From d5f67adcce34c64ed2e308b2a44cc36fd490fb10 Mon Sep 17 00:00:00 2001 From: Tanner Polley <42780259+tannerpolley@users.noreply.github.com> Date: Wed, 28 Oct 2026 14:06:23 -0600 Subject: [PATCH 1/4] feat(notebook): add bounded Jupyter execution service and contracts --- apps/server/src/ForkWsRpc.ts | 10 + apps/server/src/auth/RpcAuthorization.ts | 4 + apps/server/src/notebook/JupyterTransport.ts | 234 +++++++++ .../src/notebook/NotebookService.test.ts | 256 +++++++++ apps/server/src/notebook/NotebookService.ts | 488 ++++++++++++++++++ .../src/notebook/executionReducer.test.ts | 97 ++++ apps/server/src/notebook/executionReducer.ts | 121 +++++ apps/server/src/notebook/notebookDocument.ts | 62 +++ .../src/observability/RpcInstrumentation.ts | 8 + apps/server/src/server.ts | 2 + apps/server/src/ws.test.ts | 2 + .../contracts/src/clientRpcPermissions.ts | 14 + packages/contracts/src/index.ts | 3 + packages/contracts/src/notebook.ts | 112 ++++ packages/contracts/src/notebookRpc.ts | 41 ++ packages/contracts/src/rpc.ts | 3 + 16 files changed, 1457 insertions(+) create mode 100644 apps/server/src/notebook/JupyterTransport.ts create mode 100644 apps/server/src/notebook/NotebookService.test.ts create mode 100644 apps/server/src/notebook/NotebookService.ts create mode 100644 apps/server/src/notebook/executionReducer.test.ts create mode 100644 apps/server/src/notebook/executionReducer.ts create mode 100644 apps/server/src/notebook/notebookDocument.ts create mode 100644 packages/contracts/src/notebook.ts create mode 100644 packages/contracts/src/notebookRpc.ts diff --git a/apps/server/src/ForkWsRpc.ts b/apps/server/src/ForkWsRpc.ts index 9dd244aed3d3e..e48a8f3a906b4 100644 --- a/apps/server/src/ForkWsRpc.ts +++ b/apps/server/src/ForkWsRpc.ts @@ -1,3 +1,5 @@ +import { NOTEBOOK_METHODS } from "@t3tools/contracts"; +import * as NotebookService from "./notebook/NotebookService.ts"; /** * WebSocket handlers for the fork's own RPCs (`ForkWsRpcGroup`): issues, issue work, issue * changes and background-task output. This module joins them with upstream handlers for `ws.ts`. @@ -57,7 +59,15 @@ const layerForkWsRpc = (currentSession: EnvironmentAuth.AuthenticatedSession) => const decisions = yield* OwnerDecisions.OwnerDecisions; const inbox = yield* IssueInbox.IssueInbox; const artifacts = yield* ProjectArtifacts.ProjectArtifacts; + const notebook = yield* NotebookService.NotebookService; return { + [NOTEBOOK_METHODS.open]: (input) => notebook.open(input, currentSession), + [NOTEBOOK_METHODS.snapshot]: (input) => notebook.snapshot(input, currentSession), + [NOTEBOOK_METHODS.execute]: (input) => notebook.execute(input, currentSession), + [NOTEBOOK_METHODS.interrupt]: (input) => notebook.interrupt(input, currentSession), + [NOTEBOOK_METHODS.restart]: (input) => notebook.restart(input, currentSession), + [NOTEBOOK_METHODS.shutdown]: (input) => notebook.shutdown(input, currentSession), + [NOTEBOOK_METHODS.subscribe]: (input) => notebook.subscribe(input, currentSession), [PROJECT_ARTIFACT_METHODS.list]: (input) => artifacts.list(input.projectId), [PROJECT_ARTIFACT_METHODS.add]: (input) => artifacts.add(input, "user"), [PROJECT_ARTIFACT_METHODS.remove]: (input) => artifacts.remove(input), diff --git a/apps/server/src/auth/RpcAuthorization.ts b/apps/server/src/auth/RpcAuthorization.ts index 142c7cbac1a75..74179c6c90ef9 100644 --- a/apps/server/src/auth/RpcAuthorization.ts +++ b/apps/server/src/auth/RpcAuthorization.ts @@ -1,3 +1,4 @@ +import { NOTEBOOK_METHODS } from "@t3tools/contracts"; import { PROJECT_ARTIFACT_METHODS } from "@t3tools/contracts"; import { CLIENT_GUARDED_RPC_SCOPES, @@ -36,6 +37,9 @@ import type * as RpcGroup from "effect/rpc/RpcGroup"; type WsRpcMethod = RpcGroup.Rpcs["_tag"]; export const RPC_REQUIRED_SCOPES = { + [NOTEBOOK_METHODS.open]: AuthFilesystemReadScope, + [NOTEBOOK_METHODS.snapshot]: AuthFilesystemReadScope, + [NOTEBOOK_METHODS.subscribe]: AuthFilesystemReadScope, [PROJECT_ARTIFACT_METHODS.list]: AuthFilesystemReadScope, [PROJECT_ARTIFACT_METHODS.subscribe]: AuthFilesystemReadScope, ...CLIENT_GUARDED_RPC_SCOPES, diff --git a/apps/server/src/notebook/JupyterTransport.ts b/apps/server/src/notebook/JupyterTransport.ts new file mode 100644 index 0000000000000..4892e0fcee12e --- /dev/null +++ b/apps/server/src/notebook/JupyterTransport.ts @@ -0,0 +1,234 @@ +// @effect-diagnostics nodeBuiltinImport:off +// Node UUIDs identify kernel protocol requests at this adapter boundary. +import * as NodeCrypto from "node:crypto"; +import { fetch, WebSocket } from "undici"; +import { NOTEBOOK_LIMITS, NotebookError } from "@t3tools/contracts"; +import * as Cause from "effect/Cause"; +import * as Context from "effect/Context"; +import * as Effect from "effect/Effect"; +import * as Layer from "effect/Layer"; +import * as Queue from "effect/Queue"; +import * as Schema from "effect/Schema"; +import * as Scope from "effect/Scope"; +import * as Stream from "effect/Stream"; + +const encodeJson = Schema.encodeSync(Schema.fromJsonString(Schema.Unknown)); +const decodeJson = Schema.decodeSync(Schema.fromJsonString(Schema.Unknown)); +const Message = Schema.Struct({ + header: Schema.Struct({ msg_type: Schema.String }), + parent_header: Schema.Struct({ msg_id: Schema.optional(Schema.String) }), + channel: Schema.String, + content: Schema.Record(Schema.String, Schema.Unknown), + buffers: Schema.optional(Schema.Array(Schema.Unknown)), +}); +const mimeContent = { + data: Schema.Record(Schema.String, Schema.Unknown), + metadata: Schema.Record(Schema.String, Schema.Unknown), +}; +const contentValidators: Record boolean> = { + stream: Schema.is( + Schema.Struct({ name: Schema.Literals(["stdout", "stderr"]), text: Schema.String }), + ), + display_data: Schema.is(Schema.Struct(mimeContent)), + update_display_data: Schema.is(Schema.Struct(mimeContent)), + execute_result: Schema.is(Schema.Struct({ ...mimeContent, execution_count: Schema.Int })), + error: Schema.is( + Schema.Struct({ + ename: Schema.String, + evalue: Schema.String, + traceback: Schema.Array(Schema.String), + }), + ), + clear_output: Schema.is(Schema.Struct({ wait: Schema.Boolean })), +}; +const decodeMessage = Schema.decodeSync(Schema.fromJsonString(Message)); +const decodeKernel = Schema.decodeUnknownEffect(Schema.Struct({ id: Schema.String })); +export type Message = typeof Message.Type; +export type Endpoint = { + readonly url: string; + readonly token: string; + readonly kernelName: string; +}; +export type Channels = { + readonly messages: Stream.Stream; + readonly execute: (id: string, code: string) => Effect.Effect; + readonly close: Effect.Effect; +}; +export class JupyterTransport extends Context.Service< + JupyterTransport, + { + readonly create: (endpoint: Endpoint, cwd: string) => Effect.Effect; + readonly control: ( + endpoint: Endpoint, + id: string, + action: "interrupt" | "restart" | "shutdown", + ) => Effect.Effect; + readonly connect: ( + endpoint: Endpoint, + id: string, + ) => Effect.Effect; + } +>()("t3/notebook/JupyterTransport") {} + +function endpointUrl(endpoint: Endpoint, suffix: string) { + const url = new URL(endpoint.url); + if ( + url.protocol !== "http:" || + url.hostname !== "127.0.0.1" || + url.username || + url.password || + url.search || + url.hash || + !endpoint.token + ) + throw new Error(); + url.pathname = `${url.pathname.replace(/\/$/, "")}/api/kernels${suffix}`; + return url; +} +const request = Effect.fn("JupyterTransport.request")(function* ( + endpoint: Endpoint, + suffix: string, + method: string, + body?: unknown, +) { + return yield* Effect.tryPromise({ + try: async (signal) => { + const response = await fetch(endpointUrl(endpoint, suffix), { + method, + headers: { Authorization: `token ${endpoint.token}`, "Content-Type": "application/json" }, + redirect: "error", + signal, + ...(body === undefined ? {} : { body: encodeJson(body) }), + }); + if (!response.ok) { + await response.body?.cancel(); + throw new Error(); + } + if (response.status === 204) return undefined; + const reader = response.body?.getReader(); + if (!reader) throw new Error(); + let text = ""; + const decoder = new TextDecoder(); + try { + for (;;) { + const chunk = await reader.read(); + if (chunk.done) break; + text += decoder.decode(chunk.value, { stream: true }); + if (Buffer.byteLength(text) > NOTEBOOK_LIMITS.frameBytes) throw new Error(); + } + return decodeJson(text) as unknown; + } finally { + await reader.cancel(); + } + }, + catch: () => new NotebookError({ reason: "kernel-failure" }), + }); +}); + +export const layer = Layer.sync(JupyterTransport, () => { + const create = Effect.fn("JupyterTransport.create")(function* (endpoint: Endpoint, cwd: string) { + const result = yield* request(endpoint, "", "POST", { name: endpoint.kernelName, path: cwd }); + return yield* decodeKernel(result).pipe( + Effect.map((r) => r.id), + Effect.mapError(() => new NotebookError({ reason: "kernel-failure" })), + ); + }); + const control = Effect.fn("JupyterTransport.control")(function* ( + endpoint: Endpoint, + id: string, + action: "interrupt" | "restart" | "shutdown", + ) { + yield* request( + endpoint, + `/${encodeURIComponent(id)}${action === "shutdown" ? "" : `/${action}`}`, + action === "shutdown" ? "DELETE" : "POST", + ); + }); + const connect = Effect.fn("JupyterTransport.connect")(function* (endpoint: Endpoint, id: string) { + const messages = yield* Queue.dropping(64); + const socket = yield* Effect.try({ + try: () => { + const url = endpointUrl(endpoint, `/${encodeURIComponent(id)}/channels`); + url.protocol = "ws:"; + url.searchParams.set("session_id", NodeCrypto.randomUUID()); + return new WebSocket(url, { headers: { Authorization: `token ${endpoint.token}` } }); + }, + catch: () => new NotebookError({ reason: "unavailable" }), + }); + const fail = () => { + Queue.failCauseUnsafe(messages, Cause.fail(new NotebookError({ reason: "disconnected" }))); + socket.close(); + }; + socket.addEventListener("message", (event) => { + try { + if ( + typeof event.data !== "string" || + Buffer.byteLength(event.data) > NOTEBOOK_LIMITS.frameBytes + ) { + fail(); + return; + } + const message = decodeMessage(event.data); + if ( + message.buffers?.length || + contentValidators[message.header.msg_type]?.(message.content) === false || + !Queue.offerUnsafe(messages, message) + ) + fail(); + } catch { + fail(); + } + }); + socket.addEventListener("close", fail); + socket.addEventListener("error", fail); + const close = Effect.sync(() => { + socket.removeEventListener("close", fail); + socket.removeEventListener("error", fail); + socket.close(); + Queue.shutdownUnsafe(messages); + }); + yield* Effect.addFinalizer(() => close); + yield* Effect.callback((resume) => { + const opened = () => resume(Effect.void); + const failed = () => resume(Effect.fail(new NotebookError({ reason: "disconnected" }))); + socket.addEventListener("open", opened, { once: true }); + socket.addEventListener("error", failed, { once: true }); + socket.addEventListener("close", failed, { once: true }); + return Effect.sync(() => { + socket.removeEventListener("open", opened); + socket.removeEventListener("error", failed); + socket.removeEventListener("close", failed); + }); + }); + const execute = Effect.fn("JupyterTransport.execute")(function* (id: string, code: string) { + yield* Effect.try({ + try: () => + socket.send( + encodeJson({ + header: { + msg_id: id, + msg_type: "execute_request", + username: "t3", + session: "t3", + version: "5.3", + }, + parent_header: {}, + metadata: {}, + channel: "shell", + content: { + code, + silent: false, + store_history: true, + user_expressions: {}, + allow_stdin: false, + stop_on_error: true, + }, + }), + ), + catch: () => new NotebookError({ reason: "disconnected" }), + }); + }); + return { messages: Stream.fromQueue(messages), execute, close }; + }); + return JupyterTransport.of({ create, control, connect }); +}); diff --git a/apps/server/src/notebook/NotebookService.test.ts b/apps/server/src/notebook/NotebookService.test.ts new file mode 100644 index 0000000000000..145d3f7dbf2ba --- /dev/null +++ b/apps/server/src/notebook/NotebookService.test.ts @@ -0,0 +1,256 @@ +import { expect, it } from "@effect/vitest"; +import * as NodeServices from "@effect/platform-node/NodeServices"; +import { NotebookError, ProjectId, type AuthEnvironmentScope } from "@t3tools/contracts"; +import * as Cause from "effect/Cause"; +import * as Effect from "effect/Effect"; +import * as Deferred from "effect/Deferred"; +import * as Fiber from "effect/Fiber"; +import * as FileSystem from "effect/FileSystem"; +import * as Layer from "effect/Layer"; +import * as Queue from "effect/Queue"; +import * as Stream from "effect/Stream"; +import * as SqlClient from "effect/sql/SqlClient"; +import * as Sqlite from "../persistence/Sqlite.ts"; +import * as ProjectStore from "../orchestration-v2/ProjectStore.ts"; +import * as NotebookService from "./NotebookService.ts"; +import * as JupyterTransport from "./JupyterTransport.ts"; +import { parseNotebook } from "./notebookDocument.ts"; + +const caller = { scopes: ["filesystem:read", "terminal:operate"] satisfies AuthEnvironmentScope[] }; +const fixture = { + nbformat: 4, + nbformat_minor: 5, + metadata: { custom: [1, { nested: true }] }, + unknown: "retain", + cells: ["a", "b"].map((id) => ({ + id, + cell_type: "code", + source: ["print(42)\n"], + metadata: { arbitrary: true }, + execution_count: null, + outputs: [], + })), +}; +const setup = Effect.gen(function* () { + const messages = yield* Queue.unbounded(); + const requests = yield* Queue.unbounded<{ id: string; code: string }>(); + const controls: string[] = []; + const transport = JupyterTransport.JupyterTransport.of({ + create: Effect.fn("Fake.create")(() => Effect.succeed("owned")), + connect: Effect.fn("Fake.connect")(() => + Effect.succeed({ + messages: Stream.fromQueue(messages), + execute: Effect.fn("Fake.execute")((id: string, code: string) => + Queue.offer(requests, { id, code }).pipe(Effect.asVoid), + ), + close: Effect.void, + }), + ), + control: Effect.fn("Fake.control")((_endpoint, _id, action) => + Effect.sync(() => { + controls.push(action); + }), + ), + }); + const fs = yield* FileSystem.FileSystem; + const sql = yield* SqlClient.SqlClient; + const root = yield* fs.makeTempDirectoryScoped(); + const projectId = ProjectId.make(root); + yield* sql`INSERT INTO projection_projects (project_id, title, workspace_root, scripts_json, created_at, updated_at) VALUES (${projectId}, 'Notebook', ${root}, '[]', '2026-10-28T00:00:00.000Z', '2026-10-28T00:00:00.000Z')`; + yield* fs.writeFileString(`${root}/test.ipynb`, JSON.stringify(fixture)); + const context = yield* Effect.context(); + const serviceContext = yield* Layer.build( + NotebookService.layer.pipe( + Layer.provide(Layer.succeed(JupyterTransport.JupyterTransport, transport)), + Layer.provide( + Layer.succeed(NotebookService.NotebookConfiguration, { + endpoint: Effect.fn("Fake.endpoint")(() => + Effect.succeed({ url: "http://127.0.0.1:1", token: "test", kernelName: "python3" }), + ), + }), + ), + Layer.provide(Layer.succeedContext(context)), + ), + ); + const service = yield* NotebookService.NotebookService.pipe(Effect.provide(serviceContext)); + const opened = yield* service.open({ projectId, relativePath: "test.ipynb" }, caller); + const handle = { projectId, handle: opened.handle }; + const send = ( + id: string, + type: string, + content: Record = {}, + channel = type === "execute_reply" ? "shell" : "iopub", + ) => + Queue.offer(messages, { + header: { msg_type: type }, + parent_header: { msg_id: id }, + channel, + content, + }); + const finished = (runId: string) => + service.subscribe(handle, caller).pipe( + Stream.mapEffect(() => service.snapshot(handle, caller)), + Stream.filter((s) => s.runId === runId && s.outcome !== "running"), + Stream.runHead, + ); + return { service, opened, handle, requests, controls, send, finished, root, fs, messages }; +}); +const layer = Layer.mergeAll(NodeServices.layer, ProjectStore.layer).pipe( + Layer.provideMerge(Sqlite.layerMemory), +); + +it.effect("slow observers receive an explicit resync instead of retaining every delta", () => + Effect.gen(function* () { + const { service, opened, handle, requests, send, finished } = yield* setup; + const ready = yield* Deferred.make(); + const release = yield* Deferred.make(); + const resync = yield* Deferred.make(); + const observer = yield* service.subscribe(handle, caller).pipe( + Stream.runForEach((event) => + Effect.gen(function* () { + if (event.type === "snapshot") { + yield* Deferred.succeed(ready, undefined); + yield* Deferred.await(release); + } + if (event.type === "resync-required") yield* Deferred.succeed(resync, undefined); + }), + ), + Effect.forkScoped, + ); + yield* Deferred.await(ready); + const runId = yield* service.execute( + { ...handle, revision: opened.revision, cellIds: ["a", "b"], requestId: "slow" }, + caller, + ); + const first = yield* Queue.take(requests); + for (let i = 0; i < 40; i++) yield* send(first.id, "stream", { name: "stdout", text: "x" }); + yield* send(first.id, "execute_reply", { status: "ok" }); + yield* send(first.id, "status", { execution_state: "idle" }); + const second = yield* Queue.take(requests); + yield* Deferred.succeed(release, undefined); + yield* Deferred.await(resync); + yield* send(second.id, "execute_reply", { status: "ok" }); + yield* send(second.id, "status", { execution_state: "idle" }); + yield* finished(runId); + expect((yield* service.snapshot(handle, caller)).outcome).toBe("ok"); + yield* Fiber.interrupt(observer); + }).pipe(Effect.scoped, Effect.provide(layer)), +); + +it.effect("disconnect marks an uncertain outcome and never resends an accepted request", () => + Effect.gen(function* () { + const { service, opened, handle, requests, messages, finished } = yield* setup; + const input = { + ...handle, + revision: opened.revision, + cellIds: ["a", "b"], + requestId: "uncertain", + }; + const runId = yield* service.execute(input, caller); + yield* Queue.take(requests); + Queue.failCauseUnsafe(messages, Cause.fail(new NotebookError({ reason: "disconnected" }))); + yield* finished(runId); + expect((yield* service.snapshot(handle, caller)).outcome).toBe("unknown"); + expect(yield* service.execute(input, caller)).toBe(runId); + expect(yield* Queue.size(requests)).toBe(0); + expect( + yield* service.execute({ ...input, requestId: "new" }, caller).pipe(Effect.flip), + ).toMatchObject({ reason: "disconnected" }); + }).pipe(Effect.scoped, Effect.provide(layer)), +); + +it("preserves unknown notebook JSON and source representation, assigning only missing IDs", () => { + const parsed = parseNotebook(JSON.stringify(fixture)); + expect(parsed.document).toEqual(fixture); + const old = structuredClone(fixture); + Reflect.deleteProperty(old.cells[0]!, "id"); + const upgraded = parseNotebook(JSON.stringify(old)); + expect(upgraded.changed).toBe(true); + expect(upgraded.document.cells[0]?.id).toBeTruthy(); + expect(() => + parseNotebook(JSON.stringify({ ...fixture, cells: [fixture.cells[0], fixture.cells[0]] })), + ).toThrow(); +}); + +it.effect( + "correlates requests, waits for reply and idle in both orders, deduplicates, and retains late output", + () => + Effect.gen(function* () { + const { service, opened, handle, requests, send, finished } = yield* setup; + const input = { ...handle, revision: opened.revision, cellIds: ["a", "b"], requestId: "run" }; + const runId = yield* service.execute(input, caller); + expect(yield* service.execute(input, caller)).toBe(runId); + const first = yield* Queue.take(requests); + yield* send("unrelated", "stream", { name: "stdout", text: "ignore" }); + yield* send(first.id, "execute_reply", { status: "ok" }); + yield* send(first.id, "status", { execution_state: "idle" }); + const second = yield* Queue.take(requests); + yield* send(first.id, "stream", { name: "stdout", text: "late" }); + yield* send(second.id, "status", { execution_state: "idle" }); + yield* send(second.id, "execute_reply", { status: "ok" }); + yield* finished(runId); + const result = yield* service.snapshot(handle, caller); + expect(result.outcome).toBe("ok"); + expect((result.document.cells as Array<{ outputs: unknown[] }>)[0]?.outputs).toEqual([ + { output_type: "stream", name: "stdout", text: "late" }, + ]); + }).pipe(Effect.scoped, Effect.provide(layer)), +); + +it.effect( + "interrupt cancels undispatched cells; restart and shutdown fence stale generations", + () => + Effect.gen(function* () { + const { service, opened, handle, requests, controls, send, finished } = yield* setup; + const runId = yield* service.execute( + { ...handle, revision: opened.revision, cellIds: ["a", "b"], requestId: "run" }, + caller, + ); + const first = yield* Queue.take(requests); + const generation = (yield* service.snapshot(handle, caller)).generation; + yield* service.interrupt({ ...handle, generation, runId }, caller); + yield* send(first.id, "execute_reply", { status: "error" }); + yield* send(first.id, "status", { execution_state: "idle" }); + yield* finished(runId); + expect(yield* Queue.size(requests)).toBe(0); + expect((yield* service.snapshot(handle, caller)).outcome).toBe("aborted"); + yield* service.restart({ ...handle, generation }, caller); + expect( + yield* service.shutdown({ ...handle, generation }, caller).pipe(Effect.flip), + ).toMatchObject({ reason: "stale-generation" }); + yield* service.shutdown({ ...handle, generation: generation + 1 }, caller); + expect(controls).toEqual(["interrupt", "restart", "shutdown"]); + }).pipe(Effect.scoped, Effect.provide(layer)), +); + +it.effect( + "enforces read and execution grants, handle ownership, and canonical path containment", + () => + Effect.gen(function* () { + const { service, opened, handle, root, fs } = yield* setup; + expect(yield* service.snapshot(handle, { scopes: [] }).pipe(Effect.flip)).toMatchObject({ + reason: "permission-denied", + }); + expect( + yield* service + .execute( + { ...handle, revision: opened.revision, cellIds: ["a"], requestId: "no" }, + { scopes: ["filesystem:read"] }, + ) + .pipe(Effect.flip), + ).toMatchObject({ reason: "permission-denied" }); + expect( + yield* service + .snapshot({ ...handle, projectId: ProjectId.make("other") }, caller) + .pipe(Effect.flip), + ).toMatchObject({ reason: "not-found" }); + const outside = yield* fs.makeTempDirectoryScoped(); + yield* fs.writeFileString(`${outside}/escape.ipynb`, JSON.stringify(fixture)); + yield* fs.symlink(`${outside}/escape.ipynb`, `${root}/escape.ipynb`); + expect( + yield* service + .open({ projectId: handle.projectId, relativePath: "escape.ipynb" }, caller) + .pipe(Effect.flip), + ).toMatchObject({ reason: "invalid-path" }); + }).pipe(Effect.scoped, Effect.provide(layer)), +); diff --git a/apps/server/src/notebook/NotebookService.ts b/apps/server/src/notebook/NotebookService.ts new file mode 100644 index 0000000000000..dd570308ea459 --- /dev/null +++ b/apps/server/src/notebook/NotebookService.ts @@ -0,0 +1,488 @@ +// @effect-diagnostics nodeBuiltinImport:off +// Canonical filesystem identities and disk hashes use Node primitives. +import * as NodeCrypto from "node:crypto"; +import * as NodePath from "node:path"; +import { + NotebookError, + NOTEBOOK_LIMITS, + type NotebookSnapshot, + type NotebookOutputPatch, + type NotebookEvent, + type NotebookIdentity, + type NotebookHandle, + type NotebookExecute, + type NotebookControl, + type NotebookInterrupt, + type AuthEnvironmentScope, +} from "@t3tools/contracts"; +import * as Context from "effect/Context"; +import * as Deferred from "effect/Deferred"; +import * as Exit from "effect/Exit"; +import * as Schema from "effect/Schema"; +import * as Effect from "effect/Effect"; +import * as FileSystem from "effect/FileSystem"; +import * as Layer from "effect/Layer"; +import * as Option from "effect/Option"; +import * as Queue from "effect/Queue"; +import * as Scope from "effect/Scope"; +import * as Semaphore from "effect/Semaphore"; +import * as Stream from "effect/Stream"; +import * as ProjectStore from "../orchestration-v2/ProjectStore.ts"; +import * as JupyterTransport from "./JupyterTransport.ts"; +import { byteLength, parseNotebook } from "./notebookDocument.ts"; +import { ExecutionReducer } from "./executionReducer.ts"; + +// Trusted composition supplies configured endpoints and secrets; RPC callers cannot choose one. +export class NotebookConfiguration extends Context.Service< + NotebookConfiguration, + { + readonly endpoint: (root: string) => Effect.Effect; + } +>()("t3/notebook/NotebookService/NotebookConfiguration") {} +const disabled = Layer.succeed(NotebookConfiguration, { + endpoint: Effect.fn("NotebookConfiguration.endpoint")(() => + Effect.fail(new NotebookError({ reason: "unavailable" })), + ), +}); +const fingerprintOf = Schema.encodeSync(Schema.fromJsonString(Schema.Unknown)); +type Caller = { readonly scopes: readonly AuthEnvironmentScope[] }; +export class NotebookService extends Context.Service< + NotebookService, + { + readonly open: ( + input: NotebookIdentity, + caller: Caller, + ) => Effect.Effect; + readonly snapshot: ( + input: NotebookHandle, + caller: Caller, + ) => Effect.Effect; + readonly execute: ( + input: NotebookExecute, + caller: Caller, + ) => Effect.Effect; + readonly interrupt: ( + input: NotebookInterrupt, + caller: Caller, + ) => Effect.Effect; + readonly restart: ( + input: NotebookControl, + caller: Caller, + ) => Effect.Effect; + readonly shutdown: ( + input: NotebookControl, + caller: Caller, + ) => Effect.Effect; + readonly subscribe: ( + input: NotebookHandle, + caller: Caller, + ) => Stream.Stream; + } +>()("t3/notebook/NotebookService") {} + +const make = Effect.gen(function* () { + const fs = yield* FileSystem.FileSystem; + const projects = yield* ProjectStore.ProjectStoreV2; + const config = yield* NotebookConfiguration; + const transport = yield* JupyterTransport.JupyterTransport; + const scope = yield* Scope.Scope; + const lock = yield* Semaphore.make(1); + const epoch = NodeCrypto.randomUUID(); + type Session = { + projectId: string; + path: string; + root: string; + snapshot: { -readonly [K in keyof NotebookSnapshot]: NotebookSnapshot[K] }; + reducer: ExecutionReducer; + kernel?: { + id: string; + endpoint: JupyterTransport.Endpoint; + channels: JupyterTransport.Channels; + scope: Scope.Closeable; + }; + pending?: { id: string; done: Deferred.Deferred }; + cancelled: boolean; + requests: Map; + observers: Set>; + }; + const sessions = new Map(); + const authorize = Effect.fn("NotebookService.authorize")(function* ( + caller: Caller, + execute = false, + ) { + if ( + !caller.scopes.includes("filesystem:read") || + (execute && !caller.scopes.includes("terminal:operate")) + ) + return yield* new NotebookError({ reason: "permission-denied" }); + }); + const resolve = Effect.fn("NotebookService.resolve")(function* (input: NotebookIdentity) { + const project = yield* projects + .get(input.projectId) + .pipe(Effect.mapError(() => new NotebookError({ reason: "not-found" }))); + if (Option.isNone(project) || project.value.deletedAt !== null) + return yield* new NotebookError({ reason: "not-found" }); + const root = yield* fs + .realPath(project.value.workspaceRoot) + .pipe(Effect.mapError(() => new NotebookError({ reason: "invalid-path" }))); + if (NodePath.isAbsolute(input.relativePath) || !input.relativePath.endsWith(".ipynb")) + return yield* new NotebookError({ reason: "invalid-path" }); + const path = yield* fs + .realPath(NodePath.resolve(root, input.relativePath)) + .pipe(Effect.mapError(() => new NotebookError({ reason: "invalid-path" }))); + const relative = NodePath.relative(root, path); + if ( + !relative || + relative === ".." || + relative.startsWith(`..${NodePath.sep}`) || + NodePath.isAbsolute(relative) + ) + return yield* new NotebookError({ reason: "invalid-path" }); + return { root, path }; + }); + const get = Effect.fn("NotebookService.get")(function* ( + input: NotebookHandle, + caller: Caller, + execute = false, + ) { + yield* authorize(caller, execute); + const session = sessions.get(input.handle); + if (!session || session.projectId !== input.projectId) + return yield* new NotebookError({ reason: "not-found" }); + const resolved = yield* resolve({ + projectId: input.projectId, + relativePath: NodePath.relative(session.root, session.path), + }); + if (resolved.path !== session.path || resolved.root !== session.root) + return yield* new NotebookError({ reason: "invalid-path" }); + return session; + }); + const snapshotOf = (session: Session) => structuredClone(session.snapshot); + const changed = ( + session: Session, + cellId: string | null = null, + patches: readonly NotebookOutputPatch[] = [], + ) => { + const s = session.snapshot; + s.sequence++; + s.revision++; + s.truncated = session.reducer.truncated; + const delta: NotebookEvent = { + type: "changed", + patches, + kernel: s.kernel, + outcome: s.outcome, + truncated: s.truncated, + epoch, + sequence: s.sequence, + revision: s.revision, + generation: s.generation, + runId: cellId === null ? s.runId : (s.executions[cellId]?.runId ?? s.runId), + cellId, + }; + const event: NotebookEvent = + byteLength(delta) > NOTEBOOK_LIMITS.frameBytes + ? { type: "resync-required", epoch, sequence: s.sequence } + : delta; + for (const queue of session.observers) { + if (!Queue.offerUnsafe(queue, event)) { + for (let i = 0; i < 32; i++) Queue.takeUnsafe(queue); + Queue.offerUnsafe(queue, { type: "resync-required", epoch, sequence: s.sequence }); + } + } + }; + const unknown = Effect.fn("NotebookService.disconnected")(function* (session: Session) { + session.cancelled = true; + session.snapshot.kernel = "unknown"; + session.snapshot.outcome = "unknown"; + if (session.pending) yield* Deferred.succeed(session.pending.done, true); + changed(session); + }); + const connectOwned = Effect.fn("NotebookService.connectOwned")(function* ( + session: Session, + endpoint: JupyterTransport.Endpoint, + id: string, + ) { + const kernelScope = yield* Scope.make(); + const channels = yield* transport.connect(endpoint, id).pipe( + Effect.provideService(Scope.Scope, kernelScope), + Effect.onError(() => + Effect.all( + [ + Scope.close(kernelScope, Exit.void), + transport.control(endpoint, id, "shutdown").pipe(Effect.ignore), + unknown(session), + ], + { discard: true }, + ), + ), + ); + session.kernel = { id, endpoint, channels, scope: kernelScope }; + const generation = session.snapshot.generation; + yield* Stream.runForEach(channels.messages, (message) => + Effect.gen(function* () { + if (session.snapshot.generation !== generation) return; + const result = session.reducer.reduce(message); + if (!result) return; + changed(session, result.cellId, result.patches); + const pending = session.pending; + if (result.complete && pending && pending.id === message.parent_header.msg_id) + yield* Deferred.succeed(pending.done, result.failed); + }), + ).pipe( + Effect.catch(() => + session.snapshot.generation === generation ? unknown(session) : Effect.void, + ), + Effect.forkIn(kernelScope), + ); + }); + const open = Effect.fn("NotebookService.open")(function* ( + input: NotebookIdentity, + caller: Caller, + ) { + yield* authorize(caller); + return yield* lock.withPermit( + Effect.gen(function* () { + const { root, path } = yield* resolve(input); + for (const session of sessions.values()) + if ( + session.path === path && + session.root === root && + session.projectId === input.projectId + ) + return snapshotOf(session); + if (sessions.size >= NOTEBOOK_LIMITS.sessions) + return yield* new NotebookError({ reason: "output-limit" }); + const stat = yield* fs + .stat(path) + .pipe(Effect.mapError(() => new NotebookError({ reason: "not-found" }))); + if (stat.size > NOTEBOOK_LIMITS.bytes) + return yield* new NotebookError({ reason: "output-limit" }); + const text = yield* fs + .readFileString(path) + .pipe(Effect.mapError(() => new NotebookError({ reason: "not-found" }))); + const parsed = yield* Effect.try({ + try: () => parseNotebook(text), + catch: () => new NotebookError({ reason: "invalid-document" }), + }); + const handle = NodeCrypto.randomUUID(); + const snapshot: NotebookSnapshot = { + handle, + epoch, + sequence: 0, + revision: parsed.changed ? 1 : 0, + diskHash: NodeCrypto.createHash("sha256").update(text).digest("hex"), + generation: 0, + executions: {}, + document: parsed.document, + kernel: "stopped", + runId: null, + outcome: "none", + truncated: false, + }; + const session: Session = { + projectId: input.projectId, + path, + root, + snapshot, + reducer: new ExecutionReducer(parsed.document), + cancelled: false, + requests: new Map(), + observers: new Set(), + }; + sessions.set(handle, session); + return snapshotOf(session); + }), + ); + }); + const snapshot = Effect.fn("NotebookService.snapshot")(function* ( + input: NotebookHandle, + caller: Caller, + ) { + return snapshotOf(yield* get(input, caller)); + }); + const execute = Effect.fn("NotebookService.execute")(function* ( + input: NotebookExecute, + caller: Caller, + ) { + return yield* lock.withPermit( + Effect.gen(function* () { + const session = yield* get(input, caller, true); + const fingerprint = fingerprintOf([input.revision, input.cellIds]); + const prior = session.requests.get(input.requestId); + if (prior) { + if (prior.fingerprint !== fingerprint) + return yield* new NotebookError({ reason: "stale-revision" }); + return prior.runId; + } + if (session.snapshot.outcome === "running") + return yield* new NotebookError({ reason: "busy" }); + if (session.snapshot.kernel === "unknown") + return yield* new NotebookError({ reason: "disconnected" }); + if (input.revision !== session.snapshot.revision) + return yield* new NotebookError({ reason: "stale-revision" }); + const cells = input.cellIds.map((id) => + session.reducer.document.cells.find((c) => c.id === id), + ); + if ( + !input.requestId || + input.requestId.length > 128 || + !cells.length || + cells.length > NOTEBOOK_LIMITS.cells || + new Set(input.cellIds).size !== cells.length || + cells.some((c) => !c || c.cell_type !== "code") + ) + return yield* new NotebookError({ reason: "invalid-document" }); + if (session.requests.size >= 128) + return yield* new NotebookError({ reason: "output-limit" }); + if (!session.kernel) { + const endpoint = yield* config.endpoint(session.root); + const id = yield* transport.create(endpoint, session.root); + session.snapshot.generation++; + yield* connectOwned(session, endpoint, id); + } + const kernel = session.kernel; + if (!kernel) return yield* new NotebookError({ reason: "kernel-failure" }); + const runId = NodeCrypto.randomUUID(); + session.snapshot.runId = runId; + session.snapshot.outcome = "running"; + session.snapshot.kernel = "busy"; + session.cancelled = false; + session.requests.set(input.requestId, { fingerprint, runId }); + changed(session); + const generation = session.snapshot.generation; + const sourceRevision = input.revision; + yield* Effect.gen(function* () { + let failed = false; + for (const cell of cells) { + if (!cell || session.cancelled || session.snapshot.generation !== generation) break; + const id = NodeCrypto.randomUUID(); + const done = yield* Deferred.make(); + session.pending = { id, done }; + session.snapshot.executions = { + ...session.snapshot.executions, + [cell.id]: { + runId, + generation, + sourceRevision, + sourceHash: NodeCrypto.createHash("sha256") + .update(typeof cell.source === "string" ? cell.source : cell.source.join("")) + .digest("hex"), + }, + }; + session.reducer.begin(id, cell); + changed(session, cell.id, [{ type: "clear", cellId: cell.id }]); + yield* kernel.channels.execute( + id, + typeof cell.source === "string" ? cell.source : cell.source.join(""), + ); + failed = yield* Deferred.await(done); + if (session.pending?.id === id) delete session.pending; + if (failed) break; + } + if ( + session.snapshot.generation === generation && + session.snapshot.outcome === "running" + ) { + session.snapshot.outcome = session.cancelled ? "aborted" : failed ? "error" : "ok"; + session.snapshot.kernel = "idle"; + changed(session); + } + }).pipe( + Effect.catch(() => + session.snapshot.generation === generation ? unknown(session) : Effect.void, + ), + Effect.forkIn(scope), + ); + return runId; + }), + ); + }); + const control = (action: "interrupt" | "restart" | "shutdown") => + Effect.fn(`NotebookService.${action}`)(function* ( + input: NotebookControl & { runId?: string }, + caller: Caller, + ) { + yield* lock.withPermit( + Effect.gen(function* () { + const session = yield* get(input, caller, true); + if (session.snapshot.generation !== input.generation) + return yield* new NotebookError({ reason: "stale-generation" }); + if (action === "interrupt" && input.runId !== session.snapshot.runId) + return yield* new NotebookError({ reason: "stale-revision" }); + const kernel = session.kernel; + if (!kernel) { + if (action !== "interrupt") { + session.snapshot.kernel = "stopped"; + session.snapshot.outcome = "aborted"; + changed(session); + } + return; + } + session.cancelled = true; + yield* transport + .control(kernel.endpoint, kernel.id, action) + .pipe(Effect.onError(() => unknown(session))); + if (action !== "interrupt") { + session.snapshot.generation++; + session.reducer.owners.clear(); + if (session.pending) yield* Deferred.succeed(session.pending.done, true); + yield* Scope.close(kernel.scope, Exit.void); + delete session.kernel; + if (action === "restart") yield* connectOwned(session, kernel.endpoint, kernel.id); + session.snapshot.kernel = action === "restart" ? "idle" : "stopped"; + session.snapshot.outcome = "aborted"; + } + changed(session); + }), + ); + }); + const subscribe = (input: NotebookHandle, caller: Caller) => + Stream.unwrap( + Effect.gen(function* () { + const session = yield* get(input, caller); + const queue = yield* Queue.dropping(32); + session.observers.add(queue); + const initial = snapshotOf(session); + yield* Effect.addFinalizer(() => + Effect.sync(() => { + session.observers.delete(queue); + Queue.shutdownUnsafe(queue); + }), + ); + return Stream.concat( + Stream.make({ type: "snapshot", snapshot: initial } as NotebookEvent), + Stream.fromQueue(queue), + ); + }), + ); + yield* Effect.addFinalizer(() => + Effect.forEach( + sessions.values(), + (session) => + Effect.gen(function* () { + const kernel = session.kernel; + if (!kernel) return; + session.snapshot.generation++; + yield* Scope.close(kernel.scope, Exit.void); + yield* transport.control(kernel.endpoint, kernel.id, "shutdown").pipe(Effect.ignore); + }), + { discard: true }, + ), + ); + + return NotebookService.of({ + open, + snapshot, + execute, + interrupt: control("interrupt"), + restart: control("restart"), + shutdown: control("shutdown"), + subscribe, + }); +}); +export const layer = Layer.effect(NotebookService, make); +export const layerDisabled = layer.pipe( + Layer.provide(disabled), + Layer.provide(JupyterTransport.layer), +); diff --git a/apps/server/src/notebook/executionReducer.test.ts b/apps/server/src/notebook/executionReducer.test.ts new file mode 100644 index 0000000000000..4629c8e311d35 --- /dev/null +++ b/apps/server/src/notebook/executionReducer.test.ts @@ -0,0 +1,97 @@ +import { expect, it } from "@effect/vitest"; +import { NOTEBOOK_LIMITS } from "@t3tools/contracts"; +import { ExecutionReducer } from "./executionReducer.ts"; +import { parseNotebook, byteLength } from "./notebookDocument.ts"; + +it.each([true, false])("requires both reply and idle, reply first: %s", (replyFirst) => { + const document = parseNotebook( + JSON.stringify({ + nbformat: 4, + nbformat_minor: 5, + metadata: {}, + cells: [ + { + id: "a", + cell_type: "code", + source: "1", + metadata: {}, + execution_count: null, + outputs: [], + }, + ], + }), + ).document; + const reducer = new ExecutionReducer(document); + reducer.begin("request", document.cells[0]!); + const reply = { + header: { msg_type: "execute_reply" }, + parent_header: { msg_id: "request" }, + channel: "shell", + content: { status: "error" }, + }; + const idle = { + header: { msg_type: "status" }, + parent_header: { msg_id: "request" }, + channel: "iopub", + content: { execution_state: "idle" }, + }; + expect(reducer.reduce({ ...idle, parent_header: { msg_id: "other" } })).toBeUndefined(); + expect(reducer.reduce(replyFirst ? reply : idle)?.complete).toBe(false); + expect(reducer.reduce(replyFirst ? idle : reply)).toMatchObject({ complete: true, failed: true }); +}); + +it("updates displays across cells, defers clear, strips transient IDs, and bounds floods", () => { + const document = parseNotebook( + JSON.stringify({ + nbformat: 4, + nbformat_minor: 5, + metadata: {}, + cells: ["a", "b"].map((id) => ({ + id, + cell_type: "code", + source: "", + metadata: {}, + execution_count: null, + outputs: [], + })), + }), + ).document; + const reducer = new ExecutionReducer(document); + const send = (id: string, type: string, content: Record) => + reducer.reduce({ + header: { msg_type: type }, + parent_header: { msg_id: id }, + channel: "iopub", + content, + }); + reducer.begin("one", document.cells[0]!); + reducer.begin("two", document.cells[1]!); + for (const id of ["one", "two"]) + send(id, "display_data", { + data: { "text/plain": "old" }, + metadata: { keep: true }, + transient: { display_id: "display" }, + }); + send("two", "update_display_data", { + data: { "text/plain": "new" }, + metadata: {}, + transient: { display_id: "display" }, + }); + for (const cell of document.cells) + expect(cell.outputs).toEqual([ + { output_type: "display_data", data: { "text/plain": "new" }, metadata: {} }, + ]); + send("one", "clear_output", { wait: true }); + expect(document.cells[0]?.outputs).toHaveLength(1); + send("one", "stream", { name: "stdout", text: "next" }); + expect(document.cells[0]?.outputs).toEqual([ + { output_type: "stream", name: "stdout", text: "next" }, + ]); + for (let i = 0; i < 12; i++) + send("one", "stream", { name: "stdout", text: "x".repeat(NOTEBOOK_LIMITS.frameBytes) }); + expect(byteLength(document)).toBeLessThanOrEqual(NOTEBOOK_LIMITS.bytes); + expect(reducer.truncated).toBe(true); + expect(document.cells[0]?.outputs?.at(-1)).toMatchObject({ t3_truncated: true }); + reducer.begin("replacement", document.cells[0]!); + expect(send("one", "stream", { name: "stdout", text: "obsolete" })).toBeUndefined(); +}); diff --git a/apps/server/src/notebook/executionReducer.ts b/apps/server/src/notebook/executionReducer.ts new file mode 100644 index 0000000000000..3f32e8216d5eb --- /dev/null +++ b/apps/server/src/notebook/executionReducer.ts @@ -0,0 +1,121 @@ +import { NOTEBOOK_LIMITS, type NotebookOutputPatch } from "@t3tools/contracts"; +import type { Message } from "./JupyterTransport.ts"; +import { byteLength, type Document, type NotebookCell } from "./notebookDocument.ts"; + +type Owner = { cell: NotebookCell; reply: boolean; idle: boolean; failed: boolean; clear: boolean }; +export class ExecutionReducer { + readonly owners = new Map(); + readonly displays = new WeakMap(); + truncated = false; + readonly document: Document; + constructor(document: Document) { + this.document = document; + } + begin(id: string, cell: NotebookCell) { + for (const [key, owner] of this.owners) if (owner.cell === cell) this.owners.delete(key); + cell.outputs = []; + cell.execution_count = null; + this.owners.set(id, { cell, reply: false, idle: false, failed: false, clear: false }); + } + reduce(message: Message) { + const owner = this.owners.get(message.parent_header.msg_id ?? ""); + if (!owner) return undefined; + const { cell } = owner; + const patches: NotebookOutputPatch[] = []; + const content = message.content; + const type = message.header.msg_type; + if (type === "execute_reply" && message.channel === "shell") { + owner.reply = true; + owner.failed = content.status !== "ok"; + } + if (type === "status" && message.channel === "iopub" && content.execution_state === "idle") + owner.idle = true; + if (message.channel === "iopub") { + if (type === "execute_input" && typeof content.execution_count === "number") { + cell.execution_count = content.execution_count; + patches.push({ type: "count", cellId: cell.id, count: content.execution_count }); + } + if (type === "clear_output") { + if (content.wait === true) owner.clear = true; + else { + cell.outputs = []; + patches.push({ type: "clear", cellId: cell.id }); + owner.clear = false; + } + } + if ( + ["stream", "display_data", "execute_result", "error", "update_display_data"].includes(type) + ) { + if (owner.clear) { + cell.outputs = []; + patches.push({ type: "clear", cellId: cell.id }); + owner.clear = false; + } + const transient = content.transient as { display_id?: unknown } | undefined; + const displayId = + typeof transient?.display_id === "string" ? transient.display_id : undefined; + const { transient: _transient, ...saved } = content; + if (type === "update_display_data") { + if (displayId) + for (const other of this.document.cells) + for (const output of other.outputs ?? []) { + if (this.displays.get(output) !== displayId) continue; + const prior = { ...output }; + Object.assign(output, { data: saved.data, metadata: saved.metadata }); + if (byteLength(this.document) > NOTEBOOK_LIMITS.bytes - 256) { + Object.assign(output, prior); + this.truncated = true; + } else + patches.push({ + type: "replace", + cellId: other.id, + index: other.outputs!.indexOf(output), + output: structuredClone(output), + }); + } + } else { + const outputs = (cell.outputs ??= []); + const previous = outputs.at(-1); + if ( + type === "stream" && + typeof saved.text === "string" && + previous?.output_type === "stream" && + previous.name === saved.name + ) { + const text = previous.text; + previous.text = String(text) + saved.text; + if (byteLength(this.document) > NOTEBOOK_LIMITS.bytes - 256) { + previous.text = text; + this.truncated = true; + } else patches.push({ type: "stream", cellId: cell.id, text: saved.text }); + } else { + const output = { ...saved, output_type: type }; + outputs.push(output); + if ( + outputs.length > NOTEBOOK_LIMITS.outputs - 1 || + byteLength(this.document) > NOTEBOOK_LIMITS.bytes - 256 + ) { + outputs.pop(); + this.truncated = true; + } else { + patches.push({ type: "append", cellId: cell.id, output: structuredClone(output) }); + if (displayId) this.displays.set(output, displayId); + } + } + if (this.truncated && !outputs.some((o) => o.t3_truncated === true)) { + const marker = { + output_type: "stream", + name: "stderr", + text: "[T3 output truncated]\n", + t3_truncated: true, + }; + outputs.push(marker); + if (byteLength(this.document) > NOTEBOOK_LIMITS.bytes) outputs.pop(); + else patches.push({ type: "append", cellId: cell.id, output: marker }); + } + } + } + } + return { patches, cellId: cell.id, complete: owner.reply && owner.idle, failed: owner.failed }; + } +} diff --git a/apps/server/src/notebook/notebookDocument.ts b/apps/server/src/notebook/notebookDocument.ts new file mode 100644 index 0000000000000..f056e10154a08 --- /dev/null +++ b/apps/server/src/notebook/notebookDocument.ts @@ -0,0 +1,62 @@ +// @effect-diagnostics nodeBuiltinImport:off +// Missing nbformat cell IDs are assigned without altering existing JSON fields. +import * as NodeCrypto from "node:crypto"; +import { NOTEBOOK_LIMITS, NotebookError, type NotebookDocument } from "@t3tools/contracts"; +import * as Schema from "effect/Schema"; + +const source = Schema.Union([Schema.String, Schema.Array(Schema.String)]); +const cellSchema = Schema.Struct({ + cell_type: Schema.Literals(["code", "markdown", "raw"]), + source, + metadata: Schema.Record(Schema.String, Schema.Unknown), + id: Schema.optional(Schema.String), + outputs: Schema.optional(Schema.Array(Schema.Record(Schema.String, Schema.Unknown))), +}); +const documentSchema = Schema.Struct({ + nbformat: Schema.Literal(4), + nbformat_minor: Schema.Int, + metadata: Schema.Record(Schema.String, Schema.Unknown), + cells: Schema.Array(cellSchema), +}); +export type NotebookCell = Omit & { + id: string; + outputs?: Array>; + execution_count?: number | null; +}; +export type Document = NotebookDocument & { cells: NotebookCell[] }; +export const byteLength = (value: unknown) => Buffer.byteLength(JSON.stringify(value)); +export function parseNotebook(text: string): { document: Document; changed: boolean } { + try { + if (Buffer.byteLength(text) > NOTEBOOK_LIMITS.bytes) throw new Error(); + const value: unknown = JSON.parse(text); + const checked = Schema.decodeUnknownSync(documentSchema)(value); + if (checked.cells.length > NOTEBOOK_LIMITS.cells || checked.nbformat_minor < 0) + throw new Error(); + const document = value as Document; + const ids = new Set(); + let changed = false; + for (const cell of document.cells) { + if (cell.id === undefined) { + cell.id = NodeCrypto.randomUUID(); + changed = true; + } + if ( + !/^[a-zA-Z0-9_-]{1,64}$/.test(cell.id) || + ids.has(cell.id) || + (cell.outputs?.length ?? 0) > NOTEBOOK_LIMITS.outputs + ) + throw new Error(); + ids.add(cell.id); + if ( + cell.cell_type === "code" && + (!cell.outputs || + !(cell.execution_count === null || Number.isInteger(cell.execution_count))) + ) + throw new Error(); + } + if (byteLength(document) > NOTEBOOK_LIMITS.bytes) throw new Error(); + return { document, changed }; + } catch { + throw new NotebookError({ reason: "invalid-document" }); + } +} diff --git a/apps/server/src/observability/RpcInstrumentation.ts b/apps/server/src/observability/RpcInstrumentation.ts index 803e0a11f722d..6993afe142cf5 100644 --- a/apps/server/src/observability/RpcInstrumentation.ts +++ b/apps/server/src/observability/RpcInstrumentation.ts @@ -1,3 +1,4 @@ +import { NOTEBOOK_METHODS } from "@t3tools/contracts"; import { PROJECT_ARTIFACT_METHODS } from "@t3tools/contracts"; import { ORCHESTRATION_V2_WS_METHODS, WS_METHODS, type WsRpcGroup } from "@t3tools/contracts"; import * as Effect from "effect/Effect"; @@ -12,6 +13,13 @@ type WsRpcMethod = RpcGroup.Rpcs["_tag"]; // Trace queries and dashboards depend on these labels, including deviations from method prefixes. const RPC_AGGREGATES = { + [NOTEBOOK_METHODS.open]: "notebook", + [NOTEBOOK_METHODS.snapshot]: "notebook", + [NOTEBOOK_METHODS.execute]: "notebook", + [NOTEBOOK_METHODS.interrupt]: "notebook", + [NOTEBOOK_METHODS.restart]: "notebook", + [NOTEBOOK_METHODS.shutdown]: "notebook", + [NOTEBOOK_METHODS.subscribe]: "notebook", [PROJECT_ARTIFACT_METHODS.list]: "projectArtifacts", [PROJECT_ARTIFACT_METHODS.add]: "projectArtifacts", [PROJECT_ARTIFACT_METHODS.remove]: "projectArtifacts", diff --git a/apps/server/src/server.ts b/apps/server/src/server.ts index d6b384ff2fd6e..20c77a1d1329e 100644 --- a/apps/server/src/server.ts +++ b/apps/server/src/server.ts @@ -1,3 +1,4 @@ +import * as NotebookService from "./notebook/NotebookService.ts"; import { services as gitHubServices } from "@t3tools/source-control-github/server/driver"; import * as ResetCreditExpiry from "./provider/ResetCreditExpiry.ts"; import * as ProjectArtifacts from "./project/ProjectArtifacts.ts"; @@ -758,6 +759,7 @@ const layerIssueRoadmap = IssueRoadmap.layer.pipe( ); const layerRouterServices = Layer.mergeAll( ProjectArtifacts.layer, + NotebookService.layerDisabled, ThreadWorkspaceService.layer, layerIssueRoadmap, IssueCleanup.layer.pipe( diff --git a/apps/server/src/ws.test.ts b/apps/server/src/ws.test.ts index 7c6e68609817a..fb2fcdfec3d9c 100644 --- a/apps/server/src/ws.test.ts +++ b/apps/server/src/ws.test.ts @@ -1,3 +1,4 @@ +import * as NotebookService from "./notebook/NotebookService.ts"; import * as ProjectArtifacts from "./project/ProjectArtifacts.ts"; import * as NodeServices from "@effect/platform-node/NodeServices"; import { assert, it } from "@effect/vitest"; @@ -524,6 +525,7 @@ const layerUnusedConfigRpcServices = Layer.mergeAll( Layer.mock(IssueWorkService.IssueWorkService)({}), Layer.mock(SidebarLayout.SidebarLayoutService)({}), Layer.mock(ProjectArtifacts.ProjectArtifacts)({}), + Layer.mock(NotebookService.NotebookService)({}), Layer.mock(ThreadWorkspaceService.ThreadWorkspaceService)({}), Layer.mock(DocumentText.DocumentText)({}), Layer.mock(ServerSecretStore.ServerSecretStore)({}), diff --git a/packages/contracts/src/clientRpcPermissions.ts b/packages/contracts/src/clientRpcPermissions.ts index 65948c2512c04..081a85e458901 100644 --- a/packages/contracts/src/clientRpcPermissions.ts +++ b/packages/contracts/src/clientRpcPermissions.ts @@ -1,9 +1,11 @@ +import { NOTEBOOK_METHODS } from "./notebookRpc.ts"; import { PROJECT_ARTIFACT_METHODS } from "./projectArtifactRpc.ts"; import { type ServerSettingsPatch, requiredScopesForServerSettingsPatch } from "./settings.ts"; import * as Schema from "effect/Schema"; import { GitPreparePullRequestThreadInput, GitRunStackedActionInput } from "./git.ts"; import { AuthAccessWriteScope, + AuthFilesystemReadScope, AuthOrchestrationOperateScope, AuthProvidersManageScope, AuthOrchestrationReadScope, @@ -18,6 +20,10 @@ import { WS_METHODS } from "./rpc.ts"; /** Incremental client enforcement; the server still authorizes every request. */ export const CLIENT_GUARDED_RPC_SCOPES = { + [NOTEBOOK_METHODS.execute]: AuthTerminalOperateScope, + [NOTEBOOK_METHODS.interrupt]: AuthTerminalOperateScope, + [NOTEBOOK_METHODS.restart]: AuthTerminalOperateScope, + [NOTEBOOK_METHODS.shutdown]: AuthTerminalOperateScope, [PROJECT_ARTIFACT_METHODS.add]: AuthOrchestrationOperateScope, [PROJECT_ARTIFACT_METHODS.remove]: AuthOrchestrationOperateScope, [PROJECT_ARTIFACT_METHODS.move]: AuthOrchestrationOperateScope, @@ -103,6 +109,14 @@ export function clientRpcRequiredScopes( method: string, input: unknown, ): readonly AuthEnvironmentScope[] { + if ( + method === NOTEBOOK_METHODS.execute || + method === NOTEBOOK_METHODS.interrupt || + method === NOTEBOOK_METHODS.restart || + method === NOTEBOOK_METHODS.shutdown + ) { + return [AuthFilesystemReadScope, AuthTerminalOperateScope]; + } if (method === WS_METHODS.serverUpdateSettings && input !== undefined) { const payload = (typeof input === "object" && input !== null ? input : {}) as { patch?: unknown; diff --git a/packages/contracts/src/index.ts b/packages/contracts/src/index.ts index dac5e10fb6f27..a2efc92db4345 100644 --- a/packages/contracts/src/index.ts +++ b/packages/contracts/src/index.ts @@ -89,3 +89,6 @@ export * from "./ownerDecision.ts"; export * from "./repositoryOrchestrator.ts"; export * from "./ResetCreditExpiry.ts"; + +export * from "./notebook.ts"; +export * from "./notebookRpc.ts"; diff --git a/packages/contracts/src/notebook.ts b/packages/contracts/src/notebook.ts new file mode 100644 index 0000000000000..490663e19355e --- /dev/null +++ b/packages/contracts/src/notebook.ts @@ -0,0 +1,112 @@ +import * as Schema from "effect/Schema"; +import { ProjectId } from "./baseSchemas.ts"; + +export const NOTEBOOK_LIMITS = { + bytes: 8 * 1024 * 1024, + cells: 500, + outputs: 100, + frameBytes: 1024 * 1024, + sessions: 16, +} as const; +export class NotebookError extends Schema.TaggedError()("NotebookError", { + reason: Schema.Literals([ + "unavailable", + "invalid-document", + "permission-denied", + "not-found", + "busy", + "stale-revision", + "stale-generation", + "disconnected", + "kernel-failure", + "output-limit", + "invalid-path", + ]), +}) { + override get message() { + return `Notebook operation failed: ${this.reason}.`; + } +} +// The JSON record owns unknown nbformat fields; preview models must never be saved. +export const NotebookDocument = Schema.Record(Schema.String, Schema.Unknown); +export type NotebookDocument = typeof NotebookDocument.Type; +export const NotebookIdentity = Schema.Struct({ + projectId: ProjectId, + relativePath: Schema.String, +}); +export type NotebookIdentity = typeof NotebookIdentity.Type; +export const NotebookHandle = Schema.Struct({ projectId: ProjectId, handle: Schema.String }); +export type NotebookHandle = typeof NotebookHandle.Type; +export const NotebookExecute = Schema.Struct({ + ...NotebookHandle.fields, + revision: Schema.Int, + cellIds: Schema.Array(Schema.String), + requestId: Schema.String, +}); +export type NotebookExecute = typeof NotebookExecute.Type; +export const NotebookControl = Schema.Struct({ ...NotebookHandle.fields, generation: Schema.Int }); +export type NotebookControl = typeof NotebookControl.Type; +export const NotebookInterrupt = Schema.Struct({ ...NotebookControl.fields, runId: Schema.String }); +export type NotebookInterrupt = typeof NotebookInterrupt.Type; +export const NotebookSnapshot = Schema.Struct({ + handle: Schema.String, + epoch: Schema.String, + sequence: Schema.Int, + revision: Schema.Int, + diskHash: Schema.String, + executions: Schema.Record( + Schema.String, + Schema.Struct({ + runId: Schema.String, + generation: Schema.Int, + sourceHash: Schema.String, + sourceRevision: Schema.Int, + }), + ), + generation: Schema.Int, + document: NotebookDocument, + kernel: Schema.Literals(["stopped", "idle", "busy", "unknown"]), + runId: Schema.NullOr(Schema.String), + outcome: Schema.Literals(["none", "running", "ok", "error", "aborted", "unknown"]), + truncated: Schema.Boolean, +}); +export type NotebookSnapshot = typeof NotebookSnapshot.Type; +export const NotebookOutputPatch = Schema.Union([ + Schema.Struct({ type: Schema.Literal("clear"), cellId: Schema.String }), + Schema.Struct({ type: Schema.Literal("count"), cellId: Schema.String, count: Schema.Int }), + Schema.Struct({ + type: Schema.Literal("append"), + cellId: Schema.String, + output: Schema.Record(Schema.String, Schema.Unknown), + }), + Schema.Struct({ type: Schema.Literal("stream"), cellId: Schema.String, text: Schema.String }), + Schema.Struct({ + type: Schema.Literal("replace"), + cellId: Schema.String, + index: Schema.Int, + output: Schema.Record(Schema.String, Schema.Unknown), + }), +]); +export type NotebookOutputPatch = typeof NotebookOutputPatch.Type; +export const NotebookEvent = Schema.Union([ + Schema.Struct({ + type: Schema.Literal("resync-required"), + epoch: Schema.String, + sequence: Schema.Int, + }), + Schema.Struct({ type: Schema.Literal("snapshot"), snapshot: NotebookSnapshot }), + Schema.Struct({ + type: Schema.Literal("changed"), + patches: Schema.Array(NotebookOutputPatch), + kernel: NotebookSnapshot.fields.kernel, + outcome: NotebookSnapshot.fields.outcome, + truncated: Schema.Boolean, + epoch: Schema.String, + sequence: Schema.Int, + revision: Schema.Int, + generation: Schema.Int, + runId: Schema.NullOr(Schema.String), + cellId: Schema.NullOr(Schema.String), + }), +]); +export type NotebookEvent = typeof NotebookEvent.Type; diff --git a/packages/contracts/src/notebookRpc.ts b/packages/contracts/src/notebookRpc.ts new file mode 100644 index 0000000000000..b3e223991291b --- /dev/null +++ b/packages/contracts/src/notebookRpc.ts @@ -0,0 +1,41 @@ +import * as Rpc from "effect/rpc/Rpc"; +import * as Schema from "effect/Schema"; +import { EnvironmentAuthorizationError } from "./auth.ts"; +import { + NotebookIdentity, + NotebookHandle, + NotebookExecute, + NotebookControl, + NotebookInterrupt, + NotebookSnapshot, + NotebookEvent, + NotebookError, +} from "./notebook.ts"; +export const NOTEBOOK_METHODS = { + open: "notebook.open", + snapshot: "notebook.snapshot", + execute: "notebook.execute", + interrupt: "notebook.interrupt", + restart: "notebook.restart", + shutdown: "notebook.shutdown", + subscribe: "notebook.subscribe", +} as const; +const error = Schema.Union([NotebookError, EnvironmentAuthorizationError]); +export const notebookRpcs = [ + Rpc.make(NOTEBOOK_METHODS.open, { payload: NotebookIdentity, success: NotebookSnapshot, error }), + Rpc.make(NOTEBOOK_METHODS.snapshot, { + payload: NotebookHandle, + success: NotebookSnapshot, + error, + }), + Rpc.make(NOTEBOOK_METHODS.execute, { payload: NotebookExecute, success: Schema.String, error }), + Rpc.make(NOTEBOOK_METHODS.interrupt, { payload: NotebookInterrupt, success: Schema.Void, error }), + Rpc.make(NOTEBOOK_METHODS.restart, { payload: NotebookControl, success: Schema.Void, error }), + Rpc.make(NOTEBOOK_METHODS.shutdown, { payload: NotebookControl, success: Schema.Void, error }), + Rpc.make(NOTEBOOK_METHODS.subscribe, { + payload: NotebookHandle, + success: NotebookEvent, + error, + stream: true, + }), +] as const; diff --git a/packages/contracts/src/rpc.ts b/packages/contracts/src/rpc.ts index 45b1436721f26..c0b89119c53a2 100644 --- a/packages/contracts/src/rpc.ts +++ b/packages/contracts/src/rpc.ts @@ -1,3 +1,4 @@ +import { NOTEBOOK_METHODS, notebookRpcs } from "./notebookRpc.ts"; import { OwnerDecision, OwnerDecisionsListInput, @@ -2276,6 +2277,7 @@ const WsTerminalWatchInlineRunRpc = Rpc.make(WS_METHODS.terminalWatchInlineRun, }); const forkWsRpcs = [ + ...notebookRpcs, ...projectArtifactRpcs, WsSidebarLayoutUpdateRpc, WsSubscribeSidebarLayoutRpc, @@ -2527,6 +2529,7 @@ export const WsRpcGroup = RpcGroup.make( ).middleware(RpcScopeAuthorization); export const ForkWsRpcGroup = RpcGroup.make(...forkWsRpcs).middleware(RpcScopeAuthorization); export const UpstreamWsRpcGroup = WsRpcGroup.omit( + ...Object.values(NOTEBOOK_METHODS), ...Object.values(PROJECT_ARTIFACT_METHODS), WS_METHODS.sidebarLayoutUpdate, WS_METHODS.subscribeSidebarLayout, From ac177ebfd4fd1dc127eb1466fac5562d60254bf3 Mon Sep 17 00:00:00 2001 From: Tanner Polley <42780259+tannerpolley@users.noreply.github.com> Date: Wed, 28 Oct 2026 14:29:46 -0600 Subject: [PATCH 2/4] fix(notebook): fence kernel ownership and bound transport admission --- .../src/notebook/JupyterTransport.test.ts | 158 +++++++++++ apps/server/src/notebook/JupyterTransport.ts | 140 +++++++--- .../src/notebook/NotebookService.test.ts | 261 +++++++++++++++++- apps/server/src/notebook/NotebookService.ts | 171 +++++++----- .../src/notebook/executionReducer.test.ts | 41 +++ apps/server/src/notebook/executionReducer.ts | 39 ++- apps/server/src/notebook/kernelSocket.ts | 195 +++++++++++++ packages/contracts/src/notebook.ts | 22 +- 8 files changed, 893 insertions(+), 134 deletions(-) create mode 100644 apps/server/src/notebook/JupyterTransport.test.ts create mode 100644 apps/server/src/notebook/kernelSocket.ts diff --git a/apps/server/src/notebook/JupyterTransport.test.ts b/apps/server/src/notebook/JupyterTransport.test.ts new file mode 100644 index 0000000000000..eb2b63f750ea1 --- /dev/null +++ b/apps/server/src/notebook/JupyterTransport.test.ts @@ -0,0 +1,158 @@ +// @effect-diagnostics nodeBuiltinImport:off +// A loopback protocol fixture exercises framing and admission without a Jupyter process. +import * as NodeHttp from "node:http"; +import * as NodeCrypto from "node:crypto"; +import type * as NodeStream from "node:stream"; +import { expect, it } from "@effect/vitest"; +import * as Effect from "effect/Effect"; +import * as Stream from "effect/Stream"; +import * as Deferred from "effect/Deferred"; +import { NOTEBOOK_LIMITS } from "@t3tools/contracts"; +import * as JupyterTransport from "./JupyterTransport.ts"; +import { ExecutionReducer } from "./executionReducer.ts"; +import { parseNotebook, byteLength } from "./notebookDocument.ts"; + +const frame = (text: string, opcode = 1, final = true) => { + const payload = Buffer.from(text); + const header = Buffer.alloc(payload.length < 126 ? 2 : payload.length <= 65535 ? 4 : 10); + header[0] = (final ? 128 : 0) | opcode; + if (payload.length < 126) header[1] = payload.length; + else if (payload.length <= 65535) { + header[1] = 126; + header.writeUInt16BE(payload.length, 2); + } else { + header[1] = 127; + header.writeBigUInt64BE(BigInt(payload.length), 2); + } + return Buffer.concat([header, payload]); +}; + +it.live.each(["stream", "frame", "burst"] as const)( + "%s overflow retains bounded data without losing completion or control", + (mode) => + Effect.gen(function* () { + let peer: NodeStream.Duplex | undefined; + let interrupts = 0; + const server = NodeHttp.createServer((_request, response) => { + interrupts++; + response.writeHead(204); + response.end(); + }); + server.on("upgrade", (request, socket) => { + peer = socket; + const accept = NodeCrypto.createHash("sha1") + .update(request.headers["sec-websocket-key"] + "258EAFA5-E914-47DA-95CA-C5AB0DC85B11") + .digest("base64"); + socket.write( + "HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: " + + accept + + "\r\n\r\n", + ); + socket.on("data", () => {}); + }); + yield* Effect.addFinalizer(() => + Effect.promise( + () => + new Promise((resolve) => { + peer?.destroy(); + server.close(() => resolve()); + }), + ), + ); + yield* Effect.callback((resume) => { + server.listen(0, "127.0.0.1", () => resume(Effect.void)); + }); + const address = server.address(); + if (!address || typeof address === "string") throw new Error(); + const transport = yield* JupyterTransport.JupyterTransport; + const endpoint = { + url: "http://127.0.0.1:" + address.port, + token: "test", + kernelName: "python3", + }; + const channels = yield* transport.connect(endpoint, "owned"); + const document = parseNotebook( + JSON.stringify({ + nbformat: 4, + nbformat_minor: 5, + metadata: {}, + cells: [ + { + id: "a", + cell_type: "code", + source: "", + metadata: {}, + execution_count: null, + outputs: [], + }, + ], + }), + ).document; + const reducer = new ExecutionReducer(document); + const message = (type: string, content: Record) => + frame( + JSON.stringify({ + header: { msg_type: type }, + parent_header: { msg_id: "request" }, + channel: type === "execute_reply" ? "shell" : "iopub", + content, + }), + ); + reducer.begin("request", document.cells[0]!); + yield* channels.execute("request", "print('flood')"); + if (mode === "stream") + peer!.write( + message("stream", { name: "stdout", text: "x".repeat(NOTEBOOK_LIMITS.frameBytes) }), + ); + if (mode === "burst") + peer!.write( + Buffer.concat( + Array.from({ length: 1000 }, () => + message("stream", { name: "stdout", text: "burst\\n" }), + ), + ), + ); + if (mode === "frame") + peer!.write( + message("display_data", { + data: { "text/plain": "x".repeat(NOTEBOOK_LIMITS.bytes + 65536) }, + metadata: {}, + }), + ); + const reply = JSON.stringify({ + header: { msg_type: "execute_reply" }, + parent_header: { msg_id: "request" }, + channel: "shell", + content: { status: "ok" }, + }); + const split = Math.floor(reply.length / 2); + peer!.write( + Buffer.concat([ + message("status", { execution_state: "idle" }), + frame(reply.slice(0, split), 1, false), + frame("ping", 9), + frame(reply.slice(split), 0), + ]), + ); + const complete = yield* Deferred.make(); + let deliveredBytes = 0; + yield* channels.messages.pipe( + Stream.runForEach((message) => + Effect.gen(function* () { + deliveredBytes += byteLength(message); + if (reducer.reduce(message)?.complete) yield* Deferred.succeed(complete, undefined); + }), + ), + Effect.forkScoped, + ); + yield* Deferred.await(complete); + yield* transport.control(endpoint, "owned", "interrupt"); + expect(interrupts).toBe(1); + expect(reducer.truncated).toBe(true); + expect(document.cells[0]?.outputs).toContainEqual( + expect.objectContaining({ text: "[T3 output truncated]\n" }), + ); + expect(deliveredBytes).toBeLessThan(NOTEBOOK_LIMITS.frameBytes + 65536); + expect(byteLength(document)).toBeLessThanOrEqual(NOTEBOOK_LIMITS.bytes); + }).pipe(Effect.scoped, Effect.provide(JupyterTransport.layer)), +); diff --git a/apps/server/src/notebook/JupyterTransport.ts b/apps/server/src/notebook/JupyterTransport.ts index 4892e0fcee12e..50d56a1c1e08b 100644 --- a/apps/server/src/notebook/JupyterTransport.ts +++ b/apps/server/src/notebook/JupyterTransport.ts @@ -1,7 +1,7 @@ // @effect-diagnostics nodeBuiltinImport:off // Node UUIDs identify kernel protocol requests at this adapter boundary. import * as NodeCrypto from "node:crypto"; -import { fetch, WebSocket } from "undici"; +import { fetch } from "undici"; import { NOTEBOOK_LIMITS, NotebookError } from "@t3tools/contracts"; import * as Cause from "effect/Cause"; import * as Context from "effect/Context"; @@ -11,6 +11,7 @@ import * as Queue from "effect/Queue"; import * as Schema from "effect/Schema"; import * as Scope from "effect/Scope"; import * as Stream from "effect/Stream"; +import { connectKernelSocket } from "./kernelSocket.ts"; const encodeJson = Schema.encodeSync(Schema.fromJsonString(Schema.Unknown)); const decodeJson = Schema.decodeSync(Schema.fromJsonString(Schema.Unknown)); @@ -145,62 +146,99 @@ export const layer = Layer.sync(JupyterTransport, () => { ); }); const connect = Effect.fn("JupyterTransport.connect")(function* (endpoint: Endpoint, id: string) { - const messages = yield* Queue.dropping(64); - const socket = yield* Effect.try({ - try: () => { - const url = endpointUrl(endpoint, `/${encodeURIComponent(id)}/channels`); - url.protocol = "ws:"; - url.searchParams.set("session_id", NodeCrypto.randomUUID()); - return new WebSocket(url, { headers: { Authorization: `token ${endpoint.token}` } }); - }, - catch: () => new NotebookError({ reason: "unavailable" }), - }); + const wake = yield* Queue.sliding(1); + const output: Message[] = []; + const control = new Map(); + let retainedBytes = 0; + let currentId: string | undefined; + let lost = false; + const loss = () => { + lost = true; + Queue.offerUnsafe(wake, undefined); + }; const fail = () => { - Queue.failCauseUnsafe(messages, Cause.fail(new NotebookError({ reason: "disconnected" }))); - socket.close(); + Queue.failCauseUnsafe(wake, Cause.fail(new NotebookError({ reason: "disconnected" }))); }; - socket.addEventListener("message", (event) => { + const receive = (text: string | null) => { try { + if (text === null) { + loss(); + return; + } + // Admission is bounded before JSON decoding; large data is discarded while channels drain. + if (Buffer.byteLength(text) > NOTEBOOK_LIMITS.bytes + 65536) { + loss(); + return; + } + let message = decodeMessage(text); if ( - typeof event.data !== "string" || - Buffer.byteLength(event.data) > NOTEBOOK_LIMITS.frameBytes + message.buffers?.length || + contentValidators[message.header.msg_type]?.(message.content) === false ) { fail(); return; } - const message = decodeMessage(event.data); + const type = message.header.msg_type; if ( - message.buffers?.length || - contentValidators[message.header.msg_type]?.(message.content) === false || - !Queue.offerUnsafe(messages, message) - ) - fail(); + message.parent_header.msg_id === currentId && + currentId && + (type === "execute_reply" || + type === "execute_input" || + (type === "status" && message.content.execution_state === "idle")) + ) { + if (Buffer.byteLength(encodeJson(message)) > 65536) { + fail(); + return; + } + control.set(type, message); + } else { + if ( + type === "stream" && + typeof message.content.text === "string" && + Buffer.byteLength(message.content.text) > 65532 + ) { + message = { + ...message, + content: { + ...message.content, + text: Buffer.from(message.content.text).subarray(0, 65528).toString("utf8"), + t3_truncated: true, + }, + }; + } + const size = Buffer.byteLength(encodeJson(message)); + if ( + size > NOTEBOOK_LIMITS.frameBytes || + retainedBytes + size > NOTEBOOK_LIMITS.frameBytes || + output.length >= 64 + ) { + loss(); + return; + } + output.push(message); + retainedBytes += size; + } + Queue.offerUnsafe(wake, undefined); } catch { fail(); } + }; + const url = yield* Effect.try({ + try: () => { + const url = endpointUrl(endpoint, "/" + encodeURIComponent(id) + "/channels"); + url.searchParams.set("session_id", NodeCrypto.randomUUID()); + return url; + }, + catch: () => new NotebookError({ reason: "unavailable" }), }); - socket.addEventListener("close", fail); - socket.addEventListener("error", fail); + const socket = yield* connectKernelSocket(url, endpoint.token, receive, fail); const close = Effect.sync(() => { - socket.removeEventListener("close", fail); - socket.removeEventListener("error", fail); socket.close(); - Queue.shutdownUnsafe(messages); + Queue.shutdownUnsafe(wake); }); yield* Effect.addFinalizer(() => close); - yield* Effect.callback((resume) => { - const opened = () => resume(Effect.void); - const failed = () => resume(Effect.fail(new NotebookError({ reason: "disconnected" }))); - socket.addEventListener("open", opened, { once: true }); - socket.addEventListener("error", failed, { once: true }); - socket.addEventListener("close", failed, { once: true }); - return Effect.sync(() => { - socket.removeEventListener("open", opened); - socket.removeEventListener("error", failed); - socket.removeEventListener("close", failed); - }); - }); const execute = Effect.fn("JupyterTransport.execute")(function* (id: string, code: string) { + currentId = id; yield* Effect.try({ try: () => socket.send( @@ -228,7 +266,29 @@ export const layer = Layer.sync(JupyterTransport, () => { catch: () => new NotebookError({ reason: "disconnected" }), }); }); - return { messages: Stream.fromQueue(messages), execute, close }; + const messages = Stream.fromQueue(wake).pipe( + Stream.flatMap(() => { + const batch = output.splice(0); + retainedBytes = 0; + const input = control.get("execute_input"); + if (input) batch.unshift(input); + if (lost && currentId) + batch.push({ + header: { msg_type: "t3_output_lost" }, + parent_header: { msg_id: currentId }, + channel: "iopub", + content: {}, + }); + for (const type of ["execute_reply", "status"]) { + const message = control.get(type); + if (message) batch.push(message); + } + control.clear(); + lost = false; + return Stream.fromIterable(batch); + }), + ); + return { messages, execute, close }; }); return JupyterTransport.of({ create, control, connect }); }); diff --git a/apps/server/src/notebook/NotebookService.test.ts b/apps/server/src/notebook/NotebookService.test.ts index 145d3f7dbf2ba..b1356af0dda4f 100644 --- a/apps/server/src/notebook/NotebookService.test.ts +++ b/apps/server/src/notebook/NotebookService.test.ts @@ -31,7 +31,9 @@ const fixture = { outputs: [], })), }; -const setup = Effect.gen(function* () { +const setup = Effect.fn("Test.setup")(function* ( + overrides: Partial = {}, +) { const messages = yield* Queue.unbounded(); const requests = yield* Queue.unbounded<{ id: string; code: string }>(); const controls: string[] = []; @@ -51,6 +53,7 @@ const setup = Effect.gen(function* () { controls.push(action); }), ), + ...overrides, }); const fs = yield* FileSystem.FileSystem; const sql = yield* SqlClient.SqlClient; @@ -93,15 +96,194 @@ const setup = Effect.gen(function* () { Stream.filter((s) => s.runId === runId && s.outcome !== "running"), Stream.runHead, ); - return { service, opened, handle, requests, controls, send, finished, root, fs, messages }; + return { service, opened, handle, requests, controls, send, finished, root, fs, messages, sql }; }); const layer = Layer.mergeAll(NodeServices.layer, ProjectStore.layer).pipe( Layer.provideMerge(Sqlite.layerMemory), ); +it.effect("project aliases share the canonical notebook owner and revalidate each project", () => + Effect.gen(function* () { + const { service, opened, handle, sql, root, requests, send, finished } = yield* setup(); + const alias = ProjectId.make("alias"); + yield* sql`INSERT INTO projection_projects (project_id, title, workspace_root, scripts_json, created_at, updated_at) VALUES (${alias}, 'Alias', ${root}, '[]', '2026-10-28T00:00:00.000Z', '2026-10-28T00:00:00.000Z')`; + const other = yield* service.open({ projectId: alias, relativePath: "test.ipynb" }, caller); + expect(other.handle).toBe(opened.handle); + const runId = yield* service.execute( + { ...handle, revision: opened.revision, cellIds: ["a"], requestId: "alias" }, + caller, + ); + const request = yield* Queue.take(requests); + expect( + yield* service + .execute( + { + projectId: alias, + handle: other.handle, + revision: other.revision, + cellIds: ["b"], + requestId: "conflict", + }, + caller, + ) + .pipe(Effect.flip), + ).toMatchObject({ reason: "busy" }); + yield* send(request.id, "execute_reply", { status: "ok" }); + yield* send(request.id, "status", { execution_state: "idle" }); + yield* finished(runId); + yield* sql`UPDATE projection_projects SET workspace_root = '/missing' WHERE project_id = ${alias}`; + expect( + yield* service.snapshot({ projectId: alias, handle: other.handle }, caller).pipe(Effect.flip), + ).toMatchObject({ reason: "invalid-path" }); + }).pipe(Effect.scoped, Effect.provide(layer)), +); + +it.effect.each(["shutdown", "finalizer"] as const)( + "retains an acquired kernel for %s when connect and its first cleanup fail", + (cleanup) => + Effect.gen(function* () { + const deleted: string[] = []; + yield* Effect.scoped( + Effect.gen(function* () { + const { service, opened, handle } = yield* setup({ + connect: Effect.fn("Fault.connect")(() => + Effect.fail(new NotebookError({ reason: "disconnected" })), + ), + control: Effect.fn("Fault.control")((_endpoint, id) => + Effect.gen(function* () { + deleted.push(id); + if (deleted.length === 1) + return yield* new NotebookError({ reason: "kernel-failure" }); + }), + ), + }); + yield* service + .execute( + { ...handle, revision: opened.revision, cellIds: ["a"], requestId: "fault" }, + caller, + ) + .pipe(Effect.flip); + if (cleanup === "shutdown") { + const current = yield* service.snapshot(handle, caller); + yield* service.shutdown({ ...handle, generation: current.generation }, caller); + } + }), + ); + expect(deleted).toEqual(["owned", "owned"]); + }).pipe(Effect.scoped, Effect.provide(layer)), +); + +it.effect("a blocked acquisition in B does not block A interrupt", () => + Effect.gen(function* () { + const blocked = yield* Deferred.make(); + const release = yield* Deferred.make(); + const interrupted = yield* Deferred.make(); + let creates = 0; + const { service, opened, handle, requests, fs, root } = yield* setup({ + create: Effect.fn("Blocked.create")(() => + Effect.gen(function* () { + creates++; + if (creates === 2) { + yield* Deferred.succeed(blocked, undefined); + yield* Deferred.await(release); + } + return "owned"; + }), + ), + control: Effect.fn("Blocked.control")((_endpoint, _id, action) => + action === "interrupt" + ? Deferred.succeed(interrupted, undefined).pipe(Effect.asVoid) + : Effect.void, + ), + }); + const runId = yield* service.execute( + { ...handle, revision: opened.revision, cellIds: ["a"], requestId: "A" }, + caller, + ); + yield* Queue.take(requests); + yield* fs.writeFileString(`${root}/B.ipynb`, JSON.stringify(fixture)); + const b = yield* service.open({ projectId: handle.projectId, relativePath: "B.ipynb" }, caller); + const acquiring = yield* service + .execute( + { ...handle, handle: b.handle, revision: b.revision, cellIds: ["a"], requestId: "B" }, + caller, + ) + .pipe(Effect.forkScoped); + yield* Deferred.await(blocked); + const interrupt = yield* service + .interrupt({ ...handle, generation: 1, runId }, caller) + .pipe(Effect.forkScoped); + yield* Deferred.await(interrupted); + yield* Deferred.succeed(release, undefined); + yield* Fiber.join(acquiring); + yield* Fiber.join(interrupt); + }).pipe(Effect.scoped, Effect.provide(layer)), +); + +it.effect("execution deltas reset counts and carry provenance even when dispatch fails", () => + Effect.gen(function* () { + const { service, handle, fs, root } = yield* setup({ + connect: Effect.fn("FailDispatch.connect")(() => + Effect.succeed({ + messages: Stream.never, + execute: Effect.fn("FailDispatch.execute")(() => + Effect.fail(new NotebookError({ reason: "disconnected" })), + ), + close: Effect.void, + }), + ), + }); + const old = { + ...fixture, + cells: fixture.cells.map((cell) => ({ ...cell, execution_count: 7 as number | null })), + }; + yield* fs.writeFileString(`${root}/counts.ipynb`, JSON.stringify(old)); + const initial = yield* service.open( + { projectId: handle.projectId, relativePath: "counts.ipynb" }, + caller, + ); + const target = { ...handle, handle: initial.handle }; + const ready = yield* Deferred.make(); + const terminal = yield* Deferred.make(); + let count: number | null = 7; + let execution: unknown; + const observer = yield* service.subscribe(target, caller).pipe( + Stream.runForEach((event) => + Effect.gen(function* () { + if (event.type === "snapshot") yield* Deferred.succeed(ready, undefined); + if (event.type !== "changed") return; + for (const patch of event.patches) { + if (patch.type === "execution") { + count = patch.count; + execution = patch.execution; + } + if (patch.type === "count") count = patch.count; + } + if (event.outcome === "unknown") yield* Deferred.succeed(terminal, undefined); + }), + ), + Effect.forkScoped, + ); + yield* Deferred.await(ready); + yield* service.execute( + { ...target, revision: initial.revision, cellIds: ["a"], requestId: "failed-dispatch" }, + caller, + ); + yield* Deferred.await(terminal); + const final = yield* service.snapshot(target, caller); + expect(count).toBeNull(); + expect(count).toBe( + (final.document.cells as Array<{ execution_count: number | null }>)[0]?.execution_count, + ); + expect(execution).toEqual(final.executions.a); + expect(execution).toMatchObject({ sourceRevision: initial.revision, generation: 1 }); + yield* Fiber.interrupt(observer); + }).pipe(Effect.scoped, Effect.provide(layer)), +); + it.effect("slow observers receive an explicit resync instead of retaining every delta", () => Effect.gen(function* () { - const { service, opened, handle, requests, send, finished } = yield* setup; + const { service, opened, handle, requests, send, finished } = yield* setup(); const ready = yield* Deferred.make(); const release = yield* Deferred.make(); const resync = yield* Deferred.make(); @@ -139,7 +321,7 @@ it.effect("slow observers receive an explicit resync instead of retaining every it.effect("disconnect marks an uncertain outcome and never resends an accepted request", () => Effect.gen(function* () { - const { service, opened, handle, requests, messages, finished } = yield* setup; + const { service, opened, handle, requests, messages, finished } = yield* setup(); const input = { ...handle, revision: opened.revision, @@ -176,7 +358,7 @@ it.effect( "correlates requests, waits for reply and idle in both orders, deduplicates, and retains late output", () => Effect.gen(function* () { - const { service, opened, handle, requests, send, finished } = yield* setup; + const { service, opened, handle, requests, send, finished } = yield* setup(); const input = { ...handle, revision: opened.revision, cellIds: ["a", "b"], requestId: "run" }; const runId = yield* service.execute(input, caller); expect(yield* service.execute(input, caller)).toBe(runId); @@ -201,13 +383,18 @@ it.effect( "interrupt cancels undispatched cells; restart and shutdown fence stale generations", () => Effect.gen(function* () { - const { service, opened, handle, requests, controls, send, finished } = yield* setup; + const { service, opened, handle, requests, controls, send, finished } = yield* setup(); const runId = yield* service.execute( { ...handle, revision: opened.revision, cellIds: ["a", "b"], requestId: "run" }, caller, ); const first = yield* Queue.take(requests); const generation = (yield* service.snapshot(handle, caller)).generation; + yield* send(first.id, "display_data", { + data: { "text/plain": "generation-one" }, + metadata: {}, + transient: { display_id: "shared" }, + }); yield* service.interrupt({ ...handle, generation, runId }, caller); yield* send(first.id, "execute_reply", { status: "error" }); yield* send(first.id, "status", { execution_state: "idle" }); @@ -218,7 +405,67 @@ it.effect( expect( yield* service.shutdown({ ...handle, generation }, caller).pipe(Effect.flip), ).toMatchObject({ reason: "stale-generation" }); + const secondSnapshot = yield* service.snapshot(handle, caller); + const runB = yield* service.execute( + { + ...handle, + revision: secondSnapshot.revision, + cellIds: ["b"], + requestId: "after-restart", + }, + caller, + ); + const b = yield* Queue.take(requests); + yield* send(b.id, "display_data", { + data: { "text/plain": "generation-two" }, + metadata: {}, + transient: { display_id: "shared" }, + }); + yield* send(b.id, "update_display_data", { + data: { "text/plain": "updated-two" }, + metadata: {}, + transient: { display_id: "shared" }, + }); + yield* send(b.id, "execute_reply", { status: "ok" }); + yield* send(b.id, "status", { execution_state: "idle" }); + yield* finished(runB); + expect( + ( + (yield* service.snapshot(handle, caller)).document.cells as Array<{ outputs: unknown[] }> + )[0]?.outputs, + ).toEqual([ + { output_type: "display_data", data: { "text/plain": "generation-one" }, metadata: {} }, + ]); yield* service.shutdown({ ...handle, generation: generation + 1 }, caller); + const thirdSnapshot = yield* service.snapshot(handle, caller); + const runC = yield* service.execute( + { + ...handle, + revision: thirdSnapshot.revision, + cellIds: ["b"], + requestId: "after-shutdown", + }, + caller, + ); + const c = yield* Queue.take(requests); + yield* send(c.id, "display_data", { + data: { "text/plain": "generation-three" }, + metadata: {}, + transient: { display_id: "shared" }, + }); + yield* send(c.id, "update_display_data", { + data: { "text/plain": "updated-three" }, + metadata: {}, + transient: { display_id: "shared" }, + }); + yield* send(c.id, "execute_reply", { status: "ok" }); + yield* send(c.id, "status", { execution_state: "idle" }); + yield* finished(runC); + const final = yield* service.snapshot(handle, caller); + expect((final.document.cells as Array<{ outputs: unknown[] }>)[0]?.outputs).toEqual([ + { output_type: "display_data", data: { "text/plain": "generation-one" }, metadata: {} }, + ]); + expect(final.executions.a?.generation).toBe(1); expect(controls).toEqual(["interrupt", "restart", "shutdown"]); }).pipe(Effect.scoped, Effect.provide(layer)), ); @@ -227,7 +474,7 @@ it.effect( "enforces read and execution grants, handle ownership, and canonical path containment", () => Effect.gen(function* () { - const { service, opened, handle, root, fs } = yield* setup; + const { service, opened, handle, root, fs } = yield* setup(); expect(yield* service.snapshot(handle, { scopes: [] }).pipe(Effect.flip)).toMatchObject({ reason: "permission-denied", }); diff --git a/apps/server/src/notebook/NotebookService.ts b/apps/server/src/notebook/NotebookService.ts index dd570308ea459..306e3a40fa1b7 100644 --- a/apps/server/src/notebook/NotebookService.ts +++ b/apps/server/src/notebook/NotebookService.ts @@ -86,10 +86,10 @@ const make = Effect.gen(function* () { const config = yield* NotebookConfiguration; const transport = yield* JupyterTransport.JupyterTransport; const scope = yield* Scope.Scope; - const lock = yield* Semaphore.make(1); const epoch = NodeCrypto.randomUUID(); type Session = { - projectId: string; + projects: Set; + lock: Semaphore.Semaphore; path: string; root: string; snapshot: { -readonly [K in keyof NotebookSnapshot]: NotebookSnapshot[K] }; @@ -97,8 +97,8 @@ const make = Effect.gen(function* () { kernel?: { id: string; endpoint: JupyterTransport.Endpoint; - channels: JupyterTransport.Channels; - scope: Scope.Closeable; + channels?: JupyterTransport.Channels; + scope?: Scope.Closeable; }; pending?: { id: string; done: Deferred.Deferred }; cancelled: boolean; @@ -147,7 +147,7 @@ const make = Effect.gen(function* () { ) { yield* authorize(caller, execute); const session = sessions.get(input.handle); - if (!session || session.projectId !== input.projectId) + if (!session || !session.projects.has(input.projectId)) return yield* new NotebookError({ reason: "not-found" }); const resolved = yield* resolve({ projectId: input.projectId, @@ -204,13 +204,21 @@ const make = Effect.gen(function* () { id: string, ) { const kernelScope = yield* Scope.make(); + session.kernel = { id, endpoint, scope: kernelScope }; const channels = yield* transport.connect(endpoint, id).pipe( Effect.provideService(Scope.Scope, kernelScope), Effect.onError(() => Effect.all( [ Scope.close(kernelScope, Exit.void), - transport.control(endpoint, id, "shutdown").pipe(Effect.ignore), + transport.control(endpoint, id, "shutdown").pipe( + Effect.tap(() => + Effect.sync(() => { + delete session.kernel; + }), + ), + Effect.ignore, + ), unknown(session), ], { discard: true }, @@ -241,59 +249,62 @@ const make = Effect.gen(function* () { caller: Caller, ) { yield* authorize(caller); - return yield* lock.withPermit( - Effect.gen(function* () { - const { root, path } = yield* resolve(input); - for (const session of sessions.values()) - if ( - session.path === path && - session.root === root && - session.projectId === input.projectId - ) - return snapshotOf(session); - if (sessions.size >= NOTEBOOK_LIMITS.sessions) - return yield* new NotebookError({ reason: "output-limit" }); - const stat = yield* fs - .stat(path) - .pipe(Effect.mapError(() => new NotebookError({ reason: "not-found" }))); - if (stat.size > NOTEBOOK_LIMITS.bytes) - return yield* new NotebookError({ reason: "output-limit" }); - const text = yield* fs - .readFileString(path) - .pipe(Effect.mapError(() => new NotebookError({ reason: "not-found" }))); - const parsed = yield* Effect.try({ - try: () => parseNotebook(text), - catch: () => new NotebookError({ reason: "invalid-document" }), - }); - const handle = NodeCrypto.randomUUID(); - const snapshot: NotebookSnapshot = { - handle, - epoch, - sequence: 0, - revision: parsed.changed ? 1 : 0, - diskHash: NodeCrypto.createHash("sha256").update(text).digest("hex"), - generation: 0, - executions: {}, - document: parsed.document, - kernel: "stopped", - runId: null, - outcome: "none", - truncated: false, - }; - const session: Session = { - projectId: input.projectId, - path, - root, - snapshot, - reducer: new ExecutionReducer(parsed.document), - cancelled: false, - requests: new Map(), - observers: new Set(), - }; - sessions.set(handle, session); + const { root, path } = yield* resolve(input); + for (const session of sessions.values()) + if (session.path === path && session.root === root) { + session.projects.add(input.projectId); return snapshotOf(session); - }), - ); + } + if (sessions.size >= NOTEBOOK_LIMITS.sessions) + return yield* new NotebookError({ reason: "output-limit" }); + const stat = yield* fs + .stat(path) + .pipe(Effect.mapError(() => new NotebookError({ reason: "not-found" }))); + if (stat.size > NOTEBOOK_LIMITS.bytes) + return yield* new NotebookError({ reason: "output-limit" }); + const text = yield* fs + .readFileString(path) + .pipe(Effect.mapError(() => new NotebookError({ reason: "not-found" }))); + const parsed = yield* Effect.try({ + try: () => parseNotebook(text), + catch: () => new NotebookError({ reason: "invalid-document" }), + }); + for (const session of sessions.values()) { + if (session.path === path && session.root === root) { + session.projects.add(input.projectId); + return snapshotOf(session); + } + } + if (sessions.size >= NOTEBOOK_LIMITS.sessions) + return yield* new NotebookError({ reason: "output-limit" }); + const handle = NodeCrypto.randomUUID(); + const snapshot: NotebookSnapshot = { + handle, + epoch, + sequence: 0, + revision: parsed.changed ? 1 : 0, + diskHash: NodeCrypto.createHash("sha256").update(text).digest("hex"), + generation: 0, + executions: {}, + document: parsed.document, + kernel: "stopped", + runId: null, + outcome: "none", + truncated: false, + }; + const session: Session = { + projects: new Set([input.projectId]), + lock: Semaphore.makeUnsafe(1), + path, + root, + snapshot, + reducer: new ExecutionReducer(parsed.document), + cancelled: false, + requests: new Map(), + observers: new Set(), + }; + sessions.set(handle, session); + return snapshotOf(session); }); const snapshot = Effect.fn("NotebookService.snapshot")(function* ( input: NotebookHandle, @@ -305,9 +316,10 @@ const make = Effect.gen(function* () { input: NotebookExecute, caller: Caller, ) { - return yield* lock.withPermit( + const session = yield* get(input, caller, true); + return yield* session.lock.withPermit( Effect.gen(function* () { - const session = yield* get(input, caller, true); + yield* get(input, caller, true); const fingerprint = fingerprintOf([input.revision, input.cellIds]); const prior = session.requests.get(input.requestId); if (prior) { @@ -338,11 +350,14 @@ const make = Effect.gen(function* () { if (!session.kernel) { const endpoint = yield* config.endpoint(session.root); const id = yield* transport.create(endpoint, session.root); + session.kernel = { id, endpoint }; session.snapshot.generation++; + session.reducer.resetGeneration(); yield* connectOwned(session, endpoint, id); } const kernel = session.kernel; - if (!kernel) return yield* new NotebookError({ reason: "kernel-failure" }); + if (!kernel?.channels) return yield* new NotebookError({ reason: "kernel-failure" }); + const channels = kernel.channels; const runId = NodeCrypto.randomUUID(); session.snapshot.runId = runId; session.snapshot.outcome = "running"; @@ -359,6 +374,10 @@ const make = Effect.gen(function* () { const id = NodeCrypto.randomUUID(); const done = yield* Deferred.make(); session.pending = { id, done }; + yield* Effect.try({ + try: () => session.reducer.begin(id, cell), + catch: () => new NotebookError({ reason: "output-limit" }), + }); session.snapshot.executions = { ...session.snapshot.executions, [cell.id]: { @@ -370,9 +389,16 @@ const make = Effect.gen(function* () { .digest("hex"), }, }; - session.reducer.begin(id, cell); - changed(session, cell.id, [{ type: "clear", cellId: cell.id }]); - yield* kernel.channels.execute( + changed(session, cell.id, [ + { type: "clear", cellId: cell.id }, + { + type: "execution", + cellId: cell.id, + count: null, + execution: session.snapshot.executions[cell.id]!, + }, + ]); + yield* channels.execute( id, typeof cell.source === "string" ? cell.source : cell.source.join(""), ); @@ -403,9 +429,10 @@ const make = Effect.gen(function* () { input: NotebookControl & { runId?: string }, caller: Caller, ) { - yield* lock.withPermit( + const session = yield* get(input, caller, true); + yield* session.lock.withPermit( Effect.gen(function* () { - const session = yield* get(input, caller, true); + yield* get(input, caller, true); if (session.snapshot.generation !== input.generation) return yield* new NotebookError({ reason: "stale-generation" }); if (action === "interrupt" && input.runId !== session.snapshot.runId) @@ -425,11 +452,13 @@ const make = Effect.gen(function* () { .pipe(Effect.onError(() => unknown(session))); if (action !== "interrupt") { session.snapshot.generation++; - session.reducer.owners.clear(); + session.reducer.resetGeneration(); if (session.pending) yield* Deferred.succeed(session.pending.done, true); - yield* Scope.close(kernel.scope, Exit.void); - delete session.kernel; - if (action === "restart") yield* connectOwned(session, kernel.endpoint, kernel.id); + if (kernel.scope) yield* Scope.close(kernel.scope, Exit.void); + if (action === "restart") { + session.kernel = { id: kernel.id, endpoint: kernel.endpoint }; + yield* connectOwned(session, kernel.endpoint, kernel.id); + } else delete session.kernel; session.snapshot.kernel = action === "restart" ? "idle" : "stopped"; session.snapshot.outcome = "aborted"; } @@ -464,7 +493,7 @@ const make = Effect.gen(function* () { const kernel = session.kernel; if (!kernel) return; session.snapshot.generation++; - yield* Scope.close(kernel.scope, Exit.void); + if (kernel.scope) yield* Scope.close(kernel.scope, Exit.void); yield* transport.control(kernel.endpoint, kernel.id, "shutdown").pipe(Effect.ignore); }), { discard: true }, diff --git a/apps/server/src/notebook/executionReducer.test.ts b/apps/server/src/notebook/executionReducer.test.ts index 4629c8e311d35..cae52105670ab 100644 --- a/apps/server/src/notebook/executionReducer.test.ts +++ b/apps/server/src/notebook/executionReducer.test.ts @@ -81,12 +81,14 @@ it("updates displays across cells, defers clear, strips transient IDs, and bound expect(cell.outputs).toEqual([ { output_type: "display_data", data: { "text/plain": "new" }, metadata: {} }, ]); + send("one", "execute_input", { execution_count: 9 }); send("one", "clear_output", { wait: true }); expect(document.cells[0]?.outputs).toHaveLength(1); send("one", "stream", { name: "stdout", text: "next" }); expect(document.cells[0]?.outputs).toEqual([ { output_type: "stream", name: "stdout", text: "next" }, ]); + expect(document.cells[0]?.execution_count).toBe(9); for (let i = 0; i < 12; i++) send("one", "stream", { name: "stdout", text: "x".repeat(NOTEBOOK_LIMITS.frameBytes) }); expect(byteLength(document)).toBeLessThanOrEqual(NOTEBOOK_LIMITS.bytes); @@ -95,3 +97,42 @@ it("updates displays across cells, defers clear, strips transient IDs, and bound reducer.begin("replacement", document.cells[0]!); expect(send("one", "stream", { name: "stdout", text: "obsolete" })).toBeUndefined(); }); + +it("never exceeds the exact document bound while resetting or updating counts", () => { + const value = { + nbformat: 4, + nbformat_minor: 5, + metadata: { padding: "" }, + cells: [ + { id: "a", cell_type: "code", source: "", metadata: {}, execution_count: 0, outputs: [] }, + ], + }; + value.metadata.padding = "x".repeat(NOTEBOOK_LIMITS.bytes - byteLength(value)); + const document = parseNotebook(JSON.stringify(value)).document; + const reducer = new ExecutionReducer(document); + try { + reducer.begin("one", document.cells[0]!); + } catch { + /* An explicit limit rejection preserves the document. */ + } + expect(byteLength(document)).toBeLessThanOrEqual(NOTEBOOK_LIMITS.bytes); + reducer.reduce({ + header: { msg_type: "execute_input" }, + parent_header: { msg_id: "one" }, + channel: "iopub", + content: { execution_count: 123456789 }, + }); + expect(byteLength(document)).toBeLessThanOrEqual(NOTEBOOK_LIMITS.bytes); + value.metadata.padding = value.metadata.padding.slice(0, -4); + const nearLimit = parseNotebook(JSON.stringify(value)).document; + const countReducer = new ExecutionReducer(nearLimit); + countReducer.begin("two", nearLimit.cells[0]!); + countReducer.reduce({ + header: { msg_type: "execute_input" }, + parent_header: { msg_id: "two" }, + channel: "iopub", + content: { execution_count: 123456789 }, + }); + expect(byteLength(nearLimit)).toBeLessThanOrEqual(NOTEBOOK_LIMITS.bytes); + expect(countReducer.truncated).toBe(true); +}); diff --git a/apps/server/src/notebook/executionReducer.ts b/apps/server/src/notebook/executionReducer.ts index 3f32e8216d5eb..2b6571757c03c 100644 --- a/apps/server/src/notebook/executionReducer.ts +++ b/apps/server/src/notebook/executionReducer.ts @@ -1,20 +1,32 @@ -import { NOTEBOOK_LIMITS, type NotebookOutputPatch } from "@t3tools/contracts"; +import { NOTEBOOK_LIMITS, NotebookError, type NotebookOutputPatch } from "@t3tools/contracts"; import type { Message } from "./JupyterTransport.ts"; import { byteLength, type Document, type NotebookCell } from "./notebookDocument.ts"; type Owner = { cell: NotebookCell; reply: boolean; idle: boolean; failed: boolean; clear: boolean }; export class ExecutionReducer { readonly owners = new Map(); - readonly displays = new WeakMap(); + displays = new WeakMap(); truncated = false; readonly document: Document; constructor(document: Document) { this.document = document; } + resetGeneration() { + this.owners.clear(); + this.displays = new WeakMap(); + } begin(id: string, cell: NotebookCell) { - for (const [key, owner] of this.owners) if (owner.cell === cell) this.owners.delete(key); + const previous = { outputs: cell.outputs, count: cell.execution_count }; cell.outputs = []; cell.execution_count = null; + if (byteLength(this.document) > NOTEBOOK_LIMITS.bytes) { + if (previous.outputs === undefined) delete cell.outputs; + else cell.outputs = previous.outputs; + if (previous.count === undefined) delete cell.execution_count; + else cell.execution_count = previous.count; + throw new NotebookError({ reason: "output-limit" }); + } + for (const [key, owner] of this.owners) if (owner.cell === cell) this.owners.delete(key); this.owners.set(id, { cell, reply: false, idle: false, failed: false, clear: false }); } reduce(message: Message) { @@ -24,6 +36,7 @@ export class ExecutionReducer { const patches: NotebookOutputPatch[] = []; const content = message.content; const type = message.header.msg_type; + if (type === "t3_output_lost" || content.t3_truncated === true) this.truncated = true; if (type === "execute_reply" && message.channel === "shell") { owner.reply = true; owner.failed = content.status !== "ok"; @@ -32,8 +45,13 @@ export class ExecutionReducer { owner.idle = true; if (message.channel === "iopub") { if (type === "execute_input" && typeof content.execution_count === "number") { + const previous = cell.execution_count; cell.execution_count = content.execution_count; - patches.push({ type: "count", cellId: cell.id, count: content.execution_count }); + if (byteLength(this.document) > NOTEBOOK_LIMITS.bytes) { + if (previous === undefined) delete cell.execution_count; + else cell.execution_count = previous; + this.truncated = true; + } else patches.push({ type: "count", cellId: cell.id, count: content.execution_count }); } if (type === "clear_output") { if (content.wait === true) owner.clear = true; @@ -44,7 +62,14 @@ export class ExecutionReducer { } } if ( - ["stream", "display_data", "execute_result", "error", "update_display_data"].includes(type) + [ + "stream", + "display_data", + "execute_result", + "error", + "update_display_data", + "t3_output_lost", + ].includes(type) ) { if (owner.clear) { cell.outputs = []; @@ -54,7 +79,7 @@ export class ExecutionReducer { const transient = content.transient as { display_id?: unknown } | undefined; const displayId = typeof transient?.display_id === "string" ? transient.display_id : undefined; - const { transient: _transient, ...saved } = content; + const { transient: _transient, t3_truncated: _truncated, ...saved } = content; if (type === "update_display_data") { if (displayId) for (const other of this.document.cells) @@ -88,7 +113,7 @@ export class ExecutionReducer { previous.text = text; this.truncated = true; } else patches.push({ type: "stream", cellId: cell.id, text: saved.text }); - } else { + } else if (type !== "t3_output_lost") { const output = { ...saved, output_type: type }; outputs.push(output); if ( diff --git a/apps/server/src/notebook/kernelSocket.ts b/apps/server/src/notebook/kernelSocket.ts new file mode 100644 index 0000000000000..b0f604b1b8b54 --- /dev/null +++ b/apps/server/src/notebook/kernelSocket.ts @@ -0,0 +1,195 @@ +// @effect-diagnostics nodeBuiltinImport:off +// Streaming frame admission must discard oversized messages before allocating their declared length. +import * as NodeHttp from "node:http"; +import * as NodeCrypto from "node:crypto"; +import * as Effect from "effect/Effect"; +import { NotebookError, NOTEBOOK_LIMITS } from "@t3tools/contracts"; + +export const connectKernelSocket = Effect.fn("JupyterTransport.connectSocket")( + (url: URL, token: string, receive: (text: string | null) => void, fail: () => void) => + Effect.callback<{ send: (text: string) => void; close: () => void }, NotebookError>( + (resume) => { + const key = NodeCrypto.randomBytes(16).toString("base64"); + let connected = false; + const request = NodeHttp.request(url, { + headers: { + Authorization: `token ${token}`, + Connection: "Upgrade", + Upgrade: "websocket", + "Sec-WebSocket-Key": key, + "Sec-WebSocket-Version": "13", + }, + }); + const failed = () => { + fail(); + resume(Effect.fail(new NotebookError({ reason: "disconnected" }))); + }; + request.on("error", failed); + request.on("response", (response) => { + response.destroy(); + failed(); + }); + request.on("upgrade", (response, socket, head) => { + const accept = NodeCrypto.createHash("sha1") + .update(key + "258EAFA5-E914-47DA-95CA-C5AB0DC85B11") + .digest("base64"); + if ( + response.statusCode !== 101 || + response.headers["sec-websocket-accept"] !== accept || + response.headers["sec-websocket-extensions"] + ) { + socket.destroy(); + failed(); + return; + } + const capacity = NOTEBOOK_LIMITS.bytes + 65536; + const message = Buffer.allocUnsafe(capacity); + const header = Buffer.alloc(10); + let headerSize = 0, + headerNeeded = 2, + remaining = 0n, + opcode = 0, + final = false; + let messageSize = 0, + fragmented = false, + discarded = false; + const ping = Buffer.alloc(125); + let pingSize = 0; + let closed = false; + const sendFrame = (opcode: number, payload: Buffer) => { + if (closed || socket.writableLength + payload.length > 2 * capacity) throw new Error(); + const mask = NodeCrypto.randomBytes(4); + const prefix = Buffer.alloc( + payload.length < 126 ? 6 : payload.length <= 65535 ? 8 : 14, + ); + prefix[0] = 128 | opcode; + if (payload.length < 126) prefix[1] = 128 | payload.length; + else if (payload.length <= 65535) { + prefix[1] = 254; + prefix.writeUInt16BE(payload.length, 2); + } else { + prefix[1] = 255; + prefix.writeBigUInt64BE(BigInt(payload.length), 2); + } + mask.copy(prefix, prefix.length - 4); + const encoded = Buffer.allocUnsafe(payload.length); + for (let i = 0; i < payload.length; i++) encoded[i] = payload[i]! ^ mask[i % 4]!; + socket.write(prefix); + socket.write(encoded); + }; + const close = () => { + closed = true; + socket.destroy(); + }; + const badFrame = () => { + close(); + failed(); + }; + const finish = () => { + if (opcode === 8) { + badFrame(); + return; + } + if (opcode === 9) { + sendFrame(10, ping.subarray(0, pingSize)); + } + if ((opcode === 0 || opcode === 1) && final) { + if (discarded) receive(null); + else + receive( + new TextDecoder("utf-8", { fatal: true }).decode( + message.subarray(0, messageSize), + ), + ); + messageSize = 0; + discarded = false; + fragmented = false; + } + headerSize = 0; + headerNeeded = 2; + }; + const data = (chunk: Buffer) => { + try { + let offset = 0; + while (offset < chunk.length) { + if (closed) return; + if (headerSize < headerNeeded) { + const count = Math.min(headerNeeded - headerSize, chunk.length - offset); + chunk.copy(header, headerSize, offset, offset + count); + offset += count; + headerSize += count; + if (headerSize < headerNeeded) continue; + if (headerNeeded === 2) { + const length = header[1]! & 127; + if (header[0]! & 112 || header[1]! & 128) { + badFrame(); + return; + } + headerNeeded = length === 126 ? 4 : length === 127 ? 10 : 2; + if (headerSize < headerNeeded) continue; + } + opcode = header[0]! & 15; + final = Boolean(header[0]! & 128); + const length = header[1]! & 127; + remaining = + length === 126 + ? BigInt(header.readUInt16BE(2)) + : length === 127 + ? header.readBigUInt64BE(2) + : BigInt(length); + if ( + remaining >> 63n || + ![0, 1, 8, 9, 10].includes(opcode) || + (opcode >= 8 && (!final || remaining > 125n)) || + (opcode === 0 && !fragmented) || + (opcode === 1 && fragmented) + ) { + badFrame(); + return; + } + if (opcode === 1) fragmented = !final; + if (opcode >= 8) pingSize = 0; + if (opcode < 8 && remaining > BigInt(capacity - messageSize)) discarded = true; + if (remaining === 0n) { + finish(); + continue; + } + } + const count = Number( + remaining < BigInt(chunk.length - offset) + ? remaining + : BigInt(chunk.length - offset), + ); + if (opcode === 9) { + chunk.copy(ping, pingSize, offset, offset + count); + pingSize += count; + } else if (opcode < 8 && !discarded) { + chunk.copy(message, messageSize, offset, offset + count); + messageSize += count; + } + offset += count; + remaining -= BigInt(count); + if (remaining === 0n) finish(); + } + } catch { + badFrame(); + } + }; + socket.on("data", data); + socket.on("error", failed); + socket.on("close", () => { + if (!closed) failed(); + }); + connected = true; + resume(Effect.succeed({ send: (text) => sendFrame(1, Buffer.from(text)), close })); + if (head.length) data(head); + }); + request.end(); + return Effect.sync(() => { + if (!connected) { + request.destroy(); + } + }); + }, + ), +); diff --git a/packages/contracts/src/notebook.ts b/packages/contracts/src/notebook.ts index 490663e19355e..e2e43d0fd5b3e 100644 --- a/packages/contracts/src/notebook.ts +++ b/packages/contracts/src/notebook.ts @@ -48,21 +48,19 @@ export const NotebookControl = Schema.Struct({ ...NotebookHandle.fields, generat export type NotebookControl = typeof NotebookControl.Type; export const NotebookInterrupt = Schema.Struct({ ...NotebookControl.fields, runId: Schema.String }); export type NotebookInterrupt = typeof NotebookInterrupt.Type; +const NotebookExecution = Schema.Struct({ + runId: Schema.String, + generation: Schema.Int, + sourceHash: Schema.String, + sourceRevision: Schema.Int, +}); export const NotebookSnapshot = Schema.Struct({ handle: Schema.String, epoch: Schema.String, sequence: Schema.Int, revision: Schema.Int, diskHash: Schema.String, - executions: Schema.Record( - Schema.String, - Schema.Struct({ - runId: Schema.String, - generation: Schema.Int, - sourceHash: Schema.String, - sourceRevision: Schema.Int, - }), - ), + executions: Schema.Record(Schema.String, NotebookExecution), generation: Schema.Int, document: NotebookDocument, kernel: Schema.Literals(["stopped", "idle", "busy", "unknown"]), @@ -72,6 +70,12 @@ export const NotebookSnapshot = Schema.Struct({ }); export type NotebookSnapshot = typeof NotebookSnapshot.Type; export const NotebookOutputPatch = Schema.Union([ + Schema.Struct({ + type: Schema.Literal("execution"), + cellId: Schema.String, + count: Schema.NullOr(Schema.Int), + execution: NotebookExecution, + }), Schema.Struct({ type: Schema.Literal("clear"), cellId: Schema.String }), Schema.Struct({ type: Schema.Literal("count"), cellId: Schema.String, count: Schema.Int }), Schema.Struct({ From a07f431b924a40ffb01ac4128e08eb81b634ae24 Mon Sep 17 00:00:00 2001 From: Tanner Polley <42780259+tannerpolley@users.noreply.github.com> Date: Wed, 28 Oct 2026 14:47:14 -0600 Subject: [PATCH 3/4] fix(notebook): stream bounded envelopes through one admission owner --- .../src/notebook/JupyterTransport.test.ts | 325 ++++++++++++------ apps/server/src/notebook/JupyterTransport.ts | 146 +------- .../src/notebook/KernelMessageAdmission.ts | 290 ++++++++++++++++ apps/server/src/notebook/executionReducer.ts | 7 + apps/server/src/notebook/kernelSocket.ts | 30 +- 5 files changed, 541 insertions(+), 257 deletions(-) create mode 100644 apps/server/src/notebook/KernelMessageAdmission.ts diff --git a/apps/server/src/notebook/JupyterTransport.test.ts b/apps/server/src/notebook/JupyterTransport.test.ts index eb2b63f750ea1..ef947f72bb633 100644 --- a/apps/server/src/notebook/JupyterTransport.test.ts +++ b/apps/server/src/notebook/JupyterTransport.test.ts @@ -1,11 +1,12 @@ // @effect-diagnostics nodeBuiltinImport:off -// A loopback protocol fixture exercises framing and admission without a Jupyter process. +// A scoped loopback fixture drives actual frames, chunking, admission, and reducer correlation. import * as NodeHttp from "node:http"; import * as NodeCrypto from "node:crypto"; import type * as NodeStream from "node:stream"; import { expect, it } from "@effect/vitest"; import * as Effect from "effect/Effect"; import * as Stream from "effect/Stream"; +import * as Fiber from "effect/Fiber"; import * as Deferred from "effect/Deferred"; import { NOTEBOOK_LIMITS } from "@t3tools/contracts"; import * as JupyterTransport from "./JupyterTransport.ts"; @@ -26,119 +27,161 @@ const frame = (text: string, opcode = 1, final = true) => { } return Buffer.concat([header, payload]); }; +const json = (id: string, type: string, content: Record) => + JSON.stringify({ + content, + channel: type === "execute_reply" ? "shell" : "iopub", + parent_header: { msg_id: id }, + header: { msg_type: type }, + }); +const fixture = Effect.gen(function* () { + let peer: NodeStream.Duplex | undefined; + let interrupts = 0; + const server = NodeHttp.createServer((_request, response) => { + interrupts++; + response.writeHead(204); + response.end(); + }); + server.on("upgrade", (request, socket) => { + peer = socket; + const accept = NodeCrypto.createHash("sha1") + .update(request.headers["sec-websocket-key"] + "258EAFA5-E914-47DA-95CA-C5AB0DC85B11") + .digest("base64"); + socket.write( + "HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: " + + accept + + "\r\n\r\n", + ); + socket.on("data", () => {}); + }); + yield* Effect.addFinalizer(() => + Effect.promise( + () => + new Promise((resolve) => { + peer?.destroy(); + server.close(() => resolve()); + }), + ), + ); + yield* Effect.callback((resume) => { + server.listen(0, "127.0.0.1", () => resume(Effect.void)); + }); + const address = server.address(); + if (!address || typeof address === "string") throw new Error(); + const transport = yield* JupyterTransport.JupyterTransport; + const endpoint = { + url: "http://127.0.0.1:" + address.port, + token: "test", + kernelName: "python3", + }; + const channels = yield* transport.connect(endpoint, "owned"); + const document = parseNotebook( + JSON.stringify({ + nbformat: 4, + nbformat_minor: 5, + metadata: {}, + cells: ["a", "b"].map((id) => ({ + id, + cell_type: "code", + source: "", + metadata: {}, + execution_count: null, + outputs: [], + })), + }), + ).document; + const reducer = new ExecutionReducer(document); + return { + channels, + reducer, + document, + transport, + endpoint, + interrupts: () => interrupts, + write: (data: Buffer) => peer!.write(data), + }; +}); +const provide = (effect: Effect.Effect) => + effect.pipe(Effect.provide(JupyterTransport.layer)); -it.live.each(["stream", "frame", "burst"] as const)( - "%s overflow retains bounded data without losing completion or control", - (mode) => +it.live.each([ + "stream", + "frame", + "burst", + "input64", + "reply64", + "reply9", + "reply9-idleFirst", + "reply64-replyFirst", + "input64-replyFirst", +] as const)( + "%s retains bounded data and normalized control fields across reordered, fragmented JSON", + (variant) => Effect.gen(function* () { - let peer: NodeStream.Duplex | undefined; - let interrupts = 0; - const server = NodeHttp.createServer((_request, response) => { - interrupts++; - response.writeHead(204); - response.end(); - }); - server.on("upgrade", (request, socket) => { - peer = socket; - const accept = NodeCrypto.createHash("sha1") - .update(request.headers["sec-websocket-key"] + "258EAFA5-E914-47DA-95CA-C5AB0DC85B11") - .digest("base64"); - socket.write( - "HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: " + - accept + - "\r\n\r\n", - ); - socket.on("data", () => {}); - }); - yield* Effect.addFinalizer(() => - Effect.promise( - () => - new Promise((resolve) => { - peer?.destroy(); - server.close(() => resolve()); - }), - ), - ); - yield* Effect.callback((resume) => { - server.listen(0, "127.0.0.1", () => resume(Effect.void)); - }); - const address = server.address(); - if (!address || typeof address === "string") throw new Error(); - const transport = yield* JupyterTransport.JupyterTransport; - const endpoint = { - url: "http://127.0.0.1:" + address.port, - token: "test", - kernelName: "python3", - }; - const channels = yield* transport.connect(endpoint, "owned"); - const document = parseNotebook( - JSON.stringify({ - nbformat: 4, - nbformat_minor: 5, - metadata: {}, - cells: [ - { - id: "a", - cell_type: "code", - source: "", - metadata: {}, - execution_count: null, - outputs: [], - }, - ], - }), - ).document; - const reducer = new ExecutionReducer(document); - const message = (type: string, content: Record) => - frame( - JSON.stringify({ - header: { msg_type: type }, - parent_header: { msg_id: "request" }, - channel: type === "execute_reply" ? "shell" : "iopub", - content, - }), - ); + const mode = variant.split("-")[0]!; + const { channels, reducer, document, transport, endpoint, interrupts, write } = + yield* fixture; reducer.begin("request", document.cells[0]!); - yield* channels.execute("request", "print('flood')"); + yield* channels.execute("request", "code"); if (mode === "stream") - peer!.write( - message("stream", { name: "stdout", text: "x".repeat(NOTEBOOK_LIMITS.frameBytes) }), + write( + frame( + json("request", "stream", { + name: "stdout", + text: "x".repeat(NOTEBOOK_LIMITS.frameBytes), + }), + ), ); if (mode === "burst") - peer!.write( + write( Buffer.concat( Array.from({ length: 1000 }, () => - message("stream", { name: "stdout", text: "burst\\n" }), + frame(json("request", "stream", { name: "stdout", text: "burst\n" })), ), ), ); if (mode === "frame") - peer!.write( - message("display_data", { - data: { "text/plain": "x".repeat(NOTEBOOK_LIMITS.bytes + 65536) }, - metadata: {}, - }), + write( + frame( + json("request", "display_data", { + data: { "text/plain": "x".repeat(9 * 1024 * 1024) }, + metadata: {}, + }), + ), + ); + const echo = json("request", "execute_input", { + code: mode === "input64" ? "x".repeat(65536) : "", + execution_count: 7, + }); + // Both field order and JSON tokens cross WebSocket fragment boundaries. + for (let offset = 0; offset < echo.length; offset += 7) + write( + frame(echo.slice(offset, offset + 7), offset === 0 ? 1 : 0, offset + 7 >= echo.length), ); - const reply = JSON.stringify({ - header: { msg_type: "execute_reply" }, - parent_header: { msg_id: "request" }, - channel: "shell", - content: { status: "ok" }, + const idle = frame(json("request", "status", { execution_state: "idle" })); + const error = mode === "reply64" || mode === "reply9"; + const reply = json("request", "execute_reply", { + traceback: error ? ["x".repeat(mode === "reply9" ? 9 * 1024 * 1024 : 65536)] : [], + status: error ? "error" : "ok", }); const split = Math.floor(reply.length / 2); - peer!.write( - Buffer.concat([ - message("status", { execution_state: "idle" }), - frame(reply.slice(0, split), 1, false), - frame("ping", 9), - frame(reply.slice(split), 0), - ]), + const replies = Buffer.concat([ + frame(reply.slice(0, split), 1, false), + frame("ping", 9), + frame(reply.slice(split), 0), + ]); + write( + variant === "reply9" || variant.endsWith("-replyFirst") + ? Buffer.concat([replies, idle]) + : Buffer.concat([idle, replies]), ); const complete = yield* Deferred.make(); let deliveredBytes = 0; + const received: JupyterTransport.Message[] = []; yield* channels.messages.pipe( Stream.runForEach((message) => Effect.gen(function* () { + received.push(message); deliveredBytes += byteLength(message); if (reducer.reduce(message)?.complete) yield* Deferred.succeed(complete, undefined); }), @@ -147,12 +190,96 @@ it.live.each(["stream", "frame", "burst"] as const)( ); yield* Deferred.await(complete); yield* transport.control(endpoint, "owned", "interrupt"); - expect(interrupts).toBe(1); - expect(reducer.truncated).toBe(true); - expect(document.cells[0]?.outputs).toContainEqual( - expect.objectContaining({ text: "[T3 output truncated]\n" }), - ); + expect(interrupts()).toBe(1); + expect(document.cells[0]?.execution_count).toBe(7); + expect(reducer.truncated).toBe(["stream", "frame", "burst"].includes(mode)); + expect(received.find((m) => m.header.msg_type === "execute_input")?.content).toEqual({ + execution_count: 7, + }); + expect(received.find((m) => m.header.msg_type === "execute_reply")?.content).toEqual({ + status: error ? "error" : "ok", + }); expect(deliveredBytes).toBeLessThan(NOTEBOOK_LIMITS.frameBytes + 65536); expect(byteLength(document)).toBeLessThanOrEqual(NOTEBOOK_LIMITS.bytes); - }).pipe(Effect.scoped, Effect.provide(JupyterTransport.layer)), + }).pipe(Effect.scoped, provide), +); + +it.live("late A overflow during B retains A correlation and never truncates B", () => + Effect.gen(function* () { + const { channels, reducer, document, write } = yield* fixture; + reducer.begin("A", document.cells[0]!); + yield* channels.execute("A", "a"); + const aDone = yield* Deferred.make(); + const aReader = yield* channels.messages.pipe( + Stream.runForEach((message) => + Effect.gen(function* () { + if (reducer.reduce(message)?.complete) yield* Deferred.succeed(aDone, undefined); + }), + ), + Effect.forkScoped, + ); + write( + Buffer.concat([ + frame(json("A", "execute_reply", { status: "ok" })), + frame(json("A", "status", { execution_state: "idle" })), + ]), + ); + yield* Deferred.await(aDone); + yield* Fiber.interrupt(aReader); + reducer.begin("B", document.cells[1]!); + yield* channels.execute("B", "b"); + write( + Buffer.concat([ + ...Array.from({ length: 1000 }, () => + frame(json("A", "stream", { name: "stdout", text: "late\n" })), + ), + frame(json("B", "execute_reply", { status: "ok" })), + frame(json("B", "status", { execution_state: "idle" })), + ]), + ); + const bDone = yield* Deferred.make(); + const lossParents: Array = []; + yield* channels.messages.pipe( + Stream.runForEach((message) => + Effect.gen(function* () { + if (message.header.msg_type === "t3_output_lost") + lossParents.push(message.parent_header.msg_id); + if (reducer.reduce(message)?.complete && message.parent_header.msg_id === "B") + yield* Deferred.succeed(bDone, undefined); + }), + ), + Effect.forkScoped, + ); + yield* Deferred.await(bDone); + expect(lossParents).toContain("A"); + expect(lossParents).not.toContain("B"); + expect(document.cells[0]?.outputs).toContainEqual( + expect.objectContaining({ t3_truncated: true }), + ); + expect(document.cells[1]?.outputs).toEqual([]); + }).pipe(Effect.scoped, provide), +); + +it.live("incomplete oversized JSON never becomes a valid partial output or completion", () => + Effect.gen(function* () { + const { channels, write } = yield* fixture; + yield* channels.execute("request", "code"); + const received: JupyterTransport.Message[] = []; + const malformed = json("request", "execute_reply", { + traceback: ["x".repeat(9 * 1024 * 1024)], + status: "ok", + }).slice(0, -1); + write(frame(malformed)); + const error = yield* channels.messages.pipe( + Stream.tap((message) => + Effect.sync(() => { + received.push(message); + }), + ), + Stream.runDrain, + Effect.flip, + ); + expect(error).toMatchObject({ reason: "disconnected" }); + expect(received).toEqual([]); + }).pipe(Effect.scoped, provide), ); diff --git a/apps/server/src/notebook/JupyterTransport.ts b/apps/server/src/notebook/JupyterTransport.ts index 50d56a1c1e08b..847d2fdea2e28 100644 --- a/apps/server/src/notebook/JupyterTransport.ts +++ b/apps/server/src/notebook/JupyterTransport.ts @@ -3,55 +3,26 @@ import * as NodeCrypto from "node:crypto"; import { fetch } from "undici"; import { NOTEBOOK_LIMITS, NotebookError } from "@t3tools/contracts"; -import * as Cause from "effect/Cause"; import * as Context from "effect/Context"; import * as Effect from "effect/Effect"; import * as Layer from "effect/Layer"; -import * as Queue from "effect/Queue"; import * as Schema from "effect/Schema"; import * as Scope from "effect/Scope"; import * as Stream from "effect/Stream"; import { connectKernelSocket } from "./kernelSocket.ts"; +import * as KernelMessageAdmission from "./KernelMessageAdmission.ts"; const encodeJson = Schema.encodeSync(Schema.fromJsonString(Schema.Unknown)); const decodeJson = Schema.decodeSync(Schema.fromJsonString(Schema.Unknown)); -const Message = Schema.Struct({ - header: Schema.Struct({ msg_type: Schema.String }), - parent_header: Schema.Struct({ msg_id: Schema.optional(Schema.String) }), - channel: Schema.String, - content: Schema.Record(Schema.String, Schema.Unknown), - buffers: Schema.optional(Schema.Array(Schema.Unknown)), -}); -const mimeContent = { - data: Schema.Record(Schema.String, Schema.Unknown), - metadata: Schema.Record(Schema.String, Schema.Unknown), -}; -const contentValidators: Record boolean> = { - stream: Schema.is( - Schema.Struct({ name: Schema.Literals(["stdout", "stderr"]), text: Schema.String }), - ), - display_data: Schema.is(Schema.Struct(mimeContent)), - update_display_data: Schema.is(Schema.Struct(mimeContent)), - execute_result: Schema.is(Schema.Struct({ ...mimeContent, execution_count: Schema.Int })), - error: Schema.is( - Schema.Struct({ - ename: Schema.String, - evalue: Schema.String, - traceback: Schema.Array(Schema.String), - }), - ), - clear_output: Schema.is(Schema.Struct({ wait: Schema.Boolean })), -}; -const decodeMessage = Schema.decodeSync(Schema.fromJsonString(Message)); const decodeKernel = Schema.decodeUnknownEffect(Schema.Struct({ id: Schema.String })); -export type Message = typeof Message.Type; +export type { Message } from "./KernelMessageAdmission.ts"; export type Endpoint = { readonly url: string; readonly token: string; readonly kernelName: string; }; export type Channels = { - readonly messages: Stream.Stream; + readonly messages: Stream.Stream; readonly execute: (id: string, code: string) => Effect.Effect; readonly close: Effect.Effect; }; @@ -146,83 +117,7 @@ export const layer = Layer.sync(JupyterTransport, () => { ); }); const connect = Effect.fn("JupyterTransport.connect")(function* (endpoint: Endpoint, id: string) { - const wake = yield* Queue.sliding(1); - const output: Message[] = []; - const control = new Map(); - let retainedBytes = 0; - let currentId: string | undefined; - let lost = false; - const loss = () => { - lost = true; - Queue.offerUnsafe(wake, undefined); - }; - const fail = () => { - Queue.failCauseUnsafe(wake, Cause.fail(new NotebookError({ reason: "disconnected" }))); - }; - const receive = (text: string | null) => { - try { - if (text === null) { - loss(); - return; - } - // Admission is bounded before JSON decoding; large data is discarded while channels drain. - if (Buffer.byteLength(text) > NOTEBOOK_LIMITS.bytes + 65536) { - loss(); - return; - } - let message = decodeMessage(text); - if ( - message.buffers?.length || - contentValidators[message.header.msg_type]?.(message.content) === false - ) { - fail(); - return; - } - const type = message.header.msg_type; - if ( - message.parent_header.msg_id === currentId && - currentId && - (type === "execute_reply" || - type === "execute_input" || - (type === "status" && message.content.execution_state === "idle")) - ) { - if (Buffer.byteLength(encodeJson(message)) > 65536) { - fail(); - return; - } - control.set(type, message); - } else { - if ( - type === "stream" && - typeof message.content.text === "string" && - Buffer.byteLength(message.content.text) > 65532 - ) { - message = { - ...message, - content: { - ...message.content, - text: Buffer.from(message.content.text).subarray(0, 65528).toString("utf8"), - t3_truncated: true, - }, - }; - } - const size = Buffer.byteLength(encodeJson(message)); - if ( - size > NOTEBOOK_LIMITS.frameBytes || - retainedBytes + size > NOTEBOOK_LIMITS.frameBytes || - output.length >= 64 - ) { - loss(); - return; - } - output.push(message); - retainedBytes += size; - } - Queue.offerUnsafe(wake, undefined); - } catch { - fail(); - } - }; + const admission = yield* KernelMessageAdmission.make; const url = yield* Effect.try({ try: () => { const url = endpointUrl(endpoint, "/" + encodeURIComponent(id) + "/channels"); @@ -231,14 +126,11 @@ export const layer = Layer.sync(JupyterTransport, () => { }, catch: () => new NotebookError({ reason: "unavailable" }), }); - const socket = yield* connectKernelSocket(url, endpoint.token, receive, fail); - const close = Effect.sync(() => { - socket.close(); - Queue.shutdownUnsafe(wake); - }); + const socket = yield* connectKernelSocket(url, endpoint.token, admission, admission.fail); + const close = Effect.sync(() => socket.close()).pipe(Effect.andThen(admission.close)); yield* Effect.addFinalizer(() => close); const execute = Effect.fn("JupyterTransport.execute")(function* (id: string, code: string) { - currentId = id; + admission.execute(id); yield* Effect.try({ try: () => socket.send( @@ -266,29 +158,7 @@ export const layer = Layer.sync(JupyterTransport, () => { catch: () => new NotebookError({ reason: "disconnected" }), }); }); - const messages = Stream.fromQueue(wake).pipe( - Stream.flatMap(() => { - const batch = output.splice(0); - retainedBytes = 0; - const input = control.get("execute_input"); - if (input) batch.unshift(input); - if (lost && currentId) - batch.push({ - header: { msg_type: "t3_output_lost" }, - parent_header: { msg_id: currentId }, - channel: "iopub", - content: {}, - }); - for (const type of ["execute_reply", "status"]) { - const message = control.get(type); - if (message) batch.push(message); - } - control.clear(); - lost = false; - return Stream.fromIterable(batch); - }), - ); - return { messages, execute, close }; + return { messages: admission.messages, execute, close }; }); return JupyterTransport.of({ create, control, connect }); }); diff --git a/apps/server/src/notebook/KernelMessageAdmission.ts b/apps/server/src/notebook/KernelMessageAdmission.ts new file mode 100644 index 0000000000000..b59dd26f69c10 --- /dev/null +++ b/apps/server/src/notebook/KernelMessageAdmission.ts @@ -0,0 +1,290 @@ +import { NOTEBOOK_LIMITS, NotebookError } from "@t3tools/contracts"; +import * as Cause from "effect/Cause"; +import * as Effect from "effect/Effect"; +import * as Queue from "effect/Queue"; +import * as Schema from "effect/Schema"; +import * as Stream from "effect/Stream"; +import { none, type Many } from "stream-chain/defs.js"; +import * as StreamJson from "stream-json/core/parser.js"; +import type { ParserOptions, Token } from "stream-json/core/parser.js"; + +const Message = Schema.Struct({ + header: Schema.Struct({ msg_type: Schema.String }), + parent_header: Schema.Struct({ msg_id: Schema.optional(Schema.String) }), + channel: Schema.String, + content: Schema.Record(Schema.String, Schema.Unknown), + buffers: Schema.optional(Schema.Array(Schema.Unknown)), +}); +export type Message = typeof Message.Type; +const decodeMessage = Schema.decodeSync(Schema.fromJsonString(Message)); +const mime = { + data: Schema.Record(Schema.String, Schema.Unknown), + metadata: Schema.Record(Schema.String, Schema.Unknown), +}; +const validators: Record boolean> = { + stream: Schema.is( + Schema.Struct({ name: Schema.Literals(["stdout", "stderr"]), text: Schema.String }), + ), + display_data: Schema.is(Schema.Struct(mime)), + update_display_data: Schema.is(Schema.Struct(mime)), + execute_result: Schema.is(Schema.Struct({ ...mime, execution_count: Schema.Int })), + error: Schema.is( + Schema.Struct({ + ename: Schema.String, + evalue: Schema.String, + traceback: Schema.Array(Schema.String), + }), + ), + clear_output: Schema.is(Schema.Struct({ wait: Schema.Boolean })), +}; +const normalized = Schema.Struct({ + type: Schema.String, + channel: Schema.String, + parent: Schema.optional(Schema.String), + count: Schema.optional(Schema.Unknown), + status: Schema.optional(Schema.Unknown), + state: Schema.optional(Schema.Unknown), +}); +const decodeEnvelope = Schema.decodeUnknownSync(normalized); +const encodeJson = Schema.encodeSync(Schema.fromJsonString(Schema.Unknown)); + +// The synchronous tokenizer is used elsewhere in this server; 3.6.0 omits its type declaration. +const { jsonParser } = StreamJson as typeof StreamJson & { + jsonParser: ( + options: ParserOptions, + ) => (input: string | typeof none) => Many | typeof none; +}; + +function messageReader(raw: Buffer) { + const tokenize = jsonParser({ packValues: false }); + const decoder = new TextDecoder("utf-8", { fatal: true }); + let rawSize = 0, + oversized = false, + key = "", + scalar = "", + readingKey = false; + let field: string | undefined; + let scalarOverflow = false; + const envelope: Record = {}; + const stack: Array<{ container: string | undefined; key: string; array: boolean }> = []; + let rootComplete = false; + let buffers = false; + const selected = () => { + const parent = stack.at(-1); + if (stack.length === 1 && parent?.key === "channel") return "channel"; + if (stack.length !== 2) return undefined; + if (parent?.container === "header" && parent.key === "msg_type") return "type"; + if (parent?.container === "parent_header" && parent.key === "msg_id") return "parent"; + if (parent?.container !== "content") return undefined; + return ( + { execution_count: "count", status: "status", execution_state: "state" } as Record< + string, + string + > + )[parent.key]; + }; + const tokens = (input: string | typeof none) => { + const batch = tokenize(input); + if (batch === none) return; + for (const token of batch.values) { + const parent = stack.at(-1); + if (parent?.container === "buffers" && token.name !== "endArray") buffers = true; + switch (token.name) { + case "startObject": + case "startArray": { + if (rootComplete || stack.length >= 128) throw new Error(); + const container = stack.length === 1 ? parent?.key : undefined; + if (container === "header") delete envelope.type; + if (container === "parent_header") delete envelope.parent; + if (container === "content") { + delete envelope.count; + delete envelope.status; + delete envelope.state; + } + stack.push({ container, key: "", array: token.name === "startArray" }); + break; + } + case "endObject": + case "endArray": + stack.pop(); + if (!stack.length) rootComplete = true; + break; + case "startKey": + key = ""; + readingKey = true; + break; + case "endKey": + if (parent) parent.key = key; + readingKey = false; + break; + case "startString": + case "startNumber": + field = selected(); + scalar = ""; + scalarOverflow = false; + break; + case "stringChunk": + case "numberChunk": + if (readingKey) { + if (key.length <= 64) key = (key + token.value).slice(0, 65); + } else if (field && !scalarOverflow) { + scalar += token.value; + if (scalar.length > 128) { + scalar = ""; + scalarOverflow = true; + } + } + break; + case "endString": + case "endNumber": + if (field) { + if (scalarOverflow) delete envelope[field]; + else envelope[field] = token.name === "endNumber" ? Number(scalar) : scalar; + } + field = undefined; + break; + case "nullValue": + case "trueValue": + case "falseValue": { + const target = selected(); + if (target) envelope[target] = token.value; + break; + } + } + } + }; + return { + write(chunk: Uint8Array) { + if (rawSize + chunk.length > raw.length) oversized = true; + if (!oversized) { + raw.set(chunk, rawSize); + rawSize += chunk.length; + } + tokens(decoder.decode(chunk, { stream: true })); + }, + finish() { + tokens(decoder.decode()); + tokens(none); + if (!rootComplete || stack.length || buffers) throw new Error(); + const e = decodeEnvelope(envelope); + const base = { + header: { msg_type: e.type }, + parent_header: e.parent === undefined ? {} : { msg_id: e.parent }, + channel: e.channel, + }; + if (e.type === "execute_input") { + if (typeof e.count !== "number" || !Number.isInteger(e.count)) throw new Error(); + return { + message: { ...base, content: { execution_count: e.count } } as Message, + lost: false, + }; + } + if (e.type === "execute_reply") { + if (typeof e.status !== "string" || !["ok", "error", "abort"].includes(e.status)) + throw new Error(); + return { message: { ...base, content: { status: e.status } } as Message, lost: false }; + } + if (e.type === "status") { + if (typeof e.state !== "string" || !["busy", "idle", "starting"].includes(e.state)) + throw new Error(); + return { + message: { ...base, content: { execution_state: e.state } } as Message, + lost: false, + }; + } + if (oversized) return { message: { ...base, content: {} } as Message, lost: true }; + const message = decodeMessage( + new TextDecoder("utf-8", { fatal: true }).decode(raw.subarray(0, rawSize)), + ); + if (message.buffers?.length || validators[e.type]?.(message.content) === false) + throw new Error(); + return { message, lost: false }; + }, + }; +} + +export const make = Effect.gen(function* () { + const wake = yield* Queue.sliding(1); + const output: Message[] = []; + const controls = new Map(); + const losses = new Set(); + let unattributed = false, + retained = 0; + let active: string | undefined; + const raw = Buffer.allocUnsafe(NOTEBOOK_LIMITS.frameBytes); + let reader = messageReader(raw); + const loss = (parent: string | undefined) => { + if (parent === undefined || (!losses.has(parent) && losses.size === 64)) unattributed = true; + else losses.add(parent); + }; + const end = () => { + const { message, lost } = reader.finish(); + reader = messageReader(raw); + const parent = message.parent_header.msg_id; + const type = message.header.msg_type; + if (lost) loss(parent); + else if ( + parent === active && + active && + (type === "execute_input" || + type === "execute_reply" || + (type === "status" && message.content.execution_state === "idle")) + ) + controls.set(type, message); + else { + const size = Buffer.byteLength(encodeJson(message)); + if (output.length === 64 || retained + size > NOTEBOOK_LIMITS.frameBytes) loss(parent); + else { + output.push(message); + retained += size; + } + } + Queue.offerUnsafe(wake, undefined); + }; + const messages = Stream.fromQueue(wake).pipe( + Stream.flatMap(() => { + const batch = output.splice(0); + retained = 0; + const input = controls.get("execute_input"); + if (input) batch.unshift(input); + for (const parent of losses) + batch.push({ + header: { msg_type: "t3_output_lost" }, + parent_header: { msg_id: parent }, + channel: "iopub", + content: {}, + }); + if (unattributed) + batch.push({ + header: { msg_type: "t3_output_lost" }, + parent_header: {}, + channel: "iopub", + content: {}, + }); + for (const type of ["execute_reply", "status"]) { + const message = controls.get(type); + if (message) batch.push(message); + } + controls.clear(); + losses.clear(); + unattributed = false; + return Stream.fromIterable(batch); + }), + ); + return { + write: (bytes: Uint8Array) => { + for (let offset = 0; offset < bytes.length; offset += 8192) + reader.write(bytes.subarray(offset, offset + 8192)); + }, + end, + execute: (id: string) => { + active = id; + }, + messages, + fail: () => + Queue.failCauseUnsafe(wake, Cause.fail(new NotebookError({ reason: "disconnected" }))), + close: Effect.sync(() => { + Queue.shutdownUnsafe(wake); + }), + }; +}); diff --git a/apps/server/src/notebook/executionReducer.ts b/apps/server/src/notebook/executionReducer.ts index 2b6571757c03c..8adbe2b2059e6 100644 --- a/apps/server/src/notebook/executionReducer.ts +++ b/apps/server/src/notebook/executionReducer.ts @@ -30,6 +30,13 @@ export class ExecutionReducer { this.owners.set(id, { cell, reply: false, idle: false, failed: false, clear: false }); } reduce(message: Message) { + if ( + message.header.msg_type === "t3_output_lost" && + message.parent_header.msg_id === undefined + ) { + this.truncated = true; + return { patches: [], cellId: null, complete: false, failed: false }; + } const owner = this.owners.get(message.parent_header.msg_id ?? ""); if (!owner) return undefined; const { cell } = owner; diff --git a/apps/server/src/notebook/kernelSocket.ts b/apps/server/src/notebook/kernelSocket.ts index b0f604b1b8b54..a622fe1b3ecba 100644 --- a/apps/server/src/notebook/kernelSocket.ts +++ b/apps/server/src/notebook/kernelSocket.ts @@ -1,12 +1,17 @@ // @effect-diagnostics nodeBuiltinImport:off -// Streaming frame admission must discard oversized messages before allocating their declared length. +// Frame payloads stream to the message owner without buffering their declared length. import * as NodeHttp from "node:http"; import * as NodeCrypto from "node:crypto"; import * as Effect from "effect/Effect"; import { NotebookError, NOTEBOOK_LIMITS } from "@t3tools/contracts"; export const connectKernelSocket = Effect.fn("JupyterTransport.connectSocket")( - (url: URL, token: string, receive: (text: string | null) => void, fail: () => void) => + ( + url: URL, + token: string, + receive: { write: (bytes: Uint8Array) => void; end: () => void }, + fail: () => void, + ) => Effect.callback<{ send: (text: string) => void; close: () => void }, NotebookError>( (resume) => { const key = NodeCrypto.randomBytes(16).toString("base64"); @@ -43,16 +48,13 @@ export const connectKernelSocket = Effect.fn("JupyterTransport.connectSocket")( return; } const capacity = NOTEBOOK_LIMITS.bytes + 65536; - const message = Buffer.allocUnsafe(capacity); const header = Buffer.alloc(10); let headerSize = 0, headerNeeded = 2, remaining = 0n, opcode = 0, final = false; - let messageSize = 0, - fragmented = false, - discarded = false; + let fragmented = false; const ping = Buffer.alloc(125); let pingSize = 0; let closed = false; @@ -94,15 +96,7 @@ export const connectKernelSocket = Effect.fn("JupyterTransport.connectSocket")( sendFrame(10, ping.subarray(0, pingSize)); } if ((opcode === 0 || opcode === 1) && final) { - if (discarded) receive(null); - else - receive( - new TextDecoder("utf-8", { fatal: true }).decode( - message.subarray(0, messageSize), - ), - ); - messageSize = 0; - discarded = false; + receive.end(); fragmented = false; } headerSize = 0; @@ -149,7 +143,6 @@ export const connectKernelSocket = Effect.fn("JupyterTransport.connectSocket")( } if (opcode === 1) fragmented = !final; if (opcode >= 8) pingSize = 0; - if (opcode < 8 && remaining > BigInt(capacity - messageSize)) discarded = true; if (remaining === 0n) { finish(); continue; @@ -163,10 +156,7 @@ export const connectKernelSocket = Effect.fn("JupyterTransport.connectSocket")( if (opcode === 9) { chunk.copy(ping, pingSize, offset, offset + count); pingSize += count; - } else if (opcode < 8 && !discarded) { - chunk.copy(message, messageSize, offset, offset + count); - messageSize += count; - } + } else if (opcode < 8) receive.write(chunk.subarray(offset, offset + count)); offset += count; remaining -= BigInt(count); if (remaining === 0n) finish(); From d5739bd76db1f62a947efd27675dfd0f7454eb4a Mon Sep 17 00:00:00 2001 From: Tanner Polley <42780259+tannerpolley@users.noreply.github.com> Date: Wed, 28 Oct 2026 14:56:01 -0600 Subject: [PATCH 4/4] fix(notebook): preserve aborted execution replies and provenance --- .../src/notebook/JupyterTransport.test.ts | 12 +++-- .../src/notebook/KernelMessageAdmission.ts | 2 +- .../src/notebook/NotebookService.test.ts | 44 +++++++++++++++++++ apps/server/src/notebook/NotebookService.ts | 28 +++++++++++- apps/server/src/notebook/executionReducer.ts | 20 +++++++-- packages/contracts/src/notebook.ts | 1 + 6 files changed, 97 insertions(+), 10 deletions(-) diff --git a/apps/server/src/notebook/JupyterTransport.test.ts b/apps/server/src/notebook/JupyterTransport.test.ts index ef947f72bb633..5e6aba23a11de 100644 --- a/apps/server/src/notebook/JupyterTransport.test.ts +++ b/apps/server/src/notebook/JupyterTransport.test.ts @@ -8,7 +8,7 @@ import * as Effect from "effect/Effect"; import * as Stream from "effect/Stream"; import * as Fiber from "effect/Fiber"; import * as Deferred from "effect/Deferred"; -import { NOTEBOOK_LIMITS } from "@t3tools/contracts"; +import { NOTEBOOK_LIMITS, type NotebookError } from "@t3tools/contracts"; import * as JupyterTransport from "./JupyterTransport.ts"; import { ExecutionReducer } from "./executionReducer.ts"; import { parseNotebook, byteLength } from "./notebookDocument.ts"; @@ -114,6 +114,8 @@ it.live.each([ "reply9-idleFirst", "reply64-replyFirst", "input64-replyFirst", + "aborted", + "aborted-replyFirst", ] as const)( "%s retains bounded data and normalized control fields across reordered, fragmented JSON", (variant) => @@ -160,9 +162,10 @@ it.live.each([ ); const idle = frame(json("request", "status", { execution_state: "idle" })); const error = mode === "reply64" || mode === "reply9"; + const replyStatus = mode === "aborted" ? "aborted" : error ? "error" : "ok"; const reply = json("request", "execute_reply", { traceback: error ? ["x".repeat(mode === "reply9" ? 9 * 1024 * 1024 : 65536)] : [], - status: error ? "error" : "ok", + status: replyStatus, }); const split = Math.floor(reply.length / 2); const replies = Buffer.concat([ @@ -175,7 +178,7 @@ it.live.each([ ? Buffer.concat([replies, idle]) : Buffer.concat([idle, replies]), ); - const complete = yield* Deferred.make(); + const complete = yield* Deferred.make(); let deliveredBytes = 0; const received: JupyterTransport.Message[] = []; yield* channels.messages.pipe( @@ -186,6 +189,7 @@ it.live.each([ if (reducer.reduce(message)?.complete) yield* Deferred.succeed(complete, undefined); }), ), + Effect.catch((error) => Deferred.fail(complete, error)), Effect.forkScoped, ); yield* Deferred.await(complete); @@ -197,7 +201,7 @@ it.live.each([ execution_count: 7, }); expect(received.find((m) => m.header.msg_type === "execute_reply")?.content).toEqual({ - status: error ? "error" : "ok", + status: replyStatus, }); expect(deliveredBytes).toBeLessThan(NOTEBOOK_LIMITS.frameBytes + 65536); expect(byteLength(document)).toBeLessThanOrEqual(NOTEBOOK_LIMITS.bytes); diff --git a/apps/server/src/notebook/KernelMessageAdmission.ts b/apps/server/src/notebook/KernelMessageAdmission.ts index b59dd26f69c10..9d80663e0b8b0 100644 --- a/apps/server/src/notebook/KernelMessageAdmission.ts +++ b/apps/server/src/notebook/KernelMessageAdmission.ts @@ -180,7 +180,7 @@ function messageReader(raw: Buffer) { }; } if (e.type === "execute_reply") { - if (typeof e.status !== "string" || !["ok", "error", "abort"].includes(e.status)) + if (typeof e.status !== "string" || !["ok", "error", "aborted"].includes(e.status)) throw new Error(); return { message: { ...base, content: { status: e.status } } as Message, lost: false }; } diff --git a/apps/server/src/notebook/NotebookService.test.ts b/apps/server/src/notebook/NotebookService.test.ts index b1356af0dda4f..0b7f12f79fd1b 100644 --- a/apps/server/src/notebook/NotebookService.test.ts +++ b/apps/server/src/notebook/NotebookService.test.ts @@ -102,6 +102,50 @@ const layer = Layer.mergeAll(NodeServices.layer, ProjectStore.layer).pipe( Layer.provideMerge(Sqlite.layerMemory), ); +it.effect.each([true, false])( + "preserves an aborted cell reply and stops remaining cells, reply first: %s", + (replyFirst) => + Effect.gen(function* () { + const { service, opened, handle, requests, send, finished } = yield* setup(); + const ready = yield* Deferred.make(); + const terminal = yield* Deferred.make(); + let replyStatus: string | undefined; + const observer = yield* service.subscribe(handle, caller).pipe( + Stream.runForEach((event) => + Effect.gen(function* () { + if (event.type === "snapshot") yield* Deferred.succeed(ready, undefined); + if (event.type === "changed") { + for (const patch of event.patches) + if (patch.type === "execution") replyStatus = patch.execution.replyStatus; + if (event.outcome === "aborted" || event.outcome === "error") + yield* Deferred.succeed(terminal, undefined); + } + }), + ), + Effect.forkScoped, + ); + yield* Deferred.await(ready); + const runId = yield* service.execute( + { ...handle, revision: opened.revision, cellIds: ["a", "b"], requestId: "aborted" }, + caller, + ); + const request = yield* Queue.take(requests); + if (replyFirst) yield* send(request.id, "execute_reply", { status: "aborted" }); + yield* send(request.id, "status", { execution_state: "idle" }); + if (!replyFirst) yield* send(request.id, "execute_reply", { status: "aborted" }); + yield* Deferred.await(terminal); + yield* finished(runId); + const result = yield* service.snapshot(handle, caller); + expect(result.outcome).toBe("aborted"); + expect(result.executions.a?.replyStatus).toBe("aborted"); + expect(replyStatus).toBe("aborted"); + expect(result.executions.b).toBeUndefined(); + expect(yield* Queue.size(requests)).toBe(0); + expect((result.document.cells as Array<{ outputs: unknown[] }>)[0]?.outputs).toEqual([]); + yield* Fiber.interrupt(observer); + }).pipe(Effect.scoped, Effect.provide(layer)), +); + it.effect("project aliases share the canonical notebook owner and revalidate each project", () => Effect.gen(function* () { const { service, opened, handle, sql, root, requests, send, finished } = yield* setup(); diff --git a/apps/server/src/notebook/NotebookService.ts b/apps/server/src/notebook/NotebookService.ts index 306e3a40fa1b7..6e9809bd4975f 100644 --- a/apps/server/src/notebook/NotebookService.ts +++ b/apps/server/src/notebook/NotebookService.ts @@ -232,10 +232,34 @@ const make = Effect.gen(function* () { if (session.snapshot.generation !== generation) return; const result = session.reducer.reduce(message); if (!result) return; - changed(session, result.cellId, result.patches); + const patches = [...result.patches]; + if ( + message.header.msg_type === "execute_reply" && + result.cellId !== null && + result.replyStatus !== undefined + ) { + const execution = session.snapshot.executions[result.cellId]; + if (execution) { + const updated = { ...execution, replyStatus: result.replyStatus }; + session.snapshot.executions = { + ...session.snapshot.executions, + [result.cellId]: updated, + }; + const cell = session.reducer.document.cells.find((cell) => cell.id === result.cellId); + patches.push({ + type: "execution", + cellId: result.cellId, + count: cell?.execution_count ?? null, + execution: updated, + }); + } + } + changed(session, result.cellId, patches); const pending = session.pending; - if (result.complete && pending && pending.id === message.parent_header.msg_id) + if (result.complete && pending && pending.id === message.parent_header.msg_id) { + if (result.replyStatus === "aborted") session.cancelled = true; yield* Deferred.succeed(pending.done, result.failed); + } }), ).pipe( Effect.catch(() => diff --git a/apps/server/src/notebook/executionReducer.ts b/apps/server/src/notebook/executionReducer.ts index 8adbe2b2059e6..b4d61dc56ee51 100644 --- a/apps/server/src/notebook/executionReducer.ts +++ b/apps/server/src/notebook/executionReducer.ts @@ -2,7 +2,14 @@ import { NOTEBOOK_LIMITS, NotebookError, type NotebookOutputPatch } from "@t3too import type { Message } from "./JupyterTransport.ts"; import { byteLength, type Document, type NotebookCell } from "./notebookDocument.ts"; -type Owner = { cell: NotebookCell; reply: boolean; idle: boolean; failed: boolean; clear: boolean }; +type Owner = { + cell: NotebookCell; + reply: boolean; + idle: boolean; + failed: boolean; + clear: boolean; + replyStatus?: string; +}; export class ExecutionReducer { readonly owners = new Map(); displays = new WeakMap(); @@ -35,7 +42,7 @@ export class ExecutionReducer { message.parent_header.msg_id === undefined ) { this.truncated = true; - return { patches: [], cellId: null, complete: false, failed: false }; + return { patches: [], cellId: null, complete: false, failed: false, replyStatus: undefined }; } const owner = this.owners.get(message.parent_header.msg_id ?? ""); if (!owner) return undefined; @@ -47,6 +54,7 @@ export class ExecutionReducer { if (type === "execute_reply" && message.channel === "shell") { owner.reply = true; owner.failed = content.status !== "ok"; + if (typeof content.status === "string") owner.replyStatus = content.status; } if (type === "status" && message.channel === "iopub" && content.execution_state === "idle") owner.idle = true; @@ -148,6 +156,12 @@ export class ExecutionReducer { } } } - return { patches, cellId: cell.id, complete: owner.reply && owner.idle, failed: owner.failed }; + return { + patches, + cellId: cell.id, + complete: owner.reply && owner.idle, + failed: owner.failed, + replyStatus: owner.replyStatus, + }; } } diff --git a/packages/contracts/src/notebook.ts b/packages/contracts/src/notebook.ts index e2e43d0fd5b3e..9eeef1f440c50 100644 --- a/packages/contracts/src/notebook.ts +++ b/packages/contracts/src/notebook.ts @@ -49,6 +49,7 @@ export type NotebookControl = typeof NotebookControl.Type; export const NotebookInterrupt = Schema.Struct({ ...NotebookControl.fields, runId: Schema.String }); export type NotebookInterrupt = typeof NotebookInterrupt.Type; const NotebookExecution = Schema.Struct({ + replyStatus: Schema.optional(Schema.String), runId: Schema.String, generation: Schema.Int, sourceHash: Schema.String,