From db6a24444e2defe53afedd85f950fe10d9312214 Mon Sep 17 00:00:00 2001 From: dimakis Date: Fri, 29 May 2026 18:48:45 +0100 Subject: [PATCH 1/2] feat(session): persist boot context across reconnect and session switch MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Boot context pills disappeared on iOS reconnect and session switch because the boot_context WS message was only sent once (fire-and-forget) and never persisted. This adds a three-layer persistence strategy: - ManagedSession in-memory cache (hot path — instant replay on reconnect) - SQLite sessions table agent_name + boot_context columns (cold path) - Re-send on handleReconnect and handleSwitchSession Co-Authored-By: Claude Opus 4.6 --- packages/harness/src/session-registry.ts | 4 ++ packages/protocol/src/event-store.ts | 30 +++++++++++ packages/protocol/src/types.ts | 3 ++ server/chat.ts | 66 ++++++++++++++---------- server/query-loop.ts | 4 ++ server/ws-handler-v2.ts | 29 +++++++++++ 6 files changed, 110 insertions(+), 26 deletions(-) diff --git a/packages/harness/src/session-registry.ts b/packages/harness/src/session-registry.ts index f30f960a..a9bf83f4 100644 --- a/packages/harness/src/session-registry.ts +++ b/packages/harness/src/session-registry.ts @@ -46,6 +46,10 @@ export interface ManagedSession { telosTaskId?: string; /** Active subagent task IDs — task_id → tool_use_id (parent_tool_use_id). */ activeTaskIds: Map; + /** Agent name used for this session (e.g. 'mitzo-conversational'). */ + agentName?: string; + /** Cached boot_context payload for replay on reconnect/switch. */ + bootContext?: Record; } export interface ActiveSessionInfo { diff --git a/packages/protocol/src/event-store.ts b/packages/protocol/src/event-store.ts index 3735f1e2..19cac827 100644 --- a/packages/protocol/src/event-store.ts +++ b/packages/protocol/src/event-store.ts @@ -48,6 +48,8 @@ interface SessionRow { last_speaker_at: number | null; state: string | null; last_state_change: number | null; + agent_name: string | null; + boot_context: string | null; created_at: number; updated_at: number; } @@ -122,6 +124,7 @@ export class EventStore { this.migrateCloseTracking(db); this.migrateAttentionTracking(db); this.migrateSessionState(db); + this.migrateBootContext(db); this.log.info('EventStore initialized', { dbPath }); @@ -281,6 +284,19 @@ export class EventStore { } } + private migrateBootContext(db: Database.Database): void { + const columns = db.prepare("PRAGMA table_info('sessions')").all() as Array<{ name: string }>; + const columnNames = new Set(columns.map((c) => c.name)); + if (!columnNames.has('agent_name')) { + db.exec('ALTER TABLE sessions ADD COLUMN agent_name TEXT'); + this.log.info('migrated sessions table: added agent_name'); + } + if (!columnNames.has('boot_context')) { + db.exec('ALTER TABLE sessions ADD COLUMN boot_context TEXT'); + this.log.info('migrated sessions table: added boot_context'); + } + } + close(): void { if (this.db) { this.db.close(); @@ -351,6 +367,14 @@ export class EventStore { fields.push('closed_by = ?'); values.push(meta.closedBy); } + if (meta.agentName !== undefined) { + fields.push('agent_name = ?'); + values.push(meta.agentName); + } + if (meta.bootContext !== undefined) { + fields.push('boot_context = ?'); + values.push(meta.bootContext); + } if (meta.updatedAt !== undefined) { fields.push('updated_at = ?'); values.push(meta.updatedAt); @@ -374,6 +398,8 @@ export class EventStore { 'goal_id', 'telos_task_id', 'closed_by', + 'agent_name', + 'boot_context', ]; const vals: unknown[] = [ meta.sessionId, @@ -387,6 +413,8 @@ export class EventStore { meta.goalId ?? null, meta.telosTaskId ?? null, meta.closedBy ?? null, + meta.agentName ?? null, + meta.bootContext ?? null, ]; if (meta.updatedAt !== undefined) { cols.push('updated_at'); @@ -653,6 +681,8 @@ function rowToSession(row: SessionRow): SessionMeta { lastSpeakerAt: row.last_speaker_at ?? null, state: (row.state as SessionMeta['state']) ?? null, lastStateChange: row.last_state_change ?? null, + agentName: row.agent_name ?? null, + bootContext: row.boot_context ?? null, createdAt: row.created_at, updatedAt: row.updated_at, }; diff --git a/packages/protocol/src/types.ts b/packages/protocol/src/types.ts index be8a157d..e4e758c5 100644 --- a/packages/protocol/src/types.ts +++ b/packages/protocol/src/types.ts @@ -219,6 +219,9 @@ export interface SessionMeta { lastSpeakerAt: number | null; state: SessionState | null; lastStateChange: number | null; + agentName: string | null; + /** Serialized JSON of the boot_context payload (sources, tokens, sections). */ + bootContext: string | null; createdAt: number; updatedAt: number; } diff --git a/server/chat.ts b/server/chat.ts index 00a6e66c..c4891233 100644 --- a/server/chat.ts +++ b/server/chat.ts @@ -675,6 +675,9 @@ async function _startChatInner( const mcpAllowed = buildMcpAllowedTools(clientId); const extraTools = options.extraTools ? options.extraTools.split(',').map((t) => t.trim()) : []; + // Resolve agent name early — needed for registration, resume upsert, and boot context. + const agentName = options.agentName ?? DEFAULT_AGENT_NAME; + // Streaming-input queue — kept open for the session lifetime. const inputQueue = new AsyncQueue(); inputQueue.push(makeUserMessage(fullPrompt, 'now')); @@ -687,6 +690,7 @@ async function _startChatInner( wtId, sessionAllowList: new Set(), worktreePath, + agentName, // Set sessionId early so pre-assistant events are persisted (iOS reconnect). ...(options.resume ? { sessionId: options.resume } : {}), ...(options.telosTaskId ? { telosTaskId: options.telosTaskId } : {}), @@ -726,6 +730,7 @@ async function _startChatInner( ...(worktreePath ? { wtId } : {}), ...(options.telosTaskId ? { telosTaskId: options.telosTaskId } : {}), ...(existingMeta ? { updatedAt: existingMeta.updatedAt } : {}), + agentName, }); } @@ -741,7 +746,6 @@ async function _startChatInner( // Build session env with worktree paths for the agent (all repos including primary) const sessionEnv = sdkEnv(); sessionEnv.MITZO_SESSION_ID = wtId; - const agentName = options.agentName ?? DEFAULT_AGENT_NAME; sessionEnv.MITZO_AGENT_NAME = agentName; for (const [name, { path }] of repoWorktrees) { sessionEnv[`MITZO_REPO_${name.toUpperCase()}`] = path; @@ -777,16 +781,18 @@ async function _startChatInner( } catch (importErr: unknown) { const msg = importErr instanceof Error ? importErr.message : String(importErr); log.info('contexgin not available, using fallback', { error: msg }); - send(transport, { - type: 'boot_context', - source: 'local-fallback', + const fallback = { + source: 'local-fallback' as const, sourceCount: 0, tokenCount: 0, tokenBudget: DEFAULT_TOKEN_BUDGET, - sources: [], - included: [], - trimmed: [], - }); + sources: [] as Array<{ path: string; kind: string }>, + included: [] as Array<{ source: string; heading: string; tokens: number; content: string }>, + trimmed: [] as Array<{ source: string; heading: string; tokens: number; content: string }>, + }; + send(transport, { type: 'boot_context', ...fallback }); + const s = registry.get(clientId); + if (s) s.bootContext = fallback; return; } @@ -819,16 +825,18 @@ async function _startChatInner( // Validate the compiled object shape if (!compiled || typeof compiled !== 'object') { log.warn('contexgin compile() returned unexpected shape', { compiled }); - send(transport, { - type: 'boot_context', - source: 'local-fallback', + const fallback = { + source: 'local-fallback' as const, sourceCount: 0, tokenCount: 0, tokenBudget: tokenBudget, - sources: [], - included: [], - trimmed: [], - }); + sources: [] as Array<{ path: string; kind: string }>, + included: [] as Array<{ source: string; heading: string; tokens: number; content: string }>, + trimmed: [] as Array<{ source: string; heading: string; tokens: number; content: string }>, + }; + send(transport, { type: 'boot_context', ...fallback }); + const s = registry.get(clientId); + if (s) s.bootContext = fallback; return; } @@ -879,9 +887,8 @@ async function _startChatInner( const trimmed = extractSections(rawTrimmed); const fullMarkdown = typeof obj.bootPayload === 'string' ? obj.bootPayload : undefined; - send(transport, { - type: 'boot_context', - source: 'contexgin', + const bootPayload = { + source: 'contexgin' as const, sourceCount: sources.length, tokenCount: bootTokens, tokenBudget: tokenBudget, @@ -889,20 +896,27 @@ async function _startChatInner( included, trimmed, fullMarkdown, - }); + }; + send(transport, { type: 'boot_context', ...bootPayload }); + + // Cache in ManagedSession for replay on reconnect/switch + const s = registry.get(clientId); + if (s) s.bootContext = bootPayload; } catch (err: unknown) { const msg = err instanceof Error ? err.message : String(err); log.warn('boot context compilation failed', { error: msg }); - send(transport, { - type: 'boot_context', - source: 'local-fallback', + const fallbackPayload = { + source: 'local-fallback' as const, sourceCount: 0, tokenCount: 0, tokenBudget: tokenBudget, - sources: [], - included: [], - trimmed: [], - }); + sources: [] as Array<{ path: string; kind: string }>, + included: [] as Array<{ source: string; heading: string; tokens: number; content: string }>, + trimmed: [] as Array<{ source: string; heading: string; tokens: number; content: string }>, + }; + send(transport, { type: 'boot_context', ...fallbackPayload }); + const s = registry.get(clientId); + if (s) s.bootContext = fallbackPayload; } })(); capturePromptComparison(wtId, cwd, systemPromptAppend, repoWorktrees).catch(() => {}); diff --git a/server/query-loop.ts b/server/query-loop.ts index f416306d..d18e0131 100644 --- a/server/query-loop.ts +++ b/server/query-loop.ts @@ -497,6 +497,10 @@ async function _runQueryLoopInner( ...(currentSession.worktreePath ? { wtId: currentSession.wtId } : {}), ...(initialPrompt ? { initialPrompt } : {}), ...(currentSession.telosTaskId ? { telosTaskId: currentSession.telosTaskId } : {}), + ...(currentSession.agentName ? { agentName: currentSession.agentName } : {}), + ...(currentSession.bootContext + ? { bootContext: JSON.stringify(currentSession.bootContext) } + : {}), }); if (initialPrompt) { // Store the initial prompt as a user_message event so diff --git a/server/ws-handler-v2.ts b/server/ws-handler-v2.ts index df1ae187..8906c71d 100644 --- a/server/ws-handler-v2.ts +++ b/server/ws-handler-v2.ts @@ -280,6 +280,14 @@ export function handleReconnect( running, }); + // Re-send cached boot_context so pills reappear after reconnect + if (found && running && found.session?.bootContext) { + ctx.connRegistry.get(connectionId)?.transport.send({ + type: 'boot_context', + ...found.session.bootContext, + }); + } + log.info('reconnect replay', { connectionId, sessionId: entry.sessionId, @@ -371,6 +379,27 @@ export async function handleSwitchSession( }, }); + // Re-send boot_context so pills appear on session switch. + // Hot path: running session in SessionRegistry (in-memory cache). + // Cold path: ended session — read serialized JSON from EventStore. + const found = ctx.sessionRegistry.findBySessionId(msg.sessionId); + if (found?.session?.bootContext) { + ctx.connRegistry.get(connectionId)?.transport.send({ + type: 'boot_context', + ...found.session.bootContext, + }); + } else if (sessionMeta.bootContext) { + try { + const parsed = JSON.parse(sessionMeta.bootContext); + ctx.connRegistry.get(connectionId)?.transport.send({ + type: 'boot_context', + ...parsed, + }); + } catch { + // Invalid JSON — skip + } + } + log.info('switch_session', { connectionId, sessionId: msg.sessionId }); }, ); From 09f38d669615ea46a587fe794234585e4b74dca6 Mon Sep 17 00:00:00 2001 From: dimakis Date: Sat, 30 May 2026 02:19:01 +0100 Subject: [PATCH 2/2] style(session): format chat.ts with prettier Co-Authored-By: Claude Opus 4.6 --- server/chat.ts | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/server/chat.ts b/server/chat.ts index c4891233..fb754fff 100644 --- a/server/chat.ts +++ b/server/chat.ts @@ -831,8 +831,18 @@ async function _startChatInner( tokenCount: 0, tokenBudget: tokenBudget, sources: [] as Array<{ path: string; kind: string }>, - included: [] as Array<{ source: string; heading: string; tokens: number; content: string }>, - trimmed: [] as Array<{ source: string; heading: string; tokens: number; content: string }>, + included: [] as Array<{ + source: string; + heading: string; + tokens: number; + content: string; + }>, + trimmed: [] as Array<{ + source: string; + heading: string; + tokens: number; + content: string; + }>, }; send(transport, { type: 'boot_context', ...fallback }); const s = registry.get(clientId);