@tanstack/ai 0.32.0 → 0.34.0

package/src/activities/chat/index.ts

@@ -6,6 +6,7 @@
  */
 
 import { devtoolsMiddleware } from '@tanstack/ai-event-client'
+import { undoNullWidening } from '@tanstack/ai-utils'
 import { stripToSpecMiddleware } from '../../strip-to-spec-middleware'
 import { streamToText } from '../../stream-to-response.js'
 import { resolveDebugOption } from '../../logger/resolve'
@@ -18,6 +19,7 @@ import {
   executeToolCalls,
 } from './tools/tool-calls'
 import {
+  convertSchemaForStructuredOutput,
   convertSchemaToJsonSchema,
   isStandardSchema,
   parseWithStandardSchema,
@@ -33,7 +35,11 @@ import type {
   ClientToolRequest,
   ToolResult,
 } from './tools/tool-calls'
-import type { AnyTextAdapter, StructuredOutputOptions } from './adapter'
+import type {
+  AnyTextAdapter,
+  StructuredOutputOptions,
+  StructuredOutputResult,
+} from './adapter'
 import type {
   AgentLoopStrategy,
   AnyTool,
@@ -418,11 +424,21 @@ interface TextEngineConfig<
    *   (used by runStreamingStructuredOutput). When false, chunks are
    *   consumed internally for middleware visibility but not yielded
    *   (used by runAgenticStructuredOutput).
-   * - validate: optional callback invoked AFTER the structured-output result
-   *   is captured but BEFORE the terminal hook fires. If it throws, the
-   *   engine records a `finalizationError` and fires `onError` instead of
-   *   `onFinish` (per spec §7.3). On success, the returned value is stored
-   *   as the validated result and retrievable via
+   * - normalize: optional schema-aware transform applied to the captured
+   *   structured-output object the moment it enters the engine — BEFORE it is
+   *   stored, validated, or yielded. Used to undo strict-mode null-widening
+   *   (`undoNullWidening`): strict schemas widen optional fields to
+   *   `required` + nullable so the provider returns `null` for an absent
+   *   optional, and this strips exactly those synthesized nulls while keeping
+   *   the ones a `.nullable()` field genuinely allows. Applied here (not in
+   *   the adapter) because the engine is the only layer holding the original
+   *   schema's null-widening map, and applying it at capture fixes BOTH the
+   *   streaming chunk and the Promise<T> result with one transform.
+   * - validate: optional callback invoked AFTER `normalize` and AFTER the
+   *   structured-output result is captured, but BEFORE the terminal hook
+   *   fires. If it throws, the engine records a `finalizationError` and fires
+   *   `onError` instead of `onFinish` (per spec §7.3). On success, the
+   *   returned value is stored as the validated result and retrievable via
    *   `getValidatedStructuredOutput()`. Used by `runAgenticStructuredOutput`
    *   to perform Standard Schema validation inside the engine.
    * - nativeCombined: when true, the adapter declared
@@ -437,6 +453,7 @@ interface TextEngineConfig<
   finalStructuredOutput?: {
     jsonSchema: JSONSchema
     yieldChunks: boolean
+    normalize?: (data: unknown) => unknown
     validate?: (data: unknown) => unknown
     nativeCombined?: boolean
   }
@@ -542,8 +559,8 @@ class TextEngine<
   // to carry, so the client matches it to the streaming text deltas.
   private combinedStructuredMessageId: string | null = null
   // Holds the validated value when `finalStructuredOutput.validate` is provided
-  // and succeeds. Distinct from `structuredOutputResult.data` (the raw,
-  // unvalidated payload from the structured-output.complete chunk).
+  // and succeeds. Distinct from `structuredOutputResult.data` (the normalized
+  // but unvalidated payload from the structured-output.complete chunk).
   private validatedStructuredOutput: unknown = undefined
   private hasValidatedStructuredOutput = false
   private finalizationError: {
@@ -554,6 +571,7 @@ class TextEngine<
   private readonly finalStructuredOutput?: {
     jsonSchema: JSONSchema
     yieldChunks: boolean
+    normalize?: (data: unknown) => unknown
     validate?: (data: unknown) => unknown
     nativeCombined?: boolean
   }
@@ -646,6 +664,7 @@ class TextEngine<
         this.deferredPromises.push(promise)
       },
       // Provider / adapter info
+      activity: 'chat',
       provider: config.adapter.name,
       model: config.params.model,
       source: 'server',
@@ -1197,6 +1216,23 @@ class TextEngine<
     }
   }
 
+  /**
+   * Tools available for execution this turn. The discovery tool is dropped
+   * from the advertised set (`this.tools`) once every lazy tool is discovered,
+   * but a model may still re-request discovery; this widens execution lookup
+   * to include it so such calls don't fail with "Unknown tool". Centralised so
+   * both execution sites (`processToolCalls` and `checkForPendingToolCalls`)
+   * stay in sync.
+   */
+  private resolveExecutableTools(
+    toolCalls: ReadonlyArray<ToolCall>,
+  ): ReadonlyArray<AnyTool> {
+    return this.lazyToolManager.getExecutableTools(
+      this.tools,
+      toolCalls.map((tc) => tc.function.name),
+    )
+  }
+
   private async *checkForPendingToolCalls(): AsyncGenerator<
     StreamChunk,
     ToolPhaseResult,
@@ -1245,7 +1281,7 @@ class TextEngine<
 
     const generator = executeToolCalls(
       executablePendingCalls,
-      this.tools,
+      this.resolveExecutableTools(executablePendingCalls),
       approvals,
       clientToolResults,
       (eventName, data) => this.createCustomEventChunk(eventName, data),
@@ -1407,7 +1443,7 @@ class TextEngine<
 
     const generator = executeToolCalls(
       executableToolCalls,
-      this.tools,
+      this.resolveExecutableTools(executableToolCalls),
       approvals,
       clientToolResults,
       (eventName, data) => this.createCustomEventChunk(eventName, data),
@@ -2065,15 +2101,29 @@ class TextEngine<
       // All narrowing below is via the discriminated-union `chunk.type`
       // — no `as` casts.
 
+      // The chunk forwarded to middleware/consumers. Replaced below only for
+      // the structured-output.complete event, whose `object` we normalize
+      // (un-widen) so streaming consumers see the same cleaned payload the
+      // Promise<T> path validates and returns.
+      let outboundChunk: StreamChunk = chunk
+
       if (
         chunk.type === EventType.CUSTOM &&
         chunk.name === 'structured-output.complete'
       ) {
         const parsed = readStructuredOutputCompleteValue(chunk.value)
         if (parsed) {
-          this.structuredOutputResult = {
-            data: parsed.object,
-            rawText: parsed.raw,
+          const object = this.finalStructuredOutput.normalize
+            ? this.finalStructuredOutput.normalize(parsed.object)
+            : parsed.object
+          this.structuredOutputResult = { data: object, rawText: parsed.raw }
+          // Rewrite the outbound event so the yielded chunk carries the
+          // normalized object (the original `chunk.value` still holds the
+          // widened one). Preserve every other field — `raw`, `reasoning` —
+          // by spreading the original value.
+          const value = chunk.value
+          if (object !== parsed.object && value && typeof value === 'object') {
+            outboundChunk = { ...chunk, value: { ...value, object } }
           }
         }
       }
@@ -2097,7 +2147,7 @@ class TextEngine<
       // 7b. Pipe through middleware
       const outputChunks = await this.middlewareRunner.runOnChunk(
         this.middlewareCtx,
-        chunk,
+        outboundChunk,
       )
 
       // 7c. Decide consumer visibility — only yieldChunks=true callers get them.
@@ -2254,7 +2304,14 @@ class TextEngine<
     } else {
       try {
         const parsed: unknown = JSON.parse(rawText)
-        this.structuredOutputResult = { data: parsed, rawText }
+        // Normalize (un-widen) before storing so the synthesized
+        // structured-output.complete chunk and the Promise<T> result both
+        // carry the cleaned payload. JSON.parse preserves provider nulls, so
+        // this is where native-combined output gets its widening undone.
+        const data = this.finalStructuredOutput.normalize
+          ? this.finalStructuredOutput.normalize(parsed)
+          : parsed
+        this.structuredOutputResult = { data, rawText }
       } catch (err: unknown) {
         const detail =
           rawText.slice(0, 200) + (rawText.length > 200 ? '...' : '')
@@ -2711,18 +2768,31 @@ async function runAgenticStructuredOutput<
 
   // Same strict-conversion as the streaming path (`forStructuredOutput: true`)
   // so the same Zod schema produces the same JSON Schema regardless of
-  // stream mode — Promise<T> and stream:true must not diverge here.
-  const jsonSchema = convertSchemaToJsonSchema(outputSchema, {
-    forStructuredOutput: true,
-  })
+  // stream mode — Promise<T> and stream:true must not diverge here. The same
+  // pass also records a `nullWideningMap`: optional fields are widened to
+  // `required` + nullable for the provider, which then returns `null` for an
+  // absent optional — a `null` the original `.optional()` (`T | undefined`)
+  // schema would otherwise reject. The map pinpoints exactly those synthesized
+  // nulls so `undoNullWidening` can drop them while preserving the ones a
+  // `.nullable()` field genuinely allows.
+  const { jsonSchema, nullWideningMap } =
+    convertSchemaForStructuredOutput(outputSchema)
   if (!jsonSchema) {
     throw new Error('Failed to convert output schema to JSON Schema')
   }
 
+  // Un-widening runs in the engine the moment the structured output is
+  // captured (`finalStructuredOutput.normalize`), so it applies uniformly to
+  // every adapter and to both stream modes — the engine is the only layer
+  // holding the schema's `nullWideningMap`. Validation then runs on the
+  // already-normalized data, so `validate` is a plain Standard Schema parse.
+  const normalize = (data: unknown): unknown =>
+    undoNullWidening(data, nullWideningMap)
+
   // Validation runs INSIDE the engine (per spec §7.3) so validation failures
   // route through the engine's terminal-hook chooser as `onError`. We pass a
   // `validate` callback when the schema is a Standard Schema; otherwise we
-  // pass through the raw data and the engine returns it unchanged.
+  // pass through the (normalized) data and the engine returns it unchanged.
   const validate = isStandardSchema(outputSchema)
     ? (data: unknown): unknown =>
         parseWithStandardSchema<InferSchemaType<TSchema>>(outputSchema, data)
@@ -2754,6 +2824,7 @@ async function runAgenticStructuredOutput<
       finalStructuredOutput: {
         jsonSchema,
         yieldChunks: false,
+        normalize,
         ...(validate ? { validate } : {}),
         ...(nativeCombined ? { nativeCombined: true } : {}),
       },
@@ -2861,7 +2932,7 @@ async function* fallbackStructuredOutputStream(
     timestamp,
   }
 
-  let result: { data: unknown; rawText: string }
+  let result: StructuredOutputResult<unknown>
   try {
     result = await adapter.structuredOutput(options)
   } catch (error) {
@@ -2917,6 +2988,12 @@ async function* fallbackStructuredOutputStream(
     model,
     timestamp,
     finishReason: 'stop',
+    // Forward adapter-reported token usage so consumers reading
+    // `RUN_FINISHED.usage` (and the engine's `runOnUsage` middleware hook) see
+    // it on the fallback path, mirroring the native streaming path. The
+    // conditional spread avoids emitting `usage: undefined` for adapters that
+    // don't report it. See #758.
+    ...(result.usage ? { usage: result.usage } : {}),
   }
 }
 
@@ -2926,17 +3003,23 @@ async function* fallbackStructuredOutputStream(
  * RUN_STARTED/RUN_FINISHED are suppressed; the structured-output finalization
  * step's pair brackets the run for the consumer.
  *
- * Schema validation is intentionally NOT run on this path — it is the
- * consumer's responsibility. The `structured-output.complete` CUSTOM event
- * is forwarded with the adapter-produced `value.object` as-is. This is a
- * deliberate asymmetry vs. `runAgenticStructuredOutput` (Promise<T> path),
- * which DOES run Standard Schema validation inside the engine and routes
- * validation failures through `onError`. The reason for the asymmetry:
+ * Standard Schema *validation* is intentionally NOT run on this path — it is
+ * the consumer's responsibility. This is a deliberate asymmetry vs.
+ * `runAgenticStructuredOutput` (Promise<T> path), which DOES validate inside
+ * the engine and routes validation failures through `onError`. The reason:
  * streaming consumers typically render partial JSON progressively (via
  * `parsePartialJSON` or `useChat`'s `partial` slot) and validate downstream
  * after assembly. Running validation server-side would force a hard error
  * on partial-by-design payloads. See `docs/structured-outputs/overview.md`.
  *
+ * Null-widening normalization, however, IS run on both paths: the
+ * `structured-output.complete` CUSTOM event is forwarded with its `value.object`
+ * already un-widened (synthesized strict-mode nulls dropped, genuine
+ * `.nullable()` nulls kept), so a consumer validating the assembled object
+ * against the original schema doesn't choke on a `null` for an `.optional()`
+ * field. Same `convertSchemaForStructuredOutput` pass and same
+ * `undoNullWidening` map as the Promise<T> path — the two must not diverge.
+ *
  * Pre-flight validation (missing schema, unconvertible schema) throws
  * synchronously at call time rather than as a yielded RUN_ERROR mid-stream —
  * those are programmer errors, not runtime conditions.
@@ -2954,14 +3037,17 @@ function runStreamingStructuredOutput<
   }
 
   // forStructuredOutput strict-converts the schema once at the activity
-  // boundary. Adapters can re-convert if their wire format diverges, but the
-  // default flow hands them a strict-ready schema.
-  const jsonSchema = convertSchemaToJsonSchema(outputSchema, {
-    forStructuredOutput: true,
-  })
+  // boundary, capturing the null-widening map so the engine can un-widen the
+  // provider's response before it reaches the consumer. Adapters can re-convert
+  // if their wire format diverges, but the default flow hands them a
+  // strict-ready schema.
+  const { jsonSchema, nullWideningMap } =
+    convertSchemaForStructuredOutput(outputSchema)
   if (!jsonSchema) {
     throw new Error('Failed to convert output schema to JSON Schema')
   }
+  const normalize = (data: unknown): unknown =>
+    undoNullWidening(data, nullWideningMap)
 
   // The implementation generator yields the broader internal type
   // (`StreamChunk | StructuredOutputCompleteEvent<T>`) so agent-loop
@@ -2972,6 +3058,7 @@ function runStreamingStructuredOutput<
   return runStreamingStructuredOutputImpl(
     options,
     jsonSchema,
+    normalize,
   ) as StructuredOutputStream<InferSchemaType<TSchema>>
 }
 
@@ -2997,6 +3084,7 @@ async function* runStreamingStructuredOutputImpl<
 >(
   options: TextActivityOptions<AnyTextAdapter, TSchema, true, TContext>,
   jsonSchema: NonNullable<ReturnType<typeof convertSchemaToJsonSchema>>,
+  normalize: (data: unknown) => unknown,
 ): StructuredOutputStreamInternal<InferSchemaType<TSchema>> {
   const {
     adapter,
@@ -3041,6 +3129,7 @@ async function* runStreamingStructuredOutputImpl<
       finalStructuredOutput: {
         jsonSchema,
         yieldChunks: true,
+        normalize,
         ...(nativeCombined ? { nativeCombined: true } : {}),
       },
     },
@@ -3055,9 +3144,10 @@ async function* runStreamingStructuredOutputImpl<
     await mcpManager.dispose()
   }
 
-  // Schema validation for the streaming variant remains the consumer's
-  // responsibility — they read the CUSTOM 'structured-output.complete' from
-  // the yielded stream. Matches pre-fix behavior.
+  // Standard Schema validation for the streaming variant remains the
+  // consumer's responsibility — they read the CUSTOM 'structured-output.complete'
+  // from the yielded stream. (Null-widening normalization, by contrast, already
+  // ran inside the engine via `normalize`, so the object they read is un-widened.)
   void outputSchema
 }
 

package/src/activities/chat/middleware/types.ts

@@ -80,6 +80,13 @@ export interface ChatMiddlewareContext<TContext = unknown> {
 
   // --- Provider / adapter info (immutable for the lifetime of the request) ---
 
+  /**
+   * Which activity this context describes — always `'chat'`. Present so the
+   * chat context structurally satisfies the base `GenerationMiddlewareContext`,
+   * letting an observe-only middleware authored against the base (e.g.
+   * `otelMiddleware`) run on both chat and media activities.
+   */
+  activity: 'chat'
   /** Provider name (e.g., 'openai', 'anthropic') */
   provider: string
   /** Model identifier (e.g., 'gpt-4o') */

package/src/activities/chat/tools/lazy-tool-manager.ts

@@ -1,7 +1,14 @@
 import { convertSchemaToJsonSchema } from './schema-converter'
-import type { Tool } from '../../../types'
+import type { AnyTool, Tool } from '../../../types'
 
-const DISCOVERY_TOOL_NAME = '__lazy__tool__discovery__'
+/**
+ * Name of the synthetic tool the LLM calls to discover lazy tools.
+ *
+ * Exported so callers building custom message-compaction / history-trimming
+ * logic can reference the discovery tool by constant instead of hard-coding
+ * the string (which is an internal contract that could change).
+ */
+export const DISCOVERY_TOOL_NAME = '__lazy__tool__discovery__'
 
 /**
  * Manages lazy tool discovery for the chat agent loop.
@@ -87,6 +94,35 @@ export class LazyToolManager {
     return active
   }
 
+  /**
+   * Returns the tools that should be available for *execution* this turn.
+   *
+   * This is the advertised set (`getActiveTools()`, passed in as `activeTools`)
+   * plus the discovery tool when a pending call references it but it is no
+   * longer advertised. Once every lazy tool has been discovered the discovery
+   * tool is dropped from the advertised set, but a model may still re-request
+   * discovery (long context / hallucination); keeping it executable lets that
+   * call return the schemas again instead of failing with "Unknown tool".
+   *
+   * The advertised set is intentionally left unchanged — only execution lookup
+   * is widened. Operates on the already-built `activeTools`: it must NOT call
+   * `getActiveTools()`, which would reset `hasNewDiscoveries` before the
+   * post-execution refresh check in the agent loop.
+   */
+  getExecutableTools(
+    activeTools: ReadonlyArray<AnyTool>,
+    pendingToolCallNames: ReadonlyArray<string>,
+  ): ReadonlyArray<AnyTool> {
+    if (
+      this.discoveryTool &&
+      pendingToolCallNames.includes(DISCOVERY_TOOL_NAME) &&
+      !activeTools.some((t) => t.name === DISCOVERY_TOOL_NAME)
+    ) {
+      return [...activeTools, this.discoveryTool]
+    }
+    return activeTools
+  }
+
   /**
    * Returns whether new tools have been discovered since the last getActiveTools() call.
    */
@@ -221,8 +257,14 @@ export class LazyToolManager {
         for (const name of args.toolNames) {
           const tool = lazyToolMap.get(name)
           if (tool) {
-            manager.discoveredTools.add(name)
-            manager.hasNewDiscoveries = true
+            // Only flag a refresh for genuinely new discoveries. Re-requesting
+            // an already-discovered tool still returns its schema below (the
+            // model asked for it), but must not trigger a redundant tool-list
+            // refresh + continue in the agent loop.
+            if (!manager.discoveredTools.has(name)) {
+              manager.discoveredTools.add(name)
+              manager.hasNewDiscoveries = true
+            }
             const jsonSchema = tool.inputSchema
               ? convertSchemaToJsonSchema(tool.inputSchema)
               : undefined