ctb 1.1.0 → 1.2.1

package/src/session.ts

@@ -8,6 +8,7 @@
 import { readFileSync } from "node:fs";
 import {
 	type Options,
+	type Query,
 	query,
 	type SDKMessage,
 } from "@anthropic-ai/claude-agent-sdk";
@@ -31,8 +32,9 @@ import type { SessionData, StatusCallback, TokenUsage } from "./types";
 
 /**
  * Determine thinking token budget based on message keywords.
+ * Exported for testing.
  */
-function getThinkingLevel(message: string): number {
+export function getThinkingLevel(message: string): number {
 	const msgLower = message.toLowerCase();
 
 	// Check deep thinking triggers first (more specific)
@@ -78,6 +80,21 @@ class ClaudeSession {
 	lastErrorTime: Date | null = null;
 	lastUsage: TokenUsage | null = null;
 	lastMessage: string | null = null;
+	lastBotResponse: string | null = null;
+
+	// Model and mode settings
+	currentModel: "sonnet" | "opus" | "haiku" = "sonnet";
+	forceThinking: number | null = null; // Tokens for next message, then resets
+	planMode = false;
+
+	// Cumulative usage tracking
+	totalInputTokens = 0;
+	totalOutputTokens = 0;
+	totalCacheReadTokens = 0;
+
+	// File checkpointing for /undo
+	private _queryInstance: Query | null = null;
+	private _userMessageUuids: string[] = [];
 
 	// Mutable working directory (can be changed with /cd)
 	private _workingDir: string = WORKING_DIR;
@@ -192,11 +209,27 @@ class ClaudeSession {
 		}
 
 		const isNewSession = !this.isActive;
-		const thinkingTokens = getThinkingLevel(message);
+
+		// Determine thinking tokens - forceThinking overrides keyword detection
+		let thinkingTokens: number;
+		if (this.forceThinking !== null) {
+			thinkingTokens = this.forceThinking;
+			this.forceThinking = null; // Reset after use
+		} else {
+			thinkingTokens = getThinkingLevel(message);
+		}
 		const thinkingLabel =
 			{ 0: "off", 10000: "normal", 50000: "deep" }[thinkingTokens] ||
 			String(thinkingTokens);
 
+		// Determine model based on currentModel setting
+		const modelMap = {
+			sonnet: "claude-sonnet-4-5",
+			opus: "claude-opus-4-5",
+			haiku: "claude-haiku-3-5",
+		};
+		const modelId = modelMap[this.currentModel];
+
 		// Inject current date/time at session start so Claude doesn't need to call a tool for it
 		let messageToSend = message;
 		if (isNewSession) {
@@ -218,16 +251,17 @@ class ClaudeSession {
 
 		// Build SDK V1 options - supports all features
 		const options: Options = {
-			model: "claude-sonnet-4-5",
+			model: modelId,
 			cwd: this._workingDir,
 			settingSources: ["user", "project"],
-			permissionMode: "bypassPermissions",
-			allowDangerouslySkipPermissions: true,
+			permissionMode: this.planMode ? "plan" : "bypassPermissions",
+			allowDangerouslySkipPermissions: !this.planMode,
 			systemPrompt: SAFETY_PROMPT,
 			mcpServers: MCP_SERVERS,
 			maxThinkingTokens: thinkingTokens,
 			additionalDirectories: ALLOWED_PATHS,
 			resume: this.sessionId || undefined,
+			enableFileCheckpointing: true, // Enable /undo support
 		};
 
 		// Add Claude Code executable path if set (required for standalone builds)
@@ -281,6 +315,9 @@ class ClaudeSession {
 				},
 			});
 
+			// Store query instance for /undo support
+			this._queryInstance = queryInstance;
+
 			// Process streaming response
 			for await (const event of queryInstance) {
 				// Check for abort
@@ -308,6 +345,12 @@ class ClaudeSession {
 					this.saveSession();
 				}
 
+				// Capture user message UUIDs for /undo checkpoints
+				if (event.type === "user" && event.uuid) {
+					this._userMessageUuids.push(event.uuid);
+					console.log(`Checkpoint: user message ${event.uuid.slice(0, 8)}...`);
+				}
+
 				// Handle different message types
 				if (event.type === "assistant") {
 					for (const block of event.message.content) {
@@ -436,6 +479,10 @@ class ClaudeSession {
 					if ("usage" in event && event.usage) {
 						this.lastUsage = event.usage as TokenUsage;
 						const u = this.lastUsage;
+						// Accumulate totals
+						this.totalInputTokens += u.input_tokens || 0;
+						this.totalOutputTokens += u.output_tokens || 0;
+						this.totalCacheReadTokens += u.cache_read_input_tokens || 0;
 						console.log(
 							`Usage: in=${u.input_tokens} out=${u.output_tokens} cache_read=${
 								u.cache_read_input_tokens || 0
@@ -486,7 +533,9 @@ class ClaudeSession {
 
 		await statusCallback("done", "");
 
-		return responseParts.join("") || "No response from Claude.";
+		const finalResponse = responseParts.join("") || "No response from Claude.";
+		this.lastBotResponse = finalResponse;
+		return finalResponse;
 	}
 
 	/**
@@ -495,9 +544,91 @@ class ClaudeSession {
 	async kill(): Promise<void> {
 		this.sessionId = null;
 		this.lastActivity = null;
+		// Reset usage totals
+		this.totalInputTokens = 0;
+		this.totalOutputTokens = 0;
+		this.totalCacheReadTokens = 0;
+		// Reset modes
+		this.planMode = false;
+		this.forceThinking = null;
+		// Reset checkpoints
+		this._queryInstance = null;
+		this._userMessageUuids = [];
 		console.log("Session cleared");
 	}
 
+	/**
+	 * Check if undo is available (has checkpoints).
+	 */
+	get canUndo(): boolean {
+		return this._queryInstance !== null && this._userMessageUuids.length > 0;
+	}
+
+	/**
+	 * Get number of available undo checkpoints.
+	 */
+	get undoCheckpoints(): number {
+		return this._userMessageUuids.length;
+	}
+
+	/**
+	 * Undo file changes by rewinding to the last user message checkpoint.
+	 * Returns [success, message].
+	 */
+	async undo(): Promise<[boolean, string]> {
+		if (!this._queryInstance) {
+			return [false, "No active session to undo"];
+		}
+
+		if (this._userMessageUuids.length === 0) {
+			return [false, "No checkpoints available"];
+		}
+
+		// Get and remove the last user message UUID
+		const targetUuid = this._userMessageUuids.pop();
+		if (!targetUuid) {
+			return [false, "No checkpoints available"];
+		}
+
+		try {
+			console.log(`Rewinding files to checkpoint ${targetUuid.slice(0, 8)}...`);
+			await this._queryInstance.rewindFiles(targetUuid);
+
+			const remaining = this._userMessageUuids.length;
+			return [
+				true,
+				`✅ Reverted file changes to checkpoint \`${targetUuid.slice(0, 8)}...\`\n${remaining} checkpoint${remaining !== 1 ? "s" : ""} remaining`,
+			];
+		} catch (error) {
+			// Restore the checkpoint on failure
+			this._userMessageUuids.push(targetUuid);
+			console.error(`Undo failed: ${error}`);
+			return [false, `Failed to undo: ${error}`];
+		}
+	}
+
+	/**
+	 * Estimate cost based on current model and usage.
+	 */
+	estimateCost(): { inputCost: number; outputCost: number; total: number } {
+		// Pricing per 1M tokens (approximate as of 2024)
+		const pricing = {
+			sonnet: { input: 3, output: 15 },
+			opus: { input: 15, output: 75 },
+			haiku: { input: 0.25, output: 1.25 },
+		};
+		const rates = pricing[this.currentModel];
+
+		const inputCost = (this.totalInputTokens / 1_000_000) * rates.input;
+		const outputCost = (this.totalOutputTokens / 1_000_000) * rates.output;
+
+		return {
+			inputCost,
+			outputCost,
+			total: inputCost + outputCost,
+		};
+	}
+
 	/**
 	 * Save session to disk for resume after restart.
 	 */
@@ -561,5 +692,8 @@ class ClaudeSession {
 	}
 }
 
+// Export class for testing
+export { ClaudeSession };
+
 // Global session instance
 export const session = new ClaudeSession();