Refactor agent architecture and add session storage

Major architectural improvements:
- Renamed AgentSession → Agent (state/ → agent/)
- Removed id field from AgentState
- Fixed transport abstraction to pass messages directly instead of using callbacks
- Eliminated circular dependencies in transport creation

Transport changes:
- Changed signature: run(messages, userMessage, config, signal)
- Removed getMessages callback from ProviderTransport and AppTransport
- Transports now filter attachments internally

Session storage:
- Added SessionRepository with IndexedDB backend
- Auto-save sessions after first exchange
- Auto-generate titles from first user message
- Session list dialog with search and delete
- Persistent storage permission dialog
- Browser extension now auto-loads last session

UI improvements:
- ChatPanel creates single AgentInterface instance in setAgent()
- Added drag & drop file upload to MessageEditor
- Fixed artifacts panel auto-opening on session load
- Added "Drop files here" i18n strings
- Changed "Continue Without Saving" → "Continue Anyway"

Web example:
- Complete rewrite of main.ts with clean architecture
- Added check script to package.json
- Session management with URL state
- Editable session titles

Browser extension:
- Added full session storage support
- History and new session buttons
- Auto-load most recent session on open
- Session titles in header
This commit is contained in:
Mario Zechner 2025-10-06 12:47:52 +02:00
parent c18923a8c5
commit e5cf25a267
23 changed files with 1787 additions and 289 deletions

View file

@ -0,0 +1,288 @@
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 "../components/Messages.js";
import type { Attachment } from "../utils/attachment-utils.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 AgentState {
systemPrompt: string;
model: Model<any>;
thinkingLevel: ThinkingLevel;
tools: AgentTool<any>[];
messages: AppMessage[];
isStreaming: boolean;
streamMessage: Message | null;
pendingToolCalls: Set<string>;
error?: string;
}
export type AgentEvent =
| { type: "state-update"; state: AgentState }
| { type: "error-no-model" }
| { type: "error-no-api-key"; provider: string };
export interface AgentOptions {
initialState?: Partial<AgentState>;
debugListener?: (entry: DebugLogEntry) => void;
transport: AgentTransport;
}
export class Agent {
private _state: AgentState = {
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: AgentEvent) => void>();
private abortController?: AbortController;
private transport: AgentTransport;
private debugListener?: (entry: DebugLogEntry) => void;
constructor(opts: AgentOptions) {
this._state = { ...this._state, ...opts.initialState };
this.debugListener = opts.debugListener;
this.transport = opts.transport;
}
get state(): AgentState {
return this._state;
}
subscribe(fn: (e: AgentEvent) => 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>) {
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(
this._state.messages as Message[],
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;
}
{
const { systemPrompt, model, messages } = this._state;
console.log("final state:", { systemPrompt, model, messages });
}
}
private patch(p: Partial<AgentState>): void {
this._state = { ...this._state, ...p };
this.emit({ type: "state-update", state: this._state });
}
private emit(e: AgentEvent) {
for (const listener of this.listeners) {
listener(e);
}
}
}

View file

@ -0,0 +1,370 @@
import type {
AgentContext,
Api,
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
/**
* Transport that uses an app server with user authentication tokens.
* The server manages user accounts and proxies requests to LLM providers.
*/
export class AppTransport implements AgentTransport {
// Hardcoded proxy URL for now - will be made configurable later
private readonly proxyUrl = "https://genai.mariozechner.at";
async *run(messages: Message[], 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 = <TApi extends Api>(model: Model<TApi>, context: Context, options?: SimpleStreamOptions) => {
return streamSimpleProxy(
model,
context,
{
...options,
authToken,
},
this.proxyUrl,
);
};
// Filter out attachments from messages
const filteredMessages = messages.map((m) => {
if (m.role === "user") {
const { attachments, ...rest } = m as any;
return rest;
}
return m;
});
const context: AgentContext = {
systemPrompt: cfg.systemPrompt,
messages: filteredMessages,
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 as any)) {
yield ev;
}
}
}

View file

@ -0,0 +1,56 @@
import { type AgentContext, agentLoop, type Message, type PromptConfig, type UserMessage } from "@mariozechner/pi-ai";
import { getAppStorage } from "../../storage/app-storage.js";
import type { AgentRunConfig, AgentTransport } from "./types.js";
/**
* Transport that calls LLM providers directly.
* Optionally routes calls through a CORS proxy if enabled in settings.
*/
export class ProviderTransport implements AgentTransport {
async *run(messages: Message[], userMessage: Message, cfg: AgentRunConfig, signal?: AbortSignal) {
// Get API key from storage
const apiKey = await getAppStorage().providerKeys.getKey(cfg.model.provider);
if (!apiKey) {
throw new Error("no-api-key");
}
// Check if CORS proxy is enabled
const proxyEnabled = await getAppStorage().settings.get<boolean>("proxy.enabled");
const proxyUrl = await getAppStorage().settings.get<string>("proxy.url");
// Clone model and modify baseUrl if proxy is enabled
let model = cfg.model;
if (proxyEnabled && proxyUrl && cfg.model.baseUrl) {
model = {
...cfg.model,
baseUrl: `${proxyUrl}/?url=${encodeURIComponent(cfg.model.baseUrl)}`,
};
}
// Filter out attachments from messages
const filteredMessages = messages.map((m) => {
if (m.role === "user") {
const { attachments, ...rest } = m as any;
return rest;
}
return m;
});
const context: AgentContext = {
systemPrompt: cfg.systemPrompt,
messages: filteredMessages,
tools: cfg.tools,
};
const pc: PromptConfig = {
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;
}
}
}

View file

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

View file

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

View file

@ -0,0 +1,21 @@
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(
messages: Message[],
userMessage: Message,
config: AgentRunConfig,
signal?: AbortSignal,
): AsyncIterable<AgentEvent>; // passthrough of AgentEvent from upstream
}

View file

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