diff --git a/src/db.ts b/src/db.ts index bcd5dcc..18486ff 100644 --- a/src/db.ts +++ b/src/db.ts @@ -30,6 +30,39 @@ let respondToPromptStmt: ReturnType; let getPendingPromptsStmt: ReturnType; let appendOutputStmt: ReturnType; let getSessionOutputStmt: ReturnType; +let updateSessionStatsStmt: ReturnType; + +// Migration function to add new columns to existing tables +function runMigrations(): void { + // Add Phase 2.3 session state and stats columns + const migrations = [ + // Stats columns + "ALTER TABLE sessions ADD COLUMN state TEXT DEFAULT 'ready'", + "ALTER TABLE sessions ADD COLUMN prompts INTEGER DEFAULT 0", + "ALTER TABLE sessions ADD COLUMN completions INTEGER DEFAULT 0", + "ALTER TABLE sessions ADD COLUMN tools INTEGER DEFAULT 0", + "ALTER TABLE sessions ADD COLUMN compressions INTEGER DEFAULT 0", + "ALTER TABLE sessions ADD COLUMN thinking_seconds INTEGER DEFAULT 0", + "ALTER TABLE sessions ADD COLUMN work_seconds INTEGER DEFAULT 0", + "ALTER TABLE sessions ADD COLUMN mode TEXT DEFAULT 'normal'", + "ALTER TABLE sessions ADD COLUMN model TEXT", + "ALTER TABLE sessions ADD COLUMN idle_since INTEGER", + // Git columns (for later phase) + "ALTER TABLE sessions ADD COLUMN git_branch TEXT", + "ALTER TABLE sessions ADD COLUMN git_files_json TEXT", + ]; + + for (const migration of migrations) { + try { + db.exec(migration); + } catch (error: any) { + // Ignore "duplicate column" errors - column already exists + if (!error.message?.includes("duplicate column")) { + throw error; + } + } + } +} export function initDb(path = "claude-remote.db"): Database { db = new Database(path); @@ -74,6 +107,9 @@ export function initDb(path = "claude-remote.db"): Database { ); `); + // Run migrations to add new columns + runMigrations(); + // Prepare all statements once createDeviceStmt = db.prepare( "INSERT INTO devices (secret, name, created_at, last_seen) VALUES (?, ?, ?, ?) RETURNING *", @@ -106,6 +142,12 @@ export function initDb(path = "claude-remote.db"): Database { getSessionOutputStmt = db.prepare( "SELECT * FROM output_log WHERE session_id = ? ORDER BY timestamp ASC", ); + updateSessionStatsStmt = db.prepare(` + UPDATE sessions + SET state = ?, prompts = ?, completions = ?, tools = ?, compressions = ?, + thinking_seconds = ?, work_seconds = ?, mode = ?, model = ?, idle_since = ? + WHERE id = ? + `); return db; } @@ -226,3 +268,35 @@ export function appendOutput(sessionId: number, line: string): void { export function getSessionOutput(sessionId: number): OutputLog[] { return getSessionOutputStmt.all(sessionId) as OutputLog[]; } + +// Session state functions + +export function updateSessionStats( + sessionId: number, + state: { + state: string; + prompts: number; + completions: number; + tools: number; + compressions: number; + thinking_seconds: number; + work_seconds: number; + mode: string; + model: string | null; + idle_since: number | null; + }, +): void { + updateSessionStatsStmt.run( + state.state, + state.prompts, + state.completions, + state.tools, + state.compressions, + state.thinking_seconds, + state.work_seconds, + state.mode, + state.model, + state.idle_since, + sessionId, + ); +} diff --git a/src/server.ts b/src/server.ts index 56845cf..1337d51 100644 --- a/src/server.ts +++ b/src/server.ts @@ -13,23 +13,53 @@ import { initDb, respondToPrompt, updateLastSeen, + updateSessionStats, } from "./db"; import type { AnswerResponse, ClientMessage, ServerMessage, + SessionState, SSEEvent, } from "./types"; // Server state const sseClients = new Set>(); const sessionWebSockets = new Map>(); +const sessionStates = new Map(); interface SessionData { deviceId: number; sessionId: number | null; } +// Helper function to create default SessionState +function createDefaultSessionState(): SessionState { + return { + state: "ready", + prompts: 0, + completions: 0, + tools: 0, + compressions: 0, + thinking_seconds: 0, + work_seconds: 0, + mode: "normal", + model: null, + idle_since: null, + dirty: false, + }; +} + +// Persist dirty sessions to database (throttled, called every 5s) +function persistDirtySessions(): void { + for (const [sessionId, state] of sessionStates.entries()) { + if (state.dirty) { + updateSessionStats(sessionId, state); + state.dirty = false; + } + } +} + // Broadcast SSE event to all connected dashboards function broadcastSSE(event: SSEEvent): void { const eventStr = `event: ${event.type}\ndata: ${JSON.stringify(event)}\n\n`; @@ -337,6 +367,9 @@ const server = Bun.serve({ ws.data.sessionId = session.id; sessionWebSockets.set(session.id, ws); + // Initialize in-memory session state + sessionStates.set(session.id, createDefaultSessionState()); + console.debug( `Session ${session.id} started for device ${device.id}`, ); @@ -349,6 +382,30 @@ const server = Bun.serve({ command: session.command, }); + // Broadcast initial state and stats + const initialState = sessionStates.get(session.id); + if (initialState) { + broadcastSSE({ + type: "state", + session_id: session.id, + state: initialState.state, + timestamp: Date.now(), + }); + broadcastSSE({ + type: "stats", + session_id: session.id, + prompts: initialState.prompts, + completions: initialState.completions, + tools: initialState.tools, + compressions: initialState.compressions, + thinking_seconds: initialState.thinking_seconds, + work_seconds: initialState.work_seconds, + mode: initialState.mode, + model: initialState.model, + idle_since: initialState.idle_since, + }); + } + ws.send( JSON.stringify({ type: "authenticated", session_id: session.id }), ); @@ -387,6 +444,60 @@ const server = Bun.serve({ return; } + // Handle state message + if (msg.type === "state") { + const sessionState = sessionStates.get(ws.data.sessionId); + if (sessionState && msg.state) { + sessionState.state = msg.state; + sessionState.dirty = true; + // Broadcast SSE state event + broadcastSSE({ + type: "state", + session_id: ws.data.sessionId, + state: msg.state, + timestamp: msg.timestamp || Date.now(), + }); + } + return; + } + + // Handle stats message + if (msg.type === "stats") { + const sessionState = sessionStates.get(ws.data.sessionId); + if (sessionState) { + // Update all stats fields from the message + sessionState.prompts = msg.prompts ?? sessionState.prompts; + sessionState.completions = + msg.completions ?? sessionState.completions; + sessionState.tools = msg.tools ?? sessionState.tools; + sessionState.compressions = + msg.compressions ?? sessionState.compressions; + sessionState.thinking_seconds = + msg.thinking_seconds ?? sessionState.thinking_seconds; + sessionState.work_seconds = + msg.work_seconds ?? sessionState.work_seconds; + sessionState.mode = msg.mode ?? sessionState.mode; + sessionState.model = msg.model ?? sessionState.model; + sessionState.idle_since = msg.idle_since ?? sessionState.idle_since; + sessionState.dirty = true; + // Broadcast SSE stats event + broadcastSSE({ + type: "stats", + session_id: ws.data.sessionId, + prompts: sessionState.prompts, + completions: sessionState.completions, + tools: sessionState.tools, + compressions: sessionState.compressions, + thinking_seconds: sessionState.thinking_seconds, + work_seconds: sessionState.work_seconds, + mode: sessionState.mode, + model: sessionState.model, + idle_since: sessionState.idle_since, + }); + } + return; + } + // Handle resize message if (msg.type === "resize") { // Store resize info if needed (not yet implemented) @@ -401,6 +512,13 @@ const server = Bun.serve({ endSession(ws.data.sessionId); sessionWebSockets.delete(ws.data.sessionId); + // Persist final state before cleanup + const state = sessionStates.get(ws.data.sessionId); + if (state?.dirty) { + updateSessionStats(ws.data.sessionId, state); + } + sessionStates.delete(ws.data.sessionId); + console.debug( `Session ${ws.data.sessionId} ended with code ${msg.code}`, ); @@ -423,6 +541,14 @@ const server = Bun.serve({ close(ws) { if (ws.data.sessionId) { sessionWebSockets.delete(ws.data.sessionId); + + // Persist final state before cleanup + const state = sessionStates.get(ws.data.sessionId); + if (state?.dirty) { + updateSessionStats(ws.data.sessionId, state); + } + sessionStates.delete(ws.data.sessionId); + // Mark session as ended if not already const sessionId = ws.data.sessionId; endSession(sessionId); @@ -446,4 +572,7 @@ setInterval(() => { } }, 30_000); +// Periodic persistence of dirty session states (every 5s) +setInterval(persistDirtySessions, 5000); + export { server };