also no idea, from monorepo
This commit is contained in:
@@ -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 };
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
},
|
||||
|
||||
|
||||
Reference in New Issue
Block a user