ctb 1.0.0 → 1.2.1

package/src/handlers/text.ts

@@ -2,127 +2,205 @@
  * Text message handler for Claude Telegram Bot.
  */
 
+import { spawn } from "node:child_process";
 import type { Context } from "grammy";
-import { session } from "../session";
 import { ALLOWED_USERS } from "../config";
 import { isAuthorized, rateLimiter } from "../security";
+import { session } from "../session";
 import {
-  auditLog,
-  auditLogRateLimit,
-  checkInterrupt,
-  startTypingIndicator,
+	auditLog,
+	auditLogRateLimit,
+	checkInterrupt,
+	startTypingIndicator,
 } from "../utils";
-import { StreamingState, createStatusCallback } from "./streaming";
+import { createStatusCallback, StreamingState } from "./streaming";
+
+/**
+ * Execute a shell command and return output.
+ */
+async function execShellCommand(
+	command: string,
+	cwd: string,
+): Promise<{ stdout: string; stderr: string; exitCode: number }> {
+	return new Promise((resolve) => {
+		const proc = spawn("bash", ["-c", command], {
+			cwd,
+			timeout: 30000, // 30s timeout
+		});
+
+		let stdout = "";
+		let stderr = "";
+
+		proc.stdout.on("data", (data) => {
+			stdout += data.toString();
+		});
+
+		proc.stderr.on("data", (data) => {
+			stderr += data.toString();
+		});
+
+		proc.on("close", (code) => {
+			resolve({ stdout, stderr, exitCode: code ?? 0 });
+		});
+
+		proc.on("error", (err) => {
+			resolve({ stdout, stderr: err.message, exitCode: 1 });
+		});
+	});
+}
 
 /**
  * Handle incoming text messages.
  */
 export async function handleText(ctx: Context): Promise<void> {
-  const userId = ctx.from?.id;
-  const username = ctx.from?.username || "unknown";
-  const chatId = ctx.chat?.id;
-  let message = ctx.message?.text;
-
-  if (!userId || !message || !chatId) {
-    return;
-  }
-
-  // 1. Authorization check
-  if (!isAuthorized(userId, ALLOWED_USERS)) {
-    await ctx.reply("Unauthorized. Contact the bot owner for access.");
-    return;
-  }
-
-  // 2. Check for interrupt prefix
-  message = await checkInterrupt(message);
-  if (!message.trim()) {
-    return;
-  }
-
-  // 3. Rate limit check
-  const [allowed, retryAfter] = rateLimiter.check(userId);
-  if (!allowed) {
-    await auditLogRateLimit(userId, username, retryAfter!);
-    await ctx.reply(
-      `⏳ Rate limited. Please wait ${retryAfter!.toFixed(1)} seconds.`
-    );
-    return;
-  }
-
-  // 4. Store message for retry
-  session.lastMessage = message;
-
-  // 5. Mark processing started
-  const stopProcessing = session.startProcessing();
-
-  // 6. Start typing indicator
-  const typing = startTypingIndicator(ctx);
-
-  // 7. Create streaming state and callback
-  let state = new StreamingState();
-  let statusCallback = createStatusCallback(ctx, state);
-
-  // 8. Send to Claude with retry logic for crashes
-  const MAX_RETRIES = 1;
-
-  for (let attempt = 0; attempt <= MAX_RETRIES; attempt++) {
-    try {
-      const response = await session.sendMessageStreaming(
-        message,
-        username,
-        userId,
-        statusCallback,
-        chatId,
-        ctx
-      );
-
-      // 9. Audit log
-      await auditLog(userId, username, "TEXT", message, response);
-      break; // Success - exit retry loop
-    } catch (error) {
-      const errorStr = String(error);
-      const isClaudeCodeCrash = errorStr.includes("exited with code");
-
-      // Clean up any partial messages from this attempt
-      for (const toolMsg of state.toolMessages) {
-        try {
-          await ctx.api.deleteMessage(toolMsg.chat.id, toolMsg.message_id);
-        } catch {
-          // Ignore cleanup errors
-        }
-      }
-
-      // Retry on Claude Code crash (not user cancellation)
-      if (isClaudeCodeCrash && attempt < MAX_RETRIES) {
-        console.log(
-          `Claude Code crashed, retrying (attempt ${attempt + 2}/${MAX_RETRIES + 1})...`
-        );
-        await session.kill(); // Clear corrupted session
-        await ctx.reply(`⚠️ Claude crashed, retrying...`);
-        // Reset state for retry
-        state = new StreamingState();
-        statusCallback = createStatusCallback(ctx, state);
-        continue;
-      }
-
-      // Final attempt failed or non-retryable error
-      console.error("Error processing message:", error);
-
-      // Check if it was a cancellation
-      if (errorStr.includes("abort") || errorStr.includes("cancel")) {
-        // Only show "Query stopped" if it was an explicit stop, not an interrupt from a new message
-        const wasInterrupt = session.consumeInterruptFlag();
-        if (!wasInterrupt) {
-          await ctx.reply("🛑 Query stopped.");
-        }
-      } else {
-        await ctx.reply(`❌ Error: ${errorStr.slice(0, 200)}`);
-      }
-      break; // Exit loop after handling error
-    }
-  }
-
-  // 10. Cleanup
-  stopProcessing();
-  typing.stop();
+	const userId = ctx.from?.id;
+	const username = ctx.from?.username || "unknown";
+	const chatId = ctx.chat?.id;
+	let message = ctx.message?.text;
+
+	if (!userId || !message || !chatId) {
+		return;
+	}
+
+	// 1. Authorization check
+	if (!isAuthorized(userId, ALLOWED_USERS)) {
+		await ctx.reply("Unauthorized. Contact the bot owner for access.");
+		return;
+	}
+
+	// 2. Shell command shortcut: !command executes directly
+	if (message.startsWith("!")) {
+		const shellCmd = message.slice(1).trim();
+		if (shellCmd) {
+			const cwd = session.workingDir;
+			await ctx.reply(
+				`⚡ Running in <code>${cwd}</code>:\n<code>${shellCmd}</code>`,
+				{
+					parse_mode: "HTML",
+				},
+			);
+
+			const { stdout, stderr, exitCode } = await execShellCommand(
+				shellCmd,
+				cwd,
+			);
+			const output = (stdout + stderr).trim();
+			const maxLen = 4000;
+			const truncated =
+				output.length > maxLen
+					? `${output.slice(0, maxLen)}...(truncated)`
+					: output;
+
+			if (exitCode === 0) {
+				await ctx.reply(
+					`✅ Exit code: ${exitCode}\n<pre>${truncated || "(no output)"}</pre>`,
+					{
+						parse_mode: "HTML",
+					},
+				);
+			} else {
+				await ctx.reply(
+					`❌ Exit code: ${exitCode}\n<pre>${truncated || "(no output)"}</pre>`,
+					{
+						parse_mode: "HTML",
+					},
+				);
+			}
+			await auditLog(userId, username, "SHELL", shellCmd, `exit=${exitCode}`);
+			return;
+		}
+	}
+
+	// 3. Check for interrupt prefix
+	message = await checkInterrupt(message);
+	if (!message.trim()) {
+		return;
+	}
+
+	// 3. Rate limit check
+	const [allowed, retryAfter] = rateLimiter.check(userId);
+	if (!allowed) {
+		await auditLogRateLimit(userId, username, retryAfter!);
+		await ctx.reply(
+			`⏳ Rate limited. Please wait ${retryAfter?.toFixed(1)} seconds.`,
+		);
+		return;
+	}
+
+	// 4. Store message for retry
+	session.lastMessage = message;
+
+	// 5. Mark processing started
+	const stopProcessing = session.startProcessing();
+
+	// 6. Start typing indicator
+	const typing = startTypingIndicator(ctx);
+
+	// 7. Create streaming state and callback
+	let state = new StreamingState();
+	let statusCallback = createStatusCallback(ctx, state);
+
+	// 8. Send to Claude with retry logic for crashes
+	const MAX_RETRIES = 1;
+
+	for (let attempt = 0; attempt <= MAX_RETRIES; attempt++) {
+		try {
+			const response = await session.sendMessageStreaming(
+				message,
+				username,
+				userId,
+				statusCallback,
+				chatId,
+				ctx,
+			);
+
+			// 9. Audit log
+			await auditLog(userId, username, "TEXT", message, response);
+			break; // Success - exit retry loop
+		} catch (error) {
+			const errorStr = String(error);
+			const isClaudeCodeCrash = errorStr.includes("exited with code");
+
+			// Clean up any partial messages from this attempt
+			for (const toolMsg of state.toolMessages) {
+				try {
+					await ctx.api.deleteMessage(toolMsg.chat.id, toolMsg.message_id);
+				} catch {
+					// Ignore cleanup errors
+				}
+			}
+
+			// Retry on Claude Code crash (not user cancellation)
+			if (isClaudeCodeCrash && attempt < MAX_RETRIES) {
+				console.log(
+					`Claude Code crashed, retrying (attempt ${attempt + 2}/${MAX_RETRIES + 1})...`,
+				);
+				await session.kill(); // Clear corrupted session
+				await ctx.reply(`⚠️ Claude crashed, retrying...`);
+				// Reset state for retry
+				state = new StreamingState();
+				statusCallback = createStatusCallback(ctx, state);
+				continue;
+			}
+
+			// Final attempt failed or non-retryable error
+			console.error("Error processing message:", error);
+
+			// Check if it was a cancellation
+			if (errorStr.includes("abort") || errorStr.includes("cancel")) {
+				// Only show "Query stopped" if it was an explicit stop, not an interrupt from a new message
+				const wasInterrupt = session.consumeInterruptFlag();
+				if (!wasInterrupt) {
+					await ctx.reply("🛑 Query stopped.");
+				}
+			} else {
+				await ctx.reply(`❌ Error: ${errorStr.slice(0, 200)}`);
+			}
+			break; // Exit loop after handling error
+		}
+	}
+
+	// 10. Cleanup
+	stopProcessing();
+	typing.stop();
 }

package/src/index.ts

@@ -17,16 +17,24 @@ import {
 	handleBookmarks,
 	handleCallback,
 	handleCd,
+	handleCompact,
+	handleCost,
 	handleDocument,
+	handleFile,
+	handleModel,
 	handleNew,
 	handlePhoto,
+	handlePlan,
 	handleRestart,
 	handleResume,
 	handleRetry,
+	handleSkill,
 	handleStart,
 	handleStatus,
 	handleStop,
 	handleText,
+	handleThink,
+	handleUndo,
 	handleVoice,
 } from "./handlers";
 
@@ -59,11 +67,22 @@ bot.use(
 bot.command("start", handleStart);
 bot.command("new", handleNew);
 bot.command("stop", handleStop);
+bot.command("c", handleStop);
+bot.command("kill", handleStop);
+bot.command("dc", handleStop);
 bot.command("status", handleStatus);
 bot.command("resume", handleResume);
 bot.command("restart", handleRestart);
 bot.command("retry", handleRetry);
 bot.command("cd", handleCd);
+bot.command("skill", handleSkill);
+bot.command("file", handleFile);
+bot.command("model", handleModel);
+bot.command("cost", handleCost);
+bot.command("think", handleThink);
+bot.command("plan", handlePlan);
+bot.command("compact", handleCompact);
+bot.command("undo", handleUndo);
 bot.command("bookmarks", handleBookmarks);
 
 // ============== Message Handlers ==============

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();