From 25fdee862d7701b26ce773cd15aff518eb429b09 Mon Sep 17 00:00:00 2001 From: HeavenllyDemon <160774215+HeavenllyDemon@users.noreply.github.com> Date: Tue, 18 Aug 2026 12:17:23 +0100 Subject: [PATCH] [04/05] feat(core): add functions exec cell lifecycle --- packages/core/src/functions-exec/cells.ts | 235 ++++++++++++++++++ .../core/test/functions-exec/cells.test.ts | 155 ++++++++++++ 2 files changed, 390 insertions(+) create mode 100644 packages/core/src/functions-exec/cells.ts create mode 100644 packages/core/test/functions-exec/cells.test.ts diff --git a/packages/core/src/functions-exec/cells.ts b/packages/core/src/functions-exec/cells.ts new file mode 100644 index 00000000..f7f1a8ce --- /dev/null +++ b/packages/core/src/functions-exec/cells.ts @@ -0,0 +1,235 @@ +import { randomUUID } from "node:crypto"; +import type { CellFrame } from "./protocol"; +import type { FunctionsExecRuntimeResult } from "./runtime"; + +const MAX_CELLS_PER_SESSION = 8; +const MAX_CELLS = 32; +// A checkpoint is wrapped with a cell id and status before reaching provider input. Keep the +// combined text/error payload well below the 512-byte envelope limit from protocol.ts. +const MAX_CHECKPOINT_OUTPUT_BYTES = 256; +const MAX_CHECKPOINT_ERROR_BYTES = 96; +const encoder = new TextEncoder(); + +export type FunctionsExecCheckpoint = + | { status: "yielded"; output: string } + | { status: "completed"; output: string } + | { status: "failed"; output: string; error: string }; + +type CellStatus = "running" | "yielded" | "completed" | "failed" | "cancelled"; + +interface FunctionsExecCell { + cellId: string; + sessionId: string; + status: CellStatus; + yielded: boolean; + output: string; + checkpoint?: FunctionsExecCheckpoint; + waiters: Set<() => void>; + cancel: () => void; + onFrame?: (sessionId: string, cellId: string, frame: CellFrame) => void; + onRemoved?: (sessionId: string, cellId: string) => void; +} + +export interface StartFunctionsExecCell { + sessionId: string; + run: (cellId: string) => Promise; + cancel: (cellId: string) => void; + onFrame?: (sessionId: string, cellId: string, frame: CellFrame) => void; + onRemoved?: (sessionId: string, cellId: string) => void; +} + +function errorMessage(error: unknown): string { + return error instanceof Error ? error.message : String(error); +} + +function cap(value: string, maxBytes: number): string { + if (encoder.encode(value).byteLength <= maxBytes) return value; + let end = 0; + for (const character of value) { + const next = end + character.length; + if (encoder.encode(value.slice(0, next)).byteLength > maxBytes) break; + end = next; + } + return value.slice(0, end); +} + +function outputAfter(current: string, next: string): string { + return cap(current ? `${current}\n${next}` : next, MAX_CHECKPOINT_OUTPUT_BYTES); +} + +/** + * Engine-owned, in-memory lifecycle for a single functions.exec invocation. + * + * It deliberately has no source field, no durable media field, and no unbounded transcript. The + * engine may forward frames to a live client, but only the compact text checkpoint is available to + * the model through functions.exec/functions.wait. + */ +export class FunctionsExecCells { + private readonly cells = new Map(); + + start(input: StartFunctionsExecCell): string { + if (this.cells.size >= MAX_CELLS) { + throw new Error("too many live functions.exec cells"); + } + if (this.countForSession(input.sessionId) >= MAX_CELLS_PER_SESSION) { + throw new Error("too many live functions.exec cells"); + } + const cellId = `fx_${randomUUID()}`; + const cell: FunctionsExecCell = { + cellId, + sessionId: input.sessionId, + status: "running", + yielded: false, + output: "", + waiters: new Set(), + cancel: () => input.cancel(cellId), + onFrame: input.onFrame, + onRemoved: input.onRemoved, + }; + this.cells.set(cellId, cell); + // Defer the runner one microtask so the parent can retain the returned cell id and bind its + // nested-tool dispatcher before a fast worker makes its first bridge call. + void Promise.resolve().then(() => input.run(cellId)).then( + (result) => this.finish(cellId, result), + (error) => this.finish(cellId, { + status: "failed", + cellId, + frames: [], + error: cap(errorMessage(error), MAX_CHECKPOINT_ERROR_BYTES), + }), + ); + return cellId; + } + + recordFrame(sessionId: string, cellId: string, frame: CellFrame): void { + const cell = this.cells.get(cellId); + if (!cell || cell.sessionId !== sessionId || cell.status === "cancelled") return; + cell.onFrame?.(sessionId, cellId, frame); + switch (frame.type) { + case "text": + cell.output = outputAfter(cell.output, frame.text); + return; + case "error": + cell.output = outputAfter(cell.output, `${frame.code}: ${frame.message}`); + return; + case "yield": + cell.status = "yielded"; + cell.yielded = true; + this.publish(cell, { status: "yielded", output: cell.output }); + return; + case "image": + case "audio": + case "notification": + // These frames are intentionally forwarded only. Keeping them here would make media or + // progress survive beyond the current live turn. + return; + } + } + + canWait(sessionId: string, cellId: string): boolean { + const cell = this.cells.get(cellId); + return cell?.sessionId === sessionId && (cell.yielded || cell.checkpoint !== undefined); + } + + hasWaitable(sessionId: string): boolean { + for (const cell of this.cells.values()) { + if (cell.sessionId === sessionId && (cell.yielded || cell.checkpoint !== undefined)) return true; + } + return false; + } + + async next(sessionId: string, cellId: string, signal?: AbortSignal): Promise { + const cell = this.cellFor(sessionId, cellId); + while (cell.checkpoint === undefined) { + if (cell.status === "cancelled") throw new Error("functions execution interrupted"); + if (signal?.aborted) throw new Error("functions execution interrupted"); + await this.waitForCheckpoint(cell, signal); + } + const checkpoint = cell.checkpoint; + cell.checkpoint = undefined; + if (checkpoint.status === "yielded") cell.status = "running"; + else this.remove(cell); + return checkpoint; + } + + cancel(sessionId: string, cellId: string): boolean { + const cell = this.cells.get(cellId); + if (!cell || cell.sessionId !== sessionId) return false; + cell.status = "cancelled"; + cell.cancel(); + this.remove(cell); + return true; + } + + cancelSession(sessionId: string): void { + for (const cell of [...this.cells.values()]) { + if (cell.sessionId === sessionId) this.cancel(sessionId, cell.cellId); + } + } + + private finish(cellId: string, result: FunctionsExecRuntimeResult): void { + const cell = this.cells.get(cellId); + if (!cell || cell.status === "cancelled") return; + if (result.cellId !== cellId) { + cell.status = "failed"; + this.publish(cell, { + status: "failed", + output: cell.output, + error: "functions.exec runtime completed another cell", + }); + return; + } + if (result.status === "completed") { + cell.status = "completed"; + this.publish(cell, { status: "completed", output: cell.output }); + return; + } + cell.status = "failed"; + this.publish(cell, { + status: "failed", + output: cell.output, + error: cap(result.error, MAX_CHECKPOINT_ERROR_BYTES), + }); + } + + private countForSession(sessionId: string): number { + let count = 0; + for (const cell of this.cells.values()) if (cell.sessionId === sessionId) count += 1; + return count; + } + + private cellFor(sessionId: string, cellId: string): FunctionsExecCell { + const cell = this.cells.get(cellId); + if (!cell) throw new Error(`unknown functions.exec cell: ${cellId}`); + if (cell.sessionId !== sessionId) throw new Error(`functions.exec cell ${cellId} belongs to another session`); + return cell; + } + + private publish(cell: FunctionsExecCell, checkpoint: FunctionsExecCheckpoint): void { + cell.checkpoint = checkpoint; + for (const wake of cell.waiters) wake(); + cell.waiters.clear(); + } + + private remove(cell: FunctionsExecCell): void { + this.cells.delete(cell.cellId); + cell.onRemoved?.(cell.sessionId, cell.cellId); + for (const wake of cell.waiters) wake(); + cell.waiters.clear(); + } + + private async waitForCheckpoint(cell: FunctionsExecCell, signal?: AbortSignal): Promise { + await new Promise((resolve, reject) => { + const wake = () => { + signal?.removeEventListener("abort", abort); + resolve(); + }; + const abort = () => { + cell.waiters.delete(wake); + reject(new Error("functions execution interrupted")); + }; + cell.waiters.add(wake); + signal?.addEventListener("abort", abort, { once: true }); + }); + } +} diff --git a/packages/core/test/functions-exec/cells.test.ts b/packages/core/test/functions-exec/cells.test.ts new file mode 100644 index 00000000..85f8b0a2 --- /dev/null +++ b/packages/core/test/functions-exec/cells.test.ts @@ -0,0 +1,155 @@ +import { describe, expect, test } from "bun:test"; +import { FunctionsExecCells } from "../../src/functions-exec/cells"; +import { ImageDetail, type CellFrame } from "../../src/functions-exec/protocol"; +import type { FunctionsExecRuntimeResult } from "../../src/functions-exec/runtime"; + +function completed(cellId: string): FunctionsExecRuntimeResult { + return { status: "completed", cellId, frames: [] }; +} + +describe("functions-exec cell lifecycle", () => { + test("delivers a yield checkpoint then the terminal checkpoint without retaining media or notifications", async () => { + const cells = new FunctionsExecCells(); + let complete!: (result: FunctionsExecRuntimeResult) => void; + const live: CellFrame[] = []; + const cellId = cells.start({ + sessionId: "s1", + run: () => new Promise((resolve) => { complete = resolve; }), + cancel: () => {}, + onFrame: (_sessionId, _cellId, frame) => live.push(frame), + }); + + cells.recordFrame("s1", cellId, { type: "text", text: "first" }); + cells.recordFrame("s1", cellId, { type: "image", dataUrl: "data:image/png;base64,AA==", detail: ImageDetail.Auto }); + cells.recordFrame("s1", cellId, { type: "notification", text: "working" }); + cells.recordFrame("s1", cellId, { type: "yield" }); + expect(await cells.next("s1", cellId)).toEqual({ status: "yielded", output: "first" }); + + cells.recordFrame("s1", cellId, { type: "text", text: "second" }); + complete(completed(cellId)); + expect(await cells.next("s1", cellId)).toEqual({ status: "completed", output: "first\nsecond" }); + expect(live).toEqual([ + { type: "text", text: "first" }, + { type: "image", dataUrl: "data:image/png;base64,AA==", detail: ImageDetail.Auto }, + { type: "notification", text: "working" }, + { type: "yield" }, + { type: "text", text: "second" }, + ]); + await expect(cells.next("s1", cellId)).rejects.toThrow("unknown functions.exec cell"); + }); + + test("keeps the next checkpoint bounded even when a worker emits many text frames", async () => { + const cells = new FunctionsExecCells(); + let complete!: (result: FunctionsExecRuntimeResult) => void; + const cellId = cells.start({ + sessionId: "s1", + run: () => new Promise((resolve) => { complete = resolve; }), + cancel: () => {}, + }); + for (let index = 0; index < 8; index += 1) { + cells.recordFrame("s1", cellId, { type: "text", text: "x".repeat(256) }); + } + cells.recordFrame("s1", cellId, { type: "yield" }); + const checkpoint = await cells.next("s1", cellId); + expect(checkpoint.status).toBe("yielded"); + expect(Buffer.byteLength(checkpoint.output, "utf8")).toBeLessThanOrEqual(256); + complete(completed(cellId)); + await cells.next("s1", cellId); + }); + + test("makes wait available only after a yield or terminal checkpoint", async () => { + const cells = new FunctionsExecCells(); + let complete!: (result: FunctionsExecRuntimeResult) => void; + const cellId = cells.start({ + sessionId: "s1", + run: () => new Promise((resolve) => { complete = resolve; }), + cancel: () => {}, + }); + expect(cells.canWait("s1", cellId)).toBe(false); + cells.recordFrame("s1", cellId, { type: "yield" }); + expect(cells.canWait("s1", cellId)).toBe(true); + await cells.next("s1", cellId); + expect(cells.canWait("s1", cellId)).toBe(true); + complete(completed(cellId)); + await cells.next("s1", cellId); + }); + + test("cancels and removes a live cell when its turn is interrupted", async () => { + const cells = new FunctionsExecCells(); + let cancelled = 0; + const cellId = cells.start({ + sessionId: "s1", + run: () => new Promise(() => {}), + cancel: () => { cancelled += 1; }, + }); + const waiting = cells.next("s1", cellId); + expect(cells.cancel("s1", cellId)).toBe(true); + expect(cancelled).toBe(1); + expect(cells.cancel("s1", cellId)).toBe(false); + await expect(waiting).rejects.toThrow("interrupted"); + await expect(cells.next("s1", cellId)).rejects.toThrow("unknown functions.exec cell"); + }); + + test("fails closed if a runner tries to complete a different cell", async () => { + const cells = new FunctionsExecCells(); + const cellId = cells.start({ + sessionId: "s1", + run: async () => completed("other-cell"), + cancel: () => {}, + }); + await expect(cells.next("s1", cellId)).resolves.toEqual({ + status: "failed", + output: "", + error: "functions.exec runtime completed another cell", + }); + }); + + test("does not let a session observe another session's cell", async () => { + const cells = new FunctionsExecCells(); + const cellId = cells.start({ + sessionId: "s1", + run: async (id) => completed(id), + cancel: () => {}, + }); + await expect(cells.next("s2", cellId)).rejects.toThrow("belongs to another session"); + await cells.next("s1", cellId); + }); + + test("cancels every live cell in an interrupted session and clears removal callbacks", () => { + const cells = new FunctionsExecCells(); + let cancelled = 0; + const removed: string[] = []; + for (const sessionId of ["s1", "s1", "s2"]) { + cells.start({ + sessionId, + run: () => new Promise(() => {}), + cancel: () => { cancelled += 1; }, + onRemoved: (_sessionId, cellId) => removed.push(cellId), + }); + } + cells.cancelSession("s1"); + expect(cancelled).toBe(2); + expect(removed).toHaveLength(2); + }); + + test("caps live cells across sessions as well as within one session", () => { + const cells = new FunctionsExecCells(); + const ids: Array<{ sessionId: string; cellId: string }> = []; + for (let index = 0; index < 32; index += 1) { + const sessionId = `session-${Math.floor(index / 8)}`; + const cellId = cells.start({ + sessionId, + run: async () => await new Promise(() => {}), + cancel: () => {}, + }); + ids.push({ sessionId, cellId }); + } + + expect(() => cells.start({ + sessionId: "overflow", + run: async () => await new Promise(() => {}), + cancel: () => {}, + })).toThrow(/too many live/i); + for (const { sessionId, cellId } of ids) cells.cancel(sessionId, cellId); + }); +});