diff --git a/src/agent.ts b/src/agent.ts index d93e75a..60cd1d9 100644 --- a/src/agent.ts +++ b/src/agent.ts @@ -84,10 +84,10 @@ export class Agent { * Yields progress events; returns the model's final text answer. */ async *chat(userInput: string, signal?: AbortSignal): AsyncGenerator { - this.messages.push({ role: 'user', content: userInput }) - await this.context.maybeCompress(this.messages, this.llm, signal) - try { + this.messages.push({ role: 'user', content: userInput }) + await this.context.maybeCompress(this.messages, this.llm, signal) + // yield* forwards every event and evaluates to rounds()'s return value return yield* this.rounds(signal) } catch (e) { diff --git a/tests/abort.test.ts b/tests/abort.test.ts index 592f2e4..de0a5d0 100644 --- a/tests/abort.test.ts +++ b/tests/abort.test.ts @@ -90,6 +90,47 @@ test('abort propagates through the sub-agent tool to the sub-agent\'s LLM call', assert.equal(toolMsg!.content, '[interrupted]') }) +test('abort during context compression marks the user turn interrupted', async () => { + let compressionStarted!: () => void + const started = new Promise(resolve => { + compressionStarted = resolve + }) + const llm: LLMClient = { + model: 'slow-summarizer', + totalPromptTokens: 0, + totalCompletionTokens: 0, + usageSeen: true, + async *chat(_messages, _tools, signal) { + compressionStarted() + await new Promise((_, reject) => { + if (signal?.aborted) { + reject(new DOMException('Aborted', 'AbortError')) + return + } + signal?.addEventListener( + 'abort', + () => reject(new DOMException('Aborted', 'AbortError')), + { once: true }, + ) + }) + return new LLMResponse('unreachable') + }, + } + const agent = new Agent({ llm, maxContextTokens: 10 }) + for (let i = 0; i < 12; i++) { + agent.messages.push({ role: i % 2 === 0 ? 'user' : 'assistant', content: 'old context' }) + } + const ac = new AbortController() + + const turn = drain(agent.chat('cancel during compression', ac.signal)) + await started + ac.abort() + + await assert.rejects(turn, isAbort) + assert.equal(agent.messages.at(-1)?.role, 'assistant') + assert.equal(agent.messages.at(-1)?.content, '[interrupted by user]') +}) + test('pre-aborted signals make glob and grep reject immediately', async () => { const ac = new AbortController() ac.abort()