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
1111 status: "completed",
1212 result: { result: "ok", logs: [] },
1313 }),
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),
1424 resume: () =>
1525 Effect.succeed({
1626 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
1616 engine
1717 .executeWithPause(code)
1818 .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),
1925 // resume doesn't count as usage
2026 resume: (executionId, response) => engine.resume(executionId, response),
2127 getPausedExecution: (executionId) => engine.getPausedExecution(executionId),
e2e/scenarios/codemode-persistent-cells.test.ts+259−0
@@ -0,0 +1,259 @@
1+import { expect } from "@effect/vitest";
2+import { Effect } from "effect";
3+
4+import type { McpCallResult } from "../src/surfaces/mcp";
5+import { scenario } from "../src/scenario";
6+import { Mcp, Target } from "../src/services";
7+
8+type 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+
23+type 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+
34+const 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+
42+const eventTypes = (observation: CellObservation): readonly unknown[] =>
43+ observation.events?.map((event) => event.type) ?? [];
44+
45+const eventsOf = (observation: CellObservation): readonly CellEvent[] => observation.events ?? [];
46+
47+const 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+
57+const 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+
66+const requireString = (value: unknown, label: string): string => {
67+ expect(value, label).toEqual(expect.any(String));
68+ return value as string;
69+};
70+
71+const requireNumber = (value: unknown, label: string): number => {
72+ expect(value, label).toEqual(expect.any(Number));
73+ return value as number;
74+};
75+
76+scenario(
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
1111 code: Schema.String,
1212 });
1313
14+const StartCellRequest = Schema.Struct({
15+ code: Schema.String,
16+ yieldAfterMs: Schema.optional(Schema.Number),
17+});
18+
1419 const CompletedResult = Schema.Struct({
1520 status: Schema.Literal("completed"),
1621 text: Schema.String,
@@ -26,6 +31,15 @@ const PausedResult = Schema.Struct({ · 9 unmodified lines
2631
2732 const ExecuteResponse = Schema.Union([CompletedResult, PausedResult]);
2833
34+const 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+
2943 const ResumeRequest = Schema.Struct({
3044 action: Schema.Literals(["accept", "decline", "cancel"]),
3145 content: Schema.optional(Schema.Unknown),
@@ -47,19 +61,17 @@ const ExecutionNotFoundError = Schema.TaggedStruct("ExecutionNotFoundError", { · 15 unmodified lines
4761 // ---------------------------------------------------------------------------
4862
4963 const ExecutionParams = { executionId: Schema.String };
64+const CellParams = { cellId: Schema.String };
65+const CellWaitQuery = Schema.Struct({
66+ after: Schema.optional(Schema.String),
67+ timeoutMs: Schema.optional(Schema.String),
68+});
5069
5170 // ---------------------------------------------------------------------------
5271 // Group
5372 // ---------------------------------------------------------------------------
5473
5574 export 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− )
6375 .add(
6476 HttpApiEndpoint.post("execute", "/executions", {
6577 payload: ExecuteRequest,
@@ -67,6 +79,35 @@ export const ExecutionsApi = HttpApiGroup.make("executions") · 1 unmodified line
6779 error: InternalError,
6880 }),
6981 )
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+ )
70111 .add(
71112 HttpApiEndpoint.post("resume", "/executions/:executionId/resume", {
72113 params: ExecutionParams,
packages/core/api/src/handlers/executions.ts+92−0
@@ -14,6 +14,53 @@ class ExecutionNotFoundError extends Schema.TaggedErrorClass<ExecutionNotFoundEr · 13 unmodified lines
1414 },
1515 ) {}
1616
17+const 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+
23+const isJsonRecord = (value: unknown): value is Record<string, unknown> =>
24+ typeof value === "object" && value !== null && !Array.isArray(value);
25+
26+const 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+
62+const jsonSafe = <A>(value: A): A => toJsonCompatible(value) as A;
63+
1764 export const ExecutionsHandlers = HttpApiBuilder.group(ExecutorApi, "executions", (handlers) =>
1865 handlers
1966 .handle("getPaused", ({ params: path }) =>
@@ -55,6 +102,51 @@ export const ExecutionsHandlers = HttpApiBuilder.group(ExecutorApi, "executions" · 35 unmodified lines
55102 }),
56103 ),
57104 )
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+ )
58150 .handle("resume", ({ params: path, payload }) =>
59151 capture(
60152 Effect.gen(function* () {
packages/core/execution/src/engine.ts+281−1
@@ -1,5 +1,5 @@
11 import { Deferred, Effect, Fiber, Predicate, Queue } from "effect";
2−import type * as Cause from "effect/Cause";
2+import * as Cause from "effect/Cause";
33 import * as Exit from "effect/Exit";
44
55 import type {
@@ -12,6 +12,7 @@ import type { · 6 unmodified lines
1212 import { CodeExecutionError } from "@executor-js/codemode-core";
1313 import type {
1414 CodeExecutor,
15+ ExecuteOutputItem,
1516 ExecuteNotification,
1617 ExecuteResult,
1718 SandboxToolInvoker,
@@ -41,6 +42,50 @@ export type ExecutionResult = · 23 unmodified lines
4142 | { readonly status: "completed"; readonly result: ExecuteResult }
4243 | { readonly status: "paused"; readonly execution: PausedExecution };
4344
45+export type ExecutionCellStatus = "running" | "completed" | "failed" | "terminated";
46+
47+type 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+
67+export type ExecutionCellEvent = ExecutionCellEventPayload & {
68+ readonly id: number;
69+};
70+
71+export 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+
80+export type StartExecutionCellOptions = {
81+ readonly yieldAfterMs?: number;
82+};
83+
84+export type WaitExecutionCellOptions = {
85+ readonly after?: number;
86+ readonly timeoutMs?: number;
87+};
88+
4489 export type PausedExecution = {
4590 readonly id: string;
4691 readonly elicitationContext: ElicitationContext;
@@ -53,6 +98,25 @@ type InternalPausedExecution<E> = PausedExecution & { · 6 unmodified lines
5398 readonly pauseQueue: Queue.Queue<InternalPausedExecution<E>>;
5499 };
55100
101+type 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+
116+type CellObservationMode = "any-event" | "yield-or-terminal";
117+
118+const CELL_TERMINATED_MESSAGE = "Execution cell terminated";
119+
56120 export type ResumeResponse = {
57121 readonly action: "accept" | "decline" | "cancel";
58122 readonly content?: Record<string, unknown>;
@@ -405,6 +469,30 @@ export type ExecutionEngine<E extends Cause.YieldableError = CodeExecutionError> · 346 unmodified lines
405469 */
406470 readonly getPausedExecution: (executionId: string) => Effect.Effect<PausedExecution | null>;
407471
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+
408496 /**
409497 * Get the dynamic tool description (workflow + namespaces).
410498 */
@@ -416,6 +504,8 @@ export const createExecutionEngine = <E extends Cause.YieldableError = CodeExecu · 5 unmodified lines
416504 ): ExecutionEngine<E> => {
417505 const { executor, codeExecutor, toolDiscoveryProvider = defaultToolDiscoveryProvider } = config;
418506 const pausedExecutions = new Map<string, InternalPausedExecution<E>>();
507+ const executionCells = new Map<string, InternalExecutionCell<E>>();
508+ const runSync = Effect.runSync;
419509 // Outcomes of executions that already settled (resumed to completion, hit a
420510 // new pause, or died while paused). MCP clients retry `resume` when a
421511 // 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
466556 ),
467557 );
468558
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+
469746 /**
470747 * Start an execution in pause/resume mode.
471748 *
@@ -625,6 +902,9 @@ export const createExecutionEngine = <E extends Cause.YieldableError = CodeExecu · 153 unmodified lines
625902 resume: resumeExecution,
626903 getPausedExecution: (executionId) =>
627904 Effect.sync(() => pausedExecutions.get(executionId) ?? null),
905+ startCell: startExecutionCell,
906+ waitCell: waitExecutionCell,
907+ terminateCell: terminateExecutionCell,
628908 getDescription: buildExecuteDescription(executor),
629909 };
630910 };
packages/core/execution/src/index.ts+5−0
@@ -2,11 +2,16 @@ export { · 1 unmodified line
22 createExecutionEngine,
33 formatExecuteResult,
44 formatPausedExecution,
5+ type ExecutionCellEvent,
6+ type ExecutionCellObservation,
7+ type ExecutionCellStatus,
58 type ExecutionEngine,
69 type ExecutionEngineConfig,
710 type ExecutionResult,
811 type PausedExecution,
912 type ResumeResponse,
13+ type StartExecutionCellOptions,
14+ type WaitExecutionCellOptions,
1015 } from "./engine";
1116
1217 export { 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
2323
2424 import {
2525 createExecutionEngine as createEffectExecutionEngine,
26+ type ExecutionCellObservation,
2627 type ExecutionEngine as EffectExecutionEngine,
2728 type ExecutionResult,
2829 type PausedExecution,
2930 type ResumeResponse,
31+ type StartExecutionCellOptions,
32+ type WaitExecutionCellOptions,
3033 } from "./engine";
3134
3235 export type ElicitationHandler = (ctx: ElicitationContext) => Promise<ElicitationResponse>;
@@ -47,6 +50,15 @@ export type ExecutionEngine = { · 14 unmodified lines
4750 response: ResumeResponse,
4851 ) => Promise<ExecutionResult | null>;
4952 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>;
5062 readonly getDescription: () => Promise<string>;
5163 };
5264
@@ -152,6 +164,9 @@ export const toPromiseExecutionEngine = <E extends Cause.YieldableError>( · 99 unmodified lines
152164 executeWithPause: (code) => Effect.runPromise(engine.executeWithPause(code)),
153165 resume: (executionId, response) => Effect.runPromise(engine.resume(executionId, response)),
154166 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)),
155170 getDescription: () => Effect.runPromise(engine.getDescription),
156171 });
157172
@@ -171,7 +186,15 @@ export const createExecutionEngine = <E extends Cause.YieldableError = CodeExecu · 13 unmodified lines
171186
172187 export { formatExecuteResult, formatPausedExecution } from "./engine";
173188
174−export type { ExecutionResult, PausedExecution, ResumeResponse, ExecuteResult };
189+export type {
190+ ExecutionCellObservation,
191+ ExecutionResult,
192+ PausedExecution,
193+ ResumeResponse,
194+ ExecuteResult,
195+ StartExecutionCellOptions,
196+ WaitExecutionCellOptions,
197+};
175198
176199 export { buildExecuteDescription } from "./description";
177200 export { 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
2929 const makeStubEngine = <E extends Cause.YieldableError = never>(overrides: {
3030 execute?: ExecutionEngine<E>["execute"];
3131 executeWithPause?: ExecutionEngine<E>["executeWithPause"];
32+ startCell?: ExecutionEngine<E>["startCell"];
33+ waitCell?: ExecutionEngine<E>["waitCell"];
34+ terminateCell?: ExecutionEngine<E>["terminateCell"];
3235 resume?: ExecutionEngine<E>["resume"];
3336 description?: string;
3437 }): ExecutionEngine<E> => ({
@@ -36,6 +39,18 @@ const makeStubEngine = <E extends Cause.YieldableError = never>(overrides: { · 1 unmodified line
3639 executeWithPause:
3740 overrides.executeWithPause ??
3841 (() => 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)),
3954 resume: overrides.resume ?? (() => Effect.succeed(null)),
4055 getPausedExecution: () => Effect.succeed(null),
4156 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
2121 import type * as Tracer from "effect/Tracer";
2222 import {
2323 createExecutionEngine,
24+ type ExecutionCellEvent,
25+ type ExecutionCellObservation,
2426 formatExecuteResult,
2527 formatPausedExecution,
2628 type ExecutionEngine,
@@ -437,6 +439,43 @@ const toMcpResult = (result: FormattedExecuteInput): McpToolResult => { · 410 unmodified lines
437439 };
438440 };
439441
442+const 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+
456+const 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+
440479 const toMcpPausedResult = (formatted: ReturnType<typeof formatPausedExecution>): McpToolResult => ({
441480 content: [{ type: "text", text: formatted.text }],
442481 structuredContent: formatted.structured,
@@ -526,6 +565,20 @@ const missingExecutionResult = (executionId: string): McpToolResult => ({ · 83 unmodified lines
526565 isError: true,
527566 });
528567
568+const 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+
529582 const JsonObjectFromString = Schema.fromJsonString(Schema.Record(Schema.String, Schema.Unknown));
530583 const decodeJsonObjectString = Schema.decodeUnknownOption(JsonObjectFromString);
531584
@@ -631,6 +684,54 @@ export const createExecutorMcpServer = <E extends Cause.YieldableError>( · 99 unmodified lines
631684 }),
632685 );
633686
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+
634735 const resumeExecution = (
635736 executionId: string,
636737 action: "accept" | "decline" | "cancel",
@@ -749,6 +850,65 @@ export const createExecutorMcpServer = <E extends Cause.YieldableError>( · 112 unmodified lines
749850 }),
750851 );
751852
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+
752912 yield* Effect.sync(() => {
753913 if (elicitationMode.mode === "native") {
754914 return undefined;