From f797c6395cdedc16b773b79bfc4a54c59e3d2ce3 Mon Sep 17 00:00:00 2001 From: Rolando Santamaria Maso Date: Sat, 25 Jul 2026 23:34:02 +0200 Subject: [PATCH 1/2] feat(loop): graduated context trimming, margin calibration, rolling compaction MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Reworks context-window management in the agent loop: Correctness fixes: - estimateToolDefs now counts JSON parameter schemas (the bulk of every tool definition) instead of ignoring them - estimateMessages now counts ReasoningContent (echoed back to reasoning models, often the largest content in a turn) - trimToSurvival keeps the original task in multi-turn sessions, preserves the compaction digest, and no longer appends a zero-value message when no user message exists Trimming robustness: - budget re-checked after memory/skill/episode injections, right before the LLM call — injected blocks no longer escape the budget for a full iteration - graduated trimming: old tool results >2 KiB (except the 4 most recent) are replaced with a marker before any turn group is dropped - trim warning moved before the most recent user message (cache-stable head preserved), updated in place with cumulative totals, and names up to 5 dropped tools; new headLen protects system prompt + memory block + digest + original task - disk-cap session trimming persists a [Session storage limit: ...] marker into the transcript so resumed sessions know history was removed New features: - self-calibrating safety margin: tightens 0.75 -> 0.65 when provider- reported input tokens exceed the local estimate by >15% (margin_calibrated signal) - optional rolling compaction (compaction / ODEK_COMPACTION / --compaction, default off): dropped turn groups are summarized by the model into a rolling, untrusted-wrapped digest message instead of being lost; summarizer failures never break trimming All new/changed loop, session, and config code is at 100% function coverage. Full unit suite, race detector, golangci-lint, and CLI E2E tests pass. --- AGENTS.md | 4 +- cmd/odek/main.go | 16 + cmd/odek/repl.go | 2 + cmd/odek/schedule.go | 1 + cmd/odek/serve.go | 4 + docs/CLI.md | 1 + docs/CONFIG.md | 9 + internal/config/loader.go | 20 + internal/config/loader_test.go | 50 ++ internal/loop/loop.go | 506 +++++++++++++--- internal/loop/loop_test.go | 16 +- internal/loop/loop_trim_test.go | 721 +++++++++++++++++++++++ internal/loop/signal.go | 6 +- internal/session/session.go | 43 +- internal/session/session_savecap_test.go | 56 +- internal/session/session_test.go | 17 +- odek.go | 8 + 17 files changed, 1373 insertions(+), 107 deletions(-) create mode 100644 internal/loop/loop_trim_test.go diff --git a/AGENTS.md b/AGENTS.md index 60438df8..59436b39 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -87,11 +87,11 @@ ReAct cycle: observe → think → act → repeat. - **Parallel tool execution** — multiple independent tool calls run concurrently (`max_tool_parallel`, default: 4). - **Batch approval gate** (`internal/loop/loop.go`) — multiple risky tools shown at once in a single prompt. `classifyToolCall` now classifies every command inside `parallel_shell`, every path inside `batch_patch`, and the modern `browser` tool, shows full (untruncated) commands, and withholds blanket `SetTrustAll` when unclassifiable tools remain in the iteration. - **Tool-failure recovery** — systematic recovery from tool call failures: retry transient errors, skip permanently failed tools, and continue the loop without crashing. -- **Context-limit protection** — `trimToSurvival` drops oldest messages when approaching the model's context window, keeping the agent functional under extended sessions. Tool messages stay grouped with their parent assistant message. +- **Context-limit protection** — graduated trimming when approaching the model's context window: old, large tool results are first replaced with a marker (the 4 most recent are always kept intact), then the oldest turn groups are dropped atomically (tool messages stay grouped with their parent assistant message). The protected head (system prompt, memory block, compaction digest, original task) is never dropped. A trim warning is injected before the most recent user message and updated in place with cumulative totals (including which tools lost earlier results). The token estimator counts tool-def parameter schemas and reasoning content, and the safety margin self-tightens (0.75 → 0.65) when provider-reported input tokens exceed the estimate by >15% (`margin_calibrated` signal). `trimToSurvival` (on provider context-length errors) keeps the system prompt, digest, original task, last 2 turn groups, and latest user message. Optional **rolling compaction** (`compaction` config / `--compaction` / `ODEK_COMPACTION`) summarizes dropped groups into a rolling digest system message (untrusted-wrapped) instead of losing them — one extra LLM call per trim. - **Interaction modes** — engaging (narrated), enhance (persistent), verbose (raw), off. - Max 300 iterations by default. - **Post-response async processing** — skill learning, episode extraction, and per-turn extended-memory extraction run in background goroutines tracked by `MemoryManager.RunBackground`; `Agent.Close` drains them via `WaitForBackground` (bounded, ~15s) so the work survives CLI exit without hanging `odek run`. -- **Storage maintenance janitor** — `maintenance.Start` (internal/maintenance) runs a periodic sweep of `~/.odek` (expired sessions/audit/plans, oversized-log rotation, media sweep, skill skip-list GC) inside `odek telegram`, `odek serve`, and `odek schedule daemon` when `maintenance.enabled` is set; `odek cleanup [--dry-run]` runs the same sweep on demand. Session files are also trimmed at write time (oldest groups first, system message kept) when they would exceed `MaxSessionFileBytes`, so a session can never grow past the load cap. Operator-only config — project `./odek.json` cannot set it. See docs/MAINTENANCE.md. +- **Storage maintenance janitor** — `maintenance.Start` (internal/maintenance) runs a periodic sweep of `~/.odek` (expired sessions/audit/plans, oversized-log rotation, media sweep, skill skip-list GC) inside `odek telegram`, `odek serve`, and `odek schedule daemon` when `maintenance.enabled` is set; `odek cleanup [--dry-run]` runs the same sweep on demand. Session files are also trimmed at write time (oldest groups first, system message kept, `[Session storage limit: ...]` marker persisted into the transcript) when they would exceed `MaxSessionFileBytes`, so a session can never grow past the load cap. Operator-only config — project `./odek.json` cannot set it. See docs/MAINTENANCE.md. - **Artifact-aware file search** — `search_files` and `multi_grep` skip build/artifact directories (`node_modules`, `vendor`, `.git`, `__pycache__`, `.venv`, etc.) automatically, reducing noise and speeding scans. - **Semantic session search** — the `session_search` tool uses go-vector RandomProjections + k-NN for semantic similarity search through session content, with a two-tier pipeline: vector index (fast, ~1ms) → deepSearch fallback (exhaustive, slower). - **Security-first defaults** — the latest hardening closes the high/medium/low findings tracked in `sec_findings.md`: default `non_interactive` is `deny`, project-level `odek.json` cannot redirect backends or hijack delivery, project-level sandbox knobs require explicit operator approval, `~/.odek` trust anchors are protected, WebSocket upgrades require a per-instance CSRF token, and all untrusted content is wrapped before reaching the model. See Security Architecture below for the full list. diff --git a/cmd/odek/main.go b/cmd/odek/main.go index 5abcf59a..b9339ac8 100644 --- a/cmd/odek/main.go +++ b/cmd/odek/main.go @@ -274,6 +274,7 @@ type runFlags struct { NoColor *bool // nil = not set NoAgents *bool // nil = not set PromptCaching *bool // nil = not set; true = enable prompt caching + Compaction *bool // nil = not set; true = enable rolling compaction Session *bool // nil = not set; true = save session after run Learn *bool // nil = not set; true = enable skills learning mode Task string @@ -412,6 +413,9 @@ func parseRunFlags(args []string) (runFlags, error) { case "--prompt-caching": f.PromptCaching = boolPtr(true) i++ + case "--compaction": + f.Compaction = boolPtr(true) + i++ case "--session": f.Session = boolPtr(true) i++ @@ -655,6 +659,10 @@ done: f.PromptCaching = boolPtr(true) taskArgs = append(taskArgs[:j], taskArgs[j+1:]...) j-- + case "--compaction": + f.Compaction = boolPtr(true) + taskArgs = append(taskArgs[:j], taskArgs[j+1:]...) + j-- case "--sandbox-readonly": f.SandboxReadonly = boolPtr(true) taskArgs = append(taskArgs[:j], taskArgs[j+1:]...) @@ -740,6 +748,7 @@ type replFlags struct { ThinkingBudget int // 0 = not set; use default Sandbox *bool // nil = not set PromptCaching *bool // nil = not set; true = enable prompt caching + Compaction *bool // nil = not set; true = enable rolling compaction InteractionMode string // Sandbox-specific CLI flags @@ -809,6 +818,9 @@ func parseReplFlags(args []string) (replFlags, error) { case "--prompt-caching": f.PromptCaching = boolPtr(true) i++ + case "--compaction": + f.Compaction = boolPtr(true) + i++ case "--interaction-mode": f.InteractionMode = args[i+1] i += 2 @@ -888,6 +900,7 @@ Run flags: --no-color Disable colored terminal output --no-agents Skip loading AGENTS.md from working directory --prompt-caching Enable prompt caching markers (Anthropic/DeepSeek/OpenAI) + --compaction Enable LLM-based rolling compaction of trimmed context --session Save conversation as a multi-turn session --learn Enable skill learning mode — on by default, no flag needed --no-learn Disable skill learning mode (overrides config/default) @@ -1154,6 +1167,7 @@ func run(args []string) error { NoColor: f.NoColor, NoAgents: f.NoAgents, PromptCaching: f.PromptCaching, + Compaction: f.Compaction, Learn: f.Learn, System: f.System, Task: f.Task, @@ -1326,6 +1340,7 @@ func run(args []string) error { Skills: skillsCfg, SkillManager: sm, PromptCaching: resolved.PromptCaching, + Compaction: resolved.Compaction, MemoryDir: expandHome("~/.odek/memory"), MemoryConfig: resolved.Memory, Guard: injectionGuard, @@ -2380,6 +2395,7 @@ func continueCmd(args []string) error { Skills: skillsCfg, SkillManager: sm, PromptCaching: resolved.PromptCaching, + Compaction: resolved.Compaction, MemoryDir: expandHome("~/.odek/memory"), MemoryConfig: resolved.Memory, Guard: injectionGuard, diff --git a/cmd/odek/repl.go b/cmd/odek/repl.go index cd839928..2688efc3 100644 --- a/cmd/odek/repl.go +++ b/cmd/odek/repl.go @@ -53,6 +53,7 @@ func replCmd(args []string) error { Thinking: f.Thinking, Sandbox: f.Sandbox, PromptCaching: f.PromptCaching, + Compaction: f.Compaction, InteractionMode: f.InteractionMode, SandboxImage: f.SandboxImage, @@ -165,6 +166,7 @@ func replCmd(args []string) error { MemoryConfig: resolved.Memory, MemoryDir: expandHome("~/.odek/memory"), PromptCaching: resolved.PromptCaching, + Compaction: resolved.Compaction, Guard: injectionGuard, GuardConfig: resolved.Guard, }) diff --git a/cmd/odek/schedule.go b/cmd/odek/schedule.go index fc80fbc9..1d563922 100644 --- a/cmd/odek/schedule.go +++ b/cmd/odek/schedule.go @@ -721,6 +721,7 @@ func runTaskHeadless(ctx context.Context, resolved config.ResolvedConfig, system Renderer: render.New(io.Discard, false), // silent: unattended InteractionMode: "off", PromptCaching: resolved.PromptCaching, + Compaction: resolved.Compaction, IterationCallback: func(info loop.IterationInfo) { lastInfo = info }, Guard: injectionGuard, GuardConfig: resolved.Guard, diff --git a/cmd/odek/serve.go b/cmd/odek/serve.go index f21eb4ad..3b2237d7 100644 --- a/cmd/odek/serve.go +++ b/cmd/odek/serve.go @@ -185,6 +185,7 @@ func serveCmd(args []string) error { var sandbox *bool var sandboxReadonly *bool var promptCaching *bool + var compaction *bool var sandboxImage, sandboxNetwork, sandboxMemory, sandboxCPUs, sandboxUser string var toolsEnabled, toolsDisabled, trustedProxies []string @@ -234,6 +235,8 @@ func serveCmd(args []string) error { } case "--prompt-caching": promptCaching = boolPtr(true) + case "--compaction": + compaction = boolPtr(true) case "--tool": i++ if i >= len(args) { @@ -263,6 +266,7 @@ func serveCmd(args []string) error { resolved := config.LoadConfig(config.CLIFlags{ Sandbox: sandbox, PromptCaching: promptCaching, + Compaction: compaction, SandboxImage: sandboxImage, SandboxNetwork: sandboxNetwork, SandboxReadonly: sandboxReadonly, diff --git a/docs/CLI.md b/docs/CLI.md index 1a255950..c24d4ade 100644 --- a/docs/CLI.md +++ b/docs/CLI.md @@ -48,6 +48,7 @@ | `--interaction-mode ` | string | `engaging` | Tool-call rendering: `engaging` (emoji narration) or `verbose` (raw tool output) | | `--no-color` | bool | false | Disable colored terminal output | | `--prompt-caching` | bool | false | Enable Anthropic/OpenAI/DeepSeek prompt caching markers | +| `--compaction` | bool | false | Enable LLM-based rolling compaction of trimmed context | | `--no-agents` | bool | false | Skip loading AGENTS.md | | `--session` | bool | false | Save conversation as a multi-turn session | | `--learn` | bool | `true` | Enable skill learning mode (detects patterns, saves skills). On by default | diff --git a/docs/CONFIG.md b/docs/CONFIG.md index d15e84dd..35a0061b 100644 --- a/docs/CONFIG.md +++ b/docs/CONFIG.md @@ -118,6 +118,7 @@ Every config knob has a `ODEK_*` counterpart: | `ODEK_SYSTEM` | `--system` | string | | `ODEK_SKILLS_LEARN` | `skills.learn` | bool | | `ODEK_PROMPT_CACHING` | `prompt_caching` | bool | +| `ODEK_COMPACTION` | `compaction` | bool | | `ODEK_TOOL_PROGRESS` | `tool_progress` | string (all\|new\|verbose\|off) | | `ODEK_SANDBOX_IMAGE` | `--sandbox-image` | string | | `ODEK_SANDBOX_NETWORK` | `--sandbox-network` | string | @@ -243,6 +244,14 @@ I/O-bound tools (read_file, search_files, shell) benefit most — latency drops **Approval gate:** When an approver is configured and the LLM returns multiple tool calls, a single batch approval prompt is shown before any tool executes. If approved, all tools run in parallel. If denied, no tools run. +## Rolling compaction (`compaction`) + +When context trimming drops old conversation turns to stay within the model's context window, those turns are normally lost. With `compaction` enabled, the dropped turns are instead summarized by the model into a rolling digest message, preserving a compressed history of the session. + +| Field | Default | Env var | CLI flag | Description | +|-------|---------|---------|----------|-------------| +| `compaction` | `false` | `ODEK_COMPACTION` | `--compaction` | Enable LLM-based rolling compaction of trimmed context. Each compaction costs one extra LLM call per trim. | + ## Concurrency and reverse-proxy trust | Field | Default | Env var | Description | diff --git a/internal/config/loader.go b/internal/config/loader.go index bcaa35fe..4df76970 100644 --- a/internal/config/loader.go +++ b/internal/config/loader.go @@ -72,6 +72,10 @@ type CLIFlags struct { // Config: prompt_caching, ODEK_PROMPT_CACHING, --prompt-caching. PromptCaching *bool // nil = not set + // Compaction enables LLM-based rolling compaction of trimmed context. + // Config: compaction, ODEK_COMPACTION, --compaction. + Compaction *bool // nil = not set + // Sandbox-specific SandboxImage string SandboxNetwork string @@ -240,6 +244,9 @@ type FileConfig struct { // PromptCaching enables prompt caching markers for supported providers. PromptCaching *bool `json:"prompt_caching,omitempty"` + // Compaction enables LLM-based rolling compaction of trimmed context. + Compaction *bool `json:"compaction,omitempty"` + System string `json:"system,omitempty"` // Sandbox-specific fields. @@ -376,6 +383,7 @@ type ResolvedConfig struct { NoColor bool NoAgents bool PromptCaching bool + Compaction bool System string // SandboxImage is the Docker image for the sandbox container. @@ -1078,6 +1086,9 @@ func LoadConfig(cli CLIFlags) ResolvedConfig { if v := envBool("PROMPT_CACHING"); v != nil { cfg.PromptCaching = v } + if v := envBool("COMPACTION"); v != nil { + cfg.Compaction = v + } if v := envString("SYSTEM"); v != "" { cfg.System = v } @@ -1431,6 +1442,9 @@ func LoadConfig(cli CLIFlags) ResolvedConfig { if cli.PromptCaching != nil { cfg.PromptCaching = cli.PromptCaching } + if cli.Compaction != nil { + cfg.Compaction = cli.Compaction + } if cli.Learn != nil { if cfg.Skills == nil { cfg.Skills = &SkillsConfig{} @@ -1702,6 +1716,9 @@ func LoadConfig(cli CLIFlags) ResolvedConfig { if cfg.PromptCaching != nil { resolved.PromptCaching = *cfg.PromptCaching } + if cfg.Compaction != nil { + resolved.Compaction = *cfg.Compaction + } if cfg.SandboxReadonly != nil { resolved.SandboxReadonly = *cfg.SandboxReadonly } @@ -2273,6 +2290,9 @@ func overlayFile(base, override FileConfig) FileConfig { if override.PromptCaching != nil { base.PromptCaching = override.PromptCaching } + if override.Compaction != nil { + base.Compaction = override.Compaction + } if override.MaxConcurrency > 0 { base.MaxConcurrency = override.MaxConcurrency } diff --git a/internal/config/loader_test.go b/internal/config/loader_test.go index 2e5d6828..d8d80e6d 100644 --- a/internal/config/loader_test.go +++ b/internal/config/loader_test.go @@ -1238,6 +1238,32 @@ func TestGlobalOverlay_PromptCaching(t *testing.T) { } } +// TestGlobalOverlay_Compaction verifies that Compaction from global config +// survives the merge. +func TestGlobalOverlay_Compaction(t *testing.T) { + t.Setenv("HOME", t.TempDir()) + globalDir := filepath.Join(os.Getenv("HOME"), ".odek") + os.MkdirAll(globalDir, 0755) + + if err := os.WriteFile(filepath.Join(globalDir, "config.json"), []byte(`{ + "compaction": true + }`), 0644); err != nil { + t.Fatal(err) + } + + t.Chdir(t.TempDir()) + if err := os.WriteFile("odek.json", []byte(`{ + "model": "project-model" + }`), 0644); err != nil { + t.Fatal(err) + } + + cfg := LoadConfig(CLIFlags{}) + if !cfg.Compaction { + t.Error("Compaction should be true (global value should survive project merge)") + } +} + // TestGlobalOverlay_MCPServers verifies that MCPServers from global config // survive the merge. BUG: overlayFile doesn't transfer MCPServers. func TestGlobalOverlay_MCPServers(t *testing.T) { @@ -1678,3 +1704,27 @@ func TestIsVarCont(t *testing.T) { } } } + +// TestEnvVar_Compaction verifies ODEK_COMPACTION enables compaction. +func TestEnvVar_Compaction(t *testing.T) { + t.Setenv("HOME", t.TempDir()) + t.Chdir(t.TempDir()) + t.Setenv("ODEK_COMPACTION", "true") + + cfg := LoadConfig(CLIFlags{}) + if !cfg.Compaction { + t.Error("Compaction should be true when ODEK_COMPACTION=true") + } +} + +// TestCLIFlags_Compaction verifies the --compaction CLI flag enables compaction. +func TestCLIFlags_Compaction(t *testing.T) { + t.Setenv("HOME", t.TempDir()) + t.Chdir(t.TempDir()) + + enabled := true + cfg := LoadConfig(CLIFlags{Compaction: &enabled}) + if !cfg.Compaction { + t.Error("Compaction should be true when CLIFlags.Compaction is set") + } +} diff --git a/internal/loop/loop.go b/internal/loop/loop.go index 2a53bc1b..6f3ee05e 100644 --- a/internal/loop/loop.go +++ b/internal/loop/loop.go @@ -7,6 +7,7 @@ import ( "encoding/hex" "encoding/json" "fmt" + "sort" "strings" "time" @@ -192,6 +193,23 @@ type Engine struct { TotalCacheReadTokens int // Anthropic: tokens read from cache TotalCachedTokens int // OpenAI: cached prompt tokens TotalCacheReported bool // provider returned cache metrics at least once + + // Context trimming state, cumulative for the engine lifetime so repeated + // trims update a single warning message instead of stacking new ones. + trimGroupsTotal int // total turn groups dropped by trimming + trimTruncTotal int // total tool results truncated by trimming + trimDroppedTools map[string]int // tool names seen in dropped groups + + // Self-calibrating context margin: when the provider reports substantially + // more input tokens than the local estimate, the margin tightens once. + lastReportedInputTokens int + lastEstimatedTotal int + tightMargin bool + + // compaction enables LLM-based rolling summarization of dropped turn + // groups (Config.Compaction). compactDigest holds the current summary. + compaction bool + compactDigest string } // New creates a new loop Engine. @@ -206,6 +224,7 @@ func New(client *llm.Client, registry *tool.Registry, maxIterations int, systemM system: systemMessage, maxContext: maxContext, maxConsecutiveToolErrors: make(map[string]int), + trimDroppedTools: make(map[string]int), } } @@ -280,6 +299,13 @@ func (e *Engine) SetApprover(a danger.Approver) { e.approver = a } // risk classes require approval and would skip pre-checking. func (e *Engine) SetDangerousConfig(cfg *danger.DangerousConfig) { e.dangerousCfg = cfg } +// SetCompaction enables or disables LLM-based rolling compaction of dropped +// context. When enabled, turn groups dropped by context trimming are +// summarized into a rolling digest system message instead of vanishing +// entirely. The digest is derived from (potentially untrusted) tool output, +// so it is wrapped with the engine's untrusted-content wrapper when set. +func (e *Engine) SetCompaction(enabled bool) { e.compaction = enabled } + // ── Token Estimation ───────────────────────────────────────────────── // // Zero-dependency heuristic: 1 token ≈ 4 chars for English text. @@ -296,6 +322,30 @@ const toolCallOverhead = 30 // Input (messages + tools) should not exceed this fraction. const contextSafetyMargin = 0.75 +// contextSafetyMarginTight is the tightened margin applied once token +// calibration detects that the local heuristic is underestimating real +// usage (dense code, reasoning tokens, large tool schemas). +const contextSafetyMarginTight = 0.65 + +// toolTruncateMinBytes is the minimum size of a tool result body eligible +// for graduated truncation during trimming. Small results are kept intact — +// truncating them saves little and destroys information. +const toolTruncateMinBytes = 2000 + +// keepRecentToolResults is the number of most recent tool result messages +// that graduated truncation never touches, so the agent always sees its +// latest tool output in full. +const keepRecentToolResults = 4 + +// digestMsgPrefix marks the rolling compaction digest system message so +// trimming can recognize, preserve, and update it. +const digestMsgPrefix = "[Compacted earlier context:" + +// isDigestMessage reports whether m is the rolling compaction digest. +func isDigestMessage(m llm.Message) bool { + return m.Role == "system" && strings.HasPrefix(m.Content, digestMsgPrefix) +} + // estimateTokens returns a rough upper-bound token count for a string. // Conservative: ~4 chars per token. Dense content (code, JSON) is // closer to 2-3 chars/token; this is safe for both. @@ -309,6 +359,7 @@ func estimateMessages(messages []llm.Message) int { for _, m := range messages { total += messageOverhead total += estimateTokens(m.Content) + total += estimateTokens(m.ReasoningContent) total += estimateTokens(m.Name) total += estimateTokens(m.ToolCallID) for _, tc := range m.ToolCalls { @@ -323,6 +374,8 @@ func estimateMessages(messages []llm.Message) int { // estimateToolDefs returns the estimated tokens for tool definitions. // These are sent with every request and count toward the context budget. +// The parameter schema is the bulk of every definition, so it is marshaled +// and counted; an unmarshalable schema falls back to a flat allowance. func estimateToolDefs(defs []llm.ToolDef) int { total := 0 for _, d := range defs { @@ -330,6 +383,13 @@ func estimateToolDefs(defs []llm.ToolDef) int { total += estimateTokens(d.Type) total += estimateTokens(d.Function.Name) total += estimateTokens(d.Function.Description) + if d.Function.Parameters != nil { + if schemaJSON, err := json.Marshal(d.Function.Parameters); err == nil { + total += estimateTokens(string(schemaJSON)) + } else { + total += 200 // fallback allowance + } + } } return total } @@ -344,69 +404,134 @@ func contextBudget(maxContext int) int { // ── Context Trimming ───────────────────────────────────────────────── +// headLen returns the number of leading messages that trimming must never +// drop: the base system prompt plus any other leading system messages +// (volatile memory block) and the first user message (the original task). +// After the task, only the rolling compaction digest is protected — other +// system messages that land there (skill/episode injections, trim warnings) +// remain droppable. +func headLen(messages []llm.Message) int { + start := 0 + seenTask := false + for start < len(messages) { + m := messages[start] + switch { + case m.Role == "system" && !seenTask: + start++ + case m.Role == "user" && !seenTask: + seenTask = true + start++ + case seenTask && isDigestMessage(m): + start++ + default: + return start + } + } + return start +} + // trimContext trims the message history to stay within the context budget. -// It preserves: -// - System message (always first, if present) -// - The first user message (the original task) // -// It drops the oldest non-essential message triples (assistant tool-call -// message + its tool result(s)) to avoid orphaning tool results without -// their preceding tool_calls — DeepSeek rejects orphaned tool messages. +// It preserves the protected head (see headLen): system prompt, leading +// system messages (memory block, compaction digest), and the first user +// message (the original task). +// +// Trimming is graduated: +// 1. Old, large tool result bodies (the token hogs) are replaced with a +// short marker, preserving the assistant's reasoning and the fact that +// the tool ran. The most recent tool results are never truncated. +// 2. If still over budget, the oldest complete turn groups (assistant +// tool-call message + its tool result(s)) are dropped atomically to +// avoid orphaning tool results — DeepSeek rejects orphaned tool messages. +// When compaction is enabled, dropped groups are first summarized into +// a rolling digest message (see refreshDigest). // -// When trimming occurs, a system message is injected to warn the agent -// that context was lost, preventing it from confidently operating on -// incomplete information. +// When trimming occurs, a system message is injected (before the most recent +// user message, keeping the cache-stable head untouched) to warn the agent +// that context was lost. Repeated trims update the single existing warning +// in place with cumulative totals. // // Performance: uses a running token total to avoid O(n²) re-scanning of -// the full message list on every iteration. Previously, estimateMessages -// was called at the top of the loop, re-summing ALL messages each time -// a single group was dropped. For large conversations near the context -// limit, this was O(n²) — now it's O(n). -func (e *Engine) trimContext(messages []llm.Message, toolDefs []llm.ToolDef) []llm.Message { +// the full message list on every iteration. +func (e *Engine) trimContext(ctx context.Context, messages []llm.Message, toolDefs []llm.ToolDef) []llm.Message { budget := contextBudget(e.maxContext) if budget <= 0 { return messages } + // Self-calibration: when the provider's reported input-token count for + // the previous call exceeded the local estimate by more than 15%, the + // heuristic is underestimating real usage. Tighten the safety margin + // for the rest of this engine's lifetime so proactive trimming kicks in + // earlier instead of relying on provider context-length errors. + if !e.tightMargin && e.lastReportedInputTokens > 0 && e.lastEstimatedTotal > 0 && + float64(e.lastReportedInputTokens) > 1.15*float64(e.lastEstimatedTotal) { + e.tightMargin = true + e.emitSignal(SignalEvent{Type: "context_trimmed", Detail: "margin_calibrated"}) + } + if e.tightMargin { + budget = int(float64(e.maxContext) * contextSafetyMarginTight) + } + // Estimate tool definitions once (they don't change between iterations) defTokens := estimateToolDefs(toolDefs) - // Compute the running total ONCE — each group drop then subtracts only - // the dropped group's tokens instead of re-scanning all messages. + // Compute the running total ONCE — each truncation/drop then subtracts + // only the affected tokens instead of re-scanning all messages. totalTokens := estimateMessages(messages) + defTokens - droppedGroups := 0 - droppedTools := make(map[string]int) - - for { - if totalTokens <= budget { - break + head := headLen(messages) + + // Pass 1 — graduated truncation: replace old, large tool results with a + // short marker before resorting to deleting whole turn groups. The most + // recent tool results and the protected head are never touched. + truncated := 0 + if totalTokens > budget { + protected := make(map[int]struct{}, keepRecentToolResults) + for i, n := len(messages)-1, 0; i >= 0 && n < keepRecentToolResults; i-- { + if messages[i].Role == "tool" { + protected[i] = struct{}{} + n++ + } } - if len(messages) <= 2 { - break // can't trim further (need system + task at minimum) + for i := head; i < len(messages) && totalTokens > budget; i++ { + if messages[i].Role != "tool" { + continue + } + if _, ok := protected[i]; ok { + continue + } + if len(messages[i].Content) <= toolTruncateMinBytes { + continue + } + oldEst := estimateTokens(messages[i].Content) + messages[i].Content = fmt.Sprintf( + "[tool output trimmed: %d bytes dropped to fit context budget]", + len(messages[i].Content), + ) + truncated++ + totalTokens -= oldEst - estimateTokens(messages[i].Content) } + } - // Find the first droppable index. - // Keep messages[0] if it's the system message. - // Keep the next message too (first user message = the task). - start := 0 - if messages[0].Role == "system" { - start = 1 // keep system - } - start++ // keep system + task - if start >= len(messages) { - break + // Pass 2 — drop the oldest complete turn groups. A group is either: + // - A standalone message (user text, assistant text) + // - An assistant tool_calls message + all following tool results + if e.trimDroppedTools == nil { + e.trimDroppedTools = make(map[string]int) + } + droppedGroups := 0 + var droppedForDigest []llm.Message + for totalTokens > budget { + if len(messages) <= head { + break // can't trim further — only the protected head remains } - - // Find the first complete droppable group starting at `start`. - // A group is either: - // - A standalone message (user text, assistant text) - // - An assistant tool_calls message + all following tool results + start := head groupEnd := start + 1 if messages[start].Role == "assistant" && len(messages[start].ToolCalls) > 0 { // Track which tools were called in dropped groups for _, tc := range messages[start].ToolCalls { - droppedTools[tc.Function.Name]++ + e.trimDroppedTools[tc.Function.Name]++ } // Include all following tool result messages for groupEnd < len(messages) && messages[groupEnd].Role == "tool" { @@ -414,6 +539,9 @@ func (e *Engine) trimContext(messages []llm.Message, toolDefs []llm.ToolDef) []l } } droppedGroups++ + if e.compaction { + droppedForDigest = append(droppedForDigest, messages[start:groupEnd]...) + } // Subtract the dropped group's tokens from the running total. // This avoids O(n²) behavior: we only scan the N messages being @@ -425,21 +553,17 @@ func (e *Engine) trimContext(messages []llm.Message, toolDefs []llm.ToolDef) []l messages = append(messages[:start], messages[groupEnd:]...) } - // Inject context trim warning if we dropped messages - if droppedGroups > 0 && len(messages) > 1 { - warning := fmt.Sprintf( - "[Context trimmed: %d prior message group(s) dropped to stay within token budget. "+ - "Some earlier tool calls and their results are no longer available. "+ - "If the user references earlier work, ask them to summarize what was done.]", - droppedGroups, - ) - // Insert after system message (index 0), before task (index 1) - trimMsg := llm.Message{Role: "system", Content: warning} - newMsgs := make([]llm.Message, 0, len(messages)+1) - newMsgs = append(newMsgs, messages[0]) - newMsgs = append(newMsgs, trimMsg) - newMsgs = append(newMsgs, messages[1:]...) - messages = newMsgs + // Rolling compaction: summarize the dropped groups into a digest system + // message so the information survives in compressed form. + if e.compaction && len(droppedForDigest) > 0 { + messages = e.refreshDigest(ctx, messages, droppedForDigest) + } + + // Inject or update the context trim warning if we trimmed anything. + if (droppedGroups > 0 || truncated > 0) && len(messages) > 1 { + e.trimGroupsTotal += droppedGroups + e.trimTruncTotal += truncated + messages = upsertTrimWarning(messages, e.buildTrimWarning()) e.emitSignal(SignalEvent{ Type: "context_trimmed", @@ -448,9 +572,80 @@ func (e *Engine) trimContext(messages []llm.Message, toolDefs []llm.ToolDef) []l }) } + // Record the final estimate so the next call can calibrate the margin + // against the provider's reported input tokens. + e.lastEstimatedTotal = estimateMessages(messages) + defTokens + return messages } +// buildTrimWarning renders the cumulative trim warning text, including how +// much was truncated/dropped and which tools lost their earlier results. +func (e *Engine) buildTrimWarning() string { + var sb strings.Builder + sb.WriteString("[Context trimmed: ") + parts := make([]string, 0, 2) + if e.trimTruncTotal > 0 { + parts = append(parts, fmt.Sprintf("%d tool output(s) truncated", e.trimTruncTotal)) + } + if e.trimGroupsTotal > 0 { + parts = append(parts, fmt.Sprintf("%d prior message group(s) dropped", e.trimGroupsTotal)) + } + sb.WriteString(strings.Join(parts, ", ")) + sb.WriteString(" to stay within the token budget. Some earlier tool calls and their results are no longer available.") + if len(e.trimDroppedTools) > 0 { + names := make([]string, 0, len(e.trimDroppedTools)) + for name := range e.trimDroppedTools { + names = append(names, name) + } + sort.Strings(names) + if len(names) > 5 { + names = names[:5] + } + sb.WriteString(" Earlier calls to these tools were dropped: ") + sb.WriteString(strings.Join(names, ", ")) + sb.WriteString(".") + } + if e.compaction && e.compactDigest != "" { + sb.WriteString(" A model-generated summary of the dropped turns is available in the '" + digestMsgPrefix + "...' system message.") + } + sb.WriteString(" If the user references earlier work, ask them to summarize what was done.]") + return sb.String() +} + +// upsertTrimWarning inserts the trim warning immediately before the most +// recent user message — keeping the cache-stable head untouched — or updates +// the existing warning in place when one is already present. The warning is +// never placed at index 0 so a session without a system prompt still starts +// with the task. +func upsertTrimWarning(messages []llm.Message, warning string) []llm.Message { + for i := range messages { + if messages[i].Role == "system" && strings.HasPrefix(messages[i].Content, "[Context trimmed:") { + messages[i].Content = warning + return messages + } + } + insertIdx := 1 + for i := len(messages) - 1; i >= 0; i-- { + if messages[i].Role == "user" { + insertIdx = i + break + } + } + if insertIdx < 1 { + insertIdx = 1 + } + if insertIdx > len(messages) { + insertIdx = len(messages) + } + trimMsg := llm.Message{Role: "system", Content: warning} + newMsgs := make([]llm.Message, 0, len(messages)+1) + newMsgs = append(newMsgs, messages[:insertIdx]...) + newMsgs = append(newMsgs, trimMsg) + newMsgs = append(newMsgs, messages[insertIdx:]...) + return newMsgs +} + // isContextLengthError returns true for API errors that indicate the // input exceeded the model's context window. These errors are retryable // with aggressive trimming rather than killing the session. @@ -475,9 +670,11 @@ func isContextLengthError(err error) bool { strings.Contains(msg, "reduce the length") } -// trimToSurvival drops all but the system prompt, first user message, -// and the most recent 2 complete turn groups. This is the nuclear option -// used when the API rejects the request as context-length-exceeded. +// trimToSurvival drops all but the system prompt, the rolling compaction +// digest (if present), the first user message (the original task, when it +// differs from the latest one), the most recent 2 complete turn groups, and +// the last user message. This is the nuclear option used when the API +// rejects the request as context-length-exceeded. // Unlike trimContext which gives up when it can't stay under budget, // trimToSurvival always produces a drastically reduced message list // that nearly every model can handle. @@ -489,12 +686,22 @@ func trimToSurvival(msgs []llm.Message) []llm.Message { if msgs[0].Role == "system" { start = 1 // keep system } + + // First user message (the original task) — kept when it differs from the + // last user message, so a long multi-turn session does not silently lose + // what it was asked to do. + firstUserIdx := -1 + for i := start; i < len(msgs); i++ { + if msgs[i].Role == "user" { + firstUserIdx = i + break + } + } + // Last user message (the current task/input) — always keep it. - var lastUser llm.Message lastUserIdx := -1 for i := len(msgs) - 1; i >= 0; i-- { if msgs[i].Role == "user" { - lastUser = msgs[i] lastUserIdx = i break } @@ -502,9 +709,13 @@ func trimToSurvival(msgs []llm.Message) []llm.Message { // Collect the last 2 complete assistant→tool groups before the user msg. // Each group is a sub-slice in correct internal order: [system*, assistant, tool*]. + scanFrom := lastUserIdx - 1 + if lastUserIdx < 0 { + scanFrom = len(msgs) - 1 + } var groups [][]llm.Message seen := 0 - for i := lastUserIdx - 1; i > start && seen < 2; i-- { + for i := scanFrom; i > start && seen < 2; i-- { if msgs[i].Role == "assistant" && len(msgs[i].ToolCalls) > 0 { var group []llm.Message @@ -531,13 +742,21 @@ func trimToSurvival(msgs []llm.Message) []llm.Message { } } - // Build survival set: system + task + recent groups + last user - // Calculate total messages across all groups for capacity hint + // Preserve the rolling compaction digest if one is present in the head. + digestIdx := -1 + for i := start; i < len(msgs) && i < start+4; i++ { + if isDigestMessage(msgs[i]) { + digestIdx = i + break + } + } + + // Build survival set: system + warning + digest + task + recent groups + last user totalGroupMsgs := 0 for _, g := range groups { totalGroupMsgs += len(g) } - survival := make([]llm.Message, 0, start+1+totalGroupMsgs+1) + survival := make([]llm.Message, 0, start+3+totalGroupMsgs+1) if start > 0 { survival = append(survival, msgs[0]) // system message } @@ -545,6 +764,15 @@ func trimToSurvival(msgs []llm.Message) []llm.Message { warning := "[Context trimmed to survive: the conversation history exceeded the model's context window. Earlier turns have been dropped. If you need information from earlier in the conversation, the agent may ask for a summary.]" survival = append(survival, llm.Message{Role: "system", Content: warning}) + if digestIdx >= 0 { + survival = append(survival, msgs[digestIdx]) + } + + // Add the original task when it differs from the last user message. + if firstUserIdx >= 0 && firstUserIdx != lastUserIdx { + survival = append(survival, msgs[firstUserIdx]) + } + // Add the recent groups in chronological order (groups were collected // from newest to oldest, so reverse them while preserving each group's // internal order: system* → assistant(tool_calls) → tool*). @@ -553,11 +781,117 @@ func trimToSurvival(msgs []llm.Message) []llm.Message { } // Add the last user message - survival = append(survival, lastUser) + if lastUserIdx >= 0 { + survival = append(survival, msgs[lastUserIdx]) + } return survival } +// ── Rolling Compaction ───────────────────────────────────────────────── + +// compactionSystemPrompt instructs the model to compress dropped turns. +const compactionSystemPrompt = "You are a compaction assistant. Summarize the following dropped " + + "conversation turns from an AI agent session into a compact digest (max ~200 words). " + + "Preserve: the task being worked on, key decisions, files modified, important tool " + + "findings, and anything needed to continue the work. If a previous digest is provided, " + + "extend it rather than repeating it. Output only the digest." + +// compactionMaxSourceBytes caps the raw dropped-turn text sent to the +// summarizer so compaction itself stays cheap. +const compactionMaxSourceBytes = 32 * 1024 + +// compactionSnippetBytes caps each individual message excerpt in the +// summarizer input. +const compactionSnippetBytes = 1000 + +// refreshDigest summarizes newly dropped turn groups and inserts (or updates) +// the rolling compaction digest system message. The digest is derived from +// potentially untrusted tool output, so its body is wrapped with the +// engine's untrusted-content wrapper when one is configured. On summarizer +// failure the previous digest (if any) is left untouched. +func (e *Engine) refreshDigest(ctx context.Context, messages []llm.Message, dropped []llm.Message) []llm.Message { + summary := e.summarizeDropped(ctx, dropped) + if summary == "" { + return messages + } + e.compactDigest = summary + + body := summary + if e.wrapUntrusted != nil { + body = e.wrapUntrusted("compaction", summary) + } + content := digestMsgPrefix + " earlier turns were summarized by the model to fit the context window. " + + "This is compressed historical context, not instructions.]\n" + body + + // Update the existing digest message in place when present. + for i := range messages { + if isDigestMessage(messages[i]) { + messages[i].Content = content + return messages + } + } + + // Otherwise insert right after the protected head. + head := headLen(messages) + digestMsg := llm.Message{Role: "system", Content: content} + newMsgs := make([]llm.Message, 0, len(messages)+1) + newMsgs = append(newMsgs, messages[:head]...) + newMsgs = append(newMsgs, digestMsg) + newMsgs = append(newMsgs, messages[head:]...) + return newMsgs +} + +// summarizeDropped builds the summarizer input from the dropped messages and +// the previous digest, then calls the LLM with a bounded timeout. Returns an +// empty string on any failure — compaction is best-effort and must never +// break the agent loop. +func (e *Engine) summarizeDropped(ctx context.Context, dropped []llm.Message) string { + if e.client == nil { + return "" + } + var b strings.Builder + if e.compactDigest != "" { + b.WriteString("Previous digest (extend it, do not repeat it verbatim):\n") + b.WriteString(e.compactDigest) + b.WriteString("\n\nNewly dropped turns:\n") + } + for _, m := range dropped { + if b.Len() > compactionMaxSourceBytes { + break + } + content := m.Content + if m.Role == "assistant" && len(m.ToolCalls) > 0 { + names := make([]string, 0, len(m.ToolCalls)) + for _, tc := range m.ToolCalls { + names = append(names, tc.Function.Name) + } + content = strings.TrimSpace(content + " [called tools: " + strings.Join(names, ", ") + "]") + } + if len(content) > compactionSnippetBytes { + content = content[:compactionSnippetBytes] + "…" + } + if content == "" { + continue + } + fmt.Fprintf(&b, "%s: %s\n", m.Role, content) + } + if b.Len() == 0 { + return "" + } + + callCtx, cancel := context.WithTimeout(ctx, 30*time.Second) + defer cancel() + res, err := e.client.Call(callCtx, []llm.Message{ + {Role: "system", Content: compactionSystemPrompt}, + {Role: "user", Content: b.String()}, + }, nil, nil) + if err != nil || res == nil { + return "" + } + return strings.TrimSpace(res.Content) +} + // ── Loop ────────────────────────────────────────────────────────────── // Run executes the loop for a given task and returns the final response. @@ -636,29 +970,15 @@ func (e *Engine) runLoop(ctx context.Context, messages []llm.Message) (string, [ } // Trim context to stay within model's context window - messages = e.trimContext(messages, tools) - - // Resync memMsgIdx after trimContext — when trimContext injects a - // context-warning message at index 1, all subsequent messages shift - // right by 1, making e.memMsgIdx point to the warning instead of - // memory. Detect this by checking if the message at our tracked - // position is a trim warning. - if e.memMsgIdx > 0 && e.memMsgIdx < len(messages) { - if strings.Contains(messages[e.memMsgIdx].Content, "[Context trimmed:") { - // Trim warning was injected before our memory message. - // The actual memory message is now at memMsgIdx + 1. - e.memMsgIdx++ - } - } - - // After trimContext, verify the memory message still exists at the - // tracked position. trimContext starts dropping from index 2 onward, - // and the memory message can shift into that range after a trim - // warning is injected (shifting it from index 1 to 2). Once at index - // 2, it can be dropped by subsequent trimContext calls, leaving - // memMsgIdx pointing to the wrong message (a user/assistant message - // that shifted into that slot). When this happens, reset memMsgIdx - // so the memory is re-inserted at the correct position. + messages = e.trimContext(ctx, messages, tools) + + // Verify the memory message still exists at the tracked position. + // trimContext protects the leading run of system messages (base + // prompt + memory block) via headLen, so the memory message can no + // longer be dropped or shifted by trimming — but trimToSurvival (on + // a provider context-length error) still removes it. When that + // happens, reset memMsgIdx so memory is re-inserted at the correct + // position. if e.memMsgIdx >= 0 && e.memMsgIdx < len(messages) { if messages[e.memMsgIdx].Role != "system" { // Memory message was dropped — re-insert on next update. @@ -820,6 +1140,12 @@ func (e *Engine) runLoop(ctx context.Context, messages []llm.Message) (string, [ // THINK (timed) start := time.Now() + // Re-check the budget after all context injections (memory block, + // skills, episodes, extended memory) — those are added after the + // top-of-loop trim and can push an already-near-budget request over + // the model's context window on this very call. + messages = e.trimContext(ctx, messages, tools) + // Apply prompt caching markers when enabled — but only for Anthropic // endpoints. OpenAI rejects the Anthropic request shape (top-level // "system" field) with a 400, and DeepSeek caches automatically, @@ -870,6 +1196,10 @@ func (e *Engine) runLoop(ctx context.Context, messages []llm.Message) (string, [ e.TotalInputTokens += result.InputTokens e.TotalOutputTokens += result.OutputTokens + // Feed the margin calibration in trimContext: provider-reported input + // tokens are ground truth for how accurate the local estimate is. + e.lastReportedInputTokens = result.InputTokens + // Accumulate cache metrics // Accumulate cache metrics across iterations e.TotalCacheCreationTokens += result.CacheCreationTokens diff --git a/internal/loop/loop_test.go b/internal/loop/loop_test.go index 0d4e5657..96e7634c 100644 --- a/internal/loop/loop_test.go +++ b/internal/loop/loop_test.go @@ -543,7 +543,7 @@ func TestTrimContext_NoLimit(t *testing.T) { {Role: "system", Content: "You are a bot."}, {Role: "user", Content: "hello"}, } - result := engine.trimContext(msgs, nil) + result := engine.trimContext(context.Background(), msgs, nil) if len(result) != 2 { t.Errorf("trimContext with no limit should not change messages, got %d", len(result)) } @@ -558,7 +558,7 @@ func TestTrimContext_UnderBudget(t *testing.T) { {Role: "assistant", Content: "Hi there", ToolCalls: nil}, {Role: "tool", Content: "result", ToolCallID: "call_1"}, } - result := engine.trimContext(msgs, nil) + result := engine.trimContext(context.Background(), msgs, nil) if len(result) != 4 { t.Errorf("trimContext under budget should keep all messages, got %d", len(result)) } @@ -577,7 +577,7 @@ func TestTrimContext_OverBudget(t *testing.T) { {Role: "assistant", Content: strings.Repeat("final reasoning... ", 20)}, {Role: "tool", Content: strings.Repeat("final data ", 20), ToolCallID: "call_3"}, } - result := engine.trimContext(msgs, nil) + result := engine.trimContext(context.Background(), msgs, nil) // Should have preserved system + task (first user) if len(result) < 2 { @@ -608,7 +608,7 @@ func TestTrimContext_VeryTightBudget(t *testing.T) { {Role: "assistant", Content: strings.Repeat("data ", 50)}, {Role: "tool", Content: strings.Repeat("result ", 50), ToolCallID: "call_1"}, } - result := engine.trimContext(msgs, nil) + result := engine.trimContext(context.Background(), msgs, nil) // Must keep system + task at minimum if len(result) < 2 { @@ -632,7 +632,7 @@ func TestTrimContext_NoSystemMessage(t *testing.T) { {Role: "assistant", Content: strings.Repeat("data ", 30)}, {Role: "tool", Content: strings.Repeat("result ", 30), ToolCallID: "call_1"}, } - result := engine.trimContext(msgs, nil) + result := engine.trimContext(context.Background(), msgs, nil) // Without system, keep at least the task if len(result) < 1 { @@ -680,7 +680,7 @@ func TestTrimContext_IncludesToolDefTokens(t *testing.T) { }, }} - result := engine.trimContext(msgs, defs) + result := engine.trimContext(context.Background(), msgs, defs) if len(result) >= len(msgs) { t.Errorf("trimContext with tool defs should trim, got %d >= %d", len(result), len(msgs)) } @@ -1281,7 +1281,7 @@ func BenchmarkTrimContext(b *testing.B) { // Copy messages each iteration to avoid modifying shared state. cp := make([]llm.Message, len(msgs)) copy(cp, msgs) - engine.trimContext(cp, nil) + engine.trimContext(context.Background(), cp, nil) } }) } @@ -1298,7 +1298,7 @@ func BenchmarkTrimContext_NoTrim(b *testing.B) { b.ReportAllocs() b.ResetTimer() for range b.N { - engine.trimContext(msgs, nil) + engine.trimContext(context.Background(), msgs, nil) } } diff --git a/internal/loop/loop_trim_test.go b/internal/loop/loop_trim_test.go new file mode 100644 index 00000000..bb5b0fdd --- /dev/null +++ b/internal/loop/loop_trim_test.go @@ -0,0 +1,721 @@ +package loop + +import ( + "context" + "fmt" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/BackendStack21/odek/internal/llm" + "github.com/BackendStack21/odek/internal/tool" +) + +// ── Estimators ───────────────────────────────────────────────────────── + +func TestEstimateToolDefs_IncludesParameters(t *testing.T) { + without := estimateToolDefs([]llm.ToolDef{{ + Type: "function", + Function: llm.FunctionDef{Name: "shell", Description: "run a command"}, + }}) + with := estimateToolDefs([]llm.ToolDef{{ + Type: "function", + Function: llm.FunctionDef{ + Name: "shell", + Description: "run a command", + Parameters: map[string]any{ + "type": "object", + "properties": map[string]any{ + "command": map[string]any{ + "type": "string", + "description": strings.Repeat("the command to execute ", 20), + }, + }, + }, + }, + }}) + if with <= without { + t.Errorf("estimateToolDefs with schema = %d, want > %d (schema must be counted)", with, without) + } +} + +func TestEstimateMessages_CountsReasoningContent(t *testing.T) { + plain := estimateMessages([]llm.Message{{Role: "assistant", Content: "answer"}}) + withReasoning := estimateMessages([]llm.Message{{ + Role: "assistant", + Content: "answer", + ReasoningContent: strings.Repeat("thinking step by step ", 50), + }}) + if withReasoning <= plain { + t.Errorf("estimateMessages with reasoning = %d, want > %d", withReasoning, plain) + } +} + +// ── Graduated truncation ─────────────────────────────────────────────── + +// buildToolConversation returns system + task + n groups of +// (assistant text, tool result of toolBytes bytes). +func buildToolConversation(n, toolBytes int) []llm.Message { + msgs := []llm.Message{ + {Role: "system", Content: "sys"}, + {Role: "user", Content: "task"}, + } + for i := 0; i < n; i++ { + msgs = append(msgs, + llm.Message{Role: "assistant", Content: fmt.Sprintf("thinking %d", i)}, + llm.Message{Role: "tool", Content: strings.Repeat("x", toolBytes), ToolCallID: fmt.Sprintf("c%d", i)}, + ) + } + return msgs +} + +func TestTrimContext_TruncatesOldToolResults(t *testing.T) { + msgs := buildToolConversation(6, 4000) + origLen := len(msgs) + // Budget: over with all results intact, under once the two oldest + // (unprotected) tool results are truncated. + engine := &Engine{maxContext: 7500} + result := engine.trimContext(context.Background(), msgs, nil) + + // No group may be dropped — only the warning is added. + if len(result) != origLen+1 { + t.Fatalf("expected %d messages (no drops + warning), got %d", origLen+1, len(result)) + } + // The two oldest tool results are truncated; the 4 most recent are not. + truncated := map[string]bool{} + for _, m := range result { + if m.Role == "tool" && strings.Contains(m.Content, "[tool output trimmed:") { + truncated[m.ToolCallID] = true + } + } + if !truncated["c0"] || !truncated["c1"] || len(truncated) != 2 { + t.Errorf("expected exactly c0,c1 truncated, got %v", truncated) + } + for _, m := range result { + if m.Role == "tool" && !truncated[m.ToolCallID] && len(m.Content) != 4000 { + t.Errorf("recent tool result %s changed length: %d", m.ToolCallID, len(m.Content)) + } + } + // Warning reports truncation, not group drops. + var warning string + for _, m := range result { + if strings.HasPrefix(m.Content, "[Context trimmed:") { + warning = m.Content + } + } + if !strings.Contains(warning, "2 tool output(s) truncated") { + t.Errorf("warning should mention 2 truncated outputs, got %q", warning) + } + if strings.Contains(warning, "group(s) dropped") { + t.Errorf("warning should not mention dropped groups, got %q", warning) + } +} + +func TestTrimContext_TruncationInsufficient_DropsGroups(t *testing.T) { + msgs := buildToolConversation(6, 4000) + engine := &Engine{maxContext: 200} // tiny budget — truncation alone can't fit + result := engine.trimContext(context.Background(), msgs, nil) + + if len(result) >= len(msgs) { + t.Errorf("expected groups to be dropped, got %d >= %d messages", len(result), len(msgs)) + } + // No orphaned tool messages: every tool message must follow an assistant. + for i, m := range result { + if m.Role == "tool" && i > 0 && result[i-1].Role != "assistant" && result[i-1].Role != "tool" { + t.Errorf("orphaned tool message at index %d", i) + } + } + var warning string + for _, m := range result { + if strings.HasPrefix(m.Content, "[Context trimmed:") { + warning = m.Content + } + } + if !strings.Contains(warning, "group(s) dropped") { + t.Errorf("warning should mention dropped groups, got %q", warning) + } +} + +// ── Warning content / placement ──────────────────────────────────────── + +func TestTrimContext_WarningIncludesDroppedToolNames(t *testing.T) { + tc := func(id, name string) []llm.ToolCall { + return []llm.ToolCall{{ID: id, Type: "function", Function: struct { + Name string `json:"name"` + Arguments string `json:"arguments"` + }{Name: name, Arguments: "{}"}}} + } + msgs := []llm.Message{ + {Role: "system", Content: "sys"}, + {Role: "user", Content: "task"}, + {Role: "assistant", ToolCalls: tc("c1", "read_file")}, + {Role: "tool", Content: "small", ToolCallID: "c1"}, + {Role: "assistant", Content: strings.Repeat("pad ", 2000)}, + {Role: "user", Content: "latest input"}, + } + engine := &Engine{maxContext: 400} + result := engine.trimContext(context.Background(), msgs, nil) + + var warning string + for _, m := range result { + if strings.HasPrefix(m.Content, "[Context trimmed:") { + warning = m.Content + } + } + if warning == "" { + t.Fatal("expected a trim warning") + } + if !strings.Contains(warning, "read_file") { + t.Errorf("warning should name dropped tools, got %q", warning) + } + // Warning sits before the most recent user message, not at the head. + for i, m := range result { + if strings.HasPrefix(m.Content, "[Context trimmed:") { + if i+1 >= len(result) || result[i+1].Role != "user" { + t.Errorf("warning at index %d should be followed by a user message", i) + } + } + } +} + +func TestTrimContext_WarningUpdatesInPlace(t *testing.T) { + engine := &Engine{maxContext: 600} + msgs := []llm.Message{ + {Role: "system", Content: "sys"}, + {Role: "user", Content: "task"}, + {Role: "assistant", Content: strings.Repeat("a", 3000)}, + {Role: "assistant", Content: strings.Repeat("b", 3000)}, + } + result := engine.trimContext(context.Background(), msgs, nil) + + // Second trim on the result with a new oversized message appended. + result = append(result, llm.Message{Role: "assistant", Content: strings.Repeat("c", 3000)}) + result = engine.trimContext(context.Background(), result, nil) + + count := 0 + var warning string + for _, m := range result { + if strings.HasPrefix(m.Content, "[Context trimmed:") { + count++ + warning = m.Content + } + } + if count != 1 { + t.Fatalf("expected exactly 1 trim warning, got %d", count) + } + if !strings.Contains(warning, "3 prior message group(s) dropped") { + t.Errorf("warning should carry cumulative drop count 3, got %q", warning) + } +} + +// ── Margin calibration ───────────────────────────────────────────────── + +func TestTrimContext_MarginCalibration(t *testing.T) { + var signals []SignalEvent + engine := &Engine{ + maxContext: 100_000, + lastEstimatedTotal: 10_000, + lastReportedInputTokens: 20_000, // 2x the estimate → > 15% off + maxConsecutiveToolErrors: map[string]int{}, + trimDroppedTools: map[string]int{}, + } + engine.SetSignalHandler(func(ev SignalEvent) { signals = append(signals, ev) }) + + // ~68k estimated tokens: fits the default 75% margin (75k) but not the + // tightened 65% margin (65k). + msgs := []llm.Message{ + {Role: "system", Content: "sys"}, + {Role: "user", Content: "task"}, + {Role: "assistant", Content: strings.Repeat("x", 268_000)}, + } + result := engine.trimContext(context.Background(), msgs, nil) + + if !engine.tightMargin { + t.Error("margin should have tightened after provider reported 2x the estimate") + } + calibrated := false + for _, ev := range signals { + if ev.Type == "context_trimmed" && ev.Detail == "margin_calibrated" { + calibrated = true + } + } + if !calibrated { + t.Error("expected a margin_calibrated signal") + } + if len(result) >= len(msgs)+1 { + t.Errorf("expected trimming under the tightened margin, got %d messages", len(result)) + } +} + +// ── Post-injection budget re-check ───────────────────────────────────── + +func TestTrimContext_PostInjectionBudget(t *testing.T) { + var bodies []string + callCount := 0 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + data, _ := io.ReadAll(r.Body) + bodies = append(bodies, string(data)) + callCount++ + if callCount == 1 { + fmt.Fprint(w, `{"choices":[{"message":{"content":"","tool_calls":[{"id":"c1","function":{"name":"echo","arguments":"{}"}}]}}]}`) + } else { + fmt.Fprint(w, `{"choices":[{"message":{"content":"done"}}]}`) + } + })) + defer server.Close() + + echoTool := &fakeTool{name: "echo", description: "echo", output: "ok"} + registry := tool.NewRegistry([]tool.Tool{echoTool}) + client := llm.New(server.URL, "sk-test", "test-model", "", 0, 0) + // Budget 1500 tokens — the injected skill block (~2500 tokens) must be + // dropped before the very first API call, not one iteration later. + engine := New(client, registry, 5, "sys", nil, 2000) + engine.SetSkillLoader(func(string) string { return strings.Repeat("SKILLDATA ", 1000) }) + + if _, err := engine.Run(context.Background(), "do it"); err != nil { + t.Fatalf("Run() error: %v", err) + } + if len(bodies) == 0 { + t.Fatal("no requests recorded") + } + if strings.Contains(bodies[0], "SKILLDATA") { + t.Error("first API call carried the oversized injected skill block — post-injection trim missing") + } +} + +// ── trimToSurvival ───────────────────────────────────────────────────── + +func survivalTC(id, name string) []llm.ToolCall { + return []llm.ToolCall{{ID: id, Type: "function", Function: struct { + Name string `json:"name"` + Arguments string `json:"arguments"` + }{Name: name, Arguments: "{}"}}} +} + +func TestTrimToSurvival_KeepsOriginalTask(t *testing.T) { + msgs := []llm.Message{ + {Role: "system", Content: "sys"}, + {Role: "user", Content: "original task"}, + {Role: "assistant", ToolCalls: survivalTC("c1", "read_file")}, + {Role: "tool", Content: "r1", ToolCallID: "c1"}, + {Role: "user", Content: "follow-up question"}, + {Role: "assistant", ToolCalls: survivalTC("c2", "shell")}, + {Role: "tool", Content: "r2", ToolCallID: "c2"}, + {Role: "user", Content: "latest input"}, + } + got := trimToSurvival(msgs) + + foundTask := false + for _, m := range got { + if m.Role == "user" && m.Content == "original task" { + foundTask = true + } + } + if !foundTask { + t.Error("original task message must survive the survival trim") + } + last := got[len(got)-1] + if last.Role != "user" || last.Content != "latest input" { + t.Errorf("last message = %q/%q, want user/latest input", last.Role, last.Content) + } + // Task must appear before the last user message. + taskIdx, latestIdx := -1, -1 + for i, m := range got { + if m.Role == "user" && m.Content == "original task" { + taskIdx = i + } + if m.Role == "user" && m.Content == "latest input" { + latestIdx = i + } + } + if taskIdx < 0 || latestIdx < 0 || taskIdx > latestIdx { + t.Errorf("task at %d, latest at %d — bad ordering", taskIdx, latestIdx) + } +} + +func TestTrimToSurvival_NoUserMessage(t *testing.T) { + msgs := []llm.Message{ + {Role: "system", Content: "sys"}, + {Role: "assistant", ToolCalls: survivalTC("c1", "echo")}, + {Role: "tool", Content: "r1", ToolCallID: "c1"}, + {Role: "assistant", Content: "thinking out loud"}, + } + got := trimToSurvival(msgs) + for i, m := range got { + if m.Role == "" { + t.Errorf("message %d has empty role — zero-value message leaked", i) + } + } +} + +func TestTrimToSurvival_PreservesDigest(t *testing.T) { + msgs := []llm.Message{ + {Role: "system", Content: "sys"}, + {Role: "system", Content: digestMsgPrefix + " summary of old work]\ndigest body"}, + {Role: "user", Content: "task"}, + {Role: "assistant", ToolCalls: survivalTC("c1", "echo")}, + {Role: "tool", Content: "r1", ToolCallID: "c1"}, + {Role: "user", Content: "latest"}, + } + got := trimToSurvival(msgs) + found := false + for _, m := range got { + if isDigestMessage(m) { + found = true + } + } + if !found { + t.Error("compaction digest must survive the survival trim") + } +} + +// ── Rolling compaction ───────────────────────────────────────────────── + +func TestTrimContext_CompactionCreatesDigest(t *testing.T) { + summaryCalls := 0 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + summaryCalls++ + fmt.Fprint(w, `{"choices":[{"message":{"content":"SUMMARY: earlier work condensed"}}]}`) + })) + defer server.Close() + + client := llm.New(server.URL, "sk-test", "test-model", "", 0, 0) + engine := New(client, tool.NewRegistry(nil), 10, "", nil, 200) + engine.SetCompaction(true) + + msgs := []llm.Message{ + {Role: "system", Content: "sys"}, + {Role: "user", Content: "task"}, + {Role: "assistant", Content: strings.Repeat("a", 3000)}, + {Role: "assistant", Content: strings.Repeat("b", 3000)}, + {Role: "assistant", Content: strings.Repeat("c", 3000)}, + } + result := engine.trimContext(context.Background(), msgs, nil) + + if engine.compactDigest != "SUMMARY: earlier work condensed" { + t.Errorf("compactDigest = %q, want summary", engine.compactDigest) + } + digestCount := 0 + for _, m := range result { + if isDigestMessage(m) { + digestCount++ + if !strings.Contains(m.Content, "SUMMARY: earlier work condensed") { + t.Errorf("digest message missing summary: %q", m.Content) + } + } + } + if digestCount != 1 { + t.Fatalf("expected exactly 1 digest message, got %d", digestCount) + } + if summaryCalls == 0 { + t.Error("summarizer was never called") + } + + // A second trim updates the existing digest in place. + result = append(result, llm.Message{Role: "assistant", Content: strings.Repeat("d", 3000)}) + result = engine.trimContext(context.Background(), result, nil) + digestCount = 0 + for _, m := range result { + if isDigestMessage(m) { + digestCount++ + } + } + if digestCount != 1 { + t.Errorf("after second trim: expected 1 digest message, got %d", digestCount) + } +} + +func TestTrimContext_CompactionFailureStillTrims(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusInternalServerError) + })) + defer server.Close() + + client := llm.New(server.URL, "sk-test", "test-model", "", 0, 0) + engine := New(client, tool.NewRegistry(nil), 10, "", nil, 200) + engine.SetCompaction(true) + + msgs := []llm.Message{ + {Role: "system", Content: "sys"}, + {Role: "user", Content: "task"}, + {Role: "assistant", Content: strings.Repeat("a", 3000)}, + {Role: "assistant", Content: strings.Repeat("b", 3000)}, + } + result := engine.trimContext(context.Background(), msgs, nil) + + for _, m := range result { + if isDigestMessage(m) { + t.Error("no digest should be inserted when the summarizer fails") + } + } + if len(result) >= len(msgs) { + t.Errorf("trimming must still happen on summarizer failure, got %d messages", len(result)) + } +} + +func TestTrimContext_CompactionWrapsUntrusted(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + fmt.Fprint(w, `{"choices":[{"message":{"content":"digest"}}]}`) + })) + defer server.Close() + + client := llm.New(server.URL, "sk-test", "test-model", "", 0, 0) + engine := New(client, tool.NewRegistry(nil), 10, "", nil, 200) + engine.SetCompaction(true) + engine.SetUntrustedWrapper(func(source, content string) string { + return "" + content + "" + }) + + msgs := []llm.Message{ + {Role: "system", Content: "sys"}, + {Role: "user", Content: "task"}, + {Role: "assistant", Content: strings.Repeat("a", 3000)}, + } + result := engine.trimContext(context.Background(), msgs, nil) + + found := false + for _, m := range result { + if isDigestMessage(m) { + found = true + if !strings.Contains(m.Content, "") { + t.Errorf("digest body not wrapped as untrusted: %q", m.Content) + } + } + } + if !found { + t.Error("expected a digest message") + } +} + +// ── Coverage: estimator fallback ─────────────────────────────────────── + +func TestEstimateToolDefs_UnmarshalableSchema(t *testing.T) { + base := estimateToolDefs([]llm.ToolDef{{ + Type: "function", + Function: llm.FunctionDef{Name: "shell", Description: "run a command"}, + }}) + withBad := estimateToolDefs([]llm.ToolDef{{ + Type: "function", + Function: llm.FunctionDef{ + Name: "shell", + Description: "run a command", + Parameters: map[string]any{"bad": func() {}}, // json.Marshal fails + }, + }}) + if withBad != base+200 { + t.Errorf("unmarshalable schema should add the 200-token fallback: got %d, want %d", withBad, base+200) + } +} + +// ── Coverage: small tool results are never truncated ─────────────────── + +func TestTrimContext_SmallToolResultNotTruncated(t *testing.T) { + msgs := []llm.Message{ + {Role: "system", Content: "sys"}, + {Role: "user", Content: "task"}, + } + // 6 tool messages → the oldest 2 are unprotected: one small (skipped by + // the truncation pass), one big (truncated). + sizes := []int{500, 4000, 4000, 4000, 4000, 4000} + for i, sz := range sizes { + msgs = append(msgs, + llm.Message{Role: "assistant", Content: fmt.Sprintf("thinking %d", i)}, + llm.Message{Role: "tool", Content: strings.Repeat("x", sz), ToolCallID: fmt.Sprintf("c%d", i)}, + ) + } + origLen := len(msgs) + // Budget forces exactly the big unprotected result to be truncated. + engine := &Engine{maxContext: 7000} + result := engine.trimContext(context.Background(), msgs, nil) + + if len(result) != origLen+1 { + t.Fatalf("expected no drops (only warning added), got %d != %d", len(result), origLen+1) + } + for _, m := range result { + if m.Role != "tool" { + continue + } + switch m.ToolCallID { + case "c0": + if len(m.Content) != 500 { + t.Errorf("small unprotected result c0 must stay intact, got len=%d", len(m.Content)) + } + case "c1": + if !strings.Contains(m.Content, "[tool output trimmed:") { + t.Errorf("big unprotected result c1 should be truncated") + } + } + } +} + +// ── Coverage: warning caps dropped tool names ────────────────────────── + +func TestBuildTrimWarning_CapsToolNames(t *testing.T) { + engine := &Engine{ + trimGroupsTotal: 1, + trimDroppedTools: map[string]int{}, + } + for _, name := range []string{"aaa", "bbb", "ccc", "ddd", "eee", "fff", "ggg"} { + engine.trimDroppedTools[name] = 1 + } + warning := engine.buildTrimWarning() + if !strings.Contains(warning, "eee") { + t.Errorf("warning should include the first 5 sorted names, got %q", warning) + } + if strings.Contains(warning, "fff") || strings.Contains(warning, "ggg") { + t.Errorf("warning should cap dropped tool names at 5, got %q", warning) + } +} + +// ── Coverage: warning upsert edge cases ──────────────────────────────── + +func TestUpsertTrimWarning_EdgeCases(t *testing.T) { + // No user message — warning goes to index 1. + msgs := []llm.Message{{Role: "assistant", Content: "a"}} + got := upsertTrimWarning(msgs, "[Context trimmed: x]") + if len(got) != 2 || got[1].Content != "[Context trimmed: x]" { + t.Errorf("no-user case: got %+v", got) + } + // Empty slice — insertion index clamps to 0. + got = upsertTrimWarning(nil, "[Context trimmed: x]") + if len(got) != 1 || got[0].Content != "[Context trimmed: x]" { + t.Errorf("empty case: got %+v", got) + } + // Task at index 0 (no system prompt) — warning clamps to index 1 so the + // session still starts with the task. + got = upsertTrimWarning([]llm.Message{{Role: "user", Content: "task"}}, "[Context trimmed: x]") + if len(got) != 2 || got[0].Role != "user" || got[1].Content != "[Context trimmed: x]" { + t.Errorf("task-first case: got %+v", got) + } +} + +// ── Coverage: survival keeps preceding system messages in a group ────── + +func TestTrimToSurvival_IncludesPrecedingSystemMessages(t *testing.T) { + msgs := []llm.Message{ + {Role: "system", Content: "sys"}, + {Role: "user", Content: "task"}, + {Role: "system", Content: "correction note"}, + {Role: "assistant", ToolCalls: survivalTC("c1", "echo")}, + {Role: "tool", Content: "r1", ToolCallID: "c1"}, + {Role: "user", Content: "latest"}, + } + got := trimToSurvival(msgs) + found := false + for _, m := range got { + if m.Content == "correction note" { + found = true + } + } + if !found { + t.Error("system message preceding the assistant tool-call group must survive") + } +} + +// ── Coverage: summarizeDropped branches ──────────────────────────────── + +func TestSummarizeDropped_NilClient(t *testing.T) { + engine := &Engine{} + if got := engine.summarizeDropped(context.Background(), []llm.Message{{Role: "assistant", Content: "x"}}); got != "" { + t.Errorf("nil client must return empty, got %q", got) + } +} + +func TestSummarizeDropped_EmptyContent(t *testing.T) { + calls := 0 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + calls++ + fmt.Fprint(w, `{"choices":[{"message":{"content":"digest"}}]}`) + })) + defer server.Close() + + engine := New(llm.New(server.URL, "sk-test", "m", "", 0, 0), tool.NewRegistry(nil), 10, "", nil, 0) + got := engine.summarizeDropped(context.Background(), []llm.Message{{Role: "assistant", Content: ""}}) + if got != "" { + t.Errorf("empty dropped content must return empty, got %q", got) + } + if calls != 0 { + t.Errorf("summarizer must not be called for empty input, got %d calls", calls) + } +} + +func TestSummarizeDropped_InputBuilding(t *testing.T) { + var bodies []string + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + data, _ := io.ReadAll(r.Body) + bodies = append(bodies, string(data)) + fmt.Fprint(w, `{"choices":[{"message":{"content":"digest"}}]}`) + })) + defer server.Close() + + newEngine := func() *Engine { + return New(llm.New(server.URL, "sk-test", "m", "", 0, 0), tool.NewRegistry(nil), 10, "", nil, 0) + } + + // Assistant tool-call names are included in the summarizer input. + e := newEngine() + e.summarizeDropped(context.Background(), []llm.Message{{ + Role: "assistant", + ToolCalls: survivalTC("c1", "read_file"), + }}) + if !strings.Contains(bodies[len(bodies)-1], "[called tools: read_file]") { + t.Errorf("tool-call names missing from summarizer input: %.200s", bodies[len(bodies)-1]) + } + + // Long message content is snippet-truncated. + e = newEngine() + e.summarizeDropped(context.Background(), []llm.Message{{ + Role: "tool", + Content: strings.Repeat("y", 3000), + }}) + if strings.Contains(bodies[len(bodies)-1], strings.Repeat("y", 2000)) { + t.Error("long message content should be snippet-truncated in summarizer input") + } + + // A previous digest is included for rolling extension. + e = newEngine() + e.compactDigest = "OLD DIGEST" + e.summarizeDropped(context.Background(), []llm.Message{{Role: "assistant", Content: "new work"}}) + body := bodies[len(bodies)-1] + if !strings.Contains(body, "Previous digest") || !strings.Contains(body, "OLD DIGEST") { + t.Errorf("previous digest missing from summarizer input: %.200s", body) + } + + // The raw source is capped at compactionMaxSourceBytes. + e = newEngine() + big := make([]llm.Message, 0, 40) + for i := 0; i < 40; i++ { + big = append(big, llm.Message{Role: "assistant", Content: strings.Repeat("z", 1000)}) + } + e.summarizeDropped(context.Background(), big) + if len(bodies[len(bodies)-1]) > compactionMaxSourceBytes+4096 { + t.Errorf("summarizer input exceeds source cap: %d bytes", len(bodies[len(bodies)-1])) + } +} + +// ── Coverage: stale memMsgIdx invariant guard ────────────────────────── + +func TestRunLoop_StaleMemMsgIdxReset(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + fmt.Fprint(w, `{"choices":[{"message":{"content":"done"}}]}`) + })) + defer server.Close() + + client := llm.New(server.URL, "sk-test", "test-model", "", 0, 0) + engine := New(client, tool.NewRegistry(nil), 1, "sys", nil, 0) + // Simulate a stale memory-message index pointing at a non-system message. + engine.memMsgIdx = 1 + + _, _, err := engine.runLoop(context.Background(), []llm.Message{ + {Role: "system", Content: "s"}, + {Role: "user", Content: "task"}, + }) + if err != nil { + t.Fatalf("runLoop() error: %v", err) + } + if engine.memMsgIdx != -1 { + t.Errorf("stale memMsgIdx should be reset to -1, got %d", engine.memMsgIdx) + } +} diff --git a/internal/loop/signal.go b/internal/loop/signal.go index 26d99451..3aab54ca 100644 --- a/internal/loop/signal.go +++ b/internal/loop/signal.go @@ -12,8 +12,10 @@ type SignalEvent struct { // Type is the signal kind. One of: // "context_trimmed" — prior message groups were dropped to fit the token // budget (Count = groups dropped, Detail = "proactive" - // for the pre-call budget trim or "survival" for the - // post-error nuclear trim) + // for the pre-call budget trim, "survival" for the + // post-error nuclear trim, or "margin_calibrated" + // when the safety margin tightened after the provider + // reported more input tokens than estimated) // "tool_recovery" — a tool failed repeatedly and the engine injected a // corrective hint so the model changes approach // (Tool = failing tool, Detail = the correction) diff --git a/internal/session/session.go b/internal/session/session.go index bb08e932..a0b00bbc 100644 --- a/internal/session/session.go +++ b/internal/session/session.go @@ -445,10 +445,14 @@ func (s *Store) saveLocked(sess *Session) error { // index 0 is always kept, and an assistant tool_calls message is dropped // together with its following tool-result messages so a stored transcript // never contains orphaned tool messages (which strict providers reject). -// The turn count is recounted to match the surviving messages. If nothing -// droppable remains (a degenerate case, e.g. a single oversized system -// message), the session is written as-is — failing the save would lose data. +// The turn count is recounted to match the surviving messages. When any +// groups were dropped, a marker system message is inserted after the system +// prompt so a resumed session can see that earlier history was removed. +// If nothing droppable remains (a degenerate case, e.g. a single oversized +// system message), the session is written as-is — failing the save would +// lose data. func (s *Store) trimToFileCapLocked(sess *Session, data []byte) ([]byte, error) { + droppedGroups := 0 for len(data) > MaxSessionFileBytes { start := 0 if len(sess.Messages) > 0 && sess.Messages[0].Role == "system" { @@ -478,6 +482,7 @@ func (s *Store) trimToFileCapLocked(sess *Session, data []byte) ([]byte, error) return nil, fmt.Errorf("session: marshal trim candidate: %w", err) } freed += len(groupJSON) + droppedGroups++ dropEnd = groupEnd } sess.Messages = append(sess.Messages[:start], sess.Messages[dropEnd:]...) @@ -488,6 +493,38 @@ func (s *Store) trimToFileCapLocked(sess *Session, data []byte) ([]byte, error) return nil, fmt.Errorf("session: marshal after trim: %w", err) } } + + // Persist a marker so a resumed session knows earlier turns were removed + // (the stderr warning alone never reaches the transcript). + if droppedGroups > 0 { + marker := llm.Message{ + Role: "system", + Content: fmt.Sprintf( + "[Session storage limit: %d oldest message group(s) were removed from this transcript to stay within the %d-byte file cap. Earlier conversation context is unavailable.]", + droppedGroups, MaxSessionFileBytes, + ), + } + insertAt := 0 + if len(sess.Messages) > 0 && sess.Messages[0].Role == "system" { + insertAt = 1 + } + withMarker := make([]llm.Message, 0, len(sess.Messages)+1) + withMarker = append(withMarker, sess.Messages[:insertAt]...) + withMarker = append(withMarker, marker) + withMarker = append(withMarker, sess.Messages[insertAt:]...) + + candidate := *sess + candidate.Messages = withMarker + markerData, err := json.Marshal(&candidate) + if err != nil { + return nil, fmt.Errorf("session: marshal trim marker: %w", err) + } + // Keep the marker only if the transcript still fits the cap. + if len(markerData) <= MaxSessionFileBytes { + sess.Messages = withMarker + data = markerData + } + } return data, nil } diff --git a/internal/session/session_savecap_test.go b/internal/session/session_savecap_test.go index 27542650..5f522923 100644 --- a/internal/session/session_savecap_test.go +++ b/internal/session/session_savecap_test.go @@ -1,6 +1,7 @@ package session import ( + "encoding/json" "os" "path/filepath" "strings" @@ -110,8 +111,16 @@ func TestTrimToFileCap_DropsToolGroups(t *testing.T) { if len(data) > MaxSessionFileBytes { t.Errorf("len(data) = %d, exceeds cap %d", len(data), MaxSessionFileBytes) } - if len(sess.Messages) != 2 || sess.Messages[0].Role != "system" || sess.Messages[1].Role != "user" { - t.Errorf("system + trailing user message should survive: %+v", sess.Messages) + // system + trim marker + trailing user message should survive. + if len(sess.Messages) != 3 || sess.Messages[0].Role != "system" || sess.Messages[2].Role != "user" { + t.Errorf("system + marker + trailing user message should survive: %+v", sess.Messages) + } + marker := sess.Messages[1] + if marker.Role != "system" || !strings.Contains(marker.Content, "[Session storage limit:") { + t.Errorf("expected storage-limit trim marker at index 1, got %+v", marker) + } + if !strings.Contains(marker.Content, "1 oldest message group(s)") { + t.Errorf("marker should report 1 dropped group, got %q", marker.Content) } for _, m := range sess.Messages { if m.Role == "tool" { @@ -164,3 +173,46 @@ func TestSave_IncrementalRedaction(t *testing.T) { t.Errorf("RedactBoundary = %d, want %d after second save", reloaded.RedactBoundary, len(reloaded.Messages)) } } + +// TestTrimToFileCap_MarkerSkippedWhenItWouldExceedCap covers the edge where +// the trimmed transcript lands so close to the cap that inserting the trim +// marker would push it back over — the marker must be omitted rather than +// breaking the cap. +func TestTrimToFileCap_MarkerSkippedWhenItWouldExceedCap(t *testing.T) { + withFileCap(t, 64<<10) + store, err := NewStoreWithDir(t.TempDir()) + if err != nil { + t.Fatalf("NewStoreWithDir() error: %v", err) + } + sess := &Session{ + ID: "20260101-dddddddddddddddddddddddddddddddd", + Messages: []llm.Message{ + {Role: "system", Content: ""}, // sized below + {Role: "user", Content: "droppable question"}, + }, + } + // Size the protected system message so the post-trim transcript lands + // just under the cap, leaving no room for the ~230-byte marker. + base, err := json.Marshal(sess) + if err != nil { + t.Fatalf("marshal base session: %v", err) + } + sess.Messages[0].Content = strings.Repeat("s", MaxSessionFileBytes-len(base)-50) + + oversized := make([]byte, MaxSessionFileBytes+1) + data, err := store.trimToFileCapLocked(sess, oversized) + if err != nil { + t.Fatalf("trimToFileCapLocked() error: %v", err) + } + if len(data) > MaxSessionFileBytes { + t.Errorf("len(data) = %d, exceeds cap %d", len(data), MaxSessionFileBytes) + } + if len(sess.Messages) != 1 || sess.Messages[0].Role != "system" { + t.Errorf("only the protected system message should survive: %+v", sess.Messages) + } + for _, m := range sess.Messages { + if strings.Contains(m.Content, "[Session storage limit:") { + t.Error("marker must be omitted when it would push the transcript over the cap") + } + } +} diff --git a/internal/session/session_test.go b/internal/session/session_test.go index a1bfe057..265e7152 100644 --- a/internal/session/session_test.go +++ b/internal/session/session_test.go @@ -1335,8 +1335,21 @@ func TestSave_TrimsOversizedSession(t *testing.T) { if last.Role != "assistant" || last.Content != "final answer" { t.Errorf("trailing message not preserved: role=%q content=%q", last.Role, last.Content) } - if len(loaded.Messages) >= len(msgs) { - t.Errorf("expected oldest messages to be dropped, still have %d of %d", len(loaded.Messages), len(msgs)) + bigLeft := 0 + markerFound := false + for _, m := range loaded.Messages { + if m.Content == big { + bigLeft++ + } + if strings.Contains(m.Content, "[Session storage limit:") { + markerFound = true + } + } + if bigLeft >= 4 { + t.Errorf("expected oldest big messages to be dropped, still have %d of 4", bigLeft) + } + if !markerFound { + t.Error("expected a [Session storage limit: ...] marker in the trimmed transcript") } if got, want := loaded.Turns, countUserTurns(loaded.Messages); got != want { t.Errorf("Turns = %d, want %d (recounted after trim)", got, want) diff --git a/odek.go b/odek.go index 10419547..50823825 100644 --- a/odek.go +++ b/odek.go @@ -213,6 +213,13 @@ type Config struct { // nil, skill/episode content is injected directly (not recommended for // production surfaces). UntrustedWrapper func(source, content string) string + + // Compaction enables LLM-based rolling compaction (default: false). When + // enabled, conversation turn groups dropped by context trimming are + // summarized by the model into a rolling digest system message instead of + // vanishing entirely, so long sessions retain a compressed memory of + // earlier work. Each compaction costs one extra LLM call per trim. + Compaction bool } // Agent is the agent loop runtime. @@ -578,6 +585,7 @@ func New(cfg Config) (*Agent, error) { engine := loop.New(client, registry, cfg.MaxIterations, cfg.SystemMessage, cfg.Renderer, maxContext) engine.PromptCaching = cfg.PromptCaching + engine.SetCompaction(cfg.Compaction) engine.SetUntrustedWrapper(cfg.UntrustedWrapper) if cfg.MaxToolParallel > 0 { engine.SetMaxToolParallel(cfg.MaxToolParallel) From 9655091634665770e99d6825004afe39e469617a Mon Sep 17 00:00:00 2001 From: Rolando Santamaria Maso Date: Sat, 25 Jul 2026 23:40:01 +0200 Subject: [PATCH 2/2] feat(docker): enable rolling compaction in bundled configs --- docker/config.godmode.json | 1 + docker/config.restricted.json | 1 + 2 files changed, 2 insertions(+) diff --git a/docker/config.godmode.json b/docker/config.godmode.json index d6facbb8..ebe303de 100644 --- a/docker/config.godmode.json +++ b/docker/config.godmode.json @@ -1,5 +1,6 @@ { "sandbox": false, + "compaction": true, "interaction_mode": "engaging", "tool_progress": "verbose", "tool_progress_cleanup": false, diff --git a/docker/config.restricted.json b/docker/config.restricted.json index 0b9549ad..4ed6fe07 100644 --- a/docker/config.restricted.json +++ b/docker/config.restricted.json @@ -1,5 +1,6 @@ { "sandbox": false, + "compaction": true, "interaction_mode": "engaging", "tool_progress": "verbose", "tool_progress_cleanup": false,