diff --git a/index.ts b/index.ts index 4bef4ad..323a0ab 100644 --- a/index.ts +++ b/index.ts @@ -101,7 +101,7 @@ import { getTelegramBotTokenInputDefault, getTelegramBotTokenPromptSpec, } from "./lib/setup.ts"; -import { buildStatusHtml } from "./lib/status.ts"; +import { buildStatusHtml, extractTurnCost, formatTurnCostLine } from "./lib/status.ts"; import { buildTelegramPromptTurn, truncateTelegramQueueSummary, @@ -1563,6 +1563,33 @@ export default function (pi: ExtensionAPI) { await sendTextReply(message.chat.id, message.message_id, "No active turn."); } + async function handleQuitCommand( + message: TelegramMessage, + ctx: ExtensionContext, + ): Promise { + await sendTextReply(message.chat.id, message.message_id, "Shutting down pi session."); + ctx.shutdown(); + } + + async function handleShellCommand( + shellCmd: string, + message: TelegramMessage, + _ctx: ExtensionContext, + ): Promise { + try { + const result = await pi.exec("sh", ["-c", shellCmd], { timeout: 30_000 }); + const output = (result.stdout + result.stderr).trim(); + const codeTag = result.code !== 0 ? ` (exit ${result.code})` : ""; + const reply = output + ? `${output.slice(0, 3900)}${codeTag}` + : `(no output)${codeTag}`; + await sendTextReply(message.chat.id, message.message_id, reply); + } catch (err) { + const msg = err instanceof Error ? err.message : String(err); + await sendTextReply(message.chat.id, message.message_id, `Shell error: ${msg}`); + } + } + async function handleCompactCommand( message: TelegramMessage, ctx: ExtensionContext, @@ -1687,7 +1714,8 @@ export default function (pi: ExtensionAPI) { ctx: ExtensionContext, ): Promise { let helpText = - "Send me a message and I will forward it to pi. Commands: /status, /trace, /model, /compact, /stop."; + "Send me a message and I will forward it to pi.\n\nLocal: /status, /trace, /model, /compact, /stop, /quit\nOther /commands and ! shell commands pass through to pi directly."; + if (commandName === "start") { try { await registerTelegramBotCommands(); @@ -1719,6 +1747,8 @@ export default function (pi: ExtensionAPI) { model: () => handleModelCommand(message, ctx), help: () => handleHelpCommand(message, commandName, ctx), start: () => handleHelpCommand(message, commandName, ctx), + quit: () => handleQuitCommand(message, ctx), + exit: () => handleQuitCommand(message, ctx), }; const handler = handlers[commandName]; if (!handler) return false; @@ -1748,9 +1778,21 @@ export default function (pi: ExtensionAPI) { const firstMessage = messages[0]; if (!firstMessage) return; const rawText = extractFirstTelegramMessageText(messages); + + // Handle ! shell commands directly via ctx.exec + const trimmedRaw = rawText.trimStart(); + if (trimmedRaw.startsWith("!")) { + const shellCmd = trimmedRaw.slice(1).trim(); + if (shellCmd) { + await handleShellCommand(shellCmd, firstMessage, ctx); + return; + } + } + const commandName = parseTelegramCommand(rawText)?.name; const handled = await handleTelegramCommand(commandName, firstMessage, ctx); if (handled) return; + await enqueueTelegramTurn(messages, ctx); } @@ -2093,9 +2135,17 @@ export default function (pi: ExtensionAPI) { const assistant = turn ? extractAssistantSummary((event as { messages: AgentMessage[] }).messages) : { blocks: [] }; - const finalText = traceVisible + let finalText = traceVisible ? buildTelegramAssistantTranscriptMarkdown(assistant.blocks, true) : assistant.text; + // Append per-turn cost/context footer when trace is on + if (traceVisible && turn && finalText) { + const turnCost = extractTurnCost((event as { messages: AgentMessage[] }).messages as any); + const usage = ctx.getContextUsage(); + if (turnCost) { + finalText += `\n\n---\n${formatTurnCostLine(turnCost, usage?.percent ?? null)}`; + } + } const endPlan = buildTelegramAgentEndPlan({ hasTurn: !!turn, stopReason: assistant.stopReason, @@ -2104,6 +2154,18 @@ export default function (pi: ExtensionAPI) { preserveQueuedTurnsAsHistory, }); if (!turn) { + // Notify about non-telegram turns when trace is on (scheduled prompts, system events, etc.) + if (traceVisible && config.allowedUserId) { + const nonTelegramAssistant = extractAssistantSummary((event as { messages: AgentMessage[] }).messages); + const summary = nonTelegramAssistant.text?.slice(0, 500); + const turnCost = extractTurnCost((event as { messages: AgentMessage[] }).messages as any); + const usage = ctx.getContextUsage(); + const costLine = turnCost ? formatTurnCostLine(turnCost, usage?.percent ?? null) : undefined; + const parts = ["[non-telegram turn]"]; + if (summary) parts.push(summary); + if (costLine) parts.push(`---\n${costLine}`); + void sendTextReply(config.allowedUserId, 0, parts.join("\n")); + } if (endPlan.shouldDispatchNext) { dispatchNextQueuedTelegramTurn(ctx); } diff --git a/lib/status.ts b/lib/status.ts index d9e4be7..3b9badc 100644 --- a/lib/status.ts +++ b/lib/status.ts @@ -84,6 +84,45 @@ function buildContextSummary( return `${percent}/${formatTokens(contextWindow)}`; } +export interface TurnCostInfo { + input: number; + output: number; + cacheRead: number; + cost: number; +} + +/** + * Extract per-turn cost/token info from the agent_end messages array. + * Returns undefined if no assistant messages with usage found. + */ +export function extractTurnCost(messages: Array<{ role?: string; usage?: { input: number; output: number; cacheRead: number; cost: { total: number } } }>): TurnCostInfo | undefined { + let input = 0, output = 0, cacheRead = 0, cost = 0; + let found = false; + for (const msg of messages) { + if (msg.role === "assistant" && msg.usage) { + input += msg.usage.input; + output += msg.usage.output; + cacheRead += msg.usage.cacheRead; + cost += msg.usage.cost.total; + found = true; + } + } + return found ? { input, output, cacheRead, cost } : undefined; +} + +/** + * Format turn cost as a one-line summary for appending to trace replies. + */ +export function formatTurnCostLine(turnCost: TurnCostInfo, contextPercent: number | null): string { + const parts: string[] = []; + parts.push(`$${turnCost.cost.toFixed(3)}`); + const tokens = [`↑${formatTokens(turnCost.input)}`, `↓${formatTokens(turnCost.output)}`]; + if (turnCost.cacheRead) tokens.push(`R${formatTokens(turnCost.cacheRead)}`); + parts.push(tokens.join(" ")); + if (contextPercent !== null) parts.push(`ctx ${contextPercent.toFixed(0)}%`); + return parts.join(" | "); +} + export function buildStatusHtml( ctx: ExtensionContext, activeModel: Model | undefined, diff --git a/lib/turns.ts b/lib/turns.ts index 6fa101a..d901bed 100644 --- a/lib/turns.ts +++ b/lib/turns.ts @@ -65,8 +65,9 @@ export function buildTelegramTurnPrompt(options: { files: DownloadedTelegramTurnFileLike[]; historyTurns?: Pick[]; }): string { - // Let pi handle `!` shell commands natively - don't prepend [telegram] prefix - let prompt = options.rawText.trimStart().startsWith("!") ? "" : options.telegramPrefix; + // Let pi handle `!` shell commands and `/` slash commands natively - don't prepend [telegram] prefix + const raw = options.rawText.trimStart(); + let prompt = (raw.startsWith("!") || raw.startsWith("/")) ? "" : options.telegramPrefix; if ((options.historyTurns?.length ?? 0) > 0) { prompt += "\n\nEarlier Telegram messages arrived after an aborted turn. Treat them as prior user messages, in order:"; diff --git a/tests/queue.test.ts b/tests/queue.test.ts index 9bab282..48cb8be 100644 --- a/tests/queue.test.ts +++ b/tests/queue.test.ts @@ -1255,6 +1255,7 @@ test("Extension runtime finalizes a drafted preview into the final Telegram repl isIdle: () => true, hasPendingMessages: () => false, abort: () => {}, + getContextUsage: () => ({ tokens: 10000, contextWindow: 200000, percent: 5.0 }), } as never; await handlers.get("session_start")?.({}, ctx); await commands.get("telegram-connect")?.handler("", ctx);