storkit/frontend/src/api/client.ts

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

export interface AgentAssignment {
	agent_name: string;
	model: string | null;
	status: string;
}

export interface PipelineStageItem {
	story_id: string;
	name: string | null;
	error: string | null;
	agent: AgentAssignment | null;
}

export interface PipelineState {
	upcoming: PipelineStageItem[];
	current: PipelineStageItem[];
	qa: PipelineStageItem[];
	merge: PipelineStageItem[];
}

export type WsResponse =
	| { type: "token"; content: string }
	| { type: "update"; messages: Message[] }
	| { type: "session_id"; session_id: string }
	| { type: "error"; message: string }
	| {
			type: "pipeline_state";
			upcoming: PipelineStageItem[];
			current: PipelineStageItem[];
			qa: PipelineStageItem[];
			merge: PipelineStageItem[];
	  };

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

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;
}

declare const __STORYKIT_PORT__: string;

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

export function resolveWsHost(
	isDev: boolean,
	envPort: string | undefined,
	locationHost: string,
): string {
	return isDev ? `127.0.0.1:${envPort || "3001"}` : locationHost;
}

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,
		);
	},
	createDirectoryAbsolute(path: string, baseUrl?: string) {
		return requestJson<boolean>(
			"/io/fs/create/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 onSessionId?: (sessionId: string) => void;
	private onError?: (message: string) => void;
	private onPipelineState?: (state: PipelineState) => void;
	private connected = false;
	private closeTimer?: number;

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

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

		const protocol = window.location.protocol === "https:" ? "wss" : "ws";
		const wsHost = resolveWsHost(
			import.meta.env.DEV,
			typeof __STORYKIT_PORT__ !== "undefined" ? __STORYKIT_PORT__ : undefined,
			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 === "session_id") this.onSessionId?.(data.session_id);
				if (data.type === "error") this.onError?.(data.message);
				if (data.type === "pipeline_state")
					this.onPipelineState?.({
						upcoming: data.upcoming,
						current: data.current,
						qa: data.qa,
						merge: data.merge,
					});
			} 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));
	}
}