diff --git a/packages/ui/tui/src/index.ts b/packages/ui/tui/src/index.ts index 3fb0c4dcec..3f13fe89e4 100644 --- a/packages/ui/tui/src/index.ts +++ b/packages/ui/tui/src/index.ts @@ -2910,38 +2910,50 @@ export function createTuiChat( // Idle: the snapshot rides the prompt's admission transaction so a // blocking hook discards both together. let cleanedUp = false + let acceptedId: AgentMessageId | undefined + let acceptedContent: ContentBlock[] | undefined + const enqueued = new Map() + const discarded = new Set() const cleanup = (): void => { - // Each trigger detaches both listeners, so a second call needs a - // future third trigger; kept so adding one cannot double-release. + // Every completion path detaches all three listeners. Keep this + // idempotent so later cleanup paths cannot double-release them. /* v8 ignore next -- unreachable idempotence guard, see above */ if (cleanedUp) return cleanedUp = true + detachEnqueue() detachSubmit() detachDiscard() } + // send() snapshots input before publishing it, and publishes enqueue + // before returning its id. Capture that snapshot by id so admission can + // use exact reference identity without depending on caller-owned input. + const detachEnqueue = ctx.on('agent/inbox/enqueue', (subject, message) => { + if (subject === agent) enqueued.set(message.id, message.content) + }) // Prepended so this wrapper is outermost: it observes the admission // whether a downstream hook allows or blocks, and detaches either way. const detachSubmit = ctx.on('agent/prompt-submit', async (subject, submitted, _source, _signal, next) => { - if (subject !== agent || submitted !== content) return next() + if (subject !== agent || submitted !== acceptedContent) return next() cleanup() const decision = await next() if (decision.kind !== 'allow') return decision return { ...decision, additionalContexts: [...decision.additionalContexts ?? [], attachedContext] } }, { prepend: true }) - // Installed BEFORE followup(): admission runs synchronously inside it on - // the common path, and a listener registered after cleanup() already ran - // would never be released. Match on the `content` reference, not the - // returned id: an enqueue listener that synchronously cancels emits - // discard before followup() returns to assign the id, and content is the - // same reference send() carries onto the message (mirrors detachSubmit). + // Installed before followup(): an enqueue listener can synchronously + // cancel and discard before followup() returns its id. const detachDiscard = ctx.on('agent/inbox/discard', (subject, messages) => { - if (subject === agent && messages.some(message => message.content === content)) cleanup() + if (subject !== agent) return + for (const message of messages) discarded.add(message.id) + if (acceptedId !== undefined && discarded.has(acceptedId)) cleanup() }) // followup() accepts any typed input and contains listener failures; // this guards a future synchronous throw so the wrapper cannot leak. /* v8 ignore start -- future-proofing guard, see above */ try { - agent.followup({ content, source: { kind: 'user' } }) + acceptedId = agent.followup({ content, source: { kind: 'user' } }) + acceptedContent = enqueued.get(acceptedId) ?? content + detachEnqueue() + if (discarded.has(acceptedId)) cleanup() } catch (error: unknown) { cleanup() throw error diff --git a/packages/ui/tui/tests/tui.spec.ts b/packages/ui/tui/tests/tui.spec.ts index 1670ef90de..446d9be871 100644 --- a/packages/ui/tui/tests/tui.spec.ts +++ b/packages/ui/tui/tests/tui.spec.ts @@ -2039,16 +2039,21 @@ describe('pi-tui chat lifecycle and transcript', () => { appendUser(source, 'source background') }, }) - // Real send() emits agent/inbox/discard when an enqueue listener cancels - // synchronously, before followup() returns to assign the message id. This - // stub reproduces that timing: the wrapper must match on content, since id - // is not yet observable at discard time. + // Real send() publishes its snapshotted message, then an enqueue listener + // may synchronously cancel and discard it before followup() returns the + // already-assigned id. This stub reproduces that ordering. + const foreign = { ...result.agent, id: SessionId('foreign') } as unknown as Agent result.agent.followup = (input) => { result.agent.sent.push(input.content) - result.ctx.emit('agent/inbox/discard', result.agent, [{ - id: AgentMessageId('unassigned'), content: input.content, source: input.source, - }]) - return AgentMessageId('stub') + const message = { + id: AgentMessageId('stub'), + content: structuredClone(input.content), + source: structuredClone(input.source), + } + result.ctx.emit('agent/inbox/enqueue', foreign, message, 'queued') + result.ctx.emit('agent/inbox/enqueue', result.agent, message, 'queued') + result.ctx.emit('agent/inbox/discard', result.agent, [message]) + return message.id } result.terminal.send('@sync-source') @@ -2058,9 +2063,9 @@ describe('pi-tui chat lifecycle and transcript', () => { result.terminal.send('\r') await vi.waitFor(() => { expect(result.agent.sent).toHaveLength(1) }) - // The synchronous discard released both listeners despite the id being - // unassigned: replaying the prompt's admission attaches no stranded - // snapshot, and nothing leaks for the TUI lifetime. + // The synchronous discard released the listeners even though followup() + // had not returned the id yet: replaying the prompt's admission attaches + // no stranded snapshot, and nothing leaks for the TUI lifetime. const replay = await agentEvents(result.ctx, result.agent).waterfall( 'agent/prompt-submit', result.agent.sent[0]!, { kind: 'user' }, new AbortController().signal, () => Promise.resolve({ kind: 'allow' as const }),