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:
+189
-119
@@ -1,35 +1,180 @@
|
||||
import type { Subprocess } from "bun";
|
||||
import type { PiEvent, MessageCost } from "./types";
|
||||
import { readApiKeys } from "../server-settings/pi-mono";
|
||||
import { logger } from "./logger";
|
||||
|
||||
export type PiEventHandler = (event: PiEvent) => void;
|
||||
|
||||
export async function spawnPi(
|
||||
cwd: string,
|
||||
model: string,
|
||||
env?: Record<string, string>
|
||||
onEvent: PiEventHandler
|
||||
): Promise<Subprocess> {
|
||||
const piProcess = Bun.spawn(
|
||||
[
|
||||
"pi",
|
||||
"--mode",
|
||||
"rpc",
|
||||
"--no-extensions",
|
||||
"--no-skills",
|
||||
"--model",
|
||||
model,
|
||||
],
|
||||
{
|
||||
cwd,
|
||||
env: {
|
||||
...process.env,
|
||||
...env,
|
||||
},
|
||||
stdin: "pipe",
|
||||
stdout: "pipe",
|
||||
stderr: "inherit",
|
||||
}
|
||||
);
|
||||
const storedKeys = await readApiKeys();
|
||||
|
||||
const args = ['pi', '--mode', 'rpc', '--no-extensions', '--no-skills', '--no-prompt-templates', '--no-themes'];
|
||||
if (model) args.push('--model', model);
|
||||
|
||||
return piProcess;
|
||||
const proc = Bun.spawn(args, {
|
||||
cwd,
|
||||
stdin: 'pipe',
|
||||
stdout: 'pipe',
|
||||
stderr: 'pipe',
|
||||
env: { ...process.env, ...storedKeys },
|
||||
});
|
||||
|
||||
// Read stdout JSON event stream (runs in background)
|
||||
const reader = proc.stdout.getReader();
|
||||
const decoder = new TextDecoder();
|
||||
let buffer = '';
|
||||
let streamBuffer = '';
|
||||
|
||||
(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) as Record<string, unknown>;
|
||||
const piEvent = parsePiEvent(event, streamBuffer);
|
||||
if (piEvent) {
|
||||
if (piEvent.type === 'delta') {
|
||||
streamBuffer += piEvent.text;
|
||||
} else if (piEvent.type === 'text' || piEvent.type === 'tool:start') {
|
||||
streamBuffer = '';
|
||||
}
|
||||
onEvent(piEvent);
|
||||
}
|
||||
} catch {
|
||||
// Skip unparseable lines
|
||||
}
|
||||
}
|
||||
}
|
||||
} catch {
|
||||
// Process ended
|
||||
}
|
||||
})();
|
||||
|
||||
// Stderr → debug log
|
||||
const stderrReader = proc.stderr.getReader();
|
||||
const stderrDecoder = new TextDecoder();
|
||||
(async () => {
|
||||
try {
|
||||
while (true) {
|
||||
const { done, value } = await stderrReader.read();
|
||||
if (done) break;
|
||||
const text = stderrDecoder.decode(value, { stream: true });
|
||||
if (text.trim()) logger.info('Pi stderr', { text: text.trim() });
|
||||
}
|
||||
} catch {
|
||||
// Process ended
|
||||
}
|
||||
})();
|
||||
|
||||
// Handle process exit
|
||||
proc.exited.then((code) => {
|
||||
logger.info('Pi process exited', { code });
|
||||
});
|
||||
|
||||
return proc;
|
||||
}
|
||||
|
||||
function parsePiEvent(event: Record<string, unknown>, currentStreamBuffer: string): PiEvent | null {
|
||||
const type = event.type as string;
|
||||
|
||||
// Handle response (success/failure for commands)
|
||||
if (type === 'response') {
|
||||
if (event.command === 'prompt' && !event.success) {
|
||||
const errorMsg = (event.error as string) ?? 'Prompt failed';
|
||||
return { type: 'error', message: errorMsg };
|
||||
}
|
||||
return null;
|
||||
}
|
||||
|
||||
switch (type) {
|
||||
case 'agent_start':
|
||||
// No event to emit, just resets state
|
||||
return null;
|
||||
|
||||
case 'message_update': {
|
||||
const ame = event.assistantMessageEvent as Record<string, unknown> | undefined;
|
||||
if (ame?.type === 'text_delta') {
|
||||
const delta = ame.delta as string;
|
||||
return { type: 'delta', text: delta };
|
||||
}
|
||||
return null;
|
||||
}
|
||||
|
||||
case 'message_end': {
|
||||
if (currentStreamBuffer) {
|
||||
return { type: 'text', text: currentStreamBuffer };
|
||||
}
|
||||
return null;
|
||||
}
|
||||
|
||||
case 'tool_execution_start': {
|
||||
const toolCallId = (event.toolCallId as string) ?? '';
|
||||
const toolName = (event.toolName as string) ?? 'unknown';
|
||||
const args = (event.args as Record<string, unknown>) ?? {};
|
||||
|
||||
return {
|
||||
type: 'tool:start',
|
||||
toolCallId,
|
||||
toolName,
|
||||
toolInput: args,
|
||||
};
|
||||
}
|
||||
|
||||
case 'tool_execution_end': {
|
||||
const toolCallId = (event.toolCallId as string) ?? '';
|
||||
const result = event.result;
|
||||
const isError = (event.isError as boolean) ?? false;
|
||||
const output = result != null ? (typeof result === 'string' ? result : JSON.stringify(result)) : '';
|
||||
|
||||
return {
|
||||
type: 'tool:result',
|
||||
toolCallId,
|
||||
output,
|
||||
isError,
|
||||
};
|
||||
}
|
||||
|
||||
case 'agent_end': {
|
||||
// Pi doesn't provide cost info in agent_end, use zeros
|
||||
const cost: MessageCost = {
|
||||
inputTokens: 0,
|
||||
outputTokens: 0,
|
||||
totalUSD: 0,
|
||||
};
|
||||
return { type: 'result', cost };
|
||||
}
|
||||
|
||||
case 'extension_ui_request': {
|
||||
// Will be handled separately
|
||||
return null;
|
||||
}
|
||||
|
||||
default:
|
||||
return null;
|
||||
}
|
||||
}
|
||||
|
||||
function writeRpcCommand(proc: Subprocess, command: Record<string, unknown>): void {
|
||||
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) {
|
||||
logger.error('writeRpcCommand error', { error: String(err) });
|
||||
}
|
||||
}
|
||||
|
||||
export function sendPrompt(
|
||||
@@ -37,113 +182,38 @@ export function sendPrompt(
|
||||
prompt: string,
|
||||
requestId: string
|
||||
): void {
|
||||
const request = {
|
||||
jsonrpc: "2.0",
|
||||
method: "chat",
|
||||
params: {
|
||||
prompt,
|
||||
},
|
||||
writeRpcCommand(process, {
|
||||
type: 'prompt',
|
||||
id: requestId,
|
||||
};
|
||||
|
||||
const writer = (process.stdin as any).getWriter();
|
||||
writer.write(
|
||||
new TextEncoder().encode(JSON.stringify(request) + "\n")
|
||||
);
|
||||
writer.releaseLock();
|
||||
message: prompt,
|
||||
});
|
||||
}
|
||||
|
||||
export function abort(
|
||||
process: Subprocess,
|
||||
requestId: string
|
||||
): void {
|
||||
const request = {
|
||||
jsonrpc: "2.0",
|
||||
method: "abort",
|
||||
params: {},
|
||||
writeRpcCommand(process, {
|
||||
type: 'abort',
|
||||
id: requestId,
|
||||
};
|
||||
|
||||
const writer = (process.stdin as any).getWriter();
|
||||
writer.write(
|
||||
new TextEncoder().encode(JSON.stringify(request) + "\n")
|
||||
);
|
||||
writer.releaseLock();
|
||||
});
|
||||
}
|
||||
|
||||
export async function* readEvents(
|
||||
process: Subprocess
|
||||
): AsyncGenerator<PiEvent> {
|
||||
if (!process.stdout) {
|
||||
throw new Error("Pi process stdout not available");
|
||||
}
|
||||
|
||||
const reader = (process.stdout as any).getReader();
|
||||
const decoder = new TextDecoder();
|
||||
let buffer = "";
|
||||
|
||||
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);
|
||||
|
||||
if (event.method === "text") {
|
||||
yield { type: "text", text: event.params.text };
|
||||
} else if (event.method === "delta") {
|
||||
yield { type: "delta", text: event.params.text };
|
||||
} else if (event.method === "tool:start") {
|
||||
yield {
|
||||
type: "tool:start",
|
||||
toolCallId: event.params.toolCallId,
|
||||
toolName: event.params.toolName,
|
||||
toolInput: event.params.toolInput,
|
||||
};
|
||||
} else if (event.method === "tool:result") {
|
||||
yield {
|
||||
type: "tool:result",
|
||||
toolCallId: event.params.toolCallId,
|
||||
output: event.params.output,
|
||||
isError: event.params.isError || false,
|
||||
};
|
||||
} else if (event.method === "result") {
|
||||
const cost: MessageCost = {
|
||||
inputTokens: event.params.cost?.inputTokens || 0,
|
||||
outputTokens: event.params.cost?.outputTokens || 0,
|
||||
totalUSD: event.params.cost?.totalUSD || 0,
|
||||
};
|
||||
yield { type: "result", cost };
|
||||
} else if (event.method === "error") {
|
||||
yield {
|
||||
type: "error",
|
||||
message: event.params.message || "Unknown error",
|
||||
};
|
||||
} else if (event.method === "stopped") {
|
||||
yield { type: "stopped" };
|
||||
}
|
||||
} catch (err) {
|
||||
logger.error("Failed to parse Pi event", { line, error: String(err) });
|
||||
}
|
||||
}
|
||||
}
|
||||
} finally {
|
||||
reader.releaseLock();
|
||||
}
|
||||
export function cancelExtensionUi(
|
||||
process: Subprocess,
|
||||
id: unknown
|
||||
): void {
|
||||
writeRpcCommand(process, {
|
||||
type: 'extension_ui_response',
|
||||
id,
|
||||
cancelled: true,
|
||||
});
|
||||
}
|
||||
|
||||
export function killPi(process: Subprocess): void {
|
||||
process.kill();
|
||||
try {
|
||||
process.kill();
|
||||
} catch {
|
||||
// Already dead
|
||||
}
|
||||
}
|
||||
|
||||
+69
-37
@@ -1,6 +1,8 @@
|
||||
import type { Context } from 'hono';
|
||||
import { createRouter } from '../../create-router';
|
||||
import * as storage from './storage';
|
||||
import { readApiKeys } from '../server-settings/pi-mono';
|
||||
import { getHomeDir } from '../../data-path';
|
||||
import type { ModelInfo } from './types';
|
||||
import { logger } from './logger';
|
||||
|
||||
@@ -12,36 +14,66 @@ export const piRestRouter = createRouter();
|
||||
|
||||
/**
|
||||
* GET /api/pi/models
|
||||
* List available models
|
||||
* List available models by running `pi --list-models` with stored API keys
|
||||
*/
|
||||
piRestRouter.get('/pi/models', async (ctx: Context) => {
|
||||
// TODO: Implement dynamic model discovery
|
||||
// For now, return hardcoded models
|
||||
const models: ModelInfo[] = [
|
||||
{
|
||||
id: 'gpt-4o',
|
||||
name: 'GPT-4o',
|
||||
provider: 'openai',
|
||||
contextWindow: 128000,
|
||||
maxTokens: 4096,
|
||||
},
|
||||
{
|
||||
id: 'claude-opus-4-5',
|
||||
name: 'Claude Opus 4.5',
|
||||
provider: 'anthropic',
|
||||
contextWindow: 200000,
|
||||
maxTokens: 4096,
|
||||
},
|
||||
{
|
||||
id: 'big-pickle',
|
||||
name: 'Big Pickle',
|
||||
provider: 'opencode-zen',
|
||||
contextWindow: 128000,
|
||||
maxTokens: 4096,
|
||||
},
|
||||
];
|
||||
try {
|
||||
const storedKeys = await readApiKeys();
|
||||
|
||||
const proc = Bun.spawn(['pi', '--list-models'], {
|
||||
stdout: 'pipe',
|
||||
stderr: 'pipe',
|
||||
env: { ...process.env, ...storedKeys },
|
||||
});
|
||||
|
||||
return ctx.json({ models });
|
||||
const output = await new Response(proc.stdout).text();
|
||||
await proc.exited;
|
||||
|
||||
if (proc.exitCode !== 0) {
|
||||
logger.error('pi --list-models failed', { exitCode: proc.exitCode });
|
||||
return ctx.json({ models: [] });
|
||||
}
|
||||
|
||||
// Parse the whitespace-separated table output:
|
||||
// provider model context max-out thinking images
|
||||
// anthropic claude-sonnet-4-6 200K 128K yes yes
|
||||
const lines = output.trim().split('\n').filter(Boolean);
|
||||
const models: ModelInfo[] = [];
|
||||
|
||||
// Skip header line (first line)
|
||||
for (let i = 1; i < lines.length; i++) {
|
||||
const cols = lines[i]!.trim().split(/\s+/);
|
||||
if (cols.length < 2) continue;
|
||||
|
||||
const [provider, model, context, maxOut] = cols;
|
||||
|
||||
// Parse context window (e.g., "200K" -> 200000)
|
||||
const parseSize = (s?: string): number => {
|
||||
if (!s) return 128000;
|
||||
const match = s.match(/^(\d+)([KMG])?$/i);
|
||||
if (!match) return 128000;
|
||||
const num = parseInt(match[1]!, 10);
|
||||
const unit = (match[2] ?? '').toUpperCase();
|
||||
if (unit === 'K') return num * 1000;
|
||||
if (unit === 'M') return num * 1000000;
|
||||
if (unit === 'G') return num * 1000000000;
|
||||
return num;
|
||||
};
|
||||
|
||||
models.push({
|
||||
id: `${provider}/${model}`,
|
||||
name: model!,
|
||||
provider: provider!,
|
||||
contextWindow: parseSize(context),
|
||||
maxTokens: parseSize(maxOut),
|
||||
});
|
||||
}
|
||||
|
||||
return ctx.json({ models });
|
||||
} catch (err) {
|
||||
logger.error('Failed to list models', { error: String(err) });
|
||||
return ctx.json({ models: [] });
|
||||
}
|
||||
});
|
||||
|
||||
/**
|
||||
@@ -55,7 +87,7 @@ piRestRouter.post('/pi/sessions', async (ctx: Context) => {
|
||||
}
|
||||
|
||||
// TODO: Implement proper user home directory resolution
|
||||
const userHome = `/home/${user.email.split('@')[0]}`;
|
||||
const userHome = getHomeDir(user.email);
|
||||
|
||||
try {
|
||||
const sessions = await storage.listUserSessions(userHome);
|
||||
@@ -81,7 +113,7 @@ piRestRouter.get('/pi/sessions/:sessionId', async (ctx: Context) => {
|
||||
return ctx.json({ error: 'Session ID required' }, 400);
|
||||
}
|
||||
|
||||
const userHome = `/home/${user.email.split('@')[0]}`;
|
||||
const userHome = getHomeDir(user.email);
|
||||
|
||||
try {
|
||||
// Try loading from root first
|
||||
@@ -140,7 +172,7 @@ piRestRouter.patch('/pi/sessions/:sessionId', async (ctx: Context) => {
|
||||
return ctx.json({ error: 'Title is required and must be a string' }, 400);
|
||||
}
|
||||
|
||||
const userHome = `/home/${user.email.split('@')[0]}`;
|
||||
const userHome = getHomeDir(user.email);
|
||||
|
||||
try {
|
||||
// Find the session (root or in group)
|
||||
@@ -192,7 +224,7 @@ piRestRouter.delete('/pi/sessions/:sessionId', async (ctx: Context) => {
|
||||
return ctx.json({ error: 'Session ID required' }, 400);
|
||||
}
|
||||
|
||||
const userHome = `/home/${user.email.split('@')[0]}`;
|
||||
const userHome = getHomeDir(user.email);
|
||||
|
||||
try {
|
||||
// Find the session (root or in group)
|
||||
@@ -251,7 +283,7 @@ piRestRouter.get('/pi/sessions/search', async (ctx: Context) => {
|
||||
return ctx.json({ error: 'Query parameter required' }, 400);
|
||||
}
|
||||
|
||||
const userHome = `/home/${user.email.split('@')[0]}`;
|
||||
const userHome = getHomeDir(user.email);
|
||||
|
||||
try {
|
||||
const results = await storage.searchSessions(userHome, query);
|
||||
@@ -282,7 +314,7 @@ piRestRouter.post('/pi/groups', async (ctx: Context) => {
|
||||
return ctx.json({ error: 'Slug is required and must be a string' }, 400);
|
||||
}
|
||||
|
||||
const userHome = `/home/${user.email.split('@')[0]}`;
|
||||
const userHome = getHomeDir(user.email);
|
||||
|
||||
try {
|
||||
// Check if group already exists
|
||||
@@ -347,7 +379,7 @@ piRestRouter.get('/pi/groups', async (ctx: Context) => {
|
||||
return ctx.json({ error: 'Unauthorized' }, 401);
|
||||
}
|
||||
|
||||
const userHome = `/home/${user.email.split('@')[0]}`;
|
||||
const userHome = getHomeDir(user.email);
|
||||
|
||||
try {
|
||||
const groups = await storage.listGroups(userHome);
|
||||
@@ -387,7 +419,7 @@ piRestRouter.patch('/pi/groups/:groupSlug', async (ctx: Context) => {
|
||||
return ctx.json({ error: 'No valid updates provided' }, 400);
|
||||
}
|
||||
|
||||
const userHome = `/home/${user.email.split('@')[0]}`;
|
||||
const userHome = getHomeDir(user.email);
|
||||
|
||||
try {
|
||||
const updatedGroup = await storage.updateGroupMeta(userHome, groupSlug, updates);
|
||||
@@ -416,7 +448,7 @@ piRestRouter.delete('/pi/groups/:groupSlug', async (ctx: Context) => {
|
||||
return ctx.json({ error: 'Group slug required' }, 400);
|
||||
}
|
||||
|
||||
const userHome = `/home/${user.email.split('@')[0]}`;
|
||||
const userHome = getHomeDir(user.email);
|
||||
|
||||
try {
|
||||
await storage.deleteGroup(userHome, groupSlug);
|
||||
@@ -449,7 +481,7 @@ piRestRouter.post('/pi/sessions/:sessionId/move', async (ctx: Context) => {
|
||||
return ctx.json({ error: 'groupSlug must be a string or null' }, 400);
|
||||
}
|
||||
|
||||
const userHome = `/home/${user.email.split('@')[0]}`;
|
||||
const userHome = getHomeDir(user.email);
|
||||
|
||||
try {
|
||||
// Find the session in root or any group
|
||||
|
||||
+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