From e9d9cc80381185a87512c5f2f5ae92d183572e8f Mon Sep 17 00:00:00 2001 From: Wang Date: Thu, 23 Jul 2026 05:33:06 +0800 Subject: [PATCH 1/3] feat(runtime-host): add session continuity --- .../canonical-session-projection.test.ts | 259 ++++++ .../src/__tests__/connection-session.test.ts | 11 + .../src/__tests__/dependency-boundary.test.ts | 3 + .../src/__tests__/execution-host.test.ts | 167 ++++ .../__tests__/fixtures/uncooperative-host.ts | 8 + .../src/__tests__/host-kernel.test.ts | 2 + .../src/__tests__/message-coordinator.test.ts | 42 +- .../src/__tests__/protocol.test.ts | 250 ++++++ .../__tests__/root-turn-coordinator.test.ts | 72 +- .../__tests__/session-admission-gate.test.ts | 58 ++ .../session-continuity-coordinator.test.ts | 289 +++++++ .../session-subscription-client.test.ts | 381 +++++++++ .../runtime-host/src/client/connection.ts | 189 ++++- packages/runtime-host/src/client/index.ts | 6 + .../src/client/session-subscription.ts | 205 +++++ packages/runtime-host/src/protocol/index.ts | 14 +- .../runtime-host/src/protocol/operations.ts | 8 +- .../src/protocol/session-continuity.ts | 627 ++++++++++++++ .../server/canonical-session-projection.ts | 78 ++ .../src/server/canonical-turn-snapshot.ts | 93 +++ .../src/server/connection-session.ts | 58 +- .../src/server/execution-composition.ts | 33 +- .../runtime-host/src/server/host-kernel.ts | 3 + .../src/server/message-coordinator.ts | 30 +- .../src/server/operation-dispatcher.ts | 8 + .../src/server/root-admission-owner.ts | 13 + .../src/server/root-turn-coordinator.ts | 75 +- .../src/server/session-admission-gate.ts | 131 ++- .../server/session-continuity-coordinator.ts | 785 ++++++++++++++++++ .../src/server/session-continuity-service.ts | 21 + 30 files changed, 3870 insertions(+), 49 deletions(-) create mode 100644 packages/runtime-host/src/__tests__/canonical-session-projection.test.ts create mode 100644 packages/runtime-host/src/__tests__/session-continuity-coordinator.test.ts create mode 100644 packages/runtime-host/src/__tests__/session-subscription-client.test.ts create mode 100644 packages/runtime-host/src/client/session-subscription.ts create mode 100644 packages/runtime-host/src/protocol/session-continuity.ts create mode 100644 packages/runtime-host/src/server/canonical-session-projection.ts create mode 100644 packages/runtime-host/src/server/canonical-turn-snapshot.ts create mode 100644 packages/runtime-host/src/server/session-continuity-coordinator.ts create mode 100644 packages/runtime-host/src/server/session-continuity-service.ts diff --git a/packages/runtime-host/src/__tests__/canonical-session-projection.test.ts b/packages/runtime-host/src/__tests__/canonical-session-projection.test.ts new file mode 100644 index 0000000000..39752abcd8 --- /dev/null +++ b/packages/runtime-host/src/__tests__/canonical-session-projection.test.ts @@ -0,0 +1,259 @@ +import assert from 'node:assert/strict'; +import { mkdtemp, readFile, rm, writeFile } from 'node:fs/promises'; +import { tmpdir } from 'node:os'; +import { join } from 'node:path'; +import { test } from 'node:test'; +import type { AgentRunHeader, RuntimeEvent } from '@maka/core'; +import { + openInteractiveExecutionStoresForWrite, + type ExecutionStoresWriter, +} from '@maka/storage/execution-stores'; +import { resolveStorageRoot, tryAcquireInteractiveRootOwner } from '@maka/storage/root-authority'; +import { CanonicalSessionProjectionReader } from '../server/canonical-session-projection.js'; +import { type HostMessageRootPort, HostMessageCoordinator } from '../server/message-coordinator.js'; +import { RootAdmissionOwner } from '../server/root-admission-owner.js'; +import { SessionAdmissionGate } from '../server/session-admission-gate.js'; + +test('projects the canonical root lifecycle and the attachment queue from real Stores', async () => { + await withStores(async (root, stores) => { + const session = await stores.sessionStore.create(sessionInput(root)); + const rootAdmissions = new RootAdmissionOwner(stores.agentRunStore); + await rootAdmissions.recoverSession(session.id); + const messages = createMessages(session.id, stores); + const reader = new CanonicalSessionProjectionReader({ + stores, + rootAdmissions, + messages, + }); + + assert.deepEqual(await reader.read(session.id), { + session: { + sessionId: session.id, + status: session.status, + createdAt: session.createdAt, + lastUsedAt: session.lastUsedAt, + isArchived: false, + }, + rootTurn: null, + queue: { hostEpoch: 'epoch-1', queueRevision: 0, steering: [], followup: [] }, + }); + + const admitted = await rootAdmissions.admitRootTurn({ + sessionId: session.id, + turnId: 'turn-1', + proposedRunId: 'run-1', + proposedUserMessageId: 'user-1', + execution: { kind: 'external_message' }, + normalizedInput: { text: 'hello' }, + sourceMessages: [], + admittedAt: 10, + }); + const admittedProjection = await reader.read(session.id); + assert.ok(admittedProjection); + assert.equal(admittedProjection.rootTurn?.status, 'admitted'); + + await stores.agentRunStore.createRun(runHeader(session.id)); + await stores.agentRunStore.appendEvent(session.id, 'run-1', { + type: 'run_started', + id: 'run-started-1', + sessionId: session.id, + turnId: 'turn-1', + runId: 'run-1', + ts: 11, + }); + await stores.agentRunStore.updateRun(session.id, 'run-1', { + status: 'running', + updatedAt: 11, + }); + + messages.reserveRootTurn({ sessionId: session.id, turnId: 'turn-1', runId: 'run-1' }); + const attachment = { + kind: 'image' as const, + name: 'evidence.png', + mimeType: 'image/png', + bytes: 12, + ref: { kind: 'workspace_file' as const, relativePath: 'evidence.png' }, + }; + const submitted = await messages.handlers['turn.message.submit']( + { + originHostEpoch: 'epoch-1', + sessionId: session.id, + messageId: 'queued-1', + content: { text: 'inspect this', attachments: [attachment] }, + placement: 'current_turn', + }, + operationContext(), + ); + assert.equal(submitted.ok, true); + const running = await reader.read(session.id); + assert.ok(running); + assert.equal(running.rootTurn?.status, 'running'); + assert.deepEqual(running.queue.steering[0]?.content.attachments, [attachment]); + + const terminal = terminalEvent(session.id); + await stores.runtimeEventStore.appendRuntimeEvent(session.id, 'run-1', terminal); + await stores.agentRunStore.updateRun(session.id, 'run-1', { + status: 'completed', + updatedAt: 12, + completedAt: 12, + }); + const completed = await reader.read(session.id); + assert.ok(completed); + assert.deepEqual(completed.rootTurn, { + sessionId: session.id, + turnId: admitted.admission.turnId, + runId: admitted.admission.runId, + status: 'completed', + terminalEventId: terminal.id, + }); + + await messages.handlers['queue.retract']( + { originHostEpoch: 'epoch-1', sessionId: session.id, retractId: 'cleanup' }, + operationContext(), + ); + messages.abandonRootReservation({ sessionId: session.id, turnId: 'turn-1', runId: 'run-1' }); + await messages.close(); + }); +}); + +test('fails closed when the owned tip durable identity changes', async () => { + await withStores(async (root, stores) => { + const session = await stores.sessionStore.create(sessionInput(root)); + const rootAdmissions = new RootAdmissionOwner(stores.agentRunStore); + await rootAdmissions.recoverSession(session.id); + await rootAdmissions.admitRootTurn({ + sessionId: session.id, + turnId: 'turn-1', + proposedRunId: 'run-1', + proposedUserMessageId: 'user-1', + execution: { kind: 'external_message' }, + normalizedInput: { text: 'hello' }, + sourceMessages: [], + admittedAt: 10, + }); + const admissionPath = join(root, 'sessions', session.id, 'turn-admissions', 'turn-1.json'); + const original = await readFile(admissionPath, 'utf8'); + const durable = JSON.parse(original) as Record; + + await rm(admissionPath); + const missingReader = new CanonicalSessionProjectionReader({ + stores, + rootAdmissions, + messages: createMessages(session.id, stores), + }); + await assert.rejects(() => missingReader.read(session.id), /missing from durable storage/); + + await writeFile(admissionPath, `${JSON.stringify({ ...durable, runId: 'run-drifted' })}\n`); + + const reader = new CanonicalSessionProjectionReader({ + stores, + rootAdmissions, + messages: createMessages(session.id, stores), + }); + await assert.rejects(() => reader.read(session.id), /identity changed/); + }); +}); + +function createMessages( + sessionId: string, + stores: ExecutionStoresWriter<'interactive'>, +): HostMessageCoordinator { + const root: HostMessageRootPort = { + readSessionHeader: async () => ({ isArchived: false }), + readRootState: () => ({ kind: 'active', sessionId, turnId: 'turn-1', runId: 'run-1' }), + startFromMessage: async () => { + throw new Error('unexpected root start'); + }, + claimStop: async () => { + throw new Error('unexpected root stop'); + }, + }; + return new HostMessageCoordinator({ + hostEpoch: 'epoch-1', + root, + durableProof: { + readRootTurnSourceMessageReceipt: (requestedSessionId, messageId) => + stores.agentRunStore.readRootTurnSourceMessageReceipt(requestedSessionId, messageId), + readImmutableSteeringMessageProof: (requestedSessionId, messageId) => + stores.runtimeEventStore.readImmutableSteeringMessageProof(requestedSessionId, messageId), + }, + receipts: stores.messageReceiptStore, + sessionAdmission: new SessionAdmissionGate(), + acquireResidency: () => ({ release: () => undefined }), + createId: () => 'entry-1', + }); +} + +function sessionInput(root: string) { + return { + cwd: root, + backend: 'fake' as const, + llmConnectionSlug: 'fake', + model: 'fake-model', + permissionMode: 'ask' as const, + }; +} + +function runHeader(sessionId: string): AgentRunHeader { + return { + runId: 'run-1', + invocationId: 'run-1', + sessionId, + turnId: 'turn-1', + status: 'created', + backendKind: 'fake', + llmConnectionSlug: 'fake', + modelId: 'fake-model', + cwd: '/private/runtime-cwd', + permissionMode: 'ask', + createdAt: 10, + updatedAt: 10, + }; +} + +function terminalEvent(sessionId: string): RuntimeEvent { + return { + id: 'terminal-1', + invocationId: 'run-1', + sessionId, + turnId: 'turn-1', + runId: 'run-1', + ts: 12, + partial: false, + status: 'completed', + role: 'model', + author: 'agent', + content: { kind: 'text', text: 'done' }, + }; +} + +function operationContext() { + return { + hostEpoch: 'epoch-1', + connectionId: 'connection-1', + surface: 'tui' as const, + principal: 'local_os_user' as const, + acquireResidency: () => ({ release: () => undefined }), + }; +} + +async function withStores( + run: (root: string, stores: ExecutionStoresWriter<'interactive'>) => Promise, +): Promise { + const base = await mkdtemp(join(tmpdir(), 'maka-canonical-session-')); + const capability = await resolveStorageRoot({ + path: join(base, 'root'), + kind: 'interactive', + }); + const owner = await tryAcquireInteractiveRootOwner(capability); + assert.ok(owner); + if (!owner) throw new Error('Unable to acquire test root'); + try { + const stores = await openInteractiveExecutionStoresForWrite(owner.lease); + await stores.messageReceiptStore.beginHostEpoch('epoch-1'); + await run(capability.canonicalPath, stores); + } finally { + await owner.close(); + await rm(base, { recursive: true, force: true }); + } +} diff --git a/packages/runtime-host/src/__tests__/connection-session.test.ts b/packages/runtime-host/src/__tests__/connection-session.test.ts index e317b2b0e4..0b98813af8 100644 --- a/packages/runtime-host/src/__tests__/connection-session.test.ts +++ b/packages/runtime-host/src/__tests__/connection-session.test.ts @@ -310,6 +310,7 @@ test('connection reset while operation admission is pending does not execute the principal: 'local_os_user', }, resolveHandlers: () => handlers, + resolveContinuity: () => undefined, beginOperation: async () => { admissionEntered.resolve(); await releaseAdmission.promise; @@ -657,6 +658,7 @@ async function openHalfClosedDispatchedSession( }; }), }), + resolveContinuity: () => undefined, beginOperation: async () => ({ acquireResidency: () => ({ release() {} }), seal() {}, @@ -727,6 +729,13 @@ function createHandlers(queryTurn: TurnQueryHandler): RuntimeHostComposition['ha message: 'not available in this test composition', }, }; + const subscriptionUnavailable = { + ok: false, + error: { + code: 'operation_unavailable', + message: 'not available in this test composition', + }, + } as const; return { ...createUnavailableDomainOperationHandlers(), 'turn.start': async (input) => ({ @@ -741,6 +750,8 @@ function createHandlers(queryTurn: TurnQueryHandler): RuntimeHostComposition['ha 'turn.message.submit': async () => unavailable, 'queue.retract': async () => unavailable, 'turn.interrupt': async () => unavailable, + 'subscription.open': async () => subscriptionUnavailable, + 'subscription.close': async () => subscriptionUnavailable, }; } diff --git a/packages/runtime-host/src/__tests__/dependency-boundary.test.ts b/packages/runtime-host/src/__tests__/dependency-boundary.test.ts index 3f00bc8639..943cd52f71 100644 --- a/packages/runtime-host/src/__tests__/dependency-boundary.test.ts +++ b/packages/runtime-host/src/__tests__/dependency-boundary.test.ts @@ -33,6 +33,7 @@ const allowedServerExternalImports = new Set([ '@maka/runtime', '@maka/storage/execution-stores', '@maka/storage/runtime-policy-stores', + 'node:async_hooks', ]); const allowedExternalImports = { client: allowedHostExternalImports, @@ -132,6 +133,7 @@ test('the production Candidate dependency graph remains non-serving', () => { 'server/execution-composition.ts', 'server/root-turn-coordinator.ts', 'server/runtime-policy-coordinator.ts', + 'server/session-continuity-coordinator.ts', ]); const violations: string[] = []; for (const path of reached) { @@ -158,6 +160,7 @@ test('the public server entrypoint does not expose the test execution compositio 'server/execution-candidate.ts', 'server/execution-composition.ts', 'server/root-turn-coordinator.ts', + 'server/session-continuity-coordinator.ts', ]); assert.deepEqual( reachableModules(serverEntrypoint, publicEntrypoints) diff --git a/packages/runtime-host/src/__tests__/execution-host.test.ts b/packages/runtime-host/src/__tests__/execution-host.test.ts index be204c2a27..8485dd306a 100644 --- a/packages/runtime-host/src/__tests__/execution-host.test.ts +++ b/packages/runtime-host/src/__tests__/execution-host.test.ts @@ -31,11 +31,14 @@ import { import { connectRuntimeHost, RuntimeHostOperationError, + RuntimeHostSubscriptionError, type RuntimeHostConnection, + type RuntimeHostSessionSubscription, } from '../client/index.js'; import { decodeHostFrame, RUNTIME_HOST_PROTOCOL_VERSION, + type SubscriptionFrame, type TurnMessageSubmitInput, type TurnSnapshot, } from '../protocol/index.js'; @@ -177,6 +180,101 @@ test('two Clients share one execution after the starting Client disconnects', as }); }); +test('subscribed Clients share one canonical queue and ordered root handoff', async () => { + await withExecutionRoot(async (fixture) => { + const host = await fixture.startHost(); + const desktop = await connectClient(fixture.root, 'desktop'); + const tui = await connectClient(fixture.root, 'tui'); + const desktopSubscription = await desktop.openSessionSubscription({ + sessionId: fixture.sessionId, + }); + const tuiSubscription = await tui.openSessionSubscription({ sessionId: fixture.sessionId }); + const desktopProbe = new SubscriptionProbe(desktopSubscription); + const tuiProbe = new SubscriptionProbe(tuiSubscription); + for (const subscription of [desktopSubscription, tuiSubscription]) { + assert.equal(subscription.hostEpoch, host.hostEpoch); + assert.equal(subscription.snapshot.rootTurn, null); + assert.equal(subscription.snapshot.projectionRevision, 1); + assert.equal(subscription.snapshot.queue.hostEpoch, host.hostEpoch); + } + + const firstTurnId = randomUUID(); + const started = await desktop.startTurn({ + sessionId: fixture.sessionId, + turnId: firstTurnId, + content: { text: `continuity root ${'x'.repeat(540)}` }, + }); + for (const probe of [desktopProbe, tuiProbe]) { + const liveDelta = await probe.waitFor( + (frame) => + frame.kind === 'subscription.session_delta' && frame.delta.turnId === firstTurnId, + 'continuity did not publish the live assistant delta', + ); + assert.equal(liveDelta.kind, 'subscription.session_delta'); + if (liveDelta.kind === 'subscription.session_delta') { + assert.equal(liveDelta.delta.runId, started.runId); + } + } + + const followupId = randomUUID(); + const followupContent = { text: 'continue after the first root completes' }; + const queued = await tui.request('turn.message.submit', { + originHostEpoch: host.hostEpoch, + sessionId: fixture.sessionId, + messageId: followupId, + content: followupContent, + placement: 'next_turn', + }); + assert.equal(queued.disposition, 'followup'); + for (const probe of [desktopProbe, tuiProbe]) { + const queueProjection = await probe.waitFor( + (frame) => + frame.kind === 'subscription.session_projection' && + frame.snapshot.queue.followup.some((entry) => entry.messageId === followupId), + 'continuity did not publish the accepted follow-up', + ); + assert.equal(queueProjection.kind, 'subscription.session_projection'); + } + + await desktop.close(); + await desktopProbe.waitForFailure('connection_closed'); + assert.equal((await tui.status()).connections, 1); + const terminal = await tuiProbe.waitFor( + (frame) => + frame.kind === 'subscription.session_projection' && + frame.snapshot.rootTurn?.turnId === firstTurnId && + frame.snapshot.rootTurn.status === 'completed', + 'continuity did not publish the terminal root cut', + ); + assert.equal(terminal.kind, 'subscription.session_projection'); + const successor = await tuiProbe.waitFor( + (frame) => + frame.kind === 'subscription.session_projection' && + frame.snapshot.rootTurn !== null && + frame.snapshot.rootTurn.turnId !== firstTurnId, + 'continuity did not publish the successor root', + ); + assert.equal(successor.kind, 'subscription.session_projection'); + if (successor.kind !== 'subscription.session_projection' || !successor.snapshot.rootTurn) { + return; + } + assert.equal(successor.snapshot.rootTurn.sessionId, fixture.sessionId); + assert.ok(tuiProbe.indexOf(terminal) < tuiProbe.indexOf(successor)); + await tuiSubscription.close(); + await tuiProbe.done; + await waitForTerminalTurn(tui, fixture.sessionId, successor.snapshot.rootTurn.turnId); + await tui.close(); + await fixture.stopHost(host); + + const chain = await fixture.readAdmissionChain(); + assert.deepEqual( + chain.map((admission) => admission.turnId), + [firstTurnId, successor.snapshot.rootTurn.turnId], + ); + assert.deepEqual(chain[1]?.normalizedInput, followupContent); + }); +}); + test('concurrent root admission for one Session has a single winner', async () => { await withExecutionRoot(async (fixture) => { const host = await fixture.startHost(); @@ -259,23 +357,45 @@ test('a killed Host is recovered exactly once before its successor becomes ready await withExecutionRoot(async (fixture) => { const firstHost = await fixture.startHost(); const first = await connectClient(fixture.root, 'desktop'); + const firstSubscription = await first.openSessionSubscription({ + sessionId: fixture.sessionId, + }); + const firstProbe = new SubscriptionProbe(firstSubscription); const turnId = randomUUID(); const started = await first.startTurn({ sessionId: fixture.sessionId, turnId, content: { text: FAKE_ASK_USER_QUESTION_PROMPT }, }); + await firstProbe.waitFor( + (frame) => + frame.kind === 'subscription.session_projection' && + frame.snapshot.rootTurn?.runId === started.runId && + frame.snapshot.rootTurn.status !== 'admitted', + 'first Host did not publish the active root projection', + ); await fixture.killHost(firstHost); await first.closed; + await firstProbe.waitForFailure('connection_closed'); const secondHost = await fixture.startHost(); const second = await connectClient(fixture.root, 'tui'); + const recoveredSubscription = await second.openSessionSubscription({ + sessionId: fixture.sessionId, + }); const recovered = await second.queryTurn({ sessionId: fixture.sessionId, turnId, }); assert.equal(recovered.status, 'failed'); if (recovered.status === 'failed') assert.equal(recovered.failureClass, 'app_restarted'); + assert.notEqual(recoveredSubscription.hostEpoch, firstSubscription.hostEpoch); + assert.equal(recoveredSubscription.snapshot.projectionRevision, 1); + assert.deepEqual(recoveredSubscription.snapshot.rootTurn, recovered); + assert.equal(recoveredSubscription.snapshot.queue.hostEpoch, recoveredSubscription.hostEpoch); + assert.deepEqual(recoveredSubscription.snapshot.queue.steering, []); + assert.deepEqual(recoveredSubscription.snapshot.queue.followup, []); + await recoveredSubscription.close(); await second.close(); await fixture.stopHost(secondHost); @@ -1625,6 +1745,53 @@ async function waitForTurn( } } +class SubscriptionProbe { + readonly frames: SubscriptionFrame[] = []; + readonly done: Promise; + #failure: unknown; + #settled = false; + + constructor(subscription: RuntimeHostSessionSubscription) { + this.done = this.#consume(subscription); + } + + async waitFor( + predicate: (frame: SubscriptionFrame) => boolean, + message: string, + ): Promise { + const deadline = Date.now() + PROCESS_TIMEOUT_MS; + while (true) { + const frame = this.frames.find(predicate); + if (frame) return frame; + if (this.#failure) throw this.#failure; + if (this.#settled) throw new Error(`${message}: subscription closed`); + if (Date.now() >= deadline) throw new Error(message); + await sleep(10); + } + } + + async waitForFailure(reason: RuntimeHostSubscriptionError['reason']): Promise { + const deadline = Date.now() + PROCESS_TIMEOUT_MS; + while (!this.#failure && !this.#settled && Date.now() < deadline) await sleep(10); + assert.ok(this.#failure instanceof RuntimeHostSubscriptionError); + assert.equal(this.#failure.reason, reason); + } + + indexOf(frame: SubscriptionFrame): number { + return this.frames.indexOf(frame); + } + + async #consume(subscription: RuntimeHostSessionSubscription): Promise { + try { + for await (const frame of subscription) this.frames.push(frame); + } catch (error) { + this.#failure = error; + } finally { + this.#settled = true; + } + } +} + async function waitForTerminalTurn( connection: RuntimeHostConnection, sessionId: string, diff --git a/packages/runtime-host/src/__tests__/fixtures/uncooperative-host.ts b/packages/runtime-host/src/__tests__/fixtures/uncooperative-host.ts index 39d3081087..eb8e209b57 100644 --- a/packages/runtime-host/src/__tests__/fixtures/uncooperative-host.ts +++ b/packages/runtime-host/src/__tests__/fixtures/uncooperative-host.ts @@ -55,6 +55,14 @@ const host = await RuntimeHostKernel.start({ ok: false, error: { code: 'operation_unavailable', message: 'Operation unavailable in test Host' }, }), + 'subscription.open': async () => ({ + ok: false, + error: { code: 'operation_unavailable', message: 'Operation unavailable in test Host' }, + }), + 'subscription.close': async () => ({ + ok: false, + error: { code: 'operation_unavailable', message: 'Operation unavailable in test Host' }, + }), }, async recover() {}, async close() {}, diff --git a/packages/runtime-host/src/__tests__/host-kernel.test.ts b/packages/runtime-host/src/__tests__/host-kernel.test.ts index 7550cf0cbe..96c6595afc 100644 --- a/packages/runtime-host/src/__tests__/host-kernel.test.ts +++ b/packages/runtime-host/src/__tests__/host-kernel.test.ts @@ -139,6 +139,8 @@ describe('non-serving Runtime Host kernel', () => { 'turn.message.submit': unavailable, 'queue.retract': unavailable, 'turn.interrupt': unavailable, + 'subscription.open': unavailable, + 'subscription.close': unavailable, }, async recover() {}, async close() {}, diff --git a/packages/runtime-host/src/__tests__/message-coordinator.test.ts b/packages/runtime-host/src/__tests__/message-coordinator.test.ts index 795a6e63cc..6612ae7d69 100644 --- a/packages/runtime-host/src/__tests__/message-coordinator.test.ts +++ b/packages/runtime-host/src/__tests__/message-coordinator.test.ts @@ -51,6 +51,45 @@ test('idle submit starts exactly one root Turn and retry identity is connection- assert.equal(fixture.liveResidencies(), 0); }); +test('invalidates the canonical projection after each observable queue mutation', async () => { + const changedSessions: string[] = []; + const fixture = createFixture((sessionId) => changedSessions.push(sessionId)); + fixture.coordinator.reserveRootTurn(ROOT); + const owner = fixture.coordinator.bindRun(ROOT); + + assert.equal((await submit(fixture, 'steering-1', 'first', 'current_turn')).ok, true); + const [lease] = owner.pull(); + assert.ok(lease); + owner.ack([lease.id]); + owner.release(); + fixture.coordinator.completeIdle(fixture.coordinator.beginTerminalTransition(ROOT)); + + assert.deepEqual( + changedSessions, + Array.from({ length: 4 }, () => ROOT.sessionId), + ); + await fixture.coordinator.close(); +}); + +test('binds the exact reserved Run after a pre-bind stop fence', async () => { + const fixture = createFixture(); + fixture.coordinator.reserveRootTurn(ROOT); + assert.equal((await submit(fixture, 'queued-before-bind', 'discard me', 'next_turn')).ok, true); + + const fence = fixture.coordinator.commitStopFence(ROOT); + assert.equal(fence.retracted.length, 1); + assert.deepEqual(fixture.coordinator.projection(ROOT.sessionId).followup, []); + assert.equal(fixture.liveResidencies(), 0); + + const owner = fixture.coordinator.bindRun(ROOT); + assert.deepEqual(owner.pull(), []); + owner.release(); + const batch = fixture.coordinator.beginTerminalTransition(ROOT); + assert.deepEqual(batch.sources, []); + fixture.coordinator.completeIdle(batch); + await fixture.coordinator.close(); +}); + test('queue projection capacity is rejected before mutation or residency acquisition', async () => { const fixture = createFixture(); fixture.coordinator.reserveRootTurn(ROOT); @@ -1098,7 +1137,7 @@ test('canonical retry omits redundant display text and empty ordered refs', asyn fixture.coordinator.completeIdle(batch); }); -function createFixture() { +function createFixture(onProjectionChanged?: (sessionId: string) => void) { let nextId = 1; let liveResidencies = 0; let startCalls = 0; @@ -1215,6 +1254,7 @@ function createFixture() { requestDrain: () => { drainRequests += 1; }, + ...(onProjectionChanged ? { onProjectionChanged } : {}), createId: () => `id-${nextId++}`, }; coordinator = new HostMessageCoordinator(options); diff --git a/packages/runtime-host/src/__tests__/protocol.test.ts b/packages/runtime-host/src/__tests__/protocol.test.ts index 8b918f6c97..d557a54c53 100644 --- a/packages/runtime-host/src/__tests__/protocol.test.ts +++ b/packages/runtime-host/src/__tests__/protocol.test.ts @@ -6,6 +6,7 @@ import { decodeHostFrame, decodeHostRegistration, decodeSessionMessageQueueProjection, + decodeSessionContinuitySnapshot, encodeProtocolFrame, HOST_OPERATION_SPECS, MESSAGE_OPERATION_RESULT_MAX_BYTES, @@ -14,6 +15,10 @@ import { ProtocolFrameDecoder, RUNTIME_HOST_MAX_FRAME_BYTES, RUNTIME_HOST_PROTOCOL_VERSION, + SESSION_CONTINUITY_SCHEMA_VERSION, + SESSION_CONTINUITY_SNAPSHOT_MAX_BYTES, + SESSION_LIVE_DELTA_MAX_BYTES, + SESSION_TOOL_NAME_MAX_BYTES, TURN_MESSAGE_CONTENT_MAX_BYTES, TURN_MESSAGE_TEXT_MAX_BYTES, RUNTIME_POLICY_OPERATION_SPECS, @@ -50,6 +55,8 @@ describe('Runtime Host bootstrap protocol', () => { 'queue.retract', 'runtime.policy.mutate', 'runtime.policy.query', + 'subscription.close', + 'subscription.open', 'turn.interrupt', 'turn.message.submit', 'turn.query', @@ -86,6 +93,223 @@ describe('Runtime Host bootstrap protocol', () => { ); }); + test('keeps subscription operations closed, ready-only, and queue Epoch correlated', () => { + assert.equal(SESSION_CONTINUITY_SCHEMA_VERSION, 1); + assert.deepEqual( + Object.fromEntries( + (['subscription.open', 'subscription.close'] as const).map((operation) => [ + operation, + { + mode: HOST_OPERATION_SPECS[operation].mode, + availability: HOST_OPERATION_SPECS[operation].availability, + errors: HOST_OPERATION_SPECS[operation].errors, + }, + ]), + ), + { + 'subscription.open': { + mode: 'control', + availability: 'ready', + errors: [ + 'host_not_ready', + 'host_draining', + 'operation_unavailable', + 'not_found', + 'operation_conflict', + 'internal_failure', + ], + }, + 'subscription.close': { + mode: 'control', + availability: 'ready', + errors: [ + 'host_not_ready', + 'host_draining', + 'operation_unavailable', + 'not_found', + 'internal_failure', + ], + }, + }, + ); + const opened = { + requestId: 'open-1', + operation: 'subscription.open', + ok: true, + result: { + hostEpoch: 'epoch-1', + subscriptionId: 'subscription-1', + nextSequence: 1, + snapshot: continuitySnapshot('epoch-1'), + }, + }; + assert.deepEqual(decodeHostFrame(opened), opened); + assert.throws( + () => + decodeHostFrame({ + ...opened, + result: { ...opened.result, snapshot: continuitySnapshot('epoch-2') }, + }), + isInvalidFrame, + ); + assert.throws( + () => + decodeSessionContinuitySnapshot({ + ...continuitySnapshot('epoch-1'), + interactions: [], + }), + isInvalidFrame, + ); + }); + + test('decodes only privacy-normalized bounded subscription live frames', () => { + const envelope = { + kind: 'subscription.session_event' as const, + hostEpoch: 'epoch-1', + subscriptionId: 'subscription-1', + sequence: 1, + sessionId: 'session-1', + runId: 'run-1', + }; + const identity = { + id: 'event-1', + turnId: 'turn-1', + ts: 1, + toolUseId: 'tool-1', + }; + for (const event of [ + { + ...identity, + type: 'tool_start', + toolName: 'read', + displayName: 'Read file', + }, + { + ...identity, + type: 'tool_output_delta', + seq: 0, + stream: 'stdout', + chunk: 'visible output', + redacted: false, + createdAt: 2, + }, + { ...identity, type: 'tool_progress', chunk: 'working' }, + { ...identity, type: 'tool_result', status: 'completed', durationMs: 3 }, + ]) { + assert.doesNotThrow(() => decodeHostFrame({ ...envelope, event })); + } + for (const event of [ + { + ...identity, + type: 'tool_start', + toolName: 'read', + args: { path: '/private' }, + }, + { + ...identity, + type: 'tool_result', + status: 'errored', + result: { secret: true }, + }, + { + ...identity, + type: 'tool_result', + status: 'errored', + error: 'raw provider error', + }, + ]) { + assert.throws(() => decodeHostFrame({ ...envelope, event }), isInvalidFrame); + } + assert.throws( + () => + decodeHostFrame({ + kind: 'subscription.session_delta', + hostEpoch: 'epoch-1', + subscriptionId: 'subscription-1', + sequence: 1, + sessionId: 'session-1', + delta: { + kind: 'thinking', + turnId: 'turn-1', + runId: 'run-1', + messageId: 'message-1', + text: 'private reasoning', + signature: 'provider-signature', + }, + }), + isInvalidFrame, + ); + }); + + test('enforces UTF-8 snapshot, live field, and whole-frame byte bounds', () => { + const snapshot = continuitySnapshot('epoch-1'); + assert.ok(Buffer.byteLength(JSON.stringify(snapshot)) < SESSION_CONTINUITY_SNAPSHOT_MAX_BYTES); + assert.throws( + () => + decodeSessionContinuitySnapshot({ + ...snapshot, + padding: 'x'.repeat(SESSION_CONTINUITY_SNAPSHOT_MAX_BYTES), + }), + isInvalidFrame, + ); + const frame = { + kind: 'subscription.session_delta' as const, + hostEpoch: 'epoch-1', + subscriptionId: 'subscription-1', + sequence: 1, + sessionId: 'session-1', + delta: { + kind: 'text' as const, + turnId: 'turn-1', + runId: 'run-1', + messageId: 'message-1', + text: '界'.repeat(Math.floor(SESSION_LIVE_DELTA_MAX_BYTES / 3) + 1), + }, + }; + assert.throws(() => decodeHostFrame(frame), isInvalidFrame); + const eventEnvelope = { + kind: 'subscription.session_event', + hostEpoch: 'epoch-1', + subscriptionId: 'subscription-1', + sequence: 1, + sessionId: 'session-1', + runId: 'run-1', + }; + const eventIdentity = { id: 'event-1', turnId: 'turn-1', ts: 1, toolUseId: 'tool-1' }; + assert.throws( + () => + decodeHostFrame({ + ...eventEnvelope, + event: { + ...eventIdentity, + type: 'tool_start', + toolName: '界'.repeat(Math.floor(SESSION_TOOL_NAME_MAX_BYTES / 3) + 1), + }, + }), + isInvalidFrame, + ); + assert.throws( + () => + decodeHostFrame({ + ...eventEnvelope, + event: { + ...eventIdentity, + type: 'tool_progress', + chunk: '界'.repeat(Math.floor(SESSION_LIVE_DELTA_MAX_BYTES / 3) + 1), + }, + }), + isInvalidFrame, + ); + assert.throws( + () => + decodeHostFrame({ + ...frame, + privatePadding: 'x'.repeat(RUNTIME_HOST_MAX_FRAME_BYTES), + }), + isInvalidFrame, + ); + }); + test('declares exactly the ten Runtime Policy operations in the current framework', () => { const queries = [ 'runtime.policy.query', @@ -728,3 +952,29 @@ function attachmentRef( ) { return { kind: 'code' as const, name: 'a.ts', mimeType: 'text/typescript', bytes: 10, ref }; } + +function continuitySnapshot(hostEpoch: string) { + return { + schemaVersion: 1 as const, + session: { + sessionId: 'session-1', + status: 'running' as const, + createdAt: 1, + lastUsedAt: 2, + isArchived: false, + }, + projectionRevision: 1, + rootTurn: { + sessionId: 'session-1', + turnId: 'turn-1', + runId: 'run-1', + status: 'running' as const, + }, + queue: { + hostEpoch, + queueRevision: 1, + steering: [], + followup: [], + }, + }; +} diff --git a/packages/runtime-host/src/__tests__/root-turn-coordinator.test.ts b/packages/runtime-host/src/__tests__/root-turn-coordinator.test.ts index 4003606d32..6a759c0143 100644 --- a/packages/runtime-host/src/__tests__/root-turn-coordinator.test.ts +++ b/packages/runtime-host/src/__tests__/root-turn-coordinator.test.ts @@ -22,11 +22,13 @@ import { type RootTurnAdmissionStore, } from '@maka/storage/execution-stores'; import { resolveStorageRoot, tryAcquireInteractiveRootOwner } from '@maka/storage/root-authority'; +import { CanonicalSessionProjectionReader } from '../server/canonical-session-projection.js'; import type { RuntimeHostResidency } from '../server/host-kernel.js'; import { type HostMessageRootPort, HostMessageCoordinator } from '../server/message-coordinator.js'; import { RootAdmissionOwner } from '../server/root-admission-owner.js'; import { RootTurnCoordinator } from '../server/root-turn-coordinator.js'; import { SessionAdmissionGate } from '../server/session-admission-gate.js'; +import { SessionContinuityCoordinator } from '../server/session-continuity-coordinator.js'; const HOLD_EXTERNAL_PROMPT = 'hold external root before follow-up'; @@ -56,11 +58,14 @@ test('hosted linked child roots share admission, message, terminal, and stop aut await rootAdmissionOwner.recoverSession(parent.id); const acquireResidency = (): RuntimeHostResidency => ({ release() {} }); let coordinator: RootTurnCoordinator | undefined; + let continuity: SessionContinuityCoordinator | undefined; + let drainRequested = false; const rootPort: HostMessageRootPort = { readSessionHeader: (sessionId) => requireCoordinator(coordinator).readSessionHeader(sessionId), readRootState: (sessionId) => requireCoordinator(coordinator).readRootState(sessionId), - startFromMessage: (input) => requireCoordinator(coordinator).startFromMessage(input), + startFromMessage: (input, admission) => + requireCoordinator(coordinator).startFromMessage(input, admission), claimStop: (input, commitQueueFence) => requireCoordinator(coordinator).claimStop(input, commitQueueFence), }; @@ -78,7 +83,25 @@ test('hosted linked child roots share admission, message, terminal, and stop aut receipts: stores.messageReceiptStore, sessionAdmission, acquireResidency, + requestDrain: () => { + drainRequested = true; + }, + onProjectionChanged: (sessionId) => + requireContinuity(continuity).enqueueCanonicalRefresh(sessionId), }); + const canonicalProjection = new CanonicalSessionProjectionReader({ + stores, + rootAdmissions: rootAdmissionOwner, + messages, + }); + continuity = new SessionContinuityCoordinator( + hostEpoch, + (sessionId) => canonicalProjection.read(sessionId), + sessionAdmission, + () => { + drainRequested = true; + }, + ); const authority: RuntimeHostedRootAuthority = { bindRun: (identity) => messages.bindRun(identity), executeRoot: (input) => requireCoordinator(coordinator).executeRoot(input), @@ -104,13 +127,13 @@ test('hosted linked child roots share admission, message, terminal, and stop aut now: Date.now, messageAuthority: authority, }); - let drainRequested = false; coordinator = new RootTurnCoordinator( manager, stores, sessionAdmission, rootAdmissionOwner, messages, + continuity, acquireResidency, () => { drainRequested = true; @@ -416,6 +439,7 @@ test('hosted linked child roots share admission, message, terminal, and stop aut }); await coordinator.close(); await messages.close(); + continuity.close(); } finally { await owner.close(); await rm(base, { recursive: true, force: true }); @@ -670,11 +694,13 @@ test('shutdown re-scans a successor created by an in-flight terminal handoff', { }; const sessionAdmission = new SessionAdmissionGate(); let coordinator: RootTurnCoordinator | undefined; + let continuity: SessionContinuityCoordinator | undefined; const rootPort: HostMessageRootPort = { readSessionHeader: (sessionId) => requireCoordinator(coordinator).readSessionHeader(sessionId), readRootState: (sessionId) => requireCoordinator(coordinator).readRootState(sessionId), - startFromMessage: (input) => requireCoordinator(coordinator).startFromMessage(input), + startFromMessage: (input, admission) => + requireCoordinator(coordinator).startFromMessage(input, admission), claimStop: (input, commitQueueFence) => requireCoordinator(coordinator).claimStop(input, commitQueueFence), }; @@ -692,7 +718,19 @@ test('shutdown re-scans a successor created by an in-flight terminal handoff', { receipts: stores.messageReceiptStore, sessionAdmission, acquireResidency, + onProjectionChanged: (sessionId) => + requireContinuity(continuity).enqueueCanonicalRefresh(sessionId), + }); + const canonicalProjection = new CanonicalSessionProjectionReader({ + stores, + rootAdmissions: rootAdmissionOwner, + messages, }); + continuity = new SessionContinuityCoordinator( + hostEpoch, + (sessionId) => canonicalProjection.read(sessionId), + sessionAdmission, + ); const backends = new BackendRegistry(); backends.register('fake', (context) => new FakeBackend(context)); const manager = new SessionManager({ @@ -711,6 +749,7 @@ test('shutdown re-scans a successor created by an in-flight terminal handoff', { sessionAdmission, rootAdmissionOwner, messages, + continuity, acquireResidency, () => { drainRequested = true; @@ -746,6 +785,7 @@ test('shutdown re-scans a successor created by an in-flight terminal handoff', { releaseFollowupAdmission.resolve(); await closing; await messages.close(); + continuity.close(); assert.deepEqual(coordinator.readRootState(session.id), { kind: 'idle' }); assert.equal(liveResidencies, 0); @@ -807,10 +847,12 @@ async function createFailureFixture(options: { }; let drainRequested = false; let coordinator: RootTurnCoordinator | undefined; + let continuity: SessionContinuityCoordinator | undefined; const rootPort: HostMessageRootPort = { readSessionHeader: (sessionId) => requireCoordinator(coordinator).readSessionHeader(sessionId), readRootState: (sessionId) => requireCoordinator(coordinator).readRootState(sessionId), - startFromMessage: (input) => requireCoordinator(coordinator).startFromMessage(input), + startFromMessage: (input, admission) => + requireCoordinator(coordinator).startFromMessage(input, admission), claimStop: (input, commitQueueFence) => requireCoordinator(coordinator).claimStop(input, commitQueueFence), }; @@ -832,7 +874,20 @@ async function createFailureFixture(options: { sessionAdmission, acquireResidency, requestDrain, + onProjectionChanged: (sessionId) => + requireContinuity(continuity).enqueueCanonicalRefresh(sessionId), + }); + const canonicalProjection = new CanonicalSessionProjectionReader({ + stores, + rootAdmissions: rootAdmissionOwner, + messages, }); + continuity = new SessionContinuityCoordinator( + hostEpoch, + (sessionId) => canonicalProjection.read(sessionId), + sessionAdmission, + requestDrain, + ); const backends = new BackendRegistry(); options.registerBackend(backends); const manager = new SessionManager({ @@ -850,6 +905,7 @@ async function createFailureFixture(options: { sessionAdmission, rootAdmissionOwner, messages, + continuity, acquireResidency, requestDrain, ); @@ -864,6 +920,7 @@ async function createFailureFixture(options: { liveResidencies: () => liveResidencies, drainRequested: () => drainRequested, dispose: async () => { + continuity.close(); await owner.close(); await rm(base, { recursive: true, force: true }); }, @@ -875,6 +932,13 @@ function requireCoordinator(coordinator: RootTurnCoordinator | undefined): RootT return coordinator; } +function requireContinuity( + continuity: SessionContinuityCoordinator | undefined, +): SessionContinuityCoordinator { + if (!continuity) throw new Error('Continuity coordinator is not bound'); + return continuity; +} + function operationContext(hostEpoch: string, acquireResidency: () => RuntimeHostResidency) { return { hostEpoch, diff --git a/packages/runtime-host/src/__tests__/session-admission-gate.test.ts b/packages/runtime-host/src/__tests__/session-admission-gate.test.ts index bd8d9efd69..e6e4c4c0e8 100644 --- a/packages/runtime-host/src/__tests__/session-admission-gate.test.ts +++ b/packages/runtime-host/src/__tests__/session-admission-gate.test.ts @@ -41,6 +41,64 @@ test('does not serialize operations for different Sessions', async () => { await first; }); +test('keeps the admission open until admitted child work settles', async () => { + const gate = new SessionAdmissionGate(); + const childEntered = deferred(); + const releaseChild = deferred(); + let outerSettled = false; + + const outer = gate + .run('session', (lease) => { + void gate.runAdmitted('session', lease, async () => { + childEntered.resolve(); + await releaseChild.promise; + }); + return 'accepted'; + }) + .then((value) => { + outerSettled = true; + return value; + }); + + await childEntered.promise; + await Promise.resolve(); + assert.equal(outerSettled, false); + releaseChild.resolve(); + assert.equal(await outer, 'accepted'); +}); + +test('queues detached publication after the active admission', async () => { + const gate = new SessionAdmissionGate(); + const release = deferred(); + const order: string[] = []; + let detached!: Promise; + + const active = gate.run('session', async () => { + order.push('active:start'); + detached = gate.enqueueDetached('session', () => { + order.push('detached'); + }); + await release.promise; + order.push('active:end'); + }); + + await Promise.resolve(); + assert.deepEqual(order, ['active:start']); + release.resolve(); + await Promise.all([active, detached]); + assert.deepEqual(order, ['active:start', 'active:end', 'detached']); +}); + +test('rejects accidental admission re-entry instead of deadlocking', async () => { + const gate = new SessionAdmissionGate(); + await gate.run('session', async () => { + await assert.rejects( + gate.run('session', () => undefined), + /reuse its lease/, + ); + }); +}); + function deferred(): { promise: Promise; resolve(): void } { let resolve!: () => void; const promise = new Promise((settle) => { diff --git a/packages/runtime-host/src/__tests__/session-continuity-coordinator.test.ts b/packages/runtime-host/src/__tests__/session-continuity-coordinator.test.ts new file mode 100644 index 0000000000..9e40a577ae --- /dev/null +++ b/packages/runtime-host/src/__tests__/session-continuity-coordinator.test.ts @@ -0,0 +1,289 @@ +import assert from 'node:assert/strict'; +import { setImmediate as delayImmediate } from 'node:timers/promises'; +import test from 'node:test'; +import type { SubscriptionFrame } from '../protocol/index.js'; +import type { ConnectionContext } from '../server/operation-dispatcher.js'; +import { + type CanonicalSessionProjection, + SessionContinuityCoordinator, +} from '../server/session-continuity-coordinator.js'; +import { SessionAdmissionGate } from '../server/session-admission-gate.js'; +import type { SessionContinuityFrameSink } from '../server/session-continuity-service.js'; + +const HOST_EPOCH = 'host-epoch'; +const SESSION_ID = 'session-1'; + +test('open is an inactive publication barrier and live sequence starts at nextSequence', async () => { + const read = deferred(); + const coordinator = new SessionContinuityCoordinator( + HOST_EPOCH, + () => read.promise, + new SessionAdmissionGate(), + ); + const sink = new RecordingSink(); + const connection = coordinator.attachConnection('connection-1', sink); + + const opening = coordinator.handlers['subscription.open']( + { sessionId: SESSION_ID }, + connectionContext('connection-1'), + ); + await delayImmediate(); + const publishing = coordinator.acceptRuntimeEvent(SESSION_ID, 'run-1', textEvent(1)); + read.resolve(canonical()); + + const outcome = await opening; + assert.equal(outcome.ok, true); + if (!outcome.ok) return; + await publishing; + assert.equal(outcome.result.nextSequence, 1); + assert.equal(Object.isFrozen(outcome.result.snapshot), true); + assert.equal(sink.frames.length, 0); + + connection.activate(outcome.result.subscriptionId); + await delayImmediate(); + assert.deepEqual( + sink.frames.map((frame) => frame.sequence), + [1], + ); + assert.equal(sink.frames[0]?.kind, 'subscription.session_delta'); + + connection.abort(outcome.result.subscriptionId); + await coordinator.acceptRuntimeEvent(SESSION_ID, 'run-1', textEvent(2)); + assert.equal(sink.frames.length, 1); + coordinator.close(); +}); + +test('terminal fence suppresses ordinary refresh until the exact terminal cut publishes', async () => { + let projection = canonical({ + rootTurn: { sessionId: SESSION_ID, turnId: 'turn-1', runId: 'run-1', status: 'running' }, + }); + const coordinator = new SessionContinuityCoordinator( + HOST_EPOCH, + async () => projection, + new SessionAdmissionGate(), + ); + const sink = new RecordingSink(); + const connection = coordinator.attachConnection('connection-1', sink); + const opened = await open(coordinator, 'connection-1'); + connection.activate(opened.subscriptionId); + + await coordinator.holdTerminalPublication(SESSION_ID, 'turn-1', 'run-1'); + projection = canonical({ + rootTurn: { + sessionId: SESSION_ID, + turnId: 'turn-1', + runId: 'run-1', + status: 'completed', + terminalEventId: 'event-terminal', + }, + }); + await coordinator.refreshCanonical(SESSION_ID); + assert.equal(sink.frames.length, 0); + + await coordinator.publishTerminalProjection(SESSION_ID, 'turn-1', 'run-1'); + await delayImmediate(); + assert.equal(sink.frames.length, 1); + const frame = sink.frames[0]; + assert.equal(frame?.kind, 'subscription.session_projection'); + if (frame?.kind === 'subscription.session_projection') { + assert.equal(frame.sequence, 1); + assert.equal(frame.snapshot.projectionRevision, 2); + assert.equal(frame.snapshot.rootTurn?.status, 'completed'); + } + coordinator.close(); +}); + +test('detached canonical refreshes coalesce before Store I/O', async () => { + let projection = canonical(); + let reads = 0; + const refreshRead = deferred(); + const refreshEntered = deferred(); + const coordinator = new SessionContinuityCoordinator( + HOST_EPOCH, + async () => { + reads += 1; + if (reads === 2) { + refreshEntered.resolve(); + await refreshRead.promise; + } + return projection; + }, + new SessionAdmissionGate(), + ); + const sink = new RecordingSink(); + const connection = coordinator.attachConnection('connection-1', sink); + const opened = await open(coordinator, 'connection-1'); + connection.activate(opened.subscriptionId); + + projection = canonical({ lastUsedAt: 2 }); + coordinator.enqueueCanonicalRefresh(SESSION_ID); + coordinator.enqueueCanonicalRefresh(SESSION_ID); + await refreshEntered.promise; + assert.equal(reads, 2); + refreshRead.resolve(); + await waitFor(() => sink.frames.length === 1); + assert.equal(reads, 2); + coordinator.close(); +}); + +test('reports a detached canonical publication failure to the Host lifecycle', async () => { + let reads = 0; + const observed = deferred(); + const coordinator = new SessionContinuityCoordinator( + HOST_EPOCH, + async () => { + reads += 1; + if (reads === 1) return canonical(); + throw new Error('canonical Store read failed'); + }, + new SessionAdmissionGate(), + (error) => observed.resolve(error), + ); + const connection = coordinator.attachConnection('connection-1', new RecordingSink()); + const opened = await open(coordinator, 'connection-1'); + connection.activate(opened.subscriptionId); + + coordinator.enqueueCanonicalRefresh(SESSION_ID); + const failure = await observed.promise; + assert.match(String(failure), /canonical Store read failed/); + coordinator.close(); +}); + +test('rejects a live event that is not owned by the canonical root', async () => { + const coordinator = new SessionContinuityCoordinator( + HOST_EPOCH, + async () => canonical(), + new SessionAdmissionGate(), + ); + const connection = coordinator.attachConnection('connection-1', new RecordingSink()); + const opened = await open(coordinator, 'connection-1'); + connection.activate(opened.subscriptionId); + + await assert.rejects( + coordinator.acceptRuntimeEvent(SESSION_ID, 'different-run', textEvent(1)), + /canonical active root Turn/, + ); + coordinator.close(); +}); + +test('slow subscriber receives a terminal eviction without delaying another subscriber', async () => { + const coordinator = new SessionContinuityCoordinator( + HOST_EPOCH, + async () => canonical(), + new SessionAdmissionGate(), + ); + const slowSink = new RecordingSink(); + const fastSink = new RecordingSink(); + const slowConnection = coordinator.attachConnection('connection-slow', slowSink); + const fastConnection = coordinator.attachConnection('connection-fast', fastSink); + const slow = await open(coordinator, 'connection-slow'); + const fast = await open(coordinator, 'connection-fast'); + fastConnection.activate(fast.subscriptionId); + + for (let index = 1; index <= 32; index += 1) { + await coordinator.acceptRuntimeEvent(SESSION_ID, 'run-1', textEvent(index)); + } + slowConnection.activate(slow.subscriptionId); + await waitFor(() => slowSink.frames.length === 1 && fastSink.frames.length === 32); + + assert.equal(slowSink.closed, 0); + assert.deepEqual(slowSink.frames[0], { + kind: 'subscription.closed', + hostEpoch: HOST_EPOCH, + subscriptionId: slow.subscriptionId, + sequence: 1, + reason: 'slow_consumer', + }); + assert.equal(fastSink.closed, 0); + assert.deepEqual( + fastSink.frames.map((frame) => frame.sequence), + Array.from({ length: 32 }, (_, index) => index + 1), + ); + coordinator.close(); +}); + +class RecordingSink implements SessionContinuityFrameSink { + readonly frames: SubscriptionFrame[] = []; + closed = 0; + + async send(frame: SubscriptionFrame): Promise { + this.frames.push(frame); + } + + close(): void { + this.closed += 1; + } +} + +async function open(coordinator: SessionContinuityCoordinator, connectionId: string) { + const outcome = await coordinator.handlers['subscription.open']( + { sessionId: SESSION_ID }, + connectionContext(connectionId), + ); + if (!outcome.ok) throw new Error(outcome.error.message); + assert.equal(outcome.ok, true); + return outcome.result; +} + +function connectionContext(connectionId: string): ConnectionContext { + return { + hostEpoch: HOST_EPOCH, + connectionId, + surface: 'tui', + principal: 'local_os_user', + acquireResidency: () => ({ release() {} }), + }; +} + +function canonical( + overrides: { lastUsedAt?: number; rootTurn?: CanonicalSessionProjection['rootTurn'] } = {}, +): CanonicalSessionProjection { + return { + session: { + sessionId: SESSION_ID, + status: 'active', + createdAt: 1, + lastUsedAt: overrides.lastUsedAt ?? 1, + isArchived: false, + }, + rootTurn: + overrides.rootTurn === undefined + ? { sessionId: SESSION_ID, turnId: 'turn-1', runId: 'run-1', status: 'running' } + : overrides.rootTurn, + queue: { + hostEpoch: HOST_EPOCH, + queueRevision: 0, + steering: [], + followup: [], + }, + }; +} + +function textEvent(index: number) { + return { + type: 'text_delta' as const, + id: `event-${index}`, + turnId: 'turn-1', + ts: index, + messageId: 'message-1', + text: `chunk-${index}`, + }; +} + +function deferred() { + let resolve!: (value: T | PromiseLike) => void; + let reject!: (reason?: unknown) => void; + const promise = new Promise((resolvePromise, rejectPromise) => { + resolve = resolvePromise; + reject = rejectPromise; + }); + return { promise, resolve, reject }; +} + +async function waitFor(predicate: () => boolean): Promise { + for (let attempt = 0; attempt < 100; attempt += 1) { + if (predicate()) return; + await delayImmediate(); + } + throw new Error('Timed out waiting for continuity state'); +} diff --git a/packages/runtime-host/src/__tests__/session-subscription-client.test.ts b/packages/runtime-host/src/__tests__/session-subscription-client.test.ts new file mode 100644 index 0000000000..bc611a86ec --- /dev/null +++ b/packages/runtime-host/src/__tests__/session-subscription-client.test.ts @@ -0,0 +1,381 @@ +import assert from 'node:assert/strict'; +import { randomUUID } from 'node:crypto'; +import { mkdtemp, rm } from 'node:fs/promises'; +import { createServer, type Server } from 'node:net'; +import { tmpdir } from 'node:os'; +import { join } from 'node:path'; +import { test } from 'node:test'; +import { + prepareStorageRootControlDirectory, + resolveStorageRoot, +} from '@maka/storage/root-authority'; +import { + connectRuntimeHost, + RuntimeHostSubscriptionError, + type RuntimeHostConnection, +} from '../client/index.js'; +import { prepareRuntimeHostEndpoint } from '../control/endpoint.js'; +import { removeHostRegistration, writeHostRegistration } from '../control/registration.js'; +import { + decodeClientFrame, + encodeProtocolFrame, + RUNTIME_HOST_PROTOCOL_VERSION, + RUNTIME_HOST_REGISTRATION_SCHEMA_VERSION, + type RequestFrame, + type SubscriptionFrame, +} from '../protocol/index.js'; +import { FramedTransport } from '../transport/framed-transport.js'; + +const PROTOCOL = { + min: RUNTIME_HOST_PROTOCOL_VERSION, + max: RUNTIME_HOST_PROTOCOL_VERSION, +} as const; + +test('registers a subscription before receiving a coalesced first frame', async () => { + await withProtocolPeer( + async (transport, hostEpoch) => { + const request = await acceptConnectionAndReadOpen(transport, hostEpoch); + const opened = openResult(hostEpoch, 'subscription-ordered'); + await transport.writeEncoded( + Buffer.concat([ + encodeProtocolFrame({ + requestId: request.requestId, + operation: 'subscription.open', + ok: true, + result: opened, + }), + encodeProtocolFrame(deltaFrame(hostEpoch, opened.subscriptionId, 1)), + ]), + ); + await answerClose(transport, opened.subscriptionId); + }, + async (connection) => { + const subscription = await connection.openSessionSubscription({ + sessionId: 'session-1', + }); + assert.deepEqual(await subscription[Symbol.asyncIterator]().next(), { + done: false, + value: deltaFrame(connection.hostEpoch, subscription.subscriptionId, 1), + }); + await subscription.close(); + }, + ); +}); + +test('isolates a sequence gap and continues requests on the same connection', async () => { + await withProtocolPeer( + async (transport, hostEpoch) => { + const request = await acceptConnectionAndReadOpen(transport, hostEpoch); + const opened = openResult(hostEpoch, 'subscription-gap'); + await transport.writeEncoded( + Buffer.concat([ + encodeProtocolFrame({ + requestId: request.requestId, + operation: 'subscription.open', + ok: true, + result: opened, + }), + encodeProtocolFrame(deltaFrame(hostEpoch, opened.subscriptionId, 2)), + ]), + ); + await answerClose(transport, opened.subscriptionId); + await answerStatus(transport, hostEpoch); + }, + async (connection) => { + const subscription = await connection.openSessionSubscription({ + sessionId: 'session-1', + }); + await assert.rejects( + () => subscription[Symbol.asyncIterator]().next(), + hasSubscriptionReason('sequence_gap'), + ); + assert.equal((await connection.status()).hostEpoch, connection.hostEpoch); + }, + ); +}); + +test('rejects epoch and Session correlation changes per subscription', async () => { + for (const changed of ['epoch', 'session'] as const) { + await withProtocolPeer( + async (transport, hostEpoch) => { + const request = await acceptConnectionAndReadOpen(transport, hostEpoch); + const opened = openResult(hostEpoch, `subscription-${changed}`); + await transport.write({ + requestId: request.requestId, + operation: 'subscription.open', + ok: true, + result: opened, + }); + await transport.write({ + ...deltaFrame( + changed === 'epoch' ? 'different-epoch' : hostEpoch, + opened.subscriptionId, + 1, + ), + ...(changed === 'session' ? { sessionId: 'session-2' } : {}), + }); + await answerClose(transport, opened.subscriptionId); + await answerStatus(transport, hostEpoch); + }, + async (connection) => { + const subscription = await connection.openSessionSubscription({ + sessionId: 'session-1', + }); + await assert.rejects( + () => subscription[Symbol.asyncIterator]().next(), + hasSubscriptionReason(changed === 'epoch' ? 'host_epoch_changed' : 'correlation_changed'), + ); + assert.equal((await connection.status()).hostEpoch, connection.hostEpoch); + }, + ); + } +}); + +test('evicts a locally slow iterator and keeps the connection usable', async () => { + const closeObserved = deferred(); + await withProtocolPeer( + async (transport, hostEpoch) => { + const request = await acceptConnectionAndReadOpen(transport, hostEpoch); + const opened = openResult(hostEpoch, 'subscription-slow'); + const frames = [ + encodeProtocolFrame({ + requestId: request.requestId, + operation: 'subscription.open', + ok: true, + result: opened, + }), + ]; + for (let sequence = 1; sequence <= 33; sequence += 1) { + frames.push(encodeProtocolFrame(deltaFrame(hostEpoch, opened.subscriptionId, sequence))); + } + await transport.writeEncoded(Buffer.concat(frames)); + await answerClose(transport, opened.subscriptionId, closeObserved.resolve); + await answerStatus(transport, hostEpoch); + }, + async (connection) => { + const subscription = await connection.openSessionSubscription({ + sessionId: 'session-1', + }); + await closeObserved.promise; + await assert.rejects( + () => subscription[Symbol.asyncIterator]().next(), + hasSubscriptionReason('slow_consumer'), + ); + assert.equal((await connection.status()).hostEpoch, connection.hostEpoch); + }, + ); +}); + +test('ends every active subscription with connection_closed on EOF', async () => { + await withProtocolPeer( + async (transport, hostEpoch) => { + const request = await acceptConnectionAndReadOpen(transport, hostEpoch); + await transport.write({ + requestId: request.requestId, + operation: 'subscription.open', + ok: true, + result: openResult(hostEpoch, 'subscription-eof'), + }); + transport.destroyAfterFlush(); + }, + async (connection) => { + const subscription = await connection.openSessionSubscription({ + sessionId: 'session-1', + }); + await assert.rejects( + () => subscription[Symbol.asyncIterator]().next(), + hasSubscriptionReason('connection_closed'), + ); + }, + ); +}); + +async function withProtocolPeer( + serve: (transport: FramedTransport, hostEpoch: string) => Promise, + run: (connection: RuntimeHostConnection) => Promise, +): Promise { + const base = await mkdtemp(join(tmpdir(), 'maka-runtime-host-subscription-')); + const capability = await resolveStorageRoot({ + path: join(base, 'root'), + kind: 'interactive', + }); + const { controlDirectory } = await prepareStorageRootControlDirectory(capability); + const hostEpoch = randomUUID(); + const endpoint = await prepareRuntimeHostEndpoint({ + rootId: capability.rootId, + hostEpoch, + }); + const serverTask = deferred(); + const server = createServer((socket) => { + void serve(new FramedTransport(socket), hostEpoch).then(serverTask.resolve, serverTask.reject); + }); + try { + await listen(server, endpoint.path); + await endpoint.prepareAfterListen(); + await writeHostRegistration(controlDirectory, { + kind: 'maka-runtime-host', + schemaVersion: RUNTIME_HOST_REGISTRATION_SCHEMA_VERSION, + rootId: capability.rootId, + hostEpoch, + endpoint: endpoint.path, + protocolMin: RUNTIME_HOST_PROTOCOL_VERSION, + protocolMax: RUNTIME_HOST_PROTOCOL_VERSION, + state: 'ready', + pid: process.pid, + createdAt: new Date().toISOString(), + }); + const connected = await connectRuntimeHost({ + rootPath: join(base, 'root'), + surface: 'tui', + protocol: PROTOCOL, + }); + assert.equal(connected.kind, 'connected'); + if (connected.kind !== 'connected') return; + try { + await run(connected.connection); + } finally { + await connected.connection.close(); + } + await serverTask.promise; + } finally { + await closeServer(server); + await removeHostRegistration(controlDirectory, hostEpoch).catch(() => undefined); + await endpoint.cleanup().catch(() => undefined); + await rm(base, { recursive: true, force: true }); + } +} + +async function acceptConnectionAndReadOpen( + transport: FramedTransport, + hostEpoch: string, +): Promise> { + const hello = decodeClientFrame(await transport.read(1_000)); + assert.ok('kind' in hello && hello.kind === 'hello'); + await transport.write({ + kind: 'accepted', + hostEpoch, + connectionId: 'connection-1', + selectedProtocol: RUNTIME_HOST_PROTOCOL_VERSION, + state: 'ready', + }); + const request = decodeClientFrame(await transport.read(1_000)); + assert.ok(!('kind' in request)); + assert.equal(request.operation, 'subscription.open'); + return request as Extract; +} + +async function answerClose( + transport: FramedTransport, + subscriptionId: string, + onObserved?: () => void, +): Promise { + const request = decodeClientFrame(await transport.read(1_000)); + assert.ok(!('kind' in request)); + assert.equal(request.operation, 'subscription.close'); + assert.deepEqual(request.input, { subscriptionId }); + onObserved?.(); + await transport.write({ + requestId: request.requestId, + operation: 'subscription.close', + ok: true, + result: { subscriptionId }, + }); +} + +async function answerStatus(transport: FramedTransport, hostEpoch: string): Promise { + const request = decodeClientFrame(await transport.read(1_000)); + assert.ok(!('kind' in request)); + assert.equal(request.operation, 'host.status'); + await transport.write({ + requestId: request.requestId, + operation: 'host.status', + ok: true, + result: { + hostEpoch, + state: 'ready', + connections: 1, + activeOperations: 1, + activeResidencies: 0, + }, + }); +} + +function openResult(hostEpoch: string, subscriptionId: string) { + return { + hostEpoch, + subscriptionId, + nextSequence: 1, + snapshot: { + schemaVersion: 1 as const, + session: { + sessionId: 'session-1', + status: 'running' as const, + createdAt: 1, + lastUsedAt: 2, + isArchived: false, + }, + projectionRevision: 1, + rootTurn: { + sessionId: 'session-1', + turnId: 'turn-1', + runId: 'run-1', + status: 'running' as const, + }, + queue: { hostEpoch, queueRevision: 1, steering: [], followup: [] }, + }, + }; +} + +function deltaFrame( + hostEpoch: string, + subscriptionId: string, + sequence: number, +): SubscriptionFrame { + return { + kind: 'subscription.session_delta', + hostEpoch, + subscriptionId, + sequence, + sessionId: 'session-1', + delta: { + kind: 'text', + turnId: 'turn-1', + runId: 'run-1', + messageId: 'message-1', + text: `chunk-${sequence}`, + }, + }; +} + +function hasSubscriptionReason(reason: RuntimeHostSubscriptionError['reason']) { + return (error: unknown) => + error instanceof RuntimeHostSubscriptionError && error.reason === reason; +} + +function listen(server: Server, path: string): Promise { + return new Promise((resolve, reject) => { + server.once('error', reject); + server.listen(path, resolve); + }); +} + +function closeServer(server: Server): Promise { + if (!server.listening) return Promise.resolve(); + return new Promise((resolve, reject) => { + server.close((error) => (error ? reject(error) : resolve())); + }); +} + +function deferred(): { + promise: Promise; + resolve(value: T | PromiseLike): void; + reject(error: unknown): void; +} { + let resolve!: (value: T | PromiseLike) => void; + let reject!: (error: unknown) => void; + const promise = new Promise((resolvePromise, rejectPromise) => { + resolve = resolvePromise; + reject = rejectPromise; + }); + return { promise, resolve, reject }; +} diff --git a/packages/runtime-host/src/client/connection.ts b/packages/runtime-host/src/client/connection.ts index 8e4c6832fc..748829e8cb 100644 --- a/packages/runtime-host/src/client/connection.ts +++ b/packages/runtime-host/src/client/connection.ts @@ -20,6 +20,8 @@ import { type ProtocolRange, type RequestFrame, type ResponseFrame, + type SubscriptionFrame, + type SubscriptionOpenInput, type TurnQueryInput, type TurnSnapshot, type TurnStartInput, @@ -28,6 +30,11 @@ import { validateProtocolRange, } from '../protocol/index.js'; import { FramedTransport, RuntimeHostTransportError } from '../transport/framed-transport.js'; +import { + ClientSessionSubscription, + RuntimeHostSubscriptionError, + type RuntimeHostSessionSubscription, +} from './session-subscription.js'; const DEFAULT_CONNECT_TIMEOUT_MS = 500; const DEFAULT_HANDSHAKE_TIMEOUT_MS = 2_000; @@ -50,10 +57,22 @@ export type RuntimeHostUnavailableReason = | 'epoch_mismatch'; export type ConnectRuntimeHostResult = - | { kind: 'connected'; connection: RuntimeHostConnection; registration: HostRegistration } - | { kind: 'incompatible'; handshake: HostIncompatible; registration: HostRegistration } + | { + kind: 'connected'; + connection: RuntimeHostConnection; + registration: HostRegistration; + } + | { + kind: 'incompatible'; + handshake: HostIncompatible; + registration: HostRegistration; + } | { kind: 'draining'; registration: HostRegistration } - | { kind: 'unavailable'; reason: RuntimeHostUnavailableReason; registration?: HostRegistration }; + | { + kind: 'unavailable'; + reason: RuntimeHostUnavailableReason; + registration?: HostRegistration; + }; type ConnectResolvedRuntimeHostResult = | ConnectRuntimeHostResult @@ -82,7 +101,7 @@ export interface RuntimeHostConnection { readonly connectionId: string; readonly selectedProtocol: number; readonly closed: Promise; - request( + request( operation: K, input: OperationInput, timeoutMs?: number, @@ -91,9 +110,18 @@ export interface RuntimeHostConnection { startTurn(input: TurnStartInput, timeoutMs?: number): Promise; queryTurn(input: TurnQueryInput, timeoutMs?: number): Promise; stopTurn(input: TurnStopInput, timeoutMs?: number): Promise; + openSessionSubscription( + input: SubscriptionOpenInput, + timeoutMs?: number, + ): Promise; close(): Promise; } +export type DirectRequestOperationKey = Exclude< + OperationKey, + 'subscription.open' | 'subscription.close' +>; + export class RuntimeHostOperationError extends Error { constructor( readonly operation: OperationKey, @@ -107,6 +135,7 @@ export class RuntimeHostOperationError extends Error { interface PendingRequest { operation: OperationKey; + accept(value: unknown): unknown; resolve(value: unknown): void; reject(error: Error): void; timer: NodeJS.Timeout; @@ -119,6 +148,8 @@ class RuntimeHostConnectionImpl implements RuntimeHostConnection { readonly closed: Promise; readonly #transport: FramedTransport; readonly #pendingRequests = new Map(); + readonly #subscriptions = new Map(); + readonly #retiredSubscriptionIds = new Set(); #terminalError: Error | undefined; constructor( @@ -137,26 +168,35 @@ class RuntimeHostConnectionImpl implements RuntimeHostConnection { void this.#readResponses(); } - request( + request( operation: K, input: OperationInput, timeoutMs = DEFAULT_HANDSHAKE_TIMEOUT_MS, ): Promise> { + return this.#requestOperation(operation, input, timeoutMs, (result) => result); + } + + #requestOperation( + operation: K, + input: OperationInput, + timeoutMs: number, + accept: (result: OperationOutput) => Result, + ): Promise { const boundedTimeoutMs = requireTimeout(timeoutMs, 'timeoutMs'); if (this.#terminalError) return Promise.reject(this.#terminalError); const requestId = randomUUID(); - const result = new Promise>((resolve, reject) => { + const result = new Promise((resolve, reject) => { const timer = setTimeout(() => { - this.#fail( - new RuntimeHostTransportError( - 'read_timeout', - `Timed out waiting for Runtime Host ${operation} response`, - ), + const error = new RuntimeHostTransportError( + 'read_timeout', + `Timed out waiting for Runtime Host ${operation} response`, ); + this.#fail(error); }, boundedTimeoutMs); this.#pendingRequests.set(requestId, { operation, - resolve: (value) => resolve(value as OperationOutput), + accept: (value) => accept(value as OperationOutput), + resolve: (value) => resolve(value as Result), reject, timer, }); @@ -188,6 +228,38 @@ class RuntimeHostConnectionImpl implements RuntimeHostConnection { return this.request('turn.stop', input, timeoutMs); } + openSessionSubscription( + input: SubscriptionOpenInput, + timeoutMs = DEFAULT_HANDSHAKE_TIMEOUT_MS, + ): Promise { + const expectedSessionId = input.sessionId; + return this.#requestOperation('subscription.open', input, timeoutMs, (result) => { + if (result.hostEpoch !== this.hostEpoch) { + throw new RuntimeHostSubscriptionError( + 'host_epoch_changed', + 'Session subscription opened for a different Host Epoch', + ); + } + if (result.snapshot.session.sessionId !== expectedSessionId) { + throw new RuntimeHostSubscriptionError( + 'correlation_changed', + 'Runtime Host opened a subscription for a different Session', + ); + } + if (this.#subscriptions.has(result.subscriptionId)) { + throw new RuntimeHostSubscriptionError( + 'correlation_changed', + 'Runtime Host returned a duplicate subscription identity', + ); + } + const subscription = new ClientSessionSubscription(result, () => + this.#closeSessionSubscription(result.subscriptionId), + ); + this.#subscriptions.set(result.subscriptionId, subscription); + return subscription; + }); + } + async close(): Promise { this.#transport.destroy(); await this.#transport.closed; @@ -197,8 +269,18 @@ class RuntimeHostConnectionImpl implements RuntimeHostConnection { try { while (true) { const frame = decodeHostFrame(await this.#transport.read(0)); - if ('kind' in frame) - throw new Error('Runtime Host returned a handshake frame after acceptance'); + if ('kind' in frame) { + switch (frame.kind) { + case 'subscription.session_projection': + case 'subscription.session_delta': + case 'subscription.session_event': + case 'subscription.closed': + this.#acceptSubscriptionFrame(frame); + continue; + default: + throw new Error('Runtime Host returned a handshake frame after acceptance'); + } + } this.#acceptResponse(frame); } } catch (error) { @@ -215,7 +297,13 @@ class RuntimeHostConnectionImpl implements RuntimeHostConnection { this.#pendingRequests.delete(frame.requestId); clearTimeout(pending.timer); if (frame.ok) { - pending.resolve(frame.result); + try { + pending.resolve(pending.accept(frame.result)); + } catch (error) { + const failure = asError(error); + pending.reject(failure); + this.#fail(failure); + } return; } pending.reject( @@ -223,6 +311,63 @@ class RuntimeHostConnectionImpl implements RuntimeHostConnection { ); } + #acceptSubscriptionFrame(frame: SubscriptionFrame): void { + const subscription = this.#subscriptions.get(frame.subscriptionId); + if (!subscription) { + if (this.#retiredSubscriptionIds.has(frame.subscriptionId)) return; + this.#fail(new Error('Runtime Host returned an unmatched subscription frame')); + return; + } + try { + subscription.accept(frame); + if (frame.kind === 'subscription.closed') { + this.#subscriptions.delete(frame.subscriptionId); + } + } catch (error) { + const failure = asError(error); + if (failure instanceof RuntimeHostSubscriptionError) { + this.#invalidateSubscription(subscription, failure); + return; + } + this.#fail(failure); + } + } + + async #closeSessionSubscription(subscriptionId: string): Promise { + const subscription = this.#subscriptions.get(subscriptionId); + if (!subscription) return; + await this.#requestOperation( + 'subscription.close', + { subscriptionId }, + DEFAULT_HANDSHAKE_TIMEOUT_MS, + (result) => { + if (result.subscriptionId !== subscriptionId) { + throw new Error('Runtime Host closed a different subscription'); + } + }, + ); + this.#subscriptions.delete(subscriptionId); + subscription.finish(); + } + + #invalidateSubscription( + subscription: ClientSessionSubscription, + error: RuntimeHostSubscriptionError, + ): void { + const { subscriptionId } = subscription; + if (this.#subscriptions.get(subscriptionId) !== subscription) return; + this.#subscriptions.delete(subscriptionId); + this.#retiredSubscriptionIds.add(subscriptionId); + subscription.fail(error); + if (this.#terminalError) return; + void this.#requestOperation( + 'subscription.close', + { subscriptionId }, + DEFAULT_HANDSHAKE_TIMEOUT_MS, + () => this.#retiredSubscriptionIds.delete(subscriptionId), + ).catch((failure: unknown) => this.#fail(asError(failure))); + } + #fail(error: Error): void { if (this.#terminalError) return; this.#terminalError = error; @@ -231,6 +376,15 @@ class RuntimeHostConnectionImpl implements RuntimeHostConnection { pending.reject(error); } this.#pendingRequests.clear(); + const subscriptionError = new RuntimeHostSubscriptionError( + 'connection_closed', + `Runtime Host connection closed: ${error.message}`, + ); + for (const subscription of this.#subscriptions.values()) { + subscription.fail(subscriptionError); + } + this.#subscriptions.clear(); + this.#retiredSubscriptionIds.clear(); this.#transport.destroy(); } } @@ -248,7 +402,10 @@ export async function connectRuntimeHost( 'handshakeTimeoutMs', ); const clientInstanceId = requireClientInstanceId(input.clientInstanceId ?? randomUUID()); - const capability = await resolveStorageRoot({ path: input.rootPath, kind: 'interactive' }); + const capability = await resolveStorageRoot({ + path: input.rootPath, + kind: 'interactive', + }); const { controlDirectory } = await prepareStorageRootControlDirectory(capability); const result = await connectResolvedRuntimeHost({ ...input, diff --git a/packages/runtime-host/src/client/index.ts b/packages/runtime-host/src/client/index.ts index 93b6e60f15..c6f6b42c0b 100644 --- a/packages/runtime-host/src/client/index.ts +++ b/packages/runtime-host/src/client/index.ts @@ -5,7 +5,13 @@ export { type ConnectRuntimeHostResult, type RuntimeHostConnection, type RuntimeHostUnavailableReason, + type DirectRequestOperationKey, } from './connection.js'; +export { + RuntimeHostSubscriptionError, + type RuntimeHostSessionSubscription, + type RuntimeHostSubscriptionFailureReason, +} from './session-subscription.js'; export { connectOrSpawnRuntimeHost, type ConnectOrSpawnRuntimeHostInput, diff --git a/packages/runtime-host/src/client/session-subscription.ts b/packages/runtime-host/src/client/session-subscription.ts new file mode 100644 index 0000000000..e063b6fa40 --- /dev/null +++ b/packages/runtime-host/src/client/session-subscription.ts @@ -0,0 +1,205 @@ +import { + encodeProtocolFrame, + type SessionContinuitySnapshot, + type SubscriptionFrame, + type SubscriptionOpenResult, +} from '../protocol/index.js'; + +const MAX_CLIENT_QUEUED_FRAMES = 32; +const MAX_CLIENT_QUEUED_BYTES = 256 * 1024; + +export type RuntimeHostSubscriptionFailureReason = + | 'sequence_gap' + | 'host_epoch_changed' + | 'correlation_changed' + | 'projection_revision_invalid' + | 'slow_consumer' + | 'connection_closed'; + +export class RuntimeHostSubscriptionError extends Error { + constructor( + readonly reason: RuntimeHostSubscriptionFailureReason, + message: string, + ) { + super(message); + this.name = 'RuntimeHostSubscriptionError'; + } +} + +export interface RuntimeHostSessionSubscription extends AsyncIterable { + readonly hostEpoch: string; + readonly subscriptionId: string; + readonly snapshot: SessionContinuitySnapshot; + close(): Promise; +} + +interface QueuedFrame { + frame: SubscriptionFrame; + encodedBytes: number; +} + +export class ClientSessionSubscription + implements RuntimeHostSessionSubscription, AsyncIterator +{ + readonly hostEpoch: string; + readonly subscriptionId: string; + readonly snapshot: SessionContinuitySnapshot; + readonly #requestClose: () => Promise; + readonly #expectedSessionId: string; + readonly #queue: QueuedFrame[] = []; + #queuedBytes = 0; + #expectedSequence: number; + #latestProjectionRevision: number; + #waiting: + | { + resolve(value: IteratorResult): void; + reject(error: Error): void; + } + | undefined; + #terminalError: Error | undefined; + #done = false; + #doneAfterQueue = false; + #closeTask: Promise | undefined; + + constructor(result: SubscriptionOpenResult, requestClose: () => Promise) { + this.hostEpoch = result.hostEpoch; + this.subscriptionId = result.subscriptionId; + this.snapshot = result.snapshot; + this.#expectedSessionId = result.snapshot.session.sessionId; + this.#expectedSequence = result.nextSequence; + this.#latestProjectionRevision = result.snapshot.projectionRevision; + this.#requestClose = requestClose; + } + + [Symbol.asyncIterator](): AsyncIterator { + return this; + } + + next(): Promise> { + const queued = this.#queue.shift(); + if (queued) { + this.#queuedBytes -= queued.encodedBytes; + if (this.#queue.length === 0 && this.#doneAfterQueue) this.#done = true; + return Promise.resolve({ done: false, value: queued.frame }); + } + if (this.#terminalError) return Promise.reject(this.#terminalError); + if (this.#done || this.#doneAfterQueue) { + this.#done = true; + return Promise.resolve({ done: true, value: undefined }); + } + if (this.#waiting) { + return Promise.reject(new Error('Session subscription already has a pending iterator read')); + } + return new Promise((resolve, reject) => { + this.#waiting = { resolve, reject }; + }); + } + + async return(): Promise> { + await this.close(); + return { done: true, value: undefined }; + } + + close(): Promise { + if (this.#done || this.#terminalError) return Promise.resolve(); + if (!this.#closeTask) this.#closeTask = this.#requestClose(); + return this.#closeTask; + } + + accept(frame: SubscriptionFrame): void { + if (this.#done || this.#terminalError) return; + if (this.#doneAfterQueue) { + throw new RuntimeHostSubscriptionError( + 'correlation_changed', + 'Session subscription received a frame after closure', + ); + } + if (frame.hostEpoch !== this.hostEpoch) { + throw new RuntimeHostSubscriptionError( + 'host_epoch_changed', + 'Session subscription Host Epoch changed', + ); + } + if (frame.subscriptionId !== this.subscriptionId) { + throw new RuntimeHostSubscriptionError( + 'correlation_changed', + 'Session subscription correlation changed', + ); + } + if (frame.sequence !== this.#expectedSequence) { + throw new RuntimeHostSubscriptionError( + 'sequence_gap', + `Session subscription expected sequence ${this.#expectedSequence} but received ${frame.sequence}`, + ); + } + this.#expectedSequence += 1; + + if (frame.kind === 'subscription.session_projection') { + if (frame.snapshot.session.sessionId !== this.#expectedSessionId) { + throw new RuntimeHostSubscriptionError( + 'correlation_changed', + 'Session subscription projection identity changed', + ); + } + if (frame.snapshot.projectionRevision <= this.#latestProjectionRevision) { + throw new RuntimeHostSubscriptionError( + 'projection_revision_invalid', + 'Session projection revision did not advance', + ); + } + this.#latestProjectionRevision = frame.snapshot.projectionRevision; + } else if ( + (frame.kind === 'subscription.session_delta' || + frame.kind === 'subscription.session_event') && + frame.sessionId !== this.#expectedSessionId + ) { + throw new RuntimeHostSubscriptionError( + 'correlation_changed', + 'Session subscription frame identity changed', + ); + } + + this.#offer(frame); + if (frame.kind === 'subscription.closed') this.#doneAfterQueue = true; + } + + finish(): void { + if (this.#done || this.#terminalError) return; + this.#doneAfterQueue = true; + if (this.#queue.length === 0) { + this.#done = true; + this.#waiting?.resolve({ done: true, value: undefined }); + this.#waiting = undefined; + } + } + + fail(error: Error): void { + if (this.#done || this.#terminalError) return; + this.#terminalError = error; + this.#queue.length = 0; + this.#queuedBytes = 0; + this.#waiting?.reject(error); + this.#waiting = undefined; + } + + #offer(frame: SubscriptionFrame): void { + if (this.#waiting) { + const waiting = this.#waiting; + this.#waiting = undefined; + waiting.resolve({ done: false, value: frame }); + return; + } + const encodedBytes = encodeProtocolFrame(frame).byteLength; + if ( + this.#queue.length >= MAX_CLIENT_QUEUED_FRAMES || + this.#queuedBytes + encodedBytes > MAX_CLIENT_QUEUED_BYTES + ) { + throw new RuntimeHostSubscriptionError( + 'slow_consumer', + 'Session subscription consumer exceeded its local queue bound', + ); + } + this.#queue.push({ frame, encodedBytes }); + this.#queuedBytes += encodedBytes; + } +} diff --git a/packages/runtime-host/src/protocol/index.ts b/packages/runtime-host/src/protocol/index.ts index 7fb3644f7c..f63b129c8a 100644 --- a/packages/runtime-host/src/protocol/index.ts +++ b/packages/runtime-host/src/protocol/index.ts @@ -2,6 +2,11 @@ import { TextDecoder } from 'node:util'; import { requireCount, requireId, requireRecord, requireString } from './codec.js'; import { invalidProtocolFrame, RuntimeHostProtocolError } from './errors.js'; import { requireHostLifecycleState } from './host-status.js'; +import { + decodeSubscriptionFrame, + isSubscriptionFrameKind, + type SubscriptionFrame, +} from './session-continuity.js'; import { decodeRequestFrame, decodeResponseFrame, @@ -13,6 +18,7 @@ import { export { RuntimeHostProtocolError } from './errors.js'; export * from './message.js'; export * from './operations.js'; +export * from './session-continuity.js'; export const RUNTIME_HOST_REGISTRATION_SCHEMA_VERSION = 1 as const; export const RUNTIME_HOST_PROTOCOL_VERSION = 0 as const; @@ -58,7 +64,7 @@ export interface HostDraining { export type HostHandshakeResult = HostAccepted | HostIncompatible | HostDraining; export type ClientFrame = ClientHello | RequestFrame; -export type HostFrame = HostHandshakeResult | ResponseFrame; +export type HostFrame = HostHandshakeResult | ResponseFrame | SubscriptionFrame; export interface HostRegistration { kind: 'maka-runtime-host'; @@ -137,8 +143,12 @@ export function decodeHostFrame(value: unknown): HostFrame { } satisfies HostIncompatible; } if (frame.kind === 'draining') { - return { kind: 'draining', hostEpoch: requireId(frame.hostEpoch, 'hostEpoch') }; + return { + kind: 'draining', + hostEpoch: requireId(frame.hostEpoch, 'hostEpoch'), + }; } + if (isSubscriptionFrameKind(frame.kind)) return decodeSubscriptionFrame(frame); return decodeResponseFrame(frame); } diff --git a/packages/runtime-host/src/protocol/operations.ts b/packages/runtime-host/src/protocol/operations.ts index d6c35dba2b..d997a18c29 100644 --- a/packages/runtime-host/src/protocol/operations.ts +++ b/packages/runtime-host/src/protocol/operations.ts @@ -9,6 +9,7 @@ import { type OperationSpec, } from './operation-spec.js'; import { RUNTIME_POLICY_OPERATION_SPECS } from './runtime-policy.js'; +import { SESSION_CONTINUITY_OPERATION_SPECS } from './session-continuity.js'; import { TURN_OPERATION_SPECS } from './turn.js'; export type { HostLifecycleState, HostStatusInput, HostStatusResult } from './host-status.js'; @@ -51,11 +52,16 @@ const CORE_OPERATION_SPECS = composeOperationSpecMaps( RUNTIME_POLICY_OPERATION_SPECS, ); -export const HOST_OPERATION_SPECS = composeOperationSpecMaps( +const CORE_AND_MESSAGE_OPERATION_SPECS = composeOperationSpecMaps( CORE_OPERATION_SPECS, MESSAGE_OPERATION_SPECS, ); +export const HOST_OPERATION_SPECS = composeOperationSpecMaps( + CORE_AND_MESSAGE_OPERATION_SPECS, + SESSION_CONTINUITY_OPERATION_SPECS, +); + export type OperationSpecMap = typeof HOST_OPERATION_SPECS; export type OperationKey = keyof OperationSpecMap; diff --git a/packages/runtime-host/src/protocol/session-continuity.ts b/packages/runtime-host/src/protocol/session-continuity.ts new file mode 100644 index 0000000000..40574720a9 --- /dev/null +++ b/packages/runtime-host/src/protocol/session-continuity.ts @@ -0,0 +1,627 @@ +import { + assertExactKeys, + requireCount, + requireEntityId, + requireExactRecord, + requireId, + requireRecord, +} from './codec.js'; +import { invalidProtocolFrame } from './errors.js'; +import { + decodeSessionMessageQueueProjection, + type SessionMessageQueueProjection, +} from './message.js'; +import { defineOperation } from './operation-spec.js'; +import { decodeTurnSnapshot, type TurnSnapshot } from './turn.js'; + +export const SESSION_CONTINUITY_SCHEMA_VERSION = 1 as const; +export const SESSION_CONTINUITY_SNAPSHOT_MAX_BYTES = 56 * 1024; +export const SESSION_LIVE_DELTA_MAX_BYTES = 16 * 1024; +export const SESSION_TOOL_NAME_MAX_BYTES = 256; +export const SESSION_SUBSCRIPTION_FRAME_MAX_BYTES = 64 * 1024 - 1; + +export type SessionLifecycleStatus = + | 'active' + | 'running' + | 'waiting_for_user' + | 'blocked' + | 'review' + | 'done' + | 'archived' + | 'aborted'; + +export interface SessionContinuityIdentity { + sessionId: string; + status: SessionLifecycleStatus; + createdAt: number; + lastUsedAt: number; + isArchived: boolean; + archivedAt?: number; +} + +export interface SessionContinuitySnapshot { + schemaVersion: typeof SESSION_CONTINUITY_SCHEMA_VERSION; + session: SessionContinuityIdentity; + projectionRevision: number; + rootTurn: TurnSnapshot | null; + queue: SessionMessageQueueProjection; +} + +export interface SubscriptionOpenInput { + sessionId: string; +} + +export interface SubscriptionOpenResult { + hostEpoch: string; + subscriptionId: string; + nextSequence: number; + snapshot: SessionContinuitySnapshot; +} + +export interface SubscriptionCloseInput { + subscriptionId: string; +} + +export interface SubscriptionCloseResult { + subscriptionId: string; +} + +interface SubscriptionEnvelope { + hostEpoch: string; + subscriptionId: string; + sequence: number; +} + +export interface SessionProjectionFrame extends SubscriptionEnvelope { + kind: 'subscription.session_projection'; + snapshot: SessionContinuitySnapshot; +} + +export interface SessionAssistantDelta { + kind: 'text' | 'thinking'; + turnId: string; + runId: string; + messageId: string; + text: string; +} + +export interface SessionDeltaFrame extends SubscriptionEnvelope { + kind: 'subscription.session_delta'; + sessionId: string; + delta: SessionAssistantDelta; +} + +interface SessionToolEventIdentity { + id: string; + turnId: string; + ts: number; + toolUseId: string; +} + +export type SessionToolEvent = + | (SessionToolEventIdentity & { + type: 'tool_start'; + toolName: string; + operationId?: string; + activityKind?: + | 'read' + | 'search' + | 'websearch' + | 'webfetch' + | 'edit' + | 'command' + | 'explore' + | 'browser' + | 'tool'; + displayName?: string; + stepId?: string; + }) + | (SessionToolEventIdentity & { + type: 'tool_output_delta'; + seq: number; + stream: 'stdout' | 'stderr'; + chunk: string; + redacted: boolean; + createdAt: number; + }) + | (SessionToolEventIdentity & { + type: 'tool_progress'; + chunk: string; + }) + | (SessionToolEventIdentity & { + type: 'tool_result'; + operationId?: string; + status: 'completed' | 'errored'; + durationMs?: number; + }); + +export interface SessionEventFrame extends SubscriptionEnvelope { + kind: 'subscription.session_event'; + sessionId: string; + runId: string; + event: SessionToolEvent; +} + +export interface SubscriptionClosedFrame extends SubscriptionEnvelope { + kind: 'subscription.closed'; + reason: 'slow_consumer' | 'session_removed'; +} + +export type SubscriptionFrame = + | SessionProjectionFrame + | SessionDeltaFrame + | SessionEventFrame + | SubscriptionClosedFrame; + +const SUBSCRIPTION_OPEN_ERRORS = [ + 'host_not_ready', + 'host_draining', + 'operation_unavailable', + 'not_found', + 'operation_conflict', + 'internal_failure', +] as const; + +const SUBSCRIPTION_CLOSE_ERRORS = [ + 'host_not_ready', + 'host_draining', + 'operation_unavailable', + 'not_found', + 'internal_failure', +] as const; + +export const SESSION_CONTINUITY_OPERATION_SPECS = { + 'subscription.open': defineOperation({ + mode: 'control', + availability: 'ready', + errors: SUBSCRIPTION_OPEN_ERRORS, + decodeInput: decodeSubscriptionOpenInput, + decodeOutput: decodeSubscriptionOpenResult, + }), + 'subscription.close': defineOperation({ + mode: 'control', + availability: 'ready', + errors: SUBSCRIPTION_CLOSE_ERRORS, + decodeInput: decodeSubscriptionCloseInput, + decodeOutput: decodeSubscriptionCloseResult, + }), +} as const; + +export function decodeSubscriptionFrame(value: unknown): SubscriptionFrame { + requireEncodedByteLimit(value, 'subscription frame', SESSION_SUBSCRIPTION_FRAME_MAX_BYTES); + const record = requireRecord(value, 'subscription frame'); + const envelope = decodeEnvelope(record); + let frame: SubscriptionFrame; + if (record.kind === 'subscription.session_projection') { + assertExactKeys(record, 'Session projection frame', [ + 'kind', + 'hostEpoch', + 'subscriptionId', + 'sequence', + 'snapshot', + ]); + const snapshot = decodeSessionContinuitySnapshot(record.snapshot); + assertQueueEpoch(snapshot, envelope.hostEpoch); + frame = { kind: record.kind, ...envelope, snapshot }; + } else if (record.kind === 'subscription.session_delta') { + assertExactKeys(record, 'Session delta frame', [ + 'kind', + 'hostEpoch', + 'subscriptionId', + 'sequence', + 'sessionId', + 'delta', + ]); + frame = { + kind: record.kind, + ...envelope, + sessionId: requireEntityId(record.sessionId, 'sessionId'), + delta: decodeAssistantDelta(record.delta), + }; + } else if (record.kind === 'subscription.session_event') { + assertExactKeys(record, 'Session event frame', [ + 'kind', + 'hostEpoch', + 'subscriptionId', + 'sequence', + 'sessionId', + 'runId', + 'event', + ]); + frame = { + kind: record.kind, + ...envelope, + sessionId: requireEntityId(record.sessionId, 'sessionId'), + runId: requireEntityId(record.runId, 'runId'), + event: decodeSessionToolEvent(record.event), + }; + } else if (record.kind === 'subscription.closed') { + assertExactKeys(record, 'subscription closed frame', [ + 'kind', + 'hostEpoch', + 'subscriptionId', + 'sequence', + 'reason', + ]); + if (record.reason !== 'slow_consumer' && record.reason !== 'session_removed') { + throw invalidProtocolFrame('Invalid subscription close reason'); + } + frame = { kind: record.kind, ...envelope, reason: record.reason }; + } else { + throw invalidProtocolFrame('Unknown subscription frame kind'); + } + return frame; +} + +export function isSubscriptionFrameKind(value: unknown): value is SubscriptionFrame['kind'] { + return ( + value === 'subscription.session_projection' || + value === 'subscription.session_delta' || + value === 'subscription.session_event' || + value === 'subscription.closed' + ); +} + +export function decodeSessionContinuitySnapshot(value: unknown): SessionContinuitySnapshot { + requireEncodedByteLimit( + value, + 'Session continuity snapshot', + SESSION_CONTINUITY_SNAPSHOT_MAX_BYTES, + ); + const record = requireExactRecord(value, 'Session continuity snapshot', [ + 'schemaVersion', + 'session', + 'projectionRevision', + 'rootTurn', + 'queue', + ]); + if (record.schemaVersion !== SESSION_CONTINUITY_SCHEMA_VERSION) { + throw invalidProtocolFrame('Unsupported Session continuity snapshot schema'); + } + const session = decodeSessionContinuityIdentity(record.session); + const rootTurn = record.rootTurn === null ? null : decodeTurnSnapshot(record.rootTurn); + if (rootTurn !== null && rootTurn.sessionId !== session.sessionId) { + throw invalidProtocolFrame('Session continuity root Turn belongs to a different Session'); + } + return { + schemaVersion: SESSION_CONTINUITY_SCHEMA_VERSION, + session, + projectionRevision: requirePositiveCount(record.projectionRevision, 'projectionRevision'), + rootTurn, + queue: decodeSessionMessageQueueProjection(record.queue), + }; +} + +function decodeSubscriptionOpenInput(value: unknown): SubscriptionOpenInput { + const record = requireExactRecord(value, 'subscription.open input', ['sessionId']); + return { sessionId: requireEntityId(record.sessionId, 'sessionId') }; +} + +function decodeSubscriptionOpenResult(value: unknown): SubscriptionOpenResult { + const record = requireExactRecord(value, 'subscription.open result', [ + 'hostEpoch', + 'subscriptionId', + 'nextSequence', + 'snapshot', + ]); + const hostEpoch = requireId(record.hostEpoch, 'hostEpoch'); + const snapshot = decodeSessionContinuitySnapshot(record.snapshot); + assertQueueEpoch(snapshot, hostEpoch); + return { + hostEpoch, + subscriptionId: requireId(record.subscriptionId, 'subscriptionId'), + nextSequence: requirePositiveCount(record.nextSequence, 'nextSequence'), + snapshot, + }; +} + +function decodeSubscriptionCloseInput(value: unknown): SubscriptionCloseInput { + const record = requireExactRecord(value, 'subscription.close input', ['subscriptionId']); + return { subscriptionId: requireId(record.subscriptionId, 'subscriptionId') }; +} + +function decodeSubscriptionCloseResult(value: unknown): SubscriptionCloseResult { + const record = requireExactRecord(value, 'subscription.close result', ['subscriptionId']); + return { subscriptionId: requireId(record.subscriptionId, 'subscriptionId') }; +} + +function decodeEnvelope(record: Record): SubscriptionEnvelope { + return { + hostEpoch: requireId(record.hostEpoch, 'hostEpoch'), + subscriptionId: requireId(record.subscriptionId, 'subscriptionId'), + sequence: requirePositiveCount(record.sequence, 'sequence'), + }; +} + +function decodeAssistantDelta(value: unknown): SessionAssistantDelta { + const record = requireExactRecord(value, 'Session assistant delta', [ + 'kind', + 'turnId', + 'runId', + 'messageId', + 'text', + ]); + if (record.kind !== 'text' && record.kind !== 'thinking') { + throw invalidProtocolFrame('Invalid Session assistant delta kind'); + } + return { + kind: record.kind, + turnId: requireEntityId(record.turnId, 'turnId'), + runId: requireEntityId(record.runId, 'runId'), + messageId: requireEntityId(record.messageId, 'messageId'), + text: requireUtf8BoundedString( + record.text, + 'Session assistant delta text', + SESSION_LIVE_DELTA_MAX_BYTES, + ), + }; +} + +function decodeSessionToolEvent(value: unknown): SessionToolEvent { + const record = requireRecord(value, 'Session tool event'); + const identity = { + id: requireId(record.id, 'Session tool event id'), + turnId: requireEntityId(record.turnId, 'turnId'), + ts: requireCount(record.ts, 'Session tool event timestamp'), + toolUseId: requireId(record.toolUseId, 'toolUseId'), + }; + if (record.type === 'tool_start') { + const allowed = [ + 'type', + 'id', + 'turnId', + 'ts', + 'toolUseId', + 'toolName', + 'operationId', + 'activityKind', + 'displayName', + 'stepId', + ]; + assertAllowedKeys(record, 'Session tool start event', allowed); + assertRequiredKeys(record, 'Session tool start event', [ + 'type', + 'id', + 'turnId', + 'ts', + 'toolUseId', + 'toolName', + ]); + return { + type: record.type, + ...identity, + toolName: requireUtf8BoundedString( + record.toolName, + 'Session tool name', + SESSION_TOOL_NAME_MAX_BYTES, + ), + ...(record.operationId === undefined + ? {} + : { operationId: requireEntityId(record.operationId, 'operationId') }), + ...(record.activityKind === undefined + ? {} + : { activityKind: requireToolActivityKind(record.activityKind) }), + ...(record.displayName === undefined + ? {} + : { + displayName: requireUtf8BoundedString( + record.displayName, + 'Session tool display name', + SESSION_TOOL_NAME_MAX_BYTES, + ), + }), + ...(record.stepId === undefined ? {} : { stepId: requireEntityId(record.stepId, 'stepId') }), + }; + } + if (record.type === 'tool_output_delta') { + assertExactKeys(record, 'Session tool output delta event', [ + 'type', + 'id', + 'turnId', + 'ts', + 'toolUseId', + 'seq', + 'stream', + 'chunk', + 'redacted', + 'createdAt', + ]); + if (record.stream !== 'stdout' && record.stream !== 'stderr') { + throw invalidProtocolFrame('Invalid Session tool output stream'); + } + if (typeof record.redacted !== 'boolean') { + throw invalidProtocolFrame('Invalid Session tool output redaction'); + } + return { + type: record.type, + ...identity, + seq: requireCount(record.seq, 'Session tool output sequence'), + stream: record.stream, + chunk: requireUtf8BoundedString( + record.chunk, + 'Session tool output chunk', + SESSION_LIVE_DELTA_MAX_BYTES, + ), + redacted: record.redacted, + createdAt: requireCount(record.createdAt, 'Session tool output timestamp'), + }; + } + if (record.type === 'tool_progress') { + assertExactKeys(record, 'Session tool progress event', [ + 'type', + 'id', + 'turnId', + 'ts', + 'toolUseId', + 'chunk', + ]); + return { + type: record.type, + ...identity, + chunk: requireUtf8BoundedString( + record.chunk, + 'Session tool progress chunk', + SESSION_LIVE_DELTA_MAX_BYTES, + ), + }; + } + if (record.type === 'tool_result') { + const allowed = [ + 'type', + 'id', + 'turnId', + 'ts', + 'toolUseId', + 'operationId', + 'status', + 'durationMs', + ]; + assertAllowedKeys(record, 'Session tool result event', allowed); + assertRequiredKeys(record, 'Session tool result event', [ + 'type', + 'id', + 'turnId', + 'ts', + 'toolUseId', + 'status', + ]); + if (record.status !== 'completed' && record.status !== 'errored') { + throw invalidProtocolFrame('Invalid Session tool result status'); + } + return { + type: record.type, + ...identity, + ...(record.operationId === undefined + ? {} + : { operationId: requireEntityId(record.operationId, 'operationId') }), + status: record.status, + ...(record.durationMs === undefined + ? {} + : { + durationMs: requireCount(record.durationMs, 'Session tool result duration'), + }), + }; + } + throw invalidProtocolFrame('Invalid Session tool event type'); +} + +function decodeSessionContinuityIdentity(value: unknown): SessionContinuityIdentity { + const record = requireRecord(value, 'Session continuity identity'); + assertAllowedKeys(record, 'Session continuity identity', [ + 'sessionId', + 'status', + 'createdAt', + 'lastUsedAt', + 'isArchived', + 'archivedAt', + ]); + assertRequiredKeys(record, 'Session continuity identity', [ + 'sessionId', + 'status', + 'createdAt', + 'lastUsedAt', + 'isArchived', + ]); + if (typeof record.isArchived !== 'boolean') { + throw invalidProtocolFrame('Invalid Session archived state'); + } + return { + sessionId: requireEntityId(record.sessionId, 'sessionId'), + status: requireSessionLifecycleStatus(record.status), + createdAt: requireCount(record.createdAt, 'createdAt'), + lastUsedAt: requireCount(record.lastUsedAt, 'lastUsedAt'), + isArchived: record.isArchived, + ...(record.archivedAt === undefined + ? {} + : { archivedAt: requireCount(record.archivedAt, 'archivedAt') }), + }; +} + +function assertQueueEpoch(snapshot: SessionContinuitySnapshot, hostEpoch: string): void { + if (snapshot.queue.hostEpoch !== hostEpoch) { + throw invalidProtocolFrame('Session queue projection belongs to a different Host Epoch'); + } +} + +function assertAllowedKeys( + record: Record, + label: string, + keys: readonly string[], +): void { + const allowed = new Set(keys); + if (Object.keys(record).some((key) => !allowed.has(key))) { + throw invalidProtocolFrame(`Unknown ${label} field`); + } +} + +function assertRequiredKeys( + record: Record, + label: string, + keys: readonly string[], +): void { + if (keys.some((key) => !Object.hasOwn(record, key))) { + throw invalidProtocolFrame(`Invalid ${label} fields`); + } +} + +function requirePositiveCount(value: unknown, label: string): number { + const count = requireCount(value, label); + if (count === 0) throw invalidProtocolFrame(`Invalid ${label}`); + return count; +} + +function requireUtf8BoundedString(value: unknown, label: string, maxBytes: number): string { + if ( + typeof value !== 'string' || + value.length === 0 || + Buffer.byteLength(value, 'utf8') > maxBytes + ) { + throw invalidProtocolFrame(`Invalid ${label}`); + } + return value; +} + +function requireEncodedByteLimit(value: unknown, label: string, maxBytes: number): void { + let encoded: string | undefined; + try { + encoded = JSON.stringify(value); + } catch { + throw invalidProtocolFrame(`Invalid ${label}`); + } + if (encoded === undefined || Buffer.byteLength(encoded, 'utf8') > maxBytes) { + throw invalidProtocolFrame(`Invalid ${label}`); + } +} + +function requireToolActivityKind( + value: unknown, +): Extract['activityKind'] { + if ( + value === 'read' || + value === 'search' || + value === 'websearch' || + value === 'webfetch' || + value === 'edit' || + value === 'command' || + value === 'explore' || + value === 'browser' || + value === 'tool' + ) + return value; + throw invalidProtocolFrame('Invalid Session tool activity kind'); +} + +function requireSessionLifecycleStatus(value: unknown): SessionLifecycleStatus { + if ( + value === 'active' || + value === 'running' || + value === 'waiting_for_user' || + value === 'blocked' || + value === 'review' || + value === 'done' || + value === 'archived' || + value === 'aborted' + ) + return value; + throw invalidProtocolFrame('Invalid Session lifecycle status'); +} diff --git a/packages/runtime-host/src/server/canonical-session-projection.ts b/packages/runtime-host/src/server/canonical-session-projection.ts new file mode 100644 index 0000000000..45b58dda26 --- /dev/null +++ b/packages/runtime-host/src/server/canonical-session-projection.ts @@ -0,0 +1,78 @@ +import type { SessionHeader } from '@maka/core/session'; +import type { ExecutionStoresWriter } from '@maka/storage/execution-stores'; +import type { + SessionContinuityIdentity, + SessionMessageQueueProjection, + TurnSnapshot, +} from '../protocol/index.js'; +import { isMissingFile, readCanonicalTurnSnapshot } from './canonical-turn-snapshot.js'; +import type { HostMessageCoordinator } from './message-coordinator.js'; +import type { RootAdmissionOwner } from './root-admission-owner.js'; + +type CanonicalSessionProjectionStores = Pick< + ExecutionStoresWriter<'interactive'>, + 'sessionStore' | 'agentRunStore' | 'runtimeEventStore' +>; + +export interface CanonicalSessionProjection { + readonly session: SessionContinuityIdentity; + readonly rootTurn: TurnSnapshot | null; + readonly queue: SessionMessageQueueProjection; +} + +export interface CanonicalSessionProjectionReaderOptions { + readonly stores: CanonicalSessionProjectionStores; + readonly rootAdmissions: RootAdmissionOwner; + readonly messages: Pick; +} + +export class CanonicalSessionProjectionReader { + readonly #stores: CanonicalSessionProjectionStores; + readonly #rootAdmissions: RootAdmissionOwner; + readonly #messages: Pick; + + constructor(options: CanonicalSessionProjectionReaderOptions) { + this.#stores = options.stores; + this.#rootAdmissions = options.rootAdmissions; + this.#messages = options.messages; + } + + async read(sessionId: string): Promise { + const admission = this.#rootAdmissions.latestAdmission(sessionId); + let header: SessionHeader; + try { + header = await this.#stores.sessionStore.readHeaderSnapshot(sessionId); + } catch (error) { + if (isMissingFile(error)) return null; + throw error; + } + if (header.id !== sessionId) { + throw new Error('Durable Session identity does not match the requested Session'); + } + + let rootTurn: TurnSnapshot | null = null; + if (admission) { + const durableAdmission = await this.#stores.agentRunStore.readRootTurnAdmission( + sessionId, + admission.turnId, + ); + if (!durableAdmission) { + throw new Error('Owned Root Turn admission is missing from durable storage'); + } + this.#rootAdmissions.assertKnownAdmission(durableAdmission); + rootTurn = await readCanonicalTurnSnapshot(this.#stores, durableAdmission); + } + + // The Session lane barrier must remain held through this final synchronous read. + const queue = this.#messages.projection(sessionId); + const session: SessionContinuityIdentity = { + sessionId: header.id, + status: header.status, + createdAt: header.createdAt, + lastUsedAt: header.lastUsedAt, + isArchived: header.isArchived, + ...(header.archivedAt !== undefined ? { archivedAt: header.archivedAt } : {}), + }; + return { session, rootTurn, queue }; + } +} diff --git a/packages/runtime-host/src/server/canonical-turn-snapshot.ts b/packages/runtime-host/src/server/canonical-turn-snapshot.ts new file mode 100644 index 0000000000..777928a642 --- /dev/null +++ b/packages/runtime-host/src/server/canonical-turn-snapshot.ts @@ -0,0 +1,93 @@ +import type { AgentRunHeader } from '@maka/core/agent-run'; +import { classifyTerminalRuntimeLedger } from '@maka/runtime'; +import type { ExecutionStoresWriter } from '@maka/storage/execution-stores'; +import type { TurnSnapshot } from '../protocol/index.js'; + +type CanonicalTurnStores = Pick< + ExecutionStoresWriter<'interactive'>, + 'agentRunStore' | 'runtimeEventStore' +>; + +export interface CanonicalTurnIdentity { + readonly sessionId: string; + readonly turnId: string; + readonly runId: string; +} + +export async function readCanonicalTurnSnapshot( + stores: CanonicalTurnStores, + identity: CanonicalTurnIdentity, + knownRun?: AgentRunHeader, +): Promise { + const { sessionId, turnId, runId } = identity; + const run = knownRun ?? (await readRunIfPresent(stores, sessionId, runId)); + if (!run) return { sessionId, turnId, runId, status: 'admitted' }; + if (run.turnId !== turnId) { + throw new Error('Admitted Turn identity does not match its Run header'); + } + + const [runEvents, runtimeEvents] = await Promise.all([ + stores.agentRunStore.readEvents(sessionId, runId), + stores.runtimeEventStore.readImmutableRuntimeEvents(sessionId, runId), + ]); + const terminal = classifyTerminalRuntimeLedger(run, runtimeEvents); + if (terminal.kind === 'fact') { + const fact = terminal.fact; + if (fact.runStatus === 'completed') { + return { + sessionId, + turnId, + runId, + status: 'completed', + terminalEventId: fact.terminalEvent.id, + }; + } + if (fact.runStatus === 'failed') { + if (!fact.failureClass) throw new Error('Failed terminal fact has no failure class'); + return { + sessionId, + turnId, + runId, + status: 'failed', + terminalEventId: fact.terminalEvent.id, + failureClass: fact.failureClass, + }; + } + if (!fact.abortSource) throw new Error('Cancelled terminal fact has no abort source'); + return { + sessionId, + turnId, + runId, + status: 'cancelled', + terminalEventId: fact.terminalEvent.id, + abortSource: fact.abortSource, + }; + } + if (terminal.kind !== 'none') { + throw new Error('Runtime ledger does not contain one canonical terminal fact'); + } + if (run.status === 'completed' || run.status === 'failed' || run.status === 'cancelled') { + throw new Error('Terminal Run header has no canonical terminal RuntimeEvent'); + } + if (run.status !== 'created' && !runEvents.some((event) => event.type === 'run_started')) { + throw new Error('Non-created Run has no durable start fact'); + } + return { sessionId, turnId, runId, status: run.status }; +} + +async function readRunIfPresent( + stores: CanonicalTurnStores, + sessionId: string, + runId: string, +): Promise { + try { + return await stores.agentRunStore.readRun(sessionId, runId); + } catch (error) { + if (isMissingFile(error)) return undefined; + throw error; + } +} + +export function isMissingFile(error: unknown): boolean { + return (error as NodeJS.ErrnoException | undefined)?.code === 'ENOENT'; +} diff --git a/packages/runtime-host/src/server/connection-session.ts b/packages/runtime-host/src/server/connection-session.ts index 25dcbae976..c1d97ba779 100644 --- a/packages/runtime-host/src/server/connection-session.ts +++ b/packages/runtime-host/src/server/connection-session.ts @@ -13,6 +13,10 @@ import { } from './operation-dispatcher.js'; import { BoundedSerialOutboundWriter } from './serial-outbound-writer.js'; import { RuntimeHostTransportError } from '../transport/framed-transport.js'; +import type { + SessionContinuityConnection, + SessionContinuityService, +} from './session-continuity-service.js'; const MAX_IN_FLIGHT_REQUESTS = 64; @@ -28,6 +32,7 @@ export interface RuntimeHostConnectionSessionOptions { transport: FramedTransport; connection: AcceptedConnectionContext; resolveHandlers(): OperationHandlerMap; + resolveContinuity(): SessionContinuityService | undefined; beginOperation(frame: RequestFrame): Promise; onTeardown(): void; } @@ -36,6 +41,9 @@ export class RuntimeHostConnectionSession { readonly #options: RuntimeHostConnectionSessionOptions; readonly #writer: BoundedSerialOutboundWriter; readonly #requests = new Map>(); + #continuityService: SessionContinuityService | undefined; + #continuity: SessionContinuityConnection | undefined; + #inputClosed = false; #closed = false; constructor(options: RuntimeHostConnectionSessionOptions) { @@ -61,6 +69,8 @@ export class RuntimeHostConnectionSession { } async #closeAfterDispatchedReplies(): Promise { + this.#inputClosed = true; + this.#detachContinuity(); const outcome = await Promise.race([ Promise.allSettled([...this.#requests.values()]).then(() => 'drained' as const), this.#options.transport.closed.then(() => 'closed' as const), @@ -117,20 +127,66 @@ export class RuntimeHostConnectionSession { try { if (this.#closed) return; + const continuity = + frame.operation === 'subscription.open' || frame.operation === 'subscription.close' + ? this.#ensureContinuity() + : undefined; const response = await dispatchOperation(frame, this.#options.resolveHandlers(), { ...this.#options.connection, acquireResidency: () => admission.acquireResidency(), }); admission.seal(); - await this.#writer.enqueue(response).flushed; + const receipt = this.#writer.enqueue(response); + const openedSubscriptionId = + response.ok && response.operation === 'subscription.open' + ? response.result.subscriptionId + : undefined; + if (openedSubscriptionId) continuity?.activate(openedSubscriptionId); + try { + await receipt.flushed; + } catch (error) { + if (openedSubscriptionId) continuity?.abort(openedSubscriptionId); + throw error; + } } finally { admission.finish(); } } + #ensureContinuity(): SessionContinuityConnection | undefined { + if (this.#closed || this.#inputClosed) return; + const service = this.#options.resolveContinuity(); + if (!service) return; + if (this.#continuityService && this.#continuityService !== service) { + throw new Error('Runtime Host continuity service changed within one connection'); + } + if (!this.#continuity) { + this.#continuityService = service; + this.#continuity = service.attachConnection(this.#options.connection.connectionId, { + send: (frame) => { + try { + return this.#writer.enqueue(frame).flushed; + } catch (error) { + return Promise.reject(error); + } + }, + close: () => this.#teardown(), + }); + } + return this.#continuity; + } + + #detachContinuity(): void { + this.#continuity?.close(); + this.#continuity = undefined; + this.#continuityService = undefined; + } + #teardown(): void { if (this.#closed) return; this.#closed = true; + this.#inputClosed = true; + this.#detachContinuity(); this.#writer.close(); this.#options.transport.destroy(); this.#options.onTeardown(); diff --git a/packages/runtime-host/src/server/execution-composition.ts b/packages/runtime-host/src/server/execution-composition.ts index f7463eda8e..4213b69236 100644 --- a/packages/runtime-host/src/server/execution-composition.ts +++ b/packages/runtime-host/src/server/execution-composition.ts @@ -7,6 +7,7 @@ import { } from '@maka/runtime'; import { openInteractiveExecutionStoresForWrite } from '@maka/storage/execution-stores'; import { openInteractiveRuntimePolicyStoresForWrite } from '@maka/storage/runtime-policy-stores'; +import { CanonicalSessionProjectionReader } from './canonical-session-projection.js'; import type { RuntimeHostComposition, RuntimeHostCompositionContext } from './host-kernel.js'; import { type HostMessageRootPort, HostMessageCoordinator } from './message-coordinator.js'; import type { DomainOperationHandlerMap } from './operation-dispatcher.js'; @@ -15,6 +16,7 @@ import { RootTurnCoordinator } from './root-turn-coordinator.js'; import { RuntimePolicyActivationGate } from './runtime-policy-activation-gate.js'; import { HostRuntimePolicyCoordinator } from './runtime-policy-coordinator.js'; import { SessionAdmissionGate } from './session-admission-gate.js'; +import { SessionContinuityCoordinator } from './session-continuity-coordinator.js'; export async function createExecutionRuntimeHostComposition( context: RuntimeHostCompositionContext, @@ -30,12 +32,14 @@ export async function createExecutionRuntimeHostComposition( const runtimePolicyActivation = new RuntimePolicyActivationGate(); const sessionAdmission = new SessionAdmissionGate(); let rootCoordinator: RootTurnCoordinator | undefined; + let continuity: SessionContinuityCoordinator | undefined; const rootPort: HostMessageRootPort = { readSessionHeader: (sessionId) => requireRootCoordinator(rootCoordinator).readSessionHeader(sessionId), readRootState: (sessionId) => requireRootCoordinator(rootCoordinator).readRootState(sessionId), - startFromMessage: (input) => requireRootCoordinator(rootCoordinator).startFromMessage(input), + startFromMessage: (input, admission) => + requireRootCoordinator(rootCoordinator).startFromMessage(input, admission), claimStop: (input, commitQueueFence) => requireRootCoordinator(rootCoordinator).claimStop(input, commitQueueFence), }; @@ -52,7 +56,22 @@ export async function createExecutionRuntimeHostComposition( sessionAdmission, acquireResidency: context.acquireResidency, requestDrain: context.requestDrain, + onProjectionChanged: (sessionId) => + requireContinuity(continuity).enqueueCanonicalRefresh(sessionId), }); + const rootAdmissionOwner = new RootAdmissionOwner(stores.agentRunStore); + const canonicalProjection = new CanonicalSessionProjectionReader({ + stores, + rootAdmissions: rootAdmissionOwner, + messages, + }); + continuity = new SessionContinuityCoordinator( + context.hostEpoch, + (sessionId) => canonicalProjection.read(sessionId), + sessionAdmission, + context.requestDrain, + ); + const continuityCoordinator = continuity; const runtimeAuthority: RuntimeHostedRootAuthority = { bindRun: (identity) => messages.bindRun(identity), executeRoot: (input) => requireRootCoordinator(rootCoordinator).executeRoot(input), @@ -71,13 +90,13 @@ export async function createExecutionRuntimeHostComposition( runBackendActivation: (operation) => runtimePolicyActivation.runBackendActivation(operation), messageAuthority: runtimeAuthority, }); - const rootAdmissionOwner = new RootAdmissionOwner(stores.agentRunStore); rootCoordinator = new RootTurnCoordinator( manager, stores, sessionAdmission, rootAdmissionOwner, messages, + continuityCoordinator, context.acquireResidency, context.requestDrain, ); @@ -98,9 +117,11 @@ export async function createExecutionRuntimeHostComposition( ...coordinator.handlers, ...messages.handlers, ...runtimePolicy.handlers, + ...continuityCoordinator.handlers, } satisfies DomainOperationHandlerMap; return { handlers, + continuity: continuityCoordinator, recover: async () => { const sessions = await stores.sessionStore.listForRecovery(); for (const session of sessions) { @@ -125,6 +146,7 @@ export async function createExecutionRuntimeHostComposition( } catch (error) { errors.push(error); } + continuityCoordinator.close(); try { await stores.sessionStore.close?.(); } catch (error) { @@ -145,3 +167,10 @@ function requireRootCoordinator(coordinator: RootTurnCoordinator | undefined): R if (!coordinator) throw new Error('Runtime Host root coordinator is not composed'); return coordinator; } + +function requireContinuity( + continuity: SessionContinuityCoordinator | undefined, +): SessionContinuityCoordinator { + if (!continuity) throw new Error('Runtime Host continuity coordinator is not composed'); + return continuity; +} diff --git a/packages/runtime-host/src/server/host-kernel.ts b/packages/runtime-host/src/server/host-kernel.ts index 67c274b093..0b62f531bf 100644 --- a/packages/runtime-host/src/server/host-kernel.ts +++ b/packages/runtime-host/src/server/host-kernel.ts @@ -32,6 +32,7 @@ import { type OperationResidency, type OperationHandlerMap, } from './operation-dispatcher.js'; +import type { SessionContinuityService } from './session-continuity-service.js'; const DEFAULT_IDLE_GRACE_MS = 30_000; const DEFAULT_HANDSHAKE_TIMEOUT_MS = 5_000; @@ -63,6 +64,7 @@ export interface RuntimeHostCompositionContext { export interface RuntimeHostComposition { readonly handlers: DomainOperationHandlerMap; + readonly continuity?: SessionContinuityService; recover(): Promise; close(): Promise; } @@ -239,6 +241,7 @@ export class RuntimeHostKernel { principal: 'local_os_user', }, resolveHandlers: () => this.#operationHandlers, + resolveContinuity: () => this.#composition?.continuity, beginOperation: (request) => this.#beginOperation(request), onTeardown: releaseConnection, }); diff --git a/packages/runtime-host/src/server/message-coordinator.ts b/packages/runtime-host/src/server/message-coordinator.ts index 60150b0b43..4a105336bb 100644 --- a/packages/runtime-host/src/server/message-coordinator.ts +++ b/packages/runtime-host/src/server/message-coordinator.ts @@ -40,7 +40,7 @@ import { } from '../protocol/index.js'; import type { RuntimeHostResidency } from './host-kernel.js'; import type { MessageOperationHandlerMap } from './operation-dispatcher.js'; -import { SessionAdmissionGate } from './session-admission-gate.js'; +import { type SessionAdmissionLease, SessionAdmissionGate } from './session-admission-gate.js'; type MessageOperationErrorCode = | 'host_draining' @@ -82,7 +82,10 @@ export interface HostMessageStopClaim { export interface HostMessageRootPort { readSessionHeader(sessionId: string): Promise; readRootState(sessionId: string): Promise | HostMessageRootState; - startFromMessage(input: HostMessageStartInput): Promise<{ readonly turnId: string }>; + startFromMessage( + input: HostMessageStartInput, + admission: SessionAdmissionLease, + ): Promise<{ readonly turnId: string }>; claimStop( input: Omit, commitQueueFence: () => QueueFenceResult, @@ -109,6 +112,7 @@ export interface HostMessageCoordinatorOptions { readonly sessionAdmission: SessionAdmissionGate; readonly acquireResidency: () => RuntimeHostResidency; readonly requestDrain?: () => void; + readonly onProjectionChanged?: (sessionId: string) => void; readonly createId?: () => string; } @@ -157,6 +161,7 @@ interface TerminalTransition { } interface SessionState { + readonly sessionId: string; revision: number; generation: number; phase: 'open' | 'closed'; @@ -208,6 +213,7 @@ export class HostMessageCoordinator implements RuntimeMessageAuthority { readonly #sessionAdmission: SessionAdmissionGate; readonly #acquireResidency: () => RuntimeHostResidency; readonly #requestDrain: () => void; + readonly #onProjectionChanged: (sessionId: string) => void; readonly #createId: () => string; readonly #sessions = new Map(); readonly #pendingSubmits = new Map(); @@ -226,6 +232,7 @@ export class HostMessageCoordinator implements RuntimeMessageAuthority { this.#sessionAdmission = options.sessionAdmission; this.#acquireResidency = options.acquireResidency; this.#requestDrain = options.requestDrain ?? (() => undefined); + this.#onProjectionChanged = options.onProjectionChanged ?? (() => undefined); this.#createId = options.createId ?? randomUUID; } @@ -427,7 +434,7 @@ export class HostMessageCoordinator implements RuntimeMessageAuthority { input: TurnMessageSubmitInput, payload: CanonicalSubmitPayload, ): Promise> { - return this.#sessionAdmission.run(input.sessionId, async () => { + return this.#sessionAdmission.run(input.sessionId, async (admission) => { if (this.#failStopped) { return failure('host_draining', 'Runtime Host message authority has failed'); } @@ -483,11 +490,14 @@ export class HostMessageCoordinator implements RuntimeMessageAuthority { placement: input.placement, disposition: 'turn_started', }; - const started = await this.#root.startFromMessage({ - sessionId: input.sessionId, - content: payload.content, - sourceMessage, - }); + const started = await this.#root.startFromMessage( + { + sessionId: input.sessionId, + content: payload.content, + sourceMessage, + }, + admission, + ); if (!isEntityId(started.turnId)) { throw new RuntimeMessageAuthorityInvariantError('Started Turn identity is not encodable'); } @@ -576,7 +586,7 @@ export class HostMessageCoordinator implements RuntimeMessageAuthority { }; if (disposition === 'steering') state.steering.push(entry); else state.followup.push(entry); - state.revision = result.queueRevision; + this.#mutated(state); try { await this.#commitReceipt('submit', input.sessionId, input.messageId, payload, result); } catch (error) { @@ -1112,6 +1122,7 @@ export class HostMessageCoordinator implements RuntimeMessageAuthority { let state = this.#sessions.get(sessionId); if (!state) { state = { + sessionId, revision: 0, generation: 0, phase: 'open', @@ -1134,6 +1145,7 @@ export class HostMessageCoordinator implements RuntimeMessageAuthority { #mutated(state: SessionState): void { state.revision += 1; + this.#onProjectionChanged(state.sessionId); } #maybeReclaim(sessionId: string, state: SessionState): void { diff --git a/packages/runtime-host/src/server/operation-dispatcher.ts b/packages/runtime-host/src/server/operation-dispatcher.ts index b10a62c566..a57b38dd6e 100644 --- a/packages/runtime-host/src/server/operation-dispatcher.ts +++ b/packages/runtime-host/src/server/operation-dispatcher.ts @@ -43,10 +43,18 @@ export type MessageOperationKey = Extract< OperationKey, 'turn.message.submit' | 'queue.retract' | 'turn.interrupt' >; +export type SessionContinuityOperationKey = Extract< + OperationKey, + 'subscription.open' | 'subscription.close' +>; export type DomainOperationHandlerMap = Pick; export type TurnOperationHandlerMap = Pick; export type RuntimePolicyOperationHandlerMap = Pick; export type MessageOperationHandlerMap = Pick; +export type SessionContinuityOperationHandlerMap = Pick< + OperationHandlerMap, + SessionContinuityOperationKey +>; export function composeOperationHandlers( ...handlerMaps: readonly Partial[] diff --git a/packages/runtime-host/src/server/root-admission-owner.ts b/packages/runtime-host/src/server/root-admission-owner.ts index 1d44b04402..ea2f6e7a58 100644 --- a/packages/runtime-host/src/server/root-admission-owner.ts +++ b/packages/runtime-host/src/server/root-admission-owner.ts @@ -13,6 +13,15 @@ import type { } from '@maka/storage/execution-stores'; type OwnedAdmitRootTurnInput = Omit; +type Immutable = T extends (...args: never[]) => unknown + ? T + : T extends readonly (infer Item)[] + ? readonly Immutable[] + : T extends object + ? { readonly [Key in keyof T]: Immutable } + : T; + +export type ValidatedRootTurnAdmission = Immutable; export class RootAdmissionOwner { readonly #admissionsBySession = new Map>(); @@ -21,6 +30,10 @@ export class RootAdmissionOwner { constructor(private readonly store: RootTurnAdmissionStore) {} + latestAdmission(sessionId: string): ValidatedRootTurnAdmission | undefined { + return this.#tips.get(sessionId); + } + assertKnownAdmission(admission: RootTurnAdmission): void { const known = this.#admissionsBySession.get(admission.sessionId)?.get(admission.turnId); if (!known || !sameRootAdmission(known, admission)) { diff --git a/packages/runtime-host/src/server/root-turn-coordinator.ts b/packages/runtime-host/src/server/root-turn-coordinator.ts index bf0e6020f4..85703fb823 100644 --- a/packages/runtime-host/src/server/root-turn-coordinator.ts +++ b/packages/runtime-host/src/server/root-turn-coordinator.ts @@ -6,6 +6,7 @@ import { messageContentsEqual, normalizeMessageContent, type MessageContent, + type SessionEvent, } from '@maka/core/events'; import { isDeepResearchSession, @@ -46,7 +47,11 @@ import { } from './message-coordinator.js'; import type { ConnectionContext, TurnOperationHandlerMap } from './operation-dispatcher.js'; import { RootAdmissionOwner } from './root-admission-owner.js'; -import { SessionAdmissionGate } from './session-admission-gate.js'; +import { type SessionAdmissionLease, SessionAdmissionGate } from './session-admission-gate.js'; +import { + type RuntimeSessionTransientEvent, + SessionContinuityCoordinator, +} from './session-continuity-coordinator.js'; interface ActiveRootTurn { turnId: string; @@ -104,6 +109,7 @@ export class RootTurnCoordinator { private readonly sessionAdmission: SessionAdmissionGate, private readonly rootAdmissionOwner: RootAdmissionOwner, private readonly messages: HostMessageCoordinator, + private readonly continuity: SessionContinuityCoordinator, private readonly acquireRecoveryResidency: () => RuntimeHostResidency, private readonly requestHostDrain: () => void, ) { @@ -278,8 +284,8 @@ export class RootTurnCoordinator { turnId: admission.turnId, content: normalizeMessageContent(admission.normalizedInput), }; - const disposition = await this.sessionAdmission.run(sessionId, () => - this.prepareAdmittedTurn(input, admission, this.acquireRecoveryResidency), + const disposition = await this.sessionAdmission.run(sessionId, (lease) => + this.prepareAdmittedTurn(input, admission, this.acquireRecoveryResidency, lease), ); const outcome = await this.resolveStartDisposition(input, disposition); if (!outcome.ok) { @@ -338,7 +344,7 @@ export class RootTurnCoordinator { ...input, content: normalizeMessageContent(input.content), }; - const disposition = await this.sessionAdmission.run(input.sessionId, async () => { + const disposition = await this.sessionAdmission.run(input.sessionId, async (lease) => { const existing = await this.stores.agentRunStore.readRootTurnAdmission( input.sessionId, input.turnId, @@ -360,6 +366,7 @@ export class RootTurnCoordinator { canonicalInput, existing, this.acquireRecoveryResidency, + lease, undefined, canonicalInput, ); @@ -402,6 +409,7 @@ export class RootTurnCoordinator { canonicalInput, admitted.admission, this.acquireRecoveryResidency, + lease, undefined, canonicalInput, ); @@ -473,7 +481,10 @@ export class RootTurnCoordinator { }); } - startFromMessage(input: HostMessageStartInput): Promise<{ readonly turnId: string }> { + startFromMessage( + input: HostMessageStartInput, + admissionLease: SessionAdmissionLease, + ): Promise<{ readonly turnId: string }> { return this.runCommand(async () => { const content = normalizeMessageContent(input.content); if ( @@ -510,6 +521,7 @@ export class RootTurnCoordinator { { sessionId: input.sessionId, turnId, content }, admitted.admission, this.acquireRecoveryResidency, + admissionLease, ); if (disposition.kind !== 'await_start') { throw new RuntimeMessageAuthorityInvariantError( @@ -556,7 +568,7 @@ export class RootTurnCoordinator { ...input, content: normalizeMessageContent(input.content), }; - const disposition = await this.sessionAdmission.run(input.sessionId, async () => { + const disposition = await this.sessionAdmission.run(input.sessionId, async (lease) => { const existing = await this.stores.agentRunStore.readRootTurnAdmission( input.sessionId, input.turnId, @@ -576,7 +588,12 @@ export class RootTurnCoordinator { operationConflict('Turn identity was already admitted with a different payload'), ); } - return this.prepareAdmittedTurn(canonicalInput, existing, context.acquireResidency); + return this.prepareAdmittedTurn( + canonicalInput, + existing, + context.acquireResidency, + lease, + ); } let header: SessionHeader; @@ -627,6 +644,7 @@ export class RootTurnCoordinator { canonicalInput, admission.admission, context.acquireResidency, + lease, ); }); return this.resolveStartDisposition(canonicalInput, disposition); @@ -715,6 +733,7 @@ export class RootTurnCoordinator { input: TurnStartInput, admission: RootTurnAdmission, acquireResidency: () => RuntimeHostResidency, + admissionLease: SessionAdmissionLease, replacing?: ActiveRootTurn, execution?: RuntimeHostedRootExecutionInput, ): Promise { @@ -759,9 +778,18 @@ export class RootTurnCoordinator { const residency = acquireResidency(); const messageIdentity = { sessionId: input.sessionId, turnId: input.turnId, runId }; + let messageReserved = false; try { this.messages.reserveRootTurn(messageIdentity); + messageReserved = true; + await this.continuity.holdTerminalPublication( + input.sessionId, + input.turnId, + runId, + admissionLease, + ); } catch (error) { + if (messageReserved) this.messages.abandonRootReservation(messageIdentity); residency.release(); throw error; } @@ -826,6 +854,7 @@ export class RootTurnCoordinator { runId: active.runId, userMessageId: active.userMessageId, onRunStarted: async () => { + await this.continuity.refreshCanonical(input.sessionId); started.resolve(); await active.execution?.onReady?.(); }, @@ -837,10 +866,11 @@ export class RootTurnCoordinator { runId: active.runId, userMessageId: active.userMessageId ?? undefined, durability: 'required', - onRunStarted: (startedRunId) => { + onRunStarted: async (startedRunId) => { if (startedRunId !== active.runId) { throw new Error('Runtime started a different Run than the admitted identity'); } + await this.continuity.refreshCanonical(input.sessionId); started.resolve(); }, }, @@ -853,7 +883,9 @@ export class RootTurnCoordinator { // Presentation observers do not participate in execution authority. } } - // The Host must consume the complete stream so Runtime finalization can commit. + if (isRuntimeSessionTransientEvent(event)) { + await this.continuity.acceptRuntimeEvent(input.sessionId, active.runId, event); + } } const snapshot = await this.readCanonicalSnapshot( input.sessionId, @@ -918,7 +950,7 @@ export class RootTurnCoordinator { } private completeTerminalTransition(sessionId: string, active: ActiveRootTurn): Promise { - return this.sessionAdmission.run(sessionId, async () => { + return this.sessionAdmission.run(sessionId, async (lease) => { if (this.#activeBySession.get(sessionId) !== active) { throw new RuntimeMessageAuthorityInvariantError( 'Terminal root Turn no longer owns the Session', @@ -926,19 +958,26 @@ export class RootTurnCoordinator { } const identity = { sessionId, turnId: active.turnId, runId: active.runId }; const batch = this.messages.beginTerminalTransition(identity); + await this.continuity.publishTerminalProjection( + sessionId, + active.turnId, + active.runId, + lease, + ); if (batch.sources.length === 0) { this.messages.completeIdle(batch); active.messageTransitionCommitted = true; this.#activeBySession.delete(sessionId); return; } - await this.startFollowupBatch(batch, active); + await this.startFollowupBatch(batch, active, lease); }); } private async startFollowupBatch( batch: RootFollowupBatch, previous: ActiveRootTurn, + admissionLease: SessionAdmissionLease, ): Promise { const turnId = randomUUID(); const admitted = await this.rootAdmissionOwner.admitRootTurn({ @@ -977,6 +1016,7 @@ export class RootTurnCoordinator { }, admitted.admission, this.acquireRecoveryResidency, + admissionLease, previous, ); if (disposition.kind !== 'await_start') { @@ -1233,6 +1273,19 @@ function isTerminalSnapshot(snapshot: TurnSnapshot): boolean { ); } +function isRuntimeSessionTransientEvent( + event: SessionEvent, +): event is RuntimeSessionTransientEvent { + return ( + event.type === 'text_delta' || + event.type === 'thinking_delta' || + event.type === 'tool_start' || + event.type === 'tool_output_delta' || + event.type === 'tool_progress' || + event.type === 'tool_result' + ); +} + function completedStart(outcome: TurnStartOutcome): TurnStartDisposition { return { kind: 'complete', outcome }; } diff --git a/packages/runtime-host/src/server/session-admission-gate.ts b/packages/runtime-host/src/server/session-admission-gate.ts index d69a4ec644..fbf43c6d9f 100644 --- a/packages/runtime-host/src/server/session-admission-gate.ts +++ b/packages/runtime-host/src/server/session-admission-gate.ts @@ -1,7 +1,86 @@ +import { AsyncLocalStorage } from 'node:async_hooks'; + +const sessionAdmissionLeaseBrand: unique symbol = Symbol('SessionAdmissionLease'); + +export interface SessionAdmissionLease { + readonly [sessionAdmissionLeaseBrand]: true; +} + +interface SessionAdmissionContext { + readonly sessionId: string; + active: boolean; +} + +interface SessionAdmissionLeaseState { + readonly sessionId: string; + readonly context: SessionAdmissionContext; + readonly tasks: Promise[]; + accepting: boolean; +} + +type SessionAdmissionTaskResult = + | { readonly ok: true } + | { readonly ok: false; readonly error: unknown }; + export class SessionAdmissionGate { readonly #tails = new Map>(); + readonly #context = new AsyncLocalStorage(); + readonly #leases = new WeakMap(); + + run( + sessionId: string, + operation: (lease: SessionAdmissionLease) => Promise | T, + ): Promise { + if (this.#context.getStore()?.active) { + return Promise.reject( + new Error( + 'Cannot enter Session admission from an active admission; reuse its lease instead', + ), + ); + } + return this.#runQueued(sessionId, operation); + } + + enqueueDetached( + sessionId: string, + operation: (lease: SessionAdmissionLease) => Promise | void, + ): Promise { + return this.#runQueued(sessionId, operation); + } - async run(sessionId: string, operation: () => Promise | T): Promise { + runAdmitted( + sessionId: string, + lease: SessionAdmissionLease, + operation: () => Promise | T, + ): Promise { + const state = this.#requireLease(sessionId, lease); + if (!state.accepting) { + return Promise.reject(new Error('Session admission lease no longer accepts tasks')); + } + const inherited = this.#context.getStore(); + if (inherited?.active && inherited.sessionId !== sessionId) { + return Promise.reject(new Error('Cannot nest Session admission across Sessions')); + } + + let task: Promise; + try { + task = Promise.resolve(this.#context.run(state.context, operation)); + } catch (error) { + task = Promise.reject(error); + } + state.tasks.push( + task.then( + (): SessionAdmissionTaskResult => ({ ok: true }), + (error): SessionAdmissionTaskResult => ({ ok: false, error }), + ), + ); + return task; + } + + async #runQueued( + sessionId: string, + operation: (lease: SessionAdmissionLease) => Promise | T, + ): Promise { const previous = this.#tails.get(sessionId) ?? Promise.resolve(); let release!: () => void; const current = new Promise((resolve) => { @@ -10,11 +89,59 @@ export class SessionAdmissionGate { const tail = previous.then(() => current); this.#tails.set(sessionId, tail); await previous; + + const context: SessionAdmissionContext = { sessionId, active: true }; + const lease: SessionAdmissionLease = Object.freeze({ + [sessionAdmissionLeaseBrand]: true as const, + }); + const state: SessionAdmissionLeaseState = { + sessionId, + context, + tasks: [], + accepting: true, + }; + this.#leases.set(lease, state); try { - return await operation(); + let result!: T; + let operationError: unknown; + let operationFailed = false; + try { + result = await this.#context.run(context, () => operation(lease)); + } catch (error) { + operationFailed = true; + operationError = error; + } finally { + state.accepting = false; + } + + const taskResults = await Promise.all(state.tasks); + const errors: unknown[] = []; + const collect = (error: unknown) => { + if (!errors.some((existing) => Object.is(existing, error))) errors.push(error); + }; + if (operationFailed) collect(operationError); + for (const taskResult of taskResults) { + if (!taskResult.ok) collect(taskResult.error); + } + if (errors.length === 1) throw errors[0]; + if (errors.length > 1) { + throw new AggregateError(errors, 'Session admission operation failed'); + } + return result; } finally { + context.active = false; + this.#leases.delete(lease); release(); if (this.#tails.get(sessionId) === tail) this.#tails.delete(sessionId); } } + + #requireLease(sessionId: string, lease: SessionAdmissionLease): SessionAdmissionLeaseState { + const state = this.#leases.get(lease); + if (!state) throw new Error('Session admission lease was not issued by this gate'); + if (state.sessionId !== sessionId) { + throw new Error('Session admission lease does not match the Session'); + } + return state; + } } diff --git a/packages/runtime-host/src/server/session-continuity-coordinator.ts b/packages/runtime-host/src/server/session-continuity-coordinator.ts new file mode 100644 index 0000000000..89462a8fdb --- /dev/null +++ b/packages/runtime-host/src/server/session-continuity-coordinator.ts @@ -0,0 +1,785 @@ +import { randomUUID } from 'node:crypto'; +import { isDeepStrictEqual } from 'node:util'; +import type { SessionEvent } from '@maka/core/events'; +import { + encodeProtocolFrame, + RUNTIME_HOST_MAX_FRAME_BYTES, + SESSION_CONTINUITY_SCHEMA_VERSION, + SESSION_LIVE_DELTA_MAX_BYTES, + SESSION_TOOL_NAME_MAX_BYTES, + type SessionAssistantDelta, + type SessionContinuitySnapshot, + type SessionDeltaFrame, + type SessionEventFrame, + type SessionToolEvent, + type SubscriptionFrame, + type SubscriptionOpenResult, + type TurnSnapshot, +} from '../protocol/index.js'; +import type { SessionContinuityOperationHandlerMap } from './operation-dispatcher.js'; +import { type SessionAdmissionLease, SessionAdmissionGate } from './session-admission-gate.js'; +import type { CanonicalSessionProjection } from './canonical-session-projection.js'; +import type { + SessionContinuityConnection, + SessionContinuityFrameSink, + SessionContinuityService, +} from './session-continuity-service.js'; + +const MAX_CONNECTION_SUBSCRIPTIONS = 16; +const MAX_SUBSCRIBER_QUEUED_FRAMES = 32; +const MAX_SUBSCRIBER_QUEUED_BYTES = 256 * 1024; + +export type { CanonicalSessionProjection } from './canonical-session-projection.js'; + +export type RuntimeSessionTransientEvent = Extract< + SessionEvent, + { + type: + | 'text_delta' + | 'thinking_delta' + | 'tool_start' + | 'tool_output_delta' + | 'tool_progress' + | 'tool_result'; + } +>; + +export type ReadCanonicalSessionProjection = ( + sessionId: string, +) => Promise; + +interface SessionProjectionState { + canonical: CanonicalSessionProjection; + revision: number; + subscribers: Map; + terminalPublicationFence?: TerminalPublicationFence; +} + +interface TerminalPublicationFence { + turnId: string; + runId: string; +} + +interface ConnectionState { + sink: SessionContinuityFrameSink; + subscriptionIds: Set; + pendingOpenCount: number; +} + +interface QueuedSubscriptionFrame { + frame: SubscriptionFrame; + encodedBytes: number; +} + +interface Subscriber { + connectionId: string; + sessionId: string; + subscriptionId: string; + sink: SessionContinuityFrameSink; + phase: 'open' | 'closing' | 'closed'; + activated: boolean; + nextSequence: number; + lastFlushedSequence: number; + queue: QueuedSubscriptionFrame[]; + queuedBytes: number; + pumping: boolean; + terminalQueued: boolean; +} + +export class SessionContinuityCoordinator implements SessionContinuityService { + readonly handlers: SessionContinuityOperationHandlerMap = { + 'subscription.open': async (input, context) => { + const result = await this.#open(context.connectionId, input.sessionId); + return result.ok + ? { ok: true, result: result.value } + : { ok: false, error: { code: result.code, message: result.message } }; + }, + 'subscription.close': async (input, context) => { + const closed = this.#closeSubscription(context.connectionId, input.subscriptionId); + return closed + ? { ok: true, result: { subscriptionId: input.subscriptionId } } + : { + ok: false, + error: { code: 'not_found', message: 'Session subscription was not found' }, + }; + }, + }; + + readonly #connections = new Map(); + readonly #sessions = new Map(); + readonly #subscriptions = new Map(); + readonly #pendingRefreshes = new Set(); + readonly #hostEpoch: string; + readonly #readCanonical: ReadCanonicalSessionProjection; + #closed = false; + + constructor( + hostEpoch: string, + readCanonical: ReadCanonicalSessionProjection, + private readonly sessionAdmission: SessionAdmissionGate, + private readonly onPublicationFailure: (error: unknown) => void = () => undefined, + ) { + this.#hostEpoch = hostEpoch; + this.#readCanonical = readCanonical; + } + + attachConnection( + connectionId: string, + sink: SessionContinuityFrameSink, + ): SessionContinuityConnection { + if (this.#closed) throw new Error('Session continuity coordinator is closed'); + if (this.#connections.has(connectionId)) { + throw new Error(`Duplicate Runtime Host connection: ${connectionId}`); + } + this.#connections.set(connectionId, { + sink, + subscriptionIds: new Set(), + pendingOpenCount: 0, + }); + let attached = true; + return { + activate: (subscriptionId) => { + if (attached) this.#activate(connectionId, subscriptionId); + }, + abort: (subscriptionId) => { + if (attached) this.#abortSubscription(connectionId, subscriptionId); + }, + close: () => { + if (!attached) return; + attached = false; + this.#closeConnection(connectionId); + }, + }; + } + + async refreshCanonical(sessionId: string, admission?: SessionAdmissionLease): Promise { + await this.#runInSessionLane( + sessionId, + async () => { + if (this.#closed) return; + const state = this.#sessions.get(sessionId); + if (!state || (state.subscribers.size === 0 && !state.terminalPublicationFence)) return; + const canonical = await this.#readCanonicalProjection(sessionId); + if (this.#closed || !canonical) return; + const committed = this.#commitCanonical(sessionId, canonical); + if (committed.changed) this.#broadcastProjection(committed.state, committed.value); + }, + admission, + ); + } + + /** Safe for synchronous commit hooks: this only schedules and coalesces lane work. */ + enqueueCanonicalRefresh(sessionId: string): void { + if (this.#closed || this.#pendingRefreshes.has(sessionId)) return; + this.#pendingRefreshes.add(sessionId); + void this.sessionAdmission + .enqueueDetached(sessionId, (lease) => this.refreshCanonical(sessionId, lease)) + .then( + () => this.#pendingRefreshes.delete(sessionId), + (error) => { + this.#pendingRefreshes.delete(sessionId); + this.onPublicationFailure(error); + }, + ); + } + + async holdTerminalPublication( + sessionId: string, + turnId: string, + runId: string, + admission?: SessionAdmissionLease, + ): Promise { + await this.#runInSessionLane( + sessionId, + async () => { + if (this.#closed) throw new Error('Session continuity coordinator is closed'); + const state = this.#sessions.get(sessionId); + const existing = state?.terminalPublicationFence; + if (existing) { + if (existing.turnId === turnId && existing.runId === runId) return; + throw new Error('Session already has a different terminal publication fence'); + } + + const canonical = await this.#readCanonicalProjection(sessionId); + if (this.#closed) throw new Error('Session continuity coordinator is closed'); + if (!canonical) throw new Error('Cannot fence a missing Session projection'); + const rootTurn = requirePublicationFenceIdentity(canonical, sessionId, { turnId, runId }); + if (isTerminalTurn(rootTurn)) { + throw new Error( + 'Terminal publication fence identity does not match a non-terminal canonical Turn', + ); + } + const committed = this.#commitCanonical(sessionId, canonical); + committed.state.terminalPublicationFence = { turnId, runId }; + if (committed.changed) this.#broadcastProjection(committed.state, committed.value); + }, + admission, + ); + } + + async publishTerminalProjection( + sessionId: string, + turnId: string, + runId: string, + admission?: SessionAdmissionLease, + ): Promise { + await this.#runInSessionLane( + sessionId, + async () => { + if (this.#closed) throw new Error('Session continuity coordinator is closed'); + const state = this.#sessions.get(sessionId); + const fence = state?.terminalPublicationFence; + if (!state || !fence || fence.turnId !== turnId || fence.runId !== runId) { + throw new Error('Terminal publication does not own the Session continuity fence'); + } + const canonical = await this.#readCanonicalProjection(sessionId); + if (this.#closed) throw new Error('Session continuity coordinator is closed'); + if (!canonical) { + throw new Error('Canonical Session projection is not terminal for the fenced Turn'); + } + const rootTurn = requirePublicationFenceIdentity(canonical, sessionId, fence); + if (!isTerminalTurn(rootTurn)) { + throw new Error('Canonical Session projection is not terminal for the fenced Turn'); + } + if (isDeepStrictEqual(state.canonical, canonical)) { + throw new Error('Fenced terminal projection was already published'); + } + + const nextRevision = state.revision + 1; + const snapshot = snapshotValue(canonical, nextRevision); + state.canonical = canonical; + state.revision = nextRevision; + delete state.terminalPublicationFence; + this.#broadcastProjection(state, snapshot); + if (state.subscribers.size === 0) this.#sessions.delete(sessionId); + }, + admission, + ); + } + + async acceptRuntimeEvent( + sessionId: string, + runId: string, + event: RuntimeSessionTransientEvent, + ): Promise { + if ( + (event.type === 'text_delta' || event.type === 'thinking_delta') && + event.text.length === 0 + ) { + return; + } + if ( + (event.type === 'tool_output_delta' && event.chunk.length === 0) || + (event.type === 'tool_progress' && + (typeof event.chunk === 'string' ? event.chunk : event.chunk.text).length === 0) + ) { + return; + } + await this.sessionAdmission.run(sessionId, () => { + const state = this.#sessions.get(sessionId); + if (!state || state.subscribers.size === 0) return; + const rootTurn = state.canonical.rootTurn; + if ( + !rootTurn || + rootTurn.sessionId !== sessionId || + rootTurn.turnId !== event.turnId || + rootTurn.runId !== runId || + isTerminalTurn(rootTurn) || + (event.type === 'tool_output_delta' && event.sessionId !== sessionId) + ) { + throw new Error('Runtime event does not belong to the canonical active root Turn'); + } + if (event.type === 'text_delta' || event.type === 'thinking_delta') { + const kind: SessionAssistantDelta['kind'] = + event.type === 'text_delta' ? 'text' : 'thinking'; + for (const subscriber of state.subscribers.values()) { + this.#enqueueAssistantDelta(subscriber, sessionId, runId, event, kind); + } + return; + } + const projected = projectToolEvent(event); + for (const subscriber of state.subscribers.values()) { + const frame: SessionEventFrame = { + kind: 'subscription.session_event', + hostEpoch: this.#hostEpoch, + subscriptionId: subscriber.subscriptionId, + sequence: subscriber.nextSequence, + sessionId, + runId, + event: projected, + }; + this.#enqueue(subscriber, frame); + } + }); + } + + close(): void { + if (this.#closed) return; + this.#closed = true; + for (const connectionId of [...this.#connections.keys()]) this.#closeConnection(connectionId); + this.#sessions.clear(); + this.#subscriptions.clear(); + this.#pendingRefreshes.clear(); + } + + async #open( + connectionId: string, + sessionId: string, + ): Promise< + | { ok: true; value: SubscriptionOpenResult } + | { ok: false; code: 'not_found' | 'operation_conflict'; message: string } + > { + const connection = this.#connections.get(connectionId); + if (!connection) throw new Error('Runtime Host connection is not attached to continuity'); + if ( + connection.subscriptionIds.size + connection.pendingOpenCount >= + MAX_CONNECTION_SUBSCRIPTIONS + ) { + return { + ok: false, + code: 'operation_conflict', + message: 'Runtime Host connection subscription limit reached', + }; + } + connection.pendingOpenCount += 1; + try { + return await this.sessionAdmission.run(sessionId, async () => { + if (this.#connections.get(connectionId) !== connection) { + throw new Error('Runtime Host connection closed during subscription open'); + } + const canonical = await this.#readCanonicalProjection(sessionId); + if (this.#connections.get(connectionId) !== connection) { + throw new Error('Runtime Host connection closed during subscription open'); + } + if (!canonical) { + return { + ok: false as const, + code: 'not_found' as const, + message: 'Session was not found', + }; + } + const committed = this.#commitCanonical(sessionId, canonical); + if (committed.changed) this.#broadcastProjection(committed.state, committed.value); + if (this.#connections.get(connectionId) !== connection) { + this.#scheduleInactiveStateCleanup(sessionId, committed.state); + throw new Error('Runtime Host connection closed during subscription open'); + } + + const subscriptionId = randomUUID(); + const subscriber: Subscriber = { + connectionId, + sessionId, + subscriptionId, + sink: connection.sink, + phase: 'open', + activated: false, + nextSequence: 1, + lastFlushedSequence: 0, + queue: [], + queuedBytes: 0, + pumping: false, + terminalQueued: false, + }; + committed.state.subscribers.set(subscriptionId, subscriber); + this.#subscriptions.set(subscriptionId, subscriber); + connection.subscriptionIds.add(subscriptionId); + return { + ok: true as const, + value: { + hostEpoch: this.#hostEpoch, + subscriptionId, + nextSequence: subscriber.nextSequence, + snapshot: committed.value, + }, + }; + }); + } finally { + connection.pendingOpenCount -= 1; + } + } + + #activate(connectionId: string, subscriptionId: string): void { + const subscriber = this.#ownedSubscriber(connectionId, subscriptionId); + if (!subscriber || subscriber.activated || subscriber.phase === 'closed') return; + subscriber.activated = true; + this.#pump(subscriber); + } + + #abortSubscription(connectionId: string, subscriptionId: string): void { + const subscriber = this.#ownedSubscriber(connectionId, subscriptionId); + if (subscriber) this.#removeSubscriber(subscriber); + } + + #closeSubscription(connectionId: string, subscriptionId: string): boolean { + const connection = this.#connections.get(connectionId); + if (!connection) return false; + const subscriber = this.#subscriptions.get(subscriptionId); + if (!subscriber) return true; + if ( + subscriber.connectionId !== connectionId || + !connection.subscriptionIds.has(subscriptionId) + ) { + return false; + } + this.#removeSubscriber(subscriber); + return true; + } + + #closeConnection(connectionId: string): void { + const connection = this.#connections.get(connectionId); + if (!connection) return; + for (const subscriptionId of [...connection.subscriptionIds]) { + const subscriber = this.#ownedSubscriber(connectionId, subscriptionId); + if (subscriber) this.#removeSubscriber(subscriber); + } + this.#connections.delete(connectionId); + } + + #enqueue(subscriber: Subscriber, frame: SubscriptionFrame): void { + if (subscriber.phase !== 'open' || subscriber.terminalQueued) return; + let encodedBytes: number; + try { + encodedBytes = encodeProtocolFrame(frame).byteLength; + } catch { + this.#evictSlowSubscriber(subscriber); + return; + } + const terminalBytes = slowConsumerFrameBytes(subscriber, this.#hostEpoch); + if ( + subscriber.queue.length >= MAX_SUBSCRIBER_QUEUED_FRAMES - 1 || + subscriber.queuedBytes + encodedBytes + terminalBytes > MAX_SUBSCRIBER_QUEUED_BYTES + ) { + this.#evictSlowSubscriber(subscriber); + return; + } + subscriber.queue.push({ frame, encodedBytes }); + subscriber.queuedBytes += encodedBytes; + subscriber.nextSequence += 1; + if (subscriber.activated) this.#pump(subscriber); + } + + #evictSlowSubscriber(subscriber: Subscriber): void { + if (subscriber.phase !== 'open') return; + subscriber.phase = 'closing'; + if (subscriber.pumping) { + subscriber.sink.close(); + this.#removeSubscriber(subscriber); + return; + } + subscriber.queue = []; + subscriber.queuedBytes = 0; + subscriber.nextSequence = subscriber.lastFlushedSequence + 1; + const frame: SubscriptionFrame = { + kind: 'subscription.closed', + hostEpoch: this.#hostEpoch, + subscriptionId: subscriber.subscriptionId, + sequence: subscriber.nextSequence, + reason: 'slow_consumer', + }; + subscriber.nextSequence += 1; + subscriber.terminalQueued = true; + const encodedBytes = encodeProtocolFrame(frame).byteLength; + subscriber.queue.push({ frame, encodedBytes }); + subscriber.queuedBytes = encodedBytes; + if (subscriber.activated) this.#pump(subscriber); + } + + #enqueueAssistantDelta( + subscriber: Subscriber, + sessionId: string, + runId: string, + event: Extract, + kind: SessionAssistantDelta['kind'], + ): void { + let chunk = ''; + let rawBytes = 0; + let wireBytes = 0; + const frame = (text: string): SessionDeltaFrame => ({ + kind: 'subscription.session_delta', + hostEpoch: this.#hostEpoch, + subscriptionId: subscriber.subscriptionId, + sequence: subscriber.nextSequence, + sessionId, + delta: { kind, turnId: event.turnId, runId, messageId: event.messageId, text }, + }); + let wireLimit = wireTextByteLimit(frame('')); + for (const character of event.text) { + const rawCharacterBytes = Buffer.byteLength(character, 'utf8'); + const wireCharacterBytes = jsonStringContentBytes(character); + if ( + chunk.length > 0 && + (rawBytes + rawCharacterBytes > SESSION_LIVE_DELTA_MAX_BYTES || + wireBytes + wireCharacterBytes > wireLimit) + ) { + this.#enqueue(subscriber, frame(chunk)); + if (subscriber.phase !== 'open') return; + chunk = ''; + rawBytes = 0; + wireBytes = 0; + wireLimit = wireTextByteLimit(frame('')); + } + if (rawCharacterBytes > SESSION_LIVE_DELTA_MAX_BYTES || wireCharacterBytes > wireLimit) { + throw new Error('Session delta character exceeds the wire frame budget'); + } + chunk += character; + rawBytes += rawCharacterBytes; + wireBytes += wireCharacterBytes; + } + if (chunk.length > 0 && subscriber.phase === 'open') this.#enqueue(subscriber, frame(chunk)); + } + + #pump(subscriber: Subscriber): void { + if (subscriber.pumping || !subscriber.activated || subscriber.phase === 'closed') return; + const queued = subscriber.queue[0]; + if (!queued) return; + subscriber.pumping = true; + let flushed: Promise; + try { + flushed = subscriber.sink.send(queued.frame); + } catch { + this.#removeSubscriber(subscriber); + return; + } + void flushed.then( + () => { + subscriber.pumping = false; + if (subscriber.phase === 'closed') return; + if (subscriber.queue[0] === queued) { + subscriber.queue.shift(); + subscriber.queuedBytes -= queued.encodedBytes; + } + subscriber.lastFlushedSequence = queued.frame.sequence; + if (queued.frame.kind === 'subscription.closed') { + this.#removeSubscriber(subscriber); + return; + } + this.#pump(subscriber); + }, + () => this.#removeSubscriber(subscriber), + ); + } + + #removeSubscriber(subscriber: Subscriber): void { + if (subscriber.phase === 'closed') return; + subscriber.phase = 'closed'; + subscriber.queue = []; + subscriber.queuedBytes = 0; + const state = this.#sessions.get(subscriber.sessionId); + const removed = state?.subscribers.delete(subscriber.subscriptionId); + this.#subscriptions.delete(subscriber.subscriptionId); + this.#connections + .get(subscriber.connectionId) + ?.subscriptionIds.delete(subscriber.subscriptionId); + if (!this.#closed && state && removed && state.subscribers.size === 0) { + this.#scheduleInactiveStateCleanup(subscriber.sessionId, state); + } + } + + #ownedSubscriber(connectionId: string, subscriptionId: string): Subscriber | undefined { + const connection = this.#connections.get(connectionId); + if (!connection?.subscriptionIds.has(subscriptionId)) return; + const subscriber = this.#subscriptions.get(subscriptionId); + if (subscriber?.connectionId === connectionId) return subscriber; + } + + #scheduleInactiveStateCleanup(sessionId: string, state: SessionProjectionState): void { + if (this.#closed) return; + void this.sessionAdmission.enqueueDetached(sessionId, () => { + if ( + this.#sessions.get(sessionId) === state && + state.subscribers.size === 0 && + !state.terminalPublicationFence + ) { + this.#sessions.delete(sessionId); + } + }); + } + + async #readCanonicalProjection(sessionId: string): Promise { + const canonical = await this.#readCanonical(sessionId); + return canonical ? immutableClone(canonical) : null; + } + + #commitCanonical( + sessionId: string, + canonical: CanonicalSessionProjection, + ): { changed: boolean; state: SessionProjectionState; value: SessionContinuitySnapshot } { + let state = this.#sessions.get(sessionId); + if (state?.terminalPublicationFence) { + const rootTurn = requirePublicationFenceIdentity( + canonical, + sessionId, + state.terminalPublicationFence, + ); + if (isTerminalTurn(rootTurn)) { + return { + changed: false, + state, + value: snapshotValue(state.canonical, state.revision), + }; + } + } + if (!state) { + state = { canonical, revision: 1, subscribers: new Map() }; + this.#sessions.set(sessionId, state); + return { changed: true, state, value: snapshotValue(canonical, 1) }; + } + const changed = !isDeepStrictEqual(state.canonical, canonical); + if (changed) { + state.canonical = canonical; + state.revision += 1; + } + return { changed, state, value: snapshotValue(state.canonical, state.revision) }; + } + + #broadcastProjection(state: SessionProjectionState, snapshot: SessionContinuitySnapshot): void { + for (const subscriber of state.subscribers.values()) { + this.#enqueue(subscriber, { + kind: 'subscription.session_projection', + hostEpoch: this.#hostEpoch, + subscriptionId: subscriber.subscriptionId, + sequence: subscriber.nextSequence, + snapshot, + }); + } + } + + #runInSessionLane( + sessionId: string, + operation: () => Promise | T, + admission?: SessionAdmissionLease, + ): Promise { + return admission + ? this.sessionAdmission.runAdmitted(sessionId, admission, operation) + : this.sessionAdmission.run(sessionId, operation); + } +} + +function slowConsumerFrameBytes(subscriber: Subscriber, hostEpoch: string): number { + return encodeProtocolFrame({ + kind: 'subscription.closed', + hostEpoch, + subscriptionId: subscriber.subscriptionId, + sequence: subscriber.nextSequence + 1, + reason: 'slow_consumer', + }).byteLength; +} + +function immutableClone(value: T): T { + return deepFreeze(structuredClone(value)); +} + +function deepFreeze(value: T): T { + if (!value || typeof value !== 'object' || Object.isFrozen(value)) return value; + for (const child of Object.values(value)) deepFreeze(child); + return Object.freeze(value); +} + +function snapshotValue( + canonical: CanonicalSessionProjection, + projectionRevision: number, +): SessionContinuitySnapshot { + return immutableClone({ + schemaVersion: SESSION_CONTINUITY_SCHEMA_VERSION, + session: canonical.session, + projectionRevision, + rootTurn: canonical.rootTurn, + queue: canonical.queue, + }); +} + +function requirePublicationFenceIdentity( + canonical: CanonicalSessionProjection, + sessionId: string, + fence: TerminalPublicationFence, +): TurnSnapshot { + const rootTurn = canonical.rootTurn; + if ( + canonical.session.sessionId !== sessionId || + !rootTurn || + rootTurn.sessionId !== sessionId || + rootTurn.turnId !== fence.turnId || + rootTurn.runId !== fence.runId + ) { + throw new Error('Canonical Session projection identity does not match its publication fence'); + } + return rootTurn; +} + +function isTerminalTurn(turn: TurnSnapshot): boolean { + return turn.status === 'completed' || turn.status === 'failed' || turn.status === 'cancelled'; +} + +function wireTextByteLimit(frame: SessionDeltaFrame): number { + return RUNTIME_HOST_MAX_FRAME_BYTES - encodeProtocolFrame(frame).byteLength; +} + +function jsonStringContentBytes(value: string): number { + const encoded = JSON.stringify(value); + return Buffer.byteLength(encoded.slice(1, -1), 'utf8'); +} + +function projectToolEvent( + event: Exclude, +): SessionToolEvent { + const identity = { + id: event.id, + turnId: event.turnId, + ts: event.ts, + toolUseId: event.toolUseId, + }; + switch (event.type) { + case 'tool_start': + return { + type: event.type, + ...identity, + toolName: boundedUtf8(event.toolName, SESSION_TOOL_NAME_MAX_BYTES), + ...(event.operationId === undefined ? {} : { operationId: event.operationId }), + ...(event.activityKind === undefined ? {} : { activityKind: event.activityKind }), + ...(event.displayName === undefined + ? {} + : { displayName: boundedUtf8(event.displayName, SESSION_TOOL_NAME_MAX_BYTES) }), + ...(event.stepId === undefined ? {} : { stepId: event.stepId }), + }; + case 'tool_output_delta': + return { + type: event.type, + ...identity, + seq: event.seq, + stream: event.stream, + chunk: boundedUtf8(event.chunk, SESSION_LIVE_DELTA_MAX_BYTES), + redacted: event.redacted, + createdAt: event.createdAt, + }; + case 'tool_progress': + return { + type: event.type, + ...identity, + chunk: boundedUtf8( + typeof event.chunk === 'string' ? event.chunk : event.chunk.text, + SESSION_LIVE_DELTA_MAX_BYTES, + ), + }; + case 'tool_result': + return { + type: event.type, + ...identity, + ...(event.operationId === undefined ? {} : { operationId: event.operationId }), + status: event.isError ? 'errored' : 'completed', + ...(event.durationMs === undefined ? {} : { durationMs: event.durationMs }), + }; + } +} + +function boundedUtf8(value: string, maxBytes: number): string { + if (Buffer.byteLength(value, 'utf8') <= maxBytes) return value; + let bounded = ''; + let bytes = 0; + for (const character of value) { + const characterBytes = Buffer.byteLength(character, 'utf8'); + if (bytes + characterBytes > maxBytes) break; + bounded += character; + bytes += characterBytes; + } + return bounded; +} diff --git a/packages/runtime-host/src/server/session-continuity-service.ts b/packages/runtime-host/src/server/session-continuity-service.ts new file mode 100644 index 0000000000..51a1fa6992 --- /dev/null +++ b/packages/runtime-host/src/server/session-continuity-service.ts @@ -0,0 +1,21 @@ +import type { SubscriptionFrame } from '../protocol/index.js'; +import type { SessionContinuityOperationHandlerMap } from './operation-dispatcher.js'; + +export interface SessionContinuityFrameSink { + send(frame: SubscriptionFrame): Promise; + close(): void; +} + +export interface SessionContinuityConnection { + activate(subscriptionId: string): void; + abort(subscriptionId: string): void; + close(): void; +} + +export interface SessionContinuityService { + readonly handlers: SessionContinuityOperationHandlerMap; + attachConnection( + connectionId: string, + sink: SessionContinuityFrameSink, + ): SessionContinuityConnection; +} From 09065d374c48557d01324d1b40761400c67def8d Mon Sep 17 00:00:00 2001 From: Wang Date: Fri, 24 Jul 2026 01:02:47 +0800 Subject: [PATCH 2/3] test(runtime-host): cover linked child continuity --- .../__tests__/root-turn-coordinator.test.ts | 36 ++++++++++++++++++- 1 file changed, 35 insertions(+), 1 deletion(-) diff --git a/packages/runtime-host/src/__tests__/root-turn-coordinator.test.ts b/packages/runtime-host/src/__tests__/root-turn-coordinator.test.ts index 6a759c0143..e666f89149 100644 --- a/packages/runtime-host/src/__tests__/root-turn-coordinator.test.ts +++ b/packages/runtime-host/src/__tests__/root-turn-coordinator.test.ts @@ -162,6 +162,8 @@ test('hosted linked child roots share admission, message, terminal, and stop aut } | undefined; let initialEventCount = 0; + const childSink = new RecordingContinuitySink(); + let closeChildContinuity: (() => void) | undefined; const child = await manager.spawnChildSession(parent.id, { spawnedBy: { parentRunId: parentStarted.result.runId, @@ -170,8 +172,19 @@ test('hosted linked child roots share admission, message, terminal, and stop aut }, agentProfile: LOCAL_READ_AGENT_PROFILE, prompt: 'initial linked child', - onReady: (ready) => { + onReady: async (ready) => { initialReady = ready; + const childConnectionId = 'connection-linked-child'; + const childContinuity = requireContinuity(continuity); + const connection = childContinuity.attachConnection(childConnectionId, childSink); + const opened = await childContinuity.handlers['subscription.open']( + { sessionId: ready.childSessionId }, + operationContext(hostEpoch, acquireResidency, childConnectionId), + ); + assert.equal(opened.ok, true); + if (!opened.ok) throw new Error('Unable to subscribe to hosted linked child'); + connection.activate(opened.result.subscriptionId); + closeChildContinuity = () => connection.close(); }, onEvent: () => { initialEventCount += 1; @@ -186,6 +199,26 @@ test('hosted linked child roots share admission, message, terminal, and stop aut agentName: child.agentName, }); assert.equal(initialEventCount, child.eventCount); + assert.ok( + childSink.frames.some( + (frame) => + frame.kind === 'subscription.session_delta' && + frame.sessionId === child.childSessionId && + frame.delta.turnId === child.turnId && + frame.delta.runId === child.runId && + frame.delta.kind === 'text' && + frame.delta.text === 'linked child complete', + ), + ); + assert.ok( + childSink.frames.some( + (frame) => + frame.kind === 'subscription.session_projection' && + frame.snapshot.rootTurn?.turnId === child.turnId && + frame.snapshot.rootTurn.runId === child.runId && + frame.snapshot.rootTurn.status === 'completed', + ), + ); const initialAdmissions = await stores.agentRunStore.listRootTurnAdmissionsForRecovery( child.childSessionId, ); @@ -439,6 +472,7 @@ test('hosted linked child roots share admission, message, terminal, and stop aut }); await coordinator.close(); await messages.close(); + closeChildContinuity?.(); continuity.close(); } finally { await owner.close(); From b0c2e84403c58b9842a1f1f6214efa7e75fe168a Mon Sep 17 00:00:00 2001 From: Wang Date: Fri, 24 Jul 2026 18:01:10 +0800 Subject: [PATCH 3/3] fix(runtime-host): close continuity review gaps --- .../src/__tests__/connection-session.test.ts | 169 ++++++++++++++++++ .../src/__tests__/protocol.test.ts | 54 ++++++ .../__tests__/root-turn-coordinator.test.ts | 137 +++++++++++++- .../session-continuity-coordinator.test.ts | 148 ++++++++++++++- .../src/protocol/session-continuity.ts | 6 +- .../src/server/connection-session.ts | 1 - .../src/server/root-turn-coordinator.ts | 2 + .../server/session-continuity-coordinator.ts | 46 +++-- .../src/server/session-continuity-service.ts | 1 - .../src/__tests__/pi-agent-backend.test.ts | 34 ++++ packages/runtime/src/pi-agent-backend.ts | 14 +- 11 files changed, 582 insertions(+), 30 deletions(-) diff --git a/packages/runtime-host/src/__tests__/connection-session.test.ts b/packages/runtime-host/src/__tests__/connection-session.test.ts index 0b98813af8..b50e8cc664 100644 --- a/packages/runtime-host/src/__tests__/connection-session.test.ts +++ b/packages/runtime-host/src/__tests__/connection-session.test.ts @@ -14,6 +14,7 @@ import { connectRuntimeHost, type RuntimeHostConnection } from '../client/index. import { decodeHostFrame, RUNTIME_HOST_PROTOCOL_VERSION, + type HostFrame, type ResponseFrame, type TurnSnapshot, } from '../protocol/index.js'; @@ -23,6 +24,11 @@ import { createUnavailableDomainOperationHandlers, type OperationHandlerMap, } from '../server/operation-dispatcher.js'; +import { SessionAdmissionGate } from '../server/session-admission-gate.js'; +import { + type CanonicalSessionProjection, + SessionContinuityCoordinator, +} from '../server/session-continuity-coordinator.js'; import { BoundedSerialOutboundWriter, RuntimeHostOutboundQueueError, @@ -506,6 +512,121 @@ test('a sixty-fifth in-flight request tears down only the overflowing connection ); }); +test('evicting one slow subscription keeps sibling subscriptions and requests usable', async () => { + const pair = await openTransportPair(); + const coordinator = new SessionContinuityCoordinator( + 'host-epoch', + async (sessionId) => canonicalProjection(sessionId), + new SessionAdmissionGate(), + ); + const handlers: OperationHandlerMap = { + 'host.status': async () => ({ + ok: true, + result: { + hostEpoch: 'host-epoch', + state: 'ready', + connections: 1, + activeOperations: 1, + activeResidencies: 0, + }, + }), + ...createHandlers(async (input) => ({ + ok: true, + result: runningSnapshot(input.sessionId, input.turnId), + })), + ...coordinator.handlers, + }; + const session = new RuntimeHostConnectionSession({ + transport: pair.serverTransport, + connection: { + hostEpoch: 'host-epoch', + connectionId: 'shared-subscription-connection', + surface: 'tui', + principal: 'local_os_user', + }, + resolveHandlers: () => handlers, + resolveContinuity: () => coordinator, + beginOperation: async () => ({ + acquireResidency: () => ({ release() {} }), + seal() {}, + finish() {}, + }), + onTeardown() {}, + }); + const run = session.run(); + const slow = await openSubscription(pair.clientTransport, 'slow-session', 'open-slow'); + const sibling = await openSubscription(pair.clientTransport, 'sibling-session', 'open-sibling'); + const originalWrite = pair.serverTransport.writeEncoded.bind(pair.serverTransport); + const writeBlocked = deferred(); + const releaseWrite = deferred(); + pair.serverTransport.writeEncoded = async (encoded) => { + writeBlocked.resolve(); + await releaseWrite.promise; + return originalWrite(encoded); + }; + + try { + for (let index = 1; index <= 32; index += 1) { + await coordinator.acceptRuntimeEvent( + 'slow-session', + 'run-slow-session', + connectionTextEvent('slow-session', index), + ); + } + await withTimeout(writeBlocked.promise, 1_000, 'slow subscription never blocked in-flight'); + releaseWrite.resolve(); + + await coordinator.acceptRuntimeEvent( + 'sibling-session', + 'run-sibling-session', + connectionTextEvent('sibling-session', 1), + ); + await pair.clientTransport.write({ + requestId: 'status-after-eviction', + operation: 'host.status', + input: {}, + }); + + const observed: HostFrame[] = []; + while ( + !observed.some( + (frame) => + 'kind' in frame && + frame.kind === 'subscription.closed' && + frame.subscriptionId === slow.subscriptionId, + ) || + !observed.some( + (frame) => + 'kind' in frame && + frame.kind === 'subscription.session_delta' && + frame.subscriptionId === sibling.subscriptionId, + ) || + !observed.some((frame) => !('kind' in frame) && frame.requestId === 'status-after-eviction') + ) { + observed.push(decodeHostFrame(await pair.clientTransport.read(1_000))); + } + + const slowClosed = observed.find( + (frame) => + 'kind' in frame && + frame.kind === 'subscription.closed' && + frame.subscriptionId === slow.subscriptionId, + ); + assert.ok(slowClosed && 'kind' in slowClosed); + if (slowClosed && 'kind' in slowClosed && slowClosed.kind === 'subscription.closed') { + assert.equal(slowClosed.reason, 'slow_consumer'); + assert.equal(slowClosed.sequence, 2); + } + assert.equal(pair.serverTransport.socket.destroyed, false); + } finally { + releaseWrite.resolve(); + pair.serverTransport.writeEncoded = originalWrite; + pair.clientTransport.destroy(); + await Promise.allSettled([run, pair.close()]); + coordinator.close(); + } +}); + interface RuntimeHostTestFixture { connectClient(): Promise; endpoint: string; @@ -791,6 +912,54 @@ function runningSnapshot(sessionId: string, turnId: string): TurnSnapshot { }; } +async function openSubscription(transport: FramedTransport, sessionId: string, requestId: string) { + await transport.write({ + requestId, + operation: 'subscription.open', + input: { sessionId }, + }); + const response = decodeHostFrame(await transport.read(1_000)); + if ('kind' in response || response.operation !== 'subscription.open' || !response.ok) { + throw new Error(`Unable to open ${sessionId} subscription`); + } + return response.result; +} + +function canonicalProjection(sessionId: string): CanonicalSessionProjection { + return { + session: { + sessionId, + status: 'running', + createdAt: 1, + lastUsedAt: 1, + isArchived: false, + }, + rootTurn: { + sessionId, + turnId: `turn-${sessionId}`, + runId: `run-${sessionId}`, + status: 'running', + }, + queue: { + hostEpoch: 'host-epoch', + queueRevision: 0, + steering: [], + followup: [], + }, + }; +} + +function connectionTextEvent(sessionId: string, index: number) { + return { + type: 'text_delta' as const, + id: `event-${sessionId}-${index}`, + turnId: `turn-${sessionId}`, + ts: index, + messageId: `message-${sessionId}`, + text: `chunk-${index}`, + }; +} + async function waitForStatus( connection: RuntimeHostConnection, predicate: (status: Awaited>) => boolean, diff --git a/packages/runtime-host/src/__tests__/protocol.test.ts b/packages/runtime-host/src/__tests__/protocol.test.ts index d557a54c53..0ab001e5c4 100644 --- a/packages/runtime-host/src/__tests__/protocol.test.ts +++ b/packages/runtime-host/src/__tests__/protocol.test.ts @@ -1,6 +1,7 @@ import assert from 'node:assert/strict'; import { describe, test } from 'node:test'; import { MAX_ATTACHMENT_BYTES, MAX_ATTACHMENT_COUNT } from '@maka/core/attachments'; +import { TOOL_OUTPUT_DELTA_MAX_CHARS } from '@maka/core/events'; import { decodeClientFrame, decodeHostFrame, @@ -18,6 +19,7 @@ import { SESSION_CONTINUITY_SCHEMA_VERSION, SESSION_CONTINUITY_SNAPSHOT_MAX_BYTES, SESSION_LIVE_DELTA_MAX_BYTES, + SESSION_TOOL_OUTPUT_DELTA_MAX_BYTES, SESSION_TOOL_NAME_MAX_BYTES, TURN_MESSAGE_CONTENT_MAX_BYTES, TURN_MESSAGE_TEXT_MAX_BYTES, @@ -401,6 +403,58 @@ describe('Runtime Host bootstrap protocol', () => { ); }); + test('encodes maximum legal tool output as one bounded frame without identity loss', () => { + const chunks = [ + ['CJK', '界'.repeat(TOOL_OUTPUT_DELTA_MAX_CHARS)], + ['NUL', '\0'.repeat(TOOL_OUTPUT_DELTA_MAX_CHARS)], + ['lone surrogate', '\ud800'.repeat(TOOL_OUTPUT_DELTA_MAX_CHARS)], + ] as const; + for (const [label, chunk] of chunks) { + assert.ok( + Buffer.byteLength(chunk, 'utf8') <= SESSION_TOOL_OUTPUT_DELTA_MAX_BYTES, + `${label} exceeds the tool output raw-byte bound`, + ); + const frame = { + kind: 'subscription.session_event' as const, + hostEpoch: 'epoch-1', + subscriptionId: 'subscription-1', + sequence: 1, + sessionId: 'session-1', + runId: 'run-1', + event: { + type: 'tool_output_delta' as const, + id: `event-${label}`, + turnId: 'turn-1', + ts: 1, + toolUseId: 'tool-1', + seq: 23, + stream: 'stdout' as const, + chunk, + redacted: false, + createdAt: 2, + }, + }; + + const encoded = encodeProtocolFrame(frame); + assert.ok( + encoded.byteLength <= RUNTIME_HOST_MAX_FRAME_BYTES, + `${label} envelope exceeds the protocol frame limit`, + ); + const decodedFrames = new ProtocolFrameDecoder().push(encoded); + assert.equal(decodedFrames.length, 1); + const decoded = decodeHostFrame(decodedFrames[0]); + assert.ok('kind' in decoded); + if (!('kind' in decoded)) continue; + assert.equal(decoded.kind, 'subscription.session_event'); + if (decoded.kind !== 'subscription.session_event') continue; + assert.equal(decoded.event.type, 'tool_output_delta'); + if (decoded.event.type !== 'tool_output_delta') continue; + assert.equal(decoded.event.id, `event-${label}`); + assert.equal(decoded.event.seq, 23); + assert.equal(decoded.event.chunk, chunk); + } + }); + test('decodes split UTF-8 and multiple newline-delimited frames without an unbounded tail', () => { const decoder = new ProtocolFrameDecoder(); const wire = Buffer.from( diff --git a/packages/runtime-host/src/__tests__/root-turn-coordinator.test.ts b/packages/runtime-host/src/__tests__/root-turn-coordinator.test.ts index e666f89149..1996027274 100644 --- a/packages/runtime-host/src/__tests__/root-turn-coordinator.test.ts +++ b/packages/runtime-host/src/__tests__/root-turn-coordinator.test.ts @@ -22,6 +22,7 @@ import { type RootTurnAdmissionStore, } from '@maka/storage/execution-stores'; import { resolveStorageRoot, tryAcquireInteractiveRootOwner } from '@maka/storage/root-authority'; +import type { SubscriptionFrame } from '../protocol/index.js'; import { CanonicalSessionProjectionReader } from '../server/canonical-session-projection.js'; import type { RuntimeHostResidency } from '../server/host-kernel.js'; import { type HostMessageRootPort, HostMessageCoordinator } from '../server/message-coordinator.js'; @@ -29,6 +30,7 @@ import { RootAdmissionOwner } from '../server/root-admission-owner.js'; import { RootTurnCoordinator } from '../server/root-turn-coordinator.js'; import { SessionAdmissionGate } from '../server/session-admission-gate.js'; import { SessionContinuityCoordinator } from '../server/session-continuity-coordinator.js'; +import type { SessionContinuityFrameSink } from '../server/session-continuity-service.js'; const HOLD_EXTERNAL_PROMPT = 'hold external root before follow-up'; @@ -112,7 +114,9 @@ test('hosted linked child roots share admission, message, terminal, and stop aut const backends = new BackendRegistry(); const linkedBackends = new Map(); backends.register('fake', (context) => { - if (!context.header.subagentRuntime) return new FakeBackend(context); + if (!context.header.subagentRuntime) { + return new PermissionWaitingBackend(context.sessionId); + } const backend = new LinkedChildAuthorityBackend(context.sessionId); linkedBackends.set(context.sessionId, backend); return backend; @@ -140,6 +144,17 @@ test('hosted linked child roots share admission, message, terminal, and stop aut }, ); + const parentSink = new RecordingContinuitySink(); + const parentConnectionId = 'connection-waiting-parent'; + const parentConnection = continuity.attachConnection(parentConnectionId, parentSink); + const parentOpened = await continuity.handlers['subscription.open']( + { sessionId: parent.id }, + operationContext(hostEpoch, acquireResidency, parentConnectionId), + ); + assert.equal(parentOpened.ok, true); + if (!parentOpened.ok) return; + parentConnection.activate(parentOpened.result.subscriptionId); + const parentTurnId = randomUUID(); const parentStarted = await coordinator.handlers['turn.start']( { @@ -151,6 +166,12 @@ test('hosted linked child roots share admission, message, terminal, and stop aut ); assert.equal(parentStarted.ok, true); if (!parentStarted.ok) return; + await waitForContinuityFrame( + parentSink, + (frame) => + frame.kind === 'subscription.session_projection' && + frame.snapshot.session.status === 'waiting_for_user', + ); let initialReady: | { @@ -472,6 +493,7 @@ test('hosted linked child roots share admission, message, terminal, and stop aut }); await coordinator.close(); await messages.close(); + parentConnection.close(); closeChildContinuity?.(); continuity.close(); } finally { @@ -973,10 +995,14 @@ function requireContinuity( return continuity; } -function operationContext(hostEpoch: string, acquireResidency: () => RuntimeHostResidency) { +function operationContext( + hostEpoch: string, + acquireResidency: () => RuntimeHostResidency, + connectionId = 'connection-close-handoff', +) { return { hostEpoch, - connectionId: 'connection-close-handoff', + connectionId, surface: 'tui' as const, principal: 'local_os_user' as const, acquireResidency, @@ -1100,6 +1126,69 @@ class LinkedChildAuthorityBackend implements AgentBackend { } } +class PermissionWaitingBackend implements AgentBackend { + readonly kind = 'fake' as const; + private stopped = false; + private releaseWait: (() => void) | undefined; + + constructor(readonly sessionId: string) {} + + async *send(input: BackendSendInput): AsyncIterable { + this.stopped = false; + yield { + type: 'permission_request', + id: randomUUID(), + turnId: input.turnId, + ts: Date.now(), + kind: 'tool_permission', + requestId: randomUUID(), + toolUseId: randomUUID(), + toolName: 'Bash', + category: 'shell_safe', + reason: 'custom', + args: {}, + rememberForTurnAllowed: true, + }; + await new Promise((resolve) => { + this.releaseWait = resolve; + if (this.stopped) resolve(); + }); + yield { + type: 'abort', + id: randomUUID(), + turnId: input.turnId, + ts: Date.now(), + reason: 'user_stop', + }; + yield { + type: 'complete', + id: randomUUID(), + turnId: input.turnId, + ts: Date.now(), + stopReason: 'user_stop', + }; + } + + async stop(): Promise { + this.stopped = true; + this.releaseWait?.(); + } + + async respondToPermission(): Promise {} + + async dispose(): Promise { + this.releaseWait?.(); + } +} + +class RecordingContinuitySink implements SessionContinuityFrameSink { + readonly frames: SubscriptionFrame[] = []; + + async send(frame: SubscriptionFrame): Promise { + this.frames.push(frame); + } +} + function testTool(name: string): MakaTool { return { name, @@ -1109,3 +1198,45 @@ function testTool(name: string): MakaTool { impl: async () => ({ ok: true }), }; } + +async function completesWithin( + promise: Promise, + timeoutMs: number, + description: string, +): Promise { + let timeout: NodeJS.Timeout | undefined; + try { + return await Promise.race([ + promise, + new Promise((_resolve, reject) => { + timeout = setTimeout( + () => reject(new Error(`Timed out waiting for ${description}`)), + timeoutMs, + ); + }), + ]); + } finally { + if (timeout) clearTimeout(timeout); + } +} + +async function waitForContinuityFrame( + sink: RecordingContinuitySink, + predicate: (frame: SubscriptionFrame) => boolean, +): Promise { + return completesWithin( + new Promise((resolve) => { + const check = (): void => { + const frame = sink.frames.find(predicate); + if (frame) { + resolve(frame); + return; + } + setImmediate(check); + }; + check(); + }), + 5_000, + 'Session continuity frame', + ); +} diff --git a/packages/runtime-host/src/__tests__/session-continuity-coordinator.test.ts b/packages/runtime-host/src/__tests__/session-continuity-coordinator.test.ts index 9e40a577ae..0aa88c10c8 100644 --- a/packages/runtime-host/src/__tests__/session-continuity-coordinator.test.ts +++ b/packages/runtime-host/src/__tests__/session-continuity-coordinator.test.ts @@ -1,7 +1,16 @@ import assert from 'node:assert/strict'; import { setImmediate as delayImmediate } from 'node:timers/promises'; import test from 'node:test'; -import type { SubscriptionFrame } from '../protocol/index.js'; +import type { SessionEvent, SessionHeader } from '@maka/core'; +import { TOOL_OUTPUT_DELTA_MAX_CHARS } from '@maka/core/events'; +import { PermissionEngine, PiAgentBackend, type PiAgentTransport } from '@maka/runtime'; +import { + decodeHostFrame, + encodeProtocolFrame, + ProtocolFrameDecoder, + RUNTIME_HOST_MAX_FRAME_BYTES, + type SubscriptionFrame, +} from '../protocol/index.js'; import type { ConnectionContext } from '../server/operation-dispatcher.js'; import { type CanonicalSessionProjection, @@ -126,6 +135,113 @@ test('detached canonical refreshes coalesce before Store I/O', async () => { coordinator.close(); }); +test('in-flight canonical refresh observes an invalidation after its first read', async () => { + let projection = canonical(); + let reads = 0; + const firstRefreshRead = deferred(); + const coordinator = new SessionContinuityCoordinator( + HOST_EPOCH, + async () => { + reads += 1; + if (reads === 2) return firstRefreshRead.promise; + return projection; + }, + new SessionAdmissionGate(), + ); + const sink = new RecordingSink(); + const connection = coordinator.attachConnection('connection-1', sink); + const opened = await open(coordinator, 'connection-1'); + connection.activate(opened.subscriptionId); + + const stale = canonical({ lastUsedAt: 2 }); + coordinator.enqueueCanonicalRefresh(SESSION_ID); + await waitFor(() => reads === 2); + firstRefreshRead.resolve(stale); + projection = canonical({ lastUsedAt: 3 }); + coordinator.enqueueCanonicalRefresh(SESSION_ID); + + await waitFor(() => reads === 3 && sink.frames.length === 2); + assert.deepEqual( + sink.frames.map((frame) => + frame.kind === 'subscription.session_projection' + ? frame.snapshot.session.lastUsedAt + : undefined, + ), + [2, 3], + ); + coordinator.close(); +}); + +test('tool output preserves one domain event and one wire frame', async () => { + const coordinator = new SessionContinuityCoordinator( + HOST_EPOCH, + async () => canonical(), + new SessionAdmissionGate(), + ); + const sink = new RecordingSink(); + const connection = coordinator.attachConnection('connection-1', sink); + const opened = await open(coordinator, 'connection-1'); + connection.activate(opened.subscriptionId); + let nextId = 0; + let now = 1; + const source = '界'.repeat(TOOL_OUTPUT_DELTA_MAX_CHARS + 1); + const transport = { + async *send() { + yield { + type: 'tool_output_delta' as const, + toolUseId: 'tool-1', + stream: 'stdout' as const, + chunk: source, + }; + yield { type: 'complete' as const }; + }, + } satisfies PiAgentTransport; + const backend = new PiAgentBackend({ + sessionId: SESSION_ID, + header: piSessionHeader(), + appendMessage: async () => undefined, + permissionEngine: new PermissionEngine({ + newId: () => `permission-${++nextId}`, + now: () => now++, + }), + transport, + newId: () => `event-${++nextId}`, + now: () => now++, + }); + const produced: SessionEvent[] = []; + for await (const event of backend.send({ turnId: 'turn-1', text: 'inspect', context: [] })) { + produced.push(event); + } + const outputs = produced.filter( + (event): event is Extract => + event.type === 'tool_output_delta', + ); + assert.equal(outputs.length, 1); + const output = outputs[0]; + assert.ok(output); + assert.ok(output.chunk.length <= TOOL_OUTPUT_DELTA_MAX_CHARS); + assert.match(output.chunk, /\n\[内容已截断\]$/); + + await coordinator.acceptRuntimeEvent(SESSION_ID, 'run-1', output); + await waitFor(() => sink.frames.length === 1); + + const encoded = encodeProtocolFrame(sink.frames[0]!); + assert.ok(encoded.byteLength <= RUNTIME_HOST_MAX_FRAME_BYTES); + const decodedFrames = new ProtocolFrameDecoder().push(encoded); + assert.equal(decodedFrames.length, 1); + const decoded = decodeHostFrame(decodedFrames[0]); + assert.ok('kind' in decoded); + if (!('kind' in decoded)) return; + assert.equal(decoded.kind, 'subscription.session_event'); + if (decoded.kind !== 'subscription.session_event') return; + assert.equal(decoded.event.type, 'tool_output_delta'); + if (decoded.event.type !== 'tool_output_delta') return; + assert.equal(decoded.event.id, output.id); + assert.equal(decoded.event.seq, output.seq); + assert.equal(decoded.event.chunk, output.chunk); + coordinator.close(); +}); + test('reports a detached canonical publication failure to the Host lifecycle', async () => { let reads = 0; const observed = deferred(); @@ -186,7 +302,6 @@ test('slow subscriber receives a terminal eviction without delaying another subs slowConnection.activate(slow.subscriptionId); await waitFor(() => slowSink.frames.length === 1 && fastSink.frames.length === 32); - assert.equal(slowSink.closed, 0); assert.deepEqual(slowSink.frames[0], { kind: 'subscription.closed', hostEpoch: HOST_EPOCH, @@ -194,7 +309,6 @@ test('slow subscriber receives a terminal eviction without delaying another subs sequence: 1, reason: 'slow_consumer', }); - assert.equal(fastSink.closed, 0); assert.deepEqual( fastSink.frames.map((frame) => frame.sequence), Array.from({ length: 32 }, (_, index) => index + 1), @@ -204,15 +318,10 @@ test('slow subscriber receives a terminal eviction without delaying another subs class RecordingSink implements SessionContinuityFrameSink { readonly frames: SubscriptionFrame[] = []; - closed = 0; async send(frame: SubscriptionFrame): Promise { this.frames.push(frame); } - - close(): void { - this.closed += 1; - } } async function open(coordinator: SessionContinuityCoordinator, connectionId: string) { @@ -259,6 +368,29 @@ function canonical( }; } +function piSessionHeader(): SessionHeader { + return { + id: SESSION_ID, + workspaceRoot: '/tmp/maka', + cwd: '/tmp/maka', + createdAt: 1, + lastUsedAt: 1, + name: 'Pi continuity test', + titleIsManual: true, + isFlagged: false, + labels: [], + isArchived: false, + status: 'active', + hasUnread: false, + backend: 'pi-agent', + llmConnectionSlug: 'pi-agent', + connectionLocked: true, + model: 'pi-test', + permissionMode: 'execute', + schemaVersion: 1, + }; +} + function textEvent(index: number) { return { type: 'text_delta' as const, diff --git a/packages/runtime-host/src/protocol/session-continuity.ts b/packages/runtime-host/src/protocol/session-continuity.ts index 40574720a9..b9941edcb6 100644 --- a/packages/runtime-host/src/protocol/session-continuity.ts +++ b/packages/runtime-host/src/protocol/session-continuity.ts @@ -1,3 +1,4 @@ +import { TOOL_OUTPUT_DELTA_MAX_CHARS } from '@maka/core/events'; import { assertExactKeys, requireCount, @@ -17,6 +18,9 @@ import { decodeTurnSnapshot, type TurnSnapshot } from './turn.js'; export const SESSION_CONTINUITY_SCHEMA_VERSION = 1 as const; export const SESSION_CONTINUITY_SNAPSHOT_MAX_BYTES = 56 * 1024; export const SESSION_LIVE_DELTA_MAX_BYTES = 16 * 1024; +// Core emits at most 8,192 UTF-16 code units per tool output event. A code unit +// needs at most three UTF-8 bytes (an astral pair needs four bytes total). +export const SESSION_TOOL_OUTPUT_DELTA_MAX_BYTES = 3 * TOOL_OUTPUT_DELTA_MAX_CHARS; export const SESSION_TOOL_NAME_MAX_BYTES = 256; export const SESSION_SUBSCRIPTION_FRAME_MAX_BYTES = 64 * 1024 - 1; @@ -440,7 +444,7 @@ function decodeSessionToolEvent(value: unknown): SessionToolEvent { chunk: requireUtf8BoundedString( record.chunk, 'Session tool output chunk', - SESSION_LIVE_DELTA_MAX_BYTES, + SESSION_TOOL_OUTPUT_DELTA_MAX_BYTES, ), redacted: record.redacted, createdAt: requireCount(record.createdAt, 'Session tool output timestamp'), diff --git a/packages/runtime-host/src/server/connection-session.ts b/packages/runtime-host/src/server/connection-session.ts index c1d97ba779..6d8fd8b0a6 100644 --- a/packages/runtime-host/src/server/connection-session.ts +++ b/packages/runtime-host/src/server/connection-session.ts @@ -170,7 +170,6 @@ export class RuntimeHostConnectionSession { return Promise.reject(error); } }, - close: () => this.#teardown(), }); } return this.#continuity; diff --git a/packages/runtime-host/src/server/root-turn-coordinator.ts b/packages/runtime-host/src/server/root-turn-coordinator.ts index 85703fb823..e151f82f28 100644 --- a/packages/runtime-host/src/server/root-turn-coordinator.ts +++ b/packages/runtime-host/src/server/root-turn-coordinator.ts @@ -885,6 +885,8 @@ export class RootTurnCoordinator { } if (isRuntimeSessionTransientEvent(event)) { await this.continuity.acceptRuntimeEvent(input.sessionId, active.runId, event); + } else if (event.type === 'permission_request') { + this.continuity.enqueueCanonicalRefresh(input.sessionId); } } const snapshot = await this.readCanonicalSnapshot( diff --git a/packages/runtime-host/src/server/session-continuity-coordinator.ts b/packages/runtime-host/src/server/session-continuity-coordinator.ts index 89462a8fdb..804724e663 100644 --- a/packages/runtime-host/src/server/session-continuity-coordinator.ts +++ b/packages/runtime-host/src/server/session-continuity-coordinator.ts @@ -86,6 +86,11 @@ interface Subscriber { terminalQueued: boolean; } +interface PendingRefresh { + dirty: boolean; + inFlight: boolean; +} + export class SessionContinuityCoordinator implements SessionContinuityService { readonly handlers: SessionContinuityOperationHandlerMap = { 'subscription.open': async (input, context) => { @@ -108,7 +113,7 @@ export class SessionContinuityCoordinator implements SessionContinuityService { readonly #connections = new Map(); readonly #sessions = new Map(); readonly #subscriptions = new Map(); - readonly #pendingRefreshes = new Set(); + readonly #pendingRefreshes = new Map(); readonly #hostEpoch: string; readonly #readCanonical: ReadCanonicalSessionProjection; #closed = false; @@ -170,12 +175,27 @@ export class SessionContinuityCoordinator implements SessionContinuityService { /** Safe for synchronous commit hooks: this only schedules and coalesces lane work. */ enqueueCanonicalRefresh(sessionId: string): void { - if (this.#closed || this.#pendingRefreshes.has(sessionId)) return; - this.#pendingRefreshes.add(sessionId); + if (this.#closed) return; + const pending = this.#pendingRefreshes.get(sessionId); + if (pending) { + if (pending.inFlight) pending.dirty = true; + return; + } + const refresh: PendingRefresh = { dirty: false, inFlight: false }; + this.#pendingRefreshes.set(sessionId, refresh); void this.sessionAdmission - .enqueueDetached(sessionId, (lease) => this.refreshCanonical(sessionId, lease)) + .enqueueDetached(sessionId, async (lease) => { + refresh.inFlight = true; + await this.refreshCanonical(sessionId, lease); + if (!refresh.dirty) return; + refresh.dirty = false; + await this.refreshCanonical(sessionId, lease); + }) .then( - () => this.#pendingRefreshes.delete(sessionId), + () => { + this.#pendingRefreshes.delete(sessionId); + if (refresh.dirty) this.enqueueCanonicalRefresh(sessionId); + }, (error) => { this.#pendingRefreshes.delete(sessionId); this.onPublicationFailure(error); @@ -461,14 +481,10 @@ export class SessionContinuityCoordinator implements SessionContinuityService { #evictSlowSubscriber(subscriber: Subscriber): void { if (subscriber.phase !== 'open') return; subscriber.phase = 'closing'; - if (subscriber.pumping) { - subscriber.sink.close(); - this.#removeSubscriber(subscriber); - return; - } + const inFlight = subscriber.pumping ? subscriber.queue[0] : undefined; subscriber.queue = []; subscriber.queuedBytes = 0; - subscriber.nextSequence = subscriber.lastFlushedSequence + 1; + subscriber.nextSequence = (inFlight?.frame.sequence ?? subscriber.lastFlushedSequence) + 1; const frame: SubscriptionFrame = { kind: 'subscription.closed', hostEpoch: this.#hostEpoch, @@ -479,8 +495,12 @@ export class SessionContinuityCoordinator implements SessionContinuityService { subscriber.nextSequence += 1; subscriber.terminalQueued = true; const encodedBytes = encodeProtocolFrame(frame).byteLength; + if (inFlight) { + subscriber.queue.push(inFlight); + subscriber.queuedBytes += inFlight.encodedBytes; + } subscriber.queue.push({ frame, encodedBytes }); - subscriber.queuedBytes = encodedBytes; + subscriber.queuedBytes += encodedBytes; if (subscriber.activated) this.#pump(subscriber); } @@ -747,7 +767,7 @@ function projectToolEvent( ...identity, seq: event.seq, stream: event.stream, - chunk: boundedUtf8(event.chunk, SESSION_LIVE_DELTA_MAX_BYTES), + chunk: event.chunk, redacted: event.redacted, createdAt: event.createdAt, }; diff --git a/packages/runtime-host/src/server/session-continuity-service.ts b/packages/runtime-host/src/server/session-continuity-service.ts index 51a1fa6992..5e91769e55 100644 --- a/packages/runtime-host/src/server/session-continuity-service.ts +++ b/packages/runtime-host/src/server/session-continuity-service.ts @@ -3,7 +3,6 @@ import type { SessionContinuityOperationHandlerMap } from './operation-dispatche export interface SessionContinuityFrameSink { send(frame: SubscriptionFrame): Promise; - close(): void; } export interface SessionContinuityConnection { diff --git a/packages/runtime/src/__tests__/pi-agent-backend.test.ts b/packages/runtime/src/__tests__/pi-agent-backend.test.ts index 6ea9f2a5db..69ca7df3f2 100644 --- a/packages/runtime/src/__tests__/pi-agent-backend.test.ts +++ b/packages/runtime/src/__tests__/pi-agent-backend.test.ts @@ -4,6 +4,7 @@ import { tmpdir } from 'node:os'; import { join } from 'node:path'; import { describe, test } from 'node:test'; import type { BackendKind, SessionEvent, SessionHeader, StoredMessage } from '@maka/core'; +import { TOOL_OUTPUT_DELTA_MAX_CHARS } from '@maka/core/events'; import { createSessionStore } from '@maka/storage'; import { PermissionEngine } from '../permission-engine.js'; @@ -65,6 +66,39 @@ describe('PiAgentBackend skeleton', () => { ); }); + test('includes the truncation marker within the Core tool output limit', async () => { + const backend = new PiAgentBackend({ + sessionId: 'session-1', + header: header({ permissionMode: 'execute' }), + appendMessage: async () => undefined, + permissionEngine: new PermissionEngine({ newId: nextId('permission'), now: nextNow(2_500) }), + transport: frames([ + { + type: 'tool_output_delta', + toolUseId: 'tool-1', + stream: 'stdout', + chunk: '界'.repeat(TOOL_OUTPUT_DELTA_MAX_CHARS + 1), + }, + { type: 'complete' }, + ]), + newId: nextId('id'), + now: nextNow(2_600), + }); + + const events = await drain(backend.send({ turnId: 'turn-1', text: 'inspect', context: [] })); + const outputs = events.filter( + (event): event is Extract => + event.type === 'tool_output_delta', + ); + assert.equal(outputs.length, 1); + const output = outputs[0]; + assert.ok(output); + assert.ok(output.chunk.length <= TOOL_OUTPUT_DELTA_MAX_CHARS); + assert.match(output.chunk, /\n\[内容已截断\]$/); + assert.match(output.id, /^id-\d+$/); + assert.equal(output.seq, 1); + }); + test('normalizes noncanonical tool payloads before strict storage recovery', async () => { const root = await mkdtemp(join(tmpdir(), 'maka-pi-canonical-')); try { diff --git a/packages/runtime/src/pi-agent-backend.ts b/packages/runtime/src/pi-agent-backend.ts index 4fc9516839..dd3048f849 100644 --- a/packages/runtime/src/pi-agent-backend.ts +++ b/packages/runtime/src/pi-agent-backend.ts @@ -11,7 +11,11 @@ import type { ToolResultMessage, TokenUsageMessage, } from '@maka/core'; -import { computerUseApprovalSummary, decodeCanonicalToolResultContent } from '@maka/core'; +import { + computerUseApprovalSummary, + decodeCanonicalToolResultContent, + TOOL_OUTPUT_DELTA_MAX_CHARS, +} from '@maka/core'; import type { BackendSendInput, PermissionDecision } from '@maka/core/backend-types'; import { redactSecrets } from '@maka/core/redaction'; import { isToolCategory, type ToolCategory } from '@maka/core/permission'; @@ -685,9 +689,13 @@ function normalizeToolResultContent(content: unknown): ToolResultContent { } } -function redactBoundedText(text: string, maxChars = 8192): string { +const TRUNCATED_TEXT_MARKER = '\n[内容已截断]'; + +function redactBoundedText(text: string, maxChars = TOOL_OUTPUT_DELTA_MAX_CHARS): string { const redacted = redactSecrets(text); - return redacted.length > maxChars ? `${redacted.slice(0, maxChars)}\n[内容已截断]` : redacted; + if (redacted.length <= maxChars) return redacted; + const marker = TRUNCATED_TEXT_MARKER.slice(0, maxChars); + return `${redacted.slice(0, maxChars - marker.length)}${marker}`; } function redactUnknown(value: unknown): unknown {