More browser extension work. Old interface fully ported. Direct transport. Small UX fixes.

packages/browser-extension/src/state/agent-session.ts

@@ -0,0 +1,306 @@
+import type { Context } from "@mariozechner/pi-ai";
+import {
+	type AgentTool,
+	type AssistantMessage as AssistantMessageType,
+	getModel,
+	type ImageContent,
+	type Message,
+	type Model,
+	type TextContent,
+} from "@mariozechner/pi-ai";
+import type { AppMessage } from "../Messages.js";
+import type { Attachment } from "../utils/attachment-utils.js";
+import { getAuthToken } from "../utils/auth-token.js";
+import { DirectTransport } from "./transports/DirectTransport.js";
+import { ProxyTransport } from "./transports/ProxyTransport.js";
+import type { AgentRunConfig, AgentTransport } from "./transports/types.js";
+import type { DebugLogEntry } from "./types.js";
+
+export type ThinkingLevel = "off" | "minimal" | "low" | "medium" | "high";
+
+export interface AgentSessionState {
+	id: string;
+	systemPrompt: string;
+	model: Model<any> | null;
+	thinkingLevel: ThinkingLevel;
+	tools: AgentTool<any>[];
+	messages: AppMessage[];
+	isStreaming: boolean;
+	streamMessage: Message | null;
+	pendingToolCalls: Set<string>;
+	error?: string;
+}
+
+export type AgentSessionEvent =
+	| { type: "state-update"; state: AgentSessionState }
+	| { type: "error-no-model" }
+	| { type: "error-no-api-key"; provider: string };
+
+export type TransportMode = "direct" | "proxy";
+
+export interface AgentSessionOptions {
+	initialState?: Partial<AgentSessionState>;
+	messagePreprocessor?: (messages: AppMessage[]) => Promise<Message[]>;
+	debugListener?: (entry: DebugLogEntry) => void;
+	transportMode?: TransportMode;
+	authTokenProvider?: () => Promise<string | undefined>;
+}
+
+export class AgentSession {
+	private _state: AgentSessionState = {
+		id: "default",
+		systemPrompt: "",
+		model: getModel("google", "gemini-2.5-flash-lite-preview-06-17"),
+		thinkingLevel: "off",
+		tools: [],
+		messages: [],
+		isStreaming: false,
+		streamMessage: null,
+		pendingToolCalls: new Set<string>(),
+		error: undefined,
+	};
+	private listeners = new Set<(e: AgentSessionEvent) => void>();
+	private abortController?: AbortController;
+	private transport: AgentTransport;
+	private messagePreprocessor?: (messages: AppMessage[]) => Promise<Message[]>;
+	private debugListener?: (entry: DebugLogEntry) => void;
+
+	constructor(opts: AgentSessionOptions = {}) {
+		this._state = { ...this._state, ...opts.initialState };
+		this.messagePreprocessor = opts.messagePreprocessor;
+		this.debugListener = opts.debugListener;
+
+		const mode = opts.transportMode || "direct";
+
+		if (mode === "proxy") {
+			this.transport = new ProxyTransport(async () => this.preprocessMessages());
+		} else {
+			this.transport = new DirectTransport(async () => this.preprocessMessages());
+		}
+	}
+
+	private async preprocessMessages(): Promise<Message[]> {
+		const filtered = this._state.messages.map((m) => {
+			if (m.role === "user") {
+				const { attachments, ...rest } = m as AppMessage & { attachments?: Attachment[] };
+				return rest;
+			}
+			return m;
+		});
+		return this.messagePreprocessor ? this.messagePreprocessor(filtered as AppMessage[]) : (filtered as Message[]);
+	}
+
+	get state(): AgentSessionState {
+		return this._state;
+	}
+
+	subscribe(fn: (e: AgentSessionEvent) => void): () => void {
+		this.listeners.add(fn);
+		fn({ type: "state-update", state: this._state });
+		return () => this.listeners.delete(fn);
+	}
+
+	// Mutators
+	setSystemPrompt(v: string) {
+		this.patch({ systemPrompt: v });
+	}
+	setModel(m: Model<any> | null) {
+		this.patch({ model: m });
+	}
+	setThinkingLevel(l: ThinkingLevel) {
+		this.patch({ thinkingLevel: l });
+	}
+	setTools(t: AgentTool<any>[]) {
+		this.patch({ tools: t });
+	}
+	replaceMessages(ms: AppMessage[]) {
+		this.patch({ messages: ms.slice() });
+	}
+	appendMessage(m: AppMessage) {
+		this.patch({ messages: [...this._state.messages, m] });
+	}
+	clearMessages() {
+		this.patch({ messages: [] });
+	}
+
+	abort() {
+		this.abortController?.abort();
+	}
+
+	async prompt(input: string, attachments?: Attachment[]) {
+		const model = this._state.model;
+		if (!model) {
+			this.emit({ type: "error-no-model" });
+			return;
+		}
+
+		// Build user message with attachments
+		const content: Array<TextContent | ImageContent> = [{ type: "text", text: input }];
+		if (attachments?.length) {
+			for (const a of attachments) {
+				if (a.type === "image") {
+					content.push({ type: "image", data: a.content, mimeType: a.mimeType });
+				} else if (a.type === "document" && a.extractedText) {
+					content.push({
+						type: "text",
+						text: `\n\n[Document: ${a.fileName}]\n${a.extractedText}`,
+						isDocument: true,
+					} as TextContent);
+				}
+			}
+		}
+
+		const userMessage: AppMessage = {
+			role: "user",
+			content,
+			attachments: attachments?.length ? attachments : undefined,
+		};
+
+		this.abortController = new AbortController();
+		this.patch({ isStreaming: true, streamMessage: null, error: undefined });
+
+		const reasoning =
+			this._state.thinkingLevel === "off"
+				? undefined
+				: this._state.thinkingLevel === "minimal"
+					? "low"
+					: this._state.thinkingLevel;
+		const cfg: AgentRunConfig = {
+			systemPrompt: this._state.systemPrompt,
+			tools: this._state.tools,
+			model,
+			reasoning,
+		};
+
+		try {
+			let partial: Message | null = null;
+			let turnDebug: DebugLogEntry | null = null;
+			let turnStart = 0;
+			for await (const ev of this.transport.run(userMessage as Message, cfg, this.abortController.signal)) {
+				switch (ev.type) {
+					case "turn_start": {
+						turnStart = performance.now();
+						// Build request context snapshot
+						const existing = this._state.messages as Message[];
+						const ctx: Context = {
+							systemPrompt: this._state.systemPrompt,
+							messages: [...existing],
+							tools: this._state.tools,
+						};
+						turnDebug = {
+							timestamp: new Date().toISOString(),
+							request: {
+								provider: cfg.model.provider,
+								model: cfg.model.id,
+								context: { ...ctx },
+							},
+							sseEvents: [],
+						};
+						break;
+					}
+					case "message_start":
+					case "message_update": {
+						partial = ev.message;
+						// Collect SSE-like events for debug (drop heavy partial)
+						if (ev.type === "message_update" && ev.assistantMessageEvent && turnDebug) {
+							const copy: any = { ...ev.assistantMessageEvent };
+							if (copy && "partial" in copy) delete copy.partial;
+							turnDebug.sseEvents.push(JSON.stringify(copy));
+							if (!turnDebug.ttft) turnDebug.ttft = performance.now() - turnStart;
+						}
+						this.patch({ streamMessage: ev.message });
+						break;
+					}
+					case "message_end": {
+						partial = null;
+						this.appendMessage(ev.message as AppMessage);
+						this.patch({ streamMessage: null });
+						if (turnDebug) {
+							if (ev.message.role !== "assistant" && ev.message.role !== "toolResult") {
+								turnDebug.request.context.messages.push(ev.message);
+							}
+							if (ev.message.role === "assistant") turnDebug.response = ev.message as any;
+						}
+						break;
+					}
+					case "tool_execution_start": {
+						const s = new Set(this._state.pendingToolCalls);
+						s.add(ev.toolCallId);
+						this.patch({ pendingToolCalls: s });
+						break;
+					}
+					case "tool_execution_end": {
+						const s = new Set(this._state.pendingToolCalls);
+						s.delete(ev.toolCallId);
+						this.patch({ pendingToolCalls: s });
+						break;
+					}
+					case "turn_end": {
+						// finalize current turn
+						if (turnDebug) {
+							turnDebug.totalTime = performance.now() - turnStart;
+							this.debugListener?.(turnDebug);
+							turnDebug = null;
+						}
+						break;
+					}
+					case "agent_end": {
+						this.patch({ streamMessage: null });
+						break;
+					}
+				}
+			}
+
+			if (partial && partial.role === "assistant" && partial.content.length > 0) {
+				const onlyEmpty = !partial.content.some(
+					(c) =>
+						(c.type === "thinking" && c.thinking.trim().length > 0) ||
+						(c.type === "text" && c.text.trim().length > 0) ||
+						(c.type === "toolCall" && c.name.trim().length > 0),
+				);
+				if (!onlyEmpty) {
+					this.appendMessage(partial as AppMessage);
+				} else {
+					if (this.abortController?.signal.aborted) {
+						throw new Error("Request was aborted");
+					}
+				}
+			}
+		} catch (err: any) {
+			if (String(err?.message || err) === "no-api-key") {
+				this.emit({ type: "error-no-api-key", provider: model.provider });
+			} else {
+				const msg: AssistantMessageType = {
+					role: "assistant",
+					content: [{ type: "text", text: "" }],
+					api: model.api,
+					provider: model.provider,
+					model: model.id,
+					usage: {
+						input: 0,
+						output: 0,
+						cacheRead: 0,
+						cacheWrite: 0,
+						cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
+					},
+					stopReason: this.abortController?.signal.aborted ? "aborted" : "error",
+					errorMessage: err?.message || String(err),
+				};
+				this.appendMessage(msg as AppMessage);
+				this.patch({ error: err?.message || String(err) });
+			}
+		} finally {
+			this.patch({ isStreaming: false, streamMessage: null, pendingToolCalls: new Set<string>() });
+			this.abortController = undefined;
+		}
+	}
+
+	private patch(p: Partial<AgentSessionState>): void {
+		this._state = { ...this._state, ...p };
+		this.emit({ type: "state-update", state: this._state });
+	}
+
+	private emit(e: AgentSessionEvent) {
+		this.listeners.forEach((l) => l(e));
+	}
+}

packages/browser-extension/src/state/transports/DirectTransport.ts

@@ -0,0 +1,32 @@
+import { type AgentContext, agentLoop, type Message, type PromptConfig, type UserMessage } from "@mariozechner/pi-ai";
+import { keyStore } from "../KeyStore.js";
+import type { AgentRunConfig, AgentTransport } from "./types.js";
+
+export class DirectTransport implements AgentTransport {
+	constructor(private readonly getMessages: () => Promise<Message[]>) {}
+
+	async *run(userMessage: Message, cfg: AgentRunConfig, signal?: AbortSignal) {
+		// Get API key from KeyStore
+		const apiKey = await keyStore.getKey(cfg.model.provider);
+		if (!apiKey) {
+			throw new Error("no-api-key");
+		}
+
+		const context: AgentContext = {
+			systemPrompt: cfg.systemPrompt,
+			messages: await this.getMessages(),
+			tools: cfg.tools,
+		};
+
+		const pc: PromptConfig = {
+			model: cfg.model,
+			reasoning: cfg.reasoning,
+			apiKey,
+		};
+
+		// Yield events from agentLoop
+		for await (const ev of agentLoop(userMessage as unknown as UserMessage, context, pc, signal)) {
+			yield ev;
+		}
+	}
+}

packages/browser-extension/src/state/transports/ProxyTransport.ts

@@ -0,0 +1,358 @@
+import type {
+	AgentContext,
+	AssistantMessage,
+	AssistantMessageEvent,
+	Context,
+	Message,
+	Model,
+	PromptConfig,
+	SimpleStreamOptions,
+	ToolCall,
+	UserMessage,
+} from "@mariozechner/pi-ai";
+import { agentLoop } from "@mariozechner/pi-ai";
+import { AssistantMessageEventStream } from "@mariozechner/pi-ai/dist/utils/event-stream.js";
+import { parseStreamingJson } from "@mariozechner/pi-ai/dist/utils/json-parse.js";
+import { clearAuthToken, getAuthToken } from "../../utils/auth-token.js";
+import { i18n } from "../../utils/i18n.js";
+import type { ProxyAssistantMessageEvent } from "./proxy-types.js";
+import type { AgentRunConfig, AgentTransport } from "./types.js";
+
+/**
+ * Stream function that proxies through a server instead of calling providers directly.
+ * The server strips the partial field from delta events to reduce bandwidth.
+ * We reconstruct the partial message client-side.
+ */
+function streamSimpleProxy(
+	model: Model<any>,
+	context: Context,
+	options: SimpleStreamOptions & { authToken: string },
+	proxyUrl: string,
+): AssistantMessageEventStream {
+	const stream = new AssistantMessageEventStream();
+
+	(async () => {
+		// Initialize the partial message that we'll build up from events
+		const partial: AssistantMessage = {
+			role: "assistant",
+			stopReason: "stop",
+			content: [],
+			api: model.api,
+			provider: model.provider,
+			model: model.id,
+			usage: {
+				input: 0,
+				output: 0,
+				cacheRead: 0,
+				cacheWrite: 0,
+				cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
+			},
+		};
+
+		let reader: ReadableStreamDefaultReader<Uint8Array> | undefined;
+
+		// Set up abort handler to cancel the reader
+		const abortHandler = () => {
+			if (reader) {
+				reader.cancel("Request aborted by user").catch(() => {});
+			}
+		};
+
+		if (options.signal) {
+			options.signal.addEventListener("abort", abortHandler);
+		}
+
+		try {
+			const response = await fetch(`${proxyUrl}/api/stream`, {
+				method: "POST",
+				headers: {
+					Authorization: `Bearer ${options.authToken}`,
+					"Content-Type": "application/json",
+				},
+				body: JSON.stringify({
+					model,
+					context,
+					options: {
+						temperature: options.temperature,
+						maxTokens: options.maxTokens,
+						reasoning: options.reasoning,
+						// Don't send apiKey or signal - those are added server-side
+					},
+				}),
+				signal: options.signal,
+			});
+
+			if (!response.ok) {
+				let errorMessage = `Proxy error: ${response.status} ${response.statusText}`;
+				try {
+					const errorData = await response.json();
+					if (errorData.error) {
+						errorMessage = `Proxy error: ${errorData.error}`;
+					}
+				} catch {
+					// Couldn't parse error response, use default message
+				}
+				throw new Error(errorMessage);
+			}
+
+			// Parse SSE stream
+			reader = response.body!.getReader();
+			const decoder = new TextDecoder();
+			let buffer = "";
+
+			while (true) {
+				const { done, value } = await reader.read();
+				if (done) break;
+
+				// Check if aborted after reading
+				if (options.signal?.aborted) {
+					throw new Error("Request aborted by user");
+				}
+
+				buffer += decoder.decode(value, { stream: true });
+				const lines = buffer.split("\n");
+				buffer = lines.pop() || "";
+
+				for (const line of lines) {
+					if (line.startsWith("data: ")) {
+						const data = line.slice(6).trim();
+						if (data) {
+							const proxyEvent = JSON.parse(data) as ProxyAssistantMessageEvent;
+							let event: AssistantMessageEvent | undefined;
+
+							// Handle different event types
+							// Server sends events with partial for non-delta events,
+							// and without partial for delta events
+							switch (proxyEvent.type) {
+								case "start":
+									event = { type: "start", partial };
+									break;
+
+								case "text_start":
+									partial.content[proxyEvent.contentIndex] = {
+										type: "text",
+										text: "",
+									};
+									event = { type: "text_start", contentIndex: proxyEvent.contentIndex, partial };
+									break;
+
+								case "text_delta": {
+									const content = partial.content[proxyEvent.contentIndex];
+									if (content?.type === "text") {
+										content.text += proxyEvent.delta;
+										event = {
+											type: "text_delta",
+											contentIndex: proxyEvent.contentIndex,
+											delta: proxyEvent.delta,
+											partial,
+										};
+									} else {
+										throw new Error("Received text_delta for non-text content");
+									}
+									break;
+								}
+								case "text_end": {
+									const content = partial.content[proxyEvent.contentIndex];
+									if (content?.type === "text") {
+										content.textSignature = proxyEvent.contentSignature;
+										event = {
+											type: "text_end",
+											contentIndex: proxyEvent.contentIndex,
+											content: content.text,
+											partial,
+										};
+									} else {
+										throw new Error("Received text_end for non-text content");
+									}
+									break;
+								}
+
+								case "thinking_start":
+									partial.content[proxyEvent.contentIndex] = {
+										type: "thinking",
+										thinking: "",
+									};
+									event = { type: "thinking_start", contentIndex: proxyEvent.contentIndex, partial };
+									break;
+
+								case "thinking_delta": {
+									const content = partial.content[proxyEvent.contentIndex];
+									if (content?.type === "thinking") {
+										content.thinking += proxyEvent.delta;
+										event = {
+											type: "thinking_delta",
+											contentIndex: proxyEvent.contentIndex,
+											delta: proxyEvent.delta,
+											partial,
+										};
+									} else {
+										throw new Error("Received thinking_delta for non-thinking content");
+									}
+									break;
+								}
+
+								case "thinking_end": {
+									const content = partial.content[proxyEvent.contentIndex];
+									if (content?.type === "thinking") {
+										content.thinkingSignature = proxyEvent.contentSignature;
+										event = {
+											type: "thinking_end",
+											contentIndex: proxyEvent.contentIndex,
+											content: content.thinking,
+											partial,
+										};
+									} else {
+										throw new Error("Received thinking_end for non-thinking content");
+									}
+									break;
+								}
+
+								case "toolcall_start":
+									partial.content[proxyEvent.contentIndex] = {
+										type: "toolCall",
+										id: proxyEvent.id,
+										name: proxyEvent.toolName,
+										arguments: {},
+										partialJson: "",
+									} satisfies ToolCall & { partialJson: string } as ToolCall;
+									event = { type: "toolcall_start", contentIndex: proxyEvent.contentIndex, partial };
+									break;
+
+								case "toolcall_delta": {
+									const content = partial.content[proxyEvent.contentIndex];
+									if (content?.type === "toolCall") {
+										(content as any).partialJson += proxyEvent.delta;
+										content.arguments = parseStreamingJson((content as any).partialJson) || {};
+										event = {
+											type: "toolcall_delta",
+											contentIndex: proxyEvent.contentIndex,
+											delta: proxyEvent.delta,
+											partial,
+										};
+										partial.content[proxyEvent.contentIndex] = { ...content }; // Trigger reactivity
+									} else {
+										throw new Error("Received toolcall_delta for non-toolCall content");
+									}
+									break;
+								}
+
+								case "toolcall_end": {
+									const content = partial.content[proxyEvent.contentIndex];
+									if (content?.type === "toolCall") {
+										delete (content as any).partialJson;
+										event = {
+											type: "toolcall_end",
+											contentIndex: proxyEvent.contentIndex,
+											toolCall: content,
+											partial,
+										};
+									}
+									break;
+								}
+
+								case "done":
+									partial.stopReason = proxyEvent.reason;
+									partial.usage = proxyEvent.usage;
+									event = { type: "done", reason: proxyEvent.reason, message: partial };
+									break;
+
+								case "error":
+									partial.stopReason = proxyEvent.reason;
+									partial.errorMessage = proxyEvent.errorMessage;
+									partial.usage = proxyEvent.usage;
+									event = { type: "error", reason: proxyEvent.reason, error: partial };
+									break;
+
+								default: {
+									// Exhaustive check
+									const _exhaustiveCheck: never = proxyEvent;
+									console.warn(`Unhandled event type: ${(proxyEvent as any).type}`);
+									break;
+								}
+							}
+
+							// Push the event to stream
+							if (event) {
+								stream.push(event);
+							} else {
+								throw new Error("Failed to create event from proxy event");
+							}
+						}
+					}
+				}
+			}
+
+			// Check if aborted after reading
+			if (options.signal?.aborted) {
+				throw new Error("Request aborted by user");
+			}
+
+			stream.end();
+		} catch (error) {
+			const errorMessage = error instanceof Error ? error.message : String(error);
+			if (errorMessage.toLowerCase().includes("proxy") && errorMessage.includes("Unauthorized")) {
+				clearAuthToken();
+			}
+			partial.stopReason = options.signal?.aborted ? "aborted" : "error";
+			partial.errorMessage = errorMessage;
+			stream.push({
+				type: "error",
+				reason: partial.stopReason,
+				error: partial,
+			} satisfies AssistantMessageEvent);
+			stream.end();
+		} finally {
+			// Clean up abort handler
+			if (options.signal) {
+				options.signal.removeEventListener("abort", abortHandler);
+			}
+		}
+	})();
+
+	return stream;
+}
+
+// Proxy transport executes the turn using a remote proxy server
+export class ProxyTransport implements AgentTransport {
+	// Hardcoded proxy URL for now - will be made configurable later
+	private readonly proxyUrl = "https://genai.mariozechner.at";
+
+	constructor(private readonly getMessages: () => Promise<Message[]>) {}
+
+	async *run(userMessage: Message, cfg: AgentRunConfig, signal?: AbortSignal) {
+		const authToken = await getAuthToken();
+		if (!authToken) {
+			throw new Error(i18n("Auth token is required for proxy transport"));
+		}
+
+		// Use proxy - no local API key needed
+		const streamFn = (model: Model<any>, context: Context, options: SimpleStreamOptions | undefined) => {
+			return streamSimpleProxy(
+				model,
+				context,
+				{
+					...options,
+					authToken,
+				},
+				this.proxyUrl,
+			);
+		};
+
+		const context: AgentContext = {
+			systemPrompt: cfg.systemPrompt,
+			messages: await this.getMessages(),
+			tools: cfg.tools,
+		};
+
+		const pc: PromptConfig = {
+			model: cfg.model,
+			reasoning: cfg.reasoning,
+		};
+
+		// Yield events from the upstream agentLoop iterator
+		// Pass streamFn as the 5th parameter to use proxy
+		for await (const ev of agentLoop(userMessage as unknown as UserMessage, context, pc, signal, streamFn)) {
+			yield ev;
+		}
+	}
+}

packages/browser-extension/src/state/transports/index.ts

@@ -0,0 +1,3 @@
+export * from "./DirectTransport.js";
+export * from "./ProxyTransport.js";
+export * from "./types.js";

packages/browser-extension/src/state/transports/proxy-types.ts

@@ -0,0 +1,15 @@
+import type { StopReason, Usage } from "@mariozechner/pi-ai";
+
+export type ProxyAssistantMessageEvent =
+	| { type: "start" }
+	| { type: "text_start"; contentIndex: number }
+	| { type: "text_delta"; contentIndex: number; delta: string }
+	| { type: "text_end"; contentIndex: number; contentSignature?: string }
+	| { type: "thinking_start"; contentIndex: number }
+	| { type: "thinking_delta"; contentIndex: number; delta: string }
+	| { type: "thinking_end"; contentIndex: number; contentSignature?: string }
+	| { type: "toolcall_start"; contentIndex: number; id: string; toolName: string }
+	| { type: "toolcall_delta"; contentIndex: number; delta: string }
+	| { type: "toolcall_end"; contentIndex: number }
+	| { type: "done"; reason: Extract<StopReason, "stop" | "length" | "toolUse">; usage: Usage }
+	| { type: "error"; reason: Extract<StopReason, "aborted" | "error">; errorMessage: string; usage: Usage };

packages/browser-extension/src/state/transports/types.ts

@@ -0,0 +1,16 @@
+import type { AgentEvent, AgentTool, Message, Model } from "@mariozechner/pi-ai";
+
+// The minimal configuration needed to run a turn.
+export interface AgentRunConfig {
+	systemPrompt: string;
+	tools: AgentTool<any>[];
+	model: Model<any>;
+	reasoning?: "low" | "medium" | "high";
+}
+
+// Events yielded by transports must match the @mariozechner/pi-ai prompt() events.
+// We re-export the Message type above; consumers should use the upstream AgentEvent type.
+
+export interface AgentTransport {
+	run(userMessage: Message, config: AgentRunConfig, signal?: AbortSignal): AsyncIterable<AgentEvent>; // passthrough of AgentEvent from upstream
+}

packages/browser-extension/src/state/types.ts

@@ -0,0 +1,11 @@
+import type { AssistantMessage, Context } from "@mariozechner/pi-ai";
+
+export interface DebugLogEntry {
+	timestamp: string;
+	request: { provider: string; model: string; context: Context };
+	response?: AssistantMessage;
+	error?: unknown;
+	sseEvents: string[];
+	ttft?: number;
+	totalTime?: number;
+}