History

gate mutating tool invocations behind policy-driven approval

- Mark non-GET/HEAD OpenAPI tools and non-query GraphQL tools as interaction: required by default - Add authorizePersistedToolInvocation that evaluates the invocation policy engine and triggers an approval elicitation when needed - Build InvocationDescriptor from OpenAPI, GraphQL, and artifact tools - Derive interaction purpose from context (tool_execution_gate, source_connect_oauth2, source_connect_secret, elicitation) - Pass accountId through workspace execution environment resolver - Add server integration test verifying default gating and policy override - Update workspace execution environment tests with onElicitation

Rhys Sullivan committed 8 months ago

Commit d38d36e1 parent 6f804836 files · +462 −8Browse files

Unified · Split

6 files
  1. apps/executor/src/server/server.test.ts +194 −0
  2. packages/control-plane/src/runtime/graphql-tools.ts +1 −1
  3. packages/control-plane/src/runtime/live-execution.ts +16 −0
  4. packages/control-plane/src/runtime/local-operations.ts +18 −5
  5. packages/control-plane/src/runtime/workspace-execution-environment.test.ts +4 −0
  6. packages/control-plane/src/runtime/workspace-execution-environment.ts +229 −2
apps/executor/src/server/server.test.ts+194−0
@@ -129,6 +129,99 @@ const startOpenApiDemoServer = async () => { · 128 unmodified lines
129129 };
130130 };
131131
132+const startMutatingOpenApiDemoServer = async () => {
133+ const createdBodies: Array<Record<string, unknown>> = [];
134+ const openApiDocument = JSON.stringify({
135+ openapi: "3.0.3",
136+ info: {
137+ title: "Executor DNS Demo API",
138+ version: "1.0.0",
139+ },
140+ paths: {
141+ "/records": {
142+ post: {
143+ operationId: "records.createRecord",
144+ tags: ["records"],
145+ summary: "Create a DNS record",
146+ requestBody: {
147+ required: true,
148+ content: {
149+ "application/json": {
150+ schema: {
151+ type: "object",
152+ properties: {
153+ type: { type: "string" },
154+ name: { type: "string" },
155+ value: { type: "string" },
156+ },
157+ required: ["type", "value"],
158+ },
159+ },
160+ },
161+ },
162+ responses: {
163+ 200: {
164+ description: "ok",
165+ },
166+ },
167+ },
168+ },
169+ },
170+ });
171+
172+ const handler = (req: IncomingMessage, res: ServerResponse) => {
173+ if (req.url === "/openapi.json") {
174+ res.statusCode = 200;
175+ res.setHeader("content-type", "application/json");
176+ res.end(openApiDocument);
177+ return;
178+ }
179+
180+ if (req.method === "POST" && req.url === "/records") {
181+ const chunks: Array<Buffer> = [];
182+ req.on("data", (chunk) => {
183+ chunks.push(Buffer.isBuffer(chunk) ? chunk : Buffer.from(chunk));
184+ });
185+ req.on("end", () => {
186+ const raw = Buffer.concat(chunks).toString("utf8");
187+ const parsed = JSON.parse(raw) as Record<string, unknown>;
188+ createdBodies.push(parsed);
189+ res.statusCode = 200;
190+ res.setHeader("content-type", "application/json");
191+ res.end(JSON.stringify({
192+ ok: true,
193+ id: `rec_${createdBodies.length}`,
194+ record: parsed,
195+ }));
196+ });
197+ return;
198+ }
199+
200+ res.statusCode = 404;
201+ res.end("not found");
202+ };
203+
204+ const server = createServer(handler);
205+ await new Promise<void>((resolve, reject) => {
206+ server.once("error", reject);
207+ server.listen(0, "127.0.0.1", () => resolve());
208+ });
209+ const address = server.address();
210+ if (!address || typeof address === "string") {
211+ throw new Error("Failed to bind mutating OpenAPI demo server");
212+ }
213+
214+ return {
215+ baseUrl: `http://127.0.0.1:${address.port}`,
216+ specUrl: `http://127.0.0.1:${address.port}/openapi.json`,
217+ createdBodies,
218+ close: () =>
219+ new Promise<void>((resolve, reject) => {
220+ server.close((error) => (error ? reject(error) : resolve()));
221+ }),
222+ };
223+};
224+
132225 const DEMO_OAUTH_ACCESS_TOKEN = "demo-access-token";
133226 const DEMO_OAUTH_REFRESH_TOKEN = "demo-refresh-token";
134227
@@ -802,6 +895,107 @@ describe("local-executor-server", () => { · 667 unmodified lines
802895 }),
803896 );
804897
898+ it.scoped("gates mutating OpenAPI tools by default and allows them via organization policy", () =>
899+ Effect.gen(function* () {
900+ const openApiServer = yield* Effect.acquireRelease(
901+ Effect.promise(() => startMutatingOpenApiDemoServer()),
902+ (server) => Effect.promise(() => server.close()).pipe(Effect.orDie),
903+ );
904+ const server = yield* createLocalExecutorServer({
905+ port: 0,
906+ localDataDir: ":memory:",
907+ });
908+ const bootstrapClient = yield* createControlPlaneClient({
909+ baseUrl: server.baseUrl,
910+ });
911+ const installation = yield* bootstrapClient.local.installation({});
912+ const client = yield* createControlPlaneClient({
913+ baseUrl: server.baseUrl,
914+ accountId: installation.accountId,
915+ });
916+
917+ const connected = yield* client.sources.connect({
918+ path: {
919+ workspaceId: installation.workspaceId,
920+ },
921+ payload: {
922+ kind: "openapi",
923+ name: "DNS",
924+ namespace: "dns",
925+ endpoint: openApiServer.baseUrl,
926+ specUrl: openApiServer.specUrl,
927+ auth: {
928+ kind: "none",
929+ },
930+ },
931+ });
932+ expect(connected.kind).toBe("connected");
933+
934+ const gated = yield* client.executions.create({
935+ path: {
936+ workspaceId: installation.workspaceId,
937+ },
938+ payload: {
939+ code: 'return await tools.dns.records.createRecord({ body: { type: "TXT", name: "", value: "hello world" } });',
940+ },
941+ });
942+
943+ expect(gated.execution.status).toBe("waiting_for_interaction");
944+ expect(gated.pendingInteraction).not.toBeNull();
945+ if (gated.pendingInteraction === null) {
946+ throw new Error("Expected pending approval interaction");
947+ }
948+ expect(gated.pendingInteraction.kind).toBe("form");
949+ expect(gated.pendingInteraction.payloadJson).toContain("Allow POST /records?");
950+ expect(gated.pendingInteraction.payloadJson).toContain("\"approve\"");
951+
952+ const approved = yield* client.executions.resume({
953+ path: {
954+ workspaceId: installation.workspaceId,
955+ executionId: gated.execution.id,
956+ },
957+ payload: {
958+ responseJson: JSON.stringify({
959+ action: "accept",
960+ content: {
961+ approve: true,
962+ },
963+ }),
964+ },
965+ });
966+
967+ expect(approved.execution.status).toBe("completed");
968+ expect(openApiServer.createdBodies).toHaveLength(1);
969+
970+ const policy = yield* client.policies.createOrganization({
971+ path: {
972+ organizationId: installation.organizationId,
973+ },
974+ payload: {
975+ resourceType: "tool_path",
976+ resourcePattern: "dns.records.createRecord",
977+ matchType: "exact",
978+ effect: "allow",
979+ approvalMode: "auto",
980+ },
981+ });
982+ expect(policy.scopeType).toBe("organization");
983+
984+ const automatic = yield* client.executions.create({
985+ path: {
986+ workspaceId: installation.workspaceId,
987+ },
988+ payload: {
989+ code: 'return await tools.dns.records.createRecord({ body: { type: "TXT", name: "", value: "hello again" } });',
990+ },
991+ });
992+
993+ expect(automatic.execution.status).toBe("completed");
994+ expect(automatic.pendingInteraction).toBeNull();
995+ expect(openApiServer.createdBodies).toHaveLength(2);
996+ }),
997+ );
998+
805999 it.scoped("starts source OAuth without creating a source and stores secrets on callback", () =>
8061000 Effect.gen(function* () {
8071001 const oauthServer = yield* Effect.acquireRelease(
packages/control-plane/src/runtime/graphql-tools.ts+1−1
@@ -1074,7 +1074,7 @@ export const graphqlToolDescriptorFromDefinition = (input: { · 1073 unmodified lines
10741074 path: asToolPath(input.path),
10751075 sourceKey: input.sourceKey,
10761076 description: input.definition.description,
1077− interaction: "auto",
1077+ interaction: input.definition.operationType === "query" ? "auto" : "required",
10781078 inputType: presentation.inputType,
10791079 outputType: presentation.outputType,
10801080 ...(input.includeSchemas && presentation.inputSchemaJson
packages/control-plane/src/runtime/live-execution.ts+16−0
@@ -93,6 +93,21 @@ export const sanitizePersistedElicitationResponse = ( · 92 unmodified lines
9393 };
9494 };
9595
96+const interactionPurposeFromInput = (input: Parameters<OnElicitation>[0]): string => {
97+ const explicitPurpose = input.context?.interactionPurpose;
98+ if (typeof explicitPurpose === "string" && explicitPurpose.length > 0) {
99+ return explicitPurpose;
100+ }
101+
102+ if (input.path === "executor.sources.add") {
103+ return input.elicitation.mode === "url"
104+ ? "source_connect_oauth2"
105+ : "source_connect_secret";
106+ }
107+
108+ return "elicitation";
109+};
110+
96111
97112 export const createLiveExecutionManager = () => {
98113 const runs = new Map<Execution["id"], LiveRunEntry>();
@@ -143,6 +158,7 @@ export const createLiveExecutionManager = () => { · 44 unmodified lines
143158 executionId,
144159 status: "pending",
145160 kind: input.elicitation.mode === "url" ? "url" : "form",
161+ purpose: interactionPurposeFromInput(input),
146162 payloadJson:
147163 serializeJson({
148164 path: input.path,
packages/control-plane/src/runtime/local-operations.ts+18−5
@@ -61,6 +61,12 @@ const decodeSourceCredentialInteraction = ( · 60 unmodified lines
6161 interaction: ExecutionInteraction,
6262 ): Omit<SourceCredentialInteraction, "sourceLabel" | "endpoint"> | null => {
6363 try {
64+ if (interaction.purpose !== "source_connect_oauth2"
65+ && interaction.purpose !== "source_connect_secret"
66+ && interaction.purpose !== "elicitation") {
67+ return null;
68+ }
69+
6470 const payload = JSON.parse(interaction.payloadJson) as unknown;
6571 if (!isRecord(payload)) {
6672 return null;
@@ -72,11 +78,18 @@ const decodeSourceCredentialInteraction = ( · 5 unmodified lines
7278 return null;
7379 }
7480
75− if (
76− payload.path !== "executor.sources.add"
77− || args.kind !== "openapi"
78− || elicitation.mode !== "url"
79− ) {
81+ if (payload.path !== "executor.sources.add" || args.kind !== "openapi") {
82+ return null;
83+ }
84+ const effectivePurpose = interaction.purpose === "elicitation"
85+ ? elicitation.mode === "url"
86+ ? "source_connect_oauth2"
87+ : "source_connect_secret"
88+ : interaction.purpose;
89+ if (effectivePurpose === "source_connect_oauth2" && elicitation.mode !== "url") {
90+ return null;
91+ }
92+ if (effectivePurpose === "source_connect_secret" && elicitation.mode !== "form") {
8093 return null;
8194 }
8295
packages/control-plane/src/runtime/workspace-execution-environment.test.ts+4−0
@@ -1229,6 +1229,10 @@ describe("workspace-execution-environment", () => { · 1228 unmodified lines
12291229 workspaceId,
12301230 accountId: AccountIdSchema.make("acc_add_graphql_fresh"),
12311231 executionId: ExecutionIdSchema.make("exec_add_graphql_fresh"),
1232+ onElicitation: () =>
1233+ Effect.succeed({
1234+ action: "accept" as const,
1235+ }),
12321236 });
12331237
12341238 const discovered = (yield* freshEnvironment.toolInvoker.invoke({
packages/control-plane/src/runtime/workspace-execution-environment.ts+229−2
@@ -29,6 +29,7 @@ import { · 28 unmodified lines
2929 type SqlControlPlaneRows,
3030 } from "#persistence";
3131 import type {
32+ AccountId,
3233 Source,
3334 StoredToolArtifactRecord,
3435 } from "#schema";
@@ -70,6 +71,10 @@ import { · 35 unmodified lines
7071 resolveSourceAuthMaterial,
7172 storedToolIdFromArtifact,
7273 } from "./tool-artifacts";
74+import {
75+ evaluateInvocationPolicy,
76+ type InvocationDescriptor,
77+} from "./invocation-policy-engine";
7378
7479 const asToolPath = (value: string): ToolPath => value as ToolPath;
7580
@@ -382,7 +387,11 @@ const loadOpenApiWorkspaceTools = (input: { · 306 unmodified lines
382387 path,
383388 sourceKey: source.id,
384389 description: definition.description,
385− interaction: "auto",
390+ interaction:
391+ definition.method.toUpperCase() === "GET"
392+ || definition.method.toUpperCase() === "HEAD"
393+ ? "auto"
394+ : "required",
386395 inputType: presentation.inputType,
387396 outputType: presentation.outputType,
388397 ...(input.includeSchemas && presentation.inputSchemaJson
@@ -519,6 +528,185 @@ const loadSourceById = (input: { · 130 unmodified lines
519528 ),
520529 );
521530
531+const approvalSchema = {
532+ type: "object",
533+ properties: {
534+ approve: {
535+ type: "boolean",
536+ description: "Whether to approve this tool execution",
537+ },
538+ },
539+ required: ["approve"],
540+ additionalProperties: false,
541+} satisfies Record<string, unknown>;
542+
543+const approvalMessageForInvocation = (descriptor: InvocationDescriptor): string => {
544+ if (descriptor.httpMethod && descriptor.httpPathTemplate) {
545+ return `Allow ${descriptor.httpMethod.toUpperCase()} ${descriptor.httpPathTemplate}?`;
546+ }
547+
548+ if (descriptor.graphqlOperationType) {
549+ return `Allow GraphQL ${descriptor.graphqlOperationType} ${descriptor.toolPath}?`;
550+ }
551+
552+ return `Allow tool call: ${descriptor.toolPath}?`;
553+};
554+
555+const toInvocationDescriptorFromOpenApiTool = (input: {
556+ tool: OpenApiWorkspaceTool;
557+}): InvocationDescriptor => {
558+ const method = input.tool.definition.method.toUpperCase();
559+ return {
560+ toolPath: input.tool.path,
561+ sourceId: input.tool.source.id,
562+ sourceName: input.tool.source.name,
563+ sourceKind: input.tool.source.kind,
564+ sourceNamespace: input.tool.source.namespace ?? namespaceFromSourceName(input.tool.source.name),
565+ operationKind:
566+ method === "GET" || method === "HEAD"
567+ ? "read"
568+ : method === "DELETE"
569+ ? "delete"
570+ : "write",
571+ httpMethod: method,
572+ httpPathTemplate: input.tool.definition.path,
573+ graphqlOperationType: null,
574+ };
575+};
576+
577+const toInvocationDescriptorFromGraphqlTool = (input: {
578+ tool: GraphqlWorkspaceTool;
579+}): InvocationDescriptor => ({
580+ toolPath: input.tool.path,
581+ sourceId: input.tool.source.id,
582+ sourceName: input.tool.source.name,
583+ sourceKind: input.tool.source.kind,
584+ sourceNamespace: input.tool.source.namespace ?? namespaceFromSourceName(input.tool.source.name),
585+ operationKind:
586+ input.tool.definition.operationType === "query"
587+ ? "read"
588+ : input.tool.definition.operationType === "mutation"
589+ ? "write"
590+ : "unknown",
591+ httpMethod: null,
592+ httpPathTemplate: null,
593+ graphqlOperationType: input.tool.definition.operationType,
594+});
595+
596+const toInvocationDescriptorFromArtifact = (input: {
597+ toolPath: string;
598+ source: Source;
599+ artifact: StoredToolArtifactRecord;
600+}): InvocationDescriptor => ({
601+ toolPath: input.toolPath,
602+ sourceId: input.source.id,
603+ sourceName: input.source.name,
604+ sourceKind: input.source.kind,
605+ sourceNamespace: input.source.namespace ?? namespaceFromSourceName(input.source.name),
606+ operationKind: "unknown",
607+ httpMethod: input.artifact.openApiMethod?.toUpperCase() ?? null,
608+ httpPathTemplate: input.artifact.openApiPathTemplate,
609+ graphqlOperationType: null,
610+});
611+
612+const authorizePersistedToolInvocation = (input: {
613+ rows: SqlControlPlaneRows;
614+ workspaceId: Source["workspaceId"];
615+ accountId: AccountId;
616+ descriptor: InvocationDescriptor;
617+ args: unknown;
618+ source: Source;
619+ context?: Record<string, unknown>;
620+ onElicitation?: Parameters<typeof makeToolInvokerFromTools>[0]["onElicitation"];
621+}): Effect.Effect<void, Error, never> =>
622+ Effect.gen(function* () {
623+ const workspace = yield* input.rows.workspaces.getById(input.workspaceId).pipe(
624+ Effect.mapError((cause) =>
625+ cause instanceof Error ? cause : new Error(String(cause)),
626+ ),
627+ );
628+ const policies = Option.isSome(workspace)
629+ ? yield* input.rows.policies.listForWorkspaceContext({
630+ organizationId: workspace.value.organizationId,
631+ workspaceId: input.workspaceId,
632+ }).pipe(
633+ Effect.mapError((cause) =>
634+ cause instanceof Error ? cause : new Error(String(cause)),
635+ ),
636+ )
637+ : [];
638+
639+ const decision = evaluateInvocationPolicy({
640+ descriptor: input.descriptor,
641+ args: input.args,
642+ policies,
643+ context: {
644+ workspaceId: input.workspaceId,
645+ organizationId: Option.isSome(workspace)
646+ ? workspace.value.organizationId
647+ : ("org_unknown" as never),
648+ accountId: input.accountId,
649+ clientId:
650+ typeof input.context?.clientId === "string"
651+ && input.context.clientId.length > 0
652+ ? input.context.clientId
653+ : null,
654+ },
655+ });
656+
657+ if (decision.kind === "allow") {
658+ return;
659+ }
660+
661+ if (decision.kind === "deny") {
662+ return yield* Effect.fail(new Error(decision.reason));
663+ }
664+
665+ if (!input.onElicitation) {
666+ return yield* Effect.fail(
667+ new Error(`Approval required for ${input.descriptor.toolPath}, but no elicitation-capable host is available`),
668+ );
669+ }
670+
671+ const interactionId = typeof input.context?.callId === "string" && input.context.callId.length > 0
672+ ? `tool_execution_gate:${input.context.callId}`
673+ : `tool_execution_gate:${crypto.randomUUID()}`;
674+ const response = yield* input.onElicitation({
675+ interactionId,
676+ path: asToolPath(input.descriptor.toolPath),
677+ sourceKey: input.source.id,
678+ args: input.args,
679+ context: {
680+ ...(input.context ?? {}),
681+ interactionPurpose: "tool_execution_gate",
682+ interactionReason: decision.reason,
683+ invocationDescriptor: {
684+ operationKind: input.descriptor.operationKind,
685+ httpMethod: input.descriptor.httpMethod,
686+ httpPathTemplate: input.descriptor.httpPathTemplate,
687+ graphqlOperationType: input.descriptor.graphqlOperationType,
688+ sourceId: input.source.id,
689+ sourceName: input.source.name,
690+ },
691+ },
692+ elicitation: {
693+ mode: "form",
694+ message: approvalMessageForInvocation(input.descriptor),
695+ requestedSchema: approvalSchema,
696+ },
697+ }).pipe(
698+ Effect.mapError((cause) =>
699+ cause instanceof Error ? cause : new Error(String(cause)),
700+ ),
701+ );
702+
703+ if (response.action !== "accept") {
704+ return yield* Effect.fail(
705+ new Error(`Tool invocation not approved for ${input.descriptor.toolPath}`),
706+ );
707+ }
708+ });
709+
522710 const createWorkspaceToolCatalog = (input: {
523711 workspaceId: Source["workspaceId"];
524712 rows: SqlControlPlaneRows;
@@ -780,6 +968,7 @@ const createWorkspaceToolCatalog = (input: { · 255 unmodified lines
780968
781969 const createWorkspaceToolInvoker = (input: {
782970 workspaceId: Source["workspaceId"];
971+ accountId: AccountId;
783972 rows: SqlControlPlaneRows;
784973 resolveSecretMaterial: ResolveSecretMaterial;
785974 sourceAuthService: RuntimeSourceAuthService;
@@ -825,6 +1014,17 @@ const createWorkspaceToolInvoker = (input: { · 39 unmodified lines
8251014 });
8261015 const openApiTool = openApiTools.find((tool) => tool.path === invocation.path);
8271016 if (openApiTool) {
1017+ yield* authorizePersistedToolInvocation({
1018+ rows: input.rows,
1019+ workspaceId: input.workspaceId,
1020+ accountId: input.accountId,
1021+ descriptor: toInvocationDescriptorFromOpenApiTool({ tool: openApiTool }),
1022+ args: invocation.args,
1023+ source: openApiTool.source,
1024+ context: invocation.context,
1025+ onElicitation: input.onElicitation,
1026+ });
1027+
8281028 const auth = yield* resolveSourceAuthMaterial({
8291029 source: openApiTool.source,
8301030 resolveSecretMaterial: input.resolveSecretMaterial,
@@ -857,6 +1057,17 @@ const createWorkspaceToolInvoker = (input: { · 26 unmodified lines
8571057 });
8581058 const graphqlTool = graphqlTools.find((tool) => tool.path === invocation.path);
8591059 if (graphqlTool) {
1060+ yield* authorizePersistedToolInvocation({
1061+ rows: input.rows,
1062+ workspaceId: input.workspaceId,
1063+ accountId: input.accountId,
1064+ descriptor: toInvocationDescriptorFromGraphqlTool({ tool: graphqlTool }),
1065+ args: invocation.args,
1066+ source: graphqlTool.source,
1067+ context: invocation.context,
1068+ onElicitation: input.onElicitation,
1069+ });
1070+
8601071 const auth = yield* resolveSourceAuthMaterial({
8611072 source: graphqlTool.source,
8621073 resolveSecretMaterial: input.resolveSecretMaterial,
@@ -907,6 +1118,21 @@ const createWorkspaceToolInvoker = (input: { · 44 unmodified lines
9071118 );
9081119 }
9091120
1121+ yield* authorizePersistedToolInvocation({
1122+ rows: input.rows,
1123+ workspaceId: input.workspaceId,
1124+ accountId: input.accountId,
1125+ descriptor: toInvocationDescriptorFromArtifact({
1126+ toolPath: invocation.path,
1127+ source,
1128+ artifact,
1129+ }),
1130+ args: invocation.args,
1131+ source,
1132+ context: invocation.context,
1133+ onElicitation: input.onElicitation,
1134+ });
1135+
9101136 const auth = yield* resolveSourceAuthMaterial({
9111137 source,
9121138 resolveSecretMaterial: input.resolveSecretMaterial,
@@ -977,10 +1203,11 @@ export const createWorkspaceExecutionEnvironmentResolver = (input: { · 64 unmodified lines
9771203 rows: input.rows,
9781204 });
9791205
980− return ({ workspaceId, onElicitation }) =>
1206+ return ({ workspaceId, accountId, onElicitation }) =>
9811207 Effect.sync(() => {
9821208 const { catalog, toolInvoker } = createWorkspaceToolInvoker({
9831209 workspaceId,
1210+ accountId,
9841211 rows: input.rows,
9851212 resolveSecretMaterial,
9861213 sourceAuthService: input.sourceAuthService,