windmill/frontend/src/lib/components/agentTrace.ts

import type { WebSearchSource } from './copilot/chat/shared'
import type { AgentMessage } from './aiAgentResult'

/**
 * One entry in the trace of an agent run, built from the envelope alone so it
 * renders without waiting on any request; a tool's child job is enrichment. The
 * prompt and the question are deliberately absent: they are the step's inputs and
 * are shown as inputs, while the trace is what the agent did with them.
 */
export type AgentTraceEntry =
	| { kind: 'assistant'; content: string }
	/** The worker records only that a search happened. Its citations ride on the assistant
	 *  turn that follows, and are moved here so the search shows what it found. */
	| { kind: 'search'; sources?: WebSearchSource[] }
	| {
			kind: 'tool'
			name: string
			args?: string
			result: string
			/** Present for a flow-module tool, which runs as its own job. */
			jobId?: string
			/** An MCP tool runs in the worker, so it names its server instead. */
			resourcePath?: string
	  }

/** `content` is a string for text messages and a part list once images are involved. */
function contentText(content: unknown): string {
	if (typeof content === 'string') {
		return content
	}
	if (Array.isArray(content)) {
		return content
			.map((part) =>
				part && typeof part === 'object' && typeof (part as { text?: unknown }).text === 'string'
					? (part as { text: string }).text
					: ''
			)
			.join('')
	}
	return ''
}

function sourcesOf(message: AgentMessage): WebSearchSource[] | undefined {
	const annotations = message.annotations
	if (!annotations?.length) {
		return undefined
	}
	return annotations.map((a) => ({ url: a.url, title: a.title }))
}

export function buildAgentTrace(messages: AgentMessage[]): AgentTraceEntry[] {
	// The arguments live on the assistant message that requested the call, while
	// the action tag and the result live on the `tool` message answering it, so
	// the two are joined by `tool_call_id`.
	const argsByCallId = new Map<string, string>()
	for (const message of messages) {
		for (const call of message.tool_calls ?? []) {
			if (call.id && typeof call.function?.arguments === 'string') {
				argsByCallId.set(call.id, call.function.arguments)
			}
		}
	}

	const entries: AgentTraceEntry[] = []
	for (const message of messages) {
		const action = message.agent_action
		if (action?.type === 'tool_call') {
			entries.push({
				kind: 'tool',
				name: action.function_name,
				args: message.tool_call_id ? argsByCallId.get(message.tool_call_id) : undefined,
				result: contentText(message.content),
				jobId: action.job_id
			})
			continue
		}
		if (action?.type === 'mcp_tool_call') {
			entries.push({
				kind: 'tool',
				name: action.function_name,
				// An MCP call records its arguments on the action itself: it never
				// became a job, so there is nowhere else for them to live.
				args: action.arguments ? JSON.stringify(action.arguments, null, 2) : undefined,
				result: contentText(message.content),
				resourcePath: action.resource_path
			})
			continue
		}
		if (action?.type === 'web_search') {
			entries.push({ kind: 'search' })
			continue
		}
		// Every message this run produced is tagged, including the agent narrating
		// its next move and its final answer. Untagged ones are the prompt, the
		// question, or a previous turn replayed out of memory — history loses its
		// tags on the way back, and attributing it to this run would credit it with
		// answers it never gave.
		if (action?.type !== 'message') {
			continue
		}
		const content = contentText(message.content)
		// The worker writes a provider call's search right before that call's answer, and
		// only a call that searched carries annotations.
		const previous = entries[entries.length - 1]
		if (message.role === 'assistant' && previous?.kind === 'search' && !previous.sources) {
			previous.sources = sourcesOf(message)
		}
		if (message.role === 'assistant' && content !== '') {
			entries.push({ kind: 'assistant', content })
		}
	}
	return entries
}

/**
 * Separates the turn that produced the output from the rest of the trace, so the
 * answer is rendered once. Found by content and searched from the end: a run whose
 * last turn returned a tool call leaves its answer mid-trace, where inspecting only
 * the final entry prints it twice.
 */
export function splitFinalAnswer(entries: AgentTraceEntry[], output: unknown): AgentTraceEntry[] {
	if (typeof output !== 'string') {
		// A schema-shaped output is rendered by the result viewer itself and matches
		// no turn, so the whole trace stands.
		return entries
	}
	for (let i = entries.length - 1; i >= 0; i--) {
		const entry = entries[i]
		if (entry.kind === 'assistant' && entry.content === output) {
			return [...entries.slice(0, i), ...entries.slice(i + 1)]
		}
	}
	return entries
}