fix(ai): restore ToolLoopAgent onFinish callbacks

.changeset/fuzzy-cars-notice.md

@@ -0,0 +1,5 @@
+---
+'ai': patch
+---
+
+Fix `ToolLoopAgent` `onFinish` callbacks on the v6 release line so per-call `onFinish` is accepted again and merged with the constructor-level callback for `generate()` and `stream()`.

packages/ai/src/agent/agent.ts

@@ -5,7 +5,10 @@ import { StreamTextTransform } from '../generate-text/stream-text';
 import { StreamTextResult } from '../generate-text/stream-text-result';
 import { ToolSet } from '../generate-text/tool-set';
 import { TimeoutConfiguration } from '../prompt/call-settings';
-import type { ToolLoopAgentOnStepFinishCallback } from './tool-loop-agent-settings';
+import type {
+  ToolLoopAgentOnFinishCallback,
+  ToolLoopAgentOnStepFinishCallback,
+} from './tool-loop-agent-settings';
 
 /**
  * Parameters for calling an agent.
@@ -61,6 +64,11 @@ export type AgentCallParameters<CALL_OPTIONS, TOOLS extends ToolSet = {}> = ([
      * Callback that is called when each step (LLM call) is finished, including intermediate steps.
      */
     onStepFinish?: ToolLoopAgentOnStepFinishCallback<TOOLS>;
+
+    /**
+     * Callback that is called when all steps are finished and the response is complete.
+     */
+    onFinish?: ToolLoopAgentOnFinishCallback<TOOLS>;
   };
 
 /**

packages/ai/src/agent/tool-loop-agent.test-d.ts

@@ -83,6 +83,20 @@ describe('ToolLoopAgent', () => {
 
       expectTypeOf<typeof output>().toEqualTypeOf<{ value: string }>();
     });
+
+    it('should allow onFinish', async () => {
+      const agent = new ToolLoopAgent({
+        model: new MockLanguageModelV3(),
+      });
+
+      await agent.generate({
+        prompt: 'Hello, world!',
+        onFinish: async event => {
+          const context: unknown = event.experimental_context;
+          context;
+        },
+      });
+    });
   });
 
   describe('stream', () => {
@@ -136,5 +150,19 @@ describe('ToolLoopAgent', () => {
         AsyncIterableStream<DeepPartial<{ value: string }>>
       >();
     });
+
+    it('should allow onFinish', async () => {
+      const agent = new ToolLoopAgent({
+        model: new MockLanguageModelV3(),
+      });
+
+      await agent.stream({
+        prompt: 'Hello, world!',
+        onFinish: async event => {
+          const context: unknown = event.experimental_context;
+          context;
+        },
+      });
+    });
   });
 });

packages/ai/src/agent/tool-loop-agent.test.ts

@@ -720,4 +720,226 @@ describe('ToolLoopAgent', () => {
       });
     });
   });
+
+  describe('onFinish', () => {
+    describe('generate', () => {
+      let mockModel: MockLanguageModelV3;
+
+      beforeEach(() => {
+        mockModel = new MockLanguageModelV3({
+          doGenerate: async () => {
+            return {
+              content: [{ type: 'text', text: 'reply' }],
+              finishReason: { unified: 'stop', raw: 'stop' },
+              usage: {
+                cachedInputTokens: undefined,
+                inputTokens: {
+                  total: 3,
+                  noCache: 3,
+                  cacheRead: undefined,
+                  cacheWrite: undefined,
+                },
+                outputTokens: {
+                  total: 10,
+                  text: 10,
+                  reasoning: undefined,
+                },
+              },
+              warnings: [],
+            };
+          },
+        });
+      });
+
+      it('should call onFinish from constructor', async () => {
+        const calls: string[] = [];
+
+        const agent = new ToolLoopAgent({
+          model: mockModel,
+          onFinish: async () => {
+            calls.push('constructor');
+          },
+        });
+
+        await agent.generate({
+          prompt: 'Hello, world!',
+        });
+
+        expect(calls).toMatchInlineSnapshot(`
+          [
+            "constructor",
+          ]
+        `);
+      });
+
+      it('should call onFinish from generate method', async () => {
+        const calls: string[] = [];
+
+        const agent = new ToolLoopAgent({
+          model: mockModel,
+        });
+
+        await agent.generate({
+          prompt: 'Hello, world!',
+          onFinish: async () => {
+            calls.push('method');
+          },
+        });
+
+        expect(calls).toMatchInlineSnapshot(`
+          [
+            "method",
+          ]
+        `);
+      });
+
+      it('should call both constructor and method onFinish in correct order', async () => {
+        const calls: string[] = [];
+
+        const agent = new ToolLoopAgent({
+          model: mockModel,
+          onFinish: async () => {
+            calls.push('constructor');
+          },
+        });
+
+        await agent.generate({
+          prompt: 'Hello, world!',
+          onFinish: async () => {
+            calls.push('method');
+          },
+        });
+
+        expect(calls).toMatchInlineSnapshot(`
+          [
+            "constructor",
+            "method",
+          ]
+        `);
+      });
+    });
+
+    describe('stream', () => {
+      let mockModel: MockLanguageModelV3;
+
+      beforeEach(() => {
+        mockModel = new MockLanguageModelV3({
+          doStream: async () => {
+            return {
+              stream: convertArrayToReadableStream([
+                {
+                  type: 'stream-start',
+                  warnings: [],
+                },
+                {
+                  type: 'response-metadata',
+                  id: 'id-0',
+                  modelId: 'mock-model-id',
+                  timestamp: new Date(0),
+                },
+                { type: 'text-start', id: '1' },
+                { type: 'text-delta', id: '1', delta: 'Hello' },
+                { type: 'text-delta', id: '1', delta: ', ' },
+                { type: 'text-delta', id: '1', delta: 'world!' },
+                { type: 'text-end', id: '1' },
+                {
+                  type: 'finish',
+                  finishReason: { unified: 'stop', raw: 'stop' },
+                  usage: {
+                    inputTokens: {
+                      total: 3,
+                      noCache: 3,
+                      cacheRead: undefined,
+                      cacheWrite: undefined,
+                    },
+                    outputTokens: {
+                      total: 10,
+                      text: 10,
+                      reasoning: undefined,
+                    },
+                  },
+                  providerMetadata: {
+                    testProvider: { testKey: 'testValue' },
+                  },
+                },
+              ]),
+            };
+          },
+        });
+      });
+
+      it('should call onFinish from constructor', async () => {
+        const calls: string[] = [];
+
+        const agent = new ToolLoopAgent({
+          model: mockModel,
+          onFinish: async () => {
+            calls.push('constructor');
+          },
+        });
+
+        const result = await agent.stream({
+          prompt: 'Hello, world!',
+        });
+
+        await result.consumeStream();
+
+        expect(calls).toMatchInlineSnapshot(`
+          [
+            "constructor",
+          ]
+        `);
+      });
+
+      it('should call onFinish from stream method', async () => {
+        const calls: string[] = [];
+
+        const agent = new ToolLoopAgent({
+          model: mockModel,
+        });
+
+        const result = await agent.stream({
+          prompt: 'Hello, world!',
+          onFinish: async () => {
+            calls.push('method');
+          },
+        });
+
+        await result.consumeStream();
+
+        expect(calls).toMatchInlineSnapshot(`
+          [
+            "method",
+          ]
+        `);
+      });
+
+      it('should call both constructor and method onFinish in correct order', async () => {
+        const calls: string[] = [];
+
+        const agent = new ToolLoopAgent({
+          model: mockModel,
+          onFinish: async () => {
+            calls.push('constructor');
+          },
+        });
+
+        const result = await agent.stream({
+          prompt: 'Hello, world!',
+          onFinish: async () => {
+            calls.push('method');
+          },
+        });
+
+        await result.consumeStream();
+
+        expect(calls).toMatchInlineSnapshot(`
+          [
+            "constructor",
+            "method",
+          ]
+        `);
+      });
+    });
+  });
 });

packages/ai/src/agent/tool-loop-agent.ts

@@ -9,6 +9,7 @@ import { ToolSet } from '../generate-text/tool-set';
 import { Prompt } from '../prompt';
 import { Agent, AgentCallParameters, AgentStreamParameters } from './agent';
 import {
+  ToolLoopAgentOnFinishCallback,
   ToolLoopAgentOnStepFinishCallback,
   ToolLoopAgentSettings,
 } from './tool-loop-agent-settings';
@@ -58,12 +59,15 @@ export class ToolLoopAgent<
   }): Promise<
     Omit<
       ToolLoopAgentSettings<CALL_OPTIONS, TOOLS, OUTPUT>,
-      'prepareCall' | 'instructions' | 'onStepFinish'
+      'prepareCall' | 'instructions' | 'onStepFinish' | 'onFinish'
     > &
       Prompt
   > {
-    const { onStepFinish: _settingsOnStepFinish, ...settingsWithoutCallback } =
-      this.settings;
+    const {
+      onStepFinish: _settingsOnStepFinish,
+      onFinish: _settingsOnFinish,
+      ...settingsWithoutCallback
+    } = this.settings;
     const baseCallArgs = {
       ...settingsWithoutCallback,
       stopWhen: this.settings.stopWhen ?? stepCountIs(20),
@@ -104,13 +108,29 @@ export class ToolLoopAgent<
     return methodCallback ?? constructorCallback;
   }
 
+  private mergeOnFinishCallbacks(
+    methodCallback: ToolLoopAgentOnFinishCallback<TOOLS> | undefined,
+  ): ToolLoopAgentOnFinishCallback<TOOLS> | undefined {
+    const constructorCallback = this.settings.onFinish;
+
+    if (methodCallback && constructorCallback) {
+      return async event => {
+        await constructorCallback(event);
+        await methodCallback(event);
+      };
+    }
+
+    return methodCallback ?? constructorCallback;
+  }
+
   /**
    * Generates an output from the agent (non-streaming).
    */
   async generate({
     abortSignal,
     timeout,
     onStepFinish,
+    onFinish,
     ...options
   }: AgentCallParameters<CALL_OPTIONS, TOOLS>): Promise<
     GenerateTextResult<TOOLS, OUTPUT>
@@ -120,6 +140,7 @@ export class ToolLoopAgent<
       abortSignal,
       timeout,
       onStepFinish: this.mergeOnStepFinishCallbacks(onStepFinish),
+      onFinish: this.mergeOnFinishCallbacks(onFinish),
     });
   }
 
@@ -131,6 +152,7 @@ export class ToolLoopAgent<
     timeout,
     experimental_transform,
     onStepFinish,
+    onFinish,
     ...options
   }: AgentStreamParameters<CALL_OPTIONS, TOOLS>): Promise<
     StreamTextResult<TOOLS, OUTPUT>
@@ -141,6 +163,7 @@ export class ToolLoopAgent<
       timeout,
       experimental_transform,
       onStepFinish: this.mergeOnStepFinishCallbacks(onStepFinish),
+      onFinish: this.mergeOnFinishCallbacks(onFinish),
     });
   }
 }