Skip to content

Commit 43674f5

Browse files
author
Frank
committed
Add AWS tool
1 parent bec4b35 commit 43674f5

File tree

6 files changed

+234
-89
lines changed

6 files changed

+234
-89
lines changed

‎app/core/package.json‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
"private": true,
66
"type": "module",
77
"dependencies": {
8+
"@aws-sdk/client-sts": "3.782.0",
89
"@rocicorp/zero": "0.17.2025031904",
910
"drizzle-orm": "0.41.0",
1011
"ulid": "3.0.0"

‎app/core/src/aws.ts‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
import { z } from "zod"
2+
import { AssumeRoleCommand, STSClient } from "@aws-sdk/client-sts"
3+
import { Actor } from "./actor"
4+
import { fn } from "./util/fn"
5+
6+
export namespace Aws {
7+
export const assumeRole = fn(
8+
z.object({
9+
accountNumber: z.string(),
10+
region: z.string(),
11+
}),
12+
async (input) => {
13+
const workspaceID = Actor.workspace()
14+
15+
const sts = new STSClient({})
16+
const result = await sts.send(
17+
new AssumeRoleCommand({
18+
RoleArn: `arn:aws:iam::${input.accountNumber}:role/opencontrol-${workspaceID}-${input.region}`,
19+
RoleSessionName: "opencontrol",
20+
ExternalId: workspaceID,
21+
DurationSeconds: 3600,
22+
}),
23+
)
24+
if (!result.Credentials) throw new Error("Failed to assume role")
25+
26+
return result.Credentials
27+
},
28+
)
29+
}

‎app/function/package.json‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,12 +6,12 @@
66
"type": "module",
77
"dependencies": {
88
"@ai-sdk/anthropic": "1.2.1",
9-
"@aws-sdk/client-sts": "3.782.0",
109
"@hono/zod-validator": "0.4.3",
1110
"@openauthjs/openauth": "0.0.0-20250322224806",
1211
"@opencontrol/core": "workspace:*",
1312
"@rocicorp/zero": "0.17.2025031904",
1413
"ai": "4.2.11",
14+
"aws-sdk": "2.1692.0",
1515
"drizzle-orm": "0.41.0",
1616
"hono": "4.7.5",
1717
"opencontrol": "0.1.0"

‎app/function/src/api.ts‎

Lines changed: 139 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import AWS from "aws-sdk"
12
import { createClient } from "@openauthjs/openauth/client"
23
import { Actor } from "@opencontrol/core/actor.js"
34
import { Log } from "@opencontrol/core/util/log.js"
@@ -11,10 +12,10 @@ import { createAnthropic } from "@ai-sdk/anthropic"
1112
import { zValidator } from "@hono/zod-validator"
1213
import { z } from "zod"
1314
import { APICallError } from "ai"
14-
import { AssumeRoleCommand, STSClient } from "@aws-sdk/client-sts"
1515
import { WorkspaceTable } from "@opencontrol/core/workspace/workspace.sql.js"
1616
import { AwsAccountTable } from "@opencontrol/core/aws.sql.js"
1717
import { Identifier } from "@opencontrol/core/identifier.js"
18+
import { Aws } from "@opencontrol/core/aws.js"
1819

1920
const model = createAnthropic({
2021
apiKey: Resource.AnthropicApiKey.value,
@@ -95,57 +96,151 @@ const app = new Hono()
9596
),
9697
async (c) => {
9798
const { workspaceID, region, role } = c.req.valid("json")
99+
return Actor.provide("system", { workspaceID }, async () => {
100+
const accountNumber = role.split(":")[4]
98101

99-
// Validate workspace id
100-
const workspace = await Database.use((tx) =>
101-
tx
102-
.select({})
103-
.from(WorkspaceTable)
104-
.where(eq(WorkspaceTable.id, workspaceID)),
105-
)
106-
if (!workspace)
107-
throw new HTTPException(500, { message: "Invalid workspace ID" })
102+
// Validate workspace id
103+
const workspace = await Database.use((tx) =>
104+
tx
105+
.select({})
106+
.from(WorkspaceTable)
107+
.where(eq(WorkspaceTable.id, workspaceID)),
108+
)
109+
if (!workspace)
110+
throw new HTTPException(500, { message: "Invalid workspace ID" })
108111

109-
// Validate role by assuming it
110-
if (!role.endsWith(`role/opencontrol-${workspaceID}-${region}`))
111-
throw new HTTPException(500, { message: "Invalid role name" })
112-
const sts = new STSClient({})
113-
await sts.send(
114-
new AssumeRoleCommand({
115-
RoleArn: role,
116-
RoleSessionName: "opencontrol",
117-
ExternalId: workspaceID,
118-
DurationSeconds: 3600,
119-
}),
120-
)
112+
// Validate role by assuming it
113+
if (!role.endsWith(`role/opencontrol-${workspaceID}-${region}`))
114+
throw new HTTPException(500, { message: "Invalid role name" })
115+
await Aws.assumeRole({ accountNumber, region })
121116

122-
const accountNumber = role.split(":")[4]
123-
await Database.use((tx) =>
124-
tx
125-
.insert(AwsAccountTable)
126-
.values({
127-
id: Identifier.create("awsAccount"),
128-
workspaceID,
129-
accountNumber,
130-
region,
131-
})
132-
.onConflictDoUpdate({
133-
target: [
134-
AwsAccountTable.workspaceID,
135-
AwsAccountTable.accountNumber,
136-
],
137-
set: {
117+
await Database.use((tx) =>
118+
tx
119+
.insert(AwsAccountTable)
120+
.values({
121+
id: Identifier.create("awsAccount"),
122+
workspaceID,
123+
accountNumber,
138124
region,
139-
timeDeleted: null,
140-
},
141-
}),
142-
)
125+
})
126+
.onConflictDoUpdate({
127+
target: [
128+
AwsAccountTable.workspaceID,
129+
AwsAccountTable.accountNumber,
130+
],
131+
set: {
132+
region,
133+
timeDeleted: null,
134+
},
135+
}),
136+
)
143137

144-
return c.json({
145-
message: "ok",
138+
return c.json({
139+
message: "ok",
140+
})
146141
})
147142
},
148143
)
144+
.post("/tools/list", async (c) => {
145+
const user = Actor.assert("user")
146+
const tools = []
147+
148+
// Look up aws tool
149+
const awsAccount = await Database.use((tx) =>
150+
tx
151+
.select({})
152+
.from(AwsAccountTable)
153+
.where(eq(AwsAccountTable.workspaceID, user.properties.workspaceID)),
154+
)
155+
if (awsAccount) {
156+
tools.push({
157+
name: "aws",
158+
description: `This uses aws sdk v2 in javascript to execute aws commands
159+
this is roughly how it works
160+
\`\`\`js
161+
import aws from "aws-sdk";
162+
aws[service][method](params)
163+
\`\`\`
164+
`,
165+
args: z.object({
166+
service: z
167+
.string()
168+
.describe(
169+
"name of the aws service in the format aws sdk v2 uses, like S3 or EC2",
170+
),
171+
method: z
172+
.string()
173+
.describe("name of the aws method in the format aws sdk v2 uses"),
174+
params: z
175+
.string()
176+
.describe("params for the aws method in json format"),
177+
}),
178+
})
179+
}
180+
181+
return c.json({ tools })
182+
})
183+
.post(
184+
"/tools/call",
185+
zValidator(
186+
"json",
187+
z.custom<{
188+
name: string
189+
service: string
190+
method: string
191+
params: string
192+
}>(),
193+
),
194+
async (c) => {
195+
const body = c.req.valid("json")
196+
const { name, service, method, params } = body
197+
198+
if (name === "aws") {
199+
const awsAccount = await Database.use((tx) =>
200+
tx
201+
.select({
202+
accountNumber: AwsAccountTable.accountNumber,
203+
region: AwsAccountTable.region,
204+
})
205+
.from(AwsAccountTable)
206+
.where(eq(AwsAccountTable.workspaceID, Actor.workspace())),
207+
)
208+
if (!awsAccount) {
209+
throw new Error(
210+
"AWS integration not found. Please connect your AWS account first.",
211+
)
212+
}
213+
214+
// Assume role
215+
const credentials = await Aws.assumeRole({
216+
accountNumber: awsAccount[0].accountNumber,
217+
region: awsAccount[0].region,
218+
})
219+
220+
/* @ts-expect-error */
221+
const client = aws.default[service]
222+
if (!client) {
223+
throw new HTTPException(500, {
224+
message: `service "${service}" is not found in aws sdk v2`,
225+
})
226+
}
227+
const instance = new client({
228+
credentials: {
229+
accessKeyId: credentials.AccessKeyId,
230+
secretAccessKey: credentials.SecretAccessKey,
231+
sessionToken: credentials.SessionToken,
232+
},
233+
})
234+
if (!instance[method]) {
235+
throw new HTTPException(500, {
236+
message: `method "${method}" is not found in on the ${service} service of aws sdk v2`,
237+
})
238+
}
239+
const response = await instance[method](JSON.parse(params)).promise()
240+
return c.json(response)
241+
}
242+
},
243+
)
149244

150245
export type ApiType = typeof app
151246
export const handler = handle(app)

0 commit comments

Comments
 (0)