Sfoglia il codice sorgente

fix(agent): preserve lifecycle recovery boundaries (PR3 round 2)

Hypatia May 3 mesi fa
parent
commit
d66c926d7a

+ 1 - 0
packages/compact/compact-basic/src/automatic.ts

@@ -45,6 +45,7 @@ export function registerAutomaticCompaction(
     _step: number,
     signal: AbortSignal,
   ) => {
+    if (signal.aborted) return
     try {
       const result = await service.compactIfNeeded(agent, 'pressure', signal)
       if (result !== null) logResult(result, 'post-step pressure')

+ 15 - 0
packages/compact/compact-basic/tests/compact-basic.spec.ts

@@ -848,6 +848,21 @@ describe('automatic listener and loader composition', () => {
     expect(compact.calls).toHaveLength(1)
   })
 
+  it('skips post-step pressure when the step signal is already aborted', async () => {
+    const ctx = createContext()
+    const compact = new TestCompactService(ctx, {
+      models: { [MODEL]: { thresholdRatio: 0.5, retainTokens: 18 } },
+    })
+    const pressured = conversation(4)
+    const compactIfNeeded = vi.spyOn(compact, 'compactIfNeeded')
+
+    await expect(postStep(ctx, agent(pressured, MODEL), AbortSignal.abort('step aborted')))
+      .resolves.toBeUndefined()
+
+    expect(compactIfNeeded).not.toHaveBeenCalled()
+    expect(pressured.events.some(event => event.type === 'compact/start')).toBe(false)
+  })
+
   it('warns and continues after operational failures, including non-Errors', async () => {
     const ctx = createContext()
     const warnings: string[] = []

+ 2 - 4
packages/core/agent-loop/src/loop.ts

@@ -650,7 +650,8 @@ async function runStep(
     session, turn, step, message.content, assembler.usage, chunkSeqs,
   )
 
-  // Tool execution stays sequential; recheck abort around each normalized result.
+  // Tool execution stays sequential; cancellation latches synthetic results for
+  // every remaining call while preserving one complete result batch.
   const toolCalls = message.content.filter(block => block.type === 'tool-call')
   // Buffer context until all results are appended to preserve call/result adjacency.
   const pendingContext: HookContext[] = []
@@ -693,9 +694,6 @@ async function runStep(
     if (signal.aborted) aborted = true
   }
 
-  /* v8 ignore next -- signal.reason always set by cancellation or disposal. */
-  if (aborted) throw new Error(String(signal.reason ?? 'aborted'))
-
   // Append buffered context after the complete result batch.
   for (const context of pendingContext) {
     agent.inject(context.content, { source: context.source })

+ 54 - 3
packages/core/agent-loop/tests/contract-regressions.spec.ts

@@ -156,7 +156,7 @@ describe('successful provider completion survives agent/step-result failure', ()
 })
 
 describe('abort during tool execution ends the turn', () => {
-  it('aborting the in-flight step inside a tool prevents both remaining tools and the next model step', async () => {
+  it('balances an aborted tool batch through context, steering, and post-step before closing', async () => {
     const adapter = new MockAdapter([
       // model asks for two tool calls in one step
       [
@@ -175,8 +175,12 @@ describe('abort during tool execution ends the turn', () => {
       name: 'aborter',
       description: '',
       parameters: {},
-      async execute() {
+      async execute(_args, exec) {
         executed.push('aborter')
+        exec.agent?.steer(
+          [{ type: 'text', text: 'steering before abort' }],
+          { source: { kind: 'plugin', plugin: 'abort-test' } },
+        )
         // Fire the in-flight step's AbortController directly (the loop registers
         // it on the agent). This is the bare step-abort path — distinct from
         // cancel(), which would also clear the inbox; here the subject is the
@@ -185,6 +189,13 @@ describe('abort during tool execution ends the turn', () => {
         return [{ type: 'text', text: 'done' }]
       },
     }))
+    ctx.on('tools/post-execute', async exec => ({
+      kind: 'accept',
+      additionalContext: {
+        content: [{ type: 'text', text: `context for ${exec.callId}` }],
+        source: { kind: 'plugin', plugin: 'abort-test' },
+      },
+    }))
     ctx.tools.register(defineTool({
       name: 'second',
       description: '',
@@ -196,13 +207,53 @@ describe('abort during tool execution ends the turn', () => {
     }))
 
     const reasons: TurnEndReason[] = []
-    ctx.on('session/event', (_s, event) => { if (event.type === 'turn/end') reasons.push(event.data.reason) })
+    const order: string[] = []
+    ctx.on('session/event', (session, event) => {
+      if (session !== agent.session) return
+      switch (event.type) {
+        case 'assistant/message': order.push('assistant/message'); break
+        case 'tool/call': order.push(`tool/call:${event.data.callId}`); break
+        case 'tool/result': {
+          const outcome = event.data.error?.code === 'ABORTED' ? 'synthetic-aborted' : 'real'
+          order.push(`tool/result:${event.data.callId}:${outcome}`)
+          break
+        }
+        case 'context/message': order.push('context/message'); break
+        case 'steering/message': order.push('steering/message'); break
+        case 'step/end': order.push('step/end'); break
+        case 'turn/end': {
+          reasons.push(event.data.reason)
+          order.push(`turn/end:${event.data.reason.kind}`)
+          break
+        }
+      }
+    })
+    let postSteps = 0
+    ctx.on('agent/post-step', (subject, turn, step, signal) => {
+      if (subject !== agent) return
+      postSteps += 1
+      expect({ turn, step, aborted: signal.aborted }).toEqual({ turn: 1, step: 1, aborted: true })
+      order.push('agent/post-step')
+    })
 
     send(agent, 'go')
     await waitForIdle(ctx, agent)
 
     expect(executed).toEqual(['aborter'])           // second tool never ran
     expect(adapter.requests).toHaveLength(1)        // no follow-up model call
+    expect(postSteps).toBe(1)
+    expect(order).toEqual([
+      'assistant/message',
+      'tool/call:c1',
+      'tool/result:c1:real',
+      'tool/call:c2',
+      'tool/result:c2:synthetic-aborted',
+      'context/message',
+      'steering/message',
+      'agent/post-step',
+      'step/end',
+      'turn/end:aborted',
+    ])
     expect(reasons).toEqual([{ kind: 'aborted', reason: 'user interrupt' }])
     const calls = agent.session.events.filter(event => event.type === 'tool/call')
     const results = agent.session.events.filter(event => event.type === 'tool/result')

+ 9 - 10
packages/llm/llm/src/index.ts

@@ -124,8 +124,9 @@ export class LlmService extends Service {
    * Final adapter boundary. It tags only failures from adapter selection,
    * synchronous dispatch, iterator construction, or iteration while preserving
    * the original Error object. Middleware outside this generator remains
-   * distinguishable as plugin work. Adapter cleanup is best-effort after an
-   * earlier failure or downstream close and never masks the winning error.
+   * distinguishable as plugin work. An iteration failure skips adapter cleanup
+   * so it cannot suppress the primary provider error. A downstream close awaits
+   * adapter cleanup, whose failures remain ordinary untagged work.
    */
   private async * adapterStream(options: GenerateOptions): AsyncGenerator<StreamChunk> {
     let iterator: AsyncIterator<StreamChunk>
@@ -137,6 +138,7 @@ export class LlmService extends Service {
     }
 
     let completed = false
+    let iterationFailed = false
     try {
       while (true) {
         let value: StreamChunk
@@ -148,6 +150,7 @@ export class LlmService extends Service {
           }
           value = item.value
         } catch (error: unknown) {
+          iterationFailed = true
           throw markLlmAdapterFailure(error)
         }
         // End the adapter-owned try before yielding: consumer/middleware
@@ -155,14 +158,10 @@ export class LlmService extends Service {
         yield value
       }
     } finally {
-      if (!completed) {
-        try {
-          const close = iterator.return?.bind(iterator)
-          if (close) await close()
-        } catch {
-          // Lookup and invocation are both adapter-owned cleanup following an
-          // existing failure/downstream close; neither can replace it.
-        }
+      // eslint-disable-next-line @typescript-eslint/no-unnecessary-condition -- the iteration catch sets its latch before entering finally.
+      if (!completed && !iterationFailed) {
+        const close = iterator.return?.bind(iterator)
+        if (close) await close()
       }
     }
   }

+ 70 - 25
packages/llm/llm/tests/service.spec.ts

@@ -72,11 +72,21 @@ describe('LlmService', () => {
     const original = new LlmError(`${field} getter failed`, 'RESULT_GETTER_FAILED')
     const result = field === 'done' ? {} : { done: false }
     Object.defineProperty(result, field, { get: () => { throw original } })
+    let cleanupLookups = 0
+    const iterator: AsyncIterator<StreamChunk> = {
+      next: () => Promise.resolve(result as unknown as IteratorResult<StreamChunk>),
+    }
+    Object.defineProperty(iterator, 'return', {
+      get: () => {
+        cleanupLookups += 1
+        throw new Error('return getter must not run after iteration fails')
+      },
+    })
     const adapter = new class extends LlmAdapter {
       stream(_options: GenerateOptions): AsyncIterable<StreamChunk> {
         return {
           [Symbol.asyncIterator](): AsyncIterator<StreamChunk> {
-            return { next: () => Promise.resolve(result as unknown as IteratorResult<StreamChunk>) }
+            return iterator
           },
         }
       }
@@ -94,6 +104,7 @@ describe('LlmService', () => {
 
     expect(caught).toBe(original)
     expect(isLlmAdapterFailure(caught)).toBe(true)
+    expect(cleanupLookups).toBe(0)
   })
 
   it.each(['dispatch', 'iterator'] as const)('tags synchronous adapter %s failures without replacing their Error', async (boundary) => {
@@ -119,7 +130,7 @@ describe('LlmService', () => {
     expect(isLlmAdapterFailure(caught)).toBe(true)
   })
 
-  it('tags adapter iteration failures without replacing the original Error or cleanup outcome', async () => {
+  it('propagates a rejected next promptly without awaiting a non-settling return', async () => {
     const original = new LlmError('provider failed', 'PROVIDER_FAILED')
     let cleanupCalls = 0
     const adapter = new class extends LlmAdapter {
@@ -130,7 +141,7 @@ describe('LlmService', () => {
               next: () => Promise.reject(original),
               return: () => {
                 cleanupCalls += 1
-                return Promise.reject(new Error('cleanup failed'))
+                return new Promise<IteratorResult<StreamChunk>>(() => {})
               },
             }
           },
@@ -141,31 +152,42 @@ describe('LlmService', () => {
     await ctx.plugin(LlmService)
     ctx.llm.registerAdapter(['test-model'], adapter)
 
-    let caught: unknown
-    try {
-      for await (const _chunk of ctx.llm.stream({ model: 'test-model', messages: [] })) { /* drain */ }
-    } catch (error: unknown) {
-      caught = error
-    }
+    const failure = (async (): Promise<unknown> => {
+      try {
+        for await (const _chunk of ctx.llm.stream({ model: 'test-model', messages: [] })) { /* drain */ }
+      } catch (error: unknown) {
+        return error
+      }
+      return new Error('expected adapter iteration to fail')
+    })()
+    let timer: ReturnType<typeof setTimeout> | undefined
+    const timeout = new Promise<Error>((resolve) => {
+      timer = setTimeout(() => { resolve(new Error('adapter failure did not settle promptly')) }, 100)
+    })
+    const caught = await Promise.race([failure, timeout])
+    if (timer !== undefined) clearTimeout(timer)
 
     expect(caught).toBe(original)
     expect(isLlmAdapterFailure(caught)).toBe(true)
-    expect(cleanupCalls).toBe(1)
+    expect(cleanupCalls).toBe(0)
   })
 
-  it('contains a throwing iterator.return getter after next fails without replacing the original Error', async () => {
-    const original = new LlmError('provider failed', 'PROVIDER_FAILED')
-    let cleanupLookups = 0
-    const iterator: AsyncIterator<StreamChunk> = { next: () => Promise.reject(original) }
-    Object.defineProperty(iterator, 'return', {
-      get: () => {
-        cleanupLookups += 1
-        throw new Error('return getter failed')
-      },
-    })
+  it('awaits one adapter return on downstream close and leaves its rejection unclassified', async () => {
+    const cleanup = new Error('cleanup failed')
+    let cleanupCalls = 0
     const adapter = new class extends LlmAdapter {
       stream(_options: GenerateOptions): AsyncIterable<StreamChunk> {
-        return { [Symbol.asyncIterator]: () => iterator }
+        return {
+          [Symbol.asyncIterator](): AsyncIterator<StreamChunk> {
+            return {
+              next: () => Promise.resolve({ done: false, value: SCRIPT[0]! }),
+              return: () => {
+                cleanupCalls += 1
+                return Promise.reject(cleanup)
+              },
+            }
+          },
+        }
       }
     }()
     const ctx = new Context()
@@ -174,14 +196,37 @@ describe('LlmService', () => {
 
     let caught: unknown
     try {
-      for await (const _chunk of ctx.llm.stream({ model: 'test-model', messages: [] })) { /* drain */ }
+      for await (const _chunk of ctx.llm.stream({ model: 'test-model', messages: [] })) break
     } catch (error: unknown) {
       caught = error
     }
 
-    expect(caught).toBe(original)
-    expect(isLlmAdapterFailure(caught)).toBe(true)
-    expect(cleanupLookups).toBe(1)
+    expect(caught).toBe(cleanup)
+    expect(isLlmAdapterFailure(caught)).toBe(false)
+    expect(cleanupCalls).toBe(1)
+  })
+
+  it('allows downstream close when the adapter iterator has no return method', async () => {
+    const adapter = new class extends LlmAdapter {
+      stream(_options: GenerateOptions): AsyncIterable<StreamChunk> {
+        return {
+          [Symbol.asyncIterator](): AsyncIterator<StreamChunk> {
+            return { next: () => Promise.resolve({ done: false, value: SCRIPT[0]! }) }
+          },
+        }
+      }
+    }()
+    const ctx = new Context()
+    await ctx.plugin(LlmService)
+    ctx.llm.registerAdapter(['test-model'], adapter)
+
+    let chunks = 0
+    for await (const _chunk of ctx.llm.stream({ model: 'test-model', messages: [] })) {
+      chunks += 1
+      break
+    }
+
+    expect(chunks).toBe(1)
   })
 
   it('normalizes and tags non-Error adapter failures once', async () => {