diff --git a/.changeset/type-safety-audit.md b/.changeset/type-safety-audit.md new file mode 100644 index 0000000..5cbde56 --- /dev/null +++ b/.changeset/type-safety-audit.md @@ -0,0 +1,11 @@ +--- +"@funkai/agents": patch +--- + +Fix type safety issues in agent lifecycle hooks and flow engine + +- Remove unsafe generic hook forwarding from parent agents to sub-agents — only fixed-type hooks (`onStepStart`, `onStepFinish`) are forwarded; generic hooks (`onStart`, `onFinish`, `onError`) stay at the parent level where their `TInput`/`TOutput` types are correct +- Wrap `buildMergedHook` in `fireHooks` for error protection — merged hooks now swallow errors like all other hooks +- Fix config spread ordering in flow agent steps — framework fields (`input`, `signal`, `logger`) can no longer be overwritten by user config +- Thread `TOutput` through `BaseGenerateParams` so `onFinish` hooks receive `GenerateResult` instead of untyped `GenerateResult` +- Fix `AnyHook` contravariance in flow engine — use properly documented `any` escape hatch for internal hook merging diff --git a/packages/agents/docs/core/hooks.md b/packages/agents/docs/core/hooks.md index 9bc9bb3..930ba01 100644 --- a/packages/agents/docs/core/hooks.md +++ b/packages/agents/docs/core/hooks.md @@ -78,6 +78,65 @@ base.onStepFinish -> overrides.onStepFinish The `stepId` for agent tool-loop steps is counter-based: `agentName:0`, `agentName:1`, etc. +## Sub-Agent Hook Forwarding + +When a parent agent has sub-agents (via the `agents` config), those sub-agents are wrapped as tools. The parent forwards a subset of its hooks to each sub-agent so internal activity is observable. + +### What gets forwarded (safe — fixed event types) + +| Hook | Event type | Why safe | +| -------------- | ----------------- | -------------------------------------- | +| `onStepStart` | `StepInfo` | Fixed type, same shape for every agent | +| `onStepFinish` | `StepFinishEvent` | Fixed type, same shape for every agent | +| `logger` | `Logger` | No event type, just a logger instance | + +These hooks are passed directly into `child.generate()` as per-call hooks. The parent's `onStepFinish` is merged (config + per-call) before forwarding, so both the config-level and call-level hooks fire for sub-agent steps. + +### What stays at the parent (not forwarded — generic event types) + +| Hook | Event type | Why not forwarded | +| ---------- | -------------------------------------------------------------- | ----------------------------------------- | +| `onStart` | `{ input: TInput }` | `TInput` differs between parent and child | +| `onFinish` | `{ input: TInput, result: GenerateResult, duration }` | Both `TInput` and `TOutput` differ | +| `onError` | `{ input: TInput, error: Error }` | `TInput` differs between parent and child | + +These hooks are parameterized by the agent's generic types (`TInput`, `TOutput`). A parent typed `Agent<{ userId: string }>` would have `onStart: (e: { input: { userId: string } }) => void`, but a sub-agent might expect `{ query: string }`. Forwarding the parent's hook to the child would cause the hook to receive the wrong event shape at runtime — the compiler cannot catch this because the type boundary is erased when hooks cross agent boundaries. + +Sub-agent lifecycle activity is still observable at the parent level through `onStepFinish`, which fires for each tool-loop step including sub-agent tool calls and their results. + +### Lifecycle diagram + +``` +Parent.generate({ input, onStepFinish }) + │ + ├── Parent fires own onStart({ input }) ← parent's TInput, type-safe + │ + ├── generateText() runs tool loop + │ │ + │ ├── Step 0: LLM calls sub-agent tool + │ │ │ + │ │ │ Passed into child.generate(): + │ │ │ logger → parent's logger + │ │ │ onStepStart → parent's onStepStart (StepInfo — fixed type) + │ │ │ onStepFinish → parent's merged onStepFinish (StepFinishEvent — fixed type) + │ │ │ + │ │ │ NOT passed: + │ │ │ onStart, onFinish, onError (generic types — would break type safety) + │ │ │ + │ │ ├── Child fires own onStart({ input }) ← child's TInput, type-safe + │ │ ├── Child runs tool loop + │ │ │ ├── Child step 0 → parent's onStepFinish fires (StepFinishEvent) + │ │ │ └── Child step 1 → parent's onStepFinish fires (StepFinishEvent) + │ │ ├── Child fires own onFinish(...) ← child's types, type-safe + │ │ └── Returns result to parent + │ │ + │ └── Parent's onStepFinish fires for step 0 (includes sub-agent tool result) + │ + ├── Parent fires own onFinish({ input, result }) ← parent's TInput/TOutput, type-safe + │ + └── Returns Result to caller +``` + ## Error Handling All hooks are executed via `attemptEachAsync`, which: diff --git a/packages/agents/src/core/agents/base/agent.ts b/packages/agents/src/core/agents/base/agent.ts index 6ad5391..8d04628 100644 --- a/packages/agents/src/core/agents/base/agent.ts +++ b/packages/agents/src/core/agents/base/agent.ts @@ -11,6 +11,7 @@ import { buildPrompt, toTokenUsage, } from "@/core/agents/base/utils.js"; +import type { ParentAgentContext } from "@/core/agents/base/utils.js"; import type { Agent, AgentConfig, @@ -99,7 +100,7 @@ export function agent< * * @private */ - function extractInput(params: GenerateParams): TInput { + function extractInput(params: GenerateParams): TInput { if (Object.hasOwn(params, "prompt") && !isNil(params.prompt)) { return params.prompt as unknown as TInput; } @@ -171,7 +172,7 @@ export function agent< * @private */ function resolveSignal( - params: GenerateParams, + params: GenerateParams, ): AbortSignal | undefined { const { timeout, signal } = params; if (signal && isNotNil(timeout)) { @@ -197,7 +198,7 @@ export function agent< async function prepareGeneration( input: TInput, log: Logger, - params: GenerateParams, + params: GenerateParams, ): Promise { const resolvedModel = params.model ?? (await resolveValue(config.model, input)); const model = await withModelMiddleware({ model: resolvedModel }); @@ -210,9 +211,23 @@ export function agent< const hasTools = Object.keys(mergedTools).length > 0; const hasAgents = Object.keys(mergedAgents).length > 0; + // Only fixed-type hooks (onStepStart, onStepFinish) are forwarded to + // Sub-agents. Generic hooks (onStart, onFinish, onError) are NOT + // Forwarded because their event types are parameterized by TInput/TOutput + // — a sub-agent has different generics, so the parent's typed hook + // Would receive the wrong event shape at runtime. Sub-agent activity + // Is still observable via onStepFinish at the parent's tool-loop level. + // See packages/agents/docs/core/hooks.md for the full lifecycle. + const parentCtx: ParentAgentContext = { + log, + onStepStart: params.onStepStart, + onStepFinish: buildMergedHook(log, config.onStepFinish, params.onStepFinish), + }; + const aiTools = buildAITools( valueOrUndefined(hasTools, mergedTools), valueOrUndefined(hasAgents, mergedAgents), + parentCtx, ); // eslint-disable-next-line @typescript-eslint/no-explicit-any -- params.system is Resolver-shaped; safe to resolve @@ -269,7 +284,7 @@ export function agent< } async function generate( - params: GenerateParams, + params: GenerateParams, ): Promise>> { const startedAt = Date.now(); let resolvedInput: TInput | undefined; @@ -329,7 +344,7 @@ export function agent< wrapHook(config.onFinish, { input, result: generateResult, duration }), wrapHook(params.onFinish, { input, - result: generateResult as GenerateResult, + result: generateResult, duration, }), ); @@ -368,7 +383,7 @@ export function agent< } async function stream( - params: GenerateParams, + params: GenerateParams, ): Promise>> { const startedAt = Date.now(); let resolvedInput: TInput | undefined; @@ -454,7 +469,7 @@ export function agent< wrapHook(config.onFinish, { input, result: generateResult, duration }), wrapHook(params.onFinish, { input, - result: generateResult as GenerateResult, + result: generateResult, duration, }), ); @@ -658,3 +673,24 @@ function pickByOutput(output: unknown, ifOutput: T, ifText: T): T { } return ifText; } + +/** + * Build a merged hook that fires config-level and per-call hooks sequentially. + * + * Returns `undefined` when both are absent so `buildParentParams` skips + * the field entirely and sub-agent defaults are preserved. + * + * @private + */ +function buildMergedHook( + log: Logger, + configHook: ((event: E) => void | Promise) | undefined, + callHook: ((event: E) => void | Promise) | undefined, +): ((event: E) => void | Promise) | undefined { + if (isNil(configHook) && isNil(callHook)) { + return undefined; + } + return async (event: E) => { + await fireHooks(log, wrapHook(configHook, event), wrapHook(callHook, event)); + }; +} diff --git a/packages/agents/src/core/agents/base/utils.ts b/packages/agents/src/core/agents/base/utils.ts index b177106..8017f61 100644 --- a/packages/agents/src/core/agents/base/utils.ts +++ b/packages/agents/src/core/agents/base/utils.ts @@ -1,16 +1,75 @@ import type { LanguageModelUsage } from "ai"; import { tool } from "ai"; -import { isFunction, isNil, isNotNil, isString } from "es-toolkit"; +import { isFunction, isNil, isNotNil, isString, omitBy } from "es-toolkit"; import { match, P } from "ts-pattern"; import type { ZodType } from "zod"; import { z } from "zod"; import type { Agent, Message, Resolver } from "@/core/agents/types.js"; +import type { Logger } from "@/core/logger.js"; import type { TokenUsage } from "@/core/provider/types.js"; import type { Tool } from "@/core/tool.js"; +import type { StepFinishEvent, StepInfo } from "@/core/types.js"; import { RUNNABLE_META } from "@/lib/runnable.js"; import type { RunnableMeta } from "@/lib/runnable.js"; +/** + * Context forwarded from a parent agent to sub-agents wrapped as tools. + * + * Includes the parent's logger and **fixed-type** lifecycle hooks so that + * sub-agent internal step activity is visible to the parent. + * + * ## Why only step hooks are forwarded + * + * `onStepStart` and `onStepFinish` use fixed event types (`StepInfo`, + * `StepFinishEvent`) that are the same for every agent — safe to pass + * from parent to child with no type mismatch. + * + * `onStart`, `onFinish`, and `onError` are **not** forwarded because their + * event types are generic over `TInput`/`TOutput`. A parent agent typed + * `Agent<{ userId: string }, ...>` would have `onStart: (e: { input: { userId: string } }) => void`, + * but the sub-agent's input is a completely different type (e.g. + * `{ query: string }`). Forwarding the parent's hook to the child would + * cause the hook to receive the wrong event shape at runtime — a silent + * type-safety violation that the compiler cannot catch. + * + * Sub-agent activity is still observable at the parent level through + * `onStepFinish`, which fires for each tool-loop step including sub-agent + * tool calls. + * + * ## Lifecycle (what gets passed down vs. what stays at the parent) + * + * ``` + * Passed into child.generate() — fixed types, safe: + * log → child creates .child({ agentId }) from it + * onStepStart → StepInfo (same shape for all agents) + * onStepFinish → StepFinishEvent (same shape for all agents) + * + * NOT passed down — generic types, would break type safety: + * onStart → { input: TInput } (differs per agent) + * onFinish → { input: TInput, result: GenerateResult } (differs per agent) + * onError → { input: TInput, error: Error } (differs per agent) + * ``` + */ +export interface ParentAgentContext { + /** Parent logger — sub-agent creates `.child({ agentId })` from it. */ + log?: Logger; + + /** + * Fires when a sub-agent step starts. + * + * Uses `StepInfo` — a fixed (non-generic) type, safe to forward. + */ + onStepStart?: (event: { step: StepInfo }) => void | Promise; + + /** + * Fires when a sub-agent step finishes. + * + * Uses `StepFinishEvent` — a fixed (non-generic) type, safe to forward. + */ + onStepFinish?: (event: StepFinishEvent) => void | Promise; +} + /** * Merge `Tool` records and wrap subagent `Runnable` objects into AI SDK * tool format for `generateText` / `streamText`. @@ -24,12 +83,15 @@ import type { RunnableMeta } from "@/lib/runnable.js"; * * @param tools - Record of named tools to include. * @param agents - Record of named sub-agents to wrap as tools. + * @param parentCtx - Parent context forwarded to sub-agent generate calls + * (logger, lifecycle hooks). * @returns The merged tool set, or `undefined` when empty. */ export function buildAITools( tools?: Record, // eslint-disable-next-line @typescript-eslint/no-explicit-any -- Agent generic params are contravariant; `unknown` breaks assignability agents?: Record>, + parentCtx?: ParentAgentContext, // eslint-disable-next-line @typescript-eslint/no-explicit-any -- ToolSet requires `any` values; `unknown` breaks assignability with AI SDK ): Record | undefined { const hasTools = isNotNil(tools) && Object.keys(tools).length > 0; @@ -39,7 +101,7 @@ export function buildAITools( return undefined; } - const agentTools: Record = buildAgentTools(agents, tools); + const agentTools: Record = buildAgentTools(agents, tools, parentCtx); return { ...tools, ...agentTools }; } @@ -253,6 +315,7 @@ function buildAgentTools( // eslint-disable-next-line @typescript-eslint/no-explicit-any -- Agent generic params are contravariant; `unknown` breaks assignability agents: Record> | undefined, tools: Record | undefined, + parentCtx: ParentAgentContext | undefined, ): Record { if (!agents) { return {}; @@ -279,6 +342,7 @@ function buildAgentTools( meta, toolName, tools, + parentCtx, ); return [agentToolName, agentTool]; @@ -298,14 +362,22 @@ function buildAgentTool( meta: RunnableMeta | undefined, toolName: string, tools: Record | undefined, + parentCtx: ParentAgentContext | undefined, // eslint-disable-next-line @typescript-eslint/no-explicit-any -- ToolSet requires `any` values; `unknown` breaks assignability with AI SDK ): ReturnType> { + const parentParams = buildParentParams(parentCtx); + if (isNotNil(meta) && isNotNil(meta.inputSchema)) { return tool({ description: `Delegate to ${toolName}`, inputSchema: meta.inputSchema, execute: async (input, { abortSignal }) => { - const r = await runnable.generate({ input, signal: abortSignal, tools }); + const r = await runnable.generate({ + input, + signal: abortSignal, + tools, + ...parentParams, + }); if (!r.ok) { throw new Error(r.error.message); } @@ -321,6 +393,7 @@ function buildAgentTool( prompt: input.prompt, signal: abortSignal, tools, + ...parentParams, }); if (!r.ok) { throw new Error(r.error.message); @@ -329,3 +402,28 @@ function buildAgentTool( }, }); } + +/** + * Build the per-call params to forward from parent context to sub-agent. + * + * Only forwards the parent logger and **fixed-type** step hooks. + * Generic hooks (`onStart`, `onFinish`, `onError`) are intentionally + * excluded — see {@link ParentAgentContext} for the rationale. + * + * Omits `undefined` values so they don't override sub-agent defaults. + * + * @private + */ +function buildParentParams(ctx: ParentAgentContext | undefined): Record { + if (isNil(ctx)) { + return {}; + } + return omitBy( + { + logger: ctx.log, + onStepStart: ctx.onStepStart, + onStepFinish: ctx.onStepFinish, + }, + isNil, + ); +} diff --git a/packages/agents/src/core/agents/flow/engine.ts b/packages/agents/src/core/agents/flow/engine.ts index d31f348..1b07890 100644 --- a/packages/agents/src/core/agents/flow/engine.ts +++ b/packages/agents/src/core/agents/flow/engine.ts @@ -200,36 +200,51 @@ function createHookCaller( return undefined; } +/** + * Internal-only hook type used by `buildMergedHook` to accept any + * lifecycle hook regardless of its specific event signature. + * + * ## Why `any` is necessary here + * + * Functions are **contravariant** in their parameter types. A hook + * typed `(event: { input: TInput }) => void` is NOT assignable to + * `(event: unknown) => void` — the subtype relationship is reversed + * for function parameters. This means no strict type (including + * `unknown` or `never`) can serve as a universal hook acceptor. + * + * `any` is the only TypeScript type that bypasses contravariance, + * allowing all hook signatures to unify in a single merge function. + * + * Type safety is enforced at the public API boundary: + * - `FlowEngineConfig` defines engine hooks with `unknown` event fields + * - `FlowAgentConfig` defines flow hooks with `TInput`/`TOutput` event fields + * - Both produce the correct runtime event shapes — the merge function + * just combines two callbacks, it never inspects or constructs events. + * + * @private + */ +// eslint-disable-next-line @typescript-eslint/no-explicit-any -- see JSDoc above: contravariance requires `any` +type AnyHook = ((event: any) => void | Promise) | undefined; + /** * Build a merged hook that runs engine and flow agent hooks sequentially. * - * The `(event: never)` constraint is the widest function type under - * strict mode — any single-argument function is assignable via - * contravariance (`never extends T` for all `T`). + * Accepts `AnyHook` (event: unknown) directly — no generic type parameter + * or unsafe casts needed. Call sites widen narrower hook types to `AnyHook` + * explicitly. * * @private */ -function buildMergedHook void | Promise>( - log: Logger, - engineHook: THook | undefined, - flowHook: THook | undefined, -): THook | undefined { +function buildMergedHook(log: Logger, engineHook: AnyHook, flowHook: AnyHook): AnyHook { if (!engineHook && !flowHook) { return undefined; } - const merged = async (event: unknown): Promise => { - const engineFn = createHookCaller( - engineHook as ((event: unknown) => void | Promise) | undefined, - event, - ); - const flowFn = createHookCaller( - flowHook as ((event: unknown) => void | Promise) | undefined, - event, - ); + return async (event: unknown): Promise => { + const engineFn = createHookCaller(engineHook, event); + const flowFn = createHookCaller(flowHook, event); await fireHooks(log, engineFn, flowFn); }; - return merged as unknown as THook; } /** diff --git a/packages/agents/src/core/agents/flow/steps/agent.test.ts b/packages/agents/src/core/agents/flow/steps/agent.test.ts index 3fccef4..0dedefe 100644 --- a/packages/agents/src/core/agents/flow/steps/agent.test.ts +++ b/packages/agents/src/core/agents/flow/steps/agent.test.ts @@ -113,7 +113,7 @@ describe("agent()", () => { expect(agent.generate).toHaveBeenCalledWith( expect.objectContaining({ input: "hello", - signal: config.signal, + signal: ctx.signal, logger: expect.any(Object), }), ); @@ -132,7 +132,7 @@ describe("agent()", () => { ); }); - it("user-provided config.signal takes precedence over ctx.signal", async () => { + it("framework ctx.signal takes precedence over user-provided config.signal", async () => { const ctxController = new AbortController(); const userController = new AbortController(); const ctx = createMockCtx({ signal: ctxController.signal }); @@ -147,7 +147,7 @@ describe("agent()", () => { }); expect(agent.generate).toHaveBeenCalledWith( - expect.objectContaining({ input: "test", signal: userController.signal }), + expect.objectContaining({ input: "test", signal: ctxController.signal }), ); }); diff --git a/packages/agents/src/core/agents/flow/steps/factory.test.ts b/packages/agents/src/core/agents/flow/steps/factory.test.ts index e9e3bed..c27063b 100644 --- a/packages/agents/src/core/agents/flow/steps/factory.test.ts +++ b/packages/agents/src/core/agents/flow/steps/factory.test.ts @@ -364,7 +364,7 @@ describe("agent()", () => { expect(agent.generate).toHaveBeenCalledWith( expect.objectContaining({ input: "hello", - signal: config.signal, + signal: ctx.signal, logger: expect.any(Object), }), ); @@ -383,7 +383,7 @@ describe("agent()", () => { ); }); - it("user-provided config.signal takes precedence over ctx.signal", async () => { + it("framework ctx.signal takes precedence over user-provided config.signal", async () => { const ctxController = new AbortController(); const userController = new AbortController(); const ctx = createMockCtx({ signal: ctxController.signal }); @@ -398,7 +398,7 @@ describe("agent()", () => { }); expect(agent.generate).toHaveBeenCalledWith( - expect.objectContaining({ input: "test", signal: userController.signal }), + expect.objectContaining({ input: "test", signal: ctxController.signal }), ); }); diff --git a/packages/agents/src/core/agents/flow/steps/factory.ts b/packages/agents/src/core/agents/flow/steps/factory.ts index bb35c53..0784c67 100644 --- a/packages/agents/src/core/agents/flow/steps/factory.ts +++ b/packages/agents/src/core/agents/flow/steps/factory.ts @@ -275,9 +275,9 @@ function createStepBuilderInternal(options: StepBuilderOptions, indexRef: IndexR input: config.input, execute: async () => { const agentParams = { + ...config.config, input: config.input, signal: ctx.signal, - ...config.config, logger: ctx.log.child({ stepId: config.id }), }; diff --git a/packages/agents/src/core/agents/flow/types.ts b/packages/agents/src/core/agents/flow/types.ts index 3a8146e..f6b3352 100644 --- a/packages/agents/src/core/agents/flow/types.ts +++ b/packages/agents/src/core/agents/flow/types.ts @@ -9,6 +9,7 @@ import type { StreamResult, } from "@/core/agents/types.js"; import type { Logger } from "@/core/logger.js"; +import type { Tool } from "@/core/tool.js"; import type { StepFinishEvent, StepInfo } from "@/core/types.js"; import type { Context } from "@/lib/context.js"; import type { TraceEntry } from "@/lib/trace.js"; @@ -334,7 +335,9 @@ export interface FlowAgent { * const result = await myFlow.generate({ input: { targetDir: '.' } }) * ``` */ - generate(params: GenerateParams): Promise>>; + generate( + params: GenerateParams, Record, TOutput>, + ): Promise>>; /** * Run the flow agent with streaming step progress. @@ -346,12 +349,16 @@ export interface FlowAgent { * @param params - Input and optional per-call overrides. * @returns A `Result` wrapping the `StreamResult`. */ - stream(params: GenerateParams): Promise>>; + stream( + params: GenerateParams, Record, TOutput>, + ): Promise>>; /** * Returns a plain function that calls `.generate()`. */ - fn(): (params: GenerateParams) => Promise>>; + fn(): ( + params: GenerateParams, Record, TOutput>, + ) => Promise>>; } /** diff --git a/packages/agents/src/core/agents/types.ts b/packages/agents/src/core/agents/types.ts index 5c0e003..e9cf52f 100644 --- a/packages/agents/src/core/agents/types.ts +++ b/packages/agents/src/core/agents/types.ts @@ -297,7 +297,7 @@ export interface StreamResult { * * @private — use `GenerateParams` instead. */ -export interface BaseGenerateParams { +export interface BaseGenerateParams { /** * Override the logger for this call. * @@ -341,7 +341,7 @@ export interface BaseGenerateParams { */ onFinish?: (event: { input: TInput; - result: GenerateResult; + result: GenerateResult; duration: number; }) => void | Promise; @@ -478,7 +478,10 @@ export type GenerateParams< TInput = unknown, TTools extends Record = Record, TSubAgents extends SubAgents = Record, -> = BaseGenerateParams & AgentGenerateOverrides & InputUnion; + TOutput = string, +> = BaseGenerateParams & + AgentGenerateOverrides & + InputUnion; /** * Configuration for creating an agent. @@ -747,7 +750,7 @@ export interface Agent< * ``` */ generate( - params: GenerateParams, + params: GenerateParams, ): Promise>>; /** @@ -763,7 +766,7 @@ export interface Agent< * `result.output` / `result.messages` after the stream ends. */ stream( - params: GenerateParams, + params: GenerateParams, ): Promise>>; /** @@ -781,6 +784,6 @@ export interface Agent< * ``` */ fn(): ( - params: GenerateParams, + params: GenerateParams, ) => Promise>>; } diff --git a/packages/agents/src/integration/lifecycle.test.ts b/packages/agents/src/integration/lifecycle.test.ts new file mode 100644 index 0000000..d221321 --- /dev/null +++ b/packages/agents/src/integration/lifecycle.test.ts @@ -0,0 +1,1719 @@ +import { simulateReadableStream } from "ai"; +/** + * Integration tests for lifecycle event propagation. + * + * These tests exercise the full stack — real `agent()`, `flowAgent()`, + * `MockLanguageModelV3` from the AI SDK — to verify that hooks fire + * at every nesting level and that stream events propagate correctly. + * + * No vi.mock() calls — everything is wired together as it would be + * in production, with mock models standing in for real LLMs. + */ +import { MockLanguageModelV3, mockValues } from "ai/test"; +import { describe, it, expect, vi, beforeEach } from "vitest"; +import { z } from "zod"; + +import { agent } from "@/core/agents/base/agent.js"; +import { flowAgent } from "@/core/agents/flow/flow-agent.js"; +import type { StepFinishEvent, StepInfo, StreamPart } from "@/core/types.js"; +import { createMockLogger } from "@/testing/index.js"; + +// --------------------------------------------------------------------------- +// Mock middleware — bypass devtools import +// --------------------------------------------------------------------------- + +vi.mock( + import("@/lib/middleware.js"), + () => + ({ + withModelMiddleware: vi.fn(async ({ model }: { model: unknown }) => model), + // oxlint-disable-next-line @typescript-eslint/no-explicit-any -- Mock factory must return partial shape + }) as any, +); + +// --------------------------------------------------------------------------- +// Helpers +// --------------------------------------------------------------------------- + +const MOCK_USAGE = { + inputTokens: { total: 10, noCache: 10, cacheRead: undefined, cacheWrite: undefined }, + outputTokens: { total: 5, text: 5, reasoning: undefined }, +}; + +const MOCK_FINISH = { unified: "stop" as const, raw: undefined }; + +function createMockModel(text = "mock response") { + return new MockLanguageModelV3({ + doGenerate: mockValues({ + content: [{ type: "text" as const, text }], + finishReason: MOCK_FINISH, + usage: MOCK_USAGE, + warnings: [], + // oxlint-disable-next-line @typescript-eslint/no-explicit-any -- mockValues returns sync fn, but MockLanguageModelV3 expects PromiseLike + }) as any, + doStream: mockValues({ + stream: simulateReadableStream({ + chunks: [ + { type: "text-start" as const, id: "t1" }, + { type: "text-delta" as const, id: "t1", delta: text }, + { type: "text-end" as const, id: "t1" }, + { + type: "finish" as const, + finishReason: MOCK_FINISH, + logprobs: undefined, + usage: MOCK_USAGE, + }, + ], + chunkDelayInMs: 0, + }), + // oxlint-disable-next-line @typescript-eslint/no-explicit-any -- mockValues returns sync fn, but MockLanguageModelV3 expects PromiseLike + }) as any, + }); +} + +interface LifecycleEvent { + type: string; + detail?: string; +} + +function createLifecycleTracker() { + const events: LifecycleEvent[] = []; + + return { + events, + onStart: vi.fn((_event: { input: unknown }) => { + events.push({ type: "onStart" }); + }), + onFinish: vi.fn((_event: { input: unknown; result: unknown; duration: number }) => { + events.push({ type: "onFinish" }); + }), + onError: vi.fn((_event: { input: unknown; error: Error }) => { + events.push({ type: "onError" }); + }), + onStepStart: vi.fn((event: { step: StepInfo }) => { + events.push({ type: "onStepStart", detail: event.step.id }); + }), + onStepFinish: vi.fn((event: StepFinishEvent) => { + const id = event.step?.id ?? event.stepId ?? "unknown"; + events.push({ type: "onStepFinish", detail: id }); + }), + }; +} + +// --------------------------------------------------------------------------- +// Agent lifecycle +// --------------------------------------------------------------------------- + +describe("Agent lifecycle (integration)", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("fires onStart, onStepFinish, onFinish in order for generate()", async () => { + const tracker = createLifecycleTracker(); + + const a = agent({ + name: "test-agent", + model: createMockModel(), + system: "You are helpful.", + logger: createMockLogger(), + onStart: tracker.onStart, + onFinish: tracker.onFinish, + onStepFinish: tracker.onStepFinish, + }); + + const result = await a.generate({ prompt: "hello" }); + + expect(result.ok).toBe(true); + expect(tracker.onStart).toHaveBeenCalledTimes(1); + expect(tracker.onFinish).toHaveBeenCalledTimes(1); + // At least one step from the tool loop + expect(tracker.onStepFinish).toHaveBeenCalled(); + + // Order: onStart fires before onFinish + const startIdx = tracker.events.findIndex((e) => e.type === "onStart"); + const finishIdx = tracker.events.findIndex((e) => e.type === "onFinish"); + expect(startIdx).toBeLessThan(finishIdx); + }); + + it("fires onStart, onStepFinish, onFinish in order for stream()", async () => { + const tracker = createLifecycleTracker(); + + const a = agent({ + name: "stream-agent", + model: createMockModel(), + system: "You are helpful.", + logger: createMockLogger(), + onStart: tracker.onStart, + onFinish: tracker.onFinish, + onStepFinish: tracker.onStepFinish, + }); + + const result = await a.stream({ prompt: "hello" }); + + expect(result.ok).toBe(true); + if (!result.ok) { + return; + } + + // Drain the stream + for await (const _part of result.fullStream) { + /* Consume */ + } + + // Wait for finish hook to settle + await result.output; + + expect(tracker.onStart).toHaveBeenCalledTimes(1); + expect(tracker.onFinish).toHaveBeenCalledTimes(1); + expect(tracker.onStepFinish).toHaveBeenCalled(); + }); + + it("fires onError (not onFinish) when model throws during generate()", async () => { + const tracker = createLifecycleTracker(); + + const failModel = new MockLanguageModelV3({ + doGenerate: async () => { + throw new Error("model exploded"); + }, + }); + + const a = agent({ + name: "fail-agent", + model: failModel, + logger: createMockLogger(), + onStart: tracker.onStart, + onFinish: tracker.onFinish, + onError: tracker.onError, + }); + + const result = await a.generate({ prompt: "hello" }); + + expect(result.ok).toBe(false); + expect(tracker.onStart).toHaveBeenCalledTimes(1); + expect(tracker.onError).toHaveBeenCalledTimes(1); + expect(tracker.onFinish).not.toHaveBeenCalled(); + }); + + it("merges config hooks and per-call hooks (config first)", async () => { + const order: string[] = []; + + const a = agent({ + name: "merge-agent", + model: createMockModel(), + logger: createMockLogger(), + onStart: () => { + order.push("config:onStart"); + }, + onFinish: () => { + order.push("config:onFinish"); + }, + }); + + await a.generate({ + prompt: "hello", + onStart: () => { + order.push("call:onStart"); + }, + onFinish: () => { + order.push("call:onFinish"); + }, + }); + + expect(order).toEqual(["config:onStart", "call:onStart", "config:onFinish", "call:onFinish"]); + }); +}); + +// --------------------------------------------------------------------------- +// FlowAgent lifecycle — direct steps +// --------------------------------------------------------------------------- + +describe("FlowAgent lifecycle — direct steps (integration)", () => { + const Input = z.object({ x: z.number() }); + const Output = z.object({ y: z.number() }); + + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("fires onStart, onStepStart/Finish for each step, then onFinish", async () => { + const tracker = createLifecycleTracker(); + + const fa = flowAgent<{ x: number }, { y: number }>( + { + name: "step-flow", + input: Input, + output: Output, + logger: createMockLogger(), + onStart: tracker.onStart, + onFinish: tracker.onFinish, + onStepStart: tracker.onStepStart, + onStepFinish: tracker.onStepFinish, + }, + async ({ input, $ }) => { + await $.step({ id: "double", execute: async () => input.x * 2 }); + await $.step({ id: "add-one", execute: async () => input.x * 2 + 1 }); + return { y: input.x * 2 }; + }, + ); + + const result = await fa.generate({ input: { x: 5 } }); + + expect(result.ok).toBe(true); + expect(tracker.events).toEqual([ + { type: "onStart" }, + { type: "onStepStart", detail: "double" }, + { type: "onStepFinish", detail: "double" }, + { type: "onStepStart", detail: "add-one" }, + { type: "onStepFinish", detail: "add-one" }, + { type: "onFinish" }, + ]); + }); + + it("fires onStepStart/Finish for nested steps inside $.step()", async () => { + const tracker = createLifecycleTracker(); + + const fa = flowAgent<{ x: number }, { y: number }>( + { + name: "nested-step-flow", + input: Input, + output: Output, + logger: createMockLogger(), + onStepStart: tracker.onStepStart, + onStepFinish: tracker.onStepFinish, + }, + async ({ input, $: $outer }) => { + await $outer.step({ + id: "outer", + execute: async ({ $ }) => { + await $.step({ id: "inner", execute: async () => input.x + 1 }); + return input.x * 2; + }, + }); + return { y: input.x * 2 }; + }, + ); + + const result = await fa.generate({ input: { x: 3 } }); + + expect(result.ok).toBe(true); + expect(tracker.events).toEqual([ + { type: "onStepStart", detail: "outer" }, + { type: "onStepStart", detail: "inner" }, + { type: "onStepFinish", detail: "inner" }, + { type: "onStepFinish", detail: "outer" }, + ]); + }); + + it("fires hooks for steps inside $.map()", async () => { + const tracker = createLifecycleTracker(); + + const fa = flowAgent<{ x: number }, { y: number }>( + { + name: "map-flow", + input: Input, + output: Output, + logger: createMockLogger(), + onStepStart: tracker.onStepStart, + onStepFinish: tracker.onStepFinish, + }, + async ({ input, $: $outer }) => { + await $outer.map({ + id: "batch", + input: [1, 2], + execute: async ({ item, $ }) => { + const r = await $.step({ + id: `process-${item}`, + execute: async () => item * input.x, + }); + if (r.ok) { + return r.value; + } + return 0; + }, + }); + return { y: input.x }; + }, + ); + + const result = await fa.generate({ input: { x: 10 } }); + + expect(result.ok).toBe(true); + + // Should have: batch start, process-1 start/finish, process-2 start/finish, batch finish + const stepIds = tracker.events.map((e) => `${e.type}:${e.detail}`); + expect(stepIds).toContain("onStepStart:batch"); + expect(stepIds).toContain("onStepStart:process-1"); + expect(stepIds).toContain("onStepFinish:process-1"); + expect(stepIds).toContain("onStepStart:process-2"); + expect(stepIds).toContain("onStepFinish:process-2"); + expect(stepIds).toContain("onStepFinish:batch"); + }); + + it("fires hooks for steps inside $.each()", async () => { + const tracker = createLifecycleTracker(); + + const fa = flowAgent<{ x: number }, { y: number }>( + { + name: "each-flow", + input: Input, + output: Output, + logger: createMockLogger(), + onStepStart: tracker.onStepStart, + onStepFinish: tracker.onStepFinish, + }, + async ({ input, $: $outer }) => { + await $outer.each({ + id: "iterate", + input: [1, 2], + execute: async ({ item, $ }) => { + await $.step({ + id: `item-${item}`, + execute: async () => item * input.x, + }); + }, + }); + return { y: input.x }; + }, + ); + + const result = await fa.generate({ input: { x: 5 } }); + + expect(result.ok).toBe(true); + + const stepIds = tracker.events.map((e) => `${e.type}:${e.detail}`); + expect(stepIds).toContain("onStepStart:iterate"); + expect(stepIds).toContain("onStepStart:item-1"); + expect(stepIds).toContain("onStepFinish:item-1"); + expect(stepIds).toContain("onStepStart:item-2"); + expect(stepIds).toContain("onStepFinish:item-2"); + expect(stepIds).toContain("onStepFinish:iterate"); + }); + + it("fires hooks for steps inside $.reduce()", async () => { + const tracker = createLifecycleTracker(); + + const fa = flowAgent<{ x: number }, { y: number }>( + { + name: "reduce-flow", + input: Input, + output: Output, + logger: createMockLogger(), + onStepStart: tracker.onStepStart, + onStepFinish: tracker.onStepFinish, + }, + async ({ input, $ }) => { + const r = await $.reduce({ + id: "sum", + input: [1, 2, 3], + initial: 0, + execute: async ({ item, accumulator }) => accumulator + item * input.x, + }); + if (r.ok) { + return { y: r.value }; + } + return { y: 0 }; + }, + ); + + const result = await fa.generate({ input: { x: 1 } }); + + expect(result.ok).toBe(true); + if (!result.ok) { + return; + } + expect(result.output).toEqual({ y: 6 }); + + const stepIds = tracker.events.map((e) => `${e.type}:${e.detail}`); + expect(stepIds).toContain("onStepStart:sum"); + expect(stepIds).toContain("onStepFinish:sum"); + }); + + it("onStepFinish fires even when a step errors", async () => { + const tracker = createLifecycleTracker(); + + const fa = flowAgent<{ x: number }, { y: number }>( + { + name: "error-step-flow", + input: Input, + output: Output, + logger: createMockLogger(), + onStepStart: tracker.onStepStart, + onStepFinish: tracker.onStepFinish, + }, + async ({ input, $ }) => { + const r = await $.step({ + id: "fail-step", + execute: async () => { + throw new Error("step boom"); + }, + }); + if (r.ok) { + return { y: 0 }; + } + return { y: -1 }; + }, + ); + + const result = await fa.generate({ input: { x: 1 } }); + + expect(result.ok).toBe(true); + if (!result.ok) { + return; + } + expect(result.output).toEqual({ y: -1 }); + + expect(tracker.events).toEqual([ + { type: "onStepStart", detail: "fail-step" }, + { type: "onStepFinish", detail: "fail-step" }, + ]); + + // Verify the onStepFinish event has result: undefined on error + const finishEvent = tracker.onStepFinish.mock.calls[0]?.[0] as StepFinishEvent; + expect(finishEvent.result).toBeUndefined(); + }); +}); + +// --------------------------------------------------------------------------- +// FlowAgent with $.agent() — agent hooks propagate +// --------------------------------------------------------------------------- + +describe("FlowAgent with $.agent() (integration)", () => { + const Input = z.object({ topic: z.string() }); + const Output = z.object({ summary: z.string() }); + + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("fires flow onStepStart/Finish for $.agent() steps", async () => { + const tracker = createLifecycleTracker(); + + const writer = agent({ + name: "writer", + model: createMockModel("written content"), + logger: createMockLogger(), + }); + + const fa = flowAgent<{ topic: string }, { summary: string }>( + { + name: "agent-flow", + input: Input, + output: Output, + logger: createMockLogger(), + onStart: tracker.onStart, + onFinish: tracker.onFinish, + onStepStart: tracker.onStepStart, + onStepFinish: tracker.onStepFinish, + }, + async ({ input, $ }) => { + const r = await $.agent({ + id: "write", + agent: writer, + input: `Write about ${input.topic}`, + }); + + if (r.ok) { + return { summary: String(r.value.output) }; + } + return { summary: "failed" }; + }, + ); + + const result = await fa.generate({ input: { topic: "TypeScript" } }); + + expect(result.ok).toBe(true); + if (!result.ok) { + return; + } + expect(result.output.summary).toBe("written content"); + + expect(tracker.events).toEqual([ + { type: "onStart" }, + { type: "onStepStart", detail: "write" }, + { type: "onStepFinish", detail: "write" }, + { type: "onFinish" }, + ]); + }); + + it("fires flow hooks for multiple $.agent() calls in sequence", async () => { + const tracker = createLifecycleTracker(); + + const researcher = agent({ + name: "researcher", + model: createMockModel("research findings"), + logger: createMockLogger(), + }); + + const writer = agent({ + name: "writer", + model: createMockModel("final article"), + logger: createMockLogger(), + }); + + const fa = flowAgent<{ topic: string }, { summary: string }>( + { + name: "pipeline", + input: Input, + output: Output, + logger: createMockLogger(), + onStepStart: tracker.onStepStart, + onStepFinish: tracker.onStepFinish, + }, + async ({ input, $ }) => { + const research = await $.agent({ + id: "research", + agent: researcher, + input: input.topic, + }); + + let researchInput = ""; + if (research.ok) { + researchInput = String(research.value.output); + } + const article = await $.agent({ + id: "write", + agent: writer, + input: researchInput, + }); + + if (article.ok) { + return { summary: String(article.value.output) }; + } + return { summary: "failed" }; + }, + ); + + const result = await fa.generate({ input: { topic: "AI" } }); + + expect(result.ok).toBe(true); + expect(tracker.events).toEqual([ + { type: "onStepStart", detail: "research" }, + { type: "onStepFinish", detail: "research" }, + { type: "onStepStart", detail: "write" }, + { type: "onStepFinish", detail: "write" }, + ]); + }); + + it("fires flow hooks for $.agent() inside $.map()", async () => { + const tracker = createLifecycleTracker(); + + const processor = agent({ + name: "processor", + model: createMockModel("processed"), + logger: createMockLogger(), + }); + + const fa = flowAgent<{ topic: string }, { summary: string }>( + { + name: "map-agent-flow", + input: Input, + output: Output, + logger: createMockLogger(), + onStepStart: tracker.onStepStart, + onStepFinish: tracker.onStepFinish, + }, + async ({ $: $outer, input }) => { + await $outer.map({ + id: "batch", + input: ["a", "b"], + execute: async ({ item, $ }) => { + const r = await $.agent({ + id: `process-${item}`, + agent: processor, + input: item, + }); + if (r.ok) { + return String(r.value.output); + } + return ""; + }, + }); + + return { summary: input.topic }; + }, + ); + + const result = await fa.generate({ input: { topic: "test" } }); + + expect(result.ok).toBe(true); + + const stepIds = tracker.events.map((e) => `${e.type}:${e.detail}`); + expect(stepIds).toContain("onStepStart:batch"); + expect(stepIds).toContain("onStepStart:process-a"); + expect(stepIds).toContain("onStepFinish:process-a"); + expect(stepIds).toContain("onStepStart:process-b"); + expect(stepIds).toContain("onStepFinish:process-b"); + expect(stepIds).toContain("onStepFinish:batch"); + }); +}); + +// --------------------------------------------------------------------------- +// FlowAgent with agents config dependency +// --------------------------------------------------------------------------- + +describe("FlowAgent agents dependency lifecycle (integration)", () => { + const Input = z.object({ text: z.string() }); + const Output = z.object({ result: z.string() }); + + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("passes agents from config and fires hooks for $.agent() using them", async () => { + const tracker = createLifecycleTracker(); + + const core = agent({ + name: "core", + model: createMockModel("core output"), + logger: createMockLogger(), + }); + + const fa = flowAgent<{ text: string }, { result: string }>( + { + name: "dep-flow", + input: Input, + output: Output, + logger: createMockLogger(), + agents: { core }, + onStepStart: tracker.onStepStart, + onStepFinish: tracker.onStepFinish, + }, + async ({ input, $, agents }) => { + const r = await $.agent({ + id: "run-core", + // oxlint-disable-next-line @typescript-eslint/no-explicit-any -- FlowSubAgents union includes FlowAgent; narrow to Agent for $.agent() + agent: agents.core as any, + input: input.text, + }); + if (r.ok) { + return { result: String(r.value.output) }; + } + return { result: "failed" }; + }, + ); + + const result = await fa.generate({ input: { text: "hello" } }); + + expect(result.ok).toBe(true); + if (!result.ok) { + return; + } + expect(result.output.result).toBe("core output"); + + expect(tracker.events).toEqual([ + { type: "onStepStart", detail: "run-core" }, + { type: "onStepFinish", detail: "run-core" }, + ]); + }); +}); + +// --------------------------------------------------------------------------- +// Deep nesting — 3+ levels +// --------------------------------------------------------------------------- + +describe("Deep nesting lifecycle (integration)", () => { + const Input = z.object({ items: z.array(z.string()) }); + const Output = z.object({ results: z.array(z.string()) }); + + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("propagates onStepStart/Finish through 3 nesting levels", async () => { + const tracker = createLifecycleTracker(); + + const processor = agent({ + name: "processor", + model: createMockModel("result"), + logger: createMockLogger(), + }); + + const fa = flowAgent<{ items: string[] }, { results: string[] }>( + { + name: "deep-flow", + input: Input, + output: Output, + logger: createMockLogger(), + onStepStart: tracker.onStepStart, + onStepFinish: tracker.onStepFinish, + }, + async ({ input, $: $flow }) => { + // Level 1: $.step() + const r = await $flow.step({ + id: "outer", + execute: async ({ $: $step }) => { + // Level 2: $.map() inside $.step() + const mapResult = await $step.map({ + id: "batch", + input: input.items, + execute: async ({ item, $ }) => { + // Level 3: $.agent() inside $.map() inside $.step() + const agentResult = await $.agent({ + id: `process-${item}`, + agent: processor, + input: item, + }); + if (agentResult.ok) { + return String(agentResult.value.output); + } + return ""; + }, + }); + if (mapResult.ok) { + return mapResult.value; + } + return []; + }, + }); + + if (r.ok) { + return { results: r.value }; + } + return { results: [] }; + }, + ); + + const result = await fa.generate({ input: { items: ["x", "y"] } }); + + expect(result.ok).toBe(true); + if (!result.ok) { + return; + } + expect(result.output.results).toEqual(["result", "result"]); + + // Verify ALL nesting levels fired hooks + const stepIds = tracker.events.map((e) => `${e.type}:${e.detail}`); + + // Level 1 + expect(stepIds).toContain("onStepStart:outer"); + expect(stepIds).toContain("onStepFinish:outer"); + + // Level 2 + expect(stepIds).toContain("onStepStart:batch"); + expect(stepIds).toContain("onStepFinish:batch"); + + // Level 3 + expect(stepIds).toContain("onStepStart:process-x"); + expect(stepIds).toContain("onStepFinish:process-x"); + expect(stepIds).toContain("onStepStart:process-y"); + expect(stepIds).toContain("onStepFinish:process-y"); + + // Verify nesting order: outer starts first, finishes last + const outerStartIdx = tracker.events.findIndex( + (e) => e.type === "onStepStart" && e.detail === "outer", + ); + const outerFinishIdx = tracker.events.findIndex( + (e) => e.type === "onStepFinish" && e.detail === "outer", + ); + const batchStartIdx = tracker.events.findIndex( + (e) => e.type === "onStepStart" && e.detail === "batch", + ); + const batchFinishIdx = tracker.events.findIndex( + (e) => e.type === "onStepFinish" && e.detail === "batch", + ); + + expect(outerStartIdx).toBeLessThan(batchStartIdx); + expect(batchFinishIdx).toBeLessThan(outerFinishIdx); + }); +}); + +// --------------------------------------------------------------------------- +// Streaming lifecycle — FlowAgent +// --------------------------------------------------------------------------- + +describe("FlowAgent streaming lifecycle (integration)", () => { + const Input = z.object({ x: z.number() }); + const Output = z.object({ y: z.number() }); + + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("fires onStart, onStepStart/Finish, onFinish during stream()", async () => { + const tracker = createLifecycleTracker(); + + const fa = flowAgent<{ x: number }, { y: number }>( + { + name: "stream-flow", + input: Input, + output: Output, + logger: createMockLogger(), + onStart: tracker.onStart, + onFinish: tracker.onFinish, + onStepStart: tracker.onStepStart, + onStepFinish: tracker.onStepFinish, + }, + async ({ input, $ }) => { + await $.step({ id: "compute", execute: async () => input.x * 2 }); + return { y: input.x * 2 }; + }, + ); + + const result = await fa.stream({ input: { x: 4 } }); + + expect(result.ok).toBe(true); + if (!result.ok) { + return; + } + + // Drain the stream + const parts: StreamPart[] = []; + for await (const part of result.fullStream) { + parts.push(part); + } + + await result.output; + + expect(tracker.events).toEqual([ + { type: "onStart" }, + { type: "onStepStart", detail: "compute" }, + { type: "onStepFinish", detail: "compute" }, + { type: "onFinish" }, + ]); + + // Stream should contain tool-call, tool-result, and finish events + const types = parts.map((p) => p.type); + expect(types).toContain("tool-call"); + expect(types).toContain("tool-result"); + expect(types).toContain("finish"); + }); + + it("fires hooks for nested steps during stream()", async () => { + const tracker = createLifecycleTracker(); + + const fa = flowAgent<{ x: number }, { y: number }>( + { + name: "nested-stream-flow", + input: Input, + output: Output, + logger: createMockLogger(), + onStepStart: tracker.onStepStart, + onStepFinish: tracker.onStepFinish, + }, + async ({ input, $: $outer }) => { + await $outer.step({ + id: "outer", + execute: async ({ $ }) => { + await $.step({ id: "inner", execute: async () => input.x + 1 }); + return input.x * 2; + }, + }); + return { y: input.x * 2 }; + }, + ); + + const result = await fa.stream({ input: { x: 3 } }); + + expect(result.ok).toBe(true); + if (!result.ok) { + return; + } + + for await (const _part of result.fullStream) { + /* Drain */ + } + await result.output; + + expect(tracker.events).toEqual([ + { type: "onStepStart", detail: "outer" }, + { type: "onStepStart", detail: "inner" }, + { type: "onStepFinish", detail: "inner" }, + { type: "onStepFinish", detail: "outer" }, + ]); + }); + + it("fires hooks for $.agent() during stream()", async () => { + const tracker = createLifecycleTracker(); + + const writer = agent({ + name: "writer", + model: createMockModel("streamed content"), + logger: createMockLogger(), + }); + + const fa = flowAgent<{ x: number }, { y: number }>( + { + name: "agent-stream-flow", + input: Input, + output: Output, + logger: createMockLogger(), + onStepStart: tracker.onStepStart, + onStepFinish: tracker.onStepFinish, + }, + async ({ input, $ }) => { + await $.agent({ + id: "write", + agent: writer, + input: `Write about ${input.x}`, + }); + return { y: input.x }; + }, + ); + + const result = await fa.stream({ input: { x: 7 } }); + + expect(result.ok).toBe(true); + if (!result.ok) { + return; + } + + for await (const _part of result.fullStream) { + /* Drain */ + } + await result.output; + + expect(tracker.events).toEqual([ + { type: "onStepStart", detail: "write" }, + { type: "onStepFinish", detail: "write" }, + ]); + }); + + it("stream emits tool-call and tool-result events for $.agent() with stream: true", async () => { + const writer = agent({ + name: "stream-writer", + model: createMockModel("streaming text"), + logger: createMockLogger(), + }); + + const fa = flowAgent<{ x: number }, { y: number }>( + { + name: "stream-agent-flow", + input: Input, + output: Output, + logger: createMockLogger(), + }, + async ({ input, $ }) => { + await $.agent({ + id: "write", + agent: writer, + input: `Write about ${input.x}`, + stream: true, + }); + return { y: input.x }; + }, + ); + + const result = await fa.stream({ input: { x: 1 } }); + + expect(result.ok).toBe(true); + if (!result.ok) { + return; + } + + const parts: StreamPart[] = []; + for await (const part of result.fullStream) { + parts.push(part); + } + + const types = parts.map((p) => p.type); + expect(types).toContain("tool-call"); + expect(types).toContain("text-delta"); + expect(types).toContain("tool-result"); + expect(types).toContain("finish"); + }); + + it("fires onError (not onFinish) when handler throws during stream()", async () => { + const tracker = createLifecycleTracker(); + + const fa = flowAgent<{ x: number }, { y: number }>( + { + name: "error-stream-flow", + input: Input, + output: Output, + logger: createMockLogger(), + onStart: tracker.onStart, + onFinish: tracker.onFinish, + onError: tracker.onError, + }, + async () => { + throw new Error("stream handler fail"); + }, + ); + + const result = await fa.stream({ input: { x: 1 } }); + + expect(result.ok).toBe(true); + if (!result.ok) { + return; + } + + // Suppress derived promise rejections + result.messages.catch(() => {}); + result.usage.catch(() => {}); + result.finishReason.catch(() => {}); + + for await (const _part of result.fullStream) { + /* Drain */ + } + + await result.output.catch(() => {}); + + expect(tracker.onStart).toHaveBeenCalledTimes(1); + expect(tracker.onError).toHaveBeenCalledTimes(1); + expect(tracker.onFinish).not.toHaveBeenCalled(); + }); +}); + +// --------------------------------------------------------------------------- +// Per-call hook override merging with flow agents +// --------------------------------------------------------------------------- + +describe("FlowAgent per-call hook merging (integration)", () => { + const Input = z.object({ x: z.number() }); + const Output = z.object({ y: z.number() }); + + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("merges config + per-call hooks for onStart, onStepStart, onStepFinish, onFinish", async () => { + const order: string[] = []; + + const fa = flowAgent<{ x: number }, { y: number }>( + { + name: "merge-flow", + input: Input, + output: Output, + logger: createMockLogger(), + onStart: () => { + order.push("config:onStart"); + }, + onFinish: () => { + order.push("config:onFinish"); + }, + onStepStart: ({ step }) => { + order.push(`config:onStepStart:${step.id}`); + }, + onStepFinish: ({ step }) => { + order.push(`config:onStepFinish:${step?.id}`); + }, + }, + async ({ input, $ }) => { + await $.step({ id: "work", execute: async () => input.x }); + return { y: input.x }; + }, + ); + + await fa.generate({ + input: { x: 1 }, + onStart: () => { + order.push("call:onStart"); + }, + onFinish: () => { + order.push("call:onFinish"); + }, + onStepStart: ({ step }) => { + order.push(`call:onStepStart:${step.id}`); + }, + onStepFinish: ({ step }) => { + order.push(`call:onStepFinish:${step?.id}`); + }, + }); + + expect(order).toEqual([ + "config:onStart", + "call:onStart", + "config:onStepStart:work", + "call:onStepStart:work", + "config:onStepFinish:work", + "call:onStepFinish:work", + "config:onFinish", + "call:onFinish", + ]); + }); +}); + +// --------------------------------------------------------------------------- +// Value forwarding — logger, signal, config +// --------------------------------------------------------------------------- + +describe("Value forwarding (integration)", () => { + const Input = z.object({ text: z.string() }); + const Output = z.object({ result: z.string() }); + + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("forwards scoped child logger to $.agent() sub-agent", async () => { + const childCalls: Record[] = []; + + // Track every .child() call and what bindings it receives + function createTrackingLogger(): ReturnType { + const logger = createMockLogger(); + const originalChild = logger.child as ReturnType; + originalChild.mockImplementation((bindings: Record) => { + childCalls.push(bindings); + return createTrackingLogger(); + }); + return logger; + } + + const trackingLogger = createTrackingLogger(); + + const core = agent({ + name: "core-agent", + model: createMockModel("core output"), + // No logger here — it should receive one from the flow + }); + + const fa = flowAgent<{ text: string }, { result: string }>( + { + name: "logger-flow", + input: Input, + output: Output, + logger: trackingLogger, + }, + async ({ input, $ }) => { + await $.agent({ + id: "run-core", + agent: core, + input: input.text, + }); + return { result: "done" }; + }, + ); + + await fa.generate({ input: { text: "hello" } }); + + // Flow agent creates child({ flowAgentId: 'logger-flow' }) + expect(childCalls).toContainEqual({ flowAgentId: "logger-flow" }); + + // $.agent() creates child({ stepId: 'run-core' }) for the step context + expect(childCalls).toContainEqual({ stepId: "run-core" }); + + // The sub-agent creates child({ agentId: 'core-agent' }) from the forwarded logger + expect(childCalls).toContainEqual({ agentId: "core-agent" }); + }); + + it("forwards scoped child logger through nested steps", async () => { + const childCalls: Record[] = []; + + function createTrackingLogger(): ReturnType { + const logger = createMockLogger(); + const originalChild = logger.child as ReturnType; + originalChild.mockImplementation((bindings: Record) => { + childCalls.push(bindings); + return createTrackingLogger(); + }); + return logger; + } + + const fa = flowAgent<{ text: string }, { result: string }>( + { + name: "nested-logger-flow", + input: Input, + output: Output, + logger: createTrackingLogger(), + }, + async ({ $: $outer, input }) => { + await $outer.step({ + id: "outer", + execute: async ({ $ }) => { + await $.step({ id: "inner", execute: async () => "done" }); + return "ok"; + }, + }); + return { result: input.text }; + }, + ); + + await fa.generate({ input: { text: "hello" } }); + + expect(childCalls).toContainEqual({ flowAgentId: "nested-logger-flow" }); + expect(childCalls).toContainEqual({ stepId: "outer" }); + expect(childCalls).toContainEqual({ stepId: "inner" }); + }); + + it("forwards ctx.signal to $.agent() sub-agent by default", async () => { + let receivedSignal: AbortSignal | undefined; + + const spyAgent = agent({ + name: "spy-agent", + model: createMockModel("ok"), + logger: createMockLogger(), + }); + + const originalGenerate = spyAgent.generate.bind(spyAgent); + spyAgent.generate = async (params) => { + receivedSignal = params.signal; + return originalGenerate(params); + }; + + const flowSignal = new AbortController().signal; + + const fa = flowAgent<{ text: string }, { result: string }>( + { + name: "signal-flow", + input: Input, + output: Output, + logger: createMockLogger(), + }, + async ({ $, input }) => { + await $.agent({ + id: "spy", + agent: spyAgent, + input: input.text, + }); + return { result: "done" }; + }, + ); + + await fa.generate({ input: { text: "test" }, signal: flowSignal }); + + // When no config.signal override, the flow's signal is forwarded + expect(receivedSignal).toBe(flowSignal); + }); + + it("framework ctx.signal takes precedence over config.signal override on $.agent()", async () => { + let receivedSignal: AbortSignal | undefined; + + const spyAgent = agent({ + name: "spy-agent", + model: createMockModel("ok"), + logger: createMockLogger(), + }); + + const originalGenerate = spyAgent.generate.bind(spyAgent); + spyAgent.generate = async (params) => { + receivedSignal = params.signal; + return originalGenerate(params); + }; + + const flowSignal = new AbortController().signal; + const overrideSignal = new AbortController().signal; + + const fa = flowAgent<{ text: string }, { result: string }>( + { + name: "signal-override-flow", + input: Input, + output: Output, + logger: createMockLogger(), + }, + async ({ $, input }) => { + await $.agent({ + id: "spy", + agent: spyAgent, + input: input.text, + config: { signal: overrideSignal }, + }); + return { result: "done" }; + }, + ); + + await fa.generate({ input: { text: "test" }, signal: flowSignal }); + + // Framework ctx.signal takes precedence over user-provided config.signal + expect(receivedSignal).toBe(flowSignal); + }); + + it("config.config model override is forwarded to $.agent() sub-agent", async () => { + const overrideModel = createMockModel("override output"); + const defaultModel = createMockModel("default output"); + + const core = agent({ + name: "model-agent", + model: defaultModel, + logger: createMockLogger(), + }); + + const fa = flowAgent<{ text: string }, { result: string }>( + { + name: "model-flow", + input: Input, + output: Output, + logger: createMockLogger(), + }, + async ({ $, input }) => { + const r = await $.agent({ + id: "run-core", + agent: core, + input: input.text, + config: { model: overrideModel }, + }); + if (r.ok) { + return { result: String(r.value.output) }; + } + return { result: "failed" }; + }, + ); + + const result = await fa.generate({ input: { text: "hello" } }); + + expect(result.ok).toBe(true); + if (!result.ok) { + return; + } + // The override model returns "override output" + expect(result.output.result).toBe("override output"); + }); + + it("per-call logger override on flowAgent.generate() is used instead of config logger", async () => { + const configChildCalls: Record[] = []; + const overrideChildCalls: Record[] = []; + + function createConfigLogger(): ReturnType { + const logger = createMockLogger(); + (logger.child as ReturnType).mockImplementation( + (bindings: Record) => { + configChildCalls.push(bindings); + return createConfigLogger(); + }, + ); + return logger; + } + + function createOverrideLogger(): ReturnType { + const logger = createMockLogger(); + (logger.child as ReturnType).mockImplementation( + (bindings: Record) => { + overrideChildCalls.push(bindings); + return createOverrideLogger(); + }, + ); + return logger; + } + + const fa = flowAgent<{ text: string }, { result: string }>( + { + name: "logger-override-flow", + input: Input, + output: Output, + logger: createConfigLogger(), + }, + async ({ $, input }) => { + await $.step({ id: "work", execute: async () => "done" }); + return { result: input.text }; + }, + ); + + await fa.generate({ input: { text: "test" }, logger: createOverrideLogger() }); + + // Config logger should NOT be used + expect(configChildCalls).toHaveLength(0); + + // Override logger should be used — creates child({ flowAgentId: ... }) + expect(overrideChildCalls).toContainEqual({ flowAgentId: "logger-override-flow" }); + }); +}); + +// --------------------------------------------------------------------------- +// BuildAgentTool logger forwarding — agent subagents (not flow) +// --------------------------------------------------------------------------- + +describe("Agent subagent logger forwarding (integration)", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("forwards parent logger to sub-agents wrapped as tools", async () => { + const childCalls: Record[] = []; + + function createTrackingLogger(): ReturnType { + const logger = createMockLogger(); + (logger.child as ReturnType).mockImplementation( + (bindings: Record) => { + childCalls.push(bindings); + return createTrackingLogger(); + }, + ); + return logger; + } + + const sub = agent({ + name: "sub-agent", + model: createMockModel("sub response"), + input: z.object({ task: z.string() }), + prompt: ({ input }) => input.task, + }); + + // Mock model that makes a tool call on the first step, then stops + const toolCallModel = new MockLanguageModelV3({ + // oxlint-disable-next-line @typescript-eslint/no-explicit-any -- mockValues sync/async mismatch + doGenerate: mockValues( + { + content: [ + { + type: "tool-call" as const, + toolCallId: "tc1", + toolName: "agent_sub", + input: JSON.stringify({ task: "do it" }), + }, + ], + finishReason: MOCK_FINISH, + usage: MOCK_USAGE, + warnings: [], + }, + { + content: [{ type: "text" as const, text: "done" }], + finishReason: MOCK_FINISH, + usage: MOCK_USAGE, + warnings: [], + }, + // oxlint-disable-next-line @typescript-eslint/no-explicit-any -- mockValues returns sync fn, MockLanguageModelV3 expects PromiseLike + ) as any, + }); + + const parent = agent({ + name: "parent-agent", + model: toolCallModel, + system: "You are a parent.", + agents: { sub }, + }); + + await parent.generate({ + prompt: "delegate this", + logger: createTrackingLogger(), + }); + + // Parent creates child({ agentId: 'parent-agent' }) + expect(childCalls).toContainEqual({ agentId: "parent-agent" }); + + // Sub-agent should receive the parent's logger and create child({ agentId: 'sub-agent' }) + // This verifies buildAgentTool forwards the logger + expect(childCalls).toContainEqual({ agentId: "sub-agent" }); + }); +}); + +// --------------------------------------------------------------------------- +// Agent subagent hook forwarding — parent hooks fire for sub-agent events +// --------------------------------------------------------------------------- + +describe("Agent subagent hook forwarding (integration)", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("parent onStepFinish fires for sub-agent internal steps", async () => { + const stepEvents: string[] = []; + + const sub = agent({ + name: "sub-agent", + model: createMockModel("sub response"), + input: z.object({ task: z.string() }), + prompt: ({ input }) => input.task, + }); + + const toolCallModel = new MockLanguageModelV3({ + // oxlint-disable-next-line @typescript-eslint/no-explicit-any -- mockValues sync/async mismatch + doGenerate: mockValues( + { + content: [ + { + type: "tool-call" as const, + toolCallId: "tc1", + toolName: "agent_sub", + input: JSON.stringify({ task: "do it" }), + }, + ], + finishReason: MOCK_FINISH, + usage: MOCK_USAGE, + warnings: [], + }, + { + content: [{ type: "text" as const, text: "done" }], + finishReason: MOCK_FINISH, + usage: MOCK_USAGE, + warnings: [], + }, + // oxlint-disable-next-line @typescript-eslint/no-explicit-any -- mockValues returns sync fn, MockLanguageModelV3 expects PromiseLike + ) as any, + }); + + const parent = agent({ + name: "parent-agent", + model: toolCallModel, + system: "Delegate to agent_sub.", + agents: { sub }, + onStepFinish: (event) => { + stepEvents.push(`config:${event.stepId}`); + }, + }); + + await parent.generate({ + prompt: "go", + logger: createMockLogger(), + onStepFinish: (event) => { + stepEvents.push(`call:${event.stepId}`); + }, + }); + + // Parent's own step (the tool call round) fires parent hooks + const parentSteps = stepEvents.filter((e) => e.includes("parent-agent")); + expect(parentSteps.length).toBeGreaterThan(0); + + // Sub-agent's internal step fires parent hooks too (forwarded) + const subSteps = stepEvents.filter((e) => e.includes("sub-agent")); + expect(subSteps.length).toBeGreaterThan(0); + + // Config hook fires before per-call hook for sub-agent steps + const subConfigIdx = stepEvents.findIndex((e) => e.startsWith("config:sub-agent")); + const subCallIdx = stepEvents.findIndex((e) => e.startsWith("call:sub-agent")); + if (subConfigIdx !== -1 && subCallIdx !== -1) { + expect(subConfigIdx).toBeLessThan(subCallIdx); + } + }); + + it("parent onStart does NOT fire for sub-agent events (type safety)", async () => { + const startEvents: unknown[] = []; + + const sub = agent({ + name: "sub-agent", + model: createMockModel("sub response"), + input: z.object({ task: z.string() }), + prompt: ({ input }) => input.task, + }); + + const toolCallModel = new MockLanguageModelV3({ + // oxlint-disable-next-line @typescript-eslint/no-explicit-any -- mockValues sync/async mismatch + doGenerate: mockValues( + { + content: [ + { + type: "tool-call" as const, + toolCallId: "tc1", + toolName: "agent_sub", + input: JSON.stringify({ task: "do it" }), + }, + ], + finishReason: MOCK_FINISH, + usage: MOCK_USAGE, + warnings: [], + }, + { + content: [{ type: "text" as const, text: "done" }], + finishReason: MOCK_FINISH, + usage: MOCK_USAGE, + warnings: [], + }, + // oxlint-disable-next-line @typescript-eslint/no-explicit-any -- mockValues returns sync fn, MockLanguageModelV3 expects PromiseLike + ) as any, + }); + + const parent = agent({ + name: "parent-agent", + model: toolCallModel, + system: "Delegate.", + agents: { sub }, + }); + + await parent.generate({ + prompt: "go", + logger: createMockLogger(), + onStart: (event) => { + startEvents.push(event.input); + }, + }); + + // OnStart fires ONLY for the parent's own generate() call. + // Generic hooks (onStart, onFinish, onError) are NOT forwarded to + // Sub-agents because their event types are parameterized by + // TInput/TOutput — forwarding would cause the parent's typed hook + // To receive the sub-agent's differently-shaped event at runtime. + expect(startEvents).toHaveLength(1); + expect(startEvents[0]).toBe("go"); + }); + + it("parent onError does NOT fire for sub-agent errors (type safety)", async () => { + const errorEvents: string[] = []; + + const failingModel = new MockLanguageModelV3({ + doGenerate: () => { + throw new Error("sub-agent model failure"); + }, + }); + + const sub = agent({ + name: "sub-agent", + model: failingModel, + input: z.object({ task: z.string() }), + prompt: ({ input }) => input.task, + }); + + const toolCallModel = new MockLanguageModelV3({ + // oxlint-disable-next-line @typescript-eslint/no-explicit-any -- mockValues sync/async mismatch + doGenerate: mockValues({ + content: [ + { + type: "tool-call" as const, + toolCallId: "tc1", + toolName: "agent_sub", + input: JSON.stringify({ task: "do it" }), + }, + ], + finishReason: MOCK_FINISH, + usage: MOCK_USAGE, + warnings: [], + // oxlint-disable-next-line @typescript-eslint/no-explicit-any -- mockValues returns sync fn, MockLanguageModelV3 expects PromiseLike + }) as any, + }); + + const parent = agent({ + name: "parent-agent", + model: toolCallModel, + system: "Delegate.", + agents: { sub }, + }); + + await parent.generate({ + prompt: "go", + logger: createMockLogger(), + onError: (event) => { + errorEvents.push(event.error.message); + }, + }); + + // Generic hooks (onError) are NOT forwarded to sub-agents to + // Preserve type safety. The sub-agent's tool execute() throws, + // But the AI SDK treats tool errors as tool results — the parent + // Never enters its own catch block for this, so onError does not fire. + expect(errorEvents).toHaveLength(0); + }); +}); + +// --------------------------------------------------------------------------- +// Step index uniqueness across nesting +// --------------------------------------------------------------------------- + +describe("Step index uniqueness (integration)", () => { + const Input = z.object({ n: z.number() }); + const Output = z.object({ count: z.number() }); + + it("step indices are globally unique across nested operations", async () => { + const indices: number[] = []; + + const fa = flowAgent<{ n: number }, { count: number }>( + { + name: "index-flow", + input: Input, + output: Output, + logger: createMockLogger(), + onStepStart: ({ step }) => { + indices.push(step.index); + }, + }, + async ({ input, $: $outer }) => { + await $outer.step({ + id: "a", + execute: async ({ $ }) => { + await $.step({ id: "b", execute: async () => 1 }); + await $.step({ id: "c", execute: async () => 2 }); + return 0; + }, + }); + await $outer.step({ id: "d", execute: async () => 3 }); + return { count: input.n }; + }, + ); + + await fa.generate({ input: { n: 4 } }); + + // All indices should be unique + const uniqueIndices = new Set(indices); + expect(uniqueIndices.size).toBe(indices.length); + + // Should be sequential: 0, 1, 2, 3 + expect(indices).toEqual([0, 1, 2, 3]); + }); +});