diff --git a/src/main/native-chat/agent-session-wire/structured-agent-session-attach-context.ts b/src/main/native-chat/agent-session-wire/structured-agent-session-attach-context.ts new file mode 100644 index 00000000000..dab5eb7db4a --- /dev/null +++ b/src/main/native-chat/agent-session-wire/structured-agent-session-attach-context.ts @@ -0,0 +1,31 @@ +// What attaching needs from the host, named explicitly. +// +// Passing the host itself would let this quietly grow new dependencies; an explicit context makes +// each one a deliberate addition and keeps the orchestration testable without constructing a host. + +import type { AgentSessionWireRefusal } from '../../../shared/agent-session-wire' +import type { AgentJournalResetReason } from '../../../shared/agent-session-journal-types' +import type { AgentSessionJournal } from '../agent-session-journal/journal-store' +import type { + StructuredAgentSessionHostDeps, + StructuredAgentSessionHostSession +} from './structured-agent-session-host-types' +import type { StructuredAgentSessionHostRuntimeState } from './structured-agent-session-host-runtime-state' + +export type StructuredAgentSessionAttachContext = { + deps: StructuredAgentSessionHostDeps + runtimeState: StructuredAgentSessionHostRuntimeState + sessions: Map + subscribers: { + reset: ( + sessionId: string, + journal: AgentSessionJournal, + reset: AgentJournalResetReason, + fence: number + ) => void + snapshot: (sessionId: string, journal: AgentSessionJournal, fence: number) => void + publish: (sessionId: string, journal: AgentSessionJournal) => void + } + reconcileLeases: (sessionId: string) => Promise + now: () => number +} diff --git a/src/main/native-chat/agent-session-wire/structured-agent-session-attach-orchestration.ts b/src/main/native-chat/agent-session-wire/structured-agent-session-attach-orchestration.ts new file mode 100644 index 00000000000..08368fb6aad --- /dev/null +++ b/src/main/native-chat/agent-session-wire/structured-agent-session-attach-orchestration.ts @@ -0,0 +1,76 @@ +// The attach orchestration, lifted out of the host. +// +// The host is a coordinator: it serializes work per session and tracks in-flight attaches. THIS is +// the part that decides what attaching means — reconcile the lease, resolve any pending recovery, +// bind an event sink, and publish the right thing to subscribers depending on whether the journal +// was recovered, re-fenced, or unchanged. Keeping it here means the host file stays a surface +// rather than an implementation, and this sequence can be read without scrolling past everything +// else a session can do. + +import { randomUUID } from 'node:crypto' +import type { + AgentSessionAttachResult, + AgentSessionMutationResult +} from '../../../shared/agent-session-wire' +import type { AgentSessionAttachParams } from './structured-agent-session-attach' +import { refuseAgentSessionMutation } from './structured-agent-session-mutation-admission' +import { performAttach } from './structured-agent-session-attach-flow' +import { pinnedAgentSessionLaunchEnv } from './structured-agent-session-launch-env' +import type { StructuredAgentSessionAttachContext } from './structured-agent-session-attach-context' + +export async function attachStructuredAgentSession( + context: StructuredAgentSessionAttachContext, + callerKey: string, + params: AgentSessionAttachParams +): Promise> { + const sessionId = params.envelope.sessionId + const unreconciled = await context.reconcileLeases(sessionId) + if (unreconciled) { + return refuseAgentSessionMutation(unreconciled) + } + await context.runtimeState.resolveRecovery(sessionId) + const eventSink = context.runtimeState.eventSinkFor(sessionId) + const attached = await performAttach({ + store: context.deps.store, + adapter: context.deps.adapter, + journalRoot: context.deps.journalRoot, + eventSink: eventSink.sink, + onAcquiring: () => eventSink.unbind(), + beforeJournalOpen: async () => { + eventSink.unbind() + await eventSink.drained() + }, + authority: { + spawnToken: context.deps.mintSpawnToken?.() ?? randomUUID(), + claimKeyId: context.deps.claimKeyId, + handoffOperationId: params.envelope.clientOperationId, + probe: await context.runtimeState.probeOwner(sessionId), + ...(await pinnedAgentSessionLaunchEnv(context.deps.resolveLaunchEnv, params)) + }, + callerKey, + params, + now: () => context.now(), + onAttached: (attached) => { + const fence = context.deps.store.getRecord(sessionId)?.lease.runtimeFence ?? 0 + const previousFence = context.sessions.get(sessionId)?.fence + context.sessions.set(sessionId, { journal: attached.journal, params, fence }) + if (attached.recovery) { + context.subscribers.reset(sessionId, attached.journal, attached.recovery.reset, fence) + } else if (previousFence !== undefined && previousFence !== fence) { + context.subscribers.snapshot(sessionId, attached.journal, fence) + } else { + context.subscribers.publish(sessionId, attached.journal) + } + eventSink.bind({ + journal: attached.journal, + fence, + publish: () => context.subscribers.publish(sessionId, attached.journal) + }) + } + }) + if (!attached.ok && !context.sessions.has(sessionId)) { + eventSink.close() + context.runtimeState.discardEventSink(sessionId) + } + return attached +} diff --git a/src/main/native-chat/agent-session-wire/structured-agent-session-eviction.test.ts b/src/main/native-chat/agent-session-wire/structured-agent-session-eviction.test.ts new file mode 100644 index 00000000000..b0be080fe0f --- /dev/null +++ b/src/main/native-chat/agent-session-wire/structured-agent-session-eviction.test.ts @@ -0,0 +1,65 @@ +import { describe, expect, it, vi } from 'vitest' +import { + evictStructuredAgentSession, + StructuredAgentSessionEvictionError, + STRUCTURED_AGENT_SESSION_EVICTION_STEPS, + type StructuredAgentSessionEvictionContext +} from './structured-agent-session-eviction' + +function context(): StructuredAgentSessionEvictionContext & { order: string[] } { + const order: string[] = [] + return { + order, + sessionId: 'session-1', + eventSink: { + unbind: vi.fn(() => order.push('unbind')), + drained: vi.fn(async () => { + order.push('drained') + }), + close: vi.fn(() => order.push('close')) + } as unknown as StructuredAgentSessionEvictionContext['eventSink'], + adapter: { + closeSession: vi.fn(async () => { + order.push('closeSession') + }) + } as unknown as StructuredAgentSessionEvictionContext['adapter'], + forget: vi.fn(() => order.push('forget')) + } +} + +describe('structured agent session eviction', () => { + // Order is the correctness property: draining after the child is gone discards rows the provider + // already sent, and forgetting before it stops strands the process. + it('drains what the session produced, then stops the child, then forgets it', async () => { + const ctx = context() + await evictStructuredAgentSession(ctx) + expect(ctx.order).toEqual(['unbind', 'drained', 'close', 'closeSession', 'forget']) + }) + + it('names every step, so a half-finished eviction says which one failed', () => { + expect(STRUCTURED_AGENT_SESSION_EVICTION_STEPS.map((step) => step.name)).toEqual([ + 'stop-publishing', + 'drain-published', + 'close-sink', + 'stop-provider-child', + 'forget-session' + ]) + }) + + // A child that refuses to stop must not also leave the session wired to a dead sink. + it('runs the remaining steps when one fails, and reports which failed', async () => { + const ctx = context() + ctx.adapter.closeSession = vi.fn(async () => { + throw new Error('child would not stop') + }) + + await expect(evictStructuredAgentSession(ctx)).rejects.toBeInstanceOf( + StructuredAgentSessionEvictionError + ) + await expect(evictStructuredAgentSession(ctx)).rejects.toMatchObject({ + step: 'stop-provider-child', + sessionId: 'session-1' + }) + expect(ctx.forget).toHaveBeenCalled() + }) +}) diff --git a/src/main/native-chat/agent-session-wire/structured-agent-session-eviction.ts b/src/main/native-chat/agent-session-wire/structured-agent-session-eviction.ts new file mode 100644 index 00000000000..1d32381ef66 --- /dev/null +++ b/src/main/native-chat/agent-session-wire/structured-agent-session-eviction.ts @@ -0,0 +1,73 @@ +// Releasing one structured session's resources. +// +// Teardown is a DATA list, not a method body, for the reason this file exists at all: the host +// tracked which sessions were live in a map, and tore them down at three unrelated call sites +// (app quit, handoff to a TUI, and error cleanup). Closing a chat was never wired to any of them, +// so a provider child outlived the chat that owned it for the whole app session. +// +// A list makes the next way a session can end cheap to support — add a caller, not a fourth +// teardown — and makes a half-finished eviction say WHICH step failed instead of failing opaquely. +// Order matters and is asserted by name in the tests: drain what the session already produced +// before stopping its child, and stop the child before forgetting it. Draining after the child is +// gone discards rows the provider already sent; forgetting first strands the process. + +import type { StructuredAgentSessionAdapter } from './structured-agent-session-adapter' +import type { DeferredStructuredAgentSessionEventSink } from './structured-agent-session-event-sink' + +export type StructuredAgentSessionEvictionContext = { + sessionId: string + eventSink: DeferredStructuredAgentSessionEventSink + adapter: StructuredAgentSessionAdapter + forget: () => void +} + +export type StructuredAgentSessionEvictionStep = { + name: string + run: (context: StructuredAgentSessionEvictionContext) => Promise | void +} + +export const STRUCTURED_AGENT_SESSION_EVICTION_STEPS: readonly StructuredAgentSessionEvictionStep[] = + [ + { name: 'stop-publishing', run: (context) => context.eventSink.unbind() }, + { name: 'drain-published', run: (context) => context.eventSink.drained() }, + { name: 'close-sink', run: (context) => context.eventSink.close() }, + { + name: 'stop-provider-child', + run: async (context) => { + await context.adapter.closeSession?.(context.sessionId) + } + }, + { name: 'forget-session', run: (context) => context.forget() } + ] + +export class StructuredAgentSessionEvictionError extends Error { + constructor( + readonly step: string, + readonly sessionId: string, + override readonly cause: unknown + ) { + super(`agent session eviction failed at step "${step}" for ${sessionId}`) + this.name = 'StructuredAgentSessionEvictionError' + } +} + +/** + * Runs every eviction step in order. A failing step is reported with its name and does NOT skip the + * rest: a child that refuses to stop must not also leave the session wired to a dead sink. + */ +export async function evictStructuredAgentSession( + context: StructuredAgentSessionEvictionContext, + steps: readonly StructuredAgentSessionEvictionStep[] = STRUCTURED_AGENT_SESSION_EVICTION_STEPS +): Promise { + let failure: StructuredAgentSessionEvictionError | null = null + for (const step of steps) { + try { + await step.run(context) + } catch (error) { + failure ??= new StructuredAgentSessionEvictionError(step.name, context.sessionId, error) + } + } + if (failure) { + throw failure + } +} diff --git a/src/main/native-chat/agent-session-wire/structured-agent-session-handoff-options.test.ts b/src/main/native-chat/agent-session-wire/structured-agent-session-handoff-options.test.ts index d0adf884b89..1457793ee34 100644 --- a/src/main/native-chat/agent-session-wire/structured-agent-session-handoff-options.test.ts +++ b/src/main/native-chat/agent-session-wire/structured-agent-session-handoff-options.test.ts @@ -299,4 +299,27 @@ describe('structured session handoff options', () => { expect(launchedOptions).toEqual([]) expect(store.getRecord(SESSION)?.lease.runtimeKind).toBe('native') }) + + // Closing a chat used to leave its provider child resident for the whole app session: the host's + // session map had no delete and the only teardown was app quit. + it('stops the provider child and forgets the session when the chat closes', async () => { + const adapterCloseSession = (host as unknown as { deps: { adapter: { closeSession: Mock } } }) + .deps.adapter.closeSession + expect(host.hasSession(SESSION)).toBe(true) + + await host.close(SESSION) + + expect(adapterCloseSession).toHaveBeenCalledWith(SESSION) + expect(host.hasSession(SESSION)).toBe(false) + }) + + it('is a no-op for a session it does not hold', async () => { + await host.close(SESSION) + const adapterCloseSession = (host as unknown as { deps: { adapter: { closeSession: Mock } } }) + .deps.adapter.closeSession + adapterCloseSession.mockClear() + + await expect(host.close(SESSION)).resolves.toBeUndefined() + expect(adapterCloseSession).not.toHaveBeenCalled() + }) }) 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 52b89dafb5a..67256242733 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 @@ -1,7 +1,6 @@ // Structured agent-session host: where the lease, journal, and provider adapter meet. // Mutations share one durable admission path and serialize per session. -import { randomUUID } from 'node:crypto' import type { AgentJournalMessageItem } from '../../../shared/agent-session-journal-types' import type { AgentSessionExecutionLocation } from '../../../shared/agent-session-record' import type { @@ -21,11 +20,9 @@ import type { AgentSessionWireRefusal } from '../../../shared/agent-session-wire' import type { AgentSessionAttachParams } from './structured-agent-session-attach' -import { performAttach } from './structured-agent-session-attach-flow' import { admitAndRunAgentSessionMutation, - AGENT_SESSION_NOT_ATTACHED, - refuseAgentSessionMutation + AGENT_SESSION_NOT_ATTACHED } from './structured-agent-session-mutation-admission' import { createRestartReconciler } from './structured-agent-session-restart-reconcile' import { @@ -48,8 +45,13 @@ import { type StructuredAgentSessionHostHandoff } from './structured-agent-session-host-handoff' import { StructuredAgentSessionHostRuntimeState } from './structured-agent-session-host-runtime-state' +import { attachStructuredAgentSession } from './structured-agent-session-attach-orchestration' +import { + evictStructuredAgentSession, + type StructuredAgentSessionEvictionContext +} from './structured-agent-session-eviction' +import type { StructuredAgentSessionAttachContext } from './structured-agent-session-attach-context' import { listStructuredAgentSessionTabs } from './structured-agent-session-host-tabs' -import { pinnedAgentSessionLaunchEnv } from './structured-agent-session-launch-env' import { StructuredAgentSessionReadableRestorer } from './structured-agent-session-readable-restorer' import type { StructuredAgentSessionCaller, @@ -120,6 +122,35 @@ export class StructuredAgentSessionHost { hasSession = (sessionId: string): boolean => this.sessions.has(sessionId) + /** The host's half of attaching, named so it cannot grow dependencies unnoticed. */ + private attachContext(): StructuredAgentSessionAttachContext { + return { + deps: this.deps, + runtimeState: this.runtimeState, + sessions: this.sessions, + subscribers: this.subscribers, + reconcileLeases: (sessionId) => this.reconcileLeases(sessionId), + now: () => this.now() + } + } + + /** Releases a session's resources without ending the conversation: the record and journal stay + * on disk, so the same session can be attached again. */ + close(sessionId: string): Promise { + return this.serialize(sessionId, async () => { + if (!this.sessions.has(sessionId)) { + return + } + const context: StructuredAgentSessionEvictionContext = { + sessionId, + eventSink: this.runtimeState.eventSinkFor(sessionId), + adapter: this.deps.adapter, + forget: () => this.sessions.delete(sessionId) + } + await evictStructuredAgentSession(context) + }) + } + supportsCreate = (location: AgentSessionExecutionLocation, agent: string): boolean => providerSupport.adapterSupportsCreate(this.deps.adapter, location, agent) @@ -145,58 +176,9 @@ export class StructuredAgentSessionHost { caller: StructuredAgentSessionCaller, params: AgentSessionAttachParams ): Promise> { - const attaching = this.serialize(params.envelope.sessionId, async () => { - const sessionId = params.envelope.sessionId - const unreconciled = await this.reconcileLeases(sessionId) - if (unreconciled) { - return refuseAgentSessionMutation(unreconciled) - } - await this.runtimeState.resolveRecovery(sessionId) - const eventSink = this.runtimeState.eventSinkFor(sessionId) - const attached = await performAttach({ - store: this.deps.store, - adapter: this.deps.adapter, - journalRoot: this.deps.journalRoot, - eventSink: eventSink.sink, - onAcquiring: () => eventSink.unbind(), - beforeJournalOpen: async () => { - eventSink.unbind() - await eventSink.drained() - }, - authority: { - spawnToken: this.deps.mintSpawnToken?.() ?? randomUUID(), - claimKeyId: this.deps.claimKeyId, - handoffOperationId: params.envelope.clientOperationId, - probe: await this.runtimeState.probeOwner(sessionId), - ...(await pinnedAgentSessionLaunchEnv(this.deps.resolveLaunchEnv, params)) - }, - callerKey: caller.callerKey, - params, - now: () => this.now(), - onAttached: (attached) => { - const fence = this.deps.store.getRecord(sessionId)?.lease.runtimeFence ?? 0 - const previousFence = this.sessions.get(sessionId)?.fence - this.sessions.set(sessionId, { journal: attached.journal, params, fence }) - if (attached.recovery) { - this.subscribers.reset(sessionId, attached.journal, attached.recovery.reset, fence) - } else if (previousFence !== undefined && previousFence !== fence) { - this.subscribers.snapshot(sessionId, attached.journal, fence) - } else { - this.subscribers.publish(sessionId, attached.journal) - } - eventSink.bind({ - journal: attached.journal, - fence, - publish: () => this.subscribers.publish(sessionId, attached.journal) - }) - } - }) - if (!attached.ok && !this.sessions.has(sessionId)) { - eventSink.close() - this.runtimeState.discardEventSink(sessionId) - } - return attached - }) + const attaching = this.serialize(params.envelope.sessionId, () => + attachStructuredAgentSession(this.attachContext(), caller.callerKey, params) + ) return this.tasks.trackAttach(attaching) } diff --git a/src/main/runtime/rpc/methods/structured-agent-session.test.ts b/src/main/runtime/rpc/methods/structured-agent-session.test.ts index 93c5c4d2480..c4a86a87a14 100644 --- a/src/main/runtime/rpc/methods/structured-agent-session.test.ts +++ b/src/main/runtime/rpc/methods/structured-agent-session.test.ts @@ -192,7 +192,9 @@ describe('capability gating', () => { for (const method of STRUCTURED_AGENT_SESSION_METHODS) { expect(names).toContain(method.name) } - expect(STRUCTURED_AGENT_SESSION_METHODS).toHaveLength(15) + // Bump deliberately: the whole agentSession.* surface is behind the structured capability, + // so an additive method is invisible to old clients and needs no protocol bump. + expect(STRUCTURED_AGENT_SESSION_METHODS).toHaveLength(16) }) it('hides the surface from a declared client that did not advertise it', async () => { diff --git a/src/main/runtime/rpc/methods/structured-agent-session.ts b/src/main/runtime/rpc/methods/structured-agent-session.ts index 3663e203dfa..b2e96cfb3ad 100644 --- a/src/main/runtime/rpc/methods/structured-agent-session.ts +++ b/src/main/runtime/rpc/methods/structured-agent-session.ts @@ -223,6 +223,16 @@ export const STRUCTURED_AGENT_SESSION_METHODS: RpcAnyMethod[] = [ params: CancelParams, handler: async (params, ctx) => requireHost(ctx).cancel(callerFor(ctx), params) }), + defineMethod({ + // Releasing a chat view, not ending a conversation: the record and journal stay on disk so the + // same session can be attached again. Only the provider child and the in-memory entry go. + name: 'agentSession.close', + params: OptionsParams, + handler: async (params, ctx) => { + await requireHost(ctx).close(params.sessionId) + return { ok: true as const } + } + }), defineMethod({ name: 'agentSession.respondToApproval', params: RespondParams, diff --git a/src/main/runtime/runtime-rpc.ts b/src/main/runtime/runtime-rpc.ts index 38a05c16b9b..286d6cd0194 100644 --- a/src/main/runtime/runtime-rpc.ts +++ b/src/main/runtime/runtime-rpc.ts @@ -191,6 +191,7 @@ const MOBILE_RPC_METHOD_ALLOWLIST = new Set([ 'agentSession.create', 'agentSession.createSupport', 'agentSession.cancel', + 'agentSession.close', 'agentSession.history', 'agentSession.options', 'agentSession.respondToApproval', diff --git a/tests/e2e/cross-version-wire/cross-version-agent-session-wire.unit.test.ts b/tests/e2e/cross-version-wire/cross-version-agent-session-wire.unit.test.ts index ea10aa02bed..63b396aee8d 100644 --- a/tests/e2e/cross-version-wire/cross-version-agent-session-wire.unit.test.ts +++ b/tests/e2e/cross-version-wire/cross-version-agent-session-wire.unit.test.ts @@ -43,6 +43,7 @@ const STRUCTURED_CALLS: { method: string; hostMethod: string | null }[] = [ { method: 'agentSession.ensure', hostMethod: 'attach' }, { method: 'agentSession.send', hostMethod: 'send' }, { method: 'agentSession.cancel', hostMethod: 'cancel' }, + { method: 'agentSession.close', hostMethod: 'close' }, { method: 'agentSession.respondToApproval', hostMethod: 'respondToPrompt' }, { method: 'agentSession.respondToQuestion', hostMethod: 'respondToPrompt' }, { method: 'agentSession.setOption', hostMethod: 'setOption' }, @@ -242,6 +243,7 @@ describe('cross-version structured agent sessions', () => { attach: vi.fn(async () => ({ ok: true, replayed: false, value: { sessionId: SESSION } })), send: vi.fn(async () => ({ ok: true, replayed: false })), cancel: vi.fn(async () => ({ ok: true, replayed: false })), + close: vi.fn(async () => undefined), respondToPrompt: vi.fn(async () => ({ ok: true, replayed: false })), setOption: vi.fn(async () => ({ ok: true, replayed: false })), requestHandoff: vi.fn(async () => ({ status: { owner: 'native' } })),