From 1b1e41bdd92d3afbd8af5f645fb34cc6684ec5bc Mon Sep 17 00:00:00 2001 From: Brennan Benson <79079362+brennanb2025@users.noreply.github.com> Date: Tue, 21 Jul 2026 12:39:38 -0700 Subject: [PATCH] fix(codex): preserve subagents across relay restarts --- src/main/agent-hooks/server.test.ts | 96 ++++++++++++++++++++++++ src/main/agent-hooks/server.ts | 43 ++++++++--- src/shared/agent-hook-listener.ts | 73 ++++++++++++++++++ src/shared/codex-subagent-roster.test.ts | 22 +++++- 4 files changed, 223 insertions(+), 11 deletions(-) diff --git a/src/main/agent-hooks/server.test.ts b/src/main/agent-hooks/server.test.ts index 275fc9b38b6..350da9e2eb2 100644 --- a/src/main/agent-hooks/server.test.ts +++ b/src/main/agent-hooks/server.test.ts @@ -24,6 +24,11 @@ import { AGENT_STATUS_STALE_AFTER_MS, parseAgentStatusPayload } from '../../shared/agent-status-types' +import { + createHookListenerState, + normalizeHookPayload, + type HookListenerState +} from '../../shared/agent-hook-listener' import { makePaneKey } from '../../shared/stable-pane-id' const { getCohortAtEmitMock, trackMock } = vi.hoisted(() => ({ @@ -89,6 +94,97 @@ afterEach(() => { }) describe('AgentHookServer listener replay', () => { + it('preserves Codex sibling and lead state across relay listener restarts', () => { + const server = new AgentHookServer() + const send = (state: HookListenerState, payload: Record): void => { + const event = normalizeHookPayload(state, 'codex', buildBody(payload), 'production') + if (!event) { + throw new Error('normalizeHookPayload rejected a known-good Codex fixture') + } + server.ingestRemote(event, 'conn-1') + } + + const initialRelay = createHookListenerState() + send(initialRelay, { + hook_event_name: 'UserPromptSubmit', + session_id: 'root-session', + prompt: 'coordinate reviewers', + model: 'gpt-5.4' + }) + for (const id of ['child-a', 'child-b']) { + send(initialRelay, { + hook_event_name: 'SubagentStart', + agent_id: id, + agent_type: 'reviewer', + model: 'gpt-5.4-mini' + }) + } + + const restartedRelay = createHookListenerState() + send(restartedRelay, { hook_event_name: 'SubagentStop', agent_id: 'child-a' }) + expect(server.getStatusSnapshot()[0]).toMatchObject({ + state: 'working', + model: 'gpt-5.4', + providerSession: { key: 'session_id', id: 'root-session' }, + subagents: [expect.objectContaining({ id: 'child-b' })] + }) + + send(restartedRelay, { + hook_event_name: 'Stop', + session_id: 'root-session', + model: 'gpt-5.4' + }) + const restartedAgain = createHookListenerState() + send(restartedAgain, { hook_event_name: 'SubagentStop', agent_id: 'child-b' }) + expect(server.getStatusSnapshot()[0]).toMatchObject({ + state: 'done', + model: 'gpt-5.4', + providerSession: { key: 'session_id', id: 'root-session' }, + subagents: undefined + }) + }) + + it('does not carry a remote Codex roster across connection cleanup', () => { + const server = new AgentHookServer() + const child = (id: string) => ({ + id, + state: 'working' as const, + startedAt: 1 + }) + server.ingestRemote( + { + paneKey: PANE, + hookEventName: 'SubagentStart', + toolAgentId: 'old-child', + payload: { + state: 'working', + prompt: '', + agentType: 'codex', + subagents: [child('old-child')] + } + }, + 'conn-1' + ) + + server.clearStatusEntriesForConnection('conn-1') + server.ingestRemote( + { + paneKey: PANE, + hookEventName: 'SubagentStart', + toolAgentId: 'new-child', + payload: { + state: 'working', + prompt: '', + agentType: 'codex', + subagents: [child('new-child')] + } + }, + 'conn-2' + ) + + expect(server.getStatusSnapshot()[0]?.subagents).toEqual([child('new-child')]) + }) + it('retains root Codex identity when relay child events omit it', () => { const server = new AgentHookServer() const providerSession = { key: 'session_id' as const, id: 'root-session' } diff --git a/src/main/agent-hooks/server.ts b/src/main/agent-hooks/server.ts index 12fc82090af..75171542a53 100644 --- a/src/main/agent-hooks/server.ts +++ b/src/main/agent-hooks/server.ts @@ -23,6 +23,7 @@ import { normalizeHookPayload, parseFormEncodedBody, readRequestBody, + reconcileRemoteCodexState, resolveHookSource, preparePendingGrokResultDiscovery, seedClaudeSubagentRosterFromSnapshots, @@ -824,27 +825,43 @@ export class AgentHookServer { this.onAgentStatus?.(enriched) return enriched } + const stateReconciledPayload = + payload.connectionId && payload.payload.agentType === 'codex' && payload.hookEventName + ? { + ...payload, + payload: reconcileRemoteCodexState( + this.state, + payload.paneKey, + payload.hookEventName, + payload.toolAgentId, + payload.payload, + previous?.payload + ) + } + : payload const previousCodexRoot = - payload.payload.agentType === 'codex' && - payload.toolAgentId && + stateReconciledPayload.payload.agentType === 'codex' && + stateReconciledPayload.toolAgentId && previous?.payload.agentType === 'codex' ? previous : undefined - const preservedProviderSession = !payload.providerSession + const preservedProviderSession = !stateReconciledPayload.providerSession ? previousCodexRoot?.providerSession : undefined - const preservedRootModel = !payload.payload.model ? previousCodexRoot?.payload.model : undefined + const preservedRootModel = !stateReconciledPayload.payload.model + ? previousCodexRoot?.payload.model + : undefined // Why: an SSH relay restart forgets root-only fields; child hooks must not erase durable resume/model identity. const rootContextPreservingPayload = preservedProviderSession || preservedRootModel ? { - ...payload, + ...stateReconciledPayload, ...(preservedProviderSession ? { providerSession: preservedProviderSession } : {}), payload: preservedRootModel - ? { ...payload.payload, model: preservedRootModel } - : payload.payload + ? { ...stateReconciledPayload.payload, model: preservedRootModel } + : stateReconciledPayload.payload } - : payload + : stateReconciledPayload const identity = resolveAgentStatusIdentity({ existing: previous ? { @@ -860,7 +877,7 @@ export class AgentHookServer { previous && shouldSuppressInheritedTerminalStatus({ inheritedFromActivePane: identity.inheritedFromActivePane, - incomingState: payload.payload.state + incomingState: rootContextPreservingPayload.payload.state }) ) { return previous @@ -1599,8 +1616,14 @@ export class AgentHookServer { if (entry.connectionId !== normalizedConnectionId) { continue } - if (this.deleteStatusEntry(paneKey)) { + const deleted = this.deleteStatusEntry(paneKey) + if (deleted) { statusChanged = true + if (deleted.payload.agentType === 'codex') { + // Why: a replacement remote process may reuse the pane; don't merge it with the lost connection's children. + this.state.codexSubagentRosterByPaneKey.delete(paneKey) + this.state.codexLeadStateByPaneKey.delete(paneKey) + } } } if (statusChanged) { diff --git a/src/shared/agent-hook-listener.ts b/src/shared/agent-hook-listener.ts index 5109c66f4af..b4664b2af7f 100644 --- a/src/shared/agent-hook-listener.ts +++ b/src/shared/agent-hook-listener.ts @@ -3064,6 +3064,79 @@ export function seedCodexStateFromSnapshot( } } +function codexLeadStateForHookEvent( + eventName: string | undefined +): CodexLeadTurnState['state'] | undefined { + if (eventName === 'Stop') { + return 'done' + } + if (eventName === 'PermissionRequest') { + return 'waiting' + } + if ( + eventName === 'SessionStart' || + eventName === 'UserPromptSubmit' || + eventName === 'PreToolUse' || + eventName === 'PostToolUse' + ) { + return 'working' + } + return undefined +} + +/** Why: relay restarts lose lead/roster state; merge child events into main's longer-lived cache. */ +export function reconcileRemoteCodexState( + state: HookListenerState, + paneKey: string, + eventName: string | undefined, + agentId: string | undefined, + payload: ParsedAgentStatusPayload, + previous: ParsedAgentStatusPayload | undefined +): ParsedAgentStatusPayload { + if (previous?.agentType === 'codex') { + seedCodexStateFromSnapshot(state, paneKey, previous) + } else { + seedCodexStateFromSnapshot(state, paneKey, payload) + } + + // Why: older relays send child identity without roster snapshots; keep their already-normalized aggregate authoritative. + if (agentId && !payload.subagents && !state.codexSubagentRosterByPaneKey.has(paneKey)) { + return payload + } + const roster = getOrCreateCodexSubagentRoster(state, paneKey) + if (payload.subagents) { + seedCodexSubagentRoster(roster, payload.subagents) + } + if (agentId) { + if (eventName === 'SubagentStop') { + finishCodexSubagent(roster, agentId) + } + } else { + const leadState = codexLeadStateForHookEvent(eventName) + if (eventName === 'SessionStart' || eventName === 'Stop') { + roster.clear() + } + if (leadState) { + const previousLead = state.codexLeadStateByPaneKey.get(paneKey) + state.codexLeadStateByPaneKey.set(paneKey, { + state: leadState, + model: payload.model ?? previousLead?.model + }) + } + } + + const lead = state.codexLeadStateByPaneKey.get(paneKey) + if (!lead) { + return payload + } + return { + ...payload, + state: codexRosterEffectiveState(roster, lead.state), + model: lead.model ?? payload.model, + subagents: codexRosterToSnapshots(roster) + } +} + function buildCodexStatusPayload( state: HookListenerState, eventName: unknown, diff --git a/src/shared/codex-subagent-roster.test.ts b/src/shared/codex-subagent-roster.test.ts index 42f074f0fd6..ebb947dedf6 100644 --- a/src/shared/codex-subagent-roster.test.ts +++ b/src/shared/codex-subagent-roster.test.ts @@ -1,5 +1,9 @@ import { describe, expect, it } from 'vitest' -import { AGENT_MODEL_MAX_LENGTH, AGENT_TYPE_MAX_LENGTH } from './agent-status-types' +import { + AGENT_MODEL_MAX_LENGTH, + AGENT_STATUS_MAX_SUBAGENTS, + AGENT_TYPE_MAX_LENGTH +} from './agent-status-types' import { codexRosterToSnapshots, finishCodexSubagent, @@ -39,4 +43,20 @@ describe('Codex subagent roster', () => { expect(roster.size).toBe(0) }) + + it('bounds live storage while admitting a replacement after one child stops', () => { + const roster: CodexSubagentRoster = new Map() + for (let index = 0; index <= AGENT_STATUS_MAX_SUBAGENTS; index += 1) { + upsertCodexSubagent(roster, `child-${index}`, { state: 'working' }, index) + } + + expect(roster.size).toBe(AGENT_STATUS_MAX_SUBAGENTS) + expect(roster.has(`child-${AGENT_STATUS_MAX_SUBAGENTS}`)).toBe(false) + + finishCodexSubagent(roster, 'child-0') + upsertCodexSubagent(roster, 'replacement', { state: 'working' }, 100) + + expect(roster.size).toBe(AGENT_STATUS_MAX_SUBAGENTS) + expect(roster.has('replacement')).toBe(true) + }) })