fix: Pi harness - dynamic models, correct RPC protocol, proper event handling
Backend: - /api/pi/models now calls 'pi --list-models' with stored API keys - pi-bridge.ts: callback-based event handling (matches pi-monorepo) - pi-bridge.ts: correct RPC format (type: 'prompt' not jsonrpc) - pi-bridge.ts: pass API keys to Pi process env - websocket.ts: event handler runs in background, no blocking - rest.ts: fix user home path (getHomeDir instead of hardcoded) Frontend: - Fix /api/ double prefix in useChatSessions, useChatGroups, useModels - Add PROVIDER_DISPLAY mapping in SystemSettings.tsx - Provider tabs show friendly names (e.g., 'OpenCode Zen') UI (from previous session): - Grouped session list with collapsible folders - CreateGroupDialog, GroupContextMenu, SessionContextMenu components
This commit is contained in:
+166
-228
@@ -1,16 +1,12 @@
|
||||
import type { ServerWebSocket } from 'bun';
|
||||
import { randomUUID } from 'crypto';
|
||||
import type { ClientMessage, ServerMessage, Message } from './types';
|
||||
import type { ClientMessage, ServerMessage, Message, PiEvent } from './types';
|
||||
import { sessionManager } from './session-manager';
|
||||
import * as storage from './storage';
|
||||
import * as piBridge from './pi-bridge';
|
||||
import { getHomeDir } from '../../../servers/data-path';
|
||||
import { logger } from './logger';
|
||||
|
||||
/**
|
||||
* WebSocket Handler — Pi chat session lifecycle
|
||||
*/
|
||||
|
||||
type WSData = {
|
||||
userId: number;
|
||||
email: string;
|
||||
@@ -20,46 +16,18 @@ type WSData = {
|
||||
|
||||
const IDLE_TIMEOUT_MS = 60 * 60 * 1000; // 1 hour
|
||||
|
||||
// Track which session is attached to which WebSocket
|
||||
const wsToSessionMap = new WeakMap<any, string>();
|
||||
|
||||
/**
|
||||
* Build system prompt with conversation history
|
||||
*/
|
||||
function buildSystemPrompt(messages: Message[], homeDir: string, skills: string): string {
|
||||
const history = messages
|
||||
.map((msg) => {
|
||||
if (msg.role === 'user') return `user: ${msg.text}`;
|
||||
if (msg.role === 'assistant') return `assistant: ${msg.text}`;
|
||||
if (msg.role === 'tool') return `tool(${msg.toolName}): ${msg.output}`;
|
||||
return '';
|
||||
})
|
||||
.filter(Boolean)
|
||||
.join('\n');
|
||||
|
||||
return `
|
||||
<system>
|
||||
User home directory: ${homeDir}
|
||||
${skills}
|
||||
|
||||
Below is the conversation history from this session:
|
||||
<conversation_history>
|
||||
${history}
|
||||
</conversation_history>
|
||||
</system>
|
||||
`;
|
||||
function sendToClient(ws: ServerWebSocket<WSData> | null, msg: ServerMessage): void {
|
||||
if (ws?.readyState === 1) {
|
||||
ws.send(JSON.stringify(msg));
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* WebSocket open handler
|
||||
*/
|
||||
export async function open(ws: ServerWebSocket<WSData>): Promise<void> {
|
||||
logger.info('WebSocket connection opened', { email: ws.data.email });
|
||||
}
|
||||
|
||||
/**
|
||||
* WebSocket message handler
|
||||
*/
|
||||
export function message(ws: ServerWebSocket<WSData>, raw: string | Buffer): void {
|
||||
const data = typeof raw === 'string' ? raw : raw.toString();
|
||||
|
||||
@@ -76,18 +44,11 @@ export function message(ws: ServerWebSocket<WSData>, raw: string | Buffer): void
|
||||
}
|
||||
} catch (err) {
|
||||
logger.error('Error handling WebSocket message', { email: ws.data.email, error: String(err) });
|
||||
const errorMsg: ServerMessage = {
|
||||
type: 'error',
|
||||
message: 'Failed to process message',
|
||||
};
|
||||
ws.send(JSON.stringify(errorMsg));
|
||||
sendToClient(ws, { type: 'error', message: 'Failed to process message' });
|
||||
}
|
||||
})();
|
||||
}
|
||||
|
||||
/**
|
||||
* WebSocket close handler
|
||||
*/
|
||||
export function close(ws: ServerWebSocket<WSData>): void {
|
||||
logger.info('WebSocket connection closed', { email: ws.data.email });
|
||||
|
||||
@@ -95,49 +56,176 @@ export function close(ws: ServerWebSocket<WSData>): void {
|
||||
if (sessionId) {
|
||||
sessionManager.detachWs(sessionId);
|
||||
sessionManager.setIdleTimeout(sessionId, IDLE_TIMEOUT_MS);
|
||||
logger.info('Started idle timeout for session', { sessionId, timeoutMs: IDLE_TIMEOUT_MS });
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Handle chat message
|
||||
*/
|
||||
function createEventHandler(sessionId: string, model: string, cwd: string) {
|
||||
return async (event: PiEvent): Promise<void> => {
|
||||
const session = sessionManager.getSession(sessionId);
|
||||
if (!session) return;
|
||||
|
||||
const ws = session.ws as ServerWebSocket<WSData> | null;
|
||||
|
||||
switch (event.type) {
|
||||
case 'delta': {
|
||||
sendToClient(ws, { type: 'assistant:delta', text: event.text });
|
||||
session.streamBuffer += event.text;
|
||||
break;
|
||||
}
|
||||
|
||||
case 'text': {
|
||||
// Flush streaming buffer as complete text
|
||||
const text = event.text || session.streamBuffer;
|
||||
if (text) {
|
||||
sendToClient(ws, { type: 'assistant:text', text });
|
||||
|
||||
const assistantMsg: Message = {
|
||||
id: randomUUID(),
|
||||
timestamp: Date.now(),
|
||||
role: 'assistant',
|
||||
text,
|
||||
model,
|
||||
};
|
||||
session.messages.push(assistantMsg);
|
||||
session.meta.messageCount += 1;
|
||||
session.streamBuffer = '';
|
||||
}
|
||||
break;
|
||||
}
|
||||
|
||||
case 'tool:start': {
|
||||
// Flush any pending streaming text first
|
||||
if (session.streamBuffer) {
|
||||
sendToClient(ws, { type: 'assistant:text', text: session.streamBuffer });
|
||||
|
||||
const assistantMsg: Message = {
|
||||
id: randomUUID(),
|
||||
timestamp: Date.now(),
|
||||
role: 'assistant',
|
||||
text: session.streamBuffer,
|
||||
model,
|
||||
};
|
||||
session.messages.push(assistantMsg);
|
||||
session.meta.messageCount += 1;
|
||||
session.streamBuffer = '';
|
||||
}
|
||||
|
||||
sendToClient(ws, {
|
||||
type: 'tool:start',
|
||||
toolCallId: event.toolCallId,
|
||||
toolName: event.toolName,
|
||||
toolInput: event.toolInput,
|
||||
});
|
||||
|
||||
const toolMsg: Message = {
|
||||
id: randomUUID(),
|
||||
timestamp: Date.now(),
|
||||
role: 'tool',
|
||||
toolCallId: event.toolCallId,
|
||||
toolName: event.toolName,
|
||||
toolInput: event.toolInput,
|
||||
};
|
||||
session.messages.push(toolMsg);
|
||||
session.meta.messageCount += 1;
|
||||
break;
|
||||
}
|
||||
|
||||
case 'tool:result': {
|
||||
sendToClient(ws, {
|
||||
type: 'tool:result',
|
||||
toolCallId: event.toolCallId,
|
||||
output: event.output,
|
||||
isError: event.isError,
|
||||
});
|
||||
|
||||
// Update existing tool message with output
|
||||
for (let i = session.messages.length - 1; i >= 0; i--) {
|
||||
const m = session.messages[i]!;
|
||||
if (m.role === 'tool' && m.toolCallId === event.toolCallId) {
|
||||
m.output = event.output;
|
||||
m.isError = event.isError;
|
||||
break;
|
||||
}
|
||||
}
|
||||
break;
|
||||
}
|
||||
|
||||
case 'result': {
|
||||
// Flush any remaining streaming buffer
|
||||
if (session.streamBuffer) {
|
||||
sendToClient(ws, { type: 'assistant:text', text: session.streamBuffer });
|
||||
|
||||
const assistantMsg: Message = {
|
||||
id: randomUUID(),
|
||||
timestamp: Date.now(),
|
||||
role: 'assistant',
|
||||
text: session.streamBuffer,
|
||||
model,
|
||||
cost: event.cost,
|
||||
};
|
||||
session.messages.push(assistantMsg);
|
||||
session.meta.messageCount += 1;
|
||||
session.streamBuffer = '';
|
||||
}
|
||||
|
||||
sendToClient(ws, { type: 'result', sessionId, cost: event.cost });
|
||||
|
||||
session.isGenerating = false;
|
||||
session.meta.cost.inputTokens += event.cost.inputTokens;
|
||||
session.meta.cost.outputTokens += event.cost.outputTokens;
|
||||
session.meta.cost.totalUSD += event.cost.totalUSD;
|
||||
session.meta.updatedAt = Date.now();
|
||||
|
||||
// Save session to disk
|
||||
try {
|
||||
await storage.saveSession(cwd, sessionId, session.meta, session.messages);
|
||||
logger.info('Session saved to disk', { sessionId, messageCount: session.messages.length });
|
||||
} catch (err) {
|
||||
logger.error('Failed to save session', { sessionId, error: String(err) });
|
||||
}
|
||||
break;
|
||||
}
|
||||
|
||||
case 'error': {
|
||||
sendToClient(ws, { type: 'error', message: event.message });
|
||||
session.isGenerating = false;
|
||||
break;
|
||||
}
|
||||
|
||||
case 'stopped': {
|
||||
sendToClient(ws, { type: 'stopped' });
|
||||
session.isGenerating = false;
|
||||
break;
|
||||
}
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
async function handleChat(
|
||||
ws: ServerWebSocket<WSData>,
|
||||
msg: { prompt: string; sessionId?: string; model?: string; cwd?: string; groupSlug?: string; attachmentIds?: string[] }
|
||||
): Promise<void> {
|
||||
const { email } = ws.data;
|
||||
const sessionId = msg.sessionId || randomUUID();
|
||||
const model = msg.model || 'gpt-4o'; // Default model
|
||||
const model = msg.model || 'opencode/big-pickle';
|
||||
const cwd = msg.cwd || getHomeDir(email);
|
||||
const groupSlug = msg.groupSlug || null;
|
||||
|
||||
// Get or create session
|
||||
const session = sessionManager.getOrCreate(sessionId, email, cwd, model, groupSlug);
|
||||
sessionManager.attachWs(sessionId, ws);
|
||||
wsToSessionMap.set(ws as any, sessionId);
|
||||
|
||||
// Send session init
|
||||
const initMsg: ServerMessage = {
|
||||
type: 'session:init',
|
||||
sessionId,
|
||||
model,
|
||||
cwd,
|
||||
};
|
||||
ws.send(JSON.stringify(initMsg));
|
||||
sendToClient(ws, { type: 'session:init', sessionId, model, cwd });
|
||||
|
||||
// Spawn Pi process if not already running
|
||||
if (!session.piProcess) {
|
||||
try {
|
||||
session.piProcess = await piBridge.spawnPi(cwd, model);
|
||||
const onEvent = createEventHandler(sessionId, model, cwd);
|
||||
session.piProcess = await piBridge.spawnPi(cwd, model, onEvent);
|
||||
logger.info('Spawned Pi process for session', { sessionId, model, cwd });
|
||||
} catch (err) {
|
||||
logger.error('Failed to spawn Pi process', { sessionId, model, error: String(err) });
|
||||
const errorMsg: ServerMessage = {
|
||||
type: 'error',
|
||||
message: 'Failed to start Pi process',
|
||||
};
|
||||
ws.send(JSON.stringify(errorMsg));
|
||||
sendToClient(ws, { type: 'error', message: 'Failed to start Pi process' });
|
||||
return;
|
||||
}
|
||||
}
|
||||
@@ -153,7 +241,6 @@ async function handleChat(
|
||||
session.meta.messageCount += 1;
|
||||
session.meta.updatedAt = Date.now();
|
||||
|
||||
// Initialize title from first message
|
||||
if (!session.meta.title) {
|
||||
session.meta.title = msg.prompt.slice(0, 100);
|
||||
}
|
||||
@@ -162,117 +249,8 @@ async function handleChat(
|
||||
const requestId = randomUUID();
|
||||
session.isGenerating = true;
|
||||
piBridge.sendPrompt(session.piProcess, msg.prompt, requestId);
|
||||
|
||||
// Track assistant message and cost
|
||||
let assistantText = '';
|
||||
let assistantMsgId = randomUUID();
|
||||
|
||||
// Stream responses from Pi
|
||||
try {
|
||||
for await (const event of piBridge.readEvents(session.piProcess)) {
|
||||
if (event.type === 'text') {
|
||||
const textMsg: ServerMessage = {
|
||||
type: 'assistant:text',
|
||||
text: event.text,
|
||||
};
|
||||
ws.send(JSON.stringify(textMsg));
|
||||
assistantText = event.text;
|
||||
} else if (event.type === 'delta') {
|
||||
const deltaMsg: ServerMessage = {
|
||||
type: 'assistant:delta',
|
||||
text: event.text,
|
||||
};
|
||||
ws.send(JSON.stringify(deltaMsg));
|
||||
session.streamBuffer += event.text;
|
||||
} else if (event.type === 'tool:start') {
|
||||
const toolStartMsg: ServerMessage = {
|
||||
type: 'tool:start',
|
||||
toolCallId: event.toolCallId,
|
||||
toolName: event.toolName,
|
||||
toolInput: event.toolInput,
|
||||
};
|
||||
ws.send(JSON.stringify(toolStartMsg));
|
||||
} else if (event.type === 'tool:result') {
|
||||
const toolResultMsg: ServerMessage = {
|
||||
type: 'tool:result',
|
||||
toolCallId: event.toolCallId,
|
||||
output: event.output,
|
||||
isError: event.isError,
|
||||
};
|
||||
ws.send(JSON.stringify(toolResultMsg));
|
||||
|
||||
// Add tool message to history
|
||||
const toolMsg: Message = {
|
||||
id: randomUUID(),
|
||||
timestamp: Date.now(),
|
||||
role: 'tool',
|
||||
toolCallId: event.toolCallId,
|
||||
toolName: event.toolCallId.split(':')[0] || 'unknown',
|
||||
output: event.output,
|
||||
isError: event.isError,
|
||||
};
|
||||
session.messages.push(toolMsg);
|
||||
session.meta.messageCount += 1;
|
||||
} else if (event.type === 'result') {
|
||||
const resultMsg: ServerMessage = {
|
||||
type: 'result',
|
||||
sessionId,
|
||||
cost: event.cost,
|
||||
};
|
||||
ws.send(JSON.stringify(resultMsg));
|
||||
session.isGenerating = false;
|
||||
|
||||
// Add assistant message with cost
|
||||
const assistantMsg: Message = {
|
||||
id: assistantMsgId,
|
||||
timestamp: Date.now(),
|
||||
role: 'assistant',
|
||||
text: assistantText || session.streamBuffer,
|
||||
model,
|
||||
cost: event.cost,
|
||||
};
|
||||
session.messages.push(assistantMsg);
|
||||
session.meta.messageCount += 1;
|
||||
session.meta.cost.inputTokens += event.cost.inputTokens;
|
||||
session.meta.cost.outputTokens += event.cost.outputTokens;
|
||||
session.meta.cost.totalUSD += event.cost.totalUSD;
|
||||
session.meta.updatedAt = Date.now();
|
||||
|
||||
// Save session to disk
|
||||
try {
|
||||
await storage.saveSession(cwd, sessionId, session.meta, session.messages);
|
||||
logger.info('Session saved to disk', { sessionId, messageCount: session.messages.length });
|
||||
} catch (err) {
|
||||
logger.error('Failed to save session', { sessionId, error: String(err) });
|
||||
}
|
||||
|
||||
// Clear streaming buffer
|
||||
session.streamBuffer = '';
|
||||
} else if (event.type === 'error') {
|
||||
const errorMsg: ServerMessage = {
|
||||
type: 'error',
|
||||
message: event.message,
|
||||
};
|
||||
ws.send(JSON.stringify(errorMsg));
|
||||
session.isGenerating = false;
|
||||
} else if (event.type === 'stopped') {
|
||||
session.isGenerating = false;
|
||||
}
|
||||
}
|
||||
} catch (err) {
|
||||
logger.error('Error streaming from Pi', { sessionId, error: String(err) });
|
||||
const errorMsg: ServerMessage = {
|
||||
type: 'error',
|
||||
message: 'Stream error',
|
||||
};
|
||||
ws.send(JSON.stringify(errorMsg));
|
||||
session.isGenerating = false;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Handle resume message
|
||||
*/
|
||||
async function handleResume(
|
||||
ws: ServerWebSocket<WSData>,
|
||||
msg: { sessionId: string }
|
||||
@@ -281,105 +259,68 @@ async function handleResume(
|
||||
const { sessionId } = msg;
|
||||
|
||||
try {
|
||||
// First, check if session exists in memory
|
||||
let session = sessionManager.getSession(sessionId);
|
||||
|
||||
if (!session) {
|
||||
// Load from disk
|
||||
// We need to try finding it - iterate through potential cwds
|
||||
const homeDir = getHomeDir(email);
|
||||
|
||||
try {
|
||||
const { meta, messages } = await storage.loadSession(homeDir, sessionId);
|
||||
|
||||
// Recreate session in memory
|
||||
session = sessionManager.getOrCreate(
|
||||
sessionId,
|
||||
email,
|
||||
meta.cwd,
|
||||
meta.model
|
||||
);
|
||||
session = sessionManager.getOrCreate(sessionId, email, meta.cwd, meta.model);
|
||||
session.messages = messages;
|
||||
session.meta = meta;
|
||||
|
||||
logger.info('Loaded session from disk', { sessionId, messageCount: messages.length });
|
||||
} catch (err) {
|
||||
logger.error('Failed to load session from disk', { sessionId, error: String(err) });
|
||||
const errorMsg: ServerMessage = {
|
||||
type: 'error',
|
||||
message: 'Session not found',
|
||||
errorCode: 'SESSION_NOT_FOUND',
|
||||
};
|
||||
ws.send(JSON.stringify(errorMsg));
|
||||
sendToClient(ws, { type: 'error', message: 'Session not found', errorCode: 'SESSION_NOT_FOUND' });
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
// Attach WebSocket to session
|
||||
sessionManager.attachWs(sessionId, ws);
|
||||
wsToSessionMap.set(ws as any, sessionId);
|
||||
|
||||
// Send session init
|
||||
const initMsg: ServerMessage = {
|
||||
type: 'session:init',
|
||||
sessionId,
|
||||
model: session.model,
|
||||
cwd: session.cwd,
|
||||
};
|
||||
ws.send(JSON.stringify(initMsg));
|
||||
sendToClient(ws, { type: 'session:init', sessionId, model: session.model, cwd: session.cwd });
|
||||
|
||||
// Spawn fresh Pi process
|
||||
// Spawn fresh Pi process if needed
|
||||
if (!session.piProcess) {
|
||||
try {
|
||||
session.piProcess = await piBridge.spawnPi(session.cwd, session.model);
|
||||
const onEvent = createEventHandler(sessionId, session.model, session.cwd);
|
||||
session.piProcess = await piBridge.spawnPi(session.cwd, session.model, onEvent);
|
||||
logger.info('Spawned fresh Pi process for resumed session', { sessionId, model: session.model });
|
||||
} catch (err) {
|
||||
logger.error('Failed to spawn Pi process for resume', { sessionId, error: String(err) });
|
||||
const errorMsg: ServerMessage = {
|
||||
type: 'error',
|
||||
message: 'Failed to start Pi process',
|
||||
};
|
||||
ws.send(JSON.stringify(errorMsg));
|
||||
sendToClient(ws, { type: 'error', message: 'Failed to start Pi process' });
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
// Send full sync with history
|
||||
const syncMsg: ServerMessage = {
|
||||
sendToClient(ws, {
|
||||
type: 'sync:messages',
|
||||
sessionId,
|
||||
messages: session.messages,
|
||||
isGenerating: session.isGenerating,
|
||||
streamingText: session.streamBuffer,
|
||||
};
|
||||
ws.send(JSON.stringify(syncMsg));
|
||||
});
|
||||
|
||||
logger.info('Session resumed successfully', { sessionId, messageCount: session.messages.length });
|
||||
} catch (err) {
|
||||
logger.error('Unexpected error in handleResume', { sessionId, error: String(err) });
|
||||
const errorMsg: ServerMessage = {
|
||||
type: 'error',
|
||||
message: 'Failed to resume session',
|
||||
};
|
||||
ws.send(JSON.stringify(errorMsg));
|
||||
sendToClient(ws, { type: 'error', message: 'Failed to resume session' });
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Handle stop message
|
||||
*/
|
||||
async function handleStop(ws: ServerWebSocket<WSData>): Promise<void> {
|
||||
const sessionId = wsToSessionMap.get(ws);
|
||||
|
||||
if (sessionId) {
|
||||
const session = sessionManager.getSession(sessionId);
|
||||
|
||||
if (session && session.piProcess) {
|
||||
if (session?.piProcess) {
|
||||
try {
|
||||
// Send abort to Pi process
|
||||
const requestId = randomUUID();
|
||||
piBridge.abort(session.piProcess, requestId);
|
||||
|
||||
piBridge.abort(session.piProcess, randomUUID());
|
||||
logger.info('Sent abort to Pi process', { sessionId });
|
||||
session.isGenerating = false;
|
||||
} catch (err) {
|
||||
@@ -388,10 +329,7 @@ async function handleStop(ws: ServerWebSocket<WSData>): Promise<void> {
|
||||
}
|
||||
}
|
||||
|
||||
const stoppedMsg: ServerMessage = {
|
||||
type: 'stopped',
|
||||
};
|
||||
ws.send(JSON.stringify(stoppedMsg));
|
||||
sendToClient(ws, { type: 'stopped' });
|
||||
}
|
||||
|
||||
export const piWebsocket = {
|
||||
|
||||
Reference in New Issue
Block a user