@dbx-tools/appkit-mastra 0.9.32 → 0.9.37

package/src/serving-sanitize.ts

@@ -1,362 +1,19 @@
 /**
- * Repairs traffic between Mastra / the AI SDK and Databricks Model Serving's
- * OpenAI-compatible `/chat/completions` route, in both directions.
- *
- * Outbound ({@link rewriteServingBody}), because the transcript Mastra
- * persists is not always a transcript the provider will accept back:
- * Databricks-hosted Claude rejects replayed extended-thinking blocks and reads
- * a trailing assistant message as a prefill request, while GPT tool calls need
- * model-specific reasoning-effort normalization on Chat Completions.
- *
- * Inbound ({@link rewriteServingResponseBody} and
- * {@link rewriteServingResponseStream}), because Databricks-hosted Gemini and
- * Claude can answer with native content-parts arrays where the OpenAI contract
- * (and therefore the AI SDK's response schema) requires a plain string.
- *
- * Every repair here is a provider quirk rather than a schema violation, so all
- * of them are applied on the wire (see the `globalThis.fetch` wrapper in
- * `model.ts`) rather than by changing what the agent stores or what the UI
- * shows.
+ * Databricks Model Serving wire repairs shared with the model gateway.
  *
  * @module
  */
 
-import { policy } from "@dbx-tools/model";
-import { json, stringUtils } from "@dbx-tools/shared-core";
-import { type ChatMessage, type ChatRole, openaiChat } from "@dbx-tools/shared-model";
-
-/**
- * A chat message as it arrives on the serving wire, plus the extended-thinking
- * fields Databricks-hosted Claude adds. The OpenAI-standard part of the shape
- * is {@link ChatMessage} from `@dbx-tools/shared-model`; only the provider
- * extensions this module strips are declared here.
- */
-export interface ServingChatMessage extends ChatMessage {
-  /** Narrowed to the roles this repair pass reasons about, plus Claude's `reasoning` turn. */
-  role: ChatRole | "reasoning";
-  reasoning?: unknown;
-  reasoning_content?: unknown;
-}
-
-/** Signed Claude reasoning blocks that cannot be replayed after mutation. */
-const REASONING_PART_TYPES: ReadonlySet<string> = new Set([
-  "reasoning",
-  "thinking",
-  "redacted_thinking",
-]);
-
-/**
- * Parse, sanitize, and re-serialize a `/serving-endpoints/...` POST
- * body. Returns the original string verbatim when the body is not
- * JSON or no rewrite was needed.
- */
-export function rewriteServingBody(body: string): string {
-  const parsed = json.parseRecord(body);
-  if (!parsed) return body;
-
-  // Runs regardless of `messages`: Databricks refuses to parse a body carrying
-  // an unknown top-level field, so this failure is not specific to a transcript.
-  const toolReasoning = applyToolReasoningCompatibility(parsed);
-  let changed = openaiChat.stripUnsupportedChatFields(parsed).length > 0 || toolReasoning;
-
-  if (Array.isArray(parsed.messages)) {
-    const messages = parsed.messages as ServingChatMessage[];
-    // Evaluated eagerly, not short-circuited: one transcript can need both
-    // reasoning blocks stripped AND a trailing assistant prefill folded back.
-    const stripped = stripReasoningFromServingMessages(messages);
-    const repaired = repairAssistantPrefill(messages);
-    changed = changed || stripped || repaired;
-  }
-
-  return changed ? JSON.stringify(parsed) : body;
-}
-
-/** Prepared fetch arguments plus the post-sanitize body used for diagnostics. */
-export interface RewrittenServingRequest {
-  input: Parameters<typeof fetch>[0];
-  init: Parameters<typeof fetch>[1];
-  body: string;
-}
-
-/**
- * Read and rewrite a serving POST whether its JSON body lives on `init` or on
- * a `Request`. A changed Request is rebuilt without stale byte headers.
- */
-export async function rewriteServingRequest(
-  input: Parameters<typeof fetch>[0],
-  init?: Parameters<typeof fetch>[1],
-): Promise<RewrittenServingRequest> {
-  const request = new Request(input, init);
-  const body = await request.clone().text();
-  const rewritten = rewriteServingBody(body);
-  if (rewritten === body) return { input, init, body };
-
-  const headers = new Headers(request.headers);
-  headers.delete("content-length");
-  headers.delete("content-encoding");
-  return {
-    input: new Request(request, { body: rewritten, headers }),
-    init: undefined,
-    body: rewritten,
-  };
-}
-
-/**
- * Normalize Chat Completions reasoning for GPT models carrying function tools.
- *
- * GPT 5.6 requires reasoning disabled for tool calls on this route, so `none`
- * wins over a caller/default effort. Responses-only models are routed away
- * from Chat Completions before this sanitizer runs.
- */
-export function applyToolReasoningCompatibility(body: Record<string, unknown>): boolean {
-  if (!Array.isArray(body.tools) || body.tools.length === 0) {
-    return false;
-  }
-  const model = stringUtils.trimToNull(body.model);
-  if (!model) return false;
-  if (policy.chatToolReasoningEffort(model) !== "none") return false;
-  if (body.reasoning_effort === "none") return false;
-  body.reasoning_effort = "none";
-  return true;
-}
-
-/**
- * Drop extended-thinking / reasoning blocks from a replayed transcript.
- *
- * Hybrid Claude endpoints (e.g. Sonnet 4.5+) may emit `reasoning` content
- * parts on the first turn. Mastra persists and replays them on the next
- * agent step, but Databricks-hosted Claude rejects those blocks on
- * multi-turn tool continuations. The UI already captured reasoning for
- * display; stripping here keeps provider replay compatible without
- * changing what users see in the chat bubble.
- */
-export function stripReasoningFromServingMessages(messages: ServingChatMessage[]): boolean {
-  let changed = false;
-  for (let i = messages.length - 1; i >= 0; i--) {
-    const msg = messages[i];
-    if (!msg) continue;
-    if (msg.role === "reasoning") {
-      messages.splice(i, 1);
-      changed = true;
-      continue;
-    }
-    if (msg.reasoning !== undefined) {
-      delete msg.reasoning;
-      changed = true;
-    }
-    if (msg.reasoning_content !== undefined) {
-      delete msg.reasoning_content;
-      changed = true;
-    }
-    const parts = openaiChat.chatContentParts(msg.content);
-    if (!parts) continue;
-    const filtered = parts.filter((part) => {
-      const type = part?.type;
-      if (typeof type === "string" && REASONING_PART_TYPES.has(type)) {
-        changed = true;
-        return false;
-      }
-      return true;
-    });
-    if (filtered.length !== parts.length) {
-      msg.content = filtered;
-    }
-    const hasToolCalls = Array.isArray(msg.tool_calls) && msg.tool_calls.length > 0;
-    if (msg.role === "assistant" && !hasToolCalls && isEmptyServingContent(msg.content)) {
-      messages.splice(i, 1);
-      changed = true;
-    }
-  }
-  return changed;
-}
-
-/**
- * Repair a Mastra/AI SDK message replay that Databricks-hosted Claude
- * rejects with `"This model does not support assistant message
- * prefill. The conversation must end with a user message."`.
- *
- * The bug pattern: when an assistant turn streams text *and* a
- * `tool_call`, the AI SDK persists them as two separate assistant
- * entries (text-only and tool-call-only). On the next agent step the
- * tool-call entry is replayed *before* the tool result and the
- * text entry is replayed *after* it, so the conversation ends with a
- * trailing assistant text message. Anthropic interprets that as a
- * prefill request and rejects it on Databricks (the upstream Bedrock
- * route disallows prefill).
- *
- * Fix: when the last message is an assistant text with no `tool_calls`
- * and the chain immediately before it is `assistant(tool_calls=...)`
- * followed only by `tool(...)` results, fold the trailing text back
- * into the `content` of that opening assistant and drop the duplicate.
- */
-export function repairAssistantPrefill(messages: ServingChatMessage[]): boolean {
-  if (messages.length < 2) return false;
-  const last = messages[messages.length - 1];
-  if (!last || last.role !== "assistant" || (last.tool_calls && last.tool_calls.length > 0)) {
-    return false;
-  }
-
-  let i = messages.length - 2;
-  while (i >= 0 && messages[i]?.role === "tool") i--;
-  if (i < 0) return false;
-  const opener = messages[i];
-  if (
-    !opener ||
-    opener.role !== "assistant" ||
-    !opener.tool_calls ||
-    opener.tool_calls.length === 0
-  ) {
-    return false;
-  }
-
-  const merged = [
-    stringUtils.trimToNull(textFromServingContent(opener.content)),
-    stringUtils.trimToNull(textFromServingContent(last.content)),
-  ]
-    .filter((s): s is string => s !== null)
-    .join("\n\n");
-  opener.content = merged;
-  messages.pop();
-  return true;
-}
-
-/**
- * Flatten content to the prose a human would read: text parts only, separated
- * by blank lines, since the result is spliced into another turn's message body.
- */
-function textFromServingContent(content: ServingChatMessage["content"]): string {
-  return openaiChat.chatContentToText(content, { separator: "\n\n", types: ["text"] });
-}
-
-/**
- * Whether a turn carries nothing worth replaying. Deliberately stricter than
- * "has no text": a part of any other type (an image, a cache marker) counts as
- * content, so a message holding one is kept even though it has no prose.
- */
-function isEmptyServingContent(content: ServingChatMessage["content"]): boolean {
-  if (content === undefined) return true;
-  if (typeof content === "string") return content.trim().length === 0;
-  const parts = openaiChat.chatContentParts(content);
-  if (!parts) return true;
-  return parts.every((part) => {
-    if (part?.type === "text") {
-      return typeof part.text !== "string" || part.text.trim().length === 0;
-    }
-    return false;
-  });
-}
-
-/**
- * Parse, sanitize, and re-serialize a `/serving-endpoints/...` non-streaming
- * JSON RESPONSE body. Returns the original string verbatim when the body is
- * not JSON or no rewrite was needed, mirroring
- * {@link rewriteServingBody} on the request side.
- */
-export function rewriteServingResponseBody(body: string): string {
-  const parsed = json.parseRecord(body);
-  if (!parsed) return body;
-  return flattenChoiceMessageContent(parsed) ? JSON.stringify(parsed) : body;
-}
-
-/**
- * Repair array-valued `choices[].delta.content` inside an OpenAI SSE stream
- * without buffering the response. The decoder retains partial lines across
- * arbitrary network chunk boundaries; each complete `data:` line is parsed and
- * re-encoded only when its content needs normalization.
- */
-export function rewriteServingResponseStream(
-  body: ReadableStream<Uint8Array>,
-): ReadableStream<Uint8Array> {
-  const decoder = new TextDecoder();
-  const encoder = new TextEncoder();
-  let buffer = "";
-
-  return body.pipeThrough(
-    new TransformStream<Uint8Array, Uint8Array>({
-      transform(chunk, controller) {
-        buffer += decoder.decode(chunk, { stream: true });
-        let newline = buffer.indexOf("\n");
-        while (newline >= 0) {
-          const line = buffer.slice(0, newline + 1);
-          buffer = buffer.slice(newline + 1);
-          controller.enqueue(encoder.encode(rewriteServingResponseStreamLine(line)));
-          newline = buffer.indexOf("\n");
-        }
-      },
-      flush(controller) {
-        buffer += decoder.decode();
-        if (buffer) controller.enqueue(encoder.encode(rewriteServingResponseStreamLine(buffer)));
-      },
-    }),
-  );
-}
-
-/** Normalize one SSE line while preserving its `data:` spacing and line ending. */
-function rewriteServingResponseStreamLine(line: string): string {
-  const lineEnding = line.endsWith("\r\n") ? "\r\n" : line.endsWith("\n") ? "\n" : "";
-  const content = lineEnding ? line.slice(0, -lineEnding.length) : line;
-  const match = /^(data:\s*)(.*)$/.exec(content);
-  if (!match || match[2] === "[DONE]") return line;
-
-  const parsed = json.parseRecord(match[2]);
-  if (!parsed || !flattenChoiceDeltaContent(parsed)) return line;
-  return `${match[1]}${JSON.stringify(parsed)}${lineEnding}`;
-}
-
-/**
- * Collapse a structured `choices[].message.content` array to the plain string
- * the OpenAI Chat Completions contract specifies.
- *
- * Databricks-hosted Gemini answers a non-streaming `/chat/completions` call
- * with the Gemini-native parts shape:
- *
- * ```json
- * "content": [{ "type": "text", "text": "...", "thoughtSignature": "..." }]
- * ```
- *
- * `@ai-sdk/openai-compatible` validates the response against the OpenAI
- * schema, where `content` is a nullable string, so it throws
- * `Type validation failed: ... expected string, received array` and the whole
- * call fails (`AI_APICallError: Invalid JSON response` on an HTTP 200). Mastra
- * uses `doGenerate` for its side calls, so the visible symptom is
- * `Error generating title` - every thread keeps its placeholder name.
- *
- * Flattening on the wire keeps the repair in one place: neither the agent's
- * stored transcript nor the UI has to know the provider emitted parts. Any
- * non-text part (a reasoning block, `thoughtSignature`-only entry, or inline
- * image) contributes nothing, matching {@link openaiChat.chatContentToText},
- * and an all-parts-empty message flattens to `""` rather than being dropped, so
- * `finish_reason` and `usage` still round-trip.
- */
-export function flattenChoiceMessageContent(payload: Record<string, unknown>): boolean {
-  return flattenChoiceContent(payload, "message");
-}
-
-/**
- * Collapse structured streaming `choices[].delta.content` arrays to strings.
- * Claude reasoning-only chunks become an empty content delta while text parts
- * are concatenated in order.
- */
-export function flattenChoiceDeltaContent(payload: Record<string, unknown>): boolean {
-  return flattenChoiceContent(payload, "delta");
-}
-
-/** Shared choice walker for buffered `message` and streaming `delta` payloads. */
-function flattenChoiceContent(
-  payload: Record<string, unknown>,
-  field: "message" | "delta",
-): boolean {
-  if (!Array.isArray(payload.choices)) return false;
-  let changed = false;
-  for (const choice of payload.choices) {
-    if (!choice || typeof choice !== "object") continue;
-    const container = (choice as Record<string, unknown>)[field];
-    if (!container || typeof container !== "object") continue;
-    const target = container as { content?: unknown };
-    const parts = openaiChat.chatContentParts(target.content);
-    if (!parts) continue;
-    target.content = openaiChat.chatContentToText(parts, { types: ["text"] });
-    changed = true;
-  }
-  return changed;
-}
+export {
+  applyToolReasoningCompatibility,
+  flattenChoiceDeltaContent,
+  flattenChoiceMessageContent,
+  repairAssistantPrefill,
+  rewriteServingBody,
+  rewriteServingRequest,
+  rewriteServingResponseBody,
+  rewriteServingResponseStream,
+  stripReasoningFromServingMessages,
+  type RewrittenServingRequest,
+  type ServingChatMessage,
+} from "@dbx-tools/model/serving-wire";

package/src/trace-io.ts

@@ -1,6 +1,6 @@
 /**
- * Copy each Mastra chat turn's request/response onto the active OTel
- * root span so Databricks-managed MLflow can display them.
+ * Copy each Mastra chat turn's request/response onto its OTel root span
+ * so Databricks-managed MLflow can display them.
  *
  * Why this exists: MLflow's UC `<target>_trace_unified` view picks the
  * one span whose `parent_span_id` is empty and reads request/response
@@ -10,16 +10,19 @@
  * the view never reads, so chat traces arrive with both columns null
  * even when the spans themselves look healthy.
  *
- * The active span inside the Mastra sub-app is still AppKit's HTTP
- * server span (OTel context is async-local). With `OTEL_PROPAGATORS=none`
- * that span is the trace root - see the package README's observability
- * section for why Apps ingress `traceparent` otherwise hides every turn.
+ * AppKit's HTTP instrumentation records the server span in OTel RPC metadata,
+ * which remains available even when Express makes a middleware-layer child
+ * active. When instrumentation is unavailable, this module creates one
+ * request-lifetime server span and runs Mastra under it.
  *
  * @module
  */
 
+import { StringDecoder } from "node:string_decoder";
+
 import { json, log, object } from "@dbx-tools/shared-core";
-import { trace } from "@opentelemetry/api";
+import { context, SpanKind, trace, type Span } from "@opentelemetry/api";
+import { getRPCMetadata, RPCType } from "@opentelemetry/core";
 import type express from "express";
 
 const logger = log.logger("mastra/trace-io");
@@ -27,6 +30,12 @@ const logger = log.logger("mastra/trace-io");
 /** Cap on each payload copied onto a span, so one turn cannot bloat the export. */
 export const TRACE_IO_LIMIT = 8_000;
 
+/** Bound for a non-streaming body captured through `res.end` instead of `res.json`. */
+const RESPONSE_BODY_LIMIT = TRACE_IO_LIMIT * 4;
+
+/** Tracer scope used only when AppKit did not create a recording HTTP span. */
+const CHAT_TURN_TRACER = "@dbx-tools/appkit-mastra/trace-io";
+
 /**
  * Mount-relative Mastra agent invoke paths that carry a chat turn body
  * (`messages`) and stream an assistant answer. Resume / approve verbs
@@ -39,6 +48,25 @@ const AGENT_TURN_ROUTE = /^(?:\/agents\/[^/]+\/(?:stream|generate)(?:\/|$)|\/cha
 export const MLFLOW_SPAN_INPUTS_ATTR = "mlflow.spanInputs";
 export const MLFLOW_SPAN_OUTPUTS_ATTR = "mlflow.spanOutputs";
 
+interface TraceTarget {
+  span: Span;
+  owned: boolean;
+}
+
+/** Return the assistant text from one SSE line, or an empty string. */
+function assistantTextFromSseLine(line: string): string {
+  if (!line.startsWith("data:")) return "";
+  const frame = json.parse(line.slice(5));
+  if (!object.isRecord(frame) || frame.type !== "text-delta") return "";
+  const text =
+    typeof frame.delta === "string"
+      ? frame.delta
+      : object.isRecord(frame.payload)
+        ? frame.payload.text
+        : undefined;
+  return typeof text === "string" ? text : "";
+}
+
 /**
  * Concatenate the assistant's answer out of an AI SDK SSE transcript.
  *
@@ -50,28 +78,109 @@ export const MLFLOW_SPAN_OUTPUTS_ATTR = "mlflow.spanOutputs";
 export function assistantTextFromSse(body: string): string {
   const parts: string[] = [];
   for (const line of body.split("\n")) {
-    if (!line.startsWith("data:")) continue;
-    // SSE carries comment/keepalive lines that are not JSON; they never hold
-    // assistant text, so skipping them is the intended path rather than an error.
-    const frame = json.parse(line.slice(5));
-    if (!object.isRecord(frame) || frame.type !== "text-delta") continue;
-    const text =
-      typeof frame.delta === "string"
-        ? frame.delta
-        : object.isRecord(frame.payload)
-          ? frame.payload.text
-          : undefined;
-    if (typeof text === "string") parts.push(text);
+    const text = assistantTextFromSseLine(line);
+    if (text) parts.push(text);
+  }
+  return parts.join("").slice(0, TRACE_IO_LIMIT);
+}
+
+/** Read the final assistant text from a non-streaming Mastra response. */
+export function assistantTextFromJson(body: unknown): string {
+  if (!object.isRecord(body) || typeof body.text !== "string") return "";
+  return body.text.slice(0, TRACE_IO_LIMIT);
+}
+
+/** Incrementally assemble UTF-8 SSE text while retaining a bounded JSON fallback. */
+class AssistantResponseCollector {
+  private readonly decoder = new StringDecoder("utf8");
+  private answer = "";
+  private body = "";
+  private pendingLine = "";
+  private finished = false;
+
+  write(chunk: unknown): void {
+    if (this.finished) return;
+    if (typeof chunk === "string") {
+      this.consume(chunk);
+      return;
+    }
+    if (ArrayBuffer.isView(chunk)) {
+      this.consume(
+        this.decoder.write(Buffer.from(chunk.buffer, chunk.byteOffset, chunk.byteLength)),
+      );
+    }
+  }
+
+  finish(): string {
+    if (this.finished) return this.answer || this.jsonAnswer();
+    this.finished = true;
+    this.consume(this.decoder.end());
+    this.consumeSseLine(this.pendingLine);
+    this.pendingLine = "";
+    return this.answer || this.jsonAnswer();
+  }
+
+  private consume(text: string): void {
+    if (!text) return;
+    if (this.body.length < RESPONSE_BODY_LIMIT) {
+      this.body += text.slice(0, RESPONSE_BODY_LIMIT - this.body.length);
+    }
+    this.pendingLine += text;
+    let newline = this.pendingLine.indexOf("\n");
+    while (newline >= 0) {
+      this.consumeSseLine(this.pendingLine.slice(0, newline).replace(/\r$/, ""));
+      this.pendingLine = this.pendingLine.slice(newline + 1);
+      newline = this.pendingLine.indexOf("\n");
+    }
+    if (this.pendingLine.length > RESPONSE_BODY_LIMIT) {
+      this.pendingLine = this.pendingLine.slice(-RESPONSE_BODY_LIMIT);
+    }
+  }
+
+  private consumeSseLine(line: string): void {
+    if (this.answer.length >= TRACE_IO_LIMIT) return;
+    const text = assistantTextFromSseLine(line);
+    this.answer += text.slice(0, TRACE_IO_LIMIT - this.answer.length);
+  }
+
+  private jsonAnswer(): string {
+    return assistantTextFromJson(json.parse(this.body));
   }
-  return parts.join("");
 }
 
 /**
- * Express middleware that stamps chat turn I/O onto the active OTel span.
+ * Select the span MLflow treats as the request root.
  *
- * Tee `res.write` / `res.end` instead of listening for `finish`: the HTTP
- * instrumentation ends the root span on finish, and an attribute set after
- * a span ends never reaches the exporter.
+ * OTel HTTP instrumentation stores its server span in RPC metadata. That span
+ * wins over an active Express child. If no recording local span exists, create
+ * one under the current context so enabled W3C propagation remains intact.
+ */
+function resolveTraceTarget(): TraceTarget | undefined {
+  const activeContext = context.active();
+  const rpc = getRPCMetadata(activeContext);
+  if (rpc?.type === RPCType.HTTP && rpc.span.isRecording()) {
+    return { span: rpc.span, owned: false };
+  }
+
+  const active = trace.getActiveSpan();
+  if (active?.isRecording()) return { span: active, owned: false };
+
+  const span = trace
+    .getTracer(CHAT_TURN_TRACER)
+    .startSpan("mastra.chat_turn", { kind: SpanKind.SERVER }, activeContext);
+  if (!span.isRecording()) {
+    span.end();
+    return undefined;
+  }
+  return { span, owned: true };
+}
+
+/**
+ * Express middleware that stamps chat turn I/O onto the exported OTel root.
+ *
+ * Tee `res.write`, `res.json`, and `res.end` instead of listening for `finish`:
+ * HTTP instrumentation ends its server span on finish, and attributes set
+ * after a span ends never reach the exporter.
  *
  * `path` is mount-relative (what the Mastra sub-app sees), e.g.
  * `/agents/support/stream`.
@@ -85,48 +194,71 @@ export function chatTurnTraceIoMiddleware(
     next();
     return;
   }
-  const span = trace.getActiveSpan();
-  if (!span) {
+  const target = resolveTraceTarget();
+  if (!target) {
     next();
     return;
   }
 
   const messages = (req.body as { messages?: unknown } | undefined)?.messages;
   if (messages !== undefined) {
-    span.setAttribute(MLFLOW_SPAN_INPUTS_ATTR, JSON.stringify(messages).slice(0, TRACE_IO_LIMIT));
+    target.span.setAttribute(
+      MLFLOW_SPAN_INPUTS_ATTR,
+      JSON.stringify(messages).slice(0, TRACE_IO_LIMIT),
+    );
   }
 
-  const chunks: string[] = [];
+  const collector = new AssistantResponseCollector();
+  let outputRecorded = false;
+  let ownedSpanEnded = false;
   const passThroughWrite = res.write.bind(res) as (...args: unknown[]) => boolean;
   const passThroughEnd = res.end.bind(res) as (...args: unknown[]) => unknown;
-  const collect = (chunk: unknown): void => {
-    if (typeof chunk === "string") chunks.push(chunk);
-    else if (Buffer.isBuffer(chunk)) chunks.push(chunk.toString("utf8"));
+  const passThroughJson = res.json.bind(res) as (body?: unknown) => express.Response;
+  const recordOutput = (answer: string): void => {
+    if (outputRecorded || !answer) return;
+    target.span.setAttribute(MLFLOW_SPAN_OUTPUTS_ATTR, answer.slice(0, TRACE_IO_LIMIT));
+    outputRecorded = true;
+  };
+  const endOwnedSpan = (): void => {
+    if (!target.owned || ownedSpanEnded) return;
+    ownedSpanEnded = true;
+    target.span.end();
   };
 
   res.write = ((chunk: unknown, ...rest: unknown[]) => {
-    collect(chunk);
+    collector.write(chunk);
     return passThroughWrite(chunk, ...rest);
   }) as typeof res.write;
 
+  res.json = ((body?: unknown) => {
+    recordOutput(assistantTextFromJson(body));
+    return passThroughJson(body);
+  }) as typeof res.json;
+
   res.end = ((chunk?: unknown, ...rest: unknown[]) => {
-    collect(chunk);
-    const answer = assistantTextFromSse(chunks.join(""));
-    if (answer) {
-      span.setAttribute(MLFLOW_SPAN_OUTPUTS_ATTR, answer.slice(0, TRACE_IO_LIMIT));
+    collector.write(chunk);
+    recordOutput(collector.finish());
+    try {
+      return passThroughEnd(chunk, ...rest);
+    } finally {
+      endOwnedSpan();
     }
-    return passThroughEnd(chunk, ...rest);
   }) as typeof res.end;
 
-  next();
+  res.once("close", () => {
+    recordOutput(collector.finish());
+    endOwnedSpan();
+  });
+
+  context.with(trace.setSpan(context.active(), target.span), next);
 }
 
 /**
  * Install {@link chatTurnTraceIoMiddleware} on a Mastra Express sub-app.
  *
  * Call before `MastraServer.init()` so the layer sits ahead of the agent
- * routes. Safe to call unconditionally: with no active span (local, no
- * OTLP) the middleware is a no-op.
+ * routes. Safe to call unconditionally: when no recording tracer is registered,
+ * the middleware is a no-op.
  */
 export function attachChatTurnTraceIo(app: express.Express): void {
   app.use(chatTurnTraceIoMiddleware);