|
| 1 | +import type { PluginInput } from "@opencode-ai/plugin" |
| 2 | +import type { ExperimentalConfig } from "../../config" |
| 3 | +import type { PreemptiveCompactionState, TokenInfo } from "./types" |
| 4 | +import { |
| 5 | + DEFAULT_THRESHOLD, |
| 6 | + MIN_TOKENS_FOR_COMPACTION, |
| 7 | + COMPACTION_COOLDOWN_MS, |
| 8 | +} from "./constants" |
| 9 | +import { log } from "../../shared/logger" |
| 10 | + |
| 11 | +export interface PreemptiveCompactionOptions { |
| 12 | + experimental?: ExperimentalConfig |
| 13 | +} |
| 14 | + |
| 15 | +interface MessageInfo { |
| 16 | + id: string |
| 17 | + role: string |
| 18 | + sessionID: string |
| 19 | + providerID?: string |
| 20 | + modelID?: string |
| 21 | + tokens?: TokenInfo |
| 22 | + summary?: boolean |
| 23 | + finish?: boolean |
| 24 | +} |
| 25 | + |
| 26 | +interface MessageWrapper { |
| 27 | + info: MessageInfo |
| 28 | +} |
| 29 | + |
| 30 | +const MODEL_CONTEXT_LIMITS: Record<string, number> = { |
| 31 | + "claude-opus-4": 200_000, |
| 32 | + "claude-sonnet-4": 200_000, |
| 33 | + "claude-haiku-4": 200_000, |
| 34 | + "gpt-4o": 128_000, |
| 35 | + "gpt-4o-mini": 128_000, |
| 36 | + "gpt-4-turbo": 128_000, |
| 37 | + "gpt-4": 8_192, |
| 38 | + "gpt-5": 1_000_000, |
| 39 | + "o1": 200_000, |
| 40 | + "o1-mini": 128_000, |
| 41 | + "o1-preview": 128_000, |
| 42 | + "o3": 200_000, |
| 43 | + "o3-mini": 200_000, |
| 44 | + "gemini-2.0-flash": 1_000_000, |
| 45 | + "gemini-2.5-flash": 1_000_000, |
| 46 | + "gemini-2.5-pro": 2_000_000, |
| 47 | + "gemini-3-pro": 2_000_000, |
| 48 | +} |
| 49 | + |
| 50 | +function getContextLimit(modelID: string): number { |
| 51 | + for (const [key, limit] of Object.entries(MODEL_CONTEXT_LIMITS)) { |
| 52 | + if (modelID.includes(key)) { |
| 53 | + return limit |
| 54 | + } |
| 55 | + } |
| 56 | + return 200_000 |
| 57 | +} |
| 58 | + |
| 59 | +function createState(): PreemptiveCompactionState { |
| 60 | + return { |
| 61 | + lastCompactionTime: new Map(), |
| 62 | + compactionInProgress: new Set(), |
| 63 | + } |
| 64 | +} |
| 65 | + |
| 66 | +export function createPreemptiveCompactionHook( |
| 67 | + ctx: PluginInput, |
| 68 | + options?: PreemptiveCompactionOptions |
| 69 | +) { |
| 70 | + const experimental = options?.experimental |
| 71 | + const enabled = experimental?.preemptive_compaction !== false |
| 72 | + const threshold = experimental?.preemptive_compaction_threshold ?? DEFAULT_THRESHOLD |
| 73 | + |
| 74 | + if (!enabled) { |
| 75 | + return { event: async () => {} } |
| 76 | + } |
| 77 | + |
| 78 | + const state = createState() |
| 79 | + |
| 80 | + const checkAndTriggerCompaction = async ( |
| 81 | + sessionID: string, |
| 82 | + lastAssistant: MessageInfo |
| 83 | + ): Promise<void> => { |
| 84 | + if (state.compactionInProgress.has(sessionID)) return |
| 85 | + |
| 86 | + const lastCompaction = state.lastCompactionTime.get(sessionID) ?? 0 |
| 87 | + if (Date.now() - lastCompaction < COMPACTION_COOLDOWN_MS) return |
| 88 | + |
| 89 | + if (lastAssistant.summary === true) return |
| 90 | + |
| 91 | + const tokens = lastAssistant.tokens |
| 92 | + if (!tokens) return |
| 93 | + |
| 94 | + const modelID = lastAssistant.modelID ?? "" |
| 95 | + const contextLimit = getContextLimit(modelID) |
| 96 | + const totalUsed = tokens.input + tokens.cache.read + tokens.output |
| 97 | + |
| 98 | + if (totalUsed < MIN_TOKENS_FOR_COMPACTION) return |
| 99 | + |
| 100 | + const usageRatio = totalUsed / contextLimit |
| 101 | + |
| 102 | + log("[preemptive-compaction] checking", { |
| 103 | + sessionID, |
| 104 | + totalUsed, |
| 105 | + contextLimit, |
| 106 | + usageRatio: usageRatio.toFixed(2), |
| 107 | + threshold, |
| 108 | + }) |
| 109 | + |
| 110 | + if (usageRatio < threshold) return |
| 111 | + |
| 112 | + state.compactionInProgress.add(sessionID) |
| 113 | + state.lastCompactionTime.set(sessionID, Date.now()) |
| 114 | + |
| 115 | + const providerID = lastAssistant.providerID |
| 116 | + if (!providerID || !modelID) { |
| 117 | + state.compactionInProgress.delete(sessionID) |
| 118 | + return |
| 119 | + } |
| 120 | + |
| 121 | + await ctx.client.tui |
| 122 | + .showToast({ |
| 123 | + body: { |
| 124 | + title: "Preemptive Compaction", |
| 125 | + message: `Context at ${(usageRatio * 100).toFixed(0)}% - compacting to prevent overflow...`, |
| 126 | + variant: "warning", |
| 127 | + duration: 3000, |
| 128 | + }, |
| 129 | + }) |
| 130 | + .catch(() => {}) |
| 131 | + |
| 132 | + log("[preemptive-compaction] triggering compaction", { sessionID, usageRatio }) |
| 133 | + |
| 134 | + try { |
| 135 | + await ctx.client.session.summarize({ |
| 136 | + path: { id: sessionID }, |
| 137 | + body: { providerID, modelID }, |
| 138 | + query: { directory: ctx.directory }, |
| 139 | + }) |
| 140 | + |
| 141 | + await ctx.client.tui |
| 142 | + .showToast({ |
| 143 | + body: { |
| 144 | + title: "Compaction Complete", |
| 145 | + message: "Session compacted successfully", |
| 146 | + variant: "success", |
| 147 | + duration: 2000, |
| 148 | + }, |
| 149 | + }) |
| 150 | + .catch(() => {}) |
| 151 | + } catch (err) { |
| 152 | + log("[preemptive-compaction] compaction failed", { sessionID, error: err }) |
| 153 | + } finally { |
| 154 | + state.compactionInProgress.delete(sessionID) |
| 155 | + } |
| 156 | + } |
| 157 | + |
| 158 | + const eventHandler = async ({ event }: { event: { type: string; properties?: unknown } }) => { |
| 159 | + const props = event.properties as Record<string, unknown> | undefined |
| 160 | + |
| 161 | + if (event.type === "session.deleted") { |
| 162 | + const sessionInfo = props?.info as { id?: string } | undefined |
| 163 | + if (sessionInfo?.id) { |
| 164 | + state.lastCompactionTime.delete(sessionInfo.id) |
| 165 | + state.compactionInProgress.delete(sessionInfo.id) |
| 166 | + } |
| 167 | + return |
| 168 | + } |
| 169 | + |
| 170 | + if (event.type === "message.updated") { |
| 171 | + const info = props?.info as MessageInfo | undefined |
| 172 | + if (!info) return |
| 173 | + |
| 174 | + if (info.role !== "assistant" || !info.finish) return |
| 175 | + |
| 176 | + const sessionID = info.sessionID |
| 177 | + if (!sessionID) return |
| 178 | + |
| 179 | + await checkAndTriggerCompaction(sessionID, info) |
| 180 | + return |
| 181 | + } |
| 182 | + |
| 183 | + if (event.type === "session.idle") { |
| 184 | + const sessionID = props?.sessionID as string | undefined |
| 185 | + if (!sessionID) return |
| 186 | + |
| 187 | + try { |
| 188 | + const resp = await ctx.client.session.messages({ |
| 189 | + path: { id: sessionID }, |
| 190 | + query: { directory: ctx.directory }, |
| 191 | + }) |
| 192 | + |
| 193 | + const messages = (resp.data ?? resp) as MessageWrapper[] |
| 194 | + const assistants = messages |
| 195 | + .filter((m) => m.info.role === "assistant") |
| 196 | + .map((m) => m.info) |
| 197 | + |
| 198 | + if (assistants.length === 0) return |
| 199 | + |
| 200 | + const lastAssistant = assistants[assistants.length - 1] |
| 201 | + await checkAndTriggerCompaction(sessionID, lastAssistant) |
| 202 | + } catch {} |
| 203 | + } |
| 204 | + } |
| 205 | + |
| 206 | + return { |
| 207 | + event: eventHandler, |
| 208 | + } |
| 209 | +} |
0 commit comments