storkit/frontend/src/api/client.ts

export type WsRequest =
  | {
      type: "chat";
      messages: Message[];
      config: ProviderConfig;
    }
  | {
      type: "cancel";
    };

export type WsResponse =
  | { type: "token"; content: string }
  | { type: "update"; messages: Message[] }
  | { type: "error"; message: string };

export interface ProviderConfig {
  provider: string;
  model: string;
  base_url?: string;
  enable_tools?: boolean;
}

export type Role = "system" | "user" | "assistant" | "tool";

export interface ToolCall {
  id?: string;
  type: string;
  function: {
    name: string;
    arguments: string;
  };
}

export interface Message {
  role: Role;
  content: string;
  tool_calls?: ToolCall[];
  tool_call_id?: string;
}

export interface FileEntry {
  name: string;
  kind: "file" | "dir";
}

export interface SearchResult {
  path: string;
  matches: number;
}

export interface CommandOutput {
  stdout: string;
  stderr: string;
  exit_code: number;
}

const DEFAULT_API_BASE = "/api";
const DEFAULT_WS_PATH = "/ws";

function buildApiUrl(path: string, baseUrl = DEFAULT_API_BASE): string {
  return `${baseUrl}${path}`;
}

async function requestJson<T>(
  path: string,
  options: RequestInit = {},
  baseUrl = DEFAULT_API_BASE,
): Promise<T> {
  const res = await fetch(buildApiUrl(path, baseUrl), {
    headers: {
      "Content-Type": "application/json",
      ...(options.headers ?? {}),
    },
    ...options,
  });

  if (!res.ok) {
    const text = await res.text();
    throw new Error(text || `Request failed (${res.status})`);
  }

  return res.json() as Promise<T>;
}

export const api = {
  getCurrentProject(baseUrl?: string) {
    return requestJson<string | null>("/project", {}, baseUrl);
  },
  getKnownProjects(baseUrl?: string) {
    return requestJson<string[]>("/projects", {}, baseUrl);
  },
  forgetKnownProject(path: string, baseUrl?: string) {
    return requestJson<boolean>(
      "/projects/forget",
      { method: "POST", body: JSON.stringify({ path }) },
      baseUrl,
    );
  },
  openProject(path: string, baseUrl?: string) {
    return requestJson<string>(
      "/project",
      { method: "POST", body: JSON.stringify({ path }) },
      baseUrl,
    );
  },
  closeProject(baseUrl?: string) {
    return requestJson<boolean>("/project", { method: "DELETE" }, baseUrl);
  },
  getModelPreference(baseUrl?: string) {
    return requestJson<string | null>("/model", {}, baseUrl);
  },
  setModelPreference(model: string, baseUrl?: string) {
    return requestJson<boolean>(
      "/model",
      { method: "POST", body: JSON.stringify({ model }) },
      baseUrl,
    );
  },
  getOllamaModels(baseUrlParam?: string, baseUrl?: string) {
    const url = new URL(
      buildApiUrl("/ollama/models", baseUrl),
      window.location.origin,
    );
    if (baseUrlParam) {
      url.searchParams.set("base_url", baseUrlParam);
    }
    return requestJson<string[]>(url.pathname + url.search, {}, "");
  },
  getAnthropicApiKeyExists(baseUrl?: string) {
    return requestJson<boolean>("/anthropic/key/exists", {}, baseUrl);
  },
  getAnthropicModels(baseUrl?: string) {
    return requestJson<string[]>("/anthropic/models", {}, baseUrl);
  },
  setAnthropicApiKey(api_key: string, baseUrl?: string) {
    return requestJson<boolean>(
      "/anthropic/key",
      { method: "POST", body: JSON.stringify({ api_key }) },
      baseUrl,
    );
  },
  readFile(path: string, baseUrl?: string) {
    return requestJson<string>(
      "/fs/read",
      { method: "POST", body: JSON.stringify({ path }) },
      baseUrl,
    );
  },
  writeFile(path: string, content: string, baseUrl?: string) {
    return requestJson<boolean>(
      "/fs/write",
      { method: "POST", body: JSON.stringify({ path, content }) },
      baseUrl,
    );
  },
  listDirectory(path: string, baseUrl?: string) {
    return requestJson<FileEntry[]>(
      "/fs/list",
      { method: "POST", body: JSON.stringify({ path }) },
      baseUrl,
    );
  },
  listDirectoryAbsolute(path: string, baseUrl?: string) {
    return requestJson<FileEntry[]>(
      "/io/fs/list/absolute",
      { method: "POST", body: JSON.stringify({ path }) },
      baseUrl,
    );
  },
  getHomeDirectory(baseUrl?: string) {
    return requestJson<string>("/io/fs/home", {}, baseUrl);
  },
  searchFiles(query: string, baseUrl?: string) {
    return requestJson<SearchResult[]>(
      "/fs/search",
      { method: "POST", body: JSON.stringify({ query }) },
      baseUrl,
    );
  },
  execShell(command: string, args: string[], baseUrl?: string) {
    return requestJson<CommandOutput>(
      "/shell/exec",
      { method: "POST", body: JSON.stringify({ command, args }) },
      baseUrl,
    );
  },
  cancelChat(baseUrl?: string) {
    return requestJson<boolean>("/chat/cancel", { method: "POST" }, baseUrl);
  },
};

export class ChatWebSocket {
  private static sharedSocket: WebSocket | null = null;
  private static refCount = 0;
  private socket?: WebSocket;
  private onToken?: (content: string) => void;
  private onUpdate?: (messages: Message[]) => void;
  private onError?: (message: string) => void;
  private connected = false;
  private closeTimer?: number;

  connect(
    handlers: {
      onToken?: (content: string) => void;
      onUpdate?: (messages: Message[]) => void;
      onError?: (message: string) => void;
    },
    wsPath = DEFAULT_WS_PATH,
  ) {
    this.onToken = handlers.onToken;
    this.onUpdate = handlers.onUpdate;
    this.onError = handlers.onError;

    if (this.connected) {
      return;
    }
    this.connected = true;
    ChatWebSocket.refCount += 1;

    const protocol = window.location.protocol === "https:" ? "wss" : "ws";
    const wsHost = import.meta.env.DEV
      ? "127.0.0.1:3001"
      : window.location.host;
    const wsUrl = `${protocol}://${wsHost}${wsPath}`;

    if (
      !ChatWebSocket.sharedSocket ||
      ChatWebSocket.sharedSocket.readyState === WebSocket.CLOSED ||
      ChatWebSocket.sharedSocket.readyState === WebSocket.CLOSING
    ) {
      ChatWebSocket.sharedSocket = new WebSocket(wsUrl);
    }
    this.socket = ChatWebSocket.sharedSocket;

    this.socket.onmessage = (event) => {
      try {
        const data = JSON.parse(event.data) as WsResponse;
        if (data.type === "token") this.onToken?.(data.content);
        if (data.type === "update") this.onUpdate?.(data.messages);
        if (data.type === "error") this.onError?.(data.message);
      } catch (err) {
        this.onError?.(String(err));
      }
    };

    this.socket.onerror = () => {
      this.onError?.("WebSocket error");
    };
  }

  sendChat(messages: Message[], config: ProviderConfig) {
    this.send({ type: "chat", messages, config });
  }

  cancel() {
    this.send({ type: "cancel" });
  }

  close() {
    if (!this.connected) return;
    this.connected = false;
    ChatWebSocket.refCount = Math.max(0, ChatWebSocket.refCount - 1);

    if (import.meta.env.DEV) {
      if (this.closeTimer) {
        window.clearTimeout(this.closeTimer);
      }
      this.closeTimer = window.setTimeout(() => {
        if (ChatWebSocket.refCount === 0) {
          ChatWebSocket.sharedSocket?.close();
          ChatWebSocket.sharedSocket = null;
        }
        this.socket = ChatWebSocket.sharedSocket ?? undefined;
        this.closeTimer = undefined;
      }, 250);
      return;
    }

    if (ChatWebSocket.refCount === 0) {
      ChatWebSocket.sharedSocket?.close();
      ChatWebSocket.sharedSocket = null;
    }
    this.socket = ChatWebSocket.sharedSocket ?? undefined;
  }

  private send(payload: WsRequest) {
    if (!this.socket || this.socket.readyState !== WebSocket.OPEN) {
      this.onError?.("WebSocket is not connected");
      return;
    }
    this.socket.send(JSON.stringify(payload));
  }
}