From 95c572e91ef358f1aaffbc6957eea35db140dca0 Mon Sep 17 00:00:00 2001 From: yetone Date: Fri, 25 Sep 2026 11:42:27 +0800 Subject: [PATCH] =?UTF-8?q?gateway:=20a=20Claude=20subscription=20conversa?= =?UTF-8?q?tion's=20next=20turn=20goes=20to=20the=20Claude=20Code=20that?= =?UTF-8?q?=20had=20its=20last,=20told=20only=20what=20the=20user=20said?= =?UTF-8?q?=20since,=20so=20the=20conversation=20Claude=20Code=20keeps=20g?= =?UTF-8?q?oes=20on=20as=20it=20began=20and=20its=20cached=20prompt=20is?= =?UTF-8?q?=20read=20again=20=E2=80=94=20each=20turn=20started=20a=20new?= =?UTF-8?q?=20run=20from=20the=20whole=20history=20in=20one=20message,=20w?= =?UTF-8?q?hich=20shared=20nothing=20with=20the=20cache=20but=20Claude=20C?= =?UTF-8?q?ode's=20own=20prompt,=20so=20every=20turn=20wrote=20some=205.7K?= =?UTF-8?q?=20tokens=20to=20the=20cache=20again=20where=20it=20now=20write?= =?UTF-8?q?s=20about=20100=20(a=20three-turn=20conversation=20went=20from?= =?UTF-8?q?=2037%=20of=20its=20input=20cached=20to=2079%,=20one=20with=20t?= =?UTF-8?q?ool=20calls=20from=2063%=20to=2087%,=20about=2099%=20past=20the?= =?UTF-8?q?=20first=20call);=20a=20run=20waits=20for=20its=20conversation?= =?UTF-8?q?=20up=20to=2020=20minutes,=20six=20at=20most,=20and=20is=20foun?= =?UTF-8?q?d=20again=20by=20what=20the=20conversation=20has=20been=20so=20?= =?UTF-8?q?far,=20else=20a=20new=20one=20starts=20as=20before;=20and=20it?= =?UTF-8?q?=20is=20set=20waiting=20before=20the=20reply's=20last=20event,?= =?UTF-8?q?=20which=20a=20quick=20agent=20answers=20at=20once.=20Anthropic?= =?UTF-8?q?'s=20usage=20leaves=20out=20what=20the=20prompt=20wrote=20to=20?= =?UTF-8?q?the=20cache,=20and=20the=20OpenAI=20and=20Gemini=20counts=20mag?= =?UTF-8?q?pie=20gave=20from=20it=20did=20too,=20so=20Codex=20was=20told?= =?UTF-8?q?=20a=20Claude=20turn=20of=2026K=20tokens=20took=202;=20and=20an?= =?UTF-8?q?=20agent=20that=20hangs=20up=20once=20the=20last=20event=20is?= =?UTF-8?q?=20in=20(Codex=20does,=20before=20the=20ChatGPT=20backend=20clo?= =?UTF-8?q?ses=20its=20end)=20had=20its=20call=20logged=20as=20canceled,?= =?UTF-8?q?=20499,=20though=20it=20was=20answered=20in=20full?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- internal/gateway/chat.go | 2 +- internal/gateway/claude_subscription.go | 267 ++++++++++++++++--- internal/gateway/claude_subscription_test.go | 70 +++++ internal/gateway/fallback.go | 28 ++ internal/gateway/gateway.go | 2 +- internal/gateway/gateway_test.go | 63 +++++ internal/gateway/gemini.go | 6 +- internal/gateway/ir.go | 7 + internal/gateway/responses.go | 2 +- 9 files changed, 402 insertions(+), 45 deletions(-) diff --git a/internal/gateway/chat.go b/internal/gateway/chat.go index f688a039..24b2250b 100644 --- a/internal/gateway/chat.go +++ b/internal/gateway/chat.go @@ -301,7 +301,7 @@ func (u cUsage) usage() Usage { } func (u Usage) chat() map[string]any { - in := u.Input + u.CacheRead + in := u.prompt() return map[string]any{"prompt_tokens": in, "completion_tokens": u.Output, "total_tokens": in + u.Output, "prompt_tokens_details": map[string]any{"cached_tokens": u.CacheRead}, "completion_tokens_details": map[string]any{"reasoning_tokens": u.Reasoning}} diff --git a/internal/gateway/claude_subscription.go b/internal/gateway/claude_subscription.go index 3e543e3b..4dd57340 100644 --- a/internal/gateway/claude_subscription.go +++ b/internal/gateway/claude_subscription.go @@ -11,11 +11,19 @@ package gateway // HTTP request supplies tool_result blocks, which unblock the same Claude Code // turn. This is the protocol-neutral equivalent of pi-claude-bridge's parked // Agent SDK query (https://github.com/elidickinson/pi-claude-bridge, MIT). +// +// A turn that ends leaves its Claude Code running, its input still open, for +// the conversation's next turn: that one is told only what the user said +// since, so the conversation Claude Code keeps goes on as it began and the +// prompt it has cached is read again. A new run would start from the whole +// history in one message, which shares nothing with what was cached but +// Claude Code's own prompt. import ( "bufio" "context" "crypto/rand" + "crypto/sha256" "encoding/hex" "encoding/json" "errors" @@ -36,9 +44,18 @@ type subscriptionBridge struct { mu sync.Mutex runs map[string]*subscriptionRun // callback token → run calls map[string]*subscriptionRun // tool_use id → run + idle map[string]*subscriptionRun // conversation so far (turnKey) → run baseURL string } +// A run left for its conversation's next turn waits idleLongest at most, +// and idleMost of them are kept, the longest waiting let go first: each is +// a Claude Code process. +const ( + idleLongest = 20 * time.Minute + idleMost = 6 +) + type subscriptionRun struct { bridge *subscriptionBridge token string @@ -53,6 +70,13 @@ type subscriptionRun struct { stderr strings.Builder timer *time.Timer + // Claude Code's input, kept open for the next turn; owner is the account + // it runs as, and idleKey and idleAt say it is waiting for one. + stdin io.WriteCloser + owner string + idleKey string + idleAt time.Time + // An agent whose stream does not carry its tool calls in full (Cursor) // learns of them here, as the MCP helper hands each one over, and opens // each answer with a message start of its own. @@ -82,7 +106,7 @@ type bridgeTool struct { } func newSubscriptionBridge() *subscriptionBridge { - return &subscriptionBridge{runs: map[string]*subscriptionRun{}, calls: map[string]*subscriptionRun{}} + return &subscriptionBridge{runs: map[string]*subscriptionRun{}, calls: map[string]*subscriptionRun{}, idle: map[string]*subscriptionRun{}} } func randomToken() string { @@ -118,7 +142,7 @@ func callbackBaseURL() string { return "http://" + host + ":" + port } -func (b *subscriptionBridge) start(ctx context.Context, req *Request, model, oauth string) (*subscriptionRun, <-chan Event, error) { +func (b *subscriptionBridge) start(ctx context.Context, req *Request, model, oauth, owner string) (*subscriptionRun, <-chan Event, error) { binary, err := claudeBinary() if err != nil { return nil, nil, err @@ -170,7 +194,7 @@ func (b *subscriptionBridge) start(ctx context.Context, req *Request, model, oau return nil, nil, err } - run := &subscriptionRun{bridge: b, token: token, model: model, cmd: cmd, tmp: tmp, pending: map[string]chan mcpToolResult{}} + run := &subscriptionRun{bridge: b, token: token, model: model, cmd: cmd, tmp: tmp, pending: map[string]chan mcpToolResult{}, stdin: stdin, owner: owner} // A caller may abandon a turn after receiving tool_use. Do not leave the // parked Claude process and MCP request alive forever. run.timer = time.AfterFunc(30*time.Minute, run.abort) @@ -202,10 +226,131 @@ func (b *subscriptionBridge) start(ctx context.Context, req *Request, model, oau run.abort() return nil, nil, err } - _ = stdin.Close() return run, segment, nil } +// resume hands a conversation's next turn to the run that had its last +// one, when one is waiting: the user's messages since, and nothing else. +func (b *subscriptionBridge) resume(req *Request, owner string) (*subscriptionRun, <-chan Event) { + j := len(req.Messages) - 1 + for j >= 0 && req.Messages[j].Role != "assistant" { + j-- + } + if j < 0 || j == len(req.Messages)-1 { + return nil, nil + } + since := req.Messages[j+1:] + for _, m := range since { + for _, p := range m.Parts { + if p.Kind == ToolResult { + return nil, nil + } + } + } + key := turnKey(owner, req, req.Messages[:j+1]) + b.mu.Lock() + run := b.idle[key] + if run != nil { + delete(b.idle, key) + run.idleKey = "" + } + b.mu.Unlock() + if run == nil { + return nil, nil + } + line, _ := json.Marshal(map[string]any{"type": "user", "message": map[string]any{"role": "user", "content": renderClaudeTurn(since)}}) + run.mu.Lock() + if run.closed { + run.mu.Unlock() + return nil, nil + } + ch := make(chan Event, 64) + run.segment = ch + run.mu.Unlock() + run.timer.Reset(30 * time.Minute) + if _, err := run.stdin.Write(append(line, '\n')); err != nil { + run.abort() + return nil, nil + } + return run, ch +} + +// ended is told how a turn's reply went. A run whose reply was whole and +// asked for nothing more waits for the conversation's next turn; one that +// failed, was cut short or went unheard is let go. +func (r *subscriptionRun) ended(req *Request, said, stop string, ok bool) { + if r.stdin == nil || stop == "tool" && ok { + return // it ends by itself, or waits on its tool calls + } + if !ok || stop != "stop" { + r.abort() + return + } + reply := Message{Role: "assistant", Parts: []Part{{Kind: Text, Text: said}}} + key := turnKey(r.owner, req, append(req.Messages[:len(req.Messages):len(req.Messages)], reply)) + b := r.bridge + var drop []*subscriptionRun + b.mu.Lock() + if old := b.idle[key]; old != nil { + drop = append(drop, old) + } + b.idle[key] = r + r.idleKey, r.idleAt = key, time.Now() + for len(b.idle) > idleMost { + var oldest *subscriptionRun + for _, run := range b.idle { + if oldest == nil || run.idleAt.Before(oldest.idleAt) { + oldest = run + } + } + delete(b.idle, oldest.idleKey) + oldest.idleKey = "" + drop = append(drop, oldest) + } + b.mu.Unlock() + r.timer.Reset(idleLongest) + for _, run := range drop { + run.abort() + } +} + +// turnKey is what a conversation has been so far, for the run that had it +// to be found again: whom it runs as, the model, its settings and tools, +// and the messages' words, tool calls and results. Whitespace, thinking and +// how a reply is split into messages are left out, as clients keep those +// differently. +func turnKey(owner string, req *Request, msgs []Message) string { + h := sha256.New() + tools, _ := json.Marshal(req.Tools) + fmt.Fprintf(h, "%s\x00%s\x00%s\x00%s\x00%s\x00%s", owner, req.Model, req.Effort, req.ToolChoice, req.System, tools) + role := "" + for _, m := range msgs { + var b strings.Builder + for _, p := range m.Parts { + switch p.Kind { + case Text: + b.WriteString(p.Text + " ") + case ToolCall: + fmt.Fprintf(&b, "\x01call %s %s ", p.Name, p.ID) + case ToolResult: + fmt.Fprintf(&b, "\x01result %s %s ", p.CallID, p.Text) + case Image: + fmt.Fprintf(&b, "\x01image %d %s ", len(p.Data), p.URL) + } + } + words := strings.Join(strings.Fields(b.String()), " ") + if words == "" { + continue + } + if m.Role != role { + role = m.Role + fmt.Fprintf(h, "\x00%s:", role) + } + h.Write([]byte(words + " ")) + } + return hex.EncodeToString(h.Sum(nil)) +} + func claudeCLIArgs(model, mcpConfig, effort string) []string { args := []string{ "-p", "--output-format", "stream-json", "--input-format", "stream-json", @@ -399,38 +544,61 @@ func renderClaudePrompt(req *Request) ([]map[string]any, error) { label = "Assistant" } text.WriteString(label + ": ") - for _, p := range m.Parts { - switch p.Kind { - case Text: - text.WriteString(p.Text) - case Thinking: - text.WriteString(p.Text) - case ToolCall: - fmt.Fprintf(&text, "\n[tool call %s id=%s args=%s]", p.Name, p.ID, argsString(p)) - case ToolResult: - fmt.Fprintf(&text, "\n[tool result id=%s%s]\n%s", p.CallID, map[bool]string{true: " error"}[p.IsError], p.Text) - case Image: - if text.Len() > 0 { - blocks = append(blocks, map[string]any{"type": "text", "text": text.String()}) - text.Reset() - } - switch { - case p.Data != "": - blocks = append(blocks, map[string]any{"type": "image", "source": map[string]any{"type": "base64", "media_type": p.MediaType, "data": p.Data}}) - case p.URL != "": - blocks = append(blocks, map[string]any{"type": "image", "source": map[string]any{"type": "url", "url": p.URL}}) - } - } - } + blocks = renderParts(blocks, &text, m.Parts) text.WriteString("\n\n") } + return closeBlocks(blocks, &text), nil +} + +// renderClaudeTurn is the user's messages in a conversation Claude Code +// already has, as it would be told them itself. +func renderClaudeTurn(msgs []Message) []map[string]any { + var blocks []map[string]any + var text strings.Builder + for i, m := range msgs { + if i > 0 { + text.WriteString("\n\n") + } + blocks = renderParts(blocks, &text, m.Parts) + } + return closeBlocks(blocks, &text) +} + +func renderParts(blocks []map[string]any, text *strings.Builder, parts []Part) []map[string]any { + for _, p := range parts { + switch p.Kind { + case Text: + text.WriteString(p.Text) + case Thinking: + text.WriteString(p.Text) + case ToolCall: + fmt.Fprintf(text, "\n[tool call %s id=%s args=%s]", p.Name, p.ID, argsString(p)) + case ToolResult: + fmt.Fprintf(text, "\n[tool result id=%s%s]\n%s", p.CallID, map[bool]string{true: " error"}[p.IsError], p.Text) + case Image: + if text.Len() > 0 { + blocks = append(blocks, map[string]any{"type": "text", "text": text.String()}) + text.Reset() + } + switch { + case p.Data != "": + blocks = append(blocks, map[string]any{"type": "image", "source": map[string]any{"type": "base64", "media_type": p.MediaType, "data": p.Data}}) + case p.URL != "": + blocks = append(blocks, map[string]any{"type": "image", "source": map[string]any{"type": "url", "url": p.URL}}) + } + } + } + return blocks +} + +func closeBlocks(blocks []map[string]any, text *strings.Builder) []map[string]any { if text.Len() > 0 { blocks = append(blocks, map[string]any{"type": "text", "text": text.String()}) } if len(blocks) == 0 { blocks = append(blocks, map[string]any{"type": "text", "text": "[continue]"}) } - return blocks, nil + return blocks } func (b *subscriptionBridge) findRun(req *Request) (*subscriptionRun, []Part) { @@ -629,6 +797,9 @@ func (r *subscriptionRun) abort() { func (b *subscriptionBridge) removeRun(run *subscriptionRun) { b.mu.Lock() delete(b.runs, run.token) + if run.idleKey != "" && b.idle[run.idleKey] == run { + delete(b.idle, run.idleKey) + } for id, owner := range b.calls { if owner == run { delete(b.calls, id) @@ -640,11 +811,15 @@ func (b *subscriptionBridge) removeRun(run *subscriptionRun) { func (s *Server) serveClaudeSubscription(w http.ResponseWriter, r *http.Request, from provider.Protocol, p provider.Provider, model string, body []byte, usage *Usage) (int, string) { start := func(ctx context.Context, req *Request) (*subscriptionRun, <-chan Event, error) { + owner := p.ID + "\x00" + p.Account.User + if run, events := s.subscription.resume(req, owner); run != nil { + return run, events, nil + } token, _, err := p.Account.Token(ctx) if err != nil { return nil, nil, err } - return s.subscription.start(ctx, req, model, token) + return s.subscription.start(ctx, req, model, token, owner) } return s.serveSubscription(w, r, from, "Claude Code", model, body, usage, start) } @@ -695,33 +870,47 @@ func (s *Server) serveSubscription(w http.ResponseWriter, r *http.Request, from return writeError(w, from, code, name+": "+msg), msg } enc := encoder(from, newSSEWriter(w), req.Model) - var failed string - for _, ev := range head { - if ev.Kind == KStart || ev.Kind == KUsage { + var failed, said, stop string + see := func(ev Event) { + switch ev.Kind { + case KError: + failed = ev.Text + case KStart, KUsage: usage.add(ev.Usage) + case KText: + said += ev.Text + case KStop: + stop = ev.Stop } enc.event(ev) } + for _, ev := range head { + see(ev) + } for ev := range events { - if ev.Kind == KError { - failed = ev.Text - } - if ev.Kind == KStart || ev.Kind == KUsage { - usage.add(ev.Usage) - } - enc.event(ev) + see(ev) } + // before the reply's last event, which the agent may answer at once + run.ended(req, said, stop, failed == "" && r.Context().Err() == nil) enc.finish() return 200, failed } var col collector + var said, stop string for ev := range events { + switch ev.Kind { + case KText: + said += ev.Text + case KStop: + stop = ev.Stop + } col.add(ev) } if col.err != "" && len(col.res.Parts) == 0 { run.abort() return writeError(w, from, 502, name+": "+col.err), col.err } + run.ended(req, said, stop, col.err == "" && r.Context().Err() == nil) res := col.finish() usage.add(res.Usage) w.Header().Set("Content-Type", "application/json") diff --git a/internal/gateway/claude_subscription_test.go b/internal/gateway/claude_subscription_test.go index d701bfde..95d5bcea 100644 --- a/internal/gateway/claude_subscription_test.go +++ b/internal/gateway/claude_subscription_test.go @@ -2,8 +2,14 @@ package gateway import ( "encoding/json" + "net/http/httptest" + "os" + "path/filepath" + "strconv" "strings" "testing" + + "github.com/yetone/magpie/internal/provider" ) func TestClaudeSubscriptionPromptKeepsForeignHarnessOutOfSystem(t *testing.T) { @@ -55,3 +61,67 @@ func TestCleanClaudeEnvRemovesGatewayOverrides(t *testing.T) { } } } + +// fakeClaude answers each line it is given with the process it runs in and +// how many turns that process has had, as Claude Code's stream-json does. +func fakeClaude(t *testing.T) { + t.Helper() + dir := t.TempDir() + script := `#!/bin/sh +n=0 +while read -r line; do + n=$((n+1)) + echo '{"type":"stream_event","event":{"type":"message_start","message":{"id":"m","model":"claude-sonnet-5","usage":{"input_tokens":1}}}}' + echo '{"type":"stream_event","event":{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"pid '$$' turn '$n'"}}}' + echo '{"type":"stream_event","event":{"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":3}}}' + echo '{"type":"stream_event","event":{"type":"message_stop"}}' + echo '{"type":"result","subtype":"success","is_error":false,"result":""}' +done +` + os.WriteFile(filepath.Join(dir, "claude"), []byte(script), 0o755) + t.Setenv("PATH", dir+string(os.PathListSeparator)+os.Getenv("PATH")) +} + +// A conversation's next turn goes to the Claude Code that had its last, +// told only what the user said since; another conversation, or the same one +// with its reply changed, gets a Claude Code of its own. +func TestClaudeRunKeptForTheNextTurn(t *testing.T) { + fakeClaude(t) + s := New() + p := provider.Provider{ID: "claude", Account: &provider.Account{Agent: "claude", User: "u"}} + ask := func(msgs string) string { + t.Helper() + body := `{"model":"claude-sonnet-5","max_tokens":100,"system":"be brief","messages":` + msgs + `}` + rec := httptest.NewRecorder() + var u Usage + if code, msg := s.serveClaudeSubscription(rec, httptest.NewRequest("POST", "/v1/messages", strings.NewReader(body)), provider.Anthropic, p, "claude-sonnet-5", []byte(body), &u); code != 200 { + t.Fatalf("%d %s", code, msg) + } + var res struct { + Content []struct{ Text string } `json:"content"` + } + json.Unmarshal(rec.Body.Bytes(), &res) + if len(res.Content) == 0 { + t.Fatalf("no answer: %s", rec.Body) + } + return res.Content[0].Text + } + msg := func(role, text string) string { return `{"role":"` + role + `","content":` + strconv.Quote(text) + `}` } + + first := ask(`[` + msg("user", "hi") + `]`) + pid, _, _ := strings.Cut(strings.TrimPrefix(first, "pid "), " ") + if !strings.HasSuffix(first, "turn 1") { + t.Fatalf("first: %q", first) + } + second := ask(`[` + msg("user", "hi") + `,` + msg("assistant", " "+first+"\n") + `,` + msg("user", "and?") + `]`) + if second != "pid "+pid+" turn 2" { + t.Fatalf("second: %q, first %q", second, first) + } + if other := ask(`[` + msg("user", "bye") + `,` + msg("assistant", second) + `,` + msg("user", "and?") + `]`); strings.Contains(other, "pid "+pid) { + t.Fatalf("another conversation: %q", other) + } + third := ask(`[` + msg("user", "hi") + `,` + msg("assistant", first) + `,` + msg("user", "and?") + `,` + msg("assistant", "edited") + `,` + msg("user", "so?") + `]`) + if strings.Contains(third, "pid "+pid) { + t.Fatalf("an edited reply: %q", third) + } +} diff --git a/internal/gateway/fallback.go b/internal/gateway/fallback.go index 79a6c421..aba6e66c 100644 --- a/internal/gateway/fallback.go +++ b/internal/gateway/fallback.go @@ -383,6 +383,9 @@ type holdWriter struct { scanned int // how much of held has been read as events failure int // the status the stream's error stands for failMsg string + + ended bool // the stream's last event was written: the reply is whole + tail []byte // the end of the last write, for a marker split across two } func newHoldWriter(w http.ResponseWriter, hold bool) *holdWriter { @@ -419,6 +422,7 @@ func (h *holdWriter) Write(b []byte) (int, error) { if h.status == 0 { h.WriteHeader(http.StatusOK) } + h.see(b) if h.passing { return h.w.Write(b) } @@ -429,6 +433,30 @@ func (h *holdWriter) Write(b []byte) (int, error) { return n, err } +// streamEnds are how each protocol's stream says it is over, as they can +// only appear outside a quoted string: text that says so is escaped. +var streamEnds = [][]byte{ + []byte(`"type":"response.completed"`), []byte(`"type":"response.incomplete"`), // Responses + []byte(`"type":"message_stop"`), []byte("event: message_stop"), // Anthropic + []byte("data: [DONE]"), // Chat Completions +} + +// see notes a stream's last event going by. An agent may hang up as soon as +// it has that — Codex does — and a reply it had whole is not canceled. +func (h *holdWriter) see(b []byte) { + if h.ended { + return + } + buf := append(h.tail, b...) + for _, m := range streamEnds { + if bytes.Contains(buf, m) { + h.ended = true + return + } + } + h.tail = append(h.tail[:0], buf[max(len(buf)-32, 0):]...) +} + // holdLongest is the longest a stream is held waiting for its first // content, and holdMost the most of it. const ( diff --git a/internal/gateway/gateway.go b/internal/gateway/gateway.go index ad45cfc4..2d2dd91d 100644 --- a/internal/gateway/gateway.go +++ b/internal/gateway/gateway.go @@ -414,7 +414,7 @@ func (s *Server) serve(w http.ResponseWriter, r *http.Request, from provider.Pro call.Status, call.Error = hw.failure, c.p.Name+": "+hw.failMsg } try := Try{ID: c.rest, Start: began, Done: true, Status: call.Status, Millis: time.Since(began).Milliseconds(), Error: call.Error} - if r.Context().Err() != nil { + if r.Context().Err() != nil && !hw.ended { // the agent went away: nobody failed, and nobody else is asked call.Status, call.Error = 499, "the agent canceled the request" try.Status, try.Error, try.Fail = call.Status, call.Error, failCanceled diff --git a/internal/gateway/gateway_test.go b/internal/gateway/gateway_test.go index f3878121..423958f5 100644 --- a/internal/gateway/gateway_test.go +++ b/internal/gateway/gateway_test.go @@ -829,3 +829,66 @@ func TestModelListSaysWhichAPI(t *testing.T) { } } } + +// An agent that hangs up once the last event is in — Codex does, while the +// ChatGPT backend is still to close its end — had its answer: the call is +// served, not canceled. One that hangs up before is canceled. +func TestHangUpAfterTheLastEvent(t *testing.T) { + for _, whole := range []bool{true, false} { + last := `data: {"type":"response.output_text.delta","delta":"hi"}` + if whole { + last = `data: {"type":"response.completed","response":{"usage":{"input_tokens":5,"output_tokens":1}}}` + } + up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + io.WriteString(w, sse(`data: {"type":"response.created","response":{}}`, last)) + w.(http.Flusher).Flush() + <-r.Context().Done() + })) + t.Setenv("XDG_CONFIG_HOME", t.TempDir()) + t.Setenv("XDG_CACHE_HOME", t.TempDir()) + provider.Save(provider.Provider{ID: "fake", Name: "Fake", Key: "k", Models: []string{"m1"}, Responses: up.URL + "/v1"}) + s := New() + gw := httptest.NewServer(s.Handler()) + ctx, cancel := context.WithCancel(context.Background()) + req, _ := http.NewRequestWithContext(ctx, "POST", gw.URL+"/v1/responses", strings.NewReader(`{"model":"fake/m1","stream":true,"input":"hi"}`)) + res, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatal(err) + } + buf := make([]byte, 4096) + for got := ""; !strings.Contains(got, strings.TrimPrefix(last, "data: ")); { + n, err := res.Body.Read(buf) + if err != nil { + t.Fatal(err) + } + got += string(buf[:n]) + } + cancel() + var calls []Call + for deadline := time.Now().Add(5 * time.Second); len(calls) == 0 && time.Now().Before(deadline); { + time.Sleep(10 * time.Millisecond) + calls = s.Recent() + } + if want := map[bool]int{true: 200, false: 499}[whole]; len(calls) != 1 || calls[0].Status != want { + t.Fatalf("whole %v: %+v", whole, calls) + } + gw.Close() + up.Close() + } +} + +// Anthropic's input count leaves out what the prompt read from its cache +// and wrote to it; OpenAI's and Gemini's count both. +func TestUsagePromptCountsCacheWrites(t *testing.T) { + u := Usage{Input: 2, Output: 30, CacheRead: 6176, CacheWrite: 20563} + if r := u.responses(); r["input_tokens"] != 26741 || r["total_tokens"] != 26771 { + t.Errorf("responses: %v", r) + } + if c := u.chat(); c["prompt_tokens"] != 26741 { + t.Errorf("chat: %v", c) + } + if g := u.gemini(); g["promptTokenCount"] != 26741 || g["totalTokenCount"] != 26771 { + t.Errorf("gemini: %v", g) + } +} diff --git a/internal/gateway/gemini.go b/internal/gateway/gemini.go index 1a4e9eeb..c554fd38 100644 --- a/internal/gateway/gemini.go +++ b/internal/gateway/gemini.go @@ -325,10 +325,10 @@ func stopToGemini(s string) string { } // gemini is a Usage as usageMetadata: the prompt count includes cached -// tokens, and thoughts are counted apart from the answer. +// tokens (see prompt), and thoughts are counted apart from the answer. func (u Usage) gemini() map[string]any { - out := map[string]any{"promptTokenCount": u.Input + u.CacheRead, "candidatesTokenCount": max(u.Output-u.Reasoning, 0), - "totalTokenCount": u.Input + u.CacheRead + u.Output} + out := map[string]any{"promptTokenCount": u.prompt(), "candidatesTokenCount": max(u.Output-u.Reasoning, 0), + "totalTokenCount": u.prompt() + u.Output} if u.CacheRead > 0 { out["cachedContentTokenCount"] = u.CacheRead } diff --git a/internal/gateway/ir.go b/internal/gateway/ir.go index 0d7dea0c..ab4e05e0 100644 --- a/internal/gateway/ir.go +++ b/internal/gateway/ir.go @@ -113,6 +113,13 @@ type Usage struct { Reasoning int `json:"reasoning"` } +// prompt is every token the prompt came to, as OpenAI's and Gemini's +// counts have it: Anthropic's leaves out what was read from its cache and +// what was written to it. +func (u Usage) prompt() int { + return u.Input + u.CacheRead + u.CacheWrite +} + func (u *Usage) add(v Usage) { if v.Input > 0 { u.Input = v.Input diff --git a/internal/gateway/responses.go b/internal/gateway/responses.go index e249664c..636ef21d 100644 --- a/internal/gateway/responses.go +++ b/internal/gateway/responses.go @@ -271,7 +271,7 @@ func (u rUsage) usage() Usage { } func (u Usage) responses() map[string]any { - in := u.Input + u.CacheRead + in := u.prompt() return map[string]any{"input_tokens": in, "output_tokens": u.Output, "total_tokens": in + u.Output, "input_tokens_details": map[string]any{"cached_tokens": u.CacheRead}, "output_tokens_details": map[string]any{"reasoning_tokens": u.Reasoning}}