frontend/src/lib/agentStream.ts

import type { AgentEvent } from "./types";

const BASE = import.meta.env.VITE_API_BASE_URL || "";

export interface AgentRequest {
  prompt: string;
  session_id?: string | null;
  mode?: "standard" | "deep";
}

export interface AgentStreamHandle {
  cancel: () => void;
  done: Promise<void>;
}

/**
 * Stream the agent chat endpoint. Parses SSE manually (instead of using the
 * browser's `EventSource`) because EventSource doesn't support POST bodies.
 *
 * The frontend can call `cancel()` to abort an in-flight stream.
 */
export function streamAgentChat(
  req: AgentRequest,
  onEvent: (ev: AgentEvent) => void
): AgentStreamHandle {
  const controller = new AbortController();

  const done = (async () => {
    let response: Response;
    try {
      response = await fetch(BASE + "/api/agent/chat", {
        method: "POST",
        headers: {
          "Content-Type": "application/json",
          Accept: "text/event-stream",
        },
        body: JSON.stringify({
          prompt: req.prompt,
          session_id: req.session_id ?? null,
          mode: req.mode ?? "standard",
        }),
        signal: controller.signal,
      });
    } catch (e) {
      if ((e as Error).name === "AbortError") return;
      onEvent({ type: "error", error: `network error: ${(e as Error).message}` });
      return;
    }

    if (!response.ok || !response.body) {
      onEvent({
        type: "error",
        error: `HTTP ${response.status} ${response.statusText}`,
      });
      return;
    }

    const reader = response.body
      .pipeThrough(new TextDecoderStream())
      .getReader();
    let buffer = "";

    try {
      while (true) {
        const { value, done: streamDone } = await reader.read();
        if (streamDone) break;
        buffer += value;

        // SSE messages are separated by double newlines
        let sepIdx: number;
        while ((sepIdx = buffer.indexOf("\n\n")) !== -1) {
          const raw = buffer.slice(0, sepIdx);
          buffer = buffer.slice(sepIdx + 2);
          const parsed = parseSseFrame(raw);
          if (parsed) onEvent(parsed);
        }
      }
      // Flush leftover (rare; usually terminating event already came through)
      if (buffer.trim()) {
        const parsed = parseSseFrame(buffer);
        if (parsed) onEvent(parsed);
      }
    } catch (e) {
      if ((e as Error).name === "AbortError") return;
      onEvent({ type: "error", error: `stream error: ${(e as Error).message}` });
    } finally {
      try {
        reader.releaseLock();
      } catch {
        /* noop */
      }
    }
  })();

  return {
    cancel: () => controller.abort(),
    done,
  };
}

function parseSseFrame(raw: string): AgentEvent | null {
  let event = "message";
  let data = "";
  for (const line of raw.split("\n")) {
    if (line.startsWith("event:")) event = line.slice(6).trim();
    else if (line.startsWith("data:")) data += line.slice(5).trim();
    // ignore comments (": ...") and id: / retry:
  }
  if (!data) {
    // 'done' event has empty data ('{}'), but might come as 'data: {}' so this rarely fires
    if (event === "done") return { type: "done" } as AgentEvent;
    return null;
  }
  try {
    const payload = JSON.parse(data) as Partial<AgentEvent> & {
      type?: string;
    };
    if (!payload.type) payload.type = event as AgentEvent["type"];
    return payload as AgentEvent;
  } catch {
    return null;
  }
}