From bd456fa51ed5e46e59441a955cd22f2ecbf27065 Mon Sep 17 00:00:00 2001 From: Haogao Gu Date: Sun, 4 Oct 2026 23:52:23 -0400 Subject: [PATCH] fix: keep Claude continuations and routing stats bound to accounts (#854) * fix(gateway): bind Claude tool-result continuations to their account * fix(gui): keep routing statistics on stable account identities --- internal/gateway/automode_test.go | 2 +- .../gateway/claude_account_handoff_test.go | 89 +++++++++++++++++ internal/gateway/claude_subscription.go | 29 +++--- .../gateway/claude_subscription_remap_test.go | 2 +- .../claude_subscription_rewrite_test.go | 2 +- internal/gateway/claude_subscription_test.go | 2 +- .../claude_subscription_toolsearch_test.go | 2 +- internal/gui/assets/routing.js | 14 ++- .../tests/routing-account-identity.test.cjs | 96 +++++++++++++++++++ 9 files changed, 219 insertions(+), 19 deletions(-) create mode 100644 internal/gateway/claude_account_handoff_test.go create mode 100644 internal/gui/tests/routing-account-identity.test.cjs diff --git a/internal/gateway/automode_test.go b/internal/gateway/automode_test.go index 30338cc1..a37e4023 100644 --- a/internal/gateway/automode_test.go +++ b/internal/gateway/automode_test.go @@ -494,7 +494,7 @@ func TestClaudeBridgeSafeguardConcurrentHandoff(t *testing.T) { done := make(chan int, 1) go func() { rec := httptest.NewRecorder() - code, _ := s.serveSubscription(rec, httptest.NewRequest("POST", "/v1/messages", nil), provider.Anthropic, "Claude Code", req.Model, body, &Usage{}, + code, _ := s.serveSubscription(rec, httptest.NewRequest("POST", "/v1/messages", nil), provider.Anthropic, "Claude Code", req.Model, "", body, &Usage{}, func(context.Context, *Request) (*subscriptionRun, <-chan Event, error) { return nil, nil, fmt.Errorf("must not start another run") }) diff --git a/internal/gateway/claude_account_handoff_test.go b/internal/gateway/claude_account_handoff_test.go new file mode 100644 index 00000000..c81c5ffe --- /dev/null +++ b/internal/gateway/claude_account_handoff_test.go @@ -0,0 +1,89 @@ +package gateway + +import ( + "encoding/json" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/yetone/magpie/internal/provider" +) + +// Tool results identify a waiting process, but the router may have picked +// another account since that process made its call. The selected account +// must answer, including when the same account moved out of the agent's home. +func TestClaudeToolResultsKeepTheSelectedAccount(t *testing.T) { + for _, ids := range []string{"exact", "remapped"} { + for _, mode := range []string{"same owner", "another account", "another home"} { + t.Run(ids+"/"+mode, func(t *testing.T) { + claudeMadeFirst(t, false) + h := newRemapHarness(t) + received := make(chan bool, 1) + h.first = func(run *subscriptionRun, req *Request) { + run.owner = "claude\x00a@example.com\x00" + ownHome + waiters := h.calls(run, true, shownCall{"call-a", "read", `{}`}) + go func() { + _, ok := <-waiters[0] + if ok { + h.says(run, "kept run") + } + received <- ok + }() + } + h.ask(remapUser("read")) + p := provider.Provider{ID: "claude", Account: &provider.Account{Agent: "claude", User: "a@example.com"}} + want := "kept run" + if mode == "another account" { + for _, saved := range p.AlsoOn() { + if saved.Account.User == "b@example.com" { + p = saved + break + } + } + if p.Account.User != "b@example.com" { + t.Fatal("saved B is missing") + } + want = "tok-b turn 1" + } else if mode == "another home" { + if err := provider.SwitchLogin("claude", "b@example.com"); err != nil { + t.Fatal(err) + } + b := provider.Provider{ID: "claude", Account: &provider.Account{Agent: "claude", User: "b@example.com"}} + for _, saved := range b.AlsoOn() { + if saved.Account.User == "a@example.com" { + p = saved + break + } + } + if p.Account.AgentsOwn() { + t.Fatal("A did not move to its saved home") + } + want = "tok-a turn 1" + } + callID := "call-a" + if ids == "remapped" { + callID = "client-a" + } + body := `{"model":"claude-sonnet-5","max_tokens":100,"tools":[{"name":"read","input_schema":{"type":"object"}}],"messages":[` + remapUser("read") + `,` + remapCalls(callID+"={}") + `,` + remapResults(callID+"=result") + `]}` + rec := httptest.NewRecorder() + code, msg := h.s.serveClaudeSubscription(rec, httptest.NewRequest("POST", "/v1/messages", strings.NewReader(body)), provider.Anthropic, p, "claude-sonnet-5", []byte(body), &Usage{}) + if code != 200 { + t.Fatalf("%d %s", code, msg) + } + var reply struct{ Content []struct{ Text string } } + if err := json.Unmarshal(rec.Body.Bytes(), &reply); err != nil || len(reply.Content) != 1 || reply.Content[0].Text != want { + t.Fatalf("selected account replied %s; want %q (decode: %v)", rec.Body, want, err) + } + select { + case got := <-received: + if got != (mode == "same owner") { + t.Fatalf("old process received the result: %v, mode %s", got, mode) + } + case <-time.After(2 * time.Second): + t.Fatal("old process was neither resumed nor released") + } + }) + } + } +} diff --git a/internal/gateway/claude_subscription.go b/internal/gateway/claude_subscription.go index 3947542f..3f8aa6aa 100644 --- a/internal/gateway/claude_subscription.go +++ b/internal/gateway/claude_subscription.go @@ -2000,6 +2000,7 @@ const ( ambiguousCalls = "ambiguous" runExpired = "process expired" noRunWaiting = "no run waiting" + accountChanged = "account changed" ) // match is findRun, and how the results found their run (byExactID, @@ -2513,18 +2514,17 @@ func ownerAccount(owner string) (user string, own bool) { } 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) { - ctx = p.Via(ctx) // the account's own proxy, its CLI run's too - owner := p.ID + "\x00" + p.Account.User + owner := p.ID + if p.Account != nil { + owner += "\x00" + p.Account.User if p.Account.AgentsOwn() { - // Claude Code's own sign-in, which a switch moves to another - // account: a run kept from before goes on as that one (it - // reads its keychain again), so the account, once saved and - // run in a config directory of its own, never resumes it - // (nil_1024: made first, a saved account's turns went on as - // the spent one in its Claude Code) + // A process in the agent's home reads that home's current sign-in. + // A saved account must not resume it after the sign-in moves. owner += "\x00" + ownHome } + } + start := func(ctx context.Context, req *Request) (*subscriptionRun, <-chan Event, error) { + ctx = p.Via(ctx) // the account's own proxy, its CLI run's too if run, events := s.subscription.resume(req, owner); run != nil { return run, events, nil } @@ -2546,13 +2546,13 @@ func (s *Server) serveClaudeSubscription(w http.ResponseWriter, r *http.Request, } return s.subscription.start(ctx, req, model, dir, owner, from) } - return s.serveSubscription(w, r, from, "Claude Code", model, body, usage, start) + return s.serveSubscription(w, r, from, "Claude Code", model, owner, body, usage, start) } // serveSubscription answers a request through an agent's own binary: a new // turn starts it, a request carrying tool results resumes the turn waiting // on them. -func (s *Server) serveSubscription(w http.ResponseWriter, r *http.Request, from provider.Protocol, name, model string, body []byte, usage *Usage, +func (s *Server) serveSubscription(w http.ResponseWriter, r *http.Request, from provider.Protocol, name, model, owner string, body []byte, usage *Usage, start func(ctx context.Context, req *Request) (*subscriptionRun, <-chan Event, error)) (int, string) { req, err := parse(from, body) if err != nil { @@ -2582,6 +2582,13 @@ func (s *Server) serveSubscription(w http.ResponseWriter, r *http.Request, from msg := "the agent's turn is already being resumed" return writeError(w, from, http.StatusConflict, msg), msg } + // Tool-call IDs find the process that made them, independently of the + // account routing selected. Continuing a different owner's process would + // spend its quota and report its reply or limit against the chosen account. + if run != nil && run.owner != owner { + run.abort() + run, how = nil, accountChanged + } // the client rewrote the conversation since the run's last reply, as Pi // does compacting it mid-turn: the run's agent holds the one from before, // would answer from it, and tell the client a context as large as ever, diff --git a/internal/gateway/claude_subscription_remap_test.go b/internal/gateway/claude_subscription_remap_test.go index 894a8e3b..253c71b9 100644 --- a/internal/gateway/claude_subscription_remap_test.go +++ b/internal/gateway/claude_subscription_remap_test.go @@ -97,7 +97,7 @@ func (h *remapHarness) ask(msgs ...string) []string { body := `{"model":"claude-sonnet-5","max_tokens":100,"tools":[{"name":"read","input_schema":{"type":"object"}}],"messages":[` + strings.Join(msgs, ",") + `]}` rec := httptest.NewRecorder() var u Usage - if code, msg := h.s.serveSubscription(rec, httptest.NewRequest("POST", "/v1/messages", strings.NewReader(body)), provider.Anthropic, "Claude Code", "claude-sonnet-5", []byte(body), &u, h.start); code != 200 { + if code, msg := h.s.serveSubscription(rec, httptest.NewRequest("POST", "/v1/messages", strings.NewReader(body)), provider.Anthropic, "Claude Code", "claude-sonnet-5", "", []byte(body), &u, h.start); code != 200 { h.t.Fatalf("%d %s", code, msg) } var res struct { diff --git a/internal/gateway/claude_subscription_rewrite_test.go b/internal/gateway/claude_subscription_rewrite_test.go index 907136aa..0232536f 100644 --- a/internal/gateway/claude_subscription_rewrite_test.go +++ b/internal/gateway/claude_subscription_rewrite_test.go @@ -66,7 +66,7 @@ func TestRunLetGoWhenTheClientRewroteTheConversation(t *testing.T) { body := `{"model":"claude-sonnet-5","max_tokens":100,"tools":[{"name":"read","input_schema":{"type":"object"}}],"messages":` + msgs + `}` rec := httptest.NewRecorder() var u Usage - if code, msg := s.serveSubscription(rec, httptest.NewRequest("POST", "/v1/messages", strings.NewReader(body)), provider.Anthropic, "Claude Code", "claude-sonnet-5", []byte(body), &u, start); code != 200 { + if code, msg := s.serveSubscription(rec, httptest.NewRequest("POST", "/v1/messages", strings.NewReader(body)), provider.Anthropic, "Claude Code", "claude-sonnet-5", "", []byte(body), &u, start); code != 200 { t.Fatalf("%d %s", code, msg) } var res struct { diff --git a/internal/gateway/claude_subscription_test.go b/internal/gateway/claude_subscription_test.go index 4830c2f1..dd30d9ed 100644 --- a/internal/gateway/claude_subscription_test.go +++ b/internal/gateway/claude_subscription_test.go @@ -205,7 +205,7 @@ func TestSubscriptionStreamErrorIsTheEnd(t *testing.T) { body := `{"model":"m","max_tokens":10,"stream":true,"messages":[{"role":"user","content":"hi"}],"input":"hi"}` rec := httptest.NewRecorder() var u Usage - code, failed := s.serveSubscription(rec, httptest.NewRequest("POST", "/", strings.NewReader(body)), from, "Agent", "m", []byte(body), &u, start) + code, failed := s.serveSubscription(rec, httptest.NewRequest("POST", "/", strings.NewReader(body)), from, "Agent", "m", "", []byte(body), &u, start) out := rec.Body.String() if code != 200 || failed != "it died" || !strings.Contains(out, "it died") { t.Fatalf("%s: %d %q\n%s", from, code, failed, out) diff --git a/internal/gateway/claude_subscription_toolsearch_test.go b/internal/gateway/claude_subscription_toolsearch_test.go index 78e8a0c7..f7a232af 100644 --- a/internal/gateway/claude_subscription_toolsearch_test.go +++ b/internal/gateway/claude_subscription_toolsearch_test.go @@ -70,7 +70,7 @@ func TestToolSearchLoadKeepsTheRun(t *testing.T) { defer close(done) w := httptest.NewRecorder() r := httptest.NewRequest("POST", "/v1/messages", strings.NewReader(body)) - s.serveSubscription(w, r, provider.Anthropic, "Claude Code", "claude-sonnet-5", []byte(body), &Usage{}, start) + s.serveSubscription(w, r, provider.Anthropic, "Claude Code", "claude-sonnet-5", "", []byte(body), &Usage{}, start) }() var got mcpToolResult diff --git a/internal/gui/assets/routing.js b/internal/gui/assets/routing.js index 590b2b23..821290e4 100644 --- a/internal/gui/assets/routing.js +++ b/internal/gui/assets/routing.js @@ -1626,16 +1626,24 @@ // tried, answered and failed in them, and how the latest found it function renderActs(rs) { const by = new Map(); + // The agent's current sign-in uses its provider's ID; after switching, + // that same account uses provider@user and the base ID names another. + // Group accounts by their identity, resolving each try in its own route. + const keyOf = (w) => w.kind === "account" && w.who + ? JSON.stringify([w.provider, w.who.trim().toLowerCase()]) : w.id; for (const r of [...rs].reverse()) { // oldest first, so the latest wins + const candidates = new Map(); r.order.forEach((w, i) => { - const a = by.get(w.id) || { w, tried: 0, ok: 0, fails: {}, last: 0, rest: null, restAt: 0, seen: 0, pos: 0, models: new Set() }; + const key = keyOf(w); + const a = by.get(key) || { w, tried: 0, ok: 0, fails: {}, last: 0, rest: null, restAt: 0, seen: 0, pos: 0, models: new Set() }; a.w = w; a.seen++; a.pos = i; a.at = r.time; // a later request found it resting, or not if (at(r.time) >= a.restAt) { a.rest = w.rest || null; a.restAt = at(r.time); } - by.set(w.id, a); + by.set(key, a); + candidates.set(w.id, a); }); for (const tr of r.tries) { - const a = by.get(tr.id); + const a = candidates.get(tr.id); if (!a || !tr.done || tr.fail === "canceled") continue; // the agent's doing, not its a.tried++; const end = at(tr.start) + (tr.ms || 0); diff --git a/internal/gui/tests/routing-account-identity.test.cjs b/internal/gui/tests/routing-account-identity.test.cjs new file mode 100644 index 00000000..907dbce4 --- /dev/null +++ b/internal/gui/tests/routing-account-identity.test.cjs @@ -0,0 +1,96 @@ +const assert = require("node:assert/strict"); +const fs = require("node:fs/promises"); +const path = require("node:path"); +const { test } = require("node:test"); +const { chromium, webkit } = require("playwright"); + +// The agent's current login changes its trace ID from provider@user to +// provider. Old requests keep both IDs, and provider used to name another +// account. Counts and cooldowns must follow each request's account identity. +function switchedRoutes(now, resting) { + const reset = new Date(now.getTime() + 4 * 3600e3).toISOString(); + const university = { provider: "claude", name: "Claude Code", kind: "account", icon: "claude-color", who: "uni@example.test · University", plan: "enterprise", model: "claude-opus-5-5", known: true, used: 100, renews: [reset] }; + const team = { ...university, who: "seat@example.test · Research Group", plan: "team", used: 1 }; + const oldTeamID = "claude@" + team.who; + const old = [{ ...university, id: "claude" }, { ...team, id: oldTeamID }]; + const rest = { why: "quota", until: new Date(now.getTime() + 15 * 60e3).toISOString(), key: oldTeamID.toLowerCase() }; + return Array.from({ length: 60 }, (_, i) => { + const time = new Date(now.getTime() - (60 - i) * 1000).toISOString(); + const done = i < 58; + const failed = i === 57; + const id = i < 40 ? "claude" : oldTeamID; + const order = i < 58 ? old : [{ ...university, id: "claude@" + university.who }, { ...team, id: "claude", ...(resting ? { rest } : {}) }]; + return { id: i + 1, seq: i + 1, time, agent: "codex", model: "claude/claude-opus-5-5", provider: "claude", order, done, + status: failed ? 429 : 200, ms: 20, tries: done ? [{ id, model: "claude-opus-5-5", start: time, done: true, status: failed ? 429 : 200, ms: 20, ...(failed ? { fail: "quota", rest } : {}) }] : [] }; + }).reverse(); +} + +function serve(lang, getRoutes, now) { + const assets = path.resolve(__dirname, "../assets"); + return async (route) => { + const url = new URL(route.request().url()); + const json = (data) => route.fulfill({ json: data }); + if (url.pathname === "/boot.js") return route.fulfill({ contentType: "text/javascript", body: `window.bootPrefs={lang:"${lang}",theme:"light",web:true};` }); + if (url.pathname === "/wails/runtime.js") return route.fulfill({ contentType: "text/javascript", body: "export const Window={};" }); + if (url.pathname === "/api/state") return json({ agents: [], profiles: [], settings: { lang, theme: "light" } }); + if (url.pathname === "/api/gateway/trace") { + if (url.searchParams.has("wait")) return; + return json({ mine: true, now: now.toISOString(), seq: 60, totals: { requests: 60, rerouted: 0, errors: 1 }, routes: getRoutes() }); + } + if (url.pathname === "/api/gateway/history") return json({ cut: false, days: [], routes: [] }); + if (url.pathname === "/api/groups") return json({ groups: [] }); + if (url.pathname === "/api/providers") return json({ providers: [], presets: [], excluded: [], gateway: { running: true, window: true } }); + if (url.pathname.startsWith("/api/")) return json({}); + const file = path.join(assets, url.pathname === "/" ? "index.html" : url.pathname); + const contentType = { ".html": "text/html", ".js": "text/javascript", ".css": "text/css", ".svg": "image/svg+xml" }[path.extname(file)]; + await route.fulfill({ body: await fs.readFile(file).catch(() => ""), contentType }); + }; +} + +for (const engine of process.env.BROWSER ? [process.env.BROWSER] : ["chromium", "webkit"]) { + for (const lang of ["en", "zh"]) { + test(`${engine} ${lang}: account statistics survive a current-login switch`, async (t) => { + const browser = await (engine === "webkit" ? webkit.launch() : chromium.launch({ channel: "chromium" })); + t.after(() => browser.close()); + const page = await browser.newPage({ viewport: { width: 1100, height: 1000 } }); + page.setDefaultTimeout(5000); + const errors = []; + page.on("pageerror", (e) => errors.push(e.message)); + const now = new Date(); + let routes = switchedRoutes(now, true); + await page.route("**/*", serve(lang, () => routes, now)); + await page.goto("http://magpie.test/?view=routing"); + await page.locator(".rt-act").first().waitFor(); + assert.equal(await page.locator(".rt-act").count(), 2, "two logins stay two rows when their current-login IDs exchange"); + const uni = page.locator(".rt-act", { hasText: "uni@example.test" }); + const team = page.locator(".rt-act", { hasText: "seat@example.test" }); + const counts = async (row) => (await row.locator(".tally > span").allInnerTexts()).slice(0, 2).map((v) => Number(v.match(/\d+/)[0])); + assert.deepEqual(await counts(uni), [40, 40], "the previous current login keeps its own replies"); + assert.deepEqual(await counts(team), [18, 17], "the new current login includes its saved-login requests"); + assert.equal(await uni.locator(".st.rest").count(), 0, "another account's failure does not transfer to the reused base ID"); + assert.equal(await team.locator(".st.rest").count(), 1, "a cooldown still present in the newest snapshot follows its account"); + assert.equal(await team.locator(".tally .bad").count(), 1); + if (process.env.ARTIFACT_DIR) { + await fs.mkdir(process.env.ARTIFACT_DIR, { recursive: true }); + await page.locator(".rt-acts").screenshot({ path: path.join(process.env.ARTIFACT_DIR, `${engine}-${lang}-account-identity.png`) }); + } + + routes = switchedRoutes(now, false); + await page.reload(); + await page.locator(".rt-act").first().waitFor(); + assert.equal(await team.locator(".st.rest").count(), 0, "a newer snapshot that cleared the cooldown wins over the old alias"); + assert.deepEqual(await counts(team), [18, 17], "clearing rest does not erase the historical failure or replies"); + + // Equal emails on different subscriptions or organizations are still + // separate identities; an API key's name is not an account identity. + routes = [{ ...routes[0], id: 61, seq: 61, order: [routes[0].order[1], + { ...routes[0].order[1], id: "codex", provider: "codex", name: "Codex" }, + { ...routes[0].order[1], id: "claude@other-org", who: "seat@example.test · Other Group" }, + { ...routes[0].order[1], id: "claude#key1", kind: "key" }] }]; + await page.reload(); + await page.locator(".rt-act").first().waitFor(); + assert.equal(await page.locator(".rt-act").count(), 4, "provider, organization, and API-key boundaries are preserved"); + assert.deepEqual(errors, []); + }); + } +}