storkit/frontend/src/api/client.ts

export type WsRequest =
	| {
			type: "chat";
			messages: Message[];
			config: ProviderConfig;
	  }
	| {
			type: "cancel";
	  }
	| {
			type: "permission_response";
			request_id: string;
			approved: boolean;
	  };

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[];
	  }
	| {
			type: "permission_request";
			request_id: string;
			tool_name: string;
			tool_input: Record<string, unknown>;
	  }
	| { type: "tool_activity"; tool_name: string }
	| {
			type: "reconciliation_progress";
			story_id: string;
			status: string;
			message: string;
	  }
	/** `.story_kit/project.toml` was modified; re-fetch the agent roster. */
	| { type: "agent_config_changed" }
	| { type: "tool_activity"; tool_name: string };

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 onPermissionRequest?: (
		requestId: string,
		toolName: string,
		toolInput: Record<string, unknown>,
	) => void;
	private onActivity?: (toolName: string) => void;
	private onReconciliationProgress?: (
		storyId: string,
		status: string,
		message: string,
	) => void;
	private onAgentConfigChanged?: () => void;
	private connected = false;
	private closeTimer?: number;
	private wsPath = DEFAULT_WS_PATH;
	private reconnectTimer?: number;
	private reconnectDelay = 1000;
	private shouldReconnect = false;

	private _buildWsUrl(): string {
		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,
		);
		return `${protocol}://${wsHost}${this.wsPath}`;
	}

	private _attachHandlers(): void {
		if (!this.socket) return;
		this.socket.onopen = () => {
			this.reconnectDelay = 1000;
		};
		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,
					});
				if (data.type === "permission_request")
					this.onPermissionRequest?.(
						data.request_id,
						data.tool_name,
						data.tool_input,
					);
				if (data.type === "tool_activity") this.onActivity?.(data.tool_name);
				if (data.type === "reconciliation_progress")
					this.onReconciliationProgress?.(
						data.story_id,
						data.status,
						data.message,
					);
				if (data.type === "agent_config_changed") this.onAgentConfigChanged?.();
			} catch (err) {
				this.onError?.(String(err));
			}
		};
		this.socket.onerror = () => {
			this.onError?.("WebSocket error");
		};
		this.socket.onclose = () => {
			if (this.shouldReconnect && this.connected) {
				this._scheduleReconnect();
			}
		};
	}

	private _scheduleReconnect(): void {
		window.clearTimeout(this.reconnectTimer);
		const delay = this.reconnectDelay;
		this.reconnectDelay = Math.min(this.reconnectDelay * 2, 30000);
		this.reconnectTimer = window.setTimeout(() => {
			this.reconnectTimer = undefined;
			const wsUrl = this._buildWsUrl();
			ChatWebSocket.sharedSocket = new WebSocket(wsUrl);
			this.socket = ChatWebSocket.sharedSocket;
			this._attachHandlers();
		}, delay);
	}

	connect(
		handlers: {
			onToken?: (content: string) => void;
			onUpdate?: (messages: Message[]) => void;
			onSessionId?: (sessionId: string) => void;
			onError?: (message: string) => void;
			onPipelineState?: (state: PipelineState) => void;
			onPermissionRequest?: (
				requestId: string,
				toolName: string,
				toolInput: Record<string, unknown>,
			) => void;
			onActivity?: (toolName: string) => void;
			onReconciliationProgress?: (
				storyId: string,
				status: string,
				message: string,
			) => void;
			onAgentConfigChanged?: () => 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;
		this.onPermissionRequest = handlers.onPermissionRequest;
		this.onActivity = handlers.onActivity;
		this.onReconciliationProgress = handlers.onReconciliationProgress;
		this.onAgentConfigChanged = handlers.onAgentConfigChanged;
		this.wsPath = wsPath;
		this.shouldReconnect = true;

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

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

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

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

	sendPermissionResponse(requestId: string, approved: boolean) {
		this.send({ type: "permission_response", request_id: requestId, approved });
	}

	close() {
		this.shouldReconnect = false;
		window.clearTimeout(this.reconnectTimer);
		this.reconnectTimer = undefined;

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