chat-provider.tsx

"use client";

import { useSearchParams } from "next/navigation";
import {
  useState,
  useEffect,
  createContext,
  PropsWithChildren,
  useContext,
} from "react";
import {toast} from "sonner";
import {getErrorMessage} from "@/lib/error-utils";

interface Message {
  id: number;
  role: string;
  content: string;
}

// Draft messages are used to optmistically update the UI
// before the server responds.
interface DraftMessage extends Omit<Message, "id"> {
  id?: number;
}

interface MessageUpdateEvent {
  id: number;
  role: string;
  message: string;
  time: string;
}

interface StatusChangeEvent {
  status: string;
  agent_type: string;
}

interface APIErrorDetail {
  location: string;
  message: string;
  value: null | string | number | boolean | object;
}

interface APIErrorModel {
  $schema: string;
  detail: string;
  errors: APIErrorDetail[];
  instance: string;
  status: number;
  title: string;
  type: string;
}

function isDraftMessage(message: Message | DraftMessage): boolean {
  return message.id === undefined;
}

type MessageType = "user" | "raw";

export type ServerStatus = "stable" | "running" | "offline" | "unknown";

export interface FileUploadResponse {
  ok: boolean;
  filePath?: string;
}

export type AgentType = "claude" | "goose" | "aider" | "gemini" | "amp" | "codex" | "cursor" | "cursor-agent" | "copilot" | "auggie" | "amazonq" | "opencode" | "custom" | "unknown";

export type AgentColorDisplayNamePair = {
  displayName: string;
}

export const AgentType: Record<Exclude<AgentType, "unknown">, AgentColorDisplayNamePair> = {
  claude: {displayName: "Claude Code"},
  goose: {displayName: "Goose"},
  aider: {displayName: "Aider"},
  gemini: { displayName: "Gemini"},
  amp: {displayName: "Amp"},
  codex: {displayName: "Codex"},
  cursor: { displayName: "Cursor Agent"},
  "cursor-agent": { displayName: "Cursor Agent"},
  copilot: {displayName: "Copilot"},
  auggie: {displayName: "Auggie"},
  amazonq: {displayName: "Amazon Q"},
  opencode: {displayName: "Opencode"},
  custom: { displayName: "Custom"}
}

interface ChatContextValue {
  messages: (Message | DraftMessage)[];
  loading: boolean;
  serverStatus: ServerStatus;
  sendMessage: (message: string, type?: MessageType) => void;
  uploadFiles: (formData: FormData) => Promise<FileUploadResponse>;
  agentType: AgentType;
}

const ChatContext = createContext<ChatContextValue | undefined>(undefined);

const useAgentAPIUrl = (): string => {
  const searchParams = useSearchParams();
  const paramsUrl = searchParams.get("url");
  if (paramsUrl) {
    return paramsUrl;
  }
  const basePath = process.env.NEXT_PUBLIC_BASE_PATH;
  if (!basePath) {
    throw new Error(
      "agentAPIUrl is not set. Please set the url query parameter to the URL of the AgentAPI or the NEXT_PUBLIC_BASE_PATH environment variable."
    );
  }
  // NOTE(cian): We use '../' here to construct the agent API URL relative
  // to the chat's location. Let's say the app is hosted on a subpath
  // `/@admin/workspace.agent/apps/ccw/`. When you visit this URL you get
  // redirected to `/@admin/workspace.agent/apps/ccw/chat/embed`. This serves
  // this React application, but it needs to know where the agent API is hosted.
  // This will be at the root of where the application is mounted e.g.
  // `/@admin/workspace.agent/apps/ccw/`. Previously we used
  // `window.location.origin` but this assumes that the application owns the
  // entire origin.
  // See: https://github.com/coder/coder/issues/18779#issuecomment-3133290494 for more context.
  let chatURL: string = new URL(basePath, window.location.origin).toString();
  // NOTE: trailing slashes and relative URLs are tricky.
  // https://developer.mozilla.org/en-US/docs/Web/API/URL_API/Resolving_relative_references#current_directory_relative
  if (!chatURL.endsWith("/")) {
    chatURL += "/";
  }
  const agentAPIURL = new URL("..", chatURL).toString();
  if (agentAPIURL.endsWith("/")) {
    return agentAPIURL.slice(0, -1);
  }
  return agentAPIURL;
};

export function ChatProvider({ children }: PropsWithChildren) {
  const [messages, setMessages] = useState<(Message | DraftMessage)[]>([]);
  const [loading, setLoading] = useState<boolean>(false);
  const [serverStatus, setServerStatus] = useState<ServerStatus>("unknown");
  const [agentType, setAgentType] = useState<AgentType>("custom");
  const agentAPIUrl = useAgentAPIUrl();

  // Set up SSE connection to the events endpoint. EventSource handles
  // reconnection automatically, so we only create it once per URL and
  // let the browser manage transient failures. Messages are NOT cleared
  // on reconnect to avoid blanking the conversation on network blips.
  useEffect(() => {
    if (!agentAPIUrl) {
      console.warn(
        "agentAPIUrl is not set, SSE connection cannot be established."
      );
      setServerStatus("offline");
      return;
    }

    const eventSource = new EventSource(`${agentAPIUrl}/events`);

    // Handle message updates
    eventSource.addEventListener("message_update", (event) => {
      const data: MessageUpdateEvent = JSON.parse(event.data);
      const confirmed: Message = {
        role: data.role,
        content: data.message,
        id: data.id,
      };

      setMessages((prevMessages) => {
        // Check if message with this ID already exists
        const existingIndex = prevMessages.findIndex(
          (m) => m.id === data.id
        );

        if (existingIndex !== -1) {
          // Update existing message
          const updated = [...prevMessages];
          updated[existingIndex] = confirmed;
          return updated;
        }

        // New confirmed message: replace any trailing draft that matches
        // the same role (the optimistic message we inserted on send).
        const last = prevMessages[prevMessages.length - 1];
        if (last && isDraftMessage(last) && last.role === confirmed.role) {
          return [...prevMessages.slice(0, -1), confirmed];
        }

        return [...prevMessages, confirmed];
      });
    });

    // Handle status changes
    eventSource.addEventListener("status_change", (event) => {
      const data: StatusChangeEvent = JSON.parse(event.data);
      if (data.status === "stable") {
        setServerStatus("stable");
      } else if (data.status === "running") {
        setServerStatus("running");
      } else {
        setServerStatus("unknown");
      }
      // Set agent type
      setAgentType(data.agent_type === "" ? "unknown" : data.agent_type as AgentType);
    });

    eventSource.onopen = () => {
      console.log("EventSource connection established");
    };

    // Mark offline on error. The browser will retry automatically.
    eventSource.onerror = () => {
      setServerStatus("offline");
    };

    return () => eventSource.close();
  }, [agentAPIUrl]);

  // Send a new message
  const sendMessage = async (
    content: string,
    type: "user" | "raw" = "user"
  ) => {
    // For user messages, require non-empty content
    if (type === "user" && !content.trim()) return;

    // For raw messages, don't set loading state as it's usually fast
    if (type === "user") {
      setMessages((prevMessages) => [
        ...prevMessages,
        { role: "user", content },
      ]);
      setLoading(true);
    }

    try {
      const response = await fetch(`${agentAPIUrl}/message`, {
        method: "POST",
        headers: {
          "Content-Type": "application/json",
        },
        body: JSON.stringify({
          content: content,
          type,
        }),
      });

      if (!response.ok) {
        const errorData = await response.json() as APIErrorModel;
        console.error("Failed to send message:", errorData);
        const detail = errorData.detail;
        const messages =
          "errors" in errorData
            ?
              errorData.errors.map((e: APIErrorDetail) => e.message).join(", ")
            : "";

        const fullDetail = `${detail}: ${messages}`;
        toast.error(`Failed to send message`, {
          description: fullDetail,
        });
        // Remove the optimistic draft since the server rejected it.
        setMessages((prev) => prev.filter((m) => !isDraftMessage(m)));
      }

    } catch (error) {
      console.error("Error sending message:", error);
      const message = getErrorMessage(error)

      toast.error(`Error sending message`, {
        description: message,
      });
      // Remove the optimistic draft since the request failed.
      setMessages((prev) => prev.filter((m) => !isDraftMessage(m)));
    } finally {
      if (type === "user") {
        setLoading(false);
      }
    }
  };

  // Upload files to workspace
  const uploadFiles = async (formData: FormData): Promise<FileUploadResponse> => {
    let result: FileUploadResponse = {ok: true};
    try{
      const response = await fetch(`${agentAPIUrl}/upload`, {
        method: 'POST',
        body: formData,
      });

      if (!response.ok) {
        result.ok = false;
        const errorData = await response.json() as APIErrorModel;
        console.error("Failed to send message:", errorData);
        const detail = errorData.detail;
        const messages =
          "errors" in errorData
            ?
            errorData.errors.map((e: APIErrorDetail) => e.message).join(", ")
            : "";

        const fullDetail = `${detail}: ${messages}`;
        toast.error(`Failed to upload files`, {
          description: fullDetail,
        });
      } else {
        result = (await response.json()) as FileUploadResponse;
      }

    } catch (error) {
      result.ok = false;
      console.error("Error uploading files:", error);
      const message = getErrorMessage(error)

      toast.error(`Error uploading files`, {
        description: message,
      });
    }
    return result;
  }

  return (
    <ChatContext.Provider
      value={{
        messages,
        loading,
        sendMessage,
        serverStatus,
        uploadFiles,
        agentType,
      }}
    >
      {children}
    </ChatContext.Provider>
  );
}

export function useChat() {
  const context = useContext(ChatContext);
  if (!context) {
    throw new Error("useChat must be used within a ChatProvider");
  }
  return context;
}