Open History to see the last commit
1import { expect, it } from "@effect/vitest";
2import {
3 EnvironmentId,
4 OrchestratorMcpEnvironmentTarget,
5 OrchestratorMcpFailure,
6 ProviderInstanceId,
7 ProviderInteractionMode,
8 RuntimeMode,
9 ThreadId,
10} from "@t3tools/contracts";
11import * as Cause from "effect/Cause";
12import * as Effect from "effect/Effect";
13import * as Exit from "effect/Exit";
14import * as Layer from "effect/Layer";
15import * as Schema from "effect/Schema";
16import { McpSchema, McpServer, Tool, Toolkit } from "effect/ai";
17
18import { OrchestratorProjectionError } from "../orchestration-v2/Orchestrator.ts";
19import * as PeerForwarding from "../peer/PeerForwarding.ts";
20import * as ThreadManagement from "../orchestration-v2/ThreadManagementService.ts";
21import * as McpHttpServer from "./McpHttpServer.ts";
22import * as McpInvocationContext from "./McpInvocationContext.ts";
23import * as McpToolAccess from "./McpToolAccess.ts";
24import { liveThreadShell } from "./McpToolAccess.testkit.ts";
25
26const supervisedThreadId = ThreadId.make("thread:supervised");
27const planThreadId = ThreadId.make("thread:plan");
28const fullAccessThreadId = ThreadId.make("thread:full-access");
29const endedThreadId = ThreadId.make("thread:ended");
30const backgroundThreadId = ThreadId.make("thread:background");
31const oldProviderThreadId = ThreadId.make("thread:old-provider-background");
32const archivedThreadId = ThreadId.make("thread:archived-background");
33const unknownOwnerThreadId = ThreadId.make("thread:unknown-background-owner");
34
35const backgroundShell = (threadId: ThreadId, providerInstanceId: ProviderInstanceId) => ({
36 ...liveThreadShell(threadId, { activeRunId: null }),
37 pendingBackgroundTasks: [
38 { taskId: "task:background", kind: "subagent" as const, providerInstanceId },
39 ],
40});
41
42const shells = new Map([
43 [supervisedThreadId, liveThreadShell(supervisedThreadId, { runtimeMode: "approval-required" })],
44 [
45 planThreadId,
46 liveThreadShell(planThreadId, { runtimeMode: "approval-required", interactionMode: "plan" }),
47 ],
48 [fullAccessThreadId, liveThreadShell(fullAccessThreadId)],
49 [endedThreadId, liveThreadShell(endedThreadId, { activeRunId: null })],
50 [backgroundThreadId, backgroundShell(backgroundThreadId, ProviderInstanceId.make("codex"))],
51 [oldProviderThreadId, backgroundShell(oldProviderThreadId, ProviderInstanceId.make("claude"))],
52 [
53 archivedThreadId,
54 {
55 ...backgroundShell(archivedThreadId, ProviderInstanceId.make("codex")),
56 archivedAt: liveThreadShell(archivedThreadId).updatedAt,
57 },
58 ],
59 [
60 unknownOwnerThreadId,
61 {
62 ...liveThreadShell(unknownOwnerThreadId, { activeRunId: null }),
63 pendingBackgroundTasks: [{ taskId: "task:unknown", kind: "subagent" as const }],
64 },
65 ],
66]);
67
68const threadCaller = (threadId: ThreadId): McpInvocationContext.McpInvocationScope => ({
69 environmentId: EnvironmentId.make("environment"),
70 requestNamespace: `provider:${threadId}`,
71 thread: {
72 threadId,
73 providerSessionId: `provider:${threadId}`,
74 providerInstanceId: ProviderInstanceId.make("codex"),
75 },
76 client: undefined,
77 capabilities: new Set(["orchestration"]),
78 issuedAt: 0,
79});
80
81const clientCaller = (
82 access: McpInvocationContext.McpClientCaller["access"],
83): McpInvocationContext.McpInvocationScope => ({
84 environmentId: EnvironmentId.make("environment"),
85 requestNamespace: "client:session",
86 thread: undefined,
87 client: { sessionId: "session", label: "Claude Code", access },
88 capabilities: new Set(["orchestration"]),
89 issuedAt: 0,
90});
91
92const mcpClient = McpSchema.McpServerClient.of({
93 clientId: 1,
94 protocolVersion: "2025-06-18",
95 clientCapabilities: {},
96 clientInfo: { name: "mcp-test", version: "1" },
97 initializePayload: {
98 protocolVersion: "2025-06-18",
99 capabilities: {},
100 clientInfo: { name: "mcp-test", version: "1" },
101 },
102 getClient: Effect.die("unused"),
103});
104
105const probe = {
106 success: Schema.Struct({ ran: Schema.String }),
107 failure: OrchestratorMcpFailure,
108 failureMode: "return" as const,
109 dependencies: [
110 McpInvocationContext.McpInvocationContext,
111 ThreadManagement.ThreadManagementService,
112 ],
113};
114const ProbeToolkit = Toolkit.make(
115 Tool.make("reads", probe),
116 Tool.make("reads_as_caller", probe),
117 Tool.make("acts_as_caller", probe),
118 Tool.make("writes", probe),
119 Tool.make("writes_with_capability", probe),
120 Tool.make("writes_threads", {
121 ...probe,
122 parameters: Schema.Struct({ threadId: Schema.optional(ThreadId) }),
123 }),
124 Tool.make("starts_threads", {
125 ...probe,
126 parameters: Schema.Struct({
127 runtimeMode: Schema.optional(RuntimeMode),
128 interactionMode: Schema.optional(ProviderInteractionMode),
129 }),
130 }),
131 Tool.make("writes_environment", probe),
132);
133const ran = Effect.succeed({ ran: "ran" });
134const probeHandlers: McpToolAccess.Handlers<typeof ProbeToolkit.tools> = {
135 reads: McpToolAccess.reads(() => ran),
136 reads_as_caller: McpToolAccess.readsAsCaller(() => ran),
137 acts_as_caller: McpToolAccess.actsAsCaller(() => ran),
138 writes: McpToolAccess.writes(() => ran),
139 writes_with_capability: McpToolAccess.writesWithCapability(() => ran),
140 writes_threads: McpToolAccess.writesThreads(
141 (input) => [input.threadId],
142 () => ran,
143 ),
144 starts_threads: McpToolAccess.startsThreads(
145 (input) => input,
146 (_, modes) => Effect.succeed({ ran: `${modes.runtimeMode}/${modes.interactionMode}` }),
147 "stamped",
148 ),
149 writes_environment: McpToolAccess.writesEnvironment(() => ran),
150};
151const layerProbeHandlers = McpToolAccess.toLayer(ProbeToolkit, probeHandlers);
152
153const unchecked = () => ran;
154
155// `@ts-expect-error` fails the build if one ever compiles. Never called.
156export const refusedAtCompileTime = () => {
157 McpToolAccess.toLayer(ProbeToolkit, {
158 ...probeHandlers,
159 // @ts-expect-error a handler that skips its declaration
160 reads: unchecked,
161 });
162 McpToolAccess.toLayer(ProbeToolkit, {
163 ...probeHandlers,
164 // @ts-expect-error a declaration's fields copied onto an unchecked handler
165 reads: Object.assign(
166 unchecked,
167 McpToolAccess.reads(() => ran),
168 ),
169 });
170 // @ts-expect-error a declaration built outside McpToolAccess
171 McpToolAccess.Declaration.make(unchecked);
172 // @ts-expect-error a declaration constructed directly
173 new McpToolAccess.Declaration(unchecked); // oxlint-disable-line eslint/no-new -- This type-contract test checks direct construction.
174 // @ts-expect-error a handlers layer built outside McpToolAccess
175 McpToolAccess.HandlersLayer.make(Layer.empty);
176 // @effect-diagnostics-next-line anyUnknownInErrorContext:off
177 McpHttpServer.toolkitRegistration(
178 ProbeToolkit,
179 // @ts-expect-error a handlers layer that skipped `McpToolAccess.toLayer`
180 ProbeToolkit.toLayer({
181 reads: unchecked,
182 reads_as_caller: unchecked,
183 acts_as_caller: unchecked,
184 writes: unchecked,
185 writes_with_capability: unchecked,
186 writes_threads: unchecked,
187 starts_threads: unchecked,
188 writes_environment: unchecked,
189 }),
190 );
191};
192
193const probeServer = (threads: Layer.Layer<ThreadManagement.ThreadManagementService>) =>
194 McpHttpServer.toolkitRegistration(ProbeToolkit, layerProbeHandlers).pipe(
195 Layer.provideMerge(McpHttpServer.layerInProcessMcpServer),
196 Layer.provideMerge(threads),
197 );
198
199const decodeOutcome = Schema.decodeUnknownEffect(
200 Schema.fromJsonString(
201 Schema.Union([Schema.Struct({ ran: Schema.String }), Schema.Struct({ code: Schema.String })]),
202 ),
203);
204
205const call = (
206 tool: string,
207 scope: McpInvocationContext.McpInvocationScope,
208 args: Record<string, unknown> = {},
209) =>
210 McpServer.McpServer.pipe(
211 Effect.flatMap((server) => server.callTool({ name: tool, arguments: args })),
212 Effect.flatMap(({ content }) =>
213 decodeOutcome(content.map((part) => (part.type === "text" ? part.text : "")).join("")),
214 ),
215 Effect.map((outcome) => ("code" in outcome ? outcome.code : outcome.ran)),
216 Effect.provideService(McpInvocationContext.McpInvocationContext, scope),
217 Effect.provideService(McpSchema.McpServerClient, mcpClient),
218 );
219
220const supervised = threadCaller(supervisedThreadId);
221const plan = threadCaller(planThreadId);
222const fullAccess = threadCaller(fullAccessThreadId);
223const ended = threadCaller(endedThreadId);
224const background = threadCaller(backgroundThreadId);
225const oldProviderBackground = threadCaller(oldProviderThreadId);
226const archivedBackground = threadCaller(archivedThreadId);
227const unknownOwnerBackground = threadCaller(unknownOwnerThreadId);
228const mismatchedProvider = {
229 ...background,
230 thread: { ...background.thread!, providerInstanceId: ProviderInstanceId.make("claude") },
231};
232const supervisedClient = clientCaller("approval-required");
233const fullAccessClient = clientCaller("full-access");
234const readOnlyClient = clientCaller("read-only");
235const linkedClient: McpInvocationContext.McpInvocationScope = {
236 ...supervisedClient,
237 client: { ...supervisedClient.client!, linked: true },
238};
239// A full-access link's agents can do anything this environment's user can.
240const fullAccessLinkedClient: McpInvocationContext.McpInvocationScope = {
241 ...fullAccessClient,
242 client: { ...fullAccessClient.client!, linked: true },
243};
244const previewOnly = { ...fullAccess, capabilities: new Set(["preview"] as const) };
245
246it.effect.each([
247 ["reads", ended, {}, "ran"],
248 ["reads", supervisedClient, {}, "ran"],
249 ["reads", readOnlyClient, {}, "ran"],
250
251 ["writes", readOnlyClient, {}, "capability_denied"],
252 ["writes_threads", readOnlyClient, { threadId: supervisedThreadId }, "capability_denied"],
253 ["starts_threads", readOnlyClient, {}, "capability_denied"],
254 ["writes_environment", readOnlyClient, {}, "capability_denied"],
255
256 ["reads_as_caller", ended, {}, "ran"],
257 ["reads_as_caller", fullAccessClient, {}, "thread_credential_required"],
258 ["acts_as_caller", supervised, {}, "ran"],
259 ["acts_as_caller", ended, {}, "parent_not_active"],
260 ["acts_as_caller", background, {}, "ran"],
261 ["acts_as_caller", oldProviderBackground, {}, "parent_not_active"],
262 ["acts_as_caller", mismatchedProvider, {}, "parent_not_active"],
263 ["acts_as_caller", archivedBackground, {}, "parent_not_active"],
264 ["acts_as_caller", unknownOwnerBackground, {}, "parent_not_active"],
265 ["acts_as_caller", fullAccessClient, {}, "thread_credential_required"],
266
267 ["writes", supervised, {}, "ran"],
268 ["writes", ended, {}, "parent_not_active"],
269 ["writes", background, {}, "ran"],
270 ["writes", oldProviderBackground, {}, "parent_not_active"],
271 ["writes", supervisedClient, {}, "ran"],
272 ["writes", previewOnly, {}, "capability_denied"],
273 ["writes_environment", previewOnly, {}, "capability_denied"],
274
275 ["writes_with_capability", previewOnly, {}, "ran"],
276 ["writes_with_capability", readOnlyClient, {}, "capability_denied"],
277 ["writes_with_capability", ended, {}, "parent_not_active"],
278 ["writes", linkedClient, {}, "capability_denied"],
279 ["writes_with_capability", linkedClient, {}, "capability_denied"],
280 ["writes", fullAccessLinkedClient, {}, "ran"],
281 ["writes_with_capability", fullAccessLinkedClient, {}, "ran"],
282
283 ["writes_threads", supervised, {}, "ran"],
284 ["writes_threads", supervised, { threadId: planThreadId }, "ran"],
285 [
286 "writes_threads",
287 supervised,
288 { threadId: fullAccessThreadId },
289 "runtime_mode_escalation_denied",
290 ],
291 ["writes_threads", plan, { threadId: supervisedThreadId }, "interaction_mode_escalation_denied"],
292 ["writes_threads", fullAccess, { threadId: supervisedThreadId }, "ran"],
293 ["writes_threads", ended, { threadId: supervisedThreadId }, "parent_not_active"],
294 ["writes_threads", supervisedClient, { threadId: supervisedThreadId }, "ran"],
295 [
296 "writes_threads",
297 supervisedClient,
298 { threadId: fullAccessThreadId },
299 "runtime_mode_escalation_denied",
300 ],
301 ["writes_threads", supervised, { threadId: "thread:missing" }, "ran"],
302
303 ["starts_threads", supervised, {}, "approval-required/default"],
304 ["starts_threads", supervised, { runtimeMode: "full-access" }, "runtime_mode_escalation_denied"],
305 ["starts_threads", plan, {}, "approval-required/plan"],
306 ["starts_threads", plan, { interactionMode: "default" }, "interaction_mode_escalation_denied"],
307 ["starts_threads", fullAccess, { runtimeMode: "auto", interactionMode: "plan" }, "auto/plan"],
308 ["starts_threads", ended, {}, "parent_not_active"],
309 ["starts_threads", supervisedClient, {}, "approval-required/default"],
310 ["starts_threads", supervisedClient, { runtimeMode: "auto" }, "runtime_mode_escalation_denied"],
311 ["starts_threads", fullAccessClient, { runtimeMode: "full-access" }, "full-access/default"],
312
313 ["writes_environment", fullAccess, {}, "ran"],
314 ["writes_environment", supervised, {}, "capability_denied"],
315 ["writes_environment", plan, {}, "capability_denied"],
316 ["writes_environment", fullAccessClient, {}, "ran"],
317 ["writes_environment", supervisedClient, {}, "capability_denied"],
318] as const)("%s from %o with %o: %s", ([tool, scope, args, expected]) =>
319 call(tool, scope, args).pipe(
320 Effect.map((outcome) => expect(outcome).toBe(expected)),
321 Effect.provide(
322 probeServer(
323 Layer.mock(ThreadManagement.ThreadManagementService)({
324 getThreadShell: (threadId) => Effect.succeed(shells.get(threadId) ?? null),
325 }),
326 ),
327 ),
328 ),
329);
330
331it("refuses a declaration or handlers layer that only wears another's fields", () => {
332 const copiedDeclaration = Object.assign(
333 unchecked,
334 McpToolAccess.reads(() => ran),
335 );
336 expect(() =>
337 McpToolAccess.toLayer(ProbeToolkit, {
338 ...probeHandlers,
339 // @ts-expect-error a declaration's fields copied onto an unchecked handler
340 reads: copiedDeclaration,
341 }),
342 ).toThrow(TypeError);
343 const copiedLayer = Object.assign(Layer.empty, layerProbeHandlers);
344 expect(() => McpToolAccess.HandlersLayer.layer(copiedLayer)).toThrow(TypeError);
345});
346
347it.effect("refuses a change when the calling thread cannot be read", () =>
348 call("writes", supervised).pipe(
349 Effect.map((outcome) => expect(outcome).toBe("orchestration_error")),
350 Effect.provide(
351 probeServer(
352 Layer.mock(ThreadManagement.ThreadManagementService)({
353 getThreadShell: (threadId) =>
354 Effect.fail(new OrchestratorProjectionError({ threadId, cause: "unreadable" })),
355 }),
356 ),
357 ),
358 ),
359);
360
361const box = EnvironmentId.make("environment-box");
362const here = EnvironmentId.make("environment");
363const linkedProbe = {
364 ...probe,
365 parameters: Schema.Struct({
366 environmentId: OrchestratorMcpEnvironmentTarget,
367 threadId: Schema.optional(ThreadId),
368 runtimeMode: Schema.optional(RuntimeMode),
369 interactionMode: Schema.optional(ProviderInteractionMode),
370 }),
371};
372const LinkedProbeToolkit = Toolkit.make(
373 Tool.make("reads", linkedProbe),
374 Tool.make("writes", linkedProbe),
375 Tool.make("writes_with_capability", linkedProbe),
376 Tool.make("writes_threads", linkedProbe),
377 Tool.make("starts_threads", linkedProbe),
378 Tool.make("writes_environment", linkedProbe),
379);
380/** Runs here, and tells whether `environmentId` reached the handler. */
381const ranHere = <P extends object>(params: P) =>
382 Effect.succeed({ ran: "environmentId" in params ? "saw environmentId" : "here" });
383
384const linkedProbeServer = McpHttpServer.toolkitRegistration(
385 LinkedProbeToolkit,
386 McpToolAccess.toLayer(LinkedProbeToolkit, {
387 reads: McpToolAccess.reads(ranHere),
388 writes: McpToolAccess.writes(ranHere),
389 writes_with_capability: McpToolAccess.writesWithCapability(ranHere),
390 writes_threads: McpToolAccess.writesThreads((input) => [input.threadId], ranHere),
391 starts_threads: McpToolAccess.startsThreads((input) => input, ranHere, "stamped"),
392 writes_environment: McpToolAccess.writesEnvironment(ranHere),
393 }),
394).pipe(
395 Layer.provideMerge(McpHttpServer.layerInProcessMcpServer),
396 Layer.provideMerge(
397 Layer.mock(ThreadManagement.ThreadManagementService)({
398 getThreadShell: (threadId) => Effect.succeed(shells.get(threadId) ?? null),
399 }),
400 ),
401 // Stands in for the linked environment, answering with where the call went.
402 Layer.provide(
403 Layer.mock(PeerForwarding.PeerForwarding)({
404 call: (_scope, tool, environmentId) =>
405 Effect.succeed({ ran: `${tool.name} in ${environmentId}` } as never),
406 }),
407 ),
408);
409
410it.effect.each([
411 ["reads", supervised, { environmentId: box }, "reads in environment-box"],
412 ["reads", supervised, { environmentId: here }, "here"],
413 ["reads", supervised, {}, "here"],
414
415 // The caller is checked here before anything goes.
416 ["writes", readOnlyClient, { environmentId: box }, "capability_denied"],
417 ["writes", ended, { environmentId: box }, "parent_not_active"],
418 ["writes", supervised, { environmentId: box }, "writes in environment-box"],
419 ["writes_with_capability", readOnlyClient, { environmentId: box }, "capability_denied"],
420 [
421 "starts_threads",
422 plan,
423 { environmentId: box, interactionMode: "default" },
424 "interaction_mode_escalation_denied",
425 ],
426 ["starts_threads", plan, { environmentId: box }, "starts_threads in environment-box"],
427 ["writes_environment", supervised, { environmentId: box }, "capability_denied"],
428 [
429 "writes_environment",
430 fullAccess,
431 { environmentId: box },
432 "writes_environment in environment-box",
433 ],
434 ["writes_environment", fullAccess, { environmentId: here }, "here"],
435
436 // A thread id names a thread where the call runs, and is checked there.
437 [
438 "writes_threads",
439 supervised,
440 { environmentId: box, threadId: fullAccessThreadId },
441 "writes_threads in environment-box",
442 ],
443 [
444 "writes_threads",
445 supervised,
446 { environmentId: here, threadId: fullAccessThreadId },
447 "runtime_mode_escalation_denied",
448 ],
449] as const)("a linked %s from %o with %o: %s", ([tool, scope, args, expected]) =>
450 call(tool, scope, args).pipe(
451 Effect.map((outcome) => expect(outcome).toBe(expected)),
452 Effect.provide(linkedProbeServer),
453 ),
454);
455
456it.effect("refuses a linked environment to a tool that acts as its calling thread", () =>
457 Effect.gen(function* () {
458 const CallerToolkit = Toolkit.make(Tool.make("acts_as_caller", linkedProbe));
459 const built = yield* Layer.build(
460 McpToolAccess.HandlersLayer.layer(
461 McpToolAccess.toLayer(CallerToolkit, {
462 acts_as_caller: McpToolAccess.actsAsCaller(ranHere),
463 }),
464 ),
465 ).pipe(Effect.scoped, Effect.exit);
466 expect(Exit.isFailure(built) && Cause.hasDies(built.cause)).toBe(true);
467 }),
468);