From ef1530b6264d308ddf108a4c8926d79fd4d6f5b7 Mon Sep 17 00:00:00 2001 From: Andre Padez Date: Fri, 20 Feb 2026 18:25:15 +0000 Subject: [PATCH] also no idea, from monorepo --- .../Screens/Dashboard/Chat/usePiMono.ts | 46 +- src/apps/officer-web/styles/prose.css | 154 ++++ src/servers/api/chat-types.ts | 3 +- src/servers/api/pi-mono/websocket.ts | 664 ++++++++---------- src/workspaces/apps/Chat/MessageBubble.tsx | 22 +- src/workspaces/apps/Chat/types.ts | 3 +- 6 files changed, 473 insertions(+), 419 deletions(-) diff --git a/src/apps/officer-web/Screens/Dashboard/Chat/usePiMono.ts b/src/apps/officer-web/Screens/Dashboard/Chat/usePiMono.ts index e179b75f..cd3e6703 100644 --- a/src/apps/officer-web/Screens/Dashboard/Chat/usePiMono.ts +++ b/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(null); const sessionIdRef = useRef(initialSessionId ?? null); - const saveTimerRef = useRef(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': - setMessages((prev) => [...prev, { role: 'system', text: msg.text }]); + case 'messages:sync': + 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 } : {}), diff --git a/src/apps/officer-web/styles/prose.css b/src/apps/officer-web/styles/prose.css index 1930e291..a45c32dc 100644 --- a/src/apps/officer-web/styles/prose.css +++ b/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)); diff --git a/src/servers/api/chat-types.ts b/src/servers/api/chat-types.ts index fc8ade6a..21246360 100644 --- a/src/servers/api/chat-types.ts +++ b/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 }; diff --git a/src/servers/api/pi-mono/websocket.ts b/src/servers/api/pi-mono/websocket.ts index 6b7761e1..680b3721 100644 --- a/src/servers/api/pi-mono/websocket.ts +++ b/src/servers/api/pi-mono/websocket.ts @@ -1,47 +1,52 @@ -import type { ServerWebSocket } from 'bun'; -import type { Subprocess } from 'bun'; +import type { ServerWebSocket, 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 { createTaskLog, appendToLog, finalizeLog } from '@@/api/task-logger'; -import { readApiKeys, readLocalProviders } from '@@/api/server-settings/pi-mono'; +import type { ClientMessage, ServerMessage } from '@@/api/chat-types'; +import { readApiKeys } 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 | null; - selectedModel: string | null; + model: string | null; cwd: string | null; - resourceChatDir: string | null; - logId: string | null; - fullText: string; + messages: unknown[]; + streamBuffer: string; + isGenerating: boolean; + systemContextSent: boolean; killTimer: ReturnType | null; + saving: boolean; + dirty: boolean; }; -// Session pool — pi processes persist across websocket reconnects const sessions = new Map(); - -// Map ws → sessionId for quick lookup on close const wsToSession = new Map, string>(); - -// Grace period before killing orphaned pi processes (ms) const ORPHAN_GRACE_MS = 30_000; -function send(ws: ServerWebSocket | null, msg: ServerMessage) { - if (ws && ws.readyState === 1) ws.send(JSON.stringify(msg)); +// --- Helpers --- + +function sendToClient(session: PiSession, msg: ServerMessage) { + if (session.ws?.readyState === 1) session.ws.send(JSON.stringify(msg)); +} + +function sendDirect(ws: ServerWebSocket, 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) { + 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 { + 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; + 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\n${history}\n`; +} + +// --- Skills --- + async function buildSkillsPrompt(email: string): Promise { const nativeSkills = await readSkillDirs(getNativeSkillsDir()); const globalSkills = await readSkillDirs(getGlobalSkillsDir()); @@ -74,180 +145,15 @@ async function buildSkillsPrompt(email: string): Promise { 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) { - 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); - } -} +// --- Pi process lifecycle --- -function getOrCreateSession(sessionId: string): PiSession { - 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) { - 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) { - 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 = { '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) { +async function spawnPi(session: PiSession, cwd: string) { const args = ['pi', '--mode', 'rpc', '--no-extensions', '--no-skills', '--no-prompt-templates', '--no-themes']; - - if (session.selectedModel) { - args.push('--model', session.selectedModel); - } + if (session.model) args.push('--model', session.model); 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, event); - } catch { - // skip unparseable lines - } + handlePiEvent(session, JSON.parse(line)); + } catch { /* skip unparseable */ } } } - } catch { - // process ended - } - }; + } catch { /* process ended */ } + })(); - readLoop(); - - // Read stderr for debugging + // Stderr → debug log 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 { - // process ended - } - }; - readStderr(); + } catch { /* process ended */ } + })(); - // 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) { const type = event.type as string; - const ws = session.ws; - // RPC responses (type === 'response') if (type === 'response') { - const command = event.command as string; - if (command === 'prompt' && !event.success) { - send(ws, { type: 'error', message: (event.error as string) ?? 'Prompt failed' }); + if (event.command === 'prompt' && !event.success) { + const errorMsg = (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 | undefined; - if (!ame) break; - - const ameType = ame.type as string; - if (ameType === 'text_delta') { + if (ame?.type === 'text_delta') { const delta = ame.delta as string; - session.fullText += delta; - send(ws, { type: 'assistant:partial', text: delta }); + session.streamBuffer += delta; + sendToClient(session, { type: 'assistant:partial', text: delta }); } break; } case 'message_end': { - if (session.fullText) { - send(ws, { type: 'assistant:text', text: session.fullText }); - if (session.logId) appendToLog(session.logId, { role: 'assistant', text: session.fullText }); - session.fullText = ''; + if (session.streamBuffer) { + const text = session.streamBuffer; + session.streamBuffer = ''; + 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) ?? {}; - if (session.fullText) { - send(ws, { type: 'assistant:text', text: session.fullText }); - if (session.logId) appendToLog(session.logId, { role: 'assistant', text: session.fullText }); - 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 }); + sendToClient(session, { type: 'tool:use', toolName, toolInput: args, toolUseId: toolCallId }); + session.messages.push({ role: 'tool', toolName, toolInput: args, toolUseId: toolCallId }); + persistMessages(session); break; } @@ -376,32 +285,32 @@ function handlePiEvent(session: PiSession, event: Record) { 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 }); - if (session.logId) - appendToLog(session.logId, { - role: 'tool', - toolName: '', - toolInput: {}, - toolUseId: toolCallId, - output, - isError, - }); + sendToClient(session, { type: 'tool:result', toolUseId: toolCallId, output, isError }); + + for (let i = session.messages.length - 1; i >= 0; i--) { + const m = session.messages[i] as Record; + if (m.role === 'tool' && m.toolUseId === toolCallId) { + m.output = output; + m.isError = isError; + break; + } + } + persistMessages(session); break; } case 'agent_end': { - if (session.fullText) { - send(ws, { type: 'assistant:text', text: session.fullText }); - if (session.logId) appendToLog(session.logId, { role: 'assistant', text: session.fullText }); - session.fullText = ''; + if (session.streamBuffer) { + const text = session.streamBuffer; + session.streamBuffer = ''; + sendToClient(session, { type: 'assistant:text', text }); + session.messages.push({ role: 'assistant', text }); } - send(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; - } + sendToClient(session, { type: 'result', costUsd: 0, durationMs: 0, numTurns: 0, isError: false }); + session.messages.push({ role: 'result', costUsd: 0, durationMs: 0, numTurns: 0, isError: false }); + session.isGenerating = false; + persistMessages(session); break; } @@ -414,117 +323,129 @@ function handlePiEvent(session: PiSession, event: Record) { } } -type HandleChatParams = { - ws: ServerWebSocket; - prompt: string; - sessionId?: string; - model?: string; - cwd?: { root?: string; path: string }; - attachmentIds?: string[]; - images?: ImageData[]; - resourceChatDir?: string; - taskInfo?: TaskInfo; -}; +// --- Session management --- -async function handleChat({ - ws, - prompt, - sessionId, - model, - cwd, - attachmentIds, - 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; +function attachWs(session: PiSession, ws: ServerWebSocket) { + if (session.killTimer) { + clearTimeout(session.killTimer); + session.killTimer = null; } + session.ws = ws; + wsToSession.set(ws, session.id); +} - // Attach this ws to the session (adopts existing pi process if any) - attachWs(sid, ws); - const session = getOrCreateSession(sid); +function detachWs(ws: ServerWebSocket) { + const sessionId = wsToSession.get(ws); + wsToSession.delete(ws); + if (!sessionId) return; - if (taskInfo && !session.logId) { - session.logId = createTaskLog(email, taskInfo, 'pi-mono', model ?? 'unknown'); - appendToLog(session.logId, { role: 'user', text: prompt }); - } + const session = sessions.get(sessionId); + if (!session || session.ws !== ws) return; - if (resourceChatDir) session.resourceChatDir = resourceChatDir; - if (model) session.selectedModel = model; + session.ws = null; - 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); - 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 - if (attachmentIds?.length) { - const tmpDir = getTmpAttachmentsDir(email); - const destDir = getAttachmentsDir(email, 'pi-mono', sid); - mkdir(destDir, { recursive: true }) - .then(() => - Promise.all(attachmentIds.map((id) => rename(join(tmpDir, id), join(destDir, id)).catch(() => {}))), - ) - .catch(() => {}); + 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); + } else { + sessions.delete(sessionId); + } +} + +// --- Handlers --- + +async function handleChat(ws: ServerWebSocket, msg: Extract) { + 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); } } - // Local provider models — bypass pi, call API directly - if (session.selectedModel?.startsWith('local:')) { - handleLocalChat(session, prompt); - return; + 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); + await mkdir(destDir, { recursive: true }); + await Promise.all(msg.attachmentIds.map((id) => rename(join(tmpDir, id), join(destDir, id)).catch(() => {}))); + } } - // Ensure pi process is running - const homeDir = getHomeDir(email); - if (cwd) session.cwd = join(resolveRootDir(email, cwd.root), cwd.path); - const workingDir = session.cwd ?? homeDir; + session.messages.push({ role: 'user', text: prompt }); + persistMessages(session); - if (!session.piProcess) { - await spawnPiProcess(session, workingDir); - // Give pi a moment to initialize + const workingDir = session.cwd ?? getHomeDir(email); + + const needsSpawn = !session.piProcess; + if (needsSpawn) { + 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 - const skillsAppend = await buildSkillsPrompt(email); - const contextAppend = `\n\nThe user's home directory is: ${homeDir}` + skillsAppend; - send(session.ws, { type: 'system:prompt', text: contextAppend }); + let fullPrompt: string; + if (!session.systemContextSent) { + const homeDir = getHomeDir(email); + const skillsPrompt = await buildSkillsPrompt(email); + let systemContext = `\nThe user's home directory is: ${homeDir}${skillsPrompt}`; - const fullPrompt = `${contextAppend}\n\n${prompt}`; + if (session.messages.length > 1) { + const historyMsgs = session.messages.slice(0, -1); + const history = buildHistoryContext(historyMsgs); + if (history) systemContext += history; + } + + fullPrompt = `${systemContext}\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, 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) { const sessionId = wsToSession.get(ws); if (!sessionId) return; @@ -540,39 +474,33 @@ function handleStop(ws: ServerWebSocket) { 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) { - // Nothing to do — session is attached when a chat message arrives - }, + open(_ws: ServerWebSocket) {}, message(ws: ServerWebSocket, 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') { - handleChat({ - ws, - prompt: msg.prompt, - sessionId: msg.sessionId, - model: typeof msg.model === 'string' ? msg.model : undefined, - cwd: msg.cwd, - attachmentIds: msg.attachmentIds, - 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); + switch (msg.type) { + case 'chat': + handleChat(ws, msg); + break; + case 'resume': + handleResume(ws, msg.sessionId); + break; + case 'stop': + handleStop(ws); + break; } }, diff --git a/src/workspaces/apps/Chat/MessageBubble.tsx b/src/workspaces/apps/Chat/MessageBubble.tsx index 7b8c2836..a816f2ee 100644 --- a/src/workspaces/apps/Chat/MessageBubble.tsx +++ b/src/workspaces/apps/Chat/MessageBubble.tsx @@ -53,10 +53,12 @@ export const MessageBubble = ({ message, onAnswer }: MessageBubbleProps) => { if (!message.text) return null; return (
-
- - {message.text} - +
+
+ + {message.text} + +
); @@ -97,11 +99,13 @@ export const StreamingBubble = ({ text }: StreamingBubbleProps) => { if (!text) return null; return (
-
- - {text} - - +
+
+ + {text} + + +
); diff --git a/src/workspaces/apps/Chat/types.ts b/src/workspaces/apps/Chat/types.ts index 47696b5d..9c775662 100644 --- a/src/workspaces/apps/Chat/types.ts +++ b/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 };