From 1e9436e3747bde46e14876cf733105d73c045fc8 Mon Sep 17 00:00:00 2001 From: Osher-Elhadad Date: Mon, 10 Aug 2026 02:52:18 +0000 Subject: [PATCH 1/2] fix(expand): stop buffering every SSE response on a marker check that always matched MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The streaming fast path never engaged. `hasMarkers` tested the WHOLE outgoing request body for a `cg:` prefix, but the expand tool we inject ourselves carries `<>` in its own description, HTML-escaped by encoding/json. So from the moment `expand.Inject` fired the check was unconditionally true, every SSE response was read to completion before a byte reached the client, and the comment promising "zero added latency" for marker-free requests never held. Scope the check to model-visible content (`messages`, `system`) via the new `expand.HasMarkersInMessages`. Requiring the full marker shape is not enough on its own — the tool description contains the full shape too — so scoping is the actual fix. Both marker spellings are matched deliberately: markers travel the wire HTML-escaped, because sjson escapes "<" whenever the value contains a newline and every marker is appended after one. That dependency was load-bearing and undocumented; it is now a named regexp with a comment and tests. Add SSE streaming health to /stats — `sse_streamed`, `sse_buffered`, `sse_buffered_pct`, `sse_ttfb_ms_avg`, `sse_ttfb_ms_avg_buffered` — so this class of regression is measured rather than inferred. Fields are added only; nothing is renamed or removed, so the harbor harnesses keep parsing. Measured on a fake 1s SSE upstream (20 events x 50ms), medians of 12 trials: marker-free before 1007ms TTFB -> after 43ms TTFB marker-bearing before 1008ms TTFB -> after 1008ms TTFB (correct: inspected) Tests: a marker-free streaming request against an upstream that withholds its tail now proves the client gets the head first (this deadlocks and fails on the old code); plus buffered-by-design coverage, expand batched with another tool, the multi-round cap, and the OpenAI raw-replay fallback. Streaming restoration remains Anthropic-only — `AggregateSSE` cannot reconstruct an OpenAI stream — now documented explicitly instead of silently falling through. Closes #26 Signed-off-by: Osher-Elhadad --- docs/design.md | 31 +++- docs/how-to/recover-context.md | 34 ++++ expand/expand.go | 34 ++++ expand/expand_test.go | 71 ++++++++ metrics/metrics.go | 45 ++++++ metrics/metrics_test.go | 22 +++ proxy/proxy.go | 43 ++++- proxy/proxy_test.go | 288 +++++++++++++++++++++++++++++++++ 8 files changed, 560 insertions(+), 8 deletions(-) diff --git a/docs/design.md b/docs/design.md index 7aad713..6f0bccb 100644 --- a/docs/design.md +++ b/docs/design.md @@ -156,6 +156,29 @@ An expired/evicted original resolves to an explicit placeholder rather than bein provider requires one `tool_result` per `tool_call_id`). A miss silently turns a lossless offload lossy — the known TTL edge. +### The loop on a streaming response + +The loop needs a whole assistant message; SSE delivers events. So the host decides per request, +from the request bytes, whether it can afford to look: + +- **no marker in `messages`/`system`** → nothing to expand → stream through untouched; +- **marker present** → buffer the stream, rebuild the message with `expand.AggregateSSE` + (Anthropic dialect only — other dialects return `ok=false` and are replayed raw), inspect, and + either continue the loop or replay the buffered bytes verbatim. + +Buffering is the one thing that turns a stream into a non-stream, so the marker test must be tight +in *both* directions: a false negative loses a real expand call, a false positive silently costs +every request its time-to-first-byte. It scans only model-visible content (`messages`, `system`) — +scanning the whole body also matched the expand tool description the host injects itself, which made +it unconditionally true (issue #26). `/stats` exposes `sse_streamed` / `sse_buffered` / +`sse_buffered_pct` and the two TTFB averages so the fast path is measured, not assumed. + +**Markers on the wire are HTML-escaped.** `encoding/json` always escapes `<`; `sjson` escapes it +whenever the value contains a newline, and markers are appended after a newline. So `<>` in the +model's view is `<>` in the bytes. Marker matching on *decoded* content +(`expand.HasPlaceholder`, used by the components) sees the plain form; matching on *raw request +bytes* (`expand.rawMarkerRe`, used by the host's streaming decision) must accept both, and does. + ## State: the Store One `Store` interface, in-memory TTL+LRU default (both hosts share it). Defaults: **1800s TTL, @@ -184,7 +207,13 @@ vocabulary), `Aggregator` (in-process rollups behind `/stats`), `Tee` (fan-out), of per-request percentages. It also reports: - `wasted_tokens` / `bounces` — content offloaded then re-served via expand (a premature offload); - `adjusted_saved` = saved − wasted (bounce-adjusted, may be negative); -- `top_passthrough` — components that ran but never changed a request: dead weight to drop. +- `top_passthrough` — components that ran but never changed a request: dead weight to drop; +- `sse_streamed` / `sse_buffered` / `sse_buffered_pct` and `sse_ttfb_ms_avg` / + `sse_ttfb_ms_avg_buffered` — streaming health: how many SSE responses had to be buffered whole to + be inspected for an expand call, and what that cost in time-to-first-byte. + +Fields are only ever **added** to `/stats`; the harbor harnesses parse it, so no field is renamed +or removed. ## Config & registry diff --git a/docs/how-to/recover-context.md b/docs/how-to/recover-context.md index a746f2c..9df77ef 100644 --- a/docs/how-to/recover-context.md +++ b/docs/how-to/recover-context.md @@ -48,6 +48,40 @@ original resolves to an explicit placeholder rather than being omitted — the p `tool_result` per `tool_call_id`. A store miss silently turns a lossless offload lossy: the known TTL edge. +### Streaming (SSE): what actually happens + +The loop reasons over a complete assistant message, but a streaming response arrives as events, so +on SSE the proxy makes a per-request choice from the request bytes: + +- **No marker in `messages`/`system`** → the model has nothing to expand, so the response is + streamed straight through, byte for byte, with no added latency. +- **A marker is present** → the response is read in full, reconstructed with `expand.AggregateSSE`, + and inspected. If it is a lone expand call the loop runs; otherwise the buffered bytes are + replayed to the client verbatim. + +Buffering costs the client its streaming for that request (time-to-first-byte becomes +time-to-last-byte), which is why the marker test is narrow. It scans **only** `messages` and +`system` — the model-visible content. It used to scan the whole request body, which also matched the +`context_guru_expand` tool description we inject ourselves ("…replaced by a `<>` marker"), +so it was always true and **every** stream was silently buffered (issue #26). + +`/stats` reports this directly: `sse_streamed`, `sse_buffered`, `sse_buffered_pct`, +`sse_ttfb_ms_avg` and `sse_ttfb_ms_avg_buffered`. On traffic that never offloads, `sse_buffered` +should be 0. + +!!! note "Markers arrive HTML-escaped" + A marker the model reads as `<>` travels on the wire as `<>`: + Go's `encoding/json` always escapes `<`, and `sjson` escapes it whenever the value contains a + newline — and markers are appended after a newline. Any check matching markers against raw + request bytes must accept both spellings; `expand.rawMarkerRe` does, deliberately. + +!!! warning "Streaming restoration is Anthropic-only" + `expand.AggregateSSE` reconstructs the Anthropic Messages event stream only. A marker-bearing + **OpenAI** streaming response cannot be reconstructed, so it is replayed raw (fail-open) and + restoration does not fire on that request. Non-streaming OpenAI restoration works normally, as + does streaming Anthropic. Every streaming coding agent in scope speaks the Anthropic dialect; + OpenAI SSE aggregation will be added when a real agent needs it. + ### 3. `GET /expand?id=` The proxy exposes `GET /expand?id=` to recover an offloaded original directly, out of band from the model loop. diff --git a/expand/expand.go b/expand/expand.go index f23d75e..eba1a14 100644 --- a/expand/expand.go +++ b/expand/expand.go @@ -15,6 +15,7 @@ import ( "strings" "github.com/rossoctl/context-guru/store" + "github.com/tidwall/gjson" ) // ToolName is the model-callable tool that retrieves offloaded content. @@ -46,6 +47,39 @@ func HasPlaceholder(s string) bool { return strings.Contains(s, SummaryMarker) || markerRe.MatchString(s) } +// rawMarkerRe matches a full <> marker in RAW, un-decoded request bytes, +// where the angle brackets may arrive HTML-escaped as < / >. +// +// BOTH spellings are load-bearing, and the escaped one is the COMMON case, not an +// exotic client quirk: Go's encoding/json HTML-escapes "<" unconditionally, and sjson +// escapes it whenever the value being set contains a newline — and every Offload +// marker is appended after a newline. So a marker the MODEL reads as <> +// exists in the bytes on the wire only as <>. Any check +// that matches markers against a raw body must accept both forms deliberately. +var rawMarkerRe = regexp.MustCompile(`(?:<|\\u003c){2}cg:([A-Za-z0-9_-]{1,64})(?:>|\\u003e){2}`) + +// HasMarkersInMessages reports whether a request body carries a context-guru +// placeholder in content the MODEL can see and reference — the messages array and +// the system prompt. +// +// It deliberately does NOT scan the whole body. The `tools` array holds our own +// injected expand tool, whose description quotes the marker syntax ("…replaced by a +// <> marker"), HTML-escaped by encoding/json. A whole-body substring check +// therefore matched the tool we had just injected and was a tautology: every request +// looked marker-bearing, so every streaming response was fully buffered and the +// documented zero-added-latency fast path never engaged. Requiring the full marker +// shape is not sufficient on its own — the tool description contains the full shape +// too. Scoping to model-visible content is what fixes it. +func HasMarkersInMessages(body []byte) bool { + for _, field := range [...]string{"messages", "system"} { + if r := gjson.GetBytes(body, field); r.Exists() && + (rawMarkerRe.MatchString(r.Raw) || strings.Contains(r.Raw, SummaryMarker)) { + return true + } + } + return false +} + // ParseMarkers returns the distinct store keys referenced by any markers in s, // in first-seen order. func ParseMarkers(s string) []string { diff --git a/expand/expand_test.go b/expand/expand_test.go index 343c15f..86decd7 100644 --- a/expand/expand_test.go +++ b/expand/expand_test.go @@ -1,9 +1,11 @@ package expand import ( + "strings" "testing" "github.com/rossoctl/context-guru/store" + "github.com/tidwall/sjson" ) func TestToolDefShape(t *testing.T) { @@ -34,6 +36,75 @@ func TestParseMarkersDistinctInOrder(t *testing.T) { } } +// escLT / escGT are the JSON \uXXXX escapes Go's encoders emit for "<" and ">". +// Built from the code points rather than written literally so the test fixtures +// cannot drift from what encoding/json and sjson actually produce. +var ( + escLT = `\u` + "003c" + escGT = `\u` + "003e" +) + +// TestHasMarkersInMessagesIgnoresOwnInjectedTool is the regression guard for the +// tautology of issue #26: the injected expand tool's own description quotes the +// marker syntax, so a whole-body marker check was always true and every streaming +// response got buffered. The check must look only at model-visible content. +func TestHasMarkersInMessagesIgnoresOwnInjectedTool(t *testing.T) { + for _, provider := range []string{"anthropic", "openai"} { + body := []byte(`{"messages":[{"role":"user","content":"hello"}],"tools":[]}`) + injected, ok := Inject(provider, InjectAlways, body, true) + if !ok { + t.Fatalf("%s: Inject should have fired", provider) + } + // The escaped marker shape IS in the injected bytes — that is exactly what used + // to make the old whole-body check unconditionally true. + if !strings.Contains(string(injected), escLT+escLT+"cg:") { + t.Fatalf("%s: expected the tool description to carry an escaped marker: %s", provider, injected) + } + if HasMarkersInMessages(injected) { + t.Fatalf("%s: our own injected tool must not count as a marker: %s", provider, injected) + } + } +} + +// TestHasMarkersInMessagesEscapedForm pins the load-bearing accident: markers reach +// the wire HTML-escaped whenever the value they were appended to contains a newline +// (sjson/encoding/json escape "<"). Both spellings must be found. +func TestHasMarkersInMessagesEscapedForm(t *testing.T) { + escMarker := escLT + escLT + "cg:ABC" + escGT + escGT + cases := map[string]string{ + "plain": `{"messages":[{"role":"user","content":"see <>"}]}`, + "escaped": `{"messages":[{"role":"user","content":"see ` + escMarker + `"}]}`, + "escaped after a newline": `{"messages":[{"role":"user","content":"line1\nline2 ` + escMarker + `"}]}`, + "summary sentinel": `{"messages":[{"role":"user","content":"` + SummaryMarker + ` compacted"}]}`, + "in the system prompt": `{"system":"context: ` + escMarker + `","messages":[]}`, + "escaped in tool_result": `{"messages":[{"role":"user","content":[{"type":"tool_result","content":"out\n` + escMarker + `"}]}]}`, + } + for name, body := range cases { + if !HasMarkersInMessages([]byte(body)) { + t.Errorf("%s: marker not found in %s", name, body) + } + } + // Real sjson output, not a hand-written string: prove the escaping happens exactly + // as offload triggers it (marker appended after a newline). + sjsonBody, _ := sjson.SetBytes([]byte(`{"messages":[{"role":"user","content":""}]}`), + "messages.0.content", "line1\nline2 "+Marker("XYZ")) + if !strings.Contains(string(sjsonBody), escLT+escLT+"cg:XYZ") { + t.Fatalf("expected sjson to escape the marker after a newline: %s", sjsonBody) + } + if !HasMarkersInMessages(sjsonBody) { + t.Fatalf("sjson-escaped marker must be found: %s", sjsonBody) + } + + for name, body := range map[string]string{ + "no markers": `{"messages":[{"role":"user","content":"nothing here"}]}`, + "prefix only but not a marker": `{"messages":[{"role":"user","content":"cg: and ` + escLT + `cg:"}]}`, + } { + if HasMarkersInMessages([]byte(body)) { + t.Errorf("%s: false positive on %s", name, body) + } + } +} + func TestResolve(t *testing.T) { st := store.NewMemory(store.Options{}) if _, ok := Resolve(st, "absent"); ok { diff --git a/metrics/metrics.go b/metrics/metrics.go index c4b345f..bb3d065 100644 --- a/metrics/metrics.go +++ b/metrics/metrics.go @@ -84,6 +84,15 @@ type Aggregator struct { upstreamMsByp float64 // upstream latency on bypassed (baseline) requests upstreamN int64 upstreamNByp int64 + // SSE time-to-first-byte accounting: how long after the upstream call started the + // client got its first response byte, split by whether we had to buffer the whole + // stream to inspect it for an expand call. Buffering is the only thing that stops a + // stream being a stream, so counting it makes that cost visible instead of inferred + // (it used to be unconditional and unmeasured — issue #26). + sseTTFBMs float64 + sseTTFBMsBuf float64 + sseStreamed int64 + sseBuffered int64 } type compStat struct { @@ -183,6 +192,21 @@ func (a *Aggregator) RecordUpstreamLatency(ms float64, bypassed bool) { a.mu.Unlock() } +// RecordSSE notes one streaming response: ms from issuing the upstream request to +// the first byte handed to the client, and whether the stream had to be fully +// buffered first (which turns TTFB into total-response time). +func (a *Aggregator) RecordSSE(ttfbMs float64, buffered bool) { + a.mu.Lock() + if buffered { + a.sseTTFBMsBuf += ttfbMs + a.sseBuffered++ + } else { + a.sseTTFBMs += ttfbMs + a.sseStreamed++ + } + a.mu.Unlock() +} + func (a *Aggregator) Run(r components.RunReport) { a.mu.Lock() defer a.mu.Unlock() @@ -219,6 +243,15 @@ type Snapshot struct { AddedLatencyMsAvg float64 `json:"cg_added_ms_avg"` UpstreamMsAvg float64 `json:"upstream_ms_avg"` UpstreamMsAvgBypassed float64 `json:"upstream_ms_avg_bypassed"` + // SSE streaming health (#26). SSEBuffered counts streams context-guru had to read + // in full before the client saw a byte (to look for an expand tool call); those + // requests lose streaming entirely, so their TTFB is reported separately. A high + // buffered share on traffic that never expands is the regression to watch. + SSEStreamed int64 `json:"sse_streamed"` + SSEBuffered int64 `json:"sse_buffered"` + SSETTFBMsAvg float64 `json:"sse_ttfb_ms_avg"` // streamed-through responses + SSETTFBMsAvgBuf float64 `json:"sse_ttfb_ms_avg_buffered"` // buffered-for-inspection responses + SSEBufferedPct float64 `json:"sse_buffered_pct"` } // Snapshot returns a point-in-time copy of the rollups. @@ -257,11 +290,23 @@ func (a *Aggregator) Snapshot() Snapshot { if a.upstreamNByp > 0 { upAvgByp = a.upstreamMsByp / float64(a.upstreamNByp) } + ttfb, ttfbBuf, bufPct := 0.0, 0.0, 0.0 + if a.sseStreamed > 0 { + ttfb = a.sseTTFBMs / float64(a.sseStreamed) + } + if a.sseBuffered > 0 { + ttfbBuf = a.sseTTFBMsBuf / float64(a.sseBuffered) + } + if n := a.sseStreamed + a.sseBuffered; n > 0 { + bufPct = float64(a.sseBuffered) / float64(n) * 100 + } return Snapshot{ Requests: a.requests, TokensBefore: a.before, TokensAfter: a.after, SavedTokens: saved, SavingsPct: pct, WastedTokens: a.wasted, Bounces: a.bounces, AdjustedSaved: saved - a.wasted, Components: comps, TopPassthrough: passthrough, AddedLatencyMsAvg: addedAvg, UpstreamMsAvg: upAvg, UpstreamMsAvgBypassed: upAvgByp, + SSEStreamed: a.sseStreamed, SSEBuffered: a.sseBuffered, + SSETTFBMsAvg: ttfb, SSETTFBMsAvgBuf: ttfbBuf, SSEBufferedPct: bufPct, } } diff --git a/metrics/metrics_test.go b/metrics/metrics_test.go index c0e66d1..0f4c661 100644 --- a/metrics/metrics_test.go +++ b/metrics/metrics_test.go @@ -73,6 +73,28 @@ func TestLatencyAverages(t *testing.T) { } } +// TestSSEBufferingStats: streamed vs buffered SSE responses average separately and +// the buffered share is reported, so a buffering regression is visible in /stats. +func TestSSEBufferingStats(t *testing.T) { + a := NewAggregator() + if s := a.Snapshot(); s.SSEBufferedPct != 0 || s.SSETTFBMsAvg != 0 { + t.Fatalf("no SSE traffic should report zeros: %+v", s) + } + a.RecordSSE(20, false) + a.RecordSSE(40, false) + a.RecordSSE(900, true) + s := a.Snapshot() + if s.SSEStreamed != 2 || s.SSEBuffered != 1 { + t.Fatalf("counts=%d/%d want 2 streamed / 1 buffered", s.SSEStreamed, s.SSEBuffered) + } + if s.SSETTFBMsAvg != 30 || s.SSETTFBMsAvgBuf != 900 { + t.Fatalf("ttfb=%v buffered=%v want 30/900", s.SSETTFBMsAvg, s.SSETTFBMsAvgBuf) + } + if got := s.SSEBufferedPct; got < 33.3 || got > 33.4 { + t.Fatalf("buffered pct=%v want ~33.33", got) + } +} + // TestMutatedZeroSavingsNotPassthrough locks the fix for cacheinject-style // components: they change the request (add cache_control) but save no content // tokens, so they must NOT be flagged as dead weight in top_passthrough. diff --git a/proxy/proxy.go b/proxy/proxy.go index 19e8fb3..f8df86d 100644 --- a/proxy/proxy.go +++ b/proxy/proxy.go @@ -410,14 +410,21 @@ var errNoUpstream = errors.New("no upstream configured") // lone expand call, the buffered SSE bytes are replayed to the client verbatim (a // one-time latency cost, no correctness change). Requests without markers — early in // a session — stream straight through with zero added latency and no possible expand. +// +// Buffering is the only thing that stops a stream being a stream, so the marker test +// is scoped to the model-visible content and both outcomes are counted (agg.RecordSSE +// → /stats sse_streamed / sse_buffered). It previously matched the expand tool +// description this proxy injects itself, so it was unconditionally true and the +// zero-added-latency promise above never held for any request (issue #26). func (h *Handler) serve(w http.ResponseWriter, r *http.Request, provider bschemas.ModelProvider, up upstream, body []byte, bypassed bool) { injectOn := h.opts.InjectExpand != expand.InjectNever // For SSE we must buffer to inspect (a latency cost), so only do it when the request // actually carries expandable markers (offload happened → the model might expand). - // Tolerate a client that HTML-escapes "<" in JSON (as <) — a false positive - // only costs one buffered response; a false negative would miss a real expand. - bodyStr := string(body) - hasMarkers := expand.HasPlaceholder(bodyStr) || strings.Contains(bodyStr, "\\u003ccg:") + // Scoped to messages+system on purpose: a whole-body check also matched the expand + // tool description we inject ourselves, which made this unconditionally true and + // silently buffered EVERY stream. Both the plain and HTML-escaped marker spellings + // count — see expand.HasMarkersInMessages. + hasMarkers := expand.HasMarkersInMessages(body) for round := 0; ; round++ { upStart := time.Now() resp, err := h.doUpstream(r, up, body) @@ -434,11 +441,19 @@ func (h *Handler) serve(w http.ResponseWriter, r *http.Request, provider bschema // buffered+inspected when markers are present (else stream through, no added latency). checkExpand := injectOn && round < maxExpandRounds && (!isSSE || hasMarkers) if !checkExpand { - h.stream(w, resp) + first := h.stream(w, resp) + if isSSE && h.agg != nil { + h.agg.RecordSSE(msSince(upStart, first), false) + } return } respBody, _ := io.ReadAll(resp.Body) resp.Body.Close() + if isSSE && h.agg != nil { + // Buffered: the client sees nothing until the whole stream has arrived, so + // its first byte lands no earlier than now. + h.agg.RecordSSE(float64(time.Since(upStart).Microseconds())/1000.0, true) + } // Reconstruct the message the loop reasons over. For SSE, aggregate the events; // if that fails, replay the raw stream unchanged (fail-open). @@ -514,8 +529,9 @@ func (h *Handler) doUpstream(r *http.Request, up upstream, body []byte) (*http.R return h.client.Do(req) } -// stream copies an upstream response through with flushing (SSE-friendly). -func (h *Handler) stream(w http.ResponseWriter, resp *http.Response) { +// stream copies an upstream response through with flushing (SSE-friendly) and +// returns the instant the client got its first byte (zero if the body was empty). +func (h *Handler) stream(w http.ResponseWriter, resp *http.Response) (firstByte time.Time) { defer resp.Body.Close() copyHeaders(w.Header(), resp.Header) w.WriteHeader(resp.StatusCode) @@ -524,6 +540,9 @@ func (h *Handler) stream(w http.ResponseWriter, resp *http.Response) { for { n, rerr := resp.Body.Read(buf) if n > 0 { + if firstByte.IsZero() { + firstByte = time.Now() + } w.Write(buf[:n]) if flush != nil { flush.Flush() @@ -533,6 +552,16 @@ func (h *Handler) stream(w http.ResponseWriter, resp *http.Response) { break } } + return firstByte +} + +// msSince returns milliseconds from start to at (falling back to now if the +// response carried no bytes). +func msSince(start, at time.Time) float64 { + if at.IsZero() { + at = time.Now() + } + return float64(at.Sub(start).Microseconds()) / 1000.0 } func (h *Handler) stats(w http.ResponseWriter, _ *http.Request) { diff --git a/proxy/proxy_test.go b/proxy/proxy_test.go index 7d6f460..c44a0bc 100644 --- a/proxy/proxy_test.go +++ b/proxy/proxy_test.go @@ -8,6 +8,7 @@ import ( "net/http/httptest" "strings" "testing" + "time" _ "github.com/rossoctl/context-guru/components/all" "github.com/rossoctl/context-guru/config" @@ -430,6 +431,293 @@ func TestExpandPartialResolutionWellFormed(t *testing.T) { } } +// anthropicSSEBody builds a streaming Anthropic request that declares a tool (so +// expand injection fires) without Go's HTML-escaping, so the caller controls exactly +// which marker spelling reaches the proxy. +func anthropicSSEBody(t *testing.T, userText string) []byte { + t.Helper() + var bb bytes.Buffer + enc := json.NewEncoder(&bb) + enc.SetEscapeHTML(false) + if err := enc.Encode(map[string]any{ + "model": "claude", + "stream": true, + "tools": []map[string]any{ + {"name": "Bash", "description": "run", "input_schema": map[string]any{"type": "object"}}, + }, + "messages": []map[string]any{{"role": "user", "content": userText}}, + }); err != nil { + t.Fatal(err) + } + return bb.Bytes() +} + +// anthropicSSEBodyHTMLEscaped is the same request encoded WITH Go's HTML escaping — +// i.e. any marker in userText reaches the proxy only as <>, which +// is how markers actually arrive on the wire (see expand.rawMarkerRe). +func anthropicSSEBodyHTMLEscaped(t *testing.T, userText string) []byte { + t.Helper() + b, err := json.Marshal(map[string]any{ + "model": "claude", + "stream": true, + "tools": []map[string]any{ + {"name": "Bash", "description": "run", "input_schema": map[string]any{"type": "object"}}, + }, + "messages": []map[string]any{{"role": "user", "content": userText}}, + }) + if err != nil { + t.Fatal(err) + } + if strings.Contains(userText, "cg:") && !strings.Contains(string(b), `u003ccg:`) { + t.Fatalf("fixture is not HTML-escaped as expected: %s", b) + } + return b +} + +// sseTextStream is a minimal Anthropic text event-stream, split so a test upstream +// can send the head, pause, and only then finish. +const ( + sseHead = "event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"role\":\"assistant\"}}\n\n" + + "event: content_block_start\ndata: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"text\",\"text\":\"\"}}\n\n" + + "event: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"first\"}}\n\n" + sseTail = "event: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"last\"}}\n\n" + + "event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n" +) + +// TestMarkerFreeSSEStreamsThrough is the failing-test proof for issue #26. The +// upstream sends the head of an event-stream, then blocks until the test says the +// client has already seen bytes. If context-guru buffers the response, nothing +// reaches the client until the upstream finishes, the upstream never gets released, +// and the test deadlocks out — which is exactly what happened before the fix, +// because hasMarkers matched the expand tool description we inject ourselves. +// +// The request carries NO marker, so per the documented contract ("Requests without +// markers stream straight through with zero added latency") it must not be buffered. +func TestMarkerFreeSSEStreamsThrough(t *testing.T) { + release := make(chan struct{}) + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + w.Write([]byte(sseHead)) + w.(http.Flusher).Flush() + select { + case <-release: + case <-time.After(5 * time.Second): // fail fast instead of hanging the suite + } + w.Write([]byte(sseTail)) + })) + defer upstream.Close() + + h, _ := buildHandler(t, "pipeline: []\n", upstream.URL) + srv := httptest.NewServer(h.Mux()) + defer srv.Close() + + body := anthropicSSEBody(t, "no markers in this request at all") + resp, err := http.Post(srv.URL+"/anthropic/v1/messages", "application/json", strings.NewReader(string(body))) + if err != nil { + t.Fatal(err) + } + defer resp.Body.Close() + + // Read the first chunk. The upstream has NOT sent the tail yet and will not until + // this test releases it, so a streaming proxy can only hand us the head. A + // buffering proxy cannot return anything here at all — it deadlocks until the + // upstream's own 5s escape hatch fires, and then delivers head+tail in one go, + // which is what the "last" assertion below detects. + buf := make([]byte, 4096) + n, rerr := resp.Body.Read(buf) + first := string(buf[:n]) + close(release) + if n == 0 { + t.Fatalf("first read returned no bytes: %v", rerr) + } + if !strings.Contains(first, "first") { + t.Fatalf("expected the first streamed delta, got %q", first) + } + if strings.Contains(first, "last") { + t.Fatal("marker-free SSE response was BUFFERED: the client received the whole " + + "stream at once, after the upstream had finished (issue #26 — hasMarkers " + + "matched the expand tool description we inject ourselves)") + } + rest, _ := io.ReadAll(resp.Body) + if !strings.Contains(string(rest), "last") { + t.Fatalf("stream did not complete, tail=%q", rest) + } + + // And the fast path must be visible in /stats, not merely inferred. + var snap metrics.Snapshot + st, _ := http.Get(srv.URL + "/stats") + json.NewDecoder(st.Body).Decode(&snap) + st.Body.Close() + if snap.SSEStreamed != 1 || snap.SSEBuffered != 0 { + t.Fatalf("stats should show one streamed, zero buffered SSE: %+v", snap) + } +} + +// TestMarkerBearingSSEIsBuffered is the other half of the contract: when the request +// really does carry a marker (in its HTML-escaped wire form, which is how markers +// actually arrive), buffering is correct and must still happen — otherwise a real +// expand call would stream past uninspected. +func TestMarkerBearingSSEIsBuffered(t *testing.T) { + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + w.Write([]byte(sseHead + sseTail)) + })) + defer upstream.Close() + + h, _ := buildHandler(t, "pipeline: []\n", upstream.URL) + srv := httptest.NewServer(h.Mux()) + defer srv.Close() + + // The escaped spelling, written after a newline exactly as offload emits it. + body := anthropicSSEBodyHTMLEscaped(t, "output\nline2 <>") + resp, err := http.Post(srv.URL+"/anthropic/v1/messages", "application/json", strings.NewReader(string(body))) + if err != nil { + t.Fatal(err) + } + out, _ := io.ReadAll(resp.Body) + resp.Body.Close() + if !strings.Contains(string(out), "first") || !strings.Contains(string(out), "last") { + t.Fatalf("normal answer must be replayed verbatim: %q", out) + } + + var snap metrics.Snapshot + st, _ := http.Get(srv.URL + "/stats") + json.NewDecoder(st.Body).Decode(&snap) + st.Body.Close() + if snap.SSEBuffered != 1 || snap.SSEStreamed != 0 { + t.Fatalf("a marker-bearing SSE request must still be buffered for inspection: %+v", snap) + } +} + +// TestExpandSSEWithOtherToolReplaysVerbatim covers the otherTools bail on the +// streaming path: the model batches expand alongside a Bash call. The proxy cannot +// answer only half a batch (the client owns Bash), so it must replay the stream +// unchanged rather than continue — and the client's stream must stay well-formed. +func TestExpandSSEWithOtherToolReplaysVerbatim(t *testing.T) { + var calls int + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + calls++ + w.Header().Set("Content-Type", "text/event-stream") + w.Write([]byte("event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"role\":\"assistant\"}}\n\n" + + "event: content_block_start\ndata: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"tool_use\",\"id\":\"call_1\",\"name\":\"context_guru_expand\"}}\n\n" + + "event: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"input_json_delta\",\"partial_json\":\"{\\\"id\\\":\\\"HASH\\\"}\"}}\n\n" + + "event: content_block_stop\ndata: {\"type\":\"content_block_stop\",\"index\":0}\n\n" + + "event: content_block_start\ndata: {\"type\":\"content_block_start\",\"index\":1,\"content_block\":{\"type\":\"tool_use\",\"id\":\"call_2\",\"name\":\"Bash\"}}\n\n" + + "event: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":1,\"delta\":{\"type\":\"input_json_delta\",\"partial_json\":\"{\\\"command\\\":\\\"ls\\\"}\"}}\n\n" + + "event: content_block_stop\ndata: {\"type\":\"content_block_stop\",\"index\":1}\n\n" + + "event: message_delta\ndata: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"tool_use\"}}\n\n" + + "event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n")) + })) + defer upstream.Close() + + h, st := buildHandler(t, "pipeline: []\n", upstream.URL) + st.Put("HASH", []byte("THE ORIGINAL CONTENT")) + srv := httptest.NewServer(h.Mux()) + defer srv.Close() + + body := anthropicSSEBody(t, "look at <> then list files") + resp, err := http.Post(srv.URL+"/anthropic/v1/messages", "application/json", strings.NewReader(string(body))) + if err != nil { + t.Fatal(err) + } + out, _ := io.ReadAll(resp.Body) + resp.Body.Close() + + if calls != 1 { + t.Fatalf("batched expand+Bash must NOT trigger a continuation, got %d upstream calls", calls) + } + // Verbatim replay: both blocks intact, indices unrenumbered, no injected content. + if !strings.Contains(string(out), `"index":1`) || !strings.Contains(string(out), `"name":"Bash"`) { + t.Fatalf("client must receive the original stream unchanged: %s", out) + } + if strings.Contains(string(out), "THE ORIGINAL CONTENT") { + t.Fatalf("proxy must not splice resolved content into a stream it declined: %s", out) + } +} + +// TestExpandSSEMultiRoundCapped drives the streaming loop past maxExpandRounds: an +// upstream that answers every request with another expand call must be cut off, and +// the client must still get a well-formed stream (the model's own last call). +func TestExpandSSEMultiRoundCapped(t *testing.T) { + var calls int + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + io.Copy(io.Discard, r.Body) + calls++ + w.Header().Set("Content-Type", "text/event-stream") + w.Write([]byte("event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"role\":\"assistant\"}}\n\n" + + "event: content_block_start\ndata: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"tool_use\",\"id\":\"call_1\",\"name\":\"context_guru_expand\"}}\n\n" + + "event: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"input_json_delta\",\"partial_json\":\"{\\\"id\\\":\\\"HASH\\\"}\"}}\n\n" + + "event: content_block_stop\ndata: {\"type\":\"content_block_stop\",\"index\":0}\n\n" + + "event: message_delta\ndata: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"tool_use\"}}\n\n" + + "event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n")) + })) + defer upstream.Close() + + h, st := buildHandler(t, "pipeline: []\n", upstream.URL) + st.Put("HASH", []byte("THE ORIGINAL CONTENT")) + srv := httptest.NewServer(h.Mux()) + defer srv.Close() + + body := anthropicSSEBody(t, "look at <>") + resp, err := http.Post(srv.URL+"/anthropic/v1/messages", "application/json", strings.NewReader(string(body))) + if err != nil { + t.Fatal(err) + } + out, _ := io.ReadAll(resp.Body) + resp.Body.Close() + + // maxExpandRounds continuations, then one final pass-through: 4 upstream calls. + if calls != 4 { + t.Fatalf("round cap not honored: %d upstream calls (want 4 = 3 rounds + terminal)", calls) + } + if !strings.Contains(string(out), "message_stop") { + t.Fatalf("client must still receive a complete stream after the cap: %s", out) + } +} + +// TestExpandOpenAISSEFallsBackToRaw documents (and pins) the OpenAI streaming +// limitation: AggregateSSE only reconstructs the Anthropic event stream, so a +// marker-bearing OpenAI SSE response is replayed raw and restoration does not fire. +// Correctness is preserved (fail-open); only the feature is absent. +func TestExpandOpenAISSEFallsBackToRaw(t *testing.T) { + var calls int + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + calls++ + w.Header().Set("Content-Type", "text/event-stream") + w.Write([]byte(`data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_1","type":"function","function":{"name":"context_guru_expand","arguments":"{\"id\":\"HASH\"}"}}]}}]}` + "\n\n" + + "data: [DONE]\n\n")) + })) + defer upstream.Close() + + h, st := buildHandler(t, "pipeline: []\n", upstream.URL) + st.Put("HASH", []byte("THE ORIGINAL CONTENT")) + srv := httptest.NewServer(h.Mux()) + defer srv.Close() + + var bb bytes.Buffer + enc := json.NewEncoder(&bb) + enc.SetEscapeHTML(false) + _ = enc.Encode(map[string]any{ + "model": "gpt-x", + "stream": true, + "tools": []map[string]any{{"type": "function", "function": map[string]any{"name": "Bash", "parameters": map[string]any{"type": "object"}}}}, + "messages": []map[string]any{{"role": "user", "content": "look at <>"}}, + }) + resp, err := http.Post(srv.URL+"/openai/v1/chat/completions", "application/json", strings.NewReader(bb.String())) + if err != nil { + t.Fatal(err) + } + out, _ := io.ReadAll(resp.Body) + resp.Body.Close() + + if calls != 1 { + t.Fatalf("OpenAI SSE cannot be aggregated, so no continuation is possible: %d calls", calls) + } + if !strings.Contains(string(out), "[DONE]") || strings.Contains(string(out), "THE ORIGINAL CONTENT") { + t.Fatalf("OpenAI SSE must be replayed raw (fail-open): %s", out) + } +} + func TestExpandRoundTrip(t *testing.T) { upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { w.Write([]byte(`{}`)) })) defer upstream.Close() From 566814962baebef2efc1303f38487e56f597fba6 Mon Sep 17 00:00:00 2001 From: Osher-Elhadad Date: Mon, 10 Aug 2026 04:18:40 +0000 Subject: [PATCH 2/2] fix(metrics): count SSE buffering per client request, not per upstream round MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review of #33 caught that both RecordSSE call sites sat inside the expand continuation loop and timed from that round's upStart. Two consequences, both undermining the metric this change exists to make trustworthy: - sse_buffered_pct was documented as a share of responses but was a share of upstream calls: one client request driving the round cap reported streamed=1 buffered=3 pct=75.0. - Worse, the terminal round after the cap recorded a *streamed* sample timed from the last round alone, so a request the client waited three round-trips for contributed a healthy-looking number to sse_ttfb_ms_avg. Hoist reqStart plus sticky sse/sseBuffered flags above the loop and record once in a defer, so every terminal return path is covered — stream-through, aggregate-failure replay, normal-answer replay, nothing-resolved replay, and the round-cap exit — and so future return paths cannot silently skip accounting. sseBuffered is sticky: once a round has been buffered the client has already lost its stream, however the request ends. Also from review: - Cover the fail-open path INSIDE aggregateAnthropicSSE (truncated input_json_delta → sse.go:132 returns nil,false). The existing OpenAI test only exercised the provider gate at sse.go:21, a different branch. - Match \uXXXX escapes case-insensitively in rawMarkerRe. < was a false negative, which is the class this regexp exists to prevent: a real expand call streamed past uninspected is worse than over-buffering. - Derive the test's escape fixtures from encoding/json itself instead of a hand-written literal that merely claimed to be drift-proof; the helper panics if the encoder ever stops escaping "<", which is the signal to revisit the raw-body matcher. - Document sse_ttfb_ms_avg_buffered as time-to-LAST-byte by construction on the field itself, not only in prose, and note the per-request counting basis. - Correct "encoding/json always escapes <" — false under SetEscapeHTML(false), which this package's own test helper uses — and note that the matcher's plain-form branch makes a doc quoting a literal marker count as marker-bearing. The round-cap test now asserts one client request yields exactly one sample; it reproduces streamed=1/buffered=3/pct=75.0 against the previous placement. Signed-off-by: Osher-Elhadad --- docs/design.md | 7 +++-- docs/how-to/recover-context.md | 29 +++++++++++------ expand/expand.go | 19 +++++++----- expand/expand_test.go | 38 ++++++++++++++++++----- metrics/metrics.go | 16 +++++++--- proxy/proxy.go | 30 +++++++++++++----- proxy/proxy_test.go | 57 ++++++++++++++++++++++++++++++++++ 7 files changed, 158 insertions(+), 38 deletions(-) diff --git a/docs/design.md b/docs/design.md index 6f0bccb..c2ee1b7 100644 --- a/docs/design.md +++ b/docs/design.md @@ -173,9 +173,10 @@ scanning the whole body also matched the expand tool description the host inject it unconditionally true (issue #26). `/stats` exposes `sse_streamed` / `sse_buffered` / `sse_buffered_pct` and the two TTFB averages so the fast path is measured, not assumed. -**Markers on the wire are HTML-escaped.** `encoding/json` always escapes `<`; `sjson` escapes it -whenever the value contains a newline, and markers are appended after a newline. So `<>` in the -model's view is `<>` in the bytes. Marker matching on *decoded* content +**Markers on the wire are usually HTML-escaped.** `encoding/json` escapes `<` by default — a caller +can opt out with `Encoder.SetEscapeHTML(false)`, and some non-Go clients never escape it — and `sjson` +escapes it whenever the value contains a newline, which is how every marker is appended. So `<>` +in the model's view is normally `<>` in the bytes. Marker matching on *decoded* content (`expand.HasPlaceholder`, used by the components) sees the plain form; matching on *raw request bytes* (`expand.rawMarkerRe`, used by the host's streaming decision) must accept both, and does. diff --git a/docs/how-to/recover-context.md b/docs/how-to/recover-context.md index 9df77ef..cbbd133 100644 --- a/docs/how-to/recover-context.md +++ b/docs/how-to/recover-context.md @@ -65,15 +65,26 @@ time-to-last-byte), which is why the marker test is narrow. It scans **only** `m `context_guru_expand` tool description we inject ourselves ("…replaced by a `<>` marker"), so it was always true and **every** stream was silently buffered (issue #26). -`/stats` reports this directly: `sse_streamed`, `sse_buffered`, `sse_buffered_pct`, -`sse_ttfb_ms_avg` and `sse_ttfb_ms_avg_buffered`. On traffic that never offloads, `sse_buffered` -should be 0. - -!!! note "Markers arrive HTML-escaped" - A marker the model reads as `<>` travels on the wire as `<>`: - Go's `encoding/json` always escapes `<`, and `sjson` escapes it whenever the value contains a - newline — and markers are appended after a newline. Any check matching markers against raw - request bytes must accept both spellings; `expand.rawMarkerRe` does, deliberately. +`/stats` reports this directly, counted **once per client request** (not per upstream round, so a +request that drove several expand rounds is one sample): `sse_streamed`, `sse_buffered`, +`sse_buffered_pct`, `sse_ttfb_ms_avg` and `sse_ttfb_ms_avg_buffered`. On traffic that never offloads, +`sse_buffered` stays 0; it starts counting from the first turn that carries a marker. Note that +`sse_ttfb_ms_avg_buffered` is time-to-*last*-byte by construction — a buffered response is read in +full before the client is written to — so it is not comparable to `sse_ttfb_ms_avg`. + +!!! note "Markers usually arrive HTML-escaped" + A marker the model reads as `<>` normally travels on the wire as + `<>`: Go's `encoding/json` escapes `<` by default (callers can opt out + with `Encoder.SetEscapeHTML(false)`, and some non-Go clients never escape it), and `sjson` + escapes it whenever the value contains a newline — which is how every marker is appended. Any + check matching markers against raw request bytes must accept both spellings, case-insensitively + (`<` is as valid as `<`); `expand.rawMarkerRe` does, deliberately. A miss there is a + false negative — a real expand call streamed past uninspected — which is worse than + over-buffering. + + Because that matcher accepts the plain form too, a document or message quoting a literal + `<>` example (like this page) counts as marker-bearing. That is the intended bias: + over-inspect rather than miss a real call. !!! warning "Streaming restoration is Anthropic-only" `expand.AggregateSSE` reconstructs the Anthropic Messages event stream only. A marker-bearing diff --git a/expand/expand.go b/expand/expand.go index eba1a14..b34907f 100644 --- a/expand/expand.go +++ b/expand/expand.go @@ -51,12 +51,17 @@ func HasPlaceholder(s string) bool { // where the angle brackets may arrive HTML-escaped as < / >. // // BOTH spellings are load-bearing, and the escaped one is the COMMON case, not an -// exotic client quirk: Go's encoding/json HTML-escapes "<" unconditionally, and sjson -// escapes it whenever the value being set contains a newline — and every Offload -// marker is appended after a newline. So a marker the MODEL reads as <> -// exists in the bytes on the wire only as <>. Any check -// that matches markers against a raw body must accept both forms deliberately. -var rawMarkerRe = regexp.MustCompile(`(?:<|\\u003c){2}cg:([A-Za-z0-9_-]{1,64})(?:>|\\u003e){2}`) +// exotic client quirk: Go's encoding/json HTML-escapes "<" by default (unless a caller +// opts out via Encoder.SetEscapeHTML(false)), and sjson escapes it whenever the value +// being set contains a newline — and every Offload marker is appended after a newline. +// So a marker the MODEL reads as <> usually exists in the bytes on the wire +// only as <>. Any check matching markers against a raw body must +// accept both forms deliberately. +// +// The escape alternatives are case-insensitive: \u003C is as valid as \u003c, and a miss +// here is a FALSE NEGATIVE — a real expand call streamed past uninspected, which is +// worse than the over-buffering this regexp exists to prevent. +var rawMarkerRe = regexp.MustCompile(`(?:<|(?i:\\u003c)){2}cg:([A-Za-z0-9_-]{1,64})(?:>|(?i:\\u003e)){2}`) // HasMarkersInMessages reports whether a request body carries a context-guru // placeholder in content the MODEL can see and reference — the messages array and @@ -64,7 +69,7 @@ var rawMarkerRe = regexp.MustCompile(`(?:<|\\u003c){2}cg:([A-Za-z0-9_-]{1,64})(? // // It deliberately does NOT scan the whole body. The `tools` array holds our own // injected expand tool, whose description quotes the marker syntax ("…replaced by a -// <> marker"), HTML-escaped by encoding/json. A whole-body substring check +// <> marker"), HTML-escaped by ToolDefRaw's encoding/json. A whole-body check // therefore matched the tool we had just injected and was a tautology: every request // looked marker-bearing, so every streaming response was fully buffered and the // documented zero-added-latency fast path never engaged. Requiring the full marker diff --git a/expand/expand_test.go b/expand/expand_test.go index 86decd7..fab6e21 100644 --- a/expand/expand_test.go +++ b/expand/expand_test.go @@ -1,6 +1,7 @@ package expand import ( + "encoding/json" "strings" "testing" @@ -36,13 +37,31 @@ func TestParseMarkersDistinctInOrder(t *testing.T) { } } -// escLT / escGT are the JSON \uXXXX escapes Go's encoders emit for "<" and ">". -// Built from the code points rather than written literally so the test fixtures -// cannot drift from what encoding/json and sjson actually produce. -var ( - escLT = `\u` + "003c" - escGT = `\u` + "003e" -) +// escLT / escGT are the JSON \uXXXX escapes Go's encoders emit for "<" and ">", +// obtained by ASKING encoding/json rather than hand-writing them, so a fixture can +// never claim an escape form the encoder does not actually produce. +var escLT, escGT = jsonEscape('<'), jsonEscape('>') + +// jsonEscape marshals a single rune with HTML escaping on and returns the escape +// sequence encoding/json chose for it (e.g. "<" for '<'). +func jsonEscape(r rune) string { + b, err := json.Marshal(string(r)) + if err != nil { + panic(err) + } + s := strings.Trim(string(b), `"`) + if !strings.HasPrefix(s, `\u`) { + panic("encoding/json no longer escapes " + string(r) + " (got " + s + "); " + + "the raw-body marker matcher's escaped-form handling needs revisiting") + } + return s +} + +// upperHex uppercases the hex digits of a \uXXXX escape but not the "u" ("<" +// -> "<"), matching what a non-Go JSON encoder may legitimately emit. +func upperHex(esc string) string { + return `\u` + strings.ToUpper(strings.TrimPrefix(esc, `\u`)) +} // TestHasMarkersInMessagesIgnoresOwnInjectedTool is the regression guard for the // tautology of issue #26: the injected expand tool's own description quotes the @@ -78,6 +97,11 @@ func TestHasMarkersInMessagesEscapedForm(t *testing.T) { "summary sentinel": `{"messages":[{"role":"user","content":"` + SummaryMarker + ` compacted"}]}`, "in the system prompt": `{"system":"context: ` + escMarker + `","messages":[]}`, "escaped in tool_result": `{"messages":[{"role":"user","content":[{"type":"tool_result","content":"out\n` + escMarker + `"}]}]}`, + // < is as valid an escape as <. Missing it would be a false negative: + // a real expand call streamed past uninspected. + "uppercase hex escape": `{"messages":[{"role":"user","content":"see ` + + upperHex(escLT) + upperHex(escLT) + `cg:ABC` + + upperHex(escGT) + upperHex(escGT) + `"}]}`, } for name, body := range cases { if !HasMarkersInMessages([]byte(body)) { diff --git a/metrics/metrics.go b/metrics/metrics.go index bb3d065..3263a13 100644 --- a/metrics/metrics.go +++ b/metrics/metrics.go @@ -246,11 +246,17 @@ type Snapshot struct { // SSE streaming health (#26). SSEBuffered counts streams context-guru had to read // in full before the client saw a byte (to look for an expand tool call); those // requests lose streaming entirely, so their TTFB is reported separately. A high - // buffered share on traffic that never expands is the regression to watch. - SSEStreamed int64 `json:"sse_streamed"` - SSEBuffered int64 `json:"sse_buffered"` - SSETTFBMsAvg float64 `json:"sse_ttfb_ms_avg"` // streamed-through responses - SSETTFBMsAvgBuf float64 `json:"sse_ttfb_ms_avg_buffered"` // buffered-for-inspection responses + // buffered share on traffic that never expands is the regression to watch. All four + // count once per CLIENT REQUEST, not per upstream round: a request that drove + // several expand rounds waited for all of them. + SSEStreamed int64 `json:"sse_streamed"` + SSEBuffered int64 `json:"sse_buffered"` + SSETTFBMsAvg float64 `json:"sse_ttfb_ms_avg"` // streamed-through requests: a real TTFB + // SSETTFBMsAvgBuf is time-to-LAST-byte by construction, not a comparable TTFB: a + // buffered response is read in full before the client is written to, so its first + // byte cannot precede the buffer completing. Read it as "what buffering cost these + // requests", not as a latency to compare against sse_ttfb_ms_avg. + SSETTFBMsAvgBuf float64 `json:"sse_ttfb_ms_avg_buffered"` SSEBufferedPct float64 `json:"sse_buffered_pct"` } diff --git a/proxy/proxy.go b/proxy/proxy.go index f8df86d..f5ad0a7 100644 --- a/proxy/proxy.go +++ b/proxy/proxy.go @@ -425,6 +425,20 @@ func (h *Handler) serve(w http.ResponseWriter, r *http.Request, provider bschema // silently buffered EVERY stream. Both the plain and HTML-escaped marker spellings // count — see expand.HasMarkersInMessages. hasMarkers := expand.HasMarkersInMessages(body) + // SSE accounting is PER CLIENT REQUEST, not per upstream round: one client request + // that drives several expand rounds waited for all of them, so timing a single + // round would report a healthy TTFB for a client that waited three round-trips. + // Recorded in a defer so every terminal return path is covered exactly once — + // stream-through, aggregate-failure replay, normal-answer replay, nothing-resolved + // replay, and the round-cap exit. + reqStart := time.Now() + sse, sseBuffered := false, false + var sseFirstByte time.Time // zero on buffered paths: the client's first byte is the write itself + defer func() { + if sse && h.agg != nil { + h.agg.RecordSSE(msSince(reqStart, sseFirstByte), sseBuffered) + } + }() for round := 0; ; round++ { upStart := time.Now() resp, err := h.doUpstream(r, up, body) @@ -441,18 +455,20 @@ func (h *Handler) serve(w http.ResponseWriter, r *http.Request, provider bschema // buffered+inspected when markers are present (else stream through, no added latency). checkExpand := injectOn && round < maxExpandRounds && (!isSSE || hasMarkers) if !checkExpand { - first := h.stream(w, resp) - if isSSE && h.agg != nil { - h.agg.RecordSSE(msSince(upStart, first), false) + // sseBuffered is sticky: if an earlier round was buffered the client already + // lost its stream, so this request counts as buffered however it ends. + sse = sse || isSSE + if first := h.stream(w, resp); !sseBuffered { + sseFirstByte = first } return } respBody, _ := io.ReadAll(resp.Body) resp.Body.Close() - if isSSE && h.agg != nil { - // Buffered: the client sees nothing until the whole stream has arrived, so - // its first byte lands no earlier than now. - h.agg.RecordSSE(float64(time.Since(upStart).Microseconds())/1000.0, true) + if isSSE { + // Buffered: the client sees nothing until the whole stream has arrived, so its + // first byte lands no earlier than the write on whichever path we return from. + sse, sseBuffered, sseFirstByte = true, true, time.Time{} } // Reconstruct the message the loop reasons over. For SSE, aggregate the events; diff --git a/proxy/proxy_test.go b/proxy/proxy_test.go index c44a0bc..99fcd92 100644 --- a/proxy/proxy_test.go +++ b/proxy/proxy_test.go @@ -673,6 +673,63 @@ func TestExpandSSEMultiRoundCapped(t *testing.T) { if !strings.Contains(string(out), "message_stop") { t.Fatalf("client must still receive a complete stream after the cap: %s", out) } + + // SSE stats are per CLIENT REQUEST, not per upstream round. This one request drove + // 4 upstream calls; recording per round would report streamed=1/buffered=3 and — + // worse — count the terminal round as a healthy "streamed" TTFB timed from that + // round alone, hiding the 3 round-trips the client actually waited for. + var snap metrics.Snapshot + stx, _ := http.Get(srv.URL + "/stats") + json.NewDecoder(stx.Body).Decode(&snap) + stx.Body.Close() + if snap.SSEBuffered != 1 || snap.SSEStreamed != 0 { + t.Fatalf("one client request must yield exactly one buffered sample, got %+v", snap) + } + if snap.SSEBufferedPct != 100 { + t.Fatalf("buffered_pct must be a share of requests (want 100), got %v", snap.SSEBufferedPct) + } +} + +// TestExpandSSEAggregateFailureReplaysRaw covers the fail-open path INSIDE +// aggregateAnthropicSSE (expand/sse.go:132): a truncated input_json_delta leaves the +// tool_use input unparseable, so AggregateSSE returns ok=false even though the +// provider IS anthropic — a different branch from the provider gate at sse.go:21. +// The client must still receive the original bytes unchanged. +func TestExpandSSEAggregateFailureReplaysRaw(t *testing.T) { + var calls int + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + calls++ + w.Header().Set("Content-Type", "text/event-stream") + // partial_json is TRUNCATED — it cannot reconstruct to valid JSON. + w.Write([]byte("event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"role\":\"assistant\"}}\n\n" + + "event: content_block_start\ndata: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"tool_use\",\"id\":\"call_1\",\"name\":\"context_guru_expand\"}}\n\n" + + "event: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"input_json_delta\",\"partial_json\":\"{\\\"id\\\":\\\"HA\"}}\n\n" + + "event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n")) + })) + defer upstream.Close() + + h, st := buildHandler(t, "pipeline: []\n", upstream.URL) + st.Put("HASH", []byte("THE ORIGINAL CONTENT")) + srv := httptest.NewServer(h.Mux()) + defer srv.Close() + + body := anthropicSSEBody(t, "look at <>") + resp, err := http.Post(srv.URL+"/anthropic/v1/messages", "application/json", strings.NewReader(string(body))) + if err != nil { + t.Fatal(err) + } + out, _ := io.ReadAll(resp.Body) + resp.Body.Close() + + if calls != 1 { + t.Fatalf("an unreconstructable stream must not drive a continuation: %d calls", calls) + } + if !strings.Contains(string(out), `\"id\":\"HA`) || !strings.Contains(string(out), "message_stop") { + t.Fatalf("client must get the raw stream back verbatim (fail-open): %s", out) + } + if strings.Contains(string(out), "THE ORIGINAL CONTENT") { + t.Fatalf("nothing may be spliced into a stream we could not parse: %s", out) + } } // TestExpandOpenAISSEFallsBackToRaw documents (and pins) the OpenAI streaming