From 934c2bc5d3a0c9667667bdee08b349f32cd4466d Mon Sep 17 00:00:00 2001 From: Mario Zechner Date: Tue, 9 Dec 2025 00:14:47 +0100 Subject: [PATCH] WP9+WP10: Add AgentSession session management and utility methods --- packages/coding-agent/docs/refactor.md | 22 +- .../coding-agent/src/core/agent-session.ts | 201 ++++++++++++++++++ packages/coding-agent/src/core/index.ts | 1 + 3 files changed, 213 insertions(+), 11 deletions(-) diff --git a/packages/coding-agent/docs/refactor.md b/packages/coding-agent/docs/refactor.md index 4f7ce79e..5a5be3f8 100644 --- a/packages/coding-agent/docs/refactor.md +++ b/packages/coding-agent/docs/refactor.md @@ -1075,13 +1075,13 @@ exportToHtml(outputPath?: string): string { **Verification:** 1. `npm run check` passes -- [ ] Add `SessionStats` interface -- [ ] Add `switchSession()` method -- [ ] Add `branch()` method -- [ ] Add `getUserMessagesForBranching()` method -- [ ] Add `getSessionStats()` method -- [ ] Add `exportToHtml()` method -- [ ] Verify with `npm run check` +- [x] Add `SessionStats` interface +- [x] Add `switchSession()` method +- [x] Add `branch()` method +- [x] Add `getUserMessagesForBranching()` method +- [x] Add `getSessionStats()` method +- [x] Add `exportToHtml()` method +- [x] Verify with `npm run check` --- @@ -1138,10 +1138,10 @@ getQueuedMessages(): readonly string[] { **Verification:** 1. `npm run check` passes -- [ ] Add `getLastAssistantText()` method -- [ ] Add `queuedMessageCount` getter -- [ ] Add `getQueuedMessages()` method -- [ ] Verify with `npm run check` +- [x] Add `getLastAssistantText()` method +- [x] Add `queuedMessageCount` getter (done in WP4) +- [x] Add `getQueuedMessages()` method (done in WP4) +- [x] Verify with `npm run check` --- diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index 8d60b905..b4b0b9f6 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -17,6 +17,7 @@ import type { Agent, AgentEvent, AgentState, AppMessage, Attachment, ThinkingLev import type { AssistantMessage, Model } from "@mariozechner/pi-ai"; import { calculateContextTokens, compact, shouldCompact } from "../compaction.js"; import { getModelsPath } from "../config.js"; +import { exportSessionToHtml } from "../export-html.js"; import type { BashExecutionMessage } from "../messages.js"; import { getApiKeyForModel, getAvailableModels } from "../model-config.js"; import { loadSessionFromEntries, type SessionManager } from "../session-manager.js"; @@ -63,6 +64,25 @@ export interface CompactionResult { summary: string; } +/** Session statistics for /session command */ +export interface SessionStats { + sessionFile: string; + sessionId: string; + userMessages: number; + assistantMessages: number; + toolCalls: number; + toolResults: number; + totalMessages: number; + tokens: { + input: number; + output: number; + cacheRead: number; + cacheWrite: number; + total: number; + }; + cost: number; +} + // ============================================================================ // AgentSession Class // ============================================================================ @@ -681,4 +701,185 @@ export class AgentSession { get isBashRunning(): boolean { return this._bashAbortController !== null; } + + // ========================================================================= + // Session Management + // ========================================================================= + + /** + * Switch to a different session file. + * Aborts current operation, loads messages, restores model/thinking. + * Listeners are preserved and will continue receiving events. + */ + async switchSession(sessionPath: string): Promise { + this._disconnectFromAgent(); + await this.abort(); + this._queuedMessages = []; + + // Set new session + this.sessionManager.setSessionFile(sessionPath); + + // Reload messages + const loaded = loadSessionFromEntries(this.sessionManager.loadEntries()); + this.agent.replaceMessages(loaded.messages); + + // Restore model if saved + const savedModel = this.sessionManager.loadModel(); + if (savedModel) { + const availableModels = (await getAvailableModels()).models; + const match = availableModels.find((m) => m.provider === savedModel.provider && m.id === savedModel.modelId); + if (match) { + this.agent.setModel(match); + } + } + + // Restore thinking level if saved + const savedThinking = this.sessionManager.loadThinkingLevel(); + if (savedThinking) { + this.agent.setThinkingLevel(savedThinking as ThinkingLevel); + } + + this._reconnectToAgent(); + } + + /** + * Create a branch from a specific entry index. + * @param entryIndex Index into session entries to branch from + * @returns The text of the selected user message (for editor pre-fill) + */ + branch(entryIndex: number): string { + const entries = this.sessionManager.loadEntries(); + const selectedEntry = entries[entryIndex]; + + if (!selectedEntry || selectedEntry.type !== "message" || selectedEntry.message.role !== "user") { + throw new Error("Invalid entry index for branching"); + } + + const selectedText = this._extractUserMessageText(selectedEntry.message.content); + + // Create branched session + const newSessionFile = this.sessionManager.createBranchedSessionFromEntries(entries, entryIndex); + this.sessionManager.setSessionFile(newSessionFile); + + // Reload + const loaded = loadSessionFromEntries(this.sessionManager.loadEntries()); + this.agent.replaceMessages(loaded.messages); + + return selectedText; + } + + /** + * Get all user messages from session for branch selector. + */ + getUserMessagesForBranching(): Array<{ entryIndex: number; text: string }> { + const entries = this.sessionManager.loadEntries(); + const result: Array<{ entryIndex: number; text: string }> = []; + + for (let i = 0; i < entries.length; i++) { + const entry = entries[i]; + if (entry.type !== "message") continue; + if (entry.message.role !== "user") continue; + + const text = this._extractUserMessageText(entry.message.content); + if (text) { + result.push({ entryIndex: i, text }); + } + } + + return result; + } + + private _extractUserMessageText(content: string | Array<{ type: string; text?: string }>): string { + if (typeof content === "string") return content; + if (Array.isArray(content)) { + return content + .filter((c): c is { type: "text"; text: string } => c.type === "text") + .map((c) => c.text) + .join(""); + } + return ""; + } + + /** + * Get session statistics. + */ + getSessionStats(): SessionStats { + const state = this.state; + const userMessages = state.messages.filter((m) => m.role === "user").length; + const assistantMessages = state.messages.filter((m) => m.role === "assistant").length; + const toolResults = state.messages.filter((m) => m.role === "toolResult").length; + + let toolCalls = 0; + let totalInput = 0; + let totalOutput = 0; + let totalCacheRead = 0; + let totalCacheWrite = 0; + let totalCost = 0; + + for (const message of state.messages) { + if (message.role === "assistant") { + const assistantMsg = message as AssistantMessage; + toolCalls += assistantMsg.content.filter((c) => c.type === "toolCall").length; + totalInput += assistantMsg.usage.input; + totalOutput += assistantMsg.usage.output; + totalCacheRead += assistantMsg.usage.cacheRead; + totalCacheWrite += assistantMsg.usage.cacheWrite; + totalCost += assistantMsg.usage.cost.total; + } + } + + return { + sessionFile: this.sessionFile, + sessionId: this.sessionId, + userMessages, + assistantMessages, + toolCalls, + toolResults, + totalMessages: state.messages.length, + tokens: { + input: totalInput, + output: totalOutput, + cacheRead: totalCacheRead, + cacheWrite: totalCacheWrite, + total: totalInput + totalOutput + totalCacheRead + totalCacheWrite, + }, + cost: totalCost, + }; + } + + /** + * Export session to HTML. + * @param outputPath Optional output path (defaults to session directory) + * @returns Path to exported file + */ + exportToHtml(outputPath?: string): string { + return exportSessionToHtml(this.sessionManager, this.state, outputPath); + } + + // ========================================================================= + // Utilities + // ========================================================================= + + /** + * Get text content of last assistant message. + * Useful for /copy command. + * @returns Text content, or null if no assistant message exists + */ + getLastAssistantText(): string | null { + const lastAssistant = this.messages + .slice() + .reverse() + .find((m) => m.role === "assistant"); + + if (!lastAssistant) return null; + + let text = ""; + for (const content of (lastAssistant as AssistantMessage).content) { + if (content.type === "text") { + text += content.text; + } + } + + return text.trim() || null; + } } diff --git a/packages/coding-agent/src/core/index.ts b/packages/coding-agent/src/core/index.ts index ca42285d..6ce21a1f 100644 --- a/packages/coding-agent/src/core/index.ts +++ b/packages/coding-agent/src/core/index.ts @@ -9,5 +9,6 @@ export { type CompactionResult, type ModelCycleResult, type PromptOptions, + type SessionStats, } from "./agent-session.js"; export { type BashExecutorOptions, type BashResult, executeBash } from "./bash-executor.js";