diff --git a/src/main/agent-hooks/server-pi-normalization.test.ts b/src/main/agent-hooks/server-pi-normalization.test.ts index 2aad513cb78..b351913de6f 100644 --- a/src/main/agent-hooks/server-pi-normalization.test.ts +++ b/src/main/agent-hooks/server-pi-normalization.test.ts @@ -450,3 +450,44 @@ describe('a Pi cancel beside a live child row', () => { expect(row?.mainAgent).toBeUndefined() }) }) + +describe('a Pi compaction that tails a finished reply', () => { + const turn_completed_at = 1_700_000_005_000 + let server: AgentHookServer + + beforeEach(async () => { + server = new AgentHookServer() + await server.start({ env: 'production' }) + }) + + afterEach(() => { + server.stop() + }) + + async function post(payload: Record): Promise { + const response = await postHookEvent(server, buildBody(payload), '/hook/pi') + expect(response.status).toBe(204) + } + + it('keeps the reply finish time from the compaction rows through settle', async () => { + await post({ hook_event_name: 'before_agent_start', prompt: 'summarize' }) + await post({ hook_event_name: 'message_end', role: 'assistant', text: 'Summary ready.' }) + await post({ hook_event_name: 'compaction_start', turn_completed_at }) + expect(server.getStatusSnapshot()[0]).toMatchObject({ + state: 'working', + prompt: 'summarize', + lastAssistantMessage: 'Summary ready.', + turnCompletedAt: turn_completed_at + }) + await post({ hook_event_name: 'agent_end', turn_completed_at }) + expect(server.getStatusSnapshot()[0]).toMatchObject({ + state: 'done', + turnCompletedAt: turn_completed_at + }) + + // The same prompt again is a new turn, not the finished one. + await post({ hook_event_name: 'before_agent_start', prompt: 'summarize' }) + expect(server.getStatusSnapshot()[0]?.state).toBe('working') + expect(server.getStatusSnapshot()[0]?.turnCompletedAt).toBeUndefined() + }) +}) diff --git a/src/main/pi/agent-status-compaction-source.ts b/src/main/pi/agent-status-compaction-source.ts new file mode 100644 index 00000000000..8d3d102c6eb --- /dev/null +++ b/src/main/pi/agent-status-compaction-source.ts @@ -0,0 +1,39 @@ +import type { PiAgentKind } from '../../shared/pi-agent-kind' + +/** Pi compacts between a finished reply's agent_end and its settle; stamping that compaction + * with the reply's finish lets the host announce it before the compaction ends. */ +export function getPiAgentStatusCompactionHandlerSourceLines(kind: PiAgentKind): string[] { + if (kind !== 'pi') { + return [] + } + + return [ + ' // Why: set only from a finished reply until its settle post, so a later /compact names no old turn.', + ' let replyFinishedAt: number | undefined', + '', + ' function takeCompletedReplyFields(): Record {', + ' const fields = replyFinishedAt === undefined ? {} : { turn_completed_at: replyFinishedAt }', + ' replyFinishedAt = undefined', + ' return fields', + ' }', + '', + ' // Why: a run cut off mid-tool, failed or aborted has no finished reply to announce.', + ' function recordReplyFinish(event): void {', + ' let reply: { stopReason?: unknown } | undefined', + ' for (const message of Array.isArray(event?.messages) ? event.messages : []) {', + " if (message?.role === 'assistant') reply = message", + ' }', + " const finished = reply?.stopReason === 'stop' || reply?.stopReason === 'length'", + ' replyFinishedAt = finished && !isOmpRuntime() ? Date.now() : undefined', + ' }', + '', + " onStatus('session_before_compact', (event, ctx) => {", + ' if (replyFinishedAt === undefined || isHeldByChildren()) return', + // Why: overflow recovery retries the turn and queued messages continue it after the compaction. + " if (event?.willRetry === true || event?.reason === 'overflow') return", + " if (typeof ctx?.hasPendingMessages === 'function' && ctx.hasPendingMessages()) return", + " post('compaction_start', { turn_completed_at: replyFinishedAt })", + ' })', + '' + ] +} diff --git a/src/main/pi/agent-status-compaction.test.ts b/src/main/pi/agent-status-compaction.test.ts new file mode 100644 index 00000000000..52db618e8b6 --- /dev/null +++ b/src/main/pi/agent-status-compaction.test.ts @@ -0,0 +1,112 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' + +import { + createAgentStatusExtensionHarness, + type AgentStatusExtensionHarness +} from './agent-status-extension-test-harness' +import { idle, startAsync } from './agent-status-subagent-event-fixtures' + +const busy = { isIdle: () => false } + +function reply(stopReason: string): { messages: unknown[] } { + return { messages: [{ role: 'user' }, { role: 'assistant', stopReason }] } +} + +// The posts after the run starts, reduced to the fields the host reads for a compaction tail. +function tail(harness: AgentStatusExtensionHarness): Record[] { + return harness.fetchMock.mock.calls + .map((call) => JSON.parse(String(call[1]?.body)).payload) + .filter((payload) => payload.hook_event_name !== 'agent_start') + .map(({ hook_event_name, turn_completed_at, is_idle }) => ({ + hook_event_name, + turn_completed_at, + is_idle + })) +} + +// Why: the post queue keeps only the latest pending snapshot, and Pi's compaction posts are minutes apart. +async function step( + harness: AgentStatusExtensionHarness, + name: string, + event?: unknown, + context?: { isIdle: () => boolean; hasPendingMessages?: () => boolean } +): Promise { + await harness.callHook(name, event, context) + await vi.advanceTimersByTimeAsync(0) +} + +describe('Pi compaction after a finished reply', () => { + beforeEach(() => vi.useFakeTimers()) + afterEach(() => vi.useRealTimers()) + + it('stamps a compaction before settle and the settled completion with the reply finish', async () => { + const harness = createAgentStatusExtensionHarness({ kind: 'pi' }) + await step(harness, 'agent_start') + await step(harness, 'agent_end', reply('stop'), busy) + await step(harness, 'session_before_compact', { reason: 'threshold', willRetry: false }) + await step(harness, 'session_compact', {}) + await step(harness, 'agent_settled', undefined, idle) + + const [start, settled] = tail(harness) + expect(start).toEqual({ + hook_event_name: 'compaction_start', + turn_completed_at: expect.any(Number) + }) + expect(settled).toEqual({ + hook_event_name: 'agent_end', + turn_completed_at: start?.turn_completed_at + }) + }) + + it.each([ + { + name: 'an overflow retry', + run: async (harness: AgentStatusExtensionHarness) => { + await step(harness, 'agent_end', reply('stop'), busy) + await step(harness, 'session_before_compact', { reason: 'overflow', willRetry: true }) + } + }, + { + name: 'an aborted reply', + run: async (harness: AgentStatusExtensionHarness) => { + await step(harness, 'agent_end', reply('aborted'), busy) + await step(harness, 'session_before_compact', { reason: 'threshold', willRetry: false }) + } + }, + { + name: 'a reply with queued messages that continue the run', + run: async (harness: AgentStatusExtensionHarness) => { + await step(harness, 'agent_end', reply('stop'), busy) + await step( + harness, + 'session_before_compact', + { reason: 'threshold', willRetry: false }, + { isIdle: () => false, hasPendingMessages: () => true } + ) + } + }, + { + name: 'a /compact after the reply settled', + run: async (harness: AgentStatusExtensionHarness) => { + await step(harness, 'agent_end', reply('stop'), busy) + await step(harness, 'agent_settled', undefined, idle) + await step(harness, 'session_before_compact', { reason: 'manual', willRetry: false }) + } + }, + { + name: 'a reply whose children still run', + run: async (harness: AgentStatusExtensionHarness) => { + startAsync(harness, 'run-scout', 'scout') + await step(harness, 'agent_end', reply('stop'), busy) + await step(harness, 'session_before_compact', { reason: 'threshold', willRetry: false }) + } + } + ])('posts no compaction tail for $name', async ({ run }) => { + const harness = createAgentStatusExtensionHarness({ kind: 'pi' }) + await step(harness, 'agent_start') + await run(harness) + await step(harness, 'session_compact', {}) + + expect(tail(harness).map((post) => post.hook_event_name)).not.toContain('compaction_start') + }) +}) diff --git a/src/main/pi/agent-status-handler-source.ts b/src/main/pi/agent-status-handler-source.ts index 2b5bf3f8f4b..7ab2e27aa81 100644 --- a/src/main/pi/agent-status-handler-source.ts +++ b/src/main/pi/agent-status-handler-source.ts @@ -5,6 +5,7 @@ import type { PiAgentKind } from '../../shared/pi-agent-kind' import { getOmpSessionOwnerHandlerSourceLines } from './omp-session-status-owner-source' import { getOmpPlanReviewHandlerSourceLines } from './omp-plan-review-source' import { getPiAgentStatusUiPromptHandlerSourceLines } from './agent-status-ui-prompt-source' +import { getPiAgentStatusCompactionHandlerSourceLines } from './agent-status-compaction-source' import { getPiSubagentRosterEventSourceLines, getPiSubagentRosterSetupSourceLines @@ -29,6 +30,7 @@ export function getPiAgentStatusHandlerSourceLines(kind: PiAgentKind): string[] " post('session_start')", ...(kind === 'pi' ? [ + ' replyFinishedAt = undefined', ' restoreParkedSubagents()', // Why: the host reads session_start as an idle session; children still hold this one. " if (lifecycleState.waiting) post('agent_start')" @@ -163,7 +165,9 @@ export function getPiAgentStatusHandlerSourceLines(kind: PiAgentKind): string[] ' lifecycleState.runGeneration += 1', // Why: a turn cannot begin under a dialog holding input focus, so this is the one // boundary that can recover a modal whose close never arrived. - ...(kind === 'pi' ? [' piUiPromptDepth = 0', ' piTurnInFlight = true'] : []), + ...(kind === 'pi' + ? [' piUiPromptDepth = 0', ' piTurnInFlight = true', ' replyFinishedAt = undefined'] + : []), " post('agent_start')", ' })', '', @@ -193,6 +197,7 @@ export function getPiAgentStatusHandlerSourceLines(kind: PiAgentKind): string[] ...approvalHandlers, ...(kind === 'prime-agent' ? [] : getOmpPlanReviewHandlerSourceLines()), ...getPiAgentStatusUiPromptHandlerSourceLines(kind), + ...getPiAgentStatusCompactionHandlerSourceLines(kind), ...modelSelectHandler, " // Why: capture the assistant's final text on each completed message", ' // so the dashboard preview reflects the most recent reply even before', @@ -245,7 +250,12 @@ export function getPiAgentStatusHandlerSourceLines(kind: PiAgentKind): string[] // and so starts clean on a pane that has not run a turn yet — that pane is idle, not busy. ...(kind === 'pi' ? [' piTurnInFlight = false'] : []), // Why: a run that ends because its session did has not completed; the host records it quietly. - " post('agent_end', final ? { session_boundary: true } : {}, final)", + ...(kind === 'pi' + ? [ + ' const replyFields = takeCompletedReplyFields()', + " post('agent_end', final ? { session_boundary: true } : replyFields, final)" + ] + : [" post('agent_end', final ? { session_boundary: true } : {}, final)"]), ' return true', ' }', '', @@ -286,6 +296,7 @@ export function getPiAgentStatusHandlerSourceLines(kind: PiAgentKind): string[] ' return', ' }', ' lifecycleState.endedRunGeneration = lifecycleState.runGeneration', + ...(kind === 'pi' ? [' recordReplyFinish(event)'] : []), ' if (isOmpRuntime()) {', ' postAgentEndOnce()', ' return', diff --git a/src/shared/agent-hook-listener/providers/pi-family-events.ts b/src/shared/agent-hook-listener/providers/pi-family-events.ts index d7bf5a4b43d..1618fdea30f 100644 --- a/src/shared/agent-hook-listener/providers/pi-family-events.ts +++ b/src/shared/agent-hook-listener/providers/pi-family-events.ts @@ -73,11 +73,14 @@ export function normalizePiCompatibleEvent( const isPiUiPrompt = agentType === 'pi' && (eventName === 'ui_prompt_start' || hookPayload.ui_prompt_active === true) const isPiUiPromptEnd = agentType === 'pi' && eventName === 'ui_prompt_end' + // Why: Pi posts this only for a compaction between a finished reply and its settle. + const isPiCompactionStart = agentType === 'pi' && eventName === 'compaction_start' let stateName = isPiCompatibleAsk || isOmpApprovalRequest || isOmpPlanReview ? 'blocked' : isOmpApprovalResolution || + isPiCompactionStart || eventName === 'before_agent_start' || eventName === 'agent_start' || eventName === 'tool_call' || @@ -125,6 +128,8 @@ export function normalizePiCompatibleEvent( lastAssistantMessage: snapshot.lastAssistantMessage, lastAssistantMessageIsToolOutput: snapshot.lastAssistantMessageIsToolOutput, subagents, + // Why: a compaction row and its settle carry one finish time, so the reply is announced once. + turnCompletedAt: hookPayload.turn_completed_at, // Why: the extension ends a run whose session was replaced; that is not a completed turn. ...(eventName === 'agent_end' && hookPayload.session_boundary === true ? { sessionBoundary: true }