From b6f3a3392c3c1455df4b1940ec8b0041d4993d95 Mon Sep 17 00:00:00 2001 From: MK Date: Fri, 3 Jul 2026 16:34:07 -0400 Subject: [PATCH 01/16] feat(compress): reversible context compression (ctxzip) + provider prompt-cache hints MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Wires github.com/initializ/ctxzip into the agent loop as an opt-in feature (compression.enabled in forge.yaml, or FORGE_COMPRESSION=true). Bulky tool outputs and conversation content are compressed before reaching the LLM; everything dropped is stored in a durable local bbolt store (.forge/ctxzip.db) behind a <> marker, retrievable via the new context_expand tool — lossy on the wire, lossless end-to-end. New package forge-core/compress: - AfterToolExecHook — compresses tool output once at production time, before it enters Memory, so historic bytes never change and provider prompt caches keep hitting. Registered after guardrail/redaction hooks; error results and small outputs are left verbatim. - WrapClient — llm.Client decorator compressing the live zone of every outbound request (frozen prefix + recent turns forwarded byte-identical). Deterministic across turns: the relevance query is pinned to the first user message, never the latest turn. - ExpandTool — context_expand builtin retrieving originals by marker hash; registered only when compression is on (memory_get pattern). Provider prompt-cache hints (ClientConfig.PromptCaching, gated by compression.cache_hints, defaulting to compression.enabled): - anthropic: cache_control ephemeral breakpoints on the last tool definition and the system block (block-form system only when caching — wire format is byte-identical to the previous contract when off). Also applies on the aws_sigv4 path, which speaks the same Messages wire format. - openai: stable prompt_cache_key derived from (model, system, tool names) for cache-shard pinning; prefix caching itself is automatic. Config: CompressionConfig (enabled / store_path / ttl / min_tool_output_chars / cache_hints) following the MemoryConfig.LongTerm opt-in pattern, wired in the runner beside initLongTermMemory. Fail-open: any init error runs the agent uncompressed. Tests: hook compression + context_expand round-trip, error/small-output verbatim guarantees, live-zone vs frozen-prefix boundaries, cross-turn determinism (the cache-safety property), marker-hash normalization, anthropic/openai wire-format assertions on and off. Co-Authored-By: Claude Opus 4.8 (1M context) --- forge-cli/go.mod | 4 +- forge-cli/go.sum | 7 +- forge-cli/runtime/runner.go | 93 ++++++++ forge-core/compress/client.go | 99 +++++++++ forge-core/compress/compress.go | 96 +++++++++ forge-core/compress/compress_test.go | 204 ++++++++++++++++++ forge-core/compress/expand_tool.go | 89 ++++++++ forge-core/compress/hook.go | 59 +++++ forge-core/go.mod | 2 + forge-core/go.sum | 6 + forge-core/llm/client.go | 18 ++ forge-core/llm/providers/anthropic.go | 76 +++++-- forge-core/llm/providers/openai.go | 59 +++-- .../llm/providers/prompt_caching_test.go | 108 ++++++++++ forge-core/types/config.go | 25 +++ go.work.sum | 2 + 16 files changed, 918 insertions(+), 29 deletions(-) create mode 100644 forge-core/compress/client.go create mode 100644 forge-core/compress/compress.go create mode 100644 forge-core/compress/compress_test.go create mode 100644 forge-core/compress/expand_tool.go create mode 100644 forge-core/compress/hook.go create mode 100644 forge-core/llm/providers/prompt_caching_test.go diff --git a/forge-cli/go.mod b/forge-cli/go.mod index 46f30fbb..25c3346e 100644 --- a/forge-cli/go.mod +++ b/forge-cli/go.mod @@ -54,6 +54,7 @@ require ( github.com/gorilla/websocket v1.5.4-0.20250319132907-e064f32e3674 // indirect github.com/grpc-ecosystem/grpc-gateway/v2 v2.29.0 // indirect github.com/inconshreveable/mousetrap v1.1.0 // indirect + github.com/initializ/ctxzip v0.0.0-20260703193602-ca7fa8003af9 // indirect github.com/josharian/intern v1.0.0 // indirect github.com/json-iterator/go v1.1.12 // indirect github.com/klauspost/compress v1.16.7 // indirect @@ -71,7 +72,7 @@ require ( github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect github.com/pkg/errors v0.9.1 // indirect github.com/rivo/uniseg v0.4.7 // indirect - github.com/spf13/pflag v1.0.9 // indirect + github.com/spf13/pflag v1.0.10 // indirect github.com/x448/float16 v0.8.4 // indirect github.com/xdg-go/pbkdf2 v1.0.0 // indirect github.com/xdg-go/scram v1.1.2 // indirect @@ -81,6 +82,7 @@ require ( github.com/xeipuuv/gojsonschema v1.2.0 // indirect github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e // indirect github.com/youmark/pkcs8 v0.0.0-20240726163527-a2c0da244d78 // indirect + go.etcd.io/bbolt v1.5.0 // indirect go.opentelemetry.io/auto/sdk v1.2.1 // indirect go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.59.0 // indirect go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.44.0 // indirect diff --git a/forge-cli/go.sum b/forge-cli/go.sum index 1878cdfe..df325a6f 100644 --- a/forge-cli/go.sum +++ b/forge-cli/go.sum @@ -77,6 +77,8 @@ github.com/grpc-ecosystem/grpc-gateway/v2 v2.29.0 h1:5VipnvEpbqr2gA2VbM+nYVbkIF2 github.com/grpc-ecosystem/grpc-gateway/v2 v2.29.0/go.mod h1:Hyl3n6Twe1hvtd9XUXDec4pTvgMSEixRuQKPTMH2bNs= github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8= github.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw= +github.com/initializ/ctxzip v0.0.0-20260703193602-ca7fa8003af9 h1:a3rJaFaP+M3hERbrqAnnA5WOKYBMYuCP12wpHAMtLt8= +github.com/initializ/ctxzip v0.0.0-20260703193602-ca7fa8003af9/go.mod h1:Qs4CpXTEeEOlZwIMPm/qaFPmil+AUIqH7Jg/K81FTmU= github.com/initializ/guardrails v0.12.0 h1:YzScnl+YPihLA1gepjxKjnKH+2EbYmGC4dUtpVRAJ2c= github.com/initializ/guardrails v0.12.0/go.mod h1:bDdHx73MF0+O09KqoXmUTTiFG4H7yEVQ0NR6juP1F3Q= github.com/josharian/intern v1.0.0 h1:vlS4z54oSdjm0bgjRigI+G1HpF+tI+9rE5LLzOg8HmY= @@ -135,8 +137,9 @@ github.com/rogpeppe/go-internal v1.14.1/go.mod h1:MaRKkUm5W0goXpeCfT7UZI6fk/L7L7 github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM= github.com/spf13/cobra v1.10.2 h1:DMTTonx5m65Ic0GOoRY2c16WCbHxOOw6xxezuLaBpcU= github.com/spf13/cobra v1.10.2/go.mod h1:7C1pvHqHw5A4vrJfjNwvOdzYu0Gml16OCs2GRiTUUS4= -github.com/spf13/pflag v1.0.9 h1:9exaQaMOCwffKiiiYk6/BndUBv+iRViNW+4lEMi0PvY= github.com/spf13/pflag v1.0.9/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg= +github.com/spf13/pflag v1.0.10 h1:4EBh2KAYBwaONj6b2Ye1GiHfwjqyROoF4RwYO+vPwFk= +github.com/spf13/pflag v1.0.10/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg= github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw= github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo= @@ -169,6 +172,8 @@ github.com/youmark/pkcs8 v0.0.0-20240726163527-a2c0da244d78/go.mod h1:aL8wCCfTfS github.com/yuin/goldmark v1.1.27/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74= github.com/yuin/goldmark v1.2.1/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74= github.com/yuin/goldmark v1.4.13/go.mod h1:6yULJ656Px+3vBD8DxQVa3kxgyrAnzto9xy5taEt/CY= +go.etcd.io/bbolt v1.5.0 h1:S7GAl7Fxv12yohbwFfIbQCGDWbQbtDGPET4P/bD4lxU= +go.etcd.io/bbolt v1.5.0/go.mod h1:mkltfYE5aUHQxUct9N9V+Kp7aSjFqjgrhcXIS70Lrdk= go.mongodb.org/mongo-driver v1.17.7 h1:a9w+U3Vt67eYzcfq3k/OAv284/uUUkL0uP75VE5rCOU= go.mongodb.org/mongo-driver v1.17.7/go.mod h1:Hy04i7O2kC4RS06ZrhPRqj/u4DTYkFDAAccj+rVKqgQ= go.opentelemetry.io/auto/sdk v1.2.1 h1:jXsnJ4Lmnqd11kwkBV2LgLoFMZKizbCi5fNZ/ipaZ64= diff --git a/forge-cli/runtime/runner.go b/forge-cli/runtime/runner.go index f2d28928..94f825a5 100644 --- a/forge-cli/runtime/runner.go +++ b/forge-cli/runtime/runner.go @@ -29,6 +29,7 @@ import ( "github.com/initializ/forge/forge-core/auth/providers/httpverifier" _ "github.com/initializ/forge/forge-core/auth/providers/oidc" "github.com/initializ/forge/forge-core/auth/providers/statictoken" + "github.com/initializ/forge/forge-core/compress" "github.com/initializ/forge/forge-core/llm" "github.com/initializ/forge/forge-core/llm/oauth" "github.com/initializ/forge/forge-core/llm/providers" @@ -890,6 +891,17 @@ func (r *Runner) Run(ctx context.Context) error { r.registerSkillGuardrailHooks(hooks, sg) } + // Reversible context compression (ctxzip). The hook is + // registered AFTER guardrail/redaction hooks so it + // compresses what redaction left; the client wrapper sits + // below the FallbackChain so it also covers retries and + // compactor summarization calls. + if comp := r.initCompression(reg); comp != nil { + defer comp.Close() //nolint:errcheck + hooks.Register(coreruntime.AfterToolExec, comp.AfterToolExecHook()) + llmClient = comp.WrapClient(llmClient) + } + // Compute model-aware character budget. charBudget := r.cfg.Config.Memory.CharBudget if charBudget == 0 { @@ -2355,6 +2367,11 @@ func (r *Runner) buildLLMClient(mc *coreruntime.ModelConfig) (llm.Client, error) // OPENAI_API_KEY for that endpoint; if it's missing, surface the // configuration error rather than tunneling to ChatGPT. func (r *Runner) createProviderClient(provider string, cfg llm.ClientConfig) (llm.Client, error) { + // Provider prompt-cache hints (anthropic cache_control breakpoints / + // openai prompt_cache_key). Off by default; opted in via forge.yaml + // compression.cache_hints (or compression.enabled). See ClientConfig. + cfg.PromptCaching = r.promptCachingEnabled() + // Check for stored OAuth credentials — but only if no real API key is // configured. The "__oauth__" sentinel means the user chose OAuth auth // during init, so we should load the actual token from the credential store. @@ -3363,6 +3380,82 @@ func (r *Runner) initLongTermMemory(ctx context.Context, mc *coreruntime.ModelCo return mgr } +// compressionEnabled reports whether reversible context compression is on. +// forge.yaml `compression.enabled` sets the default; FORGE_COMPRESSION=true / +// =false overrides it either way (matching the FORGE_MEMORY_* env pattern). +// Nil-safe: callers (createProviderClient) also run with bare Runners in tests. +func (r *Runner) compressionEnabled() bool { + enabled := false + if r.cfg.Config != nil && r.cfg.Config.Compression.Enabled != nil { + enabled = *r.cfg.Config.Compression.Enabled + } + switch os.Getenv("FORGE_COMPRESSION") { + case "true": + enabled = true + case "false": + enabled = false + } + return enabled +} + +// promptCachingEnabled reports whether provider prompt-cache hints should be +// injected. Defaults to compressionEnabled(); compression.cache_hints +// overrides explicitly in either direction. +func (r *Runner) promptCachingEnabled() bool { + if r.cfg.Config != nil && r.cfg.Config.Compression.CacheHints != nil { + return *r.cfg.Config.Compression.CacheHints + } + return r.compressionEnabled() +} + +// initCompression builds the ctxzip compression runtime and registers the +// context_expand tool. Returns nil when compression is disabled or the store +// cannot be opened (the agent then runs uncompressed — fail-open). +func (r *Runner) initCompression(reg *tools.Registry) *compress.Runtime { + if !r.compressionEnabled() { + return nil + } + + cc := r.cfg.Config.Compression + storePath := cc.StorePath + if storePath == "" { + storePath = filepath.Join(r.cfg.WorkDir, ".forge", "ctxzip.db") + } + var ttl time.Duration + if cc.TTL != "" { + if d, err := time.ParseDuration(cc.TTL); err == nil { + ttl = d + } else { + r.logger.Warn("invalid compression.ttl, using default", map[string]any{ + "ttl": cc.TTL, "error": err.Error(), + }) + } + } + + comp, err := compress.New(compress.Config{ + StorePath: storePath, + TTL: ttl, + MinToolOutputChars: cc.MinToolOutputChars, + Logger: r.logger, + }) + if err != nil { + r.logger.Warn("failed to init compression, running uncompressed", map[string]any{ + "error": err.Error(), + }) + return nil + } + + if regErr := reg.Register(comp.ExpandTool()); regErr != nil { + r.logger.Warn("failed to register context_expand tool", map[string]any{"error": regErr.Error()}) + } + + r.logger.Info("context compression enabled", map[string]any{ + "store_path": storePath, + "cache_hints": r.promptCachingEnabled(), + }) + return comp +} + // resolveEmbedder creates an embedder from config or auto-detection. // Returns nil if no embedder can be created (keyword-only mode). func (r *Runner) resolveEmbedder(mc *coreruntime.ModelConfig) llm.Embedder { diff --git a/forge-core/compress/client.go b/forge-core/compress/client.go new file mode 100644 index 00000000..a1a3b295 --- /dev/null +++ b/forge-core/compress/client.go @@ -0,0 +1,99 @@ +package compress + +import ( + "context" + + "github.com/initializ/ctxzip" + + "github.com/initializ/forge/forge-core/llm" +) + +// WrapClient decorates an llm.Client so every outbound request has its live +// zone compressed. It sits below the FallbackChain, so it also covers retry +// calls and the compactor's summarization call. +// +// Cache discipline ("passthrough is sacred"): only the live zone is touched — +// the system prompt (frozen prefix) and the most recent turns are forwarded +// byte-identical. Determinism matters just as much: the relevance query is +// pinned to the FIRST user message of the conversation, never the latest +// turn. Deriving it from the latest turn would recompress the same historic +// message to different bytes each turn and bust the provider prompt cache. +func (r *Runtime) WrapClient(inner llm.Client) llm.Client { + return &compressingClient{inner: inner, rt: r} +} + +type compressingClient struct { + inner llm.Client + rt *Runtime +} + +// Chat implements llm.Client. +func (c *compressingClient) Chat(ctx context.Context, req *llm.ChatRequest) (*llm.ChatResponse, error) { + return c.inner.Chat(ctx, c.compressRequest(req)) +} + +// ChatStream implements llm.Client. +func (c *compressingClient) ChatStream(ctx context.Context, req *llm.ChatRequest) (<-chan llm.StreamDelta, error) { + return c.inner.ChatStream(ctx, c.compressRequest(req)) +} + +// ModelID implements llm.Client. +func (c *compressingClient) ModelID() string { return c.inner.ModelID() } + +// compressRequest returns req with compressed live-zone messages, or req +// unchanged when there is nothing to gain. The caller's request is never +// mutated — the loop may reuse its message slice. +func (c *compressingClient) compressRequest(req *llm.ChatRequest) *llm.ChatRequest { + if req == nil || len(req.Messages) == 0 { + return req + } + + zmsgs := make([]ctxzip.Message, len(req.Messages)) + for i, m := range req.Messages { + zmsgs[i] = ctxzip.Message{ + Role: m.Role, + Content: m.Content, + Name: m.Name, + ToolCallID: m.ToolCallID, + } + } + + opts := ctxzip.DefaultOptions() + opts.Store = c.rt.store + opts.Query = firstUserContent(req.Messages) + + res, err := ctxzip.Compress(zmsgs, opts) + if err != nil || res == nil || res.SavedTokens() == 0 { + return req + } + + out := *req + out.Messages = make([]llm.ChatMessage, len(req.Messages)) + copy(out.Messages, req.Messages) + for _, tr := range res.Transforms { + out.Messages[tr.Index].Content = res.Messages[tr.Index].Content + } + + c.rt.debugf("compressed request", map[string]any{ + "messages": len(req.Messages), + "transformed": len(res.Transforms), + "tokens_before": res.TokensBefore, + "tokens_after": res.TokensAfter, + "saved_tokens": res.SavedTokens(), + }) + return &out +} + +// firstUserContent returns the first user message — the pinned task +// statement. It is stable for the whole session, which keeps compression +// deterministic across turns (see WrapClient). Returning a single space when +// absent suppresses ctxzip's derive-from-recent-messages fallback, which +// would reintroduce turn-varying output. +func firstUserContent(msgs []llm.ChatMessage) string { + for _, m := range msgs { + if m.Role == llm.RoleUser { + return m.Content + } + } + return " " +} diff --git a/forge-core/compress/compress.go b/forge-core/compress/compress.go new file mode 100644 index 00000000..c79c66a6 --- /dev/null +++ b/forge-core/compress/compress.go @@ -0,0 +1,96 @@ +// Package compress wires reversible context compression (ctxzip) into the +// forge agent loop. +// +// ctxzip shrinks bulky content before it reaches the LLM — tool outputs, +// logs, JSON — and offloads everything it drops to a durable local store, +// replaced inline by a retrievable "<>" marker. Compression +// is lossy on the wire but lossless end-to-end: the model can recover any +// original via the context_expand tool. +// +// Three integration seams, all owned by Runtime: +// +// - AfterToolExecHook — compresses tool output once, at production time, +// before it enters Memory. Because the compressed bytes never change +// afterwards, the conversation prefix stays byte-stable across turns and +// provider prompt caches keep hitting. +// - WrapClient — an llm.Client decorator that compresses the live zone of +// every outbound request (skipping the frozen prefix and recent turns). +// It is deliberately deterministic: the relevance query is pinned to the +// first user message of the session, never the latest turn, so the same +// historic message always compresses to the same bytes. +// - ExpandTool — the context_expand builtin that retrieves originals from +// the store by marker hash. +// +// The store is a bbolt file (default .forge/ctxzip.db) so originals survive +// process restarts; entries expire after TTL, at which point the disk or the +// original command is the source of truth (the tool's miss message says so). +package compress + +import ( + "fmt" + "time" + + "github.com/initializ/ctxzip/ccr" + + "github.com/initializ/forge/forge-core/runtime" +) + +// Defaults for Config fields left at their zero value. +const ( + DefaultTTL = 30 * time.Minute + DefaultMinToolOutputChars = 2048 +) + +// Config configures the compression runtime. +type Config struct { + // StorePath is the bbolt file for the CCR store (required), + // e.g. .forge/ctxzip.db. + StorePath string + // TTL is how long offloaded originals stay retrievable. Default 30m. + TTL time.Duration + // MinToolOutputChars is the size below which tool outputs are left + // alone by the AfterToolExec hook. Default 2048. + MinToolOutputChars int + // Logger is optional; nil disables logging. + Logger runtime.Logger +} + +// Runtime owns the shared CCR store and produces the hook, client wrapper, +// and expand tool that plug into the agent loop. +type Runtime struct { + store *ccr.BoltStore + minSize int + logger runtime.Logger +} + +// New opens the durable store and returns a Runtime. Call Close on shutdown. +func New(cfg Config) (*Runtime, error) { + if cfg.StorePath == "" { + return nil, fmt.Errorf("compress: Config.StorePath is required") + } + if cfg.TTL <= 0 { + cfg.TTL = DefaultTTL + } + if cfg.MinToolOutputChars <= 0 { + cfg.MinToolOutputChars = DefaultMinToolOutputChars + } + store, err := ccr.NewBoltStore(ccr.BoltConfig{Path: cfg.StorePath, TTL: cfg.TTL}) + if err != nil { + return nil, fmt.Errorf("compress: opening store: %w", err) + } + return &Runtime{store: store, minSize: cfg.MinToolOutputChars, logger: cfg.Logger}, nil +} + +// Close releases the underlying store. +func (r *Runtime) Close() error { + return r.store.Close() +} + +// Store exposes the CCR store (used by tests and diagnostics). +func (r *Runtime) Store() ccr.Store { return r.store } + +func (r *Runtime) debugf(msg string, fields map[string]any) { + if r.logger != nil { + r.logger.Debug(msg, fields) + } +} diff --git a/forge-core/compress/compress_test.go b/forge-core/compress/compress_test.go new file mode 100644 index 00000000..11939c7a --- /dev/null +++ b/forge-core/compress/compress_test.go @@ -0,0 +1,204 @@ +package compress + +import ( + "context" + "encoding/json" + "fmt" + "path/filepath" + "strings" + "testing" + + "github.com/initializ/ctxzip/ccr" + + "github.com/initializ/forge/forge-core/llm" + "github.com/initializ/forge/forge-core/runtime" +) + +func newRuntime(t *testing.T) *Runtime { + t.Helper() + rt, err := New(Config{StorePath: filepath.Join(t.TempDir(), "ctxzip.db")}) + if err != nil { + t.Fatalf("New: %v", err) + } + t.Cleanup(func() { _ = rt.Close() }) + return rt +} + +// bigJSON builds a large JSON-array tool output with one error row. +func bigJSON(n int) string { + items := make([]map[string]any, n) + for i := range items { + items[i] = map[string]any{"name": fmt.Sprintf("pod-%03d", i), "status": "Running"} + } + items[n/2] = map[string]any{"name": "pod-bad", "status": "CrashLoopBackOff", "error": "OOMKilled"} + b, _ := json.Marshal(items) + return string(b) +} + +func TestHook_CompressesToolOutput_AndExpandRoundTrips(t *testing.T) { + rt := newRuntime(t) + hook := rt.AfterToolExecHook() + + original := bigJSON(80) + hctx := &runtime.HookContext{ToolName: "list_pods", ToolInput: `{"ns":"default"}`, ToolOutput: original} + if err := hook(context.Background(), hctx); err != nil { + t.Fatalf("hook: %v", err) + } + if hctx.ToolOutput == original { + t.Fatal("tool output was not compressed") + } + if !strings.Contains(hctx.ToolOutput, "CrashLoopBackOff") { + t.Fatal("error row dropped — must-keep floor failed") + } + hashes := ccr.ExtractHashes(hctx.ToolOutput) + if len(hashes) == 0 { + t.Fatalf("no ctxzip marker in compressed output:\n%s", hctx.ToolOutput) + } + + // The model recovers the original through the context_expand tool. + tool := rt.ExpandTool() + args, _ := json.Marshal(map[string]string{"hash": hashes[0]}) + out, err := tool.Execute(context.Background(), args) + if err != nil { + t.Fatalf("context_expand: %v", err) + } + var rows []json.RawMessage + if err := json.Unmarshal([]byte(out), &rows); err != nil || len(rows) == 0 { + t.Fatalf("expanded content invalid: err=%v rows=%d", err, len(rows)) + } +} + +func TestHook_SkipsErrorsAndSmallOutput(t *testing.T) { + rt := newRuntime(t) + hook := rt.AfterToolExecHook() + + // Error results stay verbatim regardless of size. + errOut := bigJSON(80) + hctx := &runtime.HookContext{ToolName: "t", ToolOutput: errOut, Error: fmt.Errorf("boom")} + _ = hook(context.Background(), hctx) + if hctx.ToolOutput != errOut { + t.Error("error output must not be compressed") + } + + // Small outputs stay verbatim. + small := &runtime.HookContext{ToolName: "t", ToolOutput: "tiny result"} + _ = hook(context.Background(), small) + if small.ToolOutput != "tiny result" { + t.Error("small output must not be compressed") + } +} + +// capturingClient records the request it receives and returns a stub response. +type capturingClient struct { + lastReq *llm.ChatRequest +} + +func (c *capturingClient) Chat(_ context.Context, req *llm.ChatRequest) (*llm.ChatResponse, error) { + c.lastReq = req + return &llm.ChatResponse{Message: llm.ChatMessage{Role: llm.RoleAssistant, Content: "ok"}}, nil +} + +func (c *capturingClient) ChatStream(_ context.Context, req *llm.ChatRequest) (<-chan llm.StreamDelta, error) { + c.lastReq = req + ch := make(chan llm.StreamDelta) + close(ch) + return ch, nil +} + +func (c *capturingClient) ModelID() string { return "test-model" } + +func conversation(toolOutput string) []llm.ChatMessage { + return []llm.ChatMessage{ + {Role: llm.RoleSystem, Content: "You are a k8s assistant."}, + {Role: llm.RoleUser, Content: "check the pods"}, + {Role: llm.RoleTool, Name: "list_pods", ToolCallID: "tc1", Content: toolOutput}, + {Role: llm.RoleAssistant, Content: "Looking at the pods now."}, + {Role: llm.RoleUser, Content: "and the crashing one?"}, + } +} + +func TestWrapClient_CompressesLiveZone_PreservesPrefixAndRecent(t *testing.T) { + rt := newRuntime(t) + inner := &capturingClient{} + client := rt.WrapClient(inner) + + original := bigJSON(80) + msgs := conversation(original) + req := &llm.ChatRequest{Model: "m", Messages: msgs} + + if _, err := client.Chat(context.Background(), req); err != nil { + t.Fatal(err) + } + sent := inner.lastReq + + if sent.Messages[0].Content != msgs[0].Content { + t.Error("frozen prefix (system) was modified") + } + if sent.Messages[3].Content != msgs[3].Content || sent.Messages[4].Content != msgs[4].Content { + t.Error("protected recent turns were modified") + } + if sent.Messages[2].Content == original { + t.Error("live-zone tool output was NOT compressed") + } + if sent.Messages[2].ToolCallID != "tc1" || sent.Messages[2].Name != "list_pods" { + t.Error("tool message metadata lost in compression") + } + // Caller's request untouched. + if req.Messages[2].Content != original { + t.Error("caller's request was mutated") + } +} + +// Determinism is what keeps provider prompt caches hitting: the same +// conversation must compress to identical bytes on every call. +func TestWrapClient_Deterministic(t *testing.T) { + rt := newRuntime(t) + inner := &capturingClient{} + client := rt.WrapClient(inner) + + original := bigJSON(80) + req1 := &llm.ChatRequest{Model: "m", Messages: conversation(original)} + _, _ = client.Chat(context.Background(), req1) + first := inner.lastReq.Messages[2].Content + + // Same conversation, later turn appended (the tool message is deeper in + // history now, and the latest user message differs). + extended := append(conversation(original), + llm.ChatMessage{Role: llm.RoleAssistant, Content: "It is OOMKilled."}, + llm.ChatMessage{Role: llm.RoleUser, Content: "what about the disk usage on node-a?"}, + ) + req2 := &llm.ChatRequest{Model: "m", Messages: extended} + _, _ = client.Chat(context.Background(), req2) + second := inner.lastReq.Messages[2].Content + + if first != second { + t.Fatalf("compression not deterministic across turns — prompt cache would bust:\nfirst: %.120s\nsecond: %.120s", first, second) + } +} + +func TestNormalizeHash(t *testing.T) { + cases := map[string]string{ + "abc123def456": "abc123def456", + "<>": "abc123", + "ctxzip:abc123": "abc123", + "hash=abc123": "abc123", + " ABC123 ": "abc123", + } + for in, want := range cases { + if got := normalizeHash(in); got != want { + t.Errorf("normalizeHash(%q) = %q, want %q", in, got, want) + } + } +} + +func TestExpandTool_MissIsHelpful(t *testing.T) { + rt := newRuntime(t) + args, _ := json.Marshal(map[string]string{"hash": "deadbeefdeadbeefdeadbeef"}) + out, err := rt.ExpandTool().Execute(context.Background(), args) + if err != nil { + t.Fatalf("miss should not be an error: %v", err) + } + if !strings.Contains(out, "Re-run") { + t.Errorf("miss message should guide regeneration: %q", out) + } +} diff --git a/forge-core/compress/expand_tool.go b/forge-core/compress/expand_tool.go new file mode 100644 index 00000000..f76b3331 --- /dev/null +++ b/forge-core/compress/expand_tool.go @@ -0,0 +1,89 @@ +package compress + +import ( + "context" + "encoding/json" + "fmt" + "strings" + + "github.com/initializ/ctxzip" + + "github.com/initializ/forge/forge-core/tools" +) + +// expandTool is the context_expand builtin. When compression drops content it +// leaves a "<>" marker behind; the model calls +// this tool with the hash to get the original back. The agent loop executes +// it like any other tool — no special retrieval machinery is needed. +type expandTool struct { + rt *Runtime +} + +// ExpandTool returns the context_expand tool backed by this Runtime's store. +// Register it conditionally (like memory_get) — only when compression is on. +func (r *Runtime) ExpandTool() tools.Tool { + return &expandTool{rt: r} +} + +type expandInput struct { + Hash string `json:"hash"` +} + +func (t *expandTool) Name() string { return "context_expand" } + +func (t *expandTool) Description() string { + return "Retrieve the original content behind a <> compression marker. " + + "Earlier tool results may contain such markers where bulky content was compressed away; " + + "call this with the hash to see the full original data." +} + +func (t *expandTool) Category() tools.Category { return tools.CategoryBuiltin } + +func (t *expandTool) InputSchema() json.RawMessage { + return json.RawMessage(`{ + "type": "object", + "properties": { + "hash": {"type": "string", "description": "The hash from a <> marker (the marker text itself is also accepted)"} + }, + "required": ["hash"] + }`) +} + +func (t *expandTool) Execute(_ context.Context, args json.RawMessage) (string, error) { + var input expandInput + if err := json.Unmarshal(args, &input); err != nil { + return "", fmt.Errorf("parsing input: %w", err) + } + hash := normalizeHash(input.Hash) + if hash == "" { + return "", fmt.Errorf("hash is required") + } + + original, ok := ctxzip.Unzip(t.rt.store, hash) + if !ok { + // A miss is not a dead end — the disk or the original command is the + // source of truth. Say so instead of returning a bare error. + return fmt.Sprintf( + "No stored content for hash %s (expired or evicted). "+ + "Re-run the tool that produced the original output to regenerate it.", + hash, + ), nil + } + return string(original), nil +} + +// normalizeHash tolerates the model passing a whole marker instead of the +// bare hash: "<>", "ctxzip:abc123", or +// "hash=abc123" all normalize to "abc123". +func normalizeHash(s string) string { + s = strings.TrimSpace(s) + s = strings.TrimPrefix(s, "<<") + s = strings.TrimPrefix(s, "ctxzip:") + s = strings.TrimPrefix(s, "hash=") + s = strings.TrimSuffix(s, ">>") + // Keep only the leading token — markers carry a trailing note. + if i := strings.IndexAny(s, " ,"); i >= 0 { + s = s[:i] + } + return strings.ToLower(strings.TrimSpace(s)) +} diff --git a/forge-core/compress/hook.go b/forge-core/compress/hook.go new file mode 100644 index 00000000..1302ae01 --- /dev/null +++ b/forge-core/compress/hook.go @@ -0,0 +1,59 @@ +package compress + +import ( + "context" + + "github.com/initializ/ctxzip" + + "github.com/initializ/forge/forge-core/runtime" +) + +// AfterToolExecHook returns a hook that compresses tool output at production +// time, before the loop appends it to Memory. This is the primary compression +// seam: because the output is compressed exactly once, the bytes stored in +// history never change afterwards, keeping the conversation prefix stable for +// provider prompt caches. +// +// Register it AFTER redaction/guardrail hooks so it compresses what those +// hooks left, not what they were about to remove. +// +// The hook never fails the loop: on any problem it leaves ToolOutput as-is. +// Error results (hctx.Error != nil) are always left verbatim — dropping parts +// of an error the user is about to debug is the catastrophic failure mode. +func (r *Runtime) AfterToolExecHook() runtime.Hook { + return func(_ context.Context, hctx *runtime.HookContext) error { + if hctx.Error != nil || len(hctx.ToolOutput) < r.minSize { + return nil + } + + opts := ctxzip.DefaultOptions() + opts.Store = r.store + // A single fresh message: no prefix to freeze, nothing "recent" to + // protect — those windows exist for whole conversations. + opts.FreezePrefix = 0 + opts.ProtectRecent = 0 + // The tool-call arguments are the best available relevance signal for + // what the model wanted from this output — and they are fixed at call + // time, so recompression determinism is not a concern here. + opts.Query = hctx.ToolInput + + msgs := []ctxzip.Message{{ + Role: ctxzip.RoleTool, + Content: hctx.ToolOutput, + Name: hctx.ToolName, + }} + res, err := ctxzip.Compress(msgs, opts) + if err != nil || res == nil || res.SavedTokens() == 0 { + return nil + } + + r.debugf("compressed tool output", map[string]any{ + "tool": hctx.ToolName, + "tokens_before": res.TokensBefore, + "tokens_after": res.TokensAfter, + "saved_tokens": res.SavedTokens(), + }) + hctx.ToolOutput = res.Messages[0].Content + return nil + } +} diff --git a/forge-core/go.mod b/forge-core/go.mod index c13dc448..30fe5a68 100644 --- a/forge-core/go.mod +++ b/forge-core/go.mod @@ -4,6 +4,7 @@ go 1.25.0 require ( github.com/golang-jwt/jwt/v5 v5.3.1 + github.com/initializ/ctxzip v0.0.0-20260703193602-ca7fa8003af9 github.com/initializ/forge/forge-skills v0.0.0 github.com/xeipuuv/gojsonschema v1.2.0 go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.59.0 @@ -27,6 +28,7 @@ require ( github.com/grpc-ecosystem/grpc-gateway/v2 v2.29.0 // indirect github.com/xeipuuv/gojsonpointer v0.0.0-20180127040702-4e3ac2762d5f // indirect github.com/xeipuuv/gojsonreference v0.0.0-20180127040603-bd5ef7bd5415 // indirect + go.etcd.io/bbolt v1.5.0 // indirect go.opentelemetry.io/auto/sdk v1.2.1 // indirect go.opentelemetry.io/otel/metric v1.44.0 // indirect go.opentelemetry.io/proto/otlp v1.10.0 // indirect diff --git a/forge-core/go.sum b/forge-core/go.sum index 569e02ee..51191df4 100644 --- a/forge-core/go.sum +++ b/forge-core/go.sum @@ -22,6 +22,8 @@ github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= github.com/grpc-ecosystem/grpc-gateway/v2 v2.29.0 h1:5VipnvEpbqr2gA2VbM+nYVbkIF28c5ZQfqCBQ5g2xfk= github.com/grpc-ecosystem/grpc-gateway/v2 v2.29.0/go.mod h1:Hyl3n6Twe1hvtd9XUXDec4pTvgMSEixRuQKPTMH2bNs= +github.com/initializ/ctxzip v0.0.0-20260703193602-ca7fa8003af9 h1:a3rJaFaP+M3hERbrqAnnA5WOKYBMYuCP12wpHAMtLt8= +github.com/initializ/ctxzip v0.0.0-20260703193602-ca7fa8003af9/go.mod h1:Qs4CpXTEeEOlZwIMPm/qaFPmil+AUIqH7Jg/K81FTmU= github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE= github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk= github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= @@ -40,6 +42,8 @@ github.com/xeipuuv/gojsonreference v0.0.0-20180127040603-bd5ef7bd5415 h1:EzJWgHo github.com/xeipuuv/gojsonreference v0.0.0-20180127040603-bd5ef7bd5415/go.mod h1:GwrjFmJcFw6At/Gs6z4yjiIwzuJ1/+UwLxMQDVQXShQ= github.com/xeipuuv/gojsonschema v1.2.0 h1:LhYJRs+L4fBtjZUfuSZIKGeVu0QRy8e5Xi7D17UxZ74= github.com/xeipuuv/gojsonschema v1.2.0/go.mod h1:anYRn/JVcOK2ZgGU+IjEV4nwlhoK5sQluxsYJ78Id3Y= +go.etcd.io/bbolt v1.5.0 h1:S7GAl7Fxv12yohbwFfIbQCGDWbQbtDGPET4P/bD4lxU= +go.etcd.io/bbolt v1.5.0/go.mod h1:mkltfYE5aUHQxUct9N9V+Kp7aSjFqjgrhcXIS70Lrdk= go.opentelemetry.io/auto/sdk v1.2.1 h1:jXsnJ4Lmnqd11kwkBV2LgLoFMZKizbCi5fNZ/ipaZ64= go.opentelemetry.io/auto/sdk v1.2.1/go.mod h1:KRTj+aOaElaLi+wW1kO/DZRXwkF4C5xPbEe3ZiIhN7Y= go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.59.0 h1:CV7UdSGJt/Ao6Gp4CXckLxVRRsRgDHoI8XjbL3PDl8s= @@ -68,6 +72,8 @@ golang.org/x/crypto v0.51.0 h1:IBPXwPfKxY7cWQZ38ZCIRPI50YLeevDLlLnyC5wRGTI= golang.org/x/crypto v0.51.0/go.mod h1:8AdwkbraGNABw2kOX6YFPs3WM22XqI4EXEd8g+x7Oc8= golang.org/x/net v0.55.0 h1:bcvxaJn3e1U6InsFWt1JUq1aSjnRxLzT2rtD2KfkDF8= golang.org/x/net v0.55.0/go.mod h1:L5U2KuzuOe1lY7Z+aWVIKK6qEeJXnXV9yzGA+WCHJww= +golang.org/x/sync v0.20.0 h1:e0PTpb7pjO8GAtTs2dQ6jYa5BWYlMuX047Dco/pItO4= +golang.org/x/sync v0.20.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0= golang.org/x/sys v0.45.0 h1:dO4czNzziLiiXplLQgBCEpCvXQ3dnkn0SdaZSYdQ+FY= golang.org/x/sys v0.45.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= golang.org/x/text v0.37.0 h1:Cqjiwd9eSg8e0QAkyCaQTNHFIIzWtidPahFWR83rTrc= diff --git a/forge-core/llm/client.go b/forge-core/llm/client.go index 15493237..e995e210 100644 --- a/forge-core/llm/client.go +++ b/forge-core/llm/client.go @@ -34,4 +34,22 @@ type ClientConfig struct { // signer and skips the native header logic; APIKey is ignored. AuthScheme string AWSRegion string + + // PromptCaching opts the provider client into injecting the + // provider's prompt-cache primitives on every request: + // + // - anthropic: a cache_control ephemeral breakpoint on the last + // tool definition and on the system prompt block, caching the + // stable tools+system prefix across turns. Also honored by + // Anthropic-on-Bedrock gateways (aws_sigv4), which speak the + // same wire format. + // - openai: a stable prompt_cache_key derived from + // (model, system, tool names), pinning cache routing for the + // session. OpenAI prefix caching itself is automatic ≥1024 + // tokens; the key improves hit locality. + // + // Off by default — wire formats stay byte-identical to the + // pre-compression contract unless the operator opts in + // (compression.cache_hints / compression.enabled in forge.yaml). + PromptCaching bool } diff --git a/forge-core/llm/providers/anthropic.go b/forge-core/llm/providers/anthropic.go index d001d470..79a8c26f 100644 --- a/forge-core/llm/providers/anthropic.go +++ b/forge-core/llm/providers/anthropic.go @@ -16,11 +16,12 @@ import ( // AnthropicClient implements llm.Client for the Anthropic Messages API. type AnthropicClient struct { - apiKey string - baseURL string - model string - authScheme string - client *http.Client + apiKey string + baseURL string + model string + authScheme string + promptCaching bool + client *http.Client } // NewAnthropicClient creates a new Anthropic client. @@ -51,11 +52,12 @@ func NewAnthropicClient(cfg llm.ClientConfig) *AnthropicClient { httpClient.Transport = newBedrockSigningTransport(cfg.AWSRegion, http.DefaultTransport) } return &AnthropicClient{ - apiKey: cfg.APIKey, - baseURL: strings.TrimRight(baseURL, "/"), - model: cfg.Model, - authScheme: cfg.AuthScheme, - client: httpClient, + apiKey: cfg.APIKey, + baseURL: strings.TrimRight(baseURL, "/"), + model: cfg.Model, + authScheme: cfg.AuthScheme, + promptCaching: cfg.PromptCaching, + client: httpClient, } } @@ -158,15 +160,36 @@ func (c *AnthropicClient) setHeaders(req *http.Request) { } // Anthropic-specific request types. +// +// System is `any` because the Messages API accepts either a plain string or +// an array of content blocks. Without prompt caching it stays a string — +// byte-identical to the pre-caching wire format; with prompt caching it +// becomes []anthropicSystemBlock so a cache_control breakpoint can be +// attached. type anthropicRequest struct { Model string `json:"model"` Messages []anthropicMessage `json:"messages"` - System string `json:"system,omitempty"` + System any `json:"system,omitempty"` MaxTokens int `json:"max_tokens"` Tools []anthropicTool `json:"tools,omitempty"` Stream bool `json:"stream,omitempty"` } +// anthropicCacheControl marks a prompt-cache breakpoint. Everything up to and +// including the marked block (in Anthropic's tools → system → messages cache +// order) is cached for ~5 minutes and re-billed at ~10% on hit. +type anthropicCacheControl struct { + Type string `json:"type"` // always "ephemeral" +} + +// anthropicSystemBlock is the block form of the system prompt, used only when +// prompt caching is enabled. +type anthropicSystemBlock struct { + Type string `json:"type"` // "text" + Text string `json:"text"` + CacheControl *anthropicCacheControl `json:"cache_control,omitempty"` +} + type anthropicMessage struct { Role string `json:"role"` Content json.RawMessage `json:"content"` @@ -183,9 +206,10 @@ type anthropicContentBlock struct { } type anthropicTool struct { - Name string `json:"name"` - Description string `json:"description"` - InputSchema json.RawMessage `json:"input_schema"` + Name string `json:"name"` + Description string `json:"description"` + InputSchema json.RawMessage `json:"input_schema"` + CacheControl *anthropicCacheControl `json:"cache_control,omitempty"` } func (c *AnthropicClient) toAnthropicRequest(req *llm.ChatRequest, stream bool) anthropicRequest { @@ -206,13 +230,17 @@ func (c *AnthropicClient) toAnthropicRequest(req *llm.ChatRequest, stream bool) } // Extract system message and convert remaining messages + var system string for _, m := range req.Messages { if m.Role == llm.RoleSystem { - r.System = m.Content + system = m.Content continue } r.Messages = append(r.Messages, c.convertMessage(m)) } + if system != "" { + r.System = system + } // Convert tools for _, t := range req.Tools { @@ -223,6 +251,24 @@ func (c *AnthropicClient) toAnthropicRequest(req *llm.ChatRequest, stream bool) }) } + // Prompt-cache breakpoints (opt-in via ClientConfig.PromptCaching). + // Anthropic's cache prefix serializes tools → system → messages, so a + // breakpoint on the system block caches tools+system, and one on the + // last tool keeps the tools segment cached even when the system prompt + // churns. Forge's provider-agnostic types cannot express cache_control, + // so no caller placement can conflict with this injection. Tool order is + // already deterministic (Registry.ToolDefinitions sorts by name), which + // is what makes the cached prefix byte-stable across turns. + if c.promptCaching { + ephemeral := &anthropicCacheControl{Type: "ephemeral"} + if n := len(r.Tools); n > 0 { + r.Tools[n-1].CacheControl = ephemeral + } + if system != "" { + r.System = []anthropicSystemBlock{{Type: "text", Text: system, CacheControl: ephemeral}} + } + } + return r } diff --git a/forge-core/llm/providers/openai.go b/forge-core/llm/providers/openai.go index bc24da31..413410ea 100644 --- a/forge-core/llm/providers/openai.go +++ b/forge-core/llm/providers/openai.go @@ -5,6 +5,8 @@ import ( "bufio" "bytes" "context" + "crypto/sha256" + "encoding/hex" "encoding/json" "fmt" "io" @@ -18,12 +20,13 @@ import ( // OpenAIClient implements llm.Client for the OpenAI Chat Completions API. // Also works with Azure OpenAI and any OpenAI-compatible endpoint. type OpenAIClient struct { - apiKey string - baseURL string - model string - orgID string - authScheme string - client *http.Client + apiKey string + baseURL string + model string + orgID string + authScheme string + promptCaching bool + client *http.Client } // NewOpenAIClient creates a new OpenAI client. @@ -51,12 +54,13 @@ func NewOpenAIClient(cfg llm.ClientConfig) *OpenAIClient { httpClient.Transport = newBedrockSigningTransport(cfg.AWSRegion, http.DefaultTransport) } return &OpenAIClient{ - apiKey: cfg.APIKey, - baseURL: strings.TrimRight(baseURL, "/"), - model: cfg.Model, - orgID: cfg.OrgID, - authScheme: cfg.AuthScheme, - client: httpClient, + apiKey: cfg.APIKey, + baseURL: strings.TrimRight(baseURL, "/"), + model: cfg.Model, + orgID: cfg.OrgID, + authScheme: cfg.AuthScheme, + promptCaching: cfg.PromptCaching, + client: httpClient, } } @@ -147,6 +151,11 @@ type openaiRequest struct { MaxTokens int `json:"max_tokens,omitempty"` Stream bool `json:"stream,omitempty"` StreamOptions *streamOptions `json:"stream_options,omitempty"` + // PromptCacheKey pins OpenAI's prompt-cache routing so requests from + // this agent land on the same cache shard. Prefix caching itself is + // automatic (≥1024 tokens); the key improves hit locality. Only set + // when ClientConfig.PromptCaching is on. + PromptCacheKey string `json:"prompt_cache_key,omitempty"` } type streamOptions struct { @@ -199,9 +208,35 @@ func (c *OpenAIClient) toOpenAIRequest(req *llm.ChatRequest, stream bool) openai r.StreamOptions = &streamOptions{IncludeUsage: true} } + if c.promptCaching { + r.PromptCacheKey = derivePromptCacheKey(model, req) + } + return r } +// derivePromptCacheKey builds a stable cache-routing key from the parts of +// the request that define the cacheable prefix: model, system prompt, and +// tool names. Identical (model, system, tools) across turns → identical key +// → OpenAI routes the session to the same cache shard. +func derivePromptCacheKey(model string, req *llm.ChatRequest) string { + h := sha256.New() + h.Write([]byte(model)) + h.Write([]byte{0}) + for _, m := range req.Messages { + if m.Role == llm.RoleSystem { + h.Write([]byte(m.Content)) + break + } + } + h.Write([]byte{0}) + for _, t := range req.Tools { + h.Write([]byte(t.Function.Name)) + h.Write([]byte{0}) + } + return "forge-" + hex.EncodeToString(h.Sum(nil))[:16] +} + // openaiResponse is the OpenAI-specific response format. type openaiResponse struct { ID string `json:"id"` diff --git a/forge-core/llm/providers/prompt_caching_test.go b/forge-core/llm/providers/prompt_caching_test.go new file mode 100644 index 00000000..796da4ac --- /dev/null +++ b/forge-core/llm/providers/prompt_caching_test.go @@ -0,0 +1,108 @@ +package providers + +import ( + "encoding/json" + "strings" + "testing" + + "github.com/initializ/forge/forge-core/llm" +) + +func cacheTestRequest() *llm.ChatRequest { + return &llm.ChatRequest{ + Messages: []llm.ChatMessage{ + {Role: llm.RoleSystem, Content: "You are a helpful agent."}, + {Role: llm.RoleUser, Content: "hello"}, + }, + Tools: []llm.ToolDefinition{ + {Type: "function", Function: llm.FunctionSchema{Name: "alpha_tool", Parameters: json.RawMessage(`{"type":"object"}`)}}, + {Type: "function", Function: llm.FunctionSchema{Name: "beta_tool", Parameters: json.RawMessage(`{"type":"object"}`)}}, + }, + } +} + +func TestAnthropic_PromptCaching_InjectsBreakpoints(t *testing.T) { + c := NewAnthropicClient(llm.ClientConfig{Model: "claude-sonnet-4-6", PromptCaching: true}) + body := c.toAnthropicRequest(cacheTestRequest(), false) + + // Last tool carries the ephemeral breakpoint; earlier tools do not. + if body.Tools[0].CacheControl != nil { + t.Error("first tool must not carry cache_control") + } + if body.Tools[1].CacheControl == nil || body.Tools[1].CacheControl.Type != "ephemeral" { + t.Error("last tool must carry cache_control ephemeral") + } + + // System becomes block form with a breakpoint. + blocks, ok := body.System.([]anthropicSystemBlock) + if !ok { + t.Fatalf("system should be block array when caching, got %T", body.System) + } + if len(blocks) != 1 || blocks[0].Text != "You are a helpful agent." || blocks[0].CacheControl == nil { + t.Fatalf("system block malformed: %+v", blocks) + } + + // The serialized body must contain cache_control (wire-level check). + data, err := json.Marshal(body) + if err != nil { + t.Fatal(err) + } + if !strings.Contains(string(data), `"cache_control":{"type":"ephemeral"}`) { + t.Errorf("serialized request missing cache_control: %s", data) + } +} + +func TestAnthropic_NoPromptCaching_WireFormatUnchanged(t *testing.T) { + c := NewAnthropicClient(llm.ClientConfig{Model: "claude-sonnet-4-6"}) + body := c.toAnthropicRequest(cacheTestRequest(), false) + + // System stays a plain string — byte-identical to the pre-caching format. + if _, ok := body.System.(string); !ok { + t.Fatalf("system should remain a plain string when caching is off, got %T", body.System) + } + data, err := json.Marshal(body) + if err != nil { + t.Fatal(err) + } + if strings.Contains(string(data), "cache_control") { + t.Errorf("cache_control must not appear when caching is off: %s", data) + } + if !strings.Contains(string(data), `"system":"You are a helpful agent."`) { + t.Errorf("system should serialize as a plain string: %s", data) + } +} + +func TestOpenAI_PromptCacheKey_StableAndGated(t *testing.T) { + on := NewOpenAIClient(llm.ClientConfig{Model: "gpt-4o", APIKey: "k", PromptCaching: true}) + off := NewOpenAIClient(llm.ClientConfig{Model: "gpt-4o", APIKey: "k"}) + + b1 := on.toOpenAIRequest(cacheTestRequest(), false) + b2 := on.toOpenAIRequest(cacheTestRequest(), false) + if b1.PromptCacheKey == "" { + t.Fatal("prompt_cache_key should be set when caching is on") + } + if b1.PromptCacheKey != b2.PromptCacheKey { + t.Errorf("prompt_cache_key not stable: %s vs %s", b1.PromptCacheKey, b2.PromptCacheKey) + } + if !strings.HasPrefix(b1.PromptCacheKey, "forge-") { + t.Errorf("prompt_cache_key should be forge-prefixed: %s", b1.PromptCacheKey) + } + + // Different system prompt → different key (prefix identity changed). + altered := cacheTestRequest() + altered.Messages[0].Content = "You are a different agent." + b3 := on.toOpenAIRequest(altered, false) + if b3.PromptCacheKey == b1.PromptCacheKey { + t.Error("prompt_cache_key should change when the system prompt changes") + } + + // Gated off → absent from the wire. + b4 := off.toOpenAIRequest(cacheTestRequest(), false) + if b4.PromptCacheKey != "" { + t.Error("prompt_cache_key must be empty when caching is off") + } + data, _ := json.Marshal(b4) + if strings.Contains(string(data), "prompt_cache_key") { + t.Errorf("prompt_cache_key must not serialize when caching is off: %s", data) + } +} diff --git a/forge-core/types/config.go b/forge-core/types/config.go index 66c73db6..2793b1f9 100644 --- a/forge-core/types/config.go +++ b/forge-core/types/config.go @@ -22,6 +22,7 @@ type ForgeConfig struct { Egress EgressRef `yaml:"egress,omitempty"` Skills SkillsRef `yaml:"skills,omitempty"` Memory MemoryConfig `yaml:"memory,omitempty"` + Compression CompressionConfig `yaml:"compression,omitempty"` Secrets SecretsConfig `yaml:"secrets,omitempty"` Auth AuthConfig `yaml:"auth,omitempty"` MCP MCPConfig `yaml:"mcp,omitempty"` @@ -452,6 +453,30 @@ type MemoryConfig struct { DecayHalfLifeDays int `yaml:"decay_half_life_days,omitempty"` // default: 7 } +// CompressionConfig configures reversible context compression (ctxzip). +// +// When enabled, bulky tool outputs and conversation content are compressed +// before reaching the LLM; everything dropped is stored locally (bbolt) and +// retrievable via the context_expand tool, so compression is lossy on the +// wire but lossless end-to-end. Enable via `compression.enabled: true` in +// forge.yaml or FORGE_COMPRESSION=true (env wins; "false" forces off). +type CompressionConfig struct { + Enabled *bool `yaml:"enabled,omitempty"` // default: false + // StorePath is the bbolt file holding offloaded originals. + StorePath string `yaml:"store_path,omitempty"` // default: .forge/ctxzip.db + // TTL is how long offloaded originals stay retrievable (Go duration, + // e.g. "30m", "2h"). + TTL string `yaml:"ttl,omitempty"` // default: 30m + // MinToolOutputChars is the tool-output size below which the + // compression hook leaves the output alone. + MinToolOutputChars int `yaml:"min_tool_output_chars,omitempty"` // default: 2048 + // CacheHints controls provider prompt-cache hints (Anthropic + // cache_control breakpoints, OpenAI prompt_cache_key). Defaults to the + // value of Enabled; set explicitly to run hints without compression or + // compression without hints. + CacheHints *bool `yaml:"cache_hints,omitempty"` +} + // EgressRef configures egress security controls. type EgressRef struct { Profile string `yaml:"profile,omitempty"` // strict, standard, permissive diff --git a/go.work.sum b/go.work.sum index b41fd722..a9e6653f 100644 --- a/go.work.sum +++ b/go.work.sum @@ -33,6 +33,8 @@ github.com/peterbourgon/diskv v2.0.1+incompatible/go.mod h1:uqqh8zWWbv1HBMNONnaR github.com/planetscale/vtprotobuf v0.6.1-0.20240319094008-0393e58bdf10/go.mod h1:t/avpk3KcrXxUnYOhZhMXJlSEyie6gQbtLq5NM3loB8= github.com/rogpeppe/fastuuid v1.2.0/go.mod h1:jVj6XXZzXRy/MSR5jhDC/2q6DgLz+nrA6LYCDYWNEvQ= github.com/sahilm/fuzzy v0.1.1/go.mod h1:VFvziUEIMCrT6A6tw2RFIXPXXmzXbOsSHF0DOI8ZK9Y= +github.com/spf13/pflag v1.0.10 h1:4EBh2KAYBwaONj6b2Ye1GiHfwjqyROoF4RwYO+vPwFk= +github.com/spf13/pflag v1.0.10/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg= github.com/spiffe/go-spiffe/v2 v2.6.0/go.mod h1:gm2SeUoMZEtpnzPNs2Csc0D/gX33k1xIx7lEzqblHEs= go.opentelemetry.io/contrib/detectors/gcp v1.42.0/go.mod h1:W9zQ439utxymRrXsUOzZbFX4JhLxXU4+ZnCt8GG7yA8= go.yaml.in/yaml/v2 v2.4.3 h1:6gvOSjQoTB3vt1l+CU+tSyi/HOjfOjRLJ4YwYZGwRO0= From 3e5bc3fbdbbf4edb2d351f4f9f40de20b80fb623 Mon Sep 17 00:00:00 2001 From: MK Date: Fri, 3 Jul 2026 17:54:52 -0400 Subject: [PATCH 02/16] =?UTF-8?q?fix(compress):=20live-test=20hardening=20?= =?UTF-8?q?=E2=80=94=20no=20expand=20recompression,=20hash=20resilience,?= =?UTF-8?q?=20store=20dir?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fixes from live agent testing (gpt-4o against a 150-pod fixture): - Never recompress context_expand output, at both seams (hook skips the tool by name; client wrapper passes ctxzip SkipNames). Without this the loop chased its own tail: the model expanded a marker and the hook crushed the expansion straight back into a marker. - Hash transcription resilience: models truncate or mangle marker hashes when copying them into tool calls. The Runtime now remembers emitted marker hashes; the expand tool resolves a unique prefix (≥6 chars) on exact-miss, and normalizeHash strips a glued ":count" suffix. - Create the store's parent directory (bbolt creates the file, not the dir) — on a fresh project .forge/ doesn't exist and compression failed open with "no such file or directory". - Bump ctxzip to a8b7923→94668f4: line-mode text compression (grep/log layout preserved byte-faithfully through the CCR round trip), stop-term filtering in line dedup, 12-hex marker hashes. Live result on "status breakdown + unhealthy pod" over 150 pods: tool output crushed 1397→51 tokens (96%), model called context_expand with a correctly-transcribed hash, got intact lines back, and answered with the exact error (CrashLoopBackOff / OOMKilled 512Mi) — total session 10,982 input tokens vs 19,799 in the pre-fix run. Co-Authored-By: Claude Opus 4.8 (1M context) --- forge-cli/go.mod | 2 +- forge-cli/go.sum | 4 +- forge-core/compress/client.go | 4 ++ forge-core/compress/compress.go | 53 ++++++++++++++++++++++ forge-core/compress/compress_test.go | 66 ++++++++++++++++++++++++++++ forge-core/compress/expand_tool.go | 24 +++++++--- forge-core/compress/hook.go | 9 ++++ forge-core/go.mod | 2 +- forge-core/go.sum | 4 ++ 9 files changed, 159 insertions(+), 9 deletions(-) diff --git a/forge-cli/go.mod b/forge-cli/go.mod index 25c3346e..4caec6dd 100644 --- a/forge-cli/go.mod +++ b/forge-cli/go.mod @@ -54,7 +54,7 @@ require ( github.com/gorilla/websocket v1.5.4-0.20250319132907-e064f32e3674 // indirect github.com/grpc-ecosystem/grpc-gateway/v2 v2.29.0 // indirect github.com/inconshreveable/mousetrap v1.1.0 // indirect - github.com/initializ/ctxzip v0.0.0-20260703193602-ca7fa8003af9 // indirect + github.com/initializ/ctxzip v0.0.0-20260703215102-94668f48685c // indirect github.com/josharian/intern v1.0.0 // indirect github.com/json-iterator/go v1.1.12 // indirect github.com/klauspost/compress v1.16.7 // indirect diff --git a/forge-cli/go.sum b/forge-cli/go.sum index df325a6f..fd6ec5fd 100644 --- a/forge-cli/go.sum +++ b/forge-cli/go.sum @@ -77,8 +77,8 @@ github.com/grpc-ecosystem/grpc-gateway/v2 v2.29.0 h1:5VipnvEpbqr2gA2VbM+nYVbkIF2 github.com/grpc-ecosystem/grpc-gateway/v2 v2.29.0/go.mod h1:Hyl3n6Twe1hvtd9XUXDec4pTvgMSEixRuQKPTMH2bNs= github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8= github.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw= -github.com/initializ/ctxzip v0.0.0-20260703193602-ca7fa8003af9 h1:a3rJaFaP+M3hERbrqAnnA5WOKYBMYuCP12wpHAMtLt8= -github.com/initializ/ctxzip v0.0.0-20260703193602-ca7fa8003af9/go.mod h1:Qs4CpXTEeEOlZwIMPm/qaFPmil+AUIqH7Jg/K81FTmU= +github.com/initializ/ctxzip v0.0.0-20260703215102-94668f48685c h1:w5Mq+/LEhuqElhyErASHaEJNLRvJSqdwMYwD37nC3H0= +github.com/initializ/ctxzip v0.0.0-20260703215102-94668f48685c/go.mod h1:Qs4CpXTEeEOlZwIMPm/qaFPmil+AUIqH7Jg/K81FTmU= github.com/initializ/guardrails v0.12.0 h1:YzScnl+YPihLA1gepjxKjnKH+2EbYmGC4dUtpVRAJ2c= github.com/initializ/guardrails v0.12.0/go.mod h1:bDdHx73MF0+O09KqoXmUTTiFG4H7yEVQ0NR6juP1F3Q= github.com/josharian/intern v1.0.0 h1:vlS4z54oSdjm0bgjRigI+G1HpF+tI+9rE5LLzOg8HmY= diff --git a/forge-core/compress/client.go b/forge-core/compress/client.go index a1a3b295..f3312a43 100644 --- a/forge-core/compress/client.go +++ b/forge-core/compress/client.go @@ -61,6 +61,9 @@ func (c *compressingClient) compressRequest(req *llm.ChatRequest) *llm.ChatReque opts := ctxzip.DefaultOptions() opts.Store = c.rt.store opts.Query = firstUserContent(req.Messages) + // Expansion results in history must stay verbatim — recompressing them + // recreates the marker the model already resolved. + opts.SkipNames = map[string]bool{expandToolName: true} res, err := ctxzip.Compress(zmsgs, opts) if err != nil || res == nil || res.SavedTokens() == 0 { @@ -72,6 +75,7 @@ func (c *compressingClient) compressRequest(req *llm.ChatRequest) *llm.ChatReque copy(out.Messages, req.Messages) for _, tr := range res.Transforms { out.Messages[tr.Index].Content = res.Messages[tr.Index].Content + c.rt.rememberMarkers(tr.Markers) } c.rt.debugf("compressed request", map[string]any{ diff --git a/forge-core/compress/compress.go b/forge-core/compress/compress.go index c79c66a6..e9e78a46 100644 --- a/forge-core/compress/compress.go +++ b/forge-core/compress/compress.go @@ -28,6 +28,10 @@ package compress import ( "fmt" + "os" + "path/filepath" + "strings" + "sync" "time" "github.com/initializ/ctxzip/ccr" @@ -61,6 +65,47 @@ type Runtime struct { store *ccr.BoltStore minSize int logger runtime.Logger + + // recent remembers marker hashes this process emitted so the expand tool + // can resolve a unique prefix when the model transcribes a hash + // imperfectly (observed live: models truncate hex hashes). + mu sync.Mutex + recent map[string]struct{} +} + +// rememberMarkers records emitted marker hashes for prefix resolution. +func (r *Runtime) rememberMarkers(hashes []string) { + if len(hashes) == 0 { + return + } + r.mu.Lock() + defer r.mu.Unlock() + if r.recent == nil { + r.recent = make(map[string]struct{}) + } + for _, h := range hashes { + r.recent[h] = struct{}{} + } +} + +// resolvePrefix returns the unique remembered hash starting with prefix, or +// "" when the prefix is too short, unknown, or ambiguous. +func (r *Runtime) resolvePrefix(prefix string) string { + if len(prefix) < 6 { + return "" + } + r.mu.Lock() + defer r.mu.Unlock() + var found string + for h := range r.recent { + if strings.HasPrefix(h, prefix) { + if found != "" { + return "" // ambiguous + } + found = h + } + } + return found } // New opens the durable store and returns a Runtime. Call Close on shutdown. @@ -74,6 +119,14 @@ func New(cfg Config) (*Runtime, error) { if cfg.MinToolOutputChars <= 0 { cfg.MinToolOutputChars = DefaultMinToolOutputChars } + // bbolt creates the DB file but not its parent directory; on a fresh + // project .forge/ does not exist yet, so create it (0700 — originals can + // hold sensitive tool output). + if dir := filepath.Dir(cfg.StorePath); dir != "" && dir != "." { + if err := os.MkdirAll(dir, 0o700); err != nil { + return nil, fmt.Errorf("compress: creating store dir: %w", err) + } + } store, err := ccr.NewBoltStore(ccr.BoltConfig{Path: cfg.StorePath, TTL: cfg.TTL}) if err != nil { return nil, fmt.Errorf("compress: opening store: %w", err) diff --git a/forge-core/compress/compress_test.go b/forge-core/compress/compress_test.go index 11939c7a..6b8356ac 100644 --- a/forge-core/compress/compress_test.go +++ b/forge-core/compress/compress_test.go @@ -24,6 +24,18 @@ func newRuntime(t *testing.T) *Runtime { return rt } +// Regression: on a fresh project .forge/ does not exist yet — New must create +// the store's parent directory instead of failing open (found in live testing: +// "open .forge/ctxzip.db: no such file or directory"). +func TestNew_CreatesMissingStoreDir(t *testing.T) { + path := filepath.Join(t.TempDir(), ".forge", "nested", "ctxzip.db") + rt, err := New(Config{StorePath: path}) + if err != nil { + t.Fatalf("New should create missing parent dirs: %v", err) + } + _ = rt.Close() +} + // bigJSON builds a large JSON-array tool output with one error row. func bigJSON(n int) string { items := make([]map[string]any, n) @@ -88,6 +100,36 @@ func TestHook_SkipsErrorsAndSmallOutput(t *testing.T) { } } +// Regression (live-test find): the loop chased its own tail — context_expand +// returned the original, the hook re-crushed it back into a marker. Expansion +// output must stay verbatim at both seams. +func TestNoRecompressionOfExpandOutput(t *testing.T) { + rt := newRuntime(t) + + // Hook seam. + big := bigJSON(80) + hctx := &runtime.HookContext{ToolName: "context_expand", ToolOutput: big} + _ = rt.AfterToolExecHook()(context.Background(), hctx) + if hctx.ToolOutput != big { + t.Fatal("hook recompressed context_expand output") + } + + // Client-wrapper seam: an expansion result sitting in the live zone. + inner := &capturingClient{} + client := rt.WrapClient(inner) + msgs := []llm.ChatMessage{ + {Role: llm.RoleSystem, Content: "sys"}, + {Role: llm.RoleUser, Content: "count the pods"}, + {Role: llm.RoleTool, Name: "context_expand", ToolCallID: "tc9", Content: big}, + {Role: llm.RoleAssistant, Content: "counting"}, + {Role: llm.RoleUser, Content: "go on"}, + } + _, _ = client.Chat(context.Background(), &llm.ChatRequest{Model: "m", Messages: msgs}) + if inner.lastReq.Messages[2].Content != big { + t.Fatal("wrapper recompressed context_expand output in history") + } +} + // capturingClient records the request it receives and returns a stub response. type capturingClient struct { lastReq *llm.ChatRequest @@ -183,6 +225,7 @@ func TestNormalizeHash(t *testing.T) { "ctxzip:abc123": "abc123", "hash=abc123": "abc123", " ABC123 ": "abc123", + "ctxzip:abc123:108": "abc123", // count glued on (observed live) } for in, want := range cases { if got := normalizeHash(in); got != want { @@ -191,6 +234,29 @@ func TestNormalizeHash(t *testing.T) { } } +// Regression (live-test find): models truncate hex hashes when transcribing a +// marker into a tool call. A unique prefix of an emitted hash must resolve. +func TestExpandTool_ResolvesTruncatedHash(t *testing.T) { + rt := newRuntime(t) + hook := rt.AfterToolExecHook() + hctx := &runtime.HookContext{ToolName: "list_pods", ToolOutput: bigJSON(80)} + _ = hook(context.Background(), hctx) + + hashes := ccr.ExtractHashes(hctx.ToolOutput) + if len(hashes) == 0 { + t.Fatal("no marker emitted") + } + // The model passes only the first 8 chars of the hash. + args, _ := json.Marshal(map[string]string{"hash": hashes[0][:8]}) + out, err := rt.ExpandTool().Execute(context.Background(), args) + if err != nil { + t.Fatal(err) + } + if strings.Contains(out, "No stored content") { + t.Fatalf("truncated hash did not resolve via prefix: %q", out[:80]) + } +} + func TestExpandTool_MissIsHelpful(t *testing.T) { rt := newRuntime(t) args, _ := json.Marshal(map[string]string{"hash": "deadbeefdeadbeefdeadbeef"}) diff --git a/forge-core/compress/expand_tool.go b/forge-core/compress/expand_tool.go index f76b3331..bfdee176 100644 --- a/forge-core/compress/expand_tool.go +++ b/forge-core/compress/expand_tool.go @@ -29,7 +29,11 @@ type expandInput struct { Hash string `json:"hash"` } -func (t *expandTool) Name() string { return "context_expand" } +// expandToolName is referenced by the hook and client wrapper to exempt this +// tool's output from compression. +const expandToolName = "context_expand" + +func (t *expandTool) Name() string { return expandToolName } func (t *expandTool) Description() string { return "Retrieve the original content behind a <> compression marker. " + @@ -60,6 +64,14 @@ func (t *expandTool) Execute(_ context.Context, args json.RawMessage) (string, e } original, ok := ctxzip.Unzip(t.rt.store, hash) + if !ok { + // Models sometimes transcribe a marker hash imperfectly (truncated + // hex). If the given value uniquely prefixes a hash this process + // emitted, resolve and retry before declaring a miss. + if full := t.rt.resolvePrefix(hash); full != "" { + original, ok = ctxzip.Unzip(t.rt.store, full) + } + } if !ok { // A miss is not a dead end — the disk or the original command is the // source of truth. Say so instead of returning a bare error. @@ -73,16 +85,18 @@ func (t *expandTool) Execute(_ context.Context, args json.RawMessage) (string, e } // normalizeHash tolerates the model passing a whole marker instead of the -// bare hash: "<>", "ctxzip:abc123", or -// "hash=abc123" all normalize to "abc123". +// bare hash: "<>", "ctxzip:abc123", +// "hash=abc123", or "abc123:51" (count glued on, observed live) all +// normalize to "abc123". func normalizeHash(s string) string { s = strings.TrimSpace(s) s = strings.TrimPrefix(s, "<<") s = strings.TrimPrefix(s, "ctxzip:") s = strings.TrimPrefix(s, "hash=") s = strings.TrimSuffix(s, ">>") - // Keep only the leading token — markers carry a trailing note. - if i := strings.IndexAny(s, " ,"); i >= 0 { + // Keep only the leading token — markers carry a trailing note, and + // models sometimes glue it on with a colon. + if i := strings.IndexAny(s, " ,:"); i >= 0 { s = s[:i] } return strings.ToLower(strings.TrimSpace(s)) diff --git a/forge-core/compress/hook.go b/forge-core/compress/hook.go index 1302ae01..5c1af3cd 100644 --- a/forge-core/compress/hook.go +++ b/forge-core/compress/hook.go @@ -25,6 +25,12 @@ func (r *Runtime) AfterToolExecHook() runtime.Hook { if hctx.Error != nil || len(hctx.ToolOutput) < r.minSize { return nil } + // Never compress the expansion tool's own output — the model just + // asked for those bytes back; re-crushing them recreates the marker + // it resolved and the loop chases its own tail (observed live). + if hctx.ToolName == expandToolName { + return nil + } opts := ctxzip.DefaultOptions() opts.Store = r.store @@ -47,6 +53,9 @@ func (r *Runtime) AfterToolExecHook() runtime.Hook { return nil } + for _, tr := range res.Transforms { + r.rememberMarkers(tr.Markers) + } r.debugf("compressed tool output", map[string]any{ "tool": hctx.ToolName, "tokens_before": res.TokensBefore, diff --git a/forge-core/go.mod b/forge-core/go.mod index 30fe5a68..14566a58 100644 --- a/forge-core/go.mod +++ b/forge-core/go.mod @@ -4,7 +4,7 @@ go 1.25.0 require ( github.com/golang-jwt/jwt/v5 v5.3.1 - github.com/initializ/ctxzip v0.0.0-20260703193602-ca7fa8003af9 + github.com/initializ/ctxzip v0.0.0-20260703215102-94668f48685c github.com/initializ/forge/forge-skills v0.0.0 github.com/xeipuuv/gojsonschema v1.2.0 go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.59.0 diff --git a/forge-core/go.sum b/forge-core/go.sum index 51191df4..a5d93fc7 100644 --- a/forge-core/go.sum +++ b/forge-core/go.sum @@ -24,6 +24,10 @@ github.com/grpc-ecosystem/grpc-gateway/v2 v2.29.0 h1:5VipnvEpbqr2gA2VbM+nYVbkIF2 github.com/grpc-ecosystem/grpc-gateway/v2 v2.29.0/go.mod h1:Hyl3n6Twe1hvtd9XUXDec4pTvgMSEixRuQKPTMH2bNs= github.com/initializ/ctxzip v0.0.0-20260703193602-ca7fa8003af9 h1:a3rJaFaP+M3hERbrqAnnA5WOKYBMYuCP12wpHAMtLt8= github.com/initializ/ctxzip v0.0.0-20260703193602-ca7fa8003af9/go.mod h1:Qs4CpXTEeEOlZwIMPm/qaFPmil+AUIqH7Jg/K81FTmU= +github.com/initializ/ctxzip v0.0.0-20260703214515-a8b792373e0b h1:9E4+wyRmJIhWnMpt6VdvMqNLKyyZT/pL3fJwVjTxkF0= +github.com/initializ/ctxzip v0.0.0-20260703214515-a8b792373e0b/go.mod h1:Qs4CpXTEeEOlZwIMPm/qaFPmil+AUIqH7Jg/K81FTmU= +github.com/initializ/ctxzip v0.0.0-20260703215102-94668f48685c h1:w5Mq+/LEhuqElhyErASHaEJNLRvJSqdwMYwD37nC3H0= +github.com/initializ/ctxzip v0.0.0-20260703215102-94668f48685c/go.mod h1:Qs4CpXTEeEOlZwIMPm/qaFPmil+AUIqH7Jg/K81FTmU= github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE= github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk= github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= From e6cbe2978fcf12c8eb7ce2e9e18dd707216e870d Mon Sep 17 00:00:00 2001 From: MK Date: Fri, 3 Jul 2026 18:17:04 -0400 Subject: [PATCH 03/16] feat(compress): builder-supplied keep_patterns in forge.yaml MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit compression.keep_patterns lets the agent builder declare a domain vocabulary of case-insensitive substrings compression must never drop: compression: enabled: true keep_patterns: [CrashLoopBackOff, ImagePullBackOff, OOMKilled] Threaded through compress.Config into both seams (AfterToolExec hook and the llm.Client wrapper) as ctxzip Options.MustKeep. Union semantics with ctxzip's built-in error floor — patterns only ever add protection. Bumps ctxzip to 304962f, whose defaults also grow k8s state words (crash/backoff/oomkilled/evicted/unhealthy/degraded) after live testing showed "CrashLoopBackOff" matched nothing in the original error list. Co-Authored-By: Claude Opus 4.8 (1M context) --- forge-cli/go.mod | 2 +- forge-cli/go.sum | 4 ++-- forge-cli/runtime/runner.go | 1 + forge-core/compress/client.go | 1 + forge-core/compress/compress.go | 12 +++++++++++- forge-core/compress/compress_test.go | 29 ++++++++++++++++++++++++++++ forge-core/compress/hook.go | 1 + forge-core/go.mod | 2 +- forge-core/go.sum | 2 ++ forge-core/types/config.go | 6 ++++++ 10 files changed, 55 insertions(+), 5 deletions(-) diff --git a/forge-cli/go.mod b/forge-cli/go.mod index 4caec6dd..c1dba2c6 100644 --- a/forge-cli/go.mod +++ b/forge-cli/go.mod @@ -54,7 +54,7 @@ require ( github.com/gorilla/websocket v1.5.4-0.20250319132907-e064f32e3674 // indirect github.com/grpc-ecosystem/grpc-gateway/v2 v2.29.0 // indirect github.com/inconshreveable/mousetrap v1.1.0 // indirect - github.com/initializ/ctxzip v0.0.0-20260703215102-94668f48685c // indirect + github.com/initializ/ctxzip v0.0.0-20260703221254-304962fbf29f // indirect github.com/josharian/intern v1.0.0 // indirect github.com/json-iterator/go v1.1.12 // indirect github.com/klauspost/compress v1.16.7 // indirect diff --git a/forge-cli/go.sum b/forge-cli/go.sum index fd6ec5fd..a690ddb7 100644 --- a/forge-cli/go.sum +++ b/forge-cli/go.sum @@ -77,8 +77,8 @@ github.com/grpc-ecosystem/grpc-gateway/v2 v2.29.0 h1:5VipnvEpbqr2gA2VbM+nYVbkIF2 github.com/grpc-ecosystem/grpc-gateway/v2 v2.29.0/go.mod h1:Hyl3n6Twe1hvtd9XUXDec4pTvgMSEixRuQKPTMH2bNs= github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8= github.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw= -github.com/initializ/ctxzip v0.0.0-20260703215102-94668f48685c h1:w5Mq+/LEhuqElhyErASHaEJNLRvJSqdwMYwD37nC3H0= -github.com/initializ/ctxzip v0.0.0-20260703215102-94668f48685c/go.mod h1:Qs4CpXTEeEOlZwIMPm/qaFPmil+AUIqH7Jg/K81FTmU= +github.com/initializ/ctxzip v0.0.0-20260703221254-304962fbf29f h1:RLaBHp/SGK3uW/klYJGh3WpWr4qRQWuwHf3ZoW/tqV0= +github.com/initializ/ctxzip v0.0.0-20260703221254-304962fbf29f/go.mod h1:Qs4CpXTEeEOlZwIMPm/qaFPmil+AUIqH7Jg/K81FTmU= github.com/initializ/guardrails v0.12.0 h1:YzScnl+YPihLA1gepjxKjnKH+2EbYmGC4dUtpVRAJ2c= github.com/initializ/guardrails v0.12.0/go.mod h1:bDdHx73MF0+O09KqoXmUTTiFG4H7yEVQ0NR6juP1F3Q= github.com/josharian/intern v1.0.0 h1:vlS4z54oSdjm0bgjRigI+G1HpF+tI+9rE5LLzOg8HmY= diff --git a/forge-cli/runtime/runner.go b/forge-cli/runtime/runner.go index 94f825a5..20f763a3 100644 --- a/forge-cli/runtime/runner.go +++ b/forge-cli/runtime/runner.go @@ -3436,6 +3436,7 @@ func (r *Runner) initCompression(reg *tools.Registry) *compress.Runtime { StorePath: storePath, TTL: ttl, MinToolOutputChars: cc.MinToolOutputChars, + KeepPatterns: cc.KeepPatterns, Logger: r.logger, }) if err != nil { diff --git a/forge-core/compress/client.go b/forge-core/compress/client.go index f3312a43..0da68289 100644 --- a/forge-core/compress/client.go +++ b/forge-core/compress/client.go @@ -64,6 +64,7 @@ func (c *compressingClient) compressRequest(req *llm.ChatRequest) *llm.ChatReque // Expansion results in history must stay verbatim — recompressing them // recreates the marker the model already resolved. opts.SkipNames = map[string]bool{expandToolName: true} + opts.MustKeep = c.rt.keep res, err := ctxzip.Compress(zmsgs, opts) if err != nil || res == nil || res.SavedTokens() == 0 { diff --git a/forge-core/compress/compress.go b/forge-core/compress/compress.go index e9e78a46..672cc913 100644 --- a/forge-core/compress/compress.go +++ b/forge-core/compress/compress.go @@ -55,6 +55,10 @@ type Config struct { // MinToolOutputChars is the size below which tool outputs are left // alone by the AfterToolExec hook. Default 2048. MinToolOutputChars int + // KeepPatterns is the builder's domain vocabulary of case-insensitive + // substrings compression must never drop (forge.yaml + // compression.keep_patterns). Union with ctxzip's built-in error floor. + KeepPatterns []string // Logger is optional; nil disables logging. Logger runtime.Logger } @@ -64,6 +68,7 @@ type Config struct { type Runtime struct { store *ccr.BoltStore minSize int + keep []string logger runtime.Logger // recent remembers marker hashes this process emitted so the expand tool @@ -131,7 +136,12 @@ func New(cfg Config) (*Runtime, error) { if err != nil { return nil, fmt.Errorf("compress: opening store: %w", err) } - return &Runtime{store: store, minSize: cfg.MinToolOutputChars, logger: cfg.Logger}, nil + return &Runtime{ + store: store, + minSize: cfg.MinToolOutputChars, + keep: cfg.KeepPatterns, + logger: cfg.Logger, + }, nil } // Close releases the underlying store. diff --git a/forge-core/compress/compress_test.go b/forge-core/compress/compress_test.go index 6b8356ac..577d99a7 100644 --- a/forge-core/compress/compress_test.go +++ b/forge-core/compress/compress_test.go @@ -234,6 +234,35 @@ func TestNormalizeHash(t *testing.T) { } } +// KeepPatterns (forge.yaml compression.keep_patterns) must flow through to +// the hook so builder-flagged rows survive the compressed view. +func TestHook_KeepPatterns(t *testing.T) { + rt, err := New(Config{ + StorePath: filepath.Join(t.TempDir(), "ctxzip.db"), + KeepPatterns: []string{"Quarantined"}, + }) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = rt.Close() }) + + items := make([]map[string]any, 80) + for i := range items { + items[i] = map[string]any{"id": fmt.Sprintf("r-%03d", i), "state": "nominal", "zone": "us-east-1"} + } + items[40] = map[string]any{"id": "r-040", "state": "QUARANTINED", "zone": "us-east-1"} + blob, _ := json.Marshal(items) + + hctx := &runtime.HookContext{ToolName: "fleet_list", ToolOutput: string(blob)} + _ = rt.AfterToolExecHook()(context.Background(), hctx) + if hctx.ToolOutput == string(blob) { + t.Fatal("expected compression to occur") + } + if !strings.Contains(hctx.ToolOutput, "QUARANTINED") { + t.Fatal("keep_patterns row dropped from compressed view") + } +} + // Regression (live-test find): models truncate hex hashes when transcribing a // marker into a tool call. A unique prefix of an emitted hash must resolve. func TestExpandTool_ResolvesTruncatedHash(t *testing.T) { diff --git a/forge-core/compress/hook.go b/forge-core/compress/hook.go index 5c1af3cd..aae2d7c6 100644 --- a/forge-core/compress/hook.go +++ b/forge-core/compress/hook.go @@ -42,6 +42,7 @@ func (r *Runtime) AfterToolExecHook() runtime.Hook { // what the model wanted from this output — and they are fixed at call // time, so recompression determinism is not a concern here. opts.Query = hctx.ToolInput + opts.MustKeep = r.keep msgs := []ctxzip.Message{{ Role: ctxzip.RoleTool, diff --git a/forge-core/go.mod b/forge-core/go.mod index 14566a58..36d62062 100644 --- a/forge-core/go.mod +++ b/forge-core/go.mod @@ -4,7 +4,7 @@ go 1.25.0 require ( github.com/golang-jwt/jwt/v5 v5.3.1 - github.com/initializ/ctxzip v0.0.0-20260703215102-94668f48685c + github.com/initializ/ctxzip v0.0.0-20260703221254-304962fbf29f github.com/initializ/forge/forge-skills v0.0.0 github.com/xeipuuv/gojsonschema v1.2.0 go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.59.0 diff --git a/forge-core/go.sum b/forge-core/go.sum index a5d93fc7..31686b5a 100644 --- a/forge-core/go.sum +++ b/forge-core/go.sum @@ -28,6 +28,8 @@ github.com/initializ/ctxzip v0.0.0-20260703214515-a8b792373e0b h1:9E4+wyRmJIhWnM github.com/initializ/ctxzip v0.0.0-20260703214515-a8b792373e0b/go.mod h1:Qs4CpXTEeEOlZwIMPm/qaFPmil+AUIqH7Jg/K81FTmU= github.com/initializ/ctxzip v0.0.0-20260703215102-94668f48685c h1:w5Mq+/LEhuqElhyErASHaEJNLRvJSqdwMYwD37nC3H0= github.com/initializ/ctxzip v0.0.0-20260703215102-94668f48685c/go.mod h1:Qs4CpXTEeEOlZwIMPm/qaFPmil+AUIqH7Jg/K81FTmU= +github.com/initializ/ctxzip v0.0.0-20260703221254-304962fbf29f h1:RLaBHp/SGK3uW/klYJGh3WpWr4qRQWuwHf3ZoW/tqV0= +github.com/initializ/ctxzip v0.0.0-20260703221254-304962fbf29f/go.mod h1:Qs4CpXTEeEOlZwIMPm/qaFPmil+AUIqH7Jg/K81FTmU= github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE= github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk= github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= diff --git a/forge-core/types/config.go b/forge-core/types/config.go index 2793b1f9..40024378 100644 --- a/forge-core/types/config.go +++ b/forge-core/types/config.go @@ -475,6 +475,12 @@ type CompressionConfig struct { // value of Enabled; set explicitly to run hints without compression or // compression without hints. CacheHints *bool `yaml:"cache_hints,omitempty"` + // KeepPatterns is the agent's domain vocabulary of case-insensitive + // substrings that compression must never drop — e.g. Kubernetes state + // words ("CrashLoopBackOff", "ImagePullBackOff") or product error codes. + // Union with the built-in error floor (error/fail/panic/timeout/...): + // entries only ever add protection. + KeepPatterns []string `yaml:"keep_patterns,omitempty"` } // EgressRef configures egress security controls. From 048ed5cef4e392212aa132588ade904e8c7d315f Mon Sep 17 00:00:00 2001 From: MK Date: Fri, 3 Jul 2026 18:53:19 -0400 Subject: [PATCH 04/16] feat(compress): emit compression savings into the audit stream MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two new audit events make token savings attributable instead of living only in debug logs: - context_compressed — fired from both seams (tool_output hook, request wrapper) with seam, tool, tokens_before/after, saved_tokens, plus running totals (total_saved_tokens / total_compressions / total_expansions) so any single event shows the cumulative picture, not just the per-call delta. - context_expanded — fired on every context_expand retrieval with hash, hit, bytes and the same running totals; expansions are the cost side auditors net against savings. Events flow through AuditLogger.EmitFromContext, so correlation_id / task_id / seq are stamped like every other audit event and SIEM consumers can join savings to invocations. compress stays decoupled via a Config.Audit callback; nil disables emission. Runtime.Totals() exposes the process-lifetime snapshot. Token figures are tokenizer estimates (directionally accurate), not provider-billed counts. Co-Authored-By: Claude Opus 4.8 (1M context) --- forge-cli/runtime/runner.go | 21 ++++++- forge-core/compress/client.go | 7 ++- forge-core/compress/compress.go | 93 ++++++++++++++++++++++++++++ forge-core/compress/compress_test.go | 72 +++++++++++++++++++++ forge-core/compress/expand_tool.go | 3 +- forge-core/compress/hook.go | 3 +- 6 files changed, 192 insertions(+), 7 deletions(-) diff --git a/forge-cli/runtime/runner.go b/forge-cli/runtime/runner.go index 20f763a3..26bd736f 100644 --- a/forge-cli/runtime/runner.go +++ b/forge-cli/runtime/runner.go @@ -896,7 +896,7 @@ func (r *Runner) Run(ctx context.Context) error { // compresses what redaction left; the client wrapper sits // below the FallbackChain so it also covers retries and // compactor summarization calls. - if comp := r.initCompression(reg); comp != nil { + if comp := r.initCompression(reg, auditLogger); comp != nil { defer comp.Close() //nolint:errcheck hooks.Register(coreruntime.AfterToolExec, comp.AfterToolExecHook()) llmClient = comp.WrapClient(llmClient) @@ -3411,7 +3411,13 @@ func (r *Runner) promptCachingEnabled() bool { // initCompression builds the ctxzip compression runtime and registers the // context_expand tool. Returns nil when compression is disabled or the store // cannot be opened (the agent then runs uncompressed — fail-open). -func (r *Runner) initCompression(reg *tools.Registry) *compress.Runtime { +// +// When auditLogger is non-nil, compression emits context_compressed / +// context_expanded audit events (correlation_id/task_id stamped from ctx via +// EmitFromContext) carrying per-event savings plus running totals, so SIEM +// consumers can attribute token reduction to compression rather than +// inferring it from tool_exec result sizes. +func (r *Runner) initCompression(reg *tools.Registry, auditLogger *coreruntime.AuditLogger) *compress.Runtime { if !r.compressionEnabled() { return nil } @@ -3432,12 +3438,23 @@ func (r *Runner) initCompression(reg *tools.Registry) *compress.Runtime { } } + var auditFn compress.AuditFunc + if auditLogger != nil { + auditFn = func(ctx context.Context, event string, fields map[string]any) { + auditLogger.EmitFromContext(ctx, coreruntime.AuditEvent{ + Event: event, + Fields: fields, + }) + } + } + comp, err := compress.New(compress.Config{ StorePath: storePath, TTL: ttl, MinToolOutputChars: cc.MinToolOutputChars, KeepPatterns: cc.KeepPatterns, Logger: r.logger, + Audit: auditFn, }) if err != nil { r.logger.Warn("failed to init compression, running uncompressed", map[string]any{ diff --git a/forge-core/compress/client.go b/forge-core/compress/client.go index 0da68289..f138e1e6 100644 --- a/forge-core/compress/client.go +++ b/forge-core/compress/client.go @@ -29,12 +29,12 @@ type compressingClient struct { // Chat implements llm.Client. func (c *compressingClient) Chat(ctx context.Context, req *llm.ChatRequest) (*llm.ChatResponse, error) { - return c.inner.Chat(ctx, c.compressRequest(req)) + return c.inner.Chat(ctx, c.compressRequest(ctx, req)) } // ChatStream implements llm.Client. func (c *compressingClient) ChatStream(ctx context.Context, req *llm.ChatRequest) (<-chan llm.StreamDelta, error) { - return c.inner.ChatStream(ctx, c.compressRequest(req)) + return c.inner.ChatStream(ctx, c.compressRequest(ctx, req)) } // ModelID implements llm.Client. @@ -43,7 +43,7 @@ func (c *compressingClient) ModelID() string { return c.inner.ModelID() } // compressRequest returns req with compressed live-zone messages, or req // unchanged when there is nothing to gain. The caller's request is never // mutated — the loop may reuse its message slice. -func (c *compressingClient) compressRequest(req *llm.ChatRequest) *llm.ChatRequest { +func (c *compressingClient) compressRequest(ctx context.Context, req *llm.ChatRequest) *llm.ChatRequest { if req == nil || len(req.Messages) == 0 { return req } @@ -86,6 +86,7 @@ func (c *compressingClient) compressRequest(req *llm.ChatRequest) *llm.ChatReque "tokens_after": res.TokensAfter, "saved_tokens": res.SavedTokens(), }) + c.rt.recordCompression(ctx, "request", "", res.TokensBefore, res.TokensAfter) return &out } diff --git a/forge-core/compress/compress.go b/forge-core/compress/compress.go index 672cc913..a6440d2e 100644 --- a/forge-core/compress/compress.go +++ b/forge-core/compress/compress.go @@ -27,6 +27,7 @@ package compress import ( + "context" "fmt" "os" "path/filepath" @@ -39,6 +40,26 @@ import ( "github.com/initializ/forge/forge-core/runtime" ) +// AuditFunc receives compression audit events. The runner wires this to the +// AuditLogger (EmitFromContext, so correlation_id/task_id are stamped from +// ctx); nil disables audit emission. Token figures are tokenizer estimates, +// not provider-billed counts — directionally accurate for savings reporting. +type AuditFunc func(ctx context.Context, event string, fields map[string]any) + +// Audit event names emitted by the compression runtime. +const ( + // AuditEventCompressed fires whenever content is compressed, from either + // seam (tool_output hook or request wrapper). Fields: seam, tool, + // tokens_before, tokens_after, saved_tokens, plus running totals + // total_saved_tokens / total_compressions / total_expansions so any + // single event shows the cumulative savings picture. + AuditEventCompressed = "context_compressed" + // AuditEventExpanded fires when the model retrieves offloaded content + // via context_expand. Fields: hash, hit, bytes, plus the same running + // totals — expansions are the "cost" side auditors net against savings. + AuditEventExpanded = "context_expanded" +) + // Defaults for Config fields left at their zero value. const ( DefaultTTL = 30 * time.Minute @@ -61,6 +82,8 @@ type Config struct { KeepPatterns []string // Logger is optional; nil disables logging. Logger runtime.Logger + // Audit is optional; nil disables audit emission. See AuditFunc. + Audit AuditFunc } // Runtime owns the shared CCR store and produces the hook, client wrapper, @@ -70,12 +93,81 @@ type Runtime struct { minSize int keep []string logger runtime.Logger + audit AuditFunc // recent remembers marker hashes this process emitted so the expand tool // can resolve a unique prefix when the model transcribes a hash // imperfectly (observed live: models truncate hex hashes). mu sync.Mutex recent map[string]struct{} + totals SavingsTotals +} + +// SavingsTotals is the process-lifetime savings picture. Token figures are +// tokenizer estimates. +type SavingsTotals struct { + // Compressions is how many times content was compressed (either seam). + Compressions int64 + // SavedTokens is the cumulative estimated token reduction. + SavedTokens int64 + // Expansions / ExpansionMisses count context_expand retrievals — the + // cost side to net against SavedTokens. + Expansions int64 + ExpansionMisses int64 +} + +// Totals returns a snapshot of the cumulative savings picture. +func (r *Runtime) Totals() SavingsTotals { + r.mu.Lock() + defer r.mu.Unlock() + return r.totals +} + +// recordCompression accumulates savings and emits AuditEventCompressed with +// per-event figures plus the running totals. +func (r *Runtime) recordCompression(ctx context.Context, seam, tool string, before, after int) { + r.mu.Lock() + r.totals.Compressions++ + r.totals.SavedTokens += int64(before - after) + t := r.totals + r.mu.Unlock() + + if r.audit == nil { + return + } + r.audit(ctx, AuditEventCompressed, map[string]any{ + "seam": seam, + "tool": tool, + "tokens_before": before, + "tokens_after": after, + "saved_tokens": before - after, + "total_saved_tokens": t.SavedTokens, + "total_compressions": t.Compressions, + "total_expansions": t.Expansions, + }) +} + +// recordExpansion accumulates retrieval stats and emits AuditEventExpanded. +func (r *Runtime) recordExpansion(ctx context.Context, hash string, hit bool, bytes int) { + r.mu.Lock() + r.totals.Expansions++ + if !hit { + r.totals.ExpansionMisses++ + } + t := r.totals + r.mu.Unlock() + + if r.audit == nil { + return + } + r.audit(ctx, AuditEventExpanded, map[string]any{ + "hash": hash, + "hit": hit, + "bytes": bytes, + "total_saved_tokens": t.SavedTokens, + "total_expansions": t.Expansions, + "total_expansion_misses": t.ExpansionMisses, + }) } // rememberMarkers records emitted marker hashes for prefix resolution. @@ -141,6 +233,7 @@ func New(cfg Config) (*Runtime, error) { minSize: cfg.MinToolOutputChars, keep: cfg.KeepPatterns, logger: cfg.Logger, + audit: cfg.Audit, }, nil } diff --git a/forge-core/compress/compress_test.go b/forge-core/compress/compress_test.go index 577d99a7..15d60a7e 100644 --- a/forge-core/compress/compress_test.go +++ b/forge-core/compress/compress_test.go @@ -234,6 +234,78 @@ func TestNormalizeHash(t *testing.T) { } } +// Compression must be attributable in the audit stream: every compression +// and expansion emits an event carrying per-event savings AND running totals, +// so auditors see cumulative savings, not just per-tool-call deltas. +func TestAuditEvents_SavingsAndTotals(t *testing.T) { + type captured struct { + event string + fields map[string]any + } + var events []captured + + rt, err := New(Config{ + StorePath: filepath.Join(t.TempDir(), "ctxzip.db"), + Audit: func(_ context.Context, event string, fields map[string]any) { + events = append(events, captured{event, fields}) + }, + }) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = rt.Close() }) + + // Two hook compressions accumulate totals. + hook := rt.AfterToolExecHook() + h1 := &runtime.HookContext{ToolName: "list_pods", ToolOutput: bigJSON(80)} + _ = hook(context.Background(), h1) + h2 := &runtime.HookContext{ToolName: "list_nodes", ToolOutput: bigJSON(90)} + _ = hook(context.Background(), h2) + + if len(events) != 2 { + t.Fatalf("want 2 context_compressed events, got %d", len(events)) + } + for i, e := range events { + if e.event != AuditEventCompressed { + t.Fatalf("event %d = %s, want %s", i, e.event, AuditEventCompressed) + } + if e.fields["seam"] != "tool_output" { + t.Errorf("event %d seam = %v", i, e.fields["seam"]) + } + if saved, _ := e.fields["saved_tokens"].(int); saved <= 0 { + t.Errorf("event %d saved_tokens = %v, want > 0", i, e.fields["saved_tokens"]) + } + } + // Running total on the second event exceeds its own per-event saving. + second := events[1].fields + if second["total_compressions"].(int64) != 2 { + t.Errorf("total_compressions = %v, want 2", second["total_compressions"]) + } + if second["total_saved_tokens"].(int64) <= int64(second["saved_tokens"].(int)) { + t.Errorf("running total %v should exceed per-event %v", + second["total_saved_tokens"], second["saved_tokens"]) + } + + // An expansion emits the cost-side event with the same totals context. + hashes := ccr.ExtractHashes(h1.ToolOutput) + args, _ := json.Marshal(map[string]string{"hash": hashes[0]}) + _, _ = rt.ExpandTool().Execute(context.Background(), args) + + last := events[len(events)-1] + if last.event != AuditEventExpanded { + t.Fatalf("want %s event after expand, got %s", AuditEventExpanded, last.event) + } + if last.fields["hit"] != true || last.fields["total_expansions"].(int64) != 1 { + t.Errorf("expansion fields wrong: %+v", last.fields) + } + + // Totals snapshot agrees. + tot := rt.Totals() + if tot.Compressions != 2 || tot.Expansions != 1 || tot.SavedTokens <= 0 { + t.Errorf("Totals() = %+v", tot) + } +} + // KeepPatterns (forge.yaml compression.keep_patterns) must flow through to // the hook so builder-flagged rows survive the compressed view. func TestHook_KeepPatterns(t *testing.T) { diff --git a/forge-core/compress/expand_tool.go b/forge-core/compress/expand_tool.go index bfdee176..ba75a9eb 100644 --- a/forge-core/compress/expand_tool.go +++ b/forge-core/compress/expand_tool.go @@ -53,7 +53,7 @@ func (t *expandTool) InputSchema() json.RawMessage { }`) } -func (t *expandTool) Execute(_ context.Context, args json.RawMessage) (string, error) { +func (t *expandTool) Execute(ctx context.Context, args json.RawMessage) (string, error) { var input expandInput if err := json.Unmarshal(args, &input); err != nil { return "", fmt.Errorf("parsing input: %w", err) @@ -72,6 +72,7 @@ func (t *expandTool) Execute(_ context.Context, args json.RawMessage) (string, e original, ok = ctxzip.Unzip(t.rt.store, full) } } + t.rt.recordExpansion(ctx, hash, ok, len(original)) if !ok { // A miss is not a dead end — the disk or the original command is the // source of truth. Say so instead of returning a bare error. diff --git a/forge-core/compress/hook.go b/forge-core/compress/hook.go index aae2d7c6..0e5de30c 100644 --- a/forge-core/compress/hook.go +++ b/forge-core/compress/hook.go @@ -21,7 +21,7 @@ import ( // Error results (hctx.Error != nil) are always left verbatim — dropping parts // of an error the user is about to debug is the catastrophic failure mode. func (r *Runtime) AfterToolExecHook() runtime.Hook { - return func(_ context.Context, hctx *runtime.HookContext) error { + return func(ctx context.Context, hctx *runtime.HookContext) error { if hctx.Error != nil || len(hctx.ToolOutput) < r.minSize { return nil } @@ -63,6 +63,7 @@ func (r *Runtime) AfterToolExecHook() runtime.Hook { "tokens_after": res.TokensAfter, "saved_tokens": res.SavedTokens(), }) + r.recordCompression(ctx, "tool_output", hctx.ToolName, res.TokensBefore, res.TokensAfter) hctx.ToolOutput = res.Messages[0].Content return nil } From 2e3866187c7a5e243187afb32ab4e9705f7dd922 Mon Sep 17 00:00:00 2001 From: MK Date: Fri, 3 Jul 2026 23:01:38 -0400 Subject: [PATCH 05/16] feat(compress): per-invocation savings on invocation_complete MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit invocation_complete now carries compression_saved_tokens_total, compression_count, and (when nonzero) expansion_count alongside the existing input/output token totals, so per-invocation cost rollups show what compression saved without joining context_compressed events. Savings are accumulated per correlation ID inside compress.Runtime and popped once at the response boundary (TakeInvocationTotals), so concurrent invocations never cross-contaminate — diffing the process- lifetime totals would have. Fields are present whenever compression is enabled; zeros mean "on, but nothing was worth compressing". Token figures remain tokenizer estimates. Co-Authored-By: Claude Opus 4.8 (1M context) --- forge-cli/runtime/runner.go | 14 ++++++++ forge-core/compress/compress.go | 52 ++++++++++++++++++++++++++++ forge-core/compress/compress_test.go | 36 +++++++++++++++++++ 3 files changed, 102 insertions(+) diff --git a/forge-cli/runtime/runner.go b/forge-cli/runtime/runner.go index 26bd736f..563f458a 100644 --- a/forge-cli/runtime/runner.go +++ b/forge-cli/runtime/runner.go @@ -145,6 +145,7 @@ type Runner struct { authToken string // resolved auth token (empty if --no-auth) cancelRegistry *coreruntime.CancellationRegistry // per-Runner in-flight cancellation registry (issue #88 / FWS-4) auditSigningKey *coreruntime.LoadedKey // loaded once at startup; nil when signing is off (#213). Served on JWKS endpoint. + compression *compress.Runtime // ctxzip compression runtime; nil when compression is disabled } // NewRunner creates a Runner from the given config. @@ -900,6 +901,7 @@ func (r *Runner) Run(ctx context.Context) error { defer comp.Close() //nolint:errcheck hooks.Register(coreruntime.AfterToolExec, comp.AfterToolExecHook()) llmClient = comp.WrapClient(llmClient) + r.compression = comp // invocation_complete reads per-task savings } // Compute model-aware character budget. @@ -1535,6 +1537,18 @@ func (r *Runner) executeTask( fields["provider"] = snap.PrimaryProvider } } + // Per-invocation compression savings (tokenizer estimates), popped + // from the runtime by this invocation's correlation ID so concurrent + // tasks don't cross-contaminate. Present whenever compression is + // enabled — zero values mean "on, but nothing was worth compressing". + if r.compression != nil { + ct := r.compression.TakeInvocationTotals(ctx) + fields["compression_saved_tokens_total"] = ct.SavedTokens + fields["compression_count"] = ct.Compressions + if ct.Expansions > 0 { + fields["expansion_count"] = ct.Expansions + } + } if task.Status.State == a2a.TaskStateCanceled { auditLogger.EmitInvocationCancelled(ctx, coreruntime.CancellationReasonFromCause(ctx), diff --git a/forge-core/compress/compress.go b/forge-core/compress/compress.go index a6440d2e..c37aaee7 100644 --- a/forge-core/compress/compress.go +++ b/forge-core/compress/compress.go @@ -101,6 +101,11 @@ type Runtime struct { mu sync.Mutex recent map[string]struct{} totals SavingsTotals + // perInvocation accumulates savings keyed by correlation ID so + // invocation_complete can report this invocation's savings without + // cross-contamination from concurrent tasks. Entries are popped by + // TakeInvocationTotals at the invocation's response boundary. + perInvocation map[string]*SavingsTotals } // SavingsTotals is the process-lifetime savings picture. Token figures are @@ -123,12 +128,53 @@ func (r *Runtime) Totals() SavingsTotals { return r.totals } +// TakeInvocationTotals pops and returns the savings accumulated under the +// ctx's correlation ID. Call it exactly once, at the invocation's response +// boundary (invocation_complete emission); subsequent calls for the same +// invocation return zeros. Safe under concurrent invocations — each +// correlation ID accumulates independently. +func (r *Runtime) TakeInvocationTotals(ctx context.Context) SavingsTotals { + cid := runtime.CorrelationIDFromContext(ctx) + if cid == "" { + return SavingsTotals{} + } + r.mu.Lock() + defer r.mu.Unlock() + if t, ok := r.perInvocation[cid]; ok { + delete(r.perInvocation, cid) + return *t + } + return SavingsTotals{} +} + +// bumpInvocation accumulates into the ctx's per-invocation bucket. Caller +// holds r.mu. +func (r *Runtime) bumpInvocation(ctx context.Context, f func(*SavingsTotals)) { + cid := runtime.CorrelationIDFromContext(ctx) + if cid == "" { + return + } + if r.perInvocation == nil { + r.perInvocation = make(map[string]*SavingsTotals) + } + t, ok := r.perInvocation[cid] + if !ok { + t = &SavingsTotals{} + r.perInvocation[cid] = t + } + f(t) +} + // recordCompression accumulates savings and emits AuditEventCompressed with // per-event figures plus the running totals. func (r *Runtime) recordCompression(ctx context.Context, seam, tool string, before, after int) { r.mu.Lock() r.totals.Compressions++ r.totals.SavedTokens += int64(before - after) + r.bumpInvocation(ctx, func(t *SavingsTotals) { + t.Compressions++ + t.SavedTokens += int64(before - after) + }) t := r.totals r.mu.Unlock() @@ -154,6 +200,12 @@ func (r *Runtime) recordExpansion(ctx context.Context, hash string, hit bool, by if !hit { r.totals.ExpansionMisses++ } + r.bumpInvocation(ctx, func(t *SavingsTotals) { + t.Expansions++ + if !hit { + t.ExpansionMisses++ + } + }) t := r.totals r.mu.Unlock() diff --git a/forge-core/compress/compress_test.go b/forge-core/compress/compress_test.go index 15d60a7e..c843185a 100644 --- a/forge-core/compress/compress_test.go +++ b/forge-core/compress/compress_test.go @@ -306,6 +306,42 @@ func TestAuditEvents_SavingsAndTotals(t *testing.T) { } } +// Per-invocation savings are keyed by correlation ID so concurrent tasks +// don't cross-contaminate, and TakeInvocationTotals pops exactly once. +func TestTakeInvocationTotals_PerCorrelation(t *testing.T) { + rt := newRuntime(t) + hook := rt.AfterToolExecHook() + + ctxA := runtime.WithCorrelationID(context.Background(), "task-a") + ctxB := runtime.WithCorrelationID(context.Background(), "task-b") + + // Two compressions under A, one under B. + for _, c := range []struct { + ctx context.Context + n int + }{{ctxA, 80}, {ctxA, 90}, {ctxB, 100}} { + h := &runtime.HookContext{ToolName: "t", ToolOutput: bigJSON(c.n)} + _ = hook(c.ctx, h) + } + + a := rt.TakeInvocationTotals(ctxA) + if a.Compressions != 2 || a.SavedTokens <= 0 { + t.Fatalf("A totals = %+v, want 2 compressions", a) + } + b := rt.TakeInvocationTotals(ctxB) + if b.Compressions != 1 { + t.Fatalf("B totals = %+v, want 1 compression", b) + } + // Popped — a second take returns zeros. + if again := rt.TakeInvocationTotals(ctxA); again.Compressions != 0 { + t.Fatalf("second take should be empty, got %+v", again) + } + // No correlation ID → zeros, no accumulation. + if none := rt.TakeInvocationTotals(context.Background()); none.Compressions != 0 { + t.Fatalf("no-correlation take should be empty, got %+v", none) + } +} + // KeepPatterns (forge.yaml compression.keep_patterns) must flow through to // the hook so builder-flagged rows survive the compressed view. func TestHook_KeepPatterns(t *testing.T) { From 364e3f40bb92a063a745c620d8d4a2ebc505c366 Mon Sep 17 00:00:00 2001 From: MK Date: Fri, 3 Jul 2026 23:33:51 -0400 Subject: [PATCH 06/16] =?UTF-8?q?feat(compress):=20runtime-owned=20system?= =?UTF-8?q?=20directive=20=E2=80=94=20skills=20need=20zero=20marker=20awar?= =?UTF-8?q?eness?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Marker awareness was living in the test agent's SKILL.md, which does not scale: every skill author would have to document compression. Compression is a runtime capability, so the runtime now briefs the model itself — when compression is enabled, compress.SystemDirective is appended to the system prompt (same pattern as codeAgentDirective), explaining what <> markers are, that the visible remainder keeps errors and representative content, and when/how to call context_expand. The directive is a constant, keeping the system prompt byte-stable across turns for provider prompt caches. A guard test pins it to the real tool name and marker prefix so a rename cannot silently orphan the text. Co-Authored-By: Claude Opus 4.8 (1M context) --- forge-cli/runtime/runner.go | 7 +++++++ forge-core/compress/compress_test.go | 12 ++++++++++++ forge-core/compress/expand_tool.go | 18 ++++++++++++++++++ 3 files changed, 37 insertions(+) diff --git a/forge-cli/runtime/runner.go b/forge-cli/runtime/runner.go index 563f458a..49aa22f8 100644 --- a/forge-cli/runtime/runner.go +++ b/forge-cli/runtime/runner.go @@ -915,6 +915,13 @@ func (r *Runner) Run(ctx context.Context) error { if r.hasSkill("code-agent") { sysPrompt += "\n\n" + codeAgentDirective } + // Compression marker-awareness is a runtime concern, not a + // per-skill one: whenever compression is on, every skill's + // agent learns what <> markers are and when to + // call context_expand — skill authors need do nothing. + if r.compression != nil { + sysPrompt += "\n\n" + compress.SystemDirective + } execCfg := coreruntime.LLMExecutorConfig{ Client: llmClient, diff --git a/forge-core/compress/compress_test.go b/forge-core/compress/compress_test.go index c843185a..2e3af312 100644 --- a/forge-core/compress/compress_test.go +++ b/forge-core/compress/compress_test.go @@ -306,6 +306,18 @@ func TestAuditEvents_SavingsAndTotals(t *testing.T) { } } +// The runtime-owned system directive is what makes compression work for ANY +// skill without skill authors documenting it. Guard that it names the real +// tool and the real marker prefix, so a rename cannot silently orphan it. +func TestSystemDirective_MatchesToolAndMarker(t *testing.T) { + if !strings.Contains(SystemDirective, expandToolName) { + t.Errorf("SystemDirective does not mention the %s tool", expandToolName) + } + if !strings.Contains(SystemDirective, ccr.MarkerPrefix) { + t.Errorf("SystemDirective does not show the %q marker shape", ccr.MarkerPrefix) + } +} + // Per-invocation savings are keyed by correlation ID so concurrent tasks // don't cross-contaminate, and TakeInvocationTotals pops exactly once. func TestTakeInvocationTotals_PerCorrelation(t *testing.T) { diff --git a/forge-core/compress/expand_tool.go b/forge-core/compress/expand_tool.go index ba75a9eb..cc6e03bf 100644 --- a/forge-core/compress/expand_tool.go +++ b/forge-core/compress/expand_tool.go @@ -33,6 +33,24 @@ type expandInput struct { // tool's output from compression. const expandToolName = "context_expand" +// SystemDirective is appended to the agent's system prompt whenever +// compression is enabled, so EVERY skill gets marker-awareness from the +// runtime — skill authors never need to document compression themselves. +// The text is constant, which keeps the system prompt byte-stable across +// turns (provider prompt caches stay warm). +const SystemDirective = `## Compressed context + +Large tool outputs may be automatically compressed to fit your context. +Compressed sections are replaced inline by a marker like +<> — the note says how much was offloaded. The +visible remainder keeps errors, anomalies, and representative content, so for +many questions you will not need the offloaded part. + +When you DO need offloaded data to answer precisely (exact counts, full +listings, a specific record you cannot see), call the ` + expandToolName + ` tool with the +marker's hash to retrieve the original content. If it reports the content +expired, re-run the tool that produced the output.` + func (t *expandTool) Name() string { return expandToolName } func (t *expandTool) Description() string { From e2673d1b55fa901b6a33e99d4ff20b874e5802f1 Mon Sep 17 00:00:00 2001 From: MK Date: Sat, 4 Jul 2026 15:11:05 -0400 Subject: [PATCH 07/16] chore: pin ctxzip to released v0.1.0 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replaces the feat-branch pseudo-version with the immutable release tag — the forge PR now depends on a stable, reviewable ctxzip version. Co-Authored-By: Claude Opus 4.8 (1M context) --- forge-cli/go.mod | 2 +- forge-cli/go.sum | 4 ++-- forge-core/go.mod | 2 +- forge-core/go.sum | 2 ++ 4 files changed, 6 insertions(+), 4 deletions(-) diff --git a/forge-cli/go.mod b/forge-cli/go.mod index c1dba2c6..da476b10 100644 --- a/forge-cli/go.mod +++ b/forge-cli/go.mod @@ -54,7 +54,7 @@ require ( github.com/gorilla/websocket v1.5.4-0.20250319132907-e064f32e3674 // indirect github.com/grpc-ecosystem/grpc-gateway/v2 v2.29.0 // indirect github.com/inconshreveable/mousetrap v1.1.0 // indirect - github.com/initializ/ctxzip v0.0.0-20260703221254-304962fbf29f // indirect + github.com/initializ/ctxzip v0.1.0 // indirect github.com/josharian/intern v1.0.0 // indirect github.com/json-iterator/go v1.1.12 // indirect github.com/klauspost/compress v1.16.7 // indirect diff --git a/forge-cli/go.sum b/forge-cli/go.sum index a690ddb7..37ea49b5 100644 --- a/forge-cli/go.sum +++ b/forge-cli/go.sum @@ -77,8 +77,8 @@ github.com/grpc-ecosystem/grpc-gateway/v2 v2.29.0 h1:5VipnvEpbqr2gA2VbM+nYVbkIF2 github.com/grpc-ecosystem/grpc-gateway/v2 v2.29.0/go.mod h1:Hyl3n6Twe1hvtd9XUXDec4pTvgMSEixRuQKPTMH2bNs= github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8= github.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw= -github.com/initializ/ctxzip v0.0.0-20260703221254-304962fbf29f h1:RLaBHp/SGK3uW/klYJGh3WpWr4qRQWuwHf3ZoW/tqV0= -github.com/initializ/ctxzip v0.0.0-20260703221254-304962fbf29f/go.mod h1:Qs4CpXTEeEOlZwIMPm/qaFPmil+AUIqH7Jg/K81FTmU= +github.com/initializ/ctxzip v0.1.0 h1:GDCpfYSYUB630e1aaGpH7YatTXilu1n3pZ1m5WHtLBw= +github.com/initializ/ctxzip v0.1.0/go.mod h1:Qs4CpXTEeEOlZwIMPm/qaFPmil+AUIqH7Jg/K81FTmU= github.com/initializ/guardrails v0.12.0 h1:YzScnl+YPihLA1gepjxKjnKH+2EbYmGC4dUtpVRAJ2c= github.com/initializ/guardrails v0.12.0/go.mod h1:bDdHx73MF0+O09KqoXmUTTiFG4H7yEVQ0NR6juP1F3Q= github.com/josharian/intern v1.0.0 h1:vlS4z54oSdjm0bgjRigI+G1HpF+tI+9rE5LLzOg8HmY= diff --git a/forge-core/go.mod b/forge-core/go.mod index 36d62062..d82061cb 100644 --- a/forge-core/go.mod +++ b/forge-core/go.mod @@ -4,7 +4,7 @@ go 1.25.0 require ( github.com/golang-jwt/jwt/v5 v5.3.1 - github.com/initializ/ctxzip v0.0.0-20260703221254-304962fbf29f + github.com/initializ/ctxzip v0.1.0 github.com/initializ/forge/forge-skills v0.0.0 github.com/xeipuuv/gojsonschema v1.2.0 go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.59.0 diff --git a/forge-core/go.sum b/forge-core/go.sum index 31686b5a..e6b86c2b 100644 --- a/forge-core/go.sum +++ b/forge-core/go.sum @@ -30,6 +30,8 @@ github.com/initializ/ctxzip v0.0.0-20260703215102-94668f48685c h1:w5Mq+/LEhuqElh github.com/initializ/ctxzip v0.0.0-20260703215102-94668f48685c/go.mod h1:Qs4CpXTEeEOlZwIMPm/qaFPmil+AUIqH7Jg/K81FTmU= github.com/initializ/ctxzip v0.0.0-20260703221254-304962fbf29f h1:RLaBHp/SGK3uW/klYJGh3WpWr4qRQWuwHf3ZoW/tqV0= github.com/initializ/ctxzip v0.0.0-20260703221254-304962fbf29f/go.mod h1:Qs4CpXTEeEOlZwIMPm/qaFPmil+AUIqH7Jg/K81FTmU= +github.com/initializ/ctxzip v0.1.0 h1:GDCpfYSYUB630e1aaGpH7YatTXilu1n3pZ1m5WHtLBw= +github.com/initializ/ctxzip v0.1.0/go.mod h1:Qs4CpXTEeEOlZwIMPm/qaFPmil+AUIqH7Jg/K81FTmU= github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE= github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk= github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= From e1c916bebe7d7dc91d9398f08028c9cb869884a8 Mon Sep 17 00:00:00 2001 From: MK Date: Sat, 4 Jul 2026 16:52:03 -0400 Subject: [PATCH 08/16] feat(cli): compression option in run/serve flags and the init TUI wizard MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Compression becomes reachable from every entry point, not just forge.yaml/env: - `forge run --compression` / `--compression=false` — tri-state: absent leaves yaml/env resolution untouched; explicit values override both by setting FORGE_COMPRESSION (same pattern as --model → MODEL_NAME). - `forge serve --compression[=false]` — forwarded to the forked daemon `forge run`, only when explicitly passed. - `forge init --compression` — non-interactive scaffolding writes a commented `compression.enabled: true` block into forge.yaml. - init TUI wizard — new "Context Compression" step (SingleSelect, Enabled/Disabled with explanatory descriptions) between Skills and Auth; selection flows through WizardContext.Compression into the same forge.yaml block. Scaffold test covers both directions: --compression writes the block, default omits it (off by default). Co-Authored-By: Claude Opus 4.8 (1M context) --- forge-cli/cmd/init.go | 7 ++ forge-cli/cmd/init_test.go | 51 ++++++++++ forge-cli/cmd/run.go | 14 +++ forge-cli/cmd/serve.go | 7 ++ .../internal/tui/steps/compression_step.go | 94 +++++++++++++++++++ forge-cli/internal/tui/wizard.go | 6 ++ forge-cli/templates/init/forge.yaml.tmpl | 8 ++ 7 files changed, 187 insertions(+) create mode 100644 forge-cli/internal/tui/steps/compression_step.go diff --git a/forge-cli/cmd/init.go b/forge-cli/cmd/init.go index a284b598..6956820d 100644 --- a/forge-cli/cmd/init.go +++ b/forge-cli/cmd/init.go @@ -45,6 +45,7 @@ type initOptions struct { Force bool // overwrite existing directory CustomModel string // custom provider model name AuthMethod string // "apikey" or "oauth" + Compression bool // reversible context compression (ctxzip) — compression.enabled in forge.yaml // A2A auth chain (from the wizard's Authentication step or CLI flags). AuthMode string // "", "none", "oidc", "http_verifier", "custom" @@ -76,6 +77,7 @@ type templateData struct { EgressDomains []string EnvVars []envVarEntry HasSecrets bool + Compression bool // Auth chain rendering (see forge.yaml.tmpl). Pre-rendered as a YAML // fragment because nested maps in the settings block (e.g. claim_map) @@ -126,6 +128,7 @@ func init() { initCmd.Flags().StringSlice("channels", nil, "communication channels (e.g., slack,telegram)") initCmd.Flags().String("from-skills", "", "path to SKILL.md file to parse for tools") initCmd.Flags().Bool("non-interactive", false, "run without interactive prompts (requires all flags)") + initCmd.Flags().Bool("compression", false, "enable reversible context compression (writes compression.enabled: true to forge.yaml)") initCmd.Flags().StringSlice("tools", nil, "builtin tools to enable (e.g., web_search,http_request)") initCmd.Flags().StringSlice("skills", nil, "registry skills to include (e.g., github,weather)") initCmd.Flags().String("api-key", "", "LLM provider API key") @@ -181,6 +184,7 @@ func runInit(cmd *cobra.Command, args []string) error { opts.Skills, _ = cmd.Flags().GetStringSlice("skills") opts.APIKey, _ = cmd.Flags().GetString("api-key") opts.OrganizationID, _ = cmd.Flags().GetString("org-id") + opts.Compression, _ = cmd.Flags().GetBool("compression") fallbackProviders, _ := cmd.Flags().GetStringSlice("fallbacks") for _, p := range fallbackProviders { opts.Fallbacks = append(opts.Fallbacks, tui.FallbackProvider{Provider: p}) @@ -323,6 +327,7 @@ func collectInteractive(opts *initOptions) error { steps.NewChannelStep(styles), steps.NewToolsStep(styles, toolInfos, validateWebSearchKeyFn), steps.NewSkillsStep(styles, skillInfos), + steps.NewCompressionStep(styles), steps.NewAuthStep(styles), steps.NewEgressStep(styles, deriveEgressFn), steps.NewReviewStep(styles), // scaffold is handled by the caller after collectInteractive returns @@ -372,6 +377,7 @@ func collectInteractive(opts *initOptions) error { opts.BuiltinTools = ctx.BuiltinTools opts.Skills = ctx.Skills + opts.Compression = ctx.Compression // Store provider env var storeProviderEnvVar(opts) @@ -1115,6 +1121,7 @@ func buildTemplateData(opts *initOptions) templateData { Channels: opts.Channels, Tools: opts.Tools, BuiltinTools: opts.BuiltinTools, + Compression: opts.Compression, } // Set entrypoint based on framework (only for subprocess-based frameworks) diff --git a/forge-cli/cmd/init_test.go b/forge-cli/cmd/init_test.go index 8bcb4e29..f6463131 100644 --- a/forge-cli/cmd/init_test.go +++ b/forge-cli/cmd/init_test.go @@ -439,6 +439,57 @@ func TestScaffold_EgressInForgeYAML(t *testing.T) { } } +func TestScaffold_CompressionInForgeYAML(t *testing.T) { + tmpDir := t.TempDir() + origDir, _ := os.Getwd() + if err := os.Chdir(tmpDir); err != nil { + t.Fatalf("chdir: %v", err) + } + defer func() { _ = os.Chdir(origDir) }() + + // With --compression: the block is written enabled. + opts := &initOptions{ + Name: "comp-test", + AgentID: "comp-test", + Framework: "forge", + ModelProvider: "openai", + EnvVars: map[string]string{}, + NonInteractive: true, + Compression: true, + } + if err := scaffold(opts); err != nil { + t.Fatalf("scaffold error: %v", err) + } + content, err := os.ReadFile(filepath.Join("comp-test", "forge.yaml")) + if err != nil { + t.Fatalf("reading forge.yaml: %v", err) + } + if !strings.Contains(string(content), "compression:") || + !strings.Contains(string(content), "enabled: true") { + t.Errorf("forge.yaml missing compression block:\n%s", content) + } + + // Without it: no compression block (off by default). + opts2 := &initOptions{ + Name: "nocomp-test", + AgentID: "nocomp-test", + Framework: "forge", + ModelProvider: "openai", + EnvVars: map[string]string{}, + NonInteractive: true, + } + if err := scaffold(opts2); err != nil { + t.Fatalf("scaffold error: %v", err) + } + content2, err := os.ReadFile(filepath.Join("nocomp-test", "forge.yaml")) + if err != nil { + t.Fatalf("reading forge.yaml: %v", err) + } + if strings.Contains(string(content2), "compression:") { + t.Errorf("forge.yaml should not contain a compression block by default:\n%s", content2) + } +} + func TestScaffold_GitignoreIncludesEnv(t *testing.T) { tmpDir := t.TempDir() origDir, _ := os.Getwd() diff --git a/forge-cli/cmd/run.go b/forge-cli/cmd/run.go index 3e4a59bf..83cfdf29 100644 --- a/forge-cli/cmd/run.go +++ b/forge-cli/cmd/run.go @@ -36,6 +36,11 @@ var ( runAuthURL string runAuthOrgID string runCORSOrigins string + // runCompression toggles reversible context compression. Applied only + // when the flag is explicitly passed (cmd.Flags().Changed) so the + // yaml/env resolution is untouched otherwise; --compression=false + // force-disables even when forge.yaml enables it. + runCompression bool // FWS-7 audit export sink flags (issue #95). Default zero means // "stderr only" — fully backward-compatible with pre-FWS-7. The @@ -90,6 +95,7 @@ func init() { runCmd.Flags().BoolVar(&runNoGuardrails, "no-guardrails", false, "disable all guardrail enforcement") runCmd.Flags().StringVar(&runModel, "model", "", "override model name (sets MODEL_NAME env var)") runCmd.Flags().StringVar(&runProvider, "provider", "", "LLM provider (openai, anthropic, ollama)") + runCmd.Flags().BoolVar(&runCompression, "compression", false, "enable reversible context compression; --compression=false forces it off (overrides forge.yaml; sets FORGE_COMPRESSION)") runCmd.Flags().StringVar(&runEnvFile, "env", ".env", "path to .env file") runCmd.Flags().StringVar(&runWithChannels, "with", "", "comma-separated channel adapters to start (e.g. slack,telegram)") runCmd.Flags().BoolVar(&runNoAuth, "no-auth", false, "disable bearer token authentication (localhost only)") @@ -137,6 +143,14 @@ func runRun(cmd *cobra.Command, args []string) error { return err } + // --compression / --compression=false → FORGE_COMPRESSION env, which the + // runner resolves with highest precedence (flag > env > forge.yaml). + // Only when explicitly passed, so absent flag leaves yaml/env behavior + // untouched. Same pattern as --model → MODEL_NAME. + if cmd.Flags().Changed("compression") { + _ = os.Setenv("FORGE_COMPRESSION", strconv.FormatBool(runCompression)) + } + activeChannels := parseChannels(runWithChannels) enforceGuardrails := runEnforceGuardrails diff --git a/forge-cli/cmd/serve.go b/forge-cli/cmd/serve.go index 3d260697..84c99009 100644 --- a/forge-cli/cmd/serve.go +++ b/forge-cli/cmd/serve.go @@ -33,6 +33,7 @@ var ( serveNoGuardrails bool serveModel string serveProvider string + serveCompression bool serveEnvFile string serveWithChannels string serveNoAuth bool @@ -115,6 +116,7 @@ func registerServeFlags(cmd *cobra.Command) { cmd.Flags().BoolVar(&serveNoGuardrails, "no-guardrails", false, "disable all guardrail enforcement") cmd.Flags().StringVar(&serveModel, "model", "", "override model name (sets MODEL_NAME env var)") cmd.Flags().StringVar(&serveProvider, "provider", "", "LLM provider (openai, anthropic, ollama)") + cmd.Flags().BoolVar(&serveCompression, "compression", false, "enable reversible context compression; --compression=false forces it off (forwarded to the daemon)") cmd.Flags().StringVar(&serveEnvFile, "env", ".env", "path to .env file") cmd.Flags().StringVar(&serveWithChannels, "with", "", "comma-separated channel adapters to start (e.g. slack,telegram)") cmd.Flags().BoolVar(&serveNoAuth, "no-auth", false, "disable bearer token authentication (localhost only)") @@ -213,6 +215,11 @@ func serveStartRun(cmd *cobra.Command, args []string) error { if serveProvider != "" { runArgs = append(runArgs, "--provider", serveProvider) } + // Forward only when explicitly passed — mirrors the run flag's + // tri-state (absent = yaml/env decide, =true / =false override). + if cmd.Flags().Changed("compression") { + runArgs = append(runArgs, "--compression="+strconv.FormatBool(serveCompression)) + } if serveEnvFile != ".env" { runArgs = append(runArgs, "--env", serveEnvFile) } diff --git a/forge-cli/internal/tui/steps/compression_step.go b/forge-cli/internal/tui/steps/compression_step.go new file mode 100644 index 00000000..00b21d53 --- /dev/null +++ b/forge-cli/internal/tui/steps/compression_step.go @@ -0,0 +1,94 @@ +package steps + +import ( + tea "github.com/charmbracelet/bubbletea" + + "github.com/initializ/forge/forge-cli/internal/tui" + "github.com/initializ/forge/forge-cli/internal/tui/components" +) + +// CompressionStep asks whether to enable reversible context compression +// (ctxzip). When enabled, bulky tool outputs are compressed before reaching +// the LLM and everything dropped stays retrievable via the context_expand +// tool — the wizard writes `compression.enabled: true` into forge.yaml. +type CompressionStep struct { + styles *tui.StyleSet + selector components.SingleSelect + enabled bool + complete bool +} + +// NewCompressionStep creates the compression on/off step. +func NewCompressionStep(styles *tui.StyleSet) *CompressionStep { + items := []components.SingleSelectItem{ + { + Label: "Enabled", + Value: "enabled", + Icon: "🗜️", + Description: "Compress bulky tool outputs before they reach the LLM (typically " + + "60-95% fewer tokens on large results). Reversible: dropped content stays " + + "retrievable via the context_expand tool.", + }, + { + Label: "Disabled", + Value: "disabled", + Icon: "📄", + Description: "Send tool outputs to the LLM verbatim. Enable later with compression.enabled: true in forge.yaml or FORGE_COMPRESSION=true.", + }, + } + + selector := components.NewSingleSelect( + items, + styles.Theme.Accent, + styles.Theme.Primary, + styles.Theme.Secondary, + styles.Theme.Dim, + styles.Theme.Border, + styles.Theme.ActiveBorder, + styles.Theme.ActiveBg, + styles.KbdKey, + styles.KbdDesc, + ) + + return &CompressionStep{styles: styles, selector: selector} +} + +func (s *CompressionStep) Title() string { return "Context Compression" } +func (s *CompressionStep) Icon() string { return "🗜️" } + +func (s *CompressionStep) Init() tea.Cmd { + return s.selector.Init() +} + +func (s *CompressionStep) Update(msg tea.Msg) (tui.Step, tea.Cmd) { + if s.complete { + return s, nil + } + + updated, cmd := s.selector.Update(msg) + s.selector = updated + + if s.selector.Done() { + _, val := s.selector.Selected() + s.enabled = val == "enabled" + s.complete = true + } + return s, cmd +} + +func (s *CompressionStep) View(width int) string { + return s.selector.View(width) +} + +func (s *CompressionStep) Complete() bool { return s.complete } + +func (s *CompressionStep) Summary() string { + if s.enabled { + return "Enabled — reversible, originals retrievable via context_expand" + } + return "Disabled" +} + +func (s *CompressionStep) Apply(ctx *tui.WizardContext) { + ctx.Compression = s.enabled +} diff --git a/forge-cli/internal/tui/wizard.go b/forge-cli/internal/tui/wizard.go index d22f70c8..25e461f0 100644 --- a/forge-cli/internal/tui/wizard.go +++ b/forge-cli/internal/tui/wizard.go @@ -40,6 +40,12 @@ type WizardContext struct { CustomShape string EnvVars map[string]string + // Compression enables reversible context compression (ctxzip): bulky + // tool outputs are compressed before reaching the LLM and stay + // retrievable via the context_expand tool. Written to forge.yaml as + // compression.enabled. + Compression bool + // AuthMode is the user's selection from the auth step: // "" — wizard step did not run // "none" — anonymous access (no auth: block written) diff --git a/forge-cli/templates/init/forge.yaml.tmpl b/forge-cli/templates/init/forge.yaml.tmpl index ddd2ede3..cca69161 100644 --- a/forge-cli/templates/init/forge.yaml.tmpl +++ b/forge-cli/templates/init/forge.yaml.tmpl @@ -41,6 +41,14 @@ builtin_tools: - {{.}} {{- end}} {{- end}} +{{- if .Compression}} + +# Reversible context compression (ctxzip): bulky tool outputs are compressed +# before reaching the LLM; dropped content stays retrievable via the +# context_expand tool. Tune with keep_patterns / ttl / cache_hints. +compression: + enabled: true +{{- end}} {{- if .EgressDomains}} egress: From bc7469febb817e4848bf977326897bf2f96b86e3 Mon Sep 17 00:00:00 2001 From: MK Date: Sat, 4 Jul 2026 17:23:23 -0400 Subject: [PATCH 09/16] docs: sync documentation for context compression (/sync-docs) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit New docs/core-concepts/context-compression.md — the feature's home: problem statement, pipeline diagram, keep-floor layers, configuration and precedence, provider cache hints, observability, failure posture. Per the sync-docs mapping, updates ripple to: - forge-yaml-schema.md — compression block in the full schema plus a dedicated reference section with field table - cli-reference.md — --compression rows in the init/run/serve flag tables; wizard step order documented under forge init - runtime-engine.md — Context Compression section describing the three loop seams (hook after guardrails, client wrapper below the fallback chain, context_expand tool) and the cache-stability posture - audit-logging.md — context_compressed / context_expanded event rows; invocation_complete row gains the per-invocation compression fields - tools-and-builtins.md — context_expand in the builtin table plus a Context Expansion Tool section (hash tolerance, miss guidance) - environment-variables.md — FORGE_COMPRESSION - README.md — Context Compression row in the documentation table - .claude/skills/forge.md — swept sections 8 (memory), 13 (CLI), 14 (schema), 17 (audit reference), 19 (docs map); ToC unchanged (no new numbered sections) Link check: 0 broken links across README + 55 docs files. Co-Authored-By: Claude Opus 4.8 (1M context) --- .claude/skills/forge.md | 28 +++++- README.md | 1 + docs/core-concepts/context-compression.md | 115 ++++++++++++++++++++++ docs/core-concepts/runtime-engine.md | 4 + docs/core-concepts/tools-and-builtins.md | 5 + docs/reference/cli-reference.md | 5 +- docs/reference/environment-variables.md | 1 + docs/reference/forge-yaml-schema.md | 24 +++++ docs/security/audit-logging.md | 4 +- 9 files changed, 181 insertions(+), 6 deletions(-) create mode 100644 docs/core-concepts/context-compression.md diff --git a/.claude/skills/forge.md b/.claude/skills/forge.md index 5fe13d59..98d88187 100644 --- a/.claude/skills/forge.md +++ b/.claude/skills/forge.md @@ -351,7 +351,19 @@ caps, compaction triggers. Embedding provider auto-detects from the LLM provider (Anthropic → `voyage-3` family; OpenAI → `text-embedding-3-small`) unless `memory.embedding_provider` is explicit. -**Read**: `docs/core-concepts/memory-system.md`. +Opt-in **context compression** (ctxzip): when `compression.enabled` is +set, large tool outputs are compressed reversibly before reaching the +LLM — an `AfterToolExec` hook compresses once at production time, an +`llm.Client` wrapper compresses each request's live zone, and the +`context_expand` builtin retrieves offloaded originals by +`<>` marker from a bbolt store (`.forge/ctxzip.db`, +30-min TTL). `compression.keep_patterns` declares domain vocabulary +that is never dropped; `compression.cache_hints` injects provider +prompt-cache primitives (anthropic `cache_control`, openai +`prompt_cache_key`). Fail-open: any error runs uncompressed. + +**Read**: `docs/core-concepts/memory-system.md`, +`docs/core-concepts/context-compression.md`. --- @@ -768,10 +780,10 @@ Full reference: `docs/reference/cli-reference.md`. | Subcommand | Purpose | Key flags | |---|---|---| -| `forge init` | Scaffold a new agent: `forge.yaml`, `.env`, `SKILL.md`, `guardrails.json`. Interactive TUI by default; `--non-interactive` for CI | `--model-provider`, `--model-name`, `--channels`, `--auth`, `--from-skills` | +| `forge init` | Scaffold a new agent: `forge.yaml`, `.env`, `SKILL.md`, `guardrails.json`. Interactive TUI by default; `--non-interactive` for CI | `--model-provider`, `--model-name`, `--channels`, `--auth`, `--from-skills`, `--compression` | | `forge build` | Run the build pipeline → `.forge-output/agent.json` + container Dockerfile + K8s manifests + (optional) signature | `--output-dir`, `--sign` | | `forge validate` | Lint `forge.yaml` + SKILL.md. `--platform-policy=PATH` lints a policy file standalone | `--strict`, `--command-compat`, `--platform-policy` | -| `forge run` | Dev-mode A2A server with hot-reload | `--port`, `--host`, `--with slack,telegram`, `--mock-tools`, `--no-auth`, `--cors-origins`, `--audit-socket`, `--audit-http-endpoint`, `--rate-limit-*`, `--otel-enabled`, `--otel-endpoint`, `--otel-sampler` | +| `forge run` | Dev-mode A2A server with hot-reload | `--port`, `--host`, `--with slack,telegram`, `--mock-tools`, `--no-auth`, `--cors-origins`, `--audit-socket`, `--audit-http-endpoint`, `--rate-limit-*`, `--otel-enabled`, `--otel-endpoint`, `--otel-sampler`, `--compression[=false]` | | `forge serve start \| stop \| status \| logs` | Daemonized A2A server (forks `forge run`). Forwards CLI flags + env to the child | `--port`, `--shutdown-timeout`, `--with` | | `forge export` | Export `agent.json` for registry upload | | | `forge package` | Generate Dockerfile + Kubernetes manifests + `egress_allowlist.json`. `--prod` rejects `dev-open` egress + dev-only tools | `--registry`, `--tag`, `--base`, `--prod` | @@ -852,6 +864,11 @@ memory: long_term: false embedding_provider: openai +compression: + enabled: false # reversible context compression (ctxzip) + keep_patterns: [] # never-drop vocabulary + cache_hints: true # provider prompt-cache primitives + mcp: token_store_path: ~/.forge/mcp-tokens.enc servers: @@ -1042,7 +1059,9 @@ when OTel tracing is enabled (OTel v1 / Phase 4 / #105). Both use | `EventMCPToolConflict` | `mcp_tool_conflict` | Namespaced tool collision detected | | `EventMCPTokenRefresh` | `mcp_token_refresh` | OAuth 2.1 token refresh result | | `EventAgentCardPublished` | `agent_card_published` | Agent Card finalized at startup / hot-reload; `name`, `version`, `protocol_version`, `url`, `skill_count`, `capabilities`, `security_schemes`, `card_size_bytes`, `card_sha256` (FWS-1) | -| `AuditInvocationComplete` | `invocation_complete` | A2A invocation closed; `duration_ms`, `input_tokens_total`, `output_tokens_total`, `llm_call_count`, `model`, `provider` (FWS-3) | +| `context_compressed` | `context_compressed` | Context compression shrank content; `seam` (`tool_output` / `request`), `tool`, `tokens_before` / `tokens_after` / `saved_tokens` + running totals (tokenizer estimates) | +| `context_expanded` | `context_expanded` | Model retrieved offloaded content via `context_expand`; `hash`, `hit`, `bytes` + running totals | +| `AuditInvocationComplete` | `invocation_complete` | A2A invocation closed; `duration_ms`, `input_tokens_total`, `output_tokens_total`, `llm_call_count`, `model`, `provider` (FWS-3); with compression enabled also `compression_saved_tokens_total`, `compression_count`, `expansion_count` | | `AuditInvocationCancelled` | `invocation_cancelled` | A2A invocation cancelled via `tasks/cancel`; classified `reason` + partial token totals (FWS-4) | | `AuditTaskAdmissionDenied` | `task_admission_denied` | Inbound `tasks/send` denied by the platform admission middleware (#201; opt-in via `FORGE_ADMISSION_URL` + `FORGE_PLATFORM_TOKEN`); `reason`, `scope`, `window`, `reset_at`, `cached`. Caller sees HTTP 402 Payment Required. | | `AuditPolicyLoaded` | `policy_loaded` | One per non-empty policy layer at startup; `layer`, `source`, per-list size counters (FWS-5/6) | @@ -1103,6 +1122,7 @@ docs/ │ ├── skill-md-format.md ← SKILL.md schema │ ├── channels.md │ ├── memory-system.md +│ ├── context-compression.md ← reversible tool-output compression │ ├── scheduling.md │ └── observability-tracing.md ← OTel v1 (#108) — spans, propagation, audit cross-link ├── security/ diff --git a/README.md b/README.md index e3e44933..540ddad9 100644 --- a/README.md +++ b/README.md @@ -86,6 +86,7 @@ You write a `SKILL.md`. Forge compiles it into a secure, runnable agent with egr | [Tools](docs/core-concepts/tools-and-builtins.md) | Built-in tools, adapters, and custom tools | | [Runtime](docs/core-concepts/runtime-engine.md) | LLM providers, fallback chains, running modes | | [Memory](docs/core-concepts/memory-system.md) | Session persistence and long-term memory | +| [Context Compression](docs/core-concepts/context-compression.md) | Reversible compression of bulky tool outputs — fewer tokens, nothing lost | | [Channels](docs/core-concepts/channels.md) | Slack and Telegram adapter setup | | [Scheduling](docs/core-concepts/scheduling.md) | Cron configuration and schedule tools | | [Tracing](docs/core-concepts/observability-tracing.md) | OpenTelemetry distributed tracing — spans, propagation, audit cross-link | diff --git a/docs/core-concepts/context-compression.md b/docs/core-concepts/context-compression.md new file mode 100644 index 00000000..dc695261 --- /dev/null +++ b/docs/core-concepts/context-compression.md @@ -0,0 +1,115 @@ +--- +title: "Context Compression" +description: "Reversible compression of bulky tool outputs — fewer tokens, nothing lost." +order: 6 +--- + +Forge can compress bulky tool outputs before they reach the LLM — reversibly: everything dropped stays retrievable, so compression is lossy on the wire but lossless end-to-end. + +Powered by [ctxzip](https://github.com/initializ/ctxzip). Off by default; enable per agent in `forge.yaml`, per run with a flag, or at scaffold time in the init wizard. + +## The problem it solves + +Agent tool outputs are dominated by repetition: 149 pods that are `Running` and one that is `CrashLoopBackOff`; hundreds of log lines differing only by timestamp; JSON list responses where the model needs three rows. Without compression these outputs either flood the context window or get **truncated** — destroying whatever fell past the cut, which is frequently the one row that mattered. + +Compression inverts the tradeoff: keep what matters (errors, anomalies, query-relevant rows, boundaries), offload the rest to a local store, and let the model retrieve the original if it turns out to need it. + +## How it works + +``` +tool executes + │ + ▼ +AfterToolExec hook ──── output ≥ 2 KB? ──── compress once, at production time + │ dropped content → .forge/ctxzip.db + │ replaced by <> marker + ▼ +Memory (compressed bytes never change → provider prompt caches stay warm) + │ + ▼ +LLM client wrapper ──── compresses the live zone of each request + │ (system prompt + recent turns forwarded byte-identical) + ▼ +LLM sees: [... kept rows, errors intact ...] <> + │ + └─ needs the offloaded data? → calls context_expand(hash) → original returned +``` + +Three pieces, all automatic once enabled: + +| Piece | What it does | +|-------|--------------| +| Tool-output hook | Compresses each large tool result once, before it enters session memory. Error results and small outputs are left verbatim. | +| Client wrapper | Compresses the remaining live zone of each outbound request. Deterministic across turns so historic messages always compress to identical bytes. | +| `context_expand` tool | Registered automatically. The model calls it with a marker's hash to get the original content back. A system-prompt directive teaches every agent what markers are — skills need zero awareness. | + +## What is never dropped + +Fidelity is layered; every layer only ever adds protection: + +1. **Error floor** — content matching error vocabulary (`error`, `fail`, `panic`, `timeout`, `crash`, `backoff`, `oomkilled`, `evicted`, …) is kept verbatim. +2. **`keep_patterns`** — your domain's never-drop vocabulary (see below). +3. **Query anchors** — items matching the conversation's ask survive. +4. **Structure** — head/tail windows and one exemplar of each near-duplicate group. +5. **Reversibility** — everything else is offloaded to the store, not deleted. +6. **Source of truth** — after the store TTL (30 min), the disk or the original command still holds the data; a retrieval miss tells the model to re-run the producing tool. + +## Configuration + +```yaml +# forge.yaml +compression: + enabled: true # default: false + keep_patterns: # domain vocabulary that must never be dropped + - CrashLoopBackOff + - PAYMENT_DECLINED + # store_path: .forge/ctxzip.db # offloaded-originals store (bbolt) + # ttl: 30m # how long originals stay retrievable + # min_tool_output_chars: 2048 # hook floor; smaller outputs untouched + # cache_hints: true # provider prompt-cache hints (defaults to enabled) +``` + +Precedence (most specific wins): + +``` +forge run --compression[=false] > FORGE_COMPRESSION=true|false > compression.enabled > off +``` + +| Surface | Usage | +|---------|-------| +| `forge run --compression` | Enable for one run; `--compression=false` force-disables even when forge.yaml enables it | +| `forge serve --compression[=false]` | Forwarded to the daemon | +| `forge init --compression` | Scaffold a new agent with the block enabled | +| init TUI wizard | "Context Compression" step (between Skills and Auth) | + +## Provider prompt-cache hints + +Compressing the wrong bytes can *cost* tokens by busting the provider's prompt cache, so compression never touches the system prompt, tool definitions, or recent turns, and its output is deterministic across turns. On top of that, `cache_hints` (on by default when compression is enabled) injects each provider's native cache primitives: + +| Provider | Hint | +|----------|------| +| anthropic | `cache_control: {type: ephemeral}` breakpoints on the last tool definition and the system block — caches the stable tools+system prefix across turns. Also applies on the `aws_sigv4` Bedrock-passthrough path. | +| openai / gemini | A stable `prompt_cache_key` derived from (model, system prompt, tool names) — pins cache routing; prefix caching itself is automatic. | + +When `cache_hints` is off, provider wire formats are byte-identical to a build without compression. + +## Observability + +Savings are first-class audit events, not log noise — see [Audit Logging](../security/audit-logging.md) for the event schema: + +- `context_compressed` — per compression: seam, tool, tokens before/after/saved, plus running totals. +- `context_expanded` — per retrieval: hash, hit, bytes — the cost side to net against savings. +- `invocation_complete` gains `compression_saved_tokens_total`, `compression_count`, and `expansion_count`, accumulated per invocation (concurrent tasks never cross-contaminate). + +Token figures are tokenizer estimates (directionally accurate); billed truth remains `llm_call.input_tokens`. A surgical session that produced only small outputs correctly reports `compression_count: 0` — compression is insurance against bulk, not a tax on every call. + +## Failure posture + +Fail-open, always: if the store cannot be opened, a compressor errors, or "compression" would grow a message, the original content is used unchanged. Error tool results are never compressed. An expired retrieval is not a dead end — the model is told to re-run the tool that produced the output. + +## Related + +- [Runtime Engine](runtime-engine.md) — where the hook and client wrapper sit in the agent loop +- [Tools & Builtins](tools-and-builtins.md) — the `context_expand` tool +- [forge.yaml Schema](../reference/forge-yaml-schema.md) — the `compression` block +- [CLI Reference](../reference/cli-reference.md) — flags and wizard step diff --git a/docs/core-concepts/runtime-engine.md b/docs/core-concepts/runtime-engine.md index e58bd72d..c4c93a41 100644 --- a/docs/core-concepts/runtime-engine.md +++ b/docs/core-concepts/runtime-engine.md @@ -277,6 +277,10 @@ The `FilesDir` is set via `LLMExecutorConfig.FilesDir` and made available to too For details on session persistence, context window management, compaction, and long-term memory, see [Memory](memory-system.md). +## Context Compression + +When `compression.enabled` is set, the runner wires reversible context compression (ctxzip) into the loop at three points: an `AfterToolExec` hook compresses large tool outputs once, before they enter memory (registered after the guardrail hooks, so it compresses redacted output); the LLM client is wrapped in a compressing decorator below the fallback chain (so retries and compactor summarization calls are covered too); and the `context_expand` retrieval tool is registered so the model can recover offloaded content by marker hash. A constant system-prompt directive teaches the model what `<>` markers are — individual skills need no awareness. Compression output is deterministic across turns and never touches the system prompt or recent messages, keeping provider prompt caches warm; `compression.cache_hints` additionally injects the provider's native cache primitives (anthropic `cache_control` breakpoints, openai `prompt_cache_key`). See [Context Compression](context-compression.md). + ## Hooks The engine fires hooks at key points in the loop. See [Hooks](hooks.md) for details. diff --git a/docs/core-concepts/tools-and-builtins.md b/docs/core-concepts/tools-and-builtins.md index 6dd490fd..e18e01a4 100644 --- a/docs/core-concepts/tools-and-builtins.md +++ b/docs/core-concepts/tools-and-builtins.md @@ -30,6 +30,7 @@ Tools are capabilities that an LLM agent can invoke during execution. Forge prov | `read_skill` | Load full instructions for an available skill on demand | | `memory_search` | Search long-term memory (when enabled) | | `memory_get` | Read memory files (when enabled) | +| `context_expand` | Retrieve the original content behind a `<>` compression marker (when [compression](context-compression.md) is enabled) | | `cli_execute` | Execute pre-approved CLI binaries | | `schedule_set` | Create or update a recurring cron schedule | | `schedule_list` | List all active and inactive schedules | @@ -180,6 +181,10 @@ When [long-term memory](memory-system.md) is enabled, two additional tools are r These tools allow the agent to recall information from previous sessions. +## Context Expansion Tool + +When [context compression](context-compression.md) is enabled, the `context_expand` tool is registered. Compressed tool outputs carry inline `<>` markers; the model calls `context_expand` with the hash to retrieve the offloaded original from the local store. The tool tolerates imperfect input — a whole marker pasted as the hash, or a truncated hash that uniquely prefixes a recently emitted one — and a miss (expired/evicted entry) returns guidance to re-run the producing tool rather than an error. + ## Development Tools Development tools (`local_shell`, `local_file_browser`, `debug_console`, `test_runner`) are available during `forge run --dev` but are **automatically filtered out** in production builds by the `ToolFilterStage`. diff --git a/docs/reference/cli-reference.md b/docs/reference/cli-reference.md index c02c4e79..0046872b 100644 --- a/docs/reference/cli-reference.md +++ b/docs/reference/cli-reference.md @@ -18,7 +18,7 @@ Complete reference for all Forge CLI commands. ## `forge init` -Initialize a new agent project. +Initialize a new agent project. Without `--non-interactive`, a TUI wizard walks through: name → model provider → fallbacks → channel → tools → skills → context compression → authentication → egress review → summary. ``` forge init [name] [flags] @@ -39,6 +39,7 @@ forge init [name] [flags] | `--org-id` | | | OpenAI Organization ID (enterprise) | | `--from-skills` | | | Path to a SKILL.md file for auto-configuration | | `--non-interactive` | | `false` | Skip interactive prompts | +| `--compression` | | `false` | Enable reversible context compression — writes `compression.enabled: true` to the scaffolded forge.yaml. See [Context Compression](../core-concepts/context-compression.md) | | `--auth` | | | Auth mode: `none`, `oidc`, `http_verifier`, `aws_sigv4`, `gcp_iap`, `azure_ad`, `custom` | | `--auth-issuer` | | | OIDC issuer URL (required with `--auth=oidc`) | | `--auth-audience` | | | OIDC audience (required with `--auth=oidc`) | @@ -217,6 +218,7 @@ forge run [flags] | `--enforce-guardrails` | `false` | Enforce guardrail violations as errors | | `--model` | | Override model name (sets `MODEL_NAME` env var) | | `--provider` | | LLM provider: `openai`, `anthropic`, or `ollama` | +| `--compression` | | Enable reversible context compression; `--compression=false` forces it off. Absent = forge.yaml/env decide (sets `FORGE_COMPRESSION`). See [Context Compression](../core-concepts/context-compression.md) | | `--env` | `.env` | Path to .env file | | `--with` | | Comma-separated channel adapters (e.g., `slack,telegram`) | | `--auth-url` | | External auth provider URL for token validation | @@ -291,6 +293,7 @@ forge serve [start|stop|status|logs] [flags] | `--host` | `127.0.0.1` | Bind address (secure default) | | `--with` | | Channel adapters | | `--cors-origins` | localhost | Comma-separated CORS allowed origins | +| `--compression` | | Enable reversible context compression; `--compression=false` forces it off. Forwarded to the daemon `forge run` only when explicitly passed | ### Examples diff --git a/docs/reference/environment-variables.md b/docs/reference/environment-variables.md index 926b6cd2..f2d00732 100644 --- a/docs/reference/environment-variables.md +++ b/docs/reference/environment-variables.md @@ -12,6 +12,7 @@ order: 3 | `FORGE_MODEL_FALLBACKS` | Fallback chain (e.g., `"anthropic:claude-sonnet-4,gemini"`) | | `FORGE_MEMORY_PERSISTENCE` | Set `false` to disable session persistence | | `FORGE_MEMORY_LONG_TERM` | Set `true` to enable long-term memory | +| `FORGE_COMPRESSION` | Set `true`/`false` to override `compression.enabled` (reversible context compression); the `--compression` flag overrides both | | `FORGE_EMBEDDING_PROVIDER` | Override embedding provider | | `OPENAI_API_KEY` | OpenAI API key | | `OPENAI_ORG_ID` | OpenAI Organization ID (enterprise); overrides `organization_id` in YAML | diff --git a/docs/reference/forge-yaml-schema.md b/docs/reference/forge-yaml-schema.md index 5c55dfca..1328208a 100644 --- a/docs/reference/forge-yaml-schema.md +++ b/docs/reference/forge-yaml-schema.md @@ -156,6 +156,15 @@ memory: keyword_weight: 0.3 # Hybrid search keyword weight decay_half_life_days: 7 # Temporal decay half-life +compression: # Reversible context compression (default: off) + enabled: true # Compress bulky tool outputs (default: false) + keep_patterns: # Domain vocabulary never dropped (case-insensitive substrings) + - CrashLoopBackOff + store_path: ".forge/ctxzip.db" # Offloaded-originals store (bbolt) + ttl: "30m" # How long originals stay retrievable + min_tool_output_chars: 2048 # Hook floor; smaller outputs untouched + cache_hints: true # Provider prompt-cache hints (defaults to enabled) + guardrails_path: "guardrails.json" # Path to guardrails config (default: "guardrails.json") schedules: # Recurring scheduled tasks (optional) @@ -297,6 +306,21 @@ workflow_propagation: The matcher is consulted by a `RoundTripper` wrapper around the egress transport at runner startup, so every built-in HTTP tool (`http_request`, `webhook_call`, `web_search_*`) inherits the auto-apply without per-tool changes. Empty config = zero-overhead pass-through. +## `compression` — reversible context compression + +Compresses bulky tool outputs before they reach the LLM; dropped content is stored locally and retrievable via the `context_expand` tool, so compression is lossy on the wire but lossless end-to-end. Off by default. + +| Field | Default | Description | +|-------|---------|-------------| +| `enabled` | `false` | Master switch. Env `FORGE_COMPRESSION=true\|false` overrides; `forge run --compression[=false]` overrides both | +| `keep_patterns` | — | Case-insensitive substrings never dropped (domain error codes, state words). Union with the built-in error floor — entries only add protection | +| `store_path` | `.forge/ctxzip.db` | bbolt store for offloaded originals (created 0600) | +| `ttl` | `30m` | How long originals stay retrievable; after expiry the model is told to re-run the producing tool | +| `min_tool_output_chars` | `2048` | Tool outputs below this size are never touched | +| `cache_hints` | value of `enabled` | Inject provider prompt-cache primitives (anthropic `cache_control`, openai `prompt_cache_key`) | + +See [Context Compression](../core-concepts/context-compression.md) for how the pieces fit together. + ## `security` — build-time security knobs ```yaml diff --git a/docs/security/audit-logging.md b/docs/security/audit-logging.md index 93625eb1..bff9ff3f 100644 --- a/docs/security/audit-logging.md +++ b/docs/security/audit-logging.md @@ -19,10 +19,12 @@ All runtime security events are emitted as structured NDJSON to stderr with corr | `egress_blocked` | Outbound request blocked (with domain, mode) | | `llm_call` | LLM API call completed (with `input_tokens`, `output_tokens`, `model`, `provider`, `duration_ms`, `request_id`). See [Token usage and duration](#token-usage-and-execution-duration). | | `llm_call_cancelled` | Streaming LLM call cancelled mid-flight; carries partial token counts captured up to cancellation. | -| `invocation_complete` | A2A invocation finished (auth → dispatch → engine → response). Carries `duration_ms` (wall-clock) plus aggregated `input_tokens_total` / `output_tokens_total` / `llm_call_count` / `model` / `provider`. | +| `invocation_complete` | A2A invocation finished (auth → dispatch → engine → response). Carries `duration_ms` (wall-clock) plus aggregated `input_tokens_total` / `output_tokens_total` / `llm_call_count` / `model` / `provider`. When [context compression](../core-concepts/context-compression.md) is enabled it also carries `compression_saved_tokens_total` / `compression_count` (and `expansion_count` when nonzero), accumulated per invocation by correlation ID so concurrent tasks never cross-contaminate. | | `invocation_cancelled` | A2A invocation cancelled mid-flight via `tasks/cancel` (or internal cancellation like parent ctx deadline). Carries `fields.reason` (one of `workflow_failure` / `cost_limit_exceeded` / `timeout` / `external_signal`), `duration_ms` up to cancellation, and any partial token totals consumed before the signal. See [Cancellation](#cancellation). | | `task_admission_denied` | A new inbound `tasks/send` was rejected by the platform admission middleware (issue #201; opt-in via `FORGE_ADMISSION_URL` + `FORGE_PLATFORM_TOKEN`). Carries `fields.reason` (platform-defined: `cost_limit_exceeded`, `billing_overdue`, …), `fields.scope` (`agent` / `workspace` / `org`), `fields.window` (`hourly` / `daily` / `monthly` / `billing_cycle`), `fields.reset_at` (RFC 3339), and `fields.cached` (`true` when served from the 5s per-agent cache). Caller observes HTTP 402 Payment Required with `Retry-After`. See [Platform Admission Hook](admission.md). | | `guardrail_check` | Guardrail mask / block / warn decision. Carries `fields.gate` (`input` / `context` / `tool_call` / `output` / `stream` — sourced from the library `Result.Gate`), `fields.decision` (`masked` / `warned` / `blocked`), `fields.guardrail` + `fields.category` from the triggering violation, and `fields.violation_count`. `fields.tool` is present on `tool_call` and on `output` events for tool return text. With `FORGE_GUARDRAIL_CAPTURE_EVIDENCE=true` operators also opt into `fields.evidence` carrying the redacted + truncated triggering text. See [Guardrails — Audit Events](guardrails.md#audit-events). | +| `context_compressed` | [Context compression](../core-concepts/context-compression.md) shrank content before it reached the LLM. Carries `fields.seam` (`tool_output` from the AfterToolExec hook / `request` from the client wrapper), `fields.tool`, `tokens_before` / `tokens_after` / `saved_tokens`, plus running totals `total_saved_tokens` / `total_compressions` / `total_expansions` so any single event shows the cumulative picture. Token figures are tokenizer estimates; billed truth stays in `llm_call.input_tokens`. | +| `context_expanded` | The model retrieved offloaded content via the `context_expand` tool. Carries `fields.hash`, `hit` (`false` = expired/evicted), `bytes`, and the same running totals — expansions are the cost side auditors net against savings. | | `auth_verify` | Inbound request authenticated successfully (with `provider`, `user_id`, `org_id`, `token_kind`) | | `auth_fail` | Inbound request rejected (with `reason`, `token_kind`) | | `agent_card_published` | Agent Card finalized at startup or hot-reload (with `name`, `version`, `protocol_version`, `url`, `skill_count`, `capabilities`, `security_schemes`, `card_size_bytes`, `card_sha256`). See [Agent Card reference](../reference/a2a-agent-card.md). | From deb06fdf5f285c24421b0f80215b9f809a560e1e Mon Sep 17 00:00:00 2001 From: MK Date: Sat, 4 Jul 2026 19:33:44 -0400 Subject: [PATCH 10/16] fix(compress): pop compression totals on ALL invocation_complete paths; bound runtime maps MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Addresses the PR #241 review (blocking + secondary items): - 🔴 Streaming perInvocation leak + missing metrics: TakeInvocationTotals was only called in executeTask, but invocation_complete is emitted from THREE sites — the tasks/sendSubscribe JSON-RPC SSE and REST streaming handlers ran their own emission without popping the bucket, so every streaming invocation that compressed leaked its correlation bucket permanently AND its invocation_complete lacked the compression fields. The pop+populate block is now a shared helper (appendCompressionFields, documented as required at every emission site) called from all three. Pinned by TestAppendCompressionFields_PopsAndPopulates: fields populated, pop is one-shot (no double-count), nil-safe. - 🟡 Leak backstops: perInvocation is bounded to 1024 buckets with oldest-touched eviction (a future missed pop can no longer grow unbounded); the recent-marker prefix-resolution set is bounded to 2048 with oldest-emitted eviction. Both pinned by tests. - 🟡 Inflation guards tightened from == 0 to <= 0 at both seams — ctxzip clamps SavedTokens at zero today, but the guard must not silently apply inflated output if that contract ever changes. - 🟡 bbolt single-writer constraint documented in context-compression.md: the store holds an exclusive flock with a 5s open timeout (verified in ctxzip's NewBoltStore), so a second process fails open and runs uncompressed; each replica should get its own store_path. Co-Authored-By: Claude Opus 4.8 (1M context) --- docs/core-concepts/context-compression.md | 2 + forge-cli/runtime/runner.go | 43 ++++++++--- forge-cli/runtime/runner_compression_test.go | 68 +++++++++++++++++ forge-core/compress/client.go | 3 +- forge-core/compress/compress.go | 79 ++++++++++++++++---- forge-core/compress/compress_test.go | 48 ++++++++++++ forge-core/compress/hook.go | 5 +- 7 files changed, 218 insertions(+), 30 deletions(-) create mode 100644 forge-cli/runtime/runner_compression_test.go diff --git a/docs/core-concepts/context-compression.md b/docs/core-concepts/context-compression.md index dc695261..14160a8f 100644 --- a/docs/core-concepts/context-compression.md +++ b/docs/core-concepts/context-compression.md @@ -107,6 +107,8 @@ Token figures are tokenizer estimates (directionally accurate); billed truth rem Fail-open, always: if the store cannot be opened, a compressor errors, or "compression" would grow a message, the original content is used unchanged. Error tool results are never compressed. An expired retrieval is not a dead end — the model is told to re-run the tool that produced the output. +**Single-writer store.** The bbolt store at `store_path` holds an exclusive file lock — one store per process. A second process pointing at the same file (two replicas on a shared volume, or `forge run` alongside `forge serve` in the same directory) fails to acquire the lock after a 5-second timeout and that process runs uncompressed (fail-open, with a startup warning). Give each replica its own `store_path` — offloaded originals are only ever retrieved by the process that offloaded them, so the store has no reason to be shared. + ## Related - [Runtime Engine](runtime-engine.md) — where the hook and client wrapper sit in the agent loop diff --git a/forge-cli/runtime/runner.go b/forge-cli/runtime/runner.go index 49aa22f8..052e6d83 100644 --- a/forge-cli/runtime/runner.go +++ b/forge-cli/runtime/runner.go @@ -1293,6 +1293,9 @@ func (r *Runner) registerHandlers(srv *server.Server, executor coreruntime.Agent fields["provider"] = snap.PrimaryProvider } } + // Pops the per-correlation compression bucket — required on + // every invocation_complete path (leak otherwise). + r.appendCompressionFields(ctx, fields) auditLogger.EmitInvocationComplete(ctx, snap.InvocationDuration, fields) }() @@ -1544,18 +1547,8 @@ func (r *Runner) executeTask( fields["provider"] = snap.PrimaryProvider } } - // Per-invocation compression savings (tokenizer estimates), popped - // from the runtime by this invocation's correlation ID so concurrent - // tasks don't cross-contaminate. Present whenever compression is - // enabled — zero values mean "on, but nothing was worth compressing". - if r.compression != nil { - ct := r.compression.TakeInvocationTotals(ctx) - fields["compression_saved_tokens_total"] = ct.SavedTokens - fields["compression_count"] = ct.Compressions - if ct.Expansions > 0 { - fields["expansion_count"] = ct.Expansions - } - } + // Per-invocation compression savings — see appendCompressionFields. + r.appendCompressionFields(ctx, fields) if task.Status.State == a2a.TaskStateCanceled { auditLogger.EmitInvocationCancelled(ctx, coreruntime.CancellationReasonFromCause(ctx), @@ -1818,6 +1811,9 @@ func (r *Runner) registerRESTHandlers(srv *server.Server, executor coreruntime.A fields["provider"] = snap.PrimaryProvider } } + // Pops the per-correlation compression bucket — required on + // every invocation_complete path (leak otherwise). + r.appendCompressionFields(ctx, fields) auditLogger.EmitInvocationComplete(ctx, snap.InvocationDuration, fields) }() @@ -3401,6 +3397,29 @@ func (r *Runner) initLongTermMemory(ctx context.Context, mc *coreruntime.ModelCo return mgr } +// appendCompressionFields pops this invocation's compression savings (keyed +// by the ctx's correlation ID) and adds them to an invocation_complete / +// invocation_cancelled fields map: compression_saved_tokens_total, +// compression_count, and expansion_count (when nonzero). Values are tokenizer +// estimates; zeros mean "compression on, nothing worth compressing". +// +// MUST be called exactly once per invocation, at the emission site — there +// are THREE (executeTask plus both tasks/sendSubscribe streaming handlers, +// JSON-RPC SSE and REST). Missing a site both leaks the per-correlation +// bucket and drops the metrics from that path's invocation_complete +// (PR #241 review finding). No-op when compression is disabled. +func (r *Runner) appendCompressionFields(ctx context.Context, fields map[string]any) { + if r.compression == nil { + return + } + ct := r.compression.TakeInvocationTotals(ctx) + fields["compression_saved_tokens_total"] = ct.SavedTokens + fields["compression_count"] = ct.Compressions + if ct.Expansions > 0 { + fields["expansion_count"] = ct.Expansions + } +} + // compressionEnabled reports whether reversible context compression is on. // forge.yaml `compression.enabled` sets the default; FORGE_COMPRESSION=true / // =false overrides it either way (matching the FORGE_MEMORY_* env pattern). diff --git a/forge-cli/runtime/runner_compression_test.go b/forge-cli/runtime/runner_compression_test.go new file mode 100644 index 00000000..58724e0d --- /dev/null +++ b/forge-cli/runtime/runner_compression_test.go @@ -0,0 +1,68 @@ +package runtime + +import ( + "context" + "encoding/json" + "fmt" + "path/filepath" + "testing" + + "github.com/initializ/forge/forge-core/compress" + coreruntime "github.com/initializ/forge/forge-core/runtime" +) + +// PR #241 review: invocation_complete is emitted from THREE sites +// (executeTask + both sendSubscribe streaming handlers), and every one must +// pop the per-correlation compression bucket via appendCompressionFields — +// a missed site leaks the bucket and drops the metrics. This pins the +// helper's contract: fields populated, bucket popped exactly once, nil-safe. +func TestAppendCompressionFields_PopsAndPopulates(t *testing.T) { + comp, err := compress.New(compress.Config{ + StorePath: filepath.Join(t.TempDir(), "ctxzip.db"), + }) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = comp.Close() }) + r := &Runner{compression: comp} + + // Simulate a streaming invocation that compressed one big tool output. + ctx := coreruntime.WithCorrelationID(context.Background(), "sse-task-1") + items := make([]map[string]any, 80) + for i := range items { + items[i] = map[string]any{"id": fmt.Sprintf("r-%03d", i), "state": "nominal", "zone": "us-east-1"} + } + blob, _ := json.Marshal(items) + hctx := &coreruntime.HookContext{ToolName: "list", ToolOutput: string(blob)} + if err := comp.AfterToolExecHook()(ctx, hctx); err != nil { + t.Fatal(err) + } + if hctx.ToolOutput == string(blob) { + t.Fatal("fixture did not compress — test cannot exercise the bucket") + } + + // First call: fields populated, bucket popped. + fields := map[string]any{} + r.appendCompressionFields(ctx, fields) + if fields["compression_count"].(int64) != 1 { + t.Fatalf("compression_count = %v, want 1", fields["compression_count"]) + } + if fields["compression_saved_tokens_total"].(int64) <= 0 { + t.Fatalf("compression_saved_tokens_total = %v, want > 0", fields["compression_saved_tokens_total"]) + } + + // Second call for the same invocation: bucket already popped → zeros + // (proves no double-count if two sites ever fire for one invocation). + fields2 := map[string]any{} + r.appendCompressionFields(ctx, fields2) + if fields2["compression_count"].(int64) != 0 { + t.Fatalf("second pop should be empty, got %v", fields2["compression_count"]) + } + + // Nil compression runtime: no fields, no panic. + fields3 := map[string]any{} + (&Runner{}).appendCompressionFields(ctx, fields3) + if len(fields3) != 0 { + t.Fatalf("nil runtime should add no fields: %v", fields3) + } +} diff --git a/forge-core/compress/client.go b/forge-core/compress/client.go index f138e1e6..d6d88498 100644 --- a/forge-core/compress/client.go +++ b/forge-core/compress/client.go @@ -66,8 +66,9 @@ func (c *compressingClient) compressRequest(ctx context.Context, req *llm.ChatRe opts.SkipNames = map[string]bool{expandToolName: true} opts.MustKeep = c.rt.keep + // <= 0 for the same reason as the hook: never apply inflated output. res, err := ctxzip.Compress(zmsgs, opts) - if err != nil || res == nil || res.SavedTokens() == 0 { + if err != nil || res == nil || res.SavedTokens() <= 0 { return req } diff --git a/forge-core/compress/compress.go b/forge-core/compress/compress.go index c37aaee7..fa80a344 100644 --- a/forge-core/compress/compress.go +++ b/forge-core/compress/compress.go @@ -97,17 +97,42 @@ type Runtime struct { // recent remembers marker hashes this process emitted so the expand tool // can resolve a unique prefix when the model transcribes a hash - // imperfectly (observed live: models truncate hex hashes). - mu sync.Mutex - recent map[string]struct{} - totals SavingsTotals + // imperfectly (observed live: models truncate hex hashes). Bounded to + // maxRecentMarkers, evicting oldest-emitted first — only *recent* + // markers matter for transcription repair. + mu sync.Mutex + recent map[string]struct{} + recentOrder []string + totals SavingsTotals // perInvocation accumulates savings keyed by correlation ID so // invocation_complete can report this invocation's savings without // cross-contamination from concurrent tasks. Entries are popped by - // TakeInvocationTotals at the invocation's response boundary. - perInvocation map[string]*SavingsTotals + // TakeInvocationTotals at the invocation's response boundary; as a + // leak backstop for any emission path that misses the pop, the map is + // bounded to maxInvocationBuckets with oldest-touched eviction. + perInvocation map[string]*invocationBucket } +// invocationBucket carries one invocation's savings plus a touch time for +// oldest-first eviction at the cap. +type invocationBucket struct { + totals SavingsTotals + touched time.Time +} + +// Bounds for the Runtime's process-lifetime maps (PR #241 review). +const ( + // maxInvocationBuckets bounds perInvocation. Normal operation pops each + // bucket at invocation end, so this only matters if an emission path + // misses the pop — 1024 in-flight correlation IDs is far beyond any + // realistic concurrency. + maxInvocationBuckets = 1024 + // maxRecentMarkers bounds the prefix-resolution set. 2048 markers is + // hours of heavy compression; older hashes are no longer "recent" and + // exact-hash retrieval from the store still works without them. + maxRecentMarkers = 2048 +) + // SavingsTotals is the process-lifetime savings picture. Token figures are // tokenizer estimates. type SavingsTotals struct { @@ -140,29 +165,42 @@ func (r *Runtime) TakeInvocationTotals(ctx context.Context) SavingsTotals { } r.mu.Lock() defer r.mu.Unlock() - if t, ok := r.perInvocation[cid]; ok { + if b, ok := r.perInvocation[cid]; ok { delete(r.perInvocation, cid) - return *t + return b.totals } return SavingsTotals{} } -// bumpInvocation accumulates into the ctx's per-invocation bucket. Caller -// holds r.mu. +// bumpInvocation accumulates into the ctx's per-invocation bucket, evicting +// the oldest-touched bucket when the map exceeds maxInvocationBuckets (leak +// backstop — normal operation pops buckets at invocation end). Caller holds +// r.mu. func (r *Runtime) bumpInvocation(ctx context.Context, f func(*SavingsTotals)) { cid := runtime.CorrelationIDFromContext(ctx) if cid == "" { return } if r.perInvocation == nil { - r.perInvocation = make(map[string]*SavingsTotals) + r.perInvocation = make(map[string]*invocationBucket) } - t, ok := r.perInvocation[cid] + b, ok := r.perInvocation[cid] if !ok { - t = &SavingsTotals{} - r.perInvocation[cid] = t + if len(r.perInvocation) >= maxInvocationBuckets { + var oldestKey string + var oldest time.Time + for k, v := range r.perInvocation { + if oldestKey == "" || v.touched.Before(oldest) { + oldestKey, oldest = k, v.touched + } + } + delete(r.perInvocation, oldestKey) + } + b = &invocationBucket{} + r.perInvocation[cid] = b } - f(t) + b.touched = time.Now() + f(&b.totals) } // recordCompression accumulates savings and emits AuditEventCompressed with @@ -222,7 +260,8 @@ func (r *Runtime) recordExpansion(ctx context.Context, hash string, hit bool, by }) } -// rememberMarkers records emitted marker hashes for prefix resolution. +// rememberMarkers records emitted marker hashes for prefix resolution, +// evicting the oldest-emitted entries beyond maxRecentMarkers. func (r *Runtime) rememberMarkers(hashes []string) { if len(hashes) == 0 { return @@ -233,7 +272,15 @@ func (r *Runtime) rememberMarkers(hashes []string) { r.recent = make(map[string]struct{}) } for _, h := range hashes { + if _, ok := r.recent[h]; ok { + continue + } r.recent[h] = struct{}{} + r.recentOrder = append(r.recentOrder, h) + if len(r.recentOrder) > maxRecentMarkers { + delete(r.recent, r.recentOrder[0]) + r.recentOrder = r.recentOrder[1:] + } } } diff --git a/forge-core/compress/compress_test.go b/forge-core/compress/compress_test.go index 2e3af312..4ba1e928 100644 --- a/forge-core/compress/compress_test.go +++ b/forge-core/compress/compress_test.go @@ -354,6 +354,54 @@ func TestTakeInvocationTotals_PerCorrelation(t *testing.T) { } } +// PR #241 review: perInvocation must be bounded — an emission path that +// misses TakeInvocationTotals must not leak buckets forever. Oldest-touched +// buckets are evicted at the cap; the newest survives. +func TestPerInvocationBuckets_Bounded(t *testing.T) { + rt := newRuntime(t) + for i := 0; i < maxInvocationBuckets+50; i++ { + ctx := runtime.WithCorrelationID(context.Background(), fmt.Sprintf("task-%05d", i)) + rt.recordCompression(ctx, "tool_output", "t", 100, 10) + } + rt.mu.Lock() + n := len(rt.perInvocation) + rt.mu.Unlock() + if n > maxInvocationBuckets { + t.Fatalf("perInvocation grew past cap: %d > %d", n, maxInvocationBuckets) + } + // The most recent bucket must still be intact. + last := runtime.WithCorrelationID(context.Background(), fmt.Sprintf("task-%05d", maxInvocationBuckets+49)) + if got := rt.TakeInvocationTotals(last); got.Compressions != 1 { + t.Fatalf("newest bucket evicted: %+v", got) + } +} + +// PR #241 review: the recent-marker set must be bounded, evicting +// oldest-emitted hashes first (only recent markers matter for repairing +// imperfect transcriptions; exact hashes still resolve via the store). +func TestRecentMarkers_Bounded(t *testing.T) { + rt := newRuntime(t) + // Fixed-width distinct hashes: "mk<6 digits>xx" — no hash is a prefix of + // another, so resolvePrefix answers are about presence, not ambiguity. + hash := func(i int) string { return fmt.Sprintf("mk%06dxx", i) } + for i := 0; i < maxRecentMarkers+10; i++ { + rt.rememberMarkers([]string{hash(i)}) + } + rt.mu.Lock() + n, order := len(rt.recent), len(rt.recentOrder) + rt.mu.Unlock() + if n > maxRecentMarkers || order > maxRecentMarkers { + t.Fatalf("recent markers grew past cap: map=%d order=%d", n, order) + } + // Oldest evicted, newest still resolvable. + if got := rt.resolvePrefix(hash(0)); got != "" { + t.Errorf("oldest marker should have been evicted, resolved %q", got) + } + if got := rt.resolvePrefix(hash(maxRecentMarkers + 9)); got != hash(maxRecentMarkers+9) { + t.Errorf("newest marker should resolve, got %q", got) + } +} + // KeepPatterns (forge.yaml compression.keep_patterns) must flow through to // the hook so builder-flagged rows survive the compressed view. func TestHook_KeepPatterns(t *testing.T) { diff --git a/forge-core/compress/hook.go b/forge-core/compress/hook.go index 0e5de30c..8ac2848e 100644 --- a/forge-core/compress/hook.go +++ b/forge-core/compress/hook.go @@ -49,8 +49,11 @@ func (r *Runtime) AfterToolExecHook() runtime.Hook { Content: hctx.ToolOutput, Name: hctx.ToolName, }} + // <= 0 (not == 0): ctxzip clamps SavedTokens at zero today, but the + // guard must not silently apply inflated output if that contract + // ever changes (PR #241 review). res, err := ctxzip.Compress(msgs, opts) - if err != nil || res == nil || res.SavedTokens() == 0 { + if err != nil || res == nil || res.SavedTokens() <= 0 { return nil } From eec813ca475d60d8bc83628152107e2dc96498af Mon Sep 17 00:00:00 2001 From: MK Date: Sun, 5 Jul 2026 16:58:31 -0400 Subject: [PATCH 11/16] =?UTF-8?q?fix(tui):=20compression=20step=20never=20?= =?UTF-8?q?emitted=20StepCompleteMsg=20=E2=80=94=20wizard=20hung=20on=20se?= =?UTF-8?q?lection?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The wizard advances ONLY on tui.StepCompleteMsg (wizard.go documents "never check Complete() here"), but the compression step just set complete=true and returned — selecting an option left the wizard stuck on the step with Esc as the only exit (found live). Two fixes, following the house step patterns: - Update emits StepCompleteMsg on selection (ChannelStep pattern). - Init resets completion + selector state (SkillsStep pattern) so navigating BACK to the step re-prompts instead of stranding the user on a done selector that swallows all input. Regression tests pin both: enter emits StepCompleteMsg and Apply carries the choice; down+enter selects Disabled; back-navigation resets and accepts a fresh selection. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../internal/tui/steps/compression_step.go | 8 ++ .../tui/steps/compression_step_test.go | 89 +++++++++++++++++++ 2 files changed, 97 insertions(+) create mode 100644 forge-cli/internal/tui/steps/compression_step_test.go diff --git a/forge-cli/internal/tui/steps/compression_step.go b/forge-cli/internal/tui/steps/compression_step.go index 00b21d53..07fbf510 100644 --- a/forge-cli/internal/tui/steps/compression_step.go +++ b/forge-cli/internal/tui/steps/compression_step.go @@ -57,6 +57,11 @@ func (s *CompressionStep) Title() string { return "Context Compression" } func (s *CompressionStep) Icon() string { return "🗜️" } func (s *CompressionStep) Init() tea.Cmd { + // Init also runs when the user navigates BACK to this step; reset the + // selection state (SkillsStep pattern) or the completed selector would + // swallow all input and strand the wizard here. + s.complete = false + s.selector.Reset() return s.selector.Init() } @@ -72,6 +77,9 @@ func (s *CompressionStep) Update(msg tea.Msg) (tui.Step, tea.Cmd) { _, val := s.selector.Selected() s.enabled = val == "enabled" s.complete = true + // The wizard advances ONLY on StepCompleteMsg — setting complete + // without emitting it leaves the wizard stuck on this step. + return s, func() tea.Msg { return tui.StepCompleteMsg{} } } return s, cmd } diff --git a/forge-cli/internal/tui/steps/compression_step_test.go b/forge-cli/internal/tui/steps/compression_step_test.go new file mode 100644 index 00000000..11de8b41 --- /dev/null +++ b/forge-cli/internal/tui/steps/compression_step_test.go @@ -0,0 +1,89 @@ +package steps + +import ( + "testing" + + tea "github.com/charmbracelet/bubbletea" + + "github.com/initializ/forge/forge-cli/internal/tui" +) + +func newTestCompressionStep(t *testing.T) *CompressionStep { + t.Helper() + styles := tui.NewStyleSet(tui.DarkTheme) + s := NewCompressionStep(styles) + _ = s.Init() + return s +} + +// pressCompression injects a key and returns the updated step plus the +// message produced by the returned command (nil when there is no command). +func pressCompression(t *testing.T, s *CompressionStep, key tea.KeyType) (*CompressionStep, tea.Msg) { + t.Helper() + updated, cmd := s.Update(tea.KeyMsg{Type: key}) + step := updated.(*CompressionStep) + if cmd == nil { + return step, nil + } + return step, cmd() +} + +// Regression: selecting an option MUST emit tui.StepCompleteMsg — the wizard +// advances only on that message, so a step that merely sets complete=true +// strands the user on the step (found live: TUI hung after selecting the +// compression option; Esc was the only way out). +func TestCompressionStep_EnterEmitsStepComplete(t *testing.T) { + s := newTestCompressionStep(t) + + s, msg := pressCompression(t, s, tea.KeyEnter) // select "Enabled" (first item) + if _, ok := msg.(tui.StepCompleteMsg); !ok { + t.Fatalf("enter must emit tui.StepCompleteMsg, got %T", msg) + } + if !s.Complete() { + t.Fatal("step should be complete after selection") + } + + ctx := tui.NewWizardContext() + s.Apply(ctx) + if !ctx.Compression { + t.Fatal("selecting Enabled should set ctx.Compression") + } +} + +func TestCompressionStep_DisabledSelection(t *testing.T) { + s := newTestCompressionStep(t) + + s, _ = pressCompression(t, s, tea.KeyDown) // move to "Disabled" + s, msg := pressCompression(t, s, tea.KeyEnter) + if _, ok := msg.(tui.StepCompleteMsg); !ok { + t.Fatalf("enter must emit tui.StepCompleteMsg, got %T", msg) + } + + ctx := tui.NewWizardContext() + s.Apply(ctx) + if ctx.Compression { + t.Fatal("selecting Disabled should leave ctx.Compression false") + } +} + +// Regression: navigating BACK to the step re-runs Init, which must reset the +// selection state — otherwise the completed selector swallows all input and +// the wizard is stuck here on revisit. +func TestCompressionStep_BackNavigationResets(t *testing.T) { + s := newTestCompressionStep(t) + s, _ = pressCompression(t, s, tea.KeyEnter) + if !s.Complete() { + t.Fatal("precondition: step complete") + } + + _ = s.Init() // wizard re-Inits the step on StepBackMsg + + if s.Complete() { + t.Fatal("Init must reset completion for back-navigation") + } + // The step must accept a fresh selection after the reset. + s, msg := pressCompression(t, s, tea.KeyEnter) + if _, ok := msg.(tui.StepCompleteMsg); !ok { + t.Fatalf("step did not accept a new selection after reset, got %T", msg) + } +} From cfe7bf9864d846f043610e48bdc81d6a1899ff05 Mon Sep 17 00:00:00 2001 From: MK Date: Sun, 5 Jul 2026 17:52:18 -0400 Subject: [PATCH 12/16] =?UTF-8?q?chore:=20bump=20ctxzip=20to=20v0.1.1=20?= =?UTF-8?q?=E2=80=94=20JSON-envelope=20compression?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Live k8s triage against a real cluster showed compression_count: 0 on 28-37KB kubectl outputs: cli_execute wraps results in a single-line JSON envelope whose escaped newlines defeat content detection. ctxzip v0.1.1 compresses large string fields inside JSON-object envelopes, collapses identifier tokens in line dedup (kubectl tables now crush), and adds "warn" to the error floor (TYPE=Warning events never dropped). Measured on the live shape: 7,261 -> 95 tokens (98.7%), anomaly row kept verbatim. No forge-side code change — the adapter passes envelopes through ctxzip.Compress already. Co-Authored-By: Claude Opus 4.8 (1M context) --- forge-cli/go.mod | 2 +- forge-cli/go.sum | 4 ++-- forge-core/go.mod | 2 +- forge-core/go.sum | 2 ++ 4 files changed, 6 insertions(+), 4 deletions(-) diff --git a/forge-cli/go.mod b/forge-cli/go.mod index da476b10..34981700 100644 --- a/forge-cli/go.mod +++ b/forge-cli/go.mod @@ -54,7 +54,7 @@ require ( github.com/gorilla/websocket v1.5.4-0.20250319132907-e064f32e3674 // indirect github.com/grpc-ecosystem/grpc-gateway/v2 v2.29.0 // indirect github.com/inconshreveable/mousetrap v1.1.0 // indirect - github.com/initializ/ctxzip v0.1.0 // indirect + github.com/initializ/ctxzip v0.1.1 // indirect github.com/josharian/intern v1.0.0 // indirect github.com/json-iterator/go v1.1.12 // indirect github.com/klauspost/compress v1.16.7 // indirect diff --git a/forge-cli/go.sum b/forge-cli/go.sum index 37ea49b5..f9966b54 100644 --- a/forge-cli/go.sum +++ b/forge-cli/go.sum @@ -77,8 +77,8 @@ github.com/grpc-ecosystem/grpc-gateway/v2 v2.29.0 h1:5VipnvEpbqr2gA2VbM+nYVbkIF2 github.com/grpc-ecosystem/grpc-gateway/v2 v2.29.0/go.mod h1:Hyl3n6Twe1hvtd9XUXDec4pTvgMSEixRuQKPTMH2bNs= github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8= github.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw= -github.com/initializ/ctxzip v0.1.0 h1:GDCpfYSYUB630e1aaGpH7YatTXilu1n3pZ1m5WHtLBw= -github.com/initializ/ctxzip v0.1.0/go.mod h1:Qs4CpXTEeEOlZwIMPm/qaFPmil+AUIqH7Jg/K81FTmU= +github.com/initializ/ctxzip v0.1.1 h1:bnmfAI5F5P3Tq65/QWJKUeLq3raNRg1sY8GlbZ+728I= +github.com/initializ/ctxzip v0.1.1/go.mod h1:Qs4CpXTEeEOlZwIMPm/qaFPmil+AUIqH7Jg/K81FTmU= github.com/initializ/guardrails v0.12.0 h1:YzScnl+YPihLA1gepjxKjnKH+2EbYmGC4dUtpVRAJ2c= github.com/initializ/guardrails v0.12.0/go.mod h1:bDdHx73MF0+O09KqoXmUTTiFG4H7yEVQ0NR6juP1F3Q= github.com/josharian/intern v1.0.0 h1:vlS4z54oSdjm0bgjRigI+G1HpF+tI+9rE5LLzOg8HmY= diff --git a/forge-core/go.mod b/forge-core/go.mod index d82061cb..150bae60 100644 --- a/forge-core/go.mod +++ b/forge-core/go.mod @@ -4,7 +4,7 @@ go 1.25.0 require ( github.com/golang-jwt/jwt/v5 v5.3.1 - github.com/initializ/ctxzip v0.1.0 + github.com/initializ/ctxzip v0.1.1 github.com/initializ/forge/forge-skills v0.0.0 github.com/xeipuuv/gojsonschema v1.2.0 go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.59.0 diff --git a/forge-core/go.sum b/forge-core/go.sum index e6b86c2b..231c54fb 100644 --- a/forge-core/go.sum +++ b/forge-core/go.sum @@ -32,6 +32,8 @@ github.com/initializ/ctxzip v0.0.0-20260703221254-304962fbf29f h1:RLaBHp/SGK3uW/ github.com/initializ/ctxzip v0.0.0-20260703221254-304962fbf29f/go.mod h1:Qs4CpXTEeEOlZwIMPm/qaFPmil+AUIqH7Jg/K81FTmU= github.com/initializ/ctxzip v0.1.0 h1:GDCpfYSYUB630e1aaGpH7YatTXilu1n3pZ1m5WHtLBw= github.com/initializ/ctxzip v0.1.0/go.mod h1:Qs4CpXTEeEOlZwIMPm/qaFPmil+AUIqH7Jg/K81FTmU= +github.com/initializ/ctxzip v0.1.1 h1:bnmfAI5F5P3Tq65/QWJKUeLq3raNRg1sY8GlbZ+728I= +github.com/initializ/ctxzip v0.1.1/go.mod h1:Qs4CpXTEeEOlZwIMPm/qaFPmil+AUIqH7Jg/K81FTmU= github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE= github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk= github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= From 8d2f6296d760511b6aeebca00f0f0c34499dd5bb Mon Sep 17 00:00:00 2001 From: MK Date: Sun, 5 Jul 2026 18:12:52 -0400 Subject: [PATCH 13/16] feat(compress): report REALIZED wire savings on invocation_complete MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Live A/B exposed a metric-semantics gap: an invocation whose provider bill dropped ~31K tokens (73,017 -> 42,105) reported compression_saved_tokens_total: 1,257. The field counted compression EVENTS that occurred during the invocation, but savings are realized per RESEND — the previous invocation's crushed table rode in recovered history on every one of this invocation's four calls, saving ~5.6K each time, attributed to nothing. Accounting fix: rememberMarkers now stores each marker's saved-token delta (the recent map doubles as the prefix-resolution set); the client wrapper credits, on EVERY outbound Chat/ChatStream, the deltas of all markers riding in the request — this call's transforms and history markers alike. TTL-immune (no store reads), bounded by the existing marker cap. invocation_complete now reports: - compression_saved_tokens_total — realized wire savings (matches the provider bill; the number operators expect) - compression_event_saved_tokens — the old per-event sum (still matches the invocation's context_compressed events) Pinned by TestWireSavings_CompoundPerResend (three resends credit ~3x the one-time event saving). Docs updated (audit-logging, context-compression, forge.md knowledge skill). Co-Authored-By: Claude Opus 4.8 (1M context) --- .claude/skills/forge.md | 2 +- docs/core-concepts/context-compression.md | 2 +- docs/security/audit-logging.md | 2 +- forge-cli/runtime/runner.go | 18 ++++-- forge-cli/runtime/runner_compression_test.go | 32 +++++++++- forge-core/compress/client.go | 20 +++++- forge-core/compress/compress.go | 65 ++++++++++++++++---- forge-core/compress/compress_test.go | 42 ++++++++++++- forge-core/compress/hook.go | 2 +- 9 files changed, 161 insertions(+), 24 deletions(-) diff --git a/.claude/skills/forge.md b/.claude/skills/forge.md index 98d88187..639f1094 100644 --- a/.claude/skills/forge.md +++ b/.claude/skills/forge.md @@ -1061,7 +1061,7 @@ when OTel tracing is enabled (OTel v1 / Phase 4 / #105). Both use | `EventAgentCardPublished` | `agent_card_published` | Agent Card finalized at startup / hot-reload; `name`, `version`, `protocol_version`, `url`, `skill_count`, `capabilities`, `security_schemes`, `card_size_bytes`, `card_sha256` (FWS-1) | | `context_compressed` | `context_compressed` | Context compression shrank content; `seam` (`tool_output` / `request`), `tool`, `tokens_before` / `tokens_after` / `saved_tokens` + running totals (tokenizer estimates) | | `context_expanded` | `context_expanded` | Model retrieved offloaded content via `context_expand`; `hash`, `hit`, `bytes` + running totals | -| `AuditInvocationComplete` | `invocation_complete` | A2A invocation closed; `duration_ms`, `input_tokens_total`, `output_tokens_total`, `llm_call_count`, `model`, `provider` (FWS-3); with compression enabled also `compression_saved_tokens_total`, `compression_count`, `expansion_count` | +| `AuditInvocationComplete` | `invocation_complete` | A2A invocation closed; `duration_ms`, `input_tokens_total`, `output_tokens_total`, `llm_call_count`, `model`, `provider` (FWS-3); with compression enabled also `compression_saved_tokens_total` (realized wire savings, compounds per history resend), `compression_event_saved_tokens`, `compression_count`, `expansion_count` | | `AuditInvocationCancelled` | `invocation_cancelled` | A2A invocation cancelled via `tasks/cancel`; classified `reason` + partial token totals (FWS-4) | | `AuditTaskAdmissionDenied` | `task_admission_denied` | Inbound `tasks/send` denied by the platform admission middleware (#201; opt-in via `FORGE_ADMISSION_URL` + `FORGE_PLATFORM_TOKEN`); `reason`, `scope`, `window`, `reset_at`, `cached`. Caller sees HTTP 402 Payment Required. | | `AuditPolicyLoaded` | `policy_loaded` | One per non-empty policy layer at startup; `layer`, `source`, per-list size counters (FWS-5/6) | diff --git a/docs/core-concepts/context-compression.md b/docs/core-concepts/context-compression.md index 14160a8f..dfddac8d 100644 --- a/docs/core-concepts/context-compression.md +++ b/docs/core-concepts/context-compression.md @@ -99,7 +99,7 @@ Savings are first-class audit events, not log noise — see [Audit Logging](../s - `context_compressed` — per compression: seam, tool, tokens before/after/saved, plus running totals. - `context_expanded` — per retrieval: hash, hit, bytes — the cost side to net against savings. -- `invocation_complete` gains `compression_saved_tokens_total`, `compression_count`, and `expansion_count`, accumulated per invocation (concurrent tasks never cross-contaminate). +- `invocation_complete` gains `compression_saved_tokens_total` (**realized** savings — tokens this invocation's LLM calls did not send; a tool output compressed once but resent in history across four calls saves its delta four times), `compression_event_saved_tokens` (one-time per-compression deltas), `compression_count`, and `expansion_count`, accumulated per invocation (concurrent tasks never cross-contaminate). Token figures are tokenizer estimates (directionally accurate); billed truth remains `llm_call.input_tokens`. A surgical session that produced only small outputs correctly reports `compression_count: 0` — compression is insurance against bulk, not a tax on every call. diff --git a/docs/security/audit-logging.md b/docs/security/audit-logging.md index bff9ff3f..f4940e46 100644 --- a/docs/security/audit-logging.md +++ b/docs/security/audit-logging.md @@ -19,7 +19,7 @@ All runtime security events are emitted as structured NDJSON to stderr with corr | `egress_blocked` | Outbound request blocked (with domain, mode) | | `llm_call` | LLM API call completed (with `input_tokens`, `output_tokens`, `model`, `provider`, `duration_ms`, `request_id`). See [Token usage and duration](#token-usage-and-execution-duration). | | `llm_call_cancelled` | Streaming LLM call cancelled mid-flight; carries partial token counts captured up to cancellation. | -| `invocation_complete` | A2A invocation finished (auth → dispatch → engine → response). Carries `duration_ms` (wall-clock) plus aggregated `input_tokens_total` / `output_tokens_total` / `llm_call_count` / `model` / `provider`. When [context compression](../core-concepts/context-compression.md) is enabled it also carries `compression_saved_tokens_total` / `compression_count` (and `expansion_count` when nonzero), accumulated per invocation by correlation ID so concurrent tasks never cross-contaminate. | +| `invocation_complete` | A2A invocation finished (auth → dispatch → engine → response). Carries `duration_ms` (wall-clock) plus aggregated `input_tokens_total` / `output_tokens_total` / `llm_call_count` / `model` / `provider`. When [context compression](../core-concepts/context-compression.md) is enabled it also carries `compression_saved_tokens_total` — REALIZED savings: tokens this invocation's LLM calls did not send because compression markers rode in place of originals, compounding on every resend of compressed history (this is the number that matches the provider bill) — plus `compression_event_saved_tokens` (the one-time per-compression deltas, matching the sum of this invocation's `context_compressed` events), `compression_count`, and `expansion_count` when nonzero. Accumulated per invocation by correlation ID so concurrent tasks never cross-contaminate. | | `invocation_cancelled` | A2A invocation cancelled mid-flight via `tasks/cancel` (or internal cancellation like parent ctx deadline). Carries `fields.reason` (one of `workflow_failure` / `cost_limit_exceeded` / `timeout` / `external_signal`), `duration_ms` up to cancellation, and any partial token totals consumed before the signal. See [Cancellation](#cancellation). | | `task_admission_denied` | A new inbound `tasks/send` was rejected by the platform admission middleware (issue #201; opt-in via `FORGE_ADMISSION_URL` + `FORGE_PLATFORM_TOKEN`). Carries `fields.reason` (platform-defined: `cost_limit_exceeded`, `billing_overdue`, …), `fields.scope` (`agent` / `workspace` / `org`), `fields.window` (`hourly` / `daily` / `monthly` / `billing_cycle`), `fields.reset_at` (RFC 3339), and `fields.cached` (`true` when served from the 5s per-agent cache). Caller observes HTTP 402 Payment Required with `Retry-After`. See [Platform Admission Hook](admission.md). | | `guardrail_check` | Guardrail mask / block / warn decision. Carries `fields.gate` (`input` / `context` / `tool_call` / `output` / `stream` — sourced from the library `Result.Gate`), `fields.decision` (`masked` / `warned` / `blocked`), `fields.guardrail` + `fields.category` from the triggering violation, and `fields.violation_count`. `fields.tool` is present on `tool_call` and on `output` events for tool return text. With `FORGE_GUARDRAIL_CAPTURE_EVIDENCE=true` operators also opt into `fields.evidence` carrying the redacted + truncated triggering text. See [Guardrails — Audit Events](guardrails.md#audit-events). | diff --git a/forge-cli/runtime/runner.go b/forge-cli/runtime/runner.go index 052e6d83..0c983bbe 100644 --- a/forge-cli/runtime/runner.go +++ b/forge-cli/runtime/runner.go @@ -3399,9 +3399,18 @@ func (r *Runner) initLongTermMemory(ctx context.Context, mc *coreruntime.ModelCo // appendCompressionFields pops this invocation's compression savings (keyed // by the ctx's correlation ID) and adds them to an invocation_complete / -// invocation_cancelled fields map: compression_saved_tokens_total, -// compression_count, and expansion_count (when nonzero). Values are tokenizer -// estimates; zeros mean "compression on, nothing worth compressing". +// invocation_cancelled fields map. Values are tokenizer estimates; zeros mean +// "compression on, nothing worth compressing". +// +// - compression_saved_tokens_total — REALIZED savings: tokens this +// invocation's LLM calls did not send because markers rode in place of +// originals. Compounds per resend — a tool output compressed once but +// carried in history across four calls saves its delta four times. This +// matches the provider bill (live finding: per-event accounting reported +// 1,257 while the invocation actually avoided ~31K billed tokens). +// - compression_event_saved_tokens — the one-time, per-compression deltas +// (matches the sum of this invocation's context_compressed events). +// - compression_count / expansion_count. // // MUST be called exactly once per invocation, at the emission site — there // are THREE (executeTask plus both tasks/sendSubscribe streaming handlers, @@ -3413,7 +3422,8 @@ func (r *Runner) appendCompressionFields(ctx context.Context, fields map[string] return } ct := r.compression.TakeInvocationTotals(ctx) - fields["compression_saved_tokens_total"] = ct.SavedTokens + fields["compression_saved_tokens_total"] = ct.WireSavedTokens + fields["compression_event_saved_tokens"] = ct.SavedTokens fields["compression_count"] = ct.Compressions if ct.Expansions > 0 { fields["expansion_count"] = ct.Expansions diff --git a/forge-cli/runtime/runner_compression_test.go b/forge-cli/runtime/runner_compression_test.go index 58724e0d..31100b63 100644 --- a/forge-cli/runtime/runner_compression_test.go +++ b/forge-cli/runtime/runner_compression_test.go @@ -8,9 +8,25 @@ import ( "testing" "github.com/initializ/forge/forge-core/compress" + "github.com/initializ/forge/forge-core/llm" coreruntime "github.com/initializ/forge/forge-core/runtime" ) +// stubLLMClient satisfies llm.Client for driving the compressing wrapper. +type stubLLMClient struct{} + +func (stubLLMClient) Chat(_ context.Context, _ *llm.ChatRequest) (*llm.ChatResponse, error) { + return &llm.ChatResponse{Message: llm.ChatMessage{Role: llm.RoleAssistant, Content: "ok"}}, nil +} + +func (stubLLMClient) ChatStream(_ context.Context, _ *llm.ChatRequest) (<-chan llm.StreamDelta, error) { + ch := make(chan llm.StreamDelta) + close(ch) + return ch, nil +} + +func (stubLLMClient) ModelID() string { return "stub" } + // PR #241 review: invocation_complete is emitted from THREE sites // (executeTask + both sendSubscribe streaming handlers), and every one must // pop the per-correlation compression bucket via appendCompressionFields — @@ -41,6 +57,17 @@ func TestAppendCompressionFields_PopsAndPopulates(t *testing.T) { t.Fatal("fixture did not compress — test cannot exercise the bucket") } + // One LLM call carries the compressed output in history — realized (wire) + // savings accrue per outbound request, not per compression event. + wrapped := comp.WrapClient(stubLLMClient{}) + _, _ = wrapped.Chat(ctx, &llm.ChatRequest{Model: "m", Messages: []llm.ChatMessage{ + {Role: llm.RoleSystem, Content: "sys"}, + {Role: llm.RoleUser, Content: "check"}, + {Role: llm.RoleTool, Name: "list", ToolCallID: "t1", Content: hctx.ToolOutput}, + {Role: llm.RoleAssistant, Content: "ok"}, + {Role: llm.RoleUser, Content: "go"}, + }}) + // First call: fields populated, bucket popped. fields := map[string]any{} r.appendCompressionFields(ctx, fields) @@ -48,7 +75,10 @@ func TestAppendCompressionFields_PopsAndPopulates(t *testing.T) { t.Fatalf("compression_count = %v, want 1", fields["compression_count"]) } if fields["compression_saved_tokens_total"].(int64) <= 0 { - t.Fatalf("compression_saved_tokens_total = %v, want > 0", fields["compression_saved_tokens_total"]) + t.Fatalf("compression_saved_tokens_total (wire) = %v, want > 0", fields["compression_saved_tokens_total"]) + } + if fields["compression_event_saved_tokens"].(int64) <= 0 { + t.Fatalf("compression_event_saved_tokens = %v, want > 0", fields["compression_event_saved_tokens"]) } // Second call for the same invocation: bucket already popped → zeros diff --git a/forge-core/compress/client.go b/forge-core/compress/client.go index d6d88498..87dc8520 100644 --- a/forge-core/compress/client.go +++ b/forge-core/compress/client.go @@ -43,7 +43,25 @@ func (c *compressingClient) ModelID() string { return c.inner.ModelID() } // compressRequest returns req with compressed live-zone messages, or req // unchanged when there is nothing to gain. The caller's request is never // mutated — the loop may reuse its message slice. +// +// It also credits REALIZED wire savings for this call: every marker riding in +// the outbound messages (from this call's transforms or from history +// compressed on earlier turns) is tokens this request did not send. That +// compounding — not the one-time compression events — is where most savings +// live, and it is what invocation_complete reports. func (c *compressingClient) compressRequest(ctx context.Context, req *llm.ChatRequest) *llm.ChatRequest { + out := c.compressRequestInner(ctx, req) + if out != nil && len(out.Messages) > 0 { + contents := make([]string, 0, len(out.Messages)) + for _, m := range out.Messages { + contents = append(contents, m.Content) + } + c.rt.recordWireSavings(ctx, contents) + } + return out +} + +func (c *compressingClient) compressRequestInner(ctx context.Context, req *llm.ChatRequest) *llm.ChatRequest { if req == nil || len(req.Messages) == 0 { return req } @@ -77,7 +95,7 @@ func (c *compressingClient) compressRequest(ctx context.Context, req *llm.ChatRe copy(out.Messages, req.Messages) for _, tr := range res.Transforms { out.Messages[tr.Index].Content = res.Messages[tr.Index].Content - c.rt.rememberMarkers(tr.Markers) + c.rt.rememberMarkers(tr.Markers, int64(tr.TokensBefore-tr.TokensAfter)) } c.rt.debugf("compressed request", map[string]any{ diff --git a/forge-core/compress/compress.go b/forge-core/compress/compress.go index fa80a344..d08db3f6 100644 --- a/forge-core/compress/compress.go +++ b/forge-core/compress/compress.go @@ -95,13 +95,15 @@ type Runtime struct { logger runtime.Logger audit AuditFunc - // recent remembers marker hashes this process emitted so the expand tool - // can resolve a unique prefix when the model transcribes a hash - // imperfectly (observed live: models truncate hex hashes). Bounded to - // maxRecentMarkers, evicting oldest-emitted first — only *recent* - // markers matter for transcription repair. + // recent remembers marker hashes this process emitted — for the expand + // tool's prefix resolution when the model transcribes a hash imperfectly + // (observed live), and mapping each hash to the tokens its compression + // saved so the client wrapper can credit REALIZED savings every time the + // marker rides in an outbound request (savings compound per resend, not + // per compression). Bounded to maxRecentMarkers, evicting oldest-emitted + // first. mu sync.Mutex - recent map[string]struct{} + recent map[string]int64 recentOrder []string totals SavingsTotals // perInvocation accumulates savings keyed by correlation ID so @@ -138,10 +140,19 @@ const ( type SavingsTotals struct { // Compressions is how many times content was compressed (either seam). Compressions int64 - // SavedTokens is the cumulative estimated token reduction. + // SavedTokens is the cumulative per-EVENT token reduction — counted once + // per compression, at compression time. SavedTokens int64 + // WireSavedTokens is the cumulative REALIZED reduction: every time a + // marker rides in an outbound LLM request, that marker's saved tokens are + // tokens this request did not send. A tool output compressed once but + // resent in history across ten calls saves its delta ten times — this is + // the number that matches the provider's bill (live finding: an + // invocation reporting 1,257 event-saved tokens had actually avoided + // ~31K billed tokens through history compounding). + WireSavedTokens int64 // Expansions / ExpansionMisses count context_expand retrievals — the - // cost side to net against SavedTokens. + // cost side to net against savings. Expansions int64 ExpansionMisses int64 } @@ -260,22 +271,24 @@ func (r *Runtime) recordExpansion(ctx context.Context, hash string, hit bool, by }) } -// rememberMarkers records emitted marker hashes for prefix resolution, -// evicting the oldest-emitted entries beyond maxRecentMarkers. -func (r *Runtime) rememberMarkers(hashes []string) { +// rememberMarkers records emitted marker hashes with the tokens their +// compression saved (savedTokens is split evenly across the transform's +// markers), evicting the oldest-emitted entries beyond maxRecentMarkers. +func (r *Runtime) rememberMarkers(hashes []string, savedTokens int64) { if len(hashes) == 0 { return } + perMarker := savedTokens / int64(len(hashes)) r.mu.Lock() defer r.mu.Unlock() if r.recent == nil { - r.recent = make(map[string]struct{}) + r.recent = make(map[string]int64) } for _, h := range hashes { if _, ok := r.recent[h]; ok { continue } - r.recent[h] = struct{}{} + r.recent[h] = perMarker r.recentOrder = append(r.recentOrder, h) if len(r.recentOrder) > maxRecentMarkers { delete(r.recent, r.recentOrder[0]) @@ -284,6 +297,32 @@ func (r *Runtime) rememberMarkers(hashes []string) { } } +// recordWireSavings credits the realized savings of one outbound request: +// every marker present in the request's message contents represents tokens +// this request did not send. Called by the client wrapper on every Chat / +// ChatStream, including calls where nothing new was compressed — history +// markers keep saving on every resend, which is where most realized savings +// live. +func (r *Runtime) recordWireSavings(ctx context.Context, contents []string) { + var saved int64 + r.mu.Lock() + for _, c := range contents { + if !ccr.HasMarker(c) { + continue + } + for _, h := range ccr.ExtractHashes(c) { + saved += r.recent[h] + } + } + if saved > 0 { + r.totals.WireSavedTokens += saved + r.bumpInvocation(ctx, func(t *SavingsTotals) { + t.WireSavedTokens += saved + }) + } + r.mu.Unlock() +} + // resolvePrefix returns the unique remembered hash starting with prefix, or // "" when the prefix is too short, unknown, or ambiguous. func (r *Runtime) resolvePrefix(prefix string) string { diff --git a/forge-core/compress/compress_test.go b/forge-core/compress/compress_test.go index 4ba1e928..1b9b46f1 100644 --- a/forge-core/compress/compress_test.go +++ b/forge-core/compress/compress_test.go @@ -234,6 +234,46 @@ func TestNormalizeHash(t *testing.T) { } } +// Realized (wire) savings compound per RESEND: a tool output compressed once +// by the hook keeps saving its delta on every subsequent LLM call that +// carries it in history. This is the number invocation_complete reports — +// per-event accounting alone under-reported a live invocation 1,257 vs ~31K. +func TestWireSavings_CompoundPerResend(t *testing.T) { + rt := newRuntime(t) + ctx := runtime.WithCorrelationID(context.Background(), "task-wire") + + // Hook compresses one big tool output (creates a marker, saves S once). + hctx := &runtime.HookContext{ToolName: "list_pods", ToolOutput: bigJSON(80)} + _ = rt.AfterToolExecHook()(ctx, hctx) + if !strings.Contains(hctx.ToolOutput, "< Date: Sun, 5 Jul 2026 18:56:59 -0400 Subject: [PATCH 14/16] =?UTF-8?q?feat(compress):=20compress-then-truncate?= =?UTF-8?q?=20=E2=80=94=20tool=20caps=20move=20after=20the=20hook?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Live run 004: kubectl get pods -A -o json produced 108KB; the loop's pre-hook truncation cut it mid-JSON-string, which both destroyed data and broke the envelope so compression bailed — a mangled 12K-token blob rode through four more LLM calls. With compression enabled the loop now runs: execute -> safety ceiling (16x cap, abs max 4MB) -> hooks (guardrails -> compression) -> normal cap -> memory - The compression hook sees the FULL output; envelopes stay intact and crush instead of being destroyed. Guardrails also now scan the full output rather than a truncated prefix — strictly more correct. - The pre-hook safety ceiling bounds pathological outputs so hooks never scan unbounded payloads. - The normal 25%-budget cap still applies to the post-hook result — usually a no-op after compression; the context window still wins when it isn't. - Gated by LLMExecutorConfig.DeferToolResultTruncation, set by the runner only when compression is on: compression-off deployments keep today's order byte-identical (pinned by the existing truncation test). Bumps ctxzip to v0.1.2: envelopes that still arrive cut mid-string (safety ceiling, other runtimes) are salvaged — intact prefix compressed, valid JSON re-emitted, plus a _ctxzip_note telling the model the tail was destroyed upstream, not offloaded. Tests: hook sees full 60K output and its compressed replacement reaches the LLM; post-hook cap applies when nothing shrank; 500K output bounded by the ceiling before hooks; legacy order unchanged with the flag off. Co-Authored-By: Claude Opus 4.8 (1M context) --- docs/core-concepts/context-compression.md | 2 + forge-cli/go.mod | 2 +- forge-cli/go.sum | 4 +- forge-cli/runtime/runner.go | 5 + forge-core/go.mod | 2 +- forge-core/go.sum | 2 + forge-core/runtime/loop.go | 81 +++++++++--- .../runtime/loop_defer_truncation_test.go | 116 ++++++++++++++++++ 8 files changed, 193 insertions(+), 21 deletions(-) create mode 100644 forge-core/runtime/loop_defer_truncation_test.go diff --git a/docs/core-concepts/context-compression.md b/docs/core-concepts/context-compression.md index dfddac8d..b5ca22f7 100644 --- a/docs/core-concepts/context-compression.md +++ b/docs/core-concepts/context-compression.md @@ -107,6 +107,8 @@ Token figures are tokenizer estimates (directionally accurate); billed truth rem Fail-open, always: if the store cannot be opened, a compressor errors, or "compression" would grow a message, the original content is used unchanged. Error tool results are never compressed. An expired retrieval is not a dead end — the model is told to re-run the tool that produced the output. +**Interaction with tool-output truncation.** Forge normally hard-caps tool results at 25% of the context budget *before* anything else sees them. With compression enabled, that cut moves to **after** the compression hook (behind a safety ceiling of 16× the cap, absolute max 4MB): pre-hook truncation both destroys data and breaks the JSON envelopes compression would shrink losslessly. The compressed result is then capped as usual — normally a no-op. If an envelope still arrives cut mid-string (the safety ceiling, or another runtime), ctxzip salvages the intact prefix and adds a `_ctxzip_note` field telling the model the tail was destroyed upstream — not offloaded — so it re-runs the tool rather than calling `context_expand` for bytes that no longer exist. + **Single-writer store.** The bbolt store at `store_path` holds an exclusive file lock — one store per process. A second process pointing at the same file (two replicas on a shared volume, or `forge run` alongside `forge serve` in the same directory) fails to acquire the lock after a 5-second timeout and that process runs uncompressed (fail-open, with a startup warning). Give each replica its own `store_path` — offloaded originals are only ever retrieved by the process that offloaded them, so the store has no reason to be shared. ## Related diff --git a/forge-cli/go.mod b/forge-cli/go.mod index 34981700..dbf8261b 100644 --- a/forge-cli/go.mod +++ b/forge-cli/go.mod @@ -54,7 +54,7 @@ require ( github.com/gorilla/websocket v1.5.4-0.20250319132907-e064f32e3674 // indirect github.com/grpc-ecosystem/grpc-gateway/v2 v2.29.0 // indirect github.com/inconshreveable/mousetrap v1.1.0 // indirect - github.com/initializ/ctxzip v0.1.1 // indirect + github.com/initializ/ctxzip v0.1.2 // indirect github.com/josharian/intern v1.0.0 // indirect github.com/json-iterator/go v1.1.12 // indirect github.com/klauspost/compress v1.16.7 // indirect diff --git a/forge-cli/go.sum b/forge-cli/go.sum index f9966b54..ef5178f2 100644 --- a/forge-cli/go.sum +++ b/forge-cli/go.sum @@ -77,8 +77,8 @@ github.com/grpc-ecosystem/grpc-gateway/v2 v2.29.0 h1:5VipnvEpbqr2gA2VbM+nYVbkIF2 github.com/grpc-ecosystem/grpc-gateway/v2 v2.29.0/go.mod h1:Hyl3n6Twe1hvtd9XUXDec4pTvgMSEixRuQKPTMH2bNs= github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8= github.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw= -github.com/initializ/ctxzip v0.1.1 h1:bnmfAI5F5P3Tq65/QWJKUeLq3raNRg1sY8GlbZ+728I= -github.com/initializ/ctxzip v0.1.1/go.mod h1:Qs4CpXTEeEOlZwIMPm/qaFPmil+AUIqH7Jg/K81FTmU= +github.com/initializ/ctxzip v0.1.2 h1:vJxMtYq/kONZ7lMsEVMMt7cA6tZJtaEsu/8Jsku29RA= +github.com/initializ/ctxzip v0.1.2/go.mod h1:Qs4CpXTEeEOlZwIMPm/qaFPmil+AUIqH7Jg/K81FTmU= github.com/initializ/guardrails v0.12.0 h1:YzScnl+YPihLA1gepjxKjnKH+2EbYmGC4dUtpVRAJ2c= github.com/initializ/guardrails v0.12.0/go.mod h1:bDdHx73MF0+O09KqoXmUTTiFG4H7yEVQ0NR6juP1F3Q= github.com/josharian/intern v1.0.0 h1:vlS4z54oSdjm0bgjRigI+G1HpF+tI+9rE5LLzOg8HmY= diff --git a/forge-cli/runtime/runner.go b/forge-cli/runtime/runner.go index 0c983bbe..4571c8bb 100644 --- a/forge-cli/runtime/runner.go +++ b/forge-cli/runtime/runner.go @@ -934,6 +934,11 @@ func (r *Runner) Run(ctx context.Context) error { MaxIterations: 100, CharBudget: charBudget, FilesDir: filepath.Join(r.cfg.WorkDir, ".forge", "files"), + // With compression on, tool results are capped AFTER + // the compression hook (behind a 16x/4MB safety + // ceiling) — pre-hook truncation destroys data and + // breaks the JSON envelopes compression would shrink. + DeferToolResultTruncation: r.compression != nil, // Issue #130 — the same resolved TracingConfig // already passed to NewTracerProvider drives Phase // 3.5 span-content capture inside the executor diff --git a/forge-core/go.mod b/forge-core/go.mod index 150bae60..efdd41ac 100644 --- a/forge-core/go.mod +++ b/forge-core/go.mod @@ -4,7 +4,7 @@ go 1.25.0 require ( github.com/golang-jwt/jwt/v5 v5.3.1 - github.com/initializ/ctxzip v0.1.1 + github.com/initializ/ctxzip v0.1.2 github.com/initializ/forge/forge-skills v0.0.0 github.com/xeipuuv/gojsonschema v1.2.0 go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.59.0 diff --git a/forge-core/go.sum b/forge-core/go.sum index 231c54fb..289f43b2 100644 --- a/forge-core/go.sum +++ b/forge-core/go.sum @@ -34,6 +34,8 @@ github.com/initializ/ctxzip v0.1.0 h1:GDCpfYSYUB630e1aaGpH7YatTXilu1n3pZ1m5WHtLB github.com/initializ/ctxzip v0.1.0/go.mod h1:Qs4CpXTEeEOlZwIMPm/qaFPmil+AUIqH7Jg/K81FTmU= github.com/initializ/ctxzip v0.1.1 h1:bnmfAI5F5P3Tq65/QWJKUeLq3raNRg1sY8GlbZ+728I= github.com/initializ/ctxzip v0.1.1/go.mod h1:Qs4CpXTEeEOlZwIMPm/qaFPmil+AUIqH7Jg/K81FTmU= +github.com/initializ/ctxzip v0.1.2 h1:vJxMtYq/kONZ7lMsEVMMt7cA6tZJtaEsu/8Jsku29RA= +github.com/initializ/ctxzip v0.1.2/go.mod h1:Qs4CpXTEeEOlZwIMPm/qaFPmil+AUIqH7Jg/K81FTmU= github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE= github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk= github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= diff --git a/forge-core/runtime/loop.go b/forge-core/runtime/loop.go index aa753695..49eb8c48 100644 --- a/forge-core/runtime/loop.go +++ b/forge-core/runtime/loop.go @@ -39,6 +39,15 @@ type ToolExecutor interface { ToolDefinitions() []llm.ToolDefinition } +// Pre-hook safety ceiling for deferred tool-result truncation: hooks must +// never scan unbounded payloads, but the ceiling must be generous enough +// that real bulky outputs (e.g. kubectl get -o json on a large cluster) +// reach the compression hook intact. +const ( + toolResultCeilingMultiplier = 16 + toolResultCeilingAbsolute = 4 << 20 // 4MB +) + // LLMExecutor implements AgentExecutor using an LLM client with tool calling. type LLMExecutor struct { client llm.Client @@ -56,6 +65,13 @@ type LLMExecutor struct { filesDir string // directory for file_create output sessionMaxAge time.Duration // max age for session recovery (0 = no limit) workflowPhases []string // workflow phases from skills (edit, finalize, query) + // deferToolTruncation moves the maxToolResultChars cut to AFTER the + // AfterToolExec hooks, so a compression hook sees the full output and + // the (usually smaller) compressed result is what gets capped. A + // pre-hook safety ceiling still bounds pathological outputs. Set by the + // runner only when context compression is enabled — off, the + // pre-compression truncation order is byte-identical to before. + deferToolTruncation bool // tracingCfg governs Phase 3.5 span-attribute content capture // (issue #130). Only CaptureContent + Redact are consumed here; // the rest of the struct is honored by the cli runner's tracer @@ -80,6 +96,12 @@ type LLMExecutorConfig struct { FilesDir string // directory for file_create output (default: $TMPDIR/forge-files) SessionMaxAge time.Duration // max idle time before session recovery is skipped (0 = 30m default) WorkflowPhases []string // workflow phases from skills (edit, finalize, query) + // DeferToolResultTruncation applies the tool-result size cap after the + // AfterToolExec hooks instead of before, behind a pre-hook safety + // ceiling (16x the cap, absolute max 4MB). Enable when a compression + // hook is registered: pre-hook truncation both destroys data and breaks + // the JSON envelopes the compressor could otherwise shrink losslessly. + DeferToolResultTruncation bool // TracingConfig is the same observability.TracingConfig the cli // runner resolves and passes to NewTracerProvider. The executor // reads CaptureContent + Redact to decide whether to stamp @@ -128,22 +150,23 @@ func NewLLMExecutor(cfg LLMExecutorConfig) *LLMExecutor { } return &LLMExecutor{ - client: cfg.Client, - tools: cfg.Tools, - hooks: hooks, - systemPrompt: cfg.SystemPrompt, - maxIter: maxIter, - compactor: cfg.Compactor, - store: cfg.Store, - logger: logger, - modelName: cfg.ModelName, - provider: cfg.Provider, - charBudget: budget, - maxToolResultChars: toolLimit, - filesDir: cfg.FilesDir, - sessionMaxAge: sessionMaxAge, - workflowPhases: cfg.WorkflowPhases, - tracingCfg: cfg.TracingConfig, + client: cfg.Client, + tools: cfg.Tools, + hooks: hooks, + systemPrompt: cfg.SystemPrompt, + maxIter: maxIter, + compactor: cfg.Compactor, + store: cfg.Store, + logger: logger, + modelName: cfg.ModelName, + provider: cfg.Provider, + charBudget: budget, + maxToolResultChars: toolLimit, + filesDir: cfg.FilesDir, + sessionMaxAge: sessionMaxAge, + workflowPhases: cfg.WorkflowPhases, + deferToolTruncation: cfg.DeferToolResultTruncation, + tracingCfg: cfg.TracingConfig, } } @@ -701,7 +724,24 @@ func (e *LLMExecutor) Execute(ctx context.Context, task *a2a.Task, msg *a2a.Mess // Truncate oversized tool results to avoid LLM API errors. // Limit is proportional to model context budget (25%, floor 2K, cap 400K). - if len(result) > e.maxToolResultChars { + // + // With compression enabled (deferToolTruncation), the cut moves + // to AFTER the hooks: pre-hook truncation destroys data AND + // breaks the JSON envelopes the compression hook could shrink + // losslessly (live finding: a 108KB kubectl -o json envelope was + // cut mid-string, defeating compression entirely). A generous + // pre-hook safety ceiling still bounds pathological outputs so + // hooks never scan unbounded payloads; the normal cap is applied + // to the post-hook (usually compressed) result below. + if e.deferToolTruncation { + ceiling := e.maxToolResultChars * toolResultCeilingMultiplier + if ceiling > toolResultCeilingAbsolute { + ceiling = toolResultCeilingAbsolute + } + if len(result) > ceiling { + result = result[:ceiling] + "\n\n[OUTPUT TRUNCATED -- original length: " + strconv.Itoa(len(result)) + " chars]" + } + } else if len(result) > e.maxToolResultChars { result = result[:e.maxToolResultChars] + "\n\n[OUTPUT TRUNCATED -- original length: " + strconv.Itoa(len(result)) + " chars]" } @@ -720,6 +760,13 @@ func (e *LLMExecutor) Execute(ctx context.Context, task *a2a.Task, msg *a2a.Mess } result = afterHctx.ToolOutput // allow hooks to redact output + // Deferred mode: apply the normal cap to the post-hook result — + // compression usually brought it far below the limit, making + // this a no-op; when it didn't, the context window still wins. + if e.deferToolTruncation && len(result) > e.maxToolResultChars { + result = result[:e.maxToolResultChars] + "\n\n[OUTPUT TRUNCATED -- original length: " + strconv.Itoa(len(result)) + " chars]" + } + // Handle file_create tool: always create a file part. // For other tools with large output, detect content type. // Skip cli_execute: it's an intermediate tool — the LLM should diff --git a/forge-core/runtime/loop_defer_truncation_test.go b/forge-core/runtime/loop_defer_truncation_test.go new file mode 100644 index 00000000..22ab0185 --- /dev/null +++ b/forge-core/runtime/loop_defer_truncation_test.go @@ -0,0 +1,116 @@ +package runtime + +import ( + "context" + "encoding/json" + "strings" + "testing" + + "github.com/initializ/forge/forge-core/a2a" + "github.com/initializ/forge/forge-core/llm" +) + +// deferHarness runs one tool round-trip with DeferToolResultTruncation and +// returns what the AfterToolExec hook saw and what reached the LLM. +func deferHarness(t *testing.T, toolResult string, hook Hook) (hookSawLen int, llmSaw string) { + t.Helper() + callCount := 0 + var captured []llm.ChatMessage + client := &mockLLMClient{chatFunc: func(_ context.Context, req *llm.ChatRequest) (*llm.ChatResponse, error) { + callCount++ + captured = req.Messages + if callCount == 1 { + return &llm.ChatResponse{ + Message: llm.ChatMessage{Role: llm.RoleAssistant, ToolCalls: []llm.ToolCall{{ + ID: "c1", Type: "function", Function: llm.FunctionCall{Name: "big", Arguments: `{}`}, + }}}, + FinishReason: "tool_calls", + }, nil + } + return &llm.ChatResponse{Message: llm.ChatMessage{Role: llm.RoleAssistant, Content: "done"}, FinishReason: "stop"}, nil + }} + tools := &mockToolExecutor{ + executeFunc: func(_ context.Context, _ string, _ json.RawMessage) (string, error) { + return toolResult, nil + }, + toolDefs: []llm.ToolDefinition{{Type: "function", Function: llm.FunctionSchema{Name: "big"}}}, + } + + hooks := NewHookRegistry() + hooks.Register(AfterToolExec, func(ctx context.Context, hctx *HookContext) error { + hookSawLen = len(hctx.ToolOutput) + if hook != nil { + return hook(ctx, hctx) + } + return nil + }) + + exec := NewLLMExecutor(LLMExecutorConfig{ + Client: client, Tools: tools, Hooks: hooks, + CharBudget: 100_000, // → cap 25K, ceiling 400K + DeferToolResultTruncation: true, + }) + task := &a2a.Task{ID: "t-defer"} + if _, err := exec.Execute(context.Background(), task, &a2a.Message{ + Role: a2a.MessageRoleUser, Parts: []a2a.Part{a2a.NewTextPart("go")}, + }); err != nil { + t.Fatalf("Execute: %v", err) + } + for _, m := range captured { + if m.Role == llm.RoleTool { + llmSaw = m.Content + } + } + return hookSawLen, llmSaw +} + +// The live-run-004 fix: with deferred truncation, the AfterToolExec hooks see +// the FULL tool output (no mid-JSON cut), and a hook that shrinks it (the +// compression hook) makes the post-hook cap a no-op. +func TestDeferredTruncation_HookSeesFullOutput(t *testing.T) { + full := strings.Repeat("x", 60_000) // over the 25K cap, under the ceiling + + compress := func(_ context.Context, hctx *HookContext) error { + hctx.ToolOutput = "compressed-tiny" // stand-in for the compression hook + return nil + } + hookSaw, llmSaw := deferHarness(t, full, compress) + + if hookSaw != len(full) { + t.Fatalf("hook saw %d chars, want the full %d (pre-hook cut destroys envelopes)", hookSaw, len(full)) + } + if llmSaw != "compressed-tiny" { + t.Fatalf("LLM should see the hook's compressed output, got %d chars", len(llmSaw)) + } +} + +// Without a shrinking hook, the normal cap still protects the context window +// — applied after the hooks instead of before. +func TestDeferredTruncation_PostHookCapStillApplies(t *testing.T) { + full := strings.Repeat("x", 60_000) + hookSaw, llmSaw := deferHarness(t, full, nil) + + if hookSaw != len(full) { + t.Fatalf("hook saw %d chars, want %d", hookSaw, len(full)) + } + if !strings.Contains(llmSaw, "[OUTPUT TRUNCATED") { + t.Fatal("post-hook cap missing") + } + if len(llmSaw) > 26_000 { + t.Fatalf("LLM saw %d chars, want ~25K cap", len(llmSaw)) + } +} + +// Pathological outputs are still bounded BEFORE hooks by the safety ceiling +// (16x cap, absolute 4MB) so hooks never scan unbounded payloads. +func TestDeferredTruncation_SafetyCeiling(t *testing.T) { + monster := strings.Repeat("x", 500_000) // ceiling = 16*25K = 400K + hookSaw, _ := deferHarness(t, monster, nil) + + if hookSaw > 401_000 { + t.Fatalf("hook saw %d chars — safety ceiling (400K) not enforced", hookSaw) + } + if hookSaw < 399_000 { + t.Fatalf("hook saw %d chars — ceiling cut too aggressively", hookSaw) + } +} From 93a7b95f3e95bfb870d00e83983e29c5d96a6b96 Mon Sep 17 00:00:00 2001 From: MK Date: Sun, 5 Jul 2026 19:29:20 -0400 Subject: [PATCH 15/16] fix(lint): discard unused step reassignment in back-navigation test MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit staticcheck SA4006: the final pressCompression return of s was never read — only the emitted message matters for the assertion. Co-Authored-By: Claude Opus 4.8 (1M context) --- forge-cli/internal/tui/steps/compression_step_test.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/forge-cli/internal/tui/steps/compression_step_test.go b/forge-cli/internal/tui/steps/compression_step_test.go index 11de8b41..0fe0306b 100644 --- a/forge-cli/internal/tui/steps/compression_step_test.go +++ b/forge-cli/internal/tui/steps/compression_step_test.go @@ -82,7 +82,7 @@ func TestCompressionStep_BackNavigationResets(t *testing.T) { t.Fatal("Init must reset completion for back-navigation") } // The step must accept a fresh selection after the reset. - s, msg := pressCompression(t, s, tea.KeyEnter) + _, msg := pressCompression(t, s, tea.KeyEnter) if _, ok := msg.(tui.StepCompleteMsg); !ok { t.Fatalf("step did not accept a new selection after reset, got %T", msg) } From 8e0fec6aa50528af0a38be73e8ae000922831a71 Mon Sep 17 00:00:00 2001 From: MK Date: Sun, 5 Jul 2026 21:40:50 -0400 Subject: [PATCH 16/16] fix(compress): never split a ctxzip marker at the post-hook truncation cut MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit PR #241 review follow-up: in deferred mode the post-hook cap sliced the compressed result at a byte offset, which could bisect a <> marker — leaving the model a corrupted pointer to content that is still in the store but unreachable through that marker. Rare (fires only when even the compressed result exceeds the cap), but wrong. truncateToolResult now backs the cut up to the marker's start when the byte cut would land inside one — output carries whole markers or none. Applied at both deferred-mode cut sites (safety ceiling and post-hook cap); the legacy non-deferred path is untouched (no compression, no markers, byte-identical). The marker prefix is pinned as runtime.CompressionMarkerPrefix (runtime cannot import compress — cycle) with a guard test asserting it matches ccr.MarkerPrefix so the two can never drift. Tests: straddling marker dropped whole, complete marker kept, plain cut without markers, unterminated marker dropped, plus a loop-level case where a compression hook's output places a marker straddling the cap — the LLM never sees a partial marker. Co-Authored-By: Claude Opus 4.8 (1M context) --- forge-core/compress/compress_test.go | 10 +++ forge-core/runtime/loop.go | 33 +++++++-- .../runtime/loop_defer_truncation_test.go | 70 +++++++++++++++++++ 3 files changed, 108 insertions(+), 5 deletions(-) diff --git a/forge-core/compress/compress_test.go b/forge-core/compress/compress_test.go index 1b9b46f1..be156088 100644 --- a/forge-core/compress/compress_test.go +++ b/forge-core/compress/compress_test.go @@ -346,6 +346,16 @@ func TestAuditEvents_SavingsAndTotals(t *testing.T) { } } +// The runtime's marker-aware tool-result truncation pins the ctxzip marker +// prefix as a literal (importing this package from runtime would cycle). +// Guard the two against drift. +func TestRuntimeMarkerPrefix_MatchesCCR(t *testing.T) { + if runtime.CompressionMarkerPrefix != ccr.MarkerPrefix { + t.Fatalf("runtime.CompressionMarkerPrefix=%q != ccr.MarkerPrefix=%q — marker-aware truncation would miss markers", + runtime.CompressionMarkerPrefix, ccr.MarkerPrefix) + } +} + // The runtime-owned system directive is what makes compression work for ANY // skill without skill authors documenting it. Guard that it names the real // tool and the real marker prefix, so a rename cannot silently orphan it. diff --git a/forge-core/runtime/loop.go b/forge-core/runtime/loop.go index 49eb8c48..950c19b5 100644 --- a/forge-core/runtime/loop.go +++ b/forge-core/runtime/loop.go @@ -48,6 +48,31 @@ const ( toolResultCeilingAbsolute = 4 << 20 // 4MB ) +// CompressionMarkerPrefix mirrors ctxzip's ccr.MarkerPrefix. The runtime +// cannot import forge-core/compress (it imports runtime — cycle), so the +// literal is pinned here and a guard test in the compress package asserts +// the two stay equal. +const CompressionMarkerPrefix = "<>", leaving the model +// a corrupted marker it cannot expand even though the offloaded content is +// still in the store. If the cut would bisect a marker, it backs up to the +// marker's start so the output carries whole markers or none. +func truncateToolResult(s string, limit int) string { + if len(s) <= limit { + return s + } + cut := limit + if start := strings.LastIndex(s[:cut], CompressionMarkerPrefix); start >= 0 { + end := strings.Index(s[start:], ">>") + if end < 0 || start+end+2 > cut { + cut = start // marker straddles the cut — drop it whole + } + } + return s[:cut] + "\n\n[OUTPUT TRUNCATED -- original length: " + strconv.Itoa(len(s)) + " chars]" +} + // LLMExecutor implements AgentExecutor using an LLM client with tool calling. type LLMExecutor struct { client llm.Client @@ -738,9 +763,7 @@ func (e *LLMExecutor) Execute(ctx context.Context, task *a2a.Task, msg *a2a.Mess if ceiling > toolResultCeilingAbsolute { ceiling = toolResultCeilingAbsolute } - if len(result) > ceiling { - result = result[:ceiling] + "\n\n[OUTPUT TRUNCATED -- original length: " + strconv.Itoa(len(result)) + " chars]" - } + result = truncateToolResult(result, ceiling) } else if len(result) > e.maxToolResultChars { result = result[:e.maxToolResultChars] + "\n\n[OUTPUT TRUNCATED -- original length: " + strconv.Itoa(len(result)) + " chars]" } @@ -763,8 +786,8 @@ func (e *LLMExecutor) Execute(ctx context.Context, task *a2a.Task, msg *a2a.Mess // Deferred mode: apply the normal cap to the post-hook result — // compression usually brought it far below the limit, making // this a no-op; when it didn't, the context window still wins. - if e.deferToolTruncation && len(result) > e.maxToolResultChars { - result = result[:e.maxToolResultChars] + "\n\n[OUTPUT TRUNCATED -- original length: " + strconv.Itoa(len(result)) + " chars]" + if e.deferToolTruncation { + result = truncateToolResult(result, e.maxToolResultChars) } // Handle file_create tool: always create a file part. diff --git a/forge-core/runtime/loop_defer_truncation_test.go b/forge-core/runtime/loop_defer_truncation_test.go index 22ab0185..d718261f 100644 --- a/forge-core/runtime/loop_defer_truncation_test.go +++ b/forge-core/runtime/loop_defer_truncation_test.go @@ -101,6 +101,76 @@ func TestDeferredTruncation_PostHookCapStillApplies(t *testing.T) { } } +// PR #241 review: a byte-offset cut can bisect a <> marker, +// leaving the model a corrupted pointer it cannot expand even though the +// content is still in the store. The cut must carry whole markers or none. +func TestTruncateToolResult_MarkerBoundary(t *testing.T) { + marker := "<>" + + t.Run("under limit unchanged", func(t *testing.T) { + if got := truncateToolResult("short", 100); got != "short" { + t.Fatalf("got %q", got) + } + }) + + t.Run("no marker plain cut", func(t *testing.T) { + got := truncateToolResult(strings.Repeat("x", 200), 100) + if !strings.HasPrefix(got, strings.Repeat("x", 100)) || !strings.Contains(got, "[OUTPUT TRUNCATED") { + t.Fatalf("got %q", got) + } + }) + + t.Run("marker straddling the cut is dropped whole", func(t *testing.T) { + // Cut lands mid-marker: prefix(80) + marker starts at 80, limit 100. + s := strings.Repeat("a", 80) + marker + strings.Repeat("b", 200) + got := truncateToolResult(s, 100) + if strings.Contains(got, CompressionMarkerPrefix) { + t.Fatalf("broken/partial marker survived the cut: %q", got) + } + if !strings.HasPrefix(got, strings.Repeat("a", 80)) { + t.Fatalf("prefix lost: %q", got) + } + }) + + t.Run("marker fully before the cut is kept", func(t *testing.T) { + s := marker + strings.Repeat("b", 300) + got := truncateToolResult(s, 200) + if !strings.Contains(got, marker) { + t.Fatalf("complete marker should survive: %q", got) + } + }) + + t.Run("unterminated marker before cut is dropped", func(t *testing.T) { + s := strings.Repeat("a", 50) + "<> + got := truncateToolResult(s+strings.Repeat("b", 200), 100) + if strings.Contains(got, CompressionMarkerPrefix) { + t.Fatalf("unterminated marker survived: %q", got) + } + }) +} + +// End-to-end: a compression hook whose output still exceeds the cap, with a +// marker positioned to straddle the byte cut — the LLM must never see a +// partial marker. +func TestDeferredTruncation_NeverSplitsMarker(t *testing.T) { + // Cap is 25K (CharBudget 100_000). Place the marker straddling 25_000. + marker := "<>" + compressed := strings.Repeat("k", 24_990) + marker + strings.Repeat("t", 5_000) + + hook := func(_ context.Context, hctx *HookContext) error { + hctx.ToolOutput = compressed + return nil + } + _, llmSaw := deferHarness(t, strings.Repeat("x", 60_000), hook) + + if strings.Contains(llmSaw, CompressionMarkerPrefix) && !strings.Contains(llmSaw, marker) { + t.Fatalf("LLM saw a partial marker:\n...%s", llmSaw[len(llmSaw)-120:]) + } + if !strings.Contains(llmSaw, "[OUTPUT TRUNCATED") { + t.Fatal("cap should still have fired") + } +} + // Pathological outputs are still bounded BEFORE hooks by the safety ceiling // (16x cap, absolute 4MB) so hooks never scan unbounded payloads. func TestDeferredTruncation_SafetyCeiling(t *testing.T) {