From 5435cf7257baf970d8dad4e188f31f26e6f6dfcd Mon Sep 17 00:00:00 2001 From: Brennan Benson <79079362+brennanb2025@users.noreply.github.com> Date: Sun, 9 Aug 2026 22:42:02 -0700 Subject: [PATCH] Join structured attach teardown through journal bind --- .../structured-agent-session-host.ts | 16 ++--- .../structured-agent-session-task-queue.ts | 29 ++++++++ ...ructured-agent-session-integration.test.ts | 68 +++++++++++++++++++ 3 files changed, 104 insertions(+), 9 deletions(-) create mode 100644 src/main/native-chat/agent-session-wire/structured-agent-session-task-queue.ts 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 96f89d81aa6..ade4235bfd6 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 @@ -56,6 +56,7 @@ import { AgentSessionSubscribers, type AgentSessionSubscriberEmit } from './structured-agent-session-subscribers' +import { StructuredAgentSessionTaskQueue } from './structured-agent-session-task-queue' import { restoreStructuredAgentSessionRead } from './structured-agent-session-read-restore' export type StructuredAgentSessionCaller = { @@ -88,7 +89,7 @@ type SessionState = { export class StructuredAgentSessionHost { private readonly sessions = new Map() private readonly subscribers = new AgentSessionSubscribers() - private readonly chains = new Map>() + private readonly tasks = new StructuredAgentSessionTaskQueue() private readonly eventSinks = new Map() private readonly reconcileLeases: (sessionId: string) => Promise @@ -147,13 +148,7 @@ export class StructuredAgentSessionHost { /** Per-session single file. A rejected task must not poison the chain. */ private serialize(sessionId: string, task: () => Promise): Promise { - const prior = this.chains.get(sessionId) ?? Promise.resolve() - const next = prior.then(task, task) - this.chains.set( - sessionId, - next.catch(() => undefined) - ) - return next + return this.tasks.serialize(sessionId, task) } /** @@ -164,7 +159,7 @@ export class StructuredAgentSessionHost { caller: StructuredAgentSessionCaller, params: AgentSessionAttachParams ): Promise> { - return this.serialize(params.envelope.sessionId, async () => { + const attaching = this.serialize(params.envelope.sessionId, async () => { const sessionId = params.envelope.sessionId const unreconciled = await this.reconcileLeases(sessionId) if (unreconciled) { @@ -219,6 +214,7 @@ export class StructuredAgentSessionHost { } return attached }) + return this.tasks.trackAttach(attaching) } /** Settles every row the provider streamed. Callers that must read the @@ -229,6 +225,8 @@ export class StructuredAgentSessionHost { /** Settles final provider rows after every child has stopped producing them. */ async flushAllStreamedEvents(): Promise { + // An acquired provider can still be opening its journal; bind before draining. + await this.tasks.drainAttaches() await Promise.all([...this.eventSinks.values()].map((sink) => sink.drained())) } diff --git a/src/main/native-chat/agent-session-wire/structured-agent-session-task-queue.ts b/src/main/native-chat/agent-session-wire/structured-agent-session-task-queue.ts new file mode 100644 index 00000000000..88f74be759c --- /dev/null +++ b/src/main/native-chat/agent-session-wire/structured-agent-session-task-queue.ts @@ -0,0 +1,29 @@ +export class StructuredAgentSessionTaskQueue { + private readonly chains = new Map>() + private readonly attaching = new Set>() + + serialize(sessionId: string, task: () => Promise): Promise { + const prior = this.chains.get(sessionId) ?? Promise.resolve() + const next = prior.then(task, task) + this.chains.set( + sessionId, + next.catch(() => undefined) + ) + return next + } + + trackAttach(operation: Promise): Promise { + this.attaching.add(operation) + void operation.then( + () => this.attaching.delete(operation), + () => this.attaching.delete(operation) + ) + return operation + } + + async drainAttaches(): Promise { + while (this.attaching.size > 0) { + await Promise.allSettled(this.attaching) + } + } +} diff --git a/src/main/runtime/structured-agent-session-integration.test.ts b/src/main/runtime/structured-agent-session-integration.test.ts index ff754435a7d..2f71b1a1be6 100644 --- a/src/main/runtime/structured-agent-session-integration.test.ts +++ b/src/main/runtime/structured-agent-session-integration.test.ts @@ -15,6 +15,7 @@ import type { CodexAppServerConnectionHandlers, openCodexAppServerConnection } from '../codex/codex-app-server-connection' +import type { CodexStructuredSessionAdapter } from '../codex/codex-structured-session-adapter' import { computeAgentSessionPayloadFingerprint } from '../../shared/agent-session-mutation-envelope' import { STRUCTURED_AGENT_SESSION_RUNTIME_CAPABILITY } from '../../shared/protocol-version' import type { AgentJournalRenderItem } from '../../shared/agent-session-journal-types' @@ -565,6 +566,73 @@ describe('a structured codex session over agentSession.*', () => { ) ).toBe(false) }) + + it('joins an acquired attach through journal bind before draining final rows', async () => { + const host = await ensureStructuredAgentSessionHost({ + stateDirectory: root, + hostId: 'local', + claimKeyId: 'key-1', + resolveWorkspacePath: async (workspaceId) => `/repos/${workspaceId}`, + resolveCodexCommand: () => '/usr/local/bin/codex', + openCodexConnection: codex.openConnection + }) + const adapter = (host as unknown as { deps: { adapter: CodexStructuredSessionAdapter } }).deps + .adapter + const historyEntered = Promise.withResolvers() + const historyGate = Promise.withResolvers() + const originalHistoryFilePath = adapter.historyFilePath.bind(adapter) + vi.spyOn(adapter, 'historyFilePath').mockImplementation(async (input) => { + historyEntered.resolve() + await historyGate.promise + return originalHistoryFilePath(input) + }) + + const creating = ok<{ fence: number }>('agentSession.create', createIntentParams()) + await historyEntered.promise + codex.notify('turn/started', { threadId: THREAD, turn: { id: TURN } }) + codex.notify('item/started', { + threadId: THREAD, + turnId: TURN, + item: { type: 'agentMessage', id: 'item-bind-window', text: '' } + }) + codex.notify('item/agentMessage/delta', { + threadId: THREAD, + turnId: TURN, + itemId: 'item-bind-window', + delta: 'Buffered while the journal opens.' + }) + + let stopped = false + const stopping = stopStructuredAgentSessionRuntime().then(() => { + stopped = true + }) + await new Promise((resolve) => setImmediate(resolve)) + const waitedForJournalBind = !stopped + historyGate.resolve() + await creating + await stopping + expect(waitedForJournalBind).toBe(true) + + const identity = { + sessionId: SESSION, + workspaceId: WORKSPACE, + hostId: 'local', + agent: 'codex' as const, + providerHandle: { kind: 'codex' as const, threadId: THREAD } + } + const reopened = await openAgentSessionJournal({ + identity, + journalDir: journalDirectoryFor(root, identity) + }) + expect(reopened.snapshot().items.map(textOf)).toContain('Buffered while the journal opens.') + expect( + reopened + .snapshot() + .items.some( + (item) => item.body?.kind === 'status' && item.body.turnLifecycle?.state === 'running' + ) + ).toBe(false) + }) }) /** Every item the subscription has published, latest revision per id. */