@nodetool-ai/websocket 0.8.1 → 0.8.2

package/dist/session/chat-turn.js

@@ -1,3 +1,9 @@
+import { registerChatRunTrace, registerWorkflowRunTrace, withRegisteredRunTrace, settleRegisteredRunTrace } from "@nodetool-ai/execution";
+import { getRunTraceScope, suppressRunTraceContent } from "@nodetool-ai/runtime";
+import { TRACE_RESTRICTED_CAPABILITY_MODULES } from "@nodetool-ai/protocol";
+import { trace, SpanStatusCode } from "@opentelemetry/api";
+import { settleRunTrace, registerRunTraceParent, registerRunTraceParents, threadHasTraceContentTools } from "@nodetool-ai/models";
+import { capabilityModuleOf, listCapabilitySpecs } from "@nodetool-ai/agents";
 import { createLogger } from "@nodetool-ai/config";
 import { randomUUID } from "node:crypto";
 import { SUPERSEDED_TOOL_RESULT, repairOrphanedToolCalls } from "../chat-tool-call-repair.js";
@@ -13,7 +19,7 @@ import { ACTIVE_MODEL_CONTEXT_KEY, DIRECT_TOOL_NAMES, RUN_BUDGET_CONTEXT_KEY, cr
 import { isModelSelection, PROVIDER_IDS, NO_MODEL_SELECTED_MESSAGE, noMediaModelSelectedMessage } from "@nodetool-ai/protocol";
 import { Tool, compactAssetUris, compactResourceIds } from "@nodetool-ai/agents";
 import { createChatCodeActSession, createSandboxClock, sandboxPackagesForChat, CODEACT_RESIDENT_TOOL_NAMES, EXECUTE_CODE_TOOL_NAME } from "@nodetool-ai/agents";
-import { getBuiltinTools, getAllMcpTools, getGoogleWorkspaceTools, getApifyTools, getSerpApiTools, toolForCapabilityName, gateTools, capabilityFromTool, createCapabilityRun, contextSecretAvailability, BackgroundSubtaskRegistry, UNGATED, PERMISSION_GATE_CONTEXT_KEY, GAME_ASSET_NODE_RUNNER_CONTEXT_KEY, extractInjectableImages, registerTemporaryImageHandle } from "@nodetool-ai/agents";
+import { availableBuiltinToolNames, GOOGLE_WORKSPACE_TOOL_NAMES, APIFY_TOOL_NAMES, SERPAPI_TOOL_NAMES, capabilityForName, capabilityProviderTool, getAllMcpCapabilitySpecs, toolFromCapability, capabilityFromTool, createCapabilityRun, contextSecretAvailability, BackgroundSubtaskRegistry, PERMISSION_GATE_CONTEXT_KEY, GAME_ASSET_NODE_RUNNER_CONTEXT_KEY, extractInjectableImages, registerTemporaryImageHandle } from "@nodetool-ai/agents";
 import { mcpToolHostDeps } from "../mcp-tool-deps.js";
 import { externalMcpToolsFor } from "../external-mcp.js";
 import { formatSkillCatalogForPrompt, mergeSystemSkills, formatMemoriesForPrompt } from "@nodetool-ai/agents";
@@ -25,7 +31,41 @@ import { getSetting } from "../settings-registry.js";
 import { appendContextToLastUser, createWorkflowResponseContent, dbMessageToProviderMessage, extractTextContent, historySinceCompaction, invokedSkillsSection, toolResultDisplayText } from "./chat-history.js";
 import { chooseCompactionCut, compactionThreshold, COMPACTION_SUMMARY_MAX_TOKENS, holdsTranscriptServerSide, readCompactionSettings, summarizeForCompaction } from "./chat-compaction.js";
 import { createRelayActivityWaiter, DEFAULT_RUN_JOB_EXECUTION_OPTIONS } from "./job-execution.js";
+function restrictedTraceToolNames() {
+    const modules = new Set(TRACE_RESTRICTED_CAPABILITY_MODULES);
+    return listCapabilitySpecs().filter((spec) => modules.has(capabilityModuleOf(spec.name) ?? "")).map((spec) => spec.name);
+}
+function seedStoredToolProvenance(messages) {
+    for (const message of messages) {
+        const names = message.name ? [message.name] : [];
+        if (isCompactionMessage(message)) {
+            names.push(...(message.tools ?? []));
+        }
+        for (const call of message.tool_calls ?? []) {
+            if (!isRecord(call)) {
+                continue;
+            }
+            if (isString(call.name)) {
+                names.push(call.name);
+            }
+            if (isRecord(call.function) && isString(call.function.name)) {
+                names.push(call.function.name);
+            }
+        }
+        for (const name of names) {
+            const moduleName = capabilityModuleOf(name);
+            if (moduleName === "email" || moduleName === "google" || moduleName === "browser") {
+                suppressRunTraceContent();
+            }
+        }
+    }
+}
 const log = createLogger("nodetool.websocket.runner");
+/** Providers whose credential is an account sign-in, not an API key. */
+const SIGN_IN_PROVIDER_LABELS = {
+    [PROVIDER_IDS.CODEX]: "Codex",
+    [PROVIDER_IDS.CLAUDE_AGENT_SDK]: "Claude"
+};
 const GENERIC_CHAT_ERROR = "Something went wrong while processing your request. Please try again.";
 /**
  * Keep implementation details out of the client error envelope. Provider
@@ -105,24 +145,20 @@ function providerFailureMessage(error, providerId, model) {
     const gemini = geminiFailureSummary(error, providerId);
     const groq = groqRequestFailureMessage(error);
     const status = httpStatusFromError(error) ?? gemini?.status ?? null;
+    const statusCode = status ?? undefined;
     const rawMessage = groq ?? gemini?.message ?? safeClientErrorText(error);
-    let bodyMessage = null;
-    if (isObjectLike(error)) {
-        const body = error.body ?? error.response;
-        if (isObjectLike(body)) {
-            const bodyError = isObjectLike(body.error) ? body.error : body;
-            if (isObjectLike(bodyError) && isString(bodyError.message)) {
-                const safe = safeClientErrorText(bodyError.message);
-                bodyMessage = safe === GENERIC_CHAT_ERROR ? null : safe;
-            }
-        }
-    }
+    const body = isObjectLike(error) ? (error.body ?? error.response) : undefined;
+    const bodyError = isObjectLike(body) && isObjectLike(body.error) ? body.error : body;
+    const bodyText = isObjectLike(bodyError) && isString(bodyError.message)
+        ? safeClientErrorText(bodyError.message)
+        : null;
+    const bodyMessage = bodyText === GENERIC_CHAT_ERROR ? null : bodyText;
     if (providerId.toLowerCase() === "openai" &&
         /\borganization must be verified\b/i.test(bodyMessage ?? rawMessage)) {
         return {
             errorType: status === null ? "error" : "http_status_error",
             message: "This model requires a verified OpenAI organization. Verify your provider organization or choose a model your account can access.",
-            statusCode: status ?? undefined
+            statusCode
         };
     }
     const detail = providerFailureDetail(error);
@@ -131,28 +167,31 @@ function providerFailureMessage(error, providerId, model) {
         return {
             errorType: "error",
             message: `The ${target} request is too large for the model context window. Shorten the conversation or remove some attachments and try again.`,
-            statusCode: status ?? undefined
+            statusCode
         };
     }
     if (detail?.code === "provider_auth") {
+        const signIn = SIGN_IN_PROVIDER_LABELS[detail.provider];
+        const remedy = signIn ? `Sign in to ${signIn} again` : "Check the API key";
         return {
             errorType: "error",
-            message: `Authentication failed: ${detail.provider} rejected the configured credentials. Check the API key in Settings → Models & Providers.`,
-            statusCode: status ?? undefined
+            message: `Authentication failed: ${detail.provider} rejected the configured credentials. ${remedy} in Settings → Models & Providers.`,
+            statusCode
         };
     }
     if (status === null) {
-        return { errorType: "error", message: rawMessage, statusCode: undefined };
+        return { errorType: "error", message: rawMessage, statusCode };
     }
-    let message = rawMessage;
+    const modelName = model || "requested";
+    let message;
     if (gemini?.kind === "quota") {
         message = `Account quota exhausted for ${providerId}/${model}. Check your ${providerId} quota or billing and try again later.`;
     }
     else if (gemini?.kind === "model") {
-        message = `Model ${model || "requested"} is unavailable on ${providerId}. Choose another model or check that your account has access to it.`;
+        message = `Model ${modelName} is unavailable on ${providerId}. Choose another model or check that your account has access to it.`;
     }
     else if (gemini?.kind === "tools") {
-        message = `Model ${model || "requested"} does not support the requested tools on ${providerId}. Choose a model with tool support or disable tools.`;
+        message = `Model ${modelName} does not support the requested tools on ${providerId}. Choose a model with tool support or disable tools.`;
     }
     else if (status === 400) {
         message = `Bad request: ${groq ?? bodyMessage ?? rawMessage}`;
@@ -168,20 +207,20 @@ function providerFailureMessage(error, providerId, model) {
             "Access forbidden: You don't have permission for this resource. Check the provider key's model and account access.";
     }
     else if (status === 404) {
-        message = `Model ${model || "requested"} was not found or is unavailable on ${providerId}. Choose another model or check account access.`;
+        message = `Model ${modelName} was not found or is unavailable on ${providerId}. Choose another model or check account access.`;
     }
     else if (status === 413) {
-        const detail = groq ?? bodyMessage ?? rawMessage;
+        const bodyDetail = bodyMessage ?? rawMessage;
         message =
             groq ??
-                (detail === GENERIC_CHAT_ERROR
+                (bodyDetail === GENERIC_CHAT_ERROR
                     ? "Request is too large for the provider. Shorten the conversation or remove some attachments and try again."
-                    : `Request is too large for the provider (${detail}). Shorten the conversation or remove some attachments and try again.`);
+                    : `Request is too large for the provider (${bodyDetail}). Shorten the conversation or remove some attachments and try again.`);
     }
     else if (status === 429) {
-        message = groq
-            ? groq
-            : "Rate limited: Too many requests or insufficient provider quota. Check your provider plan and try again later.";
+        message =
+            groq ??
+                "Rate limited: Too many requests or insufficient provider quota. Check your provider plan and try again later.";
     }
     else if (status >= 500) {
         message =
@@ -192,7 +231,7 @@ function providerFailureMessage(error, providerId, model) {
     else {
         message = `HTTP error (${status}): ${rawMessage}`;
     }
-    return { errorType: "http_status_error", message, statusCode: status };
+    return { errorType: "http_status_error", message, statusCode };
 }
 /**
  * How many of a user's newest memories the turn reads to build its block. The
@@ -487,11 +526,12 @@ export class ChatTurnHandler {
         delete data.thread_id;
         const userId = this.session.requireUserId();
         delete data.user_id;
-        return Message.create({
-            thread_id: threadId,
-            user_id: userId,
-            ...data
-        });
+        const message = await Message.create({ thread_id: threadId, user_id: userId, ...data });
+        const scope = getRunTraceScope();
+        if (scope) {
+            await registerRunTraceParent(userId, scope.runId, { kind: "message", id: message.id });
+        }
+        return message;
     }
     /**
      * Persist raw image bytes carried on an assistant message (native providers
@@ -923,6 +963,7 @@ export class ChatTurnHandler {
                 : null,
             assetOutputMode: this.session.mode === "text" ? "data_uri" : "temp_url"
         });
+        context.runTraceContext = getRunTraceScope() ?? null;
         // This context is built here rather than copied from the turn, so the
         // turn's gate has to be put on it by hand. The node's own loops read it
         // with `gateFromContext`; a node run outside a chat turn carries none and
@@ -1043,9 +1084,32 @@ export class ChatTurnHandler {
         log.debug("Chat message", { threadId, model, provider: providerId });
         // Save user message to DB — matches Python's _save_message_to_db_async(data)
         const turnMessage = await this.saveMessageToDb(data);
+        const rootContext = createRuntimeContext({ jobId: turnMessage.id, threadId, workflowId, userId, workspace: null });
+        await registerChatRunTrace(rootContext, { messageId: turnMessage.id, threadId });
+        return withRegisteredRunTrace(rootContext, "chat.turn", async () => {
+            try {
+                if (await threadHasTraceContentTools(userId, threadId, restrictedTraceToolNames())) {
+                    suppressRunTraceContent();
+                }
+                await this.handlePersistedChatTurn(data, turnMessage, rootContext, requestSeq, signal);
+                await settleRegisteredRunTrace(rootContext, signal?.aborted ? "cancelled" : "completed");
+            }
+            catch (error) {
+                await settleRegisteredRunTrace(rootContext, signal?.aborted ? "cancelled" : "failed", error instanceof Error ? error.message : String(error));
+                throw error;
+            }
+        });
+    }
+    async handlePersistedChatTurn(data, turnMessage, rootContext, requestSeq, signal) {
+        const threadId = String(data.thread_id);
+        const workflowId = isString(data.workflow_id) ? data.workflow_id : null;
+        const userId = this.session.requireUserId();
+        const providerId = String(data.provider);
+        const model = String(data.model);
         if (requestSeq !== undefined && requestSeq !== this.chatRequestSeq)
             return;
         if (!this.session.resolveProvider) {
+            await settleRegisteredRunTrace(rootContext, "failed", "No provider resolver configured");
             await this.session.send({
                 type: "error",
                 message: "No provider resolver configured",
@@ -1134,9 +1198,12 @@ export class ChatTurnHandler {
             toolCalls: null,
             threadId: null
         });
+        const copiedMessageIds = new Set();
         const convertDbMessages = (rows) => {
             const out = [];
             for (const m of rows) {
+                copiedMessageIds.add(m.id);
+                seedStoredToolProvenance([m]);
                 const pm = dbMessageToProviderMessage(m, this.session.userId);
                 if (pm)
                     out.push(pm);
@@ -1160,6 +1227,7 @@ export class ChatTurnHandler {
                 reverse: true,
                 limit: SESSION_PROBE_WINDOW
             });
+            seedStoredToolProvenance(recent);
             const probeHasWholeThread = recent.length < SESSION_PROBE_WINDOW;
             // `recent` is newest-first. Walk to the first boundary marker: an
             // assistant carrying a session token, or a compaction record. Both claim
@@ -1209,6 +1277,9 @@ export class ChatTurnHandler {
                     systemHash: probeSession.systemHash,
                     checkpoint: 1
                 };
+                if (probeSession.traceContentSuppressed) {
+                    priorSession.traceContentSuppressed = true;
+                }
                 // Absolute position to persist: prior prefix + the prior assistant + the
                 // new turns — identical to what the full-load path would store.
                 sessionCheckpointOverride =
@@ -1217,6 +1288,10 @@ export class ChatTurnHandler {
                     const [rows] = await Message.paginate(threadId, { limit: 1000 });
                     const full = convertDbMessages(historySinceCompaction(rows));
                     full.unshift(systemChatMessage());
+                    const scope = getRunTraceScope();
+                    if (scope) {
+                        await registerRunTraceParents(userId, scope.runId, [...copiedMessageIds].map((id) => ({ kind: "message", id })));
+                    }
                     return full;
                 };
             }
@@ -1237,6 +1312,9 @@ export class ChatTurnHandler {
                 priorSession = lastMatchingProviderSession(kept, providerId, model);
             }
         }
+        if (priorSession?.traceContentSuppressed) {
+            suppressRunTraceContent();
+        }
         // Expose the read-only `run_search` fan-out primitive by default. A client
         // can opt out by sending `enable_read_only_search: false`.
         const enableReadOnlySearch = data.enable_read_only_search !== false;
@@ -1286,13 +1364,6 @@ export class ChatTurnHandler {
         // it starts (invariant I-2). A loop that made its own would hand every
         // nested agent a fresh allowance, which is no ceiling at all.
         const turnBudget = await createChatTurnBudget();
-        const gatedRun = (context) => createCapabilityRun({
-            context,
-            gate: chatGate,
-            projectId: chatProjectId,
-            availableSecrets: contextSecretAvailability(context),
-            budget: turnBudget
-        });
         // The user's external MCP servers (Blender, Hugging Face, …). Their tools
         // are plain belt tools named `mcp_<server>_<tool>`, so every loop reaches
         // them the same way; a server that does not answer is skipped, not fatal.
@@ -1305,40 +1376,48 @@ export class ChatTurnHandler {
                 error: err instanceof Error ? err.message : String(err)
             });
         }
-        const rawToolbelt = [
-            ...getBuiltinTools(),
-            ...(googleWorkspace ? getGoogleWorkspaceTools() : []),
-            // Apify and SerpAPI have no `nodetool.*` namespace, so the belt is how
-            // a chat discovers them (`nodetool.searchTools("apify")`) at all.
-            ...getApifyTools(gatedRun),
-            ...getSerpApiTools(gatedRun),
-            ...externalMcpTools,
-            ...getAllMcpTools({
-                registry: this.session.nodeRegistry,
-                providers: chatProviders,
-                ...mcpToolHostDeps()
-            }),
-            toolForCapabilityName("list_collections"),
-            toolForCapabilityName("query_collection"),
-            runNodeTool
+        const legacyCapabilities = [runNodeTool, ...externalMcpTools].map(capabilityFromTool);
+        const nativeSpecs = [
+            ...availableBuiltinToolNames().map((name) => capabilityForName(name).spec),
+            ...(googleWorkspace ? GOOGLE_WORKSPACE_TOOL_NAMES : []).map((name) => capabilityForName(name).spec),
+            ...[...APIFY_TOOL_NAMES, ...SERPAPI_TOOL_NAMES].map((name) => capabilityForName(name).spec),
+            ...getAllMcpCapabilitySpecs({ registry: this.session.nodeRegistry, providers: chatProviders }),
+            capabilityForName("list_collections").spec,
+            capabilityForName("query_collection").spec
         ];
-        // De-duplicate by name (builtins / mcp / extras may overlap); first wins.
-        const dedupedToolbelt = [];
-        const seenToolNames = new Set();
-        for (const tool of rawToolbelt) {
-            if (seenToolNames.has(tool.name))
-                continue;
-            seenToolNames.add(tool.name);
-            dedupedToolbelt.push(tool);
-        }
-        // Wrap the toolbelt in the permission gate. The wrapper is transparent
-        // except for `process()`, so the chat loop AND any `run_subtask` child
-        // loop inherit gating by simply calling `tool.process()`.
-        const baseTools = gateTools(dedupedToolbelt, chatGate);
+        const specsByName = new Map();
+        for (const spec of [...nativeSpecs, ...legacyCapabilities.map((entry) => entry.spec)]) {
+            if (!specsByName.has(spec.name)) {
+                specsByName.set(spec.name, spec);
+            }
+        }
+        const baseSpecs = [...specsByName.values()];
+        const legacyByName = new Map(legacyCapabilities.map((entry) => [entry.spec.name, entry]));
+        // Delegated executors still consume Tool[]. Adapt only at that boundary.
+        // Each child context owns a separate run so its memory and depth are local.
+        const childRuns = new WeakMap();
+        const childRun = (context) => {
+            let run = childRuns.get(context);
+            if (!run) {
+                run = createCapabilityRun({
+                    context, gate: chatGate, nodeRegistry: this.session.nodeRegistry,
+                    providers: chatProviders, budget: turnBudget,
+                    availableSecrets: contextSecretAvailability(context),
+                    ...mcpToolHostDeps(), capabilities: legacyCapabilities
+                });
+                childRuns.set(context, run);
+            }
+            return run;
+        };
+        const baseTools = baseSpecs.map((spec) => {
+            const entry = legacyByName.get(spec.name)
+                ?? capabilityForName(spec.name);
+            return toolFromCapability(entry.spec, entry.impl, childRun);
+        });
         // Inject the recursive-decomposition primitive (ungated — it spawns a
         // child loop whose own tools are the gated `baseTools`). Child events
         // stream back tagged with `parent_tool_call_id` so the UI can nest cards.
-        const serverTools = baseTools.slice();
+        const serverTools = baseSpecs.slice();
         // The same runtime the delegation tools take, kept for the capability run
         // below: provider, model, the parent belt, and the event forwarder.
         let subAgentRuntime;
@@ -1384,18 +1463,9 @@ export class ChatTurnHandler {
             // `parent_tool_call_id` / `subtask_depth` tagging are unchanged.
             // `start_subtask` / `wait_subtasks` share the per-turn registry above:
             // spawn returns immediately, and the parent collects on its own terms.
-            const delegationRun = (context) => createCapabilityRun({
-                context,
-                // Ungated on purpose, as before: spawning a child loop has no side
-                // effect of its own, and the child's tools are the gated `baseTools`.
-                gate: UNGATED,
-                availableSecrets: contextSecretAvailability(context),
-                subAgent: subAgentRuntime,
-                budget: turnBudget
-            });
-            serverTools.unshift(toolForCapabilityName("run_subtask", delegationRun));
-            serverTools.unshift(toolForCapabilityName("start_subtask", delegationRun));
-            serverTools.unshift(toolForCapabilityName("wait_subtasks", delegationRun));
+            serverTools.unshift(capabilityForName("run_subtask").spec);
+            serverTools.unshift(capabilityForName("start_subtask").spec);
+            serverTools.unshift(capabilityForName("wait_subtasks").spec);
             // Plan mode's own capability: `create_plan` runs the TaskPlanner over the
             // objective and streams the plan into the thread, executing nothing. It
             // is offered only here, because plan mode is the one place where a plan
@@ -1404,7 +1474,7 @@ export class ChatTurnHandler {
             // mode the turn started in — `liveMode` covers a mid-turn switch for the
             // gate, not for which tools were offered.
             if (permissionMode === "plan") {
-                serverTools.unshift(toolForCapabilityName("create_plan", delegationRun));
+                serverTools.unshift(capabilityForName("create_plan").spec);
             }
             // The other half: `execute_plan` runs a plan the user has already seen,
             // on the same ParallelTaskExecutor `Agent` uses. It is on the belt in
@@ -1412,12 +1482,12 @@ export class ChatTurnHandler {
             // `blocked_in_plan_mode`, which tells the model to have the user switch
             // out, where dropping it from the belt entirely would only say the tool
             // does not exist. It is *offered* everywhere but plan mode, below.
-            serverTools.unshift(...gateTools([toolForCapabilityName("execute_plan", delegationRun)], chatGate));
+            serverTools.unshift(capabilityForName("execute_plan").spec);
             // Read-only fan-out search (opt-in). Reuses the same runtime — the
             // capability filters the parent belt to its read-only allowlist
             // internally, so passing the full snapshot is correct.
             if (enableReadOnlySearch) {
-                serverTools.unshift(toolForCapabilityName("run_search", delegationRun));
+                serverTools.unshift(capabilityForName("run_search").spec);
             }
         }
         const serverToolMap = new Map(serverTools.map((t) => [t.name, t]));
@@ -1426,7 +1496,7 @@ export class ChatTurnHandler {
             permissionMode,
             resolved: serverTools.map((t) => t.name)
         });
-        const serverSchemas = serverTools.map((t) => t.toProviderTool());
+        const serverSchemas = serverTools.map((t) => capabilityProviderTool(t));
         // Every client tool the connected UI registered is exposed. They used to be
         // gated on an active workflow, which made the editor tools unreachable from
         // plain chat; they are deferred behind `nodetool.searchTools()` anyway, and each one now
@@ -1484,7 +1554,11 @@ export class ChatTurnHandler {
             workspace: chatWorkspace,
             authToken: this.deps.authToken()
         });
+        ctx.runTraceContext = rootContext.runTraceContext;
         ctx.set("chat_message_id", turnMessage.id);
+        if (ctx.runTraceContext) {
+            await registerRunTraceParents(userId, ctx.runTraceContext.runId, [...copiedMessageIds].map((id) => ({ kind: "message", id })));
+        }
         const detachPredictions = attachChatPredictionForwarder((listener) => ctx.addMessageListener(listener), (msg) => this.session.sendDetached(msg), { threadId: threadId || null, workflowId });
         // A chat turn that generates an image or a video spends real money without
         // ever constructing an ExecutionSession, so the ledger is attached here
@@ -1530,7 +1604,7 @@ export class ChatTurnHandler {
             budget: turnBudget,
             secretPrompt: (request) => this.requestSecretEntry(threadId, request),
             ...mcpToolHostDeps(),
-            capabilities: [capabilityFromTool(runNodeTool)]
+            capabilities: legacyCapabilities
         });
         const gameAssetRun = this.chatCapabilityRun;
         const gameAssetNodeRunner = (args) => gameAssetRun.invoke("run_node", args);
@@ -1632,9 +1706,20 @@ export class ChatTurnHandler {
         // What to persist onto the assistant message: the provider's token, but with
         // the absolute checkpoint when the fast path sent only a delta (the
         // provider's emitted checkpoint is relative to the trimmed view).
-        const sessionForPersist = () => sessionCheckpointOverride != null && capturedSession
-            ? { ...capturedSession, checkpoint: sessionCheckpointOverride }
-            : capturedSession;
+        const sessionForPersist = () => {
+            if (!capturedSession) {
+                return null;
+            }
+            const contentSuppressed = getRunTraceScope()?.policy.contentSuppressed || priorSession?.traceContentSuppressed || capturedSession.traceContentSuppressed;
+            const persisted = { ...capturedSession };
+            if (sessionCheckpointOverride != null) {
+                persisted.checkpoint = sessionCheckpointOverride;
+            }
+            if (contentSuppressed) {
+                persisted.traceContentSuppressed = true;
+            }
+            return persisted;
+        };
         // Cap on tool-calling rounds before the loop stops. Generous enough to
         // build a multi-component app UI or run a long edit session in one turn —
         // 10 was too low and cut off the app builder mid-build.
@@ -1715,9 +1800,14 @@ export class ChatTurnHandler {
             }
             else if (serverTool) {
                 try {
-                    toolResult = await Tool.executeTool(serverTool, ctx, toolCall.args, {
-                        toolCallId: toolCall.id
-                    });
+                    const args = { ...Tool.stripMessage(toolCall.args) };
+                    if (serverTool.needsToolCallId) {
+                        args["_tool_call_id"] = toolCall.id;
+                    }
+                    if (!this.chatCapabilityRun) {
+                        throw new Error("Chat capability run is unavailable");
+                    }
+                    toolResult = await this.chatCapabilityRun.invoke(serverTool.name, args);
                 }
                 catch (err) {
                     const errMsg = err instanceof Error ? err.message : String(err);
@@ -2020,6 +2110,7 @@ export class ChatTurnHandler {
                 role: "user",
                 execution_event_type: COMPACTION_EVENT_TYPE,
                 content: compactionMessageContent(summary),
+                tools: getRunTraceScope()?.policy.contentSuppressed ? restrictedTraceToolNames() : null,
                 created_at: Number.isFinite(boundary)
                     ? new Date(boundary - 1).toISOString()
                     : cut.keep[0].created_at
@@ -2260,6 +2351,7 @@ export class ChatTurnHandler {
                 return;
             }
             const errMsg = err instanceof Error ? err.message : String(err);
+            await settleRegisteredRunTrace(rootContext, "failed", errMsg);
             log.error("Chat processing error", { threadId, error: errMsg });
             let errorType = "error";
             let statusCode;
@@ -3076,6 +3168,11 @@ export class ChatTurnHandler {
             if (cancelled())
                 return;
             const errMsg = err instanceof Error ? err.message : String(err);
+            const scope = getRunTraceScope();
+            if (scope) {
+                trace.getActiveSpan()?.setStatus({ code: SpanStatusCode.ERROR });
+                await settleRunTrace(userId, scope.runId, { status: "failed", error: scope.policy.contentSuppressed ? "Execution failed" : errMsg, secretValues: scope.secretValues, contentSuppressed: scope.policy.contentSuppressed });
+            }
             log.error("Media generation error", { threadId, mode, error: errMsg });
             const failure = providerFailureMessage(err, providerId, modelId);
             const failureMessage = `Generation failed: ${failure.message}`;
@@ -3160,6 +3257,11 @@ export class ChatTurnHandler {
                 "messages";
             // Build chat history for params — matches Python
             const [dbMessages] = await Message.paginate(threadId, { limit: 1000 });
+            const inheritedTrace = getRunTraceScope();
+            seedStoredToolProvenance(dbMessages);
+            if (inheritedTrace) {
+                await registerRunTraceParents(userId, inheritedTrace.runId, dbMessages.map((m) => ({ kind: "message", id: m.id })));
+            }
             const chatHistorySerialized = dbMessages.map((m) => ({
                 role: m.role,
                 content: m.content,
@@ -3207,6 +3309,7 @@ export class ChatTurnHandler {
                 workspace,
                 assetOutputMode: this.session.mode === "text" ? "data_uri" : "temp_url"
             });
+            context.runTraceContext = getRunTraceScope() ?? null;
             // Expose executor/node-type resolution for sub-workflow nodes
             context.setResolveExecutor((node) => this.session.resolveExecutor(node));
             if (this.session.resolveNodeType) {
@@ -3215,6 +3318,8 @@ export class ChatTurnHandler {
                     : this.session.resolveNodeType;
                 context.setResolveNodeType((nodeType) => resolverObj.resolveNodeType(nodeType));
             }
+            await Job.create({ id: jobId, workflow_id: workflowId, user_id: userId, status: "running", params, graph });
+            await registerWorkflowRunTrace(context, { jobId, workflowId });
             // Create and run workflow (A5: via the ExecutionSession facade — see
             // the identical note in `startJobInner`).
             const session = await ExecutionSession.create({
@@ -3249,20 +3354,6 @@ export class ChatTurnHandler {
                 }
             };
             this.deps.jobs.registerJob(jobId, active);
-            // Persist job to DB (best-effort)
-            try {
-                await Job.create({
-                    id: jobId,
-                    workflow_id: workflowId,
-                    user_id: userId,
-                    status: "running",
-                    params,
-                    graph
-                });
-            }
-            catch (error) {
-                this.session.logError("workflow job persistence failed", error);
-            }
             // The run already started inside `ExecutionSession.create()` above.
             const executePromise = session.result;
             // Stream events, collect output_update results
@@ -3441,6 +3532,11 @@ export class ChatTurnHandler {
         }
         catch (err) {
             const errMsg = err instanceof Error ? err.message : String(err);
+            const scope = getRunTraceScope();
+            if (scope) {
+                trace.getActiveSpan()?.setStatus({ code: SpanStatusCode.ERROR });
+                await settleRunTrace(userId, scope.runId, { status: "failed", error: scope.policy.contentSuppressed ? "Execution failed" : errMsg, secretValues: scope.secretValues, contentSuppressed: scope.policy.contentSuppressed });
+            }
             log.error("Workflow message error", {
                 threadId,
                 workflowId,