diff --git a/src/main/native-chat/agent-session-wire/structured-agent-session-adapter.ts b/src/main/native-chat/agent-session-wire/structured-agent-session-adapter.ts index 9b5c9dcc40b..87a91bcf4d2 100644 --- a/src/main/native-chat/agent-session-wire/structured-agent-session-adapter.ts +++ b/src/main/native-chat/agent-session-wire/structured-agent-session-adapter.ts @@ -13,6 +13,7 @@ import type { } from '../../../shared/agent-session-journal-types' import type { AgentSessionProviderHandleLink } from '../../../shared/agent-session-provider-handle' import type { AgentSessionProcessIdentity } from '../../../shared/agent-session-record' +import type { StructuredAgentSessionEventSink } from './structured-agent-session-event-sink' /** What a reservation turns into once something is actually running under it: * the process the host can probe, and the provider handle it was minted with. */ @@ -36,6 +37,11 @@ export type StructuredAgentSessionAdapter = { identity: AgentSessionJournalIdentity fence: number spawnToken: string + /** Where to write everything the provider streams. Handed over here, not + * returned, because a provider starts streaming inside this call — before + * the journal exists. The sink buffers until it does. Optional so an + * adapter that only answers calls needs no journal at all. */ + events?: StructuredAgentSessionEventSink }): Promise dispatch(input: { sessionId: string diff --git a/src/main/native-chat/agent-session-wire/structured-agent-session-attach-flow.ts b/src/main/native-chat/agent-session-wire/structured-agent-session-attach-flow.ts index 4ce563ec9ce..d360fab80c9 100644 --- a/src/main/native-chat/agent-session-wire/structured-agent-session-attach-flow.ts +++ b/src/main/native-chat/agent-session-wire/structured-agent-session-attach-flow.ts @@ -24,6 +24,7 @@ import { } from './structured-agent-session-attach' import type { AgentSessionRecordStore } from '../../runtime/agent-session-record-store' import type { StructuredAgentSessionAdapter } from './structured-agent-session-adapter' +import type { StructuredAgentSessionEventSink } from './structured-agent-session-event-sink' export type AttachFlowInput = { store: AgentSessionRecordStore @@ -37,6 +38,9 @@ export type AttachFlowInput = { * sees the result, so no client can send against a session the host has not * finished publishing. */ onAttached: (attached: AttachedJournal) => void + /** Handed to the adapter so it can journal what the provider streams. The + * host owns it and binds it to the journal inside `onAttached`. */ + eventSink?: StructuredAgentSessionEventSink } export async function performAttach( @@ -117,7 +121,8 @@ async function acquireOwner( identity: journalIdentityFor(record.sessionId, input.params), fence, // Retries must recover the original reservation, not mint a second child. - spawnToken + spawnToken, + ...(input.eventSink ? { events: input.eventSink } : {}) }) if (record.lease.ownerProcess === null) { await input.store.commitProcessIdentity({ diff --git a/src/main/native-chat/agent-session-wire/structured-agent-session-event-sink.test.ts b/src/main/native-chat/agent-session-wire/structured-agent-session-event-sink.test.ts new file mode 100644 index 00000000000..421f7885a79 --- /dev/null +++ b/src/main/native-chat/agent-session-wire/structured-agent-session-event-sink.test.ts @@ -0,0 +1,116 @@ +import { describe, expect, it, vi } from 'vitest' +import type { + AgentJournalItemBody, + AgentJournalItemIdentity +} from '../../../shared/agent-session-journal-types' +import type { AgentSessionJournal } from '../agent-session-journal/journal-store' +import { + createDeferredStructuredAgentSessionEventSink, + type StructuredAgentSessionEventTarget +} from './structured-agent-session-event-sink' + +const BODY: AgentJournalItemBody = { + kind: 'message', + role: 'assistant', + blocks: [{ type: 'text', text: 'hi' }] +} + +function identity(ordinal: number): AgentJournalItemIdentity { + return { provider: 'codex', threadId: 'thread-1', turnId: 'turn-1', ordinal } +} + +type Recorded = { call: string; fence?: number; ordinal?: number } + +function target( + fence: number, + log: Recorded[], + failOn?: number +): StructuredAgentSessionEventTarget { + const journal = { + appendItem: vi.fn(async (id: AgentJournalItemIdentity, _body: AgentJournalItemBody) => { + const ordinal = id.provider === 'codex' ? id.ordinal : -1 + if (ordinal === failOn) { + throw new Error(`refused ${ordinal}`) + } + log.push({ call: 'appendItem', fence, ordinal }) + return { cursor: { epoch: 'e', sequence: ordinal } } + }), + appendTombstone: vi.fn(async (id: AgentJournalItemIdentity) => { + log.push({ + call: 'appendTombstone', + fence, + ordinal: id.provider === 'codex' ? id.ordinal : -1 + }) + return { epoch: 'e', sequence: 0 } + }) + } as unknown as AgentSessionJournal + return { journal, fence, publish: () => log.push({ call: 'publish', fence }) } +} + +describe('deferred structured agent-session event sink', () => { + it('buffers writes made before the journal exists and drains them in arrival order', async () => { + const log: Recorded[] = [] + const deferred = createDeferredStructuredAgentSessionEventSink() + + deferred.sink.appendItem(identity(0), BODY) + deferred.sink.appendItem(identity(1), BODY) + deferred.sink.publish() + expect(log).toEqual([]) + + deferred.bind(target(7, log)) + await deferred.drained() + + expect(log).toEqual([ + { call: 'appendItem', fence: 7, ordinal: 0 }, + { call: 'appendItem', fence: 7, ordinal: 1 }, + { call: 'publish', fence: 7 } + ]) + }) + + it('writes at the fence bound at submission time, so a rebind cannot backdate a write', async () => { + const log: Recorded[] = [] + const deferred = createDeferredStructuredAgentSessionEventSink() + deferred.bind(target(1, log)) + + deferred.sink.appendItem(identity(0), BODY) + // The re-attach that raised the fence. + deferred.bind(target(2, log)) + deferred.sink.appendItem(identity(1), BODY) + await deferred.drained() + + expect(log).toEqual([ + { call: 'appendItem', fence: 1, ordinal: 0 }, + { call: 'appendItem', fence: 2, ordinal: 1 } + ]) + }) + + it('drops buffered and later writes once closed, and refuses to rebind', async () => { + const log: Recorded[] = [] + const deferred = createDeferredStructuredAgentSessionEventSink() + + deferred.sink.appendItem(identity(0), BODY) + deferred.close() + deferred.bind(target(3, log)) + deferred.sink.appendItem(identity(1), BODY) + await deferred.drained() + + expect(log).toEqual([]) + }) + + it('reports a refused append and keeps draining the rest', async () => { + const log: Recorded[] = [] + const errors: unknown[] = [] + const deferred = createDeferredStructuredAgentSessionEventSink({ + onError: (error) => errors.push(error) + }) + deferred.bind(target(4, log, 0)) + + deferred.sink.appendItem(identity(0), BODY) + deferred.sink.appendTombstone(identity(1)) + await deferred.drained() + + expect(errors).toHaveLength(1) + expect((errors[0] as Error).message).toBe('refused 0') + expect(log).toEqual([{ call: 'appendTombstone', fence: 4, ordinal: 1 }]) + }) +}) diff --git a/src/main/native-chat/agent-session-wire/structured-agent-session-event-sink.ts b/src/main/native-chat/agent-session-wire/structured-agent-session-event-sink.ts new file mode 100644 index 00000000000..26cc28a9945 --- /dev/null +++ b/src/main/native-chat/agent-session-wire/structured-agent-session-event-sink.ts @@ -0,0 +1,115 @@ +// Where an adapter writes the provider events it did not synchronously return. +// +// A provider starts streaming the moment its process exists, and that moment is +// INSIDE `adapter.acquire` — before the journal is open and before the host has +// registered the session. So the sink an adapter receives is deferred: writes +// queue in arrival order and drain once the journal exists. +// +// One sink lives for the session, not for one acquisition: a re-attach opens a +// NEW journal object at a NEW fence, and rebinding re-points the same sink at +// it. That keeps a single identity for the adapter to hold across a re-acquire, +// and the adapter closes the superseded child, so nothing writes behind a fence +// that has already moved. + +import type { + AgentJournalItemBody, + AgentJournalItemIdentity +} from '../../../shared/agent-session-journal-types' +import type { AgentSessionJournal } from '../agent-session-journal/journal-store' + +/** The only journal surface an adapter gets: append and publish, no reads. An + * adapter that could read the journal would start reconciling against it, and + * reconciliation is the wire's job, not the provider's. */ +export type StructuredAgentSessionEventSink = { + appendItem(identity: AgentJournalItemIdentity, body: AgentJournalItemBody): void + appendTombstone(identity: AgentJournalItemIdentity): void + /** Fan the journal out to subscribers. Cheap and idempotent. */ + publish(): void +} + +export type StructuredAgentSessionEventTarget = { + journal: AgentSessionJournal + /** Fence the sink writes at. Fixed for the life of the sink: a new fence + * means a new acquisition, which gets its own sink. */ + fence: number + publish: () => void +} + +export type DeferredStructuredAgentSessionEventSink = { + sink: StructuredAgentSessionEventSink + /** Drains everything buffered so far, in order, then writes through. Called + * again on every re-attach to re-point the sink at the new journal. */ + bind(target: StructuredAgentSessionEventTarget): void + /** Permanently stops the sink. Queued writes are dropped rather than landing + * in a journal the host has already let go of. */ + close(): void + /** Resolves once every write queued so far has landed. */ + drained(): Promise +} + +type SinkOperation = (target: StructuredAgentSessionEventTarget) => Promise | void + +export function createDeferredStructuredAgentSessionEventSink( + deps: { + /** A rejected append. Unset drops it: throwing here would surface inside the + * provider's notification callback and take the connection down, and the + * lease already guarantees a stale writer's rows are refused. */ + onError?: (error: unknown) => void + } = {} +): DeferredStructuredAgentSessionEventSink { + let target: StructuredAgentSessionEventTarget | null = null + let closed = false + const buffered: SinkOperation[] = [] + let chain: Promise = Promise.resolve() + + const enqueue = (operation: SinkOperation): void => { + const bound = target + chain = chain.then(async () => { + try { + await operation(bound as StructuredAgentSessionEventTarget) + } catch (error) { + deps.onError?.(error) + } + }) + } + + const submit = (operation: SinkOperation): void => { + if (closed) { + return + } + if (!target) { + buffered.push(operation) + return + } + enqueue(operation) + } + + return { + sink: { + appendItem: (identity, body: AgentJournalItemBody) => { + submit((bound) => bound.journal.appendItem(identity, body, { fence: bound.fence })) + }, + appendTombstone: (identity) => { + submit((bound) => bound.journal.appendTombstone(identity, { fence: bound.fence })) + }, + publish: () => { + submit((bound) => bound.publish()) + } + }, + bind: (next) => { + if (closed) { + return + } + target = next + const pending = buffered.splice(0, buffered.length) + for (const operation of pending) { + enqueue(operation) + } + }, + close: () => { + closed = true + buffered.length = 0 + }, + drained: () => chain + } +} diff --git a/src/main/native-chat/agent-session-wire/structured-agent-session-host.ts b/src/main/native-chat/agent-session-wire/structured-agent-session-host.ts index 4bf5940ecea..bc20c8d2867 100644 --- a/src/main/native-chat/agent-session-wire/structured-agent-session-host.ts +++ b/src/main/native-chat/agent-session-wire/structured-agent-session-host.ts @@ -14,11 +14,6 @@ import type { } from '../../../shared/agent-session-journal-types' import type { AgentSessionOwnerProbe } from '../../../shared/agent-session-lease-adjudication' import type { AgentSessionRecord } from '../../../shared/agent-session-record' -import { - admitAgentSessionMutation, - agentSessionFingerprintConflict, - computeAgentSessionPayloadFingerprint -} from '../../../shared/agent-session-mutation-envelope' import type { AgentSessionAttachResult, AgentSessionCancelResult, @@ -37,6 +32,15 @@ import { readAgentSessionHistory } from './agent-session-history-page' import type { StructuredAgentSessionAdapter } from './structured-agent-session-adapter' import type { AgentSessionAttachParams } from './structured-agent-session-attach' import { performAttach } from './structured-agent-session-attach-flow' +import { + createDeferredStructuredAgentSessionEventSink, + type DeferredStructuredAgentSessionEventSink +} from './structured-agent-session-event-sink' +import { + admitAndRunAgentSessionMutation, + AGENT_SESSION_NOT_ATTACHED, + refuseAgentSessionMutation +} from './structured-agent-session-mutation-admission' import { createRestartReconciler } from './structured-agent-session-restart-reconcile' import { cancelPlan, @@ -45,13 +49,10 @@ import { setOptionPlan, type MutationPlan } from './structured-agent-session-mutation-plans' -import { runSettledAgentSessionMutation } from './structured-agent-session-operation-settlement' -import { resolveAgentSessionReplayOutcome } from './structured-agent-session-replay-outcome' import { AgentSessionSubscribers, type AgentSessionSubscriberEmit } from './structured-agent-session-subscribers' -import type { AgentSessionTurnContext } from './structured-agent-session-turns' export type StructuredAgentSessionCaller = { /** Stable per-client identity; scopes the operation ledger and records who @@ -68,6 +69,10 @@ export type StructuredAgentSessionHostDeps = { probeOwner?: (record: AgentSessionRecord) => Promise mintSpawnToken?: () => string now?: () => number + /** A journal append the provider streamed that could not be written. Unset + * drops it: the alternative is throwing inside the provider's notification + * callback, which would take the connection down over one lost row. */ + onEventSinkError?: (input: { sessionId: string; error: unknown }) => void } type SessionState = { @@ -76,19 +81,11 @@ type SessionState = { fence: number } -function refuse(refusal: AgentSessionWireRefusal): { ok: false; refusal: AgentSessionWireRefusal } { - return { ok: false, refusal } -} - -const NO_SESSION: AgentSessionWireRefusal = { - code: 'agent_session_ownership_unknown', - message: 'This host holds no attached session by that id.' -} - export class StructuredAgentSessionHost { private readonly sessions = new Map() private readonly subscribers = new AgentSessionSubscribers() private readonly chains = new Map>() + private readonly eventSinks = new Map() private readonly reconcileLeases: (sessionId: string) => Promise constructor(private readonly deps: StructuredAgentSessionHostDeps) { @@ -128,12 +125,14 @@ export class StructuredAgentSessionHost { const sessionId = params.envelope.sessionId const unreconciled = await this.reconcileLeases(sessionId) if (unreconciled) { - return refuse(unreconciled) + return refuseAgentSessionMutation(unreconciled) } - return performAttach({ + const eventSink = this.eventSinkFor(sessionId) + const attached = await performAttach({ store: this.deps.store, adapter: this.deps.adapter, journalRoot: this.deps.journalRoot, + eventSink: eventSink.sink, authority: { spawnToken: this.deps.mintSpawnToken?.() ?? randomUUID(), claimKeyId: this.deps.claimKeyId, @@ -154,11 +153,46 @@ export class StructuredAgentSessionHost { } else { this.subscribers.publish(sessionId, attached.journal) } + // Bound last: everything the provider streamed while the journal was + // opening lands after the attach result the client is about to read, + // never interleaved with it. + eventSink.bind({ + journal: attached.journal, + fence, + publish: () => this.subscribers.publish(sessionId, attached.journal) + }) } }) + // A refused attach that never reached a journal leaves the sink holding + // writes nothing will ever drain. + if (!attached.ok && !this.sessions.has(sessionId)) { + eventSink.close() + this.eventSinks.delete(sessionId) + } + return attached }) } + /** Settles every row the provider streamed. Callers that must read the + * journal immediately after provider activity await this first. */ + flushStreamedEvents(sessionId: string): Promise { + return this.eventSinks.get(sessionId)?.drained() ?? Promise.resolve() + } + + /** One sink per session, rebound on every attach. A single identity lets an + * adapter hold it across a re-acquire without noticing the new journal. */ + private eventSinkFor(sessionId: string): DeferredStructuredAgentSessionEventSink { + const existing = this.eventSinks.get(sessionId) + if (existing) { + return existing + } + const created = createDeferredStructuredAgentSessionEventSink({ + onError: (error) => this.deps.onEventSinkError?.({ sessionId, error }) + }) + this.eventSinks.set(sessionId, created) + return created + } + /** * Liveness of whatever currently owns the session. Never client-asserted: a * client that could claim the owner is gone could evict a live one. Without a @@ -222,105 +256,18 @@ export class StructuredAgentSessionHost { envelope: AgentSessionMutationEnvelope, plan: MutationPlan ): Promise> { - return this.serialize(envelope.sessionId, () => this.runMutation(caller, envelope, plan)) - } - - private async runMutation( - caller: StructuredAgentSessionCaller, - envelope: AgentSessionMutationEnvelope, - plan: MutationPlan - ): Promise> { - const session = this.sessions.get(envelope.sessionId) - const record = this.deps.store.getRecord(envelope.sessionId) - if (!session || !record) { - return refuse(NO_SESSION) - } - const hostFingerprint = computeAgentSessionPayloadFingerprint({ - method: plan.method, - sessionId: envelope.sessionId, - fields: plan.fields - }) - const conflict = agentSessionFingerprintConflict(envelope, hostFingerprint) - if (conflict) { - return refuse(conflict) - } - const admission = admitAgentSessionMutation({ - envelope, - hostFingerprint, - ledger: await this.deps.store.admitOperation({ + return this.serialize(envelope.sessionId, () => + admitAndRunAgentSessionMutation({ + store: this.deps.store, + adapter: this.deps.adapter, callerKey: caller.callerKey, - operationId: envelope.clientOperationId, - fingerprint: hostFingerprint, - now: this.now() - }), - lease: record.lease - }) - if (admission.decision === 'refused') { - return refuse(admission.refusal) - } - - const fence = record.lease.runtimeFence - const ctx = this.contextFor(caller, session, envelope.sessionId, fence) - if (admission.decision === 'replay') { - const replay = resolveAgentSessionReplayOutcome({ - operationId: envelope.clientOperationId, - outcome: admission.row.outcome, - reconstruct: () => plan.replay(ctx, admission.row.outcome) - }) - if (replay.decision === 'refuse') { - return refuse(replay.refusal) - } - if (replay.decision === 'replay') { - return { - ok: true, - replayed: true, - fence, - cursor: ctx.journal.cursor(), - value: replay.value - } - } - // Nothing durable landed, so this id is about to run for the first time. - // A refused call leaves its ledger row behind, and replaying past the - // lease and the fence would let a resend act under an owner that has since - // changed — so a first run pays the full admission price either way. - const rerun = admitAgentSessionMutation({ envelope, - hostFingerprint, - ledger: { decision: 'admit', row: admission.row }, - lease: record.lease + plan, + journal: this.sessions.get(envelope.sessionId)?.journal, + publish: (journal) => this.subscribers.publish(envelope.sessionId, journal), + now: () => this.now() }) - if (rerun.decision === 'refused') { - return refuse(rerun.refusal) - } - } - - const outcome = await runSettledAgentSessionMutation({ - store: this.deps.store, - callerKey: caller.callerKey, - envelope, - plan, - context: ctx - }) - return outcome.ok - ? { ok: true, replayed: false, fence, cursor: ctx.journal.cursor(), value: outcome.value } - : refuse(outcome.refusal) - } - - private contextFor( - caller: StructuredAgentSessionCaller, - session: SessionState, - sessionId: string, - fence: number - ): AgentSessionTurnContext { - return { - sessionId, - journal: session.journal, - fence, - adapter: this.deps.adapter, - resolvedBy: caller.callerKey, - publish: () => this.subscribers.publish(sessionId, session.journal), - now: () => this.now() - } + ) } /** @@ -350,7 +297,7 @@ export class StructuredAgentSessionHost { private requireSession(sessionId: string): SessionState { const session = this.sessions.get(sessionId) if (!session) { - throw new Error(NO_SESSION.code) + throw new Error(AGENT_SESSION_NOT_ATTACHED.code) } return session } diff --git a/src/main/native-chat/agent-session-wire/structured-agent-session-mutation-admission.ts b/src/main/native-chat/agent-session-wire/structured-agent-session-mutation-admission.ts new file mode 100644 index 00000000000..4b6ac6d5f05 --- /dev/null +++ b/src/main/native-chat/agent-session-wire/structured-agent-session-mutation-admission.ts @@ -0,0 +1,135 @@ +// The one route every mutating agent-session call takes: recompute the +// fingerprint, admit through the durable operation ledger, check the lease, then +// run the plan. It lives outside the host so that no method can quietly grow its +// own admission rules by sitting next to the call site. + +import { + admitAgentSessionMutation, + agentSessionFingerprintConflict, + computeAgentSessionPayloadFingerprint +} from '../../../shared/agent-session-mutation-envelope' +import type { + AgentSessionMutationEnvelope, + AgentSessionMutationResult, + AgentSessionWireRefusal +} from '../../../shared/agent-session-wire' +import type { AgentSessionRecordStore } from '../../runtime/agent-session-record-store' +import type { AgentSessionJournal } from '../agent-session-journal/journal-store' +import type { StructuredAgentSessionAdapter } from './structured-agent-session-adapter' +import type { MutationPlan } from './structured-agent-session-mutation-plans' +import { runSettledAgentSessionMutation } from './structured-agent-session-operation-settlement' +import { resolveAgentSessionReplayOutcome } from './structured-agent-session-replay-outcome' +import type { AgentSessionTurnContext } from './structured-agent-session-turns' + +export const AGENT_SESSION_NOT_ATTACHED: AgentSessionWireRefusal = { + code: 'agent_session_ownership_unknown', + message: 'This host holds no attached session by that id.' +} + +export function refuseAgentSessionMutation(refusal: AgentSessionWireRefusal): { + ok: false + refusal: AgentSessionWireRefusal +} { + return { ok: false, refusal } +} + +export type AgentSessionMutationRequest = { + store: AgentSessionRecordStore + adapter: StructuredAgentSessionAdapter + callerKey: string + envelope: AgentSessionMutationEnvelope + plan: MutationPlan + /** Journal of the attached session; absent when this host holds none. */ + journal: AgentSessionJournal | undefined + publish: (journal: AgentSessionJournal) => void + now: () => number +} + +export async function admitAndRunAgentSessionMutation( + request: AgentSessionMutationRequest +): Promise> { + const { envelope, plan, journal } = request + const record = request.store.getRecord(envelope.sessionId) + if (!journal || !record) { + return refuseAgentSessionMutation(AGENT_SESSION_NOT_ATTACHED) + } + const hostFingerprint = computeAgentSessionPayloadFingerprint({ + method: plan.method, + sessionId: envelope.sessionId, + fields: plan.fields + }) + const conflict = agentSessionFingerprintConflict(envelope, hostFingerprint) + if (conflict) { + return refuseAgentSessionMutation(conflict) + } + const admission = admitAgentSessionMutation({ + envelope, + hostFingerprint, + ledger: await request.store.admitOperation({ + callerKey: request.callerKey, + operationId: envelope.clientOperationId, + fingerprint: hostFingerprint, + now: request.now() + }), + lease: record.lease + }) + if (admission.decision === 'refused') { + return refuseAgentSessionMutation(admission.refusal) + } + + const fence = record.lease.runtimeFence + const context = turnContext(request, journal, fence) + if (admission.decision === 'replay') { + const replay = resolveAgentSessionReplayOutcome({ + operationId: envelope.clientOperationId, + outcome: admission.row.outcome, + reconstruct: () => plan.replay(context, admission.row.outcome) + }) + if (replay.decision === 'refuse') { + return refuseAgentSessionMutation(replay.refusal) + } + if (replay.decision === 'replay') { + return { ok: true, replayed: true, fence, cursor: journal.cursor(), value: replay.value } + } + // Nothing durable landed, so this id is about to run for the first time. A + // refused call leaves its ledger row behind, and replaying past the lease and + // the fence would let a resend act under an owner that has since changed — so + // a first run pays the full admission price either way. + const rerun = admitAgentSessionMutation({ + envelope, + hostFingerprint, + ledger: { decision: 'admit', row: admission.row }, + lease: record.lease + }) + if (rerun.decision === 'refused') { + return refuseAgentSessionMutation(rerun.refusal) + } + } + + const outcome = await runSettledAgentSessionMutation({ + store: request.store, + callerKey: request.callerKey, + envelope, + plan, + context + }) + return outcome.ok + ? { ok: true, replayed: false, fence, cursor: journal.cursor(), value: outcome.value } + : refuseAgentSessionMutation(outcome.refusal) +} + +function turnContext( + request: AgentSessionMutationRequest, + journal: AgentSessionJournal, + fence: number +): AgentSessionTurnContext { + return { + sessionId: request.envelope.sessionId, + journal, + fence, + adapter: request.adapter, + resolvedBy: request.callerKey, + publish: () => request.publish(journal), + now: () => request.now() + } +}