Codewalkers/apps/web/src/lib/parse-agent-output.ts

export interface ParsedMessage {
  type:
    | "text"
    | "system"
    | "tool_call"
    | "tool_result"
    | "session_end"
    | "error";
  content: string;
  timestamp?: Date;
  meta?: {
    toolName?: string;
    toolInput?: unknown;
    isError?: boolean;
    cost?: number;
    duration?: number;
  };
}

export function formatToolCall(toolUse: any): string {
  const { name, input } = toolUse;

  if (name === "Bash") {
    return `$ ${input.command}${input.description ? "\n# " + input.description : ""}`;
  }

  if (name === "Read") {
    return `Read: ${input.file_path}${input.offset ? ` (lines ${input.offset}-${input.offset + (input.limit || 10)})` : ""}`;
  }

  if (name === "Edit") {
    return `Edit: ${input.file_path}\n${input.old_string.substring(0, 100)}${input.old_string.length > 100 ? "..." : ""}\n-> ${input.new_string.substring(0, 100)}${input.new_string.length > 100 ? "..." : ""}`;
  }

  if (name === "Write") {
    return `Write: ${input.file_path} (${input.content.length} chars)`;
  }

  if (name === "Task") {
    return `${input.subagent_type}: ${input.description}\n${input.prompt?.substring(0, 200)}${input.prompt && input.prompt.length > 200 ? "..." : ""}`;
  }

  return `${name}: ${JSON.stringify(input, null, 2)}`;
}

export function getMessageStyling(type: ParsedMessage["type"]): string {
  switch (type) {
    case "system":
      return "mb-1";
    case "text":
      return "mb-1";
    case "tool_call":
      return "mb-2";
    case "tool_result":
      return "mb-2";
    case "error":
      return "mb-2";
    case "session_end":
      return "mb-2";
    default:
      return "mb-1";
  }
}

/**
 * A chunk of raw JSONL content with an optional timestamp from the DB.
 */
export interface TimestampedChunk {
  content: string;
  createdAt: string;
}

/**
 * Parse agent output. Accepts either a flat string (legacy) or timestamped chunks.
 * When chunks have timestamps, each parsed message inherits the chunk's timestamp.
 */
export function parseAgentOutput(raw: string | TimestampedChunk[]): ParsedMessage[] {
  const chunks: { content: string; timestamp?: Date }[] =
    typeof raw === "string"
      ? [{ content: raw }]
      : raw.map((c) => ({ content: c.content, timestamp: new Date(c.createdAt) }));

  const parsedMessages: ParsedMessage[] = [];
  const toolUseRegistry = new Map<string, { name: string; input: unknown }>();

  for (const chunk of chunks) {
    const lines = chunk.content.split("\n").filter(Boolean);
    for (const line of lines) {
      try {
        const event = JSON.parse(line);

        // System initialization
        if (event.type === "system" && event.session_id) {
          parsedMessages.push({
            type: "system",
            content: `Session started: ${event.session_id}`,
            timestamp: chunk.timestamp,
          });
        }

        // Assistant messages with text and tool calls
        else if (
          event.type === "assistant" &&
          Array.isArray(event.message?.content)
        ) {
          for (const block of event.message.content) {
            if (block.type === "text" && block.text) {
              parsedMessages.push({
                type: "text",
                content: block.text,
                timestamp: chunk.timestamp,
              });
            } else if (block.type === "tool_use") {
              parsedMessages.push({
                type: "tool_call",
                content: formatToolCall(block),
                timestamp: chunk.timestamp,
                meta: { toolName: block.name, toolInput: block.input },
              });
              toolUseRegistry.set(block.id, { name: block.name, input: block.input });
            } else {
              parsedMessages.push({
                type: "system",
                content: `[unsupported content block: ${block.type}]`,
                timestamp: chunk.timestamp,
              });
            }
          }
        }

        // User messages with tool results
        else if (
          event.type === "user" &&
          Array.isArray(event.message?.content)
        ) {
          for (const block of event.message.content) {
            if (block.type === "tool_result") {
              const rawContent = block.content;
              const output =
                typeof rawContent === "string"
                  ? rawContent
                  : Array.isArray(rawContent)
                    ? rawContent
                        .map((c: any) => c.text ?? JSON.stringify(c))
                        .join("\n")
                    : (event.tool_use_result?.stdout || "");
              const stderr = event.tool_use_result?.stderr;

              if (stderr) {
                parsedMessages.push({
                  type: "error",
                  content: stderr,
                  timestamp: chunk.timestamp,
                  meta: { isError: true },
                });
              } else if (output) {
                const displayOutput =
                  output.length > 1000
                    ? output.substring(0, 1000) + "\n... (truncated)"
                    : output;
                const isError = block.is_error === true;
                const originatingCall = block.tool_use_id
                  ? toolUseRegistry.get(block.tool_use_id)
                  : undefined;
                parsedMessages.push({
                  type: isError ? "error" : "tool_result",
                  content: displayOutput,
                  timestamp: chunk.timestamp,
                  meta: {
                    ...(isError ? { isError: true } : {}),
                    ...(originatingCall
                      ? { toolName: originatingCall.name, toolInput: originatingCall.input }
                      : {}),
                  },
                });
              }
            }
          }
        }

        // Legacy streaming format
        else if (event.type === "stream_event" && event.event?.delta?.text) {
          parsedMessages.push({
            type: "text",
            content: event.event.delta.text,
            timestamp: chunk.timestamp,
          });
        }

        // Session completion
        else if (event.type === "result") {
          parsedMessages.push({
            type: "session_end",
            content: event.is_error ? "Session failed" : "Session completed",
            timestamp: chunk.timestamp,
            meta: {
              isError: event.is_error,
              cost: event.total_cost_usd,
              duration: event.duration_ms,
            },
          });
        } else {
          parsedMessages.push({
            type: "system",
            content: `[unknown event: ${event.type ?? "(no type)"}]`,
            timestamp: chunk.timestamp,
          });
        }
      } catch {
        // Not JSON, display as-is
        parsedMessages.push({
          type: "error",
          content: line,
          timestamp: chunk.timestamp,
          meta: { isError: true },
        });
      }
    }
  }
  return parsedMessages;
}