diff --git a/src/hooks/grep-output-truncator.ts b/src/hooks/grep-output-truncator.ts new file mode 100644 index 0000000..7b5405a --- /dev/null +++ b/src/hooks/grep-output-truncator.ts @@ -0,0 +1,131 @@ +import type { PluginInput } from "@opencode-ai/plugin" + +const ANTHROPIC_ACTUAL_LIMIT = 200_000 +const CHARS_PER_TOKEN_ESTIMATE = 4 +const TARGET_MAX_TOKENS = 50_000 + +interface AssistantMessageInfo { + role: "assistant" + tokens: { + input: number + output: number + reasoning: number + cache: { read: number; write: number } + } +} + +interface MessageWrapper { + info: { role: string } & Partial +} + +function estimateTokens(text: string): number { + return Math.ceil(text.length / CHARS_PER_TOKEN_ESTIMATE) +} + +function truncateToTokenLimit(output: string, maxTokens: number): { result: string; truncated: boolean } { + const currentTokens = estimateTokens(output) + + if (currentTokens <= maxTokens) { + return { result: output, truncated: false } + } + + const lines = output.split("\n") + + if (lines.length <= 3) { + const maxChars = maxTokens * CHARS_PER_TOKEN_ESTIMATE + return { + result: output.slice(0, maxChars) + "\n\n[Output truncated due to context window limit]", + truncated: true, + } + } + + const headerLines = lines.slice(0, 3) + const contentLines = lines.slice(3) + + const headerText = headerLines.join("\n") + const headerTokens = estimateTokens(headerText) + const availableTokens = maxTokens - headerTokens - 50 + + if (availableTokens <= 0) { + return { + result: headerText + "\n\n[Content truncated due to context window limit]", + truncated: true, + } + } + + let resultLines: string[] = [] + let currentTokenCount = 0 + + for (const line of contentLines) { + const lineTokens = estimateTokens(line + "\n") + if (currentTokenCount + lineTokens > availableTokens) { + break + } + resultLines.push(line) + currentTokenCount += lineTokens + } + + const truncatedContent = [...headerLines, ...resultLines].join("\n") + const removedCount = contentLines.length - resultLines.length + + return { + result: truncatedContent + `\n\n[${removedCount} more lines truncated due to context window limit]`, + truncated: true, + } +} + +export function createGrepOutputTruncatorHook(ctx: PluginInput) { + const GREP_TOOLS = ["safe_grep", "Grep"] + + const toolExecuteAfter = async ( + input: { tool: string; sessionID: string; callID: string }, + output: { title: string; output: string; metadata: unknown } + ) => { + if (!GREP_TOOLS.includes(input.tool)) return + + const { sessionID } = input + + try { + const response = await ctx.client.session.messages({ + path: { id: sessionID }, + }) + + const messages = (response.data ?? response) as MessageWrapper[] + + const assistantMessages = messages + .filter((m) => m.info.role === "assistant") + .map((m) => m.info as AssistantMessageInfo) + + if (assistantMessages.length === 0) return + + const totalInputTokens = assistantMessages.reduce((sum, m) => { + const inputTokens = m.tokens?.input ?? 0 + const cacheReadTokens = m.tokens?.cache?.read ?? 0 + return sum + inputTokens + cacheReadTokens + }, 0) + + const remainingTokens = ANTHROPIC_ACTUAL_LIMIT - totalInputTokens + + const maxOutputTokens = Math.min( + remainingTokens * 0.5, + TARGET_MAX_TOKENS + ) + + if (maxOutputTokens <= 0) { + output.output = "[Output suppressed - context window exhausted]" + return + } + + const { result, truncated } = truncateToTokenLimit(output.output, maxOutputTokens) + if (truncated) { + output.output = result + } + } catch { + // Graceful degradation + } + } + + return { + "tool.execute.after": toolExecuteAfter, + } +} diff --git a/src/hooks/index.ts b/src/hooks/index.ts index b0b2958..7757364 100644 --- a/src/hooks/index.ts +++ b/src/hooks/index.ts @@ -3,3 +3,4 @@ export { createContextWindowMonitorHook } from "./context-window-monitor" export { createSessionNotification } from "./session-notification" export { createSessionRecoveryHook } from "./session-recovery" export { createCommentCheckerHooks } from "./comment-checker" +export { createGrepOutputTruncatorHook } from "./grep-output-truncator" diff --git a/src/index.ts b/src/index.ts index fc97408..2c0ac73 100644 --- a/src/index.ts +++ b/src/index.ts @@ -1,6 +1,6 @@ import type { Plugin } from "@opencode-ai/plugin" import { createBuiltinAgents } from "./agents" -import { createTodoContinuationEnforcer, createContextWindowMonitorHook, createSessionRecoveryHook, createCommentCheckerHooks } from "./hooks" +import { createTodoContinuationEnforcer, createContextWindowMonitorHook, createSessionRecoveryHook, createCommentCheckerHooks, createGrepOutputTruncatorHook } from "./hooks" import { updateTerminalTitle } from "./features/terminal" import { builtinTools } from "./tools" import { createBuiltinMcps } from "./mcp" @@ -44,6 +44,7 @@ const OhMyOpenCodePlugin: Plugin = async (ctx) => { const contextWindowMonitor = createContextWindowMonitorHook(ctx) const sessionRecovery = createSessionRecoveryHook(ctx) const commentChecker = createCommentCheckerHooks() + const grepOutputTruncator = createGrepOutputTruncatorHook(ctx) updateTerminalTitle({ sessionId: "main" }) @@ -185,6 +186,7 @@ const OhMyOpenCodePlugin: Plugin = async (ctx) => { }, "tool.execute.after": async (input, output) => { + await grepOutputTruncator["tool.execute.after"](input, output) await contextWindowMonitor["tool.execute.after"](input, output) await commentChecker["tool.execute.after"](input, output)