History

Add persistent code execution cells

Rhys Sullivan committed 4 months ago

Commit 0d0215d1 parent 6f46e1b10 files · +900 −9Browse files

Unified · Split

10 files
  1. apps/cloud/src/api/protected.test.ts +10 −0
  2. apps/cloud/src/engine/execution-usage.ts +6 −0
  3. e2e/scenarios/codemode-persistent-cells.test.ts +259 −0
  4. packages/core/api/src/executions/api.ts +48 −7
  5. packages/core/api/src/handlers/executions.ts +92 −0
  6. packages/core/execution/src/engine.ts +281 −1
  7. packages/core/execution/src/index.ts +5 −0
  8. packages/core/execution/src/promise.ts +24 −1
  9. packages/hosts/mcp/src/tool-server.test.ts +15 −0
  10. packages/hosts/mcp/src/tool-server.ts +160 −0
apps/cloud/src/api/protected.test.ts+10−0
@@ -11,6 +11,16 @@ const makeBaseEngine = (): ExecutionEngine => · 10 unmodified lines
11 status: "completed",11 status: "completed",
12 result: { result: "ok", logs: [] },12 result: { result: "ok", logs: [] },
13 }),13 }),
14 startCell: () =>
15 Effect.succeed({
16 status: "completed",
17 cellId: "cell_test",
18 cursor: 1,
19 events: [],
20 result: { result: "ok", logs: [] },
21 }),
22 waitCell: () => Effect.succeed(null),
23 terminateCell: () => Effect.succeed(null),
14 resume: () =>24 resume: () =>
15 Effect.succeed({25 Effect.succeed({
16 status: "completed",26 status: "completed",
apps/cloud/src/engine/execution-usage.ts+6−0
@@ -16,6 +16,12 @@ export const withExecutionUsageTracking = <E extends Cause.YieldableError>( · 15 unmodified lines
16 engine16 engine
17 .executeWithPause(code)17 .executeWithPause(code)
18 .pipe(Effect.tap(() => Effect.sync(() => trackUsage(organizationId)))),18 .pipe(Effect.tap(() => Effect.sync(() => trackUsage(organizationId)))),
19 startCell: (code, options) =>
20 engine
21 .startCell(code, options)
22 .pipe(Effect.tap(() => Effect.sync(() => trackUsage(organizationId)))),
23 waitCell: (cellId, options) => engine.waitCell(cellId, options),
24 terminateCell: (cellId) => engine.terminateCell(cellId),
19 // resume doesn't count as usage25 // resume doesn't count as usage
20 resume: (executionId, response) => engine.resume(executionId, response),26 resume: (executionId, response) => engine.resume(executionId, response),
21 getPausedExecution: (executionId) => engine.getPausedExecution(executionId),27 getPausedExecution: (executionId) => engine.getPausedExecution(executionId),
e2e/scenarios/codemode-persistent-cells.test.ts+259−0
@@ -0,0 +1,259 @@
1import { expect } from "@effect/vitest";
2import { Effect } from "effect";
3
4import type { McpCallResult } from "../src/surfaces/mcp";
5import { scenario } from "../src/scenario";
6import { Mcp, Target } from "../src/services";
7
8type CellEvent = {
9 readonly type?: unknown;
10 readonly item?: {
11 readonly type?: unknown;
12 readonly content?: {
13 readonly type?: unknown;
14 readonly text?: unknown;
15 };
16 readonly notification?: {
17 readonly message?: unknown;
18 readonly data?: unknown;
19 };
20 };
21};
22
23type CellObservation = {
24 readonly status?: unknown;
25 readonly cellId?: unknown;
26 readonly cursor?: unknown;
27 readonly events?: readonly CellEvent[];
28 readonly result?: {
29 readonly result?: unknown;
30 readonly output?: readonly unknown[];
31 };
32};
33
34const cellObservation = (result: McpCallResult): CellObservation => {
35 const structured = (result.raw as { readonly structuredContent?: unknown }).structuredContent;
36 expect(structured, `cell call returned structured content: ${result.text.slice(0, 300)}`).toEqual(
37 expect.any(Object),
38 );
39 return structured as CellObservation;
40};
41
42const eventTypes = (observation: CellObservation): readonly unknown[] =>
43 observation.events?.map((event) => event.type) ?? [];
44
45const eventsOf = (observation: CellObservation): readonly CellEvent[] => observation.events ?? [];
46
47const textOutput = (events: readonly CellEvent[]): readonly string[] =>
48 events.flatMap((event) =>
49 event.type === "output" &&
50 event.item?.type === "content" &&
51 event.item.content?.type === "text" &&
52 typeof event.item.content.text === "string"
53 ? [event.item.content.text]
54 : [],
55 );
56
57const notificationOutput = (events: readonly CellEvent[]): readonly string[] =>
58 events.flatMap((event) =>
59 event.type === "output" &&
60 event.item?.type === "notification" &&
61 typeof event.item.notification?.message === "string"
62 ? [event.item.notification.message]
63 : [],
64 );
65
66const requireString = (value: unknown, label: string): string => {
67 expect(value, label).toEqual(expect.any(String));
68 return value as string;
69};
70
71const requireNumber = (value: unknown, label: string): number => {
72 expect(value, label).toEqual(expect.any(Number));
73 return value as number;
74};
75
76scenario(
77 "Codemode · persistent cells yield, wait, notify, and terminate through MCP",
78 { timeout: 180_000 },
79 Effect.gen(function* () {
80 const target = yield* Target;
81 const mcp = yield* Mcp;
82 const identity = yield* target.newIdentity();
83 const session = mcp.session(identity);
84
85 const tools = yield* session.listTools();
86 expect(tools, "persistent cell tools are advertised").toEqual(
87 expect.arrayContaining(["execute_cell", "wait_cell", "terminate_cell"]),
88 );
89
90 const first = yield* session.call("execute_cell", {
91 yieldAfterMs: 250,
92 code: [
93 'text("phase 1");',
94 "await yield_control();",
95 'text("phase 2");',
96 "await yieldControl();",
97 'notify({ message: "cell almost done", data: { phase: 3 } });',
98 'text("phase 3");',
99 'return { cell: "done" };',
100 ].join("\n"),
101 });
102 expect(first.ok, `execute_cell starts and reaches the first yield: ${first.text}`).toBe(true);
103
104 const firstCell = cellObservation(first);
105 expect(firstCell.status, "the first observation leaves the cell running").toBe("running");
106 const cellId = requireString(firstCell.cellId, "execute_cell returns a reusable cell id");
107 let cursor = requireNumber(firstCell.cursor, "execute_cell returns an event cursor");
108 const firstEvents: CellEvent[] = [...eventsOf(firstCell)];
109 for (
110 let attempt = 0;
111 attempt < 2 && !firstEvents.some((event) => event.type === "yielded");
112 attempt++
113 ) {
114 const yielded = yield* session.call("wait_cell", {
115 cellId,
116 after: cursor,
117 timeoutMs: 5_000,
118 });
119 expect(yielded.ok, `wait_cell observes the first yield: ${yielded.text}`).toBe(true);
120 const yieldedCell = cellObservation(yielded);
121 cursor = requireNumber(yieldedCell.cursor, "wait_cell advances the cursor");
122 firstEvents.push(...eventsOf(yieldedCell));
123 }
124 expect(
125 firstEvents.map((event) => event.type),
126 "the first phase includes output and an explicit yield checkpoint",
127 ).toEqual(expect.arrayContaining(["output", "yielded"]));
128 expect(textOutput(firstEvents), "phase 1 output is visible immediately").toContain("phase 1");
129
130 const secondEvents: CellEvent[] = [];
131 for (let attempt = 0; attempt < 4; attempt++) {
132 const next = yield* session.call("wait_cell", {
133 cellId,
134 after: cursor,
135 timeoutMs: 5_000,
136 });
137 expect(next.ok, `wait_cell observes phase 2 progress: ${next.text}`).toBe(true);
138 const nextCell = cellObservation(next);
139 cursor = requireNumber(nextCell.cursor, "wait_cell advances the phase 2 cursor");
140 secondEvents.push(...eventsOf(nextCell));
141 if (
142 textOutput(secondEvents).includes("phase 2") &&
143 secondEvents.some((event) => event.type === "yielded")
144 ) {
145 break;
146 }
147 }
148 expect(textOutput(secondEvents), "phase 2 output is delivered incrementally").toContain(
149 "phase 2",
150 );
151 expect(
152 secondEvents.map((event) => event.type),
153 "phase 2 includes another explicit yield checkpoint",
154 ).toContain("yielded");
155
156 const finalEvents: CellEvent[] = [];
157 let completedCell: CellObservation | undefined;
158 for (let attempt = 0; attempt < 5; attempt++) {
159 const next = yield* session.call("wait_cell", {
160 cellId,
161 after: cursor,
162 timeoutMs: 5_000,
163 });
164 expect(next.ok, `wait_cell observes completion progress: ${next.text}`).toBe(true);
165 const nextCell = cellObservation(next);
166 cursor = requireNumber(nextCell.cursor, "wait_cell advances the completion cursor");
167 finalEvents.push(...eventsOf(nextCell));
168 if (nextCell.status === "completed") {
169 completedCell = nextCell;
170 break;
171 }
172 }
173 expect(completedCell, "wait_cell eventually observes the completed cell").toEqual(
174 expect.any(Object),
175 );
176
177 const observedCompletedCell = completedCell as CellObservation;
178 expect(observedCompletedCell.status, "the final observation is completed").toBe("completed");
179 expect(
180 finalEvents.map((event) => event.type),
181 "completion includes final outputs and terminal event",
182 ).toEqual(expect.arrayContaining(["output", "completed"]));
183 expect(
184 notificationOutput(finalEvents),
185 "notifications are emitted as structured events",
186 ).toContain("cell almost done");
187 expect(textOutput(finalEvents), "phase 3 output is visible before completion").toContain(
188 "phase 3",
189 );
190 expect(
191 observedCompletedCell.result?.result,
192 "the returned value is preserved on completion",
193 ).toEqual({ cell: "done" });
194
195 const timerRun = yield* session.call("execute_cell", {
196 yieldAfterMs: 10,
197 code: [
198 "await new Promise((resolve) => setTimeout(resolve, 50));",
199 'text("timer fired");',
200 'return { timer: "done" };',
201 ].join("\n"),
202 });
203 expect(timerRun.ok, `execute_cell starts a timer-backed cell: ${timerRun.text}`).toBe(true);
204 let timerCell = cellObservation(timerRun);
205 const timerCellId = requireString(timerCell.cellId, "timer cell id is reusable");
206 let timerCursor = requireNumber(timerCell.cursor, "timer cell returns an event cursor");
207 const timerEvents: CellEvent[] = [...eventsOf(timerCell)];
208 for (let attempt = 0; attempt < 5 && timerCell.status !== "completed"; attempt++) {
209 const next = yield* session.call("wait_cell", {
210 cellId: timerCellId,
211 after: timerCursor,
212 timeoutMs: 5_000,
213 });
214 expect(next.ok, `wait_cell observes timer-backed completion: ${next.text}`).toBe(true);
215 timerCell = cellObservation(next);
216 timerCursor = requireNumber(timerCell.cursor, "wait_cell advances the timer cursor");
217 timerEvents.push(...eventsOf(timerCell));
218 }
219 expect(timerCell.status, "timer-backed cells can complete after the first observation").toBe(
220 "completed",
221 );
222 expect(textOutput(timerEvents), "timer output is visible on completion").toContain(
223 "timer fired",
224 );
225 expect(timerCell.result?.result, "timer cell return value is preserved").toEqual({
226 timer: "done",
227 });
228
229 const running = yield* session.call("execute_cell", {
230 yieldAfterMs: 250,
231 code: [
232 "let i = 0;",
233 "while (true) {",
234 " text(`loop ${i}`);",
235 " i += 1;",
236 " await yield_control();",
237 "}",
238 ].join("\n"),
239 });
240 expect(running.ok, `execute_cell starts a cooperative long-running cell: ${running.text}`).toBe(
241 true,
242 );
243 const runningCell = cellObservation(running);
244 expect(runningCell.status, "the loop cell is running after its first yield").toBe("running");
245 const runningCellId = requireString(runningCell.cellId, "loop cell id is reusable");
246
247 const terminated = yield* session.call("terminate_cell", {
248 cellId: runningCellId,
249 });
250 expect(terminated.ok, `terminate_cell returns a terminal observation: ${terminated.text}`).toBe(
251 true,
252 );
253 const terminatedCell = cellObservation(terminated);
254 expect(terminatedCell.status, "terminate_cell marks the cell terminated").toBe("terminated");
255 expect(eventTypes(terminatedCell), "the termination event is visible to clients").toContain(
256 "terminated",
257 );
258 }),
259);
packages/core/api/src/executions/api.ts+48−7
@@ -11,6 +11,11 @@ const ExecuteRequest = Schema.Struct({ · 10 unmodified lines
11 code: Schema.String,11 code: Schema.String,
12});12});
13 13
14const StartCellRequest = Schema.Struct({
15 code: Schema.String,
16 yieldAfterMs: Schema.optional(Schema.Number),
17});
18
14const CompletedResult = Schema.Struct({19const CompletedResult = Schema.Struct({
15 status: Schema.Literal("completed"),20 status: Schema.Literal("completed"),
16 text: Schema.String,21 text: Schema.String,
@@ -26,6 +31,15 @@ const PausedResult = Schema.Struct({ · 9 unmodified lines
26 31
27const ExecuteResponse = Schema.Union([CompletedResult, PausedResult]);32const ExecuteResponse = Schema.Union([CompletedResult, PausedResult]);
28 33
34const CellObservation = Schema.Struct({
35 status: Schema.Literals(["running", "completed", "failed", "terminated"]),
36 cellId: Schema.String,
37 cursor: Schema.Number,
38 events: Schema.Array(Schema.Unknown),
39 result: Schema.optional(Schema.Unknown),
40 error: Schema.optional(Schema.String),
41});
42
29const ResumeRequest = Schema.Struct({43const ResumeRequest = Schema.Struct({
30 action: Schema.Literals(["accept", "decline", "cancel"]),44 action: Schema.Literals(["accept", "decline", "cancel"]),
31 content: Schema.optional(Schema.Unknown),45 content: Schema.optional(Schema.Unknown),
@@ -47,19 +61,17 @@ const ExecutionNotFoundError = Schema.TaggedStruct("ExecutionNotFoundError", { · 15 unmodified lines
47// ---------------------------------------------------------------------------61// ---------------------------------------------------------------------------
48 62
49const ExecutionParams = { executionId: Schema.String };63const ExecutionParams = { executionId: Schema.String };
64const CellParams = { cellId: Schema.String };
65const CellWaitQuery = Schema.Struct({
66 after: Schema.optional(Schema.String),
67 timeoutMs: Schema.optional(Schema.String),
68});
50 69
51// ---------------------------------------------------------------------------70// ---------------------------------------------------------------------------
52// Group71// Group
53// ---------------------------------------------------------------------------72// ---------------------------------------------------------------------------
54 73
55export const ExecutionsApi = HttpApiGroup.make("executions")74export const ExecutionsApi = HttpApiGroup.make("executions")
56 .add(
57 HttpApiEndpoint.get("getPaused", "/executions/:executionId", {
58 params: ExecutionParams,
59 success: PausedExecutionInfo,
60 error: [InternalError, ExecutionNotFoundError],
61 }),
62 )
63 .add(75 .add(
64 HttpApiEndpoint.post("execute", "/executions", {76 HttpApiEndpoint.post("execute", "/executions", {
65 payload: ExecuteRequest,77 payload: ExecuteRequest,
@@ -67,6 +79,35 @@ export const ExecutionsApi = HttpApiGroup.make("executions") · 1 unmodified line
67 error: InternalError,79 error: InternalError,
68 }),80 }),
69 )81 )
82 .add(
83 HttpApiEndpoint.post("startCell", "/execution-cells", {
84 payload: StartCellRequest,
85 success: CellObservation,
86 error: InternalError,
87 }),
88 )
89 .add(
90 HttpApiEndpoint.get("waitCell", "/execution-cells/:cellId", {
91 params: CellParams,
92 query: CellWaitQuery,
93 success: CellObservation,
94 error: [InternalError, ExecutionNotFoundError],
95 }),
96 )
97 .add(
98 HttpApiEndpoint.post("terminateCell", "/execution-cells/:cellId/terminate", {
99 params: CellParams,
100 success: CellObservation,
101 error: [InternalError, ExecutionNotFoundError],
102 }),
103 )
104 .add(
105 HttpApiEndpoint.get("getPaused", "/executions/:executionId", {
106 params: ExecutionParams,
107 success: PausedExecutionInfo,
108 error: [InternalError, ExecutionNotFoundError],
109 }),
110 )
70 .add(111 .add(
71 HttpApiEndpoint.post("resume", "/executions/:executionId/resume", {112 HttpApiEndpoint.post("resume", "/executions/:executionId/resume", {
72 params: ExecutionParams,113 params: ExecutionParams,
packages/core/api/src/handlers/executions.ts+92−0
@@ -14,6 +14,53 @@ class ExecutionNotFoundError extends Schema.TaggedErrorClass<ExecutionNotFoundEr · 13 unmodified lines
14 },14 },
15) {}15) {}
16 16
17const parseOptionalNumber = (value: string | undefined): number | undefined => {
18 if (value === undefined) return undefined;
19 const parsed = Number(value);
20 return Number.isFinite(parsed) ? parsed : undefined;
21};
22
23const isJsonRecord = (value: unknown): value is Record<string, unknown> =>
24 typeof value === "object" && value !== null && !Array.isArray(value);
25
26const toJsonCompatible = (value: unknown): unknown => {
27 if (
28 value === null ||
29 typeof value === "string" ||
30 typeof value === "number" ||
31 typeof value === "boolean"
32 ) {
33 return value;
34 }
35 if (
36 typeof value === "undefined" ||
37 typeof value === "function" ||
38 typeof value === "symbol" ||
39 typeof value === "bigint"
40 ) {
41 return undefined;
42 }
43 if (Array.isArray(value)) {
44 return value.map((item) => toJsonCompatible(item) ?? null);
45 }
46 if (value instanceof Date) {
47 return value.toISOString();
48 }
49 if (isJsonRecord(value)) {
50 const result: Record<string, unknown> = {};
51 for (const [key, item] of Object.entries(value)) {
52 const converted = toJsonCompatible(item);
53 if (converted !== undefined) {
54 result[key] = converted;
55 }
56 }
57 return result;
58 }
59 return undefined;
60};
61
62const jsonSafe = <A>(value: A): A => toJsonCompatible(value) as A;
63
17export const ExecutionsHandlers = HttpApiBuilder.group(ExecutorApi, "executions", (handlers) =>64export const ExecutionsHandlers = HttpApiBuilder.group(ExecutorApi, "executions", (handlers) =>
18 handlers65 handlers
19 .handle("getPaused", ({ params: path }) =>66 .handle("getPaused", ({ params: path }) =>
@@ -55,6 +102,51 @@ export const ExecutionsHandlers = HttpApiBuilder.group(ExecutorApi, "executions" · 35 unmodified lines
55 }),102 }),
56 ),103 ),
57 )104 )
105 .handle("startCell", ({ payload }) =>
106 capture(
107 Effect.gen(function* () {
108 const engine = yield* ExecutionEngineService;
109 return yield* captureEngineError(
110 engine.startCell(payload.code, {
111 yieldAfterMs: payload.yieldAfterMs,
112 }),
113 ).pipe(Effect.map(jsonSafe));
114 }),
115 ),
116 )
117 .handle("waitCell", ({ params: path, query }) =>
118 capture(
119 Effect.gen(function* () {
120 const engine = yield* ExecutionEngineService;
121 const result = yield* captureEngineError(
122 engine.waitCell(path.cellId, {
123 after: parseOptionalNumber(query.after),
124 timeoutMs: parseOptionalNumber(query.timeoutMs),
125 }),
126 );
127
128 if (!result) {
129 return yield* new ExecutionNotFoundError({ executionId: path.cellId });
130 }
131
132 return jsonSafe(result);
133 }),
134 ),
135 )
136 .handle("terminateCell", ({ params: path }) =>
137 capture(
138 Effect.gen(function* () {
139 const engine = yield* ExecutionEngineService;
140 const result = yield* engine.terminateCell(path.cellId);
141
142 if (!result) {
143 return yield* new ExecutionNotFoundError({ executionId: path.cellId });
144 }
145
146 return jsonSafe(result);
147 }),
148 ),
149 )
58 .handle("resume", ({ params: path, payload }) =>150 .handle("resume", ({ params: path, payload }) =>
59 capture(151 capture(
60 Effect.gen(function* () {152 Effect.gen(function* () {
packages/core/execution/src/engine.ts+281−1
@@ -1,5 +1,5 @@
1import { Deferred, Effect, Fiber, Predicate, Queue } from "effect";1import { Deferred, Effect, Fiber, Predicate, Queue } from "effect";
2import type * as Cause from "effect/Cause";2import * as Cause from "effect/Cause";
3import * as Exit from "effect/Exit";3import * as Exit from "effect/Exit";
4 4
5import type {5import type {
@@ -12,6 +12,7 @@ import type { · 6 unmodified lines
12import { CodeExecutionError } from "@executor-js/codemode-core";12import { CodeExecutionError } from "@executor-js/codemode-core";
13import type {13import type {
14 CodeExecutor,14 CodeExecutor,
15 ExecuteOutputItem,
15 ExecuteNotification,16 ExecuteNotification,
16 ExecuteResult,17 ExecuteResult,
17 SandboxToolInvoker,18 SandboxToolInvoker,
@@ -41,6 +42,50 @@ export type ExecutionResult = · 23 unmodified lines
41 | { readonly status: "completed"; readonly result: ExecuteResult }42 | { readonly status: "completed"; readonly result: ExecuteResult }
42 | { readonly status: "paused"; readonly execution: PausedExecution };43 | { readonly status: "paused"; readonly execution: PausedExecution };
43 44
45export type ExecutionCellStatus = "running" | "completed" | "failed" | "terminated";
46
47type ExecutionCellEventPayload =
48 | {
49 readonly type: "output";
50 readonly item: ExecuteOutputItem;
51 }
52 | {
53 readonly type: "yielded";
54 }
55 | {
56 readonly type: "completed";
57 readonly result: ExecuteResult;
58 }
59 | {
60 readonly type: "failed";
61 readonly error: string;
62 }
63 | {
64 readonly type: "terminated";
65 };
66
67export type ExecutionCellEvent = ExecutionCellEventPayload & {
68 readonly id: number;
69};
70
71export type ExecutionCellObservation = {
72 readonly status: ExecutionCellStatus;
73 readonly cellId: string;
74 readonly cursor: number;
75 readonly events: readonly ExecutionCellEvent[];
76 readonly result?: ExecuteResult;
77 readonly error?: string;
78};
79
80export type StartExecutionCellOptions = {
81 readonly yieldAfterMs?: number;
82};
83
84export type WaitExecutionCellOptions = {
85 readonly after?: number;
86 readonly timeoutMs?: number;
87};
88
44export type PausedExecution = {89export type PausedExecution = {
45 readonly id: string;90 readonly id: string;
46 readonly elicitationContext: ElicitationContext;91 readonly elicitationContext: ElicitationContext;
@@ -53,6 +98,25 @@ type InternalPausedExecution<E> = PausedExecution & { · 6 unmodified lines
53 readonly pauseQueue: Queue.Queue<InternalPausedExecution<E>>;98 readonly pauseQueue: Queue.Queue<InternalPausedExecution<E>>;
54};99};
55 100
101type InternalExecutionCell<E> = {
102 readonly id: string;
103 events: ExecutionCellEvent[];
104 nextEventId: number;
105 status: ExecutionCellStatus;
106 result?: ExecuteResult;
107 error?: string;
108 fiber?: Fiber.Fiber<ExecuteResult, E>;
109 readonly waiters: Set<Deferred.Deferred<void>>;
110 readonly yieldContinuations: Set<{
111 readonly resume: () => void;
112 readonly terminate: () => void;
113 }>;
114};
115
116type CellObservationMode = "any-event" | "yield-or-terminal";
117
118const CELL_TERMINATED_MESSAGE = "Execution cell terminated";
119
56export type ResumeResponse = {120export type ResumeResponse = {
57 readonly action: "accept" | "decline" | "cancel";121 readonly action: "accept" | "decline" | "cancel";
58 readonly content?: Record<string, unknown>;122 readonly content?: Record<string, unknown>;
@@ -405,6 +469,30 @@ export type ExecutionEngine<E extends Cause.YieldableError = CodeExecutionError> · 346 unmodified lines
405 */469 */
406 readonly getPausedExecution: (executionId: string) => Effect.Effect<PausedExecution | null>;470 readonly getPausedExecution: (executionId: string) => Effect.Effect<PausedExecution | null>;
407 471
472 /**
473 * Start a persistent execution cell. The cell keeps running after the first
474 * observation when it yields or exceeds `yieldAfterMs`.
475 */
476 readonly startCell: (
477 code: string,
478 options?: StartExecutionCellOptions,
479 ) => Effect.Effect<ExecutionCellObservation, E>;
480
481 /**
482 * Wait for new cell events after `after`. Returns null when the cell id is
483 * unknown.
484 */
485 readonly waitCell: (
486 cellId: string,
487 options?: WaitExecutionCellOptions,
488 ) => Effect.Effect<ExecutionCellObservation | null, E>;
489
490 /**
491 * Stop a running cell and return its current observation. Returns null when
492 * the cell id is unknown.
493 */
494 readonly terminateCell: (cellId: string) => Effect.Effect<ExecutionCellObservation | null>;
495
408 /**496 /**
409 * Get the dynamic tool description (workflow + namespaces).497 * Get the dynamic tool description (workflow + namespaces).
410 */498 */
@@ -416,6 +504,8 @@ export const createExecutionEngine = <E extends Cause.YieldableError = CodeExecu · 5 unmodified lines
416): ExecutionEngine<E> => {504): ExecutionEngine<E> => {
417 const { executor, codeExecutor, toolDiscoveryProvider = defaultToolDiscoveryProvider } = config;505 const { executor, codeExecutor, toolDiscoveryProvider = defaultToolDiscoveryProvider } = config;
418 const pausedExecutions = new Map<string, InternalPausedExecution<E>>();506 const pausedExecutions = new Map<string, InternalPausedExecution<E>>();
507 const executionCells = new Map<string, InternalExecutionCell<E>>();
508 const runSync = Effect.runSync;
419 // Outcomes of executions that already settled (resumed to completion, hit a509 // Outcomes of executions that already settled (resumed to completion, hit a
420 // new pause, or died while paused). MCP clients retry `resume` when a510 // new pause, or died while paused). MCP clients retry `resume` when a
421 // response gets lost in transit; without this cache the retry of an511 // response gets lost in transit; without this cache the retry of an
@@ -466,6 +556,193 @@ export const createExecutionEngine = <E extends Cause.YieldableError = CodeExecu · 44 unmodified lines
466 ),556 ),
467 );557 );
468 558
559 const notifyCellWaiters = (cell: InternalExecutionCell<E>): void => {
560 for (const waiter of cell.waiters) {
561 runSync(Deferred.succeed(waiter, undefined));
562 }
563 cell.waiters.clear();
564 };
565
566 const appendCellEvent = (
567 cell: InternalExecutionCell<E>,
568 event: ExecutionCellEventPayload,
569 ): ExecutionCellEvent => {
570 const recorded = { ...event, id: cell.nextEventId++ } as ExecutionCellEvent;
571 cell.events.push(recorded);
572 notifyCellWaiters(cell);
573 return recorded;
574 };
575
576 const observeCell = (cell: InternalExecutionCell<E>, after = 0): ExecutionCellObservation => {
577 const events = cell.events.filter((event) => event.id > after);
578 const yielded = events.some((event) => event.type === "yielded");
579 if (yielded) {
580 for (const continuation of cell.yieldContinuations) {
581 continuation.resume();
582 }
583 cell.yieldContinuations.clear();
584 }
585 const cursor = Math.max(after, cell.events.at(-1)?.id ?? after);
586 return {
587 status: cell.status,
588 cellId: cell.id,
589 cursor,
590 events,
591 ...(cell.result !== undefined ? { result: cell.result } : {}),
592 ...(cell.error !== undefined ? { error: cell.error } : {}),
593 };
594 };
595
596 const waitForCellObservation = (
597 cell: InternalExecutionCell<E>,
598 options: WaitExecutionCellOptions = {},
599 mode: CellObservationMode = "any-event",
600 ): Effect.Effect<ExecutionCellObservation> => {
601 const after = options.after ?? 0;
602 const timeoutMs = Math.max(0, Math.floor(options.timeoutMs ?? 30_000));
603 return Effect.gen(function* () {
604 const startedAt = Date.now();
605 while (true) {
606 const current = observeCell(cell, after);
607 const hasYieldOrTerminalEvent = current.events.some((event) => event.type !== "output");
608 const ready =
609 mode === "any-event"
610 ? current.events.length > 0 || cell.status !== "running"
611 : hasYieldOrTerminalEvent || cell.status !== "running";
612 if (ready || timeoutMs === 0) return current;
613
614 const remainingMs = timeoutMs - (Date.now() - startedAt);
615 if (remainingMs <= 0) return current;
616
617 const waiter = yield* Deferred.make<void>();
618 cell.waiters.add(waiter);
619 yield* Deferred.await(waiter).pipe(
620 Effect.timeoutOrElse({
621 duration: `${remainingMs} millis`,
622 orElse: () => Effect.void,
623 }),
624 Effect.ensuring(Effect.sync(() => cell.waiters.delete(waiter))),
625 );
626 }
627 });
628 };
629
630 const startExecutionCell = Effect.fn("mcp.execute.cell.start")(function* (
631 code: string,
632 options: StartExecutionCellOptions = {},
633 ) {
634 yield* Effect.annotateCurrentSpan({
635 "mcp.execute.mode": "cell",
636 "mcp.execute.code_length": code.length,
637 });
638
639 const cell: InternalExecutionCell<E> = {
640 id: `cell_${crypto.randomUUID()}`,
641 events: [],
642 nextEventId: 1,
643 status: "running",
644 waiters: new Set(),
645 yieldContinuations: new Set(),
646 };
647 executionCells.set(cell.id, cell);
648
649 const baseInvoker = makeFullInvoker(
650 executor,
651 {
652 onElicitation: () => Effect.succeed({ action: "cancel" as const }),
653 },
654 toolDiscoveryProvider,
655 );
656 const invoker: SandboxToolInvoker = {
657 invoke: (input) => baseInvoker.invoke(input),
658 };
659
660 const fiber = yield* Effect.forkDetach(
661 codeExecutor
662 .execute(code, invoker, {
663 onOutput: (item) => {
664 if (cell.status !== "running") return;
665 appendCellEvent(cell, { type: "output", item });
666 },
667 onYield: () =>
668 new Promise<void>((resolve, reject) => {
669 if (cell.status !== "running") {
670 // oxlint-disable-next-line executor/no-promise-reject -- boundary: sandbox runtimes observe rejected yield promises as cooperative cell termination
671 reject(CELL_TERMINATED_MESSAGE);
672 return;
673 }
674 appendCellEvent(cell, { type: "yielded" });
675 cell.yieldContinuations.add({
676 resume: resolve,
677 terminate: () => {
678 // oxlint-disable-next-line executor/no-promise-reject -- boundary: sandbox runtimes observe rejected yield promises as cooperative cell termination
679 reject(CELL_TERMINATED_MESSAGE);
680 },
681 });
682 }),
683 })
684 .pipe(Effect.withSpan("executor.code.exec")),
685 );
686 cell.fiber = fiber;
687
688 yield* Effect.forkDetach(
689 Fiber.await(fiber).pipe(
690 Effect.flatMap((exit) =>
691 Effect.sync(() => {
692 if (cell.status === "terminated") return;
693 if (Exit.isSuccess(exit)) {
694 cell.result = exit.value;
695 cell.status = exit.value.error ? "failed" : "completed";
696 appendCellEvent(cell, { type: "completed", result: exit.value });
697 return;
698 }
699 const error = Cause.pretty(exit.cause);
700 cell.error = error;
701 cell.status = "failed";
702 appendCellEvent(cell, { type: "failed", error });
703 }),
704 ),
705 ),
706 );
707
708 return yield* waitForCellObservation(
709 cell,
710 {
711 after: 0,
712 timeoutMs: options.yieldAfterMs ?? 1_000,
713 },
714 "yield-or-terminal",
715 );
716 });
717
718 const waitExecutionCell = Effect.fn("mcp.execute.cell.wait")(function* (
719 cellId: string,
720 options: WaitExecutionCellOptions = {},
721 ) {
722 const cell = executionCells.get(cellId);
723 if (!cell) return null;
724 return yield* waitForCellObservation(cell, options);
725 });
726
727 const terminateExecutionCell = Effect.fn("mcp.execute.cell.terminate")(function* (
728 cellId: string,
729 ) {
730 const cell = executionCells.get(cellId);
731 if (!cell) return null;
732 if (cell.status === "running") {
733 cell.status = "terminated";
734 appendCellEvent(cell, { type: "terminated" });
735 for (const continuation of cell.yieldContinuations) {
736 continuation.terminate();
737 }
738 cell.yieldContinuations.clear();
739 if (cell.fiber) {
740 yield* Fiber.interrupt(cell.fiber).pipe(Effect.ignore);
741 }
742 }
743 return observeCell(cell);
744 });
745
469 /**746 /**
470 * Start an execution in pause/resume mode.747 * Start an execution in pause/resume mode.
471 *748 *
@@ -625,6 +902,9 @@ export const createExecutionEngine = <E extends Cause.YieldableError = CodeExecu · 153 unmodified lines
625 resume: resumeExecution,902 resume: resumeExecution,
626 getPausedExecution: (executionId) =>903 getPausedExecution: (executionId) =>
627 Effect.sync(() => pausedExecutions.get(executionId) ?? null),904 Effect.sync(() => pausedExecutions.get(executionId) ?? null),
905 startCell: startExecutionCell,
906 waitCell: waitExecutionCell,
907 terminateCell: terminateExecutionCell,
628 getDescription: buildExecuteDescription(executor),908 getDescription: buildExecuteDescription(executor),
629 };909 };
630};910};
packages/core/execution/src/index.ts+5−0
@@ -2,11 +2,16 @@ export { · 1 unmodified line
2 createExecutionEngine,2 createExecutionEngine,
3 formatExecuteResult,3 formatExecuteResult,
4 formatPausedExecution,4 formatPausedExecution,
5 type ExecutionCellEvent,
6 type ExecutionCellObservation,
7 type ExecutionCellStatus,
5 type ExecutionEngine,8 type ExecutionEngine,
6 type ExecutionEngineConfig,9 type ExecutionEngineConfig,
7 type ExecutionResult,10 type ExecutionResult,
8 type PausedExecution,11 type PausedExecution,
9 type ResumeResponse,12 type ResumeResponse,
13 type StartExecutionCellOptions,
14 type WaitExecutionCellOptions,
10} from "./engine";15} from "./engine";
11 16
12export { buildExecuteDescription } from "./description";17export { buildExecuteDescription } from "./description";
packages/core/execution/src/promise.ts+24−1
@@ -23,10 +23,13 @@ import type { CodeExecutionError, CodeExecutor, ExecuteResult } from "@executor- · 22 unmodified lines
23 23
24import {24import {
25 createExecutionEngine as createEffectExecutionEngine,25 createExecutionEngine as createEffectExecutionEngine,
26 type ExecutionCellObservation,
26 type ExecutionEngine as EffectExecutionEngine,27 type ExecutionEngine as EffectExecutionEngine,
27 type ExecutionResult,28 type ExecutionResult,
28 type PausedExecution,29 type PausedExecution,
29 type ResumeResponse,30 type ResumeResponse,
31 type StartExecutionCellOptions,
32 type WaitExecutionCellOptions,
30} from "./engine";33} from "./engine";
31 34
32export type ElicitationHandler = (ctx: ElicitationContext) => Promise<ElicitationResponse>;35export type ElicitationHandler = (ctx: ElicitationContext) => Promise<ElicitationResponse>;
@@ -47,6 +50,15 @@ export type ExecutionEngine = { · 14 unmodified lines
47 response: ResumeResponse,50 response: ResumeResponse,
48 ) => Promise<ExecutionResult | null>;51 ) => Promise<ExecutionResult | null>;
49 readonly getPausedExecution: (executionId: string) => Promise<PausedExecution | null>;52 readonly getPausedExecution: (executionId: string) => Promise<PausedExecution | null>;
53 readonly startCell: (
54 code: string,
55 options?: StartExecutionCellOptions,
56 ) => Promise<ExecutionCellObservation>;
57 readonly waitCell: (
58 cellId: string,
59 options?: WaitExecutionCellOptions,
60 ) => Promise<ExecutionCellObservation | null>;
61 readonly terminateCell: (cellId: string) => Promise<ExecutionCellObservation | null>;
50 readonly getDescription: () => Promise<string>;62 readonly getDescription: () => Promise<string>;
51};63};
52 64
@@ -152,6 +164,9 @@ export const toPromiseExecutionEngine = <E extends Cause.YieldableError>( · 99 unmodified lines
152 executeWithPause: (code) => Effect.runPromise(engine.executeWithPause(code)),164 executeWithPause: (code) => Effect.runPromise(engine.executeWithPause(code)),
153 resume: (executionId, response) => Effect.runPromise(engine.resume(executionId, response)),165 resume: (executionId, response) => Effect.runPromise(engine.resume(executionId, response)),
154 getPausedExecution: (executionId) => Effect.runPromise(engine.getPausedExecution(executionId)),166 getPausedExecution: (executionId) => Effect.runPromise(engine.getPausedExecution(executionId)),
167 startCell: (code, options) => Effect.runPromise(engine.startCell(code, options)),
168 waitCell: (cellId, options) => Effect.runPromise(engine.waitCell(cellId, options)),
169 terminateCell: (cellId) => Effect.runPromise(engine.terminateCell(cellId)),
155 getDescription: () => Effect.runPromise(engine.getDescription),170 getDescription: () => Effect.runPromise(engine.getDescription),
156});171});
157 172
@@ -171,7 +186,15 @@ export const createExecutionEngine = <E extends Cause.YieldableError = CodeExecu · 13 unmodified lines
171 186
172export { formatExecuteResult, formatPausedExecution } from "./engine";187export { formatExecuteResult, formatPausedExecution } from "./engine";
173 188
174export type { ExecutionResult, PausedExecution, ResumeResponse, ExecuteResult };189export type {
190 ExecutionCellObservation,
191 ExecutionResult,
192 PausedExecution,
193 ResumeResponse,
194 ExecuteResult,
195 StartExecutionCellOptions,
196 WaitExecutionCellOptions,
197};
175 198
176export { buildExecuteDescription } from "./description";199export { buildExecuteDescription } from "./description";
177export { ExecutionToolError } from "./errors";200export { ExecutionToolError } from "./errors";
packages/hosts/mcp/src/tool-server.test.ts+15−0
@@ -29,6 +29,9 @@ class TestExecutionError extends Data.TaggedError("TestExecutionError")<{ · 28 unmodified lines
29const makeStubEngine = <E extends Cause.YieldableError = never>(overrides: {29const makeStubEngine = <E extends Cause.YieldableError = never>(overrides: {
30 execute?: ExecutionEngine<E>["execute"];30 execute?: ExecutionEngine<E>["execute"];
31 executeWithPause?: ExecutionEngine<E>["executeWithPause"];31 executeWithPause?: ExecutionEngine<E>["executeWithPause"];
32 startCell?: ExecutionEngine<E>["startCell"];
33 waitCell?: ExecutionEngine<E>["waitCell"];
34 terminateCell?: ExecutionEngine<E>["terminateCell"];
32 resume?: ExecutionEngine<E>["resume"];35 resume?: ExecutionEngine<E>["resume"];
33 description?: string;36 description?: string;
34}): ExecutionEngine<E> => ({37}): ExecutionEngine<E> => ({
@@ -36,6 +39,18 @@ const makeStubEngine = <E extends Cause.YieldableError = never>(overrides: { · 1 unmodified line
36 executeWithPause:39 executeWithPause:
37 overrides.executeWithPause ??40 overrides.executeWithPause ??
38 (() => Effect.succeed({ status: "completed", result: { result: "default" } })),41 (() => Effect.succeed({ status: "completed", result: { result: "default" } })),
42 startCell:
43 overrides.startCell ??
44 (() =>
45 Effect.succeed({
46 status: "completed",
47 cellId: "cell_test",
48 cursor: 1,
49 events: [],
50 result: { result: "default" },
51 })),
52 waitCell: overrides.waitCell ?? (() => Effect.succeed(null)),
53 terminateCell: overrides.terminateCell ?? (() => Effect.succeed(null)),
39 resume: overrides.resume ?? (() => Effect.succeed(null)),54 resume: overrides.resume ?? (() => Effect.succeed(null)),
40 getPausedExecution: () => Effect.succeed(null),55 getPausedExecution: () => Effect.succeed(null),
41 getDescription: Effect.succeed(overrides.description ?? "test executor"),56 getDescription: Effect.succeed(overrides.description ?? "test executor"),
packages/hosts/mcp/src/tool-server.ts+160−0
@@ -21,6 +21,8 @@ import type { · 20 unmodified lines
21import type * as Tracer from "effect/Tracer";21import type * as Tracer from "effect/Tracer";
22import {22import {
23 createExecutionEngine,23 createExecutionEngine,
24 type ExecutionCellEvent,
25 type ExecutionCellObservation,
24 formatExecuteResult,26 formatExecuteResult,
25 formatPausedExecution,27 formatPausedExecution,
26 type ExecutionEngine,28 type ExecutionEngine,
@@ -437,6 +439,43 @@ const toMcpResult = (result: FormattedExecuteInput): McpToolResult => { · 410 unmodified lines
437 };439 };
438};440};
439 441
442const cellEventContent = (event: ExecutionCellEvent): ContentBlock[] => {
443 if (event.type === "output") return outputItemContent(event.item);
444 if (event.type === "yielded") {
445 return [{ type: "text", text: `Execution yielded at cursor ${event.id}.` }];
446 }
447 if (event.type === "completed") {
448 return [{ type: "text", text: formatExecuteResult(event.result).text }];
449 }
450 if (event.type === "failed") {
451 return [{ type: "text", text: `Execution failed:\n${event.error}` }];
452 }
453 return [{ type: "text", text: "Execution terminated." }];
454};
455
456const toMcpCellResult = (observation: ExecutionCellObservation): McpToolResult => {
457 const content = observation.events.flatMap(cellEventContent);
458 if (content.length === 0) {
459 content.push({
460 type: "text",
461 text: `Execution cell ${observation.cellId} is ${observation.status}. cursor=${observation.cursor}`,
462 });
463 }
464
465 return {
466 content,
467 structuredContent: {
468 status: observation.status,
469 cellId: observation.cellId,
470 cursor: observation.cursor,
471 events: observation.events,
472 ...(observation.result !== undefined ? { result: observation.result } : {}),
473 ...(observation.error !== undefined ? { error: observation.error } : {}),
474 },
475 isError: observation.status === "failed" || undefined,
476 };
477};
478
440const toMcpPausedResult = (formatted: ReturnType<typeof formatPausedExecution>): McpToolResult => ({479const toMcpPausedResult = (formatted: ReturnType<typeof formatPausedExecution>): McpToolResult => ({
441 content: [{ type: "text", text: formatted.text }],480 content: [{ type: "text", text: formatted.text }],
442 structuredContent: formatted.structured,481 structuredContent: formatted.structured,
@@ -526,6 +565,20 @@ const missingExecutionResult = (executionId: string): McpToolResult => ({ · 83 unmodified lines
526 isError: true,565 isError: true,
527});566});
528 567
568const missingCellResult = (cellId: string): McpToolResult => ({
569 content: [
570 {
571 type: "text" as const,
572 text: `No execution cell: ${cellId}. The cell expired, was terminated, or was lost when its session restarted.`,
573 },
574 ],
575 structuredContent: {
576 status: "cell_not_found",
577 cellId,
578 },
579 isError: true,
580});
581
529const JsonObjectFromString = Schema.fromJsonString(Schema.Record(Schema.String, Schema.Unknown));582const JsonObjectFromString = Schema.fromJsonString(Schema.Record(Schema.String, Schema.Unknown));
530const decodeJsonObjectString = Schema.decodeUnknownOption(JsonObjectFromString);583const decodeJsonObjectString = Schema.decodeUnknownOption(JsonObjectFromString);
531 584
@@ -631,6 +684,54 @@ export const createExecutorMcpServer = <E extends Cause.YieldableError>( · 99 unmodified lines
631 }),684 }),
632 );685 );
633 686
687 const executeCell = (
688 code: string,
689 yieldAfterMs: number | undefined,
690 ): Effect.Effect<McpToolResult, E> =>
691 Effect.gen(function* () {
692 const observation = yield* engine.startCell(code, { yieldAfterMs });
693 return toMcpCellResult(observation);
694 }).pipe(
695 Effect.withSpan("mcp.host.tool.execute_cell", {
696 attributes: {
697 "mcp.tool.name": "execute_cell",
698 "mcp.execute.code_length": code.length,
699 },
700 }),
701 );
702
703 const waitCell = (
704 cellId: string,
705 after: number | undefined,
706 timeoutMs: number | undefined,
707 ): Effect.Effect<McpToolResult, E> =>
708 Effect.gen(function* () {
709 const observation = yield* engine.waitCell(cellId, { after, timeoutMs });
710 if (!observation) return missingCellResult(cellId);
711 return toMcpCellResult(observation);
712 }).pipe(
713 Effect.withSpan("mcp.host.tool.wait_cell", {
714 attributes: {
715 "mcp.tool.name": "wait_cell",
716 "mcp.execute.cell_id": cellId,
717 },
718 }),
719 );
720
721 const terminateCell = (cellId: string): Effect.Effect<McpToolResult, E> =>
722 Effect.gen(function* () {
723 const observation = yield* engine.terminateCell(cellId);
724 if (!observation) return missingCellResult(cellId);
725 return toMcpCellResult(observation);
726 }).pipe(
727 Effect.withSpan("mcp.host.tool.terminate_cell", {
728 attributes: {
729 "mcp.tool.name": "terminate_cell",
730 "mcp.execute.cell_id": cellId,
731 },
732 }),
733 );
734
634 const resumeExecution = (735 const resumeExecution = (
635 executionId: string,736 executionId: string,
636 action: "accept" | "decline" | "cancel",737 action: "accept" | "decline" | "cancel",
@@ -749,6 +850,65 @@ export const createExecutorMcpServer = <E extends Cause.YieldableError>( · 112 unmodified lines
749 }),850 }),
750 );851 );
751 852
853 yield* Effect.sync(() =>
854 server.registerTool(
855 "execute_cell",
856 {
857 description: [
858 "Start a persistent execution cell. Use this for code that yields, runs longer than one request, or needs incremental observation.",
859 "The code can call `await yield_control()` or `await yieldControl()` to hand back currently emitted output and continue after the next wait.",
860 "Call `wait_cell` with the returned cellId and cursor to observe new events. Call `terminate_cell` to stop it.",
861 ].join("\n"),
862 inputSchema: {
863 code: z.string().trim().min(1),
864 yieldAfterMs: z.number().optional(),
865 },
866 },
867 ({ code, yieldAfterMs }) => runToolEffect(executeCell(code, yieldAfterMs)),
868 ),
869 ).pipe(
870 Effect.withSpan("mcp.host.register_tool", {
871 attributes: { "mcp.tool.name": "execute_cell" },
872 }),
873 );
874
875 yield* Effect.sync(() =>
876 server.registerTool(
877 "wait_cell",
878 {
879 description:
880 "Wait for new events from a persistent execution cell. Pass the cursor returned by execute_cell or the previous wait_cell call.",
881 inputSchema: {
882 cellId: z.string(),
883 after: z.number().optional(),
884 timeoutMs: z.number().optional(),
885 },
886 },
887 ({ cellId, after, timeoutMs }) => runToolEffect(waitCell(cellId, after, timeoutMs)),
888 ),
889 ).pipe(
890 Effect.withSpan("mcp.host.register_tool", {
891 attributes: { "mcp.tool.name": "wait_cell" },
892 }),
893 );
894
895 yield* Effect.sync(() =>
896 server.registerTool(
897 "terminate_cell",
898 {
899 description: "Terminate a persistent execution cell.",
900 inputSchema: {
901 cellId: z.string(),
902 },
903 },
904 ({ cellId }) => runToolEffect(terminateCell(cellId)),
905 ),
906 ).pipe(
907 Effect.withSpan("mcp.host.register_tool", {
908 attributes: { "mcp.tool.name": "terminate_cell" },
909 }),
910 );
911
752 yield* Effect.sync(() => {912 yield* Effect.sync(() => {
753 if (elicitationMode.mode === "native") {913 if (elicitationMode.mode === "native") {
754 return undefined;914 return undefined;