From d20def51e6538ef4d2f26162a8ed8ce2f35ffa6a Mon Sep 17 00:00:00 2001 From: voyager14 <21mh124@queensu.ca> Date: Sun, 1 Feb 2026 20:15:08 -0500 Subject: [PATCH] feat(ai): unify onFinish callback types into shared TextOnFinishCallback --- .changeset/mighty-birds-attend.md | 5 ++ .../tool-loop-agent-on-finish-callback.ts | 24 +------- .../ai/src/generate-text/generate-text.ts | 22 +------- packages/ai/src/generate-text/index.ts | 4 ++ packages/ai/src/generate-text/stream-text.ts | 22 +------- .../text-on-finish-callback.test-d.ts | 55 +++++++++++++++++++ .../generate-text/text-on-finish-callback.ts | 36 ++++++++++++ 7 files changed, 106 insertions(+), 62 deletions(-) create mode 100644 .changeset/mighty-birds-attend.md create mode 100644 packages/ai/src/generate-text/text-on-finish-callback.test-d.ts create mode 100644 packages/ai/src/generate-text/text-on-finish-callback.ts diff --git a/.changeset/mighty-birds-attend.md b/.changeset/mighty-birds-attend.md new file mode 100644 index 000000000000..ca0e25662932 --- /dev/null +++ b/.changeset/mighty-birds-attend.md @@ -0,0 +1,5 @@ +--- +'ai': patch +--- + +Extract common callback structure from GenerateTextOnFinishCallback, StreamTextOnFinishCallback, and ToolLoopAgentOnFinishCallback into a shared TextOnFinishEvent base type. This prepares the codebase for adding output/outputError support to callbacks. diff --git a/packages/ai/src/agent/tool-loop-agent-on-finish-callback.ts b/packages/ai/src/agent/tool-loop-agent-on-finish-callback.ts index 61a0f4d94159..673b92684382 100644 --- a/packages/ai/src/agent/tool-loop-agent-on-finish-callback.ts +++ b/packages/ai/src/agent/tool-loop-agent-on-finish-callback.ts @@ -1,6 +1,5 @@ -import { StepResult } from '../generate-text/step-result'; +import { TextOnFinishEvent } from '../generate-text/text-on-finish-callback'; import { ToolSet } from '../generate-text/tool-set'; -import { LanguageModelUsage } from '../types/usage'; /** * Callback that is set using the `onFinish` option. @@ -8,24 +7,5 @@ import { LanguageModelUsage } from '../types/usage'; * @param event - The event that is passed to the callback. */ export type ToolLoopAgentOnFinishCallback = ( - event: StepResult & { - /** - * Details for all steps. - */ - readonly steps: StepResult[]; - - /** - * Total usage for all steps. This is the sum of the usage of all steps. - */ - readonly totalUsage: LanguageModelUsage; - - /** - * Context that is passed into tool calls. - * - * Experimental (can break in patch releases). - * - * @default undefined - */ - experimental_context?: unknown; - }, + event: TextOnFinishEvent, ) => PromiseLike | void; diff --git a/packages/ai/src/generate-text/generate-text.ts b/packages/ai/src/generate-text/generate-text.ts index 9208c3d02597..37a4474bec81 100644 --- a/packages/ai/src/generate-text/generate-text.ts +++ b/packages/ai/src/generate-text/generate-text.ts @@ -76,6 +76,7 @@ import { ToolCallRepairFunction } from './tool-call-repair-function'; import { TypedToolError } from './tool-error'; import { ToolOutput } from './tool-output'; import { TypedToolResult } from './tool-result'; +import { TextOnFinishEvent } from './text-on-finish-callback'; import { ToolSet } from './tool-set'; import { mergeAbortSignals } from '../util/merge-abort-signals'; @@ -99,26 +100,7 @@ export type GenerateTextOnStepFinishCallback = ( * @param event - The event that is passed to the callback. */ export type GenerateTextOnFinishCallback = ( - event: StepResult & { - /** - * Details for all steps. - */ - readonly steps: StepResult[]; - - /** - * Total usage for all steps. This is the sum of the usage of all steps. - */ - readonly totalUsage: LanguageModelUsage; - - /** - * Context that is passed into tool execution. - * - * Experimental (can break in patch releases). - * - * @default undefined - */ - experimental_context: unknown; - }, + event: TextOnFinishEvent, ) => PromiseLike | void; /** diff --git a/packages/ai/src/generate-text/index.ts b/packages/ai/src/generate-text/index.ts index ec91dcb1dde3..42c5cb2a6a69 100644 --- a/packages/ai/src/generate-text/index.ts +++ b/packages/ai/src/generate-text/index.ts @@ -55,3 +55,7 @@ export type { TypedToolResult, } from './tool-result'; export type { ToolSet } from './tool-set'; +export type { + TextOnFinishCallback, + TextOnFinishEvent, +} from './text-on-finish-callback'; diff --git a/packages/ai/src/generate-text/stream-text.ts b/packages/ai/src/generate-text/stream-text.ts index 1dd9433e6008..7900a94cea23 100644 --- a/packages/ai/src/generate-text/stream-text.ts +++ b/packages/ai/src/generate-text/stream-text.ts @@ -108,6 +108,7 @@ import { TypedToolCall } from './tool-call'; import { ToolCallRepairFunction } from './tool-call-repair-function'; import { ToolOutput } from './tool-output'; import { StaticToolOutputDenied } from './tool-output-denied'; +import { TextOnFinishEvent } from './text-on-finish-callback'; import { ToolSet } from './tool-set'; const originalGenerateId = createIdGenerator({ @@ -172,26 +173,7 @@ export type StreamTextOnChunkCallback = (event: { * @param event - The event that is passed to the callback. */ export type StreamTextOnFinishCallback = ( - event: StepResult & { - /** - * Details for all steps. - */ - readonly steps: StepResult[]; - - /** - * Total usage for all steps. This is the sum of the usage of all steps. - */ - readonly totalUsage: LanguageModelUsage; - - /** - * Context that is passed into tool execution. - * - * Experimental (can break in patch releases). - * - * @default undefined - */ - experimental_context: unknown; - }, + event: TextOnFinishEvent, ) => PromiseLike | void; /** diff --git a/packages/ai/src/generate-text/text-on-finish-callback.test-d.ts b/packages/ai/src/generate-text/text-on-finish-callback.test-d.ts new file mode 100644 index 000000000000..e0bc9a1b7004 --- /dev/null +++ b/packages/ai/src/generate-text/text-on-finish-callback.test-d.ts @@ -0,0 +1,55 @@ +import { describe, expectTypeOf, it } from 'vitest'; +import { GenerateTextOnFinishCallback } from './generate-text'; +import { StreamTextOnFinishCallback } from './stream-text'; +import { + TextOnFinishCallback, + TextOnFinishEvent, +} from './text-on-finish-callback'; +import { ToolSet } from './tool-set'; +import { ToolLoopAgentOnFinishCallback } from '../agent/tool-loop-agent-on-finish-callback'; + +describe('TextOnFinishCallback types', () => { + it('GenerateTextOnFinishCallback should be assignable to TextOnFinishCallback', () => { + type Tools = ToolSet; + expectTypeOf>().toMatchTypeOf< + TextOnFinishCallback + >(); + }); + + it('StreamTextOnFinishCallback should be assignable to TextOnFinishCallback', () => { + type Tools = ToolSet; + expectTypeOf>().toMatchTypeOf< + TextOnFinishCallback + >(); + }); + + it('ToolLoopAgentOnFinishCallback should be assignable to TextOnFinishCallback', () => { + type Tools = ToolSet; + expectTypeOf>().toMatchTypeOf< + TextOnFinishCallback + >(); + }); + + it('TextOnFinishEvent should have required experimental_context', () => { + type Tools = ToolSet; + type EventContext = TextOnFinishEvent['experimental_context']; + expectTypeOf().toEqualTypeOf(); + }); + + it('callback types should be interchangeable', () => { + type Tools = ToolSet; + const generateCallback: GenerateTextOnFinishCallback = () => {}; + const streamCallback: StreamTextOnFinishCallback = () => {}; + const agentCallback: ToolLoopAgentOnFinishCallback = () => {}; + + // All should be assignable to the base type + const base1: TextOnFinishCallback = generateCallback; + const base2: TextOnFinishCallback = streamCallback; + const base3: TextOnFinishCallback = agentCallback; + + // Suppress unused variable warnings + void base1; + void base2; + void base3; + }); +}); diff --git a/packages/ai/src/generate-text/text-on-finish-callback.ts b/packages/ai/src/generate-text/text-on-finish-callback.ts new file mode 100644 index 000000000000..3beb5c94617f --- /dev/null +++ b/packages/ai/src/generate-text/text-on-finish-callback.ts @@ -0,0 +1,36 @@ +import { LanguageModelUsage } from '../types/usage'; +import { StepResult } from './step-result'; +import { ToolSet } from './tool-set'; + +/** + * Event that is passed to the `onFinish` callback. + */ +export type TextOnFinishEvent = StepResult & { + /** + * Details for all steps. + */ + readonly steps: StepResult[]; + + /** + * Total usage for all steps. This is the sum of the usage of all steps. + */ + readonly totalUsage: LanguageModelUsage; + + /** + * Context that is passed into tool execution. + * + * Experimental (can break in patch releases). + * + * @default undefined + */ + experimental_context: unknown; +}; + +/** + * Callback that is set using the `onFinish` option. + * + * @param event - The event that is passed to the callback. + */ +export type TextOnFinishCallback = ( + event: TextOnFinishEvent, +) => PromiseLike | void;