langsmith 0.4.8 → 0.4.9

package/dist/experimental/anthropic/index.js

@@ -1,567 +1,97 @@
-import { traceable, getCurrentRunTree, isTraceableFunction, } from "../../traceable.js";
-import { convertAnthropicUsageToInputTokenDetails } from "../../utils/usage.js";
-import { getNumberProperty } from "./utils.js";
-const createQueryContext = () => ({
-    activeToolRuns: new Map(),
-    clientManagedRuns: new Map(),
-    subagentSessions: new Map(),
-    activeSubagentToolUseId: undefined,
-    currentParentRun: undefined,
-});
-/**
- * PreToolUse hook that creates a tool span when a tool execution starts.
- * This traces ALL tools including built-in tools, external MCP tools, and SDK MCP tools.
- * Skips tools that are client-managed (subagent sessions and their children).
- */
-async function preToolUseHook(input, toolUseId, context) {
-    if (!toolUseId)
-        return {};
-    // Skip if this tool run is already managed by the client (subagent or its children)
-    if (context.clientManagedRuns.has(toolUseId)) {
-        return {};
-    }
-    const toolName = input.tool_name || "unknown_tool";
-    const toolInput = input.tool_input;
-    try {
-        const parent = context.currentParentRun || getCurrentRunTree();
-        if (!parent) {
-            return {};
-        }
-        const startTime = Date.now();
-        const toolRun = await parent.createChild({
-            name: toolName,
-            run_type: "tool",
-            inputs: toolInput ? { input: toolInput } : {},
-        });
-        await toolRun.postRun();
-        context.activeToolRuns.set(toolUseId, { run: toolRun, startTime });
-    }
-    catch {
-        // Silently fail - don't interrupt tool execution
-    }
-    return {};
-}
-/**
- * PostToolUse hook that ends the tool span when a tool execution completes.
- * Handles both regular tool runs and client-managed runs (subagents and their children).
- */
-async function postToolUseHook(input, toolUseId, context) {
-    if (!toolUseId)
-        return {};
-    const toolResponse = input.tool_response;
-    // Format outputs based on response type
-    const formatOutputs = (response) => {
-        let outputs;
-        if (typeof response === "object" && response !== null) {
-            if (Array.isArray(response)) {
-                outputs = { content: response };
-            }
-            else {
-                outputs = response;
-            }
-        }
-        else {
-            outputs = response ? { output: String(response) } : {};
-        }
-        const isError = typeof response === "object" &&
-            response !== null &&
-            "is_error" in response &&
-            response.is_error === true;
-        return { outputs, isError };
-    };
-    try {
-        // Check if this is a client-managed run (subagent session or its children)
-        const clientRun = context.clientManagedRuns.get(toolUseId);
-        if (clientRun) {
-            context.clientManagedRuns.delete(toolUseId);
-            const { outputs, isError } = formatOutputs(toolResponse);
-            await clientRun.end({
-                outputs,
-                error: isError ? outputs.output?.toString() : undefined,
-            });
-            await clientRun.patchRun();
-            return {};
-        }
-        // Handle regular tool runs
-        const runInfo = context.activeToolRuns.get(toolUseId);
-        if (!runInfo) {
-            return {};
-        }
-        context.activeToolRuns.delete(toolUseId);
-        const { run: toolRun } = runInfo;
-        const { outputs, isError } = formatOutputs(toolResponse);
-        await toolRun.end({
-            outputs,
-            error: isError ? outputs.output?.toString() : undefined,
-        });
-        await toolRun.patchRun();
-    }
-    catch {
-        // Silently fail - don't interrupt tool execution
-    }
-    return {};
-}
-/**
- * Creates hook matchers for LangSmith tracing.
- * Returns PreToolUse and PostToolUse hook configurations.
- */
-function createTracingHooks(context) {
-    return {
-        PreToolUse: [
-            {
-                matcher: undefined, // Match all tools
-                hooks: [
-                    async (input, toolUseId, _options) => preToolUseHook(input, toolUseId, context),
-                ],
-            },
-        ],
-        PostToolUse: [
-            {
-                matcher: undefined, // Match all tools
-                hooks: [
-                    async (input, toolUseId, _options) => postToolUseHook(input, toolUseId, context),
-                ],
-            },
-        ],
-        SessionEnd: [
-            {
-                matcher: undefined,
-                hooks: [
-                    async (_input) => {
-                        // Clean up at end of session
-                        clearActiveToolRuns(context);
-                        return {};
-                    },
-                ],
-            },
-        ],
-        SubagentStop: [
-            {
-                matcher: undefined,
-                hooks: [
-                    async (_input, toolUseId) => {
-                        // Clean up subagent session
-                        if (toolUseId) {
-                            context.subagentSessions.delete(toolUseId);
-                            context.clientManagedRuns.delete(toolUseId);
-                        }
-                        return {};
-                    },
-                ],
-            },
-        ],
-        Stop: [
-            {
-                matcher: undefined,
-                hooks: [
-                    async (_input) => {
-                        // Clean up on stop - ensure all runs are finalized
-                        clearActiveToolRuns(context);
-                        return {};
-                    },
-                ],
-            },
-        ],
-    };
-}
-/**
- * Merges LangSmith tracing hooks with existing user hooks.
- */
-function mergeHooks(existingHooks, context) {
-    const tracingHooks = createTracingHooks(context);
-    if (!existingHooks)
-        return tracingHooks;
-    const merged = { ...existingHooks };
-    // Prepend tracing hooks so they run first
-    for (const [event, matchers] of Object.entries(tracingHooks)) {
-        merged[event] = [...matchers, ...(merged[event] ?? [])];
-    }
-    return merged;
-}
-/**
- * Type assertion to check if a tool is a Task tool
- * @param tool - The tool to check
- * @returns True if the tool is a Task tool, false otherwise
- */
-function isTaskTool(tool) {
-    return tool.type === "tool_use" && tool.name === "Task";
-}
-/**
- * Type-assertion to check for tool blocks
- */
-function isToolBlock(block) {
-    if (!block || typeof block !== "object")
-        return false;
-    return block.type === "tool_use";
-}
-/**
- * Processes tool uses in an AssistantMessage to detect and create subagent sessions.
- * This matches Python's _handle_assistant_tool_uses behavior.
- *
- * @param message - The AssistantMessage to process
- * @param parentRun - The parent run tree (main conversation chain)
- */
-async function handleAssistantToolUses(message, parentRun, context) {
-    if (!parentRun)
-        return;
-    const content = message.message?.content;
-    if (!Array.isArray(content))
-        return;
-    const parentToolUseId = message.parent_tool_use_id;
-    for (const block of content) {
-        if (!isToolBlock(block) || !block.id)
-            continue;
-        try {
-            // Check if this is a Task tool (subagent) at the top level
-            if (isTaskTool(block) && !parentToolUseId) {
-                // Extract subagent name from input
-                const subagentName = block.input.subagent_type ||
-                    block.input.agent_type ||
-                    (block.input.description
-                        ? block.input.description.split(" ")[0]
-                        : null) ||
-                    "unknown-agent";
-                const subagentSession = await parentRun.createChild({
-                    name: subagentName,
-                    run_type: "chain",
-                    inputs: block.input,
-                });
-                // Post the run to start it, but DON'T end it yet
-                // It will be ended when we receive the tool result or at cleanup
-                await subagentSession.postRun();
-                // Store in both maps
-                context.subagentSessions.set(block.id, subagentSession);
-                context.clientManagedRuns.set(block.id, subagentSession);
-            }
-            // Check if tool use is within a subagent
-            else if (parentToolUseId &&
-                context.subagentSessions.has(parentToolUseId)) {
-                // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
-                const subagentSession = context.subagentSessions.get(parentToolUseId);
-                // Create tool run as child of subagent
-                const toolRun = await subagentSession.createChild({
-                    name: block.name || "unknown_tool",
-                    run_type: "tool",
-                    inputs: block.input ? { input: block.input } : {},
-                });
-                await toolRun.postRun();
-                context.clientManagedRuns.set(block.id, toolRun);
-            }
-        }
-        catch {
-            // Silently fail - don't interrupt message processing
-        }
-    }
-}
-/**
- * Clears all active tool runs and client-managed runs. Called when a conversation ends.
- */
-function clearActiveToolRuns(context) {
-    // Clean up client-managed runs (subagents and their children)
-    for (const [, run] of context.clientManagedRuns) {
-        try {
-            run
-                .end({ error: "Run not completed (conversation ended)" })
-                .then(() => run.patchRun())
-                .catch(() => { });
-        }
-        catch {
-            // Ignore cleanup errors
-        }
-    }
-    context.clientManagedRuns.clear();
-    context.subagentSessions.clear();
-    context.activeSubagentToolUseId = undefined;
-    // Clean up regular tool runs
-    for (const [, { run }] of context.activeToolRuns) {
-        try {
-            run
-                .end({ error: "Tool run not completed (conversation ended)" })
-                .then(() => run.patchRun())
-                .catch(() => { });
-        }
-        catch {
-            // Ignore cleanup errors
-        }
-    }
-    context.activeToolRuns.clear();
-}
+import { traceable, isTraceableFunction } from "../../traceable.js";
+import { StreamManager } from "./context.js";
+import { convertFromAnthropicMessage } from "./messages.js";
 /**
  * Wraps the Claude Agent SDK's query function to add LangSmith tracing.
  * Traces the entire agent interaction including all streaming messages.
- * Internal use only - use wrapClaudeAgentSDK instead.
+ * @internal Use `wrapClaudeAgentSDK` instead.
  */
 function wrapClaudeAgentQuery(queryFn, defaultThis, baseConfig) {
-    const getModifiedArgs = (args, context) => {
-        const params = (args[0] ?? {});
-        const { prompt, options = {} } = params;
-        // Inject LangSmith tracing hooks into options
-        const mergedHooks = mergeHooks(options.hooks, context);
-        const modifiedOptions = { ...options, hooks: mergedHooks };
-        const modifiedParams = { ...params, options: modifiedOptions };
-        return {
-            prompt,
-            options: modifiedOptions,
-            modifiedArgs: [modifiedParams, ...args.slice(1)],
-        };
-    };
-    async function* generator(originalGenerator, prompt, options, context) {
-        const finalResults = [];
-        // Track assistant messages by their message ID for proper streaming handling
-        // Each message ID maps to { message, startTime } - we keep the latest streaming update
-        const pendingMessages = new Map();
-        // Track which message IDs have already had spans created
-        // This prevents creating duplicate spans when the SDK sends multiple updates
-        // for the same message ID with stop_reason set
-        const completedMessageIds = new Set();
-        // Store child run promises for proper async handling
-        const childRunEndPromises = [];
-        // Track usage from ResultMessage to add to the parent span
-        let resultUsage;
-        // Track additional metadata from the SDK
-        const extraMetadata = [];
-        // Track usage from completed assistant message spans (by model)
-        // Used to calculate remaining tokens for pending messages
-        const completedUsageByModel = new Map();
-        // Create an LLM span for a specific message ID
-        const createLLMSpanForId = async (messageId) => {
-            // Skip if we've already created a span for this message ID
-            if (completedMessageIds.has(messageId)) {
-                return;
-            }
-            const pending = pendingMessages.get(messageId);
-            if (!pending)
-                return;
-            pendingMessages.delete(messageId);
-            completedMessageIds.add(messageId);
-            // Track the usage before creating the span
-            const model = pending.message.message?.model;
-            const usage = pending.message.message?.usage;
-            if (model && usage) {
-                const existing = completedUsageByModel.get(model) || {
-                    inputTokens: 0,
-                    outputTokens: 0,
-                    cacheReadTokens: 0,
-                    cacheCreationTokens: 0,
-                };
-                existing.inputTokens += usage.input_tokens || 0;
-                existing.outputTokens += usage.output_tokens || 0;
-                existing.cacheReadTokens += usage.cache_read_input_tokens || 0;
-                existing.cacheCreationTokens += usage.cache_creation_input_tokens || 0;
-                completedUsageByModel.set(model, existing);
-            }
-            const finalMessageContent = await createLLMSpanForMessages([pending.message], pending.messageHistory, options, pending.startTime, context);
-            if (finalMessageContent)
-                finalResults.push(finalMessageContent);
-        };
+    async function* generator(originalGenerator, prompt) {
+        const streamManager = new StreamManager();
         try {
+            let systemCount = 0;
             for await (const message of originalGenerator) {
-                const currentTime = Date.now();
                 if (message.type === "system") {
-                    const content = getLatestInput(prompt);
+                    const content = getLatestInput(prompt, systemCount);
+                    systemCount += 1;
                     if (content != null)
-                        finalResults.push(content);
-                }
-                // Handle assistant messages - group by message ID for streaming
-                // Multiple messages with the same ID are streaming updates; use the last one
-                if (message.type === "assistant") {
-                    const messageId = message.message?.id;
-                    // If we have an active subagent context and this message doesn't have parent_tool_use_id,
-                    // check if this is a new main conversation message (which would end the subagent execution)
-                    if (context.activeSubagentToolUseId && !message.parent_tool_use_id) {
-                        // Check if this message contains tool uses - if it does, it's part of main conversation
-                        const content = message.message?.content;
-                        if (Array.isArray(content)) {
-                            const hasToolUse = content.some((block) => block &&
-                                typeof block === "object" &&
-                                block.type === "tool_use");
-                            // If this message has tool uses and none are within the subagent, it's a new turn
-                            if (hasToolUse) {
-                                // Clean up the subagent session
-                                context.subagentSessions.delete(context.activeSubagentToolUseId);
-                                context.activeSubagentToolUseId = undefined;
-                            }
-                        }
-                    }
-                    if (messageId) {
-                        // Check if this is a new message or an update to existing
-                        const existing = pendingMessages.get(messageId);
-                        if (!existing) {
-                            // New message arrived - finalize all OTHER pending messages first
-                            // (they must be complete if we're seeing a new message)
-                            // Finalize all other pending messages
-                            for (const [otherId] of pendingMessages) {
-                                if (otherId !== messageId) {
-                                    const spanPromise = createLLMSpanForId(otherId);
-                                    childRunEndPromises.push(spanPromise);
-                                }
-                            }
-                            pendingMessages.set(messageId, {
-                                message,
-                                messageHistory: finalResults.slice(0),
-                                startTime: currentTime,
-                            });
-                        }
-                        else {
-                            // Streaming update - keep the start time, update the message
-                            pendingMessages.set(messageId, {
-                                message,
-                                messageHistory: finalResults.slice(0),
-                                startTime: existing.startTime,
-                            });
-                        }
-                        // Push the message to the final results,
-                        // Used to create spans with the full chat history as input
-                        if ("content" in message.message && message.message.content) {
-                            finalResults.push({
-                                content: flattenContentBlocks(message.message.content),
-                                role: "assistant",
-                            });
-                        }
-                        // Check if this message has a stop_reason (meaning it's complete)
-                        // If so, create the span now (createLLMSpanForId will skip if already created)
-                        if (message.message?.stop_reason) {
-                            const spanPromise = createLLMSpanForId(messageId);
-                            childRunEndPromises.push(spanPromise);
-                        }
-                    }
-                    // Process tool uses for subagent detection (matches Python's _handle_assistant_tool_uses)
-                    await handleAssistantToolUses(message, context.currentParentRun, context);
-                }
-                // Handle UserMessage - add to conversation history (matches Python)
-                if (message.type === "user") {
-                    if ("content" in message.message && message.message.content) {
-                        finalResults.push({
-                            content: flattenContentBlocks(message.message.content),
-                            role: "user",
-                        });
-                    }
-                    // If this is a tool result for a Task tool (subagent), we're entering the subagent's execution
-                    // The subagent's assistant messages will come AFTER this result
-                    if (message.parent_tool_use_id &&
-                        context.subagentSessions.has(message.parent_tool_use_id)) {
-                        context.activeSubagentToolUseId = message.parent_tool_use_id;
-                    }
-                }
-                // Handle ResultMessage - extract usage and metadata
-                if (message.type === "result") {
-                    // If modelUsage is available, aggregate from it (includes ALL models)
-                    // Otherwise fall back to top-level usage field
-                    if (message.modelUsage) {
-                        // Aggregate usage from modelUsage (includes ALL models)
-                        resultUsage = aggregateUsageFromModelUsage(message.modelUsage);
-                        // Patch token counts for pending messages using modelUsage
-                        // This handles the SDK limitation where the last assistant message
-                        // doesn't receive final streaming updates with accurate token counts
-                        for (const [, { message: pendingMsg }] of pendingMessages) {
-                            const model = pendingMsg.message?.model;
-                            if (model &&
-                                message.modelUsage[model] &&
-                                pendingMsg.message?.usage) {
-                                const modelStats = message.modelUsage[model];
-                                const completed = completedUsageByModel.get(model) || {
-                                    inputTokens: 0,
-                                    outputTokens: 0,
-                                    cacheReadTokens: 0,
-                                    cacheCreationTokens: 0,
-                                };
-                                // Calculate remaining tokens = total - completed
-                                const remainingOutput = (modelStats.outputTokens || 0) - completed.outputTokens;
-                                const remainingInput = (modelStats.inputTokens || 0) - completed.inputTokens;
-                                const remainingCacheRead = (modelStats.cacheReadInputTokens || 0) -
-                                    completed.cacheReadTokens;
-                                const remainingCacheCreation = (modelStats.cacheCreationInputTokens || 0) -
-                                    completed.cacheCreationTokens;
-                                // Update the pending message's usage with remaining tokens
-                                pendingMsg.message.usage.output_tokens = Math.max(0, remainingOutput);
-                                pendingMsg.message.usage.input_tokens = Math.max(0, remainingInput);
-                                if (remainingCacheRead > 0) {
-                                    pendingMsg.message.usage.cache_read_input_tokens =
-                                        remainingCacheRead;
-                                }
-                                if (remainingCacheCreation > 0) {
-                                    pendingMsg.message.usage.cache_creation_input_tokens =
-                                        remainingCacheCreation;
-                                }
-                            }
-                        }
-                    }
-                    else if (message.usage) {
-                        // Fall back to top-level usage if modelUsage not available
-                        resultUsage = extractUsageFromMessage(message);
-                    }
-                    // Add total_cost if available (LangSmith standard field)
-                    if (message.total_cost_usd != null && resultUsage) {
-                        resultUsage.total_cost = message.total_cost_usd;
-                    }
-                    // Add conversation-level metadata
-                    if (message.is_error != null) {
-                        extraMetadata.push(["is_error", message.is_error]);
-                    }
-                    if (message.num_turns != null) {
-                        extraMetadata.push(["num_turns", message.num_turns]);
-                    }
-                    if (message.session_id != null) {
-                        extraMetadata.push(["session_id", message.session_id]);
-                    }
-                    if (message.duration_ms != null) {
-                        extraMetadata.push(["duration_ms", message.duration_ms]);
-                    }
-                    if (message.duration_api_ms != null) {
-                        extraMetadata.push(["duration_api_ms", message.duration_api_ms]);
-                    }
+                        streamManager.addMessage(content);
                 }
+                streamManager.addMessage(message);
                 yield message;
             }
-            // Create spans for any remaining pending messages (those without stop_reason)
-            for (const messageId of pendingMessages.keys()) {
-                const spanPromise = createLLMSpanForId(messageId);
-                childRunEndPromises.push(spanPromise);
-            }
-            // Wait for all child runs to complete
-            await Promise.all(childRunEndPromises);
-            // Apply usage metadata to the chain run using LangSmith's standard fields
-            const currentRun = getCurrentRunTree();
-            if (currentRun && (resultUsage || extraMetadata.length > 0)) {
-                // Initialize metadata object if needed
-                currentRun.extra ||= {};
-                currentRun.extra.metadata ||= {};
-                if (resultUsage) {
-                    // Add LangSmith-standard usage fields directly to metadata
-                    if (resultUsage.input_tokens !== undefined) {
-                        currentRun.extra.metadata.input_tokens = resultUsage.input_tokens;
-                    }
-                    if (resultUsage.output_tokens !== undefined) {
-                        currentRun.extra.metadata.output_tokens = resultUsage.output_tokens;
-                    }
-                    if (resultUsage.total_tokens !== undefined) {
-                        currentRun.extra.metadata.total_tokens = resultUsage.total_tokens;
-                    }
-                    if (resultUsage.input_token_details) {
-                        currentRun.extra.metadata.input_token_details =
-                            resultUsage.input_token_details;
-                    }
-                    if (resultUsage.total_cost !== undefined) {
-                        currentRun.extra.metadata.total_cost = resultUsage.total_cost;
-                    }
-                }
-                for (const [key, value] of extraMetadata) {
-                    currentRun.extra.metadata[key] = value;
-                }
-            }
         }
         finally {
-            // Clean up parent run reference and any orphaned tool runs
-            context.currentParentRun = undefined;
-            clearActiveToolRuns(context);
-        }
+            await streamManager.finish();
+        }
+    }
+    function getLatestInput(arg, systemCount) {
+        const value = (() => {
+            if (typeof arg !== "object" || arg == null)
+                return arg;
+            const toJSON = arg["toJSON"];
+            if (typeof toJSON !== "function")
+                return undefined;
+            const latest = toJSON();
+            return latest?.at(systemCount);
+        })();
+        if (value == null)
+            return undefined;
+        if (typeof value === "string") {
+            return {
+                type: "user",
+                message: { content: value, role: "user" },
+                parent_tool_use_id: null,
+                session_id: "",
+            };
+        }
+        return typeof value === "object" && value != null ? value : undefined;
+    }
+    async function processInputs(rawInputs) {
+        const inputs = rawInputs;
+        const newInputs = { ...inputs };
+        return Object.assign(newInputs, {
+            toJSON: () => {
+                const toJSON = (value) => {
+                    if (typeof value !== "object" || value == null)
+                        return value;
+                    const fn = value?.toJSON;
+                    if (typeof fn === "function")
+                        return fn();
+                    return value;
+                };
+                const prompt = toJSON(inputs.prompt);
+                const options = inputs.options != null
+                    ? { ...inputs.options }
+                    : undefined;
+                if (options?.mcpServers != null) {
+                    options.mcpServers = Object.fromEntries(Object.entries(options.mcpServers ?? {}).map(([key, value]) => [
+                        key,
+                        { name: value.name, type: value.type },
+                    ]));
+                }
+                return { messages: convertFromAnthropicMessage(prompt), options };
+            },
+        });
     }
-    const wrapped = (...args) => {
-        const context = createQueryContext();
-        context.currentParentRun = getCurrentRunTree();
-        const { prompt, options, modifiedArgs } = getModifiedArgs(args, context);
-        const actualGenerator = queryFn.call(defaultThis, ...modifiedArgs);
-        const wrappedGenerator = generator(actualGenerator, prompt, options, context);
+    function processOutputs(rawOutputs) {
+        if ("outputs" in rawOutputs && Array.isArray(rawOutputs.outputs)) {
+            const sdkMessages = rawOutputs.outputs;
+            const messages = sdkMessages
+                .filter((message) => {
+                if (!("message" in message))
+                    return true;
+                return message.parent_tool_use_id == null;
+            })
+                .flatMap(convertFromAnthropicMessage);
+            return { output: { messages } };
+        }
+        return rawOutputs;
+    }
+    return traceable((params, ...args) => {
+        const actualGenerator = queryFn.call(defaultThis, params, ...args);
+        const wrappedGenerator = generator(actualGenerator, params.prompt);
         for (const method of Object.getOwnPropertyNames(Object.getPrototypeOf(actualGenerator)).filter((method) => !["constructor", "next", "throw", "return"].includes(method))) {
             Object.defineProperty(wrappedGenerator, method, {
                 get() {
@@ -573,318 +103,14 @@ function wrapClaudeAgentQuery(queryFn, defaultThis, baseConfig) {
             });
         }
         return wrappedGenerator;
-    };
-    // Wrap in traceable
-    return traceable(wrapped, {
+    }, {
         name: "claude.conversation",
         run_type: "chain",
         ...baseConfig,
         metadata: { ...baseConfig?.metadata },
         __deferredSerializableArgOptions: { maxDepth: 1 },
-        async processInputs(rawInputs) {
-            const inputs = rawInputs;
-            const newInputs = { ...inputs };
-            return Object.assign(newInputs, {
-                toJSON: () => {
-                    const toJSON = (value) => {
-                        if (typeof value !== "object" || value == null)
-                            return value;
-                        const fn = value?.toJSON;
-                        if (typeof fn === "function")
-                            return fn();
-                        return value;
-                    };
-                    const prompt = toJSON(inputs.prompt);
-                    const options = toJSON(inputs.options);
-                    const messages = (() => {
-                        if (prompt == null)
-                            return undefined;
-                        const result = [];
-                        if (typeof prompt === "string") {
-                            result.push({ content: prompt, role: "user" });
-                        }
-                        else {
-                            for (const { message } of prompt) {
-                                if (!message)
-                                    continue;
-                                result.push({
-                                    content: flattenContentBlocks(message.content),
-                                    role: message.role,
-                                });
-                            }
-                        }
-                        return result;
-                    })();
-                    return { messages, options };
-                },
-            });
-        },
-        processOutputs(rawOutputs) {
-            if ("outputs" in rawOutputs && Array.isArray(rawOutputs.outputs)) {
-                const sdkMessages = rawOutputs.outputs;
-                const messages = sdkMessages.flatMap((sdkMessage) => {
-                    if ("message" in sdkMessage && sdkMessage.message != null) {
-                        return {
-                            role: sdkMessage.message.role,
-                            content: flattenContentBlocks(sdkMessage.message.content),
-                        };
-                    }
-                    return [];
-                });
-                return { output: { messages } };
-            }
-            return rawOutputs;
-        },
-    });
-}
-/**
- * Wraps a Claude Agent SDK tool definition to add LangSmith tracing for tool executions.
- * Internal use only - use wrapClaudeAgentSDK instead.
- */
-function wrapClaudeAgentTool(toolDef, baseConfig) {
-    return {
-        ...toolDef,
-        handler: traceable(toolDef.handler, {
-            name: toolDef.name,
-            run_type: "tool",
-            ...baseConfig,
-        }),
-    };
-}
-/**
- * Aggregates usage from modelUsage breakdown (includes all models, including hidden ones).
- * This provides accurate totals when multiple models are used.
- */
-function aggregateUsageFromModelUsage(modelUsage) {
-    const metrics = {};
-    let totalInputTokens = 0;
-    let totalOutputTokens = 0;
-    let totalCacheReadTokens = 0;
-    let totalCacheCreationTokens = 0;
-    // Aggregate across all models
-    for (const modelStats of Object.values(modelUsage)) {
-        totalInputTokens += modelStats.inputTokens || 0;
-        totalOutputTokens += modelStats.outputTokens || 0;
-        totalCacheReadTokens += modelStats.cacheReadInputTokens || 0;
-        totalCacheCreationTokens += modelStats.cacheCreationInputTokens || 0;
-    }
-    // Build input_token_details if we have cache tokens
-    if (totalCacheReadTokens > 0 || totalCacheCreationTokens > 0) {
-        metrics.input_token_details = {
-            cache_read: totalCacheReadTokens,
-            cache_creation: totalCacheCreationTokens,
-        };
-    }
-    // Sum all input tokens (new + cache read + cache creation)
-    const totalPromptTokens = totalInputTokens + totalCacheReadTokens + totalCacheCreationTokens;
-    metrics.input_tokens = totalPromptTokens;
-    metrics.output_tokens = totalOutputTokens;
-    metrics.total_tokens = totalPromptTokens + totalOutputTokens;
-    return metrics;
-}
-/**
- * Extracts and normalizes usage metrics from a Claude Agent SDK message.
- */
-function extractUsageFromMessage(message) {
-    const metrics = {};
-    // Assistant messages contain usage in message.message.usage
-    // Result messages contain usage in message.usage
-    let usage;
-    if (message.type === "assistant") {
-        usage = message.message?.usage;
-    }
-    else if (message.type === "result") {
-        usage = message.usage;
-    }
-    if (!usage || typeof usage !== "object") {
-        return metrics;
-    }
-    // Standard token counts - use LangSmith's expected field names
-    const inputTokens = getNumberProperty(usage, "input_tokens") || 0;
-    const outputTokens = getNumberProperty(usage, "output_tokens") || 0;
-    // Get cache tokens
-    const cacheRead = getNumberProperty(usage, "cache_read_input_tokens") || 0;
-    const cacheCreation = getNumberProperty(usage, "cache_creation_input_tokens") || 0;
-    // Build input_token_details if we have cache tokens
-    if (cacheRead > 0 || cacheCreation > 0) {
-        const inputTokenDetails = convertAnthropicUsageToInputTokenDetails(usage);
-        if (Object.keys(inputTokenDetails).length > 0) {
-            metrics.input_token_details = inputTokenDetails;
-        }
-    }
-    // Sum cache tokens into input_tokens total (matching Python's sum_anthropic_tokens)
-    const totalInputTokens = inputTokens + cacheRead + cacheCreation;
-    metrics.input_tokens = totalInputTokens;
-    metrics.output_tokens = outputTokens;
-    metrics.total_tokens = totalInputTokens + outputTokens;
-    return metrics;
-}
-function getLatestInput(arg) {
-    const value = (() => {
-        if (typeof arg !== "object" || arg == null)
-            return arg;
-        const toJSON = arg["toJSON"];
-        if (typeof toJSON !== "function")
-            return undefined;
-        const latest = toJSON();
-        return latest?.at(-1);
-    })();
-    if (typeof value == null)
-        return undefined;
-    if (typeof value === "string")
-        return { content: value, role: "user" };
-    const userMessage = value;
-    if (typeof userMessage === "string") {
-        return { content: userMessage, role: "user" };
-    }
-    if (typeof userMessage !== "object" || userMessage == null) {
-        return undefined;
-    }
-    return {
-        role: userMessage.message.role || "user",
-        content: flattenContentBlocks(userMessage.message.content),
-    };
-}
-/**
- * Creates an LLM span for a group of messages with the same message ID.
- * Returns the final message content to add to conversation history.
- * Handles subagent LLM turns by parenting them to the correct subagent session.
- */
-async function createLLMSpanForMessages(messages, conversationHistory, options, startTime, context) {
-    if (messages.length === 0)
-        return undefined;
-    const lastMessage = messages[messages.length - 1];
-    // Create LLM spans for all AssistantMessages, not just those with usage
-    // (matches Python's behavior)
-    if (lastMessage.type !== "assistant") {
-        return undefined;
-    }
-    // Extract model from message first, fall back to options (matches Python)
-    const model = lastMessage.message.model || options.model;
-    const usage = extractUsageFromMessage(lastMessage);
-    const input = conversationHistory.length > 0 ? conversationHistory : undefined;
-    // Flatten content blocks for proper serialization (matches Python)
-    const outputs = messages
-        .map((m) => {
-        if (!("message" in m) || !("role" in m.message))
-            return undefined;
-        return {
-            content: flattenContentBlocks(m.message.content),
-            role: m.message.role,
-        };
-    })
-        .filter((c) => c !== undefined);
-    // Check if this message belongs to a subagent
-    // First check if message has explicit parent_tool_use_id
-    const parentToolUseId = lastMessage.parent_tool_use_id;
-    let subagentParent = parentToolUseId
-        ? context.subagentSessions.get(parentToolUseId)
-        : undefined;
-    // If no explicit parent, check if we're in an active subagent context
-    if (!subagentParent && context.activeSubagentToolUseId) {
-        subagentParent = context.subagentSessions.get(context.activeSubagentToolUseId);
-    }
-    const endTime = Date.now();
-    // Format inputs: if we have a single input, use it directly; otherwise wrap as messages
-    const formattedInputs = input && input.length === 1 ? input[0] : input ? { messages: input } : {};
-    if (subagentParent) {
-        // Create LLM run as child of subagent session with proper start and end time
-        try {
-            const llmRun = await subagentParent.createChild({
-                name: "claude.assistant.turn",
-                run_type: "llm",
-                inputs: formattedInputs,
-                outputs: outputs[outputs.length - 1] || { content: outputs },
-                start_time: startTime,
-                end_time: endTime,
-                extra: {
-                    metadata: {
-                        ...(model ? { ls_model_name: model } : {}),
-                        usage_metadata: usage,
-                    },
-                },
-            });
-            await llmRun.postRun();
-        }
-        catch {
-            // Silently fail
-        }
-    }
-    else {
-        // Regular LLM turn under main conversation
-        // Note: traceable doesn't support start_time config, so we use getCurrentRunTree
-        // and manually create the child run to preserve timing
-        const currentRun = getCurrentRunTree();
-        if (currentRun) {
-            try {
-                const llmRun = await currentRun.createChild({
-                    name: "claude.assistant.turn",
-                    run_type: "llm",
-                    inputs: formattedInputs,
-                    outputs: outputs[outputs.length - 1] || { content: outputs },
-                    start_time: startTime,
-                    end_time: endTime,
-                    extra: {
-                        metadata: {
-                            ...(model ? { ls_model_name: model } : {}),
-                            usage_metadata: usage,
-                        },
-                    },
-                });
-                await llmRun.postRun();
-            }
-            catch {
-                // Silently fail
-            }
-        }
-    }
-    // Return flattened content for conversation history
-    return lastMessage.message?.content && lastMessage.message?.role
-        ? {
-            content: flattenContentBlocks(lastMessage.message.content),
-            role: lastMessage.message.role,
-        }
-        : undefined;
-}
-/**
- * Converts SDK content blocks into serializable objects.
- * Matches Python's flatten_content_blocks behavior.
- */
-function flattenContentBlocks(content) {
-    if (!Array.isArray(content)) {
-        return content;
-    }
-    return content.map((block) => {
-        if (!block || typeof block !== "object" || !("type" in block)) {
-            return block;
-        }
-        const blockType = block.type;
-        switch (blockType) {
-            case "text":
-                return { type: "text", text: block.text || "" };
-            case "thinking":
-                return {
-                    type: "thinking",
-                    thinking: block.thinking || "",
-                    signature: block.signature || "",
-                };
-            case "tool_use":
-                return {
-                    type: "tool_use",
-                    id: block.id,
-                    name: block.name,
-                    input: block.input,
-                };
-            case "tool_result":
-                return {
-                    type: "tool_result",
-                    tool_use_id: block.tool_use_id,
-                    content: block.content,
-                    is_error: block.is_error || false,
-                };
-            default:
-                return block;
-        }
+        processInputs,
+        processOutputs,
     });
 }
 /**
@@ -927,14 +153,7 @@ export function wrapClaudeAgentSDK(sdk, config) {
     }
     // Wrap the tool method if it exists
     if ("tool" in inputSdk && typeof inputSdk.tool === "function") {
-        const originalTool = inputSdk.tool;
-        wrappedSdk.tool = function (...args) {
-            const toolDef = originalTool.apply(sdk, args);
-            if (toolDef && typeof toolDef === "object" && "handler" in toolDef) {
-                return wrapClaudeAgentTool(toolDef, config);
-            }
-            return toolDef;
-        };
+        wrappedSdk.tool = inputSdk.tool.bind(inputSdk);
     }
     // Keep createSdkMcpServer and other methods as-is (bound to original SDK)
     if ("createSdkMcpServer" in inputSdk &&