gateway: a Claude subscription conversation's next turn goes to the Claude Code that had its last, told only what the user said since, so the conversation Claude Code keeps goes on as it began and its cached prompt is read again — each turn started a new run from the whole history in one message, which shared nothing with the cache but Claude Code's own prompt, so every turn wrote some 5.7K tokens to the cache again where it now writes about 100 (a three-turn conversation went from 37% of its input cached to 79%, one with tool calls from 63% to 87%, about 99% past the first call); a run waits for its conversation up to 20 minutes, six at most, and is found again by what the conversation has been so far, else a new one starts as before; and it is set waiting before the reply's last event, which a quick agent answers at once. Anthropic's usage leaves out what the prompt wrote to the cache, and the OpenAI and Gemini counts magpie gave from it did too, so Codex was told a Claude turn of 26K tokens took 2; and an agent that hangs up once the last event is in (Codex does, before the ChatGPT backend closes its end) had its call logged as canceled, 499, though it was answered in full
This commit is contained in:
9 files changed
+402
-45
No files matched your search
@@ -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}}
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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 (
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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}}
|
||||
|
||||
Reference in new issue
Block a user