import { APICallError } from "ai";
import { safeJsonStringify } from "../../../../util/log-sanitizer.js";
import { getJwtTraceTags } from "../../../llmobs/helpers.js";
import { stripLengthyContent } from "../../../util/strip-content.js";
import { getTotalInputTokens } from "../../context-v2/token-tracker.js";
import { LLMProviderError } from "../../error.js";
export function classifyError(error) {
    if (LLMProviderError.isInstance(error)) {
        return {
            category: error.type,
            code: error.statusCode,
            isRetryable: error.isRetryable,
            message: error.message,
        };
    }
    if (APICallError.isInstance(error)) {
        return {
            category: mapStatusCodeToCategory(error.statusCode),
            code: error.statusCode,
            isRetryable: error.isRetryable,
            message: error.message,
        };
    }
    if (error instanceof Error) {
        if (error.name === "AbortError" || error.name === "AbortedError") {
            return {
                category: "aborted",
                isRetryable: false,
                message: error.message,
            };
        }
        if (error.message.includes("ECONNREFUSED") ||
            error.message.includes("ETIMEDOUT") ||
            error.message.includes("ENOTFOUND") ||
            error.message.includes("fetch failed")) {
            return {
                category: "network",
                isRetryable: true,
                message: error.message,
            };
        }
        if (error.message.includes("timeout") ||
            error.message.includes("Timeout")) {
            return {
                category: "timeout",
                isRetryable: true,
                message: error.message,
            };
        }
    }
    return {
        category: "unknown",
        isRetryable: true,
        message: error instanceof Error ? error.message : String(error),
    };
}
function mapStatusCodeToCategory(statusCode) {
    if (!statusCode)
        return "unknown";
    if (statusCode === 401 || statusCode === 403)
        return "auth_error";
    if (statusCode === 429)
        return "rate_limit";
    if (statusCode === 408)
        return "timeout";
    if (statusCode >= 500 && statusCode < 600)
        return "api_error";
    if (statusCode >= 400 && statusCode < 500)
        return "validation_error";
    return "unknown";
}
const THINKING_SLOW_WARNING_MS = 15_000;
/**
 * Observer that manages LLM Observability tracing spans.
 *
 * Creates and manages:
 * - A full workflow span (`ai.streamText`) for the entire session
 * - Step spans (`ai.streamText.step`) for each LLM step
 * - Segment spans (`ai.streamText.{type}`) for thinking/reasoning/text/tool-input
 * - Annotations with usage metrics, model info, and input/output data
 *
 * @example
 * ```typescript
 * const observer = new LLMObsObserver({
 *   tracer: LLMObsTracer.instance,
 *   logger: getLogger(),
 * });
 * ```
 */
export class LLMObsObserver {
    tracer;
    logger;
    fullSpan;
    stepSpan;
    currentSegment;
    stepNumber = 0;
    tools;
    session;
    messages;
    constructor(options) {
        this.tracer = options.tracer;
        this.logger = options.logger;
    }
    attach(events, _config) {
        events.on("session:start", (session, request) => {
            this.onSessionStart(session, request);
        });
        events.on("chunk", (chunk, _session) => {
            this.onChunk(chunk);
        });
        events.on("step:finish", (step, session) => {
            this.onStepComplete(step, session);
        });
        events.on("stream:finish", (result, session) => {
            this.onStreamFinish(result, session);
        });
        events.on("retry", (error, attempt, _delayMs, session) => {
            this.onRetry(error, attempt, session);
        });
        events.on("fallback", (from, to, reason, session) => {
            this.onFallback(from, to, reason, session);
        });
        events.on("session:complete", (session) => {
            this.onSessionComplete(session);
        });
        events.on("session:error", (error, session) => {
            this.onSessionError(error, session);
        });
        events.on("session:abort", (session) => {
            this.onSessionAbort(session);
        });
    }
    onSessionStart(session, request) {
        this.session = session;
        this.fullSpan = this.tracer.startActiveSpan({
            kind: "workflow",
            name: "ai.streamText",
        });
        this.tools = request?.tools;
        const systemMessages = request?.system
            ? [{ role: "system", content: request.system }]
            : [];
        const allMessages = [
            ...systemMessages,
            ...(request?.messages ?? []),
        ];
        // we keep a running list of messages for the current session, this helps us track the costs in datadog
        this.messages = formatMessages(allMessages, this.tools);
        if (this.fullSpan) {
            this.tracer.annotate(this.fullSpan, {
                metadata: {
                    sessionId: session.id,
                    provider: session.provider.name,
                    model: session.provider.model,
                },
            });
        }
        this.stepNumber = 0;
        this.startStepSpan(this.messages);
        this.startSegment("thinking");
        this.logger.debug(`[llmobs] Started session span: ${session.id}`);
    }
    onChunk(chunk) {
        if (chunk.type === "tool-input-start") {
            // this is so we can group everything going into a tool call
            this.startSegment("tool", {
                toolCallId: chunk.id,
                tool: chunk.toolName,
                input: "",
            }, 1);
            return;
        }
        if (chunk.type === "tool-input-delta") {
            this.markSegmentChunk();
            this.appendToolInputDelta(chunk.delta);
            return;
        }
        const segmentType = this.getSegmentTypeFromChunk(chunk.type);
        if (segmentType) {
            this.ensureSegment(segmentType);
            this.markSegmentChunk();
            if (this.isTextualSegment(segmentType) && "text" in chunk) {
                this.appendTextToCurrentSegment(chunk.text);
            }
        }
    }
    onStepComplete(step, _session) {
        this.messages = this.messages?.concat(formatOutputData(step.content, this.tools) ?? []);
        this.finishSegment();
        // Note: Tool spans are created by wrapToolsForObservability() during actual
        // execution, not here. We only track step-level info.
        this.finishStepSpan(step);
        this.stepNumber++;
        this.startStepSpan(this.messages);
        this.startSegment("thinking");
        this.logger.debug(`[llmobs] Step ${this.stepNumber} complete: finishReason=${step.finishReason}`);
    }
    onStreamFinish(result, session) {
        this.finishSegment();
        this.finishStepSpan(result);
        if (this.fullSpan) {
            const metrics = this.compileMetrics(result.steps);
            // Format outputData from the final result
            const outputMessages = formatOutputData(result.content, this.tools);
            const lastMessage = outputMessages.at(-1);
            this.tracer.annotate(this.fullSpan, {
                outputData: lastMessage?.content,
                metrics,
                metadata: {
                    totalSteps: result.steps.length,
                    finishReason: result.finishReason,
                    durationMs: session.getElapsedMs(),
                },
            });
        }
        this.logger.debug(`[llmobs] Stream finished: steps=${result.steps.length}`);
    }
    onRetry(error, attempt, session) {
        if (this.fullSpan) {
            const classified = classifyError(error);
            const retryMetadata = {
                error: classified.message,
                errorCategory: classified.category,
                errorCode: classified.code,
                isRetryable: classified.isRetryable,
            };
            this.tracer.annotate(this.fullSpan, {
                model: {
                    name: session.provider.model,
                    provider: session.provider.name,
                },
                tags: {
                    hasRetries: true,
                },
                metadata: {
                    [`retry_${attempt}`]: retryMetadata,
                },
            });
        }
    }
    onFallback(from, to, reason, _session) {
        if (this.fullSpan) {
            const fallbackMetadata = {
                fallbackReason: reason,
                fromProvider: from.name,
                fromModel: from.model,
                toProvider: to.name,
                toModel: to.model,
            };
            this.tracer.annotate(this.fullSpan, {
                model: {
                    name: to.model,
                    provider: to.name,
                },
                tags: {
                    hasFallback: true,
                },
                metadata: fallbackMetadata,
            });
        }
    }
    onSessionComplete(session) {
        this.finishFullSpan(session, "complete");
    }
    onSessionError(error, session) {
        if (this.fullSpan) {
            const classified = classifyError(error);
            const errorMetadata = {
                error: classified.message,
                errorCategory: classified.category,
                errorCode: classified.code,
                isRetryable: classified.isRetryable,
            };
            this.tracer.annotate(this.fullSpan, {
                tags: {
                    hasError: true,
                    errorCategory: classified.category,
                },
                metadata: errorMetadata,
            });
        }
        this.finishFullSpan(session, "error");
    }
    onSessionAbort(session) {
        if (this.fullSpan) {
            this.tracer.annotate(this.fullSpan, {
                tags: {
                    wasAborted: true,
                },
            });
        }
        this.finishFullSpan(session, "aborted");
    }
    startSegment(type, initialContent, initialChunkCount = 0) {
        if (this.currentSegment) {
            this.finishSegment();
        }
        const span = this.tracer.startActiveSpan({
            kind: "task",
            name: `ai.streamText.${type}`,
        });
        this.currentSegment = {
            type,
            span,
            content: initialContent,
            startedAtMs: Date.now(),
            chunkCount: initialChunkCount,
        };
    }
    finishSegment() {
        const segment = this.currentSegment;
        if (!segment)
            return;
        const durationMs = Date.now() - segment.startedAtMs;
        // Parse tool-input JSON if possible
        if (segment.type === "tool-input" && typeof segment.content === "object") {
            const content = segment.content;
            try {
                content.input = JSON.parse(content.input);
            }
            catch {
                // leave as string if not valid JSON
            }
        }
        if (segment.type === "thinking") {
            const hasContent = typeof segment.content === "string"
                ? segment.content.trim().length > 0
                : !!segment.content && Object.keys(segment.content).length > 0;
            this.tracer.annotate(segment.span, {
                outputData: segment.content,
                metadata: {
                    durationMs,
                    chunkCount: segment.chunkCount,
                    hasContent,
                    stepNumber: this.stepNumber,
                },
            });
            if (durationMs >= THINKING_SLOW_WARNING_MS) {
                this.logger.warn("[llmobs] Slow thinking segment detected", {
                    durationMs,
                    thresholdMs: THINKING_SLOW_WARNING_MS,
                    stepNumber: this.stepNumber,
                    chunkCount: segment.chunkCount,
                    hasContent,
                    sessionId: this.session?.id,
                });
            }
        }
        else {
            this.tracer.annotate(segment.span, {
                outputData: segment.content,
            });
        }
        segment.span.finish();
        this.currentSegment = undefined;
    }
    ensureSegment(type) {
        if (this.currentSegment?.type !== type) {
            this.startSegment(type);
        }
    }
    appendTextToCurrentSegment(text) {
        if (!this.currentSegment)
            return;
        const existing = typeof this.currentSegment.content === "string"
            ? this.currentSegment.content
            : "";
        this.currentSegment.content = `${existing}${text}`;
    }
    appendToolInputDelta(delta) {
        if (!this.currentSegment || this.currentSegment.type !== "tool-input")
            return;
        const content = this.currentSegment.content;
        content.input += delta;
    }
    markSegmentChunk() {
        if (!this.currentSegment)
            return;
        this.currentSegment.chunkCount += 1;
    }
    isTextualSegment(type) {
        return (type === "text" ||
            type === "reasoning" ||
            type === "thinking" ||
            type === "raw" ||
            type === "source");
    }
    getSegmentTypeFromChunk(chunkType) {
        const suffix = "-delta";
        if (!chunkType.endsWith(suffix))
            return undefined;
        const base = chunkType.slice(0, -suffix.length);
        const admissible = [
            "text",
            "reasoning",
            "thinking",
            "raw",
            "source",
            "tool-call",
            "tool-result",
            "tool-input",
        ];
        return admissible.includes(base) ? base : undefined;
    }
    startStepSpan(inputMessages) {
        this.stepSpan = this.tracer.startActiveSpan({
            kind: "llm",
            name: "ai.streamText.step",
        });
        if (this.stepSpan) {
            const jwt = this.session?.getJwt?.();
            const jwtTags = getJwtTraceTags(jwt);
            this.tracer.annotate(this.stepSpan, {
                inputData: inputMessages ?? [],
                tags: { stepNumber: this.stepNumber, ...jwtTags },
            });
        }
    }
    finishStepSpan(step) {
        if (!this.stepSpan)
            return;
        const metrics = {};
        if (step.usage) {
            metrics.inputTokens = getTotalInputTokens(step.usage);
            metrics.outputTokens = step.usage.outputTokens ?? 0;
            metrics.totalTokens = step.usage.totalTokens ?? 0;
            metrics.reasoningTokens = step.usage.reasoningTokens ?? 0;
            metrics.cacheReadTokens = step.usage.cachedInputTokens ?? 0;
            metrics.cacheWriteTokens =
                step.usage.inputTokenDetails?.cacheWriteTokens ?? 0;
        }
        metrics.toolCallCount = step.toolCalls?.length ?? 0;
        // Extract model info from response headers if available
        const modelId = step.response?.headers?.["x-model-id"] ?? step.response?.modelId;
        const provider = step.response?.headers?.["x-model-provider"];
        // Format outputData from step content
        const outputData = formatOutputData(step.content, this.tools);
        this.tracer.annotate(this.stepSpan, {
            outputData: stripLengthyContent(outputData),
            metrics: metrics,
            metadata: {
                finishReason: step.finishReason,
                toolCalls: step.toolCalls?.map((tc) => tc.toolName),
                usage: step.usage,
            },
            model: modelId && provider
                ? {
                    name: modelId,
                    provider,
                }
                : undefined,
        });
        this.stepSpan.finish();
        this.stepSpan = undefined;
    }
    finishFullSpan(session, outcome) {
        this.finishSegment();
        if (this.stepSpan) {
            this.stepSpan.finish();
            this.stepSpan = undefined;
        }
        if (this.tracer.rootSpan && session.metrics.timeToFirstTokenMs) {
            this.tracer.annotate(this.tracer.rootSpan, {
                metrics: {
                    timeToFirstToken: session.metrics.timeToFirstTokenMs / 1000,
                },
            });
        }
        if (this.fullSpan) {
            this.tracer.annotate(this.fullSpan, {
                metadata: {
                    outcome,
                    totalDurationMs: session.metrics.durationMs,
                    totalChunks: session.metrics.chunkCount,
                    totalRetries: session.metrics.retryCount,
                    totalRetryDelayMs: session.metrics.totalRetryDelayMs,
                    timeToFirstTokenMs: session.metrics.timeToFirstTokenMs,
                },
            });
            this.fullSpan.finish();
            this.fullSpan = undefined;
        }
        this.session = undefined;
        this.logger.debug(`[llmobs] Session ${session.id} finished: outcome=${outcome}`);
    }
    compileMetrics(steps) {
        let inputTokens = 0;
        let outputTokens = 0;
        let totalTokens = 0;
        let reasoningTokens = 0;
        let cacheReadTokens = 0;
        let cacheWriteTokens = 0;
        let toolCallCount = 0;
        let staticToolCallCount = 0;
        let dynamicToolCallCount = 0;
        let filesGeneratedCount = 0;
        for (const step of steps) {
            const usage = step.usage;
            inputTokens += getTotalInputTokens(usage);
            outputTokens += usage.outputTokens ?? 0;
            totalTokens += usage.totalTokens ?? 0;
            reasoningTokens += usage.reasoningTokens ?? 0;
            cacheReadTokens += usage.cachedInputTokens ?? 0;
            cacheWriteTokens += usage.inputTokenDetails?.cacheWriteTokens ?? 0;
            toolCallCount += step.toolCalls?.length ?? 0;
            staticToolCallCount += step.staticToolCalls?.length ?? 0;
            dynamicToolCallCount += step.dynamicToolCalls?.length ?? 0;
            filesGeneratedCount += step.files?.length ?? 0;
        }
        return {
            inputTokens,
            outputTokens,
            totalTokens,
            reasoningTokens,
            cacheReadTokens,
            cacheWriteTokens,
            toolCallCount,
            staticToolCallCount,
            dynamicToolCallCount,
            filesGeneratedCount,
            stepCount: steps.length,
        };
    }
}
/**
 * Converts ModelMessage[] to LLMObsMessage[] format.
 */
function formatMessages(messages, toolSet) {
    const result = [];
    for (const { role, content } of messages) {
        if (role === "system") {
            result.push({
                role: "system",
                content: typeof content === "string" ? content : safeJsonStringify(content),
            });
        }
        else if (role === "user") {
            if (!Array.isArray(content)) {
                result.push({ role: "user", content: content });
            }
            else {
                // Format content parts, then force role: "user" on all results
                for (const msg of formatContent(content, toolSet)) {
                    result.push({ ...msg, role: "user" });
                }
            }
        }
        else if (role === "assistant") {
            if (!Array.isArray(content)) {
                result.push({ role: "assistant", content: content });
            }
            else {
                // Let formatContent determine roles (text, reasoning, tool, etc.)
                for (const msg of formatContent(content, toolSet)) {
                    result.push(msg);
                }
            }
        }
        else if (role === "tool") {
            for (const msg of formatContent(content, toolSet)) {
                result.push(msg);
            }
        }
    }
    return result;
}
/**
 * Converts StepResult.content to LLMObsMessage[] format.
 */
function formatOutputData(parts, toolSet) {
    return formatContent(parts, toolSet);
}
/**
 * Formats content parts into LLMObsMessage format.
 */
function formatContent(parts, toolSet) {
    const messages = [];
    for (const part of parts) {
        if (!part || typeof part !== "object")
            continue;
        const { type } = part;
        const p = part;
        if (type === "text") {
            messages.push({ role: "text", content: String(p.text ?? "") });
        }
        else if (type === "reasoning") {
            messages.push({ role: "reasoning", content: String(p.text ?? "") });
        }
        else if (type === "tool-call") {
            // Group consecutive tool calls into one message
            const last = messages.at(-1);
            const message = last?.toolCalls !== undefined
                ? last
                : { role: "assistant", content: "" };
            if (!message.toolCalls) {
                message.toolCalls = [];
            }
            const toolName = String(p.toolName ?? "");
            const [callType, name] = toolName.split("_", 2);
            message.toolCalls.push({
                name: name ?? toolName,
                arguments: asDictionary(p.args) ?? asDictionary(p.input) ?? {},
                toolId: String(p.toolCallId ?? ""),
                type: callType,
            });
            if (last !== message) {
                messages.push(message);
            }
        }
        else if (type === "tool-result") {
            const toolName = String(p.toolName ?? "");
            const tool = toolSet?.[toolName];
            const output = tool?.toModelOutput
                ? tool.toModelOutput({
                    toolCallId: String(p.toolCallId ?? ""),
                    input: asDictionary(p.args) ?? asDictionary(p.input) ?? {},
                    output: p.result ?? p.output,
                })
                : (p.result ?? p.output);
            messages.push({
                role: "tool",
                toolId: String(p.toolCallId ?? ""),
                content: stringifyToolOutput(output),
            });
        }
        else if (type === "tool-error") {
            messages.push({
                role: "tool",
                toolId: String(p.toolCallId ?? ""),
                content: safeJsonStringify({ error: p.error }),
            });
        }
        else if (type === "file") {
            messages.push({
                role: "file",
                content: "Generated file",
            });
        }
        else if (type === "source") {
            messages.push({
                role: "source",
                content: `Retrieved source: ${String(p.title ?? "unknown")}`,
            });
        }
        else if (type === "image") {
            messages.push({
                role: "image",
                content: "Generated image",
            });
        }
    }
    return messages;
}
function stringifyToolOutput(output) {
    if (typeof output !== "object" || output === null || !("type" in output)) {
        return typeof output === "string" ? output : safeJsonStringify(output);
    }
    const typed = output;
    switch (typed.type) {
        case "text":
            return typed.value;
        case "error-text":
            return `ERROR: ${typed.value}`;
        case "json":
            return safeJsonStringify(typed.value);
        case "error-json":
            return safeJsonStringify({ error: typed.value });
        case "content":
            return typed.value
                .filter((v) => v.type === "text")
                .map((v) => v.text ?? "")
                .join();
        default:
            return safeJsonStringify(output);
    }
}
function asDictionary(value) {
    if (value && typeof value === "object" && !Array.isArray(value)) {
        return value;
    }
    return undefined;
}
/**
 * Wraps tools with LLMObs tracing for real-time tool execution observability.
 *
 * This utility intercepts each tool's `execute` function and wraps it with
 * a traced span, capturing:
 * - Tool name (`tool.{name}`)
 * - Input arguments
 * - Output result
 * - Execution duration
 * - Error status
 */
export function wrapToolsForObservability(tools, tracer) {
    return Object.fromEntries(Object.entries(tools).map(([name, tool]) => [
        name,
        {
            ...tool,
            async execute(input, options) {
                return await tracer.trace({
                    kind: "tool",
                    name: `tool.${tool.name}`,
                }, async (span) => {
                    try {
                        const result = await tool.execute(input, options);
                        tracer.annotate(span, {
                            inputData: typeof input === "object"
                                ? input
                                : undefined,
                            outputData: typeof result === "object"
                                ? stripLengthyContent(result)
                                : typeof result === "string"
                                    ? stripLengthyContent(result)
                                    : undefined,
                            metadata: {
                                toolCallId: options.toolCallId,
                            },
                        });
                        return result;
                    }
                    catch (error) {
                        tracer.annotate(span, {
                            inputData: typeof input === "object"
                                ? input
                                : undefined,
                            metadata: {
                                toolCallId: options.toolCallId,
                                isError: true,
                                error: error instanceof Error ? error.message : String(error),
                            },
                        });
                        throw error;
                    }
                });
            },
        },
    ]));
}
//# sourceMappingURL=llmobs.js.map