feat: Complete Pi Harness Rebuild - Phases 1 till 7.2
This commit is contained in:
@@ -0,0 +1,402 @@
|
||||
import type { ServerWebSocket } from 'bun';
|
||||
import { randomUUID } from 'crypto';
|
||||
import type { ClientMessage, ServerMessage, Message } 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;
|
||||
role: string;
|
||||
provider: string;
|
||||
};
|
||||
|
||||
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>
|
||||
`;
|
||||
}
|
||||
|
||||
/**
|
||||
* 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();
|
||||
|
||||
(async () => {
|
||||
try {
|
||||
const clientMsg = JSON.parse(data) as ClientMessage;
|
||||
|
||||
if (clientMsg.type === 'chat') {
|
||||
await handleChat(ws, clientMsg);
|
||||
} else if (clientMsg.type === 'resume') {
|
||||
await handleResume(ws, clientMsg);
|
||||
} else if (clientMsg.type === 'stop') {
|
||||
await handleStop(ws);
|
||||
}
|
||||
} 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));
|
||||
}
|
||||
})();
|
||||
}
|
||||
|
||||
/**
|
||||
* WebSocket close handler
|
||||
*/
|
||||
export function close(ws: ServerWebSocket<WSData>): void {
|
||||
logger.info('WebSocket connection closed', { email: ws.data.email });
|
||||
|
||||
const sessionId = wsToSessionMap.get(ws);
|
||||
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
|
||||
*/
|
||||
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 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));
|
||||
|
||||
// Spawn Pi process if not already running
|
||||
if (!session.piProcess) {
|
||||
try {
|
||||
session.piProcess = await piBridge.spawnPi(cwd, model);
|
||||
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));
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
// Add user message to session
|
||||
const userMsg: Message = {
|
||||
id: randomUUID(),
|
||||
timestamp: Date.now(),
|
||||
role: 'user',
|
||||
text: msg.prompt,
|
||||
};
|
||||
session.messages.push(userMsg);
|
||||
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);
|
||||
}
|
||||
|
||||
// Send prompt to Pi
|
||||
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 }
|
||||
): Promise<void> {
|
||||
const { email } = ws.data;
|
||||
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.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));
|
||||
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));
|
||||
|
||||
// Spawn fresh Pi process
|
||||
if (!session.piProcess) {
|
||||
try {
|
||||
session.piProcess = await piBridge.spawnPi(session.cwd, session.model);
|
||||
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));
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
// Send full sync with history
|
||||
const syncMsg: ServerMessage = {
|
||||
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));
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* 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) {
|
||||
try {
|
||||
// Send abort to Pi process
|
||||
const requestId = randomUUID();
|
||||
piBridge.abort(session.piProcess, requestId);
|
||||
|
||||
logger.info('Sent abort to Pi process', { sessionId });
|
||||
session.isGenerating = false;
|
||||
} catch (err) {
|
||||
logger.error('Failed to abort Pi process', { sessionId, error: String(err) });
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
const stoppedMsg: ServerMessage = {
|
||||
type: 'stopped',
|
||||
};
|
||||
ws.send(JSON.stringify(stoppedMsg));
|
||||
}
|
||||
|
||||
export const piWebsocket = {
|
||||
open,
|
||||
message,
|
||||
close,
|
||||
drain() {},
|
||||
};
|
||||
Reference in New Issue
Block a user