|
|
|
@@ -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<WSData> | 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<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) {
|
|
|
|
|
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<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>) {
|
|
|
|
|
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<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) {
|
|
|
|
|
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<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 (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<string, unknown> | 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<string, unknown>) ?? {};
|
|
|
|
|
|
|
|
|
|
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<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 });
|
|
|
|
|
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<string, unknown>;
|
|
|
|
|
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<string, unknown>) {
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
type HandleChatParams = {
|
|
|
|
|
ws: ServerWebSocket<WSData>;
|
|
|
|
|
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<WSData>) {
|
|
|
|
|
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<WSData>) {
|
|
|
|
|
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<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);
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
// 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 = `<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>) {
|
|
|
|
|
// Nothing to do — session is attached when a chat message arrives
|
|
|
|
|
},
|
|
|
|
|
open(_ws: ServerWebSocket<WSData>) {},
|
|
|
|
|
|
|
|
|
|
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') {
|
|
|
|
|
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;
|
|
|
|
|
}
|
|
|
|
|
},
|
|
|
|
|
|
|
|
|
|
