@armoriq/sdk-dev 0.6.10 → 0.8.0

package/dist/integrations/google_adk.js

@@ -31,7 +31,27 @@
  */
 Object.defineProperty(exports, "__esModule", { value: true });
 exports.ArmorIQADKBundle = exports.ArmorIQADK = void 0;
+const tool_push_1 = require("../tool_push");
+const tool_registry_1 = require("../tool_registry");
 const client_1 = require("../client");
+const tool_name_1 = require("../tool_name");
+let warnedFunctionCallIdFallback = false;
+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.includes('timeout') || code === 'etimedout' || message.includes('timed out'))
+        return 'timeout';
+    if (name === 'aborterror' || code === 'abort_err' || code === 'err_abort')
+        return 'cancelled';
+    if (code === 'econnreset' ||
+        code === 'err_stream_premature_close' ||
+        message.includes('disconnect') ||
+        message.includes('connection reset'))
+        return 'disconnected';
+    return 'error';
+}
 /**
  * Process-wide ArmorIQ factory for ADK-style agents.
  */
@@ -106,17 +126,23 @@ class ArmorIQADK {
     async toolNameParser() {
         if (this.customParser)
             return this.customParser;
-        const bootstrap = await this.bootstrap();
-        const toolMap = bootstrap.toolMap ?? {};
+        let toolMap = {};
+        try {
+            toolMap = (await this.bootstrap()).toolMap ?? {};
+        }
+        catch {
+            // Bootstrap only enriches routing. Model-only ADK work remains usable;
+            // any later real tool still goes through fail-closed session enforcement.
+            console.warn('[armoriq] google-adk bootstrap unavailable; using safe tool-name fallback');
+        }
         const defaultMcp = this.defaultMcpName;
         return (toolName) => {
             const mcp = toolMap[toolName];
             if (mcp)
                 return { mcp, action: toolName };
-            if (toolName.includes('__')) {
-                const idx = toolName.indexOf('__');
-                return { mcp: toolName.slice(0, idx), action: toolName.slice(idx + 2) };
-            }
+            const split = (0, tool_name_1.splitPrefixedToolName)(toolName);
+            if (split)
+                return split;
             if (defaultMcp)
                 return { mcp: defaultMcp, action: toolName };
             return { mcp: 'unknown', action: toolName };
@@ -126,7 +152,6 @@ class ArmorIQADK {
         this.client.invalidateUser(userEmail);
     }
     async forUser(userEmail, opts) {
-        await this.bootstrap();
         const scope = this.client.forUser(userEmail);
         const parser = await this.toolNameParser();
         return new ArmorIQADKBundle({
@@ -140,6 +165,61 @@ class ArmorIQADK {
     }
 }
 exports.ArmorIQADK = ArmorIQADK;
+/**
+ * Owner-keyed subscriber registry for watchForLateParent(): more than one
+ * bundle can independently resolve the SAME shared owner as their
+ * `invocationOwner` — e.g. two sibling agents nested inside one inner
+ * SequentialAgent, each with its own bundle installed. Object.defineProperty
+ * only supports a single descriptor per property, so each bundle calling
+ * watchForLateParent() with its own private accessor would let the second
+ * bundle's install() silently clobber the first bundle's getter/setter — and
+ * that first bundle's finalizer would then never migrate when the shared
+ * owner is later composed into an outer workflow (it would keep closing at
+ * the inner boundary). Keep exactly one accessor per owner and fan its
+ * setter out to every subscribed bundle instead of just whichever bundle
+ * happened to call watchForLateParent() last.
+ */
+const lateParentSubscribers = new WeakMap();
+const runAsyncRegistrations = new WeakMap();
+/**
+ * The four ADK lifecycle callback slots this SDK writes.
+ */
+const ADK_CALLBACK_SLOTS = [
+    'beforeModelCallback',
+    'afterModelCallback',
+    'beforeToolCallback',
+    'afterToolCallback',
+];
+const callbackRegistrations = new WeakMap();
+/**
+ * Adopt as the new original any slot the application has taken over since our
+ * last install, so the next bundle chains behind the application's current
+ * callback instead of overwriting it with a stale snapshot.
+ */
+function recaptureReplacedSlots(agent, registration) {
+    const installed = registration.installed;
+    if (!installed)
+        return;
+    for (const slot of ADK_CALLBACK_SLOTS) {
+        if (agent[slot] === installed[slot] || agent[slot] === registration.original[slot])
+            continue;
+        registration.original[slot] = agent[slot];
+        console.warn(`[armoriq] google-adk ${slot} changed after ArmorIQ installed; chaining behind the ` +
+            'current callback. Replacing a slot that still holds ArmorIQ callbacks keeps them ' +
+            'reachable through your own wrapper.');
+    }
+}
+/** Returns false if any slot was replaced by the application and left alone. */
+function restoreCallbackSlots(agent, registration) {
+    let unwoundEverySlot = true;
+    for (const slot of ADK_CALLBACK_SLOTS) {
+        if (agent[slot] === registration.installed?.[slot])
+            agent[slot] = registration.original[slot];
+        else
+            unwoundEverySlot = false;
+    }
+    return unwoundEverySlot;
+}
 /**
  * Per-request ADK bundle — installs/uninstalls lifecycle callbacks on
  * one agent and binds them to one user's session.
@@ -153,11 +233,51 @@ class ArmorIQADKBundle {
     parser;
     planMinted = false;
     pendingPlanCapture;
-    blockedTools = new Set();
-    blockedActions = new Map();
+    /** Terminal enforcement state keyed by ADK's native functionCallId. */
+    blockedCalls = new Map();
+    /**
+     * Legacy duck-typed callers may omit functionCallId. Keep a FIFO fallback
+     * for those callers; real ADK 0.6.1 always supplies the native id.
+     */
+    anonymousBlockedCalls = new Map();
+    nextAnonymousCallOrdinal = 0;
+    modelSpans = new Map();
+    toolSpans = new Map();
+    /** Model output is scoped by ADK's native invocation key. A scalar leaks a
+     * concurrent invocation's result into another span/root. */
+    modelOutputs = new Map();
+    /**
+     * Root/session output per invocation. Defaults to the raw model output
+     * (mirrors modelOutputs) but the PAP rejection branch overwrites the entry
+     * — never modelOutputs — with the safe response ADK actually returns, so
+     * root/session telemetry exports what the caller really got back instead
+     * of the blocked content.
+     */
+    rootOutputs = new Map();
+    toolOrdinals = new Map();
+    nextToolOrdinal = 0;
+    nextInvocationOrdinal = 0;
+    closed = false;
+    closePromise;
     onEvent;
     agent;
-    saved = {};
+    /**
+     * The agent whose runAsync actually got the close-on-completion wrapper.
+     * Usually the installed agent itself; when the installed agent is a child
+     * of a composed workflow (see install()) this is the outermost ancestor,
+     * so uninstall() must restore runAsync on THIS object, not on the
+     * installed agent.
+     */
+    runAsyncOwner;
+    /**
+     * Set while `runAsyncOwner` is only a provisional owner (the installed
+     * agent had no `parentAgent` yet at install() time). Holds the agent whose
+     * `parentAgent` property we replaced with an accessor so we notice if ADK
+     * composes it into a parent *after* install() — see watchForLateParent().
+     */
+    parentWatchTarget;
+    /** The agent whose callback slots this bundle last wrote. */
+    callbackAgent;
     constructor(args) {
         this.factory = args.factory;
         this.scope = args.scope;
@@ -213,12 +333,34 @@ class ArmorIQADKBundle {
     }
     /** End and ship the request-owned plan session without closing the shared client. */
     async close(status = 'ok') {
+        if (this.closePromise)
+            return this.closePromise;
+        this.closePromise = this.closeInternal(status);
+        return this.closePromise;
+    }
+    async closeInternal(status) {
+        this.closed = true;
         const session = this.session;
         const pendingPlanCapture = this.pendingPlanCapture;
+        // A bundle normally owns one ADK invocation. If a host overlaps requests
+        // on one bundle, omit the ambiguous root output rather than export the
+        // wrong request's response; individual model spans remain complete.
+        const rootOutput = this.rootOutputs.size === 1
+            ? this.rootOutputs.values().next().value
+            : undefined;
         this.session = undefined;
         this.planMinted = false;
-        this.blockedTools.clear();
-        this.blockedActions.clear();
+        this.blockedCalls.clear();
+        this.anonymousBlockedCalls.clear();
+        const error = status === 'ok' ? undefined : new Error('terminal lifecycle');
+        const unfinishedToolOutcome = status === 'error' ? 'error' : status === 'timeout' ? 'timeout' :
+            status === 'disconnected' ? 'disconnected' : 'cancelled';
+        await Promise.all([
+            ...[...this.modelSpans.values()].filter(Boolean).map((span) => session?.otelSession?.endModel(span, {}, error, status)),
+            ...[...this.toolSpans.values()].filter(Boolean).map((span) => session?.otelSession?.endTool(span, { outcome: unfinishedToolOutcome, error })),
+        ]);
+        this.toolSpans.clear();
+        this.modelSpans.clear();
         if (session && pendingPlanCapture && pendingPlanCapture.session === session) {
             try {
                 await pendingPlanCapture.promise;
@@ -227,10 +369,25 @@ class ArmorIQADKBundle {
                 // afterModel already logs capture failures; teardown remains best-effort.
             }
         }
-        if (session)
-            await session.close(status);
+        if (session) {
+            if (rootOutput === undefined)
+                await session.close(status);
+            else
+                await session.close(status, undefined, { output: rootOutput });
+        }
+        this.modelOutputs.clear();
+        this.rootOutputs.clear();
+        // Nothing calls uninstall() on a real ADK server, so without this the
+        // owner retains every closed bundle. The agent's callback slots stay
+        // installed on purpose; see releaseFromOwners().
+        this.releaseFromOwners();
     }
     async afterModel(...adkArgs) {
+        if (this.closed)
+            return null;
+        const modelKey = this.modelKey(adkArgs[0]);
+        let modelError;
+        let modelValues = {};
         try {
             // Re-mint every model turn (mirrors strands / langchain). A multi-turn run
             // chooses new tools on each turn; minting only the first turn's tools left
@@ -250,10 +407,29 @@ class ArmorIQADKBundle {
                 llmResponse = adkArgs[1] ?? adkArgs[0];
             }
             const parts = llmResponse?.content?.parts ?? [];
+            this.toolOrdinals.clear();
+            this.nextToolOrdinal = 0;
+            modelValues = {
+                ...(llmResponse?.usageMetadata?.promptTokenCount !== undefined
+                    ? { 'gen_ai.usage.input_tokens': llmResponse.usageMetadata.promptTokenCount }
+                    : {}),
+                ...(llmResponse?.usageMetadata?.candidatesTokenCount !== undefined
+                    ? { 'gen_ai.usage.output_tokens': llmResponse.usageMetadata.candidatesTokenCount }
+                    : {}),
+            };
+            const modelOutput = llmResponse?.content ?? llmResponse?.output;
+            if (modelOutput !== undefined) {
+                this.modelOutputs.set(modelKey, modelOutput);
+                // Defaults the root/export output to the raw model response; the PAP
+                // branch below overwrites this entry if it replaces what ADK returns.
+                this.rootOutputs.set(modelKey, modelOutput);
+            }
             const toolCalls = [];
             for (const p of parts) {
                 const fc = p?.functionCall ?? p?.function_call;
                 if (fc?.name) {
+                    if (typeof fc.id === 'string' && fc.id)
+                        this.toolOrdinals.set(fc.id, this.nextToolOrdinal++);
                     toolCalls.push({ name: fc.name, args: fc.args ? { ...fc.args } : {} });
                 }
             }
@@ -299,12 +475,17 @@ class ArmorIQADKBundle {
                         console.warn(`[armoriq] PAP rejected plan user=${this.userEmail} ` +
                             `violations=[${violations}] ` +
                             `predicate_fails=${refineResult.predicateFails.length}`);
-                        return buildRejectionResponse(violations, {
+                        const rejection = buildRejectionResponse(violations, {
                             decision: refineResult.decision,
                             violations: refineResult.violations,
                             predicateFails: refineResult.predicateFails,
                             intentId: refineResult.intentId,
                         });
+                        // ADK actually hands this response back to the caller, not the
+                        // original (blocked) model output captured above — root/session
+                        // telemetry must reflect what was really returned.
+                        this.rootOutputs.set(modelKey, rejection.content);
+                        return rejection;
                     }
                     console.info(`[armoriq] PAP refine ok user=${this.userEmail} intent=${refineResult.intentId}`);
                 }
@@ -312,11 +493,13 @@ class ArmorIQADKBundle {
                     const msg = refineErr.message;
                     console.warn(`[armoriq] PAP refine failed (fail-${client.papFailMode}): ${msg}`);
                     if (client.papFailMode === 'closed') {
-                        return buildRejectionResponse(`pap_unavailable: ${msg}`, {
+                        const rejection = buildRejectionResponse(`pap_unavailable: ${msg}`, {
                             decision: 'rejected',
                             violations: ['pap_unavailable'],
                             predicateFails: [],
                         });
+                        this.rootOutputs.set(modelKey, rejection.content);
+                        return rejection;
                     }
                     // fail-open: fall through, continue with original tool calls.
                 }
@@ -326,14 +509,141 @@ class ArmorIQADKBundle {
             console.info(`[armoriq] plan minted user=${this.userEmail} tools=${toolCalls.length}`);
         }
         catch (exc) {
+            modelError = exc;
             console.warn(`[armoriq] afterModelCallback failed: ${exc.message}`);
         }
+        finally {
+            const span = this.modelSpans.get(modelKey);
+            this.modelSpans.delete(modelKey);
+            const modelOutput = this.modelOutputs.get(modelKey);
+            await this.session?.otelSession?.endModel(span ?? { span: null, name: 'gen_ai.chat' }, modelValues, modelError, undefined, modelOutput === undefined ? undefined : { output: modelOutput });
+        }
         return null;
     }
-    async beforeTool(tool, args, _toolContext) {
+    /** ADK resolves toolsets on every model step, so a changed tool set is visible here. */
+    refreshDeclarations() {
+        if (!this.agent)
+            return;
+        try {
+            const d = (0, tool_registry_1.declarationsFromAdkAgent)(this.agent);
+            (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);
+    }
+    async beforeModel(...adkArgs) {
+        if (this.closed)
+            return undefined;
+        this.refreshDeclarations();
+        const params = adkArgs[0] ?? {};
+        const model = params?.request?.model ?? params?.model ?? this.factory.llm;
+        const input = params?.request?.contents ?? params?.request?.content ?? params?.input ?? this.goal;
+        const key = this.modelKey(params);
+        if (!this.modelSpans.has(key)) {
+            const otel = this.ensureSession().otelSession;
+            if (typeof otel?.beginRoot === 'function')
+                await otel.beginRoot({ input: this.goal });
+            const span = typeof otel?.beginModel === 'function'
+                ? await otel.beginModel(String(model), { input })
+                : undefined;
+            if (span)
+                this.modelSpans.set(key, span);
+        }
+        return undefined;
+    }
+    modelKey(params) {
+        const context = params?.context ?? params?.callbackContext;
+        const invocationId = context?.invocationId ?? context?.sessionId;
+        return typeof invocationId === 'string' && invocationId
+            ? `adk:${invocationId}`
+            : `anonymous:${++this.nextInvocationOrdinal}`;
+    }
+    toolOrdinal(toolContext) {
+        const functionCallId = toolContext?.functionCallId;
+        return typeof functionCallId === 'string' ? this.toolOrdinals.get(functionCallId) : undefined;
+    }
+    frameworkMcpOperation(toolName, itemOrdinal, toolCallId) {
+        return {
+            category: 'mcp',
+            name: 'mcp.execute',
+            toolType: 'mcp',
+            toolName,
+            callId: toolCallId,
+            mcpServer: this.parser(toolName).mcp,
+            planItemOrdinal: itemOrdinal,
+        };
+    }
+    anonymousCallFingerprint(toolName, args) {
+        try {
+            return `${toolName}:${JSON.stringify(args ?? {})}`;
+        }
+        catch {
+            return `${toolName}:[unserializable]`;
+        }
+    }
+    /**
+     * The fallback is only for legacy duck-typed callers. Keep the diagnostic
+     * fixed and once per process so it cannot leak tool data or flood host logs.
+     */
+    warnFunctionCallIdFallback() {
+        if (warnedFunctionCallIdFallback)
+            return;
+        warnedFunctionCallIdFallback = true;
+        console.warn('[armoriq] google-adk functionCallId missing; using bounded compatibility fallback correlation');
+    }
+    callKeyForBefore(toolName, args, toolContext) {
+        const functionCallId = toolContext?.functionCallId;
+        if (typeof functionCallId === 'string' && functionCallId)
+            return `adk:${functionCallId}`;
+        this.warnFunctionCallIdFallback();
+        const fingerprint = this.anonymousCallFingerprint(toolName, args);
+        const key = `anonymous:${fingerprint}:${++this.nextAnonymousCallOrdinal}`;
+        const queue = this.anonymousBlockedCalls.get(fingerprint) ?? [];
+        queue.push(key);
+        this.anonymousBlockedCalls.set(fingerprint, queue);
+        return key;
+    }
+    callKeyForAfter(toolName, args, toolContext) {
+        const functionCallId = toolContext?.functionCallId;
+        if (typeof functionCallId === 'string' && functionCallId)
+            return `adk:${functionCallId}`;
+        this.warnFunctionCallIdFallback();
+        const fingerprint = this.anonymousCallFingerprint(toolName, args);
+        const queue = this.anonymousBlockedCalls.get(fingerprint);
+        const key = queue?.shift();
+        if (queue && queue.length === 0)
+            this.anonymousBlockedCalls.delete(fingerprint);
+        return key;
+    }
+    rememberBlockedCall(toolName, args, toolContext, action) {
+        this.blockedCalls.set(this.callKeyForBefore(toolName, args, toolContext), action);
+    }
+    async beforeTool(tool, args, toolContext) {
+        if (this.closed) {
+            return { error: 'ArmorIQ request is closed', armoriq_enforcement: { blocked: true, action: 'block' } };
+        }
         const toolName = tool?.name ?? String(tool);
+        const session = this.ensureSession();
+        const itemOrdinal = this.toolOrdinal(toolContext);
+        const policySpan = await session.otelSession?.beginPolicy({
+            toolName, itemOrdinal, toolCallId: toolContext?.functionCallId,
+            arguments: args ?? {},
+        });
         try {
-            const decision = await this.ensureSession().check(toolName, args ?? {}, this.userEmail);
+            const decision = await session.check(toolName, args ?? {}, this.userEmail, { emitOtel: false });
+            const policy = decision;
+            await session.otelSession?.endPolicy(policySpan ?? { span: null, name: 'armoriq.policy.evaluate' }, {
+                decision: decision.allowed ? 'allow' : decision.action === 'hold' ? 'hold' : 'block',
+                policyName: decision.matchedPolicy,
+                ...(policy.policyId ? { policyId: policy.policyId } : {}),
+                policyVersion: policy.policyVersion,
+                policySource: policy.policySource,
+                policyReasonCode: decision.reason,
+                defaultAction: policy.defaultAction,
+                matchedRuleId: policy.matchedRuleId,
+            });
             if (!decision.allowed) {
                 const policy = decision.matchedPolicy ? ` (policy: ${decision.matchedPolicy})` : '';
                 if (decision.action === 'hold') {
@@ -356,11 +666,17 @@ class ArmorIQADKBundle {
                     if (outcome === 'approved') {
                         console.info(`[armoriq] APPROVED ${toolName} user=${this.userEmail}`);
                         this.emit('approved', { tool: toolName, delegationId: decision.delegationId });
+                        this.toolSpans.set(this.callKeyForBefore(toolName, args, toolContext), await session.otelSession?.beginTool({
+                            toolName,
+                            itemOrdinal,
+                            toolCallId: toolContext?.functionCallId,
+                            arguments: args,
+                            operation: this.frameworkMcpOperation(toolName, itemOrdinal, toolContext?.functionCallId),
+                        }));
                         return null; // approved -> the tool runs
                     }
                     console.info(`[armoriq] hold ${outcome} for ${toolName} user=${this.userEmail}`);
-                    this.blockedTools.add(toolName);
-                    this.blockedActions.set(toolName, 'hold');
+                    this.rememberBlockedCall(toolName, args, toolContext, 'hold');
                     this.emit(outcome, {
                         tool: toolName,
                         delegationId: decision.delegationId,
@@ -386,8 +702,7 @@ class ArmorIQADKBundle {
                     reason: decision.reason,
                     matchedPolicy: decision.matchedPolicy,
                 });
-                this.blockedTools.add(toolName);
-                this.blockedActions.set(toolName, decision.action);
+                this.rememberBlockedCall(toolName, args, toolContext, decision.action);
                 console.info(`[armoriq] BLOCKED ${toolName} user=${this.userEmail} action=${decision.action} reason=${decision.reason}`);
                 return {
                     error: `This action is not permitted by your organization's policy${policy}. Reason: ${decision.reason ?? 'policy-blocked'}.`,
@@ -401,16 +716,25 @@ class ArmorIQADKBundle {
                     },
                 };
             }
+            this.toolSpans.set(this.callKeyForBefore(toolName, args, toolContext), await session.otelSession?.beginTool({
+                toolName,
+                itemOrdinal,
+                toolCallId: toolContext?.functionCallId,
+                arguments: args,
+                operation: this.frameworkMcpOperation(toolName, itemOrdinal, toolContext?.functionCallId),
+            }));
         }
         catch (exc) {
+            await session.otelSession?.endPolicy(policySpan ?? { span: null, name: 'armoriq.policy.evaluate' }, {
+                error: exc, policyReasonCode: 'enforcement_error',
+            });
             // Fail closed: ADK reads a null return as "proceed", so an enforcement
             // error has to come back as a payload or the tool runs unchecked. Matches
             // the strands/langchain hooks, which also refuse the tool on error.
             const msg = exc.message ?? String(exc);
             this.emit('error', { tool: toolName, error: msg });
             console.error(`[armoriq] beforeToolCallback failed (fail-closed): ${msg}`);
-            this.blockedTools.add(toolName);
-            this.blockedActions.set(toolName, 'block');
+            this.rememberBlockedCall(toolName, args, toolContext, 'block');
             return {
                 error: `ArmorIQ enforcement error (fail-closed): ${msg}`,
                 armoriq_enforcement: {
@@ -425,41 +749,121 @@ class ArmorIQADKBundle {
         }
         return null;
     }
-    async afterTool(tool, args, _toolContext, toolResponse) {
+    async afterTool(tool, args, toolContext, toolResponse) {
+        if (this.closed)
+            return null;
         const toolName = tool?.name ?? String(tool);
         try {
-            if (this.blockedTools.has(toolName)) {
-                const action = this.blockedActions.get(toolName) ?? 'block';
-                this.blockedActions.delete(toolName);
-                this.blockedTools.delete(toolName);
+            const callKey = this.callKeyForAfter(toolName, args, toolContext);
+            const action = callKey ? this.blockedCalls.get(callKey) : undefined;
+            if (action) {
+                this.blockedCalls.delete(callKey);
                 if (action !== 'hold') {
                     await this.ensureSession().report(toolName, args ?? {}, toolResponse, {
                         status: 'failed',
                         errorMessage: 'Blocked by policy',
+                        operation: this.frameworkMcpOperation(toolName, this.toolOrdinal(toolContext), toolContext?.functionCallId),
                     });
                 }
                 return null;
             }
-            await this.ensureSession().report(toolName, args ?? {}, toolResponse);
+            const toolSpan = callKey ? this.toolSpans.get(callKey) : undefined;
+            if (callKey)
+                this.toolSpans.delete(callKey);
+            await this.ensureSession().otelSession?.endTool(toolSpan ?? { span: null, name: 'armoriq.tool' }, {
+                outcome: toolResponse?.error ? 'error' : 'success',
+                error: toolResponse?.error,
+                result: toolResponse,
+            });
+            await this.ensureSession().report(toolName, args ?? {}, toolResponse, {
+                emitOtel: false,
+                operation: this.frameworkMcpOperation(toolName, this.toolOrdinal(toolContext), toolContext?.functionCallId),
+            });
         }
         catch (exc) {
             console.warn(`[armoriq] afterToolCallback failed: ${exc.message}`);
         }
         return null;
     }
-    /**
-     * Attach the three callbacks to the ADK-style agent. Save originals
-     * for uninstall().
-     */
+    /** Attach this bundle to an ADK-style agent. */
     install(agent) {
         this.agent = agent;
-        this.saved = {
-            afterModelCallback: agent.afterModelCallback,
-            beforeToolCallback: agent.beforeToolCallback,
-            afterToolCallback: agent.afterToolCallback,
-        };
-        agent.afterModelCallback = (...args) => this.afterModel(...args);
-        agent.beforeToolCallback = (...args) => {
+        (0, tool_registry_1.registerForClient)(this.factory.client, ...(() => {
+            const d = (0, tool_registry_1.declarationsFromAdkAgent)(agent);
+            return [d.tools, d.servers];
+        })());
+        this.installCallbacks(agent);
+        // Bind automatic finalization to the true invocation boundary, not to
+        // "this installed agent's own iterator finished". ADK composes agents
+        // (SequentialAgent, ParallelAgent, ...) by running each child's runAsync
+        // to completion before the parent advances to the next one, so wrapping
+        // the installed agent directly would close the whole request — and stop
+        // telemetry for every later sibling — the moment the FIRST child
+        // finishes. Walk ADK's public `parentAgent` link to the outermost
+        // ancestor and bind the close-on-completion wrapper there instead; the
+        // model/tool callbacks above stay on the installed agent, since that is
+        // where ADK actually invokes them.
+        const owner = ArmorIQADKBundle.invocationOwner(agent);
+        this.bindFinalizer(owner);
+        if (!owner.parentAgent) {
+            // The *resolved owner* has no `parentAgent` yet — watch IT, not
+            // necessarily `agent`. That distinction matters when `agent` was
+            // already composed into an intermediate workflow before install():
+            // invocationOwner() then returns that intermediate ancestor (say,
+            // an inner SequentialAgent), which itself may still be composed into
+            // an outer workflow later. Comparing `owner === agent` would miss
+            // this case entirely — the intermediate ancestor is a distinct
+            // object from `agent`, so no watcher would ever get attached to it,
+            // and the bundle would close the moment the intermediate finishes
+            // instead of when the true outermost invocation completes. ADK sets
+            // a later parent with a plain `subAgent.parentAgent = this`
+            // (base_agent.js) strictly before the parent's own runAsync is ever
+            // invoked, so watch for that assignment on `owner` and move the
+            // finalizer to the real outermost owner once it happens.
+            ArmorIQADKBundle.watchForLateParent(owner, this);
+        }
+        return this;
+    }
+    /**
+     * Chained in front of the agent's pre-ArmorIQ callbacks, not in front of
+     * whatever is in the slot now. See `callbackRegistrations`. runAsync is
+     * wrapped separately, on the invocation owner rather than on the agent.
+     */
+    installCallbacks(agent) {
+        // Installing on a different agent must release the previous one first, or
+        // that agent keeps this bundle as its current owner and its slots keep our
+        // callbacks for good. bindFinalizer() releases its old owner for the same
+        // reason when a bundle migrates.
+        if (this.callbackAgent && this.callbackAgent !== agent)
+            this.releaseCallbacks(this.callbackAgent);
+        let registration = callbackRegistrations.get(agent);
+        // ADK gives each lifecycle callback a single slot, so a second live bundle
+        // on one agent cannot be dispatched to. Replacing the slots under a running
+        // request would route its remaining tool calls through the new request's
+        // policy and session, enforcing one user's call under another user's
+        // policy. Refuse instead of doing that silently. Sequential reuse is fine:
+        // a finished request's bundle is closed, and uninstall() clears `current`.
+        if (registration?.current && registration.current !== this && !registration.current.closed) {
+            throw new Error('ArmorIQADKBundle: this agent already has a live bundle installed. ADK gives each ' +
+                'lifecycle callback one slot, so installing a second bundle would enforce the first ' +
+                "request's tool calls under this request's policy and session. Give each concurrent " +
+                'request its own agent, or close the first bundle before installing the next.');
+        }
+        if (!registration) {
+            const original = {};
+            for (const slot of ADK_CALLBACK_SLOTS)
+                original[slot] = agent[slot];
+            registration = { original };
+            callbackRegistrations.set(agent, registration);
+        }
+        else {
+            recaptureReplacedSlots(agent, registration);
+        }
+        registration.current = this;
+        this.callbackAgent = agent;
+        const ourBeforeModel = (...args) => this.beforeModel(...args);
+        const ourAfterModel = (...args) => this.afterModel(...args);
+        const ourBeforeTool = (...args) => {
             const a = args[0];
             // Require `tool` specifically — anything else means we're in legacy
             // positional mode and a is the BaseTool itself.
@@ -468,7 +872,7 @@ class ArmorIQADKBundle {
             }
             return this.beforeTool(a, args[1], args[2]);
         };
-        agent.afterToolCallback = (...args) => {
+        const ourAfterTool = (...args) => {
             const a = args[0];
             // Require `tool` specifically — anything else means we're in legacy
             // positional mode and a is the BaseTool itself.
@@ -477,15 +881,292 @@ class ArmorIQADKBundle {
             }
             return this.afterTool(a, args[1], args[2], args[3]);
         };
-        return this;
+        // ADK accepts a single callback or an ordered list. Preserve callbacks the
+        // application installed first instead of silently replacing them. ADK
+        // stops the list only when a callback returns a replacement payload, so
+        // ArmorIQ still fail-closes before any later app callback/tool body runs.
+        const chain = (callback, ours) => [
+            ours,
+            ...(callback ? (Array.isArray(callback) ? callback : [callback]) : []),
+        ];
+        const ours = {
+            beforeModelCallback: ourBeforeModel,
+            afterModelCallback: ourAfterModel,
+            beforeToolCallback: ourBeforeTool,
+            afterToolCallback: ourAfterTool,
+        };
+        const installed = {};
+        for (const slot of ADK_CALLBACK_SLOTS) {
+            installed[slot] = chain(registration.original[slot], ours[slot]);
+            agent[slot] = installed[slot];
+        }
+        registration.installed = installed;
+    }
+    /**
+     * Subscribe this bundle to the close-on-completion wrapper for `owner`,
+     * installing that wrapper if this is the first bundle to claim the owner.
+     *
+     * There is exactly one wrapper per owner (see `runAsyncRegistrations`), so
+     * a second bundle resolving the same owner joins the existing wrapper's
+     * subscriber set rather than layering a second wrapper on top of the first.
+     * That keeps the true pre-wrap `runAsync` recoverable regardless of the
+     * order bundles later unsubscribe in, and means no inert wrapper can be
+     * left behind to retain a torn-down bundle.
+     */
+    bindFinalizer(owner) {
+        // Migrating to a new owner (handleLateParent) must release the old one,
+        // or the previous owner keeps this bundle alive in its subscriber set and
+        // would finalize it at the wrong — inner — invocation boundary.
+        if (this.runAsyncOwner && this.runAsyncOwner !== owner)
+            this.unsubscribeRunAsync();
+        this.runAsyncOwner = owner;
+        if (!owner.runAsync)
+            return;
+        const existing = runAsyncRegistrations.get(owner);
+        if (existing) {
+            existing.subscribers.add(this);
+            return;
+        }
+        const original = owner.runAsync;
+        const subscribers = new Set([this]);
+        const finalizeAll = async (status) => {
+            // Snapshot before the first await: closing one bundle can mutate the
+            // live set (its own uninstall, or an application close handler tearing
+            // down a sibling), and iterating a set being mutated mid-flight would
+            // skip subscribers. Each bundle checks its OWN closed state — sharing
+            // one `closed` reading across bundles would let the first bundle's
+            // close suppress the rest.
+            for (const subscriber of [...subscribers]) {
+                if (!subscriber.closed)
+                    await subscriber.close(status);
+            }
+        };
+        const wrapped = function wrappedRunAsync(...args) {
+            const source = original.apply(this, args);
+            return (async function* () {
+                let completed = false;
+                try {
+                    for await (const event of source)
+                        yield event;
+                    completed = true;
+                }
+                catch (error) {
+                    await finalizeAll(terminalStatusFromError(error));
+                    throw error;
+                }
+                finally {
+                    await finalizeAll(completed ? 'ok' : 'cancelled');
+                }
+            })();
+        };
+        owner.runAsync = wrapped;
+        runAsyncRegistrations.set(owner, { original, wrapper: wrapped, subscribers });
+    }
+    /**
+     * Drop this bundle from its owner's wrapper subscriber set. The shared
+     * wrapper stays in place while any sibling bundle still needs it; only when
+     * the set drains does the owner get its TRUE pre-wrap `runAsync` back — the
+     * one recorded at registration time, never a snapshot of whatever happened
+     * to be installed when this particular bundle bound. Uninstall order is
+     * therefore irrelevant, and no wrapper survives an owner's last subscriber.
+     */
+    unsubscribeRunAsync() {
+        const owner = this.runAsyncOwner;
+        if (!owner)
+            return;
+        this.runAsyncOwner = undefined;
+        const registration = runAsyncRegistrations.get(owner);
+        if (!registration)
+            return;
+        registration.subscribers.delete(this);
+        if (registration.subscribers.size > 0)
+            return;
+        // Something outside this module wrapped `runAsync` after us; unwinding to
+        // `original` would silently drop that wrapper. Leave the registration in
+        // place instead — our (now subscriber-less, so inert) wrapper is still in
+        // the chain, and a future bindFinalizer() re-subscribes to it rather than
+        // stacking another layer.
+        if (owner.runAsync !== registration.wrapper)
+            return;
+        owner.runAsync = registration.original;
+        runAsyncRegistrations.delete(owner);
+    }
+    /**
+     * Release this bundle from both owner-keyed structures, so neither can be
+     * left holding it because only the other was released.
+     *
+     * The agent-keyed callback registry is deliberately NOT released here: it
+     * must survive close so a late tool call still meets `beforeTool`'s
+     * fail-closed refusal. Only uninstall() releases it.
+     */
+    releaseFromOwners() {
+        this.unsubscribeRunAsync();
+        if (this.parentWatchTarget)
+            this.restoreParentWatcher(this.parentWatchTarget);
+    }
+    /**
+     * Put the agent's four callback slots back the way they were before ArmorIQ
+     * touched them, but only while this bundle still occupies them.
+     *
+     * A later bundle that installed over this one owns the slots now, so
+     * unwinding here would tear out live enforcement; that bundle's own release
+     * restores the same recorded originals when its turn comes, which is why
+     * uninstall order cannot matter. Each slot is compared against exactly what
+     * we wrote, so a callback the application replaced afterwards is left alone
+     * (the same guard `unsubscribeRunAsync()` applies to `runAsync`).
+     *
+     * Deliberately NOT called from closeInternal(). A closed bundle's callbacks
+     * stay installed so a tool call arriving after the request ended still meets
+     * `beforeTool`'s fail-closed refusal instead of reaching the application's
+     * unguarded callbacks. That leaves at most the one current bundle reachable
+     * from the agent, which the next install() replaces.
+     */
+    releaseCallbacks(agent) {
+        const a = agent ?? this.callbackAgent;
+        if (!a)
+            return;
+        const registration = callbackRegistrations.get(a);
+        if (registration?.current !== this)
+            return;
+        const unwoundEverySlot = restoreCallbackSlots(a, registration);
+        if (this.callbackAgent === a)
+            this.callbackAgent = undefined;
+        registration.current = undefined;
+        // A slot we could not unwind still holds our callback inside the
+        // application's wrapper. Dropping the record would let the next install()
+        // capture that as "original" and chain a dead bundle back in.
+        if (unwoundEverySlot)
+            callbackRegistrations.delete(a);
+    }
+    /**
+     * Replace `agent.parentAgent` with an accessor so we notice the moment ADK
+     * (or application code) assigns a real parent to it — which happens
+     * strictly before that parent's own runAsync can be invoked, since ADK
+     * always finishes composing a tree before running it. When that happens,
+     * re-bind the finalizer to the *new* outermost owner instead, so an
+     * install()-before-composition call ends up wrapping the same object
+     * install()-after-composition would have.
+     *
+     * Static, and keyed by `owner` (via `lateParentSubscribers`) rather than by
+     * `bundle`, because more than one bundle can resolve the same owner (see
+     * the registry doc comment above) — at most one accessor is ever installed
+     * per owner, and its setter fans out to every subscribed bundle.
+     */
+    static watchForLateParent(owner, bundle) {
+        bundle.parentWatchTarget = owner;
+        const existing = lateParentSubscribers.get(owner);
+        if (existing) {
+            existing.add(bundle);
+            return;
+        }
+        const subscribers = new Set([bundle]);
+        lateParentSubscribers.set(owner, subscribers);
+        let backing = owner.parentAgent;
+        Object.defineProperty(owner, 'parentAgent', {
+            configurable: true,
+            enumerable: true,
+            get: () => backing,
+            set: (value) => {
+                backing = value;
+                if (!value)
+                    return;
+                // ADK's own setParentAgentForSubAgents throws if parentAgent is
+                // already set, so a real parent assignment fires this at most once
+                // per owner. Restore a plain data property now that the late-parent
+                // transition is resolved, then notify every subscriber so each one
+                // independently migrates its own finalizer.
+                lateParentSubscribers.delete(owner);
+                delete owner.parentAgent;
+                owner.parentAgent = value;
+                for (const subscriber of subscribers)
+                    subscriber.handleLateParent(owner, value);
+            },
+        });
+    }
+    handleLateParent(agent, parent) {
+        // A closed bundle must never rejoin any owner's registry. Normally
+        // releaseFromOwners() has already dropped this bundle out of `agent`'s
+        // lateParentSubscribers set by the time a real late parent arrives, but
+        // closeInternal() sets `closed` synchronously and then awaits several
+        // times before it reaches releaseFromOwners() — a late parent CAN land
+        // on `agent` during that window, while this (already-closed) bundle is
+        // still subscribed. Without this guard, bindFinalizer() below would
+        // subscribe a dead bundle onto the new owner's runAsync registry, where
+        // finalizeAll() skips closed subscribers forever — stranding both the
+        // bundle and the new owner's wrapper permanently (issue this fixes).
+        if (this.closed)
+            return;
+        if (this.parentWatchTarget === agent)
+            this.parentWatchTarget = undefined;
+        // The provisional finalizer was bound to `agent` (owner === agent, or a
+        // shared intermediate owner, at the time this bundle last called
+        // bindFinalizer). bindFinalizer() below unsubscribes this bundle from
+        // that old owner before subscribing it to the new one, so the migration
+        // is complete in both directions: the inner boundary can no longer
+        // finalize us, and `agent.runAsync` is restored to its true original
+        // only once the last sibling bundle has migrated off it too.
+        const owner = ArmorIQADKBundle.invocationOwner(parent);
+        this.bindFinalizer(owner);
+        if (!owner.parentAgent) {
+            // Same generalization as install(): `parent` itself may already have
+            // been composed further up the tree by the time this late-parent
+            // assignment fires, so invocationOwner(parent) can resolve to a
+            // further ancestor. The resolved `owner` has no parent of its own yet
+            // — it could still be composed into a grandparent later — so watch
+            // `owner` itself (not necessarily `parent`) and keep watching.
+            ArmorIQADKBundle.watchForLateParent(owner, this);
+        }
+    }
+    /**
+     * Stop watching `agent` for a late parent assignment on this bundle's
+     * behalf — used by uninstall() when the late parent never arrived. Removes
+     * only this bundle from the shared subscriber set; the accessor is only
+     * torn down (restored to a plain property) once no bundle is left
+     * watching it, so unwinding one bundle never disturbs another bundle still
+     * watching the same owner.
+     */
+    restoreParentWatcher(agent) {
+        if (this.parentWatchTarget !== agent)
+            return;
+        this.parentWatchTarget = undefined;
+        const subscribers = lateParentSubscribers.get(agent);
+        if (!subscribers)
+            return;
+        subscribers.delete(this);
+        if (subscribers.size > 0)
+            return;
+        lateParentSubscribers.delete(agent);
+        const current = agent.parentAgent;
+        delete agent.parentAgent;
+        agent.parentAgent = current;
+    }
+    /**
+     * Walk ADK's public `parentAgent` link to the outermost agent that owns
+     * this invocation. An agent installed directly on a root (the common case,
+     * and every existing call site) resolves to itself, so behavior there is
+     * unchanged.
+     */
+    static invocationOwner(agent) {
+        let owner = agent;
+        const seen = new Set([owner]);
+        while (owner.parentAgent && !seen.has(owner.parentAgent)) {
+            owner = owner.parentAgent;
+            seen.add(owner);
+        }
+        return owner;
     }
     uninstall(agent) {
         const a = agent ?? this.agent;
         if (!a)
             return;
-        a.afterModelCallback = this.saved.afterModelCallback;
-        a.beforeToolCallback = this.saved.beforeToolCallback;
-        a.afterToolCallback = this.saved.afterToolCallback;
+        this.releaseCallbacks(a);
+        // runAsync may have been wrapped on an ancestor of `a` (see install()),
+        // and the late-parent watcher may be installed on a different owner
+        // still. A sibling bundle can share either structure, so releasing this
+        // bundle from both is what stops it from being finalized or migrated
+        // later; each structure is only unwound once no bundle needs it (see
+        // releaseFromOwners()).
+        this.releaseFromOwners();
     }
 }
 exports.ArmorIQADKBundle = ArmorIQADKBundle;