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;
			always_allow: boolean;
	  }
	| { type: "ping" }
	| {
			type: "side_question";
			question: string;
			context_messages: Message[];
			config: ProviderConfig;
	  };

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

export interface PipelineStageItem {
	story_id: string;
	name: string | null;
	error: string | null;
	merge_failure: string | null;
	agent: AgentAssignment | null;
	review_hold: boolean | null;
	manual_qa: boolean | null;
}

export interface PipelineState {
	backlog: PipelineStageItem[];
	current: PipelineStageItem[];
	qa: PipelineStageItem[];
	merge: PipelineStageItem[];
	done: 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";
			backlog: PipelineStageItem[];
			current: PipelineStageItem[];
			qa: PipelineStageItem[];
			merge: PipelineStageItem[];
			done: 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" }
	/** An agent started, stopped, or changed state; re-fetch agent list. */
	| { type: "agent_state_changed" }
	| { type: "tool_activity"; tool_name: string }
	/** Heartbeat response confirming the connection is alive. */
	| { type: "pong" }
	/** Sent on connect when the project still needs onboarding (specs are placeholders). */
	| { type: "onboarding_status"; needs_onboarding: boolean }
	/** Streaming thinking token from an extended-thinking block, separate from regular text. */
	| { type: "thinking_token"; content: string }
	/** Streaming token from a /btw side question response. */
	| { type: "side_question_token"; content: string }
	/** Final signal that the /btw side question has been fully answered. */
	| { type: "side_question_done"; response: string }
	/** A single server log entry (bulk on connect, then live). */
	| { type: "log_entry"; timestamp: string; level: string; message: 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 WorkItemContent {
	content: string;
	stage: string;
	name: string | null;
	agent: string | null;
}

export interface TestCaseResult {
	name: string;
	status: "pass" | "fail";
	details: string | null;
}

export interface TestResultsResponse {
	unit: TestCaseResult[];
	integration: TestCaseResult[];
}

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);
	},
	listProjectFiles(baseUrl?: string) {
		return requestJson<string[]>("/io/fs/files", {}, 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);
	},
	getWorkItemContent(storyId: string, baseUrl?: string) {
		return requestJson<WorkItemContent>(
			`/work-items/${encodeURIComponent(storyId)}`,
			{},
			baseUrl,
		);
	},
	getTestResults(storyId: string, baseUrl?: string) {
		return requestJson<TestResultsResponse | null>(
			`/work-items/${encodeURIComponent(storyId)}/test-results`,
			{},
			baseUrl,
		);
	},
	/** Approve a story in QA, moving it to merge. */
	approveQa(storyId: string) {
		return callMcpTool("approve_qa", { story_id: storyId });
	},
	/** Reject a story in QA, moving it back to current with notes. */
	rejectQa(storyId: string, notes: string) {
		return callMcpTool("reject_qa", { story_id: storyId, notes });
	},
	/** Launch the QA app for a story's worktree. */
	launchQaApp(storyId: string) {
		return callMcpTool("launch_qa_app", { story_id: storyId });
	},
};

async function callMcpTool(
	toolName: string,
	args: Record<string, unknown>,
): Promise<string> {
	const res = await fetch("/mcp", {
		method: "POST",
		headers: { "Content-Type": "application/json" },
		body: JSON.stringify({
			jsonrpc: "2.0",
			id: 1,
			method: "tools/call",
			params: { name: toolName, arguments: args },
		}),
	});
	const json = await res.json();
	if (json.error) {
		throw new Error(json.error.message);
	}
	const text = json.result?.content?.[0]?.text ?? "";
	return text;
}

export class ChatWebSocket {
	private static sharedSocket: WebSocket | null = null;
	private static refCount = 0;
	private socket?: WebSocket;
	private onToken?: (content: string) => void;
	private onThinkingToken?: (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 onAgentStateChanged?: () => void;
	private onOnboardingStatus?: (needsOnboarding: boolean) => void;
	private onSideQuestionToken?: (content: string) => void;
	private onSideQuestionDone?: (response: string) => void;
	private onLogEntry?: (
		timestamp: string,
		level: string,
		message: string,
	) => void;
	private connected = false;
	private closeTimer?: number;
	private wsPath = DEFAULT_WS_PATH;
	private reconnectTimer?: number;
	private reconnectDelay = 1000;
	private shouldReconnect = false;
	private heartbeatInterval?: number;
	private heartbeatTimeout?: number;
	private static readonly HEARTBEAT_INTERVAL = 30_000;
	private static readonly HEARTBEAT_TIMEOUT = 5_000;

	private _startHeartbeat(): void {
		this._stopHeartbeat();
		this.heartbeatInterval = window.setInterval(() => {
			if (!this.socket || this.socket.readyState !== WebSocket.OPEN) return;
			const ping: WsRequest = { type: "ping" };
			this.socket.send(JSON.stringify(ping));
			this.heartbeatTimeout = window.setTimeout(() => {
				// No pong received within timeout; close socket to trigger reconnect.
				this.socket?.close();
			}, ChatWebSocket.HEARTBEAT_TIMEOUT);
		}, ChatWebSocket.HEARTBEAT_INTERVAL);
	}

	private _stopHeartbeat(): void {
		window.clearInterval(this.heartbeatInterval);
		window.clearTimeout(this.heartbeatTimeout);
		this.heartbeatInterval = undefined;
		this.heartbeatTimeout = undefined;
	}

	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._startHeartbeat();
		};
		this.socket.onmessage = (event) => {
			try {
				const data = JSON.parse(event.data) as WsResponse;
				if (data.type === "token") this.onToken?.(data.content);
				if (data.type === "thinking_token")
					this.onThinkingToken?.(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?.({
						backlog: data.backlog,
						current: data.current,
						qa: data.qa,
						merge: data.merge,
						done: data.done,
					});
				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?.();
				if (data.type === "agent_state_changed") this.onAgentStateChanged?.();
				if (data.type === "onboarding_status")
					this.onOnboardingStatus?.(data.needs_onboarding);
				if (data.type === "side_question_token")
					this.onSideQuestionToken?.(data.content);
				if (data.type === "side_question_done")
					this.onSideQuestionDone?.(data.response);
				if (data.type === "log_entry")
					this.onLogEntry?.(data.timestamp, data.level, data.message);
				if (data.type === "pong") {
					window.clearTimeout(this.heartbeatTimeout);
					this.heartbeatTimeout = undefined;
				}
			} 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;
			onThinkingToken?: (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;
			onAgentStateChanged?: () => void;
			onOnboardingStatus?: (needsOnboarding: boolean) => void;
			onSideQuestionToken?: (content: string) => void;
			onSideQuestionDone?: (response: string) => void;
			onLogEntry?: (timestamp: string, level: string, message: string) => void;
		},
		wsPath = DEFAULT_WS_PATH,
	) {
		this.onToken = handlers.onToken;
		this.onThinkingToken = handlers.onThinkingToken;
		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.onAgentStateChanged = handlers.onAgentStateChanged;
		this.onOnboardingStatus = handlers.onOnboardingStatus;
		this.onSideQuestionToken = handlers.onSideQuestionToken;
		this.onSideQuestionDone = handlers.onSideQuestionDone;
		this.onLogEntry = handlers.onLogEntry;
		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 });
	}

	sendSideQuestion(
		question: string,
		contextMessages: Message[],
		config: ProviderConfig,
	) {
		this.send({
			type: "side_question",
			question,
			context_messages: contextMessages,
			config,
		});
	}

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

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

	close() {
		this.shouldReconnect = false;
		this._stopHeartbeat();
		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));
	}
}