From ba0e987b5f0e8b994a227e2b8d0c92f24937dccf Mon Sep 17 00:00:00 2001 From: MarioCadenas Date: Thu, 3 Sep 2026 16:56:15 +0200 Subject: [PATCH 1/9] feat(agents): add output schema + generic createAgent MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add an optional `output?: z.ZodType` field to AgentDefinition and make createAgent generic so `createAgent({ output: Schema })` returns `AgentDefinition>`. Thread the type param through RunAgentResult (new `output?: TOutput`) and store the runtime schema on RegisteredAgent. No behavior yet — this is the typing surface only. Signed-off-by: MarioCadenas --- .../appkit/src/core/agent/create-agent.ts | 29 ++++++++++++--- packages/appkit/src/core/agent/run-agent.ts | 36 +++++++++++++------ packages/appkit/src/core/agent/types.ts | 21 ++++++++++- 3 files changed, 70 insertions(+), 16 deletions(-) diff --git a/packages/appkit/src/core/agent/create-agent.ts b/packages/appkit/src/core/agent/create-agent.ts index 67c589317..b946097c0 100644 --- a/packages/appkit/src/core/agent/create-agent.ts +++ b/packages/appkit/src/core/agent/create-agent.ts @@ -1,3 +1,5 @@ +import type { z } from "zod"; + import { ConfigurationError } from "../../errors"; import type { AgentDefinition } from "./types"; @@ -30,8 +32,25 @@ const AGENT_BRAND: unique symbol = Symbol.for("appkit.agent"); * }, * }); * ``` + * + * @example Structured output + * ```ts + * const classify = createAgent({ + * instructions: "Classify the ticket.", + * output: z.object({ category: z.string(), urgent: z.boolean() }), + * }); + * // In-process, the result is typed via z.infer: + * const { output } = await runAgent(classify, { messages: "..." }); + * output?.category; // string | undefined + * ``` */ -export function createAgent(def: AgentDefinition): AgentDefinition { +export function createAgent( + def: AgentDefinition> & { output: S }, +): AgentDefinition>; +export function createAgent(def: AgentDefinition): AgentDefinition; +export function createAgent( + def: AgentDefinition, +): AgentDefinition { detectCycles(def); // Non-enumerable + in-place: identity, JSON, and spread are unaffected. Object.defineProperty(def, AGENT_BRAND, { @@ -58,11 +77,11 @@ export function isCreatedAgent(value: unknown): value is AgentDefinition { * Walks the `agents: { ... }` sub-agent tree via DFS and throws if a cycle is * found. Cycles would cause infinite recursion at tool-invocation time. */ -function detectCycles(def: AgentDefinition): void { - const visiting = new Set(); - const visited = new Set(); +function detectCycles(def: AgentDefinition): void { + const visiting = new Set>(); + const visited = new Set>(); - const walk = (current: AgentDefinition, path: string[]): void => { + const walk = (current: AgentDefinition, path: string[]): void => { if (visited.has(current)) return; if (visiting.has(current)) { throw new ConfigurationError( diff --git a/packages/appkit/src/core/agent/run-agent.ts b/packages/appkit/src/core/agent/run-agent.ts index 57a659686..1e1526492 100644 --- a/packages/appkit/src/core/agent/run-agent.ts +++ b/packages/appkit/src/core/agent/run-agent.ts @@ -51,11 +51,17 @@ export interface RunAgentInput { plugins?: PluginData[]; } -export interface RunAgentResult { +export interface RunAgentResult { /** Aggregated text output from all `message_delta` events. */ text: string; /** Every event the adapter yielded, in order. Useful for inspection/tests. */ events: AgentEvent[]; + /** + * Parsed, schema-validated object — present only when the agent (or the + * per-call override) declared an `output` schema. Statically typed via + * `z.infer` when the agent was built with `createAgent({ output })`. + */ + output?: TOutput; } /** @@ -86,10 +92,10 @@ export interface RunAgentResult { * `PluginContext`) throw at standalone-init time with a clear "use * createApp instead" message — not mid-stream. */ -export async function runAgent( - def: AgentDefinition, +export async function runAgent( + def: AgentDefinition, input: RunAgentInput, -): Promise { +): Promise> { // Single shared cache for the whole call graph: parent + every nested // sub-agent dispatch share constructed plugin instances. Without this, // each nested `runAgent` would build its own cache, re-instantiate every @@ -97,14 +103,22 @@ export async function runAgent( // (e.g. query result caches, connection pools). const providerCache = new Map(); await initStandalonePlugins(input.plugins ?? [], providerCache); - return runAgentInternal(def, input, providerCache); + const { text, events } = await runAgentInternal(def, input, providerCache); + // Structured-output resolution is wired in a later commit; for now the + // typed `output` field is left undefined. + return { text, events }; +} + +interface RawRunResult { + text: string; + events: AgentEvent[]; } async function runAgentInternal( - def: AgentDefinition, + def: AgentDefinition, input: RunAgentInput, providerCache: Map, -): Promise { +): Promise { const adapter = await resolveAdapter(def); const messages = normalizeMessages(input.messages, def.instructions); const toolIndex = buildStandaloneToolIndex( @@ -263,7 +277,9 @@ async function initStandalonePlugins( } } -async function resolveAdapter(def: AgentDefinition): Promise { +async function resolveAdapter( + def: AgentDefinition, +): Promise { const { model } = def; if (!model) { const { DatabricksAdapter } = await import("../../agents/databricks"); @@ -344,7 +360,7 @@ type StandaloneEntry = * references throw a named "not registered" error via the proxy. */ function buildStandaloneToolIndex( - def: AgentDefinition, + def: AgentDefinition, plugins: PluginData[], providerCache: Map, ): Map { @@ -389,7 +405,7 @@ function buildStandaloneToolIndex( * directly (no construction at toolkit-call time). */ function resolveDefTools( - def: AgentDefinition, + def: AgentDefinition, plugins: PluginData[], providerCache: Map, ): AgentTools { diff --git a/packages/appkit/src/core/agent/types.ts b/packages/appkit/src/core/agent/types.ts index 53dae3c87..3d822b776 100644 --- a/packages/appkit/src/core/agent/types.ts +++ b/packages/appkit/src/core/agent/types.ts @@ -5,6 +5,7 @@ import type { ThreadStore, ToolAnnotations, } from "shared"; +import type { z } from "zod"; import type { GenerationParams } from "../../agents/databricks"; import type { McpHostPolicyConfig } from "../../connectors/mcp"; @@ -124,7 +125,7 @@ export type AgentTools = Record; */ export type AgentToolsFn = (plugins: Plugins) => AgentTools; -export interface AgentDefinition { +export interface AgentDefinition { /** * Stable identifier for the agent. **Optional and informational** — * when the definition is registered via `agents: { foo: def }` (code) or @@ -154,6 +155,17 @@ export interface AgentDefinition { default?: boolean; /** System prompt body. For markdown-loaded agents this is the file body. */ instructions: string; + /** + * Optional Zod schema the agent's final answer is validated against. When + * set, the agent returns a typed object instead of freeform text: the + * `/invocations` envelope gains a top-level `output_parsed` field, `/chat` + * emits a final `structured_output` SSE event, and in-process `runAgent` + * populates `RunAgentResult.output`. Prefer {@link createAgent} with an + * `output` schema so the in-process result is statically typed via + * `z.infer`. Code-config agents only — markdown `agent.md` agents cannot + * carry a schema. + */ + output?: z.ZodType; /** * Model adapter (or endpoint-name string sugar for * `DatabricksAdapter.fromServingEndpoint({ endpointName })`). Optional — @@ -399,6 +411,13 @@ export interface RegisteredAgent { generationParams?: GenerationParams; /** Mirrors `AgentDefinition.ephemeral` — skip thread persistence. */ ephemeral?: boolean; + /** + * Mirrors `AgentDefinition.output` — the Zod schema the final answer is + * validated against. Present only for agents that declared structured + * output. Untyped here (the registry `Map` is string-keyed); `z.infer` + * typing lives on `createAgent`/`runAgent`. + */ + output?: z.ZodType; /** * Resolved per-agent skill catalog (visibility + collision rules applied). * Present when any skill is visible to this agent; drives the always-on From cb12a2cf7043b012d0b56922d1541d5b0b57bcfa Mon Sep 17 00:00:00 2001 From: MarioCadenas Date: Thu, 3 Sep 2026 16:58:08 +0200 Subject: [PATCH 2/9] feat(agents): response_format + structuring pass in Databricks adapter MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add an optional `outputSchema` (JSON Schema) to AgentInput. When set and the completion is tool-free, DatabricksAdapter sends it as an OpenAI-compatible `response_format: { type: json_schema, strict: true }`, so the final text is the JSON to validate. Tool-having runs are unchanged (Claude rejects response_format + tools) and get re-formatted by a separate tool-free structuring pass — a second run() with no tools. On a 400 that names the param, strip response_format and retry once (the Zod validation upstream is the real guarantee). In structured mode the text-based tool-call fallback is skipped so array-typed JSON isn't misread as a tool call. Signed-off-by: MarioCadenas --- packages/appkit/src/agents/databricks.ts | 85 ++++++++++++++++++++++-- packages/shared/src/agent.ts | 11 +++ 2 files changed, 89 insertions(+), 7 deletions(-) diff --git a/packages/appkit/src/agents/databricks.ts b/packages/appkit/src/agents/databricks.ts index 097f4fdba..d974c2e47 100644 --- a/packages/appkit/src/agents/databricks.ts +++ b/packages/appkit/src/agents/databricks.ts @@ -70,6 +70,28 @@ function applyGenerationParams( } } +/** + * True when `err` looks like an HTTP 400 that names `response_format` / + * `json_schema`. Used to strip-and-retry when a serving endpoint (or the + * gateway in front of it) doesn't support structured output — some + * Llama/DBRX/Mistral endpoints reject the param outright. The Zod boundary + * in the structured-output resolver is the real guarantee; `response_format` + * is only the retry-rate optimization, so silently dropping it on a 400 is + * safe. Deliberately narrow: a 400 that does NOT name the param is a real + * request error and must propagate. + */ +function isResponseFormatRejection(err: unknown): boolean { + const status = + isRecord(err) && typeof err.status === "number" + ? err.status + : isRecord(err) && typeof err.statusCode === "number" + ? err.statusCode + : undefined; + const msg = err instanceof Error ? err.message : String(err); + const is400 = status === 400 || /\b400\b/.test(msg); + return is400 && /response_format|json[_ ]schema/i.test(msg); +} + function extractLlamaToolJsonSlice(text: string): string | undefined { const start = text.indexOf("[{"); if (start < 0) return undefined; @@ -478,6 +500,15 @@ export class DatabricksAdapter implements AgentAdapter { const tools = this.buildTools(input.tools, nameToWire); const messages = this.buildMessages(input.messages, nameToWire); + // Structured output is only sent inline for tool-free completions — + // Databricks Claude endpoints reject `response_format` alongside `tools`. + // Tool-having agents produce prose here and are re-formatted by a + // separate tool-free structuring pass (a second `run()` with no tools). + const structuredSchema = + input.outputSchema && input.tools.length === 0 + ? input.outputSchema + : undefined; + yield { type: "status", status: "running" }; for (let step = 0; step < this.maxSteps; step++) { @@ -487,13 +518,19 @@ export class DatabricksAdapter implements AgentAdapter { messages, tools, context, + structuredSchema, ); if (toolCalls.length === 0) { - const parsed = parseTextToolCalls(text); - if (parsed.length > 0) { - yield* this.executeToolCalls(parsed, messages, context, nameToWire); - continue; + // In structured mode the model returns raw JSON (an array-typed schema + // emits `[{...}]`, which the Llama/Python text-tool-call fallback would + // otherwise misread as a tool invocation). Skip the fallback entirely. + if (!structuredSchema) { + const parsed = parseTextToolCalls(text); + if (parsed.length > 0) { + yield* this.executeToolCalls(parsed, messages, context, nameToWire); + continue; + } } break; } @@ -563,6 +600,7 @@ export class DatabricksAdapter implements AgentAdapter { messages: OpenAIMessage[], tools: OpenAITool[], context: AgentRunContext, + structuredSchema?: Record, ): AsyncGenerator< AgentEvent, { text: string; toolCalls: OpenAIToolCall[] }, @@ -578,15 +616,48 @@ export class DatabricksAdapter implements AgentAdapter { if (tools.length > 0) { body.tools = tools; + } else if (structuredSchema) { + // OpenAI-compatible structured output. `strict: true` asks the endpoint + // to constrain generation to the schema; the caller still validates + // with Zod, so a silently-ignored `response_format` is handled upstream. + body.response_format = { + type: "json_schema", + json_schema: { + name: "structured_output", + schema: structuredSchema, + strict: true, + }, + }; } let responseBody: ReadableStream; try { responseBody = await this.streamBody(body, context.signal); } catch (err) { - const msg = err instanceof Error ? err.message : "Stream request failed"; - yield { type: "status", status: "error", error: msg }; - throw err; + // Some endpoints/gateways 400 on `response_format`. Strip it and retry + // once — the structured-output resolver's Zod validation is the real + // guarantee. A 400 that doesn't name the param is a genuine error. + if ( + body.response_format !== undefined && + isResponseFormatRejection(err) + ) { + delete body.response_format; + try { + responseBody = await this.streamBody(body, context.signal); + } catch (retryErr) { + const msg = + retryErr instanceof Error + ? retryErr.message + : "Stream request failed"; + yield { type: "status", status: "error", error: msg }; + throw retryErr; + } + } else { + const msg = + err instanceof Error ? err.message : "Stream request failed"; + yield { type: "status", status: "error", error: msg }; + throw err; + } } const reader = responseBody.getReader(); diff --git a/packages/shared/src/agent.ts b/packages/shared/src/agent.ts index 811092b84..02cd857f7 100644 --- a/packages/shared/src/agent.ts +++ b/packages/shared/src/agent.ts @@ -275,6 +275,17 @@ export interface AgentInput { tools: AgentToolDefinition[]; threadId: string; signal?: AbortSignal; + /** + * JSON Schema the adapter should constrain a tool-free completion to, when + * it supports server-side structured output (e.g. an OpenAI-compatible + * `response_format: { type: "json_schema" }`). Set by the structured-output + * path; ignored when `tools` is non-empty (Databricks Claude endpoints + * reject `response_format` combined with `tools`). Adapters that can't + * constrain output ignore it — the structured-output resolver then relies + * on prompt + Zod validation. Already stripped of the top-level `$schema` + * key by {@link toToolJSONSchema}. + */ + outputSchema?: Record; /** * Adapter-specific opaque payloads, keyed by adapter namespace. The * shared contract intentionally does not enumerate keys — see each From 2432c6cee156cc242550847608573fd99f015a7b Mon Sep 17 00:00:00 2001 From: MarioCadenas Date: Thu, 3 Sep 2026 17:02:45 +0200 Subject: [PATCH 3/9] feat(agents): validate + retry + StructuredOutputError Add resolveStructuredOutput (core/agent): parse + zod-validate the agent's answer, re-prompting a tool-free structuring pass with the flattened zod issues up to 2 times before throwing StructuredOutputError with the last raw output. Tool-free answers are validated inline; tool-having answers get a first structuring pass. Strips a stray Markdown code fence before parsing. Wire the in-process runAgent to it via a new 3rd { output } options arg (overrides the agent's own schema), reusing the adapter for the tool-free structuring pass. StructuredOutputError joins the error taxonomy; lastRaw is server-side only and never leaks via clientMessage. Signed-off-by: MarioCadenas --- packages/appkit/src/beta.ts | 1 + packages/appkit/src/core/agent/run-agent.ts | 79 +++++++- .../src/core/agent/structured-output.ts | 179 ++++++++++++++++++ packages/appkit/src/errors/index.ts | 1 + .../appkit/src/errors/structured-output.ts | 50 +++++ packages/appkit/src/index.ts | 1 + 6 files changed, 306 insertions(+), 5 deletions(-) create mode 100644 packages/appkit/src/core/agent/structured-output.ts create mode 100644 packages/appkit/src/errors/structured-output.ts diff --git a/packages/appkit/src/beta.ts b/packages/appkit/src/beta.ts index 4de7ba79c..acfdefacf 100644 --- a/packages/appkit/src/beta.ts +++ b/packages/appkit/src/beta.ts @@ -42,6 +42,7 @@ export { export { createAgent } from "./core/agent/create-agent"; export { type RunAgentInput, + type RunAgentOptions, type RunAgentResult, runAgent, } from "./core/agent/run-agent"; diff --git a/packages/appkit/src/core/agent/run-agent.ts b/packages/appkit/src/core/agent/run-agent.ts index 1e1526492..278699b0d 100644 --- a/packages/appkit/src/core/agent/run-agent.ts +++ b/packages/appkit/src/core/agent/run-agent.ts @@ -9,6 +9,7 @@ import type { PluginData, ToolProvider, } from "shared"; +import type { z } from "zod"; import { isSupervisorTool, @@ -18,6 +19,10 @@ import { import { createLogger } from "../../logging/logger"; import { consumeAdapterStream } from "./consume-adapter-stream"; import { createPluginsProxy } from "./plugins-map"; +import { + resolveStructuredOutput, + type StructuringPass, +} from "./structured-output"; import { resolveToolkitFromProvider } from "./toolkit-resolver"; import { type FunctionTool, @@ -25,6 +30,7 @@ import { isFunctionTool, } from "./tools/function-tool"; import { isHostedTool } from "./tools/hosted-tools"; +import { toToolJSONSchema } from "./tools/json-schema"; import type { AgentDefinition, AgentTool, @@ -51,6 +57,17 @@ export interface RunAgentInput { plugins?: PluginData[]; } +/** Per-call options for {@link runAgent}. */ +export interface RunAgentOptions { + /** + * Structured-output schema override for this call. Takes precedence over the + * agent's own `output` schema. When either is set, `runAgent` validates the + * final answer and populates {@link RunAgentResult.output}, throwing a + * `StructuredOutputError` if it can't produce a valid object. + */ + output?: z.ZodType; +} + export interface RunAgentResult { /** Aggregated text output from all `message_delta` events. */ text: string; @@ -95,6 +112,7 @@ export interface RunAgentResult { export async function runAgent( def: AgentDefinition, input: RunAgentInput, + options?: RunAgentOptions, ): Promise> { // Single shared cache for the whole call graph: parent + every nested // sub-agent dispatch share constructed plugin instances. Without this, @@ -103,15 +121,60 @@ export async function runAgent( // (e.g. query result caches, connection pools). const providerCache = new Map(); await initStandalonePlugins(input.plugins ?? [], providerCache); - const { text, events } = await runAgentInternal(def, input, providerCache); - // Structured-output resolution is wired in a later commit; for now the - // typed `output` field is left undefined. - return { text, events }; + const { text, events, adapter, hadTools, baseMessages } = + await runAgentInternal(def, input, providerCache); + + const schema = options?.output ?? def.output; + if (!schema) return { text, events }; + + const output = await resolveStructuredOutput({ + schema, + baseMessages, + finalText: text, + hadTools, + runStructuringPass: buildStructuringPass(adapter, schema), + signal: input.signal, + }); + return { text, events, output }; +} + +/** + * Builds a {@link StructuringPass}: one tool-free, schema-constrained + * `adapter.run()`, consumed to its final text. `executeTool` throws — a + * tool-free run never dispatches one; if it somehow does, that's a bug we + * want surfaced, not swallowed. + */ +function buildStructuringPass( + adapter: AgentAdapter, + schema: z.ZodType, +): StructuringPass { + const outputSchema = toToolJSONSchema(schema); + return (messages, signal) => + consumeAdapterStream( + adapter.run( + { messages, tools: [], threadId: randomUUID(), signal, outputSchema }, + { + executeTool: () => { + throw new Error( + "runAgent: structuring pass is tool-free and must not call a tool", + ); + }, + signal, + }, + ), + { signal }, + ); } interface RawRunResult { text: string; events: AgentEvent[]; + /** Adapter used for the run — reused for the structuring pass. */ + adapter: AgentAdapter; + /** Whether the run exposed tools (tool-having answers need a structuring pass). */ + hadTools: boolean; + /** Normalized system + input messages the run saw (structuring-pass seed). */ + baseMessages: Message[]; } async function runAgentInternal( @@ -207,7 +270,13 @@ async function runAgentInternal( }, }); - return { text, events }; + return { + text, + events, + adapter, + hadTools: tools.length > 0, + baseMessages: messages, + }; } /** diff --git a/packages/appkit/src/core/agent/structured-output.ts b/packages/appkit/src/core/agent/structured-output.ts new file mode 100644 index 000000000..c043c201f --- /dev/null +++ b/packages/appkit/src/core/agent/structured-output.ts @@ -0,0 +1,179 @@ +import { randomUUID } from "node:crypto"; + +import type { Message } from "shared"; +import type { z } from "zod"; + +import { StructuredOutputError } from "../../errors"; + +/** + * Max number of re-prompted structuring passes after the first validation + * fails. So ≤3 total validation attempts, and ≤3 structuring passes for a + * tool-having agent (1 initial + 2 retries) / ≤2 for a tool-free one (the + * first attempt is the inline `response_format` output, retries add passes). + * Tracked here, separate from the tool-call / maxSteps budget. + */ +const MAX_VALIDATION_RETRIES = 2; + +const CONVERT_INSTRUCTION = + "Convert the assistant's answer above into JSON matching the provided schema. " + + "Output only the JSON, with no prose, explanation, or code fences."; + +function retryInstruction(errors: string): string { + return ( + "Your previous response did not match the required schema. " + + `Validation errors: ${errors}. ` + + "Return only JSON matching the schema, with no prose or code fences." + ); +} + +/** + * Runs ONE tool-free, schema-constrained completion over `messages` and + * returns the raw model text. Injected by the caller so the resolver stays + * free of any adapter / MLflow dependency: the agents plugin and standalone + * `runAgent` each build this from `adapter.run({ tools: [], outputSchema })`. + */ +export type StructuringPass = ( + messages: Message[], + signal?: AbortSignal, +) => Promise; + +interface ResolveStructuredOutputParams { + /** Schema the final object is validated against. */ + schema: z.ZodType; + /** + * The conversation the main run saw (system + thread messages), WITHOUT the + * final answer. Each structuring pass appends the latest answer + an + * instruction to this. + */ + baseMessages: Message[]; + /** The main run's final assistant text (already JSON when `hadTools` is false). */ + finalText: string; + /** + * Whether the main run used tools. Tool-having runs produce prose, so the + * first attempt is a structuring pass; tool-free runs already emitted JSON + * inline (via `response_format`), so `finalText` is validated directly. + */ + hadTools: boolean; + /** Runs one tool-free, schema-constrained completion (see {@link StructuringPass}). */ + runStructuringPass: StructuringPass; + signal?: AbortSignal; +} + +/** + * Validate-and-retry loop that turns an agent's answer into a typed object. + * + * - Tool-free agent: `finalText` is already `response_format` JSON, validated + * directly; on failure a re-prompted structuring pass runs. + * - Tool-having agent: the answer is prose, so a structuring pass reformats it + * into JSON first. + * + * On every failure it re-prompts with the flattened Zod issues (up to + * {@link MAX_VALIDATION_RETRIES} times), then throws {@link StructuredOutputError} + * carrying the last raw output. Never returns partial/unvalidated data. + */ +export async function resolveStructuredOutput( + params: ResolveStructuredOutputParams, +): Promise { + const { + schema, + baseMessages, + finalText, + hadTools, + runStructuringPass, + signal, + } = params; + + let lastRaw = hadTools + ? await runStructuringPass( + structuringMessages(baseMessages, finalText, CONVERT_INSTRUCTION), + signal, + ) + : finalText; + + for (let retries = 0; ; retries++) { + const parsed = parseAndValidate(schema, lastRaw); + if (parsed.ok) return parsed.value; + + if (retries >= MAX_VALIDATION_RETRIES) { + throw new StructuredOutputError( + `Agent output did not match the required schema after ` + + `${MAX_VALIDATION_RETRIES + 1} attempts: ${parsed.error}`, + { lastRaw }, + ); + } + + lastRaw = await runStructuringPass( + structuringMessages( + baseMessages, + lastRaw, + retryInstruction(parsed.error), + ), + signal, + ); + } +} + +/** Appends the latest answer (as an assistant turn) + a user instruction. */ +function structuringMessages( + base: Message[], + assistantText: string, + instruction: string, +): Message[] { + return [ + ...base, + { + id: randomUUID(), + role: "assistant", + content: assistantText, + createdAt: new Date(), + }, + { + id: randomUUID(), + role: "user", + content: instruction, + createdAt: new Date(), + }, + ]; +} + +type ParseResult = { ok: true; value: T } | { ok: false; error: string }; + +function parseAndValidate( + schema: z.ZodType, + raw: string, +): ParseResult { + let json: unknown; + try { + json = JSON.parse(stripCodeFences(raw)); + } catch { + return { ok: false, error: "response was not valid JSON" }; + } + const result = schema.safeParse(json); + if (result.success) return { ok: true, value: result.data }; + return { ok: false, error: flattenZodError(result.error) }; +} + +/** + * Strip a single leading/trailing Markdown code fence. `response_format` + * output is bare JSON, but on the 400-strip fallback (or a model that ignores + * the param) the answer often arrives fenced (` ```json … ``` `). Cheap to + * undo and materially raises the parse rate on that path. + */ +function stripCodeFences(text: string): string { + const trimmed = text.trim(); + if (!trimmed.startsWith("```")) return trimmed; + return trimmed + .replace(/^```(?:json)?\s*\n?/i, "") + .replace(/\n?```$/, "") + .trim(); +} + +/** Compact one-line rendering of a ZodError's issues, safe across zod v4. */ +function flattenZodError(error: z.ZodError): string { + return error.issues + .map((issue) => { + const path = issue.path.join(".") || "(root)"; + return `${path}: ${issue.message}`; + }) + .join("; "); +} diff --git a/packages/appkit/src/errors/index.ts b/packages/appkit/src/errors/index.ts index a367b843c..25dbb9c69 100644 --- a/packages/appkit/src/errors/index.ts +++ b/packages/appkit/src/errors/index.ts @@ -26,5 +26,6 @@ export { ConnectionError } from "./connection"; export { ExecutionError } from "./execution"; export { InitializationError } from "./initialization"; export { ServerError } from "./server"; +export { StructuredOutputError } from "./structured-output"; export { TunnelError } from "./tunnel"; export { ValidationError } from "./validation"; diff --git a/packages/appkit/src/errors/structured-output.ts b/packages/appkit/src/errors/structured-output.ts new file mode 100644 index 000000000..ad1eebd8a --- /dev/null +++ b/packages/appkit/src/errors/structured-output.ts @@ -0,0 +1,50 @@ +import { AppKitError } from "./base"; + +/** + * Thrown when an agent with an `output` schema could not produce a value that + * validates against it, even after the structuring retries are exhausted. + * Carries the last raw model output (server-side only) for debugging — it is + * never returned to the client, which sees the generic {@link clientMessage}. + * + * The structured-output path throws this rather than returning partial or + * unvalidated data: a caller that asked for a typed object gets either a valid + * one or an error, never a half-parsed shape. + * + * @example + * ```typescript + * try { + * const { output } = await runAgent(classifier, { messages: "..." }); + * } catch (e) { + * if (e instanceof StructuredOutputError) { + * // model couldn't produce schema-valid JSON; e.lastRaw has the attempt + * } + * } + * ``` + */ +export class StructuredOutputError extends AppKitError { + readonly code = "STRUCTURED_OUTPUT_ERROR"; + readonly statusCode = 500; + readonly isRetryable = false; + + /** + * Last raw model output that failed to parse/validate. Server-side only — + * kept off {@link clientMessage} since it can echo arbitrary model text. + */ + readonly lastRaw?: string; + + constructor( + message: string, + options?: { + cause?: Error; + context?: Record; + lastRaw?: string; + }, + ) { + super(message, options); + this.lastRaw = options?.lastRaw; + } + + override get clientMessage(): string { + return this._clientMessage ?? "The agent could not produce a valid result"; + } +} diff --git a/packages/appkit/src/index.ts b/packages/appkit/src/index.ts index 127cad384..f3bcded02 100644 --- a/packages/appkit/src/index.ts +++ b/packages/appkit/src/index.ts @@ -48,6 +48,7 @@ export { ExecutionError, InitializationError, ServerError, + StructuredOutputError, TunnelError, ValidationError, } from "./errors"; From 8b379742cc5831b130ba9d793d5ce856488d1d65 Mon Sep 17 00:00:00 2001 From: MarioCadenas Date: Thu, 3 Sep 2026 17:06:21 +0200 Subject: [PATCH 4/9] feat(agents): surface output_parsed (invoke) + structured_output SSE event Store the agent's output schema on RegisteredAgent and, when set, resolve structured output after the run's final text: /invocations and /responses gain a top-level output_parsed field (spread like mlflow_trace_id), and /chat emits one final structured_output AgentEvent -> appkit.structured_output SSE event after the streamed text. Retries stay invisible to the stream. The structuring pass runs as a fresh tool-free adapter.run() constrained by the schema, wrapped in a TOOL span nested under the turn's AGENT span. Adds the new event to the shared AgentEvent/ResponseStreamEvent unions and the event translator. Signed-off-by: MarioCadenas --- packages/appkit/src/plugins/agents/agents.ts | 95 +++++++++++++++++++ .../src/plugins/agents/event-translator.ts | 8 ++ packages/shared/src/agent.ts | 25 +++++ 3 files changed, 128 insertions(+) diff --git a/packages/appkit/src/plugins/agents/agents.ts b/packages/appkit/src/plugins/agents/agents.ts index 8ebbf32c3..c924d422b 100644 --- a/packages/appkit/src/plugins/agents/agents.ts +++ b/packages/appkit/src/plugins/agents/agents.ts @@ -27,6 +27,7 @@ import type { ResolvedSkillCatalog, SkillDefinition, } from "../../core/agent/skills"; +import { resolveStructuredOutput } from "../../core/agent/structured-output"; import { resolveToolkitFromProvider } from "../../core/agent/toolkit-resolver"; import { functionToolToDefinition, @@ -34,6 +35,7 @@ import { isHostedTool, resolveHostedTools, } from "../../core/agent/tools"; +import { toToolJSONSchema } from "../../core/agent/tools/json-schema"; import type { AgentDefinition, AgentsPluginConfig, @@ -65,6 +67,7 @@ import { initAgentTracing, linkTraceToRun, traceAgent, + traceTool, } from "./mlflow"; import { composePromptForAgent } from "./prompt"; import { @@ -493,6 +496,7 @@ export class AgentsPlugin extends Plugin implements ToolProvider { maxTokens: def.maxTokens, generationParams: def.generationParams, ephemeral: def.ephemeral, + output: def.output, skills, }; } @@ -1332,6 +1336,27 @@ export class AgentsPlugin extends Plugin implements ToolProvider { outboundEvents.push(evt); } } + + // Structured output: after the visible text, coerce + validate the + // answer against the agent's schema and emit one final + // `structured_output` event. Retries (if any) happen here, invisible + // to the streamed text. On exhaustion the throw propagates to the + // driver's catch and surfaces as an error event. + if (registered.output) { + const data = await this.runStructuredOutput( + registered, + messagesWithSystem, + fullContent, + tools.length > 0, + signal, + ); + for (const evt of translator.translate({ + type: "structured_output", + data, + })) { + outboundEvents.push(evt); + } + } }, ); @@ -1428,6 +1453,8 @@ export class AgentsPlugin extends Plugin implements ToolProvider { // Assigned inside the span below (the only place the active trace id // resolves), read into the response envelope after. let mlflowTraceId: string | undefined; + // Parsed structured output, when the agent declared an `output` schema. + let outputParsed: unknown; const runState: RunState = { req, @@ -1515,6 +1542,19 @@ export class AgentsPlugin extends Plugin implements ToolProvider { }); } + // Structured output: coerce + validate against the agent's schema + // and attach as the envelope's `output_parsed`. Throws on exhaustion + // (caught below as a 500) — never returns partial data. + if (registered.output) { + outputParsed = await this.runStructuredOutput( + registered, + messagesWithSystem, + fullContent, + tools.length > 0, + signal, + ); + } + mlflowTraceId = currentTraceId(); }, ); @@ -1571,10 +1611,65 @@ export class AgentsPlugin extends Plugin implements ToolProvider { ...(runState.toolErrors.length > 0 ? { tool_errors: runState.toolErrors } : {}), + // Parsed, schema-validated object when the agent declared an `output` + // schema — the non-streaming equivalent of the `structured_output` SSE + // event. Conditionally spread, like `mlflow_trace_id` / `tool_errors`. + ...(outputParsed !== undefined ? { output_parsed: outputParsed } : {}), output: [message], }); } + /** + * Resolves an agent's structured output from the run's final text. The + * structuring pass(es) run as a fresh tool-free `adapter.run()` constrained + * by the schema, wrapped in a TOOL span so they nest under the turn's AGENT + * span. Assumes `registered.output` is set. Throws `StructuredOutputError` + * if no schema-valid object can be produced within the retry budget. + */ + private runStructuredOutput( + registered: RegisteredAgent, + baseMessages: Message[], + finalText: string, + hadTools: boolean, + signal: AbortSignal, + ): Promise { + const schema = registered.output; + if (!schema) { + throw new Error("runStructuredOutput called without an output schema"); + } + const outputSchema = toToolJSONSchema(schema); + return traceTool("structured_output", { schema: outputSchema }, () => + resolveStructuredOutput({ + schema, + baseMessages, + finalText, + hadTools, + signal, + runStructuringPass: (messages, sig) => + consumeAdapterStream( + registered.adapter.run( + { + messages, + tools: [], + threadId: randomUUID(), + signal: sig, + outputSchema, + }, + { + executeTool: () => { + throw new Error( + "structured-output structuring pass is tool-free and must not call a tool", + ); + }, + signal: sig, + }, + ), + { signal: sig }, + ), + }), + ); + } + private dispatchSkillTool( entry: Extract, args: unknown, diff --git a/packages/appkit/src/plugins/agents/event-translator.ts b/packages/appkit/src/plugins/agents/event-translator.ts index ecaa2aab5..fb9980e70 100644 --- a/packages/appkit/src/plugins/agents/event-translator.ts +++ b/packages/appkit/src/plugins/agents/event-translator.ts @@ -60,6 +60,14 @@ export class AgentEventTranslator { sequence_number: this.seqNum++, }, ]; + case "structured_output": + return [ + { + type: "appkit.structured_output", + data: event.data, + sequence_number: this.seqNum++, + }, + ]; case "approval_pending": return [ { diff --git a/packages/shared/src/agent.ts b/packages/shared/src/agent.ts index 02cd857f7..e219fd7c8 100644 --- a/packages/shared/src/agent.ts +++ b/packages/shared/src/agent.ts @@ -129,6 +129,17 @@ export type AgentEvent = error?: string; } | { type: "metadata"; data: Record } + | { + /** + * Emitted by the agents plugin (not adapters) once, after the streamed + * text, when the agent declared an `output` schema: the parsed, + * schema-validated object. Delivered on the wire as + * `appkit.structured_output`. Non-streaming surfaces (`/invocations`) + * return the same object as the envelope's `output_parsed` field. + */ + type: "structured_output"; + data: unknown; + } | { /** * Emitted by the agents plugin (not adapters) when a mutating tool call @@ -245,6 +256,19 @@ export interface AppKitMetadataEvent { * arrives before the server-side timeout, the call is auto-denied and the * agent receives a denial string as the tool output. */ +/** + * Emitted once on `/chat`, after the streamed assistant text, when the agent + * declared an `output` schema. `data` is the parsed, schema-validated object. + * The `appkit.` prefix matches the other AppKit-injected wire events + * (`appkit.thinking`, `appkit.metadata`); the equivalent non-streaming field + * is `output_parsed`. + */ +export interface AppKitStructuredOutputEvent { + type: "appkit.structured_output"; + data: unknown; + sequence_number: number; +} + export interface AppKitApprovalPendingEvent { type: "appkit.approval_pending"; approval_id: string; @@ -264,6 +288,7 @@ export type ResponseStreamEvent = | ResponseFailedEvent | AppKitThinkingEvent | AppKitMetadataEvent + | AppKitStructuredOutputEvent | AppKitApprovalPendingEvent; // --------------------------------------------------------------------------- From 88a51891825a927b5665714bfa74f45696537114 Mon Sep 17 00:00:00 2001 From: MarioCadenas Date: Thu, 3 Sep 2026 17:09:09 +0200 Subject: [PATCH 5/9] feat(agents): send response_format inline on the tool-free main run MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Thread the agent's JSON schema into the main adapter.run() (runAgent + both plugin surfaces) so a tool-free structured agent gets response_format on its own completion — the answer is JSON directly, no wasted structuring round-trip. The adapter ignores outputSchema when tools are present, so tool-having agents are unaffected and still use the separate structuring pass. Completes the tool-branch strategy from the earlier adapter commit. Signed-off-by: MarioCadenas --- packages/appkit/src/core/agent/run-agent.ts | 11 +++++++++-- packages/appkit/src/plugins/agents/agents.ts | 10 ++++++++++ 2 files changed, 19 insertions(+), 2 deletions(-) diff --git a/packages/appkit/src/core/agent/run-agent.ts b/packages/appkit/src/core/agent/run-agent.ts index 278699b0d..a9bd2d1fb 100644 --- a/packages/appkit/src/core/agent/run-agent.ts +++ b/packages/appkit/src/core/agent/run-agent.ts @@ -121,10 +121,15 @@ export async function runAgent( // (e.g. query result caches, connection pools). const providerCache = new Map(); await initStandalonePlugins(input.plugins ?? [], providerCache); - const { text, events, adapter, hadTools, baseMessages } = - await runAgentInternal(def, input, providerCache); const schema = options?.output ?? def.output; + // Pass the schema into the main run so a tool-free agent gets + // `response_format` inline (no wasted round-trip); the adapter ignores it + // when tools are present. Sub-agent recursions never receive it. + const mainOutputSchema = schema ? toToolJSONSchema(schema) : undefined; + const { text, events, adapter, hadTools, baseMessages } = + await runAgentInternal(def, input, providerCache, mainOutputSchema); + if (!schema) return { text, events }; const output = await resolveStructuredOutput({ @@ -181,6 +186,7 @@ async function runAgentInternal( def: AgentDefinition, input: RunAgentInput, providerCache: Map, + mainOutputSchema?: Record, ): Promise { const adapter = await resolveAdapter(def); const messages = normalizeMessages(input.messages, def.instructions); @@ -256,6 +262,7 @@ async function runAgentInternal( threadId: randomUUID(), signal, extensions: buildStandaloneExtensions(toolIndex), + outputSchema: mainOutputSchema, }, { executeTool, signal }, ); diff --git a/packages/appkit/src/plugins/agents/agents.ts b/packages/appkit/src/plugins/agents/agents.ts index c924d422b..8f0b180c3 100644 --- a/packages/appkit/src/plugins/agents/agents.ts +++ b/packages/appkit/src/plugins/agents/agents.ts @@ -1299,6 +1299,11 @@ export class AgentsPlugin extends Plugin implements ToolProvider { threadId: thread.id, signal, extensions: buildAdapterExtensions(registered.toolIndex), + // Tool-free structured agents get `response_format` inline; the + // adapter ignores it when tools are present. + outputSchema: registered.output + ? toToolJSONSchema(registered.output) + : undefined, }, { executeTool, signal }, ); @@ -1526,6 +1531,11 @@ export class AgentsPlugin extends Plugin implements ToolProvider { tools, threadId: thread.id, signal, + // Tool-free structured agents get `response_format` inline; the + // adapter ignores it when tools are present. + outputSchema: registered.output + ? toToolJSONSchema(registered.output) + : undefined, }, { executeTool, signal }, ); From 6d65de193acc980d2078e08ac6086eee9c512ade Mon Sep 17 00:00:00 2001 From: MarioCadenas Date: Thu, 3 Sep 2026 17:19:59 +0200 Subject: [PATCH 6/9] feat(agents): overload runAgent so per-call { output } override is typed Split runAgent into two overloads: the no-override call types the result from the agent's own schema; the per-call { output } override drives the result type via z.infer of that schema (it also takes precedence at runtime). Without this, an agent with no schema is AgentDefinition, so a { output } override couldn't retype RunAgentResult. Signed-off-by: MarioCadenas --- packages/appkit/src/core/agent/run-agent.ts | 19 +++++++++++++++---- 1 file changed, 15 insertions(+), 4 deletions(-) diff --git a/packages/appkit/src/core/agent/run-agent.ts b/packages/appkit/src/core/agent/run-agent.ts index a9bd2d1fb..1a4223565 100644 --- a/packages/appkit/src/core/agent/run-agent.ts +++ b/packages/appkit/src/core/agent/run-agent.ts @@ -109,11 +109,22 @@ export interface RunAgentResult { * `PluginContext`) throw at standalone-init time with a clear "use * createApp instead" message — not mid-stream. */ -export async function runAgent( +// Per-call schema override drives the result type via z.infer. +export function runAgent( + def: AgentDefinition, + input: RunAgentInput, + options: RunAgentOptions> & { output: S }, +): Promise>>; +// No override — the result type comes from the agent's own `output` schema. +export function runAgent( def: AgentDefinition, input: RunAgentInput, - options?: RunAgentOptions, -): Promise> { +): Promise>; +export async function runAgent( + def: AgentDefinition, + input: RunAgentInput, + options?: RunAgentOptions, +): Promise> { // Single shared cache for the whole call graph: parent + every nested // sub-agent dispatch share constructed plugin instances. Without this, // each nested `runAgent` would build its own cache, re-instantiate every @@ -132,7 +143,7 @@ export async function runAgent( if (!schema) return { text, events }; - const output = await resolveStructuredOutput({ + const output = await resolveStructuredOutput({ schema, baseMessages, finalText: text, From 2c9278d956e2185715cb2683293d60fb41fdcbbe Mon Sep 17 00:00:00 2001 From: MarioCadenas Date: Thu, 3 Sep 2026 17:20:12 +0200 Subject: [PATCH 7/9] test(agents): structured output tests, docs, and dev-playground example MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Unit tests for the resolver (validate/retry/exhaust/fence-strip/no-leak), the adapter (response_format on tool-free runs, none with tools, 400 strip-retry), runAgent end-to-end (tool-free inline, tool-having structuring pass, per-call override, throw), and the structured_output event translation. Add a 'Structured output' section to the agents plugin doc and a tool-free 'classifier' example agent to dev-playground. Not verified against a live serving endpoint — needs manual dogfood. Signed-off-by: MarioCadenas --- .../server/agents/classifier/agent.ts | 36 ++++ docs/docs/plugins/agents.md | 41 +++++ .../src/agents/tests/databricks.test.ts | 165 ++++++++++++++++++ .../agent/tests/run-agent-structured.test.ts | 118 +++++++++++++ .../agent/tests/structured-output.test.ts | 119 +++++++++++++ .../agents/tests/event-translator.test.ts | 14 ++ 6 files changed, 493 insertions(+) create mode 100644 apps/dev-playground/server/agents/classifier/agent.ts create mode 100644 packages/appkit/src/core/agent/tests/run-agent-structured.test.ts create mode 100644 packages/appkit/src/core/agent/tests/structured-output.test.ts diff --git a/apps/dev-playground/server/agents/classifier/agent.ts b/apps/dev-playground/server/agents/classifier/agent.ts new file mode 100644 index 000000000..c0505bac7 --- /dev/null +++ b/apps/dev-playground/server/agents/classifier/agent.ts @@ -0,0 +1,36 @@ +import { createAgent } from "@databricks/appkit/beta"; +import { z } from "zod"; + +// Structured-output demo: a tool-free agent whose final answer is validated +// against a Zod schema instead of being returned as freeform text. Discovered +// automatically from server/agents/classifier/ (its id is the folder name, +// "classifier"). +// +// How the schema surfaces on each path: +// • POST /api/agents/chat with { message, agent: "classifier" } streams the +// text, then emits one final `appkit.structured_output` SSE event whose +// `data` is the parsed object. +// • POST /api/agents/invocations (when this is the default agent) returns a +// top-level `output_parsed` field alongside the usual `output` text. +// • In-process, the result is statically typed via z.infer: +// import { runAgent } from "@databricks/appkit/beta"; +// const { output } = await runAgent(classifier, { messages: ticket }); +// output?.category; // "billing" | "bug" | ... | undefined +// +// Because it declares no tools, AppKit sends the schema inline as the +// endpoint's `response_format`, so the answer comes back as JSON directly. +export default createAgent({ + instructions: + "You are a support-ticket triage classifier. Read the user's message and " + + "classify it into one category, decide whether it is urgent (the user is " + + "blocked or reports an outage), and write a one-sentence summary.", + output: z.object({ + category: z + .enum(["billing", "bug", "feature_request", "how_to", "other"]) + .describe("The single best-fit category for the ticket."), + urgent: z + .boolean() + .describe("True only if the user is blocked or reports an outage."), + summary: z.string().describe("A one-sentence summary of the request."), + }), +}); diff --git a/docs/docs/plugins/agents.md b/docs/docs/plugins/agents.md index 19b0f6da2..08504b7b9 100644 --- a/docs/docs/plugins/agents.md +++ b/docs/docs/plugins/agents.md @@ -313,6 +313,47 @@ const result = await runAgent(classifier, { MCP hosted tools (`mcpServer(...)`) still require `agents()` (they need a live MCP client). Supervisor-API hosted tools (`supervisorTools.*`), by contrast, **work in standalone `runAgent`** — the adapter has everything it needs to execute them server-side. This makes batch-eval / CI use of supervisor agents possible without `createApp`. Plugin tool dispatch in standalone mode runs as the service principal (no OBO) and **bypasses the agents-plugin approval gate** — treat standalone runAgent as a trusted-prompt environment (CI, batch eval, internal scripts), not as an exposed user-facing surface. +## Structured output + +Give a code agent an `output` Zod schema and its final answer is validated against that schema instead of returned as freeform text: + +```ts +import { createAgent } from "@databricks/appkit/beta"; +import { z } from "zod"; + +export default createAgent({ + instructions: "Classify the support ticket.", + output: z.object({ + category: z.enum(["billing", "bug", "feature_request", "how_to", "other"]), + urgent: z.boolean(), + summary: z.string(), + }), +}); +``` + +The parsed, schema-valid object shows up on every non-streaming surface and on the stream: + +| Surface | Where the object appears | +| --- | --- | +| `POST /invocations`, `POST /responses` | top-level `output_parsed` field on the JSON envelope (next to `output`) | +| `POST /chat` (SSE) | one final `appkit.structured_output` event (`{ data }`) **after** the streamed text | +| in-process `runAgent` | `RunAgentResult.output`, statically typed as `z.infer` of the schema | + +```ts +import { runAgent } from "@databricks/appkit/beta"; + +const { output } = await runAgent(classifier, { messages: ticket }); +output?.category; // "billing" | "bug" | … | undefined — fully typed +``` + +A per-call override is also available for one-off shapes: `runAgent(agent, input, { output: SomeSchema })` (it takes precedence over the agent's own schema and retypes the result). + +**How it works.** A tool-free agent sends the schema inline as the serving endpoint's `response_format`, so the answer comes back as JSON directly. A tool-having agent runs its normal tool loop, then a dedicated tool-free **structuring pass** reformats the answer into JSON (Databricks Claude endpoints reject `response_format` alongside `tools`). Either way the JSON is validated with Zod; on a mismatch AppKit re-prompts with the validation errors (up to two retries, invisible to the `/chat` text stream) and then throws a `StructuredOutputError` carrying the last raw output — it never returns partial or unvalidated data. If an endpoint rejects `response_format` with a 400, AppKit strips it and retries, relying on the prompt plus Zod validation. + +:::note Scope +Structured output is for **code-config agents** (`createAgent`). Markdown `agent.md` agents can't carry a schema, and sub-agents return text into their parent's context (typing buys nothing there). See the live API reference (`npx @databricks/appkit docs "appkit API reference"`) for the exact `createAgent` / `runAgent` / `StructuredOutputError` signatures. +::: + ## Adding agents to an existing app Already have an app and want to add agents? What you touch depends on the kind: diff --git a/packages/appkit/src/agents/tests/databricks.test.ts b/packages/appkit/src/agents/tests/databricks.test.ts index 4fff9c9a0..c3c2c40b2 100644 --- a/packages/appkit/src/agents/tests/databricks.test.ts +++ b/packages/appkit/src/agents/tests/databricks.test.ts @@ -1283,3 +1283,168 @@ describe("parseTextToolCalls", () => { expect(parseTextToolCalls(`${filler}${suffix}`)).toEqual([]); }); }); + +describe("DatabricksAdapter structured output", () => { + const originalFetch = globalThis.fetch; + + afterEach(() => { + globalThis.fetch = originalFetch; + mockAuthenticate.mockClear(); + }); + + interface FakeResponse { + ok: boolean; + status?: number; + chunks?: string[]; + text?: string; + } + + /** Records each request body and replays queued responses in order. */ + function capturingFetch( + bodies: Array>, + responses: FakeResponse[], + ): typeof globalThis.fetch { + let call = 0; + return vi.fn().mockImplementation((_url, init) => { + if (init?.body) bodies.push(JSON.parse(init.body)); + const r = responses[Math.min(call, responses.length - 1)]; + call++; + return Promise.resolve({ + ok: r.ok, + status: r.status ?? (r.ok ? 200 : 400), + body: r.ok ? createReadableStream(r.chunks ?? []) : null, + text: () => Promise.resolve(r.text ?? ""), + }); + }); + } + + async function drain(gen: AsyncGenerator): Promise { + for await (const _ of gen) { + // consume + } + } + + const outputSchema = { + type: "object", + properties: { answer: { type: "string" } }, + required: ["answer"], + }; + + test("tool-free run sends response_format json_schema", async () => { + const bodies: Array> = []; + globalThis.fetch = capturingFetch(bodies, [ + { ok: true, chunks: [textDelta('{"answer":"hi"}'), sseChunk("[DONE]")] }, + ]); + + const adapter = createAdapter(); + await drain( + adapter.run( + { + messages: createTestMessages(), + tools: [], + threadId: "t1", + outputSchema, + }, + { executeTool: vi.fn() }, + ), + ); + + expect(bodies).toHaveLength(1); + expect(bodies[0].response_format).toEqual({ + type: "json_schema", + json_schema: { + name: "structured_output", + schema: outputSchema, + strict: true, + }, + }); + expect(bodies[0].tools).toBeUndefined(); + }); + + test("does NOT send response_format when tools are present", async () => { + const bodies: Array> = []; + globalThis.fetch = capturingFetch(bodies, [ + { ok: true, chunks: [textDelta("done"), sseChunk("[DONE]")] }, + ]); + + const adapter = createAdapter(); + await drain( + adapter.run( + { + messages: createTestMessages(), + tools: createTestTools(), + threadId: "t1", + outputSchema, + }, + { executeTool: vi.fn() }, + ), + ); + + expect(bodies[0].response_format).toBeUndefined(); + expect(bodies[0].tools).toBeDefined(); + }); + + test("400 naming response_format strips it and retries once", async () => { + const bodies: Array> = []; + globalThis.fetch = capturingFetch(bodies, [ + { + ok: false, + status: 400, + text: "Bad request: response_format is not supported", + }, + { ok: true, chunks: [textDelta('{"answer":"ok"}'), sseChunk("[DONE]")] }, + ]); + + const adapter = createAdapter(); + const events: AgentEvent[] = []; + for await (const ev of adapter.run( + { + messages: createTestMessages(), + tools: [], + threadId: "t1", + outputSchema, + }, + { executeTool: vi.fn() }, + )) { + events.push(ev); + } + + // First body carried response_format; the retry stripped it. + expect(bodies).toHaveLength(2); + expect(bodies[0].response_format).toBeDefined(); + expect(bodies[1].response_format).toBeUndefined(); + // The stream succeeded on retry — the model's JSON came through. + expect(events).toContainEqual({ + type: "message_delta", + content: '{"answer":"ok"}', + }); + // No error status leaked (the 400 was recovered). + expect(events).not.toContainEqual( + expect.objectContaining({ type: "status", status: "error" }), + ); + }); + + test("a 400 NOT naming response_format propagates (no strip-retry)", async () => { + const bodies: Array> = []; + globalThis.fetch = capturingFetch(bodies, [ + { ok: false, status: 400, text: "Bad request: token limit exceeded" }, + ]); + + const adapter = createAdapter(); + await expect( + drain( + adapter.run( + { + messages: createTestMessages(), + tools: [], + threadId: "t1", + outputSchema, + }, + { executeTool: vi.fn() }, + ), + ), + ).rejects.toThrow(/token limit exceeded/); + // Only the original request — no retry for an unrelated 400. + expect(bodies).toHaveLength(1); + }); +}); diff --git a/packages/appkit/src/core/agent/tests/run-agent-structured.test.ts b/packages/appkit/src/core/agent/tests/run-agent-structured.test.ts new file mode 100644 index 000000000..24fc15052 --- /dev/null +++ b/packages/appkit/src/core/agent/tests/run-agent-structured.test.ts @@ -0,0 +1,118 @@ +import type { AgentAdapter, AgentInput } from "shared"; +import { describe, expect, test } from "vitest"; +import { z } from "zod"; + +import { StructuredOutputError } from "../../../errors"; +import { createAgent } from "../create-agent"; +import { runAgent } from "../run-agent"; +import { tool } from "../tools/tool"; + +const schema = z.object({ answer: z.string(), score: z.number() }); + +/** Fake adapter that records each run() input and emits a scripted text per call. */ +function recordingAdapter( + texts: string[], +): AgentAdapter & { calls: AgentInput[] } { + const calls: AgentInput[] = []; + let i = 0; + return { + calls, + async *run(input) { + calls.push(input); + const text = texts[Math.min(i, texts.length - 1)]; + i++; + yield { type: "status", status: "running" }; + yield { type: "message_delta", content: text }; + }, + }; +} + +describe("runAgent structured output", () => { + test("tool-free agent: output is parsed; main run carries outputSchema, no tools", async () => { + const adapter = recordingAdapter(['{"answer":"hi","score":1}']); + const agent = createAgent({ + instructions: "classify", + model: adapter, + output: schema, + }); + + const result = await runAgent(agent, { messages: "hello" }); + + expect(result.output).toEqual({ answer: "hi", score: 1 }); + // Statically typed via z.infer — this line only compiles if the type is right. + const typed: { answer: string; score: number } | undefined = result.output; + expect(typed?.answer).toBe("hi"); + + expect(adapter.calls).toHaveLength(1); + expect(adapter.calls[0].outputSchema).toBeDefined(); + expect(adapter.calls[0].tools).toEqual([]); + }); + + test("tool-having agent: runs a tool-free structuring pass over the prose answer", async () => { + const adapter = recordingAdapter([ + "The answer is hi with a score of 1.", // main run: prose + '{"answer":"hi","score":1}', // structuring pass: JSON + ]); + const agent = createAgent({ + instructions: "classify", + model: adapter, + output: schema, + tools: { + noop: tool({ + name: "noop", + description: "does nothing", + schema: z.object({}), + execute: () => "ok", + }), + }, + }); + + const result = await runAgent(agent, { messages: "hello" }); + + expect(result.output).toEqual({ answer: "hi", score: 1 }); + expect(adapter.calls).toHaveLength(2); + // Main run exposed the tool; the structuring pass is tool-free + constrained. + expect(adapter.calls[0].tools).toHaveLength(1); + expect(adapter.calls[1].tools).toEqual([]); + expect(adapter.calls[1].outputSchema).toBeDefined(); + }); + + test("per-call { output } override drives structured output when the agent has none", async () => { + const adapter = recordingAdapter(['{"answer":"o","score":2}']); + const agent = createAgent({ instructions: "x", model: adapter }); + + const result = await runAgent( + agent, + { messages: "hi" }, + { output: schema }, + ); + + expect(result.output).toEqual({ answer: "o", score: 2 }); + }); + + test("throws StructuredOutputError when output never validates", async () => { + const adapter = recordingAdapter(["not json at all"]); + const agent = createAgent({ + instructions: "x", + model: adapter, + output: schema, + }); + + await expect(runAgent(agent, { messages: "hi" })).rejects.toBeInstanceOf( + StructuredOutputError, + ); + // 1 main run + 2 structuring retries. + expect(adapter.calls).toHaveLength(3); + }); + + test("no output schema: result.output is undefined, no extra runs", async () => { + const adapter = recordingAdapter(["just some prose"]); + const agent = createAgent({ instructions: "x", model: adapter }); + + const result = await runAgent(agent, { messages: "hi" }); + + expect(result.output).toBeUndefined(); + expect(result.text).toBe("just some prose"); + expect(adapter.calls).toHaveLength(1); + }); +}); diff --git a/packages/appkit/src/core/agent/tests/structured-output.test.ts b/packages/appkit/src/core/agent/tests/structured-output.test.ts new file mode 100644 index 000000000..8165dea58 --- /dev/null +++ b/packages/appkit/src/core/agent/tests/structured-output.test.ts @@ -0,0 +1,119 @@ +import type { Message } from "shared"; +import { describe, expect, test, vi } from "vitest"; +import { z } from "zod"; + +import { StructuredOutputError } from "../../../errors"; +import { + resolveStructuredOutput, + type StructuringPass, +} from "../structured-output"; + +const schema = z.object({ category: z.string(), urgent: z.boolean() }); + +function baseMessages(): Message[] { + return [ + { id: "s", role: "system", content: "classify", createdAt: new Date() }, + { id: "u", role: "user", content: "help me", createdAt: new Date() }, + ]; +} + +describe("resolveStructuredOutput", () => { + test("tool-free: validates finalText directly, no structuring pass", async () => { + const pass = vi.fn(); + const output = await resolveStructuredOutput({ + schema, + baseMessages: baseMessages(), + finalText: JSON.stringify({ category: "billing", urgent: true }), + hadTools: false, + runStructuringPass: pass, + }); + + expect(output).toEqual({ category: "billing", urgent: true }); + expect(pass).not.toHaveBeenCalled(); + }); + + test("tool-having: runs one structuring pass over the answer", async () => { + const pass = vi + .fn() + .mockResolvedValue(JSON.stringify({ category: "sales", urgent: false })); + + const output = await resolveStructuredOutput({ + schema, + baseMessages: baseMessages(), + finalText: "This looks like a sales question, not urgent.", + hadTools: true, + runStructuringPass: pass, + }); + + expect(output).toEqual({ category: "sales", urgent: false }); + expect(pass).toHaveBeenCalledTimes(1); + // The prose answer is appended as an assistant turn + a convert instruction. + const [msgs] = pass.mock.calls[0]; + expect(msgs.at(-2)).toMatchObject({ role: "assistant" }); + expect(msgs.at(-1)?.role).toBe("user"); + expect(msgs.at(-1)?.content).toMatch(/JSON matching the provided schema/i); + }); + + test("strips code fences before parsing", async () => { + const output = await resolveStructuredOutput({ + schema, + baseMessages: baseMessages(), + finalText: '```json\n{"category":"support","urgent":false}\n```', + hadTools: false, + runStructuringPass: vi.fn(), + }); + expect(output).toEqual({ category: "support", urgent: false }); + }); + + test("re-prompts with zod errors on validation failure, then succeeds", async () => { + const pass = vi + .fn() + .mockResolvedValueOnce(JSON.stringify({ category: "billing" })) // missing urgent + .mockResolvedValueOnce( + JSON.stringify({ category: "billing", urgent: true }), + ); + + const output = await resolveStructuredOutput({ + schema, + baseMessages: baseMessages(), + finalText: JSON.stringify({ category: "billing" }), // invalid attempt 0 + hadTools: false, + runStructuringPass: pass, + }); + + expect(output).toEqual({ category: "billing", urgent: true }); + // attempt 0 = finalText (invalid) -> retry 1 (invalid) -> retry 2 (valid) + expect(pass).toHaveBeenCalledTimes(2); + // The retry carries the flattened zod error text. + const [retryMsgs] = pass.mock.calls[0]; + expect(retryMsgs.at(-1)?.content).toMatch( + /did not match the required schema/i, + ); + expect(retryMsgs.at(-1)?.content).toMatch(/urgent/); + }); + + test("throws StructuredOutputError with lastRaw after retries exhausted", async () => { + const pass = vi.fn().mockResolvedValue("still not json"); + + await expect( + resolveStructuredOutput({ + schema, + baseMessages: baseMessages(), + finalText: "not json either", + hadTools: false, + runStructuringPass: pass, + }), + ).rejects.toMatchObject({ + name: "StructuredOutputError", + lastRaw: "still not json", + }); + + // attempt 0 (finalText) + 2 retries = 2 structuring passes. + expect(pass).toHaveBeenCalledTimes(2); + }); + + test("StructuredOutputError does not leak lastRaw via clientMessage", () => { + const err = new StructuredOutputError("boom", { lastRaw: "secret raw" }); + expect(err.clientMessage).not.toContain("secret raw"); + }); +}); diff --git a/packages/appkit/src/plugins/agents/tests/event-translator.test.ts b/packages/appkit/src/plugins/agents/tests/event-translator.test.ts index 28fd474b3..2a974f64d 100644 --- a/packages/appkit/src/plugins/agents/tests/event-translator.test.ts +++ b/packages/appkit/src/plugins/agents/tests/event-translator.test.ts @@ -130,6 +130,20 @@ describe("AgentEventTranslator", () => { } }); + test("translates structured_output to appkit.structured_output extension event", () => { + const translator = new AgentEventTranslator(); + const events = translator.translate({ + type: "structured_output", + data: { category: "billing", urgent: true }, + }); + + expect(events).toHaveLength(1); + expect(events[0].type).toBe("appkit.structured_output"); + if (events[0].type === "appkit.structured_output") { + expect(events[0].data).toEqual({ category: "billing", urgent: true }); + } + }); + test("status:complete triggers finalize with response.completed", () => { const translator = new AgentEventTranslator(); translator.translate({ type: "message_delta", content: "Hi" }); From a310063ffbb6fcb1532a5cb3e7fa287fd7238101 Mon Sep 17 00:00:00 2001 From: MarioCadenas Date: Thu, 3 Sep 2026 18:31:00 +0200 Subject: [PATCH 8/9] fix(agents): produce structured output via a non-streaming completion Databricks rejects response_format under stream:true ("Structured output is not currently supported with streaming"), and DatabricksAdapter is streaming-only, so the original inline/streaming approach failed on every call. Structured output now runs as a dedicated NON-streaming completion: - Add a non-streaming queryBody transport (connectors/serving query()) and a structuredCompletion() path in the adapter; run() uses it when outputSchema is set and tools are empty, emitting the JSON as one message event. - The main agent run always streams normally (revert the toolless-inline wiring), so /chat still streams the visible answer; structured output is a separate pass afterwards. - resolveStructuredOutput validates the answer as-is first (cheap pre-check), else reformats via the pass; drop the hadTools branch. - Broaden the response_format-rejection detector to catch INVALID_PARAMETER_VALUE / "structured output" phrasing so strip-and-retry actually fires. Found via live dogfood; updates tests, docs, and the example accordingly. Signed-off-by: MarioCadenas --- .../server/agents/classifier/agent.ts | 5 +- docs/docs/plugins/agents.md | 2 +- packages/appkit/src/agents/databricks.ts | 238 ++++++++++++------ .../src/agents/tests/databricks.test.ts | 75 +++--- .../appkit/src/connectors/serving/client.ts | 33 +++ packages/appkit/src/core/agent/run-agent.ts | 22 +- .../src/core/agent/structured-output.ts | 47 ++-- .../agent/tests/run-agent-structured.test.ts | 42 +++- .../agent/tests/structured-output.test.ts | 28 +-- packages/appkit/src/plugins/agents/agents.ts | 14 -- 10 files changed, 320 insertions(+), 186 deletions(-) diff --git a/apps/dev-playground/server/agents/classifier/agent.ts b/apps/dev-playground/server/agents/classifier/agent.ts index c0505bac7..d714a2f00 100644 --- a/apps/dev-playground/server/agents/classifier/agent.ts +++ b/apps/dev-playground/server/agents/classifier/agent.ts @@ -17,8 +17,9 @@ import { z } from "zod"; // const { output } = await runAgent(classifier, { messages: ticket }); // output?.category; // "billing" | "bug" | ... | undefined // -// Because it declares no tools, AppKit sends the schema inline as the -// endpoint's `response_format`, so the answer comes back as JSON directly. +// The agent answers normally, then AppKit runs a dedicated non-streaming +// completion constrained by the schema (Databricks rejects `response_format` +// under streaming) and validates it with Zod before surfacing the object. export default createAgent({ instructions: "You are a support-ticket triage classifier. Read the user's message and " + diff --git a/docs/docs/plugins/agents.md b/docs/docs/plugins/agents.md index 08504b7b9..b691e7362 100644 --- a/docs/docs/plugins/agents.md +++ b/docs/docs/plugins/agents.md @@ -348,7 +348,7 @@ output?.category; // "billing" | "bug" | … | undefined — fully typed A per-call override is also available for one-off shapes: `runAgent(agent, input, { output: SomeSchema })` (it takes precedence over the agent's own schema and retypes the result). -**How it works.** A tool-free agent sends the schema inline as the serving endpoint's `response_format`, so the answer comes back as JSON directly. A tool-having agent runs its normal tool loop, then a dedicated tool-free **structuring pass** reformats the answer into JSON (Databricks Claude endpoints reject `response_format` alongside `tools`). Either way the JSON is validated with Zod; on a mismatch AppKit re-prompts with the validation errors (up to two retries, invisible to the `/chat` text stream) and then throws a `StructuredOutputError` carrying the last raw output — it never returns partial or unvalidated data. If an endpoint rejects `response_format` with a 400, AppKit strips it and retries, relying on the prompt plus Zod validation. +**How it works.** The agent answers normally first (streaming its visible text on `/chat`). Structured output is then produced by a dedicated **non-streaming** completion constrained by the schema via `response_format` — Databricks rejects `response_format` under streaming, so the structured call can't be the streamed one. The answer is validated as-is first (a cheap pre-check for a model that already emitted JSON); otherwise the non-streaming pass reformats it. Either way the JSON is validated with Zod; on a mismatch AppKit re-prompts with the validation errors (up to two retries, invisible to the `/chat` text stream) and then throws a `StructuredOutputError` carrying the last raw output — it never returns partial or unvalidated data. If an endpoint rejects `response_format` outright, AppKit strips it and retries, relying on the prompt plus Zod validation. :::note Scope Structured output is for **code-config agents** (`createAgent`). Markdown `agent.md` agents can't carry a schema, and sub-agents return text into their parent's context (typing buys nothing there). See the live API reference (`npx @databricks/appkit docs "appkit API reference"`) for the exact `createAgent` / `runAgent` / `StructuredOutputError` signatures. diff --git a/packages/appkit/src/agents/databricks.ts b/packages/appkit/src/agents/databricks.ts index d974c2e47..8d95a97a6 100644 --- a/packages/appkit/src/agents/databricks.ts +++ b/packages/appkit/src/agents/databricks.ts @@ -7,6 +7,7 @@ import type { } from "shared"; import { + query as servingQuery, type StreamBody, stream as servingStream, } from "../connectors/serving/client"; @@ -71,25 +72,59 @@ function applyGenerationParams( } /** - * True when `err` looks like an HTTP 400 that names `response_format` / - * `json_schema`. Used to strip-and-retry when a serving endpoint (or the - * gateway in front of it) doesn't support structured output — some - * Llama/DBRX/Mistral endpoints reject the param outright. The Zod boundary - * in the structured-output resolver is the real guarantee; `response_format` - * is only the retry-rate optimization, so silently dropping it on a 400 is - * safe. Deliberately narrow: a 400 that does NOT name the param is a real - * request error and must propagate. + * True when `err` looks like a client-side (400 / INVALID_PARAMETER_VALUE) + * rejection of structured output. Used to strip `response_format` and retry — + * some endpoints don't support it at all, and Databricks Claude specifically + * rejects it (the phrasing is "Structured output is not currently supported + * with streaming", though this path is already non-streaming). The Zod + * boundary in the resolver is the real guarantee, so dropping the param on + * such an error is safe. Deliberately narrow: a client error that does NOT + * name the feature is a genuine request error and must propagate. */ function isResponseFormatRejection(err: unknown): boolean { + const rec = isRecord(err) ? err : {}; const status = - isRecord(err) && typeof err.status === "number" - ? err.status - : isRecord(err) && typeof err.statusCode === "number" - ? err.statusCode + typeof rec.status === "number" + ? rec.status + : typeof rec.statusCode === "number" + ? rec.statusCode : undefined; + const code = typeof rec.errorCode === "string" ? rec.errorCode : ""; const msg = err instanceof Error ? err.message : String(err); - const is400 = status === 400 || /\b400\b/.test(msg); - return is400 && /response_format|json[_ ]schema/i.test(msg); + const looksClientError = + status === 400 || /\b400\b/.test(msg) || code === "INVALID_PARAMETER_VALUE"; + const namesFeature = /response_format|json[_ ]schema|structured output/i.test( + msg, + ); + return looksClientError && namesFeature; +} + +/** Pull the assistant message text out of a NON-streaming chat completion. */ +function extractMessageContent(parsed: unknown): string { + if (!isRecord(parsed)) return ""; + const choices = parsed.choices; + if (!Array.isArray(choices) || choices.length === 0) return ""; + const first = choices[0]; + if (!isRecord(first)) return ""; + const message = first.message; + if (!isRecord(message)) return ""; + const content = message.content; + if (typeof content === "string") return content; + // Harmony/array content: concatenate the text parts. + if (Array.isArray(content)) { + let text = ""; + for (const part of content) { + if ( + isRecord(part) && + part.type === "text" && + typeof part.text === "string" + ) { + text += part.text; + } + } + return text; + } + return ""; } function extractLlamaToolJsonSlice(text: string): string | undefined { @@ -145,6 +180,16 @@ function reasoningPartText(part: Record): string { * adapter uses a bare `fetch()` to call it. Useful for tests and for pointing * the adapter at non-workspace endpoints (reverse proxies, mocks). */ +/** + * Non-streaming transport, mirroring {@link StreamBody}. Returns the parsed + * JSON response body. Used only for structured-output completions (Databricks + * rejects `response_format` under `stream: true`). + */ +type QueryBody = ( + body: Record, + signal?: AbortSignal, +) => Promise; + interface RawFetchAdapterOptions { endpointUrl: string; authenticate: () => Promise>; @@ -168,6 +213,8 @@ interface RawFetchAdapterOptions { */ interface StreamBodyAdapterOptions { streamBody: StreamBody; + /** Non-streaming transport for structured-output completions. */ + queryBody?: QueryBody; maxSteps?: number; maxTokens?: number; generationParams?: GenerationParams; @@ -306,6 +353,8 @@ interface DeltaToolCall { */ export class DatabricksAdapter implements AgentAdapter { private streamBody: StreamBody; + /** Non-streaming transport; present only when structured output is usable. */ + private queryBody?: QueryBody; private maxSteps: number; private maxTokens: number; private generationParams: GenerationParams; @@ -326,6 +375,7 @@ export class DatabricksAdapter implements AgentAdapter { if (isStreamBodyOptions(options)) { this.streamBody = options.streamBody; + this.queryBody = options.queryBody; } else { const { endpointUrl, authenticate } = options; this.streamBody = async (body, signal) => { @@ -351,6 +401,29 @@ export class DatabricksAdapter implements AgentAdapter { if (!response.body) throw new Error("No response body"); return response.body; }; + // Non-streaming sibling of streamBody for structured-output completions. + this.queryBody = async (body, signal) => { + const fetchSignal = + signal ?? AbortSignal.timeout(RAW_FETCH_DEFAULT_TIMEOUT_MS); + const authHeaders = await authenticate(); + const response = await fetch(endpointUrl, { + method: "POST", + headers: { + "User-Agent": APPKIT_USER_AGENT, + "Content-Type": "application/json", + ...authHeaders, + }, + body: JSON.stringify({ ...body, stream: false }), + signal: fetchSignal, + }); + if (!response.ok) { + const errorText = await response.text().catch(() => "Unknown error"); + throw new Error( + `Databricks API error (${response.status}): ${errorText}`, + ); + } + return response.json(); + }; } } @@ -386,6 +459,13 @@ export class DatabricksAdapter implements AgentAdapter { body, signal, ), + queryBody: (body, signal) => + servingQuery( + workspaceClient as unknown as Parameters[0], + endpointName, + body, + signal, + ), maxSteps, maxTokens, generationParams, @@ -500,17 +580,31 @@ export class DatabricksAdapter implements AgentAdapter { const tools = this.buildTools(input.tools, nameToWire); const messages = this.buildMessages(input.messages, nameToWire); - // Structured output is only sent inline for tool-free completions — - // Databricks Claude endpoints reject `response_format` alongside `tools`. - // Tool-having agents produce prose here and are re-formatted by a - // separate tool-free structuring pass (a second `run()` with no tools). - const structuredSchema = - input.outputSchema && input.tools.length === 0 - ? input.outputSchema - : undefined; - yield { type: "status", status: "running" }; + // Structured output is a tool-free, NON-streaming completion — Databricks + // rejects `response_format` under `stream: true`. The structured-output + // resolver drives this via `run({ tools: [], outputSchema })`; the visible + // agent answer streams normally on the tool-having / non-structured path + // below and is re-formatted into JSON by this pass afterwards. + if (input.outputSchema && input.tools.length === 0) { + let text: string; + try { + text = await this.structuredCompletion( + messages, + input.outputSchema, + context, + ); + } catch (err) { + const msg = + err instanceof Error ? err.message : "Structured request failed"; + yield { type: "status", status: "error", error: msg }; + throw err; + } + yield { type: "message", content: text }; + return; + } + for (let step = 0; step < this.maxSteps; step++) { if (context.signal?.aborted) break; @@ -518,19 +612,13 @@ export class DatabricksAdapter implements AgentAdapter { messages, tools, context, - structuredSchema, ); if (toolCalls.length === 0) { - // In structured mode the model returns raw JSON (an array-typed schema - // emits `[{...}]`, which the Llama/Python text-tool-call fallback would - // otherwise misread as a tool invocation). Skip the fallback entirely. - if (!structuredSchema) { - const parsed = parseTextToolCalls(text); - if (parsed.length > 0) { - yield* this.executeToolCalls(parsed, messages, context, nameToWire); - continue; - } + const parsed = parseTextToolCalls(text); + if (parsed.length > 0) { + yield* this.executeToolCalls(parsed, messages, context, nameToWire); + continue; } break; } @@ -596,11 +684,54 @@ export class DatabricksAdapter implements AgentAdapter { } } + /** + * One tool-free, NON-streaming completion constrained by `schema` via + * `response_format`. Databricks rejects `response_format` under + * `stream: true`, so structured output must be non-streaming. Returns the + * raw message content (the JSON text) for the caller to validate. On a 400 + * that names the param, strips `response_format` and retries once — the + * structuring prompt already asks for schema-only JSON and the caller's Zod + * validation is the real guarantee. + */ + private async structuredCompletion( + messages: OpenAIMessage[], + schema: Record, + context: AgentRunContext, + ): Promise { + if (!this.queryBody) { + throw new Error( + "DatabricksAdapter: structured output requires a non-streaming transport. " + + "Build the adapter via DatabricksAdapter.fromServingEndpoint / fromModelServing.", + ); + } + const body: Record = { + messages, + max_tokens: this.maxTokens, + response_format: { + type: "json_schema", + json_schema: { name: "structured_output", schema, strict: true }, + }, + }; + applyGenerationParams(body, this.generationParams); + + let parsed: unknown; + try { + parsed = await this.queryBody(body, context.signal); + } catch (err) { + if (isResponseFormatRejection(err)) { + delete body.response_format; + parsed = await this.queryBody(body, context.signal); + } else { + throw err; + } + } + return extractMessageContent(parsed); + } + private async *streamCompletion( messages: OpenAIMessage[], tools: OpenAITool[], context: AgentRunContext, - structuredSchema?: Record, ): AsyncGenerator< AgentEvent, { text: string; toolCalls: OpenAIToolCall[] }, @@ -616,48 +747,15 @@ export class DatabricksAdapter implements AgentAdapter { if (tools.length > 0) { body.tools = tools; - } else if (structuredSchema) { - // OpenAI-compatible structured output. `strict: true` asks the endpoint - // to constrain generation to the schema; the caller still validates - // with Zod, so a silently-ignored `response_format` is handled upstream. - body.response_format = { - type: "json_schema", - json_schema: { - name: "structured_output", - schema: structuredSchema, - strict: true, - }, - }; } let responseBody: ReadableStream; try { responseBody = await this.streamBody(body, context.signal); } catch (err) { - // Some endpoints/gateways 400 on `response_format`. Strip it and retry - // once — the structured-output resolver's Zod validation is the real - // guarantee. A 400 that doesn't name the param is a genuine error. - if ( - body.response_format !== undefined && - isResponseFormatRejection(err) - ) { - delete body.response_format; - try { - responseBody = await this.streamBody(body, context.signal); - } catch (retryErr) { - const msg = - retryErr instanceof Error - ? retryErr.message - : "Stream request failed"; - yield { type: "status", status: "error", error: msg }; - throw retryErr; - } - } else { - const msg = - err instanceof Error ? err.message : "Stream request failed"; - yield { type: "status", status: "error", error: msg }; - throw err; - } + const msg = err instanceof Error ? err.message : "Stream request failed"; + yield { type: "status", status: "error", error: msg }; + throw err; } const reader = responseBody.getReader(); diff --git a/packages/appkit/src/agents/tests/databricks.test.ts b/packages/appkit/src/agents/tests/databricks.test.ts index c3c2c40b2..4bd53febc 100644 --- a/packages/appkit/src/agents/tests/databricks.test.ts +++ b/packages/appkit/src/agents/tests/databricks.test.ts @@ -1295,7 +1295,10 @@ describe("DatabricksAdapter structured output", () => { interface FakeResponse { ok: boolean; status?: number; + /** SSE chunks for the streaming path. */ chunks?: string[]; + /** Parsed body for the non-streaming (structured) path. */ + json?: unknown; text?: string; } @@ -1313,15 +1316,23 @@ describe("DatabricksAdapter structured output", () => { ok: r.ok, status: r.status ?? (r.ok ? 200 : 400), body: r.ok ? createReadableStream(r.chunks ?? []) : null, + json: () => Promise.resolve(r.json), text: () => Promise.resolve(r.text ?? ""), }); }); } - async function drain(gen: AsyncGenerator): Promise { - for await (const _ of gen) { - // consume - } + /** A non-streaming chat-completion body with a single JSON message. */ + function completion(content: string): unknown { + return { choices: [{ message: { role: "assistant", content } }] }; + } + + async function collect( + gen: AsyncGenerator, + ): Promise { + const events: AgentEvent[] = []; + for await (const ev of gen) events.push(ev); + return events; } const outputSchema = { @@ -1330,14 +1341,14 @@ describe("DatabricksAdapter structured output", () => { required: ["answer"], }; - test("tool-free run sends response_format json_schema", async () => { + test("tool-free structured run is NON-streaming with response_format", async () => { const bodies: Array> = []; globalThis.fetch = capturingFetch(bodies, [ - { ok: true, chunks: [textDelta('{"answer":"hi"}'), sseChunk("[DONE]")] }, + { ok: true, json: completion('{"answer":"hi"}') }, ]); const adapter = createAdapter(); - await drain( + const events = await collect( adapter.run( { messages: createTestMessages(), @@ -1350,6 +1361,7 @@ describe("DatabricksAdapter structured output", () => { ); expect(bodies).toHaveLength(1); + expect(bodies[0].stream).toBe(false); expect(bodies[0].response_format).toEqual({ type: "json_schema", json_schema: { @@ -1359,16 +1371,21 @@ describe("DatabricksAdapter structured output", () => { }, }); expect(bodies[0].tools).toBeUndefined(); + // The JSON arrives as a single `message` event (not streamed deltas). + expect(events).toContainEqual({ + type: "message", + content: '{"answer":"hi"}', + }); }); - test("does NOT send response_format when tools are present", async () => { + test("does NOT send response_format when tools are present (stays streaming)", async () => { const bodies: Array> = []; globalThis.fetch = capturingFetch(bodies, [ { ok: true, chunks: [textDelta("done"), sseChunk("[DONE]")] }, ]); const adapter = createAdapter(); - await drain( + await collect( adapter.run( { messages: createTestMessages(), @@ -1382,49 +1399,47 @@ describe("DatabricksAdapter structured output", () => { expect(bodies[0].response_format).toBeUndefined(); expect(bodies[0].tools).toBeDefined(); + expect(bodies[0].stream).toBe(true); }); - test("400 naming response_format strips it and retries once", async () => { + test("400 naming structured output strips response_format and retries once", async () => { const bodies: Array> = []; globalThis.fetch = capturingFetch(bodies, [ { ok: false, status: 400, - text: "Bad request: response_format is not supported", + text: "INVALID_PARAMETER_VALUE: Structured output is not currently supported with streaming.", }, - { ok: true, chunks: [textDelta('{"answer":"ok"}'), sseChunk("[DONE]")] }, + { ok: true, json: completion('{"answer":"ok"}') }, ]); const adapter = createAdapter(); - const events: AgentEvent[] = []; - for await (const ev of adapter.run( - { - messages: createTestMessages(), - tools: [], - threadId: "t1", - outputSchema, - }, - { executeTool: vi.fn() }, - )) { - events.push(ev); - } + const events = await collect( + adapter.run( + { + messages: createTestMessages(), + tools: [], + threadId: "t1", + outputSchema, + }, + { executeTool: vi.fn() }, + ), + ); - // First body carried response_format; the retry stripped it. + // First body carried response_format; the retry stripped it. Both non-streaming. expect(bodies).toHaveLength(2); expect(bodies[0].response_format).toBeDefined(); expect(bodies[1].response_format).toBeUndefined(); - // The stream succeeded on retry — the model's JSON came through. expect(events).toContainEqual({ - type: "message_delta", + type: "message", content: '{"answer":"ok"}', }); - // No error status leaked (the 400 was recovered). expect(events).not.toContainEqual( expect.objectContaining({ type: "status", status: "error" }), ); }); - test("a 400 NOT naming response_format propagates (no strip-retry)", async () => { + test("a 400 NOT naming the param propagates (no strip-retry)", async () => { const bodies: Array> = []; globalThis.fetch = capturingFetch(bodies, [ { ok: false, status: 400, text: "Bad request: token limit exceeded" }, @@ -1432,7 +1447,7 @@ describe("DatabricksAdapter structured output", () => { const adapter = createAdapter(); await expect( - drain( + collect( adapter.run( { messages: createTestMessages(), diff --git a/packages/appkit/src/connectors/serving/client.ts b/packages/appkit/src/connectors/serving/client.ts index de9d0465c..ffe822568 100644 --- a/packages/appkit/src/connectors/serving/client.ts +++ b/packages/appkit/src/connectors/serving/client.ts @@ -101,6 +101,39 @@ export async function streamPath( return response.contents; } +/** + * POSTs `body` to a serving endpoint as a NON-streaming request and returns the + * parsed JSON response (OpenAI-compatible `{ choices: [{ message }] }` shape). + * Forces `stream: false`. Used by the structured-output path — Databricks + * rejects `response_format` under `stream: true`, so structured completions + * must be non-streaming. + * + * @internal Not part of the public AppKit surface. Like {@link streamPath}, + * the endpoint name is caller-controlled and hard-coded by internal callers; + * do not expose to user input (workspace-credentialled SSRF). + */ +export async function query( + client: ApiClientLike, + endpointName: string, + body: Record, + signal?: AbortSignal, +): Promise { + const { stream: _stream, ...cleanBody } = body; + const context = contextFromAbortSignal(signal); + return client.apiClient.request( + { + path: `/serving-endpoints/${encodeURIComponent(endpointName)}/invocations`, + method: "POST", + headers: new Headers({ + "Content-Type": "application/json", + Accept: "application/json", + }), + payload: { ...cleanBody, stream: false }, + }, + context, + ); +} + /** * Returns the raw SSE byte stream from a serving endpoint. Thin wrapper over * {@link streamPath} that handles serving-specific URL encoding and forces diff --git a/packages/appkit/src/core/agent/run-agent.ts b/packages/appkit/src/core/agent/run-agent.ts index 1a4223565..b9b10ce4b 100644 --- a/packages/appkit/src/core/agent/run-agent.ts +++ b/packages/appkit/src/core/agent/run-agent.ts @@ -133,21 +133,22 @@ export async function runAgent( const providerCache = new Map(); await initStandalonePlugins(input.plugins ?? [], providerCache); - const schema = options?.output ?? def.output; - // Pass the schema into the main run so a tool-free agent gets - // `response_format` inline (no wasted round-trip); the adapter ignores it - // when tools are present. Sub-agent recursions never receive it. - const mainOutputSchema = schema ? toToolJSONSchema(schema) : undefined; - const { text, events, adapter, hadTools, baseMessages } = - await runAgentInternal(def, input, providerCache, mainOutputSchema); + const { text, events, adapter, baseMessages } = await runAgentInternal( + def, + input, + providerCache, + ); + const schema = options?.output ?? def.output; if (!schema) return { text, events }; + // Structured output is produced by a separate tool-free, non-streaming + // structuring pass (Databricks rejects response_format under streaming), + // so the main run above always streams the visible answer normally. const output = await resolveStructuredOutput({ schema, baseMessages, finalText: text, - hadTools, runStructuringPass: buildStructuringPass(adapter, schema), signal: input.signal, }); @@ -187,8 +188,6 @@ interface RawRunResult { events: AgentEvent[]; /** Adapter used for the run — reused for the structuring pass. */ adapter: AgentAdapter; - /** Whether the run exposed tools (tool-having answers need a structuring pass). */ - hadTools: boolean; /** Normalized system + input messages the run saw (structuring-pass seed). */ baseMessages: Message[]; } @@ -197,7 +196,6 @@ async function runAgentInternal( def: AgentDefinition, input: RunAgentInput, providerCache: Map, - mainOutputSchema?: Record, ): Promise { const adapter = await resolveAdapter(def); const messages = normalizeMessages(input.messages, def.instructions); @@ -273,7 +271,6 @@ async function runAgentInternal( threadId: randomUUID(), signal, extensions: buildStandaloneExtensions(toolIndex), - outputSchema: mainOutputSchema, }, { executeTool, signal }, ); @@ -292,7 +289,6 @@ async function runAgentInternal( text, events, adapter, - hadTools: tools.length > 0, baseMessages: messages, }; } diff --git a/packages/appkit/src/core/agent/structured-output.ts b/packages/appkit/src/core/agent/structured-output.ts index c043c201f..95400eb40 100644 --- a/packages/appkit/src/core/agent/structured-output.ts +++ b/packages/appkit/src/core/agent/structured-output.ts @@ -46,14 +46,8 @@ interface ResolveStructuredOutputParams { * instruction to this. */ baseMessages: Message[]; - /** The main run's final assistant text (already JSON when `hadTools` is false). */ + /** The main run's final assistant text (the answer to reformat into JSON). */ finalText: string; - /** - * Whether the main run used tools. Tool-having runs produce prose, so the - * first attempt is a structuring pass; tool-free runs already emitted JSON - * inline (via `response_format`), so `finalText` is validated directly. - */ - hadTools: boolean; /** Runs one tool-free, schema-constrained completion (see {@link StructuringPass}). */ runStructuringPass: StructuringPass; signal?: AbortSignal; @@ -62,33 +56,28 @@ interface ResolveStructuredOutputParams { /** * Validate-and-retry loop that turns an agent's answer into a typed object. * - * - Tool-free agent: `finalText` is already `response_format` JSON, validated - * directly; on failure a re-prompted structuring pass runs. - * - Tool-having agent: the answer is prose, so a structuring pass reformats it - * into JSON first. - * - * On every failure it re-prompts with the flattened Zod issues (up to - * {@link MAX_VALIDATION_RETRIES} times), then throws {@link StructuredOutputError} - * carrying the last raw output. Never returns partial/unvalidated data. + * The agent's visible answer (`finalText`) is validated as-is first — a model + * may already emit valid JSON — as a cheap pre-check. Otherwise a tool-free, + * schema-constrained structuring pass reformats it into JSON, re-prompting with + * the flattened Zod issues on each failure (up to {@link MAX_VALIDATION_RETRIES} + * times). On exhaustion it throws {@link StructuredOutputError} carrying the + * last raw output; it never returns partial/unvalidated data. */ export async function resolveStructuredOutput( params: ResolveStructuredOutputParams, ): Promise { - const { - schema, - baseMessages, - finalText, - hadTools, - runStructuringPass, + const { schema, baseMessages, finalText, runStructuringPass, signal } = + params; + + // Cheap pre-check: the answer may already be valid JSON (no round-trip). + const direct = parseAndValidate(schema, finalText); + if (direct.ok) return direct.value; + + // Reformat the answer into JSON via a tool-free structuring pass. + let lastRaw = await runStructuringPass( + structuringMessages(baseMessages, finalText, CONVERT_INSTRUCTION), signal, - } = params; - - let lastRaw = hadTools - ? await runStructuringPass( - structuringMessages(baseMessages, finalText, CONVERT_INSTRUCTION), - signal, - ) - : finalText; + ); for (let retries = 0; ; retries++) { const parsed = parseAndValidate(schema, lastRaw); diff --git a/packages/appkit/src/core/agent/tests/run-agent-structured.test.ts b/packages/appkit/src/core/agent/tests/run-agent-structured.test.ts index 24fc15052..1cad26ef7 100644 --- a/packages/appkit/src/core/agent/tests/run-agent-structured.test.ts +++ b/packages/appkit/src/core/agent/tests/run-agent-structured.test.ts @@ -28,7 +28,7 @@ function recordingAdapter( } describe("runAgent structured output", () => { - test("tool-free agent: output is parsed; main run carries outputSchema, no tools", async () => { + test("answer already valid JSON: direct pre-check, no structuring pass", async () => { const adapter = recordingAdapter(['{"answer":"hi","score":1}']); const agent = createAgent({ instructions: "classify", @@ -39,20 +39,41 @@ describe("runAgent structured output", () => { const result = await runAgent(agent, { messages: "hello" }); expect(result.output).toEqual({ answer: "hi", score: 1 }); - // Statically typed via z.infer — this line only compiles if the type is right. + // Statically typed via z.infer — compiles only if the type is right. const typed: { answer: string; score: number } | undefined = result.output; expect(typed?.answer).toBe("hi"); + // One run only; the main run never carries outputSchema. expect(adapter.calls).toHaveLength(1); - expect(adapter.calls[0].outputSchema).toBeDefined(); - expect(adapter.calls[0].tools).toEqual([]); + expect(adapter.calls[0].outputSchema).toBeUndefined(); }); - test("tool-having agent: runs a tool-free structuring pass over the prose answer", async () => { + test("prose answer: a tool-free structuring pass produces the object", async () => { const adapter = recordingAdapter([ "The answer is hi with a score of 1.", // main run: prose '{"answer":"hi","score":1}', // structuring pass: JSON ]); + const agent = createAgent({ + instructions: "classify", + model: adapter, + output: schema, + }); + + const result = await runAgent(agent, { messages: "hello" }); + + expect(result.output).toEqual({ answer: "hi", score: 1 }); + expect(adapter.calls).toHaveLength(2); + // Main run: no outputSchema. Structuring pass: tool-free + schema-constrained. + expect(adapter.calls[0].outputSchema).toBeUndefined(); + expect(adapter.calls[1].tools).toEqual([]); + expect(adapter.calls[1].outputSchema).toBeDefined(); + }); + + test("tool-having agent structures identically (tool only on the main run)", async () => { + const adapter = recordingAdapter([ + "Looks like a billing issue, not urgent.", + '{"answer":"billing","score":0}', + ]); const agent = createAgent({ instructions: "classify", model: adapter, @@ -69,11 +90,10 @@ describe("runAgent structured output", () => { const result = await runAgent(agent, { messages: "hello" }); - expect(result.output).toEqual({ answer: "hi", score: 1 }); + expect(result.output).toEqual({ answer: "billing", score: 0 }); expect(adapter.calls).toHaveLength(2); - // Main run exposed the tool; the structuring pass is tool-free + constrained. - expect(adapter.calls[0].tools).toHaveLength(1); - expect(adapter.calls[1].tools).toEqual([]); + expect(adapter.calls[0].tools).toHaveLength(1); // main run has the tool + expect(adapter.calls[1].tools).toEqual([]); // structuring pass is tool-free expect(adapter.calls[1].outputSchema).toBeDefined(); }); @@ -101,8 +121,8 @@ describe("runAgent structured output", () => { await expect(runAgent(agent, { messages: "hi" })).rejects.toBeInstanceOf( StructuredOutputError, ); - // 1 main run + 2 structuring retries. - expect(adapter.calls).toHaveLength(3); + // main run + convert pass + 2 retry passes. + expect(adapter.calls).toHaveLength(4); }); test("no output schema: result.output is undefined, no extra runs", async () => { diff --git a/packages/appkit/src/core/agent/tests/structured-output.test.ts b/packages/appkit/src/core/agent/tests/structured-output.test.ts index 8165dea58..5699d0a9f 100644 --- a/packages/appkit/src/core/agent/tests/structured-output.test.ts +++ b/packages/appkit/src/core/agent/tests/structured-output.test.ts @@ -18,13 +18,12 @@ function baseMessages(): Message[] { } describe("resolveStructuredOutput", () => { - test("tool-free: validates finalText directly, no structuring pass", async () => { + test("validates the answer as-is when it is already valid JSON (no pass)", async () => { const pass = vi.fn(); const output = await resolveStructuredOutput({ schema, baseMessages: baseMessages(), finalText: JSON.stringify({ category: "billing", urgent: true }), - hadTools: false, runStructuringPass: pass, }); @@ -32,7 +31,7 @@ describe("resolveStructuredOutput", () => { expect(pass).not.toHaveBeenCalled(); }); - test("tool-having: runs one structuring pass over the answer", async () => { + test("reformats a prose answer via one structuring pass", async () => { const pass = vi .fn() .mockResolvedValue(JSON.stringify({ category: "sales", urgent: false })); @@ -41,7 +40,6 @@ describe("resolveStructuredOutput", () => { schema, baseMessages: baseMessages(), finalText: "This looks like a sales question, not urgent.", - hadTools: true, runStructuringPass: pass, }); @@ -59,7 +57,6 @@ describe("resolveStructuredOutput", () => { schema, baseMessages: baseMessages(), finalText: '```json\n{"category":"support","urgent":false}\n```', - hadTools: false, runStructuringPass: vi.fn(), }); expect(output).toEqual({ category: "support", urgent: false }); @@ -76,20 +73,20 @@ describe("resolveStructuredOutput", () => { const output = await resolveStructuredOutput({ schema, baseMessages: baseMessages(), - finalText: JSON.stringify({ category: "billing" }), // invalid attempt 0 - hadTools: false, + finalText: JSON.stringify({ category: "billing" }), // invalid — triggers a pass runStructuringPass: pass, }); expect(output).toEqual({ category: "billing", urgent: true }); - // attempt 0 = finalText (invalid) -> retry 1 (invalid) -> retry 2 (valid) + // pass #1 (convert) invalid -> pass #2 (retry w/ errors) valid. expect(pass).toHaveBeenCalledTimes(2); - // The retry carries the flattened zod error text. - const [retryMsgs] = pass.mock.calls[0]; - expect(retryMsgs.at(-1)?.content).toMatch( - /did not match the required schema/i, + // The FIRST pass is the convert instruction; the retry carries the zod error. + expect(pass.mock.calls[0][0].at(-1)?.content).toMatch( + /JSON matching the provided schema/i, ); - expect(retryMsgs.at(-1)?.content).toMatch(/urgent/); + const retryMsg = pass.mock.calls[1][0].at(-1)?.content ?? ""; + expect(retryMsg).toMatch(/did not match the required schema/i); + expect(retryMsg).toMatch(/urgent/); }); test("throws StructuredOutputError with lastRaw after retries exhausted", async () => { @@ -100,7 +97,6 @@ describe("resolveStructuredOutput", () => { schema, baseMessages: baseMessages(), finalText: "not json either", - hadTools: false, runStructuringPass: pass, }), ).rejects.toMatchObject({ @@ -108,8 +104,8 @@ describe("resolveStructuredOutput", () => { lastRaw: "still not json", }); - // attempt 0 (finalText) + 2 retries = 2 structuring passes. - expect(pass).toHaveBeenCalledTimes(2); + // convert pass + 2 retry passes = 3 structuring attempts. + expect(pass).toHaveBeenCalledTimes(3); }); test("StructuredOutputError does not leak lastRaw via clientMessage", () => { diff --git a/packages/appkit/src/plugins/agents/agents.ts b/packages/appkit/src/plugins/agents/agents.ts index 8f0b180c3..9e424a156 100644 --- a/packages/appkit/src/plugins/agents/agents.ts +++ b/packages/appkit/src/plugins/agents/agents.ts @@ -1299,11 +1299,6 @@ export class AgentsPlugin extends Plugin implements ToolProvider { threadId: thread.id, signal, extensions: buildAdapterExtensions(registered.toolIndex), - // Tool-free structured agents get `response_format` inline; the - // adapter ignores it when tools are present. - outputSchema: registered.output - ? toToolJSONSchema(registered.output) - : undefined, }, { executeTool, signal }, ); @@ -1352,7 +1347,6 @@ export class AgentsPlugin extends Plugin implements ToolProvider { registered, messagesWithSystem, fullContent, - tools.length > 0, signal, ); for (const evt of translator.translate({ @@ -1531,11 +1525,6 @@ export class AgentsPlugin extends Plugin implements ToolProvider { tools, threadId: thread.id, signal, - // Tool-free structured agents get `response_format` inline; the - // adapter ignores it when tools are present. - outputSchema: registered.output - ? toToolJSONSchema(registered.output) - : undefined, }, { executeTool, signal }, ); @@ -1560,7 +1549,6 @@ export class AgentsPlugin extends Plugin implements ToolProvider { registered, messagesWithSystem, fullContent, - tools.length > 0, signal, ); } @@ -1640,7 +1628,6 @@ export class AgentsPlugin extends Plugin implements ToolProvider { registered: RegisteredAgent, baseMessages: Message[], finalText: string, - hadTools: boolean, signal: AbortSignal, ): Promise { const schema = registered.output; @@ -1653,7 +1640,6 @@ export class AgentsPlugin extends Plugin implements ToolProvider { schema, baseMessages, finalText, - hadTools, signal, runStructuringPass: (messages, sig) => consumeAdapterStream( From ecb6629ccb8b8b8a33ad3b97fd8743ef75bfc732 Mon Sep 17 00:00:00 2001 From: MarioCadenas Date: Fri, 4 Sep 2026 11:15:38 +0200 Subject: [PATCH 9/9] refactor(agents): dedup structuring pass, memoize schema, trim comments Quality pass over the structured-output feature (no behavior change): - Extract one shared buildStructuringPass(adapter, outputSchema) in structured-output.ts; runAgent and the agents plugin both use it instead of near-identical inline closures. - Reuse a shared formatZodIssues() core (tools/tool.ts) for the re-prompt instead of a second zod-flatten copy. - Memoize toToolJSONSchema on RegisteredAgent.outputSchema at registration rather than re-deriving it per request. - Collapse isResponseFormatRejection to the message test (a superset of the prior status/code checks) and mark its ceiling. - Make RunAgentOptions non-generic; the runAgent overload carries inference. - Share one raw-fetch helper for streamBody/queryBody in the adapter. - Trim verbose/duplicated comments to the load-bearing fact. All gates green: typecheck, check:fix, build, docs:build, agent tests. Signed-off-by: MarioCadenas --- .../server/agents/classifier/agent.ts | 23 +--- packages/appkit/src/agents/databricks.ts | 108 ++++++------------ .../appkit/src/connectors/serving/client.ts | 8 +- packages/appkit/src/core/agent/run-agent.ts | 38 +----- .../src/core/agent/structured-output.ts | 83 +++++++------- packages/appkit/src/core/agent/tools/tool.ts | 20 +++- packages/appkit/src/core/agent/types.ts | 16 +-- .../appkit/src/errors/structured-output.ts | 13 +-- packages/appkit/src/plugins/agents/agents.ts | 63 ++++------ packages/shared/src/agent.ts | 21 ++-- 10 files changed, 147 insertions(+), 246 deletions(-) diff --git a/apps/dev-playground/server/agents/classifier/agent.ts b/apps/dev-playground/server/agents/classifier/agent.ts index d714a2f00..929067a27 100644 --- a/apps/dev-playground/server/agents/classifier/agent.ts +++ b/apps/dev-playground/server/agents/classifier/agent.ts @@ -1,25 +1,10 @@ import { createAgent } from "@databricks/appkit/beta"; import { z } from "zod"; -// Structured-output demo: a tool-free agent whose final answer is validated -// against a Zod schema instead of being returned as freeform text. Discovered -// automatically from server/agents/classifier/ (its id is the folder name, -// "classifier"). -// -// How the schema surfaces on each path: -// • POST /api/agents/chat with { message, agent: "classifier" } streams the -// text, then emits one final `appkit.structured_output` SSE event whose -// `data` is the parsed object. -// • POST /api/agents/invocations (when this is the default agent) returns a -// top-level `output_parsed` field alongside the usual `output` text. -// • In-process, the result is statically typed via z.infer: -// import { runAgent } from "@databricks/appkit/beta"; -// const { output } = await runAgent(classifier, { messages: ticket }); -// output?.category; // "billing" | "bug" | ... | undefined -// -// The agent answers normally, then AppKit runs a dedicated non-streaming -// completion constrained by the schema (Databricks rejects `response_format` -// under streaming) and validates it with Zod before surfacing the object. +// Structured-output demo: a tool-free agent whose answer is validated against +// its `output` Zod schema instead of returned as text. Discovered from the +// folder name (`classifier`). See docs/plugins/agents.md "Structured output" +// for how the object surfaces on /chat, /invocations, and in-process runAgent. export default createAgent({ instructions: "You are a support-ticket triage classifier. Read the user's message and " + diff --git a/packages/appkit/src/agents/databricks.ts b/packages/appkit/src/agents/databricks.ts index 8d95a97a6..51b7b063c 100644 --- a/packages/appkit/src/agents/databricks.ts +++ b/packages/appkit/src/agents/databricks.ts @@ -72,31 +72,20 @@ function applyGenerationParams( } /** - * True when `err` looks like a client-side (400 / INVALID_PARAMETER_VALUE) - * rejection of structured output. Used to strip `response_format` and retry — - * some endpoints don't support it at all, and Databricks Claude specifically - * rejects it (the phrasing is "Structured output is not currently supported - * with streaming", though this path is already non-streaming). The Zod - * boundary in the resolver is the real guarantee, so dropping the param on - * such an error is safe. Deliberately narrow: a client error that does NOT - * name the feature is a genuine request error and must propagate. + * True when `err`'s message names structured output / `response_format`. Used + * to strip the param and retry when an endpoint rejects it — some endpoints + * don't support it at all, and Databricks returns the coarse + * `INVALID_PARAMETER_VALUE` ("Structured output is not currently supported…") + * with no granular code to key on. The Zod boundary in the resolver is the + * real guarantee, so a stray strip is harmless (validation still runs); that + * lets this stay a simple message test rather than probing status codes. + * + * ponytail: substring match — Databricks has no granular code for a + * response_format rejection; tighten to a code check if one ever ships. */ function isResponseFormatRejection(err: unknown): boolean { - const rec = isRecord(err) ? err : {}; - const status = - typeof rec.status === "number" - ? rec.status - : typeof rec.statusCode === "number" - ? rec.statusCode - : undefined; - const code = typeof rec.errorCode === "string" ? rec.errorCode : ""; const msg = err instanceof Error ? err.message : String(err); - const looksClientError = - status === 400 || /\b400\b/.test(msg) || code === "INVALID_PARAMETER_VALUE"; - const namesFeature = /response_format|json[_ ]schema|structured output/i.test( - msg, - ); - return looksClientError && namesFeature; + return /response_format|json[_ ]schema|structured output/i.test(msg); } /** Pull the assistant message text out of a NON-streaming chat completion. */ @@ -175,21 +164,17 @@ function reasoningPartText(part: Record): string { return text; } -/** - * Escape-hatch options: provide an `endpointUrl` + `authenticate()` and the - * adapter uses a bare `fetch()` to call it. Useful for tests and for pointing - * the adapter at non-workspace endpoints (reverse proxies, mocks). - */ -/** - * Non-streaming transport, mirroring {@link StreamBody}. Returns the parsed - * JSON response body. Used only for structured-output completions (Databricks - * rejects `response_format` under `stream: true`). - */ +/** Non-streaming sibling of {@link StreamBody}; returns the parsed JSON body. */ type QueryBody = ( body: Record, signal?: AbortSignal, ) => Promise; +/** + * Escape-hatch options: provide an `endpointUrl` + `authenticate()` and the + * adapter uses a bare `fetch()` to call it. Useful for tests and for pointing + * the adapter at non-workspace endpoints (reverse proxies, mocks). + */ interface RawFetchAdapterOptions { endpointUrl: string; authenticate: () => Promise>; @@ -378,9 +363,10 @@ export class DatabricksAdapter implements AgentAdapter { this.queryBody = options.queryBody; } else { const { endpointUrl, authenticate } = options; - this.streamBody = async (body, signal) => { - const fetchSignal = - signal ?? AbortSignal.timeout(RAW_FETCH_DEFAULT_TIMEOUT_MS); + const doFetch = async ( + body: Record, + signal?: AbortSignal, + ): Promise => { const authHeaders = await authenticate(); const response = await fetch(endpointUrl, { method: "POST", @@ -390,7 +376,7 @@ export class DatabricksAdapter implements AgentAdapter { ...authHeaders, }, body: JSON.stringify(body), - signal: fetchSignal, + signal: signal ?? AbortSignal.timeout(RAW_FETCH_DEFAULT_TIMEOUT_MS), }); if (!response.ok) { const errorText = await response.text().catch(() => "Unknown error"); @@ -398,32 +384,16 @@ export class DatabricksAdapter implements AgentAdapter { `Databricks API error (${response.status}): ${errorText}`, ); } + return response; + }; + this.streamBody = async (body, signal) => { + const response = await doFetch(body, signal); if (!response.body) throw new Error("No response body"); return response.body; }; - // Non-streaming sibling of streamBody for structured-output completions. - this.queryBody = async (body, signal) => { - const fetchSignal = - signal ?? AbortSignal.timeout(RAW_FETCH_DEFAULT_TIMEOUT_MS); - const authHeaders = await authenticate(); - const response = await fetch(endpointUrl, { - method: "POST", - headers: { - "User-Agent": APPKIT_USER_AGENT, - "Content-Type": "application/json", - ...authHeaders, - }, - body: JSON.stringify({ ...body, stream: false }), - signal: fetchSignal, - }); - if (!response.ok) { - const errorText = await response.text().catch(() => "Unknown error"); - throw new Error( - `Databricks API error (${response.status}): ${errorText}`, - ); - } - return response.json(); - }; + // Non-streaming sibling for structured-output completions. + this.queryBody = async (body, signal) => + (await doFetch({ ...body, stream: false }, signal)).json(); } } @@ -582,11 +552,9 @@ export class DatabricksAdapter implements AgentAdapter { yield { type: "status", status: "running" }; - // Structured output is a tool-free, NON-streaming completion — Databricks - // rejects `response_format` under `stream: true`. The structured-output - // resolver drives this via `run({ tools: [], outputSchema })`; the visible - // agent answer streams normally on the tool-having / non-structured path - // below and is re-formatted into JSON by this pass afterwards. + // Tool-free structuring pass (the resolver drives it via + // `run({ tools: [], outputSchema })`): a non-streaming, schema-constrained + // completion. The visible answer streams on the path below. if (input.outputSchema && input.tools.length === 0) { let text: string; try { @@ -685,13 +653,11 @@ export class DatabricksAdapter implements AgentAdapter { } /** - * One tool-free, NON-streaming completion constrained by `schema` via - * `response_format`. Databricks rejects `response_format` under - * `stream: true`, so structured output must be non-streaming. Returns the - * raw message content (the JSON text) for the caller to validate. On a 400 - * that names the param, strips `response_format` and retries once — the - * structuring prompt already asks for schema-only JSON and the caller's Zod - * validation is the real guarantee. + * One tool-free, non-streaming completion constrained by `schema` via + * `response_format` (Databricks rejects `response_format` under + * `stream: true`). Returns the raw message text for the caller to validate; + * on a rejection, strips `response_format` and retries once (see + * {@link isResponseFormatRejection}). */ private async structuredCompletion( messages: OpenAIMessage[], diff --git a/packages/appkit/src/connectors/serving/client.ts b/packages/appkit/src/connectors/serving/client.ts index ffe822568..662a6941a 100644 --- a/packages/appkit/src/connectors/serving/client.ts +++ b/packages/appkit/src/connectors/serving/client.ts @@ -102,11 +102,9 @@ export async function streamPath( } /** - * POSTs `body` to a serving endpoint as a NON-streaming request and returns the - * parsed JSON response (OpenAI-compatible `{ choices: [{ message }] }` shape). - * Forces `stream: false`. Used by the structured-output path — Databricks - * rejects `response_format` under `stream: true`, so structured completions - * must be non-streaming. + * POSTs `body` to a serving endpoint as a non-streaming request (forces + * `stream: false`) and returns the parsed JSON (`{ choices: [{ message }] }`). + * Used by the structured-output path. * * @internal Not part of the public AppKit surface. Like {@link streamPath}, * the endpoint name is caller-controlled and hard-coded by internal callers; diff --git a/packages/appkit/src/core/agent/run-agent.ts b/packages/appkit/src/core/agent/run-agent.ts index b9b10ce4b..05ed4be76 100644 --- a/packages/appkit/src/core/agent/run-agent.ts +++ b/packages/appkit/src/core/agent/run-agent.ts @@ -20,8 +20,8 @@ import { createLogger } from "../../logging/logger"; import { consumeAdapterStream } from "./consume-adapter-stream"; import { createPluginsProxy } from "./plugins-map"; import { + buildStructuringPass, resolveStructuredOutput, - type StructuringPass, } from "./structured-output"; import { resolveToolkitFromProvider } from "./toolkit-resolver"; import { @@ -58,14 +58,14 @@ export interface RunAgentInput { } /** Per-call options for {@link runAgent}. */ -export interface RunAgentOptions { +export interface RunAgentOptions { /** * Structured-output schema override for this call. Takes precedence over the * agent's own `output` schema. When either is set, `runAgent` validates the * final answer and populates {@link RunAgentResult.output}, throwing a * `StructuredOutputError` if it can't produce a valid object. */ - output?: z.ZodType; + output?: z.ZodType; } export interface RunAgentResult { @@ -113,7 +113,7 @@ export interface RunAgentResult { export function runAgent( def: AgentDefinition, input: RunAgentInput, - options: RunAgentOptions> & { output: S }, + options: RunAgentOptions & { output: S }, ): Promise>>; // No override — the result type comes from the agent's own `output` schema. export function runAgent( @@ -149,40 +149,12 @@ export async function runAgent( schema, baseMessages, finalText: text, - runStructuringPass: buildStructuringPass(adapter, schema), + runStructuringPass: buildStructuringPass(adapter, toToolJSONSchema(schema)), signal: input.signal, }); return { text, events, output }; } -/** - * Builds a {@link StructuringPass}: one tool-free, schema-constrained - * `adapter.run()`, consumed to its final text. `executeTool` throws — a - * tool-free run never dispatches one; if it somehow does, that's a bug we - * want surfaced, not swallowed. - */ -function buildStructuringPass( - adapter: AgentAdapter, - schema: z.ZodType, -): StructuringPass { - const outputSchema = toToolJSONSchema(schema); - return (messages, signal) => - consumeAdapterStream( - adapter.run( - { messages, tools: [], threadId: randomUUID(), signal, outputSchema }, - { - executeTool: () => { - throw new Error( - "runAgent: structuring pass is tool-free and must not call a tool", - ); - }, - signal, - }, - ), - { signal }, - ); -} - interface RawRunResult { text: string; events: AgentEvent[]; diff --git a/packages/appkit/src/core/agent/structured-output.ts b/packages/appkit/src/core/agent/structured-output.ts index 95400eb40..d2ba34e55 100644 --- a/packages/appkit/src/core/agent/structured-output.ts +++ b/packages/appkit/src/core/agent/structured-output.ts @@ -1,17 +1,13 @@ import { randomUUID } from "node:crypto"; -import type { Message } from "shared"; +import type { AgentAdapter, Message } from "shared"; import type { z } from "zod"; import { StructuredOutputError } from "../../errors"; +import { consumeAdapterStream } from "./consume-adapter-stream"; +import { formatZodIssues } from "./tools/tool"; -/** - * Max number of re-prompted structuring passes after the first validation - * fails. So ≤3 total validation attempts, and ≤3 structuring passes for a - * tool-having agent (1 initial + 2 retries) / ≤2 for a tool-free one (the - * first attempt is the inline `response_format` output, retries add passes). - * Tracked here, separate from the tool-call / maxSteps budget. - */ +/** Re-prompted structuring passes after the first convert pass (so ≤3 total). Separate from the tool-call / maxSteps budget. */ const MAX_VALIDATION_RETRIES = 2; const CONVERT_INSTRUCTION = @@ -27,10 +23,9 @@ function retryInstruction(errors: string): string { } /** - * Runs ONE tool-free, schema-constrained completion over `messages` and - * returns the raw model text. Injected by the caller so the resolver stays - * free of any adapter / MLflow dependency: the agents plugin and standalone - * `runAgent` each build this from `adapter.run({ tools: [], outputSchema })`. + * Runs one tool-free, schema-constrained completion and returns the raw text. + * Injected so the resolver stays free of any adapter / MLflow dependency; see + * {@link buildStructuringPass}. */ export type StructuringPass = ( messages: Message[], @@ -40,13 +35,9 @@ export type StructuringPass = ( interface ResolveStructuredOutputParams { /** Schema the final object is validated against. */ schema: z.ZodType; - /** - * The conversation the main run saw (system + thread messages), WITHOUT the - * final answer. Each structuring pass appends the latest answer + an - * instruction to this. - */ + /** The conversation the main run saw (system + thread), without the final answer. */ baseMessages: Message[]; - /** The main run's final assistant text (the answer to reformat into JSON). */ + /** The main run's final assistant text — the answer to reformat into JSON. */ finalText: string; /** Runs one tool-free, schema-constrained completion (see {@link StructuringPass}). */ runStructuringPass: StructuringPass; @@ -54,14 +45,11 @@ interface ResolveStructuredOutputParams { } /** - * Validate-and-retry loop that turns an agent's answer into a typed object. - * - * The agent's visible answer (`finalText`) is validated as-is first — a model - * may already emit valid JSON — as a cheap pre-check. Otherwise a tool-free, - * schema-constrained structuring pass reformats it into JSON, re-prompting with - * the flattened Zod issues on each failure (up to {@link MAX_VALIDATION_RETRIES} - * times). On exhaustion it throws {@link StructuredOutputError} carrying the - * last raw output; it never returns partial/unvalidated data. + * Turns an agent's answer into a schema-validated object. Validates `finalText` + * as-is first (skipping a structuring round-trip when the model already emitted + * valid JSON); otherwise re-prompts a tool-free structuring pass with the + * flattened Zod issues, up to {@link MAX_VALIDATION_RETRIES} times. Throws + * {@link StructuredOutputError} on exhaustion — never returns partial data. */ export async function resolveStructuredOutput( params: ResolveStructuredOutputParams, @@ -69,11 +57,9 @@ export async function resolveStructuredOutput( const { schema, baseMessages, finalText, runStructuringPass, signal } = params; - // Cheap pre-check: the answer may already be valid JSON (no round-trip). const direct = parseAndValidate(schema, finalText); if (direct.ok) return direct.value; - // Reformat the answer into JSON via a tool-free structuring pass. let lastRaw = await runStructuringPass( structuringMessages(baseMessages, finalText, CONVERT_INSTRUCTION), signal, @@ -139,15 +125,10 @@ function parseAndValidate( } const result = schema.safeParse(json); if (result.success) return { ok: true, value: result.data }; - return { ok: false, error: flattenZodError(result.error) }; + return { ok: false, error: formatZodIssues(result.error) }; } -/** - * Strip a single leading/trailing Markdown code fence. `response_format` - * output is bare JSON, but on the 400-strip fallback (or a model that ignores - * the param) the answer often arrives fenced (` ```json … ``` `). Cheap to - * undo and materially raises the parse rate on that path. - */ +/** Strip a leading/trailing Markdown code fence — models often wrap JSON in ` ```json … ``` `. */ function stripCodeFences(text: string): string { const trimmed = text.trim(); if (!trimmed.startsWith("```")) return trimmed; @@ -157,12 +138,28 @@ function stripCodeFences(text: string): string { .trim(); } -/** Compact one-line rendering of a ZodError's issues, safe across zod v4. */ -function flattenZodError(error: z.ZodError): string { - return error.issues - .map((issue) => { - const path = issue.path.join(".") || "(root)"; - return `${path}: ${issue.message}`; - }) - .join("; "); +/** + * Builds a {@link StructuringPass}: one tool-free, schema-constrained + * `adapter.run()` consumed to its text. `executeTool` throws — a tool-free + * pass must never dispatch one. + */ +export function buildStructuringPass( + adapter: AgentAdapter, + outputSchema: Record, +): StructuringPass { + return (messages, signal) => + consumeAdapterStream( + adapter.run( + { messages, tools: [], threadId: randomUUID(), signal, outputSchema }, + { + executeTool: () => { + throw new Error( + "structured-output structuring pass is tool-free and must not call a tool", + ); + }, + signal, + }, + ), + { signal }, + ); } diff --git a/packages/appkit/src/core/agent/tools/tool.ts b/packages/appkit/src/core/agent/tools/tool.ts index 62092f217..22ff6bc92 100644 --- a/packages/appkit/src/core/agent/tools/tool.ts +++ b/packages/appkit/src/core/agent/tools/tool.ts @@ -84,15 +84,25 @@ export function tool(config: ToolConfig): FunctionTool { }; } +/** + * Compact one-line rendering of a ZodError's issues: `path: message` per issue, + * joined by `; ` (root-level issues render as `(root)`). Stable across zod v4. + * Shared by {@link formatZodError} and the structured-output re-prompt. + */ +export function formatZodIssues(error: z.ZodError): string { + return error.issues + .map((issue) => { + const field = issue.path.length > 0 ? issue.path.join(".") : "(root)"; + return `${field}: ${issue.message}`; + }) + .join("; "); +} + /** * Formats a Zod validation error into an LLM-friendly string. * * Example: `Invalid arguments for get_weather: city: Invalid input: expected string, received undefined` */ export function formatZodError(error: z.ZodError, toolName: string): string { - const parts = error.issues.map((issue) => { - const field = issue.path.length > 0 ? issue.path.join(".") : "(root)"; - return `${field}: ${issue.message}`; - }); - return `Invalid arguments for ${toolName}: ${parts.join("; ")}`; + return `Invalid arguments for ${toolName}: ${formatZodIssues(error)}`; } diff --git a/packages/appkit/src/core/agent/types.ts b/packages/appkit/src/core/agent/types.ts index 3d822b776..a816130b7 100644 --- a/packages/appkit/src/core/agent/types.ts +++ b/packages/appkit/src/core/agent/types.ts @@ -156,14 +156,12 @@ export interface AgentDefinition { /** System prompt body. For markdown-loaded agents this is the file body. */ instructions: string; /** - * Optional Zod schema the agent's final answer is validated against. When - * set, the agent returns a typed object instead of freeform text: the - * `/invocations` envelope gains a top-level `output_parsed` field, `/chat` - * emits a final `structured_output` SSE event, and in-process `runAgent` - * populates `RunAgentResult.output`. Prefer {@link createAgent} with an - * `output` schema so the in-process result is statically typed via - * `z.infer`. Code-config agents only — markdown `agent.md` agents cannot - * carry a schema. + * Optional Zod schema the agent's final answer is validated against; when + * set, the agent returns a typed object instead of text — surfaced as + * `output_parsed` on `/invocations`, a final `structured_output` SSE event + * on `/chat`, and `RunAgentResult.output` in-process. Prefer {@link + * createAgent} so the in-process result is typed via `z.infer`. Code-config + * agents only — markdown agents can't carry a schema. */ output?: z.ZodType; /** @@ -418,6 +416,8 @@ export interface RegisteredAgent { * typing lives on `createAgent`/`runAgent`. */ output?: z.ZodType; + /** The {@link output} schema converted to JSON Schema, memoized at registration. Present iff `output` is. */ + outputSchema?: Record; /** * Resolved per-agent skill catalog (visibility + collision rules applied). * Present when any skill is visible to this agent; drives the always-on diff --git a/packages/appkit/src/errors/structured-output.ts b/packages/appkit/src/errors/structured-output.ts index ad1eebd8a..4c4c38ec0 100644 --- a/packages/appkit/src/errors/structured-output.ts +++ b/packages/appkit/src/errors/structured-output.ts @@ -1,14 +1,11 @@ import { AppKitError } from "./base"; /** - * Thrown when an agent with an `output` schema could not produce a value that - * validates against it, even after the structuring retries are exhausted. - * Carries the last raw model output (server-side only) for debugging — it is - * never returned to the client, which sees the generic {@link clientMessage}. - * - * The structured-output path throws this rather than returning partial or - * unvalidated data: a caller that asked for a typed object gets either a valid - * one or an error, never a half-parsed shape. + * Thrown when an agent with an `output` schema could not produce a + * schema-valid value within the structuring retry budget — the caller gets a + * valid object or this error, never partial data. Carries the last raw model + * output (`lastRaw`) for debugging, server-side only: it is never returned to + * the client, which sees the generic {@link clientMessage}. * * @example * ```typescript diff --git a/packages/appkit/src/plugins/agents/agents.ts b/packages/appkit/src/plugins/agents/agents.ts index 9e424a156..6f4ff38d3 100644 --- a/packages/appkit/src/plugins/agents/agents.ts +++ b/packages/appkit/src/plugins/agents/agents.ts @@ -27,7 +27,10 @@ import type { ResolvedSkillCatalog, SkillDefinition, } from "../../core/agent/skills"; -import { resolveStructuredOutput } from "../../core/agent/structured-output"; +import { + buildStructuringPass, + resolveStructuredOutput, +} from "../../core/agent/structured-output"; import { resolveToolkitFromProvider } from "../../core/agent/toolkit-resolver"; import { functionToolToDefinition, @@ -497,6 +500,8 @@ export class AgentsPlugin extends Plugin implements ToolProvider { generationParams: def.generationParams, ephemeral: def.ephemeral, output: def.output, + // Convert once here — the schema is fixed per agent; re-deriving per request would be wasted. + outputSchema: def.output ? toToolJSONSchema(def.output) : undefined, skills, }; } @@ -1337,11 +1342,9 @@ export class AgentsPlugin extends Plugin implements ToolProvider { } } - // Structured output: after the visible text, coerce + validate the - // answer against the agent's schema and emit one final - // `structured_output` event. Retries (if any) happen here, invisible - // to the streamed text. On exhaustion the throw propagates to the - // driver's catch and surfaces as an error event. + // After the visible text, resolve structured output and emit one + // final `structured_output` event — any retries happen here, + // invisibly to the stream; a throw surfaces via the driver's catch. if (registered.output) { const data = await this.runStructuredOutput( registered, @@ -1541,9 +1544,8 @@ export class AgentsPlugin extends Plugin implements ToolProvider { }); } - // Structured output: coerce + validate against the agent's schema - // and attach as the envelope's `output_parsed`. Throws on exhaustion - // (caught below as a 500) — never returns partial data. + // Resolve structured output for the `output_parsed` envelope field; + // a throw is caught below as a 500. if (registered.output) { outputParsed = await this.runStructuredOutput( registered, @@ -1609,20 +1611,17 @@ export class AgentsPlugin extends Plugin implements ToolProvider { ...(runState.toolErrors.length > 0 ? { tool_errors: runState.toolErrors } : {}), - // Parsed, schema-validated object when the agent declared an `output` - // schema — the non-streaming equivalent of the `structured_output` SSE - // event. Conditionally spread, like `mlflow_trace_id` / `tool_errors`. + // Parsed structured output — the non-streaming counterpart of the + // `structured_output` SSE event. ...(outputParsed !== undefined ? { output_parsed: outputParsed } : {}), output: [message], }); } /** - * Resolves an agent's structured output from the run's final text. The - * structuring pass(es) run as a fresh tool-free `adapter.run()` constrained - * by the schema, wrapped in a TOOL span so they nest under the turn's AGENT - * span. Assumes `registered.output` is set. Throws `StructuredOutputError` - * if no schema-valid object can be produced within the retry budget. + * Resolves the agent's structured output from the run's final text, wrapped + * in a TOOL span so it nests under the turn's AGENT span. Assumes the agent + * declared an `output` schema. */ private runStructuredOutput( registered: RegisteredAgent, @@ -1630,38 +1629,20 @@ export class AgentsPlugin extends Plugin implements ToolProvider { finalText: string, signal: AbortSignal, ): Promise { - const schema = registered.output; - if (!schema) { + const { output: schema, outputSchema } = registered; + if (!schema || !outputSchema) { throw new Error("runStructuredOutput called without an output schema"); } - const outputSchema = toToolJSONSchema(schema); return traceTool("structured_output", { schema: outputSchema }, () => resolveStructuredOutput({ schema, baseMessages, finalText, signal, - runStructuringPass: (messages, sig) => - consumeAdapterStream( - registered.adapter.run( - { - messages, - tools: [], - threadId: randomUUID(), - signal: sig, - outputSchema, - }, - { - executeTool: () => { - throw new Error( - "structured-output structuring pass is tool-free and must not call a tool", - ); - }, - signal: sig, - }, - ), - { signal: sig }, - ), + runStructuringPass: buildStructuringPass( + registered.adapter, + outputSchema, + ), }), ); } diff --git a/packages/shared/src/agent.ts b/packages/shared/src/agent.ts index e219fd7c8..36978593a 100644 --- a/packages/shared/src/agent.ts +++ b/packages/shared/src/agent.ts @@ -131,11 +131,9 @@ export type AgentEvent = | { type: "metadata"; data: Record } | { /** - * Emitted by the agents plugin (not adapters) once, after the streamed - * text, when the agent declared an `output` schema: the parsed, - * schema-validated object. Delivered on the wire as - * `appkit.structured_output`. Non-streaming surfaces (`/invocations`) - * return the same object as the envelope's `output_parsed` field. + * Emitted by the agents plugin (not adapters) after the streamed text: + * the parsed, schema-validated object. Wire event: + * {@link AppKitStructuredOutputEvent}. */ type: "structured_output"; data: unknown; @@ -301,14 +299,11 @@ export interface AgentInput { threadId: string; signal?: AbortSignal; /** - * JSON Schema the adapter should constrain a tool-free completion to, when - * it supports server-side structured output (e.g. an OpenAI-compatible - * `response_format: { type: "json_schema" }`). Set by the structured-output - * path; ignored when `tools` is non-empty (Databricks Claude endpoints - * reject `response_format` combined with `tools`). Adapters that can't - * constrain output ignore it — the structured-output resolver then relies - * on prompt + Zod validation. Already stripped of the top-level `$schema` - * key by {@link toToolJSONSchema}. + * JSON Schema to constrain a tool-free completion to, for adapters that + * support server-side structured output (OpenAI-compatible `response_format`). + * Adapters that can't ignore it — the structured-output resolver then falls + * back to prompt + Zod validation. Already stripped of the top-level + * `$schema` key by {@link toToolJSONSchema}. */ outputSchema?: Record; /**