also no idea, from monorepo

src/apps/officer-web/Screens/Dashboard/Chat/usePiMono.ts

@@ -1,10 +1,7 @@
 import { useState, useEffect, useRef } from 'react';
 import { useChatWebSocket } from 'hooks/useChatWebSocket';
-import { useChatSessions } from '@/state/useChatSessions';
 import type { ChatMessage, ServerMessage, TaskInfo } from 'apps/Chat';
 
-const SAVE_DEBOUNCE_MS = 1000;
-
 type UsePiMonoOptions = {
   replaceUrl?: boolean;
   taskInfo?: TaskInfo;
@@ -22,9 +19,6 @@ export const usePiMono = (initialSessionId?: string, initialModel?: string | nul
   const streamingRef = useRef('');
   const rafRef = useRef<number | null>(null);
   const sessionIdRef = useRef<string | null>(initialSessionId ?? null);
-  const saveTimerRef = useRef<number | null>(null);
-
-  const { getMessages, saveMessages } = useChatSessions();
 
   const token = localStorage.getItem('BEARER_TOKEN');
   const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:';
@@ -56,8 +50,11 @@ export const usePiMono = (initialSessionId?: string, initialModel?: string | nul
         if (replaceUrl) window.history.replaceState(null, '', `/chat/${msg.sessionId}`);
         break;
 
-      case 'system:prompt':
+      case 'messages:sync':
-        setMessages((prev) => [...prev, { role: 'system', text: msg.text }]);
+        setMessages(msg.messages);
+        streamingRef.current = msg.streamingText;
+        setStreamingText(msg.streamingText);
+        setIsGenerating(msg.isGenerating);
         break;
 
       case 'assistant:partial':
@@ -125,37 +122,6 @@ export const usePiMono = (initialSessionId?: string, initialModel?: string | nul
     }
   }, [isConnected, initialSessionId]);
 
-  // Load messages from server on mount when resuming a session
-  useEffect(() => {
-    if (!initialSessionId) return;
-    getMessages('pi-mono', initialSessionId)
-      .then((data) => {
-        if (Array.isArray(data) && data.length > 0) setMessages(data);
-      })
-      .catch(() => {});
-  }, [initialSessionId]);
-
-  // Debounced save messages to server
-  useEffect(() => {
-    if (!sessionIdRef.current || messages.length === 0) return;
-
-    if (saveTimerRef.current !== null) clearTimeout(saveTimerRef.current);
-
-    const sid = sessionIdRef.current;
-    const snapshot = messages;
-    saveTimerRef.current = window.setTimeout(() => {
-      saveMessages('pi-mono', sid, snapshot).catch(() => {});
-      saveTimerRef.current = null;
-    }, SAVE_DEBOUNCE_MS);
-
-    return () => {
-      if (saveTimerRef.current !== null) {
-        clearTimeout(saveTimerRef.current);
-        saveTimerRef.current = null;
-      }
-    };
-  }, [messages]);
-
   // Clean up RAF on unmount
   useEffect(() => {
     return () => {
@@ -184,7 +150,7 @@ export const usePiMono = (initialSessionId?: string, initialModel?: string | nul
     send({
       type: 'chat',
       prompt: text,
-      sessionId: sessionIdRef.current,
+      ...(sessionIdRef.current ? { sessionId: sessionIdRef.current } : {}),
       ...(selectedModel ? { model: selectedModel } : {}),
       ...(cwd ? { cwd } : {}),
       ...(attachmentIds?.length ? { attachmentIds } : {}),

src/apps/officer-web/styles/prose.css

@@ -1,3 +1,157 @@
+/* Chat Bubble — Markdown Prose */
+.chat-md {
+  color: hsl(var(--foreground));
+  font-size: 0.875rem;
+  line-height: 1.7;
+}
+
+.chat-md > *:first-child {
+  margin-top: 0;
+}
+
+.chat-md > *:last-child {
+  margin-bottom: 0;
+}
+
+.chat-md h1 {
+  font-size: 1.4em;
+  font-weight: 700;
+  margin: 1.2em 0 0.4em;
+  padding-bottom: 0.25em;
+  border-bottom: 1px solid hsl(var(--border));
+  letter-spacing: -0.02em;
+}
+
+.chat-md h2 {
+  font-size: 1.2em;
+  font-weight: 600;
+  margin: 1.1em 0 0.35em;
+  padding-bottom: 0.2em;
+  border-bottom: 1px solid hsl(var(--border));
+}
+
+.chat-md h3 {
+  font-size: 1.1em;
+  font-weight: 600;
+  margin: 1em 0 0.3em;
+}
+
+.chat-md h4, .chat-md h5, .chat-md h6 {
+  font-size: 1em;
+  font-weight: 600;
+  margin: 0.8em 0 0.25em;
+}
+
+.chat-md p {
+  margin: 0.6em 0;
+}
+
+.chat-md a {
+  color: var(--duck-teal);
+  text-decoration: underline;
+  text-underline-offset: 2px;
+}
+
+.chat-md a:hover {
+  opacity: 0.8;
+}
+
+.chat-md strong {
+  font-weight: 600;
+}
+
+.chat-md blockquote {
+  margin: 0.75em 0;
+  padding: 0.4em 0.8em;
+  border-left: 3px solid var(--duck-teal);
+  background: color-mix(in srgb, var(--duck-teal) 5%, transparent);
+  border-radius: 0 0.375rem 0.375rem 0;
+  color: hsl(var(--muted-foreground));
+}
+
+.chat-md ul, .chat-md ol {
+  margin: 0.5em 0;
+  padding-left: 1.5em;
+}
+
+.chat-md li {
+  margin: 0.2em 0;
+}
+
+.chat-md li::marker {
+  color: var(--duck-teal);
+}
+
+.chat-md code {
+  padding: 0.15em 0.4em;
+  border-radius: 0.25rem;
+  background: color-mix(in srgb, var(--duck-teal) 10%, transparent);
+  color: var(--duck-teal);
+  font-size: 0.85em;
+  font-family: ui-monospace, monospace;
+}
+
+.chat-md pre {
+  margin: 0.75em 0;
+  padding: 0.75em 1em;
+  border-radius: 0.5rem;
+  background: #0d1117;
+  overflow-x: auto;
+  font-size: 0.8125rem;
+  line-height: 1.6;
+}
+
+.chat-md pre code {
+  padding: 0;
+  border-radius: 0;
+  background: none;
+  color: #e6edf3;
+  font-size: inherit;
+}
+
+.chat-md hr {
+  border: none;
+  height: 1px;
+  background: hsl(var(--border));
+  margin: 1.5em 0;
+}
+
+.chat-md table {
+  width: 100%;
+  border-collapse: collapse;
+  margin: 0.75em 0;
+  font-size: 0.85em;
+}
+
+.chat-md th {
+  color: hsl(var(--muted-foreground));
+  background: color-mix(in srgb, var(--duck-teal) 8%, transparent);
+  font-weight: 600;
+  text-align: left;
+  padding: 0.4em 0.6em;
+  border: 1px solid hsl(var(--border));
+}
+
+.chat-md td {
+  padding: 0.4em 0.6em;
+  border: 1px solid hsl(var(--border));
+}
+
+.chat-md tr:nth-child(even) {
+  background: hsl(var(--muted) / 0.3);
+}
+
+.chat-md img {
+  max-width: 100%;
+  border-radius: 0.375rem;
+  margin: 0.75em 0;
+}
+
+.chat-md input[type="checkbox"] {
+  accent-color: var(--duck-teal);
+  margin-right: 0.4em;
+}
+
 /* File Viewer — Markdown Prose */
 .file-viewer-md {
   color: hsl(var(--foreground));

src/servers/api/chat-types.ts

@@ -33,4 +33,5 @@ export type ServerMessage =
   | { type: 'tool:result'; toolUseId: string; output: string; isError: boolean }
   | { type: 'result'; costUsd: number; durationMs: number; numTurns: number; isError: boolean }
   | { type: 'error'; message: string }
-  | { type: 'stopped' };
+  | { type: 'stopped' }
+  | { type: 'messages:sync'; messages: unknown[]; streamingText: string; isGenerating: boolean };

src/servers/api/pi-mono/websocket.ts

@@ -1,47 +1,52 @@
-import type { ServerWebSocket } from 'bun';
+import type { ServerWebSocket, Subprocess } from 'bun';
-import type { Subprocess } from 'bun';
 import { mkdir, rename } from 'node:fs/promises';
 import { join } from 'node:path';
 import { homedir } from 'node:os';
 import {
   getPiMonoSessionDir,
-  getTmpAttachmentsDir,
-  getAttachmentsDir,
   getHomeDir,
   getNativeSkillsDir,
   getGlobalSkillsDir,
   getUserSkillsDir,
+  getTmpAttachmentsDir,
+  getAttachmentsDir,
 } from '@@/data-path';
 import { readSkillDirs, parseFrontmatter } from '@@/api/skills/skills';
-import type { ClientMessage, ServerMessage, ImageData, TaskInfo } from '@@/api/chat-types';
+import type { ClientMessage, ServerMessage } from '@@/api/chat-types';
-import { createTaskLog, appendToLog, finalizeLog } from '@@/api/task-logger';
+import { readApiKeys } from '@@/api/server-settings/pi-mono';
-import { readApiKeys, readLocalProviders } from '@@/api/server-settings/pi-mono';
 
 type WSData = { userId: number; email: string };
 
-// Pi process state, keyed by sessionId — survives websocket reconnects
+// --- Session state ---
+
 type PiSession = {
+  id: string;
+  email: string;
   piProcess: Subprocess | null;
   ws: ServerWebSocket<WSData> | null;
-  selectedModel: string | null;
+  model: string | null;
   cwd: string | null;
-  resourceChatDir: string | null;
+  messages: unknown[];
-  logId: string | null;
+  streamBuffer: string;
-  fullText: string;
+  isGenerating: boolean;
+  systemContextSent: boolean;
   killTimer: ReturnType<typeof setTimeout> | null;
+  saving: boolean;
+  dirty: boolean;
 };
 
-// Session pool — pi processes persist across websocket reconnects
 const sessions = new Map<string, PiSession>();
-
-// Map ws → sessionId for quick lookup on close
 const wsToSession = new Map<ServerWebSocket<WSData>, string>();
-
-// Grace period before killing orphaned pi processes (ms)
 const ORPHAN_GRACE_MS = 30_000;
 
-function send(ws: ServerWebSocket<WSData> | null, msg: ServerMessage) {
+// --- Helpers ---
-  if (ws && ws.readyState === 1) ws.send(JSON.stringify(msg));
+
+function sendToClient(session: PiSession, msg: ServerMessage) {
+  if (session.ws?.readyState === 1) session.ws.send(JSON.stringify(msg));
+}
+
+function sendDirect(ws: ServerWebSocket<WSData>, msg: ServerMessage) {
+  if (ws.readyState === 1) ws.send(JSON.stringify(msg));
 }
 
 function resolveRootDir(email: string, root?: string): string {
@@ -51,6 +56,72 @@ function resolveRootDir(email: string, root?: string): string {
   return getHomeDir(email);
 }
 
+function writeRpcCommand(proc: Subprocess, command: Record<string, unknown>) {
+  const stdin = proc.stdin;
+  if (!stdin || typeof stdin === 'number') return;
+  try {
+    const writer = stdin as { write(data: string): void; flush(): void };
+    writer.write(JSON.stringify(command) + '\n');
+    writer.flush();
+  } catch (err) {
+    console.error('[pi-mono] writeRpcCommand error:', err);
+  }
+}
+
+// --- Message persistence ---
+
+async function persistMessages(session: PiSession) {
+  if (session.saving) {
+    session.dirty = true;
+    return;
+  }
+  session.saving = true;
+  session.dirty = false;
+  try {
+    const dir = getPiMonoSessionDir(session.email, session.id);
+    await Bun.write(join(dir, 'messages.json'), JSON.stringify(session.messages));
+  } catch (err) {
+    console.error('[pi-mono] persistMessages error:', err);
+  } finally {
+    session.saving = false;
+    if (session.dirty) persistMessages(session);
+  }
+}
+
+async function loadMessages(email: string, sessionId: string): Promise<unknown[]> {
+  try {
+    const file = Bun.file(join(getPiMonoSessionDir(email, sessionId), 'messages.json'));
+    if (!(await file.exists())) return [];
+    const data = await file.json();
+    return Array.isArray(data) ? data : [];
+  } catch {
+    return [];
+  }
+}
+
+function buildHistoryContext(messages: unknown[]): string {
+  const lines: string[] = [];
+  for (const msg of messages) {
+    const m = msg as Record<string, unknown>;
+    if (m.role === 'user' && m.text) lines.push(`User: ${m.text}`);
+    else if (m.role === 'assistant' && m.text) lines.push(`Assistant: ${m.text}`);
+    else if (m.role === 'tool' && m.toolName) {
+      const output = m.output ? String(m.output).slice(0, 500) : '(no output)';
+      lines.push(`[Tool: ${m.toolName}] ${output}`);
+    }
+  }
+  if (lines.length === 0) return '';
+
+  let history = lines.join('\n');
+  if (history.length > 30_000) {
+    history = '...(truncated)\n' + history.slice(-30_000);
+    history = history.slice(history.indexOf('\n') + 1);
+  }
+  return `\n\nBelow is the conversation history from this session:\n<conversation_history>\n${history}\n</conversation_history>`;
+}
+
+// --- Skills ---
+
 async function buildSkillsPrompt(email: string): Promise<string> {
   const nativeSkills = await readSkillDirs(getNativeSkillsDir());
   const globalSkills = await readSkillDirs(getGlobalSkillsDir());
@@ -74,180 +145,15 @@ async function buildSkillsPrompt(email: string): Promise<string> {
   return `\n\nYou have access to the following skills. When a user's request matches a skill, read its SKILL.md file for detailed instructions before proceeding.\n\nAvailable skills:\n${lines.join('\n')}`;
 }
 
-function writeRpcCommand(proc: Subprocess, command: Record<string, unknown>) {
+// --- Pi process lifecycle ---
-  const stdin = proc.stdin;
-  if (!stdin || typeof stdin === 'number') return;
-  try {
-    (stdin as { write: (data: string) => void; flush: () => void }).write(JSON.stringify(command) + '\n');
-    (stdin as { flush: () => void }).flush();
-  } catch (err) {
-    console.error('[pi-mono-ws] writeRpcCommand error:', err);
-  }
-}
 
-function getOrCreateSession(sessionId: string): PiSession {
+async function spawnPi(session: PiSession, cwd: string) {
-  let session = sessions.get(sessionId);
-  if (!session) {
-    session = {
-      piProcess: null,
-      ws: null,
-      selectedModel: null,
-      cwd: null,
-      resourceChatDir: null,
-      logId: null,
-      fullText: '',
-      killTimer: null,
-    };
-    sessions.set(sessionId, session);
-  }
-  return session;
-}
-
-function attachWs(sessionId: string, ws: ServerWebSocket<WSData>) {
-  const session = getOrCreateSession(sessionId);
-
-  // Cancel any pending kill timer — the session is alive again
-  if (session.killTimer) {
-    clearTimeout(session.killTimer);
-    session.killTimer = null;
-  }
-
-  session.ws = ws;
-  wsToSession.set(ws, sessionId);
-}
-
-function detachWs(ws: ServerWebSocket<WSData>) {
-  const sessionId = wsToSession.get(ws);
-  wsToSession.delete(ws);
-  if (!sessionId) return;
-
-  const session = sessions.get(sessionId);
-  if (!session || session.ws !== ws) return;
-
-  // Detach ws but keep pi process alive for grace period
-  session.ws = null;
-
-  if (session.piProcess) {
-    session.killTimer = setTimeout(() => {
-      // If no new ws has attached, kill the process
-      if (!session.ws && session.piProcess) {
-        try {
-          session.piProcess.kill();
-        } catch {
-          // already dead
-        }
-        session.piProcess = null;
-        sessions.delete(sessionId);
-      }
-    }, ORPHAN_GRACE_MS);
-  } else {
-    sessions.delete(sessionId);
-  }
-}
-
-function resolveLocalModel(modelId: string): { providerId: string; modelName: string } | null {
-  if (!modelId.startsWith('local:')) return null;
-  const parts = modelId.split(':');
-  if (parts.length < 3) return null;
-  return { providerId: parts[1]!, modelName: parts.slice(2).join(':') };
-}
-
-async function handleLocalChat(session: PiSession, prompt: string) {
-  const parsed = session.selectedModel ? resolveLocalModel(session.selectedModel) : null;
-  if (!parsed) {
-    send(session.ws, { type: 'error', message: 'Invalid local model' });
-    return;
-  }
-
-  const providers = await readLocalProviders();
-  const provider = providers.find((p) => p.id === parsed.providerId);
-  if (!provider) {
-    send(session.ws, { type: 'error', message: 'Local provider not found' });
-    return;
-  }
-
-  const base = provider.url.replace(/\/+$/, '');
-  const url = `${base}/v1/chat/completions`;
-
-  const headers: Record<string, string> = { 'Content-Type': 'application/json' };
-  if (provider.auth?.type === 'api-key') {
-    headers['Authorization'] = `Bearer ${provider.auth.apiKey}`;
-  } else if (provider.auth?.type === 'basic') {
-    headers['Authorization'] = `Basic ${btoa(`${provider.auth.username}:${provider.auth.password}`)}`;
-  }
-
-  const body = JSON.stringify({
-    model: parsed.modelName,
-    messages: [{ role: 'user', content: prompt }],
-    stream: true,
-  });
-
-  try {
-    const res = await fetch(url, { method: 'POST', headers, body });
-    if (!res.ok) {
-      const text = await res.text();
-      send(session.ws, { type: 'error', message: `Local provider error: ${res.status} ${text}` });
-      return;
-    }
-
-    session.fullText = '';
-    const reader = res.body!.getReader();
-    const decoder = new TextDecoder();
-    let buffer = '';
-
-    while (true) {
-      const { done, value } = await reader.read();
-      if (done) break;
-
-      buffer += decoder.decode(value, { stream: true });
-      const lines = buffer.split('\n');
-      buffer = lines.pop() ?? '';
-
-      for (const line of lines) {
-        if (!line.startsWith('data: ')) continue;
-        const data = line.slice(6).trim();
-        if (data === '[DONE]') continue;
-
-        try {
-          const chunk = JSON.parse(data) as { choices?: { delta?: { content?: string } }[] };
-          const delta = chunk.choices?.[0]?.delta?.content;
-          if (delta) {
-            session.fullText += delta;
-            send(session.ws, { type: 'assistant:partial', text: delta });
-          }
-        } catch {
-          // skip unparseable chunks
-        }
-      }
-    }
-
-    if (session.fullText) {
-      send(session.ws, { type: 'assistant:text', text: session.fullText });
-      if (session.logId) appendToLog(session.logId, { role: 'assistant', text: session.fullText });
-      session.fullText = '';
-    }
-
-    send(session.ws, { type: 'result', costUsd: 0, durationMs: 0, numTurns: 0, isError: false });
-    if (session.logId) {
-      appendToLog(session.logId, { role: 'result', costUsd: 0, durationMs: 0, numTurns: 0, isError: false });
-      finalizeLog(session.logId);
-      session.logId = null;
-    }
-  } catch (err) {
-    send(session.ws, { type: 'error', message: `Local provider error: ${err}` });
-  }
-}
-
-async function spawnPiProcess(session: PiSession, workingDir: string) {
   const args = ['pi', '--mode', 'rpc', '--no-extensions', '--no-skills', '--no-prompt-templates', '--no-themes'];
-
+  if (session.model) args.push('--model', session.model);
-  if (session.selectedModel) {
-    args.push('--model', session.selectedModel);
-  }
 
   const storedKeys = await readApiKeys();
   const proc = Bun.spawn(args, {
-    cwd: workingDir,
+    cwd,
     stdin: 'pipe',
     stdout: 'pipe',
     stderr: 'pipe',
@@ -256,117 +162,120 @@ async function spawnPiProcess(session: PiSession, workingDir: string) {
 
   session.piProcess = proc;
 
-  // Read stdout line-by-line for JSON events
+  // Read stdout JSON event stream
   const reader = proc.stdout.getReader();
   const decoder = new TextDecoder();
   let buffer = '';
 
-  const readLoop = async () => {
+  (async () => {
     try {
       while (true) {
         const { done, value } = await reader.read();
         if (done) break;
-
         buffer += decoder.decode(value, { stream: true });
         const lines = buffer.split('\n');
         buffer = lines.pop() ?? '';
-
         for (const line of lines) {
           if (!line.trim()) continue;
           try {
-            const event = JSON.parse(line);
+            handlePiEvent(session, JSON.parse(line));
-            handlePiEvent(session, event);
+          } catch { /* skip unparseable */ }
-          } catch {
-            // skip unparseable lines
         }
       }
-      }
+    } catch { /* process ended */ }
-    } catch {
+  })();
-      // process ended
-    }
-  };
 
-  readLoop();
+  // Stderr → debug log
-
-  // Read stderr for debugging
   const stderrReader = proc.stderr.getReader();
   const stderrDecoder = new TextDecoder();
-  const readStderr = async () => {
+  (async () => {
     try {
       while (true) {
         const { done, value } = await stderrReader.read();
         if (done) break;
         const text = stderrDecoder.decode(value, { stream: true });
-        if (text.trim()) console.log('[pi-mono-ws] stderr:', text.trim());
+        if (text.trim()) console.log('[pi-mono] stderr:', text.trim());
       }
-    } catch {
+    } catch { /* process ended */ }
-      // process ended
+  })();
-    }
-  };
-  readStderr();
 
-  // Handle process exit
   proc.exited.then((code) => {
-    console.log(`[pi-mono-ws] pi process exited with code ${code}`);
+    console.log(`[pi-mono] process exited (code ${code}) for session ${session.id}`);
     if (session.piProcess === proc) {
       session.piProcess = null;
+      session.systemContextSent = false;
+      if (session.isGenerating) {
+        session.isGenerating = false;
+        if (session.streamBuffer) {
+          session.messages.push({ role: 'assistant', text: session.streamBuffer });
+          session.streamBuffer = '';
+        }
+        session.messages.push({ role: 'error', text: 'Pi process exited unexpectedly' });
+        persistMessages(session);
+        sendToClient(session, { type: 'error', message: 'Pi process exited unexpectedly' });
+      }
     }
   });
 }
 
+// --- Pi event handling ---
+
 function handlePiEvent(session: PiSession, event: Record<string, unknown>) {
   const type = event.type as string;
-  const ws = session.ws;
 
-  // RPC responses (type === 'response')
   if (type === 'response') {
-    const command = event.command as string;
+    if (event.command === 'prompt' && !event.success) {
-    if (command === 'prompt' && !event.success) {
+      const errorMsg = (event.error as string) ?? 'Prompt failed';
-      send(ws, { type: 'error', message: (event.error as string) ?? 'Prompt failed' });
+      sendToClient(session, { type: 'error', message: errorMsg });
+      session.messages.push({ role: 'error', text: errorMsg });
+      session.isGenerating = false;
+      persistMessages(session);
     }
     return;
   }
 
   switch (type) {
     case 'agent_start':
-      session.fullText = '';
+      session.streamBuffer = '';
+      session.isGenerating = true;
       break;
 
     case 'message_update': {
       const ame = event.assistantMessageEvent as Record<string, unknown> | undefined;
-      if (!ame) break;
+      if (ame?.type === 'text_delta') {
-
-      const ameType = ame.type as string;
-      if (ameType === 'text_delta') {
         const delta = ame.delta as string;
-        session.fullText += delta;
+        session.streamBuffer += delta;
-        send(ws, { type: 'assistant:partial', text: delta });
+        sendToClient(session, { type: 'assistant:partial', text: delta });
       }
       break;
     }
 
     case 'message_end': {
-      if (session.fullText) {
+      if (session.streamBuffer) {
-        send(ws, { type: 'assistant:text', text: session.fullText });
+        const text = session.streamBuffer;
-        if (session.logId) appendToLog(session.logId, { role: 'assistant', text: session.fullText });
+        session.streamBuffer = '';
-        session.fullText = '';
+        sendToClient(session, { type: 'assistant:text', text });
+        session.messages.push({ role: 'assistant', text });
+        persistMessages(session);
       }
       break;
     }
 
     case 'tool_execution_start': {
+      if (session.streamBuffer) {
+        const text = session.streamBuffer;
+        session.streamBuffer = '';
+        sendToClient(session, { type: 'assistant:text', text });
+        session.messages.push({ role: 'assistant', text });
+      }
+
       const toolCallId = (event.toolCallId as string) ?? '';
       const toolName = (event.toolName as string) ?? 'unknown';
       const args = (event.args as Record<string, unknown>) ?? {};
 
-      if (session.fullText) {
+      sendToClient(session, { type: 'tool:use', toolName, toolInput: args, toolUseId: toolCallId });
-        send(ws, { type: 'assistant:text', text: session.fullText });
+      session.messages.push({ role: 'tool', toolName, toolInput: args, toolUseId: toolCallId });
-        if (session.logId) appendToLog(session.logId, { role: 'assistant', text: session.fullText });
+      persistMessages(session);
-        session.fullText = '';
-      }
-
-      send(ws, { type: 'tool:use', toolName, toolInput: args, toolUseId: toolCallId });
-      if (session.logId) appendToLog(session.logId, { role: 'tool', toolName, toolInput: args, toolUseId: toolCallId });
       break;
     }
 
@@ -376,32 +285,32 @@ function handlePiEvent(session: PiSession, event: Record<string, unknown>) {
       const isError = (event.isError as boolean) ?? false;
       const output = result != null ? (typeof result === 'string' ? result : JSON.stringify(result)) : '';
 
-      send(ws, { type: 'tool:result', toolUseId: toolCallId, output, isError });
+      sendToClient(session, { type: 'tool:result', toolUseId: toolCallId, output, isError });
-      if (session.logId)
+
-        appendToLog(session.logId, {
+      for (let i = session.messages.length - 1; i >= 0; i--) {
-          role: 'tool',
+        const m = session.messages[i] as Record<string, unknown>;
-          toolName: '',
+        if (m.role === 'tool' && m.toolUseId === toolCallId) {
-          toolInput: {},
+          m.output = output;
-          toolUseId: toolCallId,
+          m.isError = isError;
-          output,
+          break;
-          isError,
+        }
-        });
+      }
+      persistMessages(session);
       break;
     }
 
     case 'agent_end': {
-      if (session.fullText) {
+      if (session.streamBuffer) {
-        send(ws, { type: 'assistant:text', text: session.fullText });
+        const text = session.streamBuffer;
-        if (session.logId) appendToLog(session.logId, { role: 'assistant', text: session.fullText });
+        session.streamBuffer = '';
-        session.fullText = '';
+        sendToClient(session, { type: 'assistant:text', text });
+        session.messages.push({ role: 'assistant', text });
       }
 
-      send(ws, { type: 'result', costUsd: 0, durationMs: 0, numTurns: 0, isError: false });
+      sendToClient(session, { type: 'result', costUsd: 0, durationMs: 0, numTurns: 0, isError: false });
-      if (session.logId) {
+      session.messages.push({ role: 'result', costUsd: 0, durationMs: 0, numTurns: 0, isError: false });
-        appendToLog(session.logId, { role: 'result', costUsd: 0, durationMs: 0, numTurns: 0, isError: false });
+      session.isGenerating = false;
-        finalizeLog(session.logId);
+      persistMessages(session);
-        session.logId = null;
-      }
       break;
     }
 
@@ -414,117 +323,129 @@ function handlePiEvent(session: PiSession, event: Record<string, unknown>) {
   }
 }
 
-type HandleChatParams = {
+// --- Session management ---
-  ws: ServerWebSocket<WSData>;
-  prompt: string;
-  sessionId?: string;
-  model?: string;
-  cwd?: { root?: string; path: string };
-  attachmentIds?: string[];
-  images?: ImageData[];
-  resourceChatDir?: string;
-  taskInfo?: TaskInfo;
-};
 
-async function handleChat({
+function attachWs(session: PiSession, ws: ServerWebSocket<WSData>) {
-  ws,
+  if (session.killTimer) {
-  prompt,
+    clearTimeout(session.killTimer);
-  sessionId,
+    session.killTimer = null;
-  model,
+  }
-  cwd,
+  session.ws = ws;
-  attachmentIds,
+  wsToSession.set(ws, session.id);
-  images,
-  resourceChatDir,
-  taskInfo,
-}: HandleChatParams) {
-  const email = ws.data.email;
-
-  // Determine or create session ID
-  let sid = sessionId ?? wsToSession.get(ws) ?? null;
-  let isNewSession = false;
-
-  if (!sid) {
-    sid = crypto.randomUUID();
-    isNewSession = true;
 }
 
-  // Attach this ws to the session (adopts existing pi process if any)
+function detachWs(ws: ServerWebSocket<WSData>) {
-  attachWs(sid, ws);
+  const sessionId = wsToSession.get(ws);
-  const session = getOrCreateSession(sid);
+  wsToSession.delete(ws);
+  if (!sessionId) return;
 
-  if (taskInfo && !session.logId) {
+  const session = sessions.get(sessionId);
-    session.logId = createTaskLog(email, taskInfo, 'pi-mono', model ?? 'unknown');
+  if (!session || session.ws !== ws) return;
-    appendToLog(session.logId, { role: 'user', text: prompt });
+
+  session.ws = null;
+
+  if (session.piProcess) {
+    session.killTimer = setTimeout(() => {
+      if (!session.ws && session.piProcess) {
+        try { session.piProcess.kill(); } catch { /* already dead */ }
+        session.piProcess = null;
+        session.systemContextSent = false;
+        sessions.delete(sessionId);
       }
-
+    }, ORPHAN_GRACE_MS);
-  if (resourceChatDir) session.resourceChatDir = resourceChatDir;
-  if (model) session.selectedModel = model;
-
-  if (isNewSession) {
-    const pendingTitle = prompt.slice(0, 100);
-
-    // Send session:init AFTER attaching ws so the pi process survives the reconnect
-    send(ws, { type: 'session:init', sessionId: sid, model: model ?? null });
-
-    if (session.resourceChatDir) {
-      const chatDir = join(session.resourceChatDir, 'chat');
-      const meta = { id: sid, model: model ?? null };
-      mkdir(chatDir, { recursive: true })
-        .then(() => Bun.write(join(chatDir, 'meta.json'), JSON.stringify(meta)))
-        .catch(() => {});
   } else {
-      const dir = getPiMonoSessionDir(email, sid);
+    sessions.delete(sessionId);
-      const meta = {
+  }
-        id: sid,
+}
-        title: pendingTitle,
-        createdAt: Date.now(),
-        model: model ?? null,
-      };
-      mkdir(dir, { recursive: true })
-        .then(() => Bun.write(join(dir, 'meta.json'), JSON.stringify(meta)))
-        .catch(() => {});
 
-      // Move tmp attachments to session dir
+// --- Handlers ---
-      if (attachmentIds?.length) {
+
+async function handleChat(ws: ServerWebSocket<WSData>, msg: Extract<ClientMessage, { type: 'chat' }>) {
+  const email = ws.data.email;
+  const prompt = msg.prompt;
+  const hasExistingSession = !!msg.sessionId;
+
+  const sid = msg.sessionId ?? crypto.randomUUID();
+
+  let session = sessions.get(sid);
+  if (!session) {
+    session = {
+      id: sid,
+      email,
+      piProcess: null,
+      ws: null,
+      model: typeof msg.model === 'string' ? msg.model : null,
+      cwd: null,
+      messages: [],
+      streamBuffer: '',
+      isGenerating: false,
+      systemContextSent: false,
+      killTimer: null,
+      saving: false,
+      dirty: false,
+    };
+    sessions.set(sid, session);
+
+    if (hasExistingSession) {
+      session.messages = await loadMessages(email, sid);
+    }
+  }
+
+  if (typeof msg.model === 'string') session.model = msg.model;
+  if (msg.cwd) session.cwd = join(resolveRootDir(email, msg.cwd.root), msg.cwd.path);
+
+  attachWs(session, ws);
+
+  if (!hasExistingSession) {
+    sendDirect(ws, { type: 'session:init', sessionId: sid, model: session.model });
+
+    const dir = getPiMonoSessionDir(email, sid);
+    const meta = { id: sid, title: prompt.slice(0, 100), createdAt: Date.now(), model: session.model };
+    await mkdir(dir, { recursive: true });
+    await Bun.write(join(dir, 'meta.json'), JSON.stringify(meta));
+
+    if (msg.attachmentIds?.length) {
       const tmpDir = getTmpAttachmentsDir(email);
       const destDir = getAttachmentsDir(email, 'pi-mono', sid);
-        mkdir(destDir, { recursive: true })
+      await mkdir(destDir, { recursive: true });
-          .then(() =>
+      await Promise.all(msg.attachmentIds.map((id) => rename(join(tmpDir, id), join(destDir, id)).catch(() => {})));
-            Promise.all(attachmentIds.map((id) => rename(join(tmpDir, id), join(destDir, id)).catch(() => {}))),
-          )
-          .catch(() => {});
-      }
     }
   }
 
-  // Local provider models — bypass pi, call API directly
+  session.messages.push({ role: 'user', text: prompt });
-  if (session.selectedModel?.startsWith('local:')) {
+  persistMessages(session);
-    handleLocalChat(session, prompt);
-    return;
-  }
 
-  // Ensure pi process is running
+  const workingDir = session.cwd ?? getHomeDir(email);
-  const homeDir = getHomeDir(email);
-  if (cwd) session.cwd = join(resolveRootDir(email, cwd.root), cwd.path);
-  const workingDir = session.cwd ?? homeDir;
 
-  if (!session.piProcess) {
+  const needsSpawn = !session.piProcess;
-    await spawnPiProcess(session, workingDir);
+  if (needsSpawn) {
-    // Give pi a moment to initialize
+    session.systemContextSent = false;
+    await spawnPi(session, workingDir);
     await new Promise((r) => setTimeout(r, 500));
   }
 
   if (!session.piProcess) {
-    send(session.ws, { type: 'error', message: 'Failed to start pi process' });
+    sendToClient(session, { type: 'error', message: 'Failed to start pi process' });
     return;
   }
 
-  // Build context and send prompt
+  let fullPrompt: string;
-  const skillsAppend = await buildSkillsPrompt(email);
+  if (!session.systemContextSent) {
-  const contextAppend = `\n\nThe user's home directory is: ${homeDir}` + skillsAppend;
+    const homeDir = getHomeDir(email);
-  send(session.ws, { type: 'system:prompt', text: contextAppend });
+    const skillsPrompt = await buildSkillsPrompt(email);
+    let systemContext = `\nThe user's home directory is: ${homeDir}${skillsPrompt}`;
 
-  const fullPrompt = `<system>${contextAppend}</system>\n\n${prompt}`;
+    if (session.messages.length > 1) {
+      const historyMsgs = session.messages.slice(0, -1);
+      const history = buildHistoryContext(historyMsgs);
+      if (history) systemContext += history;
+    }
+
+    fullPrompt = `<system>${systemContext}</system>\n\n${prompt}`;
+    session.systemContextSent = true;
+  } else {
+    fullPrompt = prompt;
+  }
 
   writeRpcCommand(session.piProcess, {
     type: 'prompt',
@@ -533,6 +454,19 @@ async function handleChat({
   });
 }
 
+function handleResume(ws: ServerWebSocket<WSData>, sessionId: string) {
+  const session = sessions.get(sessionId);
+  if (session) {
+    attachWs(session, ws);
+    sendDirect(ws, {
+      type: 'messages:sync',
+      messages: session.messages,
+      streamingText: session.streamBuffer,
+      isGenerating: session.isGenerating,
+    } as ServerMessage);
+  }
+}
+
 function handleStop(ws: ServerWebSocket<WSData>) {
   const sessionId = wsToSession.get(ws);
   if (!sessionId) return;
@@ -540,39 +474,33 @@ function handleStop(ws: ServerWebSocket<WSData>) {
   if (!session?.piProcess) return;
 
   writeRpcCommand(session.piProcess, { type: 'abort', id: `abort_${Date.now()}` });
-  send(ws, { type: 'stopped' });
+  sendToClient(session, { type: 'stopped' });
 }
 
+// --- Export ---
+
 export const piMonoWebsocket = {
-  open(ws: ServerWebSocket<WSData>) {
+  open(_ws: ServerWebSocket<WSData>) {},
-    // Nothing to do — session is attached when a chat message arrives
-  },
 
   message(ws: ServerWebSocket<WSData>, raw: string | Buffer) {
     let msg: ClientMessage;
     try {
       msg = JSON.parse(typeof raw === 'string' ? raw : raw.toString()) as ClientMessage;
     } catch {
-      send(ws, { type: 'error', message: 'Invalid JSON' });
+      sendDirect(ws, { type: 'error', message: 'Invalid JSON' });
       return;
     }
 
-    if (msg.type === 'chat') {
+    switch (msg.type) {
-      handleChat({
+      case 'chat':
-        ws,
+        handleChat(ws, msg);
-        prompt: msg.prompt,
+        break;
-        sessionId: msg.sessionId,
+      case 'resume':
-        model: typeof msg.model === 'string' ? msg.model : undefined,
+        handleResume(ws, msg.sessionId);
-        cwd: msg.cwd,
+        break;
-        attachmentIds: msg.attachmentIds,
+      case 'stop':
-        images: msg.images,
-        resourceChatDir: msg.resourceChatDir,
-        taskInfo: msg.taskInfo,
-      });
-    } else if (msg.type === 'resume') {
-      attachWs(msg.sessionId, ws);
-    } else if (msg.type === 'stop') {
         handleStop(ws);
+        break;
     }
   },
 

src/workspaces/apps/Chat/MessageBubble.tsx

@@ -53,12 +53,14 @@ export const MessageBubble = ({ message, onAnswer }: MessageBubbleProps) => {
       if (!message.text) return null;
       return (
         <div className="flex justify-start">
-          <div className="max-w-[85%] rounded-2xl rounded-tl-sm bg-background/80 border border-duck-dark/10 px-4 py-2.5 text-sm text-duck-dark prose prose-sm dark:prose-invert max-w-none prose-pre:bg-gray-900 prose-pre:text-green-400 prose-code:text-duck-teal prose-code:before:content-none prose-code:after:content-none">
+          <div className="max-w-[85%] rounded-2xl rounded-tl-sm bg-muted/50 border border-border/50 px-4 py-2.5">
+            <div className="chat-md">
               <ReactMarkdown remarkPlugins={[remarkGfm]} rehypePlugins={[rehypeRaw]}>
                 {message.text}
               </ReactMarkdown>
             </div>
           </div>
+        </div>
       );
 
     case 'tool':
@@ -97,12 +99,14 @@ export const StreamingBubble = ({ text }: StreamingBubbleProps) => {
   if (!text) return null;
   return (
     <div className="flex justify-start">
-      <div className="max-w-[85%] rounded-2xl rounded-tl-sm bg-background/80 border border-duck-dark/10 px-4 py-2.5 text-sm text-duck-dark prose prose-sm dark:prose-invert max-w-none prose-pre:bg-gray-900 prose-pre:text-green-400 prose-code:text-duck-teal prose-code:before:content-none prose-code:after:content-none">
+      <div className="max-w-[85%] rounded-2xl rounded-tl-sm bg-muted/50 border border-border/50 px-4 py-2.5">
+        <div className="chat-md">
           <ReactMarkdown remarkPlugins={[remarkGfm]} rehypePlugins={[rehypeRaw]}>
             {text}
           </ReactMarkdown>
           <span className="inline-block w-2 h-4 bg-duck-teal/60 animate-pulse ml-0.5 align-middle" />
         </div>
       </div>
+    </div>
   );
 };

src/workspaces/apps/Chat/types.ts

@@ -37,4 +37,5 @@ export type ServerMessage =
   | { type: 'tool:result'; toolUseId: string; output: string; isError: boolean }
   | { type: 'result'; costUsd: number; durationMs: number; numTurns: number; isError: boolean }
   | { type: 'error'; message: string }
-  | { type: 'stopped' };
+  | { type: 'stopped' }
+  | { type: 'messages:sync'; messages: ChatMessage[]; streamingText: string; isGenerating: boolean };