@oh-my-pi/pi-coding-agent 13.12.7 → 13.12.9

package/src/session/agent-session.ts

@@ -142,7 +142,13 @@ import {
 	type PythonExecutionMessage,
 	pythonExecutionToText,
 } from "./messages";
-import type { BranchSummaryEntry, CompactionEntry, NewSessionOptions, SessionManager } from "./session-manager";
+import type {
+	BranchSummaryEntry,
+	CompactionEntry,
+	NewSessionOptions,
+	SessionContext,
+	SessionManager,
+} from "./session-manager";
 import { getLatestCompactionEntry } from "./session-manager";
 
 /** Session-specific events that extend the core AgentEvent */
@@ -215,8 +221,10 @@ export interface AgentSessionConfig {
 	rebuildSystemPrompt?: (toolNames: string[], tools: Map<string, AgentTool>) => Promise<string>;
 	/** Enable hidden-by-default MCP tool discovery for this session. */
 	mcpDiscoveryEnabled?: boolean;
-	/** MCP tool names previously selected via discovery in this session. */
+	/** MCP tool names to activate for the current session when discovery mode is enabled. */
 	initialSelectedMCPToolNames?: string[];
+	/** MCP tool names that should seed brand-new sessions created from this AgentSession. */
+	defaultSelectedMCPToolNames?: string[];
 	/** TTSR manager for time-traveling stream rules */
 	ttsrManager?: TtsrManager;
 	/** Secret obfuscator for deobfuscating streaming edit content */
@@ -415,6 +423,8 @@ export class AgentSession {
 	#discoverableMCPTools = new Map<string, DiscoverableMCPTool>();
 	#discoverableMCPSearchIndex: DiscoverableMCPSearchIndex | null = null;
 	#selectedMCPToolNames = new Set<string>();
+	#defaultSelectedMCPToolNames = new Set<string>();
+	#sessionDefaultSelectedMCPToolNames = new Map<string, string[]>();
 
 	// TTSR manager for time-traveling stream rules
 	#ttsrManager: TtsrManager | undefined = undefined;
@@ -464,7 +474,20 @@ export class AgentSession {
 		this.#mcpDiscoveryEnabled = config.mcpDiscoveryEnabled ?? false;
 		this.#setDiscoverableMCPTools(this.#collectDiscoverableMCPToolsFromRegistry());
 		this.#selectedMCPToolNames = new Set(config.initialSelectedMCPToolNames ?? []);
+		this.#defaultSelectedMCPToolNames = new Set(config.defaultSelectedMCPToolNames ?? []);
 		this.#pruneSelectedMCPToolNames();
+		const persistedSelectedMCPToolNames = this.sessionManager.buildSessionContext().selectedMCPToolNames;
+		const currentSelectedMCPToolNames = this.getSelectedMCPToolNames();
+		if (
+			this.#mcpDiscoveryEnabled &&
+			!this.#selectedMCPToolNamesMatch(persistedSelectedMCPToolNames, currentSelectedMCPToolNames)
+		) {
+			this.sessionManager.appendMCPToolSelection(currentSelectedMCPToolNames);
+		}
+		this.#rememberSessionDefaultSelectedMCPToolNames(
+			this.sessionManager.getSessionFile(),
+			this.#defaultSelectedMCPToolNames,
+		);
 		this.#ttsrManager = config.ttsrManager;
 		this.#obfuscator = config.obfuscator;
 		this.agent.providerSessionState = this.#providerSessionState;
@@ -1632,23 +1655,43 @@ export class AgentSession {
 		this.#discoverableMCPSearchIndex = null;
 	}
 
+	#filterSelectableMCPToolNames(toolNames: Iterable<string>): string[] {
+		return Array.from(toolNames).filter(name => this.#discoverableMCPTools.has(name) && this.#toolRegistry.has(name));
+	}
+
 	#pruneSelectedMCPToolNames(): void {
-		for (const name of Array.from(this.#selectedMCPToolNames)) {
-			if (!this.#discoverableMCPTools.has(name) || !this.#toolRegistry.has(name)) {
-				this.#selectedMCPToolNames.delete(name);
-			}
-		}
+		this.#selectedMCPToolNames = new Set(this.#filterSelectableMCPToolNames(this.#selectedMCPToolNames));
 	}
 
-	#getVisibleMCPToolNames(): string[] {
-		if (!this.#mcpDiscoveryEnabled) {
-			return Array.from(this.#toolRegistry.keys()).filter(name => isMCPToolName(name));
-		}
-		return Array.from(this.#selectedMCPToolNames).filter(
-			name => this.#discoverableMCPTools.has(name) && this.#toolRegistry.has(name),
+	#selectedMCPToolNamesMatch(left: string[], right: string[]): boolean {
+		return left.length === right.length && left.every((name, index) => name === right[index]);
+	}
+
+	#rememberSessionDefaultSelectedMCPToolNames(
+		sessionFile: string | null | undefined,
+		toolNames: Iterable<string>,
+	): void {
+		if (!sessionFile) return;
+		this.#sessionDefaultSelectedMCPToolNames.set(
+			path.resolve(sessionFile),
+			this.#filterSelectableMCPToolNames(toolNames),
 		);
 	}
 
+	#getSessionDefaultSelectedMCPToolNames(sessionFile: string | null | undefined): string[] {
+		if (!sessionFile) return [];
+		return this.#sessionDefaultSelectedMCPToolNames.get(path.resolve(sessionFile)) ?? [];
+	}
+
+	#persistSelectedMCPToolNamesIfChanged(previousSelectedMCPToolNames: string[]): void {
+		if (!this.#mcpDiscoveryEnabled) return;
+		const nextSelectedMCPToolNames = this.getSelectedMCPToolNames();
+		if (this.#selectedMCPToolNamesMatch(previousSelectedMCPToolNames, nextSelectedMCPToolNames)) {
+			return;
+		}
+		this.sessionManager.appendMCPToolSelection(nextSelectedMCPToolNames);
+	}
+
 	#getActiveNonMCPToolNames(): string[] {
 		return this.getActiveToolNames().filter(name => !isMCPToolName(name) && this.#toolRegistry.has(name));
 	}
@@ -1699,35 +1742,35 @@ export class AgentSession {
 		if (!this.#mcpDiscoveryEnabled) {
 			return this.getActiveToolNames().filter(name => isMCPToolName(name) && this.#toolRegistry.has(name));
 		}
-		return Array.from(this.#selectedMCPToolNames).filter(
-			name => this.#discoverableMCPTools.has(name) && this.#toolRegistry.has(name),
-		);
+		return this.#filterSelectableMCPToolNames(this.#selectedMCPToolNames);
 	}
 
 	async activateDiscoveredMCPTools(toolNames: string[]): Promise<string[]> {
+		const nextSelectedMCPToolNames = new Set(this.#selectedMCPToolNames);
 		const activated: string[] = [];
 		for (const name of toolNames) {
 			if (!isMCPToolName(name) || !this.#discoverableMCPTools.has(name) || !this.#toolRegistry.has(name)) {
 				continue;
 			}
-			this.#selectedMCPToolNames.add(name);
+			nextSelectedMCPToolNames.add(name);
 			activated.push(name);
 		}
 		if (activated.length === 0) {
 			return [];
 		}
-		const nextActive = [...this.#getActiveNonMCPToolNames(), ...this.#getVisibleMCPToolNames()];
+		const nextActive = [
+			...this.#getActiveNonMCPToolNames(),
+			...this.#filterSelectableMCPToolNames(nextSelectedMCPToolNames),
+		];
 		await this.setActiveToolsByName(nextActive);
 		return [...new Set(activated)];
 	}
 
-	/**
-	 * Set active tools by name.
-	 * Only tools in the registry can be enabled. Unknown tool names are ignored.
-	 * Also rebuilds the system prompt to reflect the new tool set.
-	 * Changes take effect before the next model call.
-	 */
-	async setActiveToolsByName(toolNames: string[]): Promise<void> {
+	async #applyActiveToolsByName(
+		toolNames: string[],
+		options?: { persistMCPSelection?: boolean; previousSelectedMCPToolNames?: string[] },
+	): Promise<void> {
+		const previousSelectedMCPToolNames = options?.previousSelectedMCPToolNames ?? this.getSelectedMCPToolNames();
 		const tools: AgentTool[] = [];
 		const validToolNames: string[] = [];
 		for (const name of toolNames) {
@@ -1751,8 +1794,36 @@ export class AgentSession {
 			this.#baseSystemPrompt = await this.#rebuildSystemPrompt(validToolNames, this.#toolRegistry);
 			this.agent.setSystemPrompt(this.#baseSystemPrompt);
 		}
+		if (options?.persistMCPSelection !== false) {
+			this.#persistSelectedMCPToolNamesIfChanged(previousSelectedMCPToolNames);
+		}
 	}
 
+	/**
+	 * Set active tools by name.
+	 * Only tools in the registry can be enabled. Unknown tool names are ignored.
+	 * Also rebuilds the system prompt to reflect the new tool set.
+	 * Changes take effect before the next model call.
+	 */
+	async setActiveToolsByName(toolNames: string[]): Promise<void> {
+		await this.#applyActiveToolsByName(toolNames);
+	}
+
+	async #restoreMCPSelectionsForSessionContext(
+		sessionContext: SessionContext,
+		options?: { fallbackSelectedMCPToolNames?: Iterable<string> },
+	): Promise<void> {
+		if (!this.#mcpDiscoveryEnabled) return;
+		const nextActiveNonMCPToolNames = this.#getActiveNonMCPToolNames();
+		const fallbackSelectedMCPToolNames = options?.fallbackSelectedMCPToolNames ?? this.#defaultSelectedMCPToolNames;
+		const restoredMCPToolNames = sessionContext.hasPersistedMCPToolSelection
+			? this.#filterSelectableMCPToolNames(sessionContext.selectedMCPToolNames)
+			: this.#filterSelectableMCPToolNames(fallbackSelectedMCPToolNames);
+		this.#rememberSessionDefaultSelectedMCPToolNames(this.sessionFile, restoredMCPToolNames);
+		await this.#applyActiveToolsByName([...nextActiveNonMCPToolNames, ...restoredMCPToolNames], {
+			persistMCPSelection: false,
+		});
+	}
 	/** Rebuild the base system prompt using the current active tool set. */
 	async refreshBaseSystemPrompt(): Promise<void> {
 		if (!this.#rebuildSystemPrompt) return;
@@ -1766,6 +1837,7 @@ export class AgentSession {
 	 * This allows /mcp add/remove/reauth to take effect without restarting the session.
 	 */
 	async refreshMCPTools(mcpTools: CustomTool[]): Promise<void> {
+		const previousSelectedMCPToolNames = this.getSelectedMCPToolNames();
 		const existingNames = Array.from(this.#toolRegistry.keys());
 		for (const name of existingNames) {
 			if (isMCPToolName(name)) {
@@ -1796,7 +1868,7 @@ export class AgentSession {
 		this.#pruneSelectedMCPToolNames();
 
 		const nextActive = [...this.#getActiveNonMCPToolNames(), ...this.getSelectedMCPToolNames()];
-		await this.setActiveToolsByName(nextActive);
+		await this.#applyActiveToolsByName(nextActive, { previousSelectedMCPToolNames });
 	}
 
 	/** Whether auto-compaction is currently running */
@@ -2747,6 +2819,12 @@ export class AgentSession {
 	 */
 	async newSession(options?: NewSessionOptions): Promise<boolean> {
 		const previousSessionFile = this.sessionFile;
+		const nextDiscoverySessionToolNames = this.#mcpDiscoveryEnabled
+			? [
+					...this.#getActiveNonMCPToolNames(),
+					...this.#filterSelectableMCPToolNames(this.#defaultSelectedMCPToolNames),
+				]
+			: undefined;
 
 		// Emit session_before_switch event with reason "new" (can be cancelled)
 		if (this.#extensionRunner?.hasHandlers("session_before_switch")) {
@@ -2774,6 +2852,13 @@ export class AgentSession {
 
 		this.sessionManager.appendThinkingLevelChange(this.thinkingLevel);
 		this.sessionManager.appendServiceTierChange(this.serviceTier ?? null);
+		if (nextDiscoverySessionToolNames) {
+			await this.#applyActiveToolsByName(nextDiscoverySessionToolNames, { persistMCPSelection: false });
+			if (this.getSelectedMCPToolNames().length > 0) {
+				this.sessionManager.appendMCPToolSelection(this.getSelectedMCPToolNames());
+			}
+		}
+		this.#rememberSessionDefaultSelectedMCPToolNames(this.sessionFile, this.#defaultSelectedMCPToolNames);
 
 		this.#todoReminderCount = 0;
 		this.#planReferenceSent = false;
@@ -4862,6 +4947,8 @@ export class AgentSession {
 
 		// Reload messages
 		const sessionContext = this.sessionManager.buildSessionContext();
+		const fallbackSelectedMCPToolNames = this.#getSessionDefaultSelectedMCPToolNames(sessionPath);
+		await this.#restoreMCPSelectionsForSessionContext(sessionContext, { fallbackSelectedMCPToolNames });
 
 		// Emit session_switch event to hooks
 		if (this.#extensionRunner) {
@@ -4965,6 +5052,8 @@ export class AgentSession {
 		// Reload messages from entries (works for both file and in-memory mode)
 		const sessionContext = this.sessionManager.buildSessionContext();
 
+		await this.#restoreMCPSelectionsForSessionContext(sessionContext);
+
 		// Emit session_branch event to hooks (after branch completes)
 		if (this.#extensionRunner) {
 			await this.#extensionRunner.emit({
@@ -5128,6 +5217,7 @@ export class AgentSession {
 
 		// Update agent state
 		const sessionContext = this.sessionManager.buildSessionContext();
+		await this.#restoreMCPSelectionsForSessionContext(sessionContext);
 		this.agent.replaceMessages(sessionContext.messages);
 		this.#syncTodoPhasesFromBranch();
 		this.#closeCodexProviderSessionsForHistoryRewrite();

package/src/session/agent-storage.ts

@@ -18,7 +18,8 @@ type ModelUsageRow = {
 };
 
 /** Bump when schema changes require migration */
-const SCHEMA_VERSION = 4;
+const SCHEMA_VERSION = 5;
+const SQLITE_NOW_EPOCH = "CAST(strftime('%s','now') AS INTEGER)";
 
 /** Singleton instances per database path */
 const instances = new Map<string, AgentStorage>();
@@ -59,9 +60,8 @@ export class AgentStorage {
 		this.#authStore = new AuthCredentialStore(this.#db);
 
 		this.#listSettingsStmt = this.#db.prepare("SELECT key, value FROM settings");
-
 		this.#upsertModelUsageStmt = this.#db.prepare(
-			"INSERT INTO model_usage (model_key, last_used_at) VALUES (?, unixepoch()) ON CONFLICT(model_key) DO UPDATE SET last_used_at = unixepoch()",
+			`INSERT INTO model_usage (model_key, last_used_at) VALUES (?, ${SQLITE_NOW_EPOCH}) ON CONFLICT(model_key) DO UPDATE SET last_used_at = ${SQLITE_NOW_EPOCH}`,
 		);
 		this.#listModelUsageStmt = this.#db.prepare(
 			"SELECT model_key, last_used_at FROM model_usage ORDER BY last_used_at DESC",
@@ -80,7 +80,7 @@ PRAGMA busy_timeout=5000;
 
 CREATE TABLE IF NOT EXISTS model_usage (
 	model_key TEXT PRIMARY KEY,
-	last_used_at INTEGER NOT NULL DEFAULT (unixepoch())
+	last_used_at INTEGER NOT NULL DEFAULT (${SQLITE_NOW_EPOCH})
 );
 
 CREATE TABLE IF NOT EXISTS schema_version (version INTEGER PRIMARY KEY);
@@ -96,7 +96,7 @@ CREATE TABLE IF NOT EXISTS schema_version (version INTEGER PRIMARY KEY);
 CREATE TABLE settings (
 	key TEXT PRIMARY KEY,
 	value TEXT NOT NULL,
-	updated_at INTEGER NOT NULL DEFAULT (unixepoch())
+	updated_at INTEGER NOT NULL DEFAULT (${SQLITE_NOW_EPOCH})
 );
 `);
 		} else if (!hasKey || !hasValue) {
@@ -122,12 +122,12 @@ CREATE TABLE settings (
 CREATE TABLE settings (
 	key TEXT PRIMARY KEY,
 	value TEXT NOT NULL,
-	updated_at INTEGER NOT NULL DEFAULT (unixepoch())
+	updated_at INTEGER NOT NULL DEFAULT (${SQLITE_NOW_EPOCH})
 );
 `);
 				if (settings) {
 					const insert = this.#db.prepare(
-						"INSERT INTO settings (key, value, updated_at) VALUES (?, ?, unixepoch())",
+						`INSERT INTO settings (key, value, updated_at) VALUES (?, ?, ${SQLITE_NOW_EPOCH})`,
 					);
 					for (const [key, value] of Object.entries(settings)) {
 						if (value === undefined) continue;
@@ -144,12 +144,15 @@ CREATE TABLE settings (
 		const versionRow = this.#db.prepare("SELECT version FROM schema_version ORDER BY version DESC LIMIT 1").get() as
 			| { version?: number }
 			| undefined;
+		const schemaVersion = typeof versionRow?.version === "number" ? versionRow.version : 0;
 		if (versionRow?.version !== undefined && versionRow.version !== SCHEMA_VERSION) {
 			logger.warn("AgentStorage schema version mismatch", {
 				current: versionRow.version,
 				expected: SCHEMA_VERSION,
 			});
-			this.#migrateSchema(versionRow.version);
+		}
+		if (schemaVersion < SCHEMA_VERSION) {
+			this.#migrateSchema(schemaVersion);
 		}
 		this.#db.prepare("INSERT OR REPLACE INTO schema_version(version) VALUES (?)").run(SCHEMA_VERSION);
 	}
@@ -159,6 +162,43 @@ CREATE TABLE settings (
 			// v3 → v4: Add disabled column to auth_credentials (handled by AuthCredentialStore)
 			// Nothing to do here - AuthCredentialStore will handle this migration
 		}
+		if (fromVersion < 5) {
+			this.#migrateSchemaV4ToV5();
+		}
+	}
+
+	#migrateSchemaV4ToV5(): void {
+		const migrate = this.#db.transaction(() => {
+			this.#db.exec("ALTER TABLE settings RENAME TO settings_legacy");
+			this.#db.exec(`
+CREATE TABLE settings (
+	key TEXT PRIMARY KEY,
+	value TEXT NOT NULL,
+	updated_at INTEGER NOT NULL DEFAULT (${SQLITE_NOW_EPOCH})
+);
+`);
+			this.#db.exec(`
+INSERT INTO settings (key, value, updated_at)
+SELECT key, value, updated_at
+FROM settings_legacy
+`);
+			this.#db.exec("DROP TABLE settings_legacy");
+
+			this.#db.exec("ALTER TABLE model_usage RENAME TO model_usage_legacy");
+			this.#db.exec(`
+CREATE TABLE model_usage (
+	model_key TEXT PRIMARY KEY,
+	last_used_at INTEGER NOT NULL DEFAULT (${SQLITE_NOW_EPOCH})
+);
+`);
+			this.#db.exec(`
+INSERT INTO model_usage (model_key, last_used_at)
+SELECT model_key, last_used_at
+FROM model_usage_legacy
+`);
+			this.#db.exec("DROP TABLE model_usage_legacy");
+		});
+		migrate();
 	}
 
 	/**

package/src/session/history-storage.ts

@@ -17,6 +17,8 @@ type HistoryRow = {
 	cwd: string | null;
 };
 
+const SQLITE_NOW_EPOCH = "CAST(strftime('%s','now') AS INTEGER)";
+
 export class HistoryStorage {
 	#db: Database;
 	static #instance?: HistoryStorage;
@@ -45,7 +47,7 @@ PRAGMA busy_timeout=5000;
 CREATE TABLE IF NOT EXISTS history (
 	id INTEGER PRIMARY KEY AUTOINCREMENT,
 	prompt TEXT NOT NULL,
-	created_at INTEGER NOT NULL DEFAULT (unixepoch()),
+	created_at INTEGER NOT NULL DEFAULT (${SQLITE_NOW_EPOCH}),
 	cwd TEXT
 );
 CREATE INDEX IF NOT EXISTS idx_history_created_at ON history(created_at DESC);
@@ -54,8 +56,12 @@ CREATE VIRTUAL TABLE IF NOT EXISTS history_fts USING fts5(prompt, content='histo
 
 CREATE TRIGGER IF NOT EXISTS history_ai AFTER INSERT ON history BEGIN
 	INSERT INTO history_fts(rowid, prompt) VALUES (new.id, new.prompt);
-END;
-`);
+	END;
+	`);
+
+		if (this.#historySchemaUsesUnixEpoch()) {
+			this.#migrateHistorySchema();
+		}
 
 		if (!hasFts) {
 			try {
@@ -135,6 +141,41 @@ END;
 		fs.mkdirSync(dir, { recursive: true });
 	}
 
+	#historySchemaUsesUnixEpoch(): boolean {
+		const row = this.#db.prepare("SELECT sql FROM sqlite_master WHERE type = 'table' AND name = 'history'").get() as
+			| { sql?: string | null }
+			| undefined;
+		return row?.sql?.includes("unixepoch(") ?? false;
+	}
+
+	#migrateHistorySchema(): void {
+		const migrate = this.#db.transaction(() => {
+			this.#db.exec("ALTER TABLE history RENAME TO history_legacy");
+			this.#db.exec("DROP INDEX IF EXISTS idx_history_created_at");
+			this.#db.exec("DROP TRIGGER IF EXISTS history_ai");
+			this.#db.exec("DROP TABLE IF EXISTS history_fts");
+			this.#db.exec(`
+CREATE TABLE history (
+	id INTEGER PRIMARY KEY AUTOINCREMENT,
+	prompt TEXT NOT NULL,
+	created_at INTEGER NOT NULL DEFAULT (${SQLITE_NOW_EPOCH}),
+	cwd TEXT
+);
+CREATE INDEX IF NOT EXISTS idx_history_created_at ON history(created_at DESC);
+INSERT INTO history (id, prompt, created_at, cwd)
+SELECT id, prompt, created_at, cwd
+FROM history_legacy;
+DROP TABLE history_legacy;
+CREATE VIRTUAL TABLE history_fts USING fts5(prompt, content='history', content_rowid='id');
+CREATE TRIGGER history_ai AFTER INSERT ON history BEGIN
+	INSERT INTO history_fts(rowid, prompt) VALUES (new.id, new.prompt);
+END;
+			`);
+			this.#db.run("INSERT INTO history_fts(history_fts) VALUES('rebuild')");
+		});
+		migrate();
+	}
+
 	#normalizeLimit(limit: number): number {
 		if (!Number.isFinite(limit)) return 0;
 		const clamped = Math.max(0, Math.floor(limit));