"use strict";
Object.defineProperty(exports, "__esModule", { value: true });
exports.wrapManagedAgentSessionEvents = wrapManagedAgentSessionEvents;
const traceable_js_1 = require("../../traceable.cjs");
const run_trees_js_1 = require("../../run_trees.cjs");
const anthropic_js_1 = require("../anthropic.cjs");
const types_js_1 = require("../../utils/types.cjs");
function processManagedAgentStreamInputs(inputs) {
    const args = Array.isArray(inputs.args) ? inputs.args : [];
    const [sessionID, params] = args.length > 0 ? args : [inputs.input, undefined];
    return {
        session_id: sessionID,
        ...(params ? { stream_params: params } : {}),
    };
}
function getManagedAgentText(content) {
    if (!Array.isArray(content))
        return "";
    return content
        .map((block) => {
        if (typeof block === "object" &&
            block != null &&
            "type" in block &&
            block.type === "text" &&
            "text" in block &&
            typeof block.text === "string") {
            return block.text;
        }
        return "";
    })
        .join("");
}
function managedAgentSessionEventsAggregator(chunks) {
    const toolCalls = [];
    const toolResults = [];
    const errors = [];
    const modelRequests = [];
    const previews = new Map();
    let status;
    let stopReason;
    const usage = {
        input_tokens: 0,
        output_tokens: 0,
        cache_creation_input_tokens: 0,
        cache_read_input_tokens: 0,
    };
    for (const event of chunks) {
        switch (event.type) {
            case "agent.message": {
                if (typeof event.id === "string")
                    previews.delete(event.id);
                break;
            }
            case "agent.tool_use":
            case "agent.mcp_tool_use":
            case "agent.custom_tool_use": {
                toolCalls.push({ ...event });
                break;
            }
            case "agent.tool_result":
            case "agent.mcp_tool_result":
            case "user.custom_tool_result": {
                toolResults.push({ ...event });
                break;
            }
            case "span.model_request_end": {
                const modelUsage = event.model_usage;
                if (modelUsage) {
                    for (const key of Object.keys(usage)) {
                        const value = modelUsage[key];
                        if (typeof value === "number")
                            usage[key] += value;
                    }
                    modelRequests.push({
                        id: event.id,
                        is_error: event.is_error,
                        model_request_start_id: event.model_request_start_id,
                        model_usage: modelUsage,
                        processed_at: event.processed_at,
                    });
                }
                previews.clear();
                break;
            }
            case "session.error":
                errors.push({ ...event });
                break;
            case "session.status_running":
            case "session.status_idle":
            case "session.status_rescheduled":
            case "session.status_terminated":
                status = event.type;
                if (event.type === "session.status_idle") {
                    stopReason = event.stop_reason;
                }
                break;
        }
    }
    const unreconciledPreviews = [...previews.entries()].map(([eventID, byIndex]) => ({
        id: eventID,
        text: [...byIndex.entries()]
            .sort(([left], [right]) => left - right)
            .map(([, text]) => text)
            .join(""),
    }));
    const messageEvents = chunks.filter((event) => event.type === "agent.message");
    const lastMessage = getManagedAgentAssistantOutputMessages(messageEvents).at(-1);
    return {
        content: messageEvents.flatMap((message) => Array.isArray(message.content) ? message.content : []),
        ...(lastMessage ? { messages: [lastMessage] } : {}),
        tool_calls: toolCalls,
        tool_results: toolResults,
        errors,
        model_requests: modelRequests,
        status,
        stop_reason: stopReason,
        events: chunks,
        ...(unreconciledPreviews.length > 0
            ? { unreconciled_previews: unreconciledPreviews }
            : {}),
    };
}
function getManagedAgentInputEvents(chunks) {
    return chunks
        .filter((event) => typeof event.type === "string" &&
        (event.type.startsWith("user.") || event.type.startsWith("system.")))
        .map((event) => ({ ...event }));
}
function contentBlocksToChatContent(content) {
    if (!Array.isArray(content))
        return content;
    const text = getManagedAgentText(content);
    return text.length > 0 &&
        content.every((block) => {
            return (typeof block === "object" &&
                block != null &&
                "type" in block &&
                block.type === "text");
        })
        ? text
        : content;
}
function managedAgentToolUseToContentBlock(event) {
    return {
        type: "tool_use",
        id: event.id,
        name: event.name,
        input: event.input,
        ...(event.type === "agent.mcp_tool_use"
            ? { mcp_server_name: event.mcp_server_name }
            : {}),
    };
}
function getManagedAgentChatMessages(events) {
    const messages = [];
    for (const event of events) {
        switch (event.type) {
            case "user.message":
                messages.push({
                    role: "user",
                    content: contentBlocksToChatContent(event.content),
                });
                break;
            case "agent.message":
                messages.push({
                    role: "assistant",
                    content: contentBlocksToChatContent(event.content),
                });
                break;
            case "agent.tool_use":
            case "agent.mcp_tool_use":
            case "agent.custom_tool_use":
                messages.push({
                    role: "assistant",
                    content: [managedAgentToolUseToContentBlock(event)],
                });
                break;
            case "agent.tool_result":
                messages.push({
                    role: "tool",
                    tool_call_id: event.tool_use_id,
                    content: event.content,
                    ...(event.is_error != null ? { is_error: event.is_error } : {}),
                });
                break;
            case "agent.mcp_tool_result":
                messages.push({
                    role: "tool",
                    tool_call_id: event.mcp_tool_use_id,
                    content: event.content,
                    ...(event.is_error != null ? { is_error: event.is_error } : {}),
                });
                break;
            case "user.custom_tool_result":
                messages.push({
                    role: "tool",
                    tool_call_id: event.custom_tool_use_id,
                    content: event.content,
                    ...(event.is_error != null ? { is_error: event.is_error } : {}),
                });
                break;
            case "user.interrupt":
            case "user.define_outcome":
            case "user.tool_result":
            case "user.tool_confirmation":
                break;
        }
    }
    return messages;
}
function getManagedAgentAssistantOutputMessages(events) {
    return events.flatMap((event) => {
        if (event.type === "agent.message") {
            return [
                {
                    id: event.id,
                    role: "assistant",
                    content: contentBlocksToChatContent(event.content),
                    processed_at: event.processed_at,
                },
            ];
        }
        if (event.type === "agent.tool_use" ||
            event.type === "agent.mcp_tool_use" ||
            event.type === "agent.custom_tool_use") {
            return [
                {
                    id: event.id,
                    role: "assistant",
                    content: [managedAgentToolUseToContentBlock(event)],
                    processed_at: event.processed_at,
                },
            ];
        }
        return [];
    });
}
function getManagedAgentStreamError(chunks) {
    const errorEvent = chunks.find((event) => event.type === "session.error");
    const error = errorEvent?.error;
    if (typeof error === "object" && error != null && "message" in error) {
        return String(error.message);
    }
    return undefined;
}
function isCustomToolRequiresActionIdle(event, chunks) {
    if (event.type !== "session.status_idle" ||
        event.stop_reason.type !== "requires_action") {
        return false;
    }
    return event.stop_reason.event_ids.every((eventID) => chunks.some((candidate) => candidate.type === "agent.custom_tool_use" && candidate.id === eventID));
}
function getManagedAgentTurnError(chunks) {
    const idleEvent = [...chunks]
        .reverse()
        .find((event) => event.type === "session.status_idle");
    if (idleEvent?.type === "session.status_idle") {
        if (idleEvent.stop_reason.type === "requires_action") {
            return `Interrupted: requires_action (${idleEvent.stop_reason.event_ids.join(", ")})`;
        }
        if (idleEvent.stop_reason.type === "retries_exhausted") {
            return "Interrupted: retries_exhausted";
        }
    }
    return undefined;
}
function stripLangSmithExtraFromRequestOptions(options) {
    if (typeof options === "object" &&
        options != null &&
        "langsmithExtra" in options) {
        const { langsmithExtra: _langsmithExtra, ...rest } = options;
        return rest;
    }
    return options;
}
function getProcessedAtMillis(event) {
    if (!event)
        return undefined;
    return typeof event.processed_at === "string"
        ? Date.parse(event.processed_at)
        : undefined;
}
function getFirstProcessedAtMillis(events) {
    for (const event of events) {
        const processedAt = getProcessedAtMillis(event);
        if (processedAt !== undefined)
            return processedAt;
    }
    return undefined;
}
function getLastProcessedAtMillis(events) {
    for (const event of [...events].reverse()) {
        const processedAt = getProcessedAtMillis(event);
        if (processedAt !== undefined)
            return processedAt;
    }
    return undefined;
}
async function createManagedAgentChildRuns(parentRun, chunks, metadata, options) {
    const { modelConfig, systemPrompt, pendingToolUseEvents, completedChildRunIds, } = options;
    const modelName = typeof modelConfig?.id === "string" ? modelConfig.id : undefined;
    const modelRequestStarts = new Map();
    const childSpecs = [];
    const toolUseEvents = new Map();
    const toolResultEvents = new Map();
    // Aggregate tool use and tool results from all events
    // This is necessary because tool results may arrive after the model request ends (eg custom tools)
    for (const event of chunks.all) {
        if ((event.type === "agent.tool_use" ||
            event.type === "agent.mcp_tool_use" ||
            event.type === "agent.custom_tool_use") &&
            typeof event.id === "string") {
            toolUseEvents.set(event.id, event);
        }
        else if (event.type === "agent.tool_result" &&
            typeof event.tool_use_id === "string") {
            toolResultEvents.set(event.tool_use_id, event);
        }
        else if (event.type === "agent.mcp_tool_result" &&
            typeof event.mcp_tool_use_id === "string") {
            toolResultEvents.set(event.mcp_tool_use_id, event);
        }
        else if (event.type === "user.custom_tool_result" &&
            typeof event.custom_tool_use_id === "string") {
            toolResultEvents.set(event.custom_tool_use_id, event);
        }
    }
    // Handle turn specific events (model requests, tool calls)
    for (const event of chunks.turn) {
        if (event.type === "span.model_request_start" &&
            typeof event.id === "string") {
            modelRequestStarts.set(event.id, event);
        }
        else if (event.type === "span.model_request_end") {
            const startID = event.model_request_start_id;
            const startEvent = modelRequestStarts.get(startID);
            const startIndex = startEvent ? chunks.all.indexOf(startEvent) : -1;
            const endIndex = chunks.all.indexOf(event);
            const eventsBeforeRequest = startIndex >= 0 ? chunks.all.slice(0, startIndex) : [];
            const eventsInRequest = startIndex >= 0 && endIndex >= startIndex
                ? chunks.all.slice(startIndex, endIndex + 1)
                : [event];
            const messageEvents = eventsInRequest.filter((candidate) => candidate.type === "agent.message");
            const messages = getManagedAgentAssistantOutputMessages(eventsInRequest);
            const toolCalls = eventsInRequest.filter((candidate) => candidate.type === "agent.tool_use" ||
                candidate.type === "agent.mcp_tool_use" ||
                candidate.type === "agent.custom_tool_use");
            const content = messageEvents.flatMap((message) => Array.isArray(message.content) ? message.content : []);
            const usageMetadata = (0, anthropic_js_1.createUsageMetadata)(event.model_usage);
            const startTime = startEvent
                ? (getProcessedAtMillis(startEvent) ?? Date.now())
                : Date.now();
            childSpecs.push({
                startTime,
                index: startIndex >= 0 ? startIndex : endIndex,
                createAndPost: async () => {
                    const childRun = parentRun.createChild({
                        name: "ClaudeManagedAgentModelRequest",
                        run_type: "llm",
                        inputs: {
                            events: eventsBeforeRequest,
                            system: systemPrompt,
                            messages: getManagedAgentChatMessages(eventsBeforeRequest),
                            ...(startEvent ? { model_request_start: startEvent } : {}),
                        },
                        metadata: {
                            ...metadata,
                            ...(modelName ? { ls_model_name: modelName } : {}),
                            ...(modelConfig
                                ? {
                                    ls_invocation_params: {
                                        ...(metadata.ls_invocation_params ?? {}),
                                        model_config: modelConfig,
                                    },
                                }
                                : {}),
                            ...(usageMetadata ? { usage_metadata: usageMetadata } : {}),
                        },
                        start_time: startTime,
                    });
                    await childRun.end({
                        content,
                        messages,
                        tool_calls: toolCalls,
                        model_request_end: event,
                        ...(event.model_usage ? { model_usage: event.model_usage } : {}),
                        ...(usageMetadata ? { usage_metadata: usageMetadata } : {}),
                    }, event.is_error ? "Model request failed" : undefined, getProcessedAtMillis(event));
                    await childRun.postRun();
                },
            });
        }
    }
    for (const [toolUseID, toolUse] of toolUseEvents.entries()) {
        const result = toolResultEvents.get(toolUseID);
        const isCurrentTurnTool = chunks.turn.includes(toolUse) ||
            (result ? chunks.turn.includes(result) : false);
        const pendingToolUse = pendingToolUseEvents.get(toolUseID);
        if (completedChildRunIds.has(toolUseID) &&
            pendingToolUse == null &&
            !isCurrentTurnTool) {
            continue;
        }
        // The model requested a tool but execution is blocked on user confirmation
        // or an external custom-tool result. Do not create a tool child yet; the
        // execution belongs to the later turn that receives the result.
        if (result == null) {
            if (chunks.turn.includes(toolUse) &&
                !completedChildRunIds.has(toolUseID)) {
                pendingToolUseEvents.set(toolUseID, toolUse);
            }
            continue;
        }
        if (!isCurrentTurnTool && pendingToolUse == null) {
            continue;
        }
        const toolUseForRun = pendingToolUse ?? toolUse;
        const toolName = typeof toolUseForRun.name === "string"
            ? toolUseForRun.name
            : "ManagedAgentTool";
        const startTime = pendingToolUse != null
            ? (getProcessedAtMillis(result) ?? Date.now())
            : (getProcessedAtMillis(toolUseForRun) ?? Date.now());
        const index = pendingToolUse != null
            ? chunks.all.indexOf(result)
            : chunks.all.indexOf(toolUseForRun);
        childSpecs.push({
            startTime,
            index,
            createAndPost: async () => {
                const resultArgs = [
                    { event: result, content: result.content },
                    result.is_error
                        ? (getManagedAgentText(result.content) ?? "Tool execution failed")
                        : undefined,
                    getProcessedAtMillis(result),
                ];
                const childRun = parentRun.createChild({
                    name: toolName,
                    run_type: "tool",
                    inputs: {
                        id: toolUseID,
                        name: toolUseForRun.name,
                        input: toolUseForRun.input,
                        event: toolUseForRun,
                    },
                    metadata,
                    start_time: startTime,
                });
                await childRun.end(...resultArgs);
                await childRun.postRun();
                pendingToolUseEvents.delete(toolUseID);
                completedChildRunIds.add(toolUseID);
            },
        });
    }
    childSpecs.sort((left, right) => left.startTime - right.startTime || left.index - right.index);
    for (const childSpec of childSpecs) {
        await childSpec.createAndPost();
    }
}
function wrapManagedAgentSessionEvents({ originalBeta, tracedBeta, cleanedOptions, prepopulatedInvocationParams, }) {
    const createManagedAgentStreamRun = (sessionID, params, requestOptions, startTime) => {
        const runtimeConfig = requestOptions?.langsmithExtra ?? {};
        const parentRun = (0, traceable_js_1.getCurrentRunTree)(true);
        const runConfig = {
            ...cleanedOptions,
            ...runtimeConfig,
            name: runtimeConfig.name ?? cleanedOptions.name ?? "ClaudeManagedAgent",
            run_type: "chain",
            inputs: { session_id: sessionID },
            ...(startTime ? { start_time: startTime } : {}),
            tags: [
                ...new Set([
                    ...(cleanedOptions.tags ?? []),
                    ...(runtimeConfig.tags ?? []),
                ]),
            ],
            metadata: {
                ...cleanedOptions.metadata,
                ...runtimeConfig.metadata,
                ls_provider: "anthropic",
                ls_model_type: "chat",
                thread_id: sessionID,
                ls_invocation_params: {
                    ...(typeof prepopulatedInvocationParams === "object" &&
                        prepopulatedInvocationParams != null
                        ? prepopulatedInvocationParams
                        : {}),
                    session_id: sessionID,
                    ...(params ?? {}),
                },
            },
        };
        return parentRun
            ? parentRun.createChild(runConfig)
            : new run_trees_js_1.RunTree(runConfig);
    };
    // Wrap Claude Managed Agents session event methods if they exist.
    if (originalBeta.sessions?.events) {
        tracedBeta.sessions = Object.create(Object.getPrototypeOf(tracedBeta.sessions));
        Object.assign(tracedBeta.sessions, originalBeta.sessions);
        tracedBeta.sessions.events = Object.create(Object.getPrototypeOf(tracedBeta.sessions.events));
        Object.assign(tracedBeta.sessions.events, originalBeta.sessions.events);
        if (typeof originalBeta.sessions.events.stream === "function") {
            // eslint-disable-next-line @typescript-eslint/no-explicit-any
            const streamManagedAgentEvents = async function (...args) {
                const [sessionID, params, requestOptions] = args;
                const sanitizedArgs = [...args];
                sanitizedArgs[2] = stripLangSmithExtraFromRequestOptions(sanitizedArgs[2]);
                while (sanitizedArgs.length > 1 &&
                    sanitizedArgs[sanitizedArgs.length - 1] === undefined) {
                    sanitizedArgs.pop();
                }
                const sessionPromise = originalBeta.sessions?.retrieve
                    ? originalBeta.sessions.retrieve
                        .bind(originalBeta.sessions)(sessionID)
                        .catch(() => undefined)
                    : Promise.resolve(undefined);
                const stream = await originalBeta.sessions?.events?.stream.bind(originalBeta.sessions.events)(...sanitizedArgs);
                const iterator = stream[Symbol.asyncIterator]();
                const allChunks = [];
                const turnChunkLengths = new Set();
                const pendingToolUseEvents = new Map();
                const completedChildRunIds = new Set();
                let flushLength = 0;
                let observedTerminalEvent = false;
                const finalize = async (kind, userError) => {
                    const chunks = allChunks.slice(flushLength);
                    flushLength = allChunks.length;
                    if (chunks.length > 0)
                        turnChunkLengths.add(flushLength);
                    // Question: should we re-attempt to fetch when finalizing again?
                    // ie if the agent's system prompt has changed mid-stream
                    const session = await sessionPromise;
                    const modelConfig = (0, types_js_1.isRecord)(session?.agent?.model)
                        ? session.agent.model
                        : undefined;
                    const systemPrompt = session?.agent.system ?? undefined;
                    if (chunks.length > 0) {
                        const turnStartTime = getFirstProcessedAtMillis(chunks) ?? Date.now();
                        const rawTurnEndTime = getLastProcessedAtMillis(chunks);
                        const turnEndTime = rawTurnEndTime !== undefined && rawTurnEndTime >= turnStartTime
                            ? rawTurnEndTime
                            : turnStartTime;
                        const runTree = createManagedAgentStreamRun(sessionID, params, requestOptions, turnStartTime);
                        const inputEvents = getManagedAgentInputEvents(chunks);
                        runTree.inputs = {
                            ...processManagedAgentStreamInputs({
                                args: [sessionID, params, sanitizedArgs[2]],
                            }),
                            ...(inputEvents.length > 0
                                ? {
                                    events: inputEvents,
                                    messages: getManagedAgentChatMessages(inputEvents),
                                }
                                : {}),
                        };
                        const turnError = userError ??
                            getManagedAgentStreamError(chunks) ??
                            getManagedAgentTurnError(chunks);
                        await runTree.end(managedAgentSessionEventsAggregator(chunks), turnError, turnEndTime);
                        await runTree.postRun();
                        await createManagedAgentChildRuns(runTree, { turn: chunks, all: allChunks }, runTree.extra.metadata ?? {}, {
                            modelConfig,
                            systemPrompt,
                            pendingToolUseEvents,
                            completedChildRunIds,
                        });
                    }
                    if (kind === "done" || kind === "throw") {
                        // Pending tool uses without results were never executed, so no tool
                        // child run was created. Clear them when the stream ends.
                        pendingToolUseEvents.clear();
                    }
                };
                stream[Symbol.asyncIterator] = () => ({
                    async next() {
                        try {
                            const result = await iterator.next();
                            if (result.done) {
                                await finalize("done");
                                return result;
                            }
                            allChunks.push(result.value);
                            if (result.value.type === "session.status_idle" ||
                                result.value.type === "session.status_terminated" ||
                                result.value.type === "session.deleted" ||
                                result.value.type === "session.error") {
                                if (!isCustomToolRequiresActionIdle(result.value, allChunks)) {
                                    observedTerminalEvent = true;
                                    await finalize("flush");
                                }
                            }
                            return result;
                        }
                        catch (error) {
                            await finalize("throw", String(error));
                            throw error;
                        }
                    },
                    async return(value) {
                        if (allChunks.length) {
                            await iterator.return?.(value);
                            if (observedTerminalEvent) {
                                await finalize("done");
                            }
                            else {
                                await finalize("throw", "Cancelled");
                            }
                        }
                        else {
                            await iterator.return?.(value);
                        }
                        return { done: true, value };
                    },
                    async throw(error) {
                        await finalize("throw", String(error));
                        if (iterator.throw)
                            return iterator.throw(error);
                        throw error;
                    },
                });
                return stream;
            };
            Object.defineProperty(streamManagedAgentEvents, "langsmith:traceable", {
                value: { name: "ClaudeManagedAgent", run_type: "chain" },
            });
            tracedBeta.sessions.events.stream = streamManagedAgentEvents;
        }
        if (typeof originalBeta.sessions.events.send === "function") {
            // Do not trace `send` as a separate run. The API returns the submitted
            // user events, which makes the run output look like a duplicate of the
            // input. Instead, annotate the active stream run for the same session.
            // eslint-disable-next-line @typescript-eslint/no-explicit-any
            tracedBeta.sessions.events.send = function sendManagedAgentEvents(...args) {
                const sanitizedArgs = [...args];
                if (typeof sanitizedArgs[2] === "object" &&
                    sanitizedArgs[2] != null &&
                    "langsmithExtra" in sanitizedArgs[2]) {
                    const { langsmithExtra: _langsmithExtra, ...rest } = sanitizedArgs[2];
                    sanitizedArgs[2] = rest;
                }
                return originalBeta.sessions?.events?.send.bind(originalBeta.sessions.events)(...sanitizedArgs);
            };
        }
    }
}
