@armoriq/sdk-dev 0.6.10 → 0.8.0

package/dist/integrations/langchain.js

@@ -21,6 +21,14 @@
  * @langchain/core is an optional peer dep, loaded lazily so importing this module
  * never requires langchain to be installed.
  *
+ * Compatibility details, tested versions, terminal hooks, and limitations are
+ * maintained in docs/integrations/compatibility-matrix.md. Keep callback
+ * chaining by adding this handler to the existing `callbacks` array; never
+ * replace that array.
+ * Root chain/model terminal callbacks close their request automatically. Call
+ * `await handler.close()` only for host-aborted runs that emit no terminal
+ * callback; it is safe and exact-once.
+ *
  * Usage:
  *   const armoriq = new ArmorIQLangChain({ client, mode: 'sdk' });
  *   const handler = await armoriq.forUser('alice@acme.com', { goal: 'reconcile' });
@@ -30,9 +38,14 @@
 Object.defineProperty(exports, "__esModule", { value: true });
 exports.ArmorIQLangChainEnforcer = exports.ArmorIQLangChain = void 0;
 exports.toolCallsFromLLMResult = toolCallsFromLLMResult;
+exports.usageFromLLMResult = usageFromLLMResult;
+exports.outputFromLLMResult = outputFromLLMResult;
 exports.nameFromSerialized = nameFromSerialized;
 exports.coerceArgs = coerceArgs;
+const tool_registry_1 = require("../tool_registry");
+const tool_push_1 = require("../tool_push");
 const exceptions_1 = require("../exceptions");
+const tool_name_1 = require("../tool_name");
 class ArmorIQLangChain {
     client;
     mode;
@@ -49,8 +62,7 @@ class ArmorIQLangChain {
         this.defaultMcpName = opts.defaultMcpName;
         this.customParser = opts.toolNameParser;
         this.approvalWaitSeconds =
-            opts.approvalWaitSeconds ??
-                Number(process.env.ARMORIQ_APPROVAL_WAIT_SECONDS ?? '300');
+            opts.approvalWaitSeconds ?? Number(process.env.ARMORIQ_APPROVAL_WAIT_SECONDS ?? '300');
         this.approvalPollInterval = opts.approvalPollInterval ?? 5;
     }
     async bootstrap() {
@@ -61,16 +73,24 @@ class ArmorIQLangChain {
     async toolNameParser() {
         if (this.customParser)
             return this.customParser;
-        const toolMap = (await this.bootstrap()).toolMap ?? {};
         const defaultMcp = this.defaultMcpName ?? 'unknown';
+        let toolMap = {};
+        try {
+            toolMap = (await this.bootstrap()).toolMap ?? {};
+        }
+        catch {
+            // Bootstrap enriches tool-name routing only. A model-only LangChain
+            // request must still complete when that optional control-plane call is
+            // unavailable; a later tool enforcement check remains fail-closed.
+            console.warn('[armoriq] langchain bootstrap unavailable; using safe tool-name fallback');
+        }
         return (toolName) => {
             const mcp = toolMap[toolName];
             if (mcp)
                 return { mcp, action: toolName };
-            if (toolName.includes('__')) {
-                const [prefix, ...rest] = toolName.split('__');
-                return { mcp: prefix, action: rest.join('__') };
-            }
+            const split = (0, tool_name_1.splitPrefixedToolName)(toolName);
+            if (split)
+                return split;
             return { mcp: defaultMcp, action: toolName };
         };
     }
@@ -88,6 +108,7 @@ class ArmorIQLangChain {
             goal: opts?.goal,
             onEvent: opts?.onEvent,
             parser: await this.toolNameParser(),
+            session: opts?.session,
         });
         return buildCallback(enforcer);
     }
@@ -96,12 +117,36 @@ exports.ArmorIQLangChain = ArmorIQLangChain;
 /** Framework-agnostic enforcement core (no langchain import), so the allow/hold/
  * block logic is unit-testable without the langchain package. */
 class ArmorIQLangChainEnforcer {
+    /** A tool ran; record its name for the inventory (names-only until a model call declares it). */
+    observeTool(name) {
+        (0, tool_registry_1.observeForClient)(this.factory.client, name);
+    }
+    get client() {
+        return this.factory.client;
+    }
+    /**
+     * Fill the inventory from the tool definitions a chat-model call carries.
+     * LangChain never hands the callback the tool objects, but every model call
+     * carries the bound definitions, so dynamic tool sets are seen as they change.
+     */
+    async registerModelTools(tools) {
+        try {
+            const client = this.factory.client;
+            const d = (0, tool_registry_1.declarationsFromOpenAiTools)(tools, await (0, tool_registry_1.knownServerNames)(client), await (0, tool_registry_1.knownToolMap)(client));
+            (0, tool_registry_1.registerForClient)(this.factory.client, d.tools, d.servers);
+        }
+        catch {
+            /* inventory problems never stop an agent */
+        }
+        (0, tool_push_1.noteModelCall)(this.factory.client);
+    }
     factory;
     scope;
     userEmail;
     goal;
     onEvent;
     parser;
+    preboundSession;
     toolCallNames = new Map();
     pendingPlanCapture;
     session;
@@ -113,6 +158,8 @@ class ArmorIQLangChainEnforcer {
         this.goal = args.goal;
         this.onEvent = args.onEvent;
         this.parser = args.parser;
+        this.session = args.session;
+        this.preboundSession = args.session;
     }
     emit(kind, payload) {
         if (!this.onEvent)
@@ -135,6 +182,27 @@ class ArmorIQLangChainEnforcer {
         }
         return this.session;
     }
+    /** Lazily creates the request session so native OTel and ArmorIQ spans use
+     * one stable session id, including a direct model-only completion. */
+    telemetrySession() {
+        return this.ensureSession();
+    }
+    hasPreboundSession() {
+        return this.preboundSession !== undefined;
+    }
+    /** Each framework root run owns an independent session and mutable callback
+     * state. Sharing one enforcer would cross-contaminate concurrent requests. */
+    fork() {
+        return new ArmorIQLangChainEnforcer({
+            factory: this.factory,
+            scope: this.scope,
+            userEmail: this.userEmail,
+            goal: this.goal,
+            onEvent: this.onEvent,
+            parser: this.parser,
+            session: this.session,
+        });
+    }
     async capturePlan(toolCalls) {
         if (!toolCalls.length)
             return;
@@ -155,7 +223,7 @@ class ArmorIQLangChainEnforcer {
         }
     }
     /** End and ship the request-owned plan session without closing the shared client. */
-    async close(status = 'ok') {
+    async close(status = 'ok', taskOutcome, content) {
         const session = this.session;
         const pendingPlanCapture = this.pendingPlanCapture;
         this.session = undefined;
@@ -169,8 +237,16 @@ class ArmorIQLangChainEnforcer {
                 // Plan capture failures are handled by the callback; teardown remains best-effort.
             }
         }
-        if (session)
-            await session.close(status);
+        if (session) {
+            // Content is authorized, bounded, and redacted by the request-owned core
+            // session. Keep the legacy call shape when no content/outcome exists.
+            if (content !== undefined)
+                await session.close(status, taskOutcome, content);
+            else if (taskOutcome !== undefined)
+                await session.close(status, taskOutcome);
+            else
+                await session.close(status);
+        }
     }
     async captureFromLlm(toolCalls) {
         this.noteToolCallIds(toolCalls);
@@ -201,6 +277,17 @@ class ArmorIQLangChainEnforcer {
         }
         return runName || undefined;
     }
+    frameworkMcpOperation(toolName, itemOrdinal, toolCallId) {
+        return {
+            category: 'mcp',
+            name: 'mcp.execute',
+            toolType: 'mcp',
+            toolName,
+            callId: toolCallId,
+            mcpServer: this.parser(toolName).mcp,
+            planItemOrdinal: itemOrdinal,
+        };
+    }
     /** A tool we cannot name cannot be enforced: a policy keyed on the real name
      * would silently never match. Refuse it instead of guessing. */
     refuseUnnamedTool(serializedHint) {
@@ -210,6 +297,11 @@ class ArmorIQLangChainEnforcer {
     }
     /** Allow -> return; block/hold-denied -> throw (LangChain stops the tool). */
     async enforceTool(toolName, args) {
+        await this.enforceToolWithDecision(toolName, args);
+    }
+    /** Internal callback path that retains safe policy metadata without changing
+     * the public enforceTool() compatibility contract. */
+    async enforceToolWithDecision(toolName, args) {
         args = args || {};
         try {
             if (!this.planStarted) {
@@ -218,9 +310,9 @@ class ArmorIQLangChainEnforcer {
                 await this.capturePlan([{ name: toolName, args }]);
             }
             const session = this.ensureSession();
-            const decision = await session.check(toolName, args, this.userEmail);
+            const decision = await session.check(toolName, args, this.userEmail, { emitOtel: false });
             if (decision.allowed)
-                return; // explicit allow -> tool runs
+                return decision; // explicit allow -> tool runs
             if (decision.action === 'hold') {
                 this.emit('hold', {
                     tool: toolName,
@@ -234,7 +326,7 @@ class ArmorIQLangChainEnforcer {
                 });
                 if (outcome === 'approved') {
                     this.emit('approved', { tool: toolName, delegationId: decision.delegationId });
-                    return; // approved -> tool runs
+                    return { ...decision, allowed: true, action: 'allow', approvalOutcome: outcome }; // approved -> tool runs
                 }
                 this.emit(outcome, {
                     tool: toolName,
@@ -261,13 +353,15 @@ class ArmorIQLangChainEnforcer {
      * 'completed' once all declared steps have run). Only called after a tool actually
      * ran — enforcement-cancelled tools never reach here. Best-effort: audit failures
      * are logged, never allowed to break the run (mirrors the strands AfterToolCall). */
-    async reportExecution(toolName, args, result, error) {
+    async reportExecution(toolName, args, result, error, frameworkCall) {
         if (!this.session || !this.planStarted)
             return;
         try {
             await this.session.report(toolName, args ?? {}, result, {
                 status: error ? 'error' : 'success',
                 errorMessage: error,
+                emitOtel: false,
+                operation: this.frameworkMcpOperation(toolName, frameworkCall?.itemOrdinal, frameworkCall?.toolCallId),
             });
         }
         catch (exc) {
@@ -290,6 +384,35 @@ function toolCallsFromLLMResult(output) {
     }
     return calls;
 }
+/** Extract only provider-reported aggregate usage; prompt/content is never
+ * copied into telemetry attributes. LangChain provider adapters use several
+ * casing variants, so accept their stable aliases. */
+function usageFromLLMResult(output) {
+    const usage = output?.llmOutput?.tokenUsage ??
+        output?.llmOutput?.usage ??
+        output?.generations?.[0]?.[0]?.message?.usage_metadata ??
+        {};
+    const input = usage?.promptTokens ?? usage?.prompt_tokens ?? usage?.inputTokens ?? usage?.input_tokens;
+    const outputTokens = usage?.completionTokens ??
+        usage?.completion_tokens ??
+        usage?.outputTokens ??
+        usage?.output_tokens;
+    return {
+        ...(typeof input === 'number' && Number.isFinite(input)
+            ? { 'gen_ai.usage.input_tokens': input }
+            : {}),
+        ...(typeof outputTokens === 'number' && Number.isFinite(outputTokens)
+            ? { 'gen_ai.usage.output_tokens': outputTokens }
+            : {}),
+    };
+}
+/** Preserve framework payload shape. The core owns redaction, limits, and
+ * capture-mode authorization; adapters must not pre-redact or stringify away
+ * structured content before the policy gate sees it. */
+function outputFromLLMResult(output) {
+    const first = output?.generations?.[0]?.[0];
+    return first?.message?.content ?? first?.text ?? output?.output;
+}
 /** Best-effort name from LangChain's Serialized payload. `id` is deliberately
  * ignored: its last element is the tool CLASS name, which looks like a tool
  * name but never is one. */
@@ -301,6 +424,29 @@ function serializedHint(tool) {
         return tool.id[tool.id.length - 1] ?? 'unknown';
     return String(tool?.id ?? 'unknown');
 }
+/**
+ * Real model identifier (e.g. "gpt-4o-mini") from LangChain's chat-model-start
+ * or LLM-start callback params, falling back to the run name. `runName` is
+ * the model class's display name (e.g. "ChatOpenAI"), not a model id, so
+ * passing it straight to `OtelSession.beginModel` as the model would leave
+ * `gen_ai.system` undetected -- `modelSystem()` only recognizes known
+ * model-id prefixes. Both callbacks carry the same `invocation_params` shape
+ * in their `extra` argument.
+ *
+ * Some providers (e.g. `ChatGoogleGenerativeAI`) never put a model id in
+ * `invocation_params` at all -- they only report it through `getLsParams()`,
+ * which @langchain/core's `CallbackManager` folds into its inheritable
+ * metadata and passes as the callback's `metadata` argument (LangSmith's
+ * `ls_model_name` field). Fall back to that before giving up on the run name.
+ */
+function modelNameFromInvocationParams(extra, metadata, fallback) {
+    const invocationParams = extra
+        ?.invocation_params;
+    const model = invocationParams?.model ??
+        invocationParams?.model_name ??
+        metadata?.ls_model_name;
+    return typeof model === 'string' && model ? model : fallback;
+}
 function coerceArgs(input) {
     if (input && typeof input === 'object')
         return input;
@@ -315,6 +461,151 @@ function coerceArgs(input) {
     }
     return {};
 }
+/**
+ * Adapter state over the request-owned OTel session. The session owns the
+ * provider, root span, bounded flush, and terminal shutdown. This class owns
+ * only LangChain's run-id correlation, so separate handlers never share state.
+ */
+class LangChainTelemetry {
+    session;
+    runs = new Map();
+    toolOrdinals = new Map();
+    nextToolOrdinal = 0;
+    constructor(session) {
+        this.session = session;
+    }
+    noteToolCalls(calls) {
+        for (const call of calls) {
+            const ordinal = this.nextToolOrdinal++;
+            if (call.id)
+                this.toolOrdinals.set(call.id, ordinal);
+        }
+    }
+    toolOrdinal(toolCallId) {
+        return toolCallId ? this.toolOrdinals.get(toolCallId) : undefined;
+    }
+    otel() {
+        return this.session().otelSession;
+    }
+    async beginRoot(input) {
+        try {
+            await this.otel()?.beginRoot(input === undefined ? undefined : { input });
+        }
+        catch {
+            // Native telemetry must never change a LangChain run result.
+        }
+    }
+    async beginModel(runId, model = 'langchain', input) {
+        if (!runId || this.runs.has(runId))
+            return;
+        try {
+            // Direct LLM callbacks have no chain start. Ensure their root receives
+            // the same prompt instead of creating a content-less root in beginModel.
+            await this.beginRoot(input);
+            const handle = await this.otel()?.beginModel(model, input === undefined ? undefined : { input });
+            if (handle)
+                this.runs.set(runId, { kind: 'model', handle });
+        }
+        catch {
+            // Native telemetry must never change a LangChain run result.
+        }
+    }
+    async beginPolicy(runId, toolName, itemOrdinal, toolCallId, args) {
+        if (!runId || this.runs.has(runId))
+            return;
+        try {
+            const handle = await this.otel()?.beginPolicy({
+                toolName,
+                itemOrdinal,
+                toolCallId,
+                ...(args === undefined ? {} : { arguments: args }),
+            });
+            if (handle)
+                this.runs.set(runId, { kind: 'policy', handle });
+        }
+        catch {
+            // Native telemetry must never change a LangChain run result.
+        }
+    }
+    async beginTool(runId, toolName, itemOrdinal, toolCallId, args, operation) {
+        if (!runId || this.runs.has(runId))
+            return;
+        try {
+            const handle = await this.otel()?.beginTool({
+                toolName,
+                itemOrdinal,
+                toolCallId,
+                operation,
+                ...(args === undefined ? {} : { arguments: args }),
+            });
+            if (handle)
+                this.runs.set(runId, { kind: 'tool', handle });
+        }
+        catch {
+            // Native telemetry must never change a LangChain run result.
+        }
+    }
+    async end(runId, outcome = 'ok', modelValues = {}, payload = {}) {
+        if (!runId)
+            return;
+        const run = this.runs.get(runId);
+        if (!run)
+            return;
+        this.runs.delete(runId);
+        try {
+            if (run.kind === 'model') {
+                await this.otel()?.endModel(run.handle, modelValues, outcome === 'ok' ? undefined : new Error('redacted'), outcome, payload.output === undefined ? undefined : { output: payload.output });
+            }
+            else if (run.kind === 'policy') {
+                await this.otel()?.endPolicy(run.handle, {
+                    ...(outcome === 'ok'
+                        ? { decision: 'allow' }
+                        : outcome === 'denied'
+                            ? { decision: 'block' }
+                            : {}),
+                    ...(outcome === 'ok' ? {} : { error: new Error('redacted') }),
+                    terminalStatus: outcome,
+                    ...payload.policy,
+                });
+            }
+            else {
+                await this.otel()?.endTool(run.handle, {
+                    outcome: outcome === 'ok'
+                        ? 'success'
+                        : outcome === 'timeout' || outcome === 'disconnected' || outcome === 'cancelled'
+                            ? outcome
+                            : 'error',
+                    ...(outcome === 'ok' ? {} : { error: new Error('redacted') }),
+                    ...(payload.toolResult === undefined ? {} : { result: payload.toolResult }),
+                });
+            }
+        }
+        catch {
+            // Native telemetry must never change a LangChain run result.
+        }
+    }
+    async endAll(outcome = 'ok') {
+        await Promise.all([...this.runs.keys()].map((runId) => this.end(runId, outcome)));
+    }
+}
+/** Translate the framework's terminal error vocabulary without treating an
+ * expected abort/disconnect as an ordinary model failure. */
+function terminalStatusFromError(error) {
+    const value = error;
+    const name = String(value?.name ?? 'Error').toLowerCase();
+    const code = String(value?.code ?? '').toLowerCase();
+    const message = String(value?.message ?? '').toLowerCase();
+    if (name === 'aborterror' || code === 'abort_err' || code === 'err_abort')
+        return 'cancelled';
+    if (name.includes('timeout') || code === 'etimedout' || message.includes('timed out'))
+        return 'timeout';
+    if (code === 'econnreset' ||
+        code === 'err_stream_premature_close' ||
+        message.includes('disconnect') ||
+        message.includes('connection reset'))
+        return 'disconnected';
+    return 'error';
+}
 function buildCallback(enforcer) {
     let BaseCallbackHandler;
     try {
@@ -329,44 +620,248 @@ function buildCallback(enforcer) {
         // Propagate our block/hold throws so LangChain actually stops the tool.
         raiseError = true;
         awaitHandlers = true;
-        // runId -> {name, args} for tools that passed enforcement and ran, so
-        // handleToolEnd/handleToolError can report the execution (close the plan).
-        pending = {};
-        async handleLLMEnd(output) {
+        requests = new Map();
+        runRoots = new Map();
+        toolCallRoots = new Map();
+        retiredRuns = new Set();
+        preboundRootRunId;
+        handlerClosed = false;
+        createRequest(rootRunId) {
+            if (enforcer.hasPreboundSession() &&
+                this.preboundRootRunId !== undefined &&
+                this.preboundRootRunId !== rootRunId) {
+                throw new exceptions_1.PolicyBlockedException('ArmorIQ enforcement error (fail-closed): a prebound session supports one root run only');
+            }
+            if (enforcer.hasPreboundSession())
+                this.preboundRootRunId = rootRunId;
+            const requestEnforcer = enforcer.fork();
+            const request = {
+                enforcer: requestEnforcer,
+                telemetry: new LangChainTelemetry(() => requestEnforcer.telemetrySession()),
+                runs: new Set([rootRunId]),
+                pending: new Map(),
+            };
+            this.requests.set(rootRunId, request);
+            this.runRoots.set(rootRunId, rootRunId);
+            return request;
+        }
+        toolCallKey(rootRunId, toolCallId) {
+            return `${rootRunId}\u0000${toolCallId}`;
+        }
+        requestFor(runId, parentRunId) {
+            if (!runId ||
+                this.handlerClosed ||
+                this.retiredRuns.has(runId) ||
+                (parentRunId !== undefined && this.retiredRuns.has(parentRunId)))
+                return undefined;
+            const rootRunId = this.runRoots.get(runId) ?? (parentRunId ? this.runRoots.get(parentRunId) : undefined);
+            const request = rootRunId ? this.requests.get(rootRunId) : this.createRequest(runId);
+            if (!request)
+                return undefined;
+            request.runs.add(runId);
+            this.runRoots.set(runId, rootRunId ?? runId);
+            return request;
+        }
+        requestForTool(runId, parentRunId, toolCallId) {
+            if (toolCallId) {
+                const rootRunId = (parentRunId ? this.runRoots.get(parentRunId) : undefined) ??
+                    (runId ? this.runRoots.get(runId) : undefined);
+                let mappedRoot = rootRunId
+                    ? this.toolCallRoots.get(this.toolCallKey(rootRunId, toolCallId))
+                    : undefined;
+                // Older/direct tool invocations may omit parentRunId. Reuse a model
+                // plan only when that id is unambiguous across live roots; otherwise
+                // fail closed rather than cross-routing a duplicate provider id.
+                if (!mappedRoot && !rootRunId) {
+                    const matches = [...this.toolCallRoots]
+                        .filter(([key]) => key.endsWith(`\u0000${toolCallId}`))
+                        .map(([, root]) => root);
+                    if (matches.length === 1)
+                        mappedRoot = matches[0];
+                }
+                const request = mappedRoot ? this.requests.get(mappedRoot) : undefined;
+                if (request && runId && !this.retiredRuns.has(runId)) {
+                    request.runs.add(runId);
+                    this.runRoots.set(runId, mappedRoot);
+                    return request;
+                }
+            }
+            return this.requestFor(runId, parentRunId);
+        }
+        async closeRequest(rootRunId, status = 'ok', output) {
+            const request = this.requests.get(rootRunId);
+            if (!request)
+                return;
+            // Retire before awaiting teardown. Late callbacks must never create a
+            // fresh session while the old request drains its final spans.
+            this.requests.delete(rootRunId);
+            for (const runId of request.runs) {
+                this.runRoots.delete(runId);
+                this.retiredRuns.add(runId);
+            }
+            for (const [toolCallId, mappedRoot] of this.toolCallRoots) {
+                if (mappedRoot === rootRunId)
+                    this.toolCallRoots.delete(toolCallId);
+            }
+            request.pending.clear();
+            await request.telemetry.endAll(status);
+            await request.enforcer.close(status, undefined, output === undefined ? undefined : { output });
+        }
+        async closeRoot(runId, status = 'ok', output) {
+            if (!runId)
+                return;
+            const rootRunId = this.runRoots.get(runId);
+            if (rootRunId)
+                await this.closeRequest(rootRunId, status, output);
+        }
+        async handleChainStart(_chain, _inputs, runId, parentRunId, _tags, _metadata, _runType, runName) {
+            void _chain;
+            void runName;
+            const request = this.requestFor(runId, parentRunId);
+            if (request && !parentRunId)
+                await request.telemetry.beginRoot(_inputs);
+        }
+        async handleChainEnd(_output, runId) {
+            if (runId && this.runRoots.get(runId) === runId)
+                await this.closeRoot(runId, 'ok', _output);
+        }
+        async handleChainError(_error, runId) {
+            if (runId && this.runRoots.get(runId) === runId) {
+                await this.closeRoot(runId, terminalStatusFromError(_error));
+            }
+        }
+        async handleChatModelStart(_model, _messages, runId, parentRunId, extra, _tags, metadata, runName) {
+            const request = this.requestFor(runId, parentRunId);
+            if (!request)
+                return;
+            const tools = extra?.invocation_params?.tools;
+            if (Array.isArray(tools) && tools.length)
+                await request.enforcer.registerModelTools(tools);
+            else
+                (0, tool_push_1.noteModelCall)(request.enforcer.client);
+            const model = modelNameFromInvocationParams(extra, metadata, runName ?? 'langchain');
+            await request.telemetry.beginModel(runId, model, _messages);
+        }
+        /**
+         * Non-chat LLMs report an array of prompts instead of message objects.
+         * `extra` must stay in this exact slot -- @langchain/core's callback
+         * manager invokes `handleLLMStart` with 8 positional args (llm, prompts,
+         * runId, parentRunId, extraParams, tags, metadata, runName); omitting
+         * `extra` here previously shifted every argument after it by one, so
+         * `runName` was silently receiving the `metadata` object instead of the
+         * run name string.
+         */
+        async handleLLMStart(_llm, prompts, runId, parentRunId, extra, _tags, metadata, runName) {
+            const request = this.requestFor(runId, parentRunId);
+            if (request) {
+                const model = modelNameFromInvocationParams(extra, metadata, runName ?? 'langchain');
+                await request.telemetry.beginModel(runId, model, prompts);
+            }
+        }
+        async handleLLMEnd(output, runId, parentRunId) {
             const calls = toolCallsFromLLMResult(output);
+            const request = this.requestFor(runId, parentRunId);
+            if (!request)
+                return;
+            await request.telemetry.beginModel(runId);
+            request.telemetry.noteToolCalls(calls);
+            const rootRunId = runId ? this.runRoots.get(runId) : undefined;
+            for (const call of calls) {
+                if (call.id && rootRunId) {
+                    this.toolCallRoots.set(this.toolCallKey(rootRunId, call.id), rootRunId);
+                }
+            }
             if (calls.length)
-                await enforcer.captureFromLlm(calls);
+                await request.enforcer.captureFromLlm(calls);
+            // Direct model invocations do not always have a chain callback, but the
+            // model completion itself is a complete OTel lifecycle.
+            await request.telemetry.end(runId, 'ok', usageFromLLMResult(output), {
+                output: outputFromLLMResult(output),
+            });
+            if (!parentRunId && calls.length === 0) {
+                await this.closeRoot(runId, 'ok', outputFromLLMResult(output));
+            }
+        }
+        async handleLLMError(_error, runId, parentRunId) {
+            const request = this.requestFor(runId, parentRunId);
+            if (!request)
+                return;
+            const status = terminalStatusFromError(_error);
+            await request.telemetry.beginModel(runId);
+            await request.telemetry.end(runId, status);
+            if (!parentRunId)
+                await this.closeRoot(runId, status);
         }
-        async handleToolStart(tool, input, runId, _parentRunId, _tags, _metadata, runName, toolCallId) {
-            const toolName = enforcer.resolveToolName(toolCallId, runName) ??
+        async handleToolStart(tool, input, runId, parentRunId, _tags, _metadata, runName, toolCallId) {
+            const request = this.requestForTool(runId, parentRunId, toolCallId);
+            if (!request) {
+                throw new exceptions_1.PolicyBlockedException('ArmorIQ enforcement error (fail-closed): callback request is closed');
+            }
+            const toolName = request.enforcer.resolveToolName(toolCallId, runName) ??
                 nameFromSerialized(tool) ??
-                enforcer.refuseUnnamedTool(serializedHint(tool));
+                request.enforcer.refuseUnnamedTool(serializedHint(tool));
+            request.enforcer.observeTool(toolName);
             const args = coerceArgs(input);
-            await enforcer.enforceTool(toolName, args); // throws to block/hold-deny
+            await request.telemetry.beginPolicy(runId, toolName, request.telemetry.toolOrdinal(toolCallId), toolCallId, args);
+            try {
+                const decision = await request.enforcer.enforceToolWithDecision(toolName, args); // throws to block/hold-deny
+                await request.telemetry.end(runId, 'ok', {}, {
+                    policy: {
+                        policyName: decision?.matchedPolicy,
+                        ...(decision?.policyId ? { policyId: decision.policyId } : {}),
+                        policyVersion: decision?.policyVersion,
+                        policySource: decision?.policySource,
+                        policyReasonCode: decision?.reason,
+                        defaultAction: decision?.defaultAction,
+                        matchedRuleId: decision?.matchedRuleId,
+                    },
+                });
+            }
+            catch (error) {
+                await request.telemetry.end(runId, error instanceof exceptions_1.PolicyBlockedException || error instanceof exceptions_1.PolicyHoldException
+                    ? 'denied'
+                    : 'error', {}, { policy: { policyReasonCode: 'enforcement_error' } });
+                throw error;
+            }
+            await request.telemetry.beginTool(runId, toolName, request.telemetry.toolOrdinal(toolCallId), toolCallId, args, request.enforcer.frameworkMcpOperation(toolName, request.telemetry.toolOrdinal(toolCallId), toolCallId));
             // Only reached when allowed -> the tool will run; remember it so
             // handleToolEnd can report the execution and close the plan lifecycle.
-            if (runId)
-                this.pending[runId] = { name: toolName, args };
+            if (runId) {
+                request.pending.set(runId, {
+                    name: toolName,
+                    args,
+                    itemOrdinal: request.telemetry.toolOrdinal(toolCallId),
+                    toolCallId,
+                });
+            }
         }
         async handleToolEnd(output, runId) {
-            const info = runId ? this.pending[runId] : undefined;
+            const rootRunId = runId ? this.runRoots.get(runId) : undefined;
+            const request = rootRunId ? this.requests.get(rootRunId) : undefined;
+            const info = runId ? request?.pending.get(runId) : undefined;
             if (info) {
-                delete this.pending[runId];
-                await enforcer.reportExecution(info.name, info.args, output);
+                request.pending.delete(runId);
+                await request.enforcer.reportExecution(info.name, info.args, output, undefined, info);
             }
+            await request?.telemetry.end(runId, 'ok', {}, { toolResult: output });
         }
         /** LangChain has no end-of-request hook, so the host finalizes the session here. */
         async close(status = 'ok') {
-            this.pending = {};
-            await enforcer.close(status);
+            if (this.handlerClosed)
+                return;
+            this.handlerClosed = true;
+            await Promise.all([...this.requests.keys()].map((rootRunId) => this.closeRequest(rootRunId, status)));
         }
         async handleToolError(err, runId) {
-            const info = runId ? this.pending[runId] : undefined;
+            const rootRunId = runId ? this.runRoots.get(runId) : undefined;
+            const request = rootRunId ? this.requests.get(rootRunId) : undefined;
+            const info = runId ? request?.pending.get(runId) : undefined;
             if (info) {
-                delete this.pending[runId];
+                request.pending.delete(runId);
                 const msg = err?.message ?? String(err);
-                await enforcer.reportExecution(info.name, info.args, msg, msg);
+                await request.enforcer.reportExecution(info.name, info.args, msg, msg, info);
             }
+            await request?.telemetry.end(runId, terminalStatusFromError(err));
         }
     }
     return new ArmorIQLangChainCallback();