| 1 | package main |
| 2 | |
| 3 | import ( |
| 4 | "bufio" |
| 5 | "bytes" |
| 6 | "encoding/json" |
| 7 | "fmt" |
| 8 | "io" |
| 9 | "maps" |
| 10 | "net" |
| 11 | "net/http" |
| 12 | "net/url" |
| 13 | "strconv" |
| 14 | "strings" |
| 15 | "sync" |
| 16 | ) |
| 17 | |
| 18 | // meter is the harness-neutral measuring point: a proxy every harness talks to |
| 19 | // instead of the provider, so nobody in a comparison reports their own token |
| 20 | // use. It is also where LongRun injects provider faults, because the request |
| 21 | // boundary is the only place both benchmarks can reach without a harness's |
| 22 | // cooperation. |
| 23 | type meter struct { |
| 24 | upstream *url.URL |
| 25 | client *http.Client |
| 26 | faults faultScript |
| 27 | |
| 28 | mu sync.Mutex |
| 29 | lastFault int |
| 30 | meterUsage |
| 31 | } |
| 32 | |
| 33 | // meterUsage is what the proxy observed. WithoutUsage is reported rather than |
| 34 | // folded into zero: a harness whose responses carry no usage block is |
| 35 | // unmeasured, and that is a finding, not a zero. |
| 36 | type meterUsage struct { |
| 37 | Requests int `json:"requests"` |
| 38 | PromptTokens int `json:"prompt_tokens"` |
| 39 | CompletionTokens int `json:"completion_tokens"` |
| 40 | CacheHitTokens int `json:"cache_hit_tokens"` |
| 41 | CacheMissTokens int `json:"cache_miss_tokens"` |
| 42 | Injected int `json:"injected_faults,omitempty"` |
| 43 | WithoutUsage int `json:"responses_without_usage,omitempty"` |
| 44 | // RequestsAfterFault counts requests issued after the first injected |
| 45 | // failure: the harness's own evidence that it retried rather than died. |
| 46 | RequestsAfterFault int `json:"requests_after_fault,omitempty"` |
| 47 | } |
| 48 | |
| 49 | // faultScript decides which requests fail. Absolute indices pin a failure to |
| 50 | // an exact point; a cadence scales with the run, which is what a mixed-length |
| 51 | // suite needs — a task that only ever makes four requests would never reach a |
| 52 | // fixed index, and would silently join the unfaulted control group. |
| 53 | type faultScript struct { |
| 54 | at map[int]int // 1-based request index -> status |
| 55 | everyN int |
| 56 | everyStatus int |
| 57 | } |
| 58 | |
| 59 | func (f faultScript) empty() bool { return len(f.at) == 0 && f.everyN == 0 } |
| 60 | |
| 61 | // statusFor reports the status to inject instead of forwarding. An absolute |
| 62 | // index wins over the cadence so a targeted failure stays exactly where it was |
| 63 | // asked for. |
| 64 | func (f faultScript) statusFor(index int) (int, bool) { |
| 65 | if status, ok := f.at[index]; ok { |
| 66 | return status, true |
| 67 | } |
| 68 | if f.everyN > 0 && index%f.everyN == 0 { |
| 69 | return f.everyStatus, true |
| 70 | } |
| 71 | return 0, false |
| 72 | } |
| 73 | |
| 74 | // parseFaultScript reads "3:429,every:5:500": fail the 3rd request with 429 and |
| 75 | // every 5th with 500. Deterministic either way, so a LongRun arm replays the |
| 76 | // same failures across harnesses. |
| 77 | func parseFaultScript(spec string) (faultScript, error) { |
| 78 | out := faultScript{at: map[int]int{}} |
| 79 | if strings.TrimSpace(spec) == "" { |
| 80 | return faultScript{}, nil |
| 81 | } |
| 82 | for field := range strings.SplitSeq(spec, ",") { |
| 83 | field = strings.TrimSpace(field) |
| 84 | if rest, ok := strings.CutPrefix(field, "every:"); ok { |
| 85 | n, status, err := parseFaultPair(field, rest) |
| 86 | if err != nil { |
| 87 | return faultScript{}, err |
| 88 | } |
| 89 | out.everyN, out.everyStatus = n, status |
| 90 | continue |
| 91 | } |
| 92 | index, status, err := parseFaultPair(field, field) |
| 93 | if err != nil { |
| 94 | return faultScript{}, err |
| 95 | } |
| 96 | out.at[index] = status |
| 97 | } |
| 98 | return out, nil |
| 99 | } |
| 100 | |
| 101 | func parseFaultPair(field, pair string) (n, status int, err error) { |
| 102 | left, right, ok := strings.Cut(pair, ":") |
| 103 | if !ok { |
| 104 | return 0, 0, fmt.Errorf("fault %q: want <request-index>:<status> or every:<n>:<status>", field) |
| 105 | } |
| 106 | n, err = strconv.Atoi(strings.TrimSpace(left)) |
| 107 | if err != nil || n < 1 { |
| 108 | return 0, 0, fmt.Errorf("fault %q: the request count must be a positive integer", field) |
| 109 | } |
| 110 | status, err = strconv.Atoi(strings.TrimSpace(right)) |
| 111 | if err != nil || status < 400 || status > 599 { |
| 112 | return 0, 0, fmt.Errorf("fault %q: status must be 4xx or 5xx", field) |
| 113 | } |
| 114 | return n, status, nil |
| 115 | } |
| 116 | |
| 117 | func newMeter(upstream string, faults faultScript) (*meter, error) { |
| 118 | u, err := url.Parse(strings.TrimSuffix(strings.TrimSpace(upstream), "/")) |
| 119 | if err != nil || u.Scheme == "" || u.Host == "" { |
| 120 | return nil, fmt.Errorf("meter upstream %q is not an absolute URL", upstream) |
| 121 | } |
| 122 | return &meter{upstream: u, client: &http.Client{}, faults: faults}, nil |
| 123 | } |
| 124 | |
| 125 | // serve starts the proxy on an ephemeral loopback port and returns its base URL |
| 126 | // plus a stop function. Loopback only: the proxy carries provider credentials. |
| 127 | func (m *meter) serve() (base string, stop func(), err error) { |
| 128 | ln, err := net.Listen("tcp", "127.0.0.1:0") |
| 129 | if err != nil { |
| 130 | return "", nil, err |
| 131 | } |
| 132 | srv := &http.Server{Handler: m} |
| 133 | go func() { _ = srv.Serve(ln) }() |
| 134 | return "http://" + ln.Addr().String(), func() { _ = srv.Close() }, nil |
| 135 | } |
| 136 | |
| 137 | func (m *meter) snapshot() meterUsage { |
| 138 | m.mu.Lock() |
| 139 | defer m.mu.Unlock() |
| 140 | return m.meterUsage |
| 141 | } |
| 142 | |
| 143 | func (m *meter) ServeHTTP(w http.ResponseWriter, r *http.Request) { |
| 144 | m.mu.Lock() |
| 145 | m.Requests++ |
| 146 | index := m.Requests |
| 147 | status, faulted := m.faults.statusFor(index) |
| 148 | switch { |
| 149 | case faulted: |
| 150 | m.Injected++ |
| 151 | m.lastFault = index |
| 152 | case m.lastFault > 0: |
| 153 | // Evidence the harness kept going after being failed. A harness that |
| 154 | // gives up on the first 429 never reaches here, and that is the |
| 155 | // difference between "recovered" and "was never really tested". |
| 156 | m.RequestsAfterFault++ |
| 157 | } |
| 158 | m.mu.Unlock() |
| 159 | |
| 160 | if faulted { |
| 161 | w.Header().Set("Content-Type", "application/json") |
| 162 | w.WriteHeader(status) |
| 163 | fmt.Fprintf(w, `{"error":{"message":"injected by e2ebench meter at request %d","type":"bench_fault"}}`, index) |
| 164 | return |
| 165 | } |
| 166 | |
| 167 | body, err := io.ReadAll(r.Body) |
| 168 | if err != nil { |
| 169 | http.Error(w, "read body: "+err.Error(), http.StatusBadGateway) |
| 170 | return |
| 171 | } |
| 172 | body = requestUsageOptIn(body) |
| 173 | |
| 174 | target := *m.upstream |
| 175 | target.Path = strings.TrimSuffix(m.upstream.Path, "/") + r.URL.Path |
| 176 | target.RawQuery = r.URL.RawQuery |
| 177 | req, err := http.NewRequestWithContext(r.Context(), r.Method, target.String(), bytes.NewReader(body)) |
| 178 | if err != nil { |
| 179 | http.Error(w, "build request: "+err.Error(), http.StatusBadGateway) |
| 180 | return |
| 181 | } |
| 182 | for key, values := range r.Header { |
| 183 | if strings.EqualFold(key, "Host") || strings.EqualFold(key, "Content-Length") { |
| 184 | continue |
| 185 | } |
| 186 | req.Header[key] = values |
| 187 | } |
| 188 | resp, err := m.client.Do(req) |
| 189 | if err != nil { |
| 190 | http.Error(w, "upstream: "+err.Error(), http.StatusBadGateway) |
| 191 | return |
| 192 | } |
| 193 | defer resp.Body.Close() |
| 194 | |
| 195 | maps.Copy(w.Header(), resp.Header) |
| 196 | w.WriteHeader(resp.StatusCode) |
| 197 | if strings.Contains(resp.Header.Get("Content-Type"), "event-stream") { |
| 198 | m.pipeStream(w, resp.Body) |
| 199 | return |
| 200 | } |
| 201 | m.pipeJSON(w, resp.Body) |
| 202 | } |
| 203 | |
| 204 | // requestUsageOptIn asks for a usage block on streamed completions. Without it |
| 205 | // an OpenAI-compatible stream may carry none at all, and a harness that never |
| 206 | // asks would measure as free. |
| 207 | func requestUsageOptIn(body []byte) []byte { |
| 208 | var payload map[string]any |
| 209 | if json.Unmarshal(body, &payload) != nil { |
| 210 | return body |
| 211 | } |
| 212 | if stream, _ := payload["stream"].(bool); !stream { |
| 213 | return body |
| 214 | } |
| 215 | if _, ok := payload["stream_options"]; ok { |
| 216 | return body |
| 217 | } |
| 218 | payload["stream_options"] = map[string]any{"include_usage": true} |
| 219 | out, err := json.Marshal(payload) |
| 220 | if err != nil { |
| 221 | return body |
| 222 | } |
| 223 | return out |
| 224 | } |
| 225 | |
| 226 | func (m *meter) pipeJSON(w http.ResponseWriter, body io.Reader) { |
| 227 | data, err := io.ReadAll(body) |
| 228 | if err != nil { |
| 229 | return |
| 230 | } |
| 231 | _, _ = w.Write(data) |
| 232 | if !m.recordUsage(data) { |
| 233 | m.noteMissingUsage() |
| 234 | } |
| 235 | } |
| 236 | |
| 237 | // pipeStream forwards SSE frames as they arrive — a buffered stream would |
| 238 | // change the harness's observed latency, which other metrics depend on — and |
| 239 | // reads usage out of the frames on the way past. |
| 240 | func (m *meter) pipeStream(w http.ResponseWriter, body io.Reader) { |
| 241 | flusher, _ := w.(http.Flusher) |
| 242 | reader := bufio.NewReader(body) |
| 243 | seen := false |
| 244 | for { |
| 245 | line, err := reader.ReadBytes('\n') |
| 246 | if len(line) > 0 { |
| 247 | _, _ = w.Write(line) |
| 248 | if flusher != nil { |
| 249 | flusher.Flush() |
| 250 | } |
| 251 | if payload, ok := bytes.CutPrefix(bytes.TrimSpace(line), []byte("data:")); ok { |
| 252 | if m.recordUsage(bytes.TrimSpace(payload)) { |
| 253 | seen = true |
| 254 | } |
| 255 | } |
| 256 | } |
| 257 | if err != nil { |
| 258 | break |
| 259 | } |
| 260 | } |
| 261 | if !seen { |
| 262 | m.noteMissingUsage() |
| 263 | } |
| 264 | } |
| 265 | |
| 266 | func (m *meter) noteMissingUsage() { |
| 267 | m.mu.Lock() |
| 268 | m.WithoutUsage++ |
| 269 | m.mu.Unlock() |
| 270 | } |
| 271 | |
| 272 | // recordUsage folds one payload's usage block in, reporting whether it had one. |
| 273 | // Both cache spellings are read: DeepSeek's explicit hit/miss split and the |
| 274 | // OpenAI-standard prompt_tokens_details.cached_tokens. |
| 275 | func (m *meter) recordUsage(payload []byte) bool { |
| 276 | var doc struct { |
| 277 | Usage *struct { |
| 278 | PromptTokens int `json:"prompt_tokens"` |
| 279 | CompletionTokens int `json:"completion_tokens"` |
| 280 | PromptCacheHitTokens int `json:"prompt_cache_hit_tokens"` |
| 281 | PromptCacheMissToken int `json:"prompt_cache_miss_tokens"` |
| 282 | PromptTokensDetails *struct { |
| 283 | CachedTokens int `json:"cached_tokens"` |
| 284 | } `json:"prompt_tokens_details"` |
| 285 | } `json:"usage"` |
| 286 | } |
| 287 | if json.Unmarshal(payload, &doc) != nil || doc.Usage == nil { |
| 288 | return false |
| 289 | } |
| 290 | hit, miss := doc.Usage.PromptCacheHitTokens, doc.Usage.PromptCacheMissToken |
| 291 | if hit == 0 && miss == 0 && doc.Usage.PromptTokensDetails != nil { |
| 292 | hit = doc.Usage.PromptTokensDetails.CachedTokens |
| 293 | miss = doc.Usage.PromptTokens - hit |
| 294 | } |
| 295 | m.mu.Lock() |
| 296 | defer m.mu.Unlock() |
| 297 | m.PromptTokens += doc.Usage.PromptTokens |
| 298 | m.CompletionTokens += doc.Usage.CompletionTokens |
| 299 | m.CacheHitTokens += hit |
| 300 | m.CacheMissTokens += miss |
| 301 | return true |
| 302 | } |
| 303 |