diff --git a/Makefile b/Makefile index 5cbb836..51bd7a0 100644 --- a/Makefile +++ b/Makefile @@ -42,8 +42,11 @@ build: sync-fingerprint @go build $(LDFLAGS) -o bin/aikey-proxy ./cmd/aikey-proxy @cp $(CONFIG) bin/$(CONFIG) +# ./cmd/aikey-proxy added 2026-07-08: the egress system-proxy-switch +# integration tests (egress_integration_test.go) live in package main because +# they drive the REAL buildTransport — internal/... alone would skip them. test: - go test -race -v ./internal/... + go test -race -v ./internal/... ./cmd/aikey-proxy/ # Chaos experiments (缺口7/8) — build-tagged so they stay OUT of the normal # `test` suite. They drive the real newStreamDrainer / http.Server code paths diff --git a/aikey-proxy.yaml.example b/aikey-proxy.yaml.example index 7ad38a9..13bf5ad 100644 --- a/aikey-proxy.yaml.example +++ b/aikey-proxy.yaml.example @@ -31,6 +31,13 @@ listen: vault: path: "~/.aikey/data/vault.db" +# Local AiKey console base (aikey-local-server web, 8090 on a default personal +# install). Used to assemble the member-login URL shown inside claude/codex +# when a shared pool account requires sign-in +# (OAUTH_GROUP_MEMBER_LOGIN_REQUIRED → /user/team-oauth). +# Empty ⇒ the error message falls back to URL-less wording. +console_url: "http://127.0.0.1:8090" + # Virtual keys are delivered via vault (managed_virtual_keys_cache or # personal_route_token columns). Add keys with `aikey add # --provider `, then `aikey list` / `aikey route` for the bearer diff --git a/cmd/aikey-proxy/egress_integration_test.go b/cmd/aikey-proxy/egress_integration_test.go new file mode 100644 index 0000000..e9e05dd --- /dev/null +++ b/cmd/aikey-proxy/egress_integration_test.go @@ -0,0 +1,261 @@ +package main + +// Egress integration test for the 2026-07-08 system-proxy auto-refresh +// requirement: a RUNNING transport must egress through the NEW system proxy +// on the very next request after a change — no daemon restart. +// +// WHY this test exists (the gap unit tests can't cover): sysproxy unit tests +// prove the resolver returns the right URL; this test proves the REAL +// buildTransport + Watcher chain — the exact objects production wires in +// main() — actually routes live HTTP traffic through the switched proxy. +// +// Hermetic by construction: both "system proxies" are local httptest servers +// acting as plain-HTTP forward proxies (they receive absolute-form request +// URIs and answer themselves — no DNS, no real egress, no Anthropic). The +// destination host is a fake non-loopback name because ProxyFunc deliberately +// sends loopback destinations direct. + +import ( + "io" + "net/http" + "net/http/httptest" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/AiKeyLabs/aikey-proxy/internal/sysproxy" +) + +// fakeForwardProxy is a recording plain-HTTP forward proxy: it asserts it was +// reached AS a proxy (absolute-form URI) and answers with its own tag. +func fakeForwardProxy(t *testing.T, tag string) (*httptest.Server, *atomic.Int64) { + t.Helper() + var hits atomic.Int64 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if !strings.HasPrefix(r.RequestURI, "http://") { + t.Errorf("%s: expected absolute-form proxy request, got %q", tag, r.RequestURI) + } + hits.Add(1) + _, _ = io.WriteString(w, tag) + })) + t.Cleanup(srv.Close) + return srv, &hits +} + +// clearProxyEnvMain pins the proxy env empty so assertions are hermetic +// against the runner's shell (dev Macs export https_proxy; since the +// 2026-07-08 refinement an empty snapshot falls through to inherited env). +func clearProxyEnvMain(t *testing.T) { + t.Helper() + for _, k := range []string{ + "HTTP_PROXY", "http_proxy", "HTTPS_PROXY", "https_proxy", + "ALL_PROXY", "all_proxy", "NO_PROXY", "no_proxy", "AIKEY_PROXYENV_KEYS", + } { + t.Setenv(k, "") + } +} + +func TestEgress_FollowsSystemProxySwitch_NoRestart(t *testing.T) { + clearProxyEnvMain(t) + proxyA, hitsA := fakeForwardProxy(t, "via-A") + proxyB, hitsB := fakeForwardProxy(t, "via-B") + + var mu sync.Mutex + cur := sysproxy.Snapshot{HTTP: proxyA.URL, HTTPS: proxyA.URL} + setSystemProxy := func(s sysproxy.Snapshot) { mu.Lock(); cur = s; mu.Unlock() } + watcher := sysproxy.NewWatcherWithReader(func() (sysproxy.Snapshot, error) { + mu.Lock() + defer mu.Unlock() + return cur, nil + }) + + // The REAL production transport for direct mode (upstream_proxy.url empty). + transport := buildTransport("", watcher.ProxyFunc()) + client := &http.Client{Transport: transport, Timeout: 5 * time.Second} + + egress := func() string { + t.Helper() + resp, err := client.Get("http://provider.test/v1/ping") + if err != nil { + t.Fatalf("egress request failed: %v", err) + } + defer resp.Body.Close() + body, _ := io.ReadAll(resp.Body) + return string(body) + } + + // 1. Primed at construction: request #1 already goes through system proxy A. + if got := egress(); got != "via-A" { + t.Fatalf("initial egress: want via-A, got %q", got) + } + + // 2. "Clash 换端口": system proxy flips A→B; after ONE poll the very next + // request must egress through B — the core acceptance of the requirement. + setSystemProxy(sysproxy.Snapshot{HTTP: proxyB.URL, HTTPS: proxyB.URL}) + if !watcher.PollOnce() { + t.Fatal("watcher must observe the A→B switch") + } + if got := egress(); got != "via-B" { + t.Fatalf("post-switch egress: want via-B, got %q", got) + } + if hitsA.Load() != 1 || hitsB.Load() != 1 { + t.Fatalf("want exactly 1 hit per proxy (A then B), got A=%d B=%d", hitsA.Load(), hitsB.Load()) + } + + // 3. "系统代理关闭": toggle off refreshes to direct — the resolver must + // return nil (we assert the resolver, not a live dial: provider.test has + // no DNS on purpose). + setSystemProxy(sysproxy.Snapshot{}) + if !watcher.PollOnce() { + t.Fatal("watcher must observe the toggle-off") + } + req, _ := http.NewRequest(http.MethodGet, "http://provider.test/v1/ping", nil) + if u, err := transport.Proxy(req); err != nil || u != nil { + t.Fatalf("after toggle-off egress must be direct, got proxy=%v err=%v", u, err) + } +} + +// Explicit upstream_proxy.url must OUTRANK the system proxy (approved +// precedence): with the system proxy pointing at A, an explicit URL of B must +// carry the traffic. +func TestEgress_ExplicitURLOutranksSystemProxy(t *testing.T) { + proxyA, hitsA := fakeForwardProxy(t, "via-A") + proxyB, _ := fakeForwardProxy(t, "via-B") + + watcher := sysproxy.NewWatcherWithReader(func() (sysproxy.Snapshot, error) { + return sysproxy.Snapshot{HTTP: proxyA.URL, HTTPS: proxyA.URL}, nil + }) + transport := buildTransport(proxyB.URL, watcher.ProxyFunc()) + client := &http.Client{Transport: transport, Timeout: 5 * time.Second} + + resp, err := client.Get("http://provider.test/v1/ping") + if err != nil { + t.Fatalf("egress request failed: %v", err) + } + defer resp.Body.Close() + body, _ := io.ReadAll(resp.Body) + if string(body) != "via-B" { + t.Fatalf("explicit URL must win: want via-B, got %q", body) + } + if hitsA.Load() != 0 { + t.Fatalf("system proxy A must not be touched when explicit URL is set, got %d hits", hitsA.Load()) + } +} + +// Fence for egressState (the /admin/upstream-proxy "egress" block that +// `aikey env` renders): each layer must be reported, and the effective value +// must match what the transport's resolver would do. +func TestEgressState_LayeredReporting(t *testing.T) { + clearProxyEnvMain(t) + sys := sysproxy.Snapshot{HTTP: "http://127.0.0.1:7890", HTTPS: "http://127.0.0.1:7890", SOCKS: "socks5://127.0.0.1:7891"} + watcher := sysproxy.NewWatcherWithReader(func() (sysproxy.Snapshot, error) { return sys, nil }) + + // No explicit URL → system proxy wins and all layer fields are visible. + st := egressState("", watcher) + if st.EffectiveSource != "system" || st.EffectiveURL != "http://127.0.0.1:7890" { + t.Fatalf("want system/http://127.0.0.1:7890, got %s/%s", st.EffectiveSource, st.EffectiveURL) + } + if !st.SystemSupported || st.SystemHTTPS != sys.HTTPS || st.SystemSOCKS != sys.SOCKS { + t.Fatalf("system layer misreported: %+v", st) + } + + // Explicit URL outranks the same system snapshot. + st = egressState("socks5://10.0.0.1:1080", watcher) + if st.EffectiveSource != "explicit" || st.EffectiveURL != "socks5://10.0.0.1:1080" { + t.Fatalf("want explicit win, got %s/%s", st.EffectiveSource, st.EffectiveURL) + } + if st.SystemHTTPS != sys.HTTPS { + t.Fatal("lower layers must stay visible even when outranked (逐级显示)") + } + + // System proxy toggled off → direct. + empty := sysproxy.NewWatcherWithReader(func() (sysproxy.Snapshot, error) { return sysproxy.Snapshot{}, nil }) + st = egressState("", empty) + if st.EffectiveSource != "direct" || st.EffectiveURL != "" { + t.Fatalf("want direct, got %s/%s", st.EffectiveSource, st.EffectiveURL) + } +} + +// Explicit (proxy.env-marked) env: flagged authoritative, listed in the +// explicit map with credentials redacted, and effective through the real +// resolver. Uses the REAL NewWatcher (explicit env keeps it inert, so no OS +// read happens even on macOS). +func TestEgressState_ExplicitEnvReported(t *testing.T) { + clearProxyEnvMain(t) + t.Setenv("HTTPS_PROXY", "http://user:secret@127.0.0.1:7899") + t.Setenv("AIKEY_PROXYENV_KEYS", "HTTPS_PROXY") + watcher := sysproxy.NewWatcher() + st := egressState("", watcher) + if !st.EnvAuthoritative { + t.Fatal("proxy.env-marked HTTPS_PROXY must be authoritative") + } + got, ok := st.EnvVars["HTTPS_PROXY"] + if !ok { + t.Fatalf("explicit vars must list HTTPS_PROXY, got %v", st.EnvVars) + } + if strings.Contains(got, "secret") { + t.Fatalf("credentials must be redacted, got %q", got) + } + if st.EffectiveSource != "env" { + t.Fatalf("effective source must be env, got %s", st.EffectiveSource) + } + if strings.Contains(st.EffectiveURL, "secret") { + t.Fatalf("effective URL must be redacted, got %q", st.EffectiveURL) + } +} + +// Inherited (unmarked) env: NOT authoritative, listed as layer 4, and only +// effective when the system snapshot is empty (2026-07-08 refinement — the +// field case from the user's Mac .zshrc and the Windows HKCU\Environment). +func TestEgressState_InheritedEnvDemotedBelowSystem(t *testing.T) { + clearProxyEnvMain(t) + t.Setenv("https_proxy", "http://127.0.0.1:7890") // no marker → inherited + + // System proxy present → it wins; inherited env is visible but outranked. + withSys := sysproxy.NewWatcherWithReader(func() (sysproxy.Snapshot, error) { + return sysproxy.Snapshot{HTTPS: "http://127.0.0.1:9999"}, nil + }) + st := egressState("", withSys) + if st.EnvAuthoritative { + t.Fatal("inherited env must not be authoritative") + } + if st.EnvInheritedVars["https_proxy"] == "" { + t.Fatalf("inherited vars must list https_proxy, got %v", st.EnvInheritedVars) + } + if st.EffectiveSource != "system" || st.EffectiveURL != "http://127.0.0.1:9999" { + t.Fatalf("system must outrank inherited env, got %s/%s", st.EffectiveSource, st.EffectiveURL) + } + + // No system proxy → inherited env is the fallback. + noSys := sysproxy.NewWatcherWithReader(func() (sysproxy.Snapshot, error) { + return sysproxy.Snapshot{}, nil + }) + st = egressState("", noSys) + if st.EffectiveSource != "env_inherited" || st.EffectiveURL != "http://127.0.0.1:7890" { + t.Fatalf("want env_inherited fallback, got %s/%s", st.EffectiveSource, st.EffectiveURL) + } +} + +// Transport-level layer-4 proof: with no system proxy, live traffic egresses +// through the INHERITED env proxy (old headless/manual behavior preserved). +func TestEgress_InheritedEnvFallbackCarriesTraffic(t *testing.T) { + clearProxyEnvMain(t) + proxyC, hitsC := fakeForwardProxy(t, "via-C") + t.Setenv("http_proxy", proxyC.URL) // inherited: no marker + watcher := sysproxy.NewWatcherWithReader(func() (sysproxy.Snapshot, error) { + return sysproxy.Snapshot{}, nil + }) + transport := buildTransport("", watcher.ProxyFunc()) + client := &http.Client{Transport: transport, Timeout: 5 * time.Second} + resp, err := client.Get("http://provider.test/v1/ping") + if err != nil { + t.Fatalf("egress request failed: %v", err) + } + defer resp.Body.Close() + body, _ := io.ReadAll(resp.Body) + if string(body) != "via-C" || hitsC.Load() != 1 { + t.Fatalf("want inherited-env egress via-C (1 hit), got %q hits=%d", body, hitsC.Load()) + } +} diff --git a/cmd/aikey-proxy/main.go b/cmd/aikey-proxy/main.go index 8396399..577aa44 100644 --- a/cmd/aikey-proxy/main.go +++ b/cmd/aikey-proxy/main.go @@ -11,6 +11,8 @@ import ( "os" "os/signal" "path/filepath" + "sync" + "sync/atomic" "time" "golang.org/x/term" @@ -23,7 +25,9 @@ import ( proxyruntime "github.com/AiKeyLabs/aikey-proxy/internal/runtime" "github.com/AiKeyLabs/aikey-proxy/internal/server" "github.com/AiKeyLabs/aikey-proxy/internal/supervisor" + "github.com/AiKeyLabs/aikey-proxy/internal/sysproxy" "github.com/AiKeyLabs/aikey-proxy/internal/vault" + "github.com/AiKeyLabs/aikey-proxy/internal/vkeys" broker "github.com/AiKeyLabs/aikey-auth-broker" "github.com/AiKeyLabs/pkg/aikeycompat" @@ -127,6 +131,21 @@ func main() { Level: slog.LevelInfo, }))) + // Handle graceful shutdown — see pkg/aikeycompat for the per-OS signal set + // (Windows: SIGINT only; Unix: SIGINT + SIGTERM). + // + // WHY Notify is installed FIRST, before any init (2026-07-08 bugfix, e2e + // i3): a service manager may forward SIGTERM at any moment after exec + // (systemd stop right after start, restart storms), and the CLI writes the + // ownership meta while we are still initializing (vault/broker/observers — + // hundreds of ms). Before this move an early SIGTERM hit Go's default + // disposition and killed the process ("signal: 15", non-zero exit) instead + // of draining gracefully. The buffered channel captures the early signal; + // the select at the bottom of main consumes it right after init completes → + // same graceful drain path, exit 0. + sigCh := make(chan os.Signal, 1) + signal.Notify(sigCh, aikeycompat.ShutdownSignals()...) + // 1. Resolve and load config. resolvedPath, err := resolveConfigPath(*configPath) if err != nil { @@ -242,21 +261,58 @@ func main() { "listen", ln.Addr(), ) + // 5a. OS system-proxy watcher (2026-07-08 需求: 代理更新时自动刷新). + // Primes synchronously so both the broker client below and the forwarding + // transport see the CURRENT system proxy from request #1. Precedence + // (user-approved): explicit upstream_proxy.url > process env > OS system + // proxy (live). The poll loop is started after egress wiring (below) so + // its change callback can reach the live transport + broker swap. + sysWatcher := sysproxy.NewWatcher() + // effectiveBrokerEgress: static egress URL for the OAuth impersonate + // client — explicit config wins; else the system-derived URL ("" = client + // falls back to its own env handling, the pre-2026-07-08 behavior). + effectiveBrokerEgress := func(cfgURL string) string { + if cfgURL != "" { + return cfgURL + } + return sysWatcher.BrokerEgressURL() + } + // 5b. Build OAuth broker (embedded, uses vault stores). // Reuse the supervisor's vault connection to avoid double-open conflicts. brokerVault := sup.VaultReader() var oauthHandler server.RouteRegistrar + var poolHandler server.RouteRegistrar // C10: pool login (memory-store broker → writeback master) if brokerVault != nil { tokenStore := vault.NewTokenStore(brokerVault.DB(), brokerVault.DerivedKey()) accountStore := vault.NewAccountStore(brokerVault.DB()) // Inject ImpersonateChrome HTTP client for Claude token endpoint (Cloudflare bypass). - broker.SetHTTPClient(vault.NewImpersonateChromeHTTPClient()) + // Impl lives in the shared broker module (moved 2026-06-26) so aikey-control-master + // can inject the same client for its server-side OAuth flow — single source of truth. + // + // 2026-06-30 (egress fix): pass the SAME configured egress as AI forwarding + // (cfg.UpstreamProxy.URL) instead of "". Why: launchd-spawned proxies don't + // inherit the login shell's HTTP_PROXY, so a "" client went DIRECT and got + // 403 "Request not allowed" from Anthropic's edge (only Chrome-TLS THROUGH the + // egress reaches the OAuth logic). Mirrors the old master pattern (egress URL → + // impersonate proxyURL). Empty url ⇒ "" ⇒ falls back to env (no regression). + // Forwarding uses the same cfg.UpstreamProxy.URL via buildTransport below; + // control-plane clients deliberately bypass it (internal/httpx.NewDirectClient). + // 2026-07-08: "" no longer means "hope the env has it" — the system + // proxy watcher supplies the live OS value (launchd-spawned proxies + // have no login-shell env), keeping OAuth on the same egress as AI + // forwarding. + broker.SetHTTPClient(broker.NewImpersonateChromeHTTPClient(effectiveBrokerEgress(cfg.UpstreamProxy.URL))) brk := broker.NewEmbedded(tokenStore, accountStore) oauthHandler = broker.NewHandler(brk) // Wire broker into supervisor so all proxy generations can resolve OAuth credentials sup.SetBroker(&brokerAdapter{inner: brk, accounts: accountStore}) + // C10/RW8: pool login handler — its own MEMORY-store broker (per-member + // tokens never touch the vault) reusing the global impersonate client set + // above; exchanges then write back to master RW10 over the team JWT. + poolHandler = sup.NewPoolLoginHandler() slog.Info("OAuth broker initialized") } else { slog.Warn("OAuth broker disabled: vault not available") @@ -275,22 +331,139 @@ func main() { adminHandler.DebugUpstreamHeadersStateFn = proxy.UpstreamHeadersDebugState adminHandler.DebugUpstreamHeadersSetFn = proxy.SetUpstreamHeadersDebugAPIOverride adminHandler.AppHealthFn = sup.AppHealthSnapshot + // Oauth-group routing health (N9): omitted from /status unless the feature is + // on, so non-pool deployments are unchanged. Surfaces which pool accounts are + // currently cooled, for the operator monitoring the first pool batch. + adminHandler.PoolHealthFn = func() *admin.PoolRoutingHealth { + if !vkeys.OauthGroupRoutingEnabled() { + return nil + } + h := &admin.PoolRoutingHealth{Enabled: true} + for id, secs := range sup.PoolCooldownSnapshot() { + h.CooledAccounts = append(h.CooledAccounts, admin.CooledAccount{AccountID: id, CooldownSeconds: secs}) + } + return h + } + // SyncRail health (2026-07-03): per-rail control-plane sync state for /status + // (health-signal-surface rule — the release E2E asserts this, not the UI). + // Empty map (no rail ever attempted — personal installs) → field omitted. + adminHandler.SyncHealthFn = func() map[string]admin.SyncRailStatus { + snap := sup.ControlPlaneSyncSnapshot() + if len(snap) == 0 { + return nil + } + out := make(map[string]admin.SyncRailStatus, len(snap)) + for name, st := range snap { + out[name] = admin.SyncRailStatus{ + State: st.State, + ConsecutiveFailures: st.ConsecutiveFailures, + LastSuccessAt: st.LastSuccessAt, + LastError: st.LastError, + } + } + return out + } adminHandler.EffectivePacksFn = sup.EffectivePacks adminHandler.AuditStatusFn = sup.AuditStatus adminHandler.ReconcileGapsFn = sup.ReconcileGaps + // Egress (upstream) proxy config endpoints (GET/PUT /admin/upstream-proxy), the + // runtime home of the egress proxy after R25 出口收敛. Get returns the live URL; + // Set persists it to aikey-user.yaml (survives re-render) and HOT-SWAPS the + // forwarding transport + OAuth impersonate client in place — no restart. egressMu + // guards the cached URL Get reads (the transport/broker swaps are themselves + // atomic). Decision: user chose hot-swap (B) + plaintext storage (i). + var egressMu sync.Mutex + egressURL := cfg.UpstreamProxy.URL + adminHandler.GetUpstreamProxyFn = func() string { + egressMu.Lock() + defer egressMu.Unlock() + return egressURL + } + // liveTransport tracks the transport currently serving forwarding, so the + // system-proxy change callback can flush its idle pool (pooled connections + // to a dead Clash port would otherwise be reused until IdleConnTimeout). + var liveTransport atomic.Pointer[http.Transport] + installTransport := func(t *http.Transport) { + liveTransport.Store(t) + sup.SetTransport(t) + } + adminHandler.SetUpstreamProxyFn = func(rawURL string) error { + if err := config.PersistUpstreamProxyURL(resolvedPath, rawURL); err != nil { + return err + } + installTransport(buildTransport(rawURL, sysWatcher.ProxyFunc())) + broker.SetHTTPClient(broker.NewImpersonateChromeHTTPClient(effectiveBrokerEgress(rawURL))) + egressMu.Lock() + egressURL = rawURL + egressMu.Unlock() + return nil + } + // EgressStateFn (2026-07-08): layered egress decision for `aikey env`. + adminHandler.EgressStateFn = func() admin.EgressState { + egressMu.Lock() + explicit := egressURL + egressMu.Unlock() + return egressState(explicit, sysWatcher) + } + + // ProbeUpstreamProxyFn tests a CANDIDATE egress URL end-to-end (through the same + // buildTransport the live path uses) to a provider host, WITHOUT persisting it, so + // the web "Test connectivity" button can verify before Save. Any HTTP response + // (even 401/404) proves the tunnel carries traffic; a transport error means the + // proxy couldn't reach out. + adminHandler.ProbeUpstreamProxyFn = func(rawURL string) (int, int64, error) { + client := &http.Client{Transport: buildTransport(rawURL, sysWatcher.ProxyFunc()), Timeout: 10 * time.Second} + req, err := http.NewRequest(http.MethodGet, upstreamProbeTarget, nil) + if err != nil { + return 0, 0, err + } + start := time.Now() + resp, err := client.Do(req) + elapsed := time.Since(start).Milliseconds() + if err != nil { + return 0, elapsed, err + } + _ = resp.Body.Close() + return resp.StatusCode, elapsed, nil + } + // Build the outbound transport for upstream providers. Always non-nil now: // even the direct path needs MaxIdleConnsPerHost tuning (see buildTransport). dataHandler := sup.Handler() - sup.SetTransport(buildTransport(cfg.UpstreamProxy.URL)) - - // 7. Build and start the HTTP server. - srv := server.New(ln, dataHandler, adminHandler, oauthHandler) + installTransport(buildTransport(cfg.UpstreamProxy.URL, sysWatcher.ProxyFunc())) + + // Start the system-proxy poll loop (inert when env config is authoritative, + // on unsupported platforms, and — via the explicit-egress guard below — + // effectively when upstream_proxy.url is set). On change: flush the idle + // pool so no request reuses a connection through the OLD proxy, and rebuild + // the OAuth client (it takes a static URL, not a per-request func). + observability.GoSafe("sysproxy.watch", observability.Isolated, func() { + sysWatcher.Run(context.Background(), func(_, _ sysproxy.Snapshot) { + egressMu.Lock() + explicit := egressURL != "" + egressMu.Unlock() + if explicit { + return // user-configured egress outranks the system proxy + } + if t := liveTransport.Load(); t != nil { + t.CloseIdleConnections() + } + broker.SetHTTPClient(broker.NewImpersonateChromeHTTPClient(effectiveBrokerEgress(""))) + }) + }) - // Handle graceful shutdown — see pkg/aikeycompat for the per-OS - // signal set (Windows: SIGINT only; Unix: SIGINT + SIGTERM). - sigCh := make(chan os.Signal, 1) - signal.Notify(sigCh, aikeycompat.ShutdownSignals()...) + // 7. Build and start the HTTP server. Only non-nil registrars are passed + // (server.New does not nil-guard; oauthHandler/poolHandler are nil when the + // vault/broker is unavailable). + extraRegistrars := make([]server.RouteRegistrar, 0, 2) + if oauthHandler != nil { + extraRegistrars = append(extraRegistrars, oauthHandler) + } + if poolHandler != nil { + extraRegistrars = append(extraRegistrars, poolHandler) + } + srv := server.New(ln, dataHandler, adminHandler, extraRegistrars...) // Fatal: server.Serve is the process's reason for being. If it dies the // proxy cannot serve requests, so exit(2) and let the OS supervisor @@ -303,9 +476,20 @@ func main() { } }) - // Wait for shutdown signal. - sig := <-sigCh - slog.Info("received signal, shutting down", "signal", sig) + // Wait for either an OS shutdown signal OR a graceful self-restart request + // from the control-plane healer (selfheal.go): the proxy went stale after a + // host network change and confirmed master is reachable, so it wants a clean + // relaunch. BOTH paths run the identical drain below (code reuse); the restart + // path just exits non-zero afterwards so launchd/systemd relaunch it. + restartRequested := false + select { + case sig := <-sigCh: + slog.Info("received signal, shutting down", "signal", sig) + case <-sup.RestartRequested(): + restartRequested = true + slog.Warn("control-plane self-heal: graceful restart requested — draining then relaunching", + "event.name", observability.EventProxyControlPlaneSelfRestart) + } // Graceful shutdown: stop accepting new connections, drain active generation. if err := srv.Shutdown(30 * time.Second); err != nil { @@ -322,6 +506,13 @@ func main() { logWriter.Flush(3 * time.Second) _ = logWriter.Close() } + + // Self-restart requests exit non-zero (EX_TEMPFAIL) AFTER the graceful drain + // above, so the OS service manager relaunches a clean process. A plain signal + // shutdown returns normally (exit 0) and is NOT relaunched. + if restartRequested { + os.Exit(75) + } } // getVaultPassword reads the master password from AIKEY_MASTER_PASSWORD env var @@ -347,9 +538,13 @@ func getVaultPassword() (string, error) { // buildTransport returns the outbound http.Transport for upstream providers. // -// Direct (no upstream_proxy): a Clone of http.DefaultTransport — preserving -// ProxyFromEnvironment (HTTP_PROXY / HTTPS_PROXY / NO_PROXY) and HTTP/2 — with -// MaxIdleConnsPerHost raised. Go's default of 2 idle conns per host means that +// Direct (no upstream_proxy): a Clone of http.DefaultTransport — with Proxy +// swapped to sysProxy (2026-07-08): the sysproxy.Watcher's per-request +// resolver, which preserves env-var behavior (HTTP_PROXY / HTTPS_PROXY / +// NO_PROXY) when the env is set and otherwise follows the LIVE OS system +// proxy, so a Clash port change / toggle applies without a daemon restart. +// HTTP/2 is preserved, with MaxIdleConnsPerHost raised. Go's default of 2 +// idle conns per host means that // under >2 concurrent requests to one provider, finished connections are // dropped and the next request pays a fresh TCP+TLS handshake. HTTP/2 // multiplexing masks this against h2 upstreams, but HTTP/1.1 fallbacks and @@ -371,22 +566,78 @@ func cloneDefaultTransport() *http.Transport { return &http.Transport{} } -func buildTransport(proxyURL string) *http.Transport { +// upstreamProbeTarget is the provider host the "Test connectivity" probe reaches for +// through a candidate egress proxy. api.anthropic.com is the Claude-centric primary +// forwarding target; any HTTP response (even 401/404, no key sent) proves the tunnel +// carries traffic to the AI world. +const upstreamProbeTarget = "https://api.anthropic.com/" + +// egressState assembles the layered egress decision for GET /admin/upstream-proxy +// (`aikey env`, 2026-07-08 逐级显示需求). The effective value is computed through +// watcher.ProxyFunc() — the SAME resolver the forwarding transport runs per +// request — so what the CLI displays can never diverge from what traffic does. +// Pulled out of main() so the decision logic sits behind a fence test +// (egress_integration_test.go). +func egressState(explicit string, watcher *sysproxy.Watcher) admin.EgressState { + snap := watcher.Current() + envExplicit, envInherited := sysproxy.EnvProxyVarsSplit() + st := admin.EgressState{ + ExplicitURL: explicit, + EnvAuthoritative: watcher.EnvExplicit(), + EnvVars: envExplicit, + EnvInheritedVars: envInherited, + SystemSupported: watcher.Supported(), + SystemHTTP: snap.HTTP, + SystemHTTPS: snap.HTTPS, + SystemSOCKS: snap.SOCKS, + } + if explicit != "" { + st.EffectiveSource, st.EffectiveURL = "explicit", explicit + return st + } + // Resolve exactly like the transport would for a provider target. + // Request object only — nothing is dialed. + req, err := http.NewRequest(http.MethodGet, upstreamProbeTarget, nil) + if err != nil { + st.EffectiveSource = "direct" + return st + } + u, perr := watcher.ProxyFunc()(req) + switch { + case perr != nil || u == nil: + st.EffectiveSource = "direct" + case st.EnvAuthoritative: + st.EffectiveSource, st.EffectiveURL = "env", u.Redacted() + case snap.ProxyFor("https") != "": + st.EffectiveSource, st.EffectiveURL = "system", u.Redacted() + default: + // Layers 1-3 empty, ProxyFunc fell through to inherited shell env. + st.EffectiveSource, st.EffectiveURL = "env_inherited", u.Redacted() + } + return st +} + +func buildTransport(proxyURL string, sysProxy func(*http.Request) (*url.URL, error)) *http.Transport { // All providers resolve to a handful of hosts (api.anthropic.com etc.), // so a generous per-host idle pool is bounded in practice by MaxIdleConns. const perHost = 100 - if proxyURL == "" { + direct := func() *http.Transport { t := cloneDefaultTransport() + if sysProxy != nil { + t.Proxy = sysProxy + } t.MaxIdleConnsPerHost = perHost return t } + + if proxyURL == "" { + return direct() + } parsed, err := url.Parse(proxyURL) if err != nil { - slog.Warn("upstream_proxy.url is invalid, falling back to env vars", "url", proxyURL, "error", err) - t := cloneDefaultTransport() - t.MaxIdleConnsPerHost = perHost - return t + slog.Warn("upstream_proxy.url is invalid, falling back to env/system proxy", "url", proxyURL, "error", err) + return direct() } t := &http.Transport{ Proxy: http.ProxyURL(parsed), diff --git a/go.mod b/go.mod index 0c0fea1..0918e57 100644 --- a/go.mod +++ b/go.mod @@ -9,10 +9,9 @@ require ( github.com/AiKeyLabs/pkg/aikeytime v0.0.0 github.com/AiKeyLabs/pkg/buildinfo v0.0.0 github.com/AiKeyLabs/pkg/providerroutes v0.0.0 + github.com/AiKeyLabs/pkg/seatassign v0.0.0 github.com/AiKeyLabs/pkg/usagehash v0.0.0 github.com/aikeylabs/ai-degrade-detector/proxy-plugin/rhythm v0.0.0 - github.com/cespare/xxhash/v2 v2.3.0 - github.com/imroc/req/v3 v3.57.0 github.com/tidwall/gjson v1.19.0 github.com/tidwall/sjson v1.2.5 golang.org/x/crypto v0.49.0 @@ -29,6 +28,8 @@ replace github.com/AiKeyLabs/pkg/aikeycompat => ../pkg/aikeycompat replace github.com/AiKeyLabs/pkg/providerroutes => ../pkg/providerroutes +replace github.com/AiKeyLabs/pkg/seatassign => ../pkg/seatassign + replace github.com/AiKeyLabs/pkg/usagehash => ../pkg/usagehash replace github.com/AiKeyLabs/aikey-auth-broker => ../aikey-auth-broker @@ -42,11 +43,13 @@ replace github.com/AiKeyLabs/aikey-config-tool => ../aikey-config-tool replace github.com/aikeylabs/ai-degrade-detector/proxy-plugin/rhythm => ../ai-degrade-detector/proxy-plugin/rhythm require ( + github.com/AiKeyLabs/pkg/routingwire v0.0.0 github.com/andybalholm/brotli v1.2.0 // indirect github.com/dustin/go-humanize v1.0.1 // indirect github.com/google/go-querystring v1.1.0 // indirect github.com/google/uuid v1.6.0 // indirect github.com/icholy/digest v1.1.0 // indirect + github.com/imroc/req/v3 v3.57.0 // indirect github.com/klauspost/compress v1.18.2 // indirect github.com/kr/text v0.2.0 // indirect github.com/mattn/go-isatty v0.0.20 // indirect @@ -57,10 +60,12 @@ require ( github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec // indirect github.com/tidwall/match v1.1.1 // indirect github.com/tidwall/pretty v1.2.0 // indirect - golang.org/x/net v0.51.0 // indirect - golang.org/x/sys v0.42.0 // indirect + golang.org/x/net v0.51.0 + golang.org/x/sys v0.42.0 golang.org/x/text v0.35.0 // indirect modernc.org/libc v1.70.0 // indirect modernc.org/mathutil v1.7.1 // indirect modernc.org/memory v1.11.0 // indirect ) + +replace github.com/AiKeyLabs/pkg/routingwire => ../pkg/routingwire diff --git a/go.sum b/go.sum index 6541d3c..2a26def 100644 --- a/go.sum +++ b/go.sum @@ -1,7 +1,5 @@ github.com/andybalholm/brotli v1.2.0 h1:ukwgCxwYrmACq68yiUqwIWnGY0cTPox/M94sVwToPjQ= github.com/andybalholm/brotli v1.2.0/go.mod h1:rzTDkvFWvIrjDXZHkuS16NPggd91W3kUSvPlQ1pLaKY= -github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs= -github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E= github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= diff --git a/internal/admin/handlers.go b/internal/admin/handlers.go index 31a7cc5..8fc39c6 100644 --- a/internal/admin/handlers.go +++ b/internal/admin/handlers.go @@ -74,6 +74,13 @@ type Handler struct { // from "feature disabled" (503). AppHealthFn func() []apppipe.AppHealth + // PoolHealthFn returns oauth-group routing health for /status (N9). nil → the + // pool_routing field is omitted (non-pool deployments unchanged). + PoolHealthFn func() *PoolRoutingHealth + // SyncHealthFn supplies the SyncRail per-rail health map for /status. Nil or + // an empty map → the control_plane_sync field is omitted. + SyncHealthFn func() map[string]SyncRailStatus + // EffectivePacksFn returns the raw JSON report of compliance packs currently // effective in the live filter child (built-in + pulled). Returns an error // (apphook.ErrPacksUnavailable) when no filter child is active / can't report; @@ -86,6 +93,31 @@ type Handler struct { // (re-send WAL-present gaps, confirm WAL-absent gaps lost now). nil → 503. AuditStatusFn func() *events.AuditStatus ReconcileGapsFn func(ctx context.Context) (events.ReconcileResult, error) + + // GetUpstreamProxyFn / SetUpstreamProxyFn back the GET/PUT /admin/upstream-proxy + // endpoints that the local web "Settings → Upstream proxy" card relays to. Get + // returns the live egress proxy URL ("" = direct); Set persists it to + // aikey-user.yaml and HOT-SWAPS the running transport + impersonate client (no + // restart). nil → 503 (endpoint disabled — e.g. cluster node). See + // config.PersistUpstreamProxyURL + the main.go wiring. + GetUpstreamProxyFn func() string + SetUpstreamProxyFn func(url string) error + + // EgressStateFn backs the layered "egress" block in GET /admin/upstream-proxy + // (2026-07-08, `aikey env` 逐级显示需求): the daemon-truth view of the egress + // decision — explicit URL > daemon process env > OS system proxy — plus the + // effective result computed through the SAME resolution path the forwarding + // transport uses (display must never diverge from behavior). nil → the GET + // response stays the legacy {"url"} shape (older wiring / cluster nodes). + EgressStateFn func() EgressState + + // ProbeUpstreamProxyFn tests whether a CANDIDATE egress URL can actually carry + // traffic to an AI provider (built with the same buildTransport the live path + // uses), WITHOUT persisting it — so the web "Test connectivity" button can verify + // before Save. Returns (httpStatus, elapsedMs, err); err = the request never got + // a response (proxy unreachable / DNS / timeout). Any HTTP status = reachable. + // nil → 503. + ProbeUpstreamProxyFn func(url string) (status int, elapsedMs int64, err error) } // KeyCheckTarget holds decrypted credentials for one provider, used by GET /health/keys. @@ -230,6 +262,38 @@ type statusResponse struct { VirtualKeys int `json:"virtual_keys_loaded"` TotalReqs int64 `json:"total_requests"` TotalErrs int64 `json:"total_errors"` + // PoolRouting is the oauth-group routing health (N9). Omitted unless the + // feature is on, so non-pool deployments' /status is unchanged. Lets the + // operator monitoring the first pool batch see which accounts are cooled. + PoolRouting *PoolRoutingHealth `json:"pool_routing,omitempty"` + // ControlPlaneSync is the SyncRail health surface (2026-07-03): per-rail + // ok/stale/offline state with failure counts and last error. Omitted when no + // rail has ever attempted a cycle (personal installs — /status unchanged). + // Release-checklist E2E and `aikey statusline` read this to assert the + // master-sync pipeline is alive (health-signal-surface rule). + ControlPlaneSync map[string]SyncRailStatus `json:"control_plane_sync,omitempty"` +} + +// SyncRailStatus mirrors supervisor.RailSyncStatus for the /status wire (built +// by the cmd layer — admin does not import supervisor). +type SyncRailStatus struct { + State string `json:"state"` + ConsecutiveFailures int `json:"consecutive_failures"` + LastSuccessAt int64 `json:"last_success_at,omitempty"` + LastError string `json:"last_error,omitempty"` +} + +// PoolRoutingHealth is the oauth-group account-routing health surface (N9). Built +// by the cmd layer from the proxy's reactive cooldown state. +type PoolRoutingHealth struct { + Enabled bool `json:"enabled"` + CooledAccounts []CooledAccount `json:"cooled_accounts,omitempty"` +} + +// CooledAccount is one pool account currently routed around (401 / exhaustion). +type CooledAccount struct { + AccountID string `json:"account_id"` + CooldownSeconds int `json:"cooldown_seconds"` } // Status returns detailed proxy status. @@ -248,27 +312,43 @@ func (h *Handler) Status(w http.ResponseWriter, r *http.Request) { totalErrs = h.TotalErrorsFn() } + var poolRouting *PoolRoutingHealth + if h.PoolHealthFn != nil { + poolRouting = h.PoolHealthFn() + } + var syncHealth map[string]SyncRailStatus + if h.SyncHealthFn != nil { + syncHealth = h.SyncHealthFn() + } + writeJSON(w, http.StatusOK, statusResponse{ - Status: "ok", - Version: Version, - Uptime: time.Since(h.startedAt).Round(time.Second).String(), - ListenAddr: h.cfg.Listen.Addr(), - VirtualKeys: h.registry.Count(), - VaultPath: h.cfg.Vault.Path, - StartedAt: h.startedAt.Format(time.RFC3339), - TotalReqs: totalReqs, - TotalErrs: totalErrs, + Status: "ok", + Version: Version, + Uptime: time.Since(h.startedAt).Round(time.Second).String(), + ListenAddr: h.cfg.Listen.Addr(), + VirtualKeys: h.registry.Count(), + VaultPath: h.cfg.Vault.Path, + StartedAt: h.startedAt.Format(time.RFC3339), + TotalReqs: totalReqs, + TotalErrs: totalErrs, + PoolRouting: poolRouting, + ControlPlaneSync: syncHealth, }) } type metricsResponse struct { - RequestsByVKey map[string]int64 `json:"requests_by_vkey"` - RequestsByProvider map[string]int64 `json:"requests_by_provider"` - Reporter *events.ReporterMetrics `json:"reporter,omitempty"` - Collector *events.CollectorMetrics `json:"collector,omitempty"` - Canary *events.CanaryResult `json:"canary,omitempty"` - TotalRequests int64 `json:"total_requests"` - TotalErrors int64 `json:"total_errors"` + RequestsByVKey map[string]int64 `json:"requests_by_vkey"` + RequestsByProvider map[string]int64 `json:"requests_by_provider"` + // RequestsByAccount: per real serving account (oauth_group attribution). + // Counts follow pool fallback (A→B counts toward B). Empty when no group + // routing has happened. Lets local audit see "which account served" without + // the collector/ODS. + RequestsByAccount map[string]int64 `json:"requests_by_account"` + Reporter *events.ReporterMetrics `json:"reporter,omitempty"` + Collector *events.CollectorMetrics `json:"collector,omitempty"` + Canary *events.CanaryResult `json:"canary,omitempty"` + TotalRequests int64 `json:"total_requests"` + TotalErrors int64 `json:"total_errors"` } // Metrics returns aggregated usage metrics. @@ -286,6 +366,10 @@ func (h *Handler) Metrics(w http.ResponseWriter, r *http.Request) { byVKey = make(map[string]int64) byProvider = make(map[string]int64) } + byAccount, err := h.store.QueryByAccount() + if err != nil { + byAccount = make(map[string]int64) + } var reporterMetrics *events.ReporterMetrics if h.ReporterMetricsFn != nil { @@ -305,6 +389,7 @@ func (h *Handler) Metrics(w http.ResponseWriter, r *http.Request) { TotalErrors: totalErrs, RequestsByVKey: byVKey, RequestsByProvider: byProvider, + RequestsByAccount: byAccount, Reporter: reporterMetrics, Collector: collectorMetrics, Canary: canaryResult, diff --git a/internal/admin/pool_health_test.go b/internal/admin/pool_health_test.go new file mode 100644 index 0000000..d2fcca7 --- /dev/null +++ b/internal/admin/pool_health_test.go @@ -0,0 +1,36 @@ +package admin + +import ( + "encoding/json" + "strings" + "testing" +) + +// N9 health surface contract: pool_routing is omitted from /status unless set +// (so non-pool deployments are byte-unchanged), and serializes the cooled-account +// shape the operator monitoring the first pool batch reads. +func TestStatusResponse_PoolRoutingSerialization(t *testing.T) { + // nil → omitted entirely. + b, _ := json.Marshal(statusResponse{Status: "ok"}) + if strings.Contains(string(b), "pool_routing") { + t.Fatalf("nil PoolRouting must be omitted from /status, got %s", b) + } + + // set → present with the cooled-account roster. + b, _ = json.Marshal(statusResponse{Status: "ok", PoolRouting: &PoolRoutingHealth{ + Enabled: true, + CooledAccounts: []CooledAccount{{AccountID: "acc-1", CooldownSeconds: 42}}, + }}) + s := string(b) + for _, want := range []string{`"pool_routing"`, `"enabled":true`, `"account_id":"acc-1"`, `"cooldown_seconds":42`} { + if !strings.Contains(s, want) { + t.Fatalf("status missing %s: %s", want, s) + } + } + + // enabled but nothing cooled → cooled_accounts omitted (clean steady state). + b, _ = json.Marshal(statusResponse{Status: "ok", PoolRouting: &PoolRoutingHealth{Enabled: true}}) + if strings.Contains(string(b), "cooled_accounts") { + t.Fatalf("empty cooled list must be omitted, got %s", b) + } +} diff --git a/internal/admin/upstream_proxy.go b/internal/admin/upstream_proxy.go new file mode 100644 index 0000000..f5abcd2 --- /dev/null +++ b/internal/admin/upstream_proxy.go @@ -0,0 +1,142 @@ +package admin + +import ( + "encoding/json" + "log/slog" + "net/http" + + "github.com/AiKeyLabs/aikey-proxy/internal/config" + "github.com/AiKeyLabs/aikey-proxy/internal/observability" +) + +// upstreamProxyBody is the GET response / PUT request shape: {"url": "..."}. +// Empty url means "no egress proxy → direct". +type upstreamProxyBody struct { + URL string `json:"url"` + // Egress is the layered decision state (2026-07-08, read side only — + // omitted when EgressStateFn isn't wired, keeping the legacy shape for + // existing consumers like the web Settings prefill, which reads .url). + Egress *EgressState `json:"egress,omitempty"` +} + +// EgressState is the daemon-truth view of the egress proxy decision, layer by +// layer, consumed by `aikey env`. WHY daemon-truth matters: the daemon's env +// is the launchd/systemd spawn snapshot — usually NOT the user's shell env — +// so only the daemon itself can answer "which proxy do my requests use". +type EgressState struct { + // ExplicitURL: layer 1 — the user-set upstream_proxy.url ("" = not set). + ExplicitURL string `json:"explicit_url"` + // EnvAuthoritative: layer 2 — true when the daemon env carries proxy vars + // EXPLICITLY configured via ~/.aikey/proxy.env (CLI-marked at spawn). + // Since the 2026-07-08 precedence refinement, ONLY those bypass OS + // detection; inherited shell env is a below-system fallback (layer 4). + EnvAuthoritative bool `json:"env_authoritative"` + // EnvVars: the EXPLICIT (proxy.env) vars (credentials redacted). + EnvVars map[string]string `json:"env_vars,omitempty"` + // EnvInheritedVars: layer 4 — proxy vars inherited from the spawn shell + // (.zshrc exports, stale terminals); consulted only when layers 1-3 are + // all empty. + EnvInheritedVars map[string]string `json:"env_inherited_vars,omitempty"` + // System: layer 3 — the live OS system-proxy snapshot. + SystemSupported bool `json:"system_supported"` + SystemHTTP string `json:"system_http,omitempty"` + SystemHTTPS string `json:"system_https,omitempty"` + SystemSOCKS string `json:"system_socks,omitempty"` + // Effective: the resolved result for https provider targets, computed via + // the SAME ProxyFunc the forwarding transport uses. Source is one of + // "explicit" | "env" | "system" | "env_inherited" | "direct"; URL is "" + // when direct. + EffectiveSource string `json:"effective_source"` + EffectiveURL string `json:"effective_url,omitempty"` +} + +// UpstreamProxyGet serves GET /admin/upstream-proxy — the live egress proxy URL the +// local web "Settings → Upstream proxy" card reads to prefill, plus (when wired) +// the layered egress decision for `aikey env`. 503 if not wired. +func (h *Handler) UpstreamProxyGet(w http.ResponseWriter, _ *http.Request) { + if h.GetUpstreamProxyFn == nil { + writeJSON(w, http.StatusServiceUnavailable, map[string]string{"error": "upstream-proxy config not supported"}) + return + } + body := upstreamProxyBody{URL: h.GetUpstreamProxyFn()} + if h.EgressStateFn != nil { + st := h.EgressStateFn() + body.Egress = &st + } + writeJSON(w, http.StatusOK, body) +} + +// UpstreamProxySet serves PUT /admin/upstream-proxy. Validates the URL, persists it +// to aikey-user.yaml, and HOT-SWAPS the running transport + impersonate client via +// SetUpstreamProxyFn (no restart). Empty url clears the proxy (direct egress). +// +// 400 on an invalid URL (bad scheme / missing host:port), 503 if not wired, 500 if +// the persist/hot-swap fails. The body that succeeds echoes the applied url so the +// caller can confirm. +func (h *Handler) UpstreamProxySet(w http.ResponseWriter, r *http.Request) { + tc := observability.ExtractOrCreate(r) + logger := slog.With("trace_id", tc.TraceID, "request_id", tc.RequestID) + + if h.SetUpstreamProxyFn == nil { + writeJSON(w, http.StatusServiceUnavailable, map[string]string{"error": "upstream-proxy config not supported"}) + return + } + + var body upstreamProxyBody + if err := json.NewDecoder(http.MaxBytesReader(w, r.Body, 4<<10)).Decode(&body); err != nil { + writeJSON(w, http.StatusBadRequest, map[string]string{"error": "invalid request body"}) + return + } + if err := config.ValidateUpstreamProxyURL(body.URL); err != nil { + writeJSON(w, http.StatusBadRequest, map[string]string{"error": err.Error()}) + return + } + + if err := h.SetUpstreamProxyFn(body.URL); err != nil { + logger.Error("admin: upstream-proxy update failed", "error.message", err.Error()) + writeJSON(w, http.StatusInternalServerError, map[string]string{"error": err.Error()}) + return + } + + logger.Info("admin: upstream-proxy updated + hot-swapped", + "event.name", "proxy.upstream_proxy.updated", "has_value", body.URL != "") + writeJSON(w, http.StatusOK, upstreamProxyBody{URL: body.URL}) +} + +// upstreamProbeResult is the POST /admin/upstream-proxy/probe response. Ok=true with +// a Status means the candidate proxy carried a request through to the provider (any +// HTTP status counts as reachable — auth/404 still proves the tunnel works). Ok=false +// with Error means no response came back (proxy down / DNS / timeout). +type upstreamProbeResult struct { + Ok bool `json:"ok"` + Status int `json:"status,omitempty"` + ElapsedMs int64 `json:"elapsed_ms,omitempty"` + Error string `json:"error,omitempty"` +} + +// UpstreamProxyProbe serves POST /admin/upstream-proxy/probe: tests a CANDIDATE egress +// URL (from the body) end-to-end to an AI provider WITHOUT saving it, so the web can +// verify connectivity before Save. 400 on an invalid URL, 503 if not wired. A probe +// that runs but can't reach the provider is a 200 with ok=false (not an HTTP error) — +// the caller renders it as a "unreachable" result, same shape as a success. +func (h *Handler) UpstreamProxyProbe(w http.ResponseWriter, r *http.Request) { + if h.ProbeUpstreamProxyFn == nil { + writeJSON(w, http.StatusServiceUnavailable, map[string]string{"error": "upstream-proxy probe not supported"}) + return + } + var body upstreamProxyBody + if err := json.NewDecoder(http.MaxBytesReader(w, r.Body, 4<<10)).Decode(&body); err != nil { + writeJSON(w, http.StatusBadRequest, map[string]string{"error": "invalid request body"}) + return + } + if err := config.ValidateUpstreamProxyURL(body.URL); err != nil { + writeJSON(w, http.StatusBadRequest, map[string]string{"error": err.Error()}) + return + } + status, elapsedMs, err := h.ProbeUpstreamProxyFn(body.URL) + if err != nil { + writeJSON(w, http.StatusOK, upstreamProbeResult{Ok: false, ElapsedMs: elapsedMs, Error: err.Error()}) + return + } + writeJSON(w, http.StatusOK, upstreamProbeResult{Ok: true, Status: status, ElapsedMs: elapsedMs}) +} diff --git a/internal/admin/upstream_proxy_test.go b/internal/admin/upstream_proxy_test.go new file mode 100644 index 0000000..bc77196 --- /dev/null +++ b/internal/admin/upstream_proxy_test.go @@ -0,0 +1,186 @@ +package admin + +import ( + "encoding/json" + "errors" + "net/http" + "net/http/httptest" + "strings" + "testing" +) + +func TestUpstreamProxyGet(t *testing.T) { + h := &Handler{GetUpstreamProxyFn: func() string { return "http://127.0.0.1:7890" }} + w := httptest.NewRecorder() + h.UpstreamProxyGet(w, httptest.NewRequest(http.MethodGet, "/admin/upstream-proxy", nil)) + if w.Code != http.StatusOK { + t.Fatalf("status = %d, want 200", w.Code) + } + var body upstreamProxyBody + _ = json.Unmarshal(w.Body.Bytes(), &body) + if body.URL != "http://127.0.0.1:7890" { + t.Fatalf("url = %q, want http://127.0.0.1:7890", body.URL) + } +} + +func TestUpstreamProxyGet_NotWired(t *testing.T) { + h := &Handler{} // GetUpstreamProxyFn nil + w := httptest.NewRecorder() + h.UpstreamProxyGet(w, httptest.NewRequest(http.MethodGet, "/admin/upstream-proxy", nil)) + if w.Code != http.StatusServiceUnavailable { + t.Fatalf("status = %d, want 503 when not wired", w.Code) + } +} + +func TestUpstreamProxySet_ValidCallsHotSwap(t *testing.T) { + var got string + called := 0 + h := &Handler{SetUpstreamProxyFn: func(url string) error { called++; got = url; return nil }} + w := httptest.NewRecorder() + r := httptest.NewRequest(http.MethodPut, "/admin/upstream-proxy", strings.NewReader(`{"url":"socks5://127.0.0.1:7891"}`)) + h.UpstreamProxySet(w, r) + if w.Code != http.StatusOK { + t.Fatalf("status = %d, want 200; body=%s", w.Code, w.Body.String()) + } + if called != 1 || got != "socks5://127.0.0.1:7891" { + t.Fatalf("SetUpstreamProxyFn called=%d url=%q", called, got) + } +} + +func TestUpstreamProxySet_EmptyClears(t *testing.T) { + called := 0 + h := &Handler{SetUpstreamProxyFn: func(string) error { called++; return nil }} + w := httptest.NewRecorder() + r := httptest.NewRequest(http.MethodPut, "/admin/upstream-proxy", strings.NewReader(`{"url":""}`)) + h.UpstreamProxySet(w, r) + if w.Code != http.StatusOK || called != 1 { + t.Fatalf("empty url should be accepted (clear): status=%d called=%d", w.Code, called) + } +} + +func TestUpstreamProxySet_InvalidURLNoHotSwap(t *testing.T) { + called := 0 + h := &Handler{SetUpstreamProxyFn: func(string) error { called++; return nil }} + w := httptest.NewRecorder() + r := httptest.NewRequest(http.MethodPut, "/admin/upstream-proxy", strings.NewReader(`{"url":"ftp://nope"}`)) + h.UpstreamProxySet(w, r) + if w.Code != http.StatusBadRequest { + t.Fatalf("status = %d, want 400 for bad scheme", w.Code) + } + if called != 0 { + t.Fatalf("invalid URL must NOT reach the hot-swap (called=%d)", called) + } +} + +func TestUpstreamProxySet_NotWired(t *testing.T) { + h := &Handler{} // SetUpstreamProxyFn nil + w := httptest.NewRecorder() + r := httptest.NewRequest(http.MethodPut, "/admin/upstream-proxy", strings.NewReader(`{"url":"http://127.0.0.1:7890"}`)) + h.UpstreamProxySet(w, r) + if w.Code != http.StatusServiceUnavailable { + t.Fatalf("status = %d, want 503 when not wired", w.Code) + } +} + +func TestUpstreamProxyProbe_Reachable(t *testing.T) { + var got string + h := &Handler{ProbeUpstreamProxyFn: func(url string) (int, int64, error) { got = url; return 401, 42, nil }} + w := httptest.NewRecorder() + r := httptest.NewRequest(http.MethodPost, "/admin/upstream-proxy/probe", strings.NewReader(`{"url":"http://127.0.0.1:7890"}`)) + h.UpstreamProxyProbe(w, r) + if w.Code != http.StatusOK { + t.Fatalf("status = %d, want 200", w.Code) + } + var res upstreamProbeResult + _ = json.Unmarshal(w.Body.Bytes(), &res) + if !res.Ok || res.Status != 401 || res.ElapsedMs != 42 { + t.Fatalf("probe result = %+v, want ok/401/42 (any HTTP status = reachable)", res) + } + if got != "http://127.0.0.1:7890" { + t.Fatalf("probe url = %q", got) + } +} + +func TestUpstreamProxyProbe_Unreachable(t *testing.T) { + h := &Handler{ProbeUpstreamProxyFn: func(string) (int, int64, error) { + return 0, 10, errors.New("dial tcp: connect: connection refused") + }} + w := httptest.NewRecorder() + r := httptest.NewRequest(http.MethodPost, "/admin/upstream-proxy/probe", strings.NewReader(`{"url":"http://127.0.0.1:9999"}`)) + h.UpstreamProxyProbe(w, r) + if w.Code != http.StatusOK { // a failed probe is a 200 with ok=false, not an HTTP error + t.Fatalf("status = %d, want 200 (probe ran, target unreachable)", w.Code) + } + var res upstreamProbeResult + _ = json.Unmarshal(w.Body.Bytes(), &res) + if res.Ok || res.Error == "" { + t.Fatalf("probe result = %+v, want ok=false with error", res) + } +} + +func TestUpstreamProxyProbe_InvalidURLNoCall(t *testing.T) { + called := 0 + h := &Handler{ProbeUpstreamProxyFn: func(string) (int, int64, error) { called++; return 200, 1, nil }} + w := httptest.NewRecorder() + r := httptest.NewRequest(http.MethodPost, "/admin/upstream-proxy/probe", strings.NewReader(`{"url":"ftp://nope"}`)) + h.UpstreamProxyProbe(w, r) + if w.Code != http.StatusBadRequest { + t.Fatalf("status = %d, want 400 for bad scheme", w.Code) + } + if called != 0 { + t.Fatalf("invalid URL must NOT reach the probe (called=%d)", called) + } +} + +func TestUpstreamProxyProbe_NotWired(t *testing.T) { + h := &Handler{} + w := httptest.NewRecorder() + r := httptest.NewRequest(http.MethodPost, "/admin/upstream-proxy/probe", strings.NewReader(`{"url":"http://127.0.0.1:7890"}`)) + h.UpstreamProxyProbe(w, r) + if w.Code != http.StatusServiceUnavailable { + t.Fatalf("status = %d, want 503 when not wired", w.Code) + } +} + +// 2026-07-08 layered egress state (`aikey env`): when EgressStateFn is wired +// the GET response carries the full decision chain; when it isn't, the body +// must stay the legacy {"url"} shape (no "egress" key) so older consumers and +// partially-wired hosts are unaffected. +func TestUpstreamProxyGet_EgressStateWired(t *testing.T) { + h := &Handler{ + GetUpstreamProxyFn: func() string { return "" }, + EgressStateFn: func() EgressState { + return EgressState{ + SystemSupported: true, + SystemHTTPS: "http://127.0.0.1:7890", + EffectiveSource: "system", + EffectiveURL: "http://127.0.0.1:7890", + } + }, + } + w := httptest.NewRecorder() + h.UpstreamProxyGet(w, httptest.NewRequest(http.MethodGet, "/admin/upstream-proxy", nil)) + if w.Code != http.StatusOK { + t.Fatalf("status = %d, want 200", w.Code) + } + var body upstreamProxyBody + if err := json.Unmarshal(w.Body.Bytes(), &body); err != nil { + t.Fatalf("unmarshal: %v", err) + } + if body.Egress == nil { + t.Fatal("egress block missing") + } + if body.Egress.EffectiveSource != "system" || body.Egress.EffectiveURL != "http://127.0.0.1:7890" { + t.Fatalf("effective = %q %q, want system http://127.0.0.1:7890", + body.Egress.EffectiveSource, body.Egress.EffectiveURL) + } +} + +func TestUpstreamProxyGet_NoEgressFnKeepsLegacyShape(t *testing.T) { + h := &Handler{GetUpstreamProxyFn: func() string { return "http://127.0.0.1:7890" }} + w := httptest.NewRecorder() + h.UpstreamProxyGet(w, httptest.NewRequest(http.MethodGet, "/admin/upstream-proxy", nil)) + if strings.Contains(w.Body.String(), "egress") { + t.Fatalf("legacy shape must not contain an egress key, got %s", w.Body.String()) + } +} diff --git a/internal/apphook/childhook.go b/internal/apphook/childhook.go index dcefacb..5d8d8bf 100644 --- a/internal/apphook/childhook.go +++ b/internal/apphook/childhook.go @@ -33,6 +33,8 @@ import ( "sync" "sync/atomic" "time" + + "github.com/AiKeyLabs/pkg/aikeycompat" ) // ChildHookConfig configures a ChildHook. @@ -144,6 +146,11 @@ func (h *ChildHook) spawnLocked(ctx context.Context) error { } cmd := exec.Command(h.cfg.BinaryPath, h.cfg.BinaryArgs...) //nolint:gosec // operator-configured app-hook binary, path stat-verified above; args from trusted vault/config, never request input + // Never flash a console window on Windows: the proxy runs console-less, + // so a console-subsystem hook child would otherwise pop a visible + // terminal window for its whole lifetime (same class as the 2026-07-07 + // web-bridge window-flash bug). No-op on Unix; stdio pipes unaffected. + aikeycompat.HideSpawnConsole(cmd) // Inherit the proxy's env (the child relies on it for AIKEY_* config) and // append any per-app ExtraEnv the supervisor derived from vault. if len(h.cfg.ExtraEnv) > 0 { diff --git a/internal/cluster/register.go b/internal/cluster/register.go index 66a35d1..57d1768 100644 --- a/internal/cluster/register.go +++ b/internal/cluster/register.go @@ -15,6 +15,8 @@ import ( "net/http" "strings" "time" + + "github.com/AiKeyLabs/aikey-proxy/internal/httpx" ) const defaultHeartbeatInterval = 5 * time.Second @@ -25,7 +27,7 @@ var errUnknownNode = fmt.Errorf("hub does not know this node") // Registrar keeps this proxy node registered + alive in the hub name service. type Registrar struct { - client *http.Client + client *httpx.SwappableClient // control-plane→hub: rebuilt on host network change (self-heal registry) // healthFn, when set, supplies the optional `health` heartbeat field // (P0-B): node-local health (daemon sync status + proxy metrics) rides // the existing heartbeat so it becomes externally visible with zero new @@ -51,7 +53,7 @@ func NewRegistrar(hubURL, nodeID, nodeAddr string, weight int, serviceToken stri nodeAddr: nodeAddr, weight: weight, serviceToken: serviceToken, - client: &http.Client{Timeout: 5 * time.Second}, + client: httpx.NewSwappableDirect(5 * time.Second), interval: defaultHeartbeatInterval, } } @@ -123,7 +125,7 @@ func (r *Registrar) post(ctx context.Context, path string, body []byte) (*http.R if r.serviceToken != "" { req.Header.Set("Authorization", "Bearer "+r.serviceToken) } - return r.client.Do(req) + return r.client.Get().Do(req) } // Run registers, then heartbeats on the hub-provided interval until ctx is done. diff --git a/internal/config/config.go b/internal/config/config.go index 6e4b0c1..4f0a082 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -41,6 +41,24 @@ type Config struct { Observers map[string]map[string]any `yaml:"observers,omitempty"` Vault VaultConfig `yaml:"vault"` UpstreamProxy UpstreamProxyConfig `yaml:"upstream_proxy"` + // ConsoleURL is the base URL of THIS machine's local AiKey console + // (aikey-local-server web, e.g. "http://127.0.0.1:8090"). Used to assemble + // user-facing login links in structured errors (OAUTH_GROUP_MEMBER_LOGIN_ + // REQUIRED → /user/team-oauth) so the message shown inside claude / + // codex names a clickable next step. Single assembly point for login_url — + // do NOT re-derive this URL in CLI/wrapper code (see + // update/20260703-OAuth组成员登录提示-CLI显示与login_url.md 决策2). + // + // *string ON PURPOSE (absent ≠ explicit-empty): + // - key ABSENT (nil): pre-20260703 personal/trial configs are PRESERVED + // on upgrade (never re-rendered), so the whole existing install base + // would silently lose the login URL. Absent ⇒ DefaultConsoleURL, which + // is correct on every personal/trial box (console port 8090). + // - key EXPLICITLY "" : deployments with no co-installed console + // (server profile render, cluster-node heredoc) opt out ⇒ URL-less + // fallback wording. Do not "simplify" this to a plain string — that + // erases the upgrade-path default. + ConsoleURL *string `yaml:"console_url,omitempty"` // Cluster is the V3c cluster-node config block. Absent / Enabled=false ⇒ // standard local proxy (Personal/Trial) — zero behavior change. When // Enabled, this proxy is a cluster node: it registers + heartbeats to the @@ -192,7 +210,10 @@ type LogConfig struct { // UpstreamProxyConfig configures the outbound proxy used when connecting to AI providers. // Supports HTTP, HTTPS, and SOCKS5 proxy URLs. -// If empty, the standard HTTP_PROXY / HTTPS_PROXY / NO_PROXY environment variables are used. +// If empty: the standard HTTP_PROXY / HTTPS_PROXY / NO_PROXY environment variables +// apply when set; otherwise the OS system proxy (macOS/Windows) is auto-detected and +// followed LIVE (2026-07-08, internal/sysproxy) — a Clash port change / toggle takes +// effect without restarting the daemon. type UpstreamProxyConfig struct { // URL is the proxy endpoint, e.g.: // http://127.0.0.1:7890 (Clash HTTP mode) @@ -332,6 +353,22 @@ func (c *Config) applyDefaults() { c.Providers[name] = p } } + // nil (key absent — pre-20260703 preserved configs) → default console. + // Explicit "" (server / cluster) is honored as opt-out. See the + // ConsoleURL field doc for why this must stay absent-vs-empty aware. + if c.ConsoleURL == nil { + v := DefaultConsoleURL + c.ConsoleURL = &v + } +} + +// ResolvedConsoleURL returns the effective console base ("" = no co-installed +// console → URL-less error wording). Safe before applyDefaults (nil → default). +func (c *Config) ResolvedConsoleURL() string { + if c.ConsoleURL == nil { + return DefaultConsoleURL + } + return *c.ConsoleURL } func (c *Config) validate() error { diff --git a/internal/config/config_test.go b/internal/config/config_test.go index dd65761..25ebc7d 100644 --- a/internal/config/config_test.go +++ b/internal/config/config_test.go @@ -174,6 +174,41 @@ func writeTestPair(t *testing.T, system, user string) string { // Without aikey-user.yaml the system value passes through unchanged — // pre-login state for fresh Personal installs must keep working. +// console_url absent-vs-empty contract (20260703 OAuth组成员登录提示): +// ABSENT key = pre-20260703 config preserved across an upgrade → must default +// (the whole existing personal/trial install base gets the login URL without a +// config migration). EXPLICIT "" = server/cluster opt-out → must stay empty. +// A plain-string "simplification" would silently break the upgrade default. +func TestLoad_ConsoleURLAbsentDefaultsExplicitEmptyHonored(t *testing.T) { + // systemProxyYaml predates console_url — the upgrade-preserved shape. + sysPath := writeTestPair(t, systemProxyYaml, "") + cfg, err := Load(sysPath) + if err != nil { + t.Fatalf("Load: %v", err) + } + if got := cfg.ResolvedConsoleURL(); got != DefaultConsoleURL { + t.Fatalf("absent console_url must default to %q, got %q", DefaultConsoleURL, got) + } + + sysPath2 := writeTestPair(t, systemProxyYaml+"\nconsole_url: \"\"\n", "") + cfg2, err := Load(sysPath2) + if err != nil { + t.Fatalf("Load explicit-empty: %v", err) + } + if got := cfg2.ResolvedConsoleURL(); got != "" { + t.Fatalf("explicit-empty console_url must stay empty (opt-out), got %q", got) + } + + sysPath3 := writeTestPair(t, systemProxyYaml+"\nconsole_url: \"http://127.0.0.1:9191\"\n", "") + cfg3, err := Load(sysPath3) + if err != nil { + t.Fatalf("Load explicit-value: %v", err) + } + if got := cfg3.ResolvedConsoleURL(); got != "http://127.0.0.1:9191" { + t.Fatalf("explicit console_url must be honored, got %q", got) + } +} + func TestLoad_NoUserYamlPreservesSystemValues(t *testing.T) { sysPath := writeTestPair(t, systemProxyYaml, "") cfg, err := Load(sysPath) diff --git a/internal/config/defaults.go b/internal/config/defaults.go index f7e4a34..cb04103 100644 --- a/internal/config/defaults.go +++ b/internal/config/defaults.go @@ -21,4 +21,12 @@ const ( // DefaultPortDriftMax bounds runtime port drift when the configured // listen port is occupied. See 20260430-端口偏移能力修复.md. DefaultPortDriftMax = 10 + // DefaultConsoleURL applies when console_url is ABSENT from yaml (pre- + // 20260703 personal/trial configs preserved across upgrades). 8090 is the + // local-server user-console port on every default personal/trial install + // (allocator "trial" key at offset 0). Explicit "" opts out — see + // Config.ConsoleURL. Known edge: a drifted console port with a preserved + // old config yields a 404-ing default URL; the message still names the + // /user/team-oauth page so the member can find it manually. + DefaultConsoleURL = "http://127.0.0.1:8090" ) diff --git a/internal/config/upstream_proxy.go b/internal/config/upstream_proxy.go new file mode 100644 index 0000000..be756b3 --- /dev/null +++ b/internal/config/upstream_proxy.go @@ -0,0 +1,131 @@ +// upstream_proxy.go — runtime read/modify/write of the egress (upstream) proxy URL. +// +// WHY this lives here (2026-06-30, R25 出口收敛 follow-up): the OAuth/AI egress +// proxy used to be a master-side system_settings value; the egress-convergence +// refactor moved that responsibility to the PROXY NODE (cfg.UpstreamProxy.URL). +// The local web "Settings → Upstream proxy" card lets a user set it without hand- +// editing yaml; the admin endpoint persists it HERE and hot-swaps the live +// transport + impersonate client (no restart). Decision record: user picked +// hot-swap (B) + plaintext storage (i). +// +// Storage = aikey-user.yaml's `proxy.upstream_proxy.url` (the USER layer). The +// proxy's Load() already deep-merges the user file's `proxy:` subtree onto the +// system config's top level (see config.go loadAndMaybeMerge + configmerge), so a +// value written here survives `make restart-personal` re-renders of the system +// yaml. We never touch the system aikey-proxy.yaml (it's template-rendered). +package config + +import ( + "fmt" + "net/url" + "os" + "path/filepath" + "strings" + + "gopkg.in/yaml.v3" +) + +// ValidateUpstreamProxyURL accepts "" (clear → direct egress) or an http/https/ +// socks5 URL with a host:port. Mirrors the (now-removed) master systemsettings +// validator so the contract is identical across the egress-convergence move. +func ValidateUpstreamProxyURL(raw string) error { + if strings.TrimSpace(raw) == "" { + return nil // empty = clear the proxy (direct egress) + } + u, err := url.Parse(raw) + if err != nil { + return fmt.Errorf("not a valid URL: %w", err) + } + switch u.Scheme { + case "http", "https", "socks5": + default: + return fmt.Errorf("scheme must be http/https/socks5, got %q", u.Scheme) + } + if u.Hostname() == "" || u.Port() == "" { + return fmt.Errorf("host:port required (e.g. http://127.0.0.1:7890)") + } + return nil +} + +// userConfigPathFor returns the aikey-user.yaml path next to the system config. +func userConfigPathFor(systemConfigPath string) string { + return filepath.Join(filepath.Dir(systemConfigPath), "aikey-user.yaml") +} + +// PersistUpstreamProxyURL writes url into aikey-user.yaml's `proxy.upstream_proxy.url` +// via read-modify-write, preserving every other field the CLI owns there +// (collector_routes, collector_credentials, …). An empty url is written as an +// explicit "" so the user layer overrides any system-template default with "direct +// egress" (deep-merge keeps the empty string as a value, not an absence). +// +// systemConfigPath is the path the proxy was loaded from (--config); the user file +// is resolved next to it. The write is atomic (temp + rename) with 0600 perms. +func PersistUpstreamProxyURL(systemConfigPath, rawURL string) error { + userPath := userConfigPathFor(systemConfigPath) + + // Load existing user file (absent ⇒ start empty; the personal-no-team case). + root := map[string]any{} + if b, err := os.ReadFile(userPath); err == nil { + if err := yaml.Unmarshal(b, &root); err != nil { + return fmt.Errorf("parse %s: %w", userPath, err) + } + if root == nil { + root = map[string]any{} + } + } else if !os.IsNotExist(err) { + return fmt.Errorf("read %s: %w", userPath, err) + } + + // Navigate/create proxy.upstream_proxy and set url, preserving siblings. + proxySec := asStringMap(root["proxy"]) + upstream := asStringMap(proxySec["upstream_proxy"]) + upstream["url"] = rawURL + proxySec["upstream_proxy"] = upstream + root["proxy"] = proxySec + + out, err := yaml.Marshal(root) + if err != nil { + return fmt.Errorf("marshal user config: %w", err) + } + return atomicWrite(userPath, out) +} + +// asStringMap coerces a yaml-decoded value into a map[string]any we can extend. +// yaml.v3 decodes mappings into map[string]any already; a nil / wrong-typed value +// (e.g. the key was absent or held a scalar) is replaced with a fresh map so the +// caller can always set into it. +func asStringMap(v any) map[string]any { + if m, ok := v.(map[string]any); ok && m != nil { + return m + } + return map[string]any{} +} + +// atomicWrite writes data to a temp file in the target dir then renames it over +// path, so a crash mid-write can't leave a truncated config. 0600 because the URL +// may carry credentials (socks5://user:pass@host) — same posture as the rest of +// aikey-user.yaml. +func atomicWrite(path string, data []byte) error { + dir := filepath.Dir(path) + tmp, err := os.CreateTemp(dir, ".aikey-user-*.yaml.tmp") + if err != nil { + return fmt.Errorf("create temp config: %w", err) + } + tmpName := tmp.Name() + defer os.Remove(tmpName) // no-op after a successful rename + if err := tmp.Chmod(0o600); err != nil { + tmp.Close() + return fmt.Errorf("chmod temp config: %w", err) + } + if _, err := tmp.Write(data); err != nil { + tmp.Close() + return fmt.Errorf("write temp config: %w", err) + } + if err := tmp.Close(); err != nil { + return fmt.Errorf("close temp config: %w", err) + } + if err := os.Rename(tmpName, path); err != nil { + return fmt.Errorf("rename temp config into place: %w", err) + } + return nil +} diff --git a/internal/config/upstream_proxy_test.go b/internal/config/upstream_proxy_test.go new file mode 100644 index 0000000..0b75d91 --- /dev/null +++ b/internal/config/upstream_proxy_test.go @@ -0,0 +1,114 @@ +package config + +import ( + "os" + "path/filepath" + "testing" + + "gopkg.in/yaml.v3" +) + +func TestValidateUpstreamProxyURL(t *testing.T) { + cases := []struct { + name string + url string + wantErr bool + }{ + {"empty clears", "", false}, + {"http ok", "http://127.0.0.1:7890", false}, + {"https ok", "https://proxy.example.com:8443", false}, + {"socks5 ok", "socks5://127.0.0.1:7891", false}, + {"socks5 with creds ok", "socks5://user:pass@127.0.0.1:7891", false}, + {"bad scheme", "ftp://127.0.0.1:21", true}, + {"missing port", "http://127.0.0.1", true}, + {"missing host", "http://:7890", true}, + {"garbage", "://nonsense", true}, + } + for _, c := range cases { + t.Run(c.name, func(t *testing.T) { + err := ValidateUpstreamProxyURL(c.url) + if c.wantErr && err == nil { + t.Fatalf("ValidateUpstreamProxyURL(%q) = nil, want error", c.url) + } + if !c.wantErr && err != nil { + t.Fatalf("ValidateUpstreamProxyURL(%q) = %v, want nil", c.url, err) + } + }) + } +} + +// TestPersistUpstreamProxyURL_RoundTripThroughLoad: persisting the URL must land in +// aikey-user.yaml AND be read back by the real Load() merge as cfg.UpstreamProxy.URL. +// End-to-end so a future change to either the writer or the merge can't silently +// break "Settings → Upstream proxy". +func TestPersistUpstreamProxyURL_RoundTripThroughLoad(t *testing.T) { + sysPath := writeTestPair(t, systemProxyYaml, "") // no user file yet + if err := PersistUpstreamProxyURL(sysPath, "http://127.0.0.1:7890"); err != nil { + t.Fatalf("PersistUpstreamProxyURL: %v", err) + } + cfg, err := Load(sysPath) + if err != nil { + t.Fatalf("Load: %v", err) + } + if cfg.UpstreamProxy.URL != "http://127.0.0.1:7890" { + t.Fatalf("UpstreamProxy.URL = %q, want http://127.0.0.1:7890", cfg.UpstreamProxy.URL) + } + // The user file must now exist (created when absent). + if _, err := os.Stat(filepath.Join(filepath.Dir(sysPath), "aikey-user.yaml")); err != nil { + t.Fatalf("aikey-user.yaml should have been created: %v", err) + } +} + +// TestPersistUpstreamProxyURL_PreservesSiblings: the read-modify-write must keep the +// CLI-owned proxy.* fields (e.g. an events.collector_routes override written by +// `aikey login --control-url`) intact — we only touch upstream_proxy.url. Re-read the +// user file directly (not via Load) to assert exactly what was persisted. +func TestPersistUpstreamProxyURL_PreservesSiblings(t *testing.T) { + userYaml := ` +proxy: + events: + collector_routes: + team: "http://team.example.com:3000" +` + sysPath := writeTestPair(t, systemProxyYaml, userYaml) + if err := PersistUpstreamProxyURL(sysPath, "socks5://127.0.0.1:7891"); err != nil { + t.Fatalf("PersistUpstreamProxyURL: %v", err) + } + + b, err := os.ReadFile(filepath.Join(filepath.Dir(sysPath), "aikey-user.yaml")) + if err != nil { + t.Fatalf("read user file: %v", err) + } + var root map[string]any + if err := yaml.Unmarshal(b, &root); err != nil { + t.Fatalf("parse user file: %v", err) + } + proxy, _ := root["proxy"].(map[string]any) + up, _ := proxy["upstream_proxy"].(map[string]any) + if up["url"] != "socks5://127.0.0.1:7891" { + t.Fatalf("upstream_proxy.url not persisted: %v", up["url"]) + } + events, _ := proxy["events"].(map[string]any) + cr, _ := events["collector_routes"].(map[string]any) + if cr["team"] != "http://team.example.com:3000" { + t.Fatalf("events.collector_routes.team sibling clobbered: %v", cr["team"]) + } +} + +// TestPersistUpstreamProxyURL_EmptyClears: writing "" overrides any system default +// with direct egress (empty string is a value the merge keeps, not an absence). +func TestPersistUpstreamProxyURL_EmptyClears(t *testing.T) { + // System yaml that ships a default upstream_proxy, to prove the user "" wins. + sysWithProxy := systemProxyYaml + "\nupstream_proxy:\n url: \"http://10.0.0.1:1080\"\n" + sysPath := writeTestPair(t, sysWithProxy, "") + if err := PersistUpstreamProxyURL(sysPath, ""); err != nil { + t.Fatalf("PersistUpstreamProxyURL: %v", err) + } + cfg, err := Load(sysPath) + if err != nil { + t.Fatalf("Load: %v", err) + } + if cfg.UpstreamProxy.URL != "" { + t.Fatalf("empty persist should clear egress, got %q", cfg.UpstreamProxy.URL) + } +} diff --git a/internal/events/canary.go b/internal/events/canary.go index 67eb241..fded6d6 100644 --- a/internal/events/canary.go +++ b/internal/events/canary.go @@ -9,6 +9,7 @@ import ( "sync" "time" + "github.com/AiKeyLabs/aikey-proxy/internal/httpx" "github.com/AiKeyLabs/aikey-proxy/internal/observability" "github.com/AiKeyLabs/pkg/aikeytime" ) @@ -59,7 +60,7 @@ type CanaryResult struct { // CanaryProbe sends periodic synthetic events through the pipeline and verifies arrival. type CanaryProbe struct { reporter *Reporter - client *http.Client + client *httpx.SwappableClient // control-plane→diagnostics: rebuilt on host network change (self-heal registry) done chan struct{} cfg CanaryConfig lastResult CanaryResult @@ -98,7 +99,7 @@ func NewCanaryProbe(reporter *Reporter, cfg CanaryConfig) *CanaryProbe { p := &CanaryProbe{ reporter: reporter, cfg: cfg, - client: &http.Client{Timeout: 10 * time.Second}, + client: httpx.NewSwappableDirect(10 * time.Second), done: make(chan struct{}), } @@ -284,7 +285,7 @@ func (p *CanaryProbe) checkArrival(eventID string, sentAt time.Time, diagnostics slog.Debug("canary check: request build failed", "url", url, "error", reqErr) return result } - resp, err := p.client.Do(req) + resp, err := p.client.Get().Do(req) if err != nil { // Diagnostics endpoint unreachable — deployment/network issue rather // than a pipeline fault. Mark as "unavailable" (not "failed") so the @@ -352,7 +353,7 @@ func (p *CanaryProbe) checkArrival(eventID string, sentAt time.Time, diagnostics slog.Debug("canary check: query request build failed", "url", queryURL, "error", qReqErr) return result } - qResp, qErr := p.client.Do(qReq) + qResp, qErr := p.client.Get().Do(qReq) if qErr != nil || qResp.StatusCode != http.StatusOK { // Query-service unreachable or error. Don't demote the overall // status to "failed" — collector side is fine. Report "partial" diff --git a/internal/events/canary_test.go b/internal/events/canary_test.go index 33ec0ac..a68b54e 100644 --- a/internal/events/canary_test.go +++ b/internal/events/canary_test.go @@ -6,6 +6,8 @@ import ( "sync" "testing" "time" + + "github.com/AiKeyLabs/aikey-proxy/internal/httpx" ) // newTestReporter builds a real Reporter backed by a temp WAL dir. The canary @@ -35,7 +37,7 @@ func newProbeNoLoop(reporter *Reporter, diagnosticsURL string) *CanaryProbe { Interval: time.Hour, CheckDelay: time.Millisecond, // don't wait 15s in tests }, - client: &http.Client{Timeout: 2 * time.Second}, + client: httpx.NewSwappableFixed(&http.Client{Timeout: 2 * time.Second}), done: make(chan struct{}), } } diff --git a/internal/events/compliance_upload.go b/internal/events/compliance_upload.go index 660ffb9..dec8ca0 100644 --- a/internal/events/compliance_upload.go +++ b/internal/events/compliance_upload.go @@ -54,7 +54,7 @@ func (r *Reporter) UploadComplianceEvents(ctx context.Context, routeSource strin } } - resp, err := r.client.Do(req) + resp, err := r.client.Get().Do(req) if err != nil { return fmt.Errorf("compliance upload: POST %s: %w", url, err) } diff --git a/internal/events/content_reporter.go b/internal/events/content_reporter.go index 387502a..9b62e12 100644 --- a/internal/events/content_reporter.go +++ b/internal/events/content_reporter.go @@ -11,6 +11,7 @@ import ( "sync" "time" + "github.com/AiKeyLabs/aikey-proxy/internal/httpx" "github.com/AiKeyLabs/aikey-proxy/internal/observability" ) @@ -65,14 +66,18 @@ type contentBatchResponse struct { // (org,event_id) dedup absorbs the overlap). Terminal (4xx) drops the batch // (advances sentSeq, no retry) — the gap is server-detectable via source_seq. type ContentReporter struct { - nextUploadAttempt time.Time - lastOKUploadAt time.Time - wal *ContentWAL - sentSeq map[string]int64 - confirmedSeq map[string]int64 - signal chan struct{} - done chan struct{} - cfg ContentReporterConfig + nextUploadAttempt time.Time + lastOKUploadAt time.Time + wal *ContentWAL + sentSeq map[string]int64 + confirmedSeq map[string]int64 + signal chan struct{} + done chan struct{} + cfg ContentReporterConfig + // client is the swappable control-plane→collector client. An injected + // cfg.HTTPClient is wrapped fixed (not globally rebuilt); the default is + // registered so a host network change rebuilds it (self-heal registry). + client *httpx.SwappableClient wg sync.WaitGroup consecutiveFailures int mu sync.Mutex @@ -88,11 +93,17 @@ func NewContentReporter(in *ContentReporterConfig, wal *ContentWAL) *ContentRepo if cfg.BatchSize <= 0 { cfg.BatchSize = 100 } - if cfg.HTTPClient == nil { - cfg.HTTPClient = &http.Client{Timeout: 30 * time.Second} + // Registered swappable for the default client (rebuilt on network change); + // an injected client is honored as-is (fixed, not registered). + var client *httpx.SwappableClient + if cfg.HTTPClient != nil { + client = httpx.NewSwappableFixed(cfg.HTTPClient) + } else { + client = httpx.NewSwappableDirect(30 * time.Second) } return &ContentReporter{ cfg: cfg, + client: client, wal: wal, sentSeq: make(map[string]int64), confirmedSeq: make(map[string]int64), @@ -332,7 +343,7 @@ func (r *ContentReporter) doUpload(body []byte) (*contentBatchResponse, int, err } else if r.cfg.CollectorToken != "" { httpReq.Header.Set("Authorization", "Bearer "+r.cfg.CollectorToken) } - httpResp, err := r.cfg.HTTPClient.Do(httpReq) + httpResp, err := r.client.Get().Do(httpReq) if err != nil { return nil, 0, err // transport → retryable } diff --git a/internal/events/credential.go b/internal/events/credential.go index d4e22a5..ac3b9e7 100644 --- a/internal/events/credential.go +++ b/internal/events/credential.go @@ -11,6 +11,8 @@ import ( "strings" "sync" "time" + + "github.com/AiKeyLabs/aikey-proxy/internal/httpx" ) // Credential is the bearer-token source for a single collector upload. @@ -200,7 +202,7 @@ func (j *RefreshableJWT) refresh(ctx context.Context) (string, time.Time, error) client := j.HTTPClient if client == nil { - client = defaultRefreshClient + client = defaultRefreshClient.Get() } resp, err := client.Do(req) if err != nil { @@ -237,6 +239,4 @@ func (j *RefreshableJWT) refresh(ctx context.Context) (string, time.Time, error) // hot path of every upload right at the edge of the access token's // validity window — a stuck refresh would block uploads. 10s is the // same shape as the reporter's batch-upload client. -var defaultRefreshClient = &http.Client{ - Timeout: 10 * time.Second, -} +var defaultRefreshClient = httpx.NewSwappableDirect(10 * time.Second) diff --git a/internal/events/preflight.go b/internal/events/preflight.go index 236c39c..c37986f 100644 --- a/internal/events/preflight.go +++ b/internal/events/preflight.go @@ -9,6 +9,8 @@ import ( "net/http" "os" "time" + + "github.com/AiKeyLabs/aikey-proxy/internal/httpx" ) // PreflightResult holds the outcome of usage pipeline readiness checks. @@ -78,7 +80,7 @@ func RunPreflight(collectorURL, collectorToken, walDir, eventsDBPath string) Pre // checkCollector tests collector reachability and token auth by sending an empty batch. func checkCollector(collectorURL, token string) (reachable, authOK bool, errMsg string) { - client := &http.Client{Timeout: 5 * time.Second} + client := httpx.NewDirectClient(5 * time.Second) url := collectorURL + "/v1/usage-events:batch" body := []byte(`{"source":"preflight","events":[]}`) diff --git a/internal/events/reconcile.go b/internal/events/reconcile.go index 6af0b13..03df396 100644 --- a/internal/events/reconcile.go +++ b/internal/events/reconcile.go @@ -241,7 +241,7 @@ func (r *Reporter) httpGetJSON(ctx context.Context, u string, out any) error { if err != nil { return err } - resp, err := r.client.Do(req) + resp, err := r.client.Get().Do(req) if err != nil { return err } @@ -262,7 +262,7 @@ func (r *Reporter) httpPostJSON(ctx context.Context, u string, body, out any) er return err } req.Header.Set("Content-Type", "application/json") - resp, err := r.client.Do(req) + resp, err := r.client.Get().Do(req) if err != nil { return err } diff --git a/internal/events/reporter.go b/internal/events/reporter.go index 175a8f5..ee040f4 100644 --- a/internal/events/reporter.go +++ b/internal/events/reporter.go @@ -13,6 +13,7 @@ import ( "sync/atomic" "time" + "github.com/AiKeyLabs/aikey-proxy/internal/httpx" "github.com/AiKeyLabs/aikey-proxy/internal/observability" "github.com/AiKeyLabs/pkg/aikeycompat" "github.com/AiKeyLabs/pkg/aikeytime" @@ -120,6 +121,15 @@ type batchResponse struct { type Reporter struct { lastCanaryEventAt time.Time lastErrorAt time.Time + // lastAutoReplayAt gates the dead-letter AUTO-replay (2026-07-04 self-heal): + // a successful upload proves the pipe recovered, so entries dead-lettered + // during the outage are replayed automatically — at most once per + // autoReplayCooldown, guarded by autoReplayInFlight. Manual replay + // (POST /admin/replay-dead-letter) stays as the emergency channel. + // Idempotency: collector ingest dedups by event_id (INSERT OR IGNORE / + // ON CONFLICT DO NOTHING), so re-uploading already-accepted entries is safe. + lastAutoReplayAt time.Time + autoReplayInFlight atomic.Bool // nextUploadAttempt is the non-blocking backoff gate (B', 2026-06-09 缺口2): // after a drain pass hits retryable upload failures, drainOnce sets this // ahead and skips upload attempts until it passes — replacing the old in-line @@ -131,13 +141,13 @@ type Reporter struct { wal *WALWriter done chan struct{} // delivery-integrity cursors (memory only; see type doc). Guarded by mu. - sentSeq map[string]int64 // source_id → highest seq handed to upload - confirmedSeq map[string]int64 // source_id → server contiguous high-water - seenV1 map[string]bool // event_id → uploaded; A1 one-shot for v1 entries - signal chan struct{} // cap-1 wakeup poke from Report → uploadLoop - dlw *deadLetterWriter // dead letter writer for terminal failures - client *http.Client - lastUploadStatus string // "ok" | "retryable_failed" | "terminal_failed" + sentSeq map[string]int64 // source_id → highest seq handed to upload + confirmedSeq map[string]int64 // source_id → server contiguous high-water + seenV1 map[string]bool // event_id → uploaded; A1 one-shot for v1 entries + signal chan struct{} // cap-1 wakeup poke from Report → uploadLoop + dlw *deadLetterWriter // dead letter writer for terminal failures + client *httpx.SwappableClient // control-plane→collector: rebuilt on host network change (self-heal registry) + lastUploadStatus string // "ok" | "retryable_failed" | "terminal_failed" cfg ReporterConfig wg sync.WaitGroup consecutiveFailures int @@ -211,7 +221,7 @@ func NewReporter(in *ReporterConfig) (*Reporter, error) { dlw: dlw, signal: make(chan struct{}, 1), done: make(chan struct{}), - client: &http.Client{Timeout: 30 * time.Second}, + client: httpx.NewSwappableDirect(30 * time.Second), sentSeq: make(map[string]int64), confirmedSeq: make(map[string]int64), seenV1: make(map[string]bool), @@ -821,6 +831,47 @@ func (r *Reporter) onUploadSuccess(count int) { "total", total, "recovered", wasRecovery) } + r.maybeAutoReplayLocked() +} + +// autoReplayCooldown spaces automatic dead-letter replays so a flapping +// upstream can't turn the replay pass into a hot loop. 10 min ≈ one recovery +// window; entries left behind are retried on the next trigger. +const autoReplayCooldown = 10 * time.Minute + +// maybeAutoReplayLocked (r.mu held) fires the self-healing dead-letter replay +// when a successful upload shows the pipe is healthy again. Non-blocking: the +// replay itself runs on an isolated goroutine — a panic or slow pass never +// touches the upload hot path. +func (r *Reporter) maybeAutoReplayLocked() { + if r.dlw == nil || time.Since(r.lastAutoReplayAt) < autoReplayCooldown { + return + } + // Cheap gate: only bother when the dead-letter file has content. Stat runs + // at most once per cooldown window. + r.lastAutoReplayAt = time.Now() + if fi, err := os.Stat(r.dlw.path); err != nil || fi.Size() == 0 { + return + } + if !r.autoReplayInFlight.CompareAndSwap(false, true) { + return + } + observability.GoSafe("events.reporter.auto_replay", observability.Isolated, func() { + defer r.autoReplayInFlight.Store(false) + res, err := r.ReplayDeadLetter(context.Background()) + if err != nil { + slog.Warn("reporter: automatic dead-letter replay failed — entries kept for the next window", + "event.name", observability.EventReporterDeadLetterReplayed, + "error", err.Error()) + return + } + if res.EntriesScanned == 0 { + return + } + slog.Info("reporter: automatic dead-letter replay finished", + "event.name", observability.EventReporterDeadLetterReplayed, + "scanned", res.EntriesScanned, "events_replayed", res.EventsReplayedOK, "still_failing", res.EntriesStillFailing) + }) } // onUploadFail updates delivery state after a failed upload. @@ -875,7 +926,7 @@ func (r *Reporter) doUpload(url string, body []byte, cred Credential) (*batchRes } } - resp, err := r.client.Do(httpReq) + resp, err := r.client.Get().Do(httpReq) if err != nil { return nil, &uploadError{Err: fmt.Errorf("http: %w", err)} } diff --git a/internal/events/reporter_test.go b/internal/events/reporter_test.go index 81de2f0..422c6d3 100644 --- a/internal/events/reporter_test.go +++ b/internal/events/reporter_test.go @@ -830,3 +830,90 @@ func TestReporter_ReplayDeadLetter_MalformedLineSkipped(t *testing.T) { t.Errorf("re-delivered entry should be removed; file:\n%s", string(data)) } } + +// ── automatic dead-letter replay (2026-07-04 self-heal) ───────────────────── +// +// A successful upload proves the pipe recovered, so maybeAutoReplayLocked must +// fire the replay (once per cooldown window) with NO manual step. 能红: drop +// the maybeAutoReplayLocked call from onUploadSuccess and the "auto" test +// fails; drop the cooldown gate and the "cooldown" test fails. +func TestReporter_AutoReplay_FiresAfterUploadSuccess(t *testing.T) { + var accepted atomic.Int64 + var auth atomic.Value + auth.Store("") + srv := recordingServer(t, &accepted, &auth) + + dir := t.TempDir() + reporter, err := NewReporter(&ReporterConfig{ + CollectorURL: srv.URL, CollectorToken: "tok", + QueueCapacity: 10, BatchSize: 5, UploadInterval: 50 * time.Millisecond, + WALDir: dir, DBPath: filepath.Join(dir, "events.db"), + }) + if err != nil { + t.Fatal(err) + } + defer reporter.Close() + + seedDeadLetter(t, dir, deadLetterEntry{ + DeadAt: aikeytime.Now(), Reason: "terminal", ErrorCode: 401, ErrorMsg: "outage", + Events: []ReportableEvent{{EventID: "ev-dl", OrgID: "o", RouteSource: "team", EventTime: aikeytime.Now(), OccurredAt: aikeytime.Now(), RequestStatus: "success", RequestCount: 1}}, + EventIDs: []string{"ev-dl"}, + }) + + // A successful upload lands → the auto-replay trigger runs asynchronously. + reporter.onUploadSuccess(1) + deadline := time.Now().Add(3 * time.Second) + for time.Now().Before(deadline) { + if len(readDeadLetterEntries(t, dir)) == 0 { + break + } + time.Sleep(20 * time.Millisecond) + } + if left := readDeadLetterEntries(t, dir); len(left) != 0 { + t.Fatalf("auto replay must drain the dead-letter file, %d entries left", len(left)) + } + if accepted.Load() == 0 { + t.Fatal("replayed events must reach the collector") + } +} + +// Within the cooldown window a second success must NOT re-stat/re-fire; after +// the window it must fire again. Pinned via the lastAutoReplayAt gate. +func TestReporter_AutoReplay_CooldownGate(t *testing.T) { + dir := t.TempDir() + reporter, err := NewReporter(&ReporterConfig{ + CollectorURL: "http://127.0.0.1:0", CollectorToken: "tok", + QueueCapacity: 10, BatchSize: 5, UploadInterval: time.Hour, + WALDir: dir, DBPath: filepath.Join(dir, "events.db"), + }) + if err != nil { + t.Fatal(err) + } + defer reporter.Close() + + // First trigger stamps the gate even when the file is empty (stat gate). + reporter.mu.Lock() + reporter.maybeAutoReplayLocked() + first := reporter.lastAutoReplayAt + reporter.mu.Unlock() + if first.IsZero() { + t.Fatal("first trigger must stamp the cooldown gate") + } + // Second trigger inside the window: stamp unchanged (no re-fire). + reporter.mu.Lock() + reporter.maybeAutoReplayLocked() + second := reporter.lastAutoReplayAt + reporter.mu.Unlock() + if !second.Equal(first) { + t.Fatal("cooldown window must suppress a second trigger") + } + // Simulate window expiry → trigger re-arms. + reporter.mu.Lock() + reporter.lastAutoReplayAt = time.Now().Add(-autoReplayCooldown - time.Second) + reporter.maybeAutoReplayLocked() + third := reporter.lastAutoReplayAt + reporter.mu.Unlock() + if third.Equal(first) || time.Since(third) > time.Minute { + t.Fatal("expired cooldown must re-arm the trigger") + } +} diff --git a/internal/events/store.go b/internal/events/store.go index 8359b28..4959c24 100644 --- a/internal/events/store.go +++ b/internal/events/store.go @@ -6,6 +6,7 @@ import ( "log/slog" "time" + "github.com/AiKeyLabs/aikey-proxy/internal/vault" _ "modernc.org/sqlite" ) @@ -16,7 +17,14 @@ type Store struct { // OpenStore opens (or creates) the events database with WAL mode. func OpenStore(dbPath string) (*Store, error) { - db, err := sql.Open("sqlite", dbPath) + // busy_timeout via DSN (per-connection, not persisted like journal_mode=WAL, + // so it must be on the DSN not an Exec'd PRAGMA). The events DB is multi-writer + // — collector Insert vs RetentionLoop PruneOlderThan on separate pool conns, + // reload building a new generation while the old one still flushes, and the + // cross-process `aikey-proxy wal vacuum` — so a second writer would otherwise + // hit SQLITE_BUSY instantly (WAL serializes writers but doesn't wait). Reuse + // vault.WithBusyTimeoutDSN as the single sqlite busy_timeout convention. + db, err := sql.Open("sqlite", vault.WithBusyTimeoutDSN(dbPath)) if err != nil { return nil, fmt.Errorf("open events db: %w", err) } @@ -80,6 +88,12 @@ func migrate(db *sql.DB) error { // new column lives at end of the table to keep INSERT ordering // of pre-existing columns stable. {"session_id", "ALTER TABLE usage_events ADD COLUMN session_id TEXT DEFAULT ''"}, + // oauth_group account attribution (2026-06-25): the REAL account that + // actually served the request, following pool fallback (A→B). Lets + // local audit/metrics answer "which account served" without querying + // the collector/ODS. Same idempotent ADD-with-pragma-probe pattern; + // end of table to keep pre-existing INSERT ordering stable. + {"account_id", "ALTER TABLE usage_events ADD COLUMN account_id TEXT DEFAULT ''"}, } for _, c := range appCols { var count int @@ -99,6 +113,10 @@ func migrate(db *sql.DB) error { if _, err := db.Exec(`CREATE INDEX IF NOT EXISTS idx_events_app_slug ON usage_events(app_slug) WHERE app_slug != ''`); err != nil { return err } + // Per-account index for "usage by account" (oauth_group attribution view). + if _, err := db.Exec(`CREATE INDEX IF NOT EXISTS idx_events_account ON usage_events(account_id) WHERE account_id != ''`); err != nil { + return err + } return nil } @@ -119,8 +137,8 @@ func (s *Store) Insert(events []UsageEvent) error { (timestamp, virtual_key_id, provider, model, input_tokens, output_tokens, duration_ms, status_code, is_streaming, error_type, request_path, app_slug, app_key_id, app_mode, bound_via, requested_model, resolved_provider, - session_id) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + session_id, account_id) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) `) if err != nil { return fmt.Errorf("prepare insert: %w", err) @@ -152,6 +170,7 @@ func (s *Store) Insert(events []UsageEvent) error { e.RequestedModel, e.ResolvedProvider, e.SessionID, + e.AccountID, ) if err != nil { return fmt.Errorf("insert event: %w", err) @@ -211,6 +230,29 @@ func (s *Store) QueryStats() (byVKey, byProvider map[string]int64, err error) { return byVKey, byProvider, nil } +// QueryByAccount returns request counts grouped by the REAL serving account +// (account_id), excluding rows with no account dimension (legacy single-binding +// paths). This is the local "which account served how much" audit view for +// oauth_group pools — it reflects fallback (a request that switched A→B counts +// toward B). Additive to QueryStats so existing callers/mocks are untouched. +func (s *Store) QueryByAccount() (map[string]int64, error) { + byAccount := make(map[string]int64) + rows, err := s.db.Query("SELECT account_id, COUNT(*) FROM usage_events WHERE account_id != '' GROUP BY account_id") + if err != nil { + return nil, err + } + defer rows.Close() + for rows.Next() { + var acct string + var count int64 + if scanErr := rows.Scan(&acct, &count); scanErr != nil { + return nil, scanErr + } + byAccount[acct] = count + } + return byAccount, rows.Err() +} + // Close closes the database. func (s *Store) Close() error { if s.db != nil { diff --git a/internal/events/types.go b/internal/events/types.go index 170461e..7397d90 100644 --- a/internal/events/types.go +++ b/internal/events/types.go @@ -42,6 +42,14 @@ type UsageEvent struct { AppKeyID string `json:"app_key_id,omitempty"` ErrorType string `json:"error_type,omitempty"` VirtualKeyID string `json:"virtual_key_id"` + // AccountID — the REAL upstream account that actually served this request. + // For oauth_group (pool) routing this follows fallback: if the proxy + // switched A→B (quota/401), this is B (the account that sent upstream), + // NOT the VK's nominal binding — so LOCAL audit can answer "which account + // served this request" without querying the collector/ODS. The user + // attribution (VirtualKeyID) is stable across the switch; only this + // account dimension changes. Empty for legacy single-binding paths. + AccountID string `json:"account_id,omitempty"` // App pipeline audit fields (Phase 4, 主方案 §5.3 主聚合). // Empty / zero for legacy paths. AppSlug string `json:"app_slug,omitempty"` diff --git a/internal/httpx/direct.go b/internal/httpx/direct.go new file mode 100644 index 0000000..d00d07f --- /dev/null +++ b/internal/httpx/direct.go @@ -0,0 +1,32 @@ +// Package httpx holds HTTP client helpers shared across the proxy. +package httpx + +import ( + "net/http" + "time" +) + +// NewDirectClient returns an *http.Client whose transport NEVER routes through +// the environment HTTP proxy (HTTP_PROXY / HTTPS_PROXY / ALL_PROXY). It clones +// http.DefaultTransport (keeping its connection-pool / dial-timeout defaults) +// and only disables the proxy lookup. +// +// WHY (2026-06-30): the aikey-proxy process commonly inherits an egress HTTP +// proxy in its env — e.g. Clash on 127.0.0.1:7890 on a CN dev box. That proxy +// is for the USER's upstream AI egress (reaching platform.claude.com out of the +// GFW), NOT for the proxy's own CONTROL-PLANE calls. Go's default transport +// honors the env proxy for EVERY destination with no LAN/localhost exception, so +// without this the proxy routes its control-plane traffic (team-master writeback +// + polls, collector upload, CLI-JWT refresh, signals, cluster register) through +// Clash, which cannot reach the internal / LAN control plane → 502 / timeout +// (e.g. the OAuth member-token writeback "context deadline exceeded" bug). The +// control plane is always direct-reachable, so these clients must bypass the env +// proxy. +// +// Do NOT use this for AI-egress / upstream-provider / request-forwarding clients +// — those deliberately honor the env (or configured upstream) proxy to get out. +func NewDirectClient(timeout time.Duration) *http.Client { + tr := http.DefaultTransport.(*http.Transport).Clone() + tr.Proxy = nil + return &http.Client{Timeout: timeout, Transport: tr} +} diff --git a/internal/httpx/direct_test.go b/internal/httpx/direct_test.go new file mode 100644 index 0000000..df1cbef --- /dev/null +++ b/internal/httpx/direct_test.go @@ -0,0 +1,41 @@ +package httpx + +import ( + "net/http" + "testing" + "time" +) + +// TestNewDirectClient_NoProxy locks the core invariant: a control-plane client +// must have NO proxy function, so it NEVER routes through HTTP_PROXY/HTTPS_PROXY. +// This is the fence for the 2026-06-30 bug where the proxy inherited Clash on +// :7890 and misrouted its team-master writeback / collector upload through it → +// "context deadline exceeded". 防退化. +func TestNewDirectClient_NoProxy(t *testing.T) { + c := NewDirectClient(5 * time.Second) + + tr, ok := c.Transport.(*http.Transport) + if !ok { + t.Fatalf("Transport is not *http.Transport: %T", c.Transport) + } + if tr.Proxy != nil { + t.Fatal("NewDirectClient transport has a Proxy func — must be nil so control-plane calls bypass HTTP_PROXY") + } + if c.Timeout != 5*time.Second { + t.Errorf("timeout = %v, want 5s", c.Timeout) + } + // Cloned from DefaultTransport → keeps the connection-pool defaults. + if tr.MaxIdleConns == 0 { + t.Error("expected cloned http.DefaultTransport defaults (MaxIdleConns > 0)") + } +} + +// TestNewDirectClient_DoesNotMutateDefaultTransport guards against Clone() being +// dropped: setting Proxy=nil must act on a COPY, never the shared global (which +// the AI-egress / forwarding paths rely on to honor the env proxy). +func TestNewDirectClient_DoesNotMutateDefaultTransport(t *testing.T) { + _ = NewDirectClient(time.Second) + if dt, ok := http.DefaultTransport.(*http.Transport); ok && dt.Proxy == nil { + t.Fatal("NewDirectClient mutated the shared http.DefaultTransport.Proxy — must clone, not modify the global") + } +} diff --git a/internal/httpx/registry.go b/internal/httpx/registry.go new file mode 100644 index 0000000..0aebaaf --- /dev/null +++ b/internal/httpx/registry.go @@ -0,0 +1,89 @@ +package httpx + +import ( + "net/http" + "sync" + "sync/atomic" + "time" +) + +// registry.go — central self-heal registry for CONTROL-PLANE HTTP clients. +// +// Every long-lived direct client the proxy uses to reach master / collector / +// hub (group-runtime, member-token writeback, routing-override, quota, +// compliance, signals, JWT refresh, events/canary/preflight, cluster register) +// can go STALE after a host network change (WiFi switch / tether / interface +// up-down): a fresh dial returns "no route to host" until the PROCESS restarts, +// even though master is reachable from a fresh process. +// +// SwappableClient makes each such client hot-swappable behind an atomic pointer; +// registering it lets the proxy's self-heal (supervisor/selfheal.go + netmon.go) +// call RebuildAllControlPlane() to install a fresh client (new dialer + empty +// pool) for EVERY control-plane rail at once on a network change — no per-client +// wiring, no rail left stale. + +// SwappableClient is a hot-swappable *http.Client. Get() returns the live client +// (safe for concurrent use); Rebuild() atomically installs a fresh one from the +// factory. Construct via NewSwappable so it's registered for RebuildAllControlPlane. +type SwappableClient struct { + ptr atomic.Pointer[http.Client] + factory func() *http.Client +} + +// NewSwappable builds a swappable client from factory, stores the first instance, +// and registers it so RebuildAllControlPlane() rebuilds it on a network change. +func NewSwappable(factory func() *http.Client) *SwappableClient { + s := &SwappableClient{factory: factory} + s.ptr.Store(factory()) + registerControlPlane(s) + return s +} + +// NewSwappableDirect is the common-case constructor: a registered swappable +// direct control-plane client with the given timeout. One factory for every +// control-plane rail so callers don't repeat the closure. +func NewSwappableDirect(timeout time.Duration) *SwappableClient { + return NewSwappable(func() *http.Client { return NewDirectClient(timeout) }) +} + +// NewSwappableFixed wraps an already-built *http.Client (e.g. one injected by a +// test or a caller) as a SwappableClient WITHOUT registering it — Rebuild is a +// no-op reinstall of the same client. Lets a struct field be *SwappableClient +// uniformly while honoring an injected client (which must not be globally +// rebuilt out from under the injector). +func NewSwappableFixed(c *http.Client) *SwappableClient { + s := &SwappableClient{factory: func() *http.Client { return c }} + s.ptr.Store(c) + return s +} + +// Get returns the live client — always call this at request time (never cache the +// result) so a mid-flight Rebuild takes effect on the next request. +func (s *SwappableClient) Get() *http.Client { return s.ptr.Load() } + +// Rebuild installs a fresh client from the factory (new transport + empty pool). +func (s *SwappableClient) Rebuild() { s.ptr.Store(s.factory()) } + +var ( + cpMu sync.Mutex + cpClients []*SwappableClient +) + +func registerControlPlane(s *SwappableClient) { + cpMu.Lock() + cpClients = append(cpClients, s) + cpMu.Unlock() +} + +// RebuildAllControlPlane rebuilds every registered control-plane client. Called by +// the proxy self-heal on a host network change (proactive) or a persistent +// no-route error (reactive). Cheap + idempotent — a fresh client is just a cloned +// transport with an empty connection pool. Returns the count rebuilt (for logs). +func RebuildAllControlPlane() int { + cpMu.Lock() + defer cpMu.Unlock() + for _, s := range cpClients { + s.Rebuild() + } + return len(cpClients) +} diff --git a/internal/httpx/registry_test.go b/internal/httpx/registry_test.go new file mode 100644 index 0000000..fe6b1a4 --- /dev/null +++ b/internal/httpx/registry_test.go @@ -0,0 +1,60 @@ +package httpx + +import ( + "net/http" + "testing" + "time" +) + +// TestSwappableDirect_BypassesEnvProxy is the REGRESSION guard for the whole +// point of these clients: control-plane clients MUST bypass the env HTTP proxy +// (Transport.Proxy == nil). Wrapping them in a SwappableClient — and rebuilding +// them on a network change — must NOT reintroduce env-proxy routing. +func TestSwappableDirect_BypassesEnvProxy(t *testing.T) { + assertBypass := func(when string, c *http.Client) { + tr, ok := c.Transport.(*http.Transport) + if !ok { + t.Fatalf("%s: transport is %T, want *http.Transport", when, c.Transport) + } + if tr.Proxy != nil { + t.Fatalf("%s: Transport.Proxy != nil — control-plane client would route through the env proxy", when) + } + } + s := NewSwappableDirect(2 * time.Second) + assertBypass("initial", s.Get()) + s.Rebuild() + assertBypass("after Rebuild", s.Get()) +} + +// TestRebuildAllControlPlane_RebuildsRegistered proves RebuildAllControlPlane +// swaps every REGISTERED client instance (not just one). +func TestRebuildAllControlPlane_RebuildsRegistered(t *testing.T) { + a := NewSwappableDirect(time.Second) + b := NewSwappableDirect(time.Second) + a0, b0 := a.Get(), b.Get() + if n := RebuildAllControlPlane(); n < 2 { + t.Fatalf("RebuildAllControlPlane rebuilt %d, want >=2 (a,b registered)", n) + } + if a.Get() == a0 || b.Get() == b0 { + t.Fatal("a registered client was not rebuilt") + } + // The rebuilt clients must STILL bypass the proxy. + if a.Get().Transport.(*http.Transport).Proxy != nil { + t.Fatal("rebuilt client reintroduced env-proxy routing") + } +} + +// TestSwappableFixed_NotRegistered_PreservesInjected proves an injected client is +// wrapped verbatim and is NOT swapped out by a global rebuild (so a test/caller +// injection survives, and its proxy behavior is whatever the caller chose). +func TestSwappableFixed_NotRegistered_PreservesInjected(t *testing.T) { + injected := &http.Client{Timeout: 3 * time.Second} + f := NewSwappableFixed(injected) + if f.Get() != injected { + t.Fatal("NewSwappableFixed did not preserve the injected client") + } + RebuildAllControlPlane() // must not touch the unregistered fixed client + if f.Get() != injected { + t.Fatal("RebuildAllControlPlane swapped an unregistered fixed client") + } +} diff --git a/internal/observability/handler.go b/internal/observability/handler.go index e4ec6ea..bd26da5 100644 --- a/internal/observability/handler.go +++ b/internal/observability/handler.go @@ -26,6 +26,46 @@ const ( EventProxyGenerationDrainTimeout = "proxy.generation.drain_timeout" ) +// Control-plane self-heal events (2026-07-01): recovery from a host network +// change that leaves the long-lived control-plane client stalled with no-route +// errors to master. See supervisor/selfheal.go + netmon.go. +const ( + EventProxyControlPlaneClientRebuilt = "proxy.control_plane.client_rebuilt" + EventProxyControlPlaneNetChange = "proxy.control_plane.net_change_detected" + EventProxyControlPlaneSelfRestart = "proxy.control_plane.self_restart" + EventProxyControlPlaneRestartExhausted = "proxy.control_plane.restart_budget_exhausted" +) + +// Egress (upstream proxy) events (2026-07-08): OS system-proxy auto-detection +// for the direct egress path. A long-lived proxy daemon must FOLLOW the host's +// system proxy (Clash port change / toggle) instead of freezing the value seen +// at process start — see internal/sysproxy. +const ( + EventProxyEgressSysProxyChanged = "proxy.egress.sysproxy_changed" + EventProxyEgressSysProxyReadFailed = "proxy.egress.sysproxy_read_failed" + EventProxyEgressSysProxyReadRecovered = "proxy.egress.sysproxy_read_recovered" +) + +// SyncRail events (2026-07-03): control-plane sync rail state transitions. +// A rail keeps serving last-known data and retrying forever in every state — +// these events are the VISIBILITY the old silent keep-last-known loops lacked +// (bugfix 2026-07-03-routing-override-rail-silent-stall.md). See +// supervisor/railset.go. +const ( + EventProxySyncRailStale = "proxy.sync.rail_stale" + EventProxySyncRailOffline = "proxy.sync.rail_offline" + EventProxySyncRailRecovered = "proxy.sync.rail_recovered" + EventProxySyncCredentialRebuilt = "proxy.sync.credential_rebuilt" + // EventProxySyncHealthFileFailed: the statusline sync-health bypass file + // (~/.aikey/run/sync-health.json) could not be written/removed — the claude + // status bar may show a stale (or miss a fresh) sync warning. + EventProxySyncHealthFileFailed = "proxy.sync.health_file_failed" + // EventReporterDeadLetterReplayed: the automatic dead-letter replay ran + // after the upload pipe recovered (2026-07-04 self-heal) — carries + // scanned/replayed/still-failing counts, or the error when the pass failed. + EventReporterDeadLetterReplayed = "proxy.events.dead_letter_replayed" +) + // Health events. const ( EventProxyHealthOk = "proxy.health.ok" @@ -48,6 +88,40 @@ const ( // token quota floor backstops it and the server baseline catches up on // re-sync. Recurring hits signal a stale summary needing a price re-sync. EventProxyQuotaModelUnpriced = "proxy.quota.model_unpriced" + // Oauth-group routing (N8). EventProxyGroupRouteResolved: a group VK request + // picked + injected a candidate account. EventProxyGroupRouteDegraded: no + // usable candidate (no material / all expired-exhausted / key unavailable) → + // the request is failed rather than silently routed wrong. + EventProxyGroupRouteResolved = "proxy.group.route_resolved" + EventProxyGroupRouteDegraded = "proxy.group.route_degraded" + // EventProxyGroupAccountCooldown (N8c): a pool account's upstream returned a + // fallback-worthy failure (401 broken / exhaustion-429), so it was cooled down + // and subsequent requests route around it. + EventProxyGroupAccountCooldown = "proxy.group.account_cooldown" + // EventProxyGroupAccountSwitched (N9 #8): the seat's rank-0 (primary) account + // was unusable (cooled / exhausted / expired / no material) so the request + // fell back to a different candidate — an auditable account switch. + EventProxyGroupAccountSwitched = "proxy.group.account_switched" + // EventProxyGroupLoginRequired (RW2/D2): the HRW-routed account has no token + // for this member (they haven't logged into it). The proxy returns a structured + // login prompt naming the account (strict HRW — it does NOT skip to a later + // logged-in account) so the member logs into THAT account on their local node. + EventProxyGroupLoginRequired = "proxy.group.login_required" + // EventProxyGroupWindowPrecut (N10): an account's upstream utilization crossed + // its randomized window cap (window_max_util_pct), so it was pre-cut for that + // window (cooled until reset) — staying under 100% which looks like abuse. + EventProxyGroupWindowPrecut = "proxy.group.window_precut" + // EventProxyGroupSeatBlocked (§5.5): the engine left this seat UNBOUND because + // every account in its pool/segment is at the ≤3-人/号 cap, so the proxy 429s it + // (never WRH-falls-back, which would route a 4th user onto a full account). + EventProxyGroupSeatBlocked = "proxy.group.seat_blocked" + // EventProxyGroupLoginStateWriteFailed / ClearFailed: the bypass + // ~/.aikey/run/group-login-required.json state file (statusline login hint, + // 20260703 update) could not be written / removed. Best-effort by design — + // the 401 response itself is unaffected — but WARN because the statusline + // hint silently disappearing (or nagging stale) is a debugging trap. + EventProxyGroupLoginStateWriteFailed = "proxy.group.login_state_write_failed" + EventProxyGroupLoginStateClearFailed = "proxy.group.login_state_clear_failed" ) // Usage extraction events. @@ -80,6 +154,13 @@ const ( ErrCodeQuotaExceededToken = "QUOTA_EXCEEDED_TOKEN" ErrCodeQuotaExceededUSD = "QUOTA_EXCEEDED_USD" ErrCodeQuotaDegradedBlock = "QUOTA_DEGRADED_BLOCK" + // Oauth-group routing degrade codes (N8). The resolver's typed reasons + // (GROUP_NO_CANDIDATES / GROUP_NO_MATERIAL / GROUP_ALL_UNUSABLE) are surfaced + // verbatim; GROUP_KEY_UNAVAILABLE is the proxy-local "can't decrypt" case. + ErrCodeGroupKeyUnavailable = "GROUP_KEY_UNAVAILABLE" + // ErrCodeGroupPoolFull (§5.5): 429 when the seat is blocked — every pool account + // is at the per-account user cap; the user waits or the admin adds accounts. + ErrCodeGroupPoolFull = "GROUP_POOL_FULL" ) // ---- HealthSnapshot ---- diff --git a/internal/observer/conversation_audit/extract.go b/internal/observer/conversation_audit/extract.go index a26b307..85ecf37 100644 --- a/internal/observer/conversation_audit/extract.go +++ b/internal/observer/conversation_audit/extract.go @@ -80,15 +80,13 @@ func extractAssistantDelta(protocolFamily, eventType string, payload []byte) str } return "" case protoOpenAI: - var f struct { - Choices []struct { - Delta struct { - Content string `json:"content"` - } `json:"delta"` - } `json:"choices"` - } - if json.Unmarshal(payload, &f) == nil && len(f.Choices) > 0 { - return f.Choices[0].Delta.Content + // Probe table (see openaiWireFormats): the two formats' frame shapes + // are disjoint, so first non-empty wins and Chat Completions keeps + // first-probe priority. + for _, wf := range openaiWireFormats { + if txt := wf.delta(payload); txt != "" { + return txt + } } return "" case protoGemini: @@ -116,7 +114,12 @@ func isCompletionMarker(protocolFamily, eventType string, payload []byte) bool { case protoAnthropic: return eventType == "message_stop" || bytes.Contains(payload, []byte(`"type":"message_stop"`)) case protoOpenAI: - return bytes.Equal(bytes.TrimSpace(payload), []byte("[DONE]")) + for _, wf := range openaiWireFormats { + if wf.done(payload) { + return true + } + } + return false case protoGemini: return bytes.Contains(payload, []byte(`"finishReason"`)) default: @@ -150,6 +153,58 @@ func extractAnthropicPrompt(body []byte) (userText, systemText, model string) { return userText, systemText, req.Model } +// --- OpenAI-family wire formats (polymorphic probe table) ------------------- +// +// The "openai_compatible" protocol family carries MORE THAN ONE wire format: +// classic Chat Completions (`/v1/chat/completions`, messages[]) and the +// Responses API (`/v1/responses`, instructions + input[], what Codex sends +// with wire_api="responses") — the same two-format split the usage extractor +// documents in provider/openai.go. Missing the second one silently skipped +// every Codex turn from the conversation audit (live incident 2026-07-07, +// CONTENT_EMPTY_EXTRACT on the Windows box while usage recorded fine). +// +// Why a probe table instead of if/else inside each function (user 拍板 +// 2026-07-07: 多态 + 可扩展 + 不影响旧 openai KEY 的采集): each wire format +// is ONE table entry owning its three probes (prompt / delta / done); the +// shared extraction flow iterates the table and the first entry that claims +// the payload wins. Adding the next OpenAI wire variant = appending one entry +// plus its fixtures — no edits to shared flow. Chat Completions probes FIRST, +// so existing openai-key deployments keep byte-identical extraction. +type openaiWireFormat struct { + name string + // prompt parses a request body; ok=false when the body is not this format + // (the table then tries the next entry). + prompt func(body []byte) (user, system, model string, ok bool) + // delta returns the assistant text carried by one SSE frame, "" when the + // frame carries none in this format. + delta func(payload []byte) string + // done reports whether the frame is this format's end-of-stream marker. + done func(payload []byte) bool +} + +var openaiWireFormats = []openaiWireFormat{ + {name: "chat_completions", prompt: chatCompletionsPrompt, delta: chatCompletionsDelta, done: chatCompletionsDone}, + {name: "responses_api", prompt: responsesAPIPrompt, delta: responsesAPIDelta, done: responsesAPIDone}, +} + +func extractOpenAIPrompt(body []byte) (userText, systemText, model string) { + for _, wf := range openaiWireFormats { + if u, s, m, ok := wf.prompt(body); ok { + return u, s, m + } + } + // No format claimed the body (e.g. a bare {"model":...} probe). Preserve + // the legacy behavior of still surfacing model so OnRequestEnd's model + // fallback chain is unchanged. + var bare struct { + Model string `json:"model"` + } + _ = json.Unmarshal(body, &bare) + return "", "", bare.Model +} + +// --- wire format 1: Chat Completions (/v1/chat/completions) ---------------- + type openaiReq struct { Model string `json:"model"` Messages []struct { @@ -158,10 +213,10 @@ type openaiReq struct { } `json:"messages"` } -func extractOpenAIPrompt(body []byte) (userText, systemText, model string) { +func chatCompletionsPrompt(body []byte) (userText, systemText, model string, ok bool) { var req openaiReq - if json.Unmarshal(body, &req) != nil { - return "", "", "" + if json.Unmarshal(body, &req) != nil || len(req.Messages) == 0 { + return "", "", "", false } for _, m := range req.Messages { if (m.Role == "system" || m.Role == "developer") && systemText == "" { @@ -174,7 +229,81 @@ func extractOpenAIPrompt(body []byte) (userText, systemText, model string) { break } } - return userText, systemText, req.Model + return userText, systemText, req.Model, true +} + +func chatCompletionsDelta(payload []byte) string { + var f struct { + Choices []struct { + Delta struct { + Content string `json:"content"` + } `json:"delta"` + } `json:"choices"` + } + if json.Unmarshal(payload, &f) == nil && len(f.Choices) > 0 { + return f.Choices[0].Delta.Content + } + return "" +} + +func chatCompletionsDone(payload []byte) bool { + return bytes.Equal(bytes.TrimSpace(payload), []byte("[DONE]")) +} + +// --- wire format 2: Responses API (/v1/responses, Codex) -------------------- + +type responsesReq struct { + Model string `json:"model"` + // Instructions is the Responses API's system prompt slot. + Instructions string `json:"instructions"` + Input json.RawMessage `json:"input"` // string | []item +} + +func responsesAPIPrompt(body []byte) (userText, systemText, model string, ok bool) { + var req responsesReq + if json.Unmarshal(body, &req) != nil { + return "", "", "", false + } + if req.Instructions == "" && len(req.Input) == 0 { + return "", "", "", false + } + systemText = req.Instructions + // input is either one plain string (single-turn shorthand) or an item + // array whose message items mirror chat messages with content parts of + // type input_text ({"type":"input_text","text":...} — joinContent reads + // the text field regardless of the type tag). + var s string + if json.Unmarshal(req.Input, &s) == nil { + return s, systemText, req.Model, true + } + var items []struct { + Role string `json:"role"` + Content json.RawMessage `json:"content"` + } + if json.Unmarshal(req.Input, &items) == nil { + for i := len(items) - 1; i >= 0; i-- { + if items[i].Role == "user" { + userText = joinContent(items[i].Content) + break + } + } + } + return userText, systemText, req.Model, true +} + +func responsesAPIDelta(payload []byte) string { + var f struct { + Type string `json:"type"` + Delta string `json:"delta"` + } + if json.Unmarshal(payload, &f) == nil && f.Type == "response.output_text.delta" { + return f.Delta + } + return "" +} + +func responsesAPIDone(payload []byte) bool { + return bytes.Contains(payload, []byte(`"type":"response.completed"`)) } type geminiPart struct { diff --git a/internal/observer/conversation_audit/extract_test.go b/internal/observer/conversation_audit/extract_test.go index ce55733..8c79f1f 100644 --- a/internal/observer/conversation_audit/extract_test.go +++ b/internal/observer/conversation_audit/extract_test.go @@ -159,3 +159,95 @@ func TestIsCompletionMarker(t *testing.T) { } } } + +// --- Responses API wire format (Codex, /v1/responses) ----------------------- +// +// Fixtures mirror what Codex actually sends with wire_api="responses" (the +// same format provider/openai.go's usage extractor documents). Live incident +// 2026-07-07: this format extracted to empty → every Codex turn was skipped +// from the conversation audit (CONTENT_EMPTY_EXTRACT) while usage recorded +// fine. Chat Completions fixtures above must stay green untouched — the probe +// table's first entry pins the legacy behavior. + +func TestExtractPrompt_OpenAI_ResponsesAPI_ItemArray(t *testing.T) { + body := []byte(`{ + "model": "gpt-5-codex", + "instructions": "You are Codex, a coding agent.", + "stream": true, + "input": [ + {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "old turn"}]}, + {"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "old answer"}]}, + {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "写一个快排"}]} + ] + }`) + user, system, model := extractPrompt(protoOpenAI, body) + if user != "写一个快排" { + t.Errorf("user = %q, want latest user turn", user) + } + if system != "You are Codex, a coding agent." { + t.Errorf("system = %q, want instructions", system) + } + if model != "gpt-5-codex" { + t.Errorf("model = %q", model) + } +} + +func TestExtractPrompt_OpenAI_ResponsesAPI_StringInput(t *testing.T) { + body := []byte(`{"model":"gpt-5-codex","input":"hello codex"}`) + user, system, model := extractPrompt(protoOpenAI, body) + if user != "hello codex" || system != "" || model != "gpt-5-codex" { + t.Errorf("got (%q,%q,%q), want plain-string input as user text", user, system, model) + } +} + +func TestExtractPrompt_OpenAI_ChatCompletionsStillWinsWhenMessagesPresent(t *testing.T) { + // A body carrying messages[] must keep the legacy extraction even if it + // also carries stray Responses-ish fields — chat/completions probes first. + body := []byte(`{ + "model": "gpt-4o", + "instructions": "should be ignored", + "messages": [ + {"role": "system", "content": "sys"}, + {"role": "user", "content": "hi"} + ] + }`) + user, system, model := extractPrompt(protoOpenAI, body) + if user != "hi" || system != "sys" || model != "gpt-4o" { + t.Errorf("got (%q,%q,%q), legacy chat/completions extraction regressed", user, system, model) + } +} + +func TestExtractPrompt_OpenAI_BareModelBodyKeepsLegacyModelFallback(t *testing.T) { + // Neither format claims a bare probe body; legacy behavior surfaced the + // model anyway (OnRequestEnd's fallback chain depends on it). + user, system, model := extractPrompt(protoOpenAI, []byte(`{"model":"gpt-4o"}`)) + if user != "" || system != "" || model != "gpt-4o" { + t.Errorf("got (%q,%q,%q), want empty text + model passthrough", user, system, model) + } +} + +func TestExtractAssistantDelta_OpenAI_ResponsesAPI(t *testing.T) { + if got := extractAssistantDelta(protoOpenAI, "", []byte(`{"type":"response.output_text.delta","delta":"chunk"}`)); got != "chunk" { + t.Errorf("delta = %q, want chunk", got) + } + // Non-text Responses frames carry no assistant text. + if got := extractAssistantDelta(protoOpenAI, "", []byte(`{"type":"response.output_item.added","item":{}}`)); got != "" { + t.Errorf("non-text frame leaked %q", got) + } + // Chat Completions delta unchanged (probe order). + if got := extractAssistantDelta(protoOpenAI, "", []byte(`{"choices":[{"delta":{"content":"cc"}}]}`)); got != "cc" { + t.Errorf("chat delta = %q, want cc", got) + } +} + +func TestIsCompletionMarker_OpenAI_BothFormats(t *testing.T) { + if !isCompletionMarker(protoOpenAI, "", []byte(`[DONE]`)) { + t.Error("[DONE] must stay a completion marker (chat/completions)") + } + if !isCompletionMarker(protoOpenAI, "", []byte(`{"type":"response.completed","response":{"usage":{"input_tokens":10}}}`)) { + t.Error("response.completed must mark completion (Responses API)") + } + if isCompletionMarker(protoOpenAI, "", []byte(`{"type":"response.output_text.delta","delta":"x"}`)) { + t.Error("delta frame must not mark completion") + } +} diff --git a/internal/observer/conversation_audit/observer.go b/internal/observer/conversation_audit/observer.go index 4bcc535..0e333cf 100644 --- a/internal/observer/conversation_audit/observer.go +++ b/internal/observer/conversation_audit/observer.go @@ -56,9 +56,14 @@ type ConversationRecord struct { UserText string `json:"user_text,omitempty"` ProviderCode string `json:"provider_code,omitempty"` Model string `json:"model,omitempty"` - EventID string `json:"event_id"` - VirtualKeyID string `json:"virtual_key_id,omitempty"` - OwnerAccountID string `json:"owner_account_id"` + EventID string `json:"event_id"` + VirtualKeyID string `json:"virtual_key_id,omitempty"` + // SeatID: seat-dimension attribution (2026-07-07) — the org seat of the + // human at the terminal, from the same route field usage events carry. + // OwnerAccountID alone misattributed shared-pool-VK turns to the VK + // owner; audit views key on seat with owner fallback for legacy rows. + SeatID string `json:"seat_id,omitempty"` + OwnerAccountID string `json:"owner_account_id"` SessionID string `json:"session_id"` RequestStatus string `json:"request_status"` OrgID string `json:"org_id"` @@ -230,6 +235,7 @@ func (o *Observer) OnRequestEnd(_ context.Context, req *observer.RequestContext, EventID: req.TraceID, // stable per-turn idempotency key OrgID: orgOrPersonal(req.OrgID), SessionID: req.SessionID, + SeatID: req.SeatID, OwnerAccountID: req.OwnerAccountID, Model: model, ProviderCode: req.ProviderID, diff --git a/internal/observer/conversation_audit/observer_test.go b/internal/observer/conversation_audit/observer_test.go index 1e0a2c0..747552e 100644 --- a/internal/observer/conversation_audit/observer_test.go +++ b/internal/observer/conversation_audit/observer_test.go @@ -35,6 +35,7 @@ func driveAnthropicTurn(o *Observer, traceID string, withStop bool) { TraceID: traceID, OrgID: "org-1", OwnerAccountID: "acct-7", + SeatID: "seat-3", SessionID: "sess-9", ProviderID: "anthropic", StartedAt: time.Unix(1_700_000_000, 0), @@ -81,6 +82,11 @@ func TestObserver_HappyPathAssemblesRecord(t *testing.T) { if r.OrgID != "org-1" || r.OwnerAccountID != "acct-7" || r.SessionID != "sess-9" { t.Fatalf("attribution: org=%q owner=%q session=%q", r.OrgID, r.OwnerAccountID, r.SessionID) } + // Seat dimension (2026-07-07): must ride along, or shared-pool-VK turns + // re-attribute to the VK owner and file under a stranger seat row. + if r.SeatID != "seat-3" { + t.Fatalf("seat_id=%q want seat-3", r.SeatID) + } if r.Model != "claude-x" || r.ProviderCode != "anthropic" { t.Fatalf("model=%q provider=%q", r.Model, r.ProviderCode) } diff --git a/internal/proxy/build_event_canonical_test.go b/internal/proxy/build_event_canonical_test.go index 8375a8d..5fe4026 100644 --- a/internal/proxy/build_event_canonical_test.go +++ b/internal/proxy/build_event_canonical_test.go @@ -46,6 +46,21 @@ func TestBuildBaseEvent_PrefersCanonicalProviderCode(t *testing.T) { } } +// Regression (2026-07-01): buildBaseEvent must NOT panic when resp is nil — the +// upstream-no-response / transport-error path. The `if resp != nil` util-report guard +// inside the function proves resp is nilable, but StatusCode used to be dereferenced +// unguarded → latent nil panic (staticcheck SA5011). nil resp → StatusCode 0, no panic. +func TestBuildBaseEvent_NilRespNoPanic(t *testing.T) { + p := &Proxy{} + req := httptest.NewRequest(http.MethodPost, "/v1/messages", nil) + route := &vkeys.ResolvedRoute{VirtualKeyID: "vk", ProviderCode: "anthropic", ProtocolType: "anthropic"} + + ev := p.buildBaseEvent(req, nil, time.Now(), route, false) // must not panic + if ev.StatusCode != 0 { + t.Fatalf("nil resp must yield StatusCode 0 (no upstream status), got %d", ev.StatusCode) + } +} + func TestBuildBaseEvent_FallbackToProviderWhenProviderCodeEmpty(t *testing.T) { // 防御:pre-2026-05-08 fixture 或 route_builders bug 可能让 ProviderCode // 为空字符串。此时退回 Provider,避免写空 provider 字段污染聚合查询。 diff --git a/internal/proxy/codex_util_test.go b/internal/proxy/codex_util_test.go new file mode 100644 index 0000000..f6d4b67 --- /dev/null +++ b/internal/proxy/codex_util_test.go @@ -0,0 +1,109 @@ +package proxy + +import ( + "net/http" + "testing" +) + +// TestParseCodexUtil pins parseCodexUtil against the REAL wire captured 2026-07-06 +// (research/oauth-codex-ratelimit/2026-07-06-codex-ratelimit-taxonomy.md) plus the +// adversarial cases the live capture warned about (the primary/secondary name is +// NOT tied to a fixed window — classification MUST be by window-minutes). +func TestParseCodexUtil(t *testing.T) { + hdr := func(kv map[string]string) http.Header { + h := http.Header{} + for k, v := range kv { + h.Set(k, v) + } + return h + } + + cases := []struct { + name string + h http.Header + wantOK bool + want5h, want7d float64 + }{ + { + // Exact live Plus capture: primary IS the 5h window here (300min), + // secondary the 7d (10080min) — the intuitive order, and the OPPOSITE + // of sub2api's "primary=weekly" comment. + name: "live_plus_capture_primary_is_5h", + h: hdr(map[string]string{ + "X-Codex-Primary-Used-Percent": "1", + "X-Codex-Primary-Window-Minutes": "300", + "X-Codex-Secondary-Used-Percent": "0", + "X-Codex-Secondary-Window-Minutes": "10080", + }), + wantOK: true, want5h: 0.01, want7d: 0.0, + }, + { + // Adversarial: primary carries the 7d window (10080min), secondary the + // 5h (300min). Classification by DURATION must map secondary→util_5h and + // primary→util_7d. If someone reverts to trust-the-name this fails. + name: "inverted_primary_is_7d", + h: hdr(map[string]string{ + "X-Codex-Primary-Used-Percent": "50", + "X-Codex-Primary-Window-Minutes": "10080", + "X-Codex-Secondary-Used-Percent": "90", + "X-Codex-Secondary-Window-Minutes": "300", + }), + wantOK: true, want5h: 0.90, want7d: 0.50, + }, + { + // percent can exceed 100 briefly → clamp to fraction 1.0. + name: "over_100_clamped", + h: hdr(map[string]string{ + "X-Codex-Primary-Used-Percent": "150", + "X-Codex-Primary-Window-Minutes": "300", + "X-Codex-Secondary-Used-Percent": "0", + "X-Codex-Secondary-Window-Minutes": "10080", + }), + wantOK: true, want5h: 1.0, want7d: 0.0, + }, + { + // Only the 5h window present (no secondary) → util_7d stays 0. + name: "only_5h_window", + h: hdr(map[string]string{ + "X-Codex-Primary-Used-Percent": "80", + "X-Codex-Primary-Window-Minutes": "300", + }), + wantOK: true, want5h: 0.80, want7d: 0.0, + }, + { + // Anthropic response (no X-Codex-* headers) → not codex, skip. + name: "anthropic_headers_only", + h: hdr(map[string]string{ + "anthropic-ratelimit-unified-5h-utilization": "0.42", + }), + wantOK: false, + }, + { + // Empty / garbage → skip. + name: "no_headers", + h: http.Header{}, + wantOK: false, + }, + } + + for _, c := range cases { + t.Run(c.name, func(t *testing.T) { + got5h, got7d, ok := parseCodexUtil(c.h) + if ok != c.wantOK { + t.Fatalf("ok = %v, want %v", ok, c.wantOK) + } + if !ok { + return + } + if !floatEq(got5h, c.want5h) || !floatEq(got7d, c.want7d) { + t.Errorf("util5h,util7d = %v,%v; want %v,%v", got5h, got7d, c.want5h, c.want7d) + } + }) + } +} + +func floatEq(a, b float64) bool { + const eps = 1e-9 + d := a - b + return d < eps && d > -eps +} diff --git a/internal/proxy/detached.go b/internal/proxy/detached.go index 700d8e4..342aebc 100644 --- a/internal/proxy/detached.go +++ b/internal/proxy/detached.go @@ -23,18 +23,31 @@ type detachedTransport struct { } func (t *detachedTransport) RoundTrip(req *http.Request) (*http.Response, error) { - ctx, cancel := context.WithTimeout(t.proxyCtx, t.maxTimeout) + // Detach from the CLIENT's cancellation (a client disconnect must not abort the + // upstream call) but PRESERVE the request's context VALUES. Bug (2026-07-01, found + // by TestGroupServe_WindowUtilHeaderPreCutsSwitchesAndUplinks): the old + // `context.WithTimeout(t.proxyCtx, …)` derived a FRESH context from the proxy + // lifecycle, dropping every stashed value on req.Context() — including the N10 + // window-cap stash that ModifyResponse reads off resp.Request. So the pre-cut + // (防封: cool a pool account before it hits 100%) silently never fired for + // NON-streaming pool requests. WithoutCancel keeps the values while dropping the + // client's Done; AfterFunc re-ties cancellation to proxy shutdown (unchanged + // lifetime: proxy shutdown OR maxTimeout OR body close). + ctx, cancel := context.WithTimeout(context.WithoutCancel(req.Context()), t.maxTimeout) + stopOnShutdown := context.AfterFunc(t.proxyCtx, cancel) + cancelAll := func() { stopOnShutdown(); cancel() } req = req.WithContext(ctx) resp, err := t.inner.RoundTrip(req) if err != nil { - cancel() + cancelAll() return nil, err } // Wrap the body so the per-request context is canceled when it is closed, - // preventing context leaks for long-lived connections. - resp.Body = &cancelOnClose{ReadCloser: resp.Body, cancel: cancel} + // preventing context leaks for long-lived connections (also stops the + // shutdown AfterFunc so it doesn't linger past the request). + resp.Body = &cancelOnClose{ReadCloser: resp.Body, cancel: cancelAll} return resp, nil } diff --git a/internal/proxy/e9_429_classification_test.go b/internal/proxy/e9_429_classification_test.go new file mode 100644 index 0000000..2f57a7a --- /dev/null +++ b/internal/proxy/e9_429_classification_test.go @@ -0,0 +1,79 @@ +package proxy + +// E9 — 429 三分类 (能力点 #8) regression, mapping §5-E9 of +// +// roadmap20260320/技术实现/阶段7-企业生产版/20260630-动态决策引擎-实现现状审计与E2E验证清单.md +// +// The audit requires THREE distinct upstream-429 buckets to behave differently: +// - 限流 (per-minute rate_limited, transient): the account is fine, just paced — do NOT +// pull it for the full 5-min default (switching wastes a good account). +// - 限额 (5h window quota exhausted): cool the account until the window resets. +// - WAF (business rejection, no rate-limit headers): not the account's fault — do not cool. +// +// cooldownDecision today is BINARY: any 429 carrying a *ratelimit* header (and no +// Retry-After) gets the flat poolCooldownDefault, so 限流 is over-cooled exactly like 限额. +// It never reads unified-status, so it cannot split the two. The WAF + 限额 buckets already +// behave correctly (asserted live below); the 限流 bucket is the gap and is xfail-skipped +// with a pointer until the split lands — never passing by asserting the broken 300s cool. +// +// Fixture-based per wire format (logging-conventions: parser/classifier changes must be +// fixture-tested per wire shape). Reuses resp() from oauth_pool_cooldown_test.go. + +import ( + "net/http" + "testing" + "time" +) + +func TestE9_429ThreeWayClassification_GAP(t *testing.T) { + now := time.Unix(1_750_000_000, 0) + + // 限额 (LIVE PASS): a 5h-window quota exhaustion 429 with Retry-After → cool until reset. + t.Run("quota_exhaustion_cools_to_reset", func(t *testing.T) { + h := http.Header{ + "Anthropic-Ratelimit-Unified-Status": {"exhausted"}, + "Retry-After": {"3600"}, + } + until, cool := cooldownDecision(resp(429, h), now) + if !cool { + t.Fatalf("5h-quota-exhaustion 429 must cool the account down") + } + if until != now.Add(3600*time.Second) { + t.Fatalf("quota exhaustion must cool until the window reset (Retry-After), got until=%v", until) + } + }) + + // WAF (LIVE PASS): a business-rejection 429 with no rate-limit headers → do NOT cool. + t.Run("waf_business_rejection_no_cooldown", func(t *testing.T) { + if _, cool := cooldownDecision(resp(429, nil), now); cool { + t.Fatalf("WAF business-rejection 429 (no rate-limit signal) must NOT cool the account") + } + }) + + // 限流 (xfail → flips green when fixed): a transient per-minute rate_limited 429, no + // Retry-After. Desired: not cooled for the full 5-min default (透传 / short-cool only). + // Current: the *ratelimit* header trips hasRateLimitSignal → flat poolCooldownDefault. + t.Run("per_minute_rate_limited_GAP", func(t *testing.T) { + h := http.Header{ + // a transient per-minute rate-limit: status says rate_limited, requests-remaining + // is momentarily 0, but there is NO Retry-After and the 5h unified window is NOT + // exhausted. The account recovers within the minute. + "Anthropic-Ratelimit-Unified-Status": {"rate_limited"}, + "Anthropic-Ratelimit-Requests-Remaining": {"0"}, + } + until, cool := cooldownDecision(resp(429, h), now) + cooledSecs := 0 + if cool { + cooledSecs = int(until.Sub(now).Seconds()) + } + if cool && cooledSecs >= int(poolCooldownDefault.Seconds()) { + t.Skipf("GAP (audit #8 / §5-E9): per-minute rate_limited 429 (no Retry-After) is "+ + "cooled %ds (the flat %ds default) — indistinguishable from 5h quota exhaustion. "+ + "A transient minute-limit should透传/short-cool, not pull a good account for 5min. "+ + "Goes green when cooldownDecision consumes unified-status to split 限流 from 限额.", + cooledSecs, int(poolCooldownDefault.Seconds())) + } + // reached only once the split lands: a transient rate-limit must not over-cool. + t.Logf("three-way classification active: rate_limited cooled=%v secs=%d", cool, cooledSecs) + }) +} diff --git a/internal/proxy/events_and_helpers.go b/internal/proxy/events_and_helpers.go index 54a20c8..489b6c2 100644 --- a/internal/proxy/events_and_helpers.go +++ b/internal/proxy/events_and_helpers.go @@ -57,14 +57,30 @@ func (p *Proxy) buildBaseEvent(req *http.Request, resp *http.Response, startTime // 极端兜底:ResolvedRoute 没填 ProviderCode 时退回 URL-form,避免空字符串 provider = route.Provider } + // resp may be nil on an upstream-no-response path (transport error → no + // response object). Every caller today is inside ModifyResponse (resp non-nil), + // but the `if resp != nil` guard below (util reporting) proves resp is treated + // as nilable — so read StatusCode nil-safely too, keeping the two consistent and + // making buildBaseEvent safe if a future error-path caller passes nil (was a + // latent panic; staticcheck SA5011). 0 = "no upstream status". + statusCode := 0 + if resp != nil { + statusCode = resp.StatusCode + } ev := events.UsageEvent{ Timestamp: startTime, VirtualKeyID: route.VirtualKeyID, - Provider: provider, - DurationMs: time.Since(startTime).Milliseconds(), - StatusCode: resp.StatusCode, - IsStreaming: streaming, - RequestPath: req.URL.Path, + // AccountID = the REAL account that served this request. For oauth_group + // routing group_serve sets route.AccountID to the CHOSEN account (and + // re-points it on fallback A→B), so local audit attributes to the + // account that actually sent upstream. VirtualKeyID stays the request's + // VK regardless of switch (stable user attribution). + AccountID: route.AccountID, + Provider: provider, + DurationMs: time.Since(startTime).Milliseconds(), + StatusCode: statusCode, + IsStreaming: streaming, + RequestPath: req.URL.Path, } // Ephemeral trace correlation (not persisted) so the async collector can // cite trace/span/request ids if it has to drop this event under backpressure. @@ -106,20 +122,42 @@ func (p *Proxy) buildBaseEvent(req *http.Request, resp *http.Response, startTime // Trust note: spoofable signal, display-only. ev.AppSlug = route.AppSlug } + // I5 (best-effort, OFF the hot path): ship the upstream's parsed 5h + 7d + // utilization to master so the allocation engine can read both (util_7d feeds + // the weekly-squeeze target). The sample still exists only when the 5h header + // is present — the two unified headers arrive together upstream — and util_7d + // just enriches it (0 + omitempty when the 7d header is absent, so util_5h + // reporting stays byte-identical). enqueue is nil-safe (reporter off → no-op) + // and non-blocking (full buffer drops the sample, never stalls). + // Provider dispatch is by header-sniff, NOT a provider_code switch: the two + // utilization header families are mutually exclusive (an account is either + // Anthropic → anthropic-ratelimit-unified-*, or Codex → X-Codex-*, never both), + // so we try Anthropic first, then Codex. This mirrors R34's reactive layer + // (hasRateLimitSignal / cooldownDecision sniff both without branching) and + // keeps the enqueue provider-neutral (design §4B: normalize at the edge, feed + // one wire). parseCodexUtil already ÷100s and maps primary/secondary→5h/7d by + // window duration, so master/engine see the identical (util_5h, util_7d) shape. + if resp != nil { + if util5h, ok := parseUnifiedUtil5h(resp.Header); ok { + util7d, _ := parseUnifiedUtil7d(resp.Header) + p.signalReporter.enqueue(route.CredentialID, time.Now().Unix(), util5h, util7d) + } else if util5h, util7d, ok := parseCodexUtil(resp.Header); ok { + p.signalReporter.enqueue(route.CredentialID, time.Now().Unix(), util5h, util7d) + } + } return ev } // recordEvent records a usage event for error responses (no token counts). func (p *Proxy) recordEvent(req *http.Request, resp *http.Response, startTime time.Time, route *vkeys.ResolvedRoute, bearerToken string, streaming bool) { ev := p.buildBaseEvent(req, resp, startTime, route, streaming) - var errMsg string + var errMsg, errType string if resp.StatusCode >= 400 { p.errors.Add(1) // Capture the upstream error envelope (re-buffers resp.Body so the client // still receives the original payload). Returns the provider's error TYPE + // human MESSAGE parsed from the body. - errType, msg := captureUpstreamErrorBody(resp) - errMsg = msg + errType, errMsg = captureUpstreamErrorBody(resp) // error_type: prefer the provider's classification (e.g. // exceeded_current_quota_error / invalid_request_error) so the usage row // tells WHY it failed; fall back to the generic HTTP reason phrase @@ -133,10 +171,33 @@ func (p *Proxy) recordEvent(req *http.Request, resp *http.Response, startTime ti } // Probe traffic bypasses all sinks — see isAikeyProbe rationale in // middleware.go. Keep p.errors.Add(1) above so the /metrics view still - // shows "N requests returned 4xx/5xx" even for self-tests. + // shows "N requests returned 4xx/5xx" even for self-tests. Probes also skip + // the revoked-feed below: a probe may deliberately send no auth and earn a + // 401, which is NOT a real revocation. if isAikeyProbe(req) { return } + // Revoked-feed emit (best-effort, OFF the hot path, mirrors the I5 util + // enqueue in buildBaseEvent): a hard-revoked OAuth token (401 "OAuth token + // has been revoked") means this credential can't serve again until re-auth, + // so flag master to quarantine the account in the allocation engine. We do + // NOT alter the 401 — it still flows to the client unchanged. enqueueRevoked + // is nil-safe (reporter off → no-op) and non-blocking (full buffer drops). + if isHardRevoked(resp.StatusCode, errType, errMsg) { + p.signalReporter.enqueueRevoked(route.CredentialID, "revoked") + } + // Rate-limit-feed emit (best-effort, OFF the hot path, mirrors the revoked + // hook above): the proxy sees every upstream 429 (rate limit) / 403 (forbidden); + // counting them per credential lets master normalize a near-window 429/403 + // frequency into the allocation engine's §5.1 w5 "Recent429FreqNorm" risk + // signal. We do NOT alter the response — it flows to the client unchanged. The + // status is the one captureUpstreamErrorBody already read; we don't re-read the + // body. incrRateLimit is nil-safe (reporter off → no-op) and non-blocking (a + // short map lock, reset every flush bounds it); an empty CredentialID (no route) + // is dropped inside. + if resp.StatusCode == http.StatusTooManyRequests || resp.StatusCode == http.StatusForbidden { + p.signalReporter.incrRateLimit(route.CredentialID) + } p.collector.Record(&ev) // Error responses are treated as interrupted — the client never got a // usable result even though the request finished "fast" from our POV. diff --git a/internal/proxy/filter_cache_perf_test.go b/internal/proxy/filter_cache_perf_test.go index ffe709f..9ce0ea6 100644 --- a/internal/proxy/filter_cache_perf_test.go +++ b/internal/proxy/filter_cache_perf_test.go @@ -32,7 +32,7 @@ func TestFilterCache_PerfDetectCallReduction(t *testing.T) { } r := newReq(`{"messages":[` + strings.Join(parts, ",") + `]}`) r.Header.Set("X-Claude-Code-Session-Id", "perf-sess") // 固定 session → 命中同一桶 - p.applyInboundFilter(httptest.NewRecorder(), r, "m", "personal", "", "", discardLogger()) + p.applyInboundFilter(httptest.NewRecorder(), r, "m", "personal", "", "", "", "", discardLogger()) } return hook.called } @@ -77,7 +77,7 @@ func TestFilterCache_LoopCoversAllNotJustTail(t *testing.T) { } r := newReq(`{"messages":[` + strings.Join(parts, ",") + `]}`) r.Header.Set("X-Claude-Code-Session-Id", "loop-cover-sess") - p.applyInboundFilter(httptest.NewRecorder(), r, "m", "personal", "", "", discardLogger()) + p.applyInboundFilter(httptest.NewRecorder(), r, "m", "personal", "", "", "", "", discardLogger()) body := readReqBody(t, r) if strings.Contains(body, "SECRET-at-the-very-first-turn") { @@ -105,12 +105,12 @@ func BenchmarkFilterCache_FullStorageWarmTurn(b *testing.B) { warm := newReq(body) warm.Header.Set("X-Claude-Code-Session-Id", "bench-sess") - p.applyInboundFilter(httptest.NewRecorder(), warm, "m", "personal", "", "", discardLogger()) + p.applyInboundFilter(httptest.NewRecorder(), warm, "m", "personal", "", "", "", "", discardLogger()) b.ResetTimer() for i := 0; i < b.N; i++ { r := newReq(body) r.Header.Set("X-Claude-Code-Session-Id", "bench-sess") - p.applyInboundFilter(httptest.NewRecorder(), r, "m", "personal", "", "", discardLogger()) + p.applyInboundFilter(httptest.NewRecorder(), r, "m", "personal", "", "", "", "", discardLogger()) } } diff --git a/internal/proxy/filter_cache_requirements_test.go b/internal/proxy/filter_cache_requirements_test.go index 7f039ff..061198c 100644 --- a/internal/proxy/filter_cache_requirements_test.go +++ b/internal/proxy/filter_cache_requirements_test.go @@ -40,11 +40,11 @@ func TestReq_HistorySensitiveStaysMaskedAcrossTurns(t *testing.T) { // 第1轮:用户发敏感词 r1 := newReq(`{"messages":[{"role":"user","content":"secret"}]}`) - p.applyInboundFilter(httptest.NewRecorder(), r1, "m", "personal", "", "", discardLogger()) + p.applyInboundFilter(httptest.NewRecorder(), r1, "m", "personal", "", "", "", "", discardLogger()) // 第2轮:敏感词进历史,最新 turn 是普通追问 r2 := newReq(`{"messages":[{"role":"user","content":"secret"},{"role":"assistant","content":"ok"},{"role":"user","content":"more"}]}`) - p.applyInboundFilter(httptest.NewRecorder(), r2, "m", "personal", "", "", discardLogger()) + p.applyInboundFilter(httptest.NewRecorder(), r2, "m", "personal", "", "", "", "", discardLogger()) body2 := readReqBody(t, r2) if strings.Contains(body2, "secret") { @@ -68,7 +68,7 @@ func TestReq_NoCrossSessionReuse(t *testing.T) { send := func(sess string) { r := newReq(`{"messages":[{"role":"user","content":"same-content"}]}`) r.Header.Set("X-Claude-Code-Session-Id", sess) // scope = h: - p.applyInboundFilter(httptest.NewRecorder(), r, "m", "personal", "", "", discardLogger()) + p.applyInboundFilter(httptest.NewRecorder(), r, "m", "personal", "", "", "", "", discardLogger()) } send("sessA") // 首次 → 扫,called=1 @@ -91,14 +91,14 @@ func TestReq_CompactionRescansSummary(t *testing.T) { // 第1轮:原始敏感消息 r1 := newReq(`{"messages":[{"role":"user","content":"敏感原文X"}]}`) - p.applyInboundFilter(httptest.NewRecorder(), r1, "m", "personal", "", "", discardLogger()) + p.applyInboundFilter(httptest.NewRecorder(), r1, "m", "personal", "", "", "", "", discardLogger()) if hook.called != 1 { t.Fatalf("turn1 called %d, want 1", hook.called) } // 第2轮(compaction):历史被压成 summary(全新文本),原消息消失 r2 := newReq(`{"messages":[{"role":"user","content":"总结:用户提过敏感的事"}]}`) - p.applyInboundFilter(httptest.NewRecorder(), r2, "m", "personal", "", "", discardLogger()) + p.applyInboundFilter(httptest.NewRecorder(), r2, "m", "personal", "", "", "", "", discardLogger()) if hook.called != 2 { t.Errorf("需求③失败:compaction summary 没被重扫:called %d, want 2(summary 是新 hash 必 miss)", hook.called) diff --git a/internal/proxy/filter_cache_test.go b/internal/proxy/filter_cache_test.go index 8035647..52971ee 100644 --- a/internal/proxy/filter_cache_test.go +++ b/internal/proxy/filter_cache_test.go @@ -20,7 +20,7 @@ func TestFilterCache_HistoryReuse(t *testing.T) { // 第1轮:messages=[A] → 扫 A r1 := newReq(`{"messages":[{"role":"user","content":"A"}]}`) - p.applyInboundFilter(httptest.NewRecorder(), r1, "m", "personal", "", "", discardLogger()) + p.applyInboundFilter(httptest.NewRecorder(), r1, "m", "personal", "", "", "", "", discardLogger()) if hook.called != 1 { t.Fatalf("turn1: Detect called %d, want 1", hook.called) } @@ -28,7 +28,7 @@ func TestFilterCache_HistoryReuse(t *testing.T) { // 第2轮:messages=[A(历史 user), x(assistant), B(user)] → A 命中缓存不重扫; // assistant x 不扫(返回内容);只扫新 user B。 r2 := newReq(`{"messages":[{"role":"user","content":"A"},{"role":"assistant","content":"x"},{"role":"user","content":"B"}]}`) - p.applyInboundFilter(httptest.NewRecorder(), r2, "m", "personal", "", "", discardLogger()) + p.applyInboundFilter(httptest.NewRecorder(), r2, "m", "personal", "", "", "", "", discardLogger()) if hook.called != 2 { // 1(A) + 1(B);assistant x 跳过;A 命中缓存 t.Errorf("turn2: Detect 累计 %d, want 2(A 命中缓存、assistant x 不扫、只扫 B)", hook.called) } @@ -40,7 +40,7 @@ func TestFilterCache_Disabled_AlwaysScans(t *testing.T) { p := &Proxy{filterHook: hook} // 不调 SetFilterCacheEnabled → 缓存关闭 for i := 0; i < 2; i++ { r := newReq(`{"messages":[{"role":"user","content":"A"}]}`) - p.applyInboundFilter(httptest.NewRecorder(), r, "m", "personal", "", "", discardLogger()) + p.applyInboundFilter(httptest.NewRecorder(), r, "m", "personal", "", "", "", "", discardLogger()) } if hook.called != 2 { t.Errorf("缓存关闭:Detect called %d, want 2(不应缓存)", hook.called) @@ -54,9 +54,9 @@ func TestFilterCache_MaskVerdictReused(t *testing.T) { p.SetFilterCacheEnabled(true, 5) r1 := newReq(`{"messages":[{"role":"user","content":"secret"}]}`) - p.applyInboundFilter(httptest.NewRecorder(), r1, "m", "personal", "", "", discardLogger()) + p.applyInboundFilter(httptest.NewRecorder(), r1, "m", "personal", "", "", "", "", discardLogger()) r2 := newReq(`{"messages":[{"role":"user","content":"secret"}]}`) - p.applyInboundFilter(httptest.NewRecorder(), r2, "m", "personal", "", "", discardLogger()) + p.applyInboundFilter(httptest.NewRecorder(), r2, "m", "personal", "", "", "", "", discardLogger()) if hook.called != 1 { t.Errorf("Detect called %d, want 1(第2次应命中缓存)", hook.called) @@ -73,7 +73,7 @@ func TestFilterCache_DegradedNotCached(t *testing.T) { p.SetFilterCacheEnabled(true, 5) for i := 0; i < 2; i++ { r := newReq(`{"messages":[{"role":"user","content":"A"}]}`) - p.applyInboundFilter(httptest.NewRecorder(), r, "m", "personal", "", "", discardLogger()) + p.applyInboundFilter(httptest.NewRecorder(), r, "m", "personal", "", "", "", "", discardLogger()) } if hook.called != 2 { t.Errorf("degraded 被缓存了?Detect called %d, want 2(degraded 不可缓存)", hook.called) @@ -97,7 +97,7 @@ func TestFilterCache_DegradedForwardsOriginalUnmasked(t *testing.T) { p.SetFilterCacheEnabled(true, 5) r := newReq(`{"messages":[{"role":"user","content":"` + sensitive + `"}]}`) - proceed := p.applyInboundFilter(httptest.NewRecorder(), r, "m", "personal", "", "", discardLogger()) + proceed := p.applyInboundFilter(httptest.NewRecorder(), r, "m", "personal", "", "", "", "", discardLogger()) if !proceed { t.Fatal("degraded 必须放行(fail-open):不拦截用户请求") @@ -210,7 +210,7 @@ func TestFilterCache_DuplicateMessagesInOneRequest(t *testing.T) { // 两条逐字相同的敏感消息一起发 r := newReq(`{"messages":[{"role":"user","content":"secret"},{"role":"user","content":"secret"}]}`) - p.applyInboundFilter(httptest.NewRecorder(), r, "m", "personal", "", "", discardLogger()) + p.applyInboundFilter(httptest.NewRecorder(), r, "m", "personal", "", "", "", "", discardLogger()) body := readReqBody(t, r) if strings.Contains(body, "secret") { diff --git a/internal/proxy/filter_content.go b/internal/proxy/filter_content.go index 98870c4..2e045f3 100644 --- a/internal/proxy/filter_content.go +++ b/internal/proxy/filter_content.go @@ -45,7 +45,7 @@ func extractFilterableContent(body []byte) (pieces []contentPiece, parsed map[st // system inside messages[] instead, so this is simply absent there. collectContentField(m, "system", &pieces) - // messages[].content for both Anthropic and OpenAI shapes. + // messages[].content for both Anthropic and OpenAI chat shapes. if msgs, isArr := m["messages"].([]any); isArr { for _, mi := range msgs { if msg, isObj := mi.(map[string]any); isObj { @@ -54,6 +54,23 @@ func extractFilterableContent(body []byte) (pieces []contentPiece, parsed map[st } } + // OpenAI Responses API (codex): content lives in input[] / input string, + // system in `instructions` (kept in scope here as the full extractor mirrors + // system for masking symmetry). See extractUserContent for the live path. + collectContentField(m, "instructions", &pieces) + switch in := m["input"].(type) { + case string: + if in != "" { + pieces = append(pieces, contentPiece{text: in, setText: func(s string) { m["input"] = s }}) + } + case []any: + for _, ii := range in { + if item, isObj := ii.(map[string]any); isObj { + collectContentField(item, "content", &pieces) + } + } + } + return pieces, m, true } @@ -82,11 +99,40 @@ func extractUserContent(body []byte) (pieces []contentPiece, parsed map[string]a if err := json.Unmarshal(body, &m); err != nil { return nil, nil, false } - msgs, isArr := m["messages"].([]any) - if !isArr { - return nil, m, false + if msgs, isArr := m["messages"].([]any); isArr { + collectUserTurns(msgs, &pieces) + return pieces, m, true } - for _, mi := range msgs { + // OpenAI Responses API (codex, /v1/responses): the request carries no + // messages[]; user turns live in input[] and the system prompt in + // `instructions` (skipped like `system`). input can be a bare string + // (single-turn shorthand) or an item array whose message items mirror + // chat messages but with input_text content parts. Without this branch + // the whole request was "not filterable" → forwarded UNFILTERED, letting + // codex prompts bypass the compliance scan while conversation audit (a + // separate, Responses-aware path) still captured them (live incident + // 2026-07-08; same wire-format class as the audit R20 fix). + switch in := m["input"].(type) { + case string: + if in != "" { + pieces = append(pieces, contentPiece{ + text: in, + setText: func(s string) { m["input"] = s }, + }) + } + return pieces, m, true + case []any: + collectUserTurns(in, &pieces) + return pieces, m, true + } + return nil, m, false +} + +// collectUserTurns scans a chat-style item array (Anthropic/OpenAI messages[] +// or Responses API input[]) for USER-role text, applying the single skip policy +// (model responses / admin instructions / tool output are never user input). +func collectUserTurns(items []any, out *[]contentPiece) { + for _, mi := range items { msg, isObj := mi.(map[string]any) if !isObj { continue @@ -99,9 +145,8 @@ func extractUserContent(body []byte) (pieces []contentPiece, parsed map[string]a // Anthropic + OpenAI-family (Codex/Kimi) + OpenClaw roles uniformly. continue } - collectContentField(msg, "content", &pieces) + collectContentField(msg, "content", out) } - return pieces, m, true } // extractLatestUserContent extracts ONLY the latest user turn's text — the NEW @@ -165,7 +210,13 @@ func collectContentField(parent map[string]any, key string, out *[]contentPiece) if !isObj { continue } - if t, _ := block["type"].(string); t != "text" { + // "text" = Anthropic/OpenAI chat content block; "input_text" = OpenAI + // Responses API user content part (codex). Both carry human text in + // the same `text` field. Non-text blocks (image / input_image / + // tool_use / output_text=model reply) are skipped. + switch t, _ := block["type"].(string); t { + case "text", "input_text": + default: continue } txt, isStr := block["text"].(string) diff --git a/internal/proxy/filter_content_test.go b/internal/proxy/filter_content_test.go index 94b4fc2..ec99efe 100644 --- a/internal/proxy/filter_content_test.go +++ b/internal/proxy/filter_content_test.go @@ -121,3 +121,76 @@ func TestContentPieceSetter_RoundTrips(t *testing.T) { } } } + +// TestExtractUserContent_ResponsesAPI is the fence for the 2026-07-08 live +// incident: codex sends the OpenAI Responses API wire format (instructions + +// input[] with input_text content parts), which the filter extractor did NOT +// understand → "no filterable content extracted; forwarded UNFILTERED" → the +// sensitive word "fuck" reached the model unscanned while the SAME message was +// captured by the conversation audit (separate path, already Responses-aware). +// The compliance filter must scan the user turns of input[] too. +func TestExtractUserContent_ResponsesAPI(t *testing.T) { + cases := []struct { + name string + body string + want []string + ok bool + }{ + { + name: "responses_input_item_array_scans_user_skips_instructions_and_assistant", + body: `{"model":"gpt-5-codex","instructions":"be a coding agent","input":[ + {"type":"message","role":"user","content":[{"type":"input_text","text":"old turn"}]}, + {"type":"message","role":"assistant","content":[{"type":"output_text","text":"model reply"}]}, + {"type":"message","role":"user","content":[{"type":"input_text","text":"fuck"}]} + ]}`, + // instructions (system-equivalent, admin) + assistant output must be + // skipped; both user turns scanned. + want: []string{"fuck", "old turn"}, + ok: true, + }, + { + name: "responses_input_plain_string", + body: `{"model":"gpt-5-codex","input":"fuck"}`, + want: []string{"fuck"}, + ok: true, + }, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + pieces, parsed, ok := extractUserContent([]byte(tc.body)) + if ok != tc.ok { + t.Fatalf("ok = %v, want %v", ok, tc.ok) + } + got := pieceTexts(pieces) + if len(got) != len(tc.want) { + t.Fatalf("texts = %v, want %v", got, tc.want) + } + for i := range got { + if got[i] != tc.want[i] { + t.Fatalf("texts = %v, want %v", got, tc.want) + } + } + if ok && parsed == nil { + t.Error("parsed nil when ok") + } + }) + } +} + +// Mask write-back must round-trip through the Responses API input[] shape too, +// so re-marshaling yields the masked request the upstream receives. +func TestExtractUserContent_ResponsesAPI_MaskRoundTrip(t *testing.T) { + body := `{"model":"gpt-5-codex","input":[{"type":"message","role":"user","content":[{"type":"input_text","text":"fuck"}]}]}` + pieces, parsed, ok := extractUserContent([]byte(body)) + if !ok || len(pieces) != 1 { + t.Fatalf("ok=%v pieces=%d, want ok + 1 piece", ok, len(pieces)) + } + pieces[0].setText("****") + out, err := json.Marshal(parsed) + if err != nil { + t.Fatalf("marshal: %v", err) + } + if !strings.Contains(string(out), `"****"`) || strings.Contains(string(out), "fuck") { + t.Fatalf("mask not written back: %s", out) + } +} diff --git a/internal/proxy/filter_dispatch.go b/internal/proxy/filter_dispatch.go index 20add06..a3ebbfc 100644 --- a/internal/proxy/filter_dispatch.go +++ b/internal/proxy/filter_dispatch.go @@ -81,6 +81,8 @@ func (p *Proxy) applyInboundFilter( routeSource string, orgID string, virtualKeyID string, + seatID string, + sessionID string, logger *slog.Logger, ) (proceed bool) { hook := p.filterHook @@ -276,7 +278,18 @@ func (p *Proxy) applyInboundFilter( // detector has neither): tenant_id = the VK's org, virtual_key_id = the // VK itself (per-seat attribution at a centralized gateway). The VK never // crosses the detector pipe. See 20260611 集中化网关归因改造. - teamEvents = append(teamEvents, injectVirtualKey(injectTenant(resp.Event, orgID), virtualKeyID)) + // seat_id: the org seat of the human (route.SeatID, same field usage + // + conversation audit carry). Without it the master compliance-audit + // page — which resolves the seat alias/email from seat_id — falls back + // to the raw detector user_id (metadata.user_id, a Claude/session id), + // so pool-VK events showed a stranger id instead of the employee's + // alias (2026-07-08, mirrors the conversation-audit seat fix). + // session_id: the conversation session (resolveSessionID, the SAME + // source the conversation-audit observer uses), so the compliance + // audit drawer can deep-link to the exact conversation turn for this + // flagged prompt — event_id joins the turn, session_id opens its + // thread (2026-07-08 cross-audit link, decision 2a). + teamEvents = append(teamEvents, injectSession(injectSeat(injectVirtualKey(injectTenant(resp.Event, orgID), virtualKeyID), seatID), sessionID)) } switch resp.Action { @@ -442,6 +455,59 @@ func injectVirtualKey(eventJSON []byte, virtualKeyID string) []byte { return out } +// injectSeat stamps the proxy-authoritative seat_id onto the team event JSON — +// the org seat of the human at the terminal (route.SeatID), the SAME field the +// usage + conversation-audit paths carry. The master compliance-audit page +// resolves the seat alias/email from this; without it, a pool-VK event's +// user_id (the detector's metadata.user_id = a Claude/session id) never matches +// a seat and the page shows a raw UUID. Empty seat (personal key / legacy) → +// unchanged, same fail-safe as injectVirtualKey. (2026-07-08 seat attribution.) +func injectSeat(eventJSON []byte, seatID string) []byte { + if seatID == "" { + return eventJSON + } + var m map[string]json.RawMessage + if err := json.Unmarshal(eventJSON, &m); err != nil { + return eventJSON + } + q, err := json.Marshal(seatID) + if err != nil { + return eventJSON + } + m["seat_id"] = q + out, err := json.Marshal(m) + if err != nil { + return eventJSON + } + return out +} + +// injectSession stamps the conversation session_id onto the team event JSON — +// resolved by resolveSessionID, the SAME source the conversation-audit observer +// uses, so the compliance audit drawer can deep-link to the exact conversation +// thread (event_id joins the turn; session_id opens its thread). Empty session +// (no session header — e.g. codex, which the conversation-audit UI keys by +// event_id instead) → unchanged, same fail-safe as injectSeat. (2026-07-08.) +func injectSession(eventJSON []byte, sessionID string) []byte { + if sessionID == "" { + return eventJSON + } + var m map[string]json.RawMessage + if err := json.Unmarshal(eventJSON, &m); err != nil { + return eventJSON + } + q, err := json.Marshal(sessionID) + if err != nil { + return eventJSON + } + m["session_id"] = q + out, err := json.Marshal(m) + if err != nil { + return eventJSON + } + return out +} + // itoaInt64 formats an int64 without pulling strconv into a hot file (mirrors // the tiny-helper convention used elsewhere in the proxy). func itoaInt64(n int64) string { diff --git a/internal/proxy/filter_dispatch_live_test.go b/internal/proxy/filter_dispatch_live_test.go index 5e222d8..bae0ed0 100644 --- a/internal/proxy/filter_dispatch_live_test.go +++ b/internal/proxy/filter_dispatch_live_test.go @@ -64,7 +64,7 @@ func TestApplyInboundFilter_LiveDetector(t *testing.T) { r := newReq(body) w := httptest.NewRecorder() - proceed := p.applyInboundFilter(w, r, "claude-3-5-sonnet", "personal", "", "", discardLogger()) + proceed := p.applyInboundFilter(w, r, "claude-3-5-sonnet", "personal", "", "", "", "", discardLogger()) // Mask verdict forwards the request (the LLM still answers, just on redacted // input). Block would also be acceptable policy, but for PII the baseline @@ -137,7 +137,7 @@ func TestApplyInboundFilter_LiveDetector_ASCII(t *testing.T) { r := newReq(body) w := httptest.NewRecorder() - proceed := p.applyInboundFilter(w, r, "claude-3-5-sonnet", "personal", "", "", discardLogger()) + proceed := p.applyInboundFilter(w, r, "claude-3-5-sonnet", "personal", "", "", "", "", discardLogger()) if !proceed { t.Fatalf("expected proceed=true; got blocked. status=%d", w.Code) } diff --git a/internal/proxy/filter_dispatch_nsfw_live_test.go b/internal/proxy/filter_dispatch_nsfw_live_test.go index b4fdfe2..e062803 100644 --- a/internal/proxy/filter_dispatch_nsfw_live_test.go +++ b/internal/proxy/filter_dispatch_nsfw_live_test.go @@ -81,7 +81,7 @@ func TestApplyInboundFilter_LiveDetector_NSFW(t *testing.T) { for time.Now().Before(deadline) { r := newReq(body) w := httptest.NewRecorder() - proceed := p.applyInboundFilter(w, r, "claude-3-5-sonnet", "personal", "", "", discardLogger()) + proceed := p.applyInboundFilter(w, r, "claude-3-5-sonnet", "personal", "", "", "", "", discardLogger()) if !proceed { // Block is also a valid interception verdict. t.Logf("NSFW intercepted via BLOCK (status=%d) — distributed pack effective", w.Code) diff --git a/internal/proxy/filter_dispatch_test.go b/internal/proxy/filter_dispatch_test.go index 2ba104c..0351c12 100644 --- a/internal/proxy/filter_dispatch_test.go +++ b/internal/proxy/filter_dispatch_test.go @@ -61,7 +61,7 @@ func TestApplyInboundFilter_PipeCapAndSplice(t *testing.T) { r := newReq(`{"model":"m","messages":[{"role":"user","content":"` + full + `"}]}`) w := httptest.NewRecorder() - if proceed := p.applyInboundFilter(w, r, "m", "personal", "", "", discardLogger()); !proceed { + if proceed := p.applyInboundFilter(w, r, "m", "personal", "", "", "", "", discardLogger()); !proceed { t.Fatal("expected proceed") } // 1. payload over the pipe is capped — the detector never sees the huge tail. @@ -84,7 +84,7 @@ func TestApplyInboundFilter_NoHook_PassThrough(t *testing.T) { r := newReq(`{"model":"claude","messages":[{"role":"user","content":"hi"}]}`) w := httptest.NewRecorder() - proceed := p.applyInboundFilter(w, r, "claude", "personal", "", "", discardLogger()) + proceed := p.applyInboundFilter(w, r, "claude", "personal", "", "", "", "", discardLogger()) if !proceed { t.Fatal("expected proceed=true with no hook") } @@ -101,7 +101,7 @@ func TestApplyInboundFilter_Allow(t *testing.T) { r := newReq(body) w := httptest.NewRecorder() - proceed := p.applyInboundFilter(w, r, "m", "personal", "", "", discardLogger()) + proceed := p.applyInboundFilter(w, r, "m", "personal", "", "", "", "", discardLogger()) if !proceed { t.Fatal("Allow should proceed") } @@ -131,7 +131,7 @@ func TestApplyInboundFilter_Mask(t *testing.T) { r := newReq(`{"messages":[{"content":"my id is 110101199001011234"}]}`) w := httptest.NewRecorder() - proceed := p.applyInboundFilter(w, r, "m", "personal", "", "", discardLogger()) + proceed := p.applyInboundFilter(w, r, "m", "personal", "", "", "", "", discardLogger()) if !proceed { t.Fatal("Mask should proceed (forward masked)") } @@ -167,7 +167,7 @@ func TestApplyInboundFilter_MaskEmptyPayload_LeavesUnchanged(t *testing.T) { r := newReq(orig) w := httptest.NewRecorder() - proceed := p.applyInboundFilter(w, r, "m", "personal", "", "", discardLogger()) + proceed := p.applyInboundFilter(w, r, "m", "personal", "", "", "", "", discardLogger()) if !proceed { t.Fatal("should proceed") } @@ -186,7 +186,7 @@ func TestApplyInboundFilter_Block(t *testing.T) { r := newReq(`{"messages":[{"content":"key=sk-ant-secret"}]}`) w := httptest.NewRecorder() - proceed := p.applyInboundFilter(w, r, "m", "personal", "", "", discardLogger()) + proceed := p.applyInboundFilter(w, r, "m", "personal", "", "", "", "", discardLogger()) if proceed { t.Fatal("Block should NOT proceed") } @@ -215,7 +215,7 @@ func TestApplyInboundFilter_BlockDefaultMessage(t *testing.T) { r := newReq(`{"messages":[{"content":"x"}]}`) // needs a content piece to inspect w := httptest.NewRecorder() - p.applyInboundFilter(w, r, "m", "personal", "", "", discardLogger()) + p.applyInboundFilter(w, r, "m", "personal", "", "", "", "", discardLogger()) var body map[string]any _ = json.Unmarshal(w.Body.Bytes(), &body) errObj := body["error"].(map[string]any) @@ -232,7 +232,7 @@ func TestApplyInboundFilter_Warn(t *testing.T) { r := newReq(body) w := httptest.NewRecorder() - proceed := p.applyInboundFilter(w, r, "m", "personal", "", "", discardLogger()) + proceed := p.applyInboundFilter(w, r, "m", "personal", "", "", "", "", discardLogger()) if !proceed { t.Fatal("Warn should proceed") } @@ -253,7 +253,7 @@ func TestApplyInboundFilter_DegradedFailsOpen(t *testing.T) { r := newReq(body) w := httptest.NewRecorder() - proceed := p.applyInboundFilter(w, r, "m", "personal", "", "", discardLogger()) + proceed := p.applyInboundFilter(w, r, "m", "personal", "", "", "", "", discardLogger()) if !proceed { t.Fatal("degraded must fail-open (proceed), NOT fail the request (§6 #11)") } @@ -273,7 +273,7 @@ func TestApplyInboundFilter_NilBody(t *testing.T) { r.Body = nil w := httptest.NewRecorder() - proceed := p.applyInboundFilter(w, r, "m", "personal", "", "", discardLogger()) + proceed := p.applyInboundFilter(w, r, "m", "personal", "", "", "", "", discardLogger()) if !proceed { t.Fatal("nil body should proceed without calling hook") } @@ -309,7 +309,7 @@ func TestApplyInboundFilter_UnknownAction_FailsLoudDegraded(t *testing.T) { var logBuf bytes.Buffer logger := slog.New(slog.NewTextHandler(&logBuf, &slog.HandlerOptions{Level: slog.LevelWarn})) - proceed := p.applyInboundFilter(w, r, "m", "personal", "", "", logger) + proceed := p.applyInboundFilter(w, r, "m", "personal", "", "", "", "", logger) // 1. fail-OPEN: the request proceeds and content is forwarded unchanged. if !proceed { @@ -323,3 +323,52 @@ func TestApplyInboundFilter_UnknownAction_FailsLoudDegraded(t *testing.T) { t.Errorf("expected a WARN (proxy.filter.unknown_action) for the unknown action, got logs:\n%s", logs) } } + +// TestInjectSeat pins the 2026-07-08 seat-attribution stamp: the compliance +// event must carry seat_id so the master audit page resolves the employee's +// alias instead of the raw detector user_id. Mirrors injectVirtualKey's +// fail-safe contract (empty seat / bad JSON → unchanged). +func TestInjectSeat(t *testing.T) { + // stamps seat_id, preserves existing fields (incl. detector's user_id) + out := injectSeat([]byte(`{"user_id":"claude-session-x","action":"block"}`), "seat-aa4c7f87") + var m map[string]any + if err := json.Unmarshal(out, &m); err != nil { + t.Fatalf("bad json: %v", err) + } + if m["seat_id"] != "seat-aa4c7f87" { + t.Fatalf("seat_id = %v, want seat-aa4c7f87", m["seat_id"]) + } + if m["user_id"] != "claude-session-x" { + t.Fatalf("user_id clobbered = %v (must be preserved — decision A)", m["user_id"]) + } + // empty seat → unchanged (personal key / legacy) + if got := string(injectSeat([]byte(`{"action":"mask"}`), "")); got != `{"action":"mask"}` { + t.Fatalf("empty seat mutated event: %s", got) + } + // non-JSON → unchanged (fail-safe) + if got := string(injectSeat([]byte(`not json`), "s1")); got != `not json` { + t.Fatalf("bad json mutated: %s", got) + } +} + +// TestInjectSession pins the 2026-07-08 cross-audit deep-link key: the +// compliance event carries session_id (same resolveSessionID source as the +// conversation-audit observer) so the drawer can open the exact conversation +// thread. Fail-safe contract mirrors injectSeat. +func TestInjectSession(t *testing.T) { + out := injectSession([]byte(`{"event_id":"trace-1","seat_id":"s1"}`), "sess-9") + var m map[string]any + if err := json.Unmarshal(out, &m); err != nil { + t.Fatalf("bad json: %v", err) + } + if m["session_id"] != "sess-9" { + t.Fatalf("session_id = %v, want sess-9", m["session_id"]) + } + if m["event_id"] != "trace-1" || m["seat_id"] != "s1" { + t.Fatalf("existing fields clobbered: %v", m) + } + // empty session (codex / no session header) → unchanged + if got := string(injectSession([]byte(`{"event_id":"x"}`), "")); got != `{"event_id":"x"}` { + t.Fatalf("empty session mutated: %s", got) + } +} diff --git a/internal/proxy/filter_incremental_test.go b/internal/proxy/filter_incremental_test.go index 399a0eb..a5a695b 100644 --- a/internal/proxy/filter_incremental_test.go +++ b/internal/proxy/filter_incremental_test.go @@ -69,7 +69,7 @@ func TestApplyInboundFilter_DeprecatedIncrementalNowFullScans(t *testing.T) { hook := &stubHook{resp: &apphook.Response{Action: apphook.ActionAllow}} // filterIncremental=true:旧 env 仍开,但已被忽略 → 仍走全量扫。 p := &Proxy{filterHook: hook, filterIncremental: true} - if proceed := p.applyInboundFilter(nil, newReq(body), "m", "personal", "", "", discardLogger()); !proceed { + if proceed := p.applyInboundFilter(nil, newReq(body), "m", "personal", "", "", "", "", discardLogger()); !proceed { t.Fatal("expected proceed") } // 全量扫:system + 3 条消息都扫 → >1 次。增量(旧漏因)只会 1 次(只扫最新 turn)。 @@ -85,7 +85,7 @@ func TestApplyInboundFilter_IncrementalFallsBackOnAgentLoop(t *testing.T) { body := `{"messages":[{"role":"user","content":"前文"},{"role":"assistant","content":"工具调用"}]}` hook := &stubHook{resp: &apphook.Response{Action: apphook.ActionAllow}} p := &Proxy{filterHook: hook, filterIncremental: true} - if proceed := p.applyInboundFilter(nil, newReq(body), "m", "personal", "", "", discardLogger()); !proceed { + if proceed := p.applyInboundFilter(nil, newReq(body), "m", "personal", "", "", "", "", discardLogger()); !proceed { t.Fatal("expected proceed") } // Full-scan fallback ran (the user piece got scanned), not a silent skip. diff --git a/internal/proxy/filter_integration_test.go b/internal/proxy/filter_integration_test.go index 7e3eecb..b289893 100644 --- a/internal/proxy/filter_integration_test.go +++ b/internal/proxy/filter_integration_test.go @@ -85,7 +85,7 @@ func TestFilterIntegration_HistoryLeakFixThroughRealDetector(t *testing.T) { apply := func(body string) string { r := newReq(body) r.Header.Set("X-Claude-Code-Session-Id", "integ-sess") - if !p.applyInboundFilter(httptest.NewRecorder(), r, "claude-3", "personal", "", "", discardLogger()) { + if !p.applyInboundFilter(httptest.NewRecorder(), r, "claude-3", "personal", "", "", "", "", discardLogger()) { t.Fatalf("request unexpectedly blocked: %s", body) } return readReqBody(t, r) @@ -155,7 +155,7 @@ func TestFilterIntegration_PerfBaseline(t *testing.T) { r := newReq(body) r.Header.Set("X-Claude-Code-Session-Id", sess) t0 := time.Now() - p.applyInboundFilter(httptest.NewRecorder(), r, "claude-3", "personal", "", "", discardLogger()) + p.applyInboundFilter(httptest.NewRecorder(), r, "claude-3", "personal", "", "", "", "", discardLogger()) return time.Since(t0) } @@ -201,7 +201,7 @@ func TestFilterIntegration_DetectorMaskIsDeterministic(t *testing.T) { // Mix a credential (regex rule) + a phone (NER/CRF path) to exercise both maskers. r := newReq(`{"messages":[{"role":"user","content":"please store my key ` + integSecret + ` and phone 13812345678 safely for later"}]}`) r.Header.Set("X-Claude-Code-Session-Id", "determinism") - p.applyInboundFilter(httptest.NewRecorder(), r, "claude-3", "personal", "", "", discardLogger()) + p.applyInboundFilter(httptest.NewRecorder(), r, "claude-3", "personal", "", "", "", "", discardLogger()) return readReqBody(t, r) } diff --git a/internal/proxy/forward_and_resolve.go b/internal/proxy/forward_and_resolve.go index 778dc04..3b448cb 100644 --- a/internal/proxy/forward_and_resolve.go +++ b/internal/proxy/forward_and_resolve.go @@ -107,6 +107,10 @@ func (p *Proxy) serveRouteWithObserver( } return p.loggedInAccountID // personal-key fallback, as in usage }(), + // Same field usage events stamp (reportable.go SeatID) — the seat + // dimension is what keeps shared-pool-VK turns attributed to the + // employee, not the VK owner (2026-07-07 audit misattribution). + SeatID: route.SeatID, } // payload_level=full (enterprise conversation audit): when some active // observer wants the raw request body, buffer+restore it here so the @@ -204,23 +208,11 @@ func (p *Proxy) ResolveBindingCredential( } } - // Codex BaseURL override pinned by TestFence_OAuthBinding_OpenAICodexBaseURLOverride. - // Why: Codex OAuth uses chatgpt.com/backend-api/codex (Responses API), - // NOT api.openai.com/v1 (Chat Completions API). API key users hit - // api.openai.com; OAuth users hit chatgpt.com. - // Ref: workflow/CI/research/oauth-codex-test/main.go - if canonicalCode == "openai" { - out.BaseURL = "https://chatgpt.com/backend-api/codex" - // Stage the `model` field into request context. Actual - // persistence is deferred to ModifyResponse, which only - // writes the file when upstream returned 2xx — see - // codex_model_capture.go for the bug-2026-06-09 rationale. - // `captureCodexModel` returns a NEW request with the value - // in context; reassign to propagate downstream. - r = captureCodexModel(r) - } else { - out.BaseURL = providerDefaultBaseURL(canonicalCode) - } + // Per-provider OAuth upstream (base URL + any provider setup) via the shared + // resolver — same source as the group route. Codex's chatgpt.com override + + // deferred model capture live in resolveOAuthUpstream; pinned by + // TestFence_OAuthBinding_OpenAICodexBaseURLOverride. + out.BaseURL, r = resolveOAuthUpstream(canonicalCode, r) oauthInject(r, cred, canonicalCode) identityTag := cred.Identity @@ -337,7 +329,7 @@ func (p *Proxy) serveRoute(w http.ResponseWriter, r *http.Request, route *vkeys. // On Block, applyInboundFilter writes the 403 and we return without // forwarding. On Mask, r.Body is rewritten to the redacted version so the // upstream LLM never sees the raw sensitive prompt. Fail-open on degraded. - if !p.applyInboundFilter(w, r, extractModel(r), route.RouteSource, route.OrgID, route.VirtualKeyID, logger) { + if !p.applyInboundFilter(w, r, extractModel(r), route.RouteSource, route.OrgID, route.VirtualKeyID, route.SeatID, resolveSessionID(r, route.ProtocolType, route.ProviderCode), logger) { return } @@ -376,7 +368,7 @@ func (p *Proxy) serveRoute(w http.ResponseWriter, r *http.Request, route *vkeys. // Streaming: keep the client context. When the client disconnects the // upstream TCP connection is released so the provider stops generation // and stops billing. Partial token usage is still recorded by the drainer. - inner := p.transport + inner := p.currentTransport() if inner == nil { inner = http.DefaultTransport logger.Debug("using default transport (no custom transport set)") @@ -471,6 +463,43 @@ func (p *Proxy) serveRoute(w http.ResponseWriter, r *http.Request, route *vkeys. // canonicalCode == "openai" OAuth branch above. persistCodexLastModelIfSuccessful(resp.Request, resp.StatusCode) + // N8c reactive fallback: if a pool account's upstream says it is + // broken (401) or its window is exhausted (rate-limit-signal 429), + // cool it down so the resolver routes subsequent requests around it. + // This request still returns its status to the client; in-request + // retry (no client-visible failure) is N9. Non-group routes + // (route.OauthGroupID == "") skip this entirely. + if route.OauthGroupID != "" && route.AccountID != "" { + // Path Z (通道3 §14): record the observed window-reset epoch so the + // next N7c pull piggybacks it to master, which re-rolls this + // account's window_max_util_pct when a new window starts. Present on + // 200s too. Observability-only side effect; never blocks the request. + if epoch, ok := observedResetEpoch(resp.Header); ok { + p.poolObservedResets.record(route.AccountID, epoch) + } + if until, ok := cooldownDecision(resp, time.Now()); ok { + p.poolCooldown.mark(route.AccountID, until) + logger.Warn("pool account cooled down after upstream failure", + "event.name", observability.EventProxyGroupAccountCooldown, + "oauth_group_id", route.OauthGroupID, + "account_id", route.AccountID, + "status", resp.StatusCode) + } + // N10 防封 pre-cut: even on a 200, if the account's utilization + // crossed its randomized cap, cool it down for this window so it + // never hits 100% (which looks like abuse upstream). + if capPct, ok := windowCapFromContext(resp.Request); ok { + if until, hit := windowPreCutDecision(resp.Header, capPct, time.Now()); hit { + p.poolCooldown.mark(route.AccountID, until) + logger.Warn("pool account pre-cut at window cap", + "event.name", observability.EventProxyGroupWindowPrecut, + "oauth_group_id", route.OauthGroupID, + "account_id", route.AccountID, + "cap_pct", capPct) + } + } + } + // Upstream-headers debug toggle: log a "response snapshot" with // status + body + selected response headers (rate-limit + retry // hints) for ANY status code. Why every status, not just 4xx: @@ -666,6 +695,11 @@ func (p *Proxy) serveRoute(w http.ResponseWriter, r *http.Request, route *vkeys. upstreamReqID := extractUpstreamRequestID(resp) p.reportUsage(route, bearerToken, ev.Model, startTime, resp.StatusCode, breakdown, "", "", realKey, sessionID, "complete", upstreamReqID) // Phase 2: accrue token + local usd quota for this completed request. + // Backfill the model for local usd pricing when the adapter left it + // empty (mirrors the streaming path) so codex/others aren't unpriced. + if breakdown.Model == "" { + breakdown.Model = ev.Model + } p.accrueQuotaUsage(route, breakdown, logger) } } else { @@ -706,6 +740,14 @@ func (p *Proxy) serveRoute(w http.ResponseWriter, r *http.Request, route *vkeys. } // Phase 2: accrue token + local usd quota on stream completion, // independent of the reporter. + // Local usd pricing keys on br.Model, but some providers' usage + // frame omits it (Codex /responses SSE) → the request would be + // left unpriced (usd=0) even though reportUsage recorded the + // correct `model`. Backfill so the edge price lookup matches the + // model this request actually ran (2026-07-06 codex-into-pool). + if br.Model == "" { + br.Model = model + } p.accrueQuotaUsage(route, br, logger) } } @@ -752,6 +794,18 @@ func (p *Proxy) serveRoute(w http.ResponseWriter, r *http.Request, route *vkeys. FlushInterval: -1, // Flush immediately for SSE streaming. } + // ConcurrencyPeak signal (best-effort, main-link-safe): count this credential + // as in-flight for the upstream forward, release on completion. defer + // guarantees the decrement even if ServeHTTP panics, so the live counter can't + // leak. trackInflight returns a no-op closure for a nil reporter / empty + // CredentialID, so this stays a single cheap map-lock with no extra branching. + // Placed AFTER the quota gate + inbound-filter early-returns above, so only + // requests that actually reach the upstream are counted (blocked ones never + // inc). ponytail: scoped to the synchronous ServeHTTP window — for streaming + // that blocks until the SSE copy to the client finishes, a good-enough + // concurrency proxy; chasing the detached non-streaming token-drain goroutine + // would need cross-goroutine lifetime tracking for marginal accuracy. + defer p.signalReporter.trackInflight(route.CredentialID)() rp.ServeHTTP(w, r) // 10. Slow request detection (after the full response is sent). diff --git a/internal/proxy/group_login_state.go b/internal/proxy/group_login_state.go new file mode 100644 index 0000000..8589f89 --- /dev/null +++ b/internal/proxy/group_login_state.go @@ -0,0 +1,128 @@ +// group_login_state.go — bypass "member login required" state file. +// +// Why a SEPARATE file instead of extending ~/.aikey/run/proxy-runtime.json: +// the runtime snapshot is rewritten wholesale by the supervisor on lifecycle +// events; this state is written from the REQUEST path. Sharing one file would +// force read-modify-write locking across owners (last-writer-wins would drop +// fields). One concern, one file, one writer — same ownership rule as +// 20260428-proxy-状态维护统一方案.md. +// +// Why a file at all (vs usage-WAL event or an HTTP status endpoint): the CLI +// statusline is offline-by-design (reads local files only, zero RPC), and the +// usage pipeline is the revenue-critical main path — a bypass display concern +// must not add events to it. See +// update/20260703-OAuth组成员登录提示-CLI显示与login_url.md 决策3. +// +// Consumers: `aikey statusline` renders "login required → " while this +// file exists. The proxy clears it on the next successful group-credential +// resolve (the member completed login), so statusline recovery is automatic. +package proxy + +import ( + "encoding/json" + "log/slog" + "os" + "path/filepath" + "sync/atomic" + "time" + + "github.com/AiKeyLabs/aikey-proxy/internal/observability" +) + +const groupLoginStateFilename = "group-login-required.json" + +// groupLoginState mirrors the JSON shape read by aikey-cli statusline. Field +// changes are a cross-repo contract — update commands_statusline.rs together. +type groupLoginStateBody struct { + Provider string `json:"provider"` + AccountID string `json:"account_id"` + LoginURL string `json:"login_url"` + // WrittenAt lets readers ignore stale files (e.g. proxy crashed before + // clearing). Unix millis, matching the usage-WAL event_time convention. + WrittenAt int64 `json:"written_at"` +} + +// groupLoginStateStore is always non-nil on a Proxy (set in New). `dirty` +// starts true so the FIRST successful group resolve after process start +// removes any file left behind by a previous incarnation (login may have +// completed while the proxy was down); afterwards Clear() is a no-op until +// the next Write, keeping the group-route hot path syscall-free. +type groupLoginStateStore struct { + dirty atomic.Bool +} + +func newGroupLoginStateStore() *groupLoginStateStore { + s := &groupLoginStateStore{} + s.dirty.Store(true) + return s +} + +// path mirrors runtime.Path(): $AIKEY_RUN_DIR override for tests, else +// ~/.aikey/run/. +func (s *groupLoginStateStore) path() (string, error) { + if dir := os.Getenv("AIKEY_RUN_DIR"); dir != "" { + return filepath.Join(dir, groupLoginStateFilename), nil + } + home, err := os.UserHomeDir() + if err != nil { + return "", err + } + return filepath.Join(home, ".aikey", "run", groupLoginStateFilename), nil +} + +// Write persists the login-required state (temp + rename, atomic for readers). +// Best-effort: a failure must never turn the structured 401 into a 500, but it +// is WARN-logged — a silently missing statusline hint is a debugging trap. +func (s *groupLoginStateStore) Write(logger *slog.Logger, provider, accountID, loginURL string) { + path, err := s.path() + if err == nil { + err = func() error { + if mkErr := os.MkdirAll(filepath.Dir(path), 0o700); mkErr != nil { + return mkErr + } + data, mErr := json.Marshal(groupLoginStateBody{ + Provider: provider, + AccountID: accountID, + LoginURL: loginURL, + WrittenAt: time.Now().UnixMilli(), + }) + if mErr != nil { + return mErr + } + tmp := path + ".tmp" + if wErr := os.WriteFile(tmp, data, 0o600); wErr != nil { + return wErr + } + return os.Rename(tmp, path) + }() + } + if err != nil { + logger.Warn("group login state: write failed — statusline will not show the login hint", + "event.name", observability.EventProxyGroupLoginStateWriteFailed, + "error", err) + return + } + s.dirty.Store(true) +} + +// Clear removes the state file after a successful group resolve. Cheap when +// nothing was written (atomic load, no syscall). Removal failure is WARN-worthy +// too: a stale file makes statusline nag about a login that already happened. +func (s *groupLoginStateStore) Clear(logger *slog.Logger) { + if !s.dirty.Load() { + return + } + path, err := s.path() + if err == nil { + if rmErr := os.Remove(path); rmErr != nil && !os.IsNotExist(rmErr) { + err = rmErr + } + } + if err != nil { + logger.Warn("group login state: clear failed — statusline may show a stale login hint", + "event.name", observability.EventProxyGroupLoginStateClearFailed, + "error", err) + return + } + s.dirty.Store(false) +} diff --git a/internal/proxy/group_resolve.go b/internal/proxy/group_resolve.go new file mode 100644 index 0000000..65011c4 --- /dev/null +++ b/internal/proxy/group_resolve.go @@ -0,0 +1,254 @@ +// group_resolve.go — N8a: oauth-group credential resolver (pure, no I/O). +// +// Given a resolved group VK route (route.OauthGroupID != ""), this picks one +// candidate account for the request and produces the credential to inject: +// +// route.GroupAccounts (candidate set, ranking inputs + identity, NO secrets) +// route.GroupRuntime (per-account encrypted material, written by N7c-2) +// │ +// ├─ seatassign.Rank(route.SeatID, candidates) ← byte-identical to master +// ├─ first usable candidate (has material, not expired/exhausted) +// ├─ base64-decode + vault.Decrypt the secret with the vault key +// └─ build *groupResolution (OAuthCredential | plaintext key) +// +// This function is deliberately side-effect free (no vault read, no HTTP, no +// header mutation) so it is fully unit-testable. The hot-path wiring that calls +// it + mutates the request lives in N8b (handle_dispatch), behind the +// oauth-group feature flag. Direct-bind / personal routes never reach here. +// +// SECURITY: the decrypted secret exists only in the returned struct's memory; +// it is never logged. Ranking MUST match master's snapshot.GroupAccountRef +// ordering (seatassign.Account{AccountID, Priority}, Weight unset = 1). +package proxy + +import ( + "encoding/base64" + "encoding/json" + + "github.com/AiKeyLabs/aikey-proxy/internal/vault" + "github.com/AiKeyLabs/aikey-proxy/internal/vkeys" + "github.com/AiKeyLabs/pkg/seatassign" +) + +// Credential types in the group material contract (match master + N7c-2 writer). +const ( + credTypeOAuth = "oauth_account" + credTypeKey = "api_key" +) + +// Resolution failure codes (mapped to HTTP responses by the N8b caller). +const ( + groupErrNoCandidates = "GROUP_NO_CANDIDATES" // route has no parseable candidate set + groupErrNoMaterial = "GROUP_NO_MATERIAL" // group_runtime empty/unparseable (not pulled yet) + groupErrAllUnusable = "GROUP_ALL_UNUSABLE" // every candidate expired / exhausted / undecryptable + // groupErrLoginRequired (RW2, per-member): the HRW-selected account has no + // token for THIS member (they haven't logged into it). Per D2 the walk stops + // here — it does NOT skip to a later already-logged-in account (that would + // break the HRW load allocation) — and the caller returns a structured login + // prompt naming groupResolveError.Account so the member logs into THAT account. + groupErrLoginRequired = "OAUTH_GROUP_MEMBER_LOGIN_REQUIRED" +) + + +// groupResolveError is a typed resolver failure so the caller can map a precise +// HTTP status + error code without string matching. +type groupResolveError struct { + Code string + Reason string + // Account is set for groupErrLoginRequired: the account the member must log + // into (the caller builds the login URL from it). + Account string +} + +func (e *groupResolveError) Error() string { return e.Code + ": " + e.Reason } + +// groupResolution is the chosen account + its decrypted credential for one +// request. Exactly one of OAuth / PlaintextKey is meaningful, per CredentialType. +type groupResolution struct { + AccountID string + CredentialID string // real credential_id of the chosen account → route.CredentialID for I5 signal reporting (T2 uplink) + CredentialType string // "oauth_account" | "api_key" + ProviderCode string // candidate's resolved provider code (oauthInject dispatch) + Identity string // display / audit only — never sent upstream + // Primary is the seat's rank-0 account (seatassign top pick). When it differs + // from AccountID, a fallback happened (the primary was cooled / exhausted / + // expired / has no material) — the caller audits the switch (N9 #8). + Primary string + + // oauth_account: header injection via oauthInject(req, OAuth, ProviderCode). + OAuth *OAuthCredential + // api_key: realKey + optional per-account upstream base URL. + PlaintextKey string + BaseURL string + Revision string + // WindowMaxUtilPct is master's randomized pre-cut cap (95-99) for this + // account's quota window (N11). When the upstream response says utilization + // ≥ this/100, N10 pre-cuts the account before it hits 100% (which looks like + // abuse). nil → no cap delivered → proxy uses 100% (natural exhaustion only). + WindowMaxUtilPct *int +} + +// resolveGroupCredential ranks the route's group candidates for route.SeatID and +// returns the first usable account's decrypted credential. `nowUnix` is the +// caller's clock (injected for deterministic tests); `derivedKey` is the vault +// key used to decrypt the at-rest material. `skip` (may be nil) names accounts +// the caller already tried this request — used by N8c fallback to advance past a +// candidate the upstream just rejected. `overrideAccountID` (may be "") is the +// allocation engine's seat→account routing override (I-side §6.5): when the +// engine has redirected this seat off an unhealthy default, the caller passes the +// engine's healthy pick here. +// +// A candidate is skipped when: it has no material in group_runtime (not pulled +// yet), its OAuth token is expired, its quota window is exhausted, or its secret +// fails to decrypt (corrupt). If every candidate is skipped → GROUP_ALL_UNUSABLE. +func resolveGroupCredential(route *vkeys.ResolvedRoute, derivedKey []byte, nowUnix int64, skip map[string]bool, overrideAccountID string) (*groupResolution, error) { + var refs []vkeys.GroupAccountRef + if route.GroupAccounts == "" || json.Unmarshal([]byte(route.GroupAccounts), &refs) != nil || len(refs) == 0 { + return nil, &groupResolveError{Code: groupErrNoCandidates, Reason: "no parseable group candidates on route"} + } + + // Distinguish "not pulled yet" from "pulled → nothing for this seat" (2026-06-30): + // "" → the channel-③ poll hasn't landed → NO_MATERIAL (retry DOES help). + // bad JSON → treat as not-ready → NO_MATERIAL (retry). + // "{}" → the proxy DID pull and this seat's group delivered NO accounts: + // the member was unbound/removed (access gate wiped it to "{}"), or + // the group has no enabled accounts → NO_CANDIDATES: retrying will + // NOT help, the member must contact an admin. The candidate snapshot + // (group_accounts) may still be STALE with entries — the material rail + // (proxy 60s) is fresher than on-demand key sync, so trust "{}" over + // stale candidates here. WHY: a removed member was shown "credentials + // still syncing, retry shortly" forever — the message told them to + // retry when only an admin re-adding the seat can fix it. + if route.GroupRuntime == "" { + return nil, &groupResolveError{Code: groupErrNoMaterial, Reason: "group_runtime not pulled yet"} + } + var material map[string]vkeys.GroupRuntimeAccount + if err := json.Unmarshal([]byte(route.GroupRuntime), &material); err != nil { + return nil, &groupResolveError{Code: groupErrNoMaterial, Reason: "group_runtime unparseable — treat as not-ready"} + } + if len(material) == 0 { + return nil, &groupResolveError{Code: groupErrNoCandidates, Reason: "group_runtime delivered no accounts for this seat (unbound/removed member or empty group)"} + } + + // Rank exactly as master does: Account{AccountID, Priority}, Weight unset. + accounts := make([]seatassign.Account, 0, len(refs)) + refByID := make(map[string]vkeys.GroupAccountRef, len(refs)) + for _, r := range refs { + accounts = append(accounts, seatassign.Account{AccountID: r.AccountID, Priority: r.Priority}) + refByID[r.AccountID] = r + } + ordered := seatassign.Rank(route.SeatID, accounts) + primary := ordered[0].AccountID // rank-0; audited when the actual pick differs + + // SINGLE SOURCE OF TRUTH (2026-07-01): the pick — engine override first (§6.5, + // member-validity re-checked; needs_login overrides are HONORED per the owner rule + // "the engine may route a member to an account they haven't logged into"), else the + // local ranked pick — is vkeys.PickRoutedAccount, the SAME function the display + // stamp (supervisor.computeRoutedAccountID → /user/vault current_routed) uses. Do + // NOT re-derive routing here; forwarding and display must agree by construction. + // + // The only hot-path-extra gate is DECRYPT (needs the vault key, so it can't be in + // the pure picker): a corrupt secret adds the account to the skip set and re-picks, + // preserving the old "corrupt material never fails the request" resilience. + localSkip, cloned := skip, false + for { + acc, oc := vkeys.PickRoutedAccount(route.SeatID, refs, material, overrideAccountID, localSkip, nowUnix) + switch oc { + case vkeys.PickNeedsLogin: + // RW2/D2: prompt login for THE routed account (engine pick or strict-HRW + // rank stop) — never silently hop past it to a later logged-in account. + return nil, &groupResolveError{Code: groupErrLoginRequired, + Reason: "member has no token for the routed account — login required", Account: acc} + case vkeys.PickOK: + mat := material[acc] + secret, err := decryptGroupSecret(derivedKey, mat) + if err != nil { + // corrupt material — route around it, don't fail the request. Clone the + // caller's skip set once (never mutate the shared cooldown view). + if !cloned { + clone := make(map[string]bool, len(skip)+1) + for k, v := range skip { + clone[k] = v + } + localSkip, cloned = clone, true + } + localSkip[acc] = true + continue + } + res := buildGroupResolution(acc, refByID[acc], mat, secret) + res.Primary = primary + return res, nil + default: // PickNone + // R36 (2026-07-04, codex pools): before declaring the admin-facing + // ALL_UNUSABLE dead-end, check whether some candidate's ONLY blocker is an + // EXPIRED member token — that's member-fixable (re-login mints a new one; + // codex tokens live ~10 days so this is their steady-state path) → prompt + // 401 login for the highest-ranked such account instead of a 503 "contact + // your admin". Mid-scan fallback THROUGH an expired account to a usable one + // is unchanged above (availability first — ExpiredPrimaryFallsToNext). + for _, a := range ordered { + if localSkip[a.AccountID] { + continue // cooled-down / undecryptable — not a login problem + } + m, ok := material[a.AccountID] + if ok && vkeys.MaterialExpired(m, nowUnix) && m.WindowStatus != "exhausted" { + return nil, &groupResolveError{Code: groupErrLoginRequired, + Reason: "member token for the routed account expired — re-login required", Account: a.AccountID} + } + } + return nil, &groupResolveError{Code: groupErrAllUnusable, Reason: "all group candidates expired, exhausted, or undecryptable"} + } + } +} + +// (resolveCandidate / materialUsable were absorbed into vkeys.PickRoutedAccount / +// vkeys.MaterialUsable — the shared pure pick used by BOTH this hot path and the +// supervisor's display stamp. 2026-07-01 single-source-of-truth unification.) + +// decryptGroupSecret base64-decodes the nonce + ciphertext and AES-GCM decrypts +// the secret with the vault key. +func decryptGroupSecret(derivedKey []byte, mat vkeys.GroupRuntimeAccount) (string, error) { + nonce, err := base64.StdEncoding.DecodeString(mat.SecretNonce) + if err != nil { + return "", err + } + ct, err := base64.StdEncoding.DecodeString(mat.SecretCiphertext) + if err != nil { + return "", err + } + pt, err := vault.Decrypt(derivedKey, nonce, ct) + if err != nil { + return "", err + } + return string(pt), nil +} + +// buildGroupResolution assembles the injectable credential for the chosen account. +func buildGroupResolution(accountID string, ref vkeys.GroupAccountRef, mat vkeys.GroupRuntimeAccount, secret string) *groupResolution { + res := &groupResolution{ + AccountID: accountID, + CredentialID: ref.CredentialID, + CredentialType: mat.CredentialType, + ProviderCode: ref.ProviderCode, + Identity: ref.Identity, + WindowMaxUtilPct: mat.WindowMaxUtilPct, // master's pre-cut cap (N10) + } + if mat.CredentialType == credTypeKey { + res.PlaintextKey = secret + res.BaseURL = mat.BaseURL + res.Revision = mat.Revision + return res + } + // oauth_account → build the credential oauthInject consumes. The real client + // identity flows upstream unchanged (transparent proxy); the former per-account + // AccountPersona normalization was removed 2026-06-29 (see oauth_inject.go). + res.OAuth = &OAuthCredential{ + AccessToken: secret, + Provider: ref.ProviderCode, + AccountID: accountID, + ExternalID: mat.ExternalID, // Claude metadata.user_id (empty until master N7a fills it) + Identity: ref.Identity, + ExpiresAt: mat.ExpiresAt, + } + return res +} diff --git a/internal/proxy/group_resolve_test.go b/internal/proxy/group_resolve_test.go new file mode 100644 index 0000000..1d05f43 --- /dev/null +++ b/internal/proxy/group_resolve_test.go @@ -0,0 +1,336 @@ +package proxy + +import ( + "encoding/base64" + "encoding/json" + "testing" + + "github.com/AiKeyLabs/aikey-proxy/internal/vault" + "github.com/AiKeyLabs/aikey-proxy/internal/vkeys" + "github.com/AiKeyLabs/pkg/seatassign" +) + +func grKey() []byte { + k := make([]byte, 32) + for i := range k { + k[i] = byte(i*3 + 1) + } + return k +} + +// encMat encrypts a secret into the at-rest group_runtime shape for tests. +func encMat(t *testing.T, key []byte, m vkeys.GroupRuntimeAccount, secret string) vkeys.GroupRuntimeAccount { + t.Helper() + nonce, ct, err := vault.Encrypt(key, []byte(secret)) + if err != nil { + t.Fatalf("encrypt: %v", err) + } + m.SecretNonce = base64.StdEncoding.EncodeToString(nonce) + m.SecretCiphertext = base64.StdEncoding.EncodeToString(ct) + return m +} + +func mustJSON(t *testing.T, v any) string { + t.Helper() + b, err := json.Marshal(v) + if err != nil { + t.Fatalf("marshal: %v", err) + } + return string(b) +} + +// rankOrder returns the seatassign order the resolver must follow (so tests +// pin behaviour to the same lib master uses, not a guessed order). +func rankOrder(seatID string, ids ...string) []string { + accts := make([]seatassign.Account, len(ids)) + for i, id := range ids { + accts[i] = seatassign.Account{AccountID: id} + } + ordered := seatassign.Rank(seatID, accts) + out := make([]string, len(ordered)) + for i, a := range ordered { + out[i] = a.AccountID + } + return out +} + +func TestResolveGroup_OAuthPrimaryDecrypts(t *testing.T) { + key := grKey() + seat := "seat-77" + order := rankOrder(seat, "acc-a", "acc-b") + primary := order[0] + + refs := []vkeys.GroupAccountRef{ + {AccountID: "acc-a", Identity: "a@x", ProviderCode: "anthropic"}, + {AccountID: "acc-b", Identity: "b@x", ProviderCode: "anthropic"}, + } + mat := map[string]vkeys.GroupRuntimeAccount{ + "acc-a": encMat(t, key, vkeys.GroupRuntimeAccount{CredentialType: "oauth_account", ExpiresAt: 9_000_000_000, ExternalID: "uuid-a"}, "tok-a"), + "acc-b": encMat(t, key, vkeys.GroupRuntimeAccount{CredentialType: "oauth_account", ExpiresAt: 9_000_000_000, ExternalID: "uuid-b"}, "tok-b"), + } + route := &vkeys.ResolvedRoute{SeatID: seat, OauthGroupID: "grp", GroupAccounts: mustJSON(t, refs), GroupRuntime: mustJSON(t, mat)} + + res, err := resolveGroupCredential(route, key, 1_000_000, nil, "") + if err != nil { + t.Fatalf("resolve: %v", err) + } + if res.AccountID != primary { + t.Fatalf("expected primary %q (seatassign order), got %q", primary, res.AccountID) + } + if res.Primary != res.AccountID { + t.Fatalf("primary usable → no switch: Primary(%q) must equal pick(%q)", res.Primary, res.AccountID) + } + if res.CredentialType != "oauth_account" || res.OAuth == nil { + t.Fatalf("want oauth cred, got %+v", res) + } + if res.OAuth.AccessToken != "tok-"+primary[len(primary)-1:] { + t.Fatalf("decrypted wrong token: %q for %q", res.OAuth.AccessToken, primary) + } + if res.OAuth.ExternalID != "uuid-"+primary[len(primary)-1:] { + t.Fatalf("external_id not carried for Claude metadata: %+v", res.OAuth) + } + if res.PlaintextKey != "" { + t.Fatalf("oauth resolution must not set PlaintextKey") + } +} + +// TestResolveGroup_LoginRequiredNoSkip (RW2/D2): the HRW-primary account has no +// material (member hasn't logged into it) while a LATER candidate does — the +// resolver returns LOGIN_REQUIRED for the PRIMARY and does NOT skip to the +// logged-in one (preserving HRW allocation). +func TestResolveGroup_LoginRequiredNoSkip(t *testing.T) { + key := grKey() + seat := "seat-77" + order := rankOrder(seat, "acc-a", "acc-b") + primary, other := order[0], order[1] + + refs := []vkeys.GroupAccountRef{{AccountID: "acc-a", ProviderCode: "anthropic"}, {AccountID: "acc-b", ProviderCode: "anthropic"}} + // Primary carries master's explicit needs_login marker; the NON-primary has a token. + mat := map[string]vkeys.GroupRuntimeAccount{ + primary: {CredentialType: "oauth_account", NeedsLogin: true}, // master: member not logged into the primary + other: encMat(t, key, vkeys.GroupRuntimeAccount{CredentialType: "oauth_account", ExpiresAt: 9_000_000_000}, "tok"), + } + route := &vkeys.ResolvedRoute{SeatID: seat, OauthGroupID: "grp", GroupAccounts: mustJSON(t, refs), GroupRuntime: mustJSON(t, mat)} + + _, err := resolveGroupCredential(route, key, 1_000_000, nil, "") + ge, ok := err.(*groupResolveError) + if !ok || ge.Code != groupErrLoginRequired { + t.Fatalf("want LOGIN_REQUIRED, got %v", err) + } + if ge.Account != primary { + t.Fatalf("D2: must prompt login for the HRW primary %q (not skip to %q); got %q", primary, other, ge.Account) + } +} + +// TestResolveGroup_AbsentMaterialSkips (P1): an account with NO material entry at +// all (channel-③ hasn't pulled it yet — NOT a master needs_login marker) is SKIPPED +// to the next usable candidate, NOT turned into a hard LOGIN_REQUIRED. This is the +// fix that stops a pre-pull race from telling a member to re-login an account that +// is actually fine. +func TestResolveGroup_AbsentMaterialSkips(t *testing.T) { + key := grKey() + seat := "seat-77" + order := rankOrder(seat, "acc-a", "acc-b") + other := order[1] + + refs := []vkeys.GroupAccountRef{{AccountID: "acc-a", ProviderCode: "anthropic"}, {AccountID: "acc-b", ProviderCode: "anthropic"}} + // Primary is simply ABSENT from material (not pulled yet); the other is usable. + mat := map[string]vkeys.GroupRuntimeAccount{ + other: encMat(t, key, vkeys.GroupRuntimeAccount{CredentialType: "oauth_account", ExpiresAt: 9_000_000_000}, "tok"), + } + route := &vkeys.ResolvedRoute{SeatID: seat, OauthGroupID: "grp", GroupAccounts: mustJSON(t, refs), GroupRuntime: mustJSON(t, mat)} + + res, err := resolveGroupCredential(route, key, 1_000_000, nil, "") + if err != nil { + t.Fatalf("absent-material primary must SKIP to the usable candidate, not error: %v", err) + } + if res.AccountID != other { + t.Fatalf("expected fallback to the usable account %q, got %q", other, res.AccountID) + } +} + +// TestResolveGroup_LoginRequiredAfterExhausted (RW2): quota fallback still skips +// an exhausted account, but stops at the next account the member hasn't logged +// into (login required), rather than skipping further to a usable one. +func TestResolveGroup_LoginRequiredAfterExhausted(t *testing.T) { + key := grKey() + seat := "seat-9" + order := rankOrder(seat, "x", "y", "z") + first, second := order[0], order[1] + + refs := []vkeys.GroupAccountRef{{AccountID: "x", ProviderCode: "anthropic"}, {AccountID: "y", ProviderCode: "anthropic"}, {AccountID: "z", ProviderCode: "anthropic"}} + mat := map[string]vkeys.GroupRuntimeAccount{ + // rank-0 exhausted → skipped (quota fallback continues). + first: encMat(t, key, vkeys.GroupRuntimeAccount{CredentialType: "oauth_account", ExpiresAt: 9_000_000_000, WindowStatus: "exhausted"}, "tok"), + // rank-2 has a usable token, but we must NOT reach it. + order[2]: encMat(t, key, vkeys.GroupRuntimeAccount{CredentialType: "oauth_account", ExpiresAt: 9_000_000_000}, "tok-z"), + // rank-1 (second) carries master's needs_login marker → login required, stops here. + second: {CredentialType: "oauth_account", NeedsLogin: true}, + } + route := &vkeys.ResolvedRoute{SeatID: seat, OauthGroupID: "grp", GroupAccounts: mustJSON(t, refs), GroupRuntime: mustJSON(t, mat)} + + _, err := resolveGroupCredential(route, key, 1_000_000, nil, "") + ge, ok := err.(*groupResolveError) + if !ok || ge.Code != groupErrLoginRequired { + t.Fatalf("want LOGIN_REQUIRED after skipping exhausted, got %v", err) + } + if ge.Account != second { + t.Fatalf("should stop at rank-1 %q (skip exhausted rank-0, not jump to usable rank-2); got %q", second, ge.Account) + } +} + +func TestResolveGroup_APIKeyCarriesBaseURL(t *testing.T) { + key := grKey() + refs := []vkeys.GroupAccountRef{{AccountID: "acc-k", ProviderCode: "openai"}} + mat := map[string]vkeys.GroupRuntimeAccount{ + "acc-k": encMat(t, key, vkeys.GroupRuntimeAccount{CredentialType: "api_key", BaseURL: "https://up.example", Revision: "r3"}, "sk-secret"), + } + route := &vkeys.ResolvedRoute{SeatID: "s1", GroupAccounts: mustJSON(t, refs), GroupRuntime: mustJSON(t, mat)} + + res, err := resolveGroupCredential(route, key, 1_000_000, nil, "") + if err != nil { + t.Fatalf("resolve: %v", err) + } + if res.CredentialType != "api_key" || res.PlaintextKey != "sk-secret" { + t.Fatalf("want api_key secret, got %+v", res) + } + if res.BaseURL != "https://up.example" || res.Revision != "r3" { + t.Fatalf("api_key meta not carried: %+v", res) + } + if res.OAuth != nil { + t.Fatalf("api_key resolution must not set OAuth") + } +} + +func TestResolveGroup_ExpiredPrimaryFallsToNext(t *testing.T) { + key := grKey() + seat := "seat-42" + order := rankOrder(seat, "acc-a", "acc-b") + primary, secondary := order[0], order[1] + + refs := []vkeys.GroupAccountRef{ + {AccountID: "acc-a", ProviderCode: "anthropic"}, + {AccountID: "acc-b", ProviderCode: "anthropic"}, + } + now := int64(1_000_000) + mat := map[string]vkeys.GroupRuntimeAccount{ + primary: encMat(t, key, vkeys.GroupRuntimeAccount{CredentialType: "oauth_account", ExpiresAt: now - 1}, "tok-primary"), // expired + secondary: encMat(t, key, vkeys.GroupRuntimeAccount{CredentialType: "oauth_account", ExpiresAt: now + 10_000}, "tok-secondary"), // fresh + } + route := &vkeys.ResolvedRoute{SeatID: seat, GroupAccounts: mustJSON(t, refs), GroupRuntime: mustJSON(t, mat)} + + res, err := resolveGroupCredential(route, key, now, nil, "") + if err != nil { + t.Fatalf("resolve: %v", err) + } + if res.AccountID != secondary { + t.Fatalf("expired primary should fall back to %q, got %q", secondary, res.AccountID) + } + if res.OAuth.AccessToken != "tok-secondary" { + t.Fatalf("wrong fallback token: %q", res.OAuth.AccessToken) + } +} + +func TestResolveGroup_ExhaustedWindowSkipped(t *testing.T) { + key := grKey() + seat := "seat-9" + order := rankOrder(seat, "acc-a", "acc-b") + primary, secondary := order[0], order[1] + + refs := []vkeys.GroupAccountRef{{AccountID: "acc-a"}, {AccountID: "acc-b"}} + mat := map[string]vkeys.GroupRuntimeAccount{ + primary: encMat(t, key, vkeys.GroupRuntimeAccount{CredentialType: "oauth_account", ExpiresAt: 9_000_000_000, WindowStatus: "exhausted"}, "tok-p"), + secondary: encMat(t, key, vkeys.GroupRuntimeAccount{CredentialType: "oauth_account", ExpiresAt: 9_000_000_000, WindowStatus: "active"}, "tok-s"), + } + route := &vkeys.ResolvedRoute{SeatID: seat, GroupAccounts: mustJSON(t, refs), GroupRuntime: mustJSON(t, mat)} + + res, err := resolveGroupCredential(route, key, 1_000_000, nil, "") + if err != nil { + t.Fatalf("resolve: %v", err) + } + if res.AccountID != secondary { + t.Fatalf("exhausted primary should be skipped → %q, got %q", secondary, res.AccountID) + } +} + +func TestResolveGroup_SkipSetAdvances(t *testing.T) { + key := grKey() + seat := "seat-3" + order := rankOrder(seat, "acc-a", "acc-b") + primary, secondary := order[0], order[1] + + refs := []vkeys.GroupAccountRef{{AccountID: "acc-a"}, {AccountID: "acc-b"}} + mat := map[string]vkeys.GroupRuntimeAccount{ + primary: encMat(t, key, vkeys.GroupRuntimeAccount{CredentialType: "oauth_account", ExpiresAt: 9_000_000_000}, "tok-p"), + secondary: encMat(t, key, vkeys.GroupRuntimeAccount{CredentialType: "oauth_account", ExpiresAt: 9_000_000_000}, "tok-s"), + } + route := &vkeys.ResolvedRoute{SeatID: seat, GroupAccounts: mustJSON(t, refs), GroupRuntime: mustJSON(t, mat)} + + // Caller already tried the primary (e.g. upstream 401) → must advance. + res, err := resolveGroupCredential(route, key, 1_000_000, map[string]bool{primary: true}, "") + if err != nil { + t.Fatalf("resolve: %v", err) + } + if res.AccountID != secondary { + t.Fatalf("skip set should advance to %q, got %q", secondary, res.AccountID) + } + // N9 #8: the rank-0 primary is reported even when skipped, so the caller can + // audit the switch (primary != actual pick). + if res.Primary != primary { + t.Fatalf("Primary must be the rank-0 account %q (for switch audit), got %q", primary, res.Primary) + } + if res.Primary == res.AccountID { + t.Fatal("a skipped primary must differ from the actual pick (switch case)") + } +} + +func TestResolveGroup_ErrorCodes(t *testing.T) { + key := grKey() + refs := []vkeys.GroupAccountRef{{AccountID: "acc-a", ProviderCode: "anthropic"}} + + // No candidates. + if _, err := resolveGroupCredential(&vkeys.ResolvedRoute{SeatID: "s", GroupAccounts: "", GroupRuntime: "{}"}, key, 1, nil, ""); !isGroupErr(err, groupErrNoCandidates) { + t.Fatalf("want NO_CANDIDATES, got %v", err) + } + // Candidates present but no material pulled yet ("" = poll not landed) → NO_MATERIAL + // (transient, retry helps). + if _, err := resolveGroupCredential(&vkeys.ResolvedRoute{SeatID: "s", GroupAccounts: mustJSON(t, refs), GroupRuntime: ""}, key, 1, nil, ""); !isGroupErr(err, groupErrNoMaterial) { + t.Fatalf("want NO_MATERIAL for unpulled material, got %v", err) + } + // 2026-06-30: candidates present (STALE snapshot) but material explicitly "{}" — + // the proxy polled and this seat's group delivered NO accounts (member removed / + // unbound, or empty group). Must be NO_CANDIDATES ("contact admin, won't self- + // resolve"), NOT NO_MATERIAL ("still syncing, retry") — retrying never helps a + // removed member. This is the fix for the "still syncing forever" report. + if _, err := resolveGroupCredential(&vkeys.ResolvedRoute{SeatID: "s", GroupAccounts: mustJSON(t, refs), GroupRuntime: "{}"}, key, 1, nil, ""); !isGroupErr(err, groupErrNoCandidates) { + t.Fatalf("want NO_CANDIDATES for pulled-but-empty material, got %v", err) + } + // R36 (2026-07-04, codex pools): the only candidate is EXPIRED → LOGIN_REQUIRED, + // NOT ALL_UNUSABLE. An expired member token is member-fixable (re-login mints a + // new one; codex tokens live ~10 days so this is their steady-state path), so the + // resolver prompts a self-serve 401 login for that account instead of dead-ending + // in the admin-facing 503. + expMat := map[string]vkeys.GroupRuntimeAccount{ + "acc-a": encMat(t, key, vkeys.GroupRuntimeAccount{CredentialType: "oauth_account", ExpiresAt: 5}, "tok"), + } + expRoute := &vkeys.ResolvedRoute{SeatID: "s", GroupAccounts: mustJSON(t, refs), GroupRuntime: mustJSON(t, expMat)} + if _, err := resolveGroupCredential(expRoute, key, 1_000_000, nil, ""); !isGroupErr(err, groupErrLoginRequired) { + t.Fatalf("want LOGIN_REQUIRED for the expired-only case (R36), got %v", err) + } + + // Window-exhausted (NOT expired) is genuinely not member-fixable → ALL_UNUSABLE. + // Only routing around it (or waiting for the window to reset) helps, so it stays + // the admin-facing dead-end and does NOT downgrade to a login prompt. + exhMat := map[string]vkeys.GroupRuntimeAccount{ + "acc-a": encMat(t, key, vkeys.GroupRuntimeAccount{CredentialType: "oauth_account", ExpiresAt: 9_000_000_000, WindowStatus: "exhausted"}, "tok"), + } + exhRoute := &vkeys.ResolvedRoute{SeatID: "s", GroupAccounts: mustJSON(t, refs), GroupRuntime: mustJSON(t, exhMat)} + if _, err := resolveGroupCredential(exhRoute, key, 1_000_000, nil, ""); !isGroupErr(err, groupErrAllUnusable) { + t.Fatalf("want ALL_UNUSABLE for the window-exhausted-only case, got %v", err) + } +} + +func isGroupErr(err error, code string) bool { + ge, ok := err.(*groupResolveError) + return ok && ge.Code == code +} diff --git a/internal/proxy/group_routed_consistency_test.go b/internal/proxy/group_routed_consistency_test.go new file mode 100644 index 0000000..f28fd9e --- /dev/null +++ b/internal/proxy/group_routed_consistency_test.go @@ -0,0 +1,67 @@ +package proxy + +import ( + "testing" + + "github.com/AiKeyLabs/aikey-proxy/internal/vkeys" +) + +// This proves the ACTUAL forward routing that the A2 display (is_current_routed / +// current_routed) claims to reflect — the answer to "刷新后变更选中账号，proxy 是否会同步路 +// 由到最新选中的那个号". +// +// The hot path (resolveGroupCredential) and the display stamp +// (supervisor.computeRoutedAccountID) BOTH read RoutingOverrideCache.Assignment ?? seat- +// assign rank-0, so: +// (1) an engine override → the proxy FORWARDS to the override account, and it FOLLOWS +// override changes on the very next request (RoutingOverrideCache is read live). This +// is the display↔actual CONSISTENCY the user asked about — confirmed. +// (2) BUT when the routed account is COOLED (N8c skip) the hot path falls through to the +// next usable candidate, while the display stamp has no cooling input and keeps +// naming rank-0 → display can diverge from actual. Companion (display side): +// supervisor.TestComputeRoutedAccountID_DisplayIsOverrideOnly_NoCoolingAwareness. +func TestResolveGroup_RoutedFollowsOverride_AndCoolingFallsThrough(t *testing.T) { + key := grKey() + seat := "seat-routed-1" + order := rankOrder(seat, "acc-a", "acc-b") + primary, secondary := order[0], order[1] + + refs := []vkeys.GroupAccountRef{ + {AccountID: "acc-a", Identity: "a@x", ProviderCode: "anthropic"}, + {AccountID: "acc-b", Identity: "b@x", ProviderCode: "anthropic"}, + } + mat := map[string]vkeys.GroupRuntimeAccount{ + "acc-a": encMat(t, key, vkeys.GroupRuntimeAccount{CredentialType: "oauth_account", ExpiresAt: 9_000_000_000}, "tok-a"), + "acc-b": encMat(t, key, vkeys.GroupRuntimeAccount{CredentialType: "oauth_account", ExpiresAt: 9_000_000_000}, "tok-b"), + } + route := &vkeys.ResolvedRoute{SeatID: seat, OauthGroupID: "grp", GroupAccounts: mustJSON(t, refs), GroupRuntime: mustJSON(t, mat)} + + // (1a) engine override → proxy forwards to the OVERRIDE account (the displayed pick). + res, err := resolveGroupCredential(route, key, 1_000_000, nil, secondary) + if err != nil { + t.Fatalf("resolve with override: %v", err) + } + if res.AccountID != secondary { + t.Fatalf("override must route to %q, got %q", secondary, res.AccountID) + } + + // (1b) override CHANGES → the next request re-routes to the new account (the "selected + // account changed → proxy re-routes to the latest" guarantee). + res, err = resolveGroupCredential(route, key, 1_000_000, nil, primary) + if err != nil { + t.Fatalf("resolve with changed override: %v", err) + } + if res.AccountID != primary { + t.Fatalf("changed override must re-route to %q, got %q", primary, res.AccountID) + } + + // (2) rank-0 COOLED (in skip), no override → the hot path forwards to the next usable + // account. THIS is what the proxy actually serves; the display stamp does not model it. + res, err = resolveGroupCredential(route, key, 1_000_000, map[string]bool{primary: true}, "") + if err != nil { + t.Fatalf("resolve with rank-0 cooled: %v", err) + } + if res.AccountID != secondary { + t.Fatalf("rank-0 cooled → hot path must forward to next usable %q, got %q (display would still show %q → the documented gap)", secondary, res.AccountID, primary) + } +} diff --git a/internal/proxy/group_serve.go b/internal/proxy/group_serve.go new file mode 100644 index 0000000..6d31009 --- /dev/null +++ b/internal/proxy/group_serve.go @@ -0,0 +1,363 @@ +// group_serve.go — N8b: serve a oauth-group virtual key on the legacy /v1 entry. +// +// A group VK carries no static key (route.PlaintextKey == ""); its per-account +// material was pulled into route.GroupRuntime by N7c-2. handleOauthGroupRoute +// picks a candidate account (N8a resolver), injects its credential exactly the +// way the personal-OAuth / team-key paths do, and forwards through the shared +// serveRouteWithObserver. The direct-bind path in handle_dispatch is untouched — +// this whole branch is reached only when route.OauthGroupID != "" AND the feature +// flag is on, and group VKs aren't even registered when the flag is off (N7c-1). +// +// SECURITY / CORRECTNESS: +// - registry.Resolve hands out a SHARED *ResolvedRoute pointer. We mutate a +// per-request COPY (rc := *route) so concurrent requests on the same group VK +// never race on BaseURL/AccountID or corrupt the registry entry. +// - OAuth injection mirrors ResolveBindingCredential's branch byte-for-byte +// (Codex chatgpt.com base + captureCodexModel; others provider-default; +// headers via oauthInject; "__oauth__" sentinel realKey so the Director only +// rewrites the URL). refresh_token never reaches here (not in the material). +package proxy + +import ( + "encoding/json" + "fmt" + "log/slog" + "net/http" + "strings" + "time" + + "github.com/AiKeyLabs/aikey-proxy/internal/observability" + "github.com/AiKeyLabs/aikey-proxy/internal/vkeys" + "github.com/AiKeyLabs/aikey-proxy/pkg/observer" +) + +// handleOauthGroupRoute resolves + serves a group VK request. Called from Handle +// after route resolution, model-allowlist, and quota gating, in place of the +// static-key step-4 path. Always writes a response (success forward or a 503 +// degrade) — the caller returns immediately after. +func (p *Proxy) handleOauthGroupRoute( + w http.ResponseWriter, r *http.Request, + route *vkeys.ResolvedRoute, inboundBearer string, + startTime time.Time, logger *slog.Logger, traceID string, +) { + // Need the vault key to decrypt the at-rest material. nil only when the + // injected vault doesn't expose DerivedKey() (tests) — degrade, never panic. + if p.groupKey == nil { + p.degradeGroup(w, logger, route, observability.ErrCodeGroupKeyUnavailable, + "Group routing is temporarily unavailable. Please retry shortly.") + return + } + + // §5.5 hard cap: the engine left this seat UNBOUND — every account in its + // pool/segment is at the ≤3-人/号 cap. 429 here; do NOT fall through to the + // local pick, which is cap-blind and would route a 4th user onto a full account. + // (Distinct from the 503 degrade: an actionable "pool full" state, not a + // transient failure.) + if p.routingOverrides.Blocked(route.SeatID, route.OauthGroupID) { + logger.Warn("oauth-group seat blocked: pool at per-account user cap", + "event.name", observability.EventProxyGroupSeatBlocked, + "oauth_group_id", route.OauthGroupID, + "seat_id", route.SeatID) + writeJSONError(w, http.StatusTooManyRequests, "rate_limit_error", observability.ErrCodeGroupPoolFull, + "No available account: every account in your pool is at the per-account user limit. Ask your admin to add accounts.") + return + } + // Skip accounts cooling down from a recent upstream failure (N8c reactive + // fallback) so this request routes around them. The allocation engine's + // routing override for this seat (§6.5; "" when off / no redirect) is applied + // inside the resolver ONLY if it's still a valid candidate — fault-isolated, + // falls back to the local pick on any miss. + override := p.routingOverrides.lookup(route.SeatID, route.OauthGroupID) + res, err := resolveGroupCredential(route, p.groupKey.DerivedKey(), time.Now().Unix(), p.poolCooldown.skipSet(), override) + if err != nil { + code := observability.ErrCodeGroupKeyUnavailable + if ge, ok := err.(*groupResolveError); ok { + // RW2/D2: the routed account has no token for this member — return a + // structured login prompt (not a 503 degrade) so the client triggers the + // local OAuth login for THAT account. + if ge.Code == groupErrLoginRequired { + p.respondLoginRequired(w, logger, route, ge.Account) + return + } + code = ge.Code + } + p.degradeGroup(w, logger, route, code, groupDegradeMessage(code)) + return + } + + // Member has a token for the routed account ⇒ any earlier "login required" + // statusline hint is now stale — clear it (no-op unless one was written; + // see groupLoginStateStore.dirty). + p.groupLoginState.Clear(logger) + + // Per-request copy — DO NOT mutate the shared registry route (see file doc). + rc := *route + rc.AccountID = res.AccountID // usage attribution → the account actually used + rc.CredentialID = res.CredentialID // I5 signal reporting keyed by credential_id (T2 uplink; empty group route.CredentialID was dropping all signals) + // Point-in-time audit identity (2026-07-01, usage-audit "selected account" display): + // the SELECTED pool account's email rides the usage event as oauth_identity + // (reportable.go reads route.OAuthIdentity → ODS → DWD → the master usage-audit + // page). Denormalized ON PURPOSE: routing changes over time, so joining live + // tables later would misattribute history — the event must carry who served it. + rc.OAuthIdentity = res.Identity + + // A group VK is bound to a oauth_group, NOT a single provider, so the VK-level + // ProviderCode is EMPTY — the provider lives per-account in group_accounts. + // Use the RESOLVED account's provider_code (the candidate ref the resolver + // picked); fall back to the route's only for safety. Using rc.ProviderCode + // here yielded canonicalCode="" → empty upstream base URL → 502 (found by the + // live full-pipeline E2E 2026-06-25; hermetic tests set route.ProviderCode so + // never hit the empty case). + provCode := res.ProviderCode + if provCode == "" { + provCode = rc.ProviderCode + } + canonicalCode := providerCanonicalCode(provCode) + var realKey string + switch res.CredentialType { + case credTypeKey: + realKey = res.PlaintextKey + if res.BaseURL != "" { + rc.BaseURL = res.BaseURL + } + default: // oauth_account + // Per-provider OAuth upstream (base URL + any provider setup like codex's + // deferred model capture) via the shared resolver — same source as the + // legacy /v1 path. Headers injected here; the Director sees the sentinel + // and only rewrites the upstream URL. + rc.BaseURL, r = resolveOAuthUpstream(canonicalCode, r) + oauthInject(r, res.OAuth, canonicalCode) + // Stash the window cap so ModifyResponse can pre-cut this account when the + // upstream's unified-utilization crosses it (N10 防封). + if res.WindowMaxUtilPct != nil { + r = stashWindowCap(r, *res.WindowMaxUtilPct) + } + realKey = oauthSentinelKey + } + + adapterKey := rc.ProtocolType + if adapterKey == "" { + adapterKey = rc.Provider + } + prov, perr := p.providers.Get(adapterKey) + if perr != nil { + p.errors.Add(1) + logger.Error("group route: unknown provider", + "event.name", observability.EventProxyRequestUpstreamError, + "error.code", observability.ErrCodeProviderError, + "error.message", perr.Error(), + ) + writeJSONError(w, http.StatusBadGateway, "server_error", observability.ErrCodeProviderError, + "Unknown provider protocol: "+adapterKey) + return + } + + // N9 #8: audit a fallback — the seat's primary account was unusable (cooled / + // exhausted / expired / no material) so we routed to a different candidate. + if res.Primary != "" && res.Primary != res.AccountID { + logger.Info("oauth-group account switched (primary unusable)", + "event.name", observability.EventProxyGroupAccountSwitched, + "oauth_group_id", rc.OauthGroupID, + "from_account_id", res.Primary, + "to_account_id", res.AccountID, + ) + } + + logger.Info("group route resolved", + "event.name", observability.EventProxyGroupRouteResolved, + "oauth_group_id", rc.OauthGroupID, + "account_id", res.AccountID, + "credential_type", res.CredentialType, + "provider", canonicalCode, + ) + + // Group VKs leave rc.ProviderCode empty by design (the provider is per-account + // in group_accounts; the base URL above already used the resolved canonicalCode, + // NOT rc.ProviderCode — see the 502 note earlier). But the conversation-audit + // observer's protocol-specific extractor needs the provider to parse the turn: + // without it serveRouteWithObserver leaves ProtocolFamily="" → the extractor + // can't decode the messages/SSE → CONTENT_EMPTY_EXTRACT drops the turn while + // usage (protocol-agnostic) still reports. Found by the OAuth-pool E2E + // (2026-06-26). Set it to the resolved canonical provider now that the base URL + // is fixed, so the observer's ProtocolFamily fallback fires. + rc.ProviderCode = canonicalCode + + p.serveRouteWithObserver(w, r, &rc, prov, realKey, inboundBearer, startTime, logger, + observer.StreamUserChat, traceID) +} + +// groupDegradeMessage maps a resolver failure code to an actionable, end-user +// (Claude Code) message. The three modes mean very different things and need +// different guidance — collapsing them into one "retry shortly" misleads a member +// whose access is permanently gone (unbound) into retrying forever. +func groupDegradeMessage(code string) string { + switch code { + case groupErrNoCandidates: + // Permanent until an admin acts: the seat has NO account available in the + // group — either it was removed from the seat group, or the group has no + // active accounts. Covers both the empty candidate snapshot AND an empty + // channel-③ delivery ("{}", 2026-06-30). Retrying will NOT help. + return "Your seat has no available account in this credential-sharing group " + + "(it may have been removed from the group, or the group has no active accounts). " + + "Contact your administrator — this will not resolve on its own." + case groupErrNoMaterial: + // Transient: candidates exist but the proxy hasn't pulled their material + // yet (channel ③ poll in flight). Retrying shortly is the right action. + return "This group's credentials are still syncing to the proxy. Please retry shortly." + case groupErrAllUnusable: + // Every candidate is expired / quota-exhausted / undecryptable right now. + return "All accounts in this credential-sharing group are currently unavailable " + + "(rate-limited or expired). Contact your administrator if this persists." + default: + return "Group routing is temporarily unavailable. Please retry shortly." + } +} + +// groupLoginConsolePath is the local-console page where a member completes the +// pool-account sign-in (C19 rename: /user/oauth-contribute → /user/team-oauth). +// Appended to the configured console base to form login_url. The page route is +// owned by aikey-control's local web router — a rename there must update this +// constant (cross-repo contract, same as the JSON body shape below). +const groupLoginConsolePath = "/user/team-oauth" + +// respondLoginRequired returns the RW2/D2 structured login prompt: the member has +// no token for the HRW-routed account, so the client must run the local OAuth +// login for THAT account (proxy did NOT skip to a later logged-in candidate). +// Status 401: the member must authenticate to the account before the request can +// proceed. +// +// Display contract (20260703 update, spike-verified in dev2): +// - error.type is "authentication_error" — the Anthropic-standard type — NOT a +// custom string. claude/codex only render error.message verbatim for types +// they recognize; the previous custom "login_required" type fell into the +// generic "API error · Retrying" path and the user never saw the prompt. +// Machine consumers keep the precise signal via error.code and the +// X-Aikey-Error-Source header (both stay OAUTH_GROUP_MEMBER_LOGIN_REQUIRED). +// - login_url is assembled HERE from Config.ConsoleURL (决策2: single assembly +// point — the statusline state file below reuses the same URL). Empty +// ConsoleURL (cluster node / server-side proxy) degrades to URL-less wording. +func (p *Proxy) respondLoginRequired(w http.ResponseWriter, logger *slog.Logger, route *vkeys.ResolvedRoute, accountID string) { + loginURL := p.groupLoginURL() + logger.Info("group route requires member login", + "event.name", observability.EventProxyGroupLoginRequired, + "oauth_group_id", route.OauthGroupID, + "virtual_key_id", route.VirtualKeyID, + "account_id", accountID, + "login_url", loginURL, + ) + // Bypass statusline hint (决策3): best-effort, never blocks the response. + p.groupLoginState.Write(logger, route.ProviderCode, accountID, loginURL) + + message := "AiKey: log in to this shared account before use. " + + "Open your local AiKey console (" + groupLoginConsolePath + " page), complete sign-in, then retry." + if loginURL != "" { + message = "AiKey: log in to this shared account before use. " + + "Open " + loginURL + " and complete sign-in, then retry." + } + + // SyncRail truthful wording (§5.4, 2026-07-03 incident): when the engine's + // assignment rail is stale/offline, THIS pick came from the local ranked + // fallback and may contradict what the team-oauth page shows (the engine may + // have routed the seat to an account the member already signed into). Saying + // "go log in" then sends the member on a wild-goose chase — say what is + // actually wrong instead. reason stays machine-readable; error.code and the + // header keep the precise signal unchanged (additive contract only). + reason := "" + if p.routingRailHealth != nil { + if st, downSecs := p.routingRailHealth(); st == "stale" || st == "offline" { + reason = "routing_sync_unavailable" + message = "AiKey: account routing sync with your team server has been unreachable for " + + humanDuration(downSecs) + ", so this sign-in request may be misdirected. " + + "Check the connection to your team server (aikey status), or contact your admin. " + + "If the account shown on your console is already signed in, this error should clear once sync recovers." + } + } + + w.Header().Set("Content-Type", "application/json") + w.Header().Set(HeaderAikeyErrorSource, groupErrLoginRequired) + w.WriteHeader(http.StatusUnauthorized) + errObj := map[string]string{ + "message": message, + "type": "authentication_error", + "code": groupErrLoginRequired, + } + if reason != "" { + errObj["reason"] = reason + } + // json.Marshal (not string concat) so account ids / future fields can't break + // the JSON or inject — correct escaping for free. + body, _ := json.Marshal(map[string]any{ + "error": errObj, + "account": accountID, + "login_url": loginURL, + }) + _, _ = w.Write(body) +} + +// humanDuration renders seconds as a coarse "N min" / "N h" for user-facing +// error text (en-US wording per the code-and-ui-language rule). +func humanDuration(secs int64) string { + switch { + case secs >= 3600: + return fmt.Sprintf("%d h", secs/3600) + case secs >= 60: + return fmt.Sprintf("%d min", secs/60) + default: + return fmt.Sprintf("%d s", secs) + } +} + +// groupLoginURL assembles the member-login page URL from the configured local +// console base. "" when no console is co-installed (empty console_url). +func (p *Proxy) groupLoginURL() string { + base := strings.TrimRight(p.consoleURL, "/") + if base == "" { + return "" + } + return base + groupLoginConsolePath +} + +// groupDegradeStatus maps a resolver failure code to the HTTP status that gives the +// client the RIGHT retry behavior (2026-07-01). The Anthropic SDK retries 5xx (and +// 429) with exponential backoff — so a PERMANENT failure returned as 503 makes +// `claude` HANG for minutes retrying something that can never succeed, and renders a +// misleading "server-side issue, usually temporary — try again" suffix that +// contradicts our own "will not resolve on its own" message. Rule: permanent → 4xx +// (fail fast); genuinely transient → 503. +func groupDegradeStatus(code string) (status int, errType string) { + switch code { + case groupErrNoCandidates: + // Permanent until an admin acts: the seat has no account in the group + // (removed / empty group). Retrying never helps → 403 so the client fails + // FAST — no backoff hang, no "try again" framing. THIS fixes the reported + // "no available oauth → claude hangs for minutes" bug. + return http.StatusForbidden, "permission_error" + case groupErrAllUnusable: + // Every candidate is rate-limited / expired — a genuine rate-limit that + // recovers when the upstream window resets. 429 is the honest code (the + // client MAY back off and retry, which is legitimate here). + return http.StatusTooManyRequests, "rate_limit_error" + default: + // NO_MATERIAL (channel-③ poll in flight) / group key unavailable (vault + // reload) → genuinely transient; retrying shortly IS the right action → 503. + return http.StatusServiceUnavailable, "server_error" + } +} + +// degradeGroup fails a group request loudly (never silently routes it to a wrong +// key). Emits a WARN with trace context + the degrade reason code, then the status +// groupDegradeStatus picks for that code (permanent → 4xx fail-fast, transient → +// 503). N8c extends this with per-candidate fallback before giving up; today a +// resolver failure means every candidate was already unusable. +func (p *Proxy) degradeGroup(w http.ResponseWriter, logger *slog.Logger, route *vkeys.ResolvedRoute, code, clientMsg string) { + p.errors.Add(1) + status, errType := groupDegradeStatus(code) + logger.Warn("group route degraded", + "event.name", observability.EventProxyGroupRouteDegraded, + "error.code", code, + "http.status", status, + "oauth_group_id", route.OauthGroupID, + "virtual_key_id", route.VirtualKeyID, + ) + writeJSONError(w, status, errType, code, clientMsg) +} diff --git a/internal/proxy/group_serve_test.go b/internal/proxy/group_serve_test.go new file mode 100644 index 0000000..0d00a8a --- /dev/null +++ b/internal/proxy/group_serve_test.go @@ -0,0 +1,867 @@ +package proxy + +// N8b integration tests: a group VK presented at the legacy /v1 entry is +// resolved (N8a) + injected + forwarded. Uses capturingTransport (defined in +// oauth_binding_fence_test.go) to observe the outbound URL + headers without +// touching the network, and encMat/grKey (group_resolve_test.go) to build the +// encrypted at-rest material. + +import ( + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strings" + "testing" + "time" + + "github.com/AiKeyLabs/aikey-proxy/internal/events" + "github.com/AiKeyLabs/aikey-proxy/internal/vkeys" +) + +// fakeGroupKey provides the derived key for at-rest group material decryption. +type fakeGroupKey struct{ k []byte } + +func (f fakeGroupKey) DerivedKey() []byte { return f.k } + +// outboundCapture records the OUTBOUND (post-Director) request — the clone the +// upstream would see. api_key injection (prov.RewriteRequest) happens in the +// Director on this clone, so the original inbound r.Header would not show it; +// asserting on the outbound request is the correct semantics. +type outboundCapture struct { + host string + auth string + apiKey string + session string + stainlessOS string + // status (0 → 200) + respHeader let a test simulate an upstream failure for + // the N8c cooldown path. + status int + respHeader http.Header +} + +func (c *outboundCapture) RoundTrip(req *http.Request) (*http.Response, error) { + c.host = req.URL.Host + c.auth = req.Header.Get("Authorization") + c.apiKey = req.Header.Get("x-api-key") + c.session = req.Header.Get("X-Claude-Code-Session-Id") + c.stainlessOS = req.Header.Get("X-Stainless-OS") + st := c.status + if st == 0 { + st = 200 + } + h := http.Header{"Content-Type": []string{"application/json"}} + for k, v := range c.respHeader { + h[k] = v + } + return &http.Response{ + StatusCode: st, + Header: h, + Body: io.NopCloser(strings.NewReader( + `{"id":"msg","type":"message","content":[{"type":"text","text":"ok"}],"model":"c","stop_reason":"end_turn","usage":{"input_tokens":1,"output_tokens":1}}`)), + Request: req, + }, nil +} + +// setupGroupProxy seeds one group VK route + wires the key provider + the +// outbound-capturing transport, returning the capture for assertions. +func setupGroupProxy(t *testing.T, key []byte, route *vkeys.ResolvedRoute) (*Proxy, *outboundCapture) { + t.Helper() + p := setupTestProxy(t, "http://unused.invalid") + p.registry.Merge(map[string]*vkeys.ResolvedRoute{"aikey_team_grouptest": route}) + p.SetGroupKeyProvider(fakeGroupKey{k: key}) + tr := &outboundCapture{} + p.SetTransport(tr) + return p, tr +} + +func groupReq(body string) (*http.Request, *httptest.ResponseRecorder) { + req := httptest.NewRequest(http.MethodPost, "/v1/messages", strings.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer aikey_team_grouptest") + return req, httptest.NewRecorder() +} + +const groupBody = `{"model":"claude-sonnet-4-5","messages":[{"role":"user","content":"hi"}],"max_tokens":10}` + +func TestGroupServe_OAuthAccountInjectsBearer(t *testing.T) { + key := grKey() + refs := []vkeys.GroupAccountRef{{AccountID: "acc-1", ProviderCode: "anthropic"}} + mat := map[string]vkeys.GroupRuntimeAccount{ + "acc-1": encMat(t, key, vkeys.GroupRuntimeAccount{ + CredentialType: "oauth_account", ExpiresAt: 9_000_000_000, ExternalID: "uuid-1", + }, "oauth-tok-live"), + } + route := &vkeys.ResolvedRoute{ + VirtualKeyID: "vk-grp", Provider: "anthropic", ProtocolType: "anthropic", + ProviderCode: "anthropic", RouteSource: "team", + SeatID: "seat-1", OauthGroupID: "grp-1", + GroupAccounts: mustJSON(t, refs), GroupRuntime: mustJSON(t, mat), + } + p, tr := setupGroupProxy(t, key, route) + + req, w := groupReq(groupBody) + p.Handle(w, req) + + if w.Code != http.StatusOK { + t.Fatalf("group OAuth route: status=%d body=%s", w.Code, w.Body.String()) + } + // OAuth → providerDefaultBaseURL(anthropic), Bearer injected, x-api-key gone. + if tr.host != "api.anthropic.com" { + t.Fatalf("outbound host=%q want api.anthropic.com", tr.host) + } + if tr.auth != "Bearer oauth-tok-live" { + t.Fatalf("outbound Authorization=%q want decrypted OAuth Bearer (oauthInject must run)", tr.auth) + } + if tr.apiKey != "" { + t.Fatalf("x-api-key must be stripped on OAuth path, got %q", tr.apiKey) + } +} + +// Regression (live full-pipeline E2E 2026-06-25): a REAL group VK has an EMPTY +// VK-level ProviderCode/Provider — it's bound to a oauth_group, not a provider, +// so the provider lives per-account in group_accounts. The serve path must take +// the provider from the RESOLVED account, not the route, or canonicalCode="" +// yields an empty upstream base URL → 502. (handleOauthGroupRoute originally used +// rc.ProviderCode; hermetic tests set route.ProviderCode so never hit the empty +// case — only the live pipeline did.) +func TestGroupServe_EmptyRouteProviderUsesAccountProvider(t *testing.T) { + key := grKey() + refs := []vkeys.GroupAccountRef{{AccountID: "acc-1", ProviderCode: "anthropic"}} // provider only on the candidate + mat := map[string]vkeys.GroupRuntimeAccount{ + "acc-1": encMat(t, key, vkeys.GroupRuntimeAccount{ + CredentialType: "oauth_account", ExpiresAt: 9_000_000_000, ExternalID: "uuid-1", + }, "oauth-tok-live"), + } + route := &vkeys.ResolvedRoute{ + VirtualKeyID: "vk-grp", ProtocolType: "anthropic", RouteSource: "team", + // Provider + ProviderCode intentionally EMPTY — the real group-VK shape. + SeatID: "seat-1", OauthGroupID: "grp-1", + GroupAccounts: mustJSON(t, refs), GroupRuntime: mustJSON(t, mat), + } + p, tr := setupGroupProxy(t, key, route) + + req, w := groupReq(groupBody) + p.Handle(w, req) + + if w.Code != http.StatusOK { + t.Fatalf("empty-ProviderCode group route: status=%d body=%s", w.Code, w.Body.String()) + } + // Provider must come from the resolved account → real upstream + Bearer injected. + if tr.host != "api.anthropic.com" { + t.Fatalf("outbound host=%q want api.anthropic.com (provider must come from the resolved account, not the empty route)", tr.host) + } + if tr.auth != "Bearer oauth-tok-live" { + t.Fatalf("outbound Authorization=%q want OAuth Bearer", tr.auth) + } +} + +func TestGroupServe_APIKeyAccountInjectsKey(t *testing.T) { + key := grKey() + refs := []vkeys.GroupAccountRef{{AccountID: "acc-k", ProviderCode: "anthropic"}} + mat := map[string]vkeys.GroupRuntimeAccount{ + "acc-k": encMat(t, key, vkeys.GroupRuntimeAccount{ + CredentialType: "api_key", BaseURL: "https://key-upstream.example", + }, "sk-group-key"), + } + route := &vkeys.ResolvedRoute{ + VirtualKeyID: "vk-grp", Provider: "anthropic", ProtocolType: "anthropic", + ProviderCode: "anthropic", RouteSource: "team", + SeatID: "seat-1", OauthGroupID: "grp-1", + GroupAccounts: mustJSON(t, refs), GroupRuntime: mustJSON(t, mat), + } + p, tr := setupGroupProxy(t, key, route) + + req, w := groupReq(groupBody) + p.Handle(w, req) + + if w.Code != http.StatusOK { + t.Fatalf("group api_key route: status=%d body=%s", w.Code, w.Body.String()) + } + // api_key → upstream is the account's base_url, key injected via adapter. + if tr.host != "key-upstream.example" { + t.Fatalf("outbound host=%q want key-upstream.example (api_key base_url)", tr.host) + } + if tr.apiKey != "sk-group-key" { + t.Fatalf("outbound x-api-key=%q want injected group key", tr.apiKey) + } +} + +func TestGroupServe_NoMaterialDegrades503(t *testing.T) { + key := grKey() + refs := []vkeys.GroupAccountRef{{AccountID: "acc-1", ProviderCode: "anthropic"}} + route := &vkeys.ResolvedRoute{ + VirtualKeyID: "vk-grp", Provider: "anthropic", ProtocolType: "anthropic", + ProviderCode: "anthropic", RouteSource: "team", + SeatID: "seat-1", OauthGroupID: "grp-1", + GroupAccounts: mustJSON(t, refs), GroupRuntime: "", // not pulled yet + } + p, tr := setupGroupProxy(t, key, route) + + req, w := groupReq(groupBody) + p.Handle(w, req) + + if w.Code != http.StatusServiceUnavailable { + t.Fatalf("no material should degrade 503, got %d", w.Code) + } + if !strings.Contains(w.Body.String(), "GROUP_NO_MATERIAL") { + t.Fatalf("degrade body should carry GROUP_NO_MATERIAL code: %s", w.Body.String()) + } + if tr.host != "" { + t.Fatalf("degraded request must NOT reach upstream, dialed %q", tr.host) + } +} + +// ③ status-code mapping (2026-07-01): the Anthropic SDK retries 5xx/429 with backoff, +// so a PERMANENT failure sent as 503 makes claude hang for minutes. Permanent codes +// must be non-retryable 4xx (fail fast); only genuinely transient ones stay 503. +func TestGroupDegradeStatus(t *testing.T) { + cases := []struct { + code string + wantStatus int + }{ + {groupErrNoCandidates, http.StatusForbidden}, // permanent (removed / empty) → fail fast + {groupErrAllUnusable, http.StatusTooManyRequests}, // rate-limited → 429 (honest) + {groupErrNoMaterial, http.StatusServiceUnavailable}, // transient sync → retry + {"SOME_OTHER_CODE", http.StatusServiceUnavailable}, // default transient + } + for _, c := range cases { + got, _ := groupDegradeStatus(c.code) + if got != c.wantStatus { + t.Errorf("groupDegradeStatus(%s)=%d want %d", c.code, got, c.wantStatus) + } + } + // The permanent code must NOT be a retryable 5xx (that's the whole bug). + if s, _ := groupDegradeStatus(groupErrNoCandidates); s >= 500 { + t.Errorf("NO_CANDIDATES must be a non-retryable 4xx, got %d (5xx → claude backoff-hang)", s) + } +} + +// ② removed-member message fix (2026-06-30): a seat REMOVED from the group gets an +// empty channel-③ delivery — the proxy wipes group_runtime to "{}". The candidate +// snapshot (group_accounts) may still be STALE with entries. This must surface the +// "no available account — contact admin, won't self-resolve" message (NO_CANDIDATES), +// NOT the misleading "credentials still syncing, retry shortly" (NO_MATERIAL) that +// told a removed member to retry forever. +func TestGroupServe_RemovedMemberEmptyMaterialNotSyncing(t *testing.T) { + key := grKey() + // Stale snapshot still lists a candidate; material was wiped to "{}". + refs := []vkeys.GroupAccountRef{{AccountID: "acc-1", ProviderCode: "anthropic"}} + route := &vkeys.ResolvedRoute{ + VirtualKeyID: "vk-grp", Provider: "anthropic", ProtocolType: "anthropic", + ProviderCode: "anthropic", RouteSource: "team", + SeatID: "seat-1", OauthGroupID: "grp-1", + GroupAccounts: mustJSON(t, refs), GroupRuntime: "{}", // pulled → delivered nothing + } + p, tr := setupGroupProxy(t, key, route) + + req, w := groupReq(groupBody) + p.Handle(w, req) + + body := w.Body.String() + // ③ (2026-07-01): a PERMANENT failure must be a NON-retryable 4xx (403), NOT 503. + // A 503 makes the Anthropic SDK retry with exponential backoff → claude hangs for + // minutes on a condition that can never succeed (the reported bug), and renders a + // contradictory "server-side issue, try again" suffix. 403 = fail fast. + if w.Code != http.StatusForbidden { + t.Fatalf("removed member (permanent NO_CANDIDATES) must be 403 fail-fast, not %d (503 → claude retry-hang)", w.Code) + } + if strings.Contains(body, "GROUP_NO_MATERIAL") || strings.Contains(strings.ToLower(body), "still syncing") { + t.Fatalf("removed member must NOT get the transient 'still syncing' message: %s", body) + } + if !strings.Contains(body, "GROUP_NO_CANDIDATES") { + t.Fatalf("removed member (empty material) must degrade as NO_CANDIDATES: %s", body) + } + if strings.Contains(strings.ToLower(body), "retry") { + t.Fatalf("removed member must not be told to retry (won't self-resolve): %s", body) + } + if tr.host != "" { + t.Fatalf("degraded request must NOT reach upstream, dialed %q", tr.host) + } +} + +// ① login-prompt producer verification (2026-06-30): a member who has NOT logged +// into the routed pool account (master delivered the account with needs_login=true, +// so group_runtime carries the marker — NON-empty) must get a structured 401 +// OAUTH_GROUP_MEMBER_LOGIN_REQUIRED, NOT the 503 "still syncing" degrade. This is the +// user-facing side of the login flow: claude surfaces the proxy's error message, so +// this 401 body IS the "please sign in" prompt. Pairs with the resolver-level +// TestResolveGroup_LoginRequiredNoSkip. Distinguishes login-needed (marker present) +// from not-synced (empty material → TestGroupServe_NoMaterialDegrades503). +func TestGroupServe_LoginRequiredReturns401(t *testing.T) { + key := grKey() + refs := []vkeys.GroupAccountRef{{AccountID: "acc-1", ProviderCode: "anthropic"}} + mat := map[string]vkeys.GroupRuntimeAccount{ + // needs_login marker carries NO secret — master delivered the account as + // "member not logged in" (this is what makes material non-empty, so the + // resolver reaches candNeedsLogin instead of bailing on NO_MATERIAL). + "acc-1": {CredentialType: "oauth_account", NeedsLogin: true}, + } + route := &vkeys.ResolvedRoute{ + VirtualKeyID: "vk-grp", Provider: "anthropic", ProtocolType: "anthropic", + ProviderCode: "anthropic", RouteSource: "team", + SeatID: "seat-1", OauthGroupID: "grp-1", + GroupAccounts: mustJSON(t, refs), GroupRuntime: mustJSON(t, mat), + } + p, tr := setupGroupProxy(t, key, route) + + req, w := groupReq(groupBody) + p.Handle(w, req) + + if w.Code != http.StatusUnauthorized { + t.Fatalf("not-logged-in routed account must return 401 login prompt, got %d body=%s", w.Code, w.Body.String()) + } + if got := w.Header().Get(HeaderAikeyErrorSource); got != groupErrLoginRequired { + t.Fatalf("X-Aikey-Error-Source=%q want %q", got, groupErrLoginRequired) + } + body := w.Body.String() + if !strings.Contains(body, groupErrLoginRequired) { + t.Fatalf("body must carry the login-required code so the client can act: %s", body) + } + if !strings.Contains(body, "acc-1") { + t.Fatalf("body must name the account to log into: %s", body) + } + if !strings.Contains(body, "sign-in") { + t.Fatalf("body must carry the human sign-in prompt claude will display: %s", body) + } + // Must NOT be mistaken for the transient "still syncing" degrade, and must not + // reach upstream (the request is halted pending login). + if strings.Contains(body, "GROUP_NO_MATERIAL") { + t.Fatalf("login-required must NOT be a NO_MATERIAL degrade: %s", body) + } + if tr.host != "" { + t.Fatalf("login-required must not reach upstream, dialed %q", tr.host) + } +} + +// Display contract (20260703 update, spike-verified in dev2): claude only renders +// error.message verbatim for Anthropic-STANDARD error types — the previous custom +// "login_required" type fell into the generic "API error · Retrying" path (11 +// blind retries) and the member never saw the sign-in prompt. With console_url +// configured, login_url must be assembled by the PROXY (决策2: single assembly +// point) and appear both as a field and inside the human message. +func TestGroupServe_LoginRequiredDisplayContract(t *testing.T) { + key := grKey() + refs := []vkeys.GroupAccountRef{{AccountID: "acc-1", ProviderCode: "anthropic"}} + mat := map[string]vkeys.GroupRuntimeAccount{ + "acc-1": {CredentialType: "oauth_account", NeedsLogin: true}, + } + route := &vkeys.ResolvedRoute{ + VirtualKeyID: "vk-grp", Provider: "anthropic", ProtocolType: "anthropic", + ProviderCode: "anthropic", RouteSource: "team", + SeatID: "seat-1", OauthGroupID: "grp-1", + GroupAccounts: mustJSON(t, refs), GroupRuntime: mustJSON(t, mat), + } + p, _ := setupGroupProxy(t, key, route) + p.SetConsoleURL("http://127.0.0.1:8090/") // trailing slash must not double up + + req, w := groupReq(groupBody) + p.Handle(w, req) + + if w.Code != http.StatusUnauthorized { + t.Fatalf("status=%d body=%s", w.Code, w.Body.String()) + } + var resp struct { + Error struct { + Type string `json:"type"` + Code string `json:"code"` + Message string `json:"message"` + } `json:"error"` + Account string `json:"account"` + LoginURL string `json:"login_url"` + } + if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil { + t.Fatalf("401 body must stay valid JSON: %v — %s", err, w.Body.String()) + } + if resp.Error.Type != "authentication_error" { + t.Fatalf("error.type=%q want authentication_error (custom types are NOT rendered by claude)", resp.Error.Type) + } + if resp.Error.Code != groupErrLoginRequired { + t.Fatalf("error.code=%q must keep the precise machine signal %q", resp.Error.Code, groupErrLoginRequired) + } + wantURL := "http://127.0.0.1:8090/user/team-oauth" + if resp.LoginURL != wantURL { + t.Fatalf("login_url=%q want %q", resp.LoginURL, wantURL) + } + if !strings.Contains(resp.Error.Message, wantURL) { + t.Fatalf("human message must carry the clickable URL (claude shows message only): %s", resp.Error.Message) + } + + // Bypass statusline hint: the state file must exist with the SAME url + // (single assembly point) while login is pending... + statePath := filepath.Join(os.Getenv("AIKEY_RUN_DIR"), "group-login-required.json") + raw, err := os.ReadFile(statePath) + if err != nil { + t.Fatalf("state file must be written on login-required 401: %v", err) + } + var st struct { + AccountID string `json:"account_id"` + LoginURL string `json:"login_url"` + WrittenAt int64 `json:"written_at"` + } + if err := json.Unmarshal(raw, &st); err != nil { + t.Fatalf("state file must be valid JSON: %v — %s", err, raw) + } + if st.AccountID != "acc-1" || st.LoginURL != wantURL || st.WrittenAt == 0 { + t.Fatalf("state file content mismatch: %+v", st) + } + + // ...and be CLEARED by the next successful group resolve (member logged in), + // so statusline recovery is automatic — a stale hint would nag forever. + mat["acc-1"] = encMat(t, key, vkeys.GroupRuntimeAccount{ + CredentialType: "oauth_account", ExpiresAt: 9_000_000_000, ExternalID: "uuid-1", + }, "oauth-tok-live") + route.GroupRuntime = mustJSON(t, mat) + req2, w2 := groupReq(groupBody) + p.Handle(w2, req2) + if w2.Code != http.StatusOK { + t.Fatalf("post-login request must succeed, got %d body=%s", w2.Code, w2.Body.String()) + } + if _, err := os.Stat(statePath); !os.IsNotExist(err) { + t.Fatalf("state file must be cleared after a successful group resolve, stat err=%v", err) + } +} + +// Empty console_url (cluster node / server-side proxy — no co-installed local +// console) must degrade to URL-less wording, never a broken half-URL, and the +// response must stay a well-formed 401 (main-path robustness). +func TestGroupServe_LoginRequiredNoConsoleURLFallback(t *testing.T) { + key := grKey() + refs := []vkeys.GroupAccountRef{{AccountID: "acc-1", ProviderCode: "anthropic"}} + mat := map[string]vkeys.GroupRuntimeAccount{ + "acc-1": {CredentialType: "oauth_account", NeedsLogin: true}, + } + route := &vkeys.ResolvedRoute{ + VirtualKeyID: "vk-grp", Provider: "anthropic", ProtocolType: "anthropic", + ProviderCode: "anthropic", RouteSource: "team", + SeatID: "seat-1", OauthGroupID: "grp-1", + GroupAccounts: mustJSON(t, refs), GroupRuntime: mustJSON(t, mat), + } + p, _ := setupGroupProxy(t, key, route) // console URL deliberately unset + + req, w := groupReq(groupBody) + p.Handle(w, req) + + if w.Code != http.StatusUnauthorized { + t.Fatalf("status=%d body=%s", w.Code, w.Body.String()) + } + var resp struct { + Error struct { + Type string `json:"type"` + Message string `json:"message"` + } `json:"error"` + LoginURL string `json:"login_url"` + } + if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil { + t.Fatalf("401 body must stay valid JSON: %v — %s", err, w.Body.String()) + } + if resp.LoginURL != "" { + t.Fatalf("login_url must be empty without console_url, got %q", resp.LoginURL) + } + if resp.Error.Type != "authentication_error" { + t.Fatalf("fallback must keep the displayable type, got %q", resp.Error.Type) + } + if strings.Contains(resp.Error.Message, "http://") || strings.Contains(resp.Error.Message, "https://") { + t.Fatalf("URL-less fallback must not leak a half-assembled URL: %s", resp.Error.Message) + } + // Even URL-less, the message must still point at the console page path so + // the member can find it manually. + if !strings.Contains(resp.Error.Message, "/user/team-oauth") { + t.Fatalf("fallback message must name the console page: %s", resp.Error.Message) + } +} + +// Transparent proxy (2026-06-29): after AccountPersona removal, an oauth_group +// pool request forwards the REAL client identity (session/OS) upstream UNCHANGED +// — no per-account synthetic device/session. Guards against re-introducing the +// forgery this code path used to apply (see oauth_inject.go NOTE for the WHY). +func TestGroupServe_PoolPassesRealIdentity(t *testing.T) { + key := grKey() + refs := []vkeys.GroupAccountRef{{AccountID: "acc-pool", ProviderCode: "anthropic"}} + mat := map[string]vkeys.GroupRuntimeAccount{ + "acc-pool": encMat(t, key, vkeys.GroupRuntimeAccount{ + CredentialType: "oauth_account", ExpiresAt: 9_000_000_000, ExternalID: "uuid-pool", + }, "tok-pool"), + } + route := &vkeys.ResolvedRoute{ + VirtualKeyID: "vk-grp", Provider: "anthropic", ProtocolType: "anthropic", + ProviderCode: "anthropic", RouteSource: "team", + SeatID: "seat-1", OauthGroupID: "grp-1", + GroupAccounts: mustJSON(t, refs), GroupRuntime: mustJSON(t, mat), + } + p, tr := setupGroupProxy(t, key, route) + + req, w := groupReq(groupBody) + req.Header.Set("X-Claude-Code-Session-Id", "REAL-SESSION") // a real employee session + req.Header.Set("X-Stainless-OS", "Windows") + p.Handle(w, req) + + if w.Code != http.StatusOK { + t.Fatalf("status=%d body=%s", w.Code, w.Body.String()) + } + // Upstream (outbound) must carry the REAL client identity, NOT a synthetic + // per-account one — the whole point of dropping the disguise. + if tr.session != "REAL-SESSION" { + t.Fatalf("upstream session must pass through real value unchanged, got %q", tr.session) + } + if tr.stainlessOS != "Windows" { + t.Fatalf("upstream OS must pass through real value unchanged, got %q", tr.stainlessOS) + } +} + +// N8c: a pool account whose upstream returns 401 is cooled down, and a later +// request routes around it (here the only candidate → GROUP_ALL_UNUSABLE 429). +func TestGroupServe_CooldownOn401(t *testing.T) { + key := grKey() + refs := []vkeys.GroupAccountRef{{AccountID: "acc-1", ProviderCode: "anthropic"}} + mat := map[string]vkeys.GroupRuntimeAccount{ + "acc-1": encMat(t, key, vkeys.GroupRuntimeAccount{ + CredentialType: "oauth_account", ExpiresAt: 9_000_000_000, ExternalID: "uuid-1", + }, "tok-1"), + } + route := &vkeys.ResolvedRoute{ + VirtualKeyID: "vk-grp", Provider: "anthropic", ProtocolType: "anthropic", + ProviderCode: "anthropic", RouteSource: "team", + SeatID: "seat-1", OauthGroupID: "grp-1", + GroupAccounts: mustJSON(t, refs), GroupRuntime: mustJSON(t, mat), + } + p, tr := setupGroupProxy(t, key, route) + tr.status = http.StatusUnauthorized // upstream says the token is broken + + req, w := groupReq(groupBody) + p.Handle(w, req) + if w.Code != http.StatusUnauthorized { + t.Fatalf("first request should forward upstream 401, got %d", w.Code) + } + if !p.poolCooldown.skipSet()["acc-1"] { + t.Fatal("a 401 upstream must cool the account down") + } + + // Second request: the only candidate is cooling down → no usable account. + tr.status = 0 + req2, w2 := groupReq(groupBody) + p.Handle(w2, req2) + // ALL_UNUSABLE = a genuine rate-limit (recovers when the window resets) → 429, + // not 503: honest code, and it's the client's call whether to back off + retry. + if w2.Code != http.StatusTooManyRequests || !strings.Contains(w2.Body.String(), "GROUP_ALL_UNUSABLE") { + t.Fatalf("cooled-down sole account must yield GROUP_ALL_UNUSABLE 429, got %d: %s", w2.Code, w2.Body.String()) + } +} + +// N8c discrimination: a WAF business-rejection 429 (no rate-limit signal) must +// NOT cool the account down — it's the request persona, not the account. +func TestGroupServe_NoCooldownOnWAF429(t *testing.T) { + key := grKey() + refs := []vkeys.GroupAccountRef{{AccountID: "acc-1", ProviderCode: "anthropic"}} + mat := map[string]vkeys.GroupRuntimeAccount{ + "acc-1": encMat(t, key, vkeys.GroupRuntimeAccount{ + CredentialType: "oauth_account", ExpiresAt: 9_000_000_000, ExternalID: "uuid-1", + }, "tok-1"), + } + route := &vkeys.ResolvedRoute{ + VirtualKeyID: "vk-grp", Provider: "anthropic", ProtocolType: "anthropic", + ProviderCode: "anthropic", RouteSource: "team", + SeatID: "seat-1", OauthGroupID: "grp-1", + GroupAccounts: mustJSON(t, refs), GroupRuntime: mustJSON(t, mat), + } + p, tr := setupGroupProxy(t, key, route) + tr.status = http.StatusTooManyRequests // 429 with NO rate-limit header → WAF + + req, w := groupReq(groupBody) + p.Handle(w, req) + if w.Code != http.StatusTooManyRequests { + t.Fatalf("first request should forward the 429, got %d", w.Code) + } + if p.poolCooldown.skipSet()["acc-1"] { + t.Fatal("WAF 429 (no rate-limit signal) must NOT cool the account") + } +} + +// §5.5 hard cap: a seat the engine marked Blocked (every pool account at the +// ≤3-人/号 cap) gets 429 — the proxy must NOT fall back to its cap-blind local pick. +func TestGroupServe_BlockedSeatReturns429(t *testing.T) { + key := grKey() + refs := []vkeys.GroupAccountRef{{AccountID: "acc-1", ProviderCode: "anthropic"}} + mat := map[string]vkeys.GroupRuntimeAccount{ + "acc-1": encMat(t, key, vkeys.GroupRuntimeAccount{ + CredentialType: "oauth_account", ExpiresAt: 9_000_000_000, ExternalID: "uuid-1", + }, "tok-1"), + } + route := &vkeys.ResolvedRoute{ + VirtualKeyID: "vk-grp", Provider: "anthropic", ProtocolType: "anthropic", + ProviderCode: "anthropic", RouteSource: "team", + SeatID: "seat-1", OauthGroupID: "grp-1", + GroupAccounts: mustJSON(t, refs), GroupRuntime: mustJSON(t, mat), + } + p, tr := setupGroupProxy(t, key, route) + + // Engine left seat-1 unbound (pool full) → proxy must 429, never route to acc-1. + cache := NewRoutingOverrideCache() + cache.StoreAll(1, nil, map[string]bool{routeKey("seat-1", "grp-1"): true}) // composite (seat,group) key + p.SetRoutingOverrides(cache) + + req, w := groupReq(groupBody) + p.Handle(w, req) + + if w.Code != http.StatusTooManyRequests { + t.Fatalf("blocked seat must 429, got %d body=%s", w.Code, w.Body.String()) + } + if !strings.Contains(w.Body.String(), "GROUP_POOL_FULL") { + t.Fatalf("429 body must carry GROUP_POOL_FULL code: %s", w.Body.String()) + } + if tr.host != "" { + t.Fatalf("blocked request must NOT reach upstream, dialed %q", tr.host) + } +} + +// Byte-unchanged guard: a non-group (direct-bind) team route must NOT enter the +// group path — it forwards via the static-key path exactly as before. +func TestGroupServe_DirectBindUnaffected(t *testing.T) { + key := grKey() + route := &vkeys.ResolvedRoute{ + VirtualKeyID: "vk-direct", Provider: "anthropic", ProtocolType: "anthropic", + ProviderCode: "anthropic", RouteSource: "team", + BaseURL: "https://direct.example", PlaintextKey: "sk-direct-key", + // OauthGroupID empty → group branch skipped. + } + p, tr := setupGroupProxy(t, key, route) + + req, w := groupReq(groupBody) + p.Handle(w, req) + + if w.Code != http.StatusOK { + t.Fatalf("direct-bind team route: status=%d body=%s", w.Code, w.Body.String()) + } + if tr.host != "direct.example" { + t.Fatalf("direct-bind must use its own BaseURL, got %q", tr.host) + } + if tr.apiKey != "sk-direct-key" { + t.Fatalf("direct-bind key injection changed: outbound x-api-key=%q", tr.apiKey) + } +} + +// TestGroupServe_FallbackAttributesToServedAccount verifies the audit/usage +// attribution rule for account switching: when the proxy falls back A→B (account +// A returns 401 and is cooled, account B serves the next request), the RECORDED +// UsageEvent's ACCOUNT attribution must follow the actually-serving account (→B), +// while the USER attribution (VirtualKeyID) stays unchanged. Covers full-test-plan +// §2.2#8 (命中=usage attribution 归账号) + the requirement "切号切审计账号归属、用户归属不变". +// +// Live-event acceptance (e2e-acceptance-live-events): drives two real requests +// through Handle → group resolve → fallback → collector → store, then asserts the +// data the pipeline actually recorded — not just HTTP status. +func TestGroupServe_FallbackAttributesToServedAccount(t *testing.T) { + key := grKey() + refs := []vkeys.GroupAccountRef{ + {AccountID: "acc-1", ProviderCode: "anthropic", Identity: "a1@pool.test"}, + {AccountID: "acc-2", ProviderCode: "anthropic", Identity: "a2@pool.test"}, + } + identityOf := map[string]string{"acc-1": "a1@pool.test", "acc-2": "a2@pool.test"} + mat := map[string]vkeys.GroupRuntimeAccount{ + "acc-1": encMat(t, key, vkeys.GroupRuntimeAccount{CredentialType: "oauth_account", ExpiresAt: 9_000_000_000, ExternalID: "uuid-1"}, "tok-1"), + "acc-2": encMat(t, key, vkeys.GroupRuntimeAccount{CredentialType: "oauth_account", ExpiresAt: 9_000_000_000, ExternalID: "uuid-2"}, "tok-2"), + } + route := &vkeys.ResolvedRoute{ + VirtualKeyID: "vk-grp", Provider: "anthropic", ProtocolType: "anthropic", + ProviderCode: "anthropic", RouteSource: "team", + SeatID: "seat-1", OauthGroupID: "grp-1", + GroupAccounts: mustJSON(t, refs), GroupRuntime: mustJSON(t, mat), + } + // Map each account's at-rest token back to its account id so we can read the + // served account off the (synchronous) outbound capture, regardless of which + // account seatassign picks as rank-0. + tokToAcct := map[string]string{"tok-1": "acc-1", "tok-2": "acc-2"} + + store := newCapturingStore() + p := setupTestProxyWithStore(t, "http://unused.invalid", store) + p.registry.Merge(map[string]*vkeys.ResolvedRoute{"aikey_team_grouptest": route}) + p.SetGroupKeyProvider(fakeGroupKey{k: key}) + tr := &outboundCapture{} + p.SetTransport(tr) + // WAL capture of the REPORTED wire event (ReportableEvent) — the shape the + // collector → DWD → usage-audit page actually consumes. + walDir := t.TempDir() + wal, err := events.NewWALWriter(walDir) + if err != nil { + t.Fatalf("NewWALWriter: %v", err) + } + t.Cleanup(func() { _ = wal.Close() }) + p.SetWAL(wal) + p.SetReporter(nil, "proxy-grp", "test", "gen-grp", 0, "acc-grp") + + // Request 1: the primary account serves; upstream 401 cools it down. + tr.status = http.StatusUnauthorized + req1, w1 := groupReq(groupBody) + p.Handle(w1, req1) + served1 := tokToAcct[strings.TrimPrefix(tr.auth, "Bearer ")] + if served1 == "" { + t.Fatalf("req1 outbound auth %q did not map to a known account", tr.auth) + } + if !p.poolCooldown.skipSet()[served1] { + t.Fatalf("401 must cool the serving account %s", served1) + } + + // Request 2: primary cooling → the proxy switches to the OTHER account. + tr.status = 0 + req2, w2 := groupReq(groupBody) + p.Handle(w2, req2) + if w2.Code != http.StatusOK { + t.Fatalf("fallback request should succeed via the other account, got %d: %s", w2.Code, w2.Body.String()) + } + served2 := tokToAcct[strings.TrimPrefix(tr.auth, "Bearer ")] + if served2 == "" || served2 == served1 { + t.Fatalf("req2 must serve via a DIFFERENT account; served1=%s served2=%s (auth=%q)", served1, served2, tr.auth) + } + + // AUDIT assertion: poll the recorded events for the 200 (fallback) request and + // assert account归属 followed the switch while user归属 stayed stable. + var got events.UsageEvent + found := false + for i := 0; i < 500 && !found; i++ { + store.mu.Lock() + for j := range store.events { + if store.events[j].StatusCode == http.StatusOK && store.events[j].AccountID == served2 { + got = store.events[j] + found = true + } + } + store.mu.Unlock() + if found { + break + } + time.Sleep(10 * time.Millisecond) + } + if !found { + t.Fatalf("no 200 usage event attributed to switched-to account %s was recorded", served2) + } + // 账号归属 = the account that actually served (B), NOT the cooled primary (A). + if got.AccountID == served1 { + t.Fatalf("audit account归属 wrongly stuck on cooled account %s instead of served %s", served1, served2) + } + // 用户归属 stable across the switch. + if got.VirtualKeyID != "vk-grp" { + t.Fatalf("user归属 must stay stable across switch: VirtualKeyID=%q want vk-grp", got.VirtualKeyID) + } + // Point-in-time audit identity (2026-07-01, usage-audit "selected account"): the + // REPORTED wire event (ReportableEvent → collector → DWD → usage-audit page) must + // carry the SERVING account's email as oauth_identity — denormalized at event time + // so the audit page shows who served even after rename/removal, and it must follow + // the switch (B's identity, not the cooled A's). 能红: drop the group_serve + // `rc.OAuthIdentity = res.Identity` stamp → this is "" → fails. + _ = wal.Close() + entry := readLastWALEntry(t, walDir) + if entry.EventJSON.OAuthIdentity != identityOf[served2] { + t.Fatalf("wire oauth_identity=%q want the SERVING account's identity %q (selected-account audit display)", + entry.EventJSON.OAuthIdentity, identityOf[served2]) + } + if entry.EventJSON.AccountID != served2 { + t.Fatalf("wire account_id=%q want served account %q", entry.EventJSON.AccountID, served2) + } +} + +// TestGroupDegradeMessage locks the per-error-code 503 guidance (finding: all +// group failures collapsed into one "retry shortly" message, misleading a member +// whose access is permanently gone). NO_CANDIDATES must NOT tell the user to +// retry; NO_MATERIAL is the only transient/retryable one; all three differ. +func TestGroupDegradeMessage(t *testing.T) { + noCand := groupDegradeMessage(groupErrNoCandidates) + noMat := groupDegradeMessage(groupErrNoMaterial) + allUnusable := groupDegradeMessage(groupErrAllUnusable) + + // NO_CANDIDATES = permanent until an admin acts → must not say "retry". + if strings.Contains(strings.ToLower(noCand), "retry") { + t.Errorf("NO_CANDIDATES message must not tell the user to retry: %q", noCand) + } + // NO_MATERIAL = transient → should invite a retry. + if !strings.Contains(strings.ToLower(noMat), "retry") { + t.Errorf("NO_MATERIAL message should invite a retry: %q", noMat) + } + // All three must be distinct (no collapse to a single generic line). + if noCand == noMat || noCand == allUnusable || noMat == allUnusable { + t.Errorf("group degrade messages must differ per code:\n noCand=%q\n noMat=%q\n allUnusable=%q", noCand, noMat, allUnusable) + } +} + +// SyncRail truthful wording (§5.4, 2026-07-03 incident): when the engine's +// assignment rail is STALE/OFFLINE, the pick behind this 401 came from the +// LOCAL ranked fallback and may contradict the engine (the member may already +// be signed into the account the engine actually routed them to). The 401 must +// then say "routing sync unreachable" — not direct the member to sign into a +// possibly-wrong account — and carry the machine-readable reason. A healthy +// rail keeps the normal sign-in prompt with NO reason field (additive contract). +func TestGroupServe_LoginRequiredRailStateWording(t *testing.T) { + key := grKey() + refs := []vkeys.GroupAccountRef{{AccountID: "acc-1", ProviderCode: "anthropic"}} + mat := map[string]vkeys.GroupRuntimeAccount{ + "acc-1": {CredentialType: "oauth_account", NeedsLogin: true}, + } + route := &vkeys.ResolvedRoute{ + VirtualKeyID: "vk-grp", Provider: "anthropic", ProtocolType: "anthropic", + ProviderCode: "anthropic", RouteSource: "team", + SeatID: "seat-1", OauthGroupID: "grp-1", + GroupAccounts: mustJSON(t, refs), GroupRuntime: mustJSON(t, mat), + } + + decode := func(w *httptest.ResponseRecorder) (string, string) { + var resp struct { + Error struct { + Message string `json:"message"` + Reason string `json:"reason"` + } `json:"error"` + } + if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil { + t.Fatalf("401 body must stay valid JSON: %v — %s", err, w.Body.String()) + } + return resp.Error.Message, resp.Error.Reason + } + + // Healthy rail (probe says ok): normal sign-in prompt, no reason field. + p, _ := setupGroupProxy(t, key, route) + p.SetConsoleURL("http://127.0.0.1:8090") + p.SetRoutingRailHealth(func() (string, int64) { return "ok", 0 }) + req, w := groupReq(groupBody) + p.Handle(w, req) + if w.Code != http.StatusUnauthorized { + t.Fatalf("status=%d body=%s", w.Code, w.Body.String()) + } + msg, reason := decode(w) + if !strings.Contains(msg, "sign-in") || reason != "" { + t.Fatalf("healthy rail must keep the sign-in prompt with no reason: msg=%q reason=%q", msg, reason) + } + + // Degraded rail: truthful wording + machine-readable reason; code and header + // unchanged (additive-only contract). + p2, _ := setupGroupProxy(t, key, route) + p2.SetConsoleURL("http://127.0.0.1:8090") + p2.SetRoutingRailHealth(func() (string, int64) { return "offline", 23 * 60 }) + req2, w2 := groupReq(groupBody) + p2.Handle(w2, req2) + if w2.Code != http.StatusUnauthorized { + t.Fatalf("degraded-rail status=%d body=%s", w2.Code, w2.Body.String()) + } + msg2, reason2 := decode(w2) + if reason2 != "routing_sync_unavailable" { + t.Fatalf("degraded rail must set reason=routing_sync_unavailable, got %q", reason2) + } + if !strings.Contains(msg2, "unreachable for 23 min") { + t.Fatalf("degraded wording must state the outage duration: %q", msg2) + } + if strings.Contains(msg2, "complete sign-in, then retry") { + t.Fatalf("degraded wording must NOT blindly direct a sign-in: %q", msg2) + } + if got := w2.Header().Get(HeaderAikeyErrorSource); got != groupErrLoginRequired { + t.Fatalf("header signal must stay %q, got %q", groupErrLoginRequired, got) + } + + // nil probe (framework off / older wiring): behaves as healthy. + p3, _ := setupGroupProxy(t, key, route) + p3.SetConsoleURL("http://127.0.0.1:8090") + req3, w3 := groupReq(groupBody) + p3.Handle(w3, req3) + msg3, reason3 := decode(w3) + if !strings.Contains(msg3, "sign-in") || reason3 != "" { + t.Fatalf("nil probe must behave as healthy: msg=%q reason=%q", msg3, reason3) + } +} diff --git a/internal/proxy/handle_dispatch.go b/internal/proxy/handle_dispatch.go index fbd8b14..cf3ebcc 100644 --- a/internal/proxy/handle_dispatch.go +++ b/internal/proxy/handle_dispatch.go @@ -224,6 +224,19 @@ func (p *Proxy) Handle(w http.ResponseWriter, r *http.Request) { return } + // 3c. Oauth-group routing (N8). A group VK carries no static key — its + // per-account material is in route.GroupRuntime; pick a candidate account + + // inject its credential via the dedicated handler. Gated on the field, not a + // re-read of the feature flag: group VKs are only ever registered when the + // flag is on (N7c-1), so route.OauthGroupID is empty in the flag-off build and + // the direct-bind path below stays byte-identical. A registered group VK MUST + // be served as a group — never fall through to the static-key path (that + // misroute is exactly what the registration gate prevents). + if route.OauthGroupID != "" { + p.handleOauthGroupRoute(w, r, route, token, startTime, logger, tc.TraceID) + return + } + // 4. Get real key — either from the pre-decrypted managed cache or from vault. var realKey string if route.PlaintextKey != "" { diff --git a/internal/proxy/main_test.go b/internal/proxy/main_test.go new file mode 100644 index 0000000..ce006d7 --- /dev/null +++ b/internal/proxy/main_test.go @@ -0,0 +1,24 @@ +package proxy + +import ( + "os" + "testing" +) + +// TestMain sandboxes AIKEY_RUN_DIR for the WHOLE package. Proxy instances +// write/remove the bypass ~/.aikey/run/group-login-required.json state file +// (group_login_state.go) on group-route 401s/successes; without this override +// every `go test` on a developer machine would transiently touch the REAL +// ~/.aikey/run — racing a live proxy's statusline hint (a success-path test +// would silently delete a genuine login prompt). Package-wide (not per-test +// t.Setenv) so future tests can't forget it. +func TestMain(m *testing.M) { + dir, err := os.MkdirTemp("", "aikey-proxy-test-run-*") + if err != nil { + panic(err) + } + _ = os.Setenv("AIKEY_RUN_DIR", dir) + code := m.Run() + _ = os.RemoveAll(dir) + os.Exit(code) +} diff --git a/internal/proxy/middleware.go b/internal/proxy/middleware.go index a7a5878..018e2c9 100644 --- a/internal/proxy/middleware.go +++ b/internal/proxy/middleware.go @@ -3,6 +3,7 @@ package proxy import ( "context" "net/http" + "os" "strings" "github.com/AiKeyLabs/aikey-proxy/internal/observability" @@ -44,6 +45,11 @@ const ( // stashExtractedFields), so the security-critical first read is never // skipped; only repeat calls hit this cache. ctxKeyExtractedModel + // ctxKeyPoolWindowCap carries the chosen pool account's window_max_util_pct + // (int, N11's randomized pre-cut cap), stashed at resolution and read in + // serveRoute's ModifyResponse to pre-cut the account when the upstream's + // unified-utilization header crosses the cap (N10 防封). Absent → no pre-cut. + ctxKeyPoolWindowCap ) // traceFromContext retrieves the request's TraceContext. Returns the zero value @@ -185,9 +191,59 @@ func providerToProtocol(providerCode string) string { // // Last synced (2026-04-24): added P0 (groq / xai / openrouter / perplexity) // + P1 (zhipu / qwen / doubao / siliconflow) alongside the original 6. +// codexUpstreamBaseURL is the Codex OAuth upstream — chatgpt.com/backend-api/codex +// (Responses API), NOT api.openai.com/v1 (that's the API-key path). Both OAuth +// dispatch sites hardcoded this string; centralized here (single source) so they +// can't drift AND so an E2E can redirect it to a local mock. +// +// Test-only hook (loopback-gated, same posture as providerDefaultBaseURL's +// AIKEY_PROXY_TEST_ANTHROPIC_BASE_URL): Codex OAuth carries no configurable +// base_url, so the codex-account routing E2E can't exercise the inject path +// against a mock without this. The loopback guard means a prod misconfig can +// never reroute real traffic. See aikey-test/oauthgroup/codex_account_routing_test.go. +func codexUpstreamBaseURL() string { + if o := os.Getenv("AIKEY_PROXY_TEST_CODEX_BASE_URL"); o != "" && + (strings.HasPrefix(o, "http://127.0.0.1:") || strings.HasPrefix(o, "http://localhost:")) { + return o + } + return "https://chatgpt.com/backend-api/codex" +} + +// resolveOAuthUpstream selects the upstream base URL for an OAuth-credential +// request AND applies any provider-specific request setup, returning the +// (possibly re-wrapped) request. Centralizes the per-provider OAuth-upstream +// policy so the two dispatch sites (legacy /v1 forward_and_resolve + the group +// route in group_serve) share ONE source instead of each carrying its own +// `if canonicalCode == "openai"` branch — a new provider whose OAuth upstream +// differs from its API-key default (like codex) adds one case here, not two. +// +// codex is the one provider whose OAuth base ≠ its API-key base: OAuth hits +// chatgpt.com/backend-api/codex (Responses API), API keys hit api.openai.com/v1. +// It also needs deferred model capture (captureCodexModel returns a NEW request +// carrying the model in context; the caller must use the returned request). +// Every other provider's OAuth base == providerDefaultBaseURL and needs no setup. +func resolveOAuthUpstream(canonicalCode string, r *http.Request) (baseURL string, req *http.Request) { + switch canonicalCode { + case "openai": + return codexUpstreamBaseURL(), captureCodexModel(r) + default: + return providerDefaultBaseURL(canonicalCode), r + } +} + func providerDefaultBaseURL(providerCode string) string { switch strings.ToLower(providerCode) { case "anthropic", "claude": + // Test-only hook (gated to loopback): the cross-component OAuth-account + // routing E2E points the otherwise-hardcoded Anthropic upstream at a local + // mock. OAuth accounts carry no configurable base_url (unlike api_key + // material), so without this the OAuth inject path can't be exercised + // against a mock. The loopback guard means a prod misconfig can never + // reroute real traffic. See aikey-test/oauthgroup/oauth_account_routing_test.go. + if o := os.Getenv("AIKEY_PROXY_TEST_ANTHROPIC_BASE_URL"); o != "" && + (strings.HasPrefix(o, "http://127.0.0.1:") || strings.HasPrefix(o, "http://localhost:")) { + return o + } return "https://api.anthropic.com" case "openai", "gpt", "chatgpt", "codex": // Why: OpenAI SDK clients (including Codex) treat base_url as already diff --git a/internal/proxy/oauth_inject.go b/internal/proxy/oauth_inject.go index 77d1e3e..df4829d 100644 --- a/internal/proxy/oauth_inject.go +++ b/internal/proxy/oauth_inject.go @@ -177,6 +177,13 @@ func injectClaudeOAuth(req *http.Request, cred *OAuthCredential) { // stored on req.Context(). rewriteToolNamesForward(req) } + + // NOTE: oauth_group POOL requests pass the REAL client identity upstream + // unchanged. The former AccountPersona normalization (N users → one frozen + // SHA256(accountID) identity) was removed 2026-06-29: a frozen synthetic + // fingerprint can't track Claude Code version bumps, so it drifts stale and + // becomes its OWN detection signal — worse than transparent pass-through. + // Risk is now bounded by a ≤3-users-per-account cap, not by disguise. } // claudeCodeSystemPrompt is the byte-exact 57-char marker that Anthropic's @@ -256,16 +263,36 @@ func hasFingerprint(block map[string]any) bool { // injectCodexOAuth injects Codex (ChatGPT Plus/Pro) headers. // // Verified 2026-04-15: -// - originator: opencode (required by Codex API) +// - originator: codex_cli_rs (the OFFICIAL codex CLI value; was "opencode", +// inherited from opencode's codex.ts). Source of truth: openai/codex +// codex-rs/login/src/auth/default_client.rs DEFAULT_ORIGINATOR="codex_cli_rs" +// (sent on every /responses request via add_originator_header). This is only +// a FALLBACK: setIfAbsent preserves the value a real codex CLI already sends, +// so we merely stop mislabelling originator-less traffic as the third-party +// "opencode" tool. NOTE: this is a required API header, NOT a persona/UA +// fingerprint — we deliberately do NOT inject a synthetic User-Agent/session +// here (that identity-forgery layer was removed 2026-06-29 for transparent +// proxy + ≤3-users/account; see CC账号池 requirement). // - ChatGPT-Account-Id: from JWT claims (for org subscriptions) // - API URL: chatgpt.com/backend-api/codex/responses (Responses API format) // // Why inject-if-absent: Codex CLI may set its own originator or account ID // in future versions. Overwriting would break forward compatibility. +// +// ChatGPT-Account-Id source (fixed 2026-07-04, R34 codex pools): the real +// chatgpt_account_id lives in ExternalID (the broker extracts it from the token +// JWT's auth.chatgpt_account_id claim). AccountID is OUR identifier — the +// broker's random account row id on the personal path, the pool's internal +// account id on the group path — sending it upstream was a wrong value that only +// went unnoticed because Codex CLI usually sets its own header (setIfAbsent). +// ExternalID first; AccountID kept as a legacy fallback for old vault rows that +// predate ExternalID. func injectCodexOAuth(req *http.Request, cred *OAuthCredential) { req.Header.Set("Authorization", "Bearer "+cred.AccessToken) - setIfAbsent(req, "originator", "opencode") - if cred.AccountID != "" { + setIfAbsent(req, "originator", "codex_cli_rs") + if id := cred.ExternalID; id != "" { + setIfAbsent(req, "ChatGPT-Account-Id", id) + } else if cred.AccountID != "" { setIfAbsent(req, "ChatGPT-Account-Id", cred.AccountID) } } diff --git a/internal/proxy/oauth_inject_test.go b/internal/proxy/oauth_inject_test.go index 3a140bb..8868484 100644 --- a/internal/proxy/oauth_inject_test.go +++ b/internal/proxy/oauth_inject_test.go @@ -71,7 +71,7 @@ func TestOAuthInject_DispatchesByProvider(t *testing.T) { wantValue string }{ {"anthropic", "Bearer oauth-tok", "anthropic-version", "2023-06-01"}, - {"openai", "Bearer oauth-tok", "originator", "opencode"}, + {"openai", "Bearer oauth-tok", "originator", "codex_cli_rs"}, {"kimi", "Bearer oauth-tok", "X-Msh-Platform", "kimi_cli"}, {"moonshot", "Bearer oauth-tok", "X-Msh-Platform", "kimi_cli"}, {"unknown", "Bearer oauth-tok", "", ""}, @@ -275,19 +275,32 @@ func TestInjectClaudeOAuth_MetadataUserID(t *testing.T) { // --- injectCodexOAuth --- func TestInjectCodexOAuth(t *testing.T) { - t.Run("sets required headers", func(t *testing.T) { + t.Run("sets required headers (ExternalID → ChatGPT-Account-Id)", func(t *testing.T) { req := httptest.NewRequest("POST", "/responses", nil) - cred := &OAuthCredential{AccessToken: "codex-tok", AccountID: "acct-456"} + // R34 fix (2026-07-04): the real chatgpt_account_id is ExternalID (broker + // extracts it from the token JWT). AccountID is OUR internal id — on the + // pool path it's the pool's internal account id, WRONG to send upstream. + cred := &OAuthCredential{AccessToken: "codex-tok", ExternalID: "chatgpt-uuid-789", AccountID: "internal-pool-id"} injectCodexOAuth(req, cred) if got := req.Header.Get("Authorization"); got != "Bearer codex-tok" { t.Errorf("Authorization = %q", got) } - if got := req.Header.Get("originator"); got != "opencode" { - t.Errorf("originator = %q, want %q", got, "opencode") + if got := req.Header.Get("originator"); got != "codex_cli_rs" { + t.Errorf("originator = %q, want %q", got, "codex_cli_rs") + } + if got := req.Header.Get("ChatGPT-Account-Id"); got != "chatgpt-uuid-789" { + t.Errorf("ChatGPT-Account-Id = %q, want ExternalID %q (not internal AccountID)", got, "chatgpt-uuid-789") } + }) + + t.Run("falls back to AccountID when ExternalID empty (legacy vault rows)", func(t *testing.T) { + req := httptest.NewRequest("POST", "/responses", nil) + cred := &OAuthCredential{AccessToken: "codex-tok", AccountID: "acct-456"} + injectCodexOAuth(req, cred) + if got := req.Header.Get("ChatGPT-Account-Id"); got != "acct-456" { - t.Errorf("ChatGPT-Account-Id = %q, want %q", got, "acct-456") + t.Errorf("ChatGPT-Account-Id = %q, want AccountID fallback %q", got, "acct-456") } }) diff --git a/internal/proxy/oauth_pool_cooldown.go b/internal/proxy/oauth_pool_cooldown.go new file mode 100644 index 0000000..f1b3611 --- /dev/null +++ b/internal/proxy/oauth_pool_cooldown.go @@ -0,0 +1,337 @@ +// oauth_pool_cooldown.go — N8c: reactive pool-account fallback. +// +// N8a already does PRE-request avoidance (skips accounts the delivered material +// marks exhausted). N8c adds the REACTIVE layer: when a chosen account's UPSTREAM +// response says the account is broken (401) or its window is used up (a real +// exhaustion 429), cool it down so the resolver's `skip` set routes subsequent +// requests around it. The failing request still returns its status to the client +// — in-request retry (no client-visible failure) is the heavier N9 work. +// +// 429 discrimination (minimal, research-backed; full 三分类 is N9): a 429 WITH a +// rate-limit signal (anthropic-ratelimit-* / Retry-After) is a real quota/limit +// → cool down. A 429 WITHOUT those headers is Anthropic's WAF business-rejection +// (the request persona is wrong, NOT the account's fault) → do NOT cool down, +// or we'd waste a good account. Same signature the OAuth-inject research uses. +package proxy + +import ( + "encoding/json" + "log/slog" + "net/http" + "os" + "path/filepath" + "strconv" + "strings" + "sync" + "time" + + "github.com/AiKeyLabs/aikey-proxy/internal/observability" +) + +const ( + // poolCooldownDefault is used when the upstream gives no Retry-After hint. + // Short enough to recover soon after a 5h window resets, long enough to not + // hammer a broken account every few seconds. + poolCooldownDefault = 5 * time.Minute + // poolCooldownMax caps a server-provided Retry-After so a hostile/huge value + // can't lock an account out for an unreasonable time. + poolCooldownMax = 1 * time.Hour +) + +// poolCooldownStore holds a per-account "avoid until" time. Concurrency-safe. +// Bounded by the number of distinct pool accounts; lapsed entries are dropped +// lazily on read. +type poolCooldownStore struct { + mu sync.Mutex + m map[string]time.Time // accountID → avoid-until + now func() time.Time // injectable clock (tests) +} + +func newPoolCooldownStore() *poolCooldownStore { + s := &poolCooldownStore{m: make(map[string]time.Time), now: time.Now} + // Cross-restart persistence (2026-07-04 self-heal, §S4): without it a proxy + // restart forgot every cooldown and could immediately route traffic back + // onto an account that just 401'd / rate-limited. STRICTLY an enhancement, + // never a dependency (owner constraint): a missing/corrupt/unreadable file + // falls back to an empty store and the main link proceeds untouched. + s.hydrateFromFile() + return s +} + +// ── cross-restart persistence (bypass state file, §S4 2026-07-04) ─────────── +// Same ownership pattern as sync-health.json / group-login-required.json: +// one concern, one file, one writer (this store). Writes happen on mark() +// (rare — an upstream 401/exhaustion) and are atomic (temp+rename). Reads +// happen ONCE at construction. Every failure path is best-effort: +// write → WARN and keep serving; read → empty store (fallback, never blocks). + +const poolCooldownFilename = "pool-cooldown.json" + +type poolCooldownFileBody struct { + // Accounts maps accountID → avoid-until unix seconds (only unexpired ones). + Accounts map[string]int64 `json:"accounts"` + WrittenAt int64 `json:"written_at"` // unix millis +} + +func poolCooldownPath() (string, error) { + if dir := os.Getenv("AIKEY_RUN_DIR"); dir != "" { + return filepath.Join(dir, poolCooldownFilename), nil + } + home, err := os.UserHomeDir() + if err != nil { + return "", err + } + return filepath.Join(home, ".aikey", "run", poolCooldownFilename), nil +} + +// hydrateFromFile preloads unexpired cooldowns from the state file. Fallback +// by design: any error (missing file, bad JSON, unreadable) leaves the store +// empty — the data path never depends on this file. +func (s *poolCooldownStore) hydrateFromFile() { + path, err := poolCooldownPath() + if err != nil { + return + } + raw, err := os.ReadFile(path) + if err != nil { + return // missing or unreadable → empty store (fallback) + } + var body poolCooldownFileBody + if json.Unmarshal(raw, &body) != nil { + return // corrupt → empty store (fallback); next mark() overwrites it + } + now := s.now() + loaded := 0 + for id, untilUnix := range body.Accounts { + until := time.Unix(untilUnix, 0) + if id != "" && now.Before(until) { + s.m[id] = until + loaded++ + } + } + if loaded > 0 { + slog.Info("pool cooldowns hydrated from state file (survive restart)", + "event.name", observability.EventProxyGroupAccountCooldown, "accounts", loaded) + } +} + +// persistLocked mirrors the current unexpired cooldowns to the state file +// (s.mu held by the caller). Empty set removes the file. Best-effort: a +// failure is WARN-logged and never surfaces to the request path. +func (s *poolCooldownStore) persistLocked() { + path, err := poolCooldownPath() + if err != nil { + return + } + now := s.now() + accounts := make(map[string]int64, len(s.m)) + for id, until := range s.m { + if now.Before(until) { + accounts[id] = until.Unix() + } + } + if len(accounts) == 0 { + if rmErr := os.Remove(path); rmErr != nil && !os.IsNotExist(rmErr) { + slog.Warn("pool cooldown state file remove failed", + "event.name", observability.EventProxyGroupAccountCooldown, "error", rmErr.Error()) + } + return + } + err = func() error { + if mkErr := os.MkdirAll(filepath.Dir(path), 0o700); mkErr != nil { + return mkErr + } + data, mErr := json.Marshal(poolCooldownFileBody{Accounts: accounts, WrittenAt: time.Now().UnixMilli()}) + if mErr != nil { + return mErr + } + tmp := path + ".tmp" + if wErr := os.WriteFile(tmp, data, 0o600); wErr != nil { + return wErr + } + return os.Rename(tmp, path) + }() + if err != nil { + slog.Warn("pool cooldown state file write failed — cooldowns won't survive a restart", + "event.name", observability.EventProxyGroupAccountCooldown, "error", err.Error()) + } +} + +// mark cools an account down until `until`. A no-op for an empty id or a time +// already in the past. +func (s *poolCooldownStore) mark(accountID string, until time.Time) { + if accountID == "" { + return + } + s.mu.Lock() + if cur, ok := s.m[accountID]; !ok || until.After(cur) { + s.m[accountID] = until + s.persistLocked() + } + s.mu.Unlock() +} + +// skipSet returns the accounts currently cooling down, for the resolver's `skip` +// argument. Lapsed entries are dropped. Returns nil when nothing is cooling +// (so the resolver's `skip[id]` lookups stay cheap). +func (s *poolCooldownStore) skipSet() map[string]bool { + s.mu.Lock() + defer s.mu.Unlock() + now := s.now() + var out map[string]bool + for id, until := range s.m { + if now.Before(until) { + if out == nil { + out = make(map[string]bool) + } + out[id] = true + } else { + delete(s.m, id) + } + } + return out +} + +// snapshot returns the accounts currently cooling down → seconds remaining, for +// the admin health surface (N9 组路由健康). Lapsed entries are pruned (same as +// skipSet). Returns nil when nothing is cooling. +func (s *poolCooldownStore) snapshot() map[string]int { + s.mu.Lock() + defer s.mu.Unlock() + now := s.now() + var out map[string]int + for id, until := range s.m { + if now.Before(until) { + if out == nil { + out = make(map[string]int) + } + out[id] = int(until.Sub(now).Seconds()) + } else { + delete(s.m, id) + } + } + return out +} + +// cooldownDecision classifies an upstream response: when the chosen pool account +// should be cooled down, returns (avoid-until, true). 401 → broken account. +// Exhaustion-429 (rate-limit signal present) → honor Retry-After if given, else +// default. WAF-business-429 (no signal) and everything else → (_, false). +func cooldownDecision(resp *http.Response, now time.Time) (time.Time, bool) { + switch resp.StatusCode { + case http.StatusUnauthorized: + return now.Add(poolCooldownDefault), true + case http.StatusTooManyRequests: + if !hasRateLimitSignal(resp.Header) { + return time.Time{}, false // WAF business rejection — not the account's fault + } + d := poolCooldownDefault + // Codex (ChatGPT) carries its own reset headers (x-codex-*-reset-after- + // seconds) instead of Retry-After — prefer them; else Retry-After; else + // default. See research/oauth-codex-ratelimit + R37 (2026-07-04). + if cx := codexRateLimitReset(resp.Header); cx > 0 { + d = cx + } else if ra := retryAfterDuration(resp.Header); ra > 0 { + d = ra + } + if d > poolCooldownMax { + d = poolCooldownMax + } + return now.Add(d), true + default: + return time.Time{}, false + } +} + +// hasRateLimitSignal reports whether the response carries a real rate-limit / +// exhaustion marker (vs a WAF business rejection that omits them). Covers the +// anthropic form (Retry-After / *ratelimit* headers) AND the codex form +// (x-codex-* usage headers) — a codex 429 that omitted both would otherwise be +// mis-read as a WAF rejection and the exhausted account never cooled (打死号). +func hasRateLimitSignal(h http.Header) bool { + if h.Get("Retry-After") != "" { + return true + } + for k := range h { + lk := strings.ToLower(k) + if strings.Contains(lk, "ratelimit") || strings.HasPrefix(lk, "x-codex-") { + return true + } + } + return false +} + +// codexRateLimitReset extracts the cooldown duration from codex's own rate-limit +// headers (ChatGPT backend; wire format verified live 2026-07-06, see +// research/oauth-codex-ratelimit/). We cool for the LONGEST reset among the +// EXHAUSTED (used_percent ≥ 100) windows, so we never un-cool the account into a +// window that is still full and immediately re-429. Returns 0 when no codex reset +// header is present (caller falls back to Retry-After / default). Anthropic +// responses carry no x-codex-* headers, so this is a no-op on the claude path. +// +// Why compare reset DURATIONS, not the primary/secondary NAME: the primary/ +// secondary label is NOT tied to a fixed 5h/7d window — a Plus account's primary +// IS the 5h window (verified live 2026-07-06; sub2api's "primary=weekly" comment +// is backwards). The previous code returned the primary reset first assuming +// primary=7d ("the longer wall"), so when BOTH windows were exhausted it +// under-cooled to whichever window happened to be primary and re-429'd. Comparing +// resets directly sidesteps the naming trap entirely. +// Ref: bugfix 2026-07-06-codex-ratelimit-reset-window-by-name.md. +func codexRateLimitReset(h http.Header) time.Duration { + primaryReset := codexHeaderInt(h, "x-codex-primary-reset-after-seconds") + secondaryReset := codexHeaderInt(h, "x-codex-secondary-reset-after-seconds") + primaryUsed := codexHeaderFloat(h, "x-codex-primary-used-percent") + secondaryUsed := codexHeaderFloat(h, "x-codex-secondary-used-percent") + + // Longest reset among EXHAUSTED windows wins (the bigger wall). + best := 0 + if primaryUsed >= 100 && primaryReset > best { + best = primaryReset + } + if secondaryUsed >= 100 && secondaryReset > best { + best = secondaryReset + } + if best == 0 { + // 429 with neither window flagged exhausted → cool for the longer reset we + // can see (both windows' resets ride on the response regardless). + best = primaryReset + if secondaryReset > best { + best = secondaryReset + } + } + if best > 0 { + return time.Duration(best) * time.Second + } + return 0 +} + +func codexHeaderInt(h http.Header, key string) int { + if v := strings.TrimSpace(h.Get(key)); v != "" { + if i, err := strconv.Atoi(v); err == nil { + return i + } + } + return 0 +} + +func codexHeaderFloat(h http.Header, key string) float64 { + if v := strings.TrimSpace(h.Get(key)); v != "" { + if f, err := strconv.ParseFloat(v, 64); err == nil { + return f + } + } + return 0 +} + +// retryAfterDuration parses the Retry-After header's delta-seconds form. The +// HTTP-date form is ignored (returns 0 → caller uses the default). +func retryAfterDuration(h http.Header) time.Duration { + v := h.Get("Retry-After") + if v == "" { + return 0 + } + if secs, err := strconv.Atoi(strings.TrimSpace(v)); err == nil && secs > 0 { + return time.Duration(secs) * time.Second + } + return 0 +} diff --git a/internal/proxy/oauth_pool_cooldown_persist_test.go b/internal/proxy/oauth_pool_cooldown_persist_test.go new file mode 100644 index 0000000..31a5c2e --- /dev/null +++ b/internal/proxy/oauth_pool_cooldown_persist_test.go @@ -0,0 +1,93 @@ +package proxy + +import ( + "os" + "testing" + "time" +) + +// ── cross-restart cooldown persistence (§S4, 2026-07-04 self-heal) ────────── +// Enhancement, NEVER a dependency (owner constraint): a missing / corrupt / +// unreadable state file must fall back to an empty store; the data path never +// blocks on it. 能红: drop hydrateFromFile from newPoolCooldownStore and the +// restart test fails; drop the corrupt-fallback and the corrupt test panics. + +func TestPoolCooldown_PersistAndHydrateAcrossRestart(t *testing.T) { + t.Setenv("AIKEY_RUN_DIR", t.TempDir()) + + s1 := newPoolCooldownStore() + until := time.Now().Add(30 * time.Minute) + s1.mark("acc-cooled", until) + path, _ := poolCooldownPath() + if _, err := os.Stat(path); err != nil { + t.Fatalf("mark must persist the state file: %v", err) + } + + // "Restart": a fresh store hydrates the same skip view. + s2 := newPoolCooldownStore() + skip := s2.skipSet() + if !skip["acc-cooled"] { + t.Fatalf("hydrated store must keep cooling acc-cooled, skip=%v", skip) + } +} + +func TestPoolCooldown_ExpiredEntriesNotHydrated(t *testing.T) { + t.Setenv("AIKEY_RUN_DIR", t.TempDir()) + + s1 := newPoolCooldownStore() + s1.mark("acc-expired", time.Now().Add(50*time.Millisecond)) + time.Sleep(80 * time.Millisecond) + + s2 := newPoolCooldownStore() + if skip := s2.skipSet(); skip != nil { + t.Fatalf("expired persisted entries must not hydrate, skip=%v", skip) + } +} + +func TestPoolCooldown_CorruptOrMissingFileFallsBackEmpty(t *testing.T) { + t.Setenv("AIKEY_RUN_DIR", t.TempDir()) + + // Missing file → empty store, no error. + s := newPoolCooldownStore() + if skip := s.skipSet(); skip != nil { + t.Fatalf("missing file must yield empty store, skip=%v", skip) + } + + // Corrupt file → empty store (fallback), and the store still WORKS. + path, _ := poolCooldownPath() + if err := os.WriteFile(path, []byte("{not-json"), 0o600); err != nil { + t.Fatal(err) + } + s2 := newPoolCooldownStore() + if skip := s2.skipSet(); skip != nil { + t.Fatalf("corrupt file must yield empty store, skip=%v", skip) + } + s2.mark("acc-new", time.Now().Add(time.Minute)) + if !s2.skipSet()["acc-new"] { + t.Fatal("store must keep working after a corrupt-file fallback") + } + // The next mark overwrote the corrupt file with valid content. + s3 := newPoolCooldownStore() + if !s3.skipSet()["acc-new"] { + t.Fatal("corrupt file must be replaced by the next persist") + } +} + +func TestPoolCooldown_AllExpiredRemovesFile(t *testing.T) { + t.Setenv("AIKEY_RUN_DIR", t.TempDir()) + + s := newPoolCooldownStore() + s.mark("acc-1", time.Now().Add(40*time.Millisecond)) + time.Sleep(60 * time.Millisecond) + // skipSet prunes the lapsed entry; the NEXT mark persists the (now empty → + // file removed on a later persist) view. Trigger a persist via a mark that + // immediately lapses. + _ = s.skipSet() + s.mu.Lock() + s.persistLocked() + s.mu.Unlock() + path, _ := poolCooldownPath() + if _, err := os.Stat(path); !os.IsNotExist(err) { + t.Fatalf("all-expired persist must remove the state file, stat err=%v", err) + } +} diff --git a/internal/proxy/oauth_pool_cooldown_test.go b/internal/proxy/oauth_pool_cooldown_test.go new file mode 100644 index 0000000..f56da97 --- /dev/null +++ b/internal/proxy/oauth_pool_cooldown_test.go @@ -0,0 +1,169 @@ +package proxy + +import ( + "net/http" + "testing" + "time" +) + +func resp(status int, header http.Header) *http.Response { + if header == nil { + header = http.Header{} + } + return &http.Response{StatusCode: status, Header: header} +} + +func TestCooldownDecision_Classification(t *testing.T) { + now := time.Unix(1_750_000_000, 0) + + // 401 → cool down (broken account), default window. + if until, ok := cooldownDecision(resp(401, nil), now); !ok || until != now.Add(poolCooldownDefault) { + t.Fatalf("401 must cool down for the default window, got until=%v ok=%v", until, ok) + } + + // 429 WITH a rate-limit signal → cool down. + rl := http.Header{"Anthropic-Ratelimit-Reset": {"123"}} + if _, ok := cooldownDecision(resp(429, rl), now); !ok { + t.Fatal("exhaustion 429 (rate-limit header) must cool down") + } + + // 429 with Retry-After → honor it (capped). + ra := http.Header{"Retry-After": {"30"}} + if until, ok := cooldownDecision(resp(429, ra), now); !ok || until != now.Add(30*time.Second) { + t.Fatalf("Retry-After must set the cooldown, got %v ok=%v", until, ok) + } + big := http.Header{"Retry-After": {"99999"}} + if until, _ := cooldownDecision(resp(429, big), now); until != now.Add(poolCooldownMax) { + t.Fatalf("oversized Retry-After must be capped at max, got %v", until) + } + + // 429 WITHOUT any rate-limit signal = WAF business rejection → NOT the + // account's fault, do not cool it down. + if _, ok := cooldownDecision(resp(429, nil), now); ok { + t.Fatal("WAF 429 (no rate-limit signal) must NOT cool down the account") + } + + // Success / other → no cooldown. + if _, ok := cooldownDecision(resp(200, nil), now); ok { + t.Fatal("200 must not cool down") + } + if _, ok := cooldownDecision(resp(500, nil), now); ok { + t.Fatal("500 must not cool down (transient upstream, not account-specific)") + } +} + +// TestCooldownDecision_CodexRateLimit covers the codex (ChatGPT) 429 wire format +// (sub2api-derived, R37 2026-07-04): codex carries x-codex-* usage headers instead +// of Retry-After/*ratelimit*. A codex 429 with ONLY these headers must still be +// recognized as a rate-limit (not mis-read as a WAF rejection → 打死号), and its +// reset-after-seconds drives the cooldown. +func TestCooldownDecision_CodexRateLimit(t *testing.T) { + now := time.Unix(1_750_000_000, 0) + + // Only ONE window exhausted (used≥100) → cool for THAT window's reset (300s), + // even though there is NO Retry-After and NO *ratelimit* header (only x-codex-*). + // The other window is below 100% so it does not gate. + sec := http.Header{ + "X-Codex-Secondary-Used-Percent": {"100"}, + "X-Codex-Secondary-Reset-After-Seconds": {"300"}, + "X-Codex-Primary-Used-Percent": {"40"}, + "X-Codex-Primary-Reset-After-Seconds": {"600000"}, + } + if until, ok := cooldownDecision(resp(429, sec), now); !ok || until != now.Add(300*time.Second) { + t.Fatalf("codex one-window-exhausted 429 must cool for that window's reset (300s), got until=%v ok=%v", until, ok) + } + + // Both windows exhausted, LONGER reset is on PRIMARY → cool for the longer wall. + wk := http.Header{ + "X-Codex-Primary-Used-Percent": {"100"}, + "X-Codex-Primary-Reset-After-Seconds": {"3600"}, + "X-Codex-Secondary-Used-Percent": {"100"}, + "X-Codex-Secondary-Reset-After-Seconds": {"120"}, + } + if until, ok := cooldownDecision(resp(429, wk), now); !ok || until != now.Add(3600*time.Second) { + t.Fatalf("codex both-exhausted 429 must cool for the longer reset (3600s), got until=%v ok=%v", until, ok) + } + + // D9 REGRESSION (bugfix 2026-07-06-codex-ratelimit-reset-window-by-name): both + // windows exhausted, but the LONGER wall is the SECONDARY window. The old code + // returned the PRIMARY reset first (assuming primary=7d) and under-cooled to + // the shorter wall → re-429. Must cool for the longer reset (1800s) regardless + // of which window is named primary/secondary. + bugBothExhaustedLongerSecondary := http.Header{ + "X-Codex-Primary-Used-Percent": {"100"}, + "X-Codex-Primary-Reset-After-Seconds": {"120"}, // 5h — shorter + "X-Codex-Secondary-Used-Percent": {"100"}, + "X-Codex-Secondary-Reset-After-Seconds": {"1800"}, // 7d — longer, MUST win + } + if until, ok := cooldownDecision(resp(429, bugBothExhaustedLongerSecondary), now); !ok || until != now.Add(1800*time.Second) { + t.Fatalf("D9: both-exhausted must cool for the LONGER reset (1800s) regardless of primary/secondary name, got until=%v ok=%v", until, ok) + } + + // 429 but neither window at 100% → cool for the larger visible reset. + partial := http.Header{ + "X-Codex-Primary-Used-Percent": {"80"}, + "X-Codex-Primary-Reset-After-Seconds": {"200"}, + "X-Codex-Secondary-Used-Percent": {"90"}, + "X-Codex-Secondary-Reset-After-Seconds": {"50"}, + } + if until, ok := cooldownDecision(resp(429, partial), now); !ok || until != now.Add(200*time.Second) { + t.Fatalf("codex sub-100%% 429 must cool for the larger reset (200s), got until=%v ok=%v", until, ok) + } + + // A codex 429 whose reset exceeds the cap is clamped. + huge := http.Header{ + "X-Codex-Primary-Used-Percent": {"100"}, + "X-Codex-Primary-Reset-After-Seconds": {"999999"}, + } + if until, _ := cooldownDecision(resp(429, huge), now); until != now.Add(poolCooldownMax) { + t.Fatalf("oversized codex reset must be capped at max, got %v", until) + } +} + +func TestPoolCooldownStore_Snapshot(t *testing.T) { + now := time.Unix(1_750_000_000, 0) + s := &poolCooldownStore{m: map[string]time.Time{}, now: func() time.Time { return now }} + + if s.snapshot() != nil { + t.Fatal("empty store → nil snapshot (health shows nothing cooled)") + } + s.mark("acc-a", now.Add(90*time.Second)) + if snap := s.snapshot(); snap["acc-a"] != 90 { + t.Fatalf("acc-a should show ~90s remaining, got %d", snap["acc-a"]) + } + // Advance past the cooldown → dropped from the snapshot. + now = now.Add(2 * time.Minute) + if s.snapshot() != nil { + t.Fatal("lapsed cooldown must drop from the health snapshot") + } +} + +func TestPoolCooldownStore_MarkAndExpire(t *testing.T) { + now := time.Unix(1_750_000_000, 0) + s := &poolCooldownStore{m: map[string]time.Time{}, now: func() time.Time { return now }} + + if len(s.skipSet()) != 0 { + t.Fatal("empty store → empty skip set") + } + s.mark("acc-a", now.Add(5*time.Minute)) + s.mark("", now.Add(time.Hour)) // empty id → ignored + if skip := s.skipSet(); !skip["acc-a"] || len(skip) != 1 { + t.Fatalf("acc-a must be cooling down, got %v", skip) + } + + // A longer mark extends; a shorter one does not shrink. + s.mark("acc-a", now.Add(time.Minute)) + // Advance 2 min: still within the original 5-min window. + now = now.Add(2 * time.Minute) + if !s.skipSet()["acc-a"] { + t.Fatal("longer cooldown must not be shortened by a later shorter mark") + } + // Advance past 5 min → lapses + is dropped. + now = now.Add(4 * time.Minute) + if s.skipSet()["acc-a"] { + t.Fatal("lapsed cooldown must be dropped") + } + if len(s.m) != 0 { + t.Fatalf("lapsed entry must be pruned, map still has %d", len(s.m)) + } +} diff --git a/internal/proxy/oauth_pool_reset.go b/internal/proxy/oauth_pool_reset.go new file mode 100644 index 0000000..6f064bf --- /dev/null +++ b/internal/proxy/oauth_pool_reset.go @@ -0,0 +1,56 @@ +// oauth_pool_reset.go — Path Z (per-window cap re-roll, signal-up half). +// +// The proxy observes each pool account's upstream window-reset epoch +// (anthropic-ratelimit-unified-reset) in ModifyResponse and records the latest +// per account here. The supervisor's N7c pull (fetchGroupRuntime) piggybacks the +// snapshot to master on the EXISTING GET /accounts/me/group-runtime call, and +// master re-rolls window_max_util_pct when it sees a newer reset (a new window). +// master learns the reset no other way today, so this is the signal-up half of +// the re-roll loop; the fresh cap comes back down via the same pull. Reuses the +// pull channel — no new endpoint, no usage-pipeline schema change. See +// 通道3系统设计 §14. +package proxy + +import "sync" + +// poolResetStore holds the latest observed upstream window-reset epoch per pool +// account. Concurrency-safe. Monotonic per account (resets only advance across +// windows; a stale/smaller value never overwrites), so master's +// "observed > stored" re-roll guard stays idempotent — re-sending the same reset +// every pull triggers at most one re-roll. +type poolResetStore struct { + mu sync.Mutex + m map[string]int64 // accountID → latest observed unified-reset epoch (seconds) +} + +func newPoolResetStore() *poolResetStore { + return &poolResetStore{m: make(map[string]int64)} +} + +// record stores the account's observed reset epoch, keeping the max. No-op for +// an empty id or a non-positive epoch. +func (s *poolResetStore) record(accountID string, resetEpoch int64) { + if accountID == "" || resetEpoch <= 0 { + return + } + s.mu.Lock() + if resetEpoch > s.m[accountID] { + s.m[accountID] = resetEpoch + } + s.mu.Unlock() +} + +// snapshot returns a copy of the per-account observed resets for the pull to +// piggyback. Returns nil when empty (so the pull omits the header). +func (s *poolResetStore) snapshot() map[string]int64 { + s.mu.Lock() + defer s.mu.Unlock() + if len(s.m) == 0 { + return nil + } + out := make(map[string]int64, len(s.m)) + for k, v := range s.m { + out[k] = v + } + return out +} diff --git a/internal/proxy/oauth_pool_reset_test.go b/internal/proxy/oauth_pool_reset_test.go new file mode 100644 index 0000000..1d4223c --- /dev/null +++ b/internal/proxy/oauth_pool_reset_test.go @@ -0,0 +1,61 @@ +package proxy + +import ( + "net/http" + "testing" +) + +func TestPoolResetStore_RecordKeepsMaxAndSnapshotIsCopy(t *testing.T) { + s := newPoolResetStore() + if s.snapshot() != nil { + t.Fatal("empty store must snapshot nil") + } + s.record("acc-1", 100) + s.record("acc-1", 50) // older → ignored (monotonic: resets only advance) + s.record("acc-1", 200) // newer → kept + s.record("acc-2", 300) + s.record("", 999) // empty id → ignored + s.record("acc-x", 0) // non-positive → ignored + + snap := s.snapshot() + if snap["acc-1"] != 200 { + t.Fatalf("acc-1=%d want 200 (max kept, stale 50 ignored)", snap["acc-1"]) + } + if snap["acc-2"] != 300 { + t.Fatalf("acc-2=%d want 300", snap["acc-2"]) + } + if len(snap) != 2 { + t.Fatalf("snapshot size=%d want 2 (empty id + zero epoch dropped)", len(snap)) + } + + // snapshot is a copy — mutating it must not affect the store. + snap["acc-1"] = 1 + if s.snapshot()["acc-1"] != 200 { + t.Fatal("snapshot must be a defensive copy") + } +} + +func TestObservedResetEpoch(t *testing.T) { + // Unified reset preferred. + h := http.Header{} + h.Set("anthropic-ratelimit-unified-reset", "1750000000") + if e, ok := observedResetEpoch(h); !ok || e != 1750000000 { + t.Fatalf("unified reset: e=%d ok=%v", e, ok) + } + // Per-window fallback when unified absent. + h2 := http.Header{} + h2.Set("anthropic-ratelimit-unified-5h-reset", "1750000500") + if e, ok := observedResetEpoch(h2); !ok || e != 1750000500 { + t.Fatalf("5h reset fallback: e=%d ok=%v", e, ok) + } + // No reset header → false. + if _, ok := observedResetEpoch(http.Header{}); ok { + t.Fatal("absent reset → false") + } + // Garbage → false (not recorded). + h3 := http.Header{} + h3.Set("anthropic-ratelimit-unified-reset", "notanum") + if _, ok := observedResetEpoch(h3); ok { + t.Fatal("garbage epoch → false") + } +} diff --git a/internal/proxy/oauth_pool_window.go b/internal/proxy/oauth_pool_window.go new file mode 100644 index 0000000..e84afb3 --- /dev/null +++ b/internal/proxy/oauth_pool_window.go @@ -0,0 +1,176 @@ +// oauth_pool_window.go — N10: window randomized-cap pre-cut (防封). +// +// An account that runs its quota window to 100% looks like abuse to the upstream. +// To stay under the wall, master generates a RANDOMIZED cap per account+window +// (window_max_util_pct, 95-99, N11) and delivers it in the group material. The +// proxy reads the upstream's live utilization from EVERY response (present on +// 200s too — verified for OAuth subscription, 20260623-OAuth账号池-技术方案 §435) +// and, when utilization ≥ cap/100, pre-cuts the account for that window: cools it +// down until the window resets, so subsequent requests route to another account +// BEFORE the account hits 100%. Reuses N8c's cooldown store. +// +// Headers (verified): anthropic-ratelimit-unified-5h-utilization / +// -7d-utilization (float 0~1), anthropic-ratelimit-unified-reset (epoch seconds, +// or per-window -5h-reset / -7d-reset). http.Header.Get canonicalizes the key, +// so the lowercase form below matches regardless of the wire casing. +package proxy + +import ( + "context" + "net/http" + "strconv" + "time" +) + +// stashWindowCap records the chosen pool account's window_max_util_pct on the +// request so serveRoute's ModifyResponse can pre-cut against the upstream's +// utilization header (N10). Returns the request carrying the stash. +func stashWindowCap(r *http.Request, capPct int) *http.Request { + return r.WithContext(context.WithValue(r.Context(), ctxKeyPoolWindowCap, capPct)) +} + +// windowCapFromContext reads the stashed cap (false when absent → no pre-cut). +func windowCapFromContext(req *http.Request) (int, bool) { + v, ok := req.Context().Value(ctxKeyPoolWindowCap).(int) + return v, ok +} + +const ( + hdrUtil5h = "anthropic-ratelimit-unified-5h-utilization" + hdrUtil7d = "anthropic-ratelimit-unified-7d-utilization" + hdrReset = "anthropic-ratelimit-unified-reset" + hdrReset5h = "anthropic-ratelimit-unified-5h-reset" + hdrReset7d = "anthropic-ratelimit-unified-7d-reset" +) + +// windowPreCutDecision reports whether the chosen pool account has reached its +// randomized window cap and should be pre-cut. capPct is the account's delivered +// window_max_util_pct; <=0 or >=100 means "no meaningful cap" → never pre-cut +// (let natural exhaustion / 429 handle it). Returns (cool-until, true) when +// either window's utilization ≥ cap/100. +func windowPreCutDecision(h http.Header, capPct int, now time.Time) (time.Time, bool) { + if capPct <= 0 || capPct >= 100 { + return time.Time{}, false + } + threshold := float64(capPct) / 100.0 + // Anthropic path (byte-identical to before): the two unified-utilization + // headers ride on every response including 200s. + u5h, has5h := parseUtil(h, hdrUtil5h) + u7d, has7d := parseUtil(h, hdrUtil7d) + if has5h || has7d { + if u5h < threshold && u7d < threshold { + return time.Time{}, false // both windows below cap → fine + } + return windowResetTime(h, now), true + } + // Codex path: same anti-ban cap, Codex's own X-Codex-* util/reset headers + // (verified live on the /responses 200, 2026-07-06). The two header families + // are mutually exclusive (an account is Anthropic OR Codex), so header-sniff + // dispatch mirrors R34's reactive layer — no provider_code branch. + return codexWindowPreCut(h, threshold, now) +} + +// codexWindowPreCut pre-cuts when ANY Codex quota window's utilization has +// reached the cap, cooling until that window's reset (the latest among over-cap +// windows). No 5h/7d classification is needed: we cool until the over-cap +// window's OWN reset, comparing per-window (the same insight as codexRateLimitReset +// and the D9 fix — never trust the primary/secondary name). Returns (zero,false) +// when the response carries no X-Codex-* util headers (not Codex traffic) or no +// window is at/over the cap. +func codexWindowPreCut(h http.Header, threshold float64, now time.Time) (time.Time, bool) { + var coolUntil time.Time + hit := false + consider := func(usedHdr, resetAtHdr, resetAfterHdr string) { + pct, ok := parseCodexPercent(h.Get(usedHdr)) + if !ok || pct/100.0 < threshold { + return + } + hit = true + if t := codexWindowReset(h.Get(resetAtHdr), h.Get(resetAfterHdr), now); t.After(coolUntil) { + coolUntil = t + } + } + consider("X-Codex-Primary-Used-Percent", "X-Codex-Primary-Reset-At", "X-Codex-Primary-Reset-After-Seconds") + consider("X-Codex-Secondary-Used-Percent", "X-Codex-Secondary-Reset-At", "X-Codex-Secondary-Reset-After-Seconds") + if !hit { + return time.Time{}, false + } + if coolUntil.IsZero() || !coolUntil.After(now) { + coolUntil = now.Add(poolCooldownDefault) + } + return coolUntil, true +} + +// codexWindowReset picks the cool-until for one Codex window: prefer the ABSOLUTE +// X-Codex-*-Reset-At epoch (verified present live 2026-07-06 — Codex is NOT +// relative-only as sub2api's data implied), else now + the relative +// X-Codex-*-Reset-After-Seconds, else a conservative default. Never a past time. +func codexWindowReset(resetAtRaw, resetAfterRaw string, now time.Time) time.Time { + if epoch := parseCodexInt(resetAtRaw); epoch > 0 { + if t := time.Unix(int64(epoch), 0); t.After(now) { + return t + } + } + if secs := parseCodexInt(resetAfterRaw); secs > 0 { + return now.Add(time.Duration(secs) * time.Second) + } + return now.Add(poolCooldownDefault) +} + +// parseUtil reads a unified utilization header as a fraction 0~1. Returns +// (value, true) only when the header is present and parses to a sane number. +func parseUtil(h http.Header, key string) (float64, bool) { + v := h.Get(key) + if v == "" { + return 0, false + } + f, err := strconv.ParseFloat(v, 64) + if err != nil || f < 0 { + return 0, false + } + return f, true +} + +// observedResetEpoch reads the upstream's raw window-reset epoch (unified, else +// per-window) for Path Z's re-roll signal. Unlike windowResetTime it does NOT +// clamp to the future: master compares the raw epoch to its stored value to +// detect a new window, so a just-passed reset must still report its true epoch. +func observedResetEpoch(h http.Header) (int64, bool) { + for _, key := range []string{hdrReset, hdrReset5h, hdrReset7d} { + if v := h.Get(key); v != "" { + if epoch, err := strconv.ParseInt(v, 10, 64); err == nil && epoch > 0 { + return epoch, true + } + } + } + // Codex: the absolute X-Codex-*-Reset-At (verified live 2026-07-06). Report the + // SOONEST future boundary (min of the two reset-ats) so master re-rolls the cap + // when the shorter (5h) window rolls over — no 5h/7d classification needed. + var soonest int64 + for _, key := range []string{"X-Codex-Primary-Reset-At", "X-Codex-Secondary-Reset-At"} { + if e := parseCodexInt(h.Get(key)); e > 0 { + if soonest == 0 || int64(e) < soonest { + soonest = int64(e) + } + } + } + if soonest > 0 { + return soonest, true + } + return 0, false +} + +// windowResetTime picks the cool-until time: prefer the unified reset epoch, then +// a per-window reset, else a conservative default. Never returns a past time. +func windowResetTime(h http.Header, now time.Time) time.Time { + for _, key := range []string{hdrReset, hdrReset5h, hdrReset7d} { + if v := h.Get(key); v != "" { + if epoch, err := strconv.ParseInt(v, 10, 64); err == nil && epoch > 0 { + if t := time.Unix(epoch, 0); t.After(now) { + return t + } + } + } + } + return now.Add(poolCooldownDefault) +} diff --git a/internal/proxy/oauth_pool_window_codex_test.go b/internal/proxy/oauth_pool_window_codex_test.go new file mode 100644 index 0000000..cc2e932 --- /dev/null +++ b/internal/proxy/oauth_pool_window_codex_test.go @@ -0,0 +1,113 @@ +package proxy + +import ( + "net/http" + "testing" + "time" +) + +// TestWindowPreCutDecision_Codex pins the Codex branch of the N10 pre-cut against +// the real wire captured 2026-07-06 (research/oauth-codex-ratelimit/): X-Codex-* +// used-percent + absolute reset-at ride on the /responses 200. Pre-cut fires when +// any window ≥ the randomized cap, cooling until that window's OWN reset (latest +// among over-cap windows) — no 5h/7d classification, mirroring the D9 fix. +func TestWindowPreCutDecision_Codex(t *testing.T) { + now := time.Unix(1_783_300_000, 0) + const resetAt5h = int64(1_783_341_932) // future, ~5h out (from the live capture) + const resetAt7d = int64(1_783_928_732) // further future, ~7d out + + hdr := func(kv map[string]string) http.Header { + h := http.Header{} + for k, v := range kv { + h.Set(k, v) + } + return h + } + + t.Run("below cap → no pre-cut", func(t *testing.T) { + h := hdr(map[string]string{ + "X-Codex-Primary-Used-Percent": "1", + "X-Codex-Primary-Window-Minutes": "300", + "X-Codex-Primary-Reset-At": "1783341932", + "X-Codex-Secondary-Used-Percent": "0", + }) + if _, ok := windowPreCutDecision(h, 95, now); ok { + t.Fatal("util below cap must not pre-cut") + } + }) + + t.Run("at/over cap → pre-cut until that window reset-at", func(t *testing.T) { + h := hdr(map[string]string{ + "X-Codex-Primary-Used-Percent": "96", + "X-Codex-Primary-Reset-At": "1783341932", + "X-Codex-Secondary-Used-Percent": "0", + }) + until, ok := windowPreCutDecision(h, 95, now) + if !ok || !until.Equal(time.Unix(resetAt5h, 0)) { + t.Fatalf("over-cap must pre-cut until reset-at %d, got until=%v ok=%v", resetAt5h, until, ok) + } + }) + + t.Run("both windows over cap → cool until the LATER reset", func(t *testing.T) { + h := hdr(map[string]string{ + "X-Codex-Primary-Used-Percent": "100", + "X-Codex-Primary-Reset-At": "1783341932", // 5h (sooner) + "X-Codex-Secondary-Used-Percent": "100", + "X-Codex-Secondary-Reset-At": "1783928732", // 7d (later) — must win + }) + until, ok := windowPreCutDecision(h, 95, now) + if !ok || !until.Equal(time.Unix(resetAt7d, 0)) { + t.Fatalf("both-over must cool until the later reset %d, got until=%v ok=%v", resetAt7d, until, ok) + } + }) + + t.Run("no absolute reset-at → falls back to now+reset-after-seconds", func(t *testing.T) { + h := hdr(map[string]string{ + "X-Codex-Primary-Used-Percent": "98", + "X-Codex-Primary-Reset-After-Seconds": "3600", + }) + until, ok := windowPreCutDecision(h, 95, now) + if !ok || !until.Equal(now.Add(3600*time.Second)) { + t.Fatalf("relative fallback must cool now+3600s, got until=%v ok=%v", until, ok) + } + }) + + t.Run("cap disabled (>=100) → never pre-cut", func(t *testing.T) { + h := hdr(map[string]string{"X-Codex-Primary-Used-Percent": "100", "X-Codex-Primary-Reset-At": "1783341932"}) + if _, ok := windowPreCutDecision(h, 100, now); ok { + t.Fatal("capPct>=100 means no meaningful cap") + } + }) + + t.Run("anthropic path unchanged (no codex headers)", func(t *testing.T) { + h := hdr(map[string]string{"anthropic-ratelimit-unified-5h-utilization": "0.96"}) + if _, ok := windowPreCutDecision(h, 95, now); !ok { + t.Fatal("anthropic util over cap must still pre-cut (byte-identical path)") + } + }) +} + +// TestObservedResetEpoch_Codex pins the Path-Z re-roll signal for Codex: report +// the SOONEST future reset-at so master re-rolls the cap when the shorter window +// rolls over. Anthropic responses still return the anthropic reset unchanged. +func TestObservedResetEpoch_Codex(t *testing.T) { + // Build headers with Set() so keys are canonicalized exactly as Go's HTTP + // parser does for a real upstream response (a raw map literal would store a + // non-canonical lowercase key that http.Header.Get can't find). + codex := http.Header{} + codex.Set("X-Codex-Primary-Reset-At", "1783341932") // soonest + codex.Set("X-Codex-Secondary-Reset-At", "1783928732") + if epoch, ok := observedResetEpoch(codex); !ok || epoch != 1783341932 { + t.Fatalf("codex observed reset must be the soonest reset-at, got %d ok=%v", epoch, ok) + } + + anthropic := http.Header{} + anthropic.Set("anthropic-ratelimit-unified-reset", "1783341000") + if epoch, ok := observedResetEpoch(anthropic); !ok || epoch != 1783341000 { + t.Fatalf("anthropic path must be unchanged, got %d ok=%v", epoch, ok) + } + + if _, ok := observedResetEpoch(http.Header{}); ok { + t.Fatal("no reset headers → (0,false)") + } +} diff --git a/internal/proxy/oauth_pool_window_precut_switch_test.go b/internal/proxy/oauth_pool_window_precut_switch_test.go new file mode 100644 index 0000000..7285748 --- /dev/null +++ b/internal/proxy/oauth_pool_window_precut_switch_test.go @@ -0,0 +1,178 @@ +package proxy + +// Integration (mock-upstream) — real Anthropic unified rate-limit headers drive the +// window pre-cut → account switch → success chain, AND the I5 utilization uplink to +// master (which feeds the 动态决策账号分配引擎). 2026-07-01. +// +// WHY this test (user ask): the "额度不足 → 切换账号 → 切换成功" path must be exercised +// with a mock Provider returning REAL response headers (not a fabricated WindowStatus), +// so the REAL parse path (oauth_pool_window.go parseUtil + signal_report.go +// parseUnifiedUtil5h) is under test. Both read the SAME header, so one injected header +// covers both legs: +// - proxy-local reactive pre-cut: util ≥ delivered cap → cool the account → the +// resolver routes the NEXT request to another account (switch), which succeeds. +// - engine uplink (I5): the util value is enqueued to the signal reporter, which the +// collector ships to master; the allocation engine reads the util trend from it. +// +// Header shape is the VERIFIED real Anthropic Claude-subscription "unified" rate-limit +// block (platform docs + reverse-engineering, 2026-07): +// anthropic-ratelimit-unified-5h-utilization: 0.98 (float 0..1) +// anthropic-ratelimit-unified-7d-utilization: 0.40 +// anthropic-ratelimit-unified-5h-reset: +// anthropic-ratelimit-unified-status: allowed | exceeded | rate_limited +// anthropic-ratelimit-unified-representative-claim: five_hour +// "额度剩余" = 1 − utilization; the proxy pre-cuts on utilization ≥ cap/100. + +import ( + "net/http" + "strconv" + "strings" + "testing" + "time" + + "github.com/AiKeyLabs/aikey-proxy/internal/vkeys" +) + +// unifiedRateLimitHeaders builds the real Anthropic unified rate-limit response block. +// Built via Set so the keys are canonicalized exactly like a real wire response +// (http.Header.Get canonicalizes on read, so casing must match). +func unifiedRateLimitHeaders(util5h, util7d string, resetEpoch int64) http.Header { + h := http.Header{} + h.Set("anthropic-ratelimit-unified-5h-utilization", util5h) + h.Set("anthropic-ratelimit-unified-7d-utilization", util7d) + h.Set("anthropic-ratelimit-unified-5h-reset", strconv.FormatInt(resetEpoch, 10)) + h.Set("anthropic-ratelimit-unified-status", "allowed") + h.Set("anthropic-ratelimit-unified-representative-claim", "five_hour") + return h +} + +// drainForUtil non-blockingly scans the signal channel for a sample matching +// (credID, util5h). Returns true on the first match. +func drainForUtil(in <-chan signalSample, credID string, want float64) bool { + for { + select { + case s := <-in: + if s.CredentialID == credID && s.Util5h == want { + return true + } + default: + return false + } + } +} + +// twoAccountWindowRoute builds a 2-OAuth-account group VK whose accounts each carry a +// delivered window cap (so stashWindowCap fires) + a real credential_id (so the I5 +// enqueue, which keys on credential_id, isn't dropped). +func twoAccountWindowRoute(t *testing.T, key []byte, capPct int) (*vkeys.ResolvedRoute, map[string]string, map[string]string) { + t.Helper() + cap := capPct + refs := []vkeys.GroupAccountRef{ + {AccountID: "acc-1", CredentialID: "cred-1", ProviderCode: "anthropic"}, + {AccountID: "acc-2", CredentialID: "cred-2", ProviderCode: "anthropic"}, + } + mk := func(ext string) vkeys.GroupRuntimeAccount { + return vkeys.GroupRuntimeAccount{ + CredentialType: "oauth_account", ExpiresAt: 9_000_000_000, ExternalID: ext, + WindowMaxUtilPct: &cap, WindowStatus: "active", + } + } + mat := map[string]vkeys.GroupRuntimeAccount{ + "acc-1": encMat(t, key, mk("uuid-1"), "tok-1"), + "acc-2": encMat(t, key, mk("uuid-2"), "tok-2"), + } + route := &vkeys.ResolvedRoute{ + VirtualKeyID: "vk-grp", Provider: "anthropic", ProtocolType: "anthropic", + ProviderCode: "anthropic", RouteSource: "team", + SeatID: "seat-1", OauthGroupID: "grp-1", + GroupAccounts: mustJSON(t, refs), GroupRuntime: mustJSON(t, mat), + } + tokToAcct := map[string]string{"tok-1": "acc-1", "tok-2": "acc-2"} + credOf := map[string]string{"acc-1": "cred-1", "acc-2": "cred-2"} + return route, tokToAcct, credOf +} + +// util ≥ cap (from a real header on a 200) → pre-cut the served account → next request +// SWITCHES to the other account and SUCCEEDS; the util is also uplinked (I5) to master. +func TestGroupServe_WindowUtilHeaderPreCutsSwitchesAndUplinks(t *testing.T) { + key := grKey() + route, tokToAcct, credOf := twoAccountWindowRoute(t, key, 97) + + p := setupTestProxy(t, "http://unused.invalid") + p.registry.Merge(map[string]*vkeys.ResolvedRoute{"aikey_team_grouptest": route}) + p.SetGroupKeyProvider(fakeGroupKey{k: key}) + // Inspectable reporter so we can assert the util was enqueued for the master + // uplink (no loop() goroutine — the buffer just holds the samples for the test). + rep := &signalReporter{in: make(chan signalSample, 8)} + p.signalReporter = rep + tr := &outboundCapture{} + p.SetTransport(tr) + + // ── Request 1: mock upstream 200 + util 0.98 (≥ cap 0.97) on the served account ── + reset := time.Now().Add(2 * time.Hour).Unix() + tr.respHeader = unifiedRateLimitHeaders("0.98", "0.40", reset) + req1, w1 := groupReq(groupBody) + p.Handle(w1, req1) + + if w1.Code != http.StatusOK { + t.Fatalf("the request itself must still succeed (200); pre-cut is for the NEXT one, got %d: %s", w1.Code, w1.Body.String()) + } + served1 := tokToAcct[strings.TrimPrefix(tr.auth, "Bearer ")] + if served1 == "" { + t.Fatalf("req1 outbound auth %q did not map to a known account", tr.auth) + } + // (a) LOCAL pre-cut: util ≥ cap must cool the served account for its window. + if !p.poolCooldown.skipSet()[served1] { + t.Fatalf("util 0.98 ≥ cap 0.97 must PRE-CUT (cool) the served account %s (防封, before it hits 100%%)", served1) + } + // (b) ENGINE uplink: the same util must be enqueued to the signal reporter (I5 → + // collector → master; the 动态决策引擎 reads the util trend from it). + if !drainForUtil(rep.in, credOf[served1], 0.98) { + t.Fatalf("util 0.98 must be enqueued to the signal reporter for cred %s (I5 uplink to the allocation engine)", credOf[served1]) + } + + // ── Request 2: util now low; pre-cut account is skipped → SWITCH → SUCCESS ── + tr.respHeader = unifiedRateLimitHeaders("0.10", "0.10", reset) + req2, w2 := groupReq(groupBody) + p.Handle(w2, req2) + + if w2.Code != http.StatusOK { + t.Fatalf("额度将满→切换：the switch request must SUCCEED via the other account, got %d: %s", w2.Code, w2.Body.String()) + } + served2 := tokToAcct[strings.TrimPrefix(tr.auth, "Bearer ")] + if served2 == "" || served2 == served1 { + t.Fatalf("must switch to a DIFFERENT account after the pre-cut; served1=%s served2=%s (auth=%q)", served1, served2, tr.auth) + } +} + +// Negative control: util 0.50 < cap 0.97 → NO pre-cut → the next request STAYS on the +// same account (a transient/low reading must not churn the route). Mirrors E4's matrix. +func TestGroupServe_WindowUtilBelowCapNoPreCut(t *testing.T) { + key := grKey() + route, tokToAcct, _ := twoAccountWindowRoute(t, key, 97) + + p := setupTestProxy(t, "http://unused.invalid") + p.registry.Merge(map[string]*vkeys.ResolvedRoute{"aikey_team_grouptest": route}) + p.SetGroupKeyProvider(fakeGroupKey{k: key}) + tr := &outboundCapture{} + p.SetTransport(tr) + + reset := time.Now().Add(2 * time.Hour).Unix() + tr.respHeader = unifiedRateLimitHeaders("0.50", "0.30", reset) + req1, w1 := groupReq(groupBody) + p.Handle(w1, req1) + if w1.Code != http.StatusOK { + t.Fatalf("req1 status=%d body=%s", w1.Code, w1.Body.String()) + } + served1 := tokToAcct[strings.TrimPrefix(tr.auth, "Bearer ")] + if p.poolCooldown.skipSet()[served1] { + t.Fatalf("util 0.50 < cap 0.97 must NOT pre-cut account %s (no needless churn)", served1) + } + + req2, w2 := groupReq(groupBody) + p.Handle(w2, req2) + served2 := tokToAcct[strings.TrimPrefix(tr.auth, "Bearer ")] + if served2 != served1 { + t.Fatalf("no pre-cut → must STAY on the same account (sticky); served1=%s served2=%s", served1, served2) + } +} diff --git a/internal/proxy/oauth_pool_window_test.go b/internal/proxy/oauth_pool_window_test.go new file mode 100644 index 0000000..ffafcb3 --- /dev/null +++ b/internal/proxy/oauth_pool_window_test.go @@ -0,0 +1,74 @@ +package proxy + +import ( + "net/http" + "strconv" + "testing" + "time" +) + +func hdr(kv map[string]string) http.Header { + h := http.Header{} + for k, v := range kv { + h.Set(k, v) + } + return h +} + +func TestWindowPreCutDecision(t *testing.T) { + now := time.Unix(1_750_000_000, 0) + reset := now.Add(time.Hour) + resetStr := strconv.FormatInt(reset.Unix(), 10) + + // 5h utilization at/over the cap → pre-cut, cool until the window reset. + if until, ok := windowPreCutDecision(hdr(map[string]string{hdrUtil5h: "0.97", hdrReset: resetStr}), 96, now); !ok || !until.Equal(reset) { + t.Fatalf("util 0.97 ≥ cap 96%% must pre-cut until reset, got until=%v ok=%v", until, ok) + } + // Under the cap → no pre-cut. + if _, ok := windowPreCutDecision(hdr(map[string]string{hdrUtil5h: "0.50"}), 96, now); ok { + t.Fatal("util under cap must not pre-cut") + } + // 7d window over the cap (5h fine) → pre-cut (either window triggers). + if _, ok := windowPreCutDecision(hdr(map[string]string{hdrUtil5h: "0.10", hdrUtil7d: "0.99"}), 96, now); !ok { + t.Fatal("7d util over cap must pre-cut") + } + // No utilization header at all → nothing to pre-cut on. + if _, ok := windowPreCutDecision(hdr(map[string]string{hdrReset: resetStr}), 96, now); ok { + t.Fatal("no utilization signal → no pre-cut") + } + // No meaningful cap (>=100 or <=0) → never pre-cut, even at high util. + if _, ok := windowPreCutDecision(hdr(map[string]string{hdrUtil5h: "0.99"}), 100, now); ok { + t.Fatal("cap>=100 → no pre-cut") + } + if _, ok := windowPreCutDecision(hdr(map[string]string{hdrUtil5h: "0.99"}), 0, now); ok { + t.Fatal("cap<=0 → no pre-cut") + } + // Over cap but no reset header → conservative default cooldown. + if until, ok := windowPreCutDecision(hdr(map[string]string{hdrUtil5h: "0.99"}), 96, now); !ok || !until.Equal(now.Add(poolCooldownDefault)) { + t.Fatalf("missing reset → default cooldown, got %v", until) + } + // A past reset epoch must not yield a past cool-until. + past := strconv.FormatInt(now.Add(-time.Hour).Unix(), 10) + if until, _ := windowPreCutDecision(hdr(map[string]string{hdrUtil5h: "0.99", hdrReset: past}), 96, now); !until.After(now) { + t.Fatalf("past reset must fall back to a future cool-until, got %v", until) + } + // Wire-casing variance: Go canonicalizes header keys, so a real upstream's + // casing still matches the lowercase const. + hWire := http.Header{} + hWire.Set("Anthropic-Ratelimit-Unified-5h-Utilization", "0.98") + if _, ok := windowPreCutDecision(hWire, 96, now); !ok { + t.Fatal("header key canonicalization should match real upstream casing") + } +} + +func TestParseUtil(t *testing.T) { + if _, ok := parseUtil(hdr(nil), hdrUtil5h); ok { + t.Fatal("missing header → not present") + } + if v, ok := parseUtil(hdr(map[string]string{hdrUtil5h: "0.83"}), hdrUtil5h); !ok || v != 0.83 { + t.Fatalf("0.83 → (0.83,true), got (%v,%v)", v, ok) + } + if _, ok := parseUtil(hdr(map[string]string{hdrUtil5h: "not-a-number"}), hdrUtil5h); ok { + t.Fatal("unparseable → not present (don't pre-cut on garbage)") + } +} diff --git a/internal/proxy/pipelines.go b/internal/proxy/pipelines.go index cba1a71..7b678cd 100644 --- a/internal/proxy/pipelines.go +++ b/internal/proxy/pipelines.go @@ -975,6 +975,35 @@ func (p *Proxy) handlePathPrefixRoute(w http.ResponseWriter, r *http.Request, pr return } + // Oauth-group VK: serve via the group handler — the path-prefix entry must + // wire this exactly like the legacy /v1 dispatch (handle_dispatch.go:235). + // A group VK carries NO VK-level provider (it's per-account in the group + // runtime), so without this branch it falls through to the provider- + // compatibility check below and 403s PROVIDER_MISMATCH on the empty + // ProviderCode. The connectivity-test probe targets //... (the + // path-prefix entry), so ONLY this entry was affected — real Claude Code + // uses the /v1 entry which already had the branch. Same empty-provider + // root cause as 2026-06-25-group-vk-empty-provider-code-502. group VKs are + // only registered when the oauth-group flag is on, so OauthGroupID is empty + // in flag-off builds and the direct-bind path stays byte-identical. + if route.OauthGroupID != "" { + // Strip the provider prefix BEFORE handing off to the group handler. + // The path-prefix entry normally defers the strip to below (after the + // provider-compat check: `r.URL.Path = strippedPath`), but + // handleOauthGroupRoute forwards r.URL.Path VERBATIM to the upstream — + // so an unstripped `/anthropic/v1/models` would hit + // `api.anthropic.com/anthropic/v1/models` → 404 (verified: Cf-Ray 404 + // from api.anthropic.com). The legacy /v1 entry (handle_dispatch.go) is + // unaffected: its path is already `/v1/...` with no provider prefix. + // Bugfix: 2026-06-26-group-vk-pathprefix-unstripped-404. + r.URL.Path = strippedPath + if r.URL.RawPath != "" { + r.URL.RawPath = strippedPath + } + p.handleOauthGroupRoute(w, r, route, rawAuthValue, startTime, logger, traceID) + return + } + // Provider compatibility check: token's provider must match path's provider. if !isProviderCompatible(route, canonicalCode) { p.errors.Add(1) @@ -1312,6 +1341,39 @@ func (p *Proxy) handlePathPrefixRoute(w http.ResponseWriter, r *http.Request, pr // credential records — see helper docstring for per-KeySourceType semantics. binding, _ := p.activeReader.GetProviderBinding(canonicalCode) if binding != nil { + // ── Group VK on the follow-active path (N8 / bugfix 2026-06-30) ──────── + // A group VK (OAuth account pool) carries NO static PlaintextKey — its + // per-account material lives in GroupRuntime and must be served via + // handleOauthGroupRoute (the same path the Tier1 `aikey_team_` token + // reaches at line ~989). The follow-active path (Claude Code via + // `aikey run` injects the `aikey_active_` sentinel) lands here + // instead, where ResolveBindingCredential's switch has no oauth_group + // case → it soft-fails with RealKey="" → the legacy fallback also can't + // serve a keyless VK → "no active key". + // + // The active team binding stores the vk_id (KeySourceRef). The supervisor + // registers every group VK in the registry under the Tier1 team token + // `aikey_team_` (supervisor.go:725) WITH the group fields fully + // populated — so re-resolving that token here yields the same complete + // group route the Tier1 path serves, without touching GetTeamKeyByID + // (whose query filters on `provider_key_ciphertext IS NOT NULL` and never + // reads group columns, so it can't see a keyless group VK at all). + if binding.KeySourceType == "team" { + if groute := p.registry.Resolve("aikey_team_" + binding.KeySourceRef); groute != nil && + groute.OauthGroupID != "" { + // Strip the provider prefix BEFORE handing off: handleOauthGroupRoute + // forwards r.URL.Path VERBATIM upstream, so an unstripped + // `/anthropic/v1/messages` would 404 (same rationale as the Tier1 + // branch at line ~999; bugfix 2026-06-26-group-vk-pathprefix-unstripped-404). + r.URL.Path = strippedPath + if r.URL.RawPath != "" { + r.URL.RawPath = strippedPath + } + p.handleOauthGroupRoute(w, r, groute, rawAuthValue, startTime, logger, traceID) + return + } + } + cred, bindErr := p.ResolveBindingCredential(r, binding, providerCode, canonicalCode, logger) if bindErr != nil { // OAuth path: writeJSONError + return (helper does NOT increment diff --git a/internal/proxy/probe_raw.go b/internal/proxy/probe_raw.go index 9a65971..d3fe0f4 100644 --- a/internal/proxy/probe_raw.go +++ b/internal/proxy/probe_raw.go @@ -238,7 +238,7 @@ func (p *Proxy) handleProbeRaw(w http.ResponseWriter, r *http.Request, canonical // 5. Send + time. Use p.transport so HTTP_PROXY env etc are honored // the same way as the regular upstream path. - transport := p.transport + transport := p.currentTransport() if transport == nil { transport = http.DefaultTransport } diff --git a/internal/proxy/proxy.go b/internal/proxy/proxy.go index 60e1359..f35b330 100644 --- a/internal/proxy/proxy.go +++ b/internal/proxy/proxy.go @@ -24,6 +24,13 @@ type VaultGetter interface { GetSecret(alias string) (string, error) } +// groupKeyProvider exposes the vault's derived key so the oauth-group resolver +// (N8) can decrypt per-account group material at request time. *vault.Reader +// implements it; an injected mock that doesn't disables group routing (safe). +type groupKeyProvider interface { + DerivedKey() []byte +} + // ActiveKeyReader extends VaultGetter with active-key lookups used by // path-prefix routing (/anthropic/v1/..., /openai/v1/...). // Implemented by *vault.Reader when the vault supports managed and personal keys. @@ -64,12 +71,54 @@ type OAuthCredential struct { // Proxy is the core reverse proxy that handles virtual key resolution // and request forwarding. type Proxy struct { - transport http.RoundTripper // nil → http.DefaultTransport (reads env vars) + // transport is the outbound RoundTripper for AI provider forwarding, held in an + // atomic.Pointer so the egress upstream-proxy URL can be HOT-SWAPPED at runtime + // (Settings → Upstream proxy, 2026-06-30) without racing the per-request read in + // forward_and_resolve. Nil box / nil rt → http.DefaultTransport (honors + // HTTP_PROXY env). Read via currentTransport, written via SetTransport. + transport atomic.Pointer[transportBox] activeReader ActiveKeyReader // non-nil when vault implements ActiveKeyReader appVault apppipe.VaultReader // non-nil when vault implements the App pipeline read surface (Phase 4) probeVault probepipe.VaultReader // non-nil when vault implements the Probe pipeline read surface (mode C, SPEC 2026-05-23) broker OAuthBroker // OAuth credential provider (nil = OAuth not available) vault VaultGetter + // groupKey exposes the vault derived key for oauth-group material decryption + // (N8). nil when the injected vault doesn't implement DerivedKey() (tests) → + // group routing degrades to GROUP_KEY_UNAVAILABLE rather than panicking. + groupKey groupKeyProvider + // poolCooldown holds per-account reactive fallback state (N8c): an account + // whose upstream failed (401 / exhaustion-429) is skipped by the resolver + // until its cooldown lapses. Always non-nil (set in New); the request path + // only consults it for group routes. + poolCooldown *poolCooldownStore + // signalReporter ships parsed unified-* utilization to master (I5, best-effort, + // off the forward hot path). nil = feature off. Set via EnableSignalReporting. + signalReporter *signalReporter + // routingOverrides is the allocation engine's seat→account routing-override + // cache (I-side §6.5). Shared across generations, polled by the supervisor; the + // group-route hot path reads it to redirect a seat off an unhealthy default. + // nil-safe → empty/unset means every request uses the local seatassign pick. + // Set via SetRoutingOverrides. See routing_override.go. + routingOverrides *RoutingOverrideCache + // poolObservedResets holds the latest upstream window-reset epoch observed per + // pool account (Path Z, 通道3 §14). The N7c pull piggybacks it to master so it + // re-rolls window_max_util_pct per window. Always non-nil; only written on + // group-route responses. + poolObservedResets *poolResetStore + // consoleURL is Config.ConsoleURL — the co-installed local console base used + // to assemble the member-login URL in OAUTH_GROUP_MEMBER_LOGIN_REQUIRED + // responses (20260703 update). "" ⇒ URL-less fallback wording. Set once via + // SetConsoleURL before serving; not hot-swapped. + consoleURL string + // groupLoginState is the bypass ~/.aikey/run/group-login-required.json store + // consumed by `aikey statusline`. Always non-nil (set in New); written on + // login-required 401s, cleared on the next successful group resolve. + groupLoginState *groupLoginStateStore + // routingRailHealth probes the routing_override SyncRail state (SyncRail + // §5.4, 2026-07-03): stale/offline → the login-required 401 wording says + // "routing sync unreachable" instead of a possibly-misdirected sign-in + // prompt (the incident shape: local pick ≠ engine assignment). nil = ok. + routingRailHealth func() (state string, failingSeconds int64) // filterHook is the P4 filter dispatcher — a generic apphook.Hook // (ai-compliance-detector / DLP / etc.) that inspects the inbound // request body before forwarding. Nil = no filter (the common default @@ -186,12 +235,26 @@ type Proxy struct { // Must be called before serving requests. A nil value restores the default // behavior (http.DefaultTransport, which honors HTTP_PROXY / HTTPS_PROXY env vars). func (p *Proxy) SetTransport(t http.RoundTripper) { - p.transport = t + p.transport.Store(&transportBox{rt: t}) if t != nil { slog.Info("proxy: custom transport set") } } +// transportBox boxes the RoundTripper so it can live in an atomic.Pointer (atomics +// can't hold an interface value directly). A nil rt means "use the default". +type transportBox struct{ rt http.RoundTripper } + +// currentTransport returns the live RoundTripper (nil → caller falls back to the +// default). Lock-free atomic read: safe on the hot path concurrently with a +// SetTransport hot-swap. +func (p *Proxy) currentTransport() http.RoundTripper { + if b := p.transport.Load(); b != nil { + return b.rt + } + return nil +} + // New creates a new Proxy. ctx is the proxy lifecycle context; canceling it // stops all detached upstream calls (called on proxy shutdown). // If v also implements ActiveKeyReader, path-prefix routing is enabled automatically. @@ -207,6 +270,9 @@ func New(v VaultGetter, reg *vkeys.Registry, prov *provider.Registry, coll *even VerySlowRequestMs: 10000, UpstreamTimeout: defaultUpstreamTimeout, appHealthCache: apppipe.NewHealthCache(), + poolCooldown: newPoolCooldownStore(), + poolObservedResets: newPoolResetStore(), + groupLoginState: newGroupLoginStateStore(), } if ar, ok := v.(ActiveKeyReader); ok { p.activeReader = ar @@ -225,6 +291,14 @@ func New(v VaultGetter, reg *vkeys.Registry, prov *provider.Registry, coll *even if pv, ok := v.(probepipe.VaultReader); ok { p.probeVault = pv } + // Oauth-group routing (N8) decrypts per-account group material at request + // time with the vault derived key. *vault.Reader exposes DerivedKey(); a + // mock that doesn't simply leaves p.groupKey nil → group routing degrades + // (GROUP_KEY_UNAVAILABLE) instead of panicking. Tests can swap via + // SetGroupKeyProvider. + if kp, ok := v.(groupKeyProvider); ok { + p.groupKey = kp + } return p } @@ -251,6 +325,14 @@ func (p *Proxy) SetProbeVault(pv probepipe.VaultReader) { p.probeVault = pv } +// SetGroupKeyProvider injects the vault derived-key accessor for oauth-group +// material decryption (N8). Mirrors SetProbeVault — auto-wired by New(...) when +// the VaultGetter argument implements DerivedKey(); tests that exercise group +// routing inject it explicitly. +func (p *Proxy) SetGroupKeyProvider(kp groupKeyProvider) { + p.groupKey = kp +} + // SetQuotaEnforcer injects the Phase 2 token-quota gate (Stage 3). Wired by the // supervisor from the per-process snapshot+counter. Safe to leave unset — a nil // enforcer is a no-op (no enforcement), which is the correct behavior for @@ -259,6 +341,24 @@ func (p *Proxy) SetQuotaEnforcer(e *quota.Enforcer) { p.quota = e } +// SetConsoleURL injects Config.ConsoleURL (20260703 update) — the co-installed +// local console base used to assemble the member-login URL in group +// login-required responses. Safe to leave unset: "" degrades to URL-less +// wording (cluster nodes / server-side proxies have no local console). +func (p *Proxy) SetConsoleURL(u string) { + p.consoleURL = u +} + +// SetRoutingRailHealth injects the supervisor's routing_override SyncRail +// health probe (SyncRail §5.4, 2026-07-03). respondLoginRequired consults it: +// when the assignment rail is stale/offline the local ranked pick may +// contradict the engine (the 2026-07-03 incident shape), so the 401 must say +// "routing sync unreachable" instead of directing the member to sign into a +// possibly-wrong account. nil (tests / framework off) → treated as healthy. +func (p *Proxy) SetRoutingRailHealth(fn func() (state string, failingSeconds int64)) { + p.routingRailHealth = fn +} + // AppHealthSnapshot returns the in-process record of the most recent // app-pipeline call per slug. Used by the admin /admin/apps/health endpoint // (wired through admin.Handler.AppHealthFn) to populate the Web "Connected @@ -275,6 +375,32 @@ func (p *Proxy) AppHealthSnapshot() []apppipe.AppHealth { return p.appHealthCache.Snapshot() } +// PoolCooldownSnapshot returns the oauth-group accounts currently in reactive +// cooldown (account_id → seconds remaining) for the admin /status health surface +// (N9 组路由健康). nil when nothing is cooling. The cmd layer wraps this into the +// admin DTO so the operator monitoring the first pool batch can see which +// accounts are being routed around. +func (p *Proxy) PoolCooldownSnapshot() map[string]int { + return p.poolCooldown.snapshot() +} + +// CooldownSkipSet returns the EXACT set of accounts the hot-path group resolver skips +// right now (poolCooldown.skipSet) so the supervisor's is_current_routed stamp can be +// computed with the SAME cooldown view the forward path uses — closing the display↔actual +// gap for cooling-driven switches (2026-07-01). Same function as group_serve.go's resolver +// call, so the two can't diverge in which accounts are considered cooled. +func (p *Proxy) CooldownSkipSet() map[string]bool { + return p.poolCooldown.skipSet() +} + +// ObservedResetsSnapshot returns the latest upstream window-reset epoch observed +// per pool account (account_id → epoch). The supervisor's N7c pull piggybacks +// it to master (Path Z) so master re-rolls window_max_util_pct per window. nil +// when nothing observed yet. +func (p *Proxy) ObservedResetsSnapshot() map[string]int64 { + return p.poolObservedResets.snapshot() +} + // recordAppHealth is the proxy-internal hook called from handleAppPipeline // after each request completes (success OR error path). Centralized so the // "what counts as a health observation" decision lives in one place: any @@ -379,6 +505,23 @@ func (p *Proxy) SetReporter(r *events.Reporter, instanceID, clientVersion, confi p.loggedInAccountID = loggedInAccountID } +// EnableSignalReporting wires the allocation-engine util signal reporter (I5): the +// proxy parses upstream unified-* utilization and best-effort POSTs it to master's +// /accounts/me/signals, authed with the same team account-JWT the group-runtime +// poll uses. nil controlURL/bearer → feature stays off (newSignalReporter returns nil). +func (p *Proxy) EnableSignalReporting(controlURL string, bearer func(ctx context.Context) (string, error)) { + p.signalReporter = newSignalReporter(controlURL, bearer, slog.Default()) +} + +// StopSignalReporting stops the signal reporter's upload loop (idempotent, +// nil-safe). Called from generation.close() so the per-generation reporter's +// goroutine + ticker don't leak across reloads. +func (p *Proxy) StopSignalReporting() { + if p.signalReporter != nil { + _ = p.signalReporter.Close() + } +} + // SetWAL attaches a local WAL writer for offline-mode usage events. // When a reporter is configured the WAL is shared with it (set once at // supervisor level) and the reporter performs the append. When the reporter diff --git a/internal/proxy/proxy_test.go b/internal/proxy/proxy_test.go index 802cc22..7aded0e 100644 --- a/internal/proxy/proxy_test.go +++ b/internal/proxy/proxy_test.go @@ -50,6 +50,21 @@ func setupTestProxy(t *testing.T, upstreamURL string) *Proxy { func setupTestProxyWithStore(t *testing.T, upstreamURL string, store events.EventInserter) *Proxy { t.Helper() + // Per-test run-dir isolation (2026-07-04). TestMain sandboxes AIKEY_RUN_DIR to + // ONE package-wide temp dir; the pool cooldown store persists to + hydrates from + // pool-cooldown.json under it (survive-restart, oauth_pool_cooldown.go). Sharing + // that one file across tests means a test that cools an account (acc-1, …) leaks + // it to every later test whose proxy hydrates the same file — so group_serve + // tests passed in isolation but failed together (cooled accounts → NO_CANDIDATES / + // ALL_UNUSABLE instead of the expected outcome). Giving each constructed proxy its + // OWN run dir makes every test hydrate an EMPTY cooldown file. This is the SINGLE + // construction point for all proxy tests (setupTestProxy delegates here), so one + // t.Setenv covers group + window + fallback tests. Single-proxy tests that read + // AIKEY_RUN_DIR later (group-login-required.json) stay consistent — t.Setenv holds + // for the whole test. Restart-survival is covered separately by the persist tests, + // which build the store directly (not via this helper). + t.Setenv("AIKEY_RUN_DIR", t.TempDir()) + v := &mockVault{ secrets: map[string]string{ "openai:test": "sk-real-openai-key-123", diff --git a/internal/proxy/routing_override.go b/internal/proxy/routing_override.go new file mode 100644 index 0000000..8415271 --- /dev/null +++ b/internal/proxy/routing_override.go @@ -0,0 +1,154 @@ +// routing_override.go — I-side §6.5 keystone (proxy consumption): the in-memory +// cache of the allocation engine's seat→account routing overrides pulled from +// master's GET /accounts/me/routing. +// +// This is a THIN redirect layer over the local seatassign pick. The supervisor's +// pollRoutingOverrides (mirror of pollGroupRuntime) GETs the sparse assignments +// map every ~60s and Store()s it here; handleOauthGroupRoute looks up the seat and +// hands the override account to resolveGroupCredential, which applies it ONLY when +// it is still a valid candidate (§6.5 member-validity re-check). On any doubt the +// resolver falls back to the local pick — the engine can redirect serving, never +// break it. +// +// MAIN-LINK SAFETY (架构第一优先级 = 主链路健壮): lookup is a lock-free atomic.Value +// read on the hot path; a nil receiver / unset cache / unknown seat returns "" so +// the caller always has the local pick as the default. Store is ONLY called with a +// fresh successful pull — a poll error never clears the cache (keep-last-known), +// so a master outage degrades to "no redirect", never to "no serving". +// +// The cache is SHARED across proxy generations: one instance lives on the +// supervisor and is injected into every Proxy (SetRoutingOverrides), so a vault +// reload never loses the overrides (mirrors why quotaSnapshot/quotaCounter are +// supervisor-scoped, not per-generation). +package proxy + +import ( + "sync/atomic" + + "github.com/AiKeyLabs/pkg/routingwire" +) + +// routeKey is the composite (seat, group) key of the in-memory override maps — a +// PACKAGE-PRIVATE cache implementation detail. The wire itself carries structured +// routingwire.RouteEntry objects (shared DTO module, 2026-07-02), so there is no +// cross-repo key-format contract anymore; only StoreRoutes below builds these keys. +// A multi-pool seat (member of >1 group) has a distinct binding per group. "|" is +// safe: seat/group ids are UUIDs. +func routeKey(seatID, groupID string) string { return seatID + "|" + groupID } + +// StoreRoutes ingests the wire's structured route entries (pkg/routingwire) and +// atomically replaces the cache — the ONLY place wire entries become cache keys, so +// the composite-key encoding never leaks past this package. Returns the bound/blocked +// counts for the caller's log line. +func (c *RoutingOverrideCache) StoreRoutes(version int64, routes []routingwire.RouteEntry) (bound, blocked int) { + assignments := make(map[string]string, len(routes)) + blockedSet := map[string]bool{} + for _, r := range routes { + if r.Blocked { + blockedSet[routeKey(r.SeatID, r.GroupID)] = true + continue + } + if r.AccountID != "" { + assignments[routeKey(r.SeatID, r.GroupID)] = r.AccountID + } + } + c.StoreAll(version, assignments, blockedSet) + return len(assignments), len(blockedSet) +} + +// RoutingOverrideCache holds the engine's sparse seat_id→account_id assignments +// plus the routing_version they came from. +type RoutingOverrideCache struct { + m atomic.Value // map[string]string; immutable once Stored (poll builds a fresh map each pull) + blocked atomic.Value // map[string]bool; seats the engine left UNBOUND (pool at the ≤3-人/号 cap) + version atomic.Int64 + stored atomic.Bool // false until the first Store — distinguishes "never pulled" from "pulled at version 0" +} + +// NewRoutingOverrideCache returns an empty cache (every lookup misses → local pick). +func NewRoutingOverrideCache() *RoutingOverrideCache { return &RoutingOverrideCache{} } + +// Store replaces assignments at a version with NO blocked seats — kept for callers +// that don't carry the blocked set; delegates to StoreAll. +func (c *RoutingOverrideCache) Store(version int64, assignments map[string]string) { + c.StoreAll(version, assignments, nil) +} + +// StoreAll atomically replaces the assignments map AND the blocked set and records +// the version. nil maps are normalized to empty so readers can always type-assert; +// the maps are read-only after StoreAll, so lookup/Blocked need no lock. +func (c *RoutingOverrideCache) StoreAll(version int64, assignments map[string]string, blocked map[string]bool) { + if c == nil { + return + } + if assignments == nil { + assignments = map[string]string{} // empty = engine redirects nothing + } + if blocked == nil { + blocked = map[string]bool{} // empty = no seat is pool-full-blocked + } + c.m.Store(assignments) + c.blocked.Store(blocked) + c.version.Store(version) + c.stored.Store(true) +} + +// Stored reports whether anything has ever been Stored. The poll uses it to +// distinguish "never pulled" (version 0 because the atomic is zero-valued) from +// "pulled at routing_version 0" — without it, master's first non-empty payload +// carrying routing_version:0 would match Version()==0 and be skipped forever. +func (c *RoutingOverrideCache) Stored() bool { + return c != nil && c.stored.Load() +} + +// Version returns the last Stored routing_version (0 when nothing stored yet) so +// the poll can skip re-applying an unchanged version. +func (c *RoutingOverrideCache) Version() int64 { + if c == nil { + return 0 + } + return c.version.Load() +} + +// lookup returns the engine's override account_id for (seatID, groupID), or "" when +// the feature is off / nothing stored / no override for this (seat,group). Keyed by +// the composite routeKey so a multi-pool seat gets the RIGHT group's override (not just +// its first group's). nil-safe so the hot path can call it unconditionally. +func (c *RoutingOverrideCache) lookup(seatID, groupID string) string { + if c == nil || seatID == "" || groupID == "" { + return "" + } + v := c.m.Load() + if v == nil { + return "" + } + return v.(map[string]string)[routeKey(seatID, groupID)] +} + +// Assignment returns the engine's override account_id for (seatID, groupID) (exported +// wrapper over lookup), or "" when nothing is overridden. Used by the supervisor's +// group_runtime writer to stamp IsCurrentRouted (C2 display) with the SAME override +// the hot-path resolver applies — one definition of "which account is routed". +func (c *RoutingOverrideCache) Assignment(seatID, groupID string) string { + return c.lookup(seatID, groupID) +} + +// Blocked reports whether the engine left (seatID, groupID) UNBOUND because every +// account in its pool/segment is at the ≤3-人/号 cap. The proxy 429s a blocked seat +// instead of falling back to the cap-blind local pick (which would route a 4th user +// onto a full account). nil-safe; false when the feature is off / nothing stored. +func (c *RoutingOverrideCache) Blocked(seatID, groupID string) bool { + if c == nil || seatID == "" || groupID == "" { + return false + } + v := c.blocked.Load() + if v == nil { + return false + } + return v.(map[string]bool)[routeKey(seatID, groupID)] +} + +// SetRoutingOverrides injects the shared, supervisor-owned routing-override cache +// (one instance across all generations). nil-safe: an unset cache makes every +// lookup miss → the local seatassign pick is always the default. +func (p *Proxy) SetRoutingOverrides(c *RoutingOverrideCache) { p.routingOverrides = c } diff --git a/internal/proxy/routing_override_test.go b/internal/proxy/routing_override_test.go new file mode 100644 index 0000000..34094a2 --- /dev/null +++ b/internal/proxy/routing_override_test.go @@ -0,0 +1,273 @@ +package proxy + +import ( + "testing" + + "github.com/AiKeyLabs/aikey-proxy/internal/vkeys" +) + +// freshOAuth is a fully-usable OAuth material entry (far-future expiry, active +// window) for the oauth-group resolver tests below. +func freshOAuth(t *testing.T, key []byte, token string) vkeys.GroupRuntimeAccount { + return encMat(t, key, vkeys.GroupRuntimeAccount{ + CredentialType: "oauth_account", ExpiresAt: 9_000_000_000, WindowStatus: "active", + }, token) +} + +// threeAccountRoute builds a 3-candidate group route where every account has +// fresh, usable material. Returns the route + the seatassign rank order so a test +// can name the local primary and a distinct override target. +func threeAccountRoute(t *testing.T, key []byte, seat string) (*vkeys.ResolvedRoute, []string) { + t.Helper() + refs := []vkeys.GroupAccountRef{ + {AccountID: "acc-a", ProviderCode: "anthropic"}, + {AccountID: "acc-b", ProviderCode: "anthropic"}, + {AccountID: "acc-c", ProviderCode: "anthropic"}, + } + mat := map[string]vkeys.GroupRuntimeAccount{ + "acc-a": freshOAuth(t, key, "tok-acc-a"), + "acc-b": freshOAuth(t, key, "tok-acc-b"), + "acc-c": freshOAuth(t, key, "tok-acc-c"), + } + route := &vkeys.ResolvedRoute{SeatID: seat, OauthGroupID: "grp", GroupAccounts: mustJSON(t, refs), GroupRuntime: mustJSON(t, mat)} + return route, rankOrder(seat, "acc-a", "acc-b", "acc-c") +} + +// A valid engine override (override account IS a serving candidate) redirects the +// seat off its local primary to the engine's healthy pick. +func TestResolveGroup_OverrideValidCandidateRedirects(t *testing.T) { + key := grKey() + seat := "seat-ovr-1" + route, order := threeAccountRoute(t, key, seat) + primary := order[0] + override := order[len(order)-1] // a different, valid candidate + + if override == primary { + t.Fatalf("test setup: override %q must differ from primary %q", override, primary) + } + + // Feed the override through the cache exactly like the hot path does. + cache := NewRoutingOverrideCache() + cache.Store(7, map[string]string{routeKey(seat, route.OauthGroupID): override}) + + res, err := resolveGroupCredential(route, key, 1_000_000, nil, cache.lookup(seat, route.OauthGroupID)) + if err != nil { + t.Fatalf("resolve: %v", err) + } + if res.AccountID != override { + t.Fatalf("engine override should redirect to %q, got %q", override, res.AccountID) + } + if res.OAuth == nil || res.OAuth.AccessToken != "tok-"+override { + t.Fatalf("override resolution wrong token: %+v", res.OAuth) + } + // The local rank-0 primary is still reported so the switch is audited (N9 #8). + if res.Primary != primary { + t.Fatalf("Primary must stay the local rank-0 %q for audit, got %q", primary, res.Primary) + } + if res.Primary == res.AccountID { + t.Fatal("an applied override must differ from the local primary (switch case)") + } +} + +// A STALE override (the account is no longer a candidate in this group's set) is +// ignored — the §6.5 member-validity re-check falls back to the local pick rather +// than routing to an account the proxy has no material for. +func TestResolveGroup_StaleOverrideFallsBackToLocal(t *testing.T) { + key := grKey() + seat := "seat-ovr-2" + route, order := threeAccountRoute(t, key, seat) + primary := order[0] + + cache := NewRoutingOverrideCache() + cache.Store(3, map[string]string{routeKey(seat, route.OauthGroupID): "ghost-account"}) // not in the candidate set + + res, err := resolveGroupCredential(route, key, 1_000_000, nil, cache.lookup(seat, route.OauthGroupID)) + if err != nil { + t.Fatalf("resolve: %v", err) + } + if res.AccountID != primary { + t.Fatalf("stale override must fall back to local primary %q, got %q", primary, res.AccountID) + } + if res.Primary != res.AccountID { + t.Fatalf("local pick → no switch: Primary(%q) must equal pick(%q)", res.Primary, res.AccountID) + } +} + +// OWNER RULE (2026-07-01): the engine MAY route a member to an account they have NOT +// logged into — an override naming a needs_login account is HONORED: the hot path stops +// there with LOGIN_REQUIRED for THAT account (matching what vault/virtual-keys/team-oauth +// display as the routed account), it does NOT silently fall through to a logged-in local +// pick. 能红: make the picker treat a needs_login override as unusable (fall through) → +// this resolves the local pick instead of erroring → fails. +func TestResolveGroup_NeedsLoginOverridePromptsLoginForThatAccount(t *testing.T) { + key := grKey() + seat := "seat-ovr-nl" + now := int64(1_000_000) + + refs := []vkeys.GroupAccountRef{ + {AccountID: "acc-a", ProviderCode: "anthropic"}, + {AccountID: "acc-b", ProviderCode: "anthropic"}, + } + order := rankOrder(seat, "acc-a", "acc-b") + primary := order[0] // logged in + usable — the fall-through would serve this + override := order[1] // the ENGINE's pick — member not logged in + mat := map[string]vkeys.GroupRuntimeAccount{ + primary: freshOAuth(t, key, "tok-primary"), + override: {CredentialType: "oauth_account", NeedsLogin: true}, + } + route := &vkeys.ResolvedRoute{SeatID: seat, OauthGroupID: "grp", GroupAccounts: mustJSON(t, refs), GroupRuntime: mustJSON(t, mat)} + + _, err := resolveGroupCredential(route, key, now, nil, override) + if err == nil { + t.Fatal("needs_login override must NOT silently serve the local pick — want LOGIN_REQUIRED") + } + ge, ok := err.(*groupResolveError) + if !ok || ge.Code != groupErrLoginRequired { + t.Fatalf("want LOGIN_REQUIRED, got %v", err) + } + if ge.Account != override { + t.Fatalf("login prompt must name the ENGINE-routed account %q (the one all pages display), got %q", override, ge.Account) + } +} + +// An override pointing at a candidate ref WITHOUT usable material (expired) is also +// ignored: the override path runs the SAME usability gate as the ranked loop, so it +// never routes to an unusable account — it falls back to the local pick. +func TestResolveGroup_OverrideWithoutUsableMaterialFallsBack(t *testing.T) { + key := grKey() + seat := "seat-ovr-3" + now := int64(1_000_000) + + refs := []vkeys.GroupAccountRef{ + {AccountID: "acc-a", ProviderCode: "anthropic"}, + {AccountID: "acc-b", ProviderCode: "anthropic"}, + } + order := rankOrder(seat, "acc-a", "acc-b") + primary := order[0] + // The override target is a candidate ref but its token is expired (no usable + // material) — engine picked it before it expired locally. + expired := order[1] + mat := map[string]vkeys.GroupRuntimeAccount{ + primary: freshOAuth(t, key, "tok-primary"), + expired: encMat(t, key, vkeys.GroupRuntimeAccount{CredentialType: "oauth_account", ExpiresAt: now - 1}, "tok-expired"), + } + route := &vkeys.ResolvedRoute{SeatID: seat, GroupAccounts: mustJSON(t, refs), GroupRuntime: mustJSON(t, mat)} + + res, err := resolveGroupCredential(route, key, now, nil, expired) + if err != nil { + t.Fatalf("resolve: %v", err) + } + if res.AccountID != primary { + t.Fatalf("override-without-material must fall back to local primary %q, got %q", primary, res.AccountID) + } +} + +// No override ("") leaves the local seatassign pick unchanged. +func TestResolveGroup_NoOverrideUsesLocalPick(t *testing.T) { + key := grKey() + seat := "seat-ovr-4" + route, order := threeAccountRoute(t, key, seat) + primary := order[0] + + res, err := resolveGroupCredential(route, key, 1_000_000, nil, "") + if err != nil { + t.Fatalf("resolve: %v", err) + } + if res.AccountID != primary { + t.Fatalf("no override → local primary %q, got %q", primary, res.AccountID) + } +} + +// §11.1 invariant WITH the override layer present: when the engine is down (poll +// never succeeded → empty cache), serving still resolves via local seatassign. +func TestResolveGroup_EngineDownEmptyCacheLocalPick(t *testing.T) { + key := grKey() + seat := "seat-ovr-5" + route, order := threeAccountRoute(t, key, seat) + primary := order[0] + + cache := NewRoutingOverrideCache() // never Stored — engine down + if got := cache.lookup(seat, route.OauthGroupID); got != "" { + t.Fatalf("empty cache must miss, got %q", got) + } + res, err := resolveGroupCredential(route, key, 1_000_000, nil, cache.lookup(seat, route.OauthGroupID)) + if err != nil { + t.Fatalf("resolve: %v", err) + } + if res.AccountID != primary { + t.Fatalf("engine down → local primary %q, got %q", primary, res.AccountID) + } +} + +// The cache is nil-safe and keep-last-known: a nil receiver / unset cache misses; +// Store replaces the whole map; Version gates re-apply. +func TestRoutingOverrideCache_NilSafeAndStore(t *testing.T) { + var nilCache *RoutingOverrideCache + if got := nilCache.lookup("seat", "g"); got != "" { + t.Fatalf("nil cache lookup must be empty, got %q", got) + } + if v := nilCache.Version(); v != 0 { + t.Fatalf("nil cache version must be 0, got %d", v) + } + nilCache.Store(1, map[string]string{routeKey("seat", "g"): "x"}) // must not panic + + c := NewRoutingOverrideCache() + if got := c.lookup("seat", "g"); got != "" { + t.Fatalf("fresh cache must miss, got %q", got) + } + c.Store(5, map[string]string{routeKey("seat-1", "g"): "acc-1"}) + if got := c.lookup("seat-1", "g"); got != "acc-1" { + t.Fatalf("lookup after store: got %q", got) + } + if got := c.lookup("seat-unknown", "g"); got != "" { + t.Fatalf("unknown seat must miss, got %q", got) + } + // Multi-pool (2026-07-01): the SAME seat in TWO groups keeps DISTINCT overrides — + // keying by seat alone (the old bug) would collapse them to one. + c.Store(5, map[string]string{routeKey("seat-1", "g1"): "acc-g1", routeKey("seat-1", "g2"): "acc-g2"}) + if got := c.lookup("seat-1", "g1"); got != "acc-g1" { + t.Fatalf("seat-1/g1 override: got %q", got) + } + if got := c.lookup("seat-1", "g2"); got != "acc-g2" { + t.Fatalf("seat-1/g2 override (multi-pool, distinct per group): got %q", got) + } + if v := c.Version(); v != 5 { + t.Fatalf("version: got %d", v) + } + // A later store replaces the whole map (old key gone, new key present). + c.Store(6, map[string]string{routeKey("seat-2", "g"): "acc-2"}) + if got := c.lookup("seat-1", "g1"); got != "" { + t.Fatalf("replaced map must drop seat-1/g1, got %q", got) + } + if got := c.lookup("seat-2", "g"); got != "acc-2" { + t.Fatalf("replaced map must carry seat-2, got %q", got) + } + // A nil map normalizes to empty (no panic, all lookups miss). + c.Store(7, nil) + if got := c.lookup("seat-2", "g"); got != "" { + t.Fatalf("nil-map store must clear, got %q", got) + } +} + +// TestRoutingOverrideCache_StoredDistinguishesVersionZero is the regression for the +// first-pull-at-version-0 skip hole (review HIGH): the cache's version atomic is +// zero-valued, so the poll cannot treat Version()==0 alone as "unchanged" — master's +// first non-empty payload at routing_version 0 must still be applied. Stored() makes +// the "never pulled" vs "pulled at 0" distinction the poll's skip guard now relies on. +func TestRoutingOverrideCache_StoredDistinguishesVersionZero(t *testing.T) { + c := NewRoutingOverrideCache() + if c.Stored() { + t.Fatal("fresh cache must report Stored()==false (never pulled)") + } + if c.Version() != 0 { + t.Fatalf("fresh cache Version() must be 0, got %d", c.Version()) + } + // First pull: a non-empty assignment map carrying routing_version 0. + c.Store(0, map[string]string{routeKey("seat-1", "g"): "acct-1"}) + if !c.Stored() { + t.Fatal("after Store(0,...) Stored() must be true — else the poll skips it forever") + } + if got := c.lookup("seat-1", "g"); got != "acct-1" { + t.Fatalf("version-0 non-empty payload must be applied, lookup got %q want acct-1", got) + } +} diff --git a/internal/proxy/signal_report.go b/internal/proxy/signal_report.go new file mode 100644 index 0000000..63c1733 --- /dev/null +++ b/internal/proxy/signal_report.go @@ -0,0 +1,507 @@ +package proxy + +import ( + "bytes" + "context" + "encoding/json" + "log/slog" + "net/http" + "strconv" + "strings" + "sync" + "time" + + "github.com/AiKeyLabs/aikey-proxy/internal/httpx" +) + +// ponytail: flush cadence doubles as the rate-limit count window — one constant +// so loop()'s ticker and the reported window_secs can never drift apart. +const signalFlushInterval = 30 * time.Second + +// signal_report.go — I5 proxy emit. The proxy parses the upstream's unified-* +// 5h-utilization off the response headers and ships it best-effort to master's +// POST /accounts/me/signals, where it lands in the account's engine_meta sample +// ring for the allocation engine to read. +// +// MAIN-LINK SAFETY (架构第一优先级 = 主链路健壮): this never touches the forward +// path's latency or success. enqueue() is a NON-BLOCKING channel send (a full +// buffer drops the sample — utilization is a trend signal, a lost reading is +// harmless); the upload runs in a background goroutine whose failures are only +// logged, never surfaced to the request. A nil reporter = feature off. + +// signalSample is one parsed util reading queued for delivery. Util7d is the +// 7-day-window sibling of Util5h (master reads it into engine_meta.util_7d for +// the weekly-squeeze target). It's `omitempty` so a 5h-only reading stays +// byte-identical to the pre-7d wire — ponytail: a genuine util_7d==0.0 is +// indistinguishable from "no 7d header" on the wire, which is fine for a trend +// signal (other samples in the window carry the non-zero reading); switch to +// *float64 only if master ever needs to tell zero-util from no-data. +type signalSample struct { + CredentialID string `json:"credential_id"` + TS int64 `json:"ts"` + Util5h float64 `json:"util_5h"` + Util7d float64 `json:"util_7d,omitempty"` +} + +// concurrencySample reports the PEAK number of concurrent in-flight forwarded +// requests a credential drew within one flush window. The allocation engine's +// DeriveFanout (§8.1) treats a high-concurrency account as several humans, so a +// raw peak (not an average) is what it wants. Like rateLimitSample this is a +// per-credential aggregate reset every flush, not a per-event channel. +type concurrencySample struct { + CredentialID string `json:"credential_id"` + Peak int `json:"peak"` +} + +// revokedSample flags a credential whose OAuth token the upstream hard-revoked +// (401 "OAuth token has been revoked") — master quarantines the account so the +// allocation engine stops picking it. Same best-effort, fault-isolated delivery +// as signalSample (see file header); a dropped flag just delays quarantine until +// the next 401 re-emits it, so loss is self-healing. +type revokedSample struct { + CredentialID string `json:"credential_id"` + Reason string `json:"reason"` +} + +// rateLimitSample reports how many 429/403 responses a credential drew in the +// last flush window. Master normalizes count/window into the allocation engine's +// §5.1 w5 "Recent429FreqNorm" risk signal (near-window 429/403 rhythm). Unlike +// util/revoked these are a per-credential COUNTER (a map+mutex, not a channel): +// 429s are frequent — a rate-limited account can draw a burst per second — so a +// per-event channel send would be a POST storm. The integer is reset every flush, +// which bounds the map to the credentials seen in one window. WindowSecs = the +// flush cadence so master can divide count by it without assuming the interval. +type rateLimitSample struct { + CredentialID string `json:"credential_id"` + Count int `json:"count"` + WindowSecs int `json:"window_secs"` +} + +type signalReporter struct { + url string + bearer func(ctx context.Context) (string, error) // account-JWT (reuses the group-runtime poll credential) + client *httpx.SwappableClient // control-plane: rebuilt on host network change (self-heal registry) + in chan signalSample + revokedIn chan revokedSample + rlMu sync.Mutex // guards rlCounts + rlCounts map[string]int // per-credential 429/403 tally, reset each flush + // in-flight concurrency tracking. inflCur is the LIVE count (inc on forward + // start, dec on completion — persists across windows so a request spanning a + // flush keeps being counted); inflPeak is the max inflCur seen this window, + // reset each flush. Both guarded by inflMu. + // ponytail: one global mutex for both maps — the critical section is two map + // ops; split to per-credential locks only if this lock ever shows up hot. + inflMu sync.Mutex + inflCur map[string]int + inflPeak map[string]int + logger *slog.Logger + + // stop terminates loop() on Close so the goroutine + ticker don't leak. A + // fresh signalReporter is built per generation (buildGeneration); without + // this, every reload (aikey use, filter/quota/audit toggle, /admin/reload) + // leaked one loop() + 30s ticker holding a live bearer closure over the old + // vault reader. Close is idempotent (stopOnce) — generation.close() may run + // once but defensive against double-close. + stop chan struct{} + stopOnce sync.Once +} + +// newSignalReporter returns nil (feature off) unless both a control URL and an +// auth bearer are configured. On success it starts the background upload loop. +func newSignalReporter(controlURL string, bearer func(context.Context) (string, error), logger *slog.Logger) *signalReporter { + if controlURL == "" || bearer == nil { + return nil + } + r := &signalReporter{ + url: strings.TrimRight(controlURL, "/") + "/accounts/me/signals", + bearer: bearer, + client: httpx.NewSwappableDirect(10 * time.Second), + in: make(chan signalSample, 256), + revokedIn: make(chan revokedSample, 64), + rlCounts: make(map[string]int), + inflCur: make(map[string]int), + inflPeak: make(map[string]int), + logger: logger, + stop: make(chan struct{}), + } + go r.loop() + return r +} + +// Close stops the upload loop (idempotent, nil-safe). Called from +// generation.close() on every reload so the loop() goroutine + ticker don't +// leak. It does NOT close r.in / r.revokedIn — the forward path may still send +// there (non-blocking, default-drop), and closing them would panic that send. +func (r *signalReporter) Close() error { + if r == nil { + return nil + } + r.stopOnce.Do(func() { close(r.stop) }) + return nil +} + +// enqueue is a non-blocking best-effort hand-off from the forward path. util7d +// is the optional 7-day reading (pass 0 when the upstream sent no 7d header — +// signalSample.Util7d is omitempty so a 0 drops off the wire and util_5h stays +// byte-identical to the pre-7d format). +func (r *signalReporter) enqueue(credentialID string, ts int64, util5h, util7d float64) { + if r == nil || credentialID == "" { + return + } + select { + case r.in <- signalSample{CredentialID: credentialID, TS: ts, Util5h: util5h, Util7d: util7d}: + default: // buffer full → drop (trend signal; never block forwarding) + } +} + +// enqueueRevoked is the revoked-feed sibling of enqueue: a non-blocking, +// best-effort hand-off from the forward path's 401 detection. Same MAIN-LINK +// safety contract — nil receiver / empty id are dropped silently and a full +// buffer drops rather than blocks (a hard-ban 401s many in-flight requests at +// once, so the queue can fill; a dropped flag self-heals on the next 401). +func (r *signalReporter) enqueueRevoked(credentialID, reason string) { + if r == nil || credentialID == "" { + return + } + if reason == "" { + reason = "revoked" + } + select { + case r.revokedIn <- revokedSample{CredentialID: credentialID, Reason: reason}: + default: // buffer full → drop (best-effort; never block forwarding) + } +} + +// incrRateLimit tallies one upstream 429/403 for a credential. Same MAIN-LINK +// safety contract as enqueue (see file header): nil receiver / empty id are +// dropped, and the bump is a short-held map lock that never blocks forwarding. +// A counter (not a channel) collapses a 429 burst to O(1) memory per credential; +// the map is reset every flush (see snapshotRateLimits), so it stays bounded by +// the live credential set. Lazy-inits the map so a directly-constructed reporter +// (tests, future call sites) can't nil-panic. +func (r *signalReporter) incrRateLimit(credentialID string) { + if r == nil || credentialID == "" { + return + } + r.rlMu.Lock() + if r.rlCounts == nil { + r.rlCounts = make(map[string]int) + } + r.rlCounts[credentialID]++ + r.rlMu.Unlock() +} + +// snapshotRateLimits atomically reads + clears the 429/403 counters into wire +// samples. Reset-on-read is what bounds the map: each window only retains the +// credentials seen since the previous flush. Returns nil when nothing was seen +// (so an empty window omits the rate_limits array entirely). +func (r *signalReporter) snapshotRateLimits() []rateLimitSample { + r.rlMu.Lock() + defer r.rlMu.Unlock() + if len(r.rlCounts) == 0 { + return nil + } + out := make([]rateLimitSample, 0, len(r.rlCounts)) + for id, n := range r.rlCounts { + out = append(out, rateLimitSample{ + CredentialID: id, + Count: n, + WindowSecs: int(signalFlushInterval / time.Second), + }) + } + r.rlCounts = make(map[string]int) // reset the window + return out +} + +// trackInflight marks one in-flight forwarded request for credentialID and +// returns a release func to call (via `defer …()`) on completion. It bumps the +// live count and updates this window's peak under a short map lock. Same +// MAIN-LINK safety contract as the other hooks: a nil receiver / empty id yield +// a no-op closure so the caller can `defer trackInflight(id)()` unconditionally +// without a nil check, and the lock is held only for two map ops. +// +// The returned func MUST run (use defer) — it decrements the live count, so a +// missed call would leak a credential as permanently "in flight" and pin its +// peak high forever. defer at the forward call site guarantees it fires even on +// panic. Lazy-inits both maps so a directly-constructed reporter (tests) can't +// nil-panic. +func (r *signalReporter) trackInflight(credentialID string) func() { + if r == nil || credentialID == "" { + return func() {} // no-op: safe to `defer …()` + } + r.inflMu.Lock() + if r.inflCur == nil { + r.inflCur = make(map[string]int) + r.inflPeak = make(map[string]int) + } + r.inflCur[credentialID]++ + if r.inflCur[credentialID] > r.inflPeak[credentialID] { + r.inflPeak[credentialID] = r.inflCur[credentialID] + } + r.inflMu.Unlock() + return func() { + r.inflMu.Lock() + if r.inflCur[credentialID] > 0 { + r.inflCur[credentialID]-- + } + r.inflMu.Unlock() + } +} + +// snapshotConcurrency atomically reads + clears this window's per-credential +// peak into wire samples. Reset-on-read bounds the map to credentials seen this +// window — exactly like snapshotRateLimits. inflCur is deliberately NOT reset: +// requests still forwarding at the window boundary stay counted, so the next +// inc keeps computing a truthful peak (a cross-window burst still registers +// because inflPeak starts from 0 but inflCur already carries the live count). +// Returns nil for an idle window so post() omits the concurrency array. +// +// ponytail: a single long-lived stream sitting in a window with no NEW arrivals +// won't re-bump its peak, so a quiet steady-state-1 window can report nothing +// for that credential — harmless for a "several humans?" signal where the peak +// is what matters; revisit only if steady-state concurrency becomes a target. +func (r *signalReporter) snapshotConcurrency() []concurrencySample { + r.inflMu.Lock() + defer r.inflMu.Unlock() + if len(r.inflPeak) == 0 { + return nil + } + out := make([]concurrencySample, 0, len(r.inflPeak)) + for id, peak := range r.inflPeak { + out = append(out, concurrencySample{CredentialID: id, Peak: peak}) + } + r.inflPeak = make(map[string]int) // reset the window (inflCur persists) + return out +} + +func (r *signalReporter) loop() { + ticker := time.NewTicker(signalFlushInterval) + defer ticker.Stop() + var batch []signalSample + var revoked []revokedSample + // flush takes the rate-limit + concurrency snapshots as args rather than + // reading the per-window aggregates itself: the size-triggered early flushes + // below pass nil so they DON'T disturb those windows, keeping each span + // exactly one ticker period — only the ticker case snapshots, so the reported + // window_secs / peak stay accurate even when a util/revoked burst forces an + // early flush. + flush := func(rl []rateLimitSample, conc []concurrencySample) { + if len(batch) == 0 && len(revoked) == 0 && len(rl) == 0 && len(conc) == 0 { + return + } + r.post(batch, revoked, rl, conc) + batch = batch[:0] + revoked = revoked[:0] + } + for { + select { + case s := <-r.in: + if batch = append(batch, s); len(batch) >= 64 { + flush(nil, nil) + } + case rv := <-r.revokedIn: + // Revoked rides the SAME 30s batch flush rather than an immediate + // per-event POST: a hard-ban 401s many concurrent in-flight requests + // at once, so per-event posting would be a POST storm; batching reuses + // one flush path (simplicity + main-link isolation) and the engine + // tolerating ≤30s quarantine latency is an acceptable best-effort + // trade. The >=64 cap bounds a burst so it doesn't wait the full tick. + if revoked = append(revoked, rv); len(revoked) >= 64 { + flush(nil, nil) + } + case <-ticker.C: + flush(r.snapshotRateLimits(), r.snapshotConcurrency()) + case <-r.stop: + // Final flush on shutdown so a reload doesn't drop the in-flight + // batch, then return — ends the goroutine + ticker (no leak). + flush(r.snapshotRateLimits(), r.snapshotConcurrency()) + return + } + } +} + +func (r *signalReporter) post(samples []signalSample, revoked []revokedSample, rateLimits []rateLimitSample, concurrency []concurrencySample) { + // All arrays are optional: marshal only the non-empty ones so an all-samples, + // all-revoked, all-rate_limits, or all-concurrency flush still posts a valid + // body the master can decode. + payload := make(map[string]any, 4) + if len(samples) > 0 { + payload["samples"] = samples + } + if len(revoked) > 0 { + payload["revoked"] = revoked + } + if len(rateLimits) > 0 { + payload["rate_limits"] = rateLimits + } + if len(concurrency) > 0 { + payload["concurrency"] = concurrency + } + if len(payload) == 0 { + return + } + body, err := json.Marshal(payload) + if err != nil { + return + } + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + tok, err := r.bearer(ctx) + if err != nil { + r.logger.Warn("signal report bearer unavailable", + "event.name", "proxy.signal.bearer_failed", "error", err) + return + } + req, err := http.NewRequestWithContext(ctx, http.MethodPost, r.url, bytes.NewReader(body)) + if err != nil { + return + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer "+tok) + resp, err := r.client.Get().Do(req) + if err != nil { + r.logger.Warn("signal report upload failed", + "event.name", "proxy.signal.upload_failed", "error", err) + return + } + _ = resp.Body.Close() + if resp.StatusCode >= 300 { + r.logger.Warn("signal report rejected", + "event.name", "proxy.signal.upload_rejected", "status", resp.StatusCode) + } +} + +// parseUnifiedUtil5h reads the anthropic-ratelimit-unified-5h-utilization header +// (a float 0..1) off a response, returning (value, true) only on a clean parse — +// a missing / malformed header yields (0, false) so the caller simply skips it. +func parseUnifiedUtil5h(h http.Header) (float64, bool) { + raw := h.Get("anthropic-ratelimit-unified-5h-utilization") + if raw == "" { + return 0, false + } + v, err := strconv.ParseFloat(strings.TrimSpace(raw), 64) + if err != nil || v < 0 || v > 1 { + return 0, false + } + return v, true +} + +// parseUnifiedUtil7d is the 7-day-window sibling of parseUnifiedUtil5h: it reads +// the anthropic-ratelimit-unified-7d-utilization header (a float 0..1) the +// upstream sends alongside the 5h one, with identical (value, true)-only-on- +// clean-parse semantics. ponytail: a near-mirror of the 5h reader (3 lines of +// shared parse logic, below the extract-a-helper threshold) — kept separate so +// the master-facing names map 1:1 to the two header names. +func parseUnifiedUtil7d(h http.Header) (float64, bool) { + raw := h.Get("anthropic-ratelimit-unified-7d-utilization") + if raw == "" { + return 0, false + } + v, err := strconv.ParseFloat(strings.TrimSpace(raw), 64) + if err != nil || v < 0 || v > 1 { + return 0, false + } + return v, true +} + +// codex5hThresholdMinutes: a window shorter than this is the "5h" (short/head) +// window, longer is the "7d" (weekly) window. 360 = 6h, mirroring sub2api's +// single-window disambiguation. Live Plus windows are 300 (5h) and 10080 (7d). +const codex5hThresholdMinutes = 360 + +// parseCodexUtil reads Codex's per-window utilization off an upstream response +// (the X-Codex-* headers ride on every chatgpt.com /responses 200 — verified +// live 2026-07-06, see research/oauth-codex-ratelimit/2026-07-06-codex-ratelimit- +// taxonomy.md) and NORMALIZES it to the SAME (util_5h, util_7d) fraction shape +// the allocation engine consumes for Anthropic, so master stays provider-neutral +// (design §4B/§5.1). Two gotchas baked in: +// - percent is 0..100 → ÷100 (Anthropic util is already 0..1). +// - Codex ships TWO windows named primary/secondary whose durations are NOT +// fixed to 5h/7d (a Plus account's primary is 5h, but sub2api's own comment +// had it backwards). We classify by X-Codex-*-Window-Minutes — smaller = 5h, +// larger = 7d — NEVER by the primary/secondary name. +// Returns (0,0,false) when no X-Codex-*-used-percent header is present (i.e. the +// response is not Codex traffic), so the caller simply skips the sample — the +// same "clean-parse-only" contract as parseUnifiedUtil5h. util_7d is 0 when its +// window is absent (matches the Anthropic path's omitempty-0 behavior). +func parseCodexUtil(h http.Header) (util5h, util7d float64, ok bool) { + pUsed, pOK := parseCodexPercent(h.Get("X-Codex-Primary-Used-Percent")) + sUsed, sOK := parseCodexPercent(h.Get("X-Codex-Secondary-Used-Percent")) + if !pOK && !sOK { + return 0, 0, false // not Codex traffic (or no util headers) → skip + } + pMin := parseCodexInt(h.Get("X-Codex-Primary-Window-Minutes")) + sMin := parseCodexInt(h.Get("X-Codex-Secondary-Window-Minutes")) + + // primaryIs5h: is the primary window the SHORT (5h) one? Classify by duration. + primaryIs5h := true // default when minutes are absent: primary = 5h (head signal) + switch { + case pMin > 0 && sMin > 0: + primaryIs5h = pMin <= sMin + case pMin > 0: + primaryIs5h = pMin <= codex5hThresholdMinutes + case sMin > 0: + primaryIs5h = sMin > codex5hThresholdMinutes + } + + if primaryIs5h { + util5h, util7d = pUsed/100, sUsed/100 + } else { + util5h, util7d = sUsed/100, pUsed/100 + } + return util5h, util7d, true +} + +// parseCodexPercent parses an X-Codex-*-used-percent header (float 0..100+). Codex +// can report >100 briefly; we clamp the upper bound to 100 so the normalized +// fraction stays in [0,1] for the engine, but treat a missing/negative/garbage +// value as "absent" (ok=false). +func parseCodexPercent(raw string) (float64, bool) { + raw = strings.TrimSpace(raw) + if raw == "" { + return 0, false + } + v, err := strconv.ParseFloat(raw, 64) + if err != nil || v < 0 { + return 0, false + } + if v > 100 { + v = 100 + } + return v, true +} + +// parseCodexInt parses an X-Codex-*-window-minutes / -reset-after-seconds header; +// 0 on missing/garbage (treated as "unknown" by callers). +func parseCodexInt(raw string) int { + raw = strings.TrimSpace(raw) + if raw == "" { + return 0 + } + if i, err := strconv.Atoi(raw); err == nil && i > 0 { + return i + } + return 0 +} + +// isHardRevoked reports whether a 401 upstream response means the credential's +// OAuth token was HARD-revoked (gone for good) vs a routine expiry the refresh +// path recovers from. We gate on status==401 AND a documented hard-revocation +// marker in the parsed error type/message (errMsg is the raw body), so a plain +// token-expiry 401 does NOT quarantine an otherwise-healthy account: +// - anthropic: "OAuth token has been revoked" (contains "revoked"). +// - codex/openai (sub2api-derived, see research/oauth-codex-ratelimit): the +// error code "token_revoked" (contains "revoked") or "token_invalidated", +// or the non-standard body {"detail":"Unauthorized"} = permanently invalid. +// +// ponytail: narrow keyword/shape match on purpose to avoid false-positive +// quarantines; widen the term list here if other hard-ban phrasings show up. +func isHardRevoked(statusCode int, errType, errMsg string) bool { + if statusCode != http.StatusUnauthorized { + return false + } + blob := strings.ToLower(errType + " " + errMsg) + return strings.Contains(blob, "revoked") || + strings.Contains(blob, "token_invalidated") || + strings.Contains(blob, `"detail":"unauthorized"`) +} diff --git a/internal/proxy/signal_report_test.go b/internal/proxy/signal_report_test.go new file mode 100644 index 0000000..9e3eb41 --- /dev/null +++ b/internal/proxy/signal_report_test.go @@ -0,0 +1,546 @@ +package proxy + +import ( + "context" + "encoding/json" + "io" + "log/slog" + "net/http" + "net/http/httptest" + "sync" + "sync/atomic" + "testing" +) + +// signal_report_test.go — covers the I5 emit path's pure + best-effort pieces. +// loop()'s 30s ticker is timing-driven and not deterministically testable, so we +// call post() directly (the loop only batches + invokes it) — see TestSignalPost*. + +func TestParseUnifiedUtil5h(t *testing.T) { + const hdr = "anthropic-ratelimit-unified-5h-utilization" + tests := []struct { + name string + set bool + val string + wantV float64 + wantOK bool + }{ + {"valid", true, "0.6", 0.6, true}, + {"missing", false, "", 0, false}, + {"malformed", true, "abc", 0, false}, + {"above_one", true, "1.5", 0, false}, + {"negative", true, "-0.1", 0, false}, + {"zero_ok", true, "0", 0, true}, + {"one_ok", true, "1", 1, true}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + h := http.Header{} + if tt.set { + h.Set(hdr, tt.val) + } + v, ok := parseUnifiedUtil5h(h) + if ok != tt.wantOK { + t.Fatalf("ok = %v, want %v", ok, tt.wantOK) + } + if ok && v != tt.wantV { + t.Fatalf("v = %v, want %v", v, tt.wantV) + } + }) + } +} + +func TestEnqueueNilAndEmptyAreSafe(t *testing.T) { + // nil receiver guard: feature-off reporter must not panic. + var nilR *signalReporter + nilR.enqueue("c1", 100, 0.5, 0) // must not panic + + // empty credentialID is dropped — observable via the buffered channel. + r := &signalReporter{in: make(chan signalSample, 4)} + r.enqueue("", 100, 0.5, 0) + if len(r.in) != 0 { + t.Fatalf("empty credentialID should be dropped, buffered = %d", len(r.in)) + } + r.enqueue("c1", 100, 0.5, 0) + if len(r.in) != 1 { + t.Fatalf("valid sample should be queued, buffered = %d", len(r.in)) + } +} + +func TestSignalPostSendsBatch(t *testing.T) { + type captured struct { + method string + auth string + contentType string + body []byte + } + got := make(chan captured, 1) + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) { + b, _ := io.ReadAll(req.Body) + got <- captured{req.Method, req.Header.Get("Authorization"), req.Header.Get("Content-Type"), b} + })) + defer srv.Close() + + r := newSignalReporter(srv.URL, func(context.Context) (string, error) { return "tok-123", nil }, slog.Default()) + if r == nil { + t.Fatal("newSignalReporter returned nil") + } + r.post([]signalSample{{CredentialID: "c1", TS: 100, Util5h: 0.6}}, nil, nil, nil) + + c := <-got + if c.method != http.MethodPost { + t.Errorf("method = %q, want POST", c.method) + } + if c.auth != "Bearer tok-123" { + t.Errorf("Authorization = %q, want %q", c.auth, "Bearer tok-123") + } + if c.contentType != "application/json" { + t.Errorf("Content-Type = %q, want application/json", c.contentType) + } + var decoded struct { + Samples []signalSample `json:"samples"` + } + if err := json.Unmarshal(c.body, &decoded); err != nil { + t.Fatalf("body not valid JSON: %v (raw %s)", err, c.body) + } + if len(decoded.Samples) != 1 || decoded.Samples[0] != (signalSample{CredentialID: "c1", TS: 100, Util5h: 0.6}) { + t.Fatalf("decoded samples = %+v, want one {c1,100,0.6}", decoded.Samples) + } +} + +func TestSignalPostBearerErrorDoesNotPost(t *testing.T) { + var hits int32 + srv := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) { + atomic.AddInt32(&hits, 1) + })) + defer srv.Close() + + r := newSignalReporter(srv.URL, func(context.Context) (string, error) { + return "", io.ErrUnexpectedEOF + }, slog.Default()) + if r == nil { + t.Fatal("newSignalReporter returned nil") + } + r.post([]signalSample{{CredentialID: "c1", TS: 100, Util5h: 0.6}}, nil, nil, nil) // must not panic + + if n := atomic.LoadInt32(&hits); n != 0 { + t.Fatalf("server hit %d times, want 0 (bearer error short-circuits)", n) + } +} + +func TestEnqueueRevokedNilEmptyAndNonBlocking(t *testing.T) { + // nil receiver guard: feature-off reporter must not panic. + var nilR *signalReporter + nilR.enqueueRevoked("c1", "revoked") // must not panic + + // empty credentialID is dropped; valid one is queued; empty reason defaults. + r := &signalReporter{revokedIn: make(chan revokedSample, 4)} + r.enqueueRevoked("", "revoked") + if len(r.revokedIn) != 0 { + t.Fatalf("empty credentialID should be dropped, buffered = %d", len(r.revokedIn)) + } + r.enqueueRevoked("c1", "") // empty reason → defaults to "revoked" + if len(r.revokedIn) != 1 { + t.Fatalf("valid revoked should be queued, buffered = %d", len(r.revokedIn)) + } + if rv := <-r.revokedIn; rv != (revokedSample{CredentialID: "c1", Reason: "revoked"}) { + t.Fatalf("queued = %+v, want {c1, revoked}", rv) + } + + // non-blocking: a full buffer drops rather than blocks the forward path. + full := &signalReporter{revokedIn: make(chan revokedSample, 1)} + full.enqueueRevoked("c1", "revoked") + full.enqueueRevoked("c2", "revoked") // buffer full → dropped, must not block + if len(full.revokedIn) != 1 { + t.Fatalf("full buffer should drop, buffered = %d", len(full.revokedIn)) + } +} + +func TestSignalPostSendsRevoked(t *testing.T) { + type captured struct { + method string + auth string + body []byte + } + got := make(chan captured, 2) + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) { + b, _ := io.ReadAll(req.Body) + got <- captured{req.Method, req.Header.Get("Authorization"), b} + })) + defer srv.Close() + + r := newSignalReporter(srv.URL, func(context.Context) (string, error) { return "tok-123", nil }, slog.Default()) + if r == nil { + t.Fatal("newSignalReporter returned nil") + } + + // all-revoked batch: body omits "samples" and carries only "revoked". + r.post(nil, []revokedSample{{CredentialID: "c1", Reason: "revoked"}}, nil, nil) + c := <-got + if c.method != http.MethodPost { + t.Errorf("method = %q, want POST", c.method) + } + if c.auth != "Bearer tok-123" { + t.Errorf("Authorization = %q, want %q", c.auth, "Bearer tok-123") + } + if want := `{"revoked":[{"credential_id":"c1","reason":"revoked"}]}`; string(c.body) != want { + t.Fatalf("body = %s, want %s", c.body, want) + } + + // mixed batch: both arrays serialize. + r.post([]signalSample{{CredentialID: "c1", TS: 100, Util5h: 0.6}}, + []revokedSample{{CredentialID: "c2", Reason: "revoked"}}, nil, nil) + c = <-got + var decoded struct { + Samples []signalSample `json:"samples"` + Revoked []revokedSample `json:"revoked"` + } + if err := json.Unmarshal(c.body, &decoded); err != nil { + t.Fatalf("body not valid JSON: %v (raw %s)", err, c.body) + } + if len(decoded.Samples) != 1 || decoded.Samples[0] != (signalSample{CredentialID: "c1", TS: 100, Util5h: 0.6}) { + t.Fatalf("decoded samples = %+v, want one {c1,100,0.6}", decoded.Samples) + } + if len(decoded.Revoked) != 1 || decoded.Revoked[0] != (revokedSample{CredentialID: "c2", Reason: "revoked"}) { + t.Fatalf("decoded revoked = %+v, want one {c2,revoked}", decoded.Revoked) + } +} + +func TestIncrRateLimitNilAndEmpty(t *testing.T) { + // nil receiver guard: feature-off reporter must not panic. + var nilR *signalReporter + nilR.incrRateLimit("c1") // must not panic + + // nil rlCounts map → incrRateLimit must lazy-init, not panic; empty id drops. + r := &signalReporter{} + r.incrRateLimit("") // empty credentialID dropped (no lazy-init needed) + if len(r.snapshotRateLimits()) != 0 { + t.Fatal("empty credentialID should be dropped") + } + r.incrRateLimit("c1") // lazy-inits the map + if rl := r.snapshotRateLimits(); len(rl) != 1 || rl[0].Count != 1 { + t.Fatalf("after lazy init = %+v, want one {c1, count 1}", rl) + } +} + +func TestRateLimitCounting(t *testing.T) { + r := &signalReporter{rlCounts: make(map[string]int)} + // two 429s + one 403 for c1 → count 3; c2 is never seen → stays absent. + r.incrRateLimit("c1") + r.incrRateLimit("c1") + r.incrRateLimit("c1") + rl := r.snapshotRateLimits() + if len(rl) != 1 { + t.Fatalf("snapshot = %+v, want exactly one entry", rl) + } + if rl[0] != (rateLimitSample{CredentialID: "c1", Count: 3, WindowSecs: 30}) { + t.Fatalf("snapshot[0] = %+v, want {c1, 3, 30}", rl[0]) + } +} + +func TestRateLimitResetAfterFlush(t *testing.T) { + r := &signalReporter{rlCounts: make(map[string]int)} + r.incrRateLimit("c1") + r.incrRateLimit("c1") + if rl := r.snapshotRateLimits(); len(rl) != 1 || rl[0].Count != 2 { + t.Fatalf("first window = %+v, want c1 count 2", rl) + } + // next window starts empty (counter was reset on snapshot). + if rl := r.snapshotRateLimits(); rl != nil { + t.Fatalf("after flush snapshot = %+v, want nil (reset)", rl) + } + // new increments count up from 0, not from the prior window. + r.incrRateLimit("c1") + if rl := r.snapshotRateLimits(); len(rl) != 1 || rl[0].Count != 1 { + t.Fatalf("second window = %+v, want c1 count 1", rl) + } +} + +func TestRateLimitConcurrent(t *testing.T) { + // map+mutex concurrency smoke: many goroutines incrementing the same + // credential must total exactly N*per (run under -race to catch data races). + r := &signalReporter{rlCounts: make(map[string]int)} + const goroutines, per = 8, 100 + var wg sync.WaitGroup + for i := 0; i < goroutines; i++ { + wg.Add(1) + go func() { + defer wg.Done() + for j := 0; j < per; j++ { + r.incrRateLimit("c1") + } + }() + } + wg.Wait() + rl := r.snapshotRateLimits() + if len(rl) != 1 || rl[0].Count != goroutines*per { + t.Fatalf("concurrent count = %+v, want one {c1, %d}", rl, goroutines*per) + } +} + +func TestSignalPostSendsRateLimits(t *testing.T) { + got := make(chan []byte, 2) + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) { + b, _ := io.ReadAll(req.Body) + got <- b + })) + defer srv.Close() + + r := newSignalReporter(srv.URL, func(context.Context) (string, error) { return "tok-123", nil }, slog.Default()) + if r == nil { + t.Fatal("newSignalReporter returned nil") + } + + // rate-limits-only batch: body omits samples + revoked, exact wire contract. + r.post(nil, nil, []rateLimitSample{{CredentialID: "c1", Count: 3, WindowSecs: 30}}, nil) + if want := `{"rate_limits":[{"credential_id":"c1","count":3,"window_secs":30}]}`; string(<-got) != want { + t.Fatalf("rate-limits-only body mismatch, want %s", want) + } + + // mixed batch: samples + revoked + rate_limits all serialize. + r.post([]signalSample{{CredentialID: "c1", TS: 100, Util5h: 0.6}}, + []revokedSample{{CredentialID: "c2", Reason: "revoked"}}, + []rateLimitSample{{CredentialID: "c3", Count: 5, WindowSecs: 30}}, nil) + var decoded struct { + Samples []signalSample `json:"samples"` + Revoked []revokedSample `json:"revoked"` + RateLimits []rateLimitSample `json:"rate_limits"` + } + body := <-got + if err := json.Unmarshal(body, &decoded); err != nil { + t.Fatalf("body not valid JSON: %v (raw %s)", err, body) + } + if len(decoded.Samples) != 1 || decoded.Samples[0] != (signalSample{CredentialID: "c1", TS: 100, Util5h: 0.6}) { + t.Fatalf("decoded samples = %+v, want one {c1,100,0.6}", decoded.Samples) + } + if len(decoded.Revoked) != 1 || decoded.Revoked[0] != (revokedSample{CredentialID: "c2", Reason: "revoked"}) { + t.Fatalf("decoded revoked = %+v, want one {c2,revoked}", decoded.Revoked) + } + if len(decoded.RateLimits) != 1 || decoded.RateLimits[0] != (rateLimitSample{CredentialID: "c3", Count: 5, WindowSecs: 30}) { + t.Fatalf("decoded rate_limits = %+v, want one {c3,5,30}", decoded.RateLimits) + } +} + +func TestIsHardRevoked(t *testing.T) { + tests := []struct { + name string + status int + errType string + errMsg string + want bool + }{ + {"401_revoked_msg", 401, "authentication_error", "OAuth token has been revoked", true}, + {"401_revoked_in_type", 401, "token_revoked", "", true}, + {"401_plain_expiry", 401, "authentication_error", "token expired", false}, + {"429_revoked", 429, "", "revoked", false}, + {"200_ok", 200, "", "", false}, + // codex/openai hard-revocation shapes (sub2api-derived, R37 2026-07-04): + // errMsg is the raw body. token_invalidated has NO "revoked" substring, so + // the old matcher missed it; the {"detail":"Unauthorized"} form is + // non-standard (no error envelope at all). + {"401_codex_token_invalidated", 401, "", `{"error":{"code":"token_invalidated"}}`, true}, + {"401_codex_token_revoked_code", 401, "", `{"error":{"code":"token_revoked"}}`, true}, + {"401_codex_detail_unauthorized", 401, "", `{"detail":"Unauthorized"}`, true}, + // a plain codex 401 (transient / refresh-recoverable) must NOT quarantine. + {"401_codex_plain", 401, "", `{"error":{"message":"invalid token"}}`, false}, + {"401_generic_unauthorized_word", 401, "", "request was unauthorized, please retry", false}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := isHardRevoked(tt.status, tt.errType, tt.errMsg); got != tt.want { + t.Fatalf("isHardRevoked(%d,%q,%q) = %v, want %v", tt.status, tt.errType, tt.errMsg, got, tt.want) + } + }) + } +} + +func TestParseUnifiedUtil7d(t *testing.T) { + const hdr = "anthropic-ratelimit-unified-7d-utilization" + tests := []struct { + name string + set bool + val string + wantV float64 + wantOK bool + }{ + {"valid", true, "0.42", 0.42, true}, + {"missing", false, "", 0, false}, + {"malformed", true, "xyz", 0, false}, + {"above_one", true, "1.5", 0, false}, + {"negative", true, "-0.1", 0, false}, + {"zero_ok", true, "0", 0, true}, + {"one_ok", true, "1", 1, true}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + h := http.Header{} + if tt.set { + h.Set(hdr, tt.val) + } + v, ok := parseUnifiedUtil7d(h) + if ok != tt.wantOK { + t.Fatalf("ok = %v, want %v", ok, tt.wantOK) + } + if ok && v != tt.wantV { + t.Fatalf("v = %v, want %v", v, tt.wantV) + } + }) + } +} + +func TestSignalSampleUtil7dSerialization(t *testing.T) { + // both readings present → both serialize, util_7d after util_5h. + both, _ := json.Marshal(signalSample{CredentialID: "c1", TS: 100, Util5h: 0.6, Util7d: 0.4}) + if want := `{"credential_id":"c1","ts":100,"util_5h":0.6,"util_7d":0.4}`; string(both) != want { + t.Fatalf("both = %s, want %s", both, want) + } + // 5h-only (no 7d header) → util_7d omitempty drops it, byte-identical to the + // pre-7d wire format so existing master ingest is unaffected. + only, _ := json.Marshal(signalSample{CredentialID: "c1", TS: 100, Util5h: 0.6}) + if want := `{"credential_id":"c1","ts":100,"util_5h":0.6}`; string(only) != want { + t.Fatalf("5h-only = %s, want %s", only, want) + } +} + +func TestConcurrencyPeak(t *testing.T) { + r := &signalReporter{} // lazy-init path (no maps preallocated) + // start, start → peak 2; one ends → cur 1; start → cur 2 (peak stays 2). + d1 := r.trackInflight("c1") + d2 := r.trackInflight("c1") + d1() // one completes → cur 1 + d3 := r.trackInflight("c1") // back to cur 2 + snap := r.snapshotConcurrency() + if len(snap) != 1 || snap[0] != (concurrencySample{CredentialID: "c1", Peak: 2}) { + t.Fatalf("snapshot = %+v, want one {c1, peak 2}", snap) + } + // next window: peak map reset. d2/d3 are still in flight (cur 2) but with no + // NEW arrival the idle window reports nothing — peak only bumps on inc (the + // documented ponytail steady-state-trough ceiling). Pins that behavior. + if snap2 := r.snapshotConcurrency(); snap2 != nil { + t.Fatalf("idle next window = %+v, want nil (peak reset)", snap2) + } + d2() + d3() // cur back to 0 + // a fresh request in the new window peaks at 1 (cur started from 0). + d4 := r.trackInflight("c1") + if snap3 := r.snapshotConcurrency(); len(snap3) != 1 || snap3[0].Peak != 1 { + t.Fatalf("new window = %+v, want c1 peak 1", snap3) + } + d4() +} + +func TestConcurrencyPeakNilAndEmpty(t *testing.T) { + // nil receiver guard: feature-off reporter must not panic, returns a usable + // no-op so callers can `defer trackInflight(id)()` unconditionally. + var nilR *signalReporter + nilR.trackInflight("c1")() // must not panic + + // empty credentialID → no-op, nothing tracked. + r := &signalReporter{} + r.trackInflight("")() // must not panic, no map entry + if r.snapshotConcurrency() != nil { + t.Fatal("empty credentialID should track nothing") + } +} + +func TestConcurrencyPeakConcurrent(t *testing.T) { + // Deterministic peak under real concurrency: all goroutines hold their + // in-flight slot simultaneously, so the peak is exactly `goroutines`. Run + // under -race to catch a data race on the inflCur/inflPeak maps. + r := &signalReporter{} + const goroutines = 8 + release := make(chan struct{}) + var held, wg sync.WaitGroup + held.Add(goroutines) + wg.Add(goroutines) + for i := 0; i < goroutines; i++ { + go func() { + defer wg.Done() + done := r.trackInflight("c1") + held.Done() // signal "I'm in flight" + <-release // hold until every goroutine is concurrently in flight + done() + }() + } + held.Wait() // all goroutines have incremented → cur == peak == goroutines + snap := r.snapshotConcurrency() + close(release) + wg.Wait() + if len(snap) != 1 || snap[0] != (concurrencySample{CredentialID: "c1", Peak: goroutines}) { + t.Fatalf("concurrent peak = %+v, want one {c1, %d}", snap, goroutines) + } +} + +func TestSignalPostSendsConcurrency(t *testing.T) { + got := make(chan []byte, 2) + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) { + b, _ := io.ReadAll(req.Body) + got <- b + })) + defer srv.Close() + + r := newSignalReporter(srv.URL, func(context.Context) (string, error) { return "tok-123", nil }, slog.Default()) + if r == nil { + t.Fatal("newSignalReporter returned nil") + } + + // concurrency-only batch: body omits the other three, exact wire contract. + r.post(nil, nil, nil, []concurrencySample{{CredentialID: "c1", Peak: 2}}) + if want := `{"concurrency":[{"credential_id":"c1","peak":2}]}`; string(<-got) != want { + t.Fatalf("concurrency-only body mismatch, want %s", want) + } + + // mixed all-four batch: samples (with util_7d) + revoked + rate_limits + + // concurrency all serialize. + r.post( + []signalSample{{CredentialID: "c1", TS: 100, Util5h: 0.6, Util7d: 0.4}}, + []revokedSample{{CredentialID: "c2", Reason: "revoked"}}, + []rateLimitSample{{CredentialID: "c3", Count: 5, WindowSecs: 30}}, + []concurrencySample{{CredentialID: "c4", Peak: 2}}) + var decoded struct { + Samples []signalSample `json:"samples"` + Revoked []revokedSample `json:"revoked"` + RateLimits []rateLimitSample `json:"rate_limits"` + Concurrency []concurrencySample `json:"concurrency"` + } + body := <-got + if err := json.Unmarshal(body, &decoded); err != nil { + t.Fatalf("body not valid JSON: %v (raw %s)", err, body) + } + if len(decoded.Samples) != 1 || decoded.Samples[0] != (signalSample{CredentialID: "c1", TS: 100, Util5h: 0.6, Util7d: 0.4}) { + t.Fatalf("decoded samples = %+v, want one {c1,100,0.6,0.4}", decoded.Samples) + } + if len(decoded.Revoked) != 1 || decoded.Revoked[0] != (revokedSample{CredentialID: "c2", Reason: "revoked"}) { + t.Fatalf("decoded revoked = %+v, want one {c2,revoked}", decoded.Revoked) + } + if len(decoded.RateLimits) != 1 || decoded.RateLimits[0] != (rateLimitSample{CredentialID: "c3", Count: 5, WindowSecs: 30}) { + t.Fatalf("decoded rate_limits = %+v, want one {c3,5,30}", decoded.RateLimits) + } + if len(decoded.Concurrency) != 1 || decoded.Concurrency[0] != (concurrencySample{CredentialID: "c4", Peak: 2}) { + t.Fatalf("decoded concurrency = %+v, want one {c4,2}", decoded.Concurrency) + } +} + +// TestSignalReporterCloseIdempotentAndNilSafe covers the leak-fix Close contract: +// loop() now returns on <-r.stop (so the per-generation goroutine + ticker don't +// leak across reloads), and Close must be safe to call more than once +// (generation.close paths) and on a nil receiver (feature-off). The concrete risk +// is a double close(r.stop) panic — guarded by stopOnce; this test fails if that +// guard is removed. (loop()'s timing isn't deterministically testable — see the +// file header — so the stop case itself is verified by inspection.) +func TestSignalReporterCloseIdempotentAndNilSafe(t *testing.T) { + r := newSignalReporter("http://example.invalid", func(context.Context) (string, error) { return "tok", nil }, slog.Default()) + if r == nil { + t.Fatal("newSignalReporter returned nil") + } + if err := r.Close(); err != nil { + t.Fatalf("first Close: %v", err) + } + if err := r.Close(); err != nil { // idempotent — must not panic on double close + t.Fatalf("second Close: %v", err) + } + var nilR *signalReporter + if err := nilR.Close(); err != nil { // nil-safe — feature-off passes a nil reporter + t.Fatalf("nil Close: %v", err) + } +} diff --git a/internal/quota/local_usage.go b/internal/quota/local_usage.go index 81e8d34..164a766 100644 --- a/internal/quota/local_usage.go +++ b/internal/quota/local_usage.go @@ -15,6 +15,8 @@ package quota import ( "database/sql" "fmt" + + "github.com/AiKeyLabs/aikey-proxy/internal/vault" ) // LocalUsageRow is one persisted (subject, metric, period) increment. Mirrors the @@ -75,7 +77,11 @@ func WriteLocalUsage(dbPath string, rows []LocalUsageRow) error { if len(rows) == 0 { return nil } - db, err := sql.Open("sqlite", dbPath) + // Same vault DB file as vault.WriteConfig*/WriteGroupRuntime — must carry the + // busy_timeout guard (this is the busiest vault writer, 5s ticks), else it + // fails instantly on SQLITE_BUSY under lock contention on Linux. Single source + // of truth: vault.WithBusyTimeoutDSN (do not open the vault file raw). + db, err := sql.Open("sqlite", vault.WithBusyTimeoutDSN(dbPath)) if err != nil { return fmt.Errorf("open vault db for local-usage flush: %w", err) } diff --git a/internal/quota/policy_write.go b/internal/quota/policy_write.go index c45b338..60a0392 100644 --- a/internal/quota/policy_write.go +++ b/internal/quota/policy_write.go @@ -15,6 +15,8 @@ import ( "database/sql" "encoding/json" "fmt" + + "github.com/AiKeyLabs/aikey-proxy/internal/vault" ) // PolicySubject is one subject as delivered by GET /v1/quota/policy. It mirrors @@ -37,7 +39,10 @@ type PolicySubject struct { // A missing table (a vault that predates the Phase 2 schema) is tolerated (no-op) // so the poll never errors a pre-migration vault. func WriteSubjects(dbPath string, subjects []PolicySubject) error { - db, err := sql.Open("sqlite", dbPath) + // Same vault DB file — reuse the busy_timeout guard (single source of truth, + // vault.WithBusyTimeoutDSN) so this poll's full-replace write retries on lock + // contention instead of failing an admin's quota-limit edit with SQLITE_BUSY. + db, err := sql.Open("sqlite", vault.WithBusyTimeoutDSN(dbPath)) if err != nil { return fmt.Errorf("open vault db for quota policy write: %w", err) } diff --git a/internal/server/server.go b/internal/server/server.go index 208213a..558e60f 100644 --- a/internal/server/server.go +++ b/internal/server/server.go @@ -91,6 +91,12 @@ func New(ln net.Listener, dataHandler http.Handler, adminHandler *admin.Handler, mux.HandleFunc("GET /admin/audit/status", adminHandler.AuditStatus) mux.HandleFunc("POST /admin/audit/reconcile", adminHandler.AuditReconcile) + // Egress (upstream) proxy config — read + hot-swap. Relayed by the local web + // "Settings → Upstream proxy" card (R25 出口收敛: egress lives on the proxy node). + mux.HandleFunc("GET /admin/upstream-proxy", adminHandler.UpstreamProxyGet) + mux.HandleFunc("PUT /admin/upstream-proxy", adminHandler.UpstreamProxySet) + mux.HandleFunc("POST /admin/upstream-proxy/probe", adminHandler.UpstreamProxyProbe) + // Extra route registrars (e.g., OAuth broker handler) for _, h := range extraHandlers { h.RegisterRoutes(mux) diff --git a/internal/supervisor/compliance_policy.go b/internal/supervisor/compliance_policy.go index 9fee433..39391b9 100644 --- a/internal/supervisor/compliance_policy.go +++ b/internal/supervisor/compliance_policy.go @@ -25,6 +25,7 @@ import ( "os" "time" + "github.com/AiKeyLabs/aikey-proxy/internal/httpx" "github.com/AiKeyLabs/aikey-proxy/internal/vault" ) @@ -36,7 +37,7 @@ const ( compliancePollInterval = 60 * time.Second ) -var complianceHTTPClient = &http.Client{Timeout: 10 * time.Second} +var complianceHTTPClient = httpx.NewSwappableDirect(10 * time.Second) // resolveTeamOrgID returns the org this node's team mandates follow — BOTH the // compliance master policy (this file) AND the conversation-audit capture switch @@ -129,7 +130,7 @@ func fetchComplianceMasterPolicy(ctx context.Context, masterURL, orgID string) ( if err != nil { return false, false } - resp, err := complianceHTTPClient.Do(req) + resp, err := complianceHTTPClient.Get().Do(req) if err != nil { return false, false } diff --git a/internal/supervisor/conversation_audit_policy.go b/internal/supervisor/conversation_audit_policy.go index 32d5858..d55f02a 100644 --- a/internal/supervisor/conversation_audit_policy.go +++ b/internal/supervisor/conversation_audit_policy.go @@ -92,7 +92,7 @@ func fetchConversationAuditPolicy(ctx context.Context, masterURL, orgID string) if err != nil { return false, 0, false } - resp, err := complianceHTTPClient.Do(req) // reuse the package HTTP client (10s timeout) + resp, err := complianceHTTPClient.Get().Do(req) // reuse the package HTTP client (10s timeout) if err != nil { return false, 0, false } diff --git a/internal/supervisor/group_login_handler.go b/internal/supervisor/group_login_handler.go new file mode 100644 index 0000000..d0dc0fc --- /dev/null +++ b/internal/supervisor/group_login_handler.go @@ -0,0 +1,359 @@ +// group_login_handler.go — C10/RW8 (Path α, proxy-orchestrated): the pool login +// endpoints the local contribute page relays to. Unlike the personal broker +// (/oauth/*, vault-backed), the POOL flow exchanges with a MEMORY-store broker so +// the per-member token NEVER lands in the proxy's personal vault, then writes it +// back to master (RW10) over the team JWT — the token is never returned to the +// caller (no Path-β localhost token leak). +// +// Flow: +// +// POST /oauth/pool/authorize-url {provider, credential_id} +// → broker.StartLogin → {session_id, authorize_url}; remember session→credential +// POST /oauth/pool/submit-code {session_id, code} +// → broker.SubmitCode → read token from memory store +// → postMemberToken(master RW10, team JWT, {credential_id, token…}) +// → {status:"ok"} (no token in the response) +package supervisor + +import ( + "context" + "encoding/json" + "errors" + "log/slog" + "net/http" + "sync" + "time" + + broker "github.com/AiKeyLabs/aikey-auth-broker" +) + +// poolSessionTTL bounds how long a started-but-uncompleted pool login session is +// kept. It matches the broker's own login-session TTL (15 min); an authorize-url +// that's never followed by a successful submit-code is reaped on the next +// authorize-url so the session map can't grow unbounded (P2). +const poolSessionTTL = 15 * time.Minute + +// poolSession is the value stored per started login: which account it binds to + +// when it started (for TTL reaping). +type poolSession struct { + credentialID string + createdAt time.Time +} + +// errNoTeamCredential — no team account-JWT available (not logged in to a team). +var errNoTeamCredential = errors.New("supervisor: no team credential (run aikey login)") + +// poolExchanger is the broker surface the pool login needs, abstracted so the +// handler's session-tracking + writeback logic is unit-testable without a live +// provider. brokerPoolExchanger is the production impl (memory-store broker). +type poolExchanger interface { + // StartLogin begins an OAuth login and returns the session id + authorize URL. + StartLogin(ctx context.Context, provider string) (sessionID, authURL string, err error) + // SubmitCode exchanges the pasted code and returns the member's token material + // for writeback (access/refresh/expiry + the provider account UUID) plus the + // exchanged account's identity (email) for display / mismatch-warning. It is + // IDEMPOTENT per session: a second call for an already-exchanged session replays + // the cached token instead of re-spending the one-shot code (so a retry after a + // failed writeback doesn't force the member to re-login). + SubmitCode(ctx context.Context, sessionID, code string) (accountID, access, refresh string, expiresAt int64, externalID, identity string, err error) + // Forget clears the cached exchange result (token + broker session) once the + // token has been durably written back to master, so it doesn't linger in memory. + Forget(ctx context.Context, sessionID, accountID string) + // LoginStatus reports a login session's progress (pending/success/failed/ + // expired + provider error text). Codex's auth_code flow completes via the + // broker's localhost:1455 callback — the page can't paste a code, it POLLS + // this until the callback fires, then submit-code replays the cached exchange + // (empty code, idempotent). NO token material in the result. + LoginStatus(ctx context.Context, sessionID string) (status, errText string, err error) +} + +// brokerPoolExchanger wraps a memory-store EmbeddedBroker: SubmitCode reads the +// just-exchanged token back from the memory store (same pattern master uses), so +// nothing is persisted to the proxy vault. +type brokerPoolExchanger struct { + brk *broker.EmbeddedBroker + memTok broker.TokenStore + memAcc broker.AccountStore +} + +func (e *brokerPoolExchanger) StartLogin(ctx context.Context, provider string) (string, string, error) { + s, err := e.brk.StartLogin(ctx, provider) + if err != nil { + return "", "", err + } + return s.ID, s.AuthURL, nil +} + +func (e *brokerPoolExchanger) SubmitCode(ctx context.Context, sessionID, code string) (string, string, string, int64, string, string, error) { + r, err := e.brk.SubmitCode(ctx, sessionID, code) + if err != nil { + return "", "", "", 0, "", "", err + } + access, exp, err := e.memTok.GetAccessToken(ctx, r.AccountID) + if err != nil { + return "", "", "", 0, "", "", err + } + refresh, _ := e.memTok.GetRefreshToken(ctx, r.AccountID) // optional (Claude 1y has none) + externalID := "" + if acct, aerr := e.memAcc.GetByID(ctx, r.AccountID); aerr == nil && acct != nil { + externalID = acct.ExternalID + } + // DisplayIdentity is the exchanged account's email (ExtractIdentity.Email) — used + // for display + the team-account mismatch warning, never for the token itself. + return r.AccountID, access, refresh, exp, externalID, r.DisplayIdentity, nil +} + +// Forget removes the just-completed login's token from the memory store and its +// broker session, so the per-member token doesn't sit in proxy memory past the +// successful writeback. Memory-store deletes ignore ctx, but it's threaded through +// for interface symmetry. +func (e *brokerPoolExchanger) Forget(ctx context.Context, sessionID, accountID string) { + _ = e.memTok.DeleteTokens(ctx, accountID) + e.brk.ForgetSession(sessionID) +} + +func (e *brokerPoolExchanger) LoginStatus(ctx context.Context, sessionID string) (string, string, error) { + s, err := e.brk.GetLoginStatus(ctx, sessionID) + if err != nil { + return "", "", err + } + return s.Status, s.Error, nil +} + +// poolLoginHandler serves the pool login endpoints. bearer/masterURL are injected +// (the supervisor wires them from the team credential + control-panel URL) so the +// handler stays testable. +type poolLoginHandler struct { + ex poolExchanger + masterURL func() string + bearer func(ctx context.Context) (string, error) + client *http.Client + sessions sync.Map // sessionID → credentialID (the account being logged into) +} + +func (h *poolLoginHandler) authorizeURL(w http.ResponseWriter, r *http.Request) { + var req struct { + Provider string `json:"provider"` + CredentialID string `json:"credential_id"` + } + if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + poolErr(w, http.StatusBadRequest, "BAD_BODY", "invalid request body") + return + } + if req.CredentialID == "" { + poolErr(w, http.StatusBadRequest, "MISSING_CREDENTIAL_ID", "credential_id is required") + return + } + sid, authURL, err := h.ex.StartLogin(r.Context(), req.Provider) + if err != nil { + poolErr(w, http.StatusBadGateway, "LOGIN_START_FAILED", err.Error()) + return + } + // Reap expired sessions (P2: bound the map; no background goroutine needed) then + // bind this session to the account so submit-code writes the token to the right + // credential (the account the proxy told the member to log into). + h.sweepExpiredSessions() + h.sessions.Store(sid, poolSession{credentialID: req.CredentialID, createdAt: time.Now()}) + poolJSON(w, map[string]any{"session_id": sid, "authorize_url": authURL}) +} + +func (h *poolLoginHandler) submitCode(w http.ResponseWriter, r *http.Request) { + var req struct { + SessionID string `json:"session_id"` + Code string `json:"code"` + // Confirm gates the WRITEBACK. false (step 1) = exchange only + return the + // resolved account for review; true (step 2) = write the reviewed token back. + Confirm bool `json:"confirm"` + } + if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + poolErr(w, http.StatusBadRequest, "BAD_BODY", "invalid request body") + return + } + credAny, ok := h.sessions.Load(req.SessionID) + if !ok { + poolErr(w, http.StatusBadRequest, "UNKNOWN_SESSION", "no pool login session for that id") + return + } + sess, ok := credAny.(poolSession) + if !ok || time.Since(sess.createdAt) > poolSessionTTL { + h.sessions.Delete(req.SessionID) + poolErr(w, http.StatusBadRequest, "SESSION_EXPIRED", "pool login session expired; restart the login") + return + } + credentialID := sess.credentialID + + accountID, access, refresh, exp, externalID, identity, err := h.ex.SubmitCode(r.Context(), req.SessionID, req.Code) + if err != nil { + // Surface the provider exchange failure (CLAUDE.md 失败要显眼). err carries + // the wrapped upstream response (e.g. `[http_403] {provider body}`), which is + // the only place to see WHY the OAuth provider rejected the code. + slog.Warn("broker.pool.exchange_failed", + "error", err.Error(), "session_id", req.SessionID, "credential_id", credentialID) + poolErr(w, http.StatusBadGateway, "EXCHANGE_FAILED", err.Error()) + return + } + + // Step 1 (confirm=false): the code is exchanged and the resolved account returned + // for the member to review — but NOTHING is written to master yet. The token is + // held in the memory store (idempotent replay) and the session kept, so the + // follow-up confirm call replays it WITHOUT re-spending the one-shot code. If the + // member never confirms, the 15-min session TTL reaps the held token. + if !req.Confirm { + poolJSON(w, map[string]any{"status": "pending", "identity": identity}) + return + } + + // Step 2 (confirm=true): write the reviewed token back to master. + bearer, err := h.bearer(r.Context()) + if err != nil { + poolErr(w, http.StatusUnauthorized, "NO_TEAM_CREDENTIAL", "not logged in to the team (run aikey login)") + return + } + masterURL := h.masterURL() + if masterURL == "" { + poolErr(w, http.StatusServiceUnavailable, "NO_MASTER_URL", "control-panel URL not configured") + return + } + if err := postMemberToken(r.Context(), h.writebackClientFn(), masterURL, bearer, memberTokenWriteback{ + CredentialID: credentialID, AccessToken: access, RefreshToken: refresh, ExpiresAt: exp, ExternalID: externalID, + }); err != nil { + // DELIBERATELY keep the session (do NOT delete here): the code was already + // spent, so the member must be able to retry just the writeback. SubmitCode is + // idempotent per session — a re-submit replays the cached token instead of + // re-exchanging — so the page can re-POST the same code#state once master is + // reachable, and this token lands without a fresh login. The session is reaped + // by poolSessionTTL (15 min) if the member gives up. + slog.Warn("broker.pool.writeback_failed", + "error", err.Error(), "credential_id", credentialID) + poolErr(w, http.StatusBadGateway, "WRITEBACK_FAILED", err.Error()) + return + } + // Durably on master now: drop the proxy-side binding AND the broker's cached token + // (Forget) so the per-member token doesn't linger in proxy memory. One-shot done. + h.sessions.Delete(req.SessionID) + h.ex.Forget(r.Context(), req.SessionID, accountID) + // Return the exchanged account's identity (email) — NOT the token — so the page + // shows which Claude account was actually logged into and can warn (yellow) if it + // doesn't match the team account this slot expects. Email is not a secret (already + // shown as the account identity); the token still never leaves the proxy. + poolJSON(w, map[string]any{"status": "ok", "identity": identity}) +} + +// status serves GET /oauth/pool/status?session_id= — the codex (auth_code) +// pool login's polling leg. The browser can't paste a code for codex (OpenAI +// redirects to the broker's own localhost:1455 callback, which exchanges +// in-place), so the page polls here until the session flips to success, then +// calls submit-code with an EMPTY code — the broker's idempotent replay returns +// the cached exchange without re-spending anything. Only sessions this handler +// started are visible (fail-closed vs probing the personal broker's sessions); +// the response carries status + provider error text, never token material. +func (h *poolLoginHandler) status(w http.ResponseWriter, r *http.Request) { + sid := r.URL.Query().Get("session_id") + if sid == "" { + poolErr(w, http.StatusBadRequest, "MISSING_SESSION_ID", "session_id query param is required") + return + } + sessAny, ok := h.sessions.Load(sid) + if !ok { + poolErr(w, http.StatusBadRequest, "UNKNOWN_SESSION", "no pool login session for that id") + return + } + if sess, ok := sessAny.(poolSession); !ok || time.Since(sess.createdAt) > poolSessionTTL { + h.sessions.Delete(sid) + poolErr(w, http.StatusBadRequest, "SESSION_EXPIRED", "pool login session expired; restart the login") + return + } + status, errText, err := h.ex.LoginStatus(r.Context(), sid) + if err != nil { + poolErr(w, http.StatusBadGateway, "STATUS_FAILED", err.Error()) + return + } + resp := map[string]any{"status": status} + if errText != "" { + resp["error_detail"] = errText + } + poolJSON(w, resp) +} + +// RegisterRoutes mounts the pool login endpoints (server.RouteRegistrar). They +// live alongside the personal broker's /oauth/* routes on the proxy's local mux. +func (h *poolLoginHandler) RegisterRoutes(mux *http.ServeMux) { + mux.HandleFunc("POST /oauth/pool/authorize-url", h.authorizeURL) + mux.HandleFunc("POST /oauth/pool/submit-code", h.submitCode) + mux.HandleFunc("GET /oauth/pool/status", h.status) +} + +// teamBearer returns a fresh team account-JWT Bearer (same credential channel ③ +// uses to call master). Errors when not logged in to a team. +func (s *Supervisor) teamBearer(ctx context.Context) (string, error) { + gen := s.active.Load() + if gen == nil || gen.vault == nil { + return "", errNoTeamCredential + } + creds := buildCollectorCredentials(s.cfg.Events.CollectorCredentials, gen.vault) + cred := creds["team"] + if cred == nil { + return "", errNoTeamCredential + } + return cred.Bearer(ctx) +} + +// NewPoolLoginHandler builds the pool login handler with a MEMORY-store broker +// (per-member tokens never touch the proxy vault) + the team bearer / control-panel +// URL closures. The broker reuses the package-global impersonate HTTP client set +// during the personal broker init (Cloudflare bypass). Returns nil if the broker +// dependency isn't available. +func (s *Supervisor) NewPoolLoginHandler() *poolLoginHandler { + memTok := broker.NewMemoryTokenStore() + memAcc := broker.NewMemoryAccountStore() + brk := broker.NewEmbedded(memTok, memAcc) + return &poolLoginHandler{ + ex: &brokerPoolExchanger{brk: brk, memTok: memTok, memAcc: memAcc}, + masterURL: readControlPanelURL, + bearer: s.teamBearer, + // client left nil in production ⇒ writebackClientFn() resolves the LIVE + // swappable groupRuntimeClient each retry, so a network-change rebuild + // heals the writeback. Tests inject a fixed h.client (see writebackClientFn). + } +} + +// writebackClientFn returns the *http.Client provider postMemberToken re-resolves +// every retry. When h.client is set (tests inject a fixed httptest client) it +// wins; otherwise production uses the live swappable control-plane client so a +// mid-flight rebuild (selfheal) takes effect on the next attempt. +func (h *poolLoginHandler) writebackClientFn() func() *http.Client { + if h.client != nil { + c := h.client + return func() *http.Client { return c } + } + return groupRuntimeClient +} + +// sweepExpiredSessions deletes login sessions older than poolSessionTTL. Called +// opportunistically on authorize-url so abandoned/failed logins can't accumulate +// (P2) without a dedicated reaper goroutine. +func (h *poolLoginHandler) sweepExpiredSessions() { + cutoff := time.Now().Add(-poolSessionTTL) + h.sessions.Range(func(k, v any) bool { + if s, ok := v.(poolSession); ok && s.createdAt.Before(cutoff) { + h.sessions.Delete(k) + } + return true + }) +} + +func poolJSON(w http.ResponseWriter, v any) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _ = json.NewEncoder(w).Encode(v) +} + +func poolErr(w http.ResponseWriter, status int, code, msg string) { + w.Header().Set("Content-Type", "application/json") + // Mark as an aikey-generated error (mirrors proxy.HeaderAikeyErrorSource / + // writeJSONError); literal to avoid a supervisor→proxy import. The local web + // relay keys on this to distinguish aikey errors from upstream pass-throughs. + w.Header().Set("X-Aikey-Error-Source", code) + w.WriteHeader(status) + _ = json.NewEncoder(w).Encode(map[string]any{"error": map[string]string{"code": code, "message": msg}}) +} diff --git a/internal/supervisor/group_login_handler_test.go b/internal/supervisor/group_login_handler_test.go new file mode 100644 index 0000000..895a1f5 --- /dev/null +++ b/internal/supervisor/group_login_handler_test.go @@ -0,0 +1,351 @@ +package supervisor + +import ( + "context" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "strings" + "sync/atomic" + "testing" +) + +// fakePoolExchanger stands in for the memory-store broker so the handler's +// session-tracking + writeback chain is testable without a live provider. +type fakePoolExchanger struct { + authURL string + accountID string + access string + refresh string + expiresAt int64 + externalID string + identity string + submitErr error + submitN int // # of SubmitCode calls (idempotent-retry assertions) + forgotN int // # of Forget calls (cache-clear-on-success assertions) + forgotSess string // last Forget sessionID + forgotAcct string // last Forget accountID + status string // LoginStatus result (codex polling leg); "" ⇒ pending + statusErr string // LoginStatus provider error text +} + +func (f *fakePoolExchanger) StartLogin(_ context.Context, _ string) (string, string, error) { + return "sess-1", f.authURL, nil +} +func (f *fakePoolExchanger) SubmitCode(_ context.Context, _, _ string) (string, string, string, int64, string, string, error) { + f.submitN++ + if f.submitErr != nil { + return "", "", "", 0, "", "", f.submitErr + } + return f.accountID, f.access, f.refresh, f.expiresAt, f.externalID, f.identity, nil +} +func (f *fakePoolExchanger) Forget(_ context.Context, sessionID, accountID string) { + f.forgotN++ + f.forgotSess = sessionID + f.forgotAcct = accountID +} +func (f *fakePoolExchanger) LoginStatus(_ context.Context, _ string) (string, string, error) { + if f.status == "" { + return "pending", "", nil + } + return f.status, f.statusErr, nil +} + +func newPoolHandler(t *testing.T, ex poolExchanger, masterURL string) *poolLoginHandler { + t.Helper() + return &poolLoginHandler{ + ex: ex, + masterURL: func() string { return masterURL }, + bearer: func(context.Context) (string, error) { return "JWT", nil }, + client: http.DefaultClient, + } +} + +func doJSON(h http.HandlerFunc, body string) *httptest.ResponseRecorder { + r := httptest.NewRequest(http.MethodPost, "/", strings.NewReader(body)) + w := httptest.NewRecorder() + h(w, r) + return w +} + +// TestPoolLogin_EndToEnd: authorize-url binds session→credential; submit-code +// exchanges, then writes the token back to master RW10 with the bound credential_id +// — and the token is NEVER in the submit-code response. +func TestPoolLogin_EndToEnd(t *testing.T) { + var gotWB memberTokenWriteback + master := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/accounts/me/oauth-member-token" { + t.Errorf("unexpected master path %q", r.URL.Path) + } + b, _ := io.ReadAll(r.Body) + _ = json.Unmarshal(b, &gotWB) + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`{"status":"ok"}`)) + })) + defer master.Close() + + ex := &fakePoolExchanger{authURL: "https://login", accountID: "acc-x", access: "TOK", refresh: "RT", expiresAt: 42, externalID: "uuid-x", identity: "member@team.com"} + h := newPoolHandler(t, ex, master.URL) + h.client = master.Client() + + // 1) authorize-url for credential c1. + w1 := doJSON(h.authorizeURL, `{"provider":"claude","credential_id":"c1"}`) + if w1.Code != http.StatusOK { + t.Fatalf("authorize-url: %d %s", w1.Code, w1.Body.String()) + } + var sresp struct { + SessionID string `json:"session_id"` + AuthorizeURL string `json:"authorize_url"` + } + _ = json.Unmarshal(w1.Body.Bytes(), &sresp) + if sresp.SessionID != "sess-1" || sresp.AuthorizeURL != "https://login" { + t.Fatalf("authorize-url resp: %+v", sresp) + } + + // 2) submit-code with confirm → exchange + writeback. + w2 := doJSON(h.submitCode, `{"session_id":"sess-1","code":"abc#state","confirm":true}`) + if w2.Code != http.StatusOK { + t.Fatalf("submit-code: %d %s", w2.Code, w2.Body.String()) + } + // writeback carried the SESSION'S credential_id + the exchanged token. + if gotWB.CredentialID != "c1" || gotWB.AccessToken != "TOK" || gotWB.RefreshToken != "RT" || gotWB.ExpiresAt != 42 || gotWB.ExternalID != "uuid-x" { + t.Fatalf("writeback wrong: %+v", gotWB) + } + // token must NOT be echoed to the caller. + if strings.Contains(w2.Body.String(), "TOK") || strings.Contains(w2.Body.String(), "RT") { + t.Fatalf("token leaked into submit-code response: %s", w2.Body.String()) + } + // the exchanged account's identity (email) IS returned, for display + the + // team-account mismatch warning (email is not a secret). + var okResp struct { + Status string `json:"status"` + Identity string `json:"identity"` + } + _ = json.Unmarshal(w2.Body.Bytes(), &okResp) + if okResp.Identity != "member@team.com" { + t.Fatalf("submit-code response should carry identity email, got %q", okResp.Identity) + } +} + +// TestPoolLogin_PendingThenConfirm: step 1 (confirm=false) exchanges and returns the +// resolved account for review WITHOUT writing to master; step 2 (confirm=true) writes +// the reviewed token back. Guards the two-step confirm gate (2026-06-30). +func TestPoolLogin_PendingThenConfirm(t *testing.T) { + var writebacks int32 + master := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + atomic.AddInt32(&writebacks, 1) + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`{"status":"ok"}`)) + })) + defer master.Close() + + ex := &fakePoolExchanger{authURL: "u", accountID: "acc-x", access: "T", identity: "member@team.com"} + h := newPoolHandler(t, ex, master.URL) + h.client = master.Client() + + _ = doJSON(h.authorizeURL, `{"provider":"claude","credential_id":"c1"}`) + + // Step 1: no confirm → pending + identity, NO writeback, session kept. + w1 := doJSON(h.submitCode, `{"session_id":"sess-1","code":"abc#st"}`) + if w1.Code != http.StatusOK { + t.Fatalf("step 1 (no confirm) → 200 pending, got %d %s", w1.Code, w1.Body.String()) + } + var pending struct { + Status string `json:"status"` + Identity string `json:"identity"` + } + _ = json.Unmarshal(w1.Body.Bytes(), &pending) + if pending.Status != "pending" || pending.Identity != "member@team.com" { + t.Fatalf("step 1 should return pending + identity, got %+v", pending) + } + if n := atomic.LoadInt32(&writebacks); n != 0 { + t.Fatalf("step 1 must NOT write to master, got %d writebacks", n) + } + if ex.forgotN != 0 { + t.Fatalf("step 1 must NOT Forget the session (needed for confirm), got %d", ex.forgotN) + } + + // Step 2: confirm → writeback lands, session consumed. + w2 := doJSON(h.submitCode, `{"session_id":"sess-1","code":"abc#st","confirm":true}`) + if w2.Code != http.StatusOK { + t.Fatalf("step 2 (confirm) → 200 ok, got %d %s", w2.Code, w2.Body.String()) + } + if n := atomic.LoadInt32(&writebacks); n != 1 { + t.Fatalf("step 2 should write exactly once, got %d", n) + } + if ex.forgotN != 1 { + t.Fatalf("step 2 should Forget on success, got %d", ex.forgotN) + } +} + +// TestPoolLogin_RegisterRoutesMounts: RegisterRoutes actually mounts both pool +// endpoints on a real ServeMux and they reach the handler (a missing-field 400, +// not a 404) — verifies the wiring that main.go relies on without starting the +// full proxy. +func TestPoolLogin_RegisterRoutesMounts(t *testing.T) { + h := newPoolHandler(t, &fakePoolExchanger{authURL: "u"}, "http://unused") + mux := http.NewServeMux() + h.RegisterRoutes(mux) + + // authorize-url mounted → missing credential_id → 400 (proves reachable). + r1 := httptest.NewRequest(http.MethodPost, "/oauth/pool/authorize-url", strings.NewReader(`{"provider":"claude"}`)) + w1 := httptest.NewRecorder() + mux.ServeHTTP(w1, r1) + if w1.Code == http.StatusNotFound { + t.Fatal("/oauth/pool/authorize-url not mounted (404)") + } + if w1.Code != http.StatusBadRequest { + t.Fatalf("authorize-url mounted but wrong code: %d", w1.Code) + } + + // submit-code mounted → unknown session → 400 (proves reachable). + r2 := httptest.NewRequest(http.MethodPost, "/oauth/pool/submit-code", strings.NewReader(`{"session_id":"x","code":"y"}`)) + w2 := httptest.NewRecorder() + mux.ServeHTTP(w2, r2) + if w2.Code == http.StatusNotFound { + t.Fatal("/oauth/pool/submit-code not mounted (404)") + } +} + +// TestPoolLogin_UnknownSession: submit-code with a session that was never started +// (or already consumed) is rejected — no writeback attempted. +func TestPoolLogin_UnknownSession(t *testing.T) { + h := newPoolHandler(t, &fakePoolExchanger{}, "http://unused") + w := doJSON(h.submitCode, `{"session_id":"nope","code":"x"}`) + if w.Code != http.StatusBadRequest { + t.Fatalf("unknown session → 400, got %d", w.Code) + } +} + +// TestPoolLogin_MissingCredentialID: authorize-url requires credential_id (which +// account to bind the resulting token to). +func TestPoolLogin_MissingCredentialID(t *testing.T) { + h := newPoolHandler(t, &fakePoolExchanger{authURL: "u"}, "http://unused") + if w := doJSON(h.authorizeURL, `{"provider":"claude"}`); w.Code != http.StatusBadRequest { + t.Fatalf("missing credential_id → 400, got %d", w.Code) + } +} + +// TestPoolLogin_SessionConsumedOnce: a session can't be replayed after success. +func TestPoolLogin_SessionConsumedOnce(t *testing.T) { + master := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`{"status":"ok"}`)) + })) + defer master.Close() + ex := &fakePoolExchanger{authURL: "u", accountID: "a", access: "T"} + h := newPoolHandler(t, ex, master.URL) + h.client = master.Client() + + _ = doJSON(h.authorizeURL, `{"provider":"claude","credential_id":"c1"}`) + if w := doJSON(h.submitCode, `{"session_id":"sess-1","code":"x","confirm":true}`); w.Code != http.StatusOK { + t.Fatalf("first submit: %d", w.Code) + } + // replay → session gone → 400. + if w := doJSON(h.submitCode, `{"session_id":"sess-1","code":"x","confirm":true}`); w.Code != http.StatusBadRequest { + t.Fatalf("replay should be rejected, got %d", w.Code) + } + // success cleared the cached token via Forget (so it doesn't linger in memory). + if ex.forgotN != 1 || ex.forgotSess != "sess-1" || ex.forgotAcct != "a" { + t.Fatalf("Forget(sess-1,a) expected once on success, got n=%d sess=%q acct=%q", ex.forgotN, ex.forgotSess, ex.forgotAcct) + } +} + +// TestPoolLogin_WritebackFailureKeepsSessionForRetry: the OAuth code is spent at +// exchange, so a transient master outage during writeback must NOT waste it. The +// handler keeps the session on WRITEBACK_FAILED; the page can re-POST the same +// code#state and — because SubmitCode is idempotent per session — the cached token +// is replayed and lands once master recovers. Forget runs only on the successful +// writeback. 防退化 for the 2026-06-30 idempotent-retry design. +func TestPoolLogin_WritebackFailureKeepsSessionForRetry(t *testing.T) { + fastBackoff(t) // shrink the writeback retry backoff for the test + var hits int32 + master := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + // 503 for the whole first submit (all writebackMaxAttempts), then recover. + if atomic.AddInt32(&hits, 1) <= int32(writebackMaxAttempts) { + w.WriteHeader(http.StatusServiceUnavailable) + return + } + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`{"status":"ok"}`)) + })) + defer master.Close() + + ex := &fakePoolExchanger{authURL: "u", accountID: "acc-x", access: "TOK"} + h := newPoolHandler(t, ex, master.URL) + h.client = master.Client() + + _ = doJSON(h.authorizeURL, `{"provider":"claude","credential_id":"c1"}`) + + // 1) master down for every attempt → WRITEBACK_FAILED, session KEPT, no Forget. + if w := doJSON(h.submitCode, `{"session_id":"sess-1","code":"abc#st","confirm":true}`); w.Code != http.StatusBadGateway { + t.Fatalf("first submit (master down) → 502, got %d %s", w.Code, w.Body.String()) + } + if ex.forgotN != 0 { + t.Fatalf("Forget must NOT run on writeback failure (got %d)", ex.forgotN) + } + + // 2) retry same session+code → master recovered → writeback lands, then Forget. + if w := doJSON(h.submitCode, `{"session_id":"sess-1","code":"abc#st","confirm":true}`); w.Code != http.StatusOK { + t.Fatalf("retry (master back) → 200, got %d %s", w.Code, w.Body.String()) + } + if ex.submitN != 2 { + t.Fatalf("SubmitCode called once per submit; want 2, got %d", ex.submitN) + } + if ex.forgotN != 1 || ex.forgotSess != "sess-1" || ex.forgotAcct != "acc-x" { + t.Fatalf("Forget(sess-1,acc-x) expected once on success, got n=%d sess=%q acct=%q", ex.forgotN, ex.forgotSess, ex.forgotAcct) + } + + // 3) session consumed on success → a later replay is rejected. + if w := doJSON(h.submitCode, `{"session_id":"sess-1","code":"abc#st","confirm":true}`); w.Code != http.StatusBadRequest { + t.Fatalf("post-success replay → 400, got %d", w.Code) + } +} + +// TestPoolLogin_Status: the codex polling leg. Only sessions the pool handler +// started are visible (fail-closed: probing an unknown/personal-broker session id +// → 400), and the response carries status/error text but never token material. +func TestPoolLogin_Status(t *testing.T) { + ex := &fakePoolExchanger{authURL: "u", accountID: "acc-x", access: "SECRET-TOK"} + h := newPoolHandler(t, ex, "http://unused") + + _ = doJSON(h.authorizeURL, `{"provider":"codex","credential_id":"c1"}`) + + get := func(sid string) *httptest.ResponseRecorder { + r := httptest.NewRequest(http.MethodGet, "/oauth/pool/status?session_id="+sid, nil) + w := httptest.NewRecorder() + h.status(w, r) + return w + } + + // Unknown session → 400 (no probing the broker through this handler). + if w := get("not-ours"); w.Code != http.StatusBadRequest { + t.Fatalf("unknown session → 400, got %d %s", w.Code, w.Body.String()) + } + // Missing session_id → 400. + { + r := httptest.NewRequest(http.MethodGet, "/oauth/pool/status", nil) + w := httptest.NewRecorder() + h.status(w, r) + if w.Code != http.StatusBadRequest { + t.Fatalf("missing session_id → 400, got %d", w.Code) + } + } + // Pending → {"status":"pending"}. + if w := get("sess-1"); w.Code != http.StatusOK || !strings.Contains(w.Body.String(), `"pending"`) { + t.Fatalf("pending status expected, got %d %s", w.Code, w.Body.String()) + } + // Callback fired → success surfaces, provider error text passes through on + // failure, and NO token material ever appears in the body. + ex.status = "success" + if w := get("sess-1"); w.Code != http.StatusOK || !strings.Contains(w.Body.String(), `"success"`) { + t.Fatalf("success status expected, got %d %s", w.Code, w.Body.String()) + } else if strings.Contains(w.Body.String(), "SECRET-TOK") { + t.Fatalf("status body must never contain token material: %s", w.Body.String()) + } + ex.status, ex.statusErr = "failed", "provider said no" + if w := get("sess-1"); !strings.Contains(w.Body.String(), "provider said no") { + t.Fatalf("provider error text should pass through: %s", w.Body.String()) + } +} diff --git a/internal/supervisor/group_login_writeback.go b/internal/supervisor/group_login_writeback.go new file mode 100644 index 0000000..bad9e6e --- /dev/null +++ b/internal/supervisor/group_login_writeback.go @@ -0,0 +1,146 @@ +// group_login_writeback.go — RW8 (Path α, proxy-orchestrated): after the proxy's +// broker exchanges a member's per-account OAuth login, the resulting token is +// written BACK to master (RW10 POST /accounts/me/oauth-member-token) so it lands +// in oauth_member_token — NOT in the proxy's personal vault, and NEVER returned to +// any HTTP caller. The proxy is the only local component that holds the team +// account-JWT (RefreshableJWT) + reaches master, so it owns this writeback. +// +// SECURITY: the token travels proxy→master over TLS only. There is deliberately +// no path that hands the plaintext token back to the contribute page / browser +// (that would be the Path β localhost-token-leak we rejected). refresh_token is +// included only because master is its authoritative store — it is never delivered +// back down to any proxy/client (channel ③ omits it). +package supervisor + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "log/slog" + "net/http" + "time" + + "github.com/AiKeyLabs/aikey-proxy/internal/httpx" + "github.com/AiKeyLabs/aikey-proxy/internal/observability" +) + +// memberTokenWriteback is the RW10 request body (mirrors api.MemberTokenHandler). +type memberTokenWriteback struct { + CredentialID string `json:"credential_id"` + AccessToken string `json:"access_token"` + RefreshToken string `json:"refresh_token,omitempty"` + ExpiresAt int64 `json:"expires_at,omitempty"` + // ExternalID (optional, C5): the provider account UUID from the exchange, so + // master backfills it on first login (Claude metadata.user_id). + ExternalID string `json:"external_id,omitempty"` +} + +// writebackMaxAttempts / writebackBaseBackoff / writebackMaxBackoff bound the retry +// window: 6 attempts with 500ms base doubling capped at 3s wait ~0.5+1+2+3+3 = 9.5s +// of backoff (plus per-try request time). Short enough that the member isn't left +// hanging, long enough to ride out the post-recovery ARP transient + a quick master +// restart. +const ( + writebackMaxAttempts = 6 + // writebackMaxBackoff caps the per-retry wait so the backoff doesn't grow + // unbounded; over 6 attempts the waits are 0.5,1,2,3,3 ⇒ ~9.5s total. Sized to + // ride out a master that JUST came back: a route/ARP recovery has a ~1-2s window + // where a fresh Go dial still gets EHOSTUNREACH (curl retries past it; verified + // 2026-06-30 reproduction), plus a quick container/service restart. It does NOT + // cover a multi-minute VM redeploy — that needs the master up before login. + writebackMaxBackoff = 3 * time.Second +) + +// writebackBaseBackoff is a var (not const) only so tests can shrink it; production +// keeps 500ms (×2 each retry, capped at writebackMaxBackoff). +var writebackBaseBackoff = 500 * time.Millisecond + +// postMemberToken POSTs the freshly-exchanged per-member token to master's RW10 +// endpoint with the team account-JWT Bearer. clientFn is injected for testability +// and is RE-RESOLVED every attempt (production passes groupRuntimeClient, the +// live swappable accessor) so a client rebuilt mid-retry takes effect on the +// next try. +// +// RETRY (2026-06-30): the OAuth exchange already CONSUMED the one-shot auth code, +// so a TRANSIENT master blip (the VM restarting → "no route to host", or a 5xx +// while nginx is up but the backend isn't) must NOT waste it — we retry with +// exponential backoff. A 4xx is PERMANENT (bad request / auth / validation): +// retrying can't help, so we fail fast. The write is idempotent on master via the +// (credential_id, seat) PK, so re-POSTing the same token is safe. +// +// SELF-HEAL (2026-07-01): a routing-layer dial error (no-route/EHOSTUNREACH) can +// mean the long-lived client went stale after a host network change while master +// is actually reachable. We rebuild the shared client mid-retry so subsequent +// attempts (and later polls) dial clean instead of burning all 6 tries on a dead +// transport. The guarded self-restart backstop (selfheal.go) covers the case +// where even a rebuilt in-process client stays stuck. +func postMemberToken(ctx context.Context, clientFn func() *http.Client, masterURL, bearer string, wb memberTokenWriteback) error { + body, err := json.Marshal(wb) + if err != nil { + return err + } + backoff := writebackBaseBackoff + var lastErr error + for attempt := 1; attempt <= writebackMaxAttempts; attempt++ { + status, err := writeMemberTokenOnce(ctx, clientFn(), masterURL, bearer, body) + if err == nil { + return nil + } + lastErr = err + // Permanent client error (4xx) — retrying won't help; surface immediately. + if status >= 400 && status < 500 { + return err + } + // Network-change dial error → swap in a fresh control-plane client so the + // NEXT clientFn() call dials clean (WARN, not silent — logging-conventions). + if isNetChangeDialErr(err) { + slog.Warn("control-plane client rebuilt after network-change dial error", + "event.name", observability.EventProxyControlPlaneClientRebuilt, "caller", "member_token_writeback", + "credential_id", wb.CredentialID, "error", err.Error()) + httpx.RebuildAllControlPlane() + } + if attempt == writebackMaxAttempts { + break + } + slog.Warn("broker.pool.writeback_retry", + "attempt", attempt, "max", writebackMaxAttempts, "next_backoff", backoff.String(), + "credential_id", wb.CredentialID, "error", err.Error()) + select { + case <-ctx.Done(): + return fmt.Errorf("member-token writeback aborted (attempt %d/%d): %w; last error: %v", + attempt, writebackMaxAttempts, ctx.Err(), lastErr) + case <-time.After(backoff): + } + if backoff *= 2; backoff > writebackMaxBackoff { + backoff = writebackMaxBackoff + } + } + return fmt.Errorf("member-token writeback failed after %d attempts: %w", writebackMaxAttempts, lastErr) +} + +// writeMemberTokenOnce performs ONE writeback POST. Returns (statusCode, err): +// statusCode is 0 on a transport/connection error (err set, transient); on a +// non-2xx HTTP response statusCode is the code and err carries the body; on 2xx +// err is nil. +func writeMemberTokenOnce(ctx context.Context, client *http.Client, masterURL, bearer string, body []byte) (int, error) { + req, err := http.NewRequestWithContext(ctx, http.MethodPost, masterURL+"/accounts/me/oauth-member-token", bytes.NewReader(body)) + if err != nil { + return 0, err + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer "+bearer) + resp, err := client.Do(req) + if err != nil { + return 0, err // transport / connection error (e.g. "no route to host") → transient + } + defer resp.Body.Close() + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + // Read a bounded slice of the error body for the log/return (master sends a + // structured {"error":{code,message}} — surface it, don't swallow). + snippet, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<10)) + return resp.StatusCode, fmt.Errorf("master member-token writeback failed: %d: %s", resp.StatusCode, string(snippet)) + } + return resp.StatusCode, nil +} diff --git a/internal/supervisor/group_login_writeback_test.go b/internal/supervisor/group_login_writeback_test.go new file mode 100644 index 0000000..2947b2a --- /dev/null +++ b/internal/supervisor/group_login_writeback_test.go @@ -0,0 +1,132 @@ +package supervisor + +import ( + "context" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "sync/atomic" + "testing" + "time" +) + +// fastBackoff shrinks the retry backoff for tests + restores it after. +func fastBackoff(t *testing.T) { + t.Helper() + prev := writebackBaseBackoff + writebackBaseBackoff = time.Millisecond + t.Cleanup(func() { writebackBaseBackoff = prev }) +} + +// TestPostMemberToken_PostsToMasterWithBearer: the writeback POSTs the per-member +// token to master's RW10 endpoint with the account-JWT Bearer + JSON body, and the +// token is sent ONLY in the request body (never echoed anywhere). 2xx → nil error. +func TestPostMemberToken_PostsToMasterWithBearer(t *testing.T) { + var gotAuth, gotPath, gotCT string + var gotBody memberTokenWriteback + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotAuth = r.Header.Get("Authorization") + gotPath = r.URL.Path + gotCT = r.Header.Get("Content-Type") + b, _ := io.ReadAll(r.Body) + _ = json.Unmarshal(b, &gotBody) + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`{"status":"ok"}`)) + })) + defer srv.Close() + + wb := memberTokenWriteback{CredentialID: "c1", AccessToken: "tok", RefreshToken: "rt", ExpiresAt: 100, ExternalID: "uuid-1"} + if err := postMemberToken(context.Background(), func() *http.Client { return srv.Client() }, srv.URL, "JWT123", wb); err != nil { + t.Fatalf("postMemberToken: %v", err) + } + if gotPath != "/accounts/me/oauth-member-token" { + t.Errorf("path = %q, want /accounts/me/oauth-member-token", gotPath) + } + if gotAuth != "Bearer JWT123" { + t.Errorf("auth = %q, want Bearer JWT123", gotAuth) + } + if gotCT != "application/json" { + t.Errorf("content-type = %q", gotCT) + } + if gotBody.CredentialID != "c1" || gotBody.AccessToken != "tok" || gotBody.ExternalID != "uuid-1" { + t.Errorf("body mismatch: %+v", gotBody) + } +} + +// TestPostMemberToken_Non2xxSurfaces: a master error (e.g. 403 forbidden) is +// surfaced, not swallowed (the carry-back is auditable; the member can retry). +func TestPostMemberToken_Non2xxSurfaces(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusForbidden) + _, _ = w.Write([]byte(`{"error":{"code":"BIZ_OAUTH_MEMBER_TOKEN_FORBIDDEN"}}`)) + })) + defer srv.Close() + + err := postMemberToken(context.Background(), func() *http.Client { return srv.Client() }, srv.URL, "JWT", memberTokenWriteback{CredentialID: "c1", AccessToken: "t"}) + if err == nil { + t.Fatal("non-2xx master response must surface as an error") + } +} + +// TestPostMemberToken_RetriesTransient5xxThenSucceeds: a TRANSIENT master failure +// (5xx while the VM restarts / nginx backend flaps) is retried; once the master +// recovers the writeback lands. The OAuth code was already consumed, so this is the +// whole point — a blip must not waste it. 防退化 for the 2026-06-30 retry. +func TestPostMemberToken_RetriesTransient5xxThenSucceeds(t *testing.T) { + fastBackoff(t) + var hits int32 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if atomic.AddInt32(&hits, 1) < 3 { + w.WriteHeader(http.StatusServiceUnavailable) // 503 → transient → retry + return + } + w.WriteHeader(http.StatusOK) + })) + defer srv.Close() + + if err := postMemberToken(context.Background(), func() *http.Client { return srv.Client() }, srv.URL, "JWT", memberTokenWriteback{CredentialID: "c1", AccessToken: "t"}); err != nil { + t.Fatalf("expected success after transient 5xx, got: %v", err) + } + if got := atomic.LoadInt32(&hits); got != 3 { + t.Errorf("server hit %d times, want 3 (2 transient + 1 success)", got) + } +} + +// TestPostMemberToken_4xxFailsFastNoRetry: a 4xx is PERMANENT — retrying can't help, +// so we must fail on the FIRST attempt (no wasted retries / latency). +func TestPostMemberToken_4xxFailsFastNoRetry(t *testing.T) { + fastBackoff(t) + var hits int32 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + atomic.AddInt32(&hits, 1) + w.WriteHeader(http.StatusBadRequest) // 400 → permanent + })) + defer srv.Close() + + if err := postMemberToken(context.Background(), func() *http.Client { return srv.Client() }, srv.URL, "JWT", memberTokenWriteback{CredentialID: "c1", AccessToken: "t"}); err == nil { + t.Fatal("4xx must surface as an error") + } + if got := atomic.LoadInt32(&hits); got != 1 { + t.Errorf("server hit %d times, want 1 (4xx must NOT retry)", got) + } +} + +// TestPostMemberToken_ExhaustsRetriesOnPersistent5xx: a master that stays down +// exhausts all attempts and returns an error (bounded, doesn't hang forever). +func TestPostMemberToken_ExhaustsRetriesOnPersistent5xx(t *testing.T) { + fastBackoff(t) + var hits int32 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + atomic.AddInt32(&hits, 1) + w.WriteHeader(http.StatusBadGateway) // 502 → transient, but never recovers + })) + defer srv.Close() + + if err := postMemberToken(context.Background(), func() *http.Client { return srv.Client() }, srv.URL, "JWT", memberTokenWriteback{CredentialID: "c1", AccessToken: "t"}); err == nil { + t.Fatal("persistent 5xx must surface as an error after exhausting retries") + } + if got := atomic.LoadInt32(&hits); got != writebackMaxAttempts { + t.Errorf("server hit %d times, want %d (all attempts)", got, writebackMaxAttempts) + } +} diff --git a/internal/supervisor/group_runtime_policy.go b/internal/supervisor/group_runtime_policy.go new file mode 100644 index 0000000..732cb49 --- /dev/null +++ b/internal/supervisor/group_runtime_policy.go @@ -0,0 +1,503 @@ +// group_runtime_policy.go — N7c-2: proxy-side channel-③ group-runtime follower. +// +// The always-on proxy pulls the account-level group-runtime endpoint from master +// (GET /accounts/me/group-runtime, account-JWT), encrypts each candidate +// account's token/key with the vault key, and writes it into the group VK's +// managed_virtual_keys_cache.group_runtime column. The route resolver (N8) reads +// it at request time to pick + inject an account. +// +// WHY proxy (not CLI): access_token refresh is high-frequency (Kimi 15min) and +// the CLI isn't always running — same rail compliance/quota already use +// (quota_policy.go). System design: update/20260625-通道3-组物料下发-proxy拉取系统设计.md. +// +// SECURITY: refresh_token is NEVER fetched or stored (the master response has no +// such field). access_token/key arrive plaintext over TLS and are re-encrypted +// with the vault derivedKey before they touch disk — same at-rest protection as +// provider_account_tokens. +package supervisor + +import ( + "context" + "encoding/base64" + "encoding/json" + "fmt" + "io" + "log/slog" + "net/http" + "time" + + "github.com/AiKeyLabs/aikey-proxy/internal/httpx" + "github.com/AiKeyLabs/aikey-proxy/internal/observability" + "github.com/AiKeyLabs/aikey-proxy/internal/vault" + "github.com/AiKeyLabs/aikey-proxy/internal/vkeys" + "github.com/AiKeyLabs/pkg/seatassign" +) + +func defaultGroupRuntimeClient() *http.Client { return httpx.NewDirectClient(10 * time.Second) } + +const groupRuntimePollInterval = 60 * time.Second + +// groupRuntimeRail declares this rail for the SyncRail framework (railset.go, +// 2026-07-03): gate/generation/URL/credential re-evaluated every cycle, failures +// counted into the visibility state machine. The old hand-written loop built its +// team credential once at start and early-returned forever on any startup +// precondition miss — silently (bugfix +// 2026-07-03-routing-override-rail-silent-stall.md). +func (s *Supervisor) groupRuntimeRail() railSpec { + return railSpec{ + name: "group_runtime", + interval: groupRuntimePollInterval, + needsTeamJWT: true, + gate: func(gen *generation) bool { + if !oauthGroupRoutingEnabled() { + return false // feature off → the whole rail is bypassed (direct-bind unchanged) + } + mks, _ := gen.vault.GetActiveManagedKeys() + for i := range mks { + if mks[i].OauthGroupID != "" { + return true + } + } + return false // no local group VK → nothing to pull/store (idle, not broken) + }, + sync: s.syncGroupRuntime, + } +} + +// syncGroupRuntime runs one pull of the account-level group runtime and, only +// when the material actually changed, rewrites each group VK's group_runtime +// column and reloads so the resolver (N8) picks up the fresh tokens. Any error +// keeps the last-known material (don't flap) and is COUNTED by the framework. +func (s *Supervisor) syncGroupRuntime(ctx context.Context, gen *generation, masterURL, bearer string) error { + mks, err := gen.vault.GetActiveManagedKeys() + if err != nil { + return err + } + // Path Z (通道3 §14): piggyback the proxy's observed window-reset epochs so + // master re-rolls each account's window_max_util_pct per window. + var observedResets map[string]int64 + if gen.proxy != nil { + observedResets = gen.proxy.ObservedResetsSnapshot() + } + groups, sig, err := fetchGroupRuntime(ctx, masterURL, bearer, observedResets) + if err != nil { + return err // unreachable / bad response → keep last-known (don't flap) + } + // Steady state: same material (no token refresh) → no rewrite/reload. + if prev := s.lastGroupRuntimeSig.Load(); prev != nil && *prev == sig { + return nil + } + // s.routingOverrides.Assignment is nil-safe (guards nil receiver → "" = rank-0), + // so the routed-account stamp degrades to the local pick when overrides are unset. + // Same cooldown view the hot path uses, so is_current_routed reflects cooling failover. + var grSkip map[string]bool + if gen.proxy != nil { + grSkip = gen.proxy.CooldownSkipSet() + } + if err := writeGroupRuntimeForGroups(s.cfg.Vault.Path, gen.vault.DerivedKey(), mks, groups, s.routingOverrides.Assignment, grSkip); err != nil { + slog.Warn("group_runtime write failed", + "event.name", "proxy.group_runtime.write_failed", "error", err.Error()) + return err // leave lastGroupRuntimeSig unchanged → retry next tick + } + s.lastGroupRuntimeSig.Store(&sig) + slog.Info("group runtime material changed", + "event.name", "proxy.group_runtime.changed", "groups", len(groups)) + if err := s.Reload(ctx); err != nil { + // Local reload hiccup, not a master-sync failure: the material IS stored + // (next reload picks it up), so the cycle still counts as a success. + slog.Warn("group_runtime reload failed", + "event.name", "proxy.group_runtime.reload_failed", "error", err.Error()) + } + return nil +} + +// ── master response (mirrors groupruntime.GroupDelivery / AccountMaterial) ── + +type grDeliveryResp struct { + Groups []grGroup `json:"groups"` +} + +type grGroup struct { + OauthGroupID string `json:"oauth_group_id"` + RoutingConfig string `json:"routing_config"` + Accounts []grAccount `json:"accounts"` +} + +// grAccount is one candidate account's PLAINTEXT material from master (TLS-only). +// NOTE: there is intentionally no refresh_token field — master never sends it. +type grAccount struct { + AccountID string `json:"account_id"` + CredentialID string `json:"credential_id"` + CredentialType string `json:"credential_type"` // oauth_account | api_key + // Display meta (non-secret, 2026-07-01) so the client's candidate LIST membership + // refreshes off this fast rail (a fast-rail-only account renders with its + // email/provider, not a bare UUID). Threaded straight through to group_runtime. + Identity string `json:"identity"` + ProviderCode string `json:"provider_code"` + Priority int `json:"priority"` + // NeedsLogin: master delivered this OAuth account as "member not logged in" + // (no token) — the proxy returns LOGIN_REQUIRED for it (vs an absent account = + // material not pulled yet → retryable skip). P1. + NeedsLogin bool `json:"needs_login"` + // OAuth-only: + AccessToken string `json:"access_token"` + ExpiresAt int64 `json:"expires_at"` + WindowMaxUtilPct *int `json:"window_max_util_pct"` + WindowStatus string `json:"window_status"` + WindowResetAt *int64 `json:"window_reset_at"` + // ExternalID is the OAuth provider's account UUID. Required for Claude's + // metadata.user_id injection (N8); empty for non-Claude. Master's N7a + // producer must populate it for Claude OAuth group accounts — until then it + // arrives empty and Claude OAuth group routing is degraded (see N8 finding). + ExternalID string `json:"external_id"` + // KEY-only: + Key string `json:"key"` + BaseURL string `json:"base_url"` + Revision string `json:"revision"` +} + +// The stored group_runtime contract (map[account_id]vkeys.GroupRuntimeAccount) +// lives in package vkeys so the writer here and the reader in package proxy (N8) +// share one definition. The secret (access_token for OAuth / key for KEY) is +// AES-GCM encrypted with the vault key; nonce + ciphertext are base64 in the +// JSON. N8 base64-decodes + vault.Decrypt. + +// groupRuntimeSwap is the group-runtime + member-token-writeback control-plane +// client, registered in the CENTRAL self-heal registry (internal/httpx). A host +// network change stalls long-lived direct clients with no-route errors (see +// selfheal.go); httpx.RebuildAllControlPlane() then installs a FRESH client (new +// dialer + empty pool) for THIS rail and every other control-plane rail at once. +// Both the poll (fetchGroupRuntime) and the writeback read the LIVE client via +// groupRuntimeClient(). WHY a whole new client rather than CloseIdleConnections(): +// closing idle conns alone does NOT reliably clear the stale post-network-change +// state (golang/go#23427); a fresh Transport is the reliable reset. +var groupRuntimeSwap = httpx.NewSwappable(defaultGroupRuntimeClient) + +// groupRuntimeClient returns the live (swappable) control-plane HTTP client. +func groupRuntimeClient() *http.Client { return groupRuntimeSwap.Get() } + +// fetchGroupRuntime GETs the account-level group-runtime endpoint with an +// account-JWT Bearer. Returns (groups, rawBody, err); a non-nil error means the +// caller keeps the last-known group_runtime (don't flap) and the SyncRail +// framework surfaces the cause (2026-07-03: a bare ok=false hid the +// connection-refused detail that mattered most in the incident). rawBody is the +// change signature — same plaintext response = no token change = skip the vault +// rewrite (the encrypted form can't be compared: a fresh nonce each encrypt). +// observedResetsHeader piggybacks the proxy's observed per-account window-reset +// epochs on the pull (Path Z, 通道3 §14): base64(JSON {account_id: epoch}). +// master re-rolls window_max_util_pct when an epoch is newer than its stored +// window_reset_at. Optional — master ignores it when absent (backward compatible). +const observedResetsHeader = "X-Aikey-Observed-Resets" + +func fetchGroupRuntime(ctx context.Context, masterURL, bearer string, observedResets map[string]int64) ([]grGroup, string, error) { + req, err := http.NewRequestWithContext(ctx, http.MethodGet, masterURL+"/accounts/me/group-runtime", http.NoBody) + if err != nil { + return nil, "", err + } + req.Header.Set("Authorization", "Bearer "+bearer) + if len(observedResets) > 0 { + if b, mErr := json.Marshal(observedResets); mErr == nil { + req.Header.Set(observedResetsHeader, base64.StdEncoding.EncodeToString(b)) + } + } + resp, err := groupRuntimeClient().Do(req) + if err != nil { + // Self-heal (2026-07-01): a host network change can stall the long-lived + // client with a routing-layer error while master is perfectly reachable + // from a fresh process. The poll is the steady 60s heartbeat that drives + // recovery: Tier1 rebuilds the client, and after `threshold` consecutive + // stuck cycles (with master confirmed reachable) escalates to a guarded + // self-restart. Non-net-change errors (master 5xx, timeout) don't touch + // the healer — they aren't a routing-stale symptom. + if isNetChangeDialErr(err) { + if d := controlPlaneHeal.onPollNetChange(masterURL); d == restartSkipBreaker { + slog.Error("control-plane stuck reaching master after network change; restart budget exhausted — manual `aikey proxy restart` may be needed", + "event.name", observability.EventProxyControlPlaneRestartExhausted, "error", err.Error()) + } else { + slog.Warn("control-plane client rebuilt after network-change dial error", + "event.name", observability.EventProxyControlPlaneClientRebuilt, "caller", "group_runtime_poll", + "error", err.Error()) + } + } + return nil, "", err + } + // Reached master (any HTTP response proves the routing path is alive) → clear + // the net-change stuck counter so a later transient can't ride an old tally. + controlPlaneHeal.onPollOK() + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + return nil, "", fmt.Errorf("GET /accounts/me/group-runtime: HTTP %d", resp.StatusCode) + } + body, err := io.ReadAll(io.LimitReader(resp.Body, 8<<20)) + if err != nil { + return nil, "", err + } + var out grDeliveryResp + if err := json.Unmarshal(body, &out); err != nil { + return nil, "", fmt.Errorf("decode group-runtime response: %w", err) + } + return out.Groups, string(body), nil +} + +// restampCurrentRouted recomputes IsCurrentRouted on every local group VK's EXISTING +// group_runtime after a routing-override change — WITHOUT refetching material from +// master or re-encrypting secrets (it reads the already-loaded mk.GroupRuntime and +// flips only the plaintext display flag). This couples the routing-override rail to +// the group_runtime display column (owner-approved 2026-06-30) so /user/vault reflects +// an engine redirect within one override poll, not only on the next material refresh. +// +// NO Reload: IsCurrentRouted is display-only (the hot-path resolver never reads it), +// and the CLI/web read the vault column directly — so writing the column suffices. +// The actual routing redirect already took effect via the RoutingOverrideCache the +// resolver reads at request time; this only keeps the DISPLAY in step. +func (s *Supervisor) restampCurrentRouted() { + gen := s.active.Load() + if gen == nil || gen.vault == nil { + return + } + mks, err := gen.vault.GetActiveManagedKeys() + if err != nil { + return + } + // Same cooldown view the hot-path resolver uses, so the stamp names the account the + // proxy actually forwards to (cooling-driven failover included). + var skip map[string]bool + if gen.proxy != nil { + skip = gen.proxy.CooldownSkipSet() + } + nowUnix := time.Now().Unix() + for i := range mks { + mk := mks[i] + if mk.OauthGroupID == "" || mk.GroupRuntime == "" { + continue + } + // Feed the STORED material to the shared pick so the re-stamp applies the same + // usability gates as the hot path (unparseable → nil = blind mode, nominal pick). + var material map[string]vkeys.GroupRuntimeAccount + _ = json.Unmarshal([]byte(mk.GroupRuntime), &material) + newJSON, changed, err := stampCurrentRoutedJSON(mk.GroupRuntime, computeRoutedAccountID(mk, material, s.routingOverrides.Assignment, skip, nowUnix)) + if err != nil { + slog.Warn("group_runtime restamp parse failed", + "event.name", "proxy.group_runtime.restamp_failed", + "virtual_key_id", mk.VirtualKeyID, "error", err.Error()) + continue + } + if !changed { + continue + } + if err := vault.WriteGroupRuntime(s.cfg.Vault.Path, mk.VirtualKeyID, newJSON); err != nil { + slog.Warn("group_runtime restamp write failed", + "event.name", "proxy.group_runtime.restamp_failed", + "virtual_key_id", mk.VirtualKeyID, "error", err.Error()) + } + } +} + +// buildGroupRuntimeMap encrypts each account's secret (access_token | key) with the +// vault key and returns the per-account material map for one group. It is +// ROUTED-AGNOSTIC — IsCurrentRouted (C2 display) is stamped per-VK later by +// marshalGroupRuntime, because "which account is routed" is per-seat and the same +// group's VKs belong to different seats. An account whose encryption fails is skipped +// (best-effort — one bad account must not blank the whole group). +func buildGroupRuntimeMap(derivedKey []byte, accounts []grAccount) map[string]vkeys.GroupRuntimeAccount { + out := make(map[string]vkeys.GroupRuntimeAccount, len(accounts)) + for _, a := range accounts { + // needs_login marker carries NO secret — store it as-is so the resolver can + // return LOGIN_REQUIRED for it (P1), distinct from an absent account. + if a.NeedsLogin { + out[a.AccountID] = vkeys.GroupRuntimeAccount{ + CredentialType: a.CredentialType, NeedsLogin: true, + Identity: a.Identity, ProviderCode: a.ProviderCode, Priority: a.Priority, + } + continue + } + secret := a.AccessToken + if a.CredentialType == "api_key" { + secret = a.Key + } + nonce, ct, err := vault.Encrypt(derivedKey, []byte(secret)) + if err != nil { + continue // skip this account; keep the rest + } + gra := vkeys.GroupRuntimeAccount{ + CredentialType: a.CredentialType, + SecretNonce: base64.StdEncoding.EncodeToString(nonce), + SecretCiphertext: base64.StdEncoding.EncodeToString(ct), + Identity: a.Identity, // non-secret display meta → client list refresh + ProviderCode: a.ProviderCode, // + Priority: a.Priority, // + } + if a.CredentialType == "api_key" { + gra.BaseURL = a.BaseURL + gra.Revision = a.Revision + } else { + gra.ExpiresAt = a.ExpiresAt + gra.WindowMaxUtilPct = a.WindowMaxUtilPct + gra.WindowStatus = a.WindowStatus + gra.WindowResetAt = a.WindowResetAt + gra.ExternalID = a.ExternalID + } + out[a.AccountID] = gra + } + return out +} + +// marshalGroupRuntime renders the material map, stamping IsCurrentRouted=true on the +// single routedAccountID (C2 display). routedAccountID "" (or absent from the map) → +// no account is flagged. The input map is NOT mutated (entries are copied by value). +func marshalGroupRuntime(base map[string]vkeys.GroupRuntimeAccount, routedAccountID string) (string, error) { + out := make(map[string]vkeys.GroupRuntimeAccount, len(base)) + for id, acc := range base { + acc.IsCurrentRouted = id == routedAccountID + out[id] = acc + } + b, err := json.Marshal(out) + if err != nil { + return "", fmt.Errorf("marshal group_runtime: %w", err) + } + return string(b), nil +} + +// buildGroupRuntimeJSON renders the routed-agnostic group_runtime JSON for one group +// (no IsCurrentRouted flag). Retained for callers/tests that don't need per-seat +// routing display. +func buildGroupRuntimeJSON(derivedKey []byte, accounts []grAccount) (string, error) { + return marshalGroupRuntime(buildGroupRuntimeMap(derivedKey, accounts), "") +} + +// computeRoutedAccountID returns the account this seat's traffic is routed to (C2 +// display, the is_current_routed / current_routed stamp). +// +// SINGLE SOURCE OF TRUTH (2026-07-01 unification): the pick itself is +// vkeys.PickRoutedAccount — the EXACT function the hot-path resolver +// (proxy.resolveGroupCredential) forwards with, fed the same override view (the shared +// RoutingOverrideCache via overrideFor), the same cooldown skip view +// (proxy.CooldownSkipSet), and the same material usability gates. Display == forwarding +// by construction; do NOT re-derive routing here. A needs_login pick IS the routed +// account (owner rule: the engine may route a member to an account they haven't logged +// into — the stamp shows it so the member logs into THAT one). +// +// material is the group's CURRENT runtime material map (pass the freshly built map when +// writing, or the parsed stored column when re-stamping); empty/nil → the picker's blind +// mode (rank/override-only — pre-poll display still shows a nominal pick). Residual +// (accepted): the hot path also skips undecryptable secrets per request (decrypt needs +// the vault key, can't be in the pure pick), and the stamp persists at the 60s poll + +// override-change cadence, so mid-cycle changes converge within ≤60s. skip may be nil. +// "" when the candidate list is absent/unparseable. overrideFor may be nil (→ rank-0). +func computeRoutedAccountID(mk vault.ManagedKey, material map[string]vkeys.GroupRuntimeAccount, overrideFor func(seatID, groupID string) string, skip map[string]bool, nowUnix int64) string { + var refs []vkeys.GroupAccountRef + if mk.GroupAccounts == "" || json.Unmarshal([]byte(mk.GroupAccounts), &refs) != nil || len(refs) == 0 { + return "" + } + override := "" + if overrideFor != nil { + override = overrideFor(mk.SeatID, mk.OauthGroupID) + } + if acc, oc := vkeys.PickRoutedAccount(mk.SeatID, refs, material, override, skip, nowUnix); oc != vkeys.PickNone { + return acc + } + // Nothing pickable (hot path would 429/ALL_UNUSABLE) → show the nominal seatassign + // rank-0 rather than blank (display-only fallback; no traffic is served here anyway). + accounts := make([]seatassign.Account, 0, len(refs)) + for _, r := range refs { + accounts = append(accounts, seatassign.Account{AccountID: r.AccountID, Priority: r.Priority}) + } + if ordered := seatassign.Rank(mk.SeatID, accounts); len(ordered) > 0 { + return ordered[0].AccountID + } + return "" +} + +// stampCurrentRoutedJSON rewrites an EXISTING group_runtime JSON so ONLY +// routedAccountID carries IsCurrentRouted (C2 re-stamp on a routing-override change — +// no master fetch, no re-encryption, secrets untouched). Returns (json, changed, err); +// changed=false when nothing moved (caller skips the write). Unparseable input is +// returned unchanged with the error so a corrupt column can't crash the poll. +func stampCurrentRoutedJSON(runtimeJSON, routedAccountID string) (string, bool, error) { + if runtimeJSON == "" || runtimeJSON == "{}" { + return runtimeJSON, false, nil + } + var m map[string]vkeys.GroupRuntimeAccount + if err := json.Unmarshal([]byte(runtimeJSON), &m); err != nil { + return runtimeJSON, false, err + } + changed := false + for id, acc := range m { + want := id == routedAccountID + if acc.IsCurrentRouted != want { + acc.IsCurrentRouted = want + m[id] = acc + changed = true + } + } + if !changed { + return runtimeJSON, false, nil + } + b, err := json.Marshal(m) + if err != nil { + return runtimeJSON, false, fmt.Errorf("marshal group_runtime: %w", err) + } + return string(b), true, nil +} + +// writeGroupRuntimeForGroups writes the encrypted material into every group VK's +// group_runtime column, stamping IsCurrentRouted PER-VK (per-seat) via +// computeRoutedAccountID — the same group's VKs belong to different seats and can +// route to different accounts. A VK belongs to a group when its OauthGroupID matches; +// a group with no local VK is simply skipped (the proxy only stores what it routes). +// overrideFor supplies the engine's seat→account routing override (nil → rank-0 only). +// skip is the current cooldown view (proxy.CooldownSkipSet) so the routed stamp reflects +// cooling-driven failover, matching the hot path (nil → no cooldown view). +func writeGroupRuntimeForGroups(dbPath string, derivedKey []byte, mks []vault.ManagedKey, groups []grGroup, overrideFor func(seatID, groupID string) string, skip map[string]bool) error { + // group_id → its locally-known managed keys (need the whole mk for SeatID + + // GroupAccounts to compute the per-seat routed account, not just the VK id). + mksByGroup := make(map[string][]vault.ManagedKey) + for i := range mks { + if mks[i].OauthGroupID != "" { + mksByGroup[mks[i].OauthGroupID] = append(mksByGroup[mks[i].OauthGroupID], mks[i]) + } + } + delivered := make(map[string]bool, len(groups)) + for _, g := range groups { + delivered[g.OauthGroupID] = true + groupMks := mksByGroup[g.OauthGroupID] + if len(groupMks) == 0 { + continue + } + // Encrypt the material ONCE per group (routed-agnostic), then stamp the + // per-seat routed flag per VK — avoids re-encrypting for each seat. The pick is + // fed the FRESH material map being written (not the stale stored column) so the + // stamp's usability gates see exactly what the hot path will read next. + base := buildGroupRuntimeMap(derivedKey, g.Accounts) + nowUnix := time.Now().Unix() + for _, mk := range groupMks { + jsonVal, err := marshalGroupRuntime(base, computeRoutedAccountID(mk, base, overrideFor, skip, nowUnix)) + if err != nil { + return err + } + if err := vault.WriteGroupRuntime(dbPath, mk.VirtualKeyID, jsonVal); err != nil { + return err + } + } + } + // Access gate (defense-in-depth): a local group VK whose group is NO LONGER in + // the delivery — its seat was unbound, so master stopped delivering it (channel + // ③'s oauth_group_member gate) — must have its cached token WIPED, so a stale + // secret can't keep serving. The master-side snapshot candidate-set gate already + // cuts the route on the next key sync; this clears the residual material on the + // proxy's own poll (independent of the CLI), closing the window either way. + // group_runtime is a JSON object {account_id:{...}}, so empty = "{}". + // Only reached after a SUCCESSFUL delivery fetch (caller gates on ok), so this + // never wipes on a transient master error. + for gid, groupMks := range mksByGroup { + if delivered[gid] { + continue + } + for _, mk := range groupMks { + if err := vault.WriteGroupRuntime(dbPath, mk.VirtualKeyID, "{}"); err != nil { + return err + } + } + } + return nil +} diff --git a/internal/supervisor/group_runtime_policy_test.go b/internal/supervisor/group_runtime_policy_test.go new file mode 100644 index 0000000..61bcd0a --- /dev/null +++ b/internal/supervisor/group_runtime_policy_test.go @@ -0,0 +1,405 @@ +package supervisor + +import ( + "context" + "database/sql" + "encoding/base64" + "encoding/json" + "net/http" + "net/http/httptest" + "path/filepath" + "strings" + "testing" + + "github.com/AiKeyLabs/aikey-proxy/internal/vault" + "github.com/AiKeyLabs/aikey-proxy/internal/vkeys" + _ "modernc.org/sqlite" +) + +func testKey() []byte { + k := make([]byte, 32) + for i := range k { + k[i] = byte(i + 7) + } + return k +} + +func TestBuildGroupRuntimeJSON_EncryptsBothTypesNoRefresh(t *testing.T) { + key := testKey() + pct := 97 + reset := int64(1750000000) + accts := []grAccount{ + {AccountID: "a-oauth", CredentialType: "oauth_account", AccessToken: "at-live", ExpiresAt: 200, + WindowMaxUtilPct: &pct, WindowStatus: "active", WindowResetAt: &reset}, + {AccountID: "a-key", CredentialType: "api_key", Key: "sk-real", BaseURL: "https://x", Revision: "r9"}, + } + js, err := buildGroupRuntimeJSON(key, accts) + if err != nil { + t.Fatalf("build: %v", err) + } + // No plaintext secret + no refresh anywhere on the wire. + low := strings.ToLower(js) + if strings.Contains(js, "at-live") || strings.Contains(js, "sk-real") { + t.Fatalf("plaintext secret leaked into group_runtime: %s", js) + } + if strings.Contains(low, "refresh") { + t.Fatalf("refresh token reference in group_runtime: %s", js) + } + + var m map[string]vkeys.GroupRuntimeAccount + if err := json.Unmarshal([]byte(js), &m); err != nil { + t.Fatalf("unmarshal: %v", err) + } + // OAuth: decrypts back to the access_token + carries window meta. + oa := m["a-oauth"] + if got := decryptSecret(t, key, oa); got != "at-live" { + t.Fatalf("oauth secret decrypt: %q", got) + } + if oa.ExpiresAt != 200 || oa.WindowMaxUtilPct == nil || *oa.WindowMaxUtilPct != 97 || oa.WindowStatus != "active" { + t.Fatalf("oauth meta wrong: %+v", oa) + } + if oa.BaseURL != "" || oa.Revision != "" { + t.Fatalf("oauth must not carry KEY meta: %+v", oa) + } + // KEY: decrypts back to the key + carries base_url/revision. + k := m["a-key"] + if got := decryptSecret(t, key, k); got != "sk-real" { + t.Fatalf("key secret decrypt: %q", got) + } + if k.BaseURL != "https://x" || k.Revision != "r9" || k.ExpiresAt != 0 { + t.Fatalf("key meta wrong: %+v", k) + } +} + +func decryptSecret(t *testing.T, key []byte, a vkeys.GroupRuntimeAccount) string { + t.Helper() + nonce, err := base64.StdEncoding.DecodeString(a.SecretNonce) + if err != nil { + t.Fatalf("nonce b64: %v", err) + } + ct, err := base64.StdEncoding.DecodeString(a.SecretCiphertext) + if err != nil { + t.Fatalf("ct b64: %v", err) + } + pt, err := vault.Decrypt(key, nonce, ct) + if err != nil { + t.Fatalf("decrypt: %v", err) + } + return string(pt) +} + +func TestFetchGroupRuntime_ParsesAndSendsBearer(t *testing.T) { + var gotAuth, gotObsReset string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotAuth = r.Header.Get("Authorization") + gotObsReset = r.Header.Get(observedResetsHeader) + if r.URL.Path != "/accounts/me/group-runtime" { + w.WriteHeader(404) + return + } + _, _ = w.Write([]byte(`{"groups":[{"oauth_group_id":"grp-1","routing_config":"{}","accounts":[{"account_id":"a1","credential_type":"oauth_account","access_token":"tok","expires_at":9}]}]}`)) + })) + defer srv.Close() + + groups, body, err := fetchGroupRuntime(context.Background(), srv.URL, "JWT123", map[string]int64{"acc-1": 1750000000}) + if err != nil || len(groups) != 1 || groups[0].OauthGroupID != "grp-1" || len(groups[0].Accounts) != 1 { + t.Fatalf("fetch: err=%v groups=%+v", err, groups) + } + if body == "" || !strings.Contains(body, "grp-1") { + t.Fatalf("raw body (change signature) missing: %q", body) + } + if groups[0].Accounts[0].AccessToken != "tok" { + t.Fatalf("account material not parsed: %+v", groups[0].Accounts[0]) + } + if gotAuth != "Bearer JWT123" { + t.Fatalf("bearer not sent: %q", gotAuth) + } + // Path Z: observed resets piggybacked as base64(JSON) on the pull. + if gotObsReset == "" { + t.Fatal("observed-resets header not sent") + } + raw, decErr := base64.StdEncoding.DecodeString(gotObsReset) + if decErr != nil { + t.Fatalf("observed-resets header not base64: %v", decErr) + } + var m map[string]int64 + if json.Unmarshal(raw, &m) != nil || m["acc-1"] != 1750000000 { + t.Fatalf("observed-resets header payload wrong: %q → %+v", string(raw), m) + } + + // Non-200 → ok=false (keep last-known). Nil resets → header omitted. + var gotObs2 string + bad := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotObs2 = r.Header.Get(observedResetsHeader) + w.WriteHeader(401) + })) + defer bad.Close() + if _, _, err := fetchGroupRuntime(context.Background(), bad.URL, "x", nil); err == nil { + t.Fatal("401 must yield a non-nil error (keep-last-known + rail failure count)") + } + if gotObs2 != "" { + t.Fatalf("nil observed-resets must omit the header, got %q", gotObs2) + } +} + +func TestWriteGroupRuntimeForGroups_PerVKEncrypted(t *testing.T) { + dbPath := filepath.Join(t.TempDir(), "vault.db") + db, err := sql.Open("sqlite", dbPath) + if err != nil { + t.Fatalf("open: %v", err) + } + if _, err := db.Exec(`CREATE TABLE managed_virtual_keys_cache ( + virtual_key_id TEXT PRIMARY KEY, oauth_group_id TEXT, group_runtime TEXT)`); err != nil { + t.Fatalf("create: %v", err) + } + // vk-g1 → grp-1, vk-direct → no group. + db.Exec(`INSERT INTO managed_virtual_keys_cache (virtual_key_id, oauth_group_id) VALUES ('vk-g1','grp-1')`) + db.Exec(`INSERT INTO managed_virtual_keys_cache (virtual_key_id, oauth_group_id) VALUES ('vk-direct','')`) + db.Close() + + key := testKey() + mks := []vault.ManagedKey{ + {VirtualKeyID: "vk-g1", OauthGroupID: "grp-1"}, + {VirtualKeyID: "vk-direct"}, + } + groups := []grGroup{{OauthGroupID: "grp-1", Accounts: []grAccount{ + {AccountID: "a1", CredentialType: "oauth_account", AccessToken: "tok-1", ExpiresAt: 5}, + }}} + + if err := writeGroupRuntimeForGroups(dbPath, key, mks, groups, nil, nil); err != nil { + t.Fatalf("write: %v", err) + } + + // vk-g1 got an encrypted group_runtime decrypting to tok-1; vk-direct untouched. + db, _ = sql.Open("sqlite", dbPath) + defer db.Close() + var gr sql.NullString + db.QueryRow(`SELECT group_runtime FROM managed_virtual_keys_cache WHERE virtual_key_id='vk-g1'`).Scan(&gr) + if !gr.Valid || gr.String == "" { + t.Fatal("vk-g1 group_runtime not written") + } + var m map[string]vkeys.GroupRuntimeAccount + if err := json.Unmarshal([]byte(gr.String), &m); err != nil { + t.Fatalf("parse stored: %v", err) + } + if got := decryptSecret(t, key, m["a1"]); got != "tok-1" { + t.Fatalf("stored secret decrypt: %q", got) + } + + var grD sql.NullString + db.QueryRow(`SELECT group_runtime FROM managed_virtual_keys_cache WHERE virtual_key_id='vk-direct'`).Scan(&grD) + if grD.Valid && grD.String != "" { + t.Fatalf("direct-bind VK group_runtime must stay empty, got %q", grD.String) + } +} + +// TestWriteGroupRuntimeForGroups_ClearsUndeliveredGroupVK: a local group VK whose +// group is NO LONGER in the delivery (its seat was unbound → master stopped +// delivering it) gets its cached token WIPED to "{}" (access gate, defense-in- +// depth), while a still-delivered group VK keeps fresh material. +func TestWriteGroupRuntimeForGroups_ClearsUndeliveredGroupVK(t *testing.T) { + dbPath := filepath.Join(t.TempDir(), "vault.db") + db, err := sql.Open("sqlite", dbPath) + if err != nil { + t.Fatalf("open: %v", err) + } + if _, err := db.Exec(`CREATE TABLE managed_virtual_keys_cache ( + virtual_key_id TEXT PRIMARY KEY, oauth_group_id TEXT, group_runtime TEXT)`); err != nil { + t.Fatalf("create: %v", err) + } + // vk-gone (grp-gone) carries a STALE cached token; vk-stay (grp-stay) is still a member. + db.Exec(`INSERT INTO managed_virtual_keys_cache (virtual_key_id, oauth_group_id, group_runtime) + VALUES ('vk-gone','grp-gone','{"a-old":{"secret_ciphertext":"stale","secret_nonce":"x"}}')`) + db.Exec(`INSERT INTO managed_virtual_keys_cache (virtual_key_id, oauth_group_id, group_runtime) + VALUES ('vk-stay','grp-stay','{"a-old":{}}')`) + db.Close() + + key := testKey() + mks := []vault.ManagedKey{ + {VirtualKeyID: "vk-gone", OauthGroupID: "grp-gone"}, + {VirtualKeyID: "vk-stay", OauthGroupID: "grp-stay"}, + } + // Delivery includes grp-stay ONLY — grp-gone dropped out (seat unbound). + groups := []grGroup{{OauthGroupID: "grp-stay", Accounts: []grAccount{ + {AccountID: "a1", CredentialType: "oauth_account", AccessToken: "tok-1", ExpiresAt: 5}, + }}} + + if err := writeGroupRuntimeForGroups(dbPath, key, mks, groups, nil, nil); err != nil { + t.Fatalf("write: %v", err) + } + + db, _ = sql.Open("sqlite", dbPath) + defer db.Close() + // vk-gone: stale token WIPED to "{}". + var grGone sql.NullString + db.QueryRow(`SELECT group_runtime FROM managed_virtual_keys_cache WHERE virtual_key_id='vk-gone'`).Scan(&grGone) + if grGone.String != "{}" { + t.Fatalf("undelivered group VK token must be wiped to {}, got %q", grGone.String) + } + // vk-stay: fresh material delivered (still a member). + var grStay sql.NullString + db.QueryRow(`SELECT group_runtime FROM managed_virtual_keys_cache WHERE virtual_key_id='vk-stay'`).Scan(&grStay) + var m map[string]vkeys.GroupRuntimeAccount + if err := json.Unmarshal([]byte(grStay.String), &m); err != nil || decryptSecret(t, key, m["a1"]) != "tok-1" { + t.Fatalf("still-member group VK must keep fresh material, got %q", grStay.String) + } +} + +// twoCandGroupAccounts is a seat's candidate list JSON (mirrors the CLI-synced +// group_accounts column) with two equal-priority accounts, so seatassign.Rank +// decides the rank-0 default deterministically. +const twoCandGroupAccounts = `[{"account_id":"a1","priority":1},{"account_id":"a2","priority":1}]` + +// TestComputeRoutedAccountID_OverrideBeatsRankZero (C2): the routed account = the +// engine override when it still names a candidate, else the seatassign rank-0 pick. +// Asserted RELATIVE to the deterministic default (HRW output isn't hand-predicted): +// picking the OTHER candidate as override MUST flip the result; a non-candidate +// override MUST be ignored (fall back to rank-0). 能红: if computeRoutedAccountID +// ignored the override, the "override flips it" assertion fails. +func TestComputeRoutedAccountID_OverrideBeatsRankZero(t *testing.T) { + mk := vault.ManagedKey{SeatID: "seat-x", GroupAccounts: twoCandGroupAccounts} + + def := computeRoutedAccountID(mk, nil, nil, nil, 1_000_000) // rank-0, no override, no cooldown + if def != "a1" && def != "a2" { + t.Fatalf("rank-0 default must be a candidate, got %q", def) + } + other := "a1" + if def == "a1" { + other = "a2" + } + // Override naming the OTHER candidate must win. + if got := computeRoutedAccountID(mk, nil, func(string, string) string { return other }, nil, 1_000_000); got != other { + t.Fatalf("override should route to %q, got %q", other, got) + } + // Override naming a NON-candidate must be ignored → rank-0 default. + if got := computeRoutedAccountID(mk, nil, func(string, string) string { return "ghost" }, nil, 1_000_000); got != def { + t.Fatalf("non-candidate override must fall back to rank-0 %q, got %q", def, got) + } + // No parseable candidates → "". + if got := computeRoutedAccountID(vault.ManagedKey{SeatID: "s"}, nil, nil, nil, 1_000_000); got != "" { + t.Fatalf("no candidates → \"\", got %q", got) + } +} + +// CONVERGENCE (2026-07-01): computeRoutedAccountID (the is_current_routed / current_routed +// DISPLAY stamp read by /user/vault + /user/virtual-keys after A2) is now cooldown-aware — +// it takes the SAME skip view the hot-path resolver uses (proxy.CooldownSkipSet), so the +// displayed account MATCHES what the proxy actually forwards to under cooling-driven +// failover. Mirrors proxy.TestResolveGroup_RoutedFollowsOverride_AndCoolingFallsThrough +// (the actual forward side). 能红: drop the `!skip` guards in computeRoutedAccountID → the +// cooled-rank-0 case below reverts to rank-0 and diverges from the hot path → fails. +func TestComputeRoutedAccountID_CoolingAware_MatchesHotPath(t *testing.T) { + mk := vault.ManagedKey{SeatID: "seat-cool", GroupAccounts: twoCandGroupAccounts} + def := computeRoutedAccountID(mk, nil, nil, nil, 1_000_000) // rank-0, nothing cooled + other := "a1" + if def == "a1" { + other = "a2" + } + + // rank-0 (def) COOLED, no override → the stamp moves to the next non-cooled account, + // exactly like the hot path's ranked-loop fall-through. + if got := computeRoutedAccountID(mk, nil, nil, map[string]bool{def: true}, 1_000_000); got != other { + t.Fatalf("cooled rank-0 → stamp must move to next non-cooled %q, got %q", other, got) + } + // Override account COOLED → the override is NOT honored; fall through to non-cooled + // (same gate as the hot path: `override != "" && !skip[override]`). + if got := computeRoutedAccountID(mk, nil, func(string, string) string { return def }, map[string]bool{def: true}, 1_000_000); got != other { + t.Fatalf("cooled override must fall through to %q, got %q", other, got) + } + // ALL cooled → nominal rank-0 (hot path would 429; display shows the engine pick). + if got := computeRoutedAccountID(mk, nil, nil, map[string]bool{"a1": true, "a2": true}, 1_000_000); got != def { + t.Fatalf("all cooled → nominal rank-0 %q, got %q", def, got) + } +} + +// TestWriteGroupRuntimeForGroups_StampsCurrentRouted (C2): the writer flags exactly +// ONE account per VK as IsCurrentRouted, honoring the override, and the same group's +// two seats get DIFFERENT routed accounts (per-VK, not per-group-shared). +func TestWriteGroupRuntimeForGroups_StampsCurrentRouted(t *testing.T) { + dbPath := filepath.Join(t.TempDir(), "vault.db") + db, err := sql.Open("sqlite", dbPath) + if err != nil { + t.Fatalf("open: %v", err) + } + if _, err := db.Exec(`CREATE TABLE managed_virtual_keys_cache ( + virtual_key_id TEXT PRIMARY KEY, oauth_group_id TEXT, group_runtime TEXT)`); err != nil { + t.Fatalf("create: %v", err) + } + db.Exec(`INSERT INTO managed_virtual_keys_cache (virtual_key_id, oauth_group_id) VALUES ('vk-s1','grp-1')`) + db.Exec(`INSERT INTO managed_virtual_keys_cache (virtual_key_id, oauth_group_id) VALUES ('vk-s2','grp-1')`) + db.Close() + + key := testKey() + // Two seats on the SAME group → must get independently-stamped routed accounts. + mks := []vault.ManagedKey{ + {VirtualKeyID: "vk-s1", OauthGroupID: "grp-1", SeatID: "seat-1", GroupAccounts: twoCandGroupAccounts}, + {VirtualKeyID: "vk-s2", OauthGroupID: "grp-1", SeatID: "seat-2", GroupAccounts: twoCandGroupAccounts}, + } + groups := []grGroup{{OauthGroupID: "grp-1", Accounts: []grAccount{ + {AccountID: "a1", CredentialType: "oauth_account", AccessToken: "t1", ExpiresAt: 9}, + {AccountID: "a2", CredentialType: "oauth_account", AccessToken: "t2", ExpiresAt: 9}, + }}} + // Force seat-1 → a1 via override; seat-2 gets no override (rank-0). + overrideFor := func(seat, _ string) string { + if seat == "seat-1" { + return "a1" + } + return "" + } + if err := writeGroupRuntimeForGroups(dbPath, key, mks, groups, overrideFor, nil); err != nil { + t.Fatalf("write: %v", err) + } + + routedOf := func(vk string) string { + db, _ := sql.Open("sqlite", dbPath) + defer db.Close() + var gr string + db.QueryRow(`SELECT group_runtime FROM managed_virtual_keys_cache WHERE virtual_key_id=?`, vk).Scan(&gr) + var m map[string]vkeys.GroupRuntimeAccount + if err := json.Unmarshal([]byte(gr), &m); err != nil { + t.Fatalf("parse %s: %v", vk, err) + } + routed := "" + for id, a := range m { + if a.IsCurrentRouted { + if routed != "" { + t.Fatalf("%s: more than one account flagged current-routed", vk) + } + routed = id + } + } + return routed + } + if got := routedOf("vk-s1"); got != "a1" { + t.Fatalf("seat-1 override → a1 current-routed, got %q", got) + } + // seat-2 has no override → rank-0; just assert exactly one is flagged and it's a candidate. + if got := routedOf("vk-s2"); got != "a1" && got != "a2" { + t.Fatalf("seat-2 current-routed must be a candidate, got %q", got) + } +} + +// TestStampCurrentRoutedJSON_FlipsFlagInPlace (C2 coupling): the override-change +// re-stamp path moves IsCurrentRouted to the new account WITHOUT touching secrets, +// and reports changed=false when nothing moves. +func TestStampCurrentRoutedJSON_FlipsFlagInPlace(t *testing.T) { + orig := `{"a1":{"credential_type":"oauth_account","secret_ciphertext":"ZZ","is_current_routed":true},"a2":{"credential_type":"oauth_account","secret_ciphertext":"YY"}}` + // Re-route to a2. + out, changed, err := stampCurrentRoutedJSON(orig, "a2") + if err != nil || !changed { + t.Fatalf("expected change, got changed=%v err=%v", changed, err) + } + var m map[string]vkeys.GroupRuntimeAccount + json.Unmarshal([]byte(out), &m) + if m["a1"].IsCurrentRouted || !m["a2"].IsCurrentRouted { + t.Fatalf("flag should move a1→a2, got a1=%v a2=%v", m["a1"].IsCurrentRouted, m["a2"].IsCurrentRouted) + } + // Secret untouched (never decrypted / re-encrypted). + if m["a1"].SecretCiphertext != "ZZ" || m["a2"].SecretCiphertext != "YY" { + t.Fatalf("secrets must be untouched, got %q / %q", m["a1"].SecretCiphertext, m["a2"].SecretCiphertext) + } + // Idempotent: re-stamping the same routed account reports no change. + if _, changed2, _ := stampCurrentRoutedJSON(out, "a2"); changed2 { + t.Fatalf("re-stamp of unchanged routed must report changed=false") + } +} diff --git a/internal/supervisor/netmon.go b/internal/supervisor/netmon.go new file mode 100644 index 0000000..d282547 --- /dev/null +++ b/internal/supervisor/netmon.go @@ -0,0 +1,121 @@ +// netmon.go — Stage 2 of control-plane self-heal: PROACTIVE host-network-change +// detection (dependency-free). +// +// Stage 1 (selfheal.go) recovers AFTER a control-plane call fails. This monitor +// recovers the moment the host's addresses change (WiFi switch, interface up/down, +// USB/phone tether plug) — the exact events that leave a long-lived process's +// direct client stalled with "no route to host". On a change it rebuilds the +// control-plane client so the next poll/writeback already dials clean, instead of +// burning a failure + waiting the 60s poll cadence. +// +// WHY a poll-a-fingerprint approach (not OS route/netlink events): a fingerprint +// of net.Interfaces() is pure Go stdlib and identical on macOS/Linux/Windows, so +// it needs NO cgo and NO third-party dependency. The reference event-driven +// implementation (Tailscale's tailscale.com/net/netmon: PF_ROUTE / netlink / +// NotifyAddrChange) reacts faster and catches pure default-route changes with +// unchanged IPs, but adds a heavy dependency; it's a drop-in upgrade if ever +// needed (swap fingerprint-poll for an event channel — same onChange callback). +package supervisor + +import ( + "context" + "log/slog" + "net" + "sort" + "strings" + "time" + + "github.com/AiKeyLabs/aikey-proxy/internal/httpx" + "github.com/AiKeyLabs/aikey-proxy/internal/observability" +) + +const netChangePollInterval = 20 * time.Second + +// interfaceFingerprint is a cheap, order-stable signature of the host's usable +// network addresses (up, non-loopback, global-unicast). It flips exactly when +// the set of interface IPs changes — the trigger for control-plane staleness. +func interfaceFingerprint() string { + ifaces, err := net.Interfaces() + if err != nil { + return "" + } + var ips []string + for _, ifc := range ifaces { + if ifc.Flags&net.FlagUp == 0 || ifc.Flags&net.FlagLoopback != 0 { + continue + } + addrs, err := ifc.Addrs() + if err != nil { + continue + } + for _, a := range addrs { + if ipn, ok := a.(*net.IPNet); ok && ipn.IP.IsGlobalUnicast() { + ips = append(ips, ipn.IP.String()) + } + } + } + sort.Strings(ips) + return strings.Join(ips, ",") +} + +// changeDetector holds the last-seen fingerprint. The FIRST observation only +// primes the baseline (no change reported), so a fresh monitor never fires on +// startup. Pulled out of the loop so the flip logic is deterministically testable +// without a ticker/clock. +type changeDetector struct { + primed bool + last string +} + +// observe records cur and reports whether it differs from the previously-seen +// value, plus the previous value (for from→to logging). First call primes the +// baseline and returns (false, ""). +func (d *changeDetector) observe(cur string) (changed bool, prev string) { + if !d.primed { + d.primed, d.last = true, cur + return false, "" + } + if cur != d.last { + prev, d.last = d.last, cur + return true, prev + } + return false, d.last +} + +// watchNetworkChanges polls `fingerprint` every `interval` and invokes onChange +// each time it flips. fingerprint/onChange are injected so tests drive it without +// real interfaces; production passes interfaceFingerprint. +func watchNetworkChanges(ctx context.Context, interval time.Duration, fingerprint func() string, onChange func(old, cur string)) { + var det changeDetector + det.observe(fingerprint()) // prime baseline; never fire on start + t := time.NewTicker(interval) + defer t.Stop() + for { + select { + case <-ctx.Done(): + return + case <-t.C: + cur := fingerprint() + if changed, prev := det.observe(cur); changed { + onChange(prev, cur) + } + } + } +} + +// onNetworkChange is the production reaction to a detected host network change: +// rebuild EVERY registered control-plane client (so all rails dial clean, not just +// group-runtime) and reset the self-heal streak (the change is a fresh start). +// Named (not an inline closure) so integration tests drive the REAL callback. +func onNetworkChange(old, cur string) { + n := httpx.RebuildAllControlPlane() + controlPlaneHeal.onPollOK() + slog.Info("host network change detected; rebuilt control-plane clients", + "event.name", observability.EventProxyControlPlaneNetChange, + "from", old, "to", cur, "clients_rebuilt", n) +} + +// runNetChangeMonitor is the production loop. +func runNetChangeMonitor(ctx context.Context) { + watchNetworkChanges(ctx, netChangePollInterval, interfaceFingerprint, onNetworkChange) +} diff --git a/internal/supervisor/netmon_test.go b/internal/supervisor/netmon_test.go new file mode 100644 index 0000000..cc01bb6 --- /dev/null +++ b/internal/supervisor/netmon_test.go @@ -0,0 +1,78 @@ +package supervisor + +import ( + "context" + "testing" + "time" +) + +func TestChangeDetector(t *testing.T) { + var d changeDetector + // first observation primes the baseline — never a change + if changed, _ := d.observe("A"); changed { + t.Fatal("first observe reported a change (should prime baseline)") + } + // same value → no change + if changed, _ := d.observe("A"); changed { + t.Fatal("same fingerprint reported a change") + } + // flip → change, prev is the old value + changed, prev := d.observe("B") + if !changed || prev != "A" { + t.Fatalf("flip A→B: changed=%v prev=%q, want true/\"A\"", changed, prev) + } + // stays at B → no change + if changed, _ := d.observe("B"); changed { + t.Fatal("stable B reported a change") + } + // flip back → change, prev=B + if changed, prev := d.observe("A"); !changed || prev != "B" { + t.Fatalf("flip B→A: changed=%v prev=%q, want true/\"B\"", changed, prev) + } +} + +func TestWatchNetworkChanges_FiresOnlyOnFlip(t *testing.T) { + // scripted fingerprint sequence: baseline X, then X (no fire), Y (fire), + // Y (no fire), Z (fire). Driven by a counter so it's deterministic. + seq := []string{"X", "X", "Y", "Y", "Z"} + i := 0 + fp := func() string { + v := seq[i] + if i < len(seq)-1 { + i++ + } + return v + } + var fires [][2]string + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan struct{}) + onChange := func(old, cur string) { + fires = append(fires, [2]string{old, cur}) + if len(fires) == 2 { // both flips seen → stop + cancel() + } + } + go func() { watchNetworkChanges(ctx, time.Millisecond, fp, onChange); close(done) }() + + select { + case <-done: + case <-time.After(2 * time.Second): + cancel() + t.Fatal("watchNetworkChanges did not report the expected flips in time") + } + if len(fires) != 2 || fires[0] != [2]string{"X", "Y"} || fires[1] != [2]string{"Y", "Z"} { + t.Fatalf("fires=%v, want [X→Y, Y→Z]", fires) + } +} + +func TestInterfaceFingerprint_StableAndBounded(t *testing.T) { + a := interfaceFingerprint() + b := interfaceFingerprint() + if a != b { + t.Fatalf("fingerprint not stable across calls: %q vs %q", a, b) + } + // Loopback must be excluded — it is never a routing change signal. + if a == "127.0.0.1" || a == "::1" { + t.Fatalf("fingerprint leaked loopback: %q", a) + } +} diff --git a/internal/supervisor/oauth_group.go b/internal/supervisor/oauth_group.go new file mode 100644 index 0000000..c6f894f --- /dev/null +++ b/internal/supervisor/oauth_group.go @@ -0,0 +1,17 @@ +package supervisor + +import "github.com/AiKeyLabs/aikey-proxy/internal/vkeys" + +// Oauth-group (channel ③) proxy-side gating (N7c). The env gate's single source +// of truth is vkeys.OauthGroupRoutingEnabled (read by both supervisor and the +// proxy data plane); this thin wrapper keeps supervisor call sites unchanged. +// +// Group virtual keys carry no PlaintextKey — their per-account material lives in +// managed_virtual_keys_cache.group_runtime, pulled from master by the proxy, and +// a request is routed by picking a candidate account (seatassign) + injecting its +// token (resolver = N8). Until that resolver + the group-runtime pull loop are +// complete, group VKs MUST NOT enter the route registry (they'd fall to the +// personal-key path and 401). Default OFF → the direct-bind path is byte-unchanged. + +// oauthGroupRoutingEnabled reports whether proxy-side group VK routing is on. +func oauthGroupRoutingEnabled() bool { return vkeys.OauthGroupRoutingEnabled() } diff --git a/internal/supervisor/quota_policy.go b/internal/supervisor/quota_policy.go index 8a39375..609cc70 100644 --- a/internal/supervisor/quota_policy.go +++ b/internal/supervisor/quota_policy.go @@ -29,13 +29,14 @@ import ( "strings" "time" + "github.com/AiKeyLabs/aikey-proxy/internal/httpx" "github.com/AiKeyLabs/aikey-proxy/internal/quota" "github.com/AiKeyLabs/aikey-proxy/internal/vault" ) const quotaPollInterval = 60 * time.Second -var quotaHTTPClient = &http.Client{Timeout: 10 * time.Second} +var quotaHTTPClient = httpx.NewSwappableDirect(10 * time.Second) // pollQuotaPolicy runs until ctx is canceled, refreshing the org quota policy // every quotaPollInterval (plus once immediately). @@ -114,7 +115,7 @@ func fetchQuotaPolicy(ctx context.Context, masterURL, orgID string, seats []stri if err != nil { return nil, "", false } - resp, err := quotaHTTPClient.Do(req) + resp, err := quotaHTTPClient.Get().Do(req) if err != nil { return nil, "", false } diff --git a/internal/supervisor/railset.go b/internal/supervisor/railset.go new file mode 100644 index 0000000..e786fea --- /dev/null +++ b/internal/supervisor/railset.go @@ -0,0 +1,489 @@ +// railset.go — the SyncRail framework: ONE declarative driver for every +// control-plane sync rail (master-poll loop) in the proxy. +// +// WHY (2026-07-03 incident, bugfix 2026-07-03-routing-override-rail-silent-stall.md): +// the proxy grew six hand-written "pull from master" loops whose behaviors +// drifted apart — the routing-override and group-runtime polls built their +// team credential ONCE at goroutine start (baking a possibly-stale control URL +// into the refresh path) and early-returned FOREVER on any startup precondition +// miss, all silently. A server IP drift then starved both rails for 7+ hours +// with zero logs while the resolver kept demanding login for an account the +// engine had routed the seat away from. +// +// The framework inverts every "evaluate once at start" into "evaluate every +// cycle" and centralizes the failure-visibility rules so a rail cannot opt out: +// - gate / generation / credential / control URL: re-checked each cycle +// - failure: counted per rail; OK → STALE (3) → OFFLINE (20) transitions log +// with the underlying error; recovery logs the outage duration +// - state affects VISIBILITY ONLY (/status, statusline): the data path keeps +// serving last-known caches and retries never stop (offline-first, owner +// decision 2026-07-03 — OFFLINE is a label, not a terminal state) +// +// Design doc: update/20260703-控制面同步框架SyncRail-技术方案.md. +// Phase 1 rails: routing_override, group_runtime. quota/compliance/audit stay +// on their legacy loops until Phase 2 (scope fidelity — do not migrate here). +package supervisor + +import ( + "context" + "encoding/json" + "errors" + "log/slog" + "os" + "path/filepath" + "sync" + "time" + + "github.com/AiKeyLabs/aikey-proxy/internal/events" + "github.com/AiKeyLabs/aikey-proxy/internal/observability" +) + +// State thresholds (§2.3): consecutive failures before the rail is announced +// STALE (first WARN) and OFFLINE (ERROR + /status red). Centralized so every +// rail escalates identically; tune here, not per rail. +const ( + railStaleAfterFailures = 3 + railOfflineAfterFailures = 20 + // railReWarnEveryFailures re-emits the offline WARN once an hour (60 × 60s + // cycles) so a long outage stays visible in logs without per-cycle spam. + railReWarnEveryFailures = 60 +) + +// cliTokenRefreshPath is the control-service token-refresh endpoint, appended +// to the CURRENT control URL when (re)building the team credential. Deriving it +// here — instead of trusting the refresh_url baked into aikey-proxy.yaml at +// process start — is the core incident fix: the yaml value goes stale when the +// server address changes while the proxy is running. +const cliTokenRefreshPath = "/v1/auth/cli/token/refresh" + +// railState is the §2.3 visibility state machine. It never gates serving. +type railState int32 + +const ( + railInit railState = iota + railOK + railStale + railOffline +) + +func (s railState) String() string { + switch s { + case railOK: + return "ok" + case railStale: + return "stale" + case railOffline: + return "offline" + default: + return "init" + } +} + +// railSpec declares one control-plane sync rail (§3.1). Everything else — +// ticker, credential, state machine, logging, /status — is the framework's. +type railSpec struct { + // name keys the /status entry and the goroutine label. Keep the historical + // poll names (e.g. "routing_override") so dashboards/log filters carry over. + name string + interval time.Duration + // needsTeamJWT: true → the framework resolves a Bearer via the shared + // teamCredentialSource each cycle and skips the cycle (counted as failure, + // visible) when auth fails. false → sync receives bearer "". + needsTeamJWT bool + // gate: local preconditions (feature flag, "this vault has group VKs", …), + // re-evaluated EVERY cycle. false → the cycle is skipped WITHOUT touching + // the failure counter (a personal install without group VKs is idle, not + // broken). The generation is the CURRENT one (re-loaded each cycle). + gate func(gen *generation) bool + // hydrate, when non-nil, runs once before the first cycle to preload the + // rail's last-known persisted state (survives restarts, §5.3). Best-effort. + hydrate func(gen *generation) + // sync performs one pull+apply. A nil error is a success (state → OK) — + // including "fetched, nothing changed". Any error is counted and surfaces + // in transition logs and /status last_error. + sync func(ctx context.Context, gen *generation, masterURL, bearer string) error +} + +// RailSyncStatus is one rail's /status snapshot entry (§3.3). Exported so the +// cmd layer can hand it to the admin handler (same wiring shape as +// PoolCooldownSnapshot → PoolRoutingHealth). +type RailSyncStatus struct { + State string `json:"state"` + ConsecutiveFailures int `json:"consecutive_failures"` + LastSuccessAt int64 `json:"last_success_at,omitempty"` + LastError string `json:"last_error,omitempty"` + // Attempted distinguishes "never had anything to do" (gate always false — + // omitted from /status) from "tried and is in trouble". + Attempted bool `json:"-"` +} + +// railRunner is the per-rail runtime: counters + state, guarded by mu (written +// only by the rail's own goroutine; read by the /status snapshot). +type railRunner struct { + spec railSpec + + mu sync.Mutex + state railState + failures int + lastSuccessAt int64 + lastError string + attempted bool + failedSince int64 // unix of the first failure in the current streak (recovery log) + + kick chan struct{} // buffered(1): Reload/settings nudge → immediate cycle +} + +// railSet drives all declared rails. One goroutine per rail (GoSafe/Isolated — +// a panic in one rail never touches the data path or sibling rails). +type railSet struct { + rails []*railRunner +} + +func newRailSet(specs ...railSpec) *railSet { + rs := &railSet{} + for _, sp := range specs { + rs.rails = append(rs.rails, &railRunner{spec: sp, kick: make(chan struct{}, 1)}) + } + return rs +} + +// start launches every rail loop. Called once from supervisor.New. +func (rs *railSet) start(s *Supervisor) { + for _, r := range rs.rails { + runner := r + observability.GoSafe("supervisor."+runner.spec.name+"_poll", observability.Isolated, func() { + runner.loop(s) + }) + } +} + +// kickAll nudges every rail to run a cycle NOW (non-blocking; a rail already +// mid-cycle coalesces the nudge via the buffered channel). Used by Reload so a +// control-URL change converges in seconds instead of one poll interval. +func (rs *railSet) kickAll() { + for _, r := range rs.rails { + select { + case r.kick <- struct{}{}: + default: + } + } +} + +// snapshot returns the /status view. Rails that never had anything to do +// (gate never passed) are omitted so personal installs don't render noise. +func (rs *railSet) snapshot() map[string]RailSyncStatus { + out := make(map[string]RailSyncStatus, len(rs.rails)) + for _, r := range rs.rails { + r.mu.Lock() + if r.attempted { + out[r.spec.name] = RailSyncStatus{ + State: r.state.String(), + ConsecutiveFailures: r.failures, + LastSuccessAt: r.lastSuccessAt, + LastError: r.lastError, + Attempted: true, + } + } + r.mu.Unlock() + } + return out +} + +// railHealthFor returns one rail's (state, failingSeconds) — the resolver's +// truthful-wording input (§5.4): a stale/offline routing rail means the local +// pick may contradict the engine, so the 401 must not blindly demand a login. +// Unknown rail / unwired railset → ("init", 0), treated as healthy. +func (s *Supervisor) railHealthFor(name string) (string, int64) { + if s.railset == nil { + return railInit.String(), 0 + } + for _, r := range s.railset.rails { + if r.spec.name != name { + continue + } + r.mu.Lock() + defer r.mu.Unlock() + secs := int64(0) + if r.failedSince > 0 { + secs = time.Now().Unix() - r.failedSince + } + return r.state.String(), secs + } + return railInit.String(), 0 +} + +// ControlPlaneSyncSnapshot returns each attempted rail's visibility state for +// the admin /status surface (§3.3) and the statusline sync-health file. Rails +// that never had anything to do (gate never passed) are omitted so personal +// installs render no noise. nil-safe on a not-yet-wired railset. +func (s *Supervisor) ControlPlaneSyncSnapshot() map[string]RailSyncStatus { + if s.railset == nil { + return nil + } + return s.railset.snapshot() +} + +// loop is the rail's lifetime: hydrate once, then cycle on ticker/kick until +// the supervisor context ends. Every cycle re-derives gen/gate/URL/credential — +// nothing is baked at start (the incident's root cause). +func (r *railRunner) loop(s *Supervisor) { + if r.spec.hydrate != nil { + if gen := s.active.Load(); gen != nil && gen.vault != nil { + r.spec.hydrate(gen) + } + } + r.cycle(s) + ticker := time.NewTicker(r.spec.interval) + defer ticker.Stop() + for { + select { + case <-s.ctx.Done(): + return + case <-ticker.C: + r.cycle(s) + case <-r.kick: + r.cycle(s) + } + } +} + +var ( + errRailNoGeneration = errors.New("no active generation/vault yet") + errRailNoControlURL = errors.New("control panel URL not configured") +) + +func (r *railRunner) cycle(s *Supervisor) { + gen := s.active.Load() + if gen == nil || gen.vault == nil { + // Local not-ready (mid-reload window / early start): neither success nor + // failure — the master isn't being blamed for a local swap. + return + } + if r.spec.gate != nil && !r.spec.gate(gen) { + return // idle by design (feature off / nothing local to sync) — not a failure + } + masterURL := readControlPanelURL() + if masterURL == "" { + // A team rail with local work but no control URL is a real broken state + // (e.g. config wiped) — count it so it surfaces, don't silently idle. + r.finish(s, errRailNoControlURL) + return + } + bearer := "" + if r.spec.needsTeamJWT { + b, err := s.teamCred.bearer(s.ctx, gen.vault, masterURL) + if err != nil { + r.finish(s, err) + return + } + bearer = b + } + r.finish(s, r.spec.sync(s.ctx, gen, masterURL, bearer)) +} + +// finish records the cycle outcome and, on a state TRANSITION, refreshes the +// statusline sync-health bypass file (§5.5) — written outside the rail mutex +// (writeSyncHealth snapshots every rail). +func (r *railRunner) finish(s *Supervisor, err error) { + if r.observe(err) && s.railset != nil { + s.railset.writeSyncHealth() + } +} + +// observe feeds one cycle outcome into the state machine and emits the +// transition logs (§2.3). Mandatory-WARN rule: every fallback-to-last-known +// path lands here — a rail can no longer fail silently. Returns true when the +// visibility STATE changed (the caller then refreshes the sync-health file). +func (r *railRunner) observe(err error) bool { + r.mu.Lock() + defer r.mu.Unlock() + r.attempted = true + now := time.Now().Unix() + if err == nil { + prev := r.state + outageSecs := int64(0) + if r.failedSince > 0 { + outageSecs = now - r.failedSince + } + r.state, r.failures, r.lastError, r.failedSince = railOK, 0, "", 0 + r.lastSuccessAt = now + if prev == railStale || prev == railOffline { + slog.Info("control-plane sync rail recovered", + "event.name", observability.EventProxySyncRailRecovered, + "rail", r.spec.name, "outage_seconds", outageSecs) + } + return prev != railOK + } + r.failures++ + r.lastError = err.Error() + if r.failedSince == 0 { + r.failedSince = now + } + switch { + case r.failures == railStaleAfterFailures: + r.state = railStale + slog.Warn("control-plane sync rail is stale — serving last-known data, retrying every cycle", + "event.name", observability.EventProxySyncRailStale, + "rail", r.spec.name, "consecutive_failures", r.failures, "error", r.lastError) + return true + case r.failures == railOfflineAfterFailures: + r.state = railOffline + slog.Error("control-plane sync rail is offline — serving last-known data, retrying every cycle (will not give up)", + "event.name", observability.EventProxySyncRailOffline, + "rail", r.spec.name, "consecutive_failures", r.failures, "error", r.lastError) + return true + case r.state == railOffline && r.failures%railReWarnEveryFailures == 0: + slog.Warn("control-plane sync rail still offline", + "event.name", observability.EventProxySyncRailOffline, + "rail", r.spec.name, "consecutive_failures", r.failures, "error", r.lastError) + } + return false +} + +// ── statusline sync-health bypass file (§5.5) ─────────────────────────────── + +// syncHealthFilename is the statusline's zero-RPC view of degraded rails — +// same ownership pattern as the group-login state file (one concern, one file, +// one writer; the writer here is the SUPERVISOR's rail framework, transitions +// only, so the hot path never touches it). Cross-repo contract with +// aikey-cli commands_statusline.rs — change shapes together. +const syncHealthFilename = "sync-health.json" + +type syncHealthRail struct { + State string `json:"state"` + // FailedSince (unix seconds) lets the reader render a LIVE outage duration + // without the writer re-writing the file every cycle. + FailedSince int64 `json:"failed_since"` +} + +type syncHealthBody struct { + Rails map[string]syncHealthRail `json:"rails"` + WrittenAt int64 `json:"written_at"` // unix millis +} + +func syncHealthPath() (string, error) { + if dir := os.Getenv("AIKEY_RUN_DIR"); dir != "" { + return filepath.Join(dir, syncHealthFilename), nil + } + home, err := os.UserHomeDir() + if err != nil { + return "", err + } + return filepath.Join(home, ".aikey", "run", syncHealthFilename), nil +} + +// writeSyncHealth snapshots the degraded rails into the bypass file, or removes +// the file when every rail is healthy (statusline recovery is automatic). +// Called on state TRANSITIONS only (never per-cycle) and outside rail mutexes. +// Best-effort with a WARN on failure — a silently missing statusline hint is a +// debugging trap (same rule as the group-login state file). +func (rs *railSet) writeSyncHealth() { + degraded := map[string]syncHealthRail{} + for _, r := range rs.rails { + r.mu.Lock() + if r.state == railStale || r.state == railOffline { + degraded[r.spec.name] = syncHealthRail{State: r.state.String(), FailedSince: r.failedSince} + } + r.mu.Unlock() + } + path, err := syncHealthPath() + if err == nil { + if len(degraded) == 0 { + if rmErr := os.Remove(path); rmErr != nil && !os.IsNotExist(rmErr) { + err = rmErr + } + } else { + err = func() error { + if mkErr := os.MkdirAll(filepath.Dir(path), 0o700); mkErr != nil { + return mkErr + } + data, mErr := json.Marshal(syncHealthBody{Rails: degraded, WrittenAt: time.Now().UnixMilli()}) + if mErr != nil { + return mErr + } + tmp := path + ".tmp" + if wErr := os.WriteFile(tmp, data, 0o600); wErr != nil { + return wErr + } + return os.Rename(tmp, path) + }() + } + } + if err != nil { + slog.Warn("sync-health state file update failed — statusline may show a stale sync warning", + "event.name", observability.EventProxySyncHealthFileFailed, "error", err.Error()) + } +} + +// ── shared team credential source (§3.2) ──────────────────────────────────── + +// teamCredentialSource builds the team account-JWT on demand and rebuilds it +// whenever the control URL changes, a refresh fails, or a Reload invalidates it +// — the per-cycle inversion of the old "build once at poll start" pattern. +// +// The refresh URL is DERIVED from the current control URL (cliTokenRefreshPath) +// and the refresh_token is read from the vault at build time; the yaml +// collector_credentials bundle (loaded once at process start) is deliberately +// not consulted — it is exactly the stale-URL source the incident exposed. +type teamCredentialSource struct { + mu sync.Mutex + cred *events.RefreshableJWT + builtURL string +} + +var errRailNoTeamCredential = errors.New("no team credential in vault (not logged in?)") + +// bearer returns a valid team Bearer for masterURL, (re)building the underlying +// credential when needed. A refresh failure drops the credential so the next +// cycle rebuilds from the CURRENT vault refresh_token + control URL. +// vaultReader is the minimal contract (refreshTokenSource) so tests stub it. +func (c *teamCredentialSource) bearer(ctx context.Context, vaultReader refreshTokenSource, masterURL string) (string, error) { + c.mu.Lock() + if c.cred == nil || c.builtURL != masterURL { + refreshToken, err := vaultReader.GetPlatformRefreshToken() + if err != nil || refreshToken == "" { + c.mu.Unlock() + if err != nil { + return "", err + } + return "", errRailNoTeamCredential + } + if c.cred != nil && c.builtURL != masterURL { + slog.Info("control-plane sync: team credential rebuilt for new control URL", + "event.name", observability.EventProxySyncCredentialRebuilt, + "old_url", c.builtURL, "new_url", masterURL) + } + // Zero ExpiresAt → the first Bearer() refreshes immediately against the + // derived URL, minting a fresh access token. PersistFn nil on purpose + // (same trust boundary as buildCollectorCredentials: the proxy never + // writes user.yaml). + c.cred = &events.RefreshableJWT{ + RefreshToken: refreshToken, + RefreshURL: masterURL + cliTokenRefreshPath, + } + c.builtURL = masterURL + } + cred := c.cred + c.mu.Unlock() + + b, err := cred.Bearer(ctx) + if err != nil { + // Drop so the next cycle rebuilds (fresh vault refresh_token + URL) — + // a rotated refresh_token or moved server heals within one cycle. + c.mu.Lock() + if c.cred == cred { + c.cred = nil + } + c.mu.Unlock() + return "", err + } + return b, nil +} + +// invalidate drops the cached credential; the next cycle rebuilds from the +// current vault + control URL. Called by Reload (settings/set-url convergence). +func (c *teamCredentialSource) invalidate() { + c.mu.Lock() + c.cred, c.builtURL = nil, "" + c.mu.Unlock() +} diff --git a/internal/supervisor/railset_integration_test.go b/internal/supervisor/railset_integration_test.go new file mode 100644 index 0000000..e60f57a --- /dev/null +++ b/internal/supervisor/railset_integration_test.go @@ -0,0 +1,214 @@ +package supervisor + +import ( + "context" + "database/sql" + "fmt" + "net/http" + "net/http/httptest" + "os" + "strings" + "sync/atomic" + "testing" + + "github.com/AiKeyLabs/aikey-proxy/internal/config" + "github.com/AiKeyLabs/aikey-proxy/internal/proxy" + + _ "modernc.org/sqlite" +) + +// ── cycle-level integration: the full production path per rail cycle ──────── +// gate → control URL → credential (vault refresh_token + derived refresh URL) +// → fetch (real HTTP) → apply (cache) → persist (vault column) → state machine +// → sync-health file. Everything real except the master (httptest). + +// addPlatformAccount gives the openable vault a platform_account row so the +// teamCredentialSource can mint (same table GetPlatformRefreshToken reads). +func addPlatformAccount(t *testing.T, dbPath, refreshToken string) { + t.Helper() + db, err := sql.Open("sqlite", dbPath) + if err != nil { + t.Fatalf("open: %v", err) + } + defer db.Close() + if _, err := db.Exec(`CREATE TABLE platform_account ( + id INTEGER PRIMARY KEY, refresh_token TEXT)`); err != nil { + t.Fatalf("create platform_account: %v", err) + } + if _, err := db.Exec(`INSERT INTO platform_account (id, refresh_token) VALUES (1, ?)`, + refreshToken); err != nil { + t.Fatalf("insert platform_account: %v", err) + } +} + +// mockMaster serves the token-refresh + routing endpoints like a real control +// service. `down` (atomic) simulates an outage without changing the URL. +func mockMaster(t *testing.T, name string, version int64, accountID string, down *atomic.Bool) *httptest.Server { + t.Helper() + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if down != nil && down.Load() { + w.WriteHeader(http.StatusBadGateway) + return + } + switch r.URL.Path { + case cliTokenRefreshPath: + _, _ = fmt.Fprintf(w, `{"access_token":"tok-%s","expires_in":3600}`, name) + case "/accounts/me/routing": + if r.Header.Get("Authorization") != "Bearer tok-"+name { + w.WriteHeader(http.StatusUnauthorized) + return + } + _, _ = fmt.Fprintf(w, + `{"routing_version":%d,"routes":[{"seat_id":"seat-1","group_id":"grp-1","account_id":"%s"}]}`, + version, accountID) + default: + w.WriteHeader(http.StatusNotFound) + } + })) +} + +func newCycleHarness(t *testing.T) (*Supervisor, *railRunner, string) { + t.Helper() + t.Setenv("AIKEY_PROXY_OAUTH_GROUP_ENABLED", "1") + t.Setenv("AIKEY_RUN_DIR", t.TempDir()) + dbPath, reader := newOpenableVault(t, []map[string]string{ + {"vk": "vk-a", "seat": "seat-1", "group": "grp-1", "override": ""}, + }) + addPlatformAccount(t, dbPath, "refresh-token-1") + + cfg := &config.Config{} + cfg.Vault.Path = dbPath + s := &Supervisor{ + cfg: cfg, + routingOverrides: proxy.NewRoutingOverrideCache(), + teamCred: &teamCredentialSource{}, + ctx: context.Background(), + } + s.active.Store(&generation{vault: reader}) + s.railset = newRailSet(s.routingOverrideRail()) + return s, s.railset.rails[0], dbPath +} + +// THE INCIDENT, as an integration test: the control URL drifts between cycles +// (server A dies, config now points at server B). The very next cycle must +// rebuild the credential against B, pull B's newer assignment, apply AND +// persist it — no restart, no reload, no manual step. 能红: bake the +// credential/URL at rail start (the old behavior) and cycle 2 still dials A → +// this test fails on the stale assignment. +func TestRailCycleIntegration_URLDriftSelfHeals(t *testing.T) { + s, runner, dbPath := newCycleHarness(t) + + serverA := mockMaster(t, "A", 41, "acc-from-A", nil) + t.Setenv("AIKEY_HUB_CONTROL_URL", serverA.URL) + + runner.cycle(s) + if got := s.routingOverrides.Assignment("seat-1", "grp-1"); got != "acc-from-A" { + t.Fatalf("cycle 1 assignment=%q want acc-from-A", got) + } + + // Drift: A is gone; the config now names B (fresh URL, newer version). + serverB := mockMaster(t, "B", 42, "acc-from-B", nil) + defer serverB.Close() + serverA.Close() + t.Setenv("AIKEY_HUB_CONTROL_URL", serverB.URL) + + runner.cycle(s) + if got := s.routingOverrides.Assignment("seat-1", "grp-1"); got != "acc-from-B" { + t.Fatalf("post-drift assignment=%q want acc-from-B (rail did not self-heal onto the new URL)", got) + } + if v := s.routingOverrides.Version(); v != 42 { + t.Fatalf("post-drift version=%d want 42", v) + } + // Persisted for the next restart, from the NEW server. + if col := readOverrideColumn(t, dbPath, "vk-a"); !strings.Contains(col, "acc-from-B") { + t.Fatalf("persisted column must carry the new assignment: %q", col) + } + // The rail is healthy again — visible state agrees. + if st, _ := s.railHealthFor("routing_override"); st != "ok" { + t.Fatalf("rail state=%q want ok after self-heal", st) + } +} + +// Outage semantics end-to-end: the master goes down mid-life → the rail turns +// STALE (visible in /status AND the statusline sync-health file) while the +// cache keeps serving the last-known assignment (offline-first); the master +// returns → one cycle recovers everything and removes the file. +func TestRailCycleIntegration_OutageKeepsLastKnownAndRecovers(t *testing.T) { + s, runner, _ := newCycleHarness(t) + + var down atomic.Bool + master := mockMaster(t, "M", 7, "acc-live", &down) + defer master.Close() + t.Setenv("AIKEY_HUB_CONTROL_URL", master.URL) + + runner.cycle(s) + if got := s.routingOverrides.Assignment("seat-1", "grp-1"); got != "acc-live" { + t.Fatalf("baseline assignment=%q want acc-live", got) + } + + down.Store(true) + for i := 0; i < railStaleAfterFailures; i++ { + runner.cycle(s) + } + if st, secs := s.railHealthFor("routing_override"); st != "stale" || secs < 0 { + t.Fatalf("after %d failed cycles state=%q want stale", railStaleAfterFailures, st) + } + // Offline-first: the data path still serves the last-known assignment. + if got := s.routingOverrides.Assignment("seat-1", "grp-1"); got != "acc-live" { + t.Fatalf("outage must keep last-known assignment, got %q", got) + } + // The statusline bypass file appeared on the transition. + healthPath, _ := syncHealthPath() + if _, err := os.Stat(healthPath); err != nil { + t.Fatalf("stale transition must write the sync-health file: %v", err) + } + // /status snapshot carries the failure detail. + snap := s.ControlPlaneSyncSnapshot() + if st := snap["routing_override"]; st.State != "stale" || st.ConsecutiveFailures != railStaleAfterFailures || st.LastError == "" { + t.Fatalf("/status snapshot wrong: %+v", st) + } + + // Recovery: one healthy cycle heals state and removes the file. + down.Store(false) + runner.cycle(s) + if st, _ := s.railHealthFor("routing_override"); st != "ok" { + t.Fatalf("post-recovery state=%q want ok", st) + } + if _, err := os.Stat(healthPath); !os.IsNotExist(err) { + t.Fatalf("recovery must remove the sync-health file, stat err=%v", err) + } +} + +// Reload convergence (§5.2): invalidate + kickAll make the NEXT cycle rebuild +// the credential even when the URL string is unchanged (e.g. same host, new +// refresh_token after re-login). Exercised at the same integration level. +func TestRailCycleIntegration_InvalidateRebuildsCredential(t *testing.T) { + s, runner, _ := newCycleHarness(t) + + var mints atomic.Int32 + master := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case cliTokenRefreshPath: + mints.Add(1) + _, _ = fmt.Fprintf(w, `{"access_token":"tok-M","expires_in":3600}`) + case "/accounts/me/routing": + _, _ = fmt.Fprintf(w, `{"routing_version":7,"routes":[{"seat_id":"seat-1","group_id":"grp-1","account_id":"acc-live"}]}`) + default: + w.WriteHeader(http.StatusNotFound) + } + })) + defer master.Close() + t.Setenv("AIKEY_HUB_CONTROL_URL", master.URL) + + runner.cycle(s) + runner.cycle(s) + if got := mints.Load(); got != 1 { + t.Fatalf("steady state must reuse the minted token (1 refresh), got %d", got) + } + // Reload's convergence hook: drop the credential → next cycle re-mints. + s.teamCred.invalidate() + runner.cycle(s) + if got := mints.Load(); got != 2 { + t.Fatalf("invalidate must force a rebuild on the next cycle, got %d refreshes", got) + } +} diff --git a/internal/supervisor/railset_test.go b/internal/supervisor/railset_test.go new file mode 100644 index 0000000..f57d004 --- /dev/null +++ b/internal/supervisor/railset_test.go @@ -0,0 +1,227 @@ +package supervisor + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "net/http" + "net/http/httptest" + "os" + "strings" + "testing" + "time" +) + +// The §2.3 state machine: OK → STALE (3 consecutive failures) → OFFLINE (20), +// any success → OK with counters reset. OFFLINE never stops anything — the +// runner keeps observing (owner decision 2026-07-03: offline-first, no give-up). +// 能红: change the thresholds or the reset-on-success and this goes red. +func TestRailRunner_StateMachineTransitions(t *testing.T) { + r := &railRunner{spec: railSpec{name: "test_rail"}} + boom := errors.New("dial tcp: connection refused") + + r.observe(nil) + if got := r.state; got != railOK { + t.Fatalf("after success: state=%v want ok", got) + } + for i := 0; i < railStaleAfterFailures-1; i++ { + r.observe(boom) + } + if r.state != railOK { + t.Fatalf("below stale threshold must keep prior state, got %v", r.state) + } + r.observe(boom) // 3rd consecutive + if r.state != railStale { + t.Fatalf("at %d consecutive failures: state=%v want stale", railStaleAfterFailures, r.state) + } + for i := railStaleAfterFailures; i < railOfflineAfterFailures; i++ { + r.observe(boom) + } + if r.state != railOffline { + t.Fatalf("at %d consecutive failures: state=%v want offline", railOfflineAfterFailures, r.state) + } + if r.failures != railOfflineAfterFailures { + t.Fatalf("failure count=%d want %d", r.failures, railOfflineAfterFailures) + } + // Offline is a LABEL: a success at any point recovers fully. + r.observe(nil) + if r.state != railOK || r.failures != 0 || r.lastError != "" { + t.Fatalf("recovery must reset: state=%v failures=%d lastError=%q", r.state, r.failures, r.lastError) + } + if r.lastSuccessAt == 0 { + t.Fatal("recovery must stamp lastSuccessAt") + } +} + +// last_error must carry the underlying cause — the 2026-07-03 incident's +// connection-refused detail was invisible for 7 hours because the old loop +// swallowed it into a bare early return. +func TestRailRunner_LastErrorSurfaced(t *testing.T) { + r := &railRunner{spec: railSpec{name: "test_rail"}} + r.observe(fmt.Errorf("GET /accounts/me/routing: %w", errors.New("dial tcp 192.168.0.120:3000: connect: connection refused"))) + if r.lastError == "" || r.state != railInit && r.failures != 1 { + t.Fatalf("failure must record lastError, got %q", r.lastError) + } + if want := "connection refused"; !strings.Contains(r.lastError, want) { + t.Fatalf("lastError %q must contain %q", r.lastError, want) + } +} + +// /status omits rails that never attempted a cycle (gate never passed): a +// personal install without group VKs renders no control_plane_sync noise. +func TestRailSet_SnapshotOmitsNeverAttempted(t *testing.T) { + rs := newRailSet( + railSpec{name: "idle_rail"}, + railSpec{name: "busy_rail"}, + ) + rs.rails[1].observe(errors.New("x")) + + snap := rs.snapshot() + if _, ok := snap["idle_rail"]; ok { + t.Fatal("never-attempted rail must be omitted from the snapshot") + } + st, ok := snap["busy_rail"] + if !ok || st.ConsecutiveFailures != 1 || st.State != "init" { + t.Fatalf("attempted rail must be present with counters: %+v", st) + } +} + +// kickAll must never block, even when nobody is draining the kick channel. +func TestRailSet_KickAllNonBlocking(t *testing.T) { + rs := newRailSet(railSpec{name: "r1"}, railSpec{name: "r2"}) + done := make(chan struct{}) + go func() { + rs.kickAll() + rs.kickAll() // second kick coalesces into the buffered slot + close(done) + }() + select { + case <-done: + case <-time.After(2 * time.Second): + t.Fatal("kickAll blocked") + } +} + +type stubRefreshTokenSource struct { + token string + err error +} + +func (s stubRefreshTokenSource) GetPlatformRefreshToken() (string, error) { return s.token, s.err } + +// The incident fix in one test: the credential must be REBUILT when the control +// URL changes between cycles — the refresh POST must land on the NEW server, +// not the one baked at first use. 能红: cache the credential without comparing +// builtURL and the second Bearer still hits the old server. +func TestTeamCredentialSource_RebuildsOnURLChange(t *testing.T) { + mint := func(name string) *httptest.Server { + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != cliTokenRefreshPath { + w.WriteHeader(404) + return + } + _, _ = fmt.Fprintf(w, `{"access_token":"tok-from-%s","expires_in":3600}`, name) + })) + } + oldSrv, newSrv := mint("old"), mint("new") + defer oldSrv.Close() + defer newSrv.Close() + + src := &teamCredentialSource{} + vaultStub := stubRefreshTokenSource{token: "refresh-1"} + + b1, err := src.bearer(context.Background(), vaultStub, oldSrv.URL) + if err != nil || b1 != "tok-from-old" { + t.Fatalf("first bearer: %q err=%v", b1, err) + } + // Control URL drifts (the 2026-07-03 incident: .120 → .121). The very next + // cycle must mint against the NEW URL — no restart, no reload required. + b2, err := src.bearer(context.Background(), vaultStub, newSrv.URL) + if err != nil || b2 != "tok-from-new" { + t.Fatalf("post-drift bearer: %q err=%v (credential not rebuilt for new URL)", b2, err) + } +} + +// A refresh failure must drop the cached credential so the next cycle rebuilds +// from the current vault + URL (a rotated refresh_token heals in one cycle). +func TestTeamCredentialSource_DropsCredentialOnRefreshFailure(t *testing.T) { + fail := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(401) + })) + defer fail.Close() + ok := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _, _ = w.Write([]byte(`{"access_token":"tok-recovered","expires_in":3600}`)) + })) + defer ok.Close() + + src := &teamCredentialSource{} + if _, err := src.bearer(context.Background(), stubRefreshTokenSource{token: "rt"}, fail.URL); err == nil { + t.Fatal("401 refresh must surface an error") + } + // Same source, working server → rebuilt credential succeeds. + b, err := src.bearer(context.Background(), stubRefreshTokenSource{token: "rt"}, ok.URL) + if err != nil || b != "tok-recovered" { + t.Fatalf("post-failure rebuild: %q err=%v", b, err) + } +} + +// No refresh token in the vault (pre-login) is a countable, visible failure — +// the old loop early-returned FOREVER here with zero logs. +func TestTeamCredentialSource_NoTokenIsAnError(t *testing.T) { + src := &teamCredentialSource{} + if _, err := src.bearer(context.Background(), stubRefreshTokenSource{token: ""}, "http://127.0.0.1:1"); err == nil { + t.Fatal("empty refresh token must be an error, not a silent skip") + } + // invalidate() is idempotent and safe on an empty source. + src.invalidate() +} + +// The §5.5 sync-health bypass file lifecycle: a rail transitioning into STALE +// writes the file (with failed_since so the reader renders a live duration); +// recovery back to OK removes it — statusline recovery is automatic. Uses the +// same AIKEY_RUN_DIR override as the group-login state file. +func TestWriteSyncHealth_FileLifecycle(t *testing.T) { + t.Setenv("AIKEY_RUN_DIR", t.TempDir()) + rs := newRailSet(railSpec{name: "routing_override"}, railSpec{name: "group_runtime"}) + boom := errors.New("dial tcp: connection refused") + + // Drive routing_override into STALE via the real observe path. + transitioned := false + for i := 0; i < railStaleAfterFailures; i++ { + transitioned = rs.rails[0].observe(boom) + } + if !transitioned { + t.Fatal("reaching the stale threshold must report a transition") + } + rs.writeSyncHealth() + + path, _ := syncHealthPath() + raw, err := os.ReadFile(path) + if err != nil { + t.Fatalf("degraded rail must write the sync-health file: %v", err) + } + var body syncHealthBody + if err := json.Unmarshal(raw, &body); err != nil { + t.Fatalf("sync-health file must be valid JSON: %v — %s", err, raw) + } + entry, ok := body.Rails["routing_override"] + if !ok || entry.State != "stale" || entry.FailedSince == 0 { + t.Fatalf("sync-health entry wrong: %+v", body.Rails) + } + if _, ok := body.Rails["group_runtime"]; ok { + t.Fatal("healthy rail must not appear in the sync-health file") + } + if body.WrittenAt == 0 { + t.Fatal("written_at must be stamped") + } + + // Recovery removes the file. + if !rs.rails[0].observe(nil) { + t.Fatal("recovery from stale must report a transition") + } + rs.writeSyncHealth() + if _, err := os.Stat(path); !os.IsNotExist(err) { + t.Fatalf("all-healthy must remove the sync-health file, stat err=%v", err) + } +} diff --git a/internal/supervisor/route_builders.go b/internal/supervisor/route_builders.go index d0ee817..68219a1 100644 --- a/internal/supervisor/route_builders.go +++ b/internal/supervisor/route_builders.go @@ -55,6 +55,13 @@ func managedKeyToRoute(mk *vault.ManagedKey) *vkeys.ResolvedRoute { CredentialRevision: mk.CredentialRevision, VirtualKeyRevision: mk.VirtualKeyRevision, RouteSource: "team", + // Oauth-group fields (N7c): empty for direct-bind VKs (PlaintextKey set, + // existing path unchanged); populated for group VKs (PlaintextKey empty, + // resolver picks an account from GroupRuntime). See dispatch N8. + OauthGroupID: mk.OauthGroupID, + GroupAccounts: mk.GroupAccounts, + GroupRuntime: mk.GroupRuntime, + RoutingConfig: mk.RoutingConfig, } } diff --git a/internal/supervisor/routing_override_persist_test.go b/internal/supervisor/routing_override_persist_test.go new file mode 100644 index 0000000..fe77b5a --- /dev/null +++ b/internal/supervisor/routing_override_persist_test.go @@ -0,0 +1,242 @@ +package supervisor + +import ( + "database/sql" + "encoding/binary" + "encoding/json" + "path/filepath" + "testing" + + "github.com/AiKeyLabs/aikey-proxy/internal/config" + "github.com/AiKeyLabs/aikey-proxy/internal/proxy" + "github.com/AiKeyLabs/aikey-proxy/internal/vault" + "github.com/AiKeyLabs/pkg/routingwire" + + _ "modernc.org/sqlite" +) + +// newOpenableVault builds a REAL on-disk vault that vault.Open can open (config +// table with salt/kdf/password_hash) plus a managed_virtual_keys_cache in the +// SAME column shape the CLI migration produces (all five group columns land in +// one batch — schema-code coherence). Low Argon2 cost keeps the test fast. +// Returns (dbPath, reader). +func newOpenableVault(t *testing.T, rows []map[string]string) (string, *vault.Reader) { + t.Helper() + dbPath := filepath.Join(t.TempDir(), "vault.db") + db, err := sql.Open("sqlite", dbPath) + if err != nil { + t.Fatalf("open: %v", err) + } + if _, err := db.Exec(`CREATE TABLE config (key TEXT PRIMARY KEY, value BLOB)`); err != nil { + t.Fatalf("create config: %v", err) + } + salt := []byte("0123456789abcdef") + u32 := func(v uint32) []byte { b := make([]byte, 4); binary.LittleEndian.PutUint32(b, v); return b } + // Cheap KDF params (test-only): 8 KiB / 1 iter / 1 lane. + for k, v := range map[string][]byte{ + "master_salt": salt, "kdf_m_cost": u32(8), "kdf_t_cost": u32(1), "kdf_p_cost": u32(1), + } { + if _, err := db.Exec(`INSERT INTO config (key, value) VALUES (?,?)`, k, v); err != nil { + t.Fatalf("config %s: %v", k, err) + } + } + derived := vault.DeriveKeyWithParams([]byte("pw"), salt, 8, 1, 1) + if _, err := db.Exec(`INSERT INTO config (key, value) VALUES ('password_hash',?)`, derived); err != nil { + t.Fatalf("password_hash: %v", err) + } + if _, err := db.Exec(`CREATE TABLE managed_virtual_keys_cache ( + virtual_key_id TEXT PRIMARY KEY, alias TEXT NOT NULL, local_alias TEXT, + provider_code TEXT, protocol_type TEXT, base_url TEXT, + provider_key_nonce BLOB, provider_key_ciphertext BLOB, provider_base_urls TEXT, + org_id TEXT, seat_id TEXT, credential_id TEXT, credential_revision TEXT, + virtual_key_revision TEXT, owner_account_id TEXT, + key_status TEXT, local_state TEXT, + oauth_group_id TEXT, group_accounts TEXT, group_runtime TEXT, + routing_config TEXT, my_assignment_override TEXT)`); err != nil { + t.Fatalf("create mvk: %v", err) + } + for _, r := range rows { + if _, err := db.Exec(`INSERT INTO managed_virtual_keys_cache + (virtual_key_id, alias, provider_code, protocol_type, base_url, + org_id, seat_id, credential_id, credential_revision, virtual_key_revision, + oauth_group_id, my_assignment_override, key_status) + VALUES (?,?, 'anthropic', 'anthropic', '', 'org', ?, 'cred', 'r', 'vr', ?, ?, 'active')`, + r["vk"], r["vk"], r["seat"], r["group"], r["override"]); err != nil { + t.Fatalf("insert %s: %v", r["vk"], err) + } + } + db.Close() + + reader, err := vault.Open(dbPath, "pw") + if err != nil { + t.Fatalf("vault.Open: %v", err) + } + t.Cleanup(func() { reader.Close() }) + return dbPath, reader +} + +func newPersistSupervisor(dbPath string, reader *vault.Reader) *Supervisor { + cfg := &config.Config{} + cfg.Vault.Path = dbPath + s := &Supervisor{cfg: cfg, routingOverrides: proxy.NewRoutingOverrideCache()} + s.active.Store(&generation{vault: reader}) + return s +} + +func readOverrideColumn(t *testing.T, dbPath, vk string) string { + t.Helper() + db, _ := sql.Open("sqlite", dbPath) + defer db.Close() + var v sql.NullString + if err := db.QueryRow(`SELECT my_assignment_override FROM managed_virtual_keys_cache + WHERE virtual_key_id=?`, vk).Scan(&v); err != nil { + t.Fatalf("read %s: %v", vk, err) + } + return v.String +} + +// The §5.3 restart contract, exercised through the REAL production functions: +// persistAssignmentOverrides writes the applied cache to the columns, then a +// FRESH supervisor (new cache, same vault) hydrates and resolves the same +// (seat,group)→account — including blocked seats (losing that negative +// directive would re-admit a seat the engine capped). 能红: drop the persist +// call, the blocked branch, or hydrate's StoreRoutes and this reds. +func TestPersistThenHydrate_RestartConsistency(t *testing.T) { + t.Setenv("AIKEY_PROXY_OAUTH_GROUP_ENABLED", "1") + dbPath, reader := newOpenableVault(t, []map[string]string{ + {"vk": "vk-a", "seat": "seat-1", "group": "grp-1", "override": ""}, + {"vk": "vk-b", "seat": "seat-2", "group": "grp-1", "override": ""}, + }) + + // Process 1: a live pull stored assignments; persist mirrors them to vault. + s1 := newPersistSupervisor(dbPath, reader) + s1.routingOverrides.StoreRoutes(42, []routingwire.RouteEntry{ + {SeatID: "seat-1", GroupID: "grp-1", AccountID: "acc-engine"}, + {SeatID: "seat-2", GroupID: "grp-1", Blocked: true}, + }) + s1.persistAssignmentOverrides(s1.active.Load(), 42) + + var pa persistedAssignment + if err := json.Unmarshal([]byte(readOverrideColumn(t, dbPath, "vk-a")), &pa); err != nil { + t.Fatalf("vk-a column not valid JSON: %v", err) + } + if pa.AccountID != "acc-engine" || pa.RoutingVersion != 42 || pa.SyncedAt == 0 { + t.Fatalf("vk-a persisted payload wrong: %+v", pa) + } + var pb persistedAssignment + if err := json.Unmarshal([]byte(readOverrideColumn(t, dbPath, "vk-b")), &pb); err != nil { + t.Fatalf("vk-b column not valid JSON: %v", err) + } + if !pb.Blocked { + t.Fatalf("vk-b must persist the blocked directive: %+v", pb) + } + + // Process 2 (simulated restart): fresh cache, same vault → hydrate. + s2 := newPersistSupervisor(dbPath, reader) + s2.hydrateRoutingOverrides(s2.active.Load()) + + if got := s2.routingOverrides.Assignment("seat-1", "grp-1"); got != "acc-engine" { + t.Fatalf("post-restart assignment=%q want acc-engine (hydrate lost it)", got) + } + if !s2.routingOverrides.Blocked("seat-2", "grp-1") { + t.Fatal("post-restart blocked seat lost — engine cap would be re-admitted") + } + if v := s2.routingOverrides.Version(); v != 42 { + t.Fatalf("hydrated version=%d want 42", v) + } + if !s2.routingOverrides.Stored() { + t.Fatal("hydrate must mark Stored (first-pull version-skip contract)") + } +} + +// A withdrawn override must CLEAR the column, and a later hydrate must not +// resurrect it (§5.3 "hydrate can't resurrect a withdrawn assignment"). +func TestPersistAssignmentOverrides_ClearsWithdrawn(t *testing.T) { + t.Setenv("AIKEY_PROXY_OAUTH_GROUP_ENABLED", "1") + pa, _ := json.Marshal(persistedAssignment{AccountID: "acc-old", RoutingVersion: 41, SyncedAt: 90}) + dbPath, reader := newOpenableVault(t, []map[string]string{ + {"vk": "vk-a", "seat": "seat-1", "group": "grp-1", "override": string(pa)}, + }) + + s := newPersistSupervisor(dbPath, reader) + // Engine now overrides NOTHING (empty route set at a newer version). + s.routingOverrides.StoreRoutes(50, nil) + s.persistAssignmentOverrides(s.active.Load(), 50) + + if got := readOverrideColumn(t, dbPath, "vk-a"); got != "" { + t.Fatalf("withdrawn override must clear the column, still %q", got) + } + // Restart: nothing to hydrate → cache stays unstored (local pick default). + s2 := newPersistSupervisor(dbPath, reader) + s2.hydrateRoutingOverrides(s2.active.Load()) + if s2.routingOverrides.Stored() { + t.Fatal("cleared columns must not hydrate a stale assignment") + } +} + +// Corrupt column content is skipped (never fails hydrate), and a healthy row +// alongside it still hydrates. +func TestHydrateRoutingOverrides_SkipsCorruptRow(t *testing.T) { + t.Setenv("AIKEY_PROXY_OAUTH_GROUP_ENABLED", "1") + good, _ := json.Marshal(persistedAssignment{AccountID: "acc-ok", RoutingVersion: 7, SyncedAt: 9}) + dbPath, reader := newOpenableVault(t, []map[string]string{ + {"vk": "vk-bad", "seat": "seat-9", "group": "grp-1", "override": "{not-json"}, + {"vk": "vk-good", "seat": "seat-1", "group": "grp-1", "override": string(good)}, + }) + s := newPersistSupervisor(dbPath, reader) + s.hydrateRoutingOverrides(s.active.Load()) + if got := s.routingOverrides.Assignment("seat-1", "grp-1"); got != "acc-ok" { + t.Fatalf("healthy row must hydrate despite corrupt sibling, got %q", got) + } + if got := s.routingOverrides.Assignment("seat-9", "grp-1"); got != "" { + t.Fatalf("corrupt row must be skipped, got %q", got) + } +} + +// Steady state writes nothing: persisting the same assignment at the same +// version twice must not churn the vault (synced_at-only diffs are ignored). +func TestPersistAssignmentOverrides_NoChurnOnSteadyState(t *testing.T) { + t.Setenv("AIKEY_PROXY_OAUTH_GROUP_ENABLED", "1") + dbPath, reader := newOpenableVault(t, []map[string]string{ + {"vk": "vk-a", "seat": "seat-1", "group": "grp-1", "override": ""}, + }) + s := newPersistSupervisor(dbPath, reader) + s.routingOverrides.StoreRoutes(42, []routingwire.RouteEntry{ + {SeatID: "seat-1", GroupID: "grp-1", AccountID: "acc-engine"}, + }) + s.persistAssignmentOverrides(s.active.Load(), 42) + first := readOverrideColumn(t, dbPath, "vk-a") + + s.persistAssignmentOverrides(s.active.Load(), 42) + second := readOverrideColumn(t, dbPath, "vk-a") + if first != second { + t.Fatalf("steady-state persist must be byte-identical (no synced_at churn): %q vs %q", first, second) + } +} + +// sameAssignmentPayload ignores synced_at (a timestamp-only diff is churn) but +// distinguishes account / blocked / version changes. +func TestSameAssignmentPayload(t *testing.T) { + a := `{"account_id":"x","routing_version":1,"synced_at":100}` + b := `{"account_id":"x","routing_version":1,"synced_at":999}` + if !sameAssignmentPayload(a, b) { + t.Fatal("timestamp-only diff must compare equal (no write churn)") + } + c := `{"account_id":"y","routing_version":1,"synced_at":100}` + if sameAssignmentPayload(a, c) { + t.Fatal("account change must compare different") + } + d := `{"account_id":"x","routing_version":2,"synced_at":100}` + if sameAssignmentPayload(a, d) { + t.Fatal("version change must compare different") + } + if sameAssignmentPayload("", a) || sameAssignmentPayload(a, "") { + t.Fatal("empty vs non-empty must compare different (set/clear transitions)") + } + if !sameAssignmentPayload("", "") { + t.Fatal("both empty must compare equal") + } + if sameAssignmentPayload("not-json", a) { + t.Fatal("corrupt existing must trigger a rewrite (repair path)") + } +} diff --git a/internal/supervisor/routing_override_policy.go b/internal/supervisor/routing_override_policy.go new file mode 100644 index 0000000..e7a5664 --- /dev/null +++ b/internal/supervisor/routing_override_policy.go @@ -0,0 +1,299 @@ +// routing_override_policy.go — I-side §6.5: the allocation engine's seat→account +// routing-override rail. Pulls GET /accounts/me/routing (account-JWT) and stores +// the sparse assignments in the shared in-memory RoutingOverrideCache the +// group-route resolver reads at request time. +// +// Since 2026-07-03 this rail is DRIVEN BY THE SYNCRAIL FRAMEWORK (railset.go): +// gate, generation, control URL and team credential are re-evaluated every +// cycle, failures are counted into the OK→STALE→OFFLINE visibility state +// machine, and Reload kicks an immediate cycle. The old hand-written loop built +// its credential once at start and early-returned forever on any precondition +// miss — silently (bugfix 2026-07-03-routing-override-rail-silent-stall.md). +// +// WHY proxy (not CLI): the engine recomputes overrides as account health drifts; +// the CLI isn't always running. Same master-poll rail compliance/quota/group-runtime +// already use. +// +// MAIN-LINK SAFETY: on master error / non-200 / timeout / no-auth the cycle +// returns an error WITHOUT touching the cache (keep-last-known, never flap) — +// the framework makes the failure visible, the cache is a pure redirect layer +// the resolver always falls back from. +package supervisor + +import ( + "context" + "encoding/json" + "fmt" + "io" + "log/slog" + "net/http" + "time" + + "github.com/AiKeyLabs/aikey-proxy/internal/httpx" + "github.com/AiKeyLabs/aikey-proxy/internal/vault" + "github.com/AiKeyLabs/pkg/routingwire" +) + +const routingOverridePollInterval = 60 * time.Second + +var routingOverrideHTTPClient = httpx.NewSwappableDirect(10 * time.Second) + +// routingOverrideRail declares this rail for the SyncRail framework. Gate: the +// oauth-group feature is on AND this vault actually has group VKs — a personal +// install without them is idle by design (no failure noise), and the old +// always-pull behavior fetched overrides no resolver would ever read. +func (s *Supervisor) routingOverrideRail() railSpec { + return railSpec{ + name: "routing_override", + interval: routingOverridePollInterval, + needsTeamJWT: true, + gate: func(gen *generation) bool { + return oauthGroupRoutingEnabled() && len(s.localSeatGroupsFor(gen)) > 0 + }, + hydrate: s.hydrateRoutingOverrides, + sync: s.syncRoutingOverrides, + } +} + +// persistedAssignment is the my_assignment_override column payload — the +// engine's last-known assignment for one VK's (seat, group), written on each +// applied routing_version so a proxy restart resumes with the SAME pick the +// engine last issued (N12's persistence leg; without it the restart window fell +// back to the local ranked pick, the 2026-07-03 incident amplifier). +type persistedAssignment struct { + AccountID string `json:"account_id,omitempty"` + Blocked bool `json:"blocked,omitempty"` + // RoutingVersion audits which engine version issued this assignment and + // lets hydrate pick the newest across rows written at different times. + RoutingVersion int64 `json:"routing_version"` + // SyncedAt (unix) makes staleness inspectable (SELECT-level debugging). + SyncedAt int64 `json:"synced_at"` +} + +// hydrateRoutingOverrides preloads the in-memory override cache from the +// persisted my_assignment_override columns — runs once before the rail's first +// pull so the restart window serves the engine's last-known assignments instead +// of falling back to the local ranked pick (§5.3). A later successful pull is +// authoritative and overwrites both the cache and the columns. +func (s *Supervisor) hydrateRoutingOverrides(gen *generation) { + if !oauthGroupRoutingEnabled() || s.routingOverrides == nil || s.routingOverrides.Stored() { + return + } + mks, err := gen.vault.GetActiveManagedKeys() + if err != nil { + return // best-effort: the first pull populates the cache anyway + } + // Rebuild wire entries and ingest through StoreRoutes — the ONLY public key + // builder — so the (seat,group) cache-key encoding stays private to the + // proxy package (its file-doc contract). + var entries []routingwire.RouteEntry + var version int64 + for i := range mks { + mk := mks[i] + if mk.OauthGroupID == "" || mk.SeatID == "" || mk.MyAssignmentOverride == "" { + continue + } + var pa persistedAssignment + if json.Unmarshal([]byte(mk.MyAssignmentOverride), &pa) != nil { + continue // corrupt row: skip it, never fail the hydrate + } + if !pa.Blocked && pa.AccountID == "" { + continue + } + entries = append(entries, routingwire.RouteEntry{ + SeatID: mk.SeatID, GroupID: mk.OauthGroupID, + AccountID: pa.AccountID, Blocked: pa.Blocked, + }) + if pa.RoutingVersion > version { + version = pa.RoutingVersion + } + } + if len(entries) == 0 { + return + } + s.routingOverrides.StoreRoutes(version, entries) + slog.Info("routing overrides hydrated from vault (last-known, pre-pull)", + "event.name", "proxy.routing_override.hydrated", + "routing_version", version, "entries", len(entries)) +} + +// persistAssignmentOverrides mirrors the freshly-stored cache into each group +// VK's my_assignment_override column. Write-on-change only (steady state is +// zero writes); a VK the engine no longer overrides gets its column cleared so +// hydrate can't resurrect a withdrawn assignment. The live pull is +// authoritative — it overwrites regardless of version direction (a reinstalled +// master may legitimately restart its version counter); RoutingVersion in the +// payload exists for hydrate ordering and auditability, not as a write gate. +func (s *Supervisor) persistAssignmentOverrides(gen *generation, version int64) { + mks, err := gen.vault.GetActiveManagedKeys() + if err != nil { + return + } + now := time.Now().Unix() + for i := range mks { + mk := mks[i] + if mk.OauthGroupID == "" || mk.SeatID == "" { + continue + } + acct := s.routingOverrides.Assignment(mk.SeatID, mk.OauthGroupID) + blk := s.routingOverrides.Blocked(mk.SeatID, mk.OauthGroupID) + desired := "" + if acct != "" || blk { + b, mErr := json.Marshal(persistedAssignment{ + AccountID: acct, Blocked: blk, RoutingVersion: version, SyncedAt: now, + }) + if mErr != nil { + continue + } + desired = string(b) + } + if sameAssignmentPayload(mk.MyAssignmentOverride, desired) { + continue // steady state: same assignment at same version → no write + } + if err := vault.WriteAssignmentOverride(s.cfg.Vault.Path, mk.VirtualKeyID, desired); err != nil { + // Persistence is an enhancement, not a dependency: the in-memory cache + // already serves this cycle; WARN (mandatory-visibility) and move on. + slog.Warn("my_assignment_override write failed — restart hydrate will miss this assignment", + "event.name", "proxy.routing_override.persist_failed", + "virtual_key_id", mk.VirtualKeyID, "error", err.Error()) + } + } +} + +// sameAssignmentPayload compares two persisted payloads ignoring synced_at (a +// re-write that only bumps the timestamp is churn, not information). +func sameAssignmentPayload(existing, desired string) bool { + if existing == desired { + return true + } + if existing == "" || desired == "" { + return false + } + var a, b persistedAssignment + if json.Unmarshal([]byte(existing), &a) != nil || json.Unmarshal([]byte(desired), &b) != nil { + return false + } + return a.AccountID == b.AccountID && a.Blocked == b.Blocked && a.RoutingVersion == b.RoutingVersion +} + +// syncRoutingOverrides runs one pull. Only when the routing_version actually +// advanced does it replace the shared cache; a fetch error keeps the last-known +// assignments in place (don't flap) and is COUNTED by the framework (no more +// silent returns). Mirrors syncGroupRuntime's version-skip + keep-last-known. +func (s *Supervisor) syncRoutingOverrides(ctx context.Context, gen *generation, masterURL, bearer string) error { + if s.routingOverrides == nil { + return nil + } + version, routes, err := fetchRoutingOverrides(ctx, masterURL, bearer) + if err != nil { + return err // keep last-known; framework counts + surfaces + } + // Skip only once we've ALREADY stored at this version. The Stored() guard + // closes the first-pull-at-version-0 hole: the cache's version atomic starts + // at 0, so without it master's first non-empty assignments carrying + // routing_version:0 would match Version()==0 and be skipped forever (the + // override never applied, no signal). After the first Store, steady-state + // re-pulls of the same version skip as before (no churn, no log spam). + if s.routingOverrides.Stored() && s.routingOverrides.Version() == version { + return nil // unchanged → no churn (a successful no-op cycle IS a success) + } + bound, blocked := s.routingOverrides.StoreRoutes(version, routes) + slog.Info("routing overrides updated", + "event.name", "proxy.routing_override.changed", + "routing_version", version, "route_entries", bound, "blocked_entries", blocked) + // Format-mismatch fingerprint (2026-07-02, review F1): a NON-EMPTY route set in + // which not a single entry matches any of this vault's local (seat,group) pairs + // means the two sides disagree about the wire (or the payload is for the wrong + // account) — a failure that is otherwise indistinguishable from "no overrides". + // WARN once per routing_version (the ticker would repeat it every 60s otherwise). + // Mandatory-WARN rule: fallback-to-default paths must log. + if len(routes) > 0 && !routesMatchAnyLocal(routes, s.localSeatGroupsFor(gen)) { + if s.lastRoutingMismatchVersion.Swap(version) != version { + slog.Warn("routing overrides matched no local group VK — wire format / account mismatch?", + "event.name", "proxy.routing_override.format_mismatch", + "routing_version", version, "route_entries", len(routes)) + } + } + // C2 display coupling (owner-approved 2026-06-30): the engine just redirected one + // or more seats. Re-stamp IsCurrentRouted on the affected group VKs' group_runtime + // so /user/vault shows the new routed account within this override poll, not only on + // the next material refresh. Display-only + network-free (RMW on the cached column); + // the routing redirect itself already took effect via the override cache above. + s.restampCurrentRouted() + // N12 persistence leg (2026-07-03): mirror the applied assignments to the + // my_assignment_override columns so a restart hydrates the same picks (§5.3). + s.persistAssignmentOverrides(gen, version) + return nil +} + +// fetchRoutingOverrides GETs the routing endpoint with an account-JWT Bearer and +// decodes the SHARED wire DTO (pkg/routingwire — the same structs master emits, so +// the two repos cannot drift). Returns a non-nil error on ANY failure so the +// caller keeps the last-known cache AND the framework can surface the underlying +// cause (a bare ok=false hid the connection-refused detail that mattered most in +// the 2026-07-03 incident). nil routes normalize to an empty slice so a +// successful pull always carries a concrete (possibly empty) set. +func fetchRoutingOverrides(ctx context.Context, masterURL, bearer string) (int64, []routingwire.RouteEntry, error) { + req, err := http.NewRequestWithContext(ctx, http.MethodGet, masterURL+"/accounts/me/routing", http.NoBody) + if err != nil { + return 0, nil, err + } + req.Header.Set("Authorization", "Bearer "+bearer) + resp, err := routingOverrideHTTPClient.Get().Do(req) + if err != nil { + return 0, nil, err + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + return 0, nil, fmt.Errorf("GET /accounts/me/routing: HTTP %d", resp.StatusCode) + } + body, err := io.ReadAll(io.LimitReader(resp.Body, 1<<20)) + if err != nil { + return 0, nil, err + } + var out routingwire.RoutingResponse + if err := json.Unmarshal(body, &out); err != nil { + return 0, nil, fmt.Errorf("decode routing response: %w", err) + } + if out.Routes == nil { + out.Routes = []routingwire.RouteEntry{} + } + return out.RoutingVersion, out.Routes, nil +} + +// routesMatchAnyLocal reports whether at least one wire entry targets one of this +// vault's local (seat, group) pairs. Pure — unit-tested as the format-mismatch +// fingerprint's decision core. +func routesMatchAnyLocal(routes []routingwire.RouteEntry, local map[[2]string]bool) bool { + for _, r := range routes { + if local[[2]string{r.SeatID, r.GroupID}] { + return true + } + } + return false +} + +// localSeatGroupsFor collects the (seat, group) pairs of a generation's group VKs +// — the set a healthy routing payload should intersect, and this rail's gate +// ("is there anything to route locally?"). Empty on any read problem. +func (s *Supervisor) localSeatGroupsFor(gen *generation) map[[2]string]bool { + if gen == nil || gen.vault == nil { + return nil + } + mks, err := gen.vault.GetActiveManagedKeys() + if err != nil { + return nil + } + out := map[[2]string]bool{} + for i := range mks { + if mks[i].OauthGroupID != "" && mks[i].SeatID != "" { + out[[2]string{mks[i].SeatID, mks[i].OauthGroupID}] = true + } + } + return out +} + +// localSeatGroups keeps the historical accessor shape (current generation). +func (s *Supervisor) localSeatGroups() map[[2]string]bool { + return s.localSeatGroupsFor(s.active.Load()) +} diff --git a/internal/supervisor/routing_override_policy_test.go b/internal/supervisor/routing_override_policy_test.go new file mode 100644 index 0000000..46b6956 --- /dev/null +++ b/internal/supervisor/routing_override_policy_test.go @@ -0,0 +1,126 @@ +package supervisor + +import ( + "context" + "net/http" + "net/http/httptest" + "testing" + + "github.com/AiKeyLabs/aikey-proxy/internal/proxy" + "github.com/AiKeyLabs/pkg/routingwire" +) + +// Mirrors TestFetchGroupRuntime_ParsesAndSendsBearer: the routing-override pull +// decodes the SHARED routingwire wire (structured route entries), sends the +// account-JWT bearer, and yields ok=false on a non-200 so the caller keeps the +// last-known cache. +func TestFetchRoutingOverrides_ParsesAndSendsBearer(t *testing.T) { + var gotAuth string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotAuth = r.Header.Get("Authorization") + if r.URL.Path != "/accounts/me/routing" { + w.WriteHeader(404) + return + } + _, _ = w.Write([]byte(`{"routing_version":42,"routes":[ + {"seat_id":"seat-1","group_id":"g1","account_id":"acc-9"}, + {"seat_id":"seat-1","group_id":"g2","account_id":"acc-3"}, + {"seat_id":"seat-9","group_id":"g1","blocked":true}]}`)) + })) + defer srv.Close() + + version, routes, err := fetchRoutingOverrides(context.Background(), srv.URL, "JWT123") + if err != nil || version != 42 || len(routes) != 3 { + t.Fatalf("fetch: err=%v version=%d routes=%d", err, version, len(routes)) + } + if routes[0].AccountID != "acc-9" || routes[1].GroupID != "g2" || !routes[2].Blocked { + t.Fatalf("routes not parsed: %+v", routes) + } + if gotAuth != "Bearer JWT123" { + t.Fatalf("bearer not sent: %q", gotAuth) + } + + // The cache ingest keeps multi-pool entries distinct and flags blocked pairs. + cache := proxy.NewRoutingOverrideCache() + bound, blocked := cache.StoreRoutes(version, routes) + if bound != 2 || blocked != 1 { + t.Fatalf("StoreRoutes counts: bound=%d blocked=%d", bound, blocked) + } + if got := cache.Assignment("seat-1", "g1"); got != "acc-9" { + t.Fatalf("seat-1/g1 override: %q", got) + } + if got := cache.Assignment("seat-1", "g2"); got != "acc-3" { + t.Fatalf("seat-1/g2 override (multi-pool distinct): %q", got) + } + if !cache.Blocked("seat-9", "g1") { + t.Fatal("seat-9/g1 must be blocked") + } + + // Empty routes (engine redirects nothing) → ok=true, empty slice. + empty := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _, _ = w.Write([]byte(`{"routing_version":7}`)) + })) + defer empty.Close() + v, rts, err := fetchRoutingOverrides(context.Background(), empty.URL, "x") + if err != nil || v != 7 || rts == nil || len(rts) != 0 { + t.Fatalf("empty-routes pull: err=%v v=%d routes=%+v", err, v, rts) + } + + // Non-200 → error (caller keeps last-known; framework surfaces the cause). + bad := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(401) + })) + defer bad.Close() + if _, _, err := fetchRoutingOverrides(context.Background(), bad.URL, "x"); err == nil { + t.Fatal("401 must yield a non-nil error") + } +} + +// A failed pull must keep the last-known cache (never clear/flap): we Store a v1 +// assignment, then simulate a poll that does NOT re-Store (the ok=false path), and +// assert the prior assignment is still served. This pins the keep-last-known +// contract syncRoutingOverrides relies on. +func TestRoutingOverride_KeepLastKnownOnFailure(t *testing.T) { + cache := proxy.NewRoutingOverrideCache() + cache.StoreRoutes(1, []routingwire.RouteEntry{{SeatID: "seat-1", GroupID: "g1", AccountID: "acc-1"}}) + + // fetch against a dead endpoint → error → syncRoutingOverrides returns it + // WITHOUT touching the cache. Assert by simulating that control flow. + _, _, err := fetchRoutingOverrides(context.Background(), "http://127.0.0.1:0", "x") + if err == nil { + t.Fatal("unreachable endpoint must yield a non-nil error") + } + // Cache untouched. + if v := cache.Version(); v != 1 { + t.Fatalf("version must stay last-known 1, got %d", v) + } +} + +// routesMatchAnyLocal is the decision core of the proxy.routing_override. +// format_mismatch WARN (review F1: a wire/format drift makes every entry miss the +// local (seat,group) set — otherwise indistinguishable from "no overrides"). 能红: +// break the pair comparison → the mismatch case below stops returning false → red. +func TestRoutesMatchAnyLocal(t *testing.T) { + local := map[[2]string]bool{{"seat-1", "g1"}: true, {"seat-2", "g1"}: true} + + if !routesMatchAnyLocal([]routingwire.RouteEntry{{SeatID: "seat-1", GroupID: "g1", AccountID: "a"}}, local) { + t.Fatal("an entry targeting a local (seat,group) must match") + } + if !routesMatchAnyLocal([]routingwire.RouteEntry{ + {SeatID: "ghost", GroupID: "gX", AccountID: "a"}, + {SeatID: "seat-2", GroupID: "g1", Blocked: true}, + }, local) { + t.Fatal("a blocked entry targeting a local pair must match too") + } + // The F1 fingerprint: non-empty payload, zero local intersection → mismatch. + if routesMatchAnyLocal([]routingwire.RouteEntry{ + {SeatID: "seat-1", GroupID: "OTHER-GROUP", AccountID: "a"}, + {SeatID: "unknown-seat", GroupID: "g1", AccountID: "b"}, + }, local) { + t.Fatal("entries matching no local (seat,group) must report mismatch (→ WARN)") + } + // No local group VKs → nothing to mismatch (WARN stays silent). + if routesMatchAnyLocal([]routingwire.RouteEntry{{SeatID: "s", GroupID: "g", AccountID: "a"}}, nil) { + t.Fatal("empty local set must not match") + } +} diff --git a/internal/supervisor/selfheal.go b/internal/supervisor/selfheal.go new file mode 100644 index 0000000..1f52db3 --- /dev/null +++ b/internal/supervisor/selfheal.go @@ -0,0 +1,253 @@ +// selfheal.go — control-plane self-healing after a host NETWORK CHANGE. +// +// PROBLEM (2026-07-01, field-diagnosed): the aikey-proxy is a long-running +// daemon. When the host network changes UNDER a running proxy — a WiFi switch, +// an interface IP change, or a VPN/proxy (e.g. Clash TUN) starting up AFTER the +// proxy — the proxy's outbound control-plane calls to the LAN team-master can get +// stuck failing with "no route to host" / EHOSTUNREACH, and STAY stuck until the +// PROCESS is restarted. A freshly-launched process (curl) reaches the same master +// fine, which proves the stale state is process-scoped, not a real outage. +// +// INDUSTRY-ALIGNED FIX (deep-research 2026-07-01): production network daemons +// (Tailscale's link monitor + socket rebind is the reference) recover by +// (1) REBUILDING their client/transport on a network change, and (2) keeping a +// service-manager RESTART as a guarded backstop. Crucially, Go's +// Transport.CloseIdleConnections() alone does NOT reliably clear the stale state +// (golang/go#23427), so we rebuild a FRESH *http.Client (new dialer + empty pool) +// rather than just closing idle conns. +// +// This file implements the failure-triggered tier (Stage 1): on a routing-layer +// error to master we (Tier1) swap in a fresh control-plane client, and if that +// still doesn't clear it AND a fresh out-of-process probe proves master is +// actually reachable (i.e. WE are the stale one, not the network), we (Tier3) +// exit non-zero so launchd/systemd relaunches a clean process — guarded by a +// cooldown + a circuit breaker so a genuine outage can't cause a restart storm. +// +// MAIN-LINK SAFETY: this is control-plane only. It never touches the request +// forwarding path, and Tier3 only fires when master is confirmed reachable, so a +// real upstream/GFW outage (which the env proxy egress handles) can't trip it. +package supervisor + +import ( + "errors" + "net" + "net/http" + "strings" + "sync" + "syscall" + "time" + + "github.com/AiKeyLabs/aikey-proxy/internal/httpx" +) + +// isNetChangeDialErr reports whether err is a routing-layer dial failure of the +// kind a host network change produces — "no route to host" (EHOSTUNREACH), +// "network is unreachable" (ENETUNREACH), or "network is down" (ENETDOWN). These +// are the errno's the kernel returns when a long-running process's routing/source +// state has gone stale relative to the current interfaces. +// +// WHY string-match as well as errors.Is: syscall errno constants are defined +// per-GOOS (Windows uses WSAE* values), so a portable classifier can't rely on +// errors.Is(err, syscall.EHOSTUNREACH) alone on every platform. We check the +// unix errno's where they exist (cheap, exact) AND fall back to the stable +// error-string tails Go formats them into, which are identical across platforms. +func isNetChangeDialErr(err error) bool { + if err == nil { + return false + } + // Only dial-time failures qualify — a mid-stream reset is a different animal + // (a fresh client wouldn't necessarily help there). net.OpError{Op:"dial"} + // wraps the syscall error for connect() failures. + var opErr *net.OpError + if errors.As(err, &opErr) && opErr.Op != "dial" { + return false + } + if errors.Is(err, syscall.EHOSTUNREACH) || + errors.Is(err, syscall.ENETUNREACH) || + errors.Is(err, syscall.ENETDOWN) { + return true + } + s := strings.ToLower(err.Error()) + return strings.Contains(s, "no route to host") || + strings.Contains(s, "network is unreachable") || + strings.Contains(s, "host is unreachable") || + strings.Contains(s, "network is down") +} + +// controlPlaneHealer holds the cooldown + circuit-breaker state that decides +// whether a persistent no-route condition should escalate to a process restart. +// It is intentionally tiny and pure (no I/O): callers feed it "we just hit a +// net-change error and a fresh probe says master IS reachable" and it answers +// whether to restart now, applying: +// +// - cooldown: at most one restart per restartCooldown (a fresh process needs +// time to re-establish; back-to-back restarts help nothing). +// - circuit breaker: at most restartBudget restarts per breakerWindow; once +// tripped it stops restarting and the caller logs loudly, so a mis-detection +// or a genuinely broken host can't spin the process forever +// (systemd StartLimitBurst analogue; Azure retry-storm antipattern). +// +// `now` and `exit` are injected so the decision + effect are unit-testable +// without a clock or an actual os.Exit. +type controlPlaneHealer struct { + mu sync.Mutex + restarts []time.Time // timestamps of restarts inside the current window + last time.Time // last restart (for cooldown) + consecutive int // consecutive poll cycles failing with a net-change error + + // tunables (vars-in-struct so tests can shrink them) + threshold int // escalate to a restart only after this many consecutive failures + cooldown time.Duration + window time.Duration + budget int + + now func() time.Time + exit func(code int) + probe func(masterURL string) bool // "is master reachable from a FRESH client?" +} + +func newControlPlaneHealer() *controlPlaneHealer { + return &controlPlaneHealer{ + threshold: 3, // ~3 poll cycles (~3 min) stuck before we consider a restart + cooldown: 2 * time.Minute, + window: 30 * time.Minute, + budget: 4, + now: time.Now, + // Production: request a GRACEFUL restart — main() runs the same + // srv.Shutdown + sup.Shutdown drain path as SIGTERM, then exits non-zero + // so launchd KeepAlive{SuccessfulExit:false} / systemd Restart=on-failure + // relaunch a clean process. In-flight forwarded requests drain first + // (up to the 30s/streaming timeout) instead of being cut. The `code` arg + // is ignored here (main owns the non-zero exit); tests inject their own + // exit to assert the decision without draining. + exit: func(int) { requestGracefulRestart() }, + probe: func(masterURL string) bool { return masterReachableFresh(defaultGroupRuntimeClient(), masterURL) }, + } +} + +// restartRequestCh carries a graceful-self-restart request from the control-plane +// healer to main(). Buffered + non-blocking send so repeated requests (a stuck +// streak firing every poll) coalesce into one; the cooldown in shouldRestart also +// prevents re-signalling within the window. +var restartRequestCh = make(chan struct{}, 1) + +// RestartRequested is selected by main() alongside the OS signal channel: it runs +// the SAME graceful drain as SIGTERM, then exits non-zero to trigger a relaunch. +func (s *Supervisor) RestartRequested() <-chan struct{} { return restartRequestCh } + +// requestGracefulRestart asks main() to drain + relaunch. Idempotent. +func requestGracefulRestart() { + select { + case restartRequestCh <- struct{}{}: + default: + } +} + +// controlPlaneHeal is the process-wide healer driven by the group-runtime poll +// (the steady 60s heartbeat). The writeback path only does the cheap Tier-1 +// client rebuild; escalation to a restart is owned here so a 6-try writeback +// burst can't inflate the counter. +var controlPlaneHeal = newControlPlaneHealer() + +// onPollOK resets the consecutive-failure counter after any successful poll. +func (h *controlPlaneHealer) onPollOK() { + h.mu.Lock() + h.consecutive = 0 + h.mu.Unlock() +} + +// onPollNetChange handles a routing-layer poll failure: it always rebuilds the +// shared client (Tier1), and once we've failed `threshold` consecutive poll +// cycles it escalates — but ONLY if a FRESH client can reach master (proving the +// staleness is ours, not a real outage) AND the cooldown/budget allow — by +// exiting non-zero so the service manager relaunches a clean process (Tier3). +// Returns the decision (for logging + tests); callers log with their vocabulary. +func (h *controlPlaneHealer) onPollNetChange(masterURL string) restartDecision { + httpx.RebuildAllControlPlane() // Tier1: always cheap + safe + h.mu.Lock() + h.consecutive++ + n := h.consecutive + h.mu.Unlock() + if n < h.threshold { + return restartSkipCooldown // not escalating yet (reuse enum: "not now") + } + // Tier3 gate: is it US (stale) or the network (real outage)? A fresh client + // reaching master means we're the stale one → a restart will clear it. + if !h.probe(masterURL) { + return restartSkipCooldown // genuine outage — never restart into a dead net + } + d := h.shouldRestart() + if d == restartNow { + h.exitNow() + } + return d +} + +// restartDecision is the outcome of asking the healer to escalate. It's returned +// (instead of restarting inline) so the caller can log the reason with its own +// event vocabulary before the process goes down. +type restartDecision int + +const ( + restartSkipCooldown restartDecision = iota // too soon since the last restart + restartSkipBreaker // budget exhausted — stop restarting, log loudly + restartNow // clear to restart +) + +// shouldRestart records the intent and returns whether a restart is warranted +// right now. The caller must have already confirmed (out-of-process) that master +// is reachable — this method only enforces the rate limits. On restartNow it +// stamps the clock (so the cooldown starts) but does NOT exit; call exitNow to +// actually go down, keeping "decide" and "act" separable for logging + tests. +func (h *controlPlaneHealer) shouldRestart() restartDecision { + h.mu.Lock() + defer h.mu.Unlock() + t := h.now() + if !h.last.IsZero() && t.Sub(h.last) < h.cooldown { + return restartSkipCooldown + } + // prune restarts older than the window, then check the budget + cut := t.Add(-h.window) + kept := h.restarts[:0] + for _, r := range h.restarts { + if r.After(cut) { + kept = append(kept, r) + } + } + h.restarts = kept + if len(h.restarts) >= h.budget { + return restartSkipBreaker + } + h.restarts = append(h.restarts, t) + h.last = t + return restartNow +} + +// exitNow performs the actual restart (via the injected exit). Split from +// shouldRestart so the caller can log first and tests can assert the code. +func (h *controlPlaneHealer) exitNow() { h.exit(75) } // EX_TEMPFAIL + +// masterReachableFresh probes masterURL from a BRAND-NEW *http.Client (fresh +// transport + dialer + empty pool). It is the "is it US or the network?" check: +// if this fresh client reaches master while the long-lived control-plane client +// keeps getting no-route, the staleness is ours and a restart will clear it; if +// this ALSO fails, it's a real outage and we must NOT restart. +// +// NOTE (honest): a fresh client in the SAME process is a weaker signal than a +// fresh PROCESS — if the staleness turns out to be process-global rather than +// transport-local, this probe could also succeed-or-fail with the live client. +// It is still the best in-process discriminator available and errs safe: any +// probe error (including net-change errors) returns false → we do NOT restart. +func masterReachableFresh(client *http.Client, masterURL string) bool { + req, err := http.NewRequest(http.MethodHead, strings.TrimRight(masterURL, "/")+"/", nil) + if err != nil { + return false + } + resp, err := client.Do(req) + if err != nil { + return false + } + _ = resp.Body.Close() + // Any HTTP response (even 404/405) proves the TCP+routing path is alive. + return true +} diff --git a/internal/supervisor/selfheal_integration_test.go b/internal/supervisor/selfheal_integration_test.go new file mode 100644 index 0000000..39c717e --- /dev/null +++ b/internal/supervisor/selfheal_integration_test.go @@ -0,0 +1,108 @@ +package supervisor + +import ( + "context" + "net" + "net/http" + "os" + "syscall" + "testing" + "time" + + "github.com/AiKeyLabs/aikey-proxy/internal/httpx" +) + +// These are INTEGRATION tests: they drive the REAL production self-heal wiring +// (the real onNetworkChange callback, the real postMemberToken retry loop, the +// real isNetChangeDialErr classifier, the real httpx control-plane registry) and +// assert observable effects — no copied/re-implemented logic. Only the trigger +// (a network-change fingerprint flip / a routing-layer dial error) is injected. + +// netChangeRoundTripper makes every request fail with a genuine routing-layer +// dial error — the exact error type the kernel produces after a host network +// change — so the REAL postMemberToken path classifies it and self-heals. +type netChangeRoundTripper struct{} + +func (netChangeRoundTripper) RoundTrip(*http.Request) (*http.Response, error) { + return nil, &net.OpError{Op: "dial", Net: "tcp", + Err: os.NewSyscallError("connect", syscall.EHOSTUNREACH)} +} + +// TestIntegration_NetChangeMonitor_RebuildsAllRegisteredClients drives the REAL +// watchNetworkChanges loop with the REAL onNetworkChange callback and asserts the +// central registry rebuilt EVERY registered control-plane client (group-runtime + +// an extra registered probe client), and reset the self-heal streak. +func TestIntegration_NetChangeMonitor_RebuildsAllRegisteredClients(t *testing.T) { + // A second registered control-plane client, to prove RebuildAll hits more than + // just group-runtime (i.e. the registry, not a single client). + probe := httpx.NewSwappable(func() *http.Client { return &http.Client{} }) + + grBefore := groupRuntimeClient() + probeBefore := probe.Get() + + // Prime the healer streak so we can prove onNetworkChange resets it. + controlPlaneHeal.mu.Lock() + controlPlaneHeal.consecutive = 2 + controlPlaneHeal.mu.Unlock() + + // Drive the REAL monitor loop: a fingerprint that flips once (A→B) triggers the + // REAL onNetworkChange (rebuild-all + reset). fp returns A, then B forever. + calls := 0 + fp := func() string { + calls++ + if calls <= 1 { + return "A" + } + return "B" + } + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + go watchNetworkChanges(ctx, time.Millisecond, fp, onNetworkChange) + + // Wait for the rebuild to land (both clients swapped). + deadline := time.After(3 * time.Second) + for { + if groupRuntimeClient() != grBefore && probe.Get() != probeBefore { + break + } + select { + case <-deadline: + t.Fatal("onNetworkChange did not rebuild all registered clients in time") + case <-time.After(5 * time.Millisecond): + } + } + + controlPlaneHeal.mu.Lock() + streak := controlPlaneHeal.consecutive + controlPlaneHeal.mu.Unlock() + if streak != 0 { + t.Fatalf("onNetworkChange did not reset the self-heal streak: consecutive=%d, want 0", streak) + } +} + +// TestIntegration_Writeback_SelfHealsOnNetChange drives the REAL postMemberToken +// retry loop with a client whose transport returns a genuine EHOSTUNREACH. It +// asserts the real path classified it and rebuilt the registry (a registered +// client instance changed) while still surfacing the final error. +func TestIntegration_Writeback_SelfHealsOnNetChange(t *testing.T) { + // Speed up the 6-attempt backoff for the test. + orig := writebackBaseBackoff + writebackBaseBackoff = time.Millisecond + defer func() { writebackBaseBackoff = orig }() + + probe := httpx.NewSwappable(func() *http.Client { return &http.Client{} }) + probeBefore := probe.Get() + + // The writeback's client always returns a net-change dial error. + clientFn := func() *http.Client { return &http.Client{Transport: netChangeRoundTripper{}} } + + err := postMemberToken(context.Background(), clientFn, "http://master.invalid", "JWT", + memberTokenWriteback{CredentialID: "c1", AccessToken: "tok"}) + if err == nil { + t.Fatal("postMemberToken returned nil despite persistent no-route errors") + } + // The REAL net-change branch must have invoked the REAL registry rebuild. + if probe.Get() == probeBefore { + t.Fatal("postMemberToken did not rebuild control-plane clients on the net-change dial error") + } +} diff --git a/internal/supervisor/selfheal_test.go b/internal/supervisor/selfheal_test.go new file mode 100644 index 0000000..bc6c0c7 --- /dev/null +++ b/internal/supervisor/selfheal_test.go @@ -0,0 +1,168 @@ +package supervisor + +import ( + "errors" + "net" + "net/http" + "net/http/httptest" + "os" + "syscall" + "testing" + "time" + + "github.com/AiKeyLabs/aikey-proxy/internal/httpx" +) + +func TestIsNetChangeDialErr(t *testing.T) { + cases := []struct { + name string + err error + want bool + }{ + {"nil", nil, false}, + {"dial EHOSTUNREACH", &net.OpError{Op: "dial", Err: os.NewSyscallError("connect", syscall.EHOSTUNREACH)}, true}, + {"dial ENETUNREACH", &net.OpError{Op: "dial", Err: os.NewSyscallError("connect", syscall.ENETUNREACH)}, true}, + {"string no route to host", errors.New(`Post "http://x/y": dial tcp 1.2.3.4:3000: connect: no route to host`), true}, + {"string network is unreachable", errors.New("dial tcp: connect: network is unreachable"), true}, + {"read op (not dial) is not a net-change", &net.OpError{Op: "read", Err: syscall.EHOSTUNREACH}, false}, + {"connection refused is not net-change", errors.New("dial tcp 1.2.3.4:3000: connect: connection refused"), false}, + {"timeout is not net-change", errors.New("context deadline exceeded"), false}, + } + for _, c := range cases { + t.Run(c.name, func(t *testing.T) { + if got := isNetChangeDialErr(c.err); got != c.want { + t.Fatalf("isNetChangeDialErr(%v) = %v, want %v", c.err, got, c.want) + } + }) + } +} + +// newTestHealer builds a healer with tiny/injected knobs — a fake clock, a +// recorded exit, and a scripted reachability probe — so the escalation logic is +// deterministic without a real os.Exit or a real clock. +func newTestHealer(reachable bool) (*controlPlaneHealer, *int, *time.Time) { + exits := 0 + clk := time.Unix(1_700_000_000, 0) + h := &controlPlaneHealer{ + threshold: 3, + cooldown: 2 * time.Minute, + window: 30 * time.Minute, + budget: 4, + now: func() time.Time { return clk }, + exit: func(int) { exits++ }, + probe: func(string) bool { return reachable }, + } + return h, &exits, &clk +} + +func TestHealer_BelowThreshold_NoRestart(t *testing.T) { + h, exits, _ := newTestHealer(true) + for i := 0; i < h.threshold-1; i++ { + if d := h.onPollNetChange("http://master:3000"); d == restartNow { + t.Fatalf("restarted before threshold at i=%d", i) + } + } + if *exits != 0 { + t.Fatalf("exit called %d times below threshold, want 0", *exits) + } +} + +func TestHealer_AtThreshold_ReachableMaster_Restarts(t *testing.T) { + h, exits, _ := newTestHealer(true) // fresh probe says master IS reachable → it's us + var last restartDecision + for i := 0; i < h.threshold; i++ { + last = h.onPollNetChange("http://master:3000") + } + if last != restartNow { + t.Fatalf("decision at threshold = %v, want restartNow", last) + } + if *exits != 1 { + t.Fatalf("exit called %d times, want 1", *exits) + } +} + +func TestHealer_MasterUnreachable_NeverRestarts(t *testing.T) { + h, exits, _ := newTestHealer(false) // real outage: fresh probe ALSO fails + for i := 0; i < h.threshold+5; i++ { + if d := h.onPollNetChange("http://master:3000"); d == restartNow { + t.Fatalf("restarted during a genuine outage at i=%d", i) + } + } + if *exits != 0 { + t.Fatalf("exit called %d times during outage, want 0", *exits) + } +} + +func TestHealer_PollOK_ResetsCounter(t *testing.T) { + h, exits, _ := newTestHealer(true) + h.onPollNetChange("http://master:3000") + h.onPollNetChange("http://master:3000") + h.onPollOK() // master came back → counter resets + // two more failures should NOT reach the threshold now + h.onPollNetChange("http://master:3000") + if d := h.onPollNetChange("http://master:3000"); d == restartNow { + t.Fatal("restarted after onPollOK reset — counter did not clear") + } + if *exits != 0 { + t.Fatalf("exit called %d times, want 0", *exits) + } +} + +func TestHealer_Cooldown_And_Breaker(t *testing.T) { + h, exits, clk := newTestHealer(true) + fire := func() restartDecision { // drive one full escalation (threshold cycles) + h.onPollOK() // reset counter between escalations + var d restartDecision + for i := 0; i < h.threshold; i++ { + d = h.onPollNetChange("http://master:3000") + } + return d + } + // 1st escalation → restarts + if d := fire(); d != restartNow { + t.Fatalf("1st = %v, want restartNow", d) + } + // immediately again (same clock) → cooldown blocks it + if d := fire(); d != restartSkipCooldown { + t.Fatalf("2nd within cooldown = %v, want restartSkipCooldown", d) + } + // advance past cooldown 3 more times → restarts up to the budget (4 total) + for n := 2; n <= h.budget; n++ { + *clk = clk.Add(h.cooldown + time.Second) + if d := fire(); d != restartNow { + t.Fatalf("escalation %d past cooldown = %v, want restartNow", n, d) + } + } + // budget now exhausted → breaker trips even past cooldown + *clk = clk.Add(h.cooldown + time.Second) + if d := fire(); d != restartSkipBreaker { + t.Fatalf("past budget = %v, want restartSkipBreaker", d) + } + if *exits != h.budget { + t.Fatalf("exit called %d times, want %d (budget)", *exits, h.budget) + } +} + +func TestMasterReachableFresh(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusNotFound) // even a 404 proves the path is alive + })) + defer srv.Close() + if !masterReachableFresh(srv.Client(), srv.URL) { + t.Fatal("reachable server reported unreachable") + } + // A dead address must report unreachable (short timeout so the test is fast). + if masterReachableFresh(&http.Client{Timeout: time.Second}, "http://127.0.0.1:1") { + t.Fatal("dead address reported reachable") + } +} + +func TestRebuildAllControlPlane_SwapsGroupRuntimeClient(t *testing.T) { + before := groupRuntimeClient() + if n := httpx.RebuildAllControlPlane(); n == 0 { + t.Fatal("RebuildAllControlPlane rebuilt 0 clients — group-runtime not registered?") + } + if after := groupRuntimeClient(); before == after { + t.Fatal("RebuildAllControlPlane did not install a new group-runtime client") + } +} diff --git a/internal/supervisor/supervisor.go b/internal/supervisor/supervisor.go index f7868ed..f067129 100644 --- a/internal/supervisor/supervisor.go +++ b/internal/supervisor/supervisor.go @@ -43,6 +43,7 @@ import ( "github.com/AiKeyLabs/aikey-proxy/internal/cluster" "github.com/AiKeyLabs/aikey-proxy/internal/config" "github.com/AiKeyLabs/aikey-proxy/internal/events" + "github.com/AiKeyLabs/aikey-proxy/internal/httpx" "github.com/AiKeyLabs/aikey-proxy/internal/observability" "github.com/AiKeyLabs/aikey-proxy/internal/observer/conversation_audit" "github.com/AiKeyLabs/aikey-proxy/internal/provider" @@ -189,6 +190,13 @@ func (g *generation) close() { if g.reporter != nil { _ = g.reporter.Close() } + // Stop the allocation-engine signal reporter (lives on this generation's + // proxy, started per generation in buildGeneration via EnableSignalReporting). + // Without this its loop() goroutine + 30s ticker leak on every reload, each + // holding a live bearer closure over the old vault reader. + if g.proxy != nil { + g.proxy.StopSignalReporting() + } // Standalone WAL (collector_url empty): generation owns the writer and // must close it explicitly, otherwise every reload leaks a file handle // on offline-mode deployments. Only non-nil when reporter is nil. @@ -265,7 +273,10 @@ func (g *generation) drain(timeout time.Duration, reloadID string) { // Supervisor manages the proxy lifecycle and exposes the data-plane handler. type Supervisor struct { startedAt time.Time - transport http.RoundTripper // optional upstream proxy transport; nil = default + // transport is the optional upstream-proxy RoundTripper applied to every + // generation's proxy (nil = default). atomic.Pointer so an egress hot-swap + // (SetTransport, 2026-06-30) can't race the gen-build read in applyToProxy. + transport atomic.Pointer[transportBox] // ctx / cancel bound the lifetime of all detached upstream calls. // Canceled in Shutdown() to stop any in-flight upstream requests. ctx context.Context @@ -290,6 +301,29 @@ type Supervisor struct { // takes effect within the 60s poll WITHOUT the employee running any command, // and steady state adds no churn. Mirrors lastFilterSig / masterCompliance. lastQuotaSig atomic.Pointer[string] + // lastGroupRuntimeSig is the signature (raw response body) of the last group + // runtime this node pulled (N7c-2). pollGroupRuntime rewrites group_runtime + + // Reloads ONLY when the material actually changed (a token refreshed), so a + // 60s poll over unchanged tokens adds no churn. Mirrors lastQuotaSig. + lastGroupRuntimeSig atomic.Pointer[string] + // routingOverrides is the allocation engine's seat→account routing-override + // cache (I-side §6.5). Supervisor-scoped (not per-generation) so a reload never + // loses it — the same instance is injected into every Proxy. Populated by + // pollRoutingOverrides; read on the group-route hot path to redirect a seat off + // an unhealthy account. nil-safe everywhere → empty means "use the local pick". + routingOverrides *proxy.RoutingOverrideCache + // lastRoutingMismatchVersion throttles the proxy.routing_override.format_mismatch + // WARN (non-empty routes, zero matching a local (seat,group)) to once per + // routing_version — the 60s ticker would otherwise repeat it every cycle. + lastRoutingMismatchVersion atomic.Int64 + // railset drives the SyncRail control-plane sync rails (railset.go, + // 2026-07-03): routing_override + group_runtime in Phase 1. Supervisor-scoped + // so Reload can kick an immediate re-sync and /status can snapshot rail health. + railset *railSet + // teamCred is the shared on-demand team account-JWT source for the rails — + // rebuilt whenever the control URL changes or a refresh fails (the fix for the + // 2026-07-03 "credential baked at start with a stale URL" incident). + teamCred *teamCredentialSource // quotaHeartbeat is the traffic-independent server-reachability probe behind // budget-mode staleness (D-U7/P9). nil unless enforce_mode=budget AND a // collector URL is configured — so the default availability path (and Personal) @@ -349,17 +383,20 @@ func (s *Supervisor) VaultReader() *vault.Reader { func New(cfg *config.Config, configPath, password, version string) (*Supervisor, error) { ctx, cancel := context.WithCancel(context.Background()) s := &Supervisor{ - cfg: cfg, - configPath: configPath, - password: password, - version: version, - startedAt: time.Now(), - ctx: ctx, - cancel: cancel, - quotaEnabled: quotaEnabledFromEnv(), - quotaSnapshot: quota.NewSnapshot(), - quotaCounter: quota.NewCounter(), - } + cfg: cfg, + configPath: configPath, + password: password, + version: version, + startedAt: time.Now(), + ctx: ctx, + cancel: cancel, + quotaEnabled: quotaEnabledFromEnv(), + quotaSnapshot: quota.NewSnapshot(), + quotaCounter: quota.NewCounter(), + routingOverrides: proxy.NewRoutingOverrideCache(), + teamCred: &teamCredentialSource{}, + } + s.railset = newRailSet(s.groupRuntimeRail(), s.routingOverrideRail()) gen, err := s.buildGeneration() if err != nil { return nil, fmt.Errorf("initial generation failed: %w", err) @@ -413,6 +450,23 @@ func New(cfg *config.Config, configPath, password, version string) (*Supervisor, // or no active seats (Personal), or when quota is disabled. observability.GoSafe("supervisor.quota_policy_poll", observability.Isolated, func() { s.pollQuotaPolicy(s.ctx) }) + // SyncRail (2026-07-03): the group_runtime (N7c-2 channel ③ material) and + // routing_override (I-side §6.5 engine assignments) rails, driven by the + // declarative framework in railset.go — gate/URL/credential re-evaluated every + // cycle, failures counted into the OK→STALE→OFFLINE visibility state machine, + // Reload kicks an immediate re-sync. One GoSafe/Isolated goroutine per rail + // (names kept from the old hand-written polls so log filters carry over). + // quota/compliance/audit stay on their legacy loops until Phase 2. + s.railset.start(s) + + // Control-plane self-heal, Stage 2 (2026-07-01): proactively rebuild the + // control-plane client the moment the host's network changes (WiFi switch / + // tether / interface up-down), so control-plane calls to master dial clean + // without waiting for a failure. Dependency-free (net.Interfaces fingerprint). + // Isolated + cheap: a 20s poll; a panic here must never touch the data path. + // See netmon.go / selfheal.go. + observability.GoSafe("supervisor.net_change_monitor", observability.Isolated, func() { runNetChangeMonitor(s.ctx) }) + // Cluster mode (V3c): register this node with the hub name service + heartbeat // so clients discover it via /cluster/resolve. Inert for non-cluster proxies — // Personal/Trial never set Cluster.Enabled, so this is a no-op there. Isolated @@ -511,7 +565,7 @@ func (s *Supervisor) startQuotaHeartbeat() { if interval > 120*time.Second { interval = 120 * time.Second } - client := &http.Client{Timeout: 5 * time.Second} + client := httpx.NewDirectClient(5 * time.Second) s.quotaHeartbeat = heartbeat.New(interval, func(ctx context.Context) error { req, err := http.NewRequestWithContext(ctx, http.MethodGet, healthURL, http.NoBody) if err != nil { @@ -544,13 +598,18 @@ func (s *Supervisor) SetBroker(b proxy.OAuthBroker) { } func (s *Supervisor) SetTransport(t http.RoundTripper) { - s.transport = t - // Also apply to the already-running initial generation. + s.transport.Store(&transportBox{rt: t}) + // Also apply to the already-running initial generation. proxy.SetTransport is + // itself atomic, so this hot-swap is safe while the generation serves requests. if gen := s.active.Load(); gen != nil { gen.proxy.SetTransport(t) } } +// transportBox boxes the RoundTripper so it can live in an atomic.Pointer (atomics +// can't hold an interface value directly). nil rt = use the default transport. +type transportBox struct{ rt http.RoundTripper } + // managedKeySyncLoop runs in a background goroutine and periodically merges // newly-active managed keys into the live registry without a full reload. // @@ -679,6 +738,15 @@ func (s *Supervisor) syncManagedKeys() { } for i := range managedKeys { mk := &managedKeys[i] + // N7c safety gate: a group VK (OauthGroupID != "") carries NO PlaintextKey — + // its per-account material lives in GroupRuntime and routing requires the + // group resolver (N8). Until that's enabled, do NOT register group VKs; + // otherwise the hot path's `PlaintextKey != ""` check fails and the request + // falls to the personal-key path and 401s. Gated by + // AIKEY_PROXY_OAUTH_GROUP_ENABLED (default off) → direct-bind path unchanged. + if mk.OauthGroupID != "" && !oauthGroupRoutingEnabled() { + continue + } // 2026-04-29 prefix rename: team token = aikey_team_. // Use NormalizeTeamToken so historical-prefix dirty data in // mk.VirtualKeyID gets stripped+rebuilt (defense-in-depth: should @@ -939,6 +1007,13 @@ func (s *Supervisor) AppHealthSnapshot() []apppipe.AppHealth { return s.active.Load().proxy.AppHealthSnapshot() } +// PoolCooldownSnapshot returns oauth-group accounts currently cooling down +// (account_id → seconds remaining) from the active generation's proxy, for the +// admin /status pool-routing health surface (N9). +func (s *Supervisor) PoolCooldownSnapshot() map[string]int { + return s.active.Load().proxy.PoolCooldownSnapshot() +} + // ReporterMetrics returns usage reporter counters from the active generation. // Returns nil if reporter is not configured (no collector_url). func (s *Supervisor) ReporterMetrics() *events.ReporterMetrics { @@ -1245,6 +1320,20 @@ func (s *Supervisor) Reload(ctx context.Context) error { ) }) + // SyncRail convergence (§5.2, 2026-07-03): a reload usually follows a config / + // vault change (`aikey account set-url`, login, settings save). Drop the + // cached team credential so the next cycle rebuilds against the CURRENT + // control URL, and nudge every rail to run that cycle NOW — the settings + // change converges in seconds instead of one poll interval. Both calls are + // non-blocking (kick channels are buffered); the 60s per-cycle URL re-check + // remains the bottom-line self-heal when no reload ever fires. + if s.teamCred != nil { + s.teamCred.invalidate() + } + if s.railset != nil { + s.railset.kickAll() + } + return nil } @@ -1341,10 +1430,22 @@ func (s *Supervisor) buildGeneration() (*generation, error) { // Build the proxy handler with configured thresholds. p := proxy.New(vaultReader, registry, providers, collector, s.ctx) + // Inject the shared, supervisor-owned routing-override cache (I-side §6.5) so + // the group-route hot path can read the engine's seat→account redirects. + // Unconditional + nil-safe: an empty cache (no team cred / control URL / poll + // not landed) just means every request uses the local seatassign pick. + p.SetRoutingOverrides(s.routingOverrides) + // Local console base for member-login URLs in group login-required 401s + // (20260703 update). Explicitly-empty (cluster/server configs) → URL-less + // fallback; absent key (pre-20260703 preserved configs) → default 8090. + p.SetConsoleURL(s.cfg.ResolvedConsoleURL()) + // SyncRail §5.4: let the 401 wording distinguish "you need to sign in" from + // "the assignment rail is unreachable so this pick may be misdirected". + p.SetRoutingRailHealth(func() (string, int64) { return s.railHealthFor("routing_override") }) p.SlowRequestMs = int64(s.cfg.Log.SlowRequestMs) p.VerySlowRequestMs = int64(s.cfg.Log.VerySlowRequestMs) - if s.transport != nil { - p.SetTransport(s.transport) + if b := s.transport.Load(); b != nil && b.rt != nil { + p.SetTransport(b.rt) } if s.broker != nil { p.SetBroker(s.broker) @@ -1530,6 +1631,12 @@ func (s *Supervisor) buildGeneration() (*generation, error) { loadedSeq = int64(seq) } p.SetReporter(reporter, fmt.Sprintf("proxy-%d", id), s.version, fmt.Sprintf("gen-%d", id), loadedSeq, vaultReader.GetLoggedInAccountID()) + // I5: wire the allocation-engine util signal reporter with the team + // account-JWT (same credential the group-runtime poll uses) → master + // /accounts/me/signals. Off unless a team credential + control URL exist. + if teamCred := buildCollectorCredentials(s.cfg.Events.CollectorCredentials, vaultReader)["team"]; teamCred != nil { + p.EnableSignalReporting(readControlPanelURL(), teamCred.Bearer) + } slog.Info("usage reporter enabled", "collector_url", s.cfg.Events.CollectorURL) // Start canary probe. As of 2026-04-17 diagnostics live on the diff --git a/internal/sysproxy/parse.go b/internal/sysproxy/parse.go new file mode 100644 index 0000000..7827108 --- /dev/null +++ b/internal/sysproxy/parse.go @@ -0,0 +1,87 @@ +// parse.go — pure parsers for the two platform proxy-setting wire formats. +// Kept build-tag-free so fixture tests run on every development platform +// (test-fixture rule: each wire format gets fixture-based coverage). +package sysproxy + +import ( + "net" + "strings" +) + +// parseScutilProxy parses `scutil --proxy` output (macOS). Relevant keys: +// +// HTTPEnable : 1 HTTPProxy : 127.0.0.1 HTTPPort : 7890 +// HTTPSEnable : 1 HTTPSProxy : 127.0.0.1 HTTPSPort : 7890 +// SOCKSEnable : 1 SOCKSProxy : 127.0.0.1 SOCKSPort : 7891 +// +// An entry contributes only when its *Enable flag is 1 AND a host is present. +// PAC (ProxyAutoConfigEnable) is NOT supported — evaluating a PAC script needs +// a JS engine; a PAC-only setup yields an empty snapshot (direct), same as the +// pre-detection behavior, so nothing regresses. +func parseScutilProxy(out string) Snapshot { + kv := map[string]string{} + for _, line := range strings.Split(out, "\n") { + k, v, ok := strings.Cut(line, ":") + if !ok { + continue + } + kv[strings.TrimSpace(k)] = strings.TrimSpace(v) + } + entry := func(prefix, scheme string) string { + if kv[prefix+"Enable"] != "1" { + return "" + } + host, port := kv[prefix+"Proxy"], kv[prefix+"Port"] + if host == "" || port == "" { + return "" + } + return scheme + "://" + net.JoinHostPort(host, port) + } + return Snapshot{ + HTTP: entry("HTTP", "http"), + HTTPS: entry("HTTPS", "http"), // HTTPS entry = proxy for https targets; the proxy itself speaks plain HTTP CONNECT + SOCKS: entry("SOCKS", "socks5"), + } +} + +// parseWindowsProxyServer parses the registry `ProxyServer` value (used only +// when `ProxyEnable` is 1). Two wire formats exist: +// +// "host:port" → one proxy for all protocols +// "http=h:p;https=h:p;ftp=h:p;socks=h:p" → per-protocol entries +// +// A leading scheme ("http://host:port") is tolerated and stripped. +func parseWindowsProxyServer(val string) Snapshot { + val = strings.TrimSpace(val) + if val == "" { + return Snapshot{} + } + if !strings.Contains(val, "=") { + u := "http://" + stripScheme(val) + return Snapshot{HTTP: u, HTTPS: u} + } + var s Snapshot + for _, part := range strings.Split(val, ";") { + proto, hp, ok := strings.Cut(strings.TrimSpace(part), "=") + if !ok || hp == "" { + continue + } + hp = stripScheme(hp) + switch strings.ToLower(strings.TrimSpace(proto)) { + case "http": + s.HTTP = "http://" + hp + case "https": + s.HTTPS = "http://" + hp + case "socks": + s.SOCKS = "socks5://" + hp + } + } + return s +} + +func stripScheme(hp string) string { + if _, rest, ok := strings.Cut(hp, "://"); ok { + return rest + } + return hp +} diff --git a/internal/sysproxy/read_darwin.go b/internal/sysproxy/read_darwin.go new file mode 100644 index 0000000..19e3f0f --- /dev/null +++ b/internal/sysproxy/read_darwin.go @@ -0,0 +1,25 @@ +//go:build darwin + +package sysproxy + +import ( + "fmt" + "os/exec" + "time" +) + +const platformSupported = true + +// readSystemProxy shells out to `scutil --proxy` — the canonical, cgo-free way +// to read the primary network service's proxy settings. It reflects Clash / +// System Settings changes immediately. Cost ~10ms, run every pollInterval. +func readSystemProxy() (Snapshot, error) { + cmd := exec.Command("/usr/sbin/scutil", "--proxy") + // WaitDelay guards against a wedged scutil pinning the poll goroutine. + cmd.WaitDelay = 5 * time.Second + out, err := cmd.Output() + if err != nil { + return Snapshot{}, fmt.Errorf("scutil --proxy: %w", err) + } + return parseScutilProxy(string(out)), nil +} diff --git a/internal/sysproxy/read_darwin_test.go b/internal/sysproxy/read_darwin_test.go new file mode 100644 index 0000000..0c4b680 --- /dev/null +++ b/internal/sysproxy/read_darwin_test.go @@ -0,0 +1,16 @@ +//go:build darwin + +package sysproxy + +import "testing" + +// Live exec smoke: `scutil --proxy` must run and parse on any macOS box — +// whatever it returns (proxy on/off) must not error. Guards the exec wiring +// the fixture tests can't cover. +func TestReadSystemProxy_ExecSmoke(t *testing.T) { + snap, err := readSystemProxy() + if err != nil { + t.Fatalf("scutil exec path failed: %v", err) + } + t.Logf("live system proxy snapshot: %+v", snap) +} diff --git a/internal/sysproxy/read_other.go b/internal/sysproxy/read_other.go new file mode 100644 index 0000000..6a49d4a --- /dev/null +++ b/internal/sysproxy/read_other.go @@ -0,0 +1,12 @@ +//go:build !darwin && !windows + +package sysproxy + +// Linux/other: no OS system-proxy detection (user-approved scope 2026-07-08). +// Server editions (Cluster nodes, CI) run here — the watcher stays inert and +// egress behavior is byte-identical to the pre-detection daemon (env vars via +// http.ProxyFromEnvironment). Desktop-Linux gsettings support is a future +// drop-in: implement readSystemProxy + flip platformSupported. +const platformSupported = false + +func readSystemProxy() (Snapshot, error) { return Snapshot{}, nil } diff --git a/internal/sysproxy/read_windows.go b/internal/sysproxy/read_windows.go new file mode 100644 index 0000000..dc35975 --- /dev/null +++ b/internal/sysproxy/read_windows.go @@ -0,0 +1,35 @@ +//go:build windows + +package sysproxy + +import ( + "fmt" + + "golang.org/x/sys/windows/registry" +) + +const platformSupported = true + +// readSystemProxy reads the per-user WinINET proxy settings — the registry +// location Windows' "Proxy" settings page (and Clash for Windows) writes. +// AutoConfigURL (PAC) is not evaluated (needs a JS engine); PAC-only setups +// yield an empty snapshot = direct, identical to pre-detection behavior. +func readSystemProxy() (Snapshot, error) { + k, err := registry.OpenKey(registry.CURRENT_USER, + `Software\Microsoft\Windows\CurrentVersion\Internet Settings`, registry.QUERY_VALUE) + if err != nil { + return Snapshot{}, fmt.Errorf("open Internet Settings key: %w", err) + } + defer k.Close() + + enabled, _, err := k.GetIntegerValue("ProxyEnable") + if err != nil || enabled == 0 { + // Missing value or disabled both mean "no system proxy" — not an error. + return Snapshot{}, nil + } + server, _, err := k.GetStringValue("ProxyServer") + if err != nil { + return Snapshot{}, nil + } + return parseWindowsProxyServer(server), nil +} diff --git a/internal/sysproxy/sysproxy.go b/internal/sysproxy/sysproxy.go new file mode 100644 index 0000000..32a5060 --- /dev/null +++ b/internal/sysproxy/sysproxy.go @@ -0,0 +1,426 @@ +// Package sysproxy detects the OS-level system proxy (macOS scutil, Windows +// registry) and keeps a live snapshot the egress transport reads per request. +// +// WHY this package exists (2026-07-08 需求: 代理更新时 proxy 自动刷新): +// aikey-proxy is a long-lived daemon. Before this package its egress proxy was +// frozen three times over: (1) env vars are snapshotted at process spawn, +// (2) Go's http.ProxyFromEnvironment permanently caches the env via sync.Once +// on the first request, (3) there was NO system-proxy detection at all — a +// launchd/systemd-spawned daemon never sees the login shell's HTTP_PROXY, so a +// user driving egress through Clash's *system proxy* was invisible to us, and +// any Clash port change / toggle stranded requests until `aikey proxy restart`. +// +// Design (方案 A, user-approved 2026-07-08): the direct-mode Transport.Proxy is +// a FUNCTION reading an atomic snapshot refreshed by a poll loop — never +// http.ProxyFromEnvironment's cached global. Precedence (single source of +// truth, user-approved; refined same day after field evidence from two +// machines — see below): +// +// explicit upstream_proxy.url > proxy.env EXPLICIT env > OS system proxy (live) > inherited shell env > direct +// +// WHY the env layer is split in two (2026-07-08 refinement, user-approved): +// "process env" conflates the user's EXPLICIT aikey config (~/.aikey/proxy.env, +// which the CLI injects at spawn) with ACCIDENTAL shell inheritance (.zshrc +// exports, stale terminal sessions). Clash-style users practically always have +// shell exports, so an undivided env layer permanently masked the system-proxy +// auto-follow this package exists for — and made behavior differ between +// launchd auto-start (no shell env) and manual restart (shell env). The CLI +// marks its proxy.env injections via AIKEY_PROXYENV_KEYS; only marked proxy +// vars outrank OS detection, inherited ones fall below it as a last resort +// (still covering headless/Linux manual starts). +// +// The explicit-URL layer stays OUTSIDE this package (buildTransport keeps +// using http.ProxyURL for it); this package resolves the rest. +// +// WHY poll (not events): netmon's fingerprint is the set of interface IPs — +// toggling/re-porting a system proxy does NOT change interface IPs, so the +// existing net-change monitor can never observe it. Polling `scutil --proxy` / +// the registry every pollInterval is the only reliable, dependency-free signal. +// +// Failure posture (增强非依赖): detection is a bypass. Read errors keep the +// last-known snapshot and WARN once per failure streak; unsupported platforms +// (Linux servers) return an inert watcher — behavior is byte-identical to the +// pre-2026-07-08 daemon there. +package sysproxy + +import ( + "context" + "log/slog" + "net" + "net/http" + "net/url" + "os" + "strings" + "sync" + "time" + + "golang.org/x/net/http/httpproxy" + + "github.com/AiKeyLabs/aikey-proxy/internal/observability" +) + +// pollInterval matches supervisor's netChangePollInterval cadence: fast enough +// that a Clash toggle heals well inside a user's "retry the request" window, +// cheap enough (~10ms scutil exec / registry read) to be negligible. +const pollInterval = 20 * time.Second + +// Snapshot is one observation of the OS system proxy. Empty string = that +// protocol has no system proxy. Values are normalized absolute URLs +// (http://host:port, socks5://host:port). +type Snapshot struct { + HTTP string + HTTPS string + SOCKS string +} + +// Empty reports whether no system proxy is configured at all. +func (s Snapshot) Empty() bool { return s.HTTP == "" && s.HTTPS == "" && s.SOCKS == "" } + +// ProxyFor returns the proxy URL string for a request scheme, or "" for +// direct. https prefers the HTTPS entry, http the HTTP entry; SOCKS is the +// last resort for both (Clash-style setups usually set all three identically). +func (s Snapshot) ProxyFor(scheme string) string { + if scheme == "https" { + return firstNonEmpty(s.HTTPS, s.HTTP, s.SOCKS) + } + return firstNonEmpty(s.HTTP, s.SOCKS) +} + +func firstNonEmpty(vals ...string) string { + for _, v := range vals { + if v != "" { + return v + } + } + return "" +} + +// Watcher owns the live snapshot. Construct with NewWatcher (primes +// synchronously so the first outbound request already sees the system proxy), +// then start Run in a goroutine. +type Watcher struct { + mu sync.Mutex + cur Snapshot + read func() (Snapshot, error) // platform reader; injectable in tests + + // supported mirrors the platform const as a field so tests exercise the + // full watcher on any development/CI OS. + supported bool + + // envProxy resolves proxy env vars (both layers 2 and 4). Built ONCE at + // construction from the frozen process env via x/net/httpproxy — the same + // implementation net/http delegates to, but WITHOUT the process-global + // sync.Once cache, so tests can drive it deterministically with t.Setenv. + envProxy func(*url.URL) (*url.URL, error) + + // envExplicit: the process env carries proxy vars that the CLI marked as + // coming from ~/.aikey/proxy.env (AIKEY_PROXYENV_KEYS). That is EXPLICIT + // user config, so it outranks OS detection — including its NO_PROXY + // exclusions, which must not "fall through" to the system layer. + // Inherited (unmarked) proxy vars do NOT set this; they participate only + // as the below-system fallback inside ProxyFunc. + envExplicit bool + + // readFailing dedups WARN logs: one WARN per failure streak, one INFO on + // recovery, never a 20s WARN drumbeat. + readFailing bool +} + +// NewWatcher builds and synchronously primes a watcher using the platform +// reader. On unsupported platforms or with authoritative env config the +// watcher is inert (Run returns immediately, ProxyFunc delegates to env). +func NewWatcher() *Watcher { + w := &Watcher{ + read: readSystemProxy, + supported: platformSupported, + envExplicit: envProxyExplicit(os.Getenv), + envProxy: httpproxy.FromEnvironment().ProxyFunc(), + } + w.prime() + return w +} + +// NewWatcherWithReader is the HARNESS constructor: a Watcher driven by an +// injected reader instead of the OS. Exported (internal/ only) so integration +// tests outside this package — e.g. cmd/aikey-proxy's egress switch test — +// can drive the REAL Watcher + buildTransport chain hermetically on any OS. +// Always "supported" and env-independent; never use in production wiring +// (production is NewWatcher, which honors platform + env precedence). +func NewWatcherWithReader(read func() (Snapshot, error)) *Watcher { + w := &Watcher{ + read: read, + supported: true, + envProxy: httpproxy.FromEnvironment().ProxyFunc(), + } + w.prime() + return w +} + +// PollOnce performs exactly one poll outside Run's ticker and reports whether +// the snapshot changed — the harness seam that advances the watcher +// deterministically (tests must not wait out the 20s production cadence). +func (w *Watcher) PollOnce() bool { + _, _, changed := w.observe() + return changed +} + +// newWatcherForTest injects a fake reader and explicit-env flag; always +// "supported" so the full poll/observe path runs on any test OS. +func newWatcherForTest(read func() (Snapshot, error), envExplicit bool) *Watcher { + w := &Watcher{ + read: read, + supported: true, + envExplicit: envExplicit, + envProxy: httpproxy.FromEnvironment().ProxyFunc(), + } + w.prime() + return w +} + +func (w *Watcher) prime() { + if !w.active() { + return + } + snap, err := w.read() + if err != nil { + w.readFailing = true + slog.Warn("system proxy read failed at startup; assuming none", + "event.name", observability.EventProxyEgressSysProxyReadFailed, + "error", err) + return + } + w.cur = snap + if !snap.Empty() { + slog.Info("system proxy detected", + "event.name", observability.EventProxyEgressSysProxyChanged, + "http", snap.HTTP, "https", snap.HTTPS, "socks", snap.SOCKS) + } +} + +// active reports whether OS detection participates at all. Only EXPLICIT +// (proxy.env-marked) env config disables it — inherited shell env does not, +// because the system proxy now outranks inheritance. +func (w *Watcher) active() bool { return w.supported && !w.envExplicit } + +// Current returns the last-known snapshot (zero value when inactive). +func (w *Watcher) Current() Snapshot { + w.mu.Lock() + defer w.mu.Unlock() + return w.cur +} + +// EnvExplicit reports whether proxy.env-marked (explicit) proxy vars are +// present and therefore outrank OS detection. Read-only observability +// accessor (admin egress state / `aikey env`). +func (w *Watcher) EnvExplicit() bool { return w.envExplicit } + +// Supported reports whether this platform has OS system-proxy detection. +// Read-only observability accessor. +func (w *Watcher) Supported() bool { return w.supported } + +// EnvProxyVarsSplit returns the proxy-relevant environment variables THIS +// process sees (the daemon's frozen spawn env — deliberately not the user's +// shell), split by provenance: explicit = injected from ~/.aikey/proxy.env +// (AIKEY_PROXYENV_KEYS-marked), inherited = everything else (shell exports, +// stale terminals). Credentials in URLs are redacted. Observability only: +// the egress decision itself goes through ProxyFunc, never these maps. +func EnvProxyVarsSplit() (explicit, inherited map[string]string) { + return envProxyVarsSplitFrom(os.Getenv) +} + +func envProxyVarsSplitFrom(get func(string) string) (explicit, inherited map[string]string) { + marked := explicitEnvKeySet(get) + explicit, inherited = map[string]string{}, map[string]string{} + for k, v := range envProxyVarsFrom(get) { + if marked[strings.ToUpper(k)] { + explicit[k] = v + } else { + inherited[k] = v + } + } + return explicit, inherited +} + +// envProxyVarsFrom is the injectable core (tests fake the getter). WHY the +// twin-dedupe (2026-07-08 Windows field report): Windows env keys are +// case-INsensitive — Getenv("http_proxy") returns the SAME variable as +// Getenv("HTTP_PROXY"), so the naive 8-key loop displayed every variable +// twice in `aikey env`. The lowercase twin is listed only when it's a +// genuinely distinct value (possible on Unix, where env is case-sensitive). +func envProxyVarsFrom(get func(string) string) map[string]string { + out := map[string]string{} + for _, pair := range [][2]string{ + {"HTTP_PROXY", "http_proxy"}, + {"HTTPS_PROXY", "https_proxy"}, + {"ALL_PROXY", "all_proxy"}, + {"NO_PROXY", "no_proxy"}, + } { + upper, lower := get(pair[0]), get(pair[1]) + if upper != "" { + out[pair[0]] = redactURLCredentials(upper) + } + if lower != "" && lower != upper { + out[pair[1]] = redactURLCredentials(lower) + } + } + return out +} + +// redactURLCredentials masks user:password inside a proxy URL (rare but legal, +// e.g. http://user:pass@host:port) so the value is safe to expose over the +// admin API and in CLI output. Non-URL values pass through unchanged. +func redactURLCredentials(raw string) string { + u, err := url.Parse(raw) + if err != nil || u.User == nil { + return raw + } + return u.Redacted() +} + +// Run polls the OS until ctx is done, invoking onChange(old, cur) after the +// internal snapshot flips. Inert (returns immediately) when env config is +// authoritative or the platform is unsupported. +func (w *Watcher) Run(ctx context.Context, onChange func(old, cur Snapshot)) { + if !w.active() { + return + } + t := time.NewTicker(pollInterval) + defer t.Stop() + for { + select { + case <-ctx.Done(): + return + case <-t.C: + if old, cur, changed := w.observe(); changed { + onChange(old, cur) + } + } + } +} + +// observe performs one poll: reads the OS, updates the snapshot, logs +// transitions. Split from Run so tests drive it without a ticker (same +// pattern as supervisor's changeDetector). +func (w *Watcher) observe() (old, cur Snapshot, changed bool) { + snap, err := w.read() + w.mu.Lock() + defer w.mu.Unlock() + if err != nil { + // Keep last-known: a transient scutil/registry hiccup must not flap + // the egress path to direct and back. + if !w.readFailing { + w.readFailing = true + slog.Warn("system proxy read failed; keeping last-known value", + "event.name", observability.EventProxyEgressSysProxyReadFailed, + "error", err) + } + return w.cur, w.cur, false + } + if w.readFailing { + w.readFailing = false + slog.Info("system proxy read recovered", + "event.name", observability.EventProxyEgressSysProxyReadRecovered) + } + if snap == w.cur { + return w.cur, w.cur, false + } + old, w.cur = w.cur, snap + slog.Info("system proxy changed; egress refreshed", + "event.name", observability.EventProxyEgressSysProxyChanged, + "http", snap.HTTP, "https", snap.HTTPS, "socks", snap.SOCKS, + "old_http", old.HTTP, "old_https", old.HTTPS, "old_socks", old.SOCKS) + return old, snap, true +} + +// ProxyFunc returns the Transport.Proxy for the DIRECT egress mode +// (upstream_proxy.url empty). Layering per the approved precedence: +// - explicit proxy.env env (marked) → env resolution with full NO_PROXY +// handling (x/net httpproxy — identical semantics to +// http.ProxyFromEnvironment, minus its process-global cache); +// - otherwise → the live system-proxy snapshot, re-read on every request so +// a mid-flight Clash change applies to the very next request; +// - no system proxy → INHERITED shell env as last-resort fallback (covers +// headless/Linux manual starts and unsupported platforms, where this +// equals the pre-detection behavior); +// - nothing anywhere → direct. +func (w *Watcher) ProxyFunc() func(*http.Request) (*url.URL, error) { + if !w.active() { + return func(req *http.Request) (*url.URL, error) { return w.envProxy(req.URL) } + } + return func(req *http.Request) (*url.URL, error) { + // Never proxy loopback destinations (local relays, self-probes) — + // parity with ProxyFromEnvironment's localhost rule. + if isLoopbackHost(req.URL.Hostname()) { + return nil, nil + } + if raw := w.Current().ProxyFor(req.URL.Scheme); raw != "" { + return url.Parse(raw) + } + // Layer 4: inherited shell env (unmarked vars) as fallback. + return w.envProxy(req.URL) + } +} + +// BrokerEgressURL is the layered egress for clients that take a static proxy +// URL string (the OAuth ImpersonateChrome client). "" when explicit env +// config applies or nothing is detected — in both cases the client's own env +// fallback yields the same result (explicit or inherited env respectively). +// Callers must re-invoke this on change (main's onChange rebuilds the broker +// client) — the returned string is a point-in-time value by design. +func (w *Watcher) BrokerEgressURL() string { + if !w.active() { + return "" + } + return w.Current().ProxyFor("https") +} + +// proxyEnvNames is the variable set golang.org/x/net/http/httpproxy consults. +// NO_PROXY counts too: it expresses "I curated env-level proxy rules". +var proxyEnvNames = []string{ + "HTTP_PROXY", "http_proxy", + "HTTPS_PROXY", "https_proxy", + "ALL_PROXY", "all_proxy", + "NO_PROXY", "no_proxy", +} + +// explicitEnvKeyMarker is set by aikey-cli at spawn: comma-separated env keys +// it injected from ~/.aikey/proxy.env (the user's EXPLICIT aikey config). +const explicitEnvKeyMarker = "AIKEY_PROXYENV_KEYS" + +// explicitEnvKeySet parses the marker into an upper-cased key set. Upper-cased +// on both sides because Windows env keys are case-insensitive and proxy.env +// keys are user-typed in either case. +func explicitEnvKeySet(get func(string) string) map[string]bool { + out := map[string]bool{} + for _, k := range strings.Split(get(explicitEnvKeyMarker), ",") { + if k = strings.TrimSpace(k); k != "" { + out[strings.ToUpper(k)] = true + } + } + return out +} + +// envProxyExplicit reports whether any proxy var is BOTH set and marked as +// coming from proxy.env — the only condition that outranks OS detection. +func envProxyExplicit(get func(string) string) bool { + marked := explicitEnvKeySet(get) + if len(marked) == 0 { + return false + } + for _, k := range proxyEnvNames { + if get(k) != "" && marked[strings.ToUpper(k)] { + return true + } + } + return false +} + +func isLoopbackHost(host string) bool { + if host == "localhost" { + return true + } + if ip := net.ParseIP(host); ip != nil { + return ip.IsLoopback() + } + return false +} diff --git a/internal/sysproxy/sysproxy_test.go b/internal/sysproxy/sysproxy_test.go new file mode 100644 index 0000000..38334f8 --- /dev/null +++ b/internal/sysproxy/sysproxy_test.go @@ -0,0 +1,351 @@ +package sysproxy + +import ( + "context" + "errors" + "net/http" + "os" + "reflect" + "strings" + "testing" + "time" +) + +// --- fixture parsers (one test per wire format, per test-fixture rule) --- + +// Real `scutil --proxy` output shape on macOS with Clash-style system proxy. +const scutilClashFixture = ` { + ExceptionsList : { + 0 : 192.168.0.0/16 + 1 : localhost + } + FTPPassive : 1 + HTTPEnable : 1 + HTTPPort : 7890 + HTTPProxy : 127.0.0.1 + HTTPSEnable : 1 + HTTPSPort : 7890 + HTTPSProxy : 127.0.0.1 + SOCKSEnable : 1 + SOCKSPort : 7891 + SOCKSProxy : 127.0.0.1 +}` + +func TestParseScutilProxy_ClashAllOn(t *testing.T) { + got := parseScutilProxy(scutilClashFixture) + want := Snapshot{ + HTTP: "http://127.0.0.1:7890", + HTTPS: "http://127.0.0.1:7890", + SOCKS: "socks5://127.0.0.1:7891", + } + if got != want { + t.Fatalf("got %+v want %+v", got, want) + } +} + +func TestParseScutilProxy_DisabledAndPACOnly(t *testing.T) { + // Toggled off: Enable flags 0 (or absent) must yield empty even when + // stale Proxy/Port keys linger — macOS keeps them after disable. + off := ` { + HTTPEnable : 0 + HTTPPort : 7890 + HTTPProxy : 127.0.0.1 + ProxyAutoConfigEnable : 1 + ProxyAutoConfigURLString : http://127.0.0.1:7890/pac +}` + if got := parseScutilProxy(off); !got.Empty() { + t.Fatalf("disabled/PAC-only must parse as empty (direct), got %+v", got) + } + if got := parseScutilProxy(" {\n}"); !got.Empty() { + t.Fatalf("empty dictionary must parse as empty, got %+v", got) + } +} + +func TestParseWindowsProxyServer(t *testing.T) { + cases := []struct { + name string + in string + want Snapshot + }{ + {"single host:port applies to http+https", "127.0.0.1:7890", + Snapshot{HTTP: "http://127.0.0.1:7890", HTTPS: "http://127.0.0.1:7890"}}, + {"scheme prefix tolerated", "http://127.0.0.1:7890", + Snapshot{HTTP: "http://127.0.0.1:7890", HTTPS: "http://127.0.0.1:7890"}}, + {"per-protocol list", "http=127.0.0.1:7890;https=127.0.0.1:7891;ftp=1.2.3.4:21;socks=127.0.0.1:7892", + Snapshot{HTTP: "http://127.0.0.1:7890", HTTPS: "http://127.0.0.1:7891", SOCKS: "socks5://127.0.0.1:7892"}}, + {"empty", "", Snapshot{}}, + } + for _, c := range cases { + if got := parseWindowsProxyServer(c.in); got != c.want { + t.Errorf("%s: got %+v want %+v", c.name, got, c.want) + } + } +} + +// --- watcher behavior --- + +func reqTo(t *testing.T, rawurl string) *http.Request { + t.Helper() + req, err := http.NewRequest(http.MethodGet, rawurl, nil) + if err != nil { + t.Fatal(err) + } + return req +} + +// clearProxyEnv makes the test hermetic against the RUNNER's shell env (a dev +// Mac usually exports https_proxy): since the 2026-07-08 refinement, an empty +// snapshot falls through to inherited env — so tests asserting "direct" must +// pin the env to empty BEFORE constructing the watcher (envProxy is built at +// construction from the frozen env). +func clearProxyEnv(t *testing.T) { + t.Helper() + for _, k := range proxyEnvNames { + t.Setenv(k, "") + } + t.Setenv(explicitEnvKeyMarker, "") +} + +func TestProxyFunc_FollowsLiveSnapshot(t *testing.T) { + clearProxyEnv(t) + cur := Snapshot{HTTP: "http://127.0.0.1:7890", HTTPS: "http://127.0.0.1:7890"} + var readErr error + w := newWatcherForTest(func() (Snapshot, error) { return cur, readErr }, false) + fn := w.ProxyFunc() + + // Primed at construction: request #1 already proxied. + u, err := fn(reqTo(t, "https://api.anthropic.com/v1/messages")) + if err != nil || u == nil || u.String() != "http://127.0.0.1:7890" { + t.Fatalf("want primed proxy, got %v err %v", u, err) + } + + // Clash port change: next observe flips, next request uses the NEW port + // — the core 2026-07-08 requirement. + cur = Snapshot{HTTP: "http://127.0.0.1:9999", HTTPS: "http://127.0.0.1:9999"} + if _, _, changed := w.observe(); !changed { + t.Fatal("port change must be observed") + } + if u, _ := fn(reqTo(t, "https://api.anthropic.com/v1/messages")); u == nil || u.Port() != "9999" { + t.Fatalf("want refreshed proxy :9999, got %v", u) + } + + // System proxy toggled OFF: refresh to direct (env cleared above, so the + // layer-4 inherited fallback has nothing to offer). + cur = Snapshot{} + if _, _, changed := w.observe(); !changed { + t.Fatal("toggle-off must be observed") + } + if u, _ := fn(reqTo(t, "https://api.anthropic.com/v1/messages")); u != nil { + t.Fatalf("want direct after toggle-off, got %v", u) + } +} + +// Layer 4 (2026-07-08 refinement): with no system proxy, INHERITED shell env +// (unmarked) is the fallback; when a system proxy appears it takes over — +// inherited env is outranked. +func TestProxyFunc_InheritedEnvIsBelowSystemProxy(t *testing.T) { + clearProxyEnv(t) + t.Setenv("HTTPS_PROXY", "http://127.0.0.1:7890") // inherited: no marker + cur := Snapshot{} + w := newWatcherForTest(func() (Snapshot, error) { return cur, nil }, false) + fn := w.ProxyFunc() + + // No system proxy → inherited env fallback engages (old Linux/manual + // behavior preserved). + if u, _ := fn(reqTo(t, "https://api.anthropic.com/v1/messages")); u == nil || u.Port() != "7890" { + t.Fatalf("want inherited env fallback :7890, got %v", u) + } + + // System proxy turns on → it outranks the inherited env (the whole point + // of the refinement: Clash users' .zshrc exports must not pin the port). + cur = Snapshot{HTTPS: "http://127.0.0.1:9999"} + if _, _, changed := w.observe(); !changed { + t.Fatal("system proxy appearance must be observed") + } + if u, _ := fn(reqTo(t, "https://api.anthropic.com/v1/messages")); u == nil || u.Port() != "9999" { + t.Fatalf("system proxy must outrank inherited env, got %v", u) + } +} + +// Explicit proxy.env config (marker) still outranks everything below it, +// including the system proxy — the enterprise no-GUI contract is unchanged. +func TestProxyFunc_ExplicitEnvOutranksSystemProxy(t *testing.T) { + clearProxyEnv(t) + t.Setenv("HTTPS_PROXY", "http://127.0.0.1:7890") + t.Setenv(explicitEnvKeyMarker, "HTTPS_PROXY") + if !envProxyExplicit(os.Getenv) { + t.Fatal("marked HTTPS_PROXY must be explicit") + } + w := newWatcherForTest(func() (Snapshot, error) { + return Snapshot{HTTPS: "http://127.0.0.1:9999"}, nil + }, envProxyExplicit(os.Getenv)) + if w.active() { + t.Fatal("explicit env must keep the watcher inert") + } + u, err := w.ProxyFunc()(reqTo(t, "https://api.anthropic.com/v1/messages")) + if err != nil || u == nil || u.Port() != "7890" { + t.Fatalf("explicit env must win over system proxy, got %v err %v", u, err) + } +} + +func TestProxyFunc_SkipsLoopbackDestinations(t *testing.T) { + w := newWatcherForTest(func() (Snapshot, error) { + return Snapshot{HTTP: "http://127.0.0.1:7890", HTTPS: "http://127.0.0.1:7890"}, nil + }, false) + fn := w.ProxyFunc() + for _, dest := range []string{"http://localhost:8080/x", "http://127.0.0.1:9100/x", "http://[::1]:9100/x"} { + if u, _ := fn(reqTo(t, dest)); u != nil { + t.Errorf("loopback dest %s must go direct, got %v", dest, u) + } + } +} + +func TestObserve_ReadFailureKeepsLastKnown(t *testing.T) { + snap := Snapshot{HTTPS: "http://127.0.0.1:7890"} + fail := false + w := newWatcherForTest(func() (Snapshot, error) { + if fail { + return Snapshot{}, errors.New("scutil exploded") + } + return snap, nil + }, false) + + fail = true + if _, cur, changed := w.observe(); changed || cur != snap { + t.Fatalf("read failure must keep last-known, got changed=%v cur=%+v", changed, cur) + } + if !w.readFailing { + t.Fatal("failure streak flag must be set (WARN dedup)") + } + fail = false + if _, _, changed := w.observe(); changed { + t.Fatal("recovery to identical snapshot must not report change") + } + if w.readFailing { + t.Fatal("recovery must clear the failure streak flag") + } +} + +func TestEnvAuthoritative_WatcherInert(t *testing.T) { + reads := 0 + w := newWatcherForTest(func() (Snapshot, error) { reads++; return Snapshot{HTTP: "http://x:1"}, nil }, true) + if reads != 0 { + t.Fatal("env-authoritative watcher must not even prime from the OS") + } + if got := w.BrokerEgressURL(); got != "" { + t.Fatalf("env-authoritative BrokerEgressURL must be empty, got %q", got) + } + // Run must return immediately (no poll loop) even with a live context. + done := make(chan struct{}) + go func() { w.Run(context.Background(), nil); close(done) }() + select { + case <-done: + case <-time.After(2 * time.Second): + t.Fatal("Run must be inert when env config is authoritative") + } +} + +func TestSnapshotProxyForFallbacks(t *testing.T) { + socksOnly := Snapshot{SOCKS: "socks5://127.0.0.1:7891"} + if got := socksOnly.ProxyFor("https"); got != "socks5://127.0.0.1:7891" { + t.Fatalf("https must fall back to SOCKS, got %q", got) + } + httpOnly := Snapshot{HTTP: "http://127.0.0.1:7890"} + if got := httpOnly.ProxyFor("https"); got != "http://127.0.0.1:7890" { + t.Fatalf("https must fall back to HTTP entry, got %q", got) + } +} + +// 2026-07-08 precedence refinement: only proxy vars MARKED as coming from +// proxy.env (AIKEY_PROXYENV_KEYS) count as explicit; inherited shell exports +// must NOT disable OS detection. +func TestEnvProxyExplicit_MarkerGated(t *testing.T) { + get := func(env map[string]string) func(string) string { + return func(k string) string { return env[k] } + } + // Inherited only (the .zshrc-export field case): NOT explicit. + if envProxyExplicit(get(map[string]string{"https_proxy": "http://127.0.0.1:7890"})) { + t.Fatal("inherited shell env must not be explicit") + } + // Marked via proxy.env (case-insensitive match): explicit. + if !envProxyExplicit(get(map[string]string{ + "https_proxy": "http://127.0.0.1:7890", + "AIKEY_PROXYENV_KEYS": "HTTPS_PROXY,DEGRADE_DETECTOR_PROXY_TOKEN", + })) { + t.Fatal("proxy.env-marked https_proxy must be explicit") + } + // Marker lists non-proxy keys only: NOT explicit. + if envProxyExplicit(get(map[string]string{ + "https_proxy": "http://127.0.0.1:7890", + "AIKEY_PROXYENV_KEYS": "DEGRADE_DETECTOR_PROXY_TOKEN", + })) { + t.Fatal("marker without proxy keys must not be explicit") + } +} + +// Split by provenance: marked vars → explicit map, unmarked → inherited map. +func TestEnvProxyVarsSplit_ByMarker(t *testing.T) { + get := func(k string) string { + switch k { + case "HTTPS_PROXY": + return "http://127.0.0.1:7890" + case "all_proxy": + return "socks5://127.0.0.1:7890" + case "AIKEY_PROXYENV_KEYS": + return "https_proxy" + } + return "" + } + explicit, inherited := envProxyVarsSplitFrom(get) + if len(explicit) != 1 || explicit["HTTPS_PROXY"] == "" { + t.Fatalf("HTTPS_PROXY must be explicit (marker is case-insensitive), got %v", explicit) + } + if len(inherited) != 1 || inherited["all_proxy"] == "" { + t.Fatalf("all_proxy must be inherited, got %v", inherited) + } +} + +// Layer-4 fallback: system snapshot empty → ProxyFunc falls through to the +// inherited env (ProxyFromEnvironment) instead of going direct. +// NOTE: not asserted via a live ProxyFromEnvironment call here — its +// process-global sync.Once cache would make the assertion order-dependent +// across the package's tests. The fall-through branch is asserted by +// TestEgress_InheritedEnvFallback in cmd/aikey-proxy (fresh binary run). + +// 2026-07-08 Windows field report: env keys are case-insensitive on Windows, +// so both Getenv("HTTP_PROXY") and Getenv("http_proxy") return the same +// variable — the display must not list it twice. +func TestEnvProxyVars_WindowsCaseInsensitiveDedupe(t *testing.T) { + winGet := func(k string) string { // case-insensitive like Windows + switch strings.ToUpper(k) { + case "HTTP_PROXY", "HTTPS_PROXY": + return "http://127.0.0.1:7890" + } + return "" + } + got := envProxyVarsFrom(winGet) + want := map[string]string{ + "HTTP_PROXY": "http://127.0.0.1:7890", + "HTTPS_PROXY": "http://127.0.0.1:7890", + } + if !reflect.DeepEqual(got, want) { + t.Fatalf("want deduped %v, got %v", want, got) + } +} + +// Unix: case-sensitive env — genuinely distinct lowercase twins stay visible. +func TestEnvProxyVars_UnixDistinctCasesBothShown(t *testing.T) { + unixGet := func(k string) string { + switch k { + case "HTTPS_PROXY": + return "http://127.0.0.1:7890" + case "https_proxy": + return "http://127.0.0.1:9999" + } + return "" + } + got := envProxyVarsFrom(unixGet) + if got["HTTPS_PROXY"] != "http://127.0.0.1:7890" || got["https_proxy"] != "http://127.0.0.1:9999" { + t.Fatalf("distinct twins must both be shown, got %v", got) + } +} diff --git a/internal/vault/busy_timeout_fence_test.go b/internal/vault/busy_timeout_fence_test.go new file mode 100644 index 0000000..e5720b2 --- /dev/null +++ b/internal/vault/busy_timeout_fence_test.go @@ -0,0 +1,42 @@ +package vault + +import ( + "os" + "path/filepath" + "regexp" + "strings" + "testing" +) + +// TestNoRawSQLOpenInVault is a source-scan fence for the busy_timeout mandate. +// +// Why: modernc's default busy_timeout is 0, so any raw sql.Open against the +// shared vault fails instantly with SQLITE_BUSY under a concurrent reader and +// the write is silently dropped. This regressed twice by hand-audit misses +// (2026-07-03 R2: five raw opens; 2026-07-07 parity audit P1-4: +// WriteAssignmentOverride left behind). A structural scan cannot miss the +// eighth site the way a human can. +func TestNoRawSQLOpenInVault(t *testing.T) { + openCall := regexp.MustCompile(`sql\.Open\(\s*"sqlite"\s*,\s*([^)]*)\)`) + + entries, err := os.ReadDir(".") + if err != nil { + t.Fatalf("read package dir: %v", err) + } + for _, e := range entries { + name := e.Name() + if !strings.HasSuffix(name, ".go") || strings.HasSuffix(name, "_test.go") { + continue + } + src, err := os.ReadFile(filepath.Clean(name)) + if err != nil { + t.Fatalf("read %s: %v", name, err) + } + for _, m := range openCall.FindAllStringSubmatch(string(src), -1) { + arg := strings.TrimSpace(m[1]) + if !strings.HasPrefix(arg, "WithBusyTimeoutDSN(") { + t.Errorf("%s: raw sql.Open DSN %q — wrap it with WithBusyTimeoutDSN (busy_timeout is mandatory for all vault opens, see WriteGroupRuntime comment)", name, arg) + } + } + } +} diff --git a/internal/vault/get_secret_test.go b/internal/vault/get_secret_test.go index 947bb85..265792d 100644 --- a/internal/vault/get_secret_test.go +++ b/internal/vault/get_secret_test.go @@ -90,8 +90,8 @@ func TestWithBusyTimeoutDSN(t *testing.T) { {"/path/vault.db?cache=shared", "/path/vault.db?cache=shared&_pragma=busy_timeout(5000)"}, } for _, c := range cases { - if got := withBusyTimeoutDSN(c.in); got != c.want { - t.Errorf("withBusyTimeoutDSN(%q) = %q, want %q", c.in, got, c.want) + if got := WithBusyTimeoutDSN(c.in); got != c.want { + t.Errorf("WithBusyTimeoutDSN(%q) = %q, want %q", c.in, got, c.want) } } } diff --git a/internal/vault/managed_keys_group_test.go b/internal/vault/managed_keys_group_test.go new file mode 100644 index 0000000..fd4ba64 --- /dev/null +++ b/internal/vault/managed_keys_group_test.go @@ -0,0 +1,116 @@ +package vault + +import ( + "database/sql" + "testing" + + _ "modernc.org/sqlite" +) + +// managedKeysSchema builds an in-memory managed_virtual_keys_cache. withGroup +// controls whether the oauth-group columns exist (false = older vault, exercises +// GetActiveManagedKeys' legacy fallback). +func newManagedKeysReader(t *testing.T, withGroup bool) *Reader { + t.Helper() + db, err := sql.Open("sqlite", ":memory:") + if err != nil { + t.Fatalf("open sqlite: %v", err) + } + t.Cleanup(func() { db.Close() }) + groupCols := "" + if withGroup { + // Mirrors the CLI migration (migrations.rs), which adds all five group + // columns in ONE batch — a vault with oauth_group_id always has + // my_assignment_override too, so the group-aware query may select both. + groupCols = ", oauth_group_id TEXT, group_accounts TEXT, group_runtime TEXT, routing_config TEXT, my_assignment_override TEXT" + } + if _, err := db.Exec(`CREATE TABLE managed_virtual_keys_cache ( + virtual_key_id TEXT PRIMARY KEY, alias TEXT NOT NULL, local_alias TEXT, + provider_code TEXT, protocol_type TEXT, base_url TEXT, + provider_key_nonce BLOB, provider_key_ciphertext BLOB, provider_base_urls TEXT, + org_id TEXT, seat_id TEXT, credential_id TEXT, credential_revision TEXT, + virtual_key_revision TEXT, owner_account_id TEXT, + key_status TEXT, local_state TEXT` + groupCols + `)`); err != nil { + t.Fatalf("create table: %v", err) + } + key := make([]byte, 32) + for i := range key { + key[i] = byte(i + 1) + } + return &Reader{db: db, derivedKey: key} +} + +func insertDirectBind(t *testing.T, r *Reader, vkID, plaintext string) { + t.Helper() + nonce, ct, err := Encrypt(r.derivedKey, []byte(plaintext)) + if err != nil { + t.Fatalf("encrypt: %v", err) + } + if _, err := r.db.Exec(`INSERT INTO managed_virtual_keys_cache + (virtual_key_id, alias, provider_code, protocol_type, base_url, + provider_key_nonce, provider_key_ciphertext, org_id, seat_id, + credential_id, credential_revision, virtual_key_revision, owner_account_id, key_status) + VALUES (?, ?, 'anthropic', 'anthropic', 'https://x', ?, ?, 'org', 'seat', 'cred', 'r', 'vr', 'acct', 'active')`, + vkID, vkID, nonce, ct); err != nil { + t.Fatalf("insert direct-bind: %v", err) + } +} + +func TestGetActiveManagedKeys_GroupAndDirectCoexist(t *testing.T) { + r := newManagedKeysReader(t, true) + insertDirectBind(t, r, "vk-direct", "sk-real-key") + // A group VK: NO provider_key_ciphertext; material is in group_runtime. + if _, err := r.db.Exec(`INSERT INTO managed_virtual_keys_cache + (virtual_key_id, alias, provider_code, protocol_type, base_url, org_id, seat_id, + credential_id, credential_revision, virtual_key_revision, owner_account_id, key_status, + oauth_group_id, group_accounts, group_runtime, routing_config) + VALUES ('vk-group', 'vk-group', 'anthropic', 'anthropic', '', 'org', 'seat-g', + '', '', 'vr', 'acct', 'active', + 'grp-1', '[{"account_id":"a1"}]', '{"a1":{"access_token":"enc:x"}}', '{"warn_ratio":2}')`); err != nil { + t.Fatalf("insert group VK: %v", err) + } + + keys, err := r.GetActiveManagedKeys() + if err != nil { + t.Fatalf("GetActiveManagedKeys: %v", err) + } + if len(keys) != 2 { + t.Fatalf("want 2 keys (direct + group), got %d: %+v", len(keys), keys) + } + byID := map[string]ManagedKey{} + for _, k := range keys { + byID[k.VirtualKeyID] = k + } + d := byID["vk-direct"] + if d.OauthGroupID != "" || d.PlaintextKey != "sk-real-key" { + t.Fatalf("direct-bind wrong: oauth_group=%q plaintext=%q", d.OauthGroupID, d.PlaintextKey) + } + g := byID["vk-group"] + if g.OauthGroupID != "grp-1" { + t.Fatalf("group oauth_group_id not read: %q", g.OauthGroupID) + } + if g.PlaintextKey != "" { + t.Fatalf("group VK must have EMPTY PlaintextKey (material in group_runtime): %q", g.PlaintextKey) + } + if g.GroupAccounts == "" || g.GroupRuntime == "" || g.RoutingConfig != `{"warn_ratio":2}` { + t.Fatalf("group fields not read: %+v", g) + } +} + +func TestGetActiveManagedKeys_LegacyVaultFallback(t *testing.T) { + // A vault WITHOUT the oauth-group columns: the group-aware query errors and + // GetActiveManagedKeys must fall back to the legacy query (never lose keys). + r := newManagedKeysReader(t, false) + insertDirectBind(t, r, "vk-direct", "sk-real-key") + + keys, err := r.GetActiveManagedKeys() + if err != nil { + t.Fatalf("GetActiveManagedKeys: %v", err) + } + if len(keys) != 1 || keys[0].VirtualKeyID != "vk-direct" || keys[0].PlaintextKey != "sk-real-key" { + t.Fatalf("legacy fallback failed to return direct-bind key: %+v", keys) + } + if keys[0].OauthGroupID != "" { + t.Fatalf("legacy vault must yield empty group fields: %q", keys[0].OauthGroupID) + } +} diff --git a/internal/vault/oauth_http_client.go b/internal/vault/oauth_http_client.go deleted file mode 100644 index 45a5808..0000000 --- a/internal/vault/oauth_http_client.go +++ /dev/null @@ -1,54 +0,0 @@ -package vault - -import ( - "context" - "fmt" - "net/url" - "time" - - broker "github.com/AiKeyLabs/aikey-auth-broker" - "github.com/imroc/req/v3" -) - -// ImpersonateChromeHTTPClient implements broker.OAuthHTTPClient with Chrome TLS -// fingerprint impersonation for Cloudflare bypass. -// -// Required by Claude's token endpoint (platform.claude.com/v1/oauth/token). -// Also works for Codex/Kimi endpoints (standard TLS is fine, but Chrome TLS is safe). -// -// Verified 2026-04-15: -// - Go net/http → 403/429 on Claude token endpoint (Cloudflare TLS detection) -// - ImpersonateChrome() → 200 (bypass) -// - Token endpoint requires MINIMAL headers: Content-Type + Accept + User-Agent: axios/1.13.6 -// - Adding Sec-Fetch-*/Origin/Referer → 429 rate limiting -type ImpersonateChromeHTTPClient struct { - client *req.Client -} - -// NewImpersonateChromeHTTPClient creates an HTTP client with Chrome TLS fingerprint. -func NewImpersonateChromeHTTPClient() *ImpersonateChromeHTTPClient { - c := req.C(). - SetTimeout(60 * time.Second). - ImpersonateChrome(). - SetCookieJar(nil) - return &ImpersonateChromeHTTPClient{client: c} -} - -func (c *ImpersonateChromeHTTPClient) PostForm(_ context.Context, targetURL string, values url.Values, _ ...broker.HTTPOption) (body []byte, statusCode int, err error) { - // Always use minimal headers for token endpoints (verified 2026-04-15). - // User-Agent: axios/1.13.6 matches the verified working pattern. - // No Sec-Fetch, no Origin, no Referer. - resp, err := c.client.R(). - SetHeader("Accept", "application/json, text/plain, */*"). - SetHeader("Content-Type", "application/x-www-form-urlencoded"). - SetHeader("User-Agent", "axios/1.13.6"). - SetFormDataFromValues(values). - Post(targetURL) - if err != nil { - return nil, 0, fmt.Errorf("http post: %w", err) - } - return resp.Bytes(), resp.StatusCode, nil -} - -// Compile-time interface check. -var _ broker.OAuthHTTPClient = (*ImpersonateChromeHTTPClient)(nil) diff --git a/internal/vault/vault.go b/internal/vault/vault.go index f081fad..d11a134 100644 --- a/internal/vault/vault.go +++ b/internal/vault/vault.go @@ -43,12 +43,17 @@ type Reader struct { derivedKey []byte } -// withBusyTimeoutDSN appends the busy_timeout pragma to a vault DB path, +// WithBusyTimeoutDSN appends the busy_timeout pragma to a vault DB path. +// Exported so other packages that open the SAME vault DB file by path (e.g. +// internal/quota WriteLocalUsage/WriteSubjects, internal/events) reuse the one +// busy_timeout convention — single source of truth, no raw opens of the vault file. +// (original doc continues below) +// WithBusyTimeoutDSN appends the busy_timeout pragma to a vault DB path, // preserving any query params the path already carries. modernc.org/sqlite // accepts `?_pragma=busy_timeout(5000)` (and `&_pragma=...` when other params // are present). The vault path is normally a bare filesystem path, but we // detect an existing `?` so callers passing a DSN don't get a malformed string. -func withBusyTimeoutDSN(dbPath string) string { +func WithBusyTimeoutDSN(dbPath string) string { const pragma = "_pragma=busy_timeout(5000)" if strings.Contains(dbPath, "?") { return dbPath + "&" + pragma @@ -67,7 +72,7 @@ func Open(dbPath, password string) (*Reader, error) { // lock. The default modernc busy_timeout is 0 (fail immediately, no retry); // 5s lets a momentary lock retry instead of bubbling up as a spurious // vault error. modernc.org/sqlite honors `_pragma` DSN query params. - db, err := sql.Open("sqlite", withBusyTimeoutDSN(dbPath)) + db, err := sql.Open("sqlite", WithBusyTimeoutDSN(dbPath)) if err != nil { return nil, fmt.Errorf("open vault db: %w", err) } @@ -257,6 +262,35 @@ type ManagedKey struct { CredentialRevision string VirtualKeyRevision string OwnerAccountID string + + // ── Oauth-group (channel ③) fields (N7c) ─────────────────────────────── + // Populated only for group VKs (binding.target = oauth_group). For these, + // PlaintextKey is EMPTY — the real per-account material lives in GroupRuntime + // and the route resolver (N8) picks a candidate account + injects its token. + // + // OauthGroupID != "" is the discriminator: empty = direct-bind VK (existing + // path, byte-unchanged); non-empty = group VK. + OauthGroupID string + // GroupAccounts: the candidate account list JSON (structural, CLI-synced) — + // [{account_id, identity, provider_code, priority, assigned}]. + GroupAccounts string + // GroupRuntime: per-account token/key material JSON (channel ③, proxy-pulled, + // encrypted at rest) — {account_id:{access_token ciphertext, expires, window}/ + // {key ciphertext, base_url, revision}}. NEVER refresh_token. + GroupRuntime string + // RoutingConfig: the group's routing knobs JSON (exhaustion_signals/util_cap/ + // ratios), structural-synced; the resolver reads it for classification. + RoutingConfig string + // MyAssignmentOverride: the engine's persisted seat→account routing assignment + // for THIS VK's (seat, group) — {"account_id","blocked","routing_version", + // "synced_at"} JSON, written by the routing_override SyncRail after each + // successful pull and hydrated into the in-memory override cache at startup + // (N12's persistence leg, completed 2026-07-03: without it a proxy restart + // forgot the engine assignment and the local ranked pick could demand login + // for an account the engine had routed the seat away from). Empty = no + // persisted assignment (local ranked pick). CLI never touches this column + // (fenced in its structural sync). + MyAssignmentOverride string } // GetActiveManagedKeys reads all team keys from managed_virtual_keys_cache that @@ -280,31 +314,51 @@ type ManagedKey struct { // different master password or corrupted data) so a single bad entry does not // block the proxy from starting. func (r *Reader) GetActiveManagedKeys() ([]ManagedKey, error) { - // 2026-05-09 alias surfacing: COALESCE(local_alias, alias) — the cache - // has two alias-like columns: - // - `alias` (NOT NULL): server-side canonical alias (e.g. - // "key-335923591-0011-1") set when the key was provisioned. Always - // populated. - // - `local_alias` (nullable, retrofit): per-client rename overlay - // populated by `aikey key rename`. NULL until the user explicitly - // renames. - // Use local_alias when set, fall back to alias. Either way the WAL / - // receipt now shows a meaningful name instead of the vk_id tail. + // Try the group-aware query first (reads the oauth_group columns added by + // N0/N6). Fall back to the legacy query on error — an older vault that lacks + // those columns must still load its direct-bind keys, never lose all keys + // over a missing column. Mirrors the CLI's column-cascade (storage_platform.rs). + keys, err := r.queryManagedKeys(true) + if err != nil { + keys, err = r.queryManagedKeys(false) + if err != nil { + // Table missing on very old vaults — treat as empty, not an error. + return nil, nil //nolint:nilerr // missing table → empty result, not a failure + } + } + return keys, nil +} + +// queryManagedKeys reads active managed keys. withGroup=true also selects the +// oauth-group columns (group VKs); false is the legacy shape for older vaults. +// +// 2026-05-09 alias surfacing: COALESCE(local_alias, alias) — the cache has two +// alias-like columns (`alias` server-canonical NOT NULL, `local_alias` nullable +// per-client rename overlay). Use local_alias when set, fall back to alias. +func (r *Reader) queryManagedKeys(withGroup bool) ([]ManagedKey, error) { + groupCols := "" + // Direct-bind keys require provider_key_ciphertext; group VKs have NONE + // (material rides group_runtime), so when reading group columns we also admit + // rows whose only target is a oauth_group. + targetFilter := "provider_key_ciphertext IS NOT NULL" + if withGroup { + groupCols = ", oauth_group_id, group_accounts, group_runtime, routing_config, my_assignment_override" + targetFilter = "(provider_key_ciphertext IS NOT NULL OR oauth_group_id IS NOT NULL)" + } rows, err := r.db.Query(` SELECT virtual_key_id, COALESCE(NULLIF(local_alias, ''), alias) AS effective_alias, provider_code, protocol_type, base_url, provider_key_nonce, provider_key_ciphertext, provider_base_urls, org_id, seat_id, credential_id, credential_revision, - virtual_key_revision, owner_account_id + virtual_key_revision, owner_account_id` + groupCols + ` FROM managed_virtual_keys_cache WHERE key_status = 'active' - AND provider_key_ciphertext IS NOT NULL + AND ` + targetFilter + ` AND COALESCE(local_state, '') != 'stale' AND COALESCE(local_state, '') NOT LIKE 'disabled_by_%' `) if err != nil { - // Table may not exist on older vaults — treat as empty, not an error. - return nil, nil //nolint:nilerr // older vaults lack this table; missing-table error means empty result, not a failure + return nil, err } defer rows.Close() @@ -316,56 +370,70 @@ func (r *Reader) GetActiveManagedKeys() ([]ManagedKey, error) { var providerBaseURLsJSON *string var orgID, seatID, credID, credRev, vkRev string var ownerAccountID *string - if err := rows.Scan(&vkID, &localAlias, &provCode, &protType, &baseURL, &nonce, &ciphertext, &providerBaseURLsJSON, - &orgID, &seatID, &credID, &credRev, &vkRev, &ownerAccountID); err != nil { + var oauthGroupID, groupAccounts, groupRuntime, routingConfig, myAssignmentOverride *string // group cols (nullable) + + dest := []any{&vkID, &localAlias, &provCode, &protType, &baseURL, &nonce, &ciphertext, &providerBaseURLsJSON, + &orgID, &seatID, &credID, &credRev, &vkRev, &ownerAccountID} + if withGroup { + dest = append(dest, &oauthGroupID, &groupAccounts, &groupRuntime, &routingConfig, &myAssignmentOverride) + } + if err := rows.Scan(dest...); err != nil { slog.Warn("managed key: scan error, skipping", "error", err) continue } - plaintext, err := Decrypt(r.derivedKey, nonce, ciphertext) - if err != nil { - // 2026-05-11: was WARN, upgraded to ERROR. A decrypt failure on - // the proxy's hot path means this team key was written with a - // vault_key that no longer matches the current password_hash — - // downstream requests routed via aikey_team_ will surface - // as 401 "Route token not found in registry" with no actionable - // signal. ERROR + an actionable hint forces the operator to see - // it (no silent registry shrink). See - // workflow/CI/bugfix/2026-05-11-team-key-decrypt-inconsistent.md. - slog.Error( - "managed key ciphertext was written with an inconsistent vault_key — proxy cannot decrypt. Recovery: aikey key sync --force-reencrypt", - "event.name", "vault.managed_key.decrypt_inconsistent", - "error.code", "VAULT_MANAGED_KEY_DECRYPT_INCONSISTENT", - "vk_id", vkID, - "error.message", err.Error(), - ) + + sgID := derefStr(oauthGroupID) + // Direct-bind keys decrypt their provider_key_ciphertext. Group VKs have + // NO ciphertext (material rides group_runtime, per-account) → skip decrypt, + // leave PlaintextKey empty; the route resolver (N8) picks an account. + var plaintextKey string + if len(ciphertext) > 0 { + plaintext, err := Decrypt(r.derivedKey, nonce, ciphertext) + if err != nil { + // 2026-05-11: ERROR (not WARN). A decrypt failure means this team + // key was written with a vault_key that no longer matches the + // current password_hash — requests via aikey_team_ 401 with + // no actionable signal. See + // workflow/CI/bugfix/2026-05-11-team-key-decrypt-inconsistent.md. + slog.Error( + "managed key ciphertext was written with an inconsistent vault_key — proxy cannot decrypt. Recovery: aikey key sync --force-reencrypt", + "event.name", "vault.managed_key.decrypt_inconsistent", + "error.code", "VAULT_MANAGED_KEY_DECRYPT_INCONSISTENT", + "vk_id", vkID, + "error.message", err.Error(), + ) + continue + } + plaintextKey = string(plaintext) + } else if sgID == "" { + // No ciphertext AND not a group VK — shouldn't pass the WHERE, but + // guard: skip rather than register an empty-key route. continue } + providerBaseURLs := make(map[string]string) if providerBaseURLsJSON != nil && *providerBaseURLsJSON != "" { _ = json.Unmarshal([]byte(*providerBaseURLsJSON), &providerBaseURLs) } - var accountID string - if ownerAccountID != nil { - accountID = *ownerAccountID - } - var aliasStr string - if localAlias != nil { - aliasStr = *localAlias - } keys = append(keys, ManagedKey{ VirtualKeyID: vkID, - LocalAlias: aliasStr, + LocalAlias: derefStr(localAlias), ProviderCode: provCode, ProtocolType: protType, BaseURL: baseURL, - PlaintextKey: string(plaintext), + PlaintextKey: plaintextKey, ProviderBaseURLs: providerBaseURLs, OrgID: orgID, SeatID: seatID, CredentialID: credID, CredentialRevision: credRev, VirtualKeyRevision: vkRev, - OwnerAccountID: accountID, + OwnerAccountID: derefStr(ownerAccountID), + OauthGroupID: sgID, + GroupAccounts: derefStr(groupAccounts), + GroupRuntime: derefStr(groupRuntime), + RoutingConfig: derefStr(routingConfig), + MyAssignmentOverride: derefStr(myAssignmentOverride), }) } if err := rows.Err(); err != nil { @@ -374,6 +442,14 @@ func (r *Reader) GetActiveManagedKeys() ([]ManagedKey, error) { return keys, nil } +// derefStr returns the pointed-at string, or "" when nil (nullable columns). +func derefStr(p *string) string { + if p == nil { + return "" + } + return *p +} + // GetPersonalKeyByAlias returns the plaintext key, provider_code, and custom base_url // for a personal key stored in the `entries` table. Uses the derivedKey already held // in memory from vault.Open — no additional authentication is needed. @@ -888,7 +964,7 @@ func readU32LEOrDefault(db *sql.DB, key string, defaultVal uint32) uint32 { // not yet exist. Opens a fresh read-write connection so callers do not need // an existing Reader (e.g. the Supervisor calling this before vault.Open). func ReadConfigU64LE(dbPath, key string) (uint64, error) { - db, err := sql.Open("sqlite", dbPath) + db, err := sql.Open("sqlite", WithBusyTimeoutDSN(dbPath)) if err != nil { return 0, fmt.Errorf("open vault db for config read: %w", err) } @@ -917,7 +993,7 @@ func ReadConfigU64LE(dbPath, key string) (uint64, error) { // runtime.source_identity for delivery-integrity event stamping — written by // the CLI side (storage.rs SOURCE_IDENTITY_KEY). func ReadConfigString(dbPath, key string) (string, error) { - db, err := sql.Open("sqlite", dbPath) + db, err := sql.Open("sqlite", WithBusyTimeoutDSN(dbPath)) if err != nil { return "", fmt.Errorf("open vault db for config read: %w", err) } @@ -1580,7 +1656,7 @@ func (r *Reader) GetActiveAppKeysForSlug(slug string) (int, error) { // vault config table (INSERT OR REPLACE). Opens its own connection so it can // be called independently of an existing Reader. func WriteConfigU64LE(dbPath, key string, value uint64) error { - db, err := sql.Open("sqlite", dbPath) + db, err := sql.Open("sqlite", WithBusyTimeoutDSN(dbPath)) if err != nil { return fmt.Errorf("open vault db for config write: %w", err) } @@ -1601,7 +1677,7 @@ func WriteConfigU64LE(dbPath, key string, value uint64) error { // table (INSERT OR REPLACE). Mirror of WriteConfigU64LE for non-numeric config // such as compliance.master_policy JSON. Opens its own connection. func WriteConfigString(dbPath, key, value string) error { - db, err := sql.Open("sqlite", dbPath) + db, err := sql.Open("sqlite", WithBusyTimeoutDSN(dbPath)) if err != nil { return fmt.Errorf("open vault db for config write: %w", err) } @@ -1613,3 +1689,57 @@ func WriteConfigString(dbPath, key, value string) error { } return nil } + +// WriteGroupRuntime updates the group_runtime column for one group VK (N7c-2). +// jsonValue is the proxy-built per-account material JSON (token/key ciphertext + +// window/meta). Only managed_virtual_keys_cache.group_runtime is touched — never +// the CLI-owned structural columns (which the CLI's sync fences this out of). A +// separate write connection mirrors WriteConfigString's pattern. +// +// busy_timeout is MANDATORY here (2026-07-03, Ubuntu client-leg regression): +// modernc's default busy_timeout is 0 — the UPDATE fails instantly with +// SQLITE_BUSY whenever any reader (CLI vault poll, second proxy connection) +// briefly holds the lock. The supervisor's 60s pull cycle then aborts and +// group VKs 503 "credentials are still syncing" for minutes on Linux, where +// the lock timing loses far more often than on macOS. Same guard applied to +// the config read/write helpers above — do not reintroduce raw dbPath opens. +func WriteGroupRuntime(dbPath, virtualKeyID, jsonValue string) error { + db, err := sql.Open("sqlite", WithBusyTimeoutDSN(dbPath)) + if err != nil { + return fmt.Errorf("open vault db for group_runtime write: %w", err) + } + defer db.Close() + + _, err = db.Exec( + "UPDATE managed_virtual_keys_cache SET group_runtime = ? WHERE virtual_key_id = ?", + jsonValue, virtualKeyID) + if err != nil { + return fmt.Errorf("write group_runtime for vk %q: %w", virtualKeyID, err) + } + return nil +} + +// WriteAssignmentOverride updates the my_assignment_override column for one +// group VK (the routing_override SyncRail's persistence leg, N12 completed +// 2026-07-03). jsonValue "" clears it (engine no longer overrides this seat). +// Only this proxy-owned column is touched — the CLI's structural sync fences it +// out, mirroring WriteGroupRuntime's ownership split. +// busy_timeout is MANDATORY (same 2026-07-03 regression class as WriteGroupRuntime +// above): a raw open here made this UPDATE fail instantly with SQLITE_BUSY under +// any concurrent reader, silently dropping the routing override (2026-07-07 parity +// audit P1-4). Do not reintroduce raw dbPath opens. +func WriteAssignmentOverride(dbPath, virtualKeyID, jsonValue string) error { + db, err := sql.Open("sqlite", WithBusyTimeoutDSN(dbPath)) + if err != nil { + return fmt.Errorf("open vault db for my_assignment_override write: %w", err) + } + defer db.Close() + + _, err = db.Exec( + "UPDATE managed_virtual_keys_cache SET my_assignment_override = ? WHERE virtual_key_id = ?", + jsonValue, virtualKeyID) + if err != nil { + return fmt.Errorf("write my_assignment_override for vk %q: %w", virtualKeyID, err) + } + return nil +} diff --git a/internal/vkeys/group_runtime.go b/internal/vkeys/group_runtime.go new file mode 100644 index 0000000..a9dcb1d --- /dev/null +++ b/internal/vkeys/group_runtime.go @@ -0,0 +1,92 @@ +// group_runtime.go — shared oauth-group routing contracts (N8). +// +// Two JSON shapes the oauth-group feature passes between components live here, +// in vkeys (the bottom of the proxy dependency graph), so the WRITER +// (supervisor.group_runtime_policy — pulls + encrypts) and the READER +// (proxy.group_resolve — ranks + decrypts + injects) share ONE definition +// instead of two structs that can silently drift. Per project rule +// "source-of-truth 不要分裂". +// +// Neither struct carries logic or non-stdlib imports — they are pure wire/at- +// rest data, so keeping them in vkeys does not pull new dependencies into the +// bottom layer. +package vkeys + +// GroupAccountRef is one candidate account in a seat group's routing set. It is +// the proxy-side mirror of master's snapshot.GroupAccountRef (JSON tags MUST +// match byte-for-byte) and arrives in ResolvedRoute.GroupAccounts, projected by +// master's RefreshSnapshot (N5). It carries only ranking inputs + display +// identity — NEVER secrets. The matching material (token/key) is in +// GroupRuntime, keyed by AccountID. +// +// NOTE on Weight: master ranks with seatassign.Account{AccountID, Priority} +// and leaves Weight unset (treated as 1 by seatassign). The proxy MUST do the +// same so its local ranking is byte-identical to master's — hence there is no +// Weight field to carry here. +type GroupAccountRef struct { + AccountID string `json:"account_id"` + Identity string `json:"identity"` // email / alias (display + audit only) + ProviderCode string `json:"provider_code"` // resolved provider code (injection dispatch) + Priority int `json:"priority"` // deterministic tie-break (lower wins) + Assigned bool `json:"assigned"` // master's rank-0 pick (advisory; proxy re-ranks) + // CredentialID is the account's real credential_id (same id the engine, + // auth-gate, and static-key track use). The proxy stamps it onto the resolved + // route so I5 usage signals enqueue keyed by credential_id instead of being + // dropped at the empty-id guard. json tag MUST stay byte-identical to master's + // snapshot.GroupAccountRef. Empty when talking to an old master (back-compat). + // Bugfix 2026-06-30: T2 uplink credential_id identity split. + CredentialID string `json:"credential_id,omitempty"` +} + +// GroupRuntimeAccount is one account's AT-REST material inside a group VK's +// managed_virtual_keys_cache.group_runtime column. The group_runtime value is a +// JSON map[account_id]GroupRuntimeAccount. The secret (access_token for OAuth / +// key for an API key) is AES-GCM encrypted with the vault derivedKey; nonce + +// ciphertext are base64 in the JSON. +// +// SECURITY: there is intentionally NO refresh_token field — master never +// delivers it, and the at-rest secret is always encrypted (the plaintext only +// exists transiently in memory between TLS receipt and re-encryption). +// +// WRITER: supervisor.buildGroupRuntimeJSON (N7c-2). READER: +// proxy.resolveGroupCredential (N8a). +type GroupRuntimeAccount struct { + CredentialType string `json:"credential_type"` // oauth_account | api_key + SecretNonce string `json:"secret_nonce"` // base64(nonce) + SecretCiphertext string `json:"secret_ciphertext"` // base64(enc(access_token|key)) + // Display meta (non-secret, 2026-07-01): identity/provider_code/priority carried on + // the fast rail so the client's candidate LIST membership refreshes here (a fast-rail- + // only account renders with its email/provider, not a bare UUID) — the CLI's + // merge_group_accounts_live treats this map as authoritative for membership. Priority + // is NOT omitempty (0 is a valid priority). + Identity string `json:"identity,omitempty"` + ProviderCode string `json:"provider_code,omitempty"` + Priority int `json:"priority"` + // NeedsLogin: the member has no token for this OAuth account (master said so) → + // the resolver returns LOGIN_REQUIRED for it. No secret when true. P1. + NeedsLogin bool `json:"needs_login,omitempty"` + // IsCurrentRouted: DISPLAY-ONLY (C2, 2026-06-30) — true on the ONE account this + // seat's traffic is currently routed to in steady state = routing-override ?? + // seatassign rank-0. NOT read by the resolver (which re-decides per request incl. + // cooldown failover); it exists only so /user/vault can show "current routed" + // distinct from the static "assigned" default. Recomputed on every material + // refresh AND on every routing-override change (couples the two 60s rails, owner- + // approved 2026-06-30). Per-VK (per-seat): the same group's VKs carry different + // flags. Deliberately excludes transient per-request cooldown failover (would + // flap; owner chose the stable pick). + IsCurrentRouted bool `json:"is_current_routed,omitempty"` + // OAuth meta: + ExpiresAt int64 `json:"expires_at,omitempty"` + WindowMaxUtilPct *int `json:"window_max_util_pct,omitempty"` + WindowStatus string `json:"window_status,omitempty"` + WindowResetAt *int64 `json:"window_reset_at,omitempty"` + // ExternalID is the OAuth provider's account UUID (e.g. Claude account.uuid). + // Claude's OAuth injection needs it to build a valid metadata.user_id; + // without it Claude returns a 429 business rejection. Codex uses AccountID + // instead and Kimi needs neither, so this is only load-bearing for Claude + // OAuth group accounts. Empty until master's N7a producer populates it. + ExternalID string `json:"external_id,omitempty"` + // KEY meta: + BaseURL string `json:"base_url,omitempty"` + Revision string `json:"revision,omitempty"` +} diff --git a/internal/vkeys/oauth_group_flag.go b/internal/vkeys/oauth_group_flag.go new file mode 100644 index 0000000..42f06b3 --- /dev/null +++ b/internal/vkeys/oauth_group_flag.go @@ -0,0 +1,34 @@ +package vkeys + +import "os" + +// Oauth-group (channel ③) routing feature flag — the single source of truth for +// the env gate, read by BOTH the supervisor (registry gate + group-runtime pull +// loop, N7c) and the proxy data plane (hot-path group resolver, N8). Keeping the +// env key in one place (here, the bottom of the dependency graph) avoids two +// packages drifting on the variable name. +// +// Group virtual keys carry no static PlaintextKey — their per-account material +// lives in managed_virtual_keys_cache.group_runtime and a request is routed by +// picking a candidate account (seatassign) + injecting its token. Until that +// path is complete + trusted, group VKs MUST NOT serve. Default OFF → the +// direct-bind path is byte-unchanged. + +// OauthGroupEnvKey enables proxy-side group VK routing. +const OauthGroupEnvKey = "AIKEY_PROXY_OAUTH_GROUP_ENABLED" + +// OauthGroupRoutingEnabled reports whether proxy-side group VK routing is on. +// +// Default ON since 2026-06-26 (user decision): oauth-group is the supported way an +// enterprise shares a credential pool, so a freshly-installed proxy must route +// group VKs without an extra env toggle — otherwise a user who was issued a group +// VK can't use it after a plain install. Set AIKEY_PROXY_OAUTH_GROUP_ENABLED=0 (or +// "false") to force it off. Builds with no group VKs are unaffected: the registry +// has no group VK to register and the group-runtime pull returns empty (no-op). +func OauthGroupRoutingEnabled() bool { + v := os.Getenv(OauthGroupEnvKey) + if v == "" { + return true + } + return v != "0" && v != "false" +} diff --git a/internal/vkeys/routed_pick.go b/internal/vkeys/routed_pick.go new file mode 100644 index 0000000..85105c7 --- /dev/null +++ b/internal/vkeys/routed_pick.go @@ -0,0 +1,126 @@ +// routed_pick.go — the SINGLE SOURCE OF TRUTH for "which pool account does this seat +// route to right now" on the PROXY side (2026-07-01, owner-approved unification). +// +// Consumers (both import this ONE function — never re-derive the pick): +// - proxy.resolveGroupCredential — the HOT PATH (what the request actually forwards to) +// - supervisor.computeRoutedAccountID — the is_current_routed DISPLAY stamp +// (/user/vault + /user/virtual-keys read it as current_routed) +// +// Together with the master side (poolroute.Resolve → the engine ledger feeding both the +// team-oauth page AND GET /accounts/me/routing → RoutingOverrideCache → the `override` +// input here), all four surfaces — proxy forwarding, vault, virtual-keys, team-oauth — +// resolve the SAME account by construction. +// +// OWNER RULE (2026-07-01): the engine is ALLOWED to route a member to an account they +// have NOT logged into. So an override naming a needs_login account is HONORED +// (PickNeedsLogin → the hot path returns LOGIN_REQUIRED for THAT account, the pages show +// that account with a "log in" prompt) — it is NOT treated as a stale redirect to fall +// through. Only a genuinely unusable override (not a candidate / no material / expired / +// window-exhausted / cooling) falls through to the local ranked pick. +package vkeys + +import "github.com/AiKeyLabs/pkg/seatassign" + +// PickOutcome is the 3-way result of picking a routed account. +type PickOutcome int + +const ( + // PickNone: no candidate is usable (all skipped/expired/exhausted/absent). + PickNone PickOutcome = iota + // PickOK: the returned account is usable right now (has deliverable material). + PickOK + // PickNeedsLogin: the returned account IS the routed account but the member has + // no token for it — callers prompt login for it (hot path: LOGIN_REQUIRED; display: + // stamp it, the login badge explains). + PickNeedsLogin +) + +// PickRoutedAccount resolves the ONE account a seat routes to within a group. +// +// refs — the group's candidate set (key-sync snapshot shape). +// material — the group_runtime material map (per-account, PLAINTEXT flags only are +// read; secrets untouched). EMPTY map = "material unknown" (the proxy +// hasn't pulled yet) → gates degrade to rank/override-only (blind mode), +// so the display can still stamp a nominal pick pre-poll. The hot path +// never passes an empty map (it errors NO_MATERIAL before picking). +// override — the engine's (seat,group) routing override ("" = none). Honored when +// the account is a candidate and not skipped — INCLUDING needs_login +// (owner rule above). Falls through when genuinely unusable. +// skip — accounts to route around (cooldown / this-request retries). nil ok. +// nowUnix — clock for the expiry gate (injected for deterministic tests). +// +// The ranked loop STOPS at the first non-skipped needs_login candidate (strict HRW — +// never silently hop past an account the member merely needs to log into; RW2/D2). +func PickRoutedAccount(seatID string, refs []GroupAccountRef, material map[string]GroupRuntimeAccount, override string, skip map[string]bool, nowUnix int64) (string, PickOutcome) { + if len(refs) == 0 { + return "", PickNone + } + accounts := make([]seatassign.Account, 0, len(refs)) + inSet := make(map[string]bool, len(refs)) + for _, r := range refs { + accounts = append(accounts, seatassign.Account{AccountID: r.AccountID, Priority: r.Priority}) + inSet[r.AccountID] = true + } + blind := len(material) == 0 // pre-poll display mode: no usability info yet + + gate := func(accountID string) PickOutcome { + if !inSet[accountID] || skip[accountID] { + return PickNone + } + if blind { + return PickOK + } + mat, ok := material[accountID] + if !ok { + return PickNone // material not delivered (yet) — retryable skip, not a login prompt + } + if mat.NeedsLogin { + return PickNeedsLogin + } + if !MaterialUsable(mat, nowUnix) { + return PickNone // expired / window-exhausted + } + return PickOK + } + + // Engine override first (owner rule: needs_login is a valid, honored destination). + if override != "" { + if oc := gate(override); oc != PickNone { + return override, oc + } + } + // Local ranked pick — stop at the first usable OR needs_login candidate. + for _, a := range seatassign.Rank(seatID, accounts) { + if oc := gate(a.AccountID); oc != PickNone { + return a.AccountID, oc + } + } + return "", PickNone +} + +// MaterialUsable reports whether an account's material can serve a request now. +// OAuth: not past expiry and quota window not exhausted. API key: always usable if +// present (no expiry/window in the contract). Shared by the pick gate above and the +// hot path's post-pick sanity. +func MaterialUsable(mat GroupRuntimeAccount, nowUnix int64) bool { + if mat.CredentialType == "api_key" { + return true + } + if mat.ExpiresAt > 0 && mat.ExpiresAt <= nowUnix { + return false // access_token expired (refresh is master's job — N7b) + } + if mat.WindowStatus == "exhausted" { + return false // oauth-group quota window used up — route around it + } + return true +} + +// MaterialExpired reports whether an OAuth account's material is stale +// SPECIFICALLY because the member's access token passed its expiry. R36 +// (2026-07-04): expiry is MEMBER-fixable — re-logging in mints a new token — so +// the resolver's dead-end classification uses this to prompt a 401 re-login +// instead of the admin-facing ALL_UNUSABLE 503. Window-exhausted is deliberately +// NOT this: only routing around (or waiting) fixes it. API keys never expire. +func MaterialExpired(mat GroupRuntimeAccount, nowUnix int64) bool { + return mat.CredentialType != "api_key" && mat.ExpiresAt > 0 && mat.ExpiresAt <= nowUnix +} diff --git a/internal/vkeys/routed_pick_test.go b/internal/vkeys/routed_pick_test.go new file mode 100644 index 0000000..ab4e1a1 --- /dev/null +++ b/internal/vkeys/routed_pick_test.go @@ -0,0 +1,89 @@ +package vkeys + +import ( + "testing" + + "github.com/AiKeyLabs/pkg/seatassign" +) + +// PickRoutedAccount is the proxy-side single source of truth for "which account does +// this seat route to" — shared by the hot-path resolver AND the display stamp. These +// cases lock its contract (esp. the 2026-07-01 owner rule: an engine override naming a +// needs_login account is HONORED, not treated as stale). +func TestPickRoutedAccount(t *testing.T) { + refs := []GroupAccountRef{ + {AccountID: "acc-a", Priority: 1}, + {AccountID: "acc-b", Priority: 1}, + } + fresh := GroupRuntimeAccount{CredentialType: "oauth_account", ExpiresAt: 9_000_000_000} + needsLogin := GroupRuntimeAccount{CredentialType: "oauth_account", NeedsLogin: true} + now := int64(1_000_000) + hrw := seatassign.Primary("seat-1", []seatassign.Account{{AccountID: "acc-a", Priority: 1}, {AccountID: "acc-b", Priority: 1}}) + other := "acc-a" + if hrw == "acc-a" { + other = "acc-b" + } + + t.Run("owner rule: needs_login override is HONORED (PickNeedsLogin, that account)", func(t *testing.T) { + mat := map[string]GroupRuntimeAccount{hrw: fresh, other: needsLogin} + acc, oc := PickRoutedAccount("seat-1", refs, mat, other, nil, now) + if acc != other || oc != PickNeedsLogin { + t.Fatalf("engine may route to a not-logged-in account: want (%q, NeedsLogin), got (%q, %v)", other, acc, oc) + } + }) + + t.Run("usable override redirects off the HRW pick", func(t *testing.T) { + mat := map[string]GroupRuntimeAccount{hrw: fresh, other: fresh} + if acc, oc := PickRoutedAccount("seat-1", refs, mat, other, nil, now); acc != other || oc != PickOK { + t.Fatalf("want (%q, OK), got (%q, %v)", other, acc, oc) + } + }) + + t.Run("genuinely unusable override falls through (expired / not-a-candidate / cooled)", func(t *testing.T) { + expired := GroupRuntimeAccount{CredentialType: "oauth_account", ExpiresAt: now - 1} + mat := map[string]GroupRuntimeAccount{hrw: fresh, other: expired} + if acc, oc := PickRoutedAccount("seat-1", refs, mat, other, nil, now); acc != hrw || oc != PickOK { + t.Fatalf("expired override must fall through to HRW %q, got (%q, %v)", hrw, acc, oc) + } + if acc, _ := PickRoutedAccount("seat-1", refs, mat, "ghost", nil, now); acc != hrw { + t.Fatalf("non-candidate override must fall through, got %q", acc) + } + mat[other] = fresh + if acc, _ := PickRoutedAccount("seat-1", refs, mat, other, map[string]bool{other: true}, now); acc != hrw { + t.Fatalf("cooled override must fall through, got %q", acc) + } + }) + + t.Run("ranked loop stops at first needs_login (strict HRW, RW2)", func(t *testing.T) { + mat := map[string]GroupRuntimeAccount{hrw: needsLogin, other: fresh} + if acc, oc := PickRoutedAccount("seat-1", refs, mat, "", nil, now); acc != hrw || oc != PickNeedsLogin { + t.Fatalf("must stop at needs_login rank-0 %q (not hop to %q), got (%q, %v)", hrw, other, acc, oc) + } + }) + + t.Run("no-material candidate is a retryable skip, not a login prompt", func(t *testing.T) { + mat := map[string]GroupRuntimeAccount{other: fresh} // hrw's material not delivered yet + if acc, oc := PickRoutedAccount("seat-1", refs, mat, "", nil, now); acc != other || oc != PickOK { + t.Fatalf("undelivered rank-0 must be skipped to %q, got (%q, %v)", other, acc, oc) + } + }) + + t.Run("blind mode (empty material): rank/override-only, for pre-poll display", func(t *testing.T) { + if acc, oc := PickRoutedAccount("seat-1", refs, nil, "", nil, now); acc != hrw || oc != PickOK { + t.Fatalf("blind rank-0: want %q, got (%q, %v)", hrw, acc, oc) + } + if acc, _ := PickRoutedAccount("seat-1", refs, nil, other, nil, now); acc != other { + t.Fatalf("blind override: want %q, got %q", other, acc) + } + }) + + t.Run("all unusable → PickNone", func(t *testing.T) { + mat := map[string]GroupRuntimeAccount{ + "acc-a": {CredentialType: "oauth_account", ExpiresAt: now - 1}, + "acc-b": {CredentialType: "oauth_account", WindowStatus: "exhausted"}, + } + if acc, oc := PickRoutedAccount("seat-1", refs, mat, "", nil, now); oc != PickNone { + t.Fatalf("want PickNone, got (%q, %v)", acc, oc) + } + }) +} diff --git a/internal/vkeys/types.go b/internal/vkeys/types.go index bd0ca37..db3d99c 100644 --- a/internal/vkeys/types.go +++ b/internal/vkeys/types.go @@ -170,6 +170,19 @@ type ResolvedRoute struct { // own active profile). first-party only; see AppKind above for the // defense-in-depth check. FollowUserActive bool + + // ── Oauth-group (channel ③) route fields (N7c) ───────────────────────── + // OauthGroupID != "" marks a GROUP virtual key: PlaintextKey is empty and the + // dispatch resolver (N8) picks a candidate account from GroupAccounts (ranked + // via pkg/seatassign for SeatID), reads its token from GroupRuntime, and + // injects it. Empty for direct-bind VKs → existing path, byte-unchanged. + OauthGroupID string + // GroupAccounts: candidate account list JSON (structural). GroupRuntime: + // per-account token/key material JSON (channel ③, encrypted, proxy-pulled; + // NEVER refresh_token). RoutingConfig: the group's routing knobs JSON. + GroupAccounts string + GroupRuntime string + RoutingConfig string } // IsModelAllowed checks if the given model is permitted by this route. diff --git a/pkg/observer/registry.go b/pkg/observer/registry.go index 16fbcb7..1ea7c40 100644 --- a/pkg/observer/registry.go +++ b/pkg/observer/registry.go @@ -158,6 +158,15 @@ type RequestContext struct { // Empty OrgID means a personal (non-org) request. OrgID string OwnerAccountID string + // SeatID is the org seat this request's HUMAN belongs to, from + // route.SeatID — the same field usage events carry (reportable.go). Why + // it exists SEPARATELY from OwnerAccountID (2026-07-07, org 624a2488 + // live incident): for a shared OAuth-pool VK the route's AccountID is + // the VK OWNER (pool creator), not the employee at the terminal — + // audit records keyed only on owner filed under a stranger seat while + // usage (seat-keyed UI) attributed correctly. Empty for personal keys + // and legacy vault rows that pre-date seat stamping. + SeatID string // Stream is the SPEC §1.4.1 stream name this request emits under // (StreamUserChat / StreamAppPipeline / StreamProbe). Set by the // Registry on the way into NotifyStart so observers can read it