diff --git a/.gitignore b/.gitignore index 25906042..3dbd4394 100644 --- a/.gitignore +++ b/.gitignore @@ -22,3 +22,9 @@ Thumbs.db # Local backups / scratch backup*/ PR-*.md +# /test.py: local scratch script holding a live API key (never commit) +/test.py +# editor backup files (e.g. kiro-go.exe~) +*~ +# local smoke-test output (dispatch runs, cached prompts, r-*.out) +.smoke/ diff --git a/Dockerfile b/Dockerfile index b2c56422..ea8b01ab 100644 --- a/Dockerfile +++ b/Dockerfile @@ -23,6 +23,8 @@ COPY --from=builder /app/web ./web RUN mkdir -p /app/data EXPOSE 8080 +# Enterprise SSO (Microsoft 365) loopback callback port — see docker-compose.yml. +EXPOSE 3128 VOLUME /app/data CMD ["./kiro-go"] diff --git a/README.md b/README.md index 67e96093..54ebe0b6 100644 --- a/README.md +++ b/README.md @@ -117,6 +117,10 @@ The setting takes effect immediately without restarting. Friendly discussion is welcome. If you run into issues, try asking Claude Code, Codex, or similar tools for help first — most problems can be solved that way. PRs are even better. +## Contact + +Telegram: [@tutua16888](https://t.me/tutua16888) + ## Friend Links - [LINUX DO](https://linux.do) diff --git a/README_CN.md b/README_CN.md index 5f51fe38..88a53bcf 100644 --- a/README_CN.md +++ b/README_CN.md @@ -117,6 +117,10 @@ curl http://localhost:8080/v1/chat/completions \ 欢迎友好交流。遇到问题时，建议先让 Claude Code、Codex 等工具帮忙排查一下，大部分问题都能自己解决。如果能直接提个 PR 就更好了。 +## 联系方式 + +Telegram：[@tutua16888](https://t.me/tutua16888) + ## 友情链接 - [LINUX DO](https://linux.do) diff --git a/auth/kiro_sso.go b/auth/kiro_sso.go new file mode 100644 index 00000000..e968d747 --- /dev/null +++ b/auth/kiro_sso.go @@ -0,0 +1,906 @@ +package auth + +// kiro_sso.go implements the Kiro hosted browser sign-in flow — the same flow +// the Kiro IDE uses at https://app.kiro.dev/signin. Unlike Builder ID / IAM +// Identity Center (AWS SSO OIDC), this portal federates Google, GitHub, AND +// enterprise identity providers (e.g. a Microsoft 365 / Entra ID / Azure AD +// tenant) behind a single PKCE authorization-code flow. It is the only way an +// enterprise Azure-tenant account — which is neither an AWS Builder ID nor an +// AWS IAM Identity Center account — can sign in to Kiro. +// +// The flow has two possible legs, both captured by one transient loopback +// listener bound on the fixed redirect port: +// +// - Social (Google/GitHub): the portal authenticates via its Cognito backend +// and redirects the authorization code straight back to the loopback +// redirect. The code is exchanged at the Kiro social token endpoint. +// +// - Enterprise / external IdP (Azure AD): the portal detects the email belongs +// to an external IdP and redirects to /signin/callback with the IdP +// descriptor (issuer_url, client_id, scopes) instead of a code. We then drive +// a SECOND OIDC authorization-code+PKCE flow directly against that IdP +// (loopback redirect to /oauth/callback) and exchange the code at the IdP +// token endpoint. The resulting access token is an IdP-issued token scoped +// for CodeWhisperer; it is used as the runtime bearer and refreshed against +// the IdP token endpoint (see refreshExternalIdpToken in oidc.go). +// +// The login is exposed through the admin panel with the same Start/Poll session +// pattern as Builder ID: StartKiroSsoLogin binds the listener and returns the +// sign-in URL; the operator opens it in a browser ON THE SAME HOST (the redirect +// targets 127.0.0.1:3128); PollKiroSsoAuth reports pending until the listener +// captures the code, then exchanges it and returns the credential. + +import ( + "bytes" + "encoding/base64" + "encoding/json" + "fmt" + "io" + "kiro-go/config" + "kiro-go/logger" + "net" + "net/http" + "net/url" + "os" + "strings" + "sync" + "time" + + "github.com/google/uuid" +) + +const ( + // kiroSignInBaseURL is the Kiro hosted sign-in page opened in the browser. + kiroSignInBaseURL = "https://app.kiro.dev/signin" + // kiroRedirectURI is the fixed loopback redirect the portal validates and + // redirects back to once sign-in succeeds. The host is "localhost" (the value + // the portal expects) while the listener binds the 127.0.0.1 / [::1] literals; + // the browser resolving "localhost" to either loopback address is what bridges + // the two, so the operator's host must resolve localhost to loopback. + kiroRedirectURI = "http://localhost:3128" + // kiroRedirectPort is the loopback port embedded in kiroRedirectURI. + kiroRedirectPort = "3128" + // kiroRedirectFrom mirrors the Kiro IDE client tag the portal expects. + kiroRedirectFrom = "KiroIDE" + // kiroOAuthCallbackPath is the loopback path the enterprise (external IdP) + // leg redirects the authorization code back to. It is distinct from the + // portal's /signin/callback so the listener can tell the two legs apart. + kiroOAuthCallbackPath = "/oauth/callback" + // kiroSocialTokenURL is the Cognito-backed social code-exchange endpoint. Note + // this is deliberately a DIFFERENT path from the social refresh endpoint + // (socialTokenURL() in oidc.go -> /refreshToken): the Kiro IDE exchanges the + // login code at /oauth/token and refreshes at /refreshToken. Do not unify them. + kiroSocialTokenURL = "https://prod.us-east-1.auth.desktop.kiro.dev/oauth/token" + // kiroSsoLoginTimeout bounds how long the listener waits for the user. + kiroSsoLoginTimeout = 10 * time.Minute +) + +// allowedExternalIdpIssuerSuffixes restricts which IdP issuer/endpoint hosts the +// enterprise leg will discover and redirect to. The issuer arrives in an +// attacker-influenceable portal callback query, so it is constrained to known +// enterprise IdP hosts (Microsoft Entra / Azure AD — the supported provider). +// This is the primary control against SSRF, open-redirect, and forced-auth abuse +// via a forged /signin/callback. The leading dot anchors each suffix to a real +// subdomain boundary so "evil-microsoftonline.com" cannot match. Extend this +// list to onboard additional enterprise IdPs. +var allowedExternalIdpIssuerSuffixes = []string{ + ".microsoftonline.com", + ".microsoftonline.us", + ".microsoftonline.cn", +} + +// KiroSsoSession holds the transient state for one hosted-portal sign-in attempt. +type KiroSsoSession struct { + ID string + Verifier string // social-leg PKCE verifier (sent at social code exchange) + State string // portal anti-CSRF state echoed on the social redirect + Region string + ProxyURL string + ExpiresAt time.Time + + srv *http.Server + resultCh chan kiroSsoCapture + once sync.Once + closeOnce sync.Once + timer *time.Timer // deadline self-teardown; freed in close() + + mu sync.Mutex + leg2 *kiroLeg2 // set when the enterprise descriptor arrives +} + +// kiroLeg2 is the per-attempt state captured when the enterprise descriptor +// arrives at /signin/callback and consumed when the IdP redirects the code back. +type kiroLeg2 struct { + state string + verifier string + tokenEndpoint string + issuerURL string + clientID string + scopes string + redirectURI string +} + +// kiroSsoCapture is the raw outcome delivered by the loopback listener: either a +// social authorization code, an enterprise (external IdP) code plus its leg-2 +// context, or an error. +type kiroSsoCapture struct { + kind string // "social" | "external_idp" + code string + err error + + tokenEndpoint string + issuerURL string + clientID string + scopes string + redirectURI string + codeVerifier string +} + +// KiroSsoResult is the resolved credential returned to the admin handler once the +// captured code has been exchanged for tokens. +type KiroSsoResult struct { + AccessToken string + RefreshToken string + AuthMethod string // "external_idp" | "social" + Provider string // "AzureAD" | "Kiro SSO" + ClientID string // external IdP client id (refresh material) + TokenEndpoint string // external IdP token endpoint (refresh material) + IssuerURL string + Scopes string + ProfileArn string // social exchange may return it; external IdP resolves lazily + Region string + ExpiresIn int + Email string +} + +var ( + kiroSsoSessions = make(map[string]*KiroSsoSession) + kiroSsoSessionsMu sync.RWMutex +) + +// StartKiroSsoLogin generates PKCE codes, binds the loopback listener, and +// returns the session plus the hosted sign-in URL the operator must open. +func StartKiroSsoLogin(region string) (*KiroSsoSession, string, error) { + // Free the callback port from any previous abandoned session — only one + // sign-in can bind 127.0.0.1:3128 at a time. + cancelAllKiroSsoSessions() + + if region == "" { + region = "us-east-1" + } + + verifier := generateCodeVerifier() + challenge := generateCodeChallenge(verifier) + state := uuid.New().String() + + session := &KiroSsoSession{ + ID: uuid.New().String(), + Verifier: verifier, + State: state, + Region: region, + ProxyURL: config.GetProxyURL(), + ExpiresAt: time.Now().Add(kiroSsoLoginTimeout), + resultCh: make(chan kiroSsoCapture, 1), + } + + if err := session.startListener(); err != nil { + return nil, "", err + } + + params := url.Values{} + params.Set("state", state) + params.Set("code_challenge", challenge) + params.Set("code_challenge_method", "S256") + params.Set("redirect_uri", kiroRedirectURI) + params.Set("redirect_from", kiroRedirectFrom) + signInURL := kiroSignInBaseURL + "?" + params.Encode() + + kiroSsoSessionsMu.Lock() + kiroSsoSessions[session.ID] = session + kiroSsoSessionsMu.Unlock() + + // Self-teardown at the deadline: free the loopback listener and drop the + // session even if the operator abandons the sign-in and the front end stops + // polling. Without this an abandoned login would hold 127.0.0.1:3128 until the + // process restarts and block every subsequent SSO login (the redirect port is + // fixed, so only one sign-in can use it at a time). + session.timer = time.AfterFunc(kiroSsoLoginTimeout, func() { + session.close() + removeKiroSsoSession(session.ID) + }) + + return session, signInURL, nil +} + +// PollKiroSsoAuth reports the login status. It returns ("pending", nil) until the +// listener captures a code, then exchanges it and returns the resolved +// credential with status "completed". A terminal error (timeout, exchange +// failure) is returned as a non-nil error. +func PollKiroSsoAuth(sessionID string) (*KiroSsoResult, string, error) { + kiroSsoSessionsMu.RLock() + session, ok := kiroSsoSessions[sessionID] + kiroSsoSessionsMu.RUnlock() + if !ok { + return nil, "", fmt.Errorf("session not found or expired") + } + + select { + case capture := <-session.resultCh: + // Terminal: a code (or error) was captured. Tear the listener down and + // drop the session regardless of exchange outcome. + session.close() + removeKiroSsoSession(sessionID) + if capture.err != nil { + return nil, "", capture.err + } + return session.exchange(capture) + default: + if time.Now().After(session.ExpiresAt) { + session.close() + removeKiroSsoSession(sessionID) + return nil, "", fmt.Errorf("SSO login timed out after %s", kiroSsoLoginTimeout) + } + return nil, "pending", nil + } +} + +// exchange swaps a captured authorization code for tokens and assembles the +// resolved credential. +func (s *KiroSsoSession) exchange(capture kiroSsoCapture) (*KiroSsoResult, string, error) { + client := GetAuthClientForProxy(s.ProxyURL) + + if capture.kind == "external_idp" { + access, refresh, expiresIn, err := exchangeExternalIdpCode( + client, capture.tokenEndpoint, capture.clientID, capture.code, + capture.codeVerifier, capture.redirectURI, capture.scopes, + ) + if err != nil { + return nil, "", fmt.Errorf("enterprise SSO token exchange failed: %w", err) + } + return &KiroSsoResult{ + AccessToken: access, + RefreshToken: refresh, + AuthMethod: "external_idp", + Provider: "AzureAD", + ClientID: capture.clientID, + TokenEndpoint: capture.tokenEndpoint, + IssuerURL: capture.issuerURL, + Scopes: capture.scopes, + Region: s.Region, + ExpiresIn: expiresIn, + Email: ExtractEmailFromJWT(access), + }, "completed", nil + } + + access, refresh, expiresIn, profileArn, err := exchangeSocialCode(client, capture.code, s.Verifier) + if err != nil { + return nil, "", fmt.Errorf("SSO token exchange failed: %w", err) + } + return &KiroSsoResult{ + AccessToken: access, + RefreshToken: refresh, + AuthMethod: "social", + Provider: "Kiro SSO", + ProfileArn: profileArn, + Region: s.Region, + ExpiresIn: expiresIn, + Email: ExtractEmailFromJWT(access), + }, "completed", nil +} + +// --- Loopback callback listener (state machine across the legs) ------------- + +// kiroCallbackBindAddrs returns the address(es) the SSO callback listener binds. +// +// By default it binds loopback only — IPv4 127.0.0.1 plus, best-effort, IPv6 ::1 +// (a browser resolving "localhost" may use either) — which is the secure default: +// only the same host can reach the transient callback. Set KIRO_SSO_CALLBACK_BIND +// to override the bind host; this is needed when the proxy runs in a container and +// the operator's browser reaches the published port on a non-loopback interface +// (e.g. KIRO_SSO_CALLBACK_BIND=0.0.0.0 with `3128:3128` published in compose). The +// callback is transient (closed at the deadline) and every leg is anti-CSRF +// state-matched, but a non-loopback bind does expose it on the network for the +// login window, so only set it in a trusted/containerized network. +func kiroCallbackBindAddrs() []string { + if bind := strings.TrimSpace(os.Getenv("KIRO_SSO_CALLBACK_BIND")); bind != "" { + return []string{net.JoinHostPort(bind, kiroRedirectPort)} + } + return []string{"127.0.0.1:" + kiroRedirectPort, "[::1]:" + kiroRedirectPort} +} + +// startListener binds the SSO callback listener(s) on the fixed redirect port and +// serves the redirect state machine. The first address is mandatory (its bind +// failure aborts the login); any remaining addresses are best-effort (e.g. the IPv6 +// loopback when IPv6 is unavailable). +func (s *KiroSsoSession) startListener() error { + addrs := kiroCallbackBindAddrs() + + ln, err := net.Listen("tcp", addrs[0]) + if err != nil { + return fmt.Errorf("cannot bind %s for the SSO callback (is the port already in use?): %w", addrs[0], err) + } + + mux := http.NewServeMux() + mux.HandleFunc("/", s.handleCallback) + // ReadHeaderTimeout bounds a stalled local client (slowloris-style). + s.srv = &http.Server{Handler: mux, ReadHeaderTimeout: 5 * time.Second} + + serve := func(l net.Listener) { + go func() { + if errServe := s.srv.Serve(l); errServe != nil && errServe != http.ErrServerClosed { + logger.Debugf("[KiroSSO] callback listener (%s) stopped: %v", l.Addr(), errServe) + } + }() + } + serve(ln) + for _, addr := range addrs[1:] { + if extra, errExtra := net.Listen("tcp", addr); errExtra == nil { + serve(extra) + } else { + logger.Debugf("[KiroSSO] secondary callback bind %s skipped: %v", addr, errExtra) + } + } + return nil +} + +// close shuts the loopback listener(s) down and stops the deadline timer. Safe to +// call more than once and from multiple goroutines (Poll, the cancel endpoint, and +// the deadline AfterFunc may all race to tear the session down). +func (s *KiroSsoSession) close() { + s.closeOnce.Do(func() { + if s.timer != nil { + s.timer.Stop() + } + if s.srv != nil { + _ = s.srv.Close() + } + }) +} + +// cancelAllKiroSsoSessions tears down every in-flight session, freeing the +// loopback callback port. Called whenever a new SSO login starts so that a +// stale/abandoned session never blocks the next attempt. +func cancelAllKiroSsoSessions() { + kiroSsoSessionsMu.Lock() + defer kiroSsoSessionsMu.Unlock() + for id, session := range kiroSsoSessions { + session.close() + delete(kiroSsoSessions, id) + } +} + +// CancelKiroSsoLogin tears an in-flight session down immediately (operator +// cancelled in the admin panel), freeing the loopback port without waiting for the +// deadline. A no-op for an unknown or already-finished session. +func CancelKiroSsoLogin(sessionID string) { + kiroSsoSessionsMu.RLock() + session, ok := kiroSsoSessions[sessionID] + kiroSsoSessionsMu.RUnlock() + if !ok { + return + } + session.close() + removeKiroSsoSession(sessionID) +} + +// FeedCallbackURL lets an operator paste a callback URL that the browser could not +// deliver (e.g. in a remote deployment where localhost isn't the proxy host). Parses +// the URL, constructs a synthetic GET request, and feeds it through the session's +// callback state machine. For enterprise SSO leg-1 (external IdP descriptor) the +// return value is the IdP authorize URL the operator should open next; for leg-2 and +// social logins it returns "". +func FeedCallbackURL(sessionID, rawURL string) (string, error) { + kiroSsoSessionsMu.RLock() + session, ok := kiroSsoSessions[sessionID] + kiroSsoSessionsMu.RUnlock() + if !ok { + return "", fmt.Errorf("session %q not found (expired or never started)", sessionID) + } + u, err := url.Parse(rawURL) + if err != nil { + return "", fmt.Errorf("invalid callback URL: %w", err) + } + req := &http.Request{ + Method: http.MethodGet, + URL: u, + Header: http.Header{}, + } + rec := &redirectCapturingWriter{} + session.handleCallback(rec, req) + return rec.redirectURL, nil +} + +// redirectCapturingWriter records the Location header from a 302 redirect and +// discards everything else. Used by FeedCallbackURL so that enterprise SSO leg-1 +// can return the IdP authorize URL to the operator without a real browser redirect. +type redirectCapturingWriter struct { + hdr http.Header + redirectURL string +} + +func (w *redirectCapturingWriter) Header() http.Header { + if w.hdr == nil { + w.hdr = http.Header{} + } + return w.hdr +} +func (w *redirectCapturingWriter) Write(b []byte) (int, error) { + return len(b), nil +} +func (w *redirectCapturingWriter) WriteHeader(statusCode int) { + if statusCode == http.StatusFound && w.hdr != nil { + w.redirectURL = w.hdr.Get("Location") + } +} + +// discardResponseWriter is an http.ResponseWriter that silently discards everything. +// Used by FeedCallbackURL so the callback state machine doesn't try to write redirects +// or HTML pages to a real HTTP response. +type discardResponseWriter struct{} + +func (discardResponseWriter) Header() http.Header { return http.Header{} } +func (discardResponseWriter) Write(b []byte) (int, error) { return len(b), nil } +func (discardResponseWriter) WriteHeader(statusCode int) {} + +// deliver pushes the first (and only) capture onto the result channel. +func (s *KiroSsoSession) deliver(capture kiroSsoCapture) { + s.once.Do(func() { s.resultCh <- capture }) +} + +// handleCallback implements the redirect state machine: enterprise leg-1 +// descriptor -> 302 to the IdP; enterprise leg-2 code at /oauth/callback; +// otherwise the social code. +func (s *KiroSsoSession) handleCallback(w http.ResponseWriter, req *http.Request) { + // Only browser GET redirects are expected; reject other methods to shrink + // the local attack surface. + if req.Method != http.MethodGet { + w.WriteHeader(http.StatusMethodNotAllowed) + return + } + q := req.URL.Query() + + // --- Enterprise leg-1: external IdP descriptor (no code) --- + // Gate on path != /oauth/callback so a forged /oauth/callback?issuer_url=... + // cannot be routed here and reset an in-flight leg-2. + if req.URL.Path != kiroOAuthCallbackPath && + (strings.EqualFold(strings.TrimSpace(q.Get("login_option")), "external_idp") || strings.TrimSpace(q.Get("issuer_url")) != "") { + // Single-shot: once a leg-2 is in flight, ignore further descriptors so a + // stray or forged local request cannot reset/hijack the active login. + s.mu.Lock() + alreadyStarted := s.leg2 != nil + s.mu.Unlock() + if alreadyStarted { + w.WriteHeader(http.StatusNoContent) + return + } + issuerURL := strings.TrimSpace(q.Get("issuer_url")) + clientID := strings.TrimSpace(q.Get("client_id")) + scopes := strings.TrimSpace(q.Get("scopes")) + loginHint := strings.TrimSpace(q.Get("login_hint")) + if clientID == "" { + writeKiroCallbackPage(w, false) + s.deliver(kiroSsoCapture{err: fmt.Errorf("invalid external IdP descriptor (missing client_id)")}) + return + } + // oidcDiscover validates the issuer + both discovered endpoints against + // the IdP host allow-list, so the issuer here is not trusted blindly. + authEndpoint, tokenEndpoint, errDisc := oidcDiscover(GetAuthClientForProxy(s.ProxyURL), issuerURL, s.ProxyURL) + if errDisc != nil { + writeKiroCallbackPage(w, false) + s.deliver(kiroSsoCapture{err: errDisc}) + return + } + verifier := generateCodeVerifier() + state2 := uuid.New().String() + redirectURI := kiroRedirectURI + kiroOAuthCallbackPath + s.mu.Lock() + // Re-check under the lock to resolve a race between concurrent + // descriptors: only the first sets leg2 and is redirected. + if s.leg2 != nil { + s.mu.Unlock() + w.WriteHeader(http.StatusNoContent) + return + } + s.leg2 = &kiroLeg2{ + state: state2, + verifier: verifier, + tokenEndpoint: tokenEndpoint, + issuerURL: issuerURL, + clientID: clientID, + scopes: scopes, + redirectURI: redirectURI, + } + s.mu.Unlock() + authURL := externalIdpAuthorizeURL(authEndpoint, clientID, redirectURI, scopes, generateCodeChallenge(verifier), state2, loginHint) + // Redirect the SAME browser tab on to the IdP login page. + http.Redirect(w, req, authURL, http.StatusFound) + return + } + + // --- Enterprise leg-2: IdP authorization code at /oauth/callback --- + if req.URL.Path == kiroOAuthCallbackPath { + s.mu.Lock() + ctx2 := s.leg2 + s.mu.Unlock() + code := strings.TrimSpace(q.Get("code")) + state := strings.TrimSpace(q.Get("state")) + errParam := strings.TrimSpace(q.Get("error")) + // Ignore callbacks that don't match the in-flight leg-2 state. + if ctx2 == nil || state == "" || state != ctx2.state { + w.WriteHeader(http.StatusNoContent) + return + } + if errParam != "" { + desc := strings.TrimSpace(q.Get("error_description")) + writeKiroCallbackPage(w, false) + s.deliver(kiroSsoCapture{err: fmt.Errorf("external IdP authorization error: %s %s", errParam, desc)}) + return + } + if code == "" { + w.WriteHeader(http.StatusNoContent) + return + } + writeKiroCallbackPage(w, true) + s.deliver(kiroSsoCapture{ + kind: "external_idp", + code: code, + tokenEndpoint: ctx2.tokenEndpoint, + issuerURL: ctx2.issuerURL, + clientID: ctx2.clientID, + scopes: ctx2.scopes, + redirectURI: ctx2.redirectURI, + codeVerifier: ctx2.verifier, + }) + return + } + + // --- Social leg-1: Cognito authorization code --- + code := strings.TrimSpace(q.Get("code")) + errParam := strings.TrimSpace(q.Get("error")) + state := strings.TrimSpace(q.Get("state")) + // Ignore stray hits with neither a code nor an error, and any callback whose + // state does not match — without consuming the one-shot. + if code == "" && errParam == "" { + w.WriteHeader(http.StatusNoContent) + return + } + if s.State == "" || state != s.State { + w.WriteHeader(http.StatusNoContent) + return + } + if errParam != "" { + desc := strings.TrimSpace(q.Get("error_description")) + writeKiroCallbackPage(w, false) + s.deliver(kiroSsoCapture{err: fmt.Errorf("SSO authorization error: %s %s", errParam, desc)}) + return + } + writeKiroCallbackPage(w, true) + s.deliver(kiroSsoCapture{kind: "social", code: code}) +} + +// --- OIDC discovery + token exchange (enterprise / external IdP leg) --------- + +// validateExternalIdpEndpoint verifies rawURL is an https URL whose host is a +// non-IP, allow-listed enterprise IdP host. It gates the issuer (before +// discovery) and BOTH discovered endpoints (the authorize URL the browser is +// 302'd to, and the token endpoint the code is exchanged at). +func validateExternalIdpEndpoint(rawURL string) error { + u, err := url.Parse(strings.TrimSpace(rawURL)) + if err != nil { + return fmt.Errorf("invalid external IdP URL: %w", err) + } + if !strings.EqualFold(u.Scheme, "https") { + return fmt.Errorf("external IdP URL must be https") + } + host := strings.ToLower(u.Hostname()) + if host == "" { + return fmt.Errorf("external IdP URL has no host") + } + // Reject IP-literal hosts outright; only named, allow-listed hosts pass. + if net.ParseIP(host) != nil { + return fmt.Errorf("external IdP host must not be an IP literal") + } + for _, suffix := range allowedExternalIdpIssuerSuffixes { + if strings.HasSuffix(host, suffix) { + return nil + } + } + return fmt.Errorf("external IdP host %q is not allow-listed", host) +} + +// externalIdpEndpointValidator is the function ValidateExternalIdpEndpoint delegates +// to. Tests override it via SetExternalIdpValidatorForTest so a happy-path import +// test can POST against an httptest server (http + 127.0.0.1) that the real +// allow-list would reject. +var externalIdpEndpointValidator = validateExternalIdpEndpoint + +// ValidateExternalIdpEndpoint is the exported entry point for validating a user- or +// discovery-supplied external IdP endpoint URL. The credential-import path +// (package proxy) uses this to guard against SSRF / refresh-token exfiltration: a +// pasted tokenEndpoint pointing at an internal or attacker-controlled host would +// otherwise cause the server to POST the account's refresh token there. +func ValidateExternalIdpEndpoint(rawURL string) error { + return externalIdpEndpointValidator(rawURL) +} + +// issuerFromAccessTokenJWT decodes an unverified Azure AD access token's payload +// and returns its iss claim (e.g. https://login.microsoftonline.com//v2.0). +// No signature verification: this is used only to classify a pasted credential and +// recover the Azure tenant, never to trust the token itself. Returns "" if +// accessToken is not a JWT or has no iss. +func issuerFromAccessTokenJWT(accessToken string) string { + raw := strings.TrimSpace(accessToken) + if raw == "" { + return "" + } + parts := strings.Split(raw, ".") + if len(parts) < 2 { + return "" + } + payload, err := base64.RawURLEncoding.DecodeString(parts[1]) + if err != nil { + return "" + } + var claims struct { + Iss string `json:"iss"` + } + if err := json.Unmarshal(payload, &claims); err != nil { + return "" + } + return claims.Iss +} + +// ExpFromAccessTokenJWT decodes an unverified access token's payload and returns +// its exp claim (Unix seconds). No signature verification: this is used only to +// set a trustworthy ExpiresAt when importing an external_idp credential WITHOUT a +// live refresh round-trip (trust-on-import), never to authenticate the token. +// Returns 0 if accessToken is not a JWT or has no exp. +func ExpFromAccessTokenJWT(accessToken string) int64 { + raw := strings.TrimSpace(accessToken) + if raw == "" { + return 0 + } + parts := strings.Split(raw, ".") + if len(parts) < 2 { + return 0 + } + payload, err := base64.RawURLEncoding.DecodeString(parts[1]) + if err != nil { + return 0 + } + var claims struct { + Exp int64 `json:"exp"` + } + if err := json.Unmarshal(payload, &claims); err != nil { + return 0 + } + return claims.Exp +} + +// DeriveExternalIdpEndpoints reconstructs the Microsoft / Azure AD token endpoint, +// OIDC issuer, and default scopes for an external-IdP credential. The Azure tenant +// is recovered from userId (Kiro Account Manager exports carry it at account +// level) or, failing that, from the accessToken JWT's issuer (bare blobs with only +// clientId + accessToken + refreshToken). This lets the credential-import path +// accept those shapes even though they omit tokenEndpoint/issuerUrl/scopes. +// +// userId / iss look like: https://login.microsoftonline.com//v2.0. +// Returns empty strings if neither source yields a usable tenant, so the caller +// can fall back to its "requires clientId and tokenEndpoint" error. The derived +// tokenEndpoint is re-validated against the IdP allow-list by the caller, so a +// non-allow-listed host (or the test's http+127.0.0.1 fake) is still gated. +func DeriveExternalIdpEndpoints(userId, clientID, accessToken string) (tokenEndpoint, issuerURL, scopes string) { + src := strings.TrimSpace(userId) + if src == "" { + // Bare credential blobs carry only clientId + accessToken: recover the + // tenant from the access token's JWT issuer. + src = strings.TrimSpace(issuerFromAccessTokenJWT(accessToken)) + } + if src == "" { + return "", "", "" + } + u, err := url.Parse(src) + if err != nil || u.Host == "" { + return "", "", "" + } + segments := strings.Split(strings.Trim(u.Path, "/"), "/") + if len(segments) == 0 || segments[0] == "" { + return "", "", "" + } + tenant := segments[0] + scheme := u.Scheme + if scheme == "" { + scheme = "https" + } + tokenEndpoint = fmt.Sprintf("%s://%s/%s/oauth2/v2.0/token", scheme, u.Host, tenant) + issuerURL = fmt.Sprintf("%s://%s/%s/v2.0", scheme, u.Host, tenant) + if clientID != "" { + scopes = fmt.Sprintf("api://%s/codewhisperer:conversations api://%s/codewhisperer:completions offline_access", clientID, clientID) + } + return tokenEndpoint, issuerURL, scopes +} + +// oidcDiscover fetches the OpenID Connect discovery document for issuerURL and +// returns its authorization and token endpoints. The issuer and BOTH discovered +// endpoints are validated against the IdP host allow-list; redirects are NOT +// followed (so a discovery host cannot bounce the fetch to an internal target); +// and no response body is echoed into errors. +func oidcDiscover(client *http.Client, issuerURL, proxyURL string) (authEndpoint, tokenEndpoint string, err error) { + if err = validateExternalIdpEndpoint(issuerURL); err != nil { + return "", "", err + } + docURL := strings.TrimRight(strings.TrimSpace(issuerURL), "/") + "/.well-known/openid-configuration" + req, err := http.NewRequest(http.MethodGet, docURL, nil) + if err != nil { + return "", "", fmt.Errorf("failed to build OIDC discovery request: %w", err) + } + req.Header.Set("Accept", "application/json") + // Do not follow redirects: the allow-listed issuer host must answer directly, + // so a 3xx (which could point at an internal/link-local target) is a failure. + noRedirect := &http.Client{ + Timeout: 30 * time.Second, + Transport: buildAuthTransport(proxyURL), + CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse }, + } + resp, err := noRedirect.Do(req) + if err != nil { + return "", "", fmt.Errorf("OIDC discovery request failed: %w", err) + } + defer resp.Body.Close() + body, err := io.ReadAll(io.LimitReader(resp.Body, 1<<20)) + if err != nil { + return "", "", fmt.Errorf("failed to read OIDC discovery response: %w", err) + } + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + // Deliberately omit the body to avoid exfiltrating an internal response. + return "", "", fmt.Errorf("OIDC discovery failed (status %d)", resp.StatusCode) + } + var doc struct { + AuthorizationEndpoint string `json:"authorization_endpoint"` + TokenEndpoint string `json:"token_endpoint"` + } + if err = json.Unmarshal(body, &doc); err != nil { + return "", "", fmt.Errorf("failed to parse OIDC discovery document: %w", err) + } + if doc.AuthorizationEndpoint == "" || doc.TokenEndpoint == "" { + return "", "", fmt.Errorf("OIDC discovery document missing authorization_endpoint or token_endpoint") + } + // Both endpoints must themselves be https + allow-listed. The allow-list (the + // trust boundary) is the only check applied here, not equality with the issuer + // host: within Microsoft's own *.microsoftonline.com infrastructure the + // discovery doc legitimately points authorize/token at sibling hosts. If the + // allow-list is ever broadened to multiple distinct IdPs, tighten this to also + // pin the endpoints to the issuer's registrable host. + if err = validateExternalIdpEndpoint(doc.AuthorizationEndpoint); err != nil { + return "", "", fmt.Errorf("discovered authorization_endpoint rejected: %w", err) + } + if err = validateExternalIdpEndpoint(doc.TokenEndpoint); err != nil { + return "", "", fmt.Errorf("discovered token_endpoint rejected: %w", err) + } + return doc.AuthorizationEndpoint, doc.TokenEndpoint, nil +} + +// externalIdpAuthorizeURL builds the IdP authorization-code+PKCE URL the browser +// is redirected to for the enterprise leg. scopes is passed through verbatim from +// the portal (already a space-separated list). +func externalIdpAuthorizeURL(authEndpoint, clientID, redirectURI, scopes, challenge, state, loginHint string) string { + q := url.Values{} + q.Set("client_id", clientID) + q.Set("response_type", "code") + q.Set("redirect_uri", redirectURI) + q.Set("scope", scopes) + q.Set("code_challenge", challenge) + q.Set("code_challenge_method", "S256") + q.Set("response_mode", "query") + q.Set("state", state) + if strings.TrimSpace(loginHint) != "" { + q.Set("login_hint", loginHint) + } + return authEndpoint + "?" + q.Encode() +} + +// exchangeExternalIdpCode exchanges an IdP authorization code (with its PKCE +// verifier) for IdP tokens at the discovered token endpoint. Standard OAuth2 +// authorization_code grant for a public client (PKCE, no client secret); +// request is form-encoded and the response is snake_case. +func exchangeExternalIdpCode(client *http.Client, tokenEndpoint, clientID, code, codeVerifier, redirectURI, scopes string) (accessToken, refreshToken string, expiresIn int, err error) { + form := url.Values{} + form.Set("client_id", clientID) + form.Set("grant_type", "authorization_code") + form.Set("code", strings.TrimSpace(code)) + form.Set("redirect_uri", redirectURI) + form.Set("code_verifier", codeVerifier) + if strings.TrimSpace(scopes) != "" { + form.Set("scope", scopes) + } + return postExternalIdpToken(client, tokenEndpoint, form) +} + +// exchangeSocialCode exchanges a Cognito authorization code (with its PKCE +// verifier) for Kiro tokens at the social token endpoint. Request body matches +// the Kiro IDE client — {code, code_verifier, redirect_uri} — and the response +// is camelCase (accessToken/refreshToken/profileArn/expiresIn). +func exchangeSocialCode(client *http.Client, code, codeVerifier string) (accessToken, refreshToken string, expiresIn int, profileArn string, err error) { + payload := map[string]string{ + "code": strings.TrimSpace(code), + "code_verifier": codeVerifier, + "redirect_uri": kiroRedirectURI, + } + body, _ := json.Marshal(payload) + req, err := http.NewRequest(http.MethodPost, kiroSocialTokenURL, bytes.NewReader(body)) + if err != nil { + return "", "", 0, "", fmt.Errorf("failed to build social token request: %w", err) + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "application/json") + resp, err := client.Do(req) + if err != nil { + return "", "", 0, "", err + } + defer resp.Body.Close() + respBody, _ := io.ReadAll(resp.Body) + var out struct { + AccessToken string `json:"accessToken"` + RefreshToken string `json:"refreshToken"` + ProfileArn string `json:"profileArn"` + ExpiresIn int `json:"expiresIn"` + } + _ = json.Unmarshal(respBody, &out) + if resp.StatusCode < 200 || resp.StatusCode >= 300 || out.AccessToken == "" { + return "", "", 0, "", fmt.Errorf("social token exchange failed (status %d): %s", resp.StatusCode, string(respBody)) + } + return out.AccessToken, out.RefreshToken, out.ExpiresIn, out.ProfileArn, nil +} + +// writeKiroCallbackPage renders a minimal HTML page shown in the browser after +// the final redirect. +func writeKiroCallbackPage(w http.ResponseWriter, ok bool) { + w.Header().Set("Content-Type", "text/html; charset=utf-8") + w.WriteHeader(http.StatusOK) + msg := "Kiro sign-in complete. You can close this tab and return to the admin panel." + if !ok { + msg = "Kiro sign-in failed. Return to the admin panel and try again." + } + _, _ = fmt.Fprintf(w, "Kiro Sign-In

%s

", msg) +} + +// ExtractEmailFromJWT decodes the JWT payload of an access token (best-effort) +// and returns the account email. Azure AD v2.0 tokens frequently omit the +// "email" claim and carry the sign-in name in "preferred_username"/"upn", so +// those are used as fallbacks. +func ExtractEmailFromJWT(accessToken string) string { + parts := strings.Split(strings.TrimSpace(accessToken), ".") + if len(parts) < 2 { + return "" + } + payload, err := base64.RawURLEncoding.DecodeString(parts[1]) + if err != nil { + // Some tokens use standard (padded) base64; retry leniently. + if payload, err = base64.StdEncoding.DecodeString(parts[1]); err != nil { + return "" + } + } + var claims struct { + Email string `json:"email"` + PreferredUsername string `json:"preferred_username"` + Upn string `json:"upn"` + UniqueName string `json:"unique_name"` + } + if err = json.Unmarshal(payload, &claims); err != nil { + return "" + } + for _, candidate := range []string{claims.Email, claims.PreferredUsername, claims.Upn, claims.UniqueName} { + if v := strings.TrimSpace(candidate); v != "" { + return v + } + } + return "" +} + +// removeKiroSsoSession deletes a session from the registry. +func removeKiroSsoSession(sessionID string) { + kiroSsoSessionsMu.Lock() + delete(kiroSsoSessions, sessionID) + kiroSsoSessionsMu.Unlock() +} + diff --git a/auth/kiro_sso_test.go b/auth/kiro_sso_test.go new file mode 100644 index 00000000..d2ba8dc5 --- /dev/null +++ b/auth/kiro_sso_test.go @@ -0,0 +1,314 @@ +package auth + +import ( + "encoding/base64" + "encoding/json" + "net/http" + "net/http/httptest" + "net/url" + "reflect" + "strings" + "testing" +) + +// TestKiroCallbackBindAddrs locks in the secure default (loopback-only) and the +// KIRO_SSO_CALLBACK_BIND override used for containerized deployments. +func TestKiroCallbackBindAddrs(t *testing.T) { + // Unset/blank -> loopback v4 (mandatory) + v6 (best-effort). + t.Setenv("KIRO_SSO_CALLBACK_BIND", "") + if got, want := kiroCallbackBindAddrs(), []string{"127.0.0.1:3128", "[::1]:3128"}; !reflect.DeepEqual(got, want) { + t.Fatalf("default bind addrs = %v, want %v", got, want) + } + // Whitespace is treated as unset (still the secure default). + t.Setenv("KIRO_SSO_CALLBACK_BIND", " ") + if got := kiroCallbackBindAddrs(); len(got) != 2 { + t.Fatalf("whitespace should fall back to loopback default, got %v", got) + } + // IPv4 wildcard override -> single mandatory bind. + t.Setenv("KIRO_SSO_CALLBACK_BIND", "0.0.0.0") + if got, want := kiroCallbackBindAddrs(), []string{"0.0.0.0:3128"}; !reflect.DeepEqual(got, want) { + t.Fatalf("0.0.0.0 bind addrs = %v, want %v", got, want) + } + // IPv6 wildcard override -> bracketed host:port. + t.Setenv("KIRO_SSO_CALLBACK_BIND", "::") + if got, want := kiroCallbackBindAddrs(), []string{"[::]:3128"}; !reflect.DeepEqual(got, want) { + t.Fatalf(":: bind addrs = %v, want %v", got, want) + } +} + +// makeJWT builds an unsigned JWT-shaped string whose payload encodes claims, for +// testing the best-effort claim extraction (the signature is never verified). +func makeJWT(claims map[string]string) string { + header := base64.RawURLEncoding.EncodeToString([]byte(`{"alg":"none"}`)) + payloadBytes, _ := json.Marshal(claims) + payload := base64.RawURLEncoding.EncodeToString(payloadBytes) + return header + "." + payload + ".sig" +} + +func TestExtractEmailFromJWT(t *testing.T) { + cases := []struct { + name string + claims map[string]string + want string + }{ + {"email claim", map[string]string{"email": "a@b.com"}, "a@b.com"}, + // Azure AD v2.0 tokens usually omit "email" and carry preferred_username. + {"preferred_username fallback", map[string]string{"preferred_username": "user@tenant.onmicrosoft.com"}, "user@tenant.onmicrosoft.com"}, + {"upn fallback", map[string]string{"upn": "u@corp.com"}, "u@corp.com"}, + {"none", map[string]string{"sub": "xyz"}, ""}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + if got := ExtractEmailFromJWT(makeJWT(tc.claims)); got != tc.want { + t.Fatalf("ExtractEmailFromJWT = %q, want %q", got, tc.want) + } + }) + } + if got := ExtractEmailFromJWT("not-a-jwt"); got != "" { + t.Fatalf("malformed token should yield empty email, got %q", got) + } +} + +func TestValidateExternalIdpEndpoint(t *testing.T) { + valid := []string{ + "https://login.microsoftonline.com/5fbc183e/v2.0", + "https://login.microsoftonline.us/tenant/v2.0", + "https://login.microsoftonline.cn/tenant/oauth2/v2.0/token", + } + for _, u := range valid { + if err := validateExternalIdpEndpoint(u); err != nil { + t.Fatalf("expected %q to be allowed, got %v", u, err) + } + } + invalid := []string{ + "http://login.microsoftonline.com/x", // not https + "https://evil-microsoftonline.com/x", // suffix not anchored to a subdomain boundary + "https://login.microsoftonline.com.evil.co", // not an allowed suffix + "https://10.0.0.5/x", // IP literal + "https://accounts.google.com/x", // not allow-listed + "https:///x", // no host + } + for _, u := range invalid { + if err := validateExternalIdpEndpoint(u); err == nil { + t.Fatalf("expected %q to be rejected", u) + } + } +} + +func TestExternalIdpAuthorizeURL(t *testing.T) { + raw := externalIdpAuthorizeURL( + "https://login.microsoftonline.com/t/oauth2/v2.0/authorize", + "client-123", + "http://localhost:3128/oauth/callback", + "api://client-123/codewhisperer:conversations offline_access", + "challenge-abc", + "state-xyz", + "user@corp.com", + ) + u, err := url.Parse(raw) + if err != nil { + t.Fatalf("parse: %v", err) + } + q := u.Query() + checks := map[string]string{ + "client_id": "client-123", + "response_type": "code", + "redirect_uri": "http://localhost:3128/oauth/callback", + "code_challenge": "challenge-abc", + "code_challenge_method": "S256", + "response_mode": "query", + "state": "state-xyz", + "login_hint": "user@corp.com", + } + for k, want := range checks { + if got := q.Get(k); got != want { + t.Fatalf("authorize url param %q = %q, want %q", k, got, want) + } + } + if !strings.Contains(q.Get("scope"), "offline_access") { + t.Fatalf("expected offline_access in scope, got %q", q.Get("scope")) + } +} + +// TestExternalIdpAuthorizeURLOmitsEmptyLoginHint ensures we don't emit an empty +// login_hint parameter when the portal supplied none. +func TestExternalIdpAuthorizeURLOmitsEmptyLoginHint(t *testing.T) { + raw := externalIdpAuthorizeURL("https://login.microsoftonline.com/t/authorize", "c", "http://localhost:3128/oauth/callback", "s", "ch", "st", "") + u, _ := url.Parse(raw) + if _, ok := u.Query()["login_hint"]; ok { + t.Fatalf("login_hint should be omitted when empty") + } +} + +// TestRefreshExternalIdpToken drives the refresh_token grant against a stub IdP +// token endpoint and asserts the form encoding and response mapping. +func TestRefreshExternalIdpToken(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if err := r.ParseForm(); err != nil { + t.Fatalf("parse form: %v", err) + } + if r.Form.Get("grant_type") != "refresh_token" { + t.Fatalf("grant_type = %q", r.Form.Get("grant_type")) + } + if r.Form.Get("client_id") != "azure-client" { + t.Fatalf("client_id = %q", r.Form.Get("client_id")) + } + if r.Form.Get("refresh_token") != "old-refresh" { + t.Fatalf("refresh_token = %q", r.Form.Get("refresh_token")) + } + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"access_token":"new-access","refresh_token":"new-refresh","expires_in":3600}`)) + })) + defer srv.Close() + + // The POST boundary re-validates tokenEndpoint against the allow-list, which rejects + // the httptest URL (http + 127.0.0.1). Install a permissive validator for the test. + restore := SetExternalIdpValidatorForTest(func(string) error { return nil }) + defer SetExternalIdpValidatorForTest(restore) + + access, refresh, expiresAt, profileArn, err := refreshExternalIdpToken( + "old-refresh", "azure-client", srv.URL, "api://x/codewhisperer:conversations offline_access", srv.Client(), + ) + if err != nil { + t.Fatalf("refreshExternalIdpToken: %v", err) + } + if access != "new-access" { + t.Fatalf("access = %q", access) + } + if refresh != "new-refresh" { + t.Fatalf("refresh = %q", refresh) + } + if profileArn != "" { + t.Fatalf("external IdP refresh should not return a profileArn, got %q", profileArn) + } + if expiresAt == 0 { + t.Fatalf("expected a non-zero absolute expiry") + } +} + +// TestRefreshExternalIdpTokenKeepsRefreshTokenWhenOmitted verifies the existing +// refresh token is retained when the IdP response omits a rotated one. +func TestRefreshExternalIdpTokenKeepsRefreshTokenWhenOmitted(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"access_token":"a2","expires_in":1200}`)) + })) + defer srv.Close() + + // See TestRefreshExternalIdpToken: relax the allow-list for the httptest endpoint. + restore := SetExternalIdpValidatorForTest(func(string) error { return nil }) + defer SetExternalIdpValidatorForTest(restore) + + _, refresh, _, _, err := refreshExternalIdpToken("keep-me", "c", srv.URL, "", srv.Client()) + if err != nil { + t.Fatalf("refreshExternalIdpToken: %v", err) + } + if refresh != "keep-me" { + t.Fatalf("expected refresh token to be retained, got %q", refresh) + } +} + +// TestRefreshExternalIdpTokenRequiresClientAndEndpoint guards the precondition +// that distinguishes the external-IdP branch from the AWS OIDC branch. +func TestRefreshExternalIdpTokenRequiresClientAndEndpoint(t *testing.T) { + if _, _, _, _, err := refreshExternalIdpToken("r", "", "https://login.microsoftonline.com/t/token", "", http.DefaultClient); err == nil { + t.Fatalf("expected error when clientID is empty") + } + if _, _, _, _, err := refreshExternalIdpToken("r", "c", "", "", http.DefaultClient); err == nil { + t.Fatalf("expected error when tokenEndpoint is empty") + } +} + +// TestValidateExternalIdpEndpointAcceptsAllowListed verifies the exported validator +// accepts real Azure / Microsoft 365 token endpoints (commercial, us-gov, china). +func TestValidateExternalIdpEndpointAcceptsAllowListed(t *testing.T) { + for _, raw := range []string{ + "https://login.microsoftonline.com/tenant/oauth2/v2.0/token", + "https://login.microsoftonline.us/tenant/v2.0", + "https://login.partner.microsoftonline.cn/tenant/oauth2/v2.0/token", + } { + if err := ValidateExternalIdpEndpoint(raw); err != nil { + t.Errorf("expected %q accepted, got %v", raw, err) + } + } +} + +// TestValidateExternalIdpEndpointRejectsUnsafe verifies the validator rejects the +// SSRF shapes a pasted credential JSON could carry: cleartext http, IP literals, +// and non-allow-listed hosts. +func TestValidateExternalIdpEndpointRejectsUnsafe(t *testing.T) { + for _, raw := range []string{ + "http://login.microsoftonline.com/x", // not https + "https://127.0.0.1/oauth/token", // IP literal + "https://evil.example.com/oauth/token", // not allow-listed + } { + if err := ValidateExternalIdpEndpoint(raw); err == nil { + t.Errorf("expected %q rejected, got nil", raw) + } + } +} + +// TestSetExternalIdpValidatorForTestSwapsAndRestores verifies the test seam lets a +// test override (and restore) the validator so happy-path import tests can POST +// against an httptest server (http + 127.0.0.1) that the real allow-list rejects. +func TestSetExternalIdpValidatorForTestSwapsAndRestores(t *testing.T) { + restore := SetExternalIdpValidatorForTest(func(string) error { return nil }) + defer SetExternalIdpValidatorForTest(restore) + if err := ValidateExternalIdpEndpoint("https://evil.example.com/x"); err != nil { + t.Fatalf("expected swapped no-op validator to accept, got %v", err) + } +} + +// TestDeriveExternalIdpEndpoints verifies the endpoints+scopes are reconstructed +// from userId (Kiro export) OR the accessToken JWT issuer (bare blobs), plus the +// Kiro client ID. This is what lets the import path accept shapes that omit +// tokenEndpoint/issuerUrl/scopes. +func TestDeriveExternalIdpEndpoints(t *testing.T) { + const userID = "https://login.microsoftonline.com/5fbc183e-3d09-4043-b36f-0c49d3665977/v2.0.8db0e2eb-d491-4a1a-98f1-cbdc12bb60a0" + const clientID = "fa6d79bf-cdaa-495e-8359-78aab7c7cd9b" + const wantTE = "https://login.microsoftonline.com/5fbc183e-3d09-4043-b36f-0c49d3665977/oauth2/v2.0/token" + const wantIss = "https://login.microsoftonline.com/5fbc183e-3d09-4043-b36f-0c49d3665977/v2.0" + + // From userId (Kiro export carries it at account level). + te, iss, sc := DeriveExternalIdpEndpoints(userID, clientID, "") + if te != wantTE || iss != wantIss { + t.Fatalf("from userId: te=%q iss=%q (want %q / %q)", te, iss, wantTE, wantIss) + } + if !strings.Contains(sc, "api://"+clientID+"/codewhisperer:conversations") || !strings.Contains(sc, "offline_access") { + t.Fatalf("scopes: got %q", sc) + } + + // From the accessToken JWT issuer (bare blobs carry no userId). + jwt := "eyJhbGciOiJub25lIn0." + base64.RawURLEncoding.EncodeToString([]byte(`{"iss":"`+userID+`"}`)) + "." + te2, iss2, sc2 := DeriveExternalIdpEndpoints("", clientID, jwt) + if te2 != wantTE || iss2 != wantIss { + t.Fatalf("from accessToken JWT: te=%q iss=%q (want %q / %q)", te2, iss2, wantTE, wantIss) + } + if sc2 == "" { + t.Fatalf("scopes from JWT path: got empty") + } + + // Neither source → all-empty (caller falls back to its 400). + if te3, iss3, sc3 := DeriveExternalIdpEndpoints("", clientID, ""); te3 != "" || iss3 != "" || sc3 != "" { + t.Fatalf("empty sources should yield all-empty, got %q %q %q", te3, iss3, sc3) + } + // userId takes precedence over accessToken. + if te4, _, _ := DeriveExternalIdpEndpoints(userID, clientID, jwt); te4 != wantTE { + t.Fatalf("userId should take precedence over accessToken, got %q", te4) + } +} + +// TestExpFromAccessTokenJWT pins the exp extraction used for trust-on-import. +func TestExpFromAccessTokenJWT(t *testing.T) { + payload := base64.RawURLEncoding.EncodeToString([]byte(`{"iss":"x","exp":2000000000}`)) + jwt := "eyJhbGciOiJub25lIn0." + payload + "." + if got := ExpFromAccessTokenJWT(jwt); got != 2000000000 { + t.Fatalf("ExpFromAccessTokenJWT: got %d, want 2000000000", got) + } + if got := ExpFromAccessTokenJWT(""); got != 0 { + t.Fatalf("empty → 0, got %d", got) + } + if got := ExpFromAccessTokenJWT("not-a-jwt"); got != 0 { + t.Fatalf("non-JWT → 0, got %d", got) + } +} diff --git a/auth/oidc.go b/auth/oidc.go index 4e12d171..ec0b1d5e 100644 --- a/auth/oidc.go +++ b/auth/oidc.go @@ -7,6 +7,8 @@ import ( "io" "kiro-go/config" "net/http" + "net/url" + "strings" "time" ) @@ -30,12 +32,101 @@ func RefreshToken(account *config.Account) (string, string, int64, string, error } client := GetAuthClientForProxy(proxyURL) + // External IdP (enterprise SSO, e.g. Azure AD) tokens are refreshed against the + // IdP token endpoint (refresh_token grant, public client), NOT the AWS SSO OIDC + // endpoint. Selecting it on AuthMethod (rather than letting it fall through to the + // OIDC branch, which requires clientSecret) is what makes these accounts refresh. + if account.AuthMethod == "external_idp" { + return refreshExternalIdpToken(account.RefreshToken, account.ClientID, account.TokenEndpoint, account.Scopes, client) + } if account.AuthMethod == "social" { return refreshSocialToken(account.RefreshToken, client) } return refreshOIDCToken(account.RefreshToken, account.ClientID, account.ClientSecret, account.Region, client) } +// refreshExternalIdpToken refreshes an external-IdP (enterprise SSO) access token +// through the IdP token endpoint using the OAuth2 refresh_token grant for a public +// client (no client secret). offline_access in the original scopes is what makes a +// refresh token available. The IdP issues no profileArn (it is resolved separately +// via ListAvailableProfiles using the EXTERNAL_IDP token type), so "" is returned +// for the profileArn. +func refreshExternalIdpToken(refreshToken, clientID, tokenEndpoint, scopes string, client *http.Client) (string, string, int64, string, error) { + if clientID == "" || tokenEndpoint == "" { + return "", "", 0, "", fmt.Errorf("external IdP refresh requires clientId and tokenEndpoint") + } + form := url.Values{} + form.Set("client_id", clientID) + form.Set("grant_type", "refresh_token") + form.Set("refresh_token", refreshToken) + if scopes != "" { + form.Set("scope", scopes) + } + accessToken, newRefreshToken, expiresIn, err := postExternalIdpToken(client, tokenEndpoint, form) + if err != nil { + return "", "", 0, "", err + } + // Some IdPs (Azure AD) rotate refresh tokens; others omit it on refresh. Keep the + // existing refresh token when the response does not carry a new one. + if newRefreshToken == "" { + newRefreshToken = refreshToken + } + expiresAt := time.Now().Unix() + int64(expiresIn) + return accessToken, newRefreshToken, expiresAt, "", nil +} + +// postExternalIdpToken performs a form-encoded POST to an external-IdP token +// endpoint and maps the snake_case OAuth2 token response onto the standard return +// shape. Shared by the authorization-code exchange (login) and the refresh_token +// grant (renewal). +func postExternalIdpToken(client *http.Client, tokenEndpoint string, form url.Values) (accessToken, refreshToken string, expiresIn int, err error) { + if strings.TrimSpace(tokenEndpoint) == "" { + return "", "", 0, fmt.Errorf("external IdP token endpoint is empty") + } + // Defense-in-depth: re-validate the endpoint at the outbound-POST boundary so the + // refresh token is never sent to a non-allow-listed host — even if a persisted + // account's TokenEndpoint was set out-of-band (backup restore, an external file + // write, or a future caller that stores an endpoint without validating). This makes + // allow-list validation an invariant of the exfiltration-sensitive operation itself + // rather than of every caller. Uses the exported ValidateExternalIdpEndpoint so the + // test seam (SetExternalIdpValidatorForTest) still relaxes it for httptest servers. + if err := ValidateExternalIdpEndpoint(tokenEndpoint); err != nil { + return "", "", 0, fmt.Errorf("external IdP token endpoint rejected: %w", err) + } + req, err := http.NewRequest("POST", tokenEndpoint, strings.NewReader(form.Encode())) + if err != nil { + return "", "", 0, fmt.Errorf("failed to build external IdP token request: %w", err) + } + req.Header.Set("Content-Type", "application/x-www-form-urlencoded") + req.Header.Set("Accept", "application/json") + + resp, err := client.Do(req) + if err != nil { + return "", "", 0, fmt.Errorf("external IdP token request failed: %w", err) + } + defer resp.Body.Close() + + body, err := io.ReadAll(resp.Body) + if err != nil { + return "", "", 0, fmt.Errorf("failed to read external IdP token response: %w", err) + } + var out struct { + AccessToken string `json:"access_token"` + RefreshToken string `json:"refresh_token"` + ExpiresIn int `json:"expires_in"` + Error string `json:"error"` + ErrorDesc string `json:"error_description"` + } + _ = json.Unmarshal(body, &out) + if resp.StatusCode < 200 || resp.StatusCode >= 300 || out.AccessToken == "" { + if out.Error != "" { + return "", "", 0, fmt.Errorf("external IdP token exchange failed (status %d): %s: %s", resp.StatusCode, out.Error, out.ErrorDesc) + } + return "", "", 0, fmt.Errorf("external IdP token exchange failed (status %d): %s", resp.StatusCode, string(body)) + } + return out.AccessToken, out.RefreshToken, out.ExpiresIn, nil +} + // refreshOIDCToken IdC/Builder ID token 刷新 func refreshOIDCToken(refreshToken, clientID, clientSecret, region string, client *http.Client) (string, string, int64, string, error) { if clientID == "" || clientSecret == "" { diff --git a/auth/sso_token.go b/auth/sso_token.go index dee05405..c549c517 100644 --- a/auth/sso_token.go +++ b/auth/sso_token.go @@ -274,8 +274,17 @@ func pollForToken(oidcBase, clientID, clientSecret, deviceCode string, interval RefreshToken string `json:"refreshToken"` ExpiresIn int `json:"expiresIn"` } - json.NewDecoder(resp.Body).Decode(&result) + if err := json.NewDecoder(resp.Body).Decode(&result); err != nil { + resp.Body.Close() + return "", "", 0, fmt.Errorf("parse token response failed: %w", err) + } resp.Body.Close() + // A corporate filtering proxy that substitutes an HTML block-page + // (HTTP 200, non-JSON body) would decode into empty fields. Refuse + // it rather than silently creating an account with empty tokens. + if result.AccessToken == "" { + return "", "", 0, fmt.Errorf("token response missing accessToken") + } return result.AccessToken, result.RefreshToken, result.ExpiresIn, nil } diff --git a/auth/sso_token_test.go b/auth/sso_token_test.go new file mode 100644 index 00000000..7edb033a --- /dev/null +++ b/auth/sso_token_test.go @@ -0,0 +1,45 @@ +package auth + +import ( + "io" + "net/http" + "net/http/httptest" + "testing" + "time" +) + +// TestPollForTokenRejectsNonJSON200 guards against the empty-token account bug: +// a corporate filtering proxy that substitutes an HTML block-page (HTTP 200, +// non-JSON body) between the proxy and AWS must NOT be decoded into empty +// credentials and returned as a nil-error success. Previously pollForToken +// discarded the decode error, returned ("","",0,nil), and ImportFromSsoToken / +// PollBuilderIdAuth propagated that as success — silently creating an account +// with empty AccessToken/RefreshToken treated as valid until first use fails. +// See Bug G. +func TestPollForTokenRejectsNonJSON200(t *testing.T) { + // Install a client whose Transport does not consult ProxyFromEnvironment so + // the POST reliably reaches the in-process httptest server regardless of the + // ambient proxy env. + prevClient := SetGlobalAuthClientForTest(&http.Client{ + Timeout: 5 * time.Second, + Transport: &http.Transport{}, + }) + defer SetGlobalAuthClientForTest(prevClient) + + // Fake token endpoint substitutes a corporate block-page: HTTP 200 + HTML. + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/html") + _, _ = io.WriteString(w, "Blocked by corporate policy") + })) + defer server.Close() + + // interval=1 keeps the first poll ~1s out; the 200 branch returns on it. + accessToken, refreshToken, expiresIn, err := pollForToken(server.URL, "cid", "secret", "device-code", 1) + + if err == nil { + t.Fatalf("pollForToken must return an error when the token endpoint returns a non-JSON 200 (e.g. a corporate block-page); got empty tokens (access=%q refresh=%q expiresIn=%d) and nil error", accessToken, refreshToken, expiresIn) + } + if accessToken != "" || refreshToken != "" { + t.Fatalf("on a non-JSON 200, pollForToken must not return token values; got access=%q refresh=%q", accessToken, refreshToken) + } +} diff --git a/auth/testhooks.go b/auth/testhooks.go index 5e84b0b6..6c911292 100644 --- a/auth/testhooks.go +++ b/auth/testhooks.go @@ -28,3 +28,16 @@ func SetGlobalAuthClientForTest(c *http.Client) *http.Client { } return old } + +// SetExternalIdpValidatorForTest swaps the validator behind ValidateExternalIdpEndpoint +// and returns the previous one so callers can restore it. Tests POST against httptest +// servers (http + 127.0.0.1), which the real allow-list validator rejects, so tests +// install a no-op validator here. Mirrors SetGlobalAuthClientForTest's swap-and-restore +// shape. Test-only. +func SetExternalIdpValidatorForTest(fn func(string) error) func(string) error { + old := externalIdpEndpointValidator + if fn != nil { + externalIdpEndpointValidator = fn + } + return old +} diff --git a/config/client_profile_test.go b/config/client_profile_test.go new file mode 100644 index 00000000..e8f59ba2 --- /dev/null +++ b/config/client_profile_test.go @@ -0,0 +1,95 @@ +package config + +import ( + "strings" + "testing" +) + +// TestDeriveClientProfileNeverReturnsLinux là tính chất quan trọng nhất của pool: một tài khoản +// không bao giờ được gán dấu vân tay Linux — dấu hiệu nhận diện mạnh nhất cho việc "không phải +// người dùng Kiro IDE thực sự" khi triển khai trên máy chủ Linux. Bất kể accountID là gì (kể cả +// rỗng), kết quả phải luôn là mac/win. +func TestDeriveClientProfileNeverReturnsLinux(t *testing.T) { + for _, id := range []string{"", "a", "acct-1", "acct-2", "some-long-account-id-xyz"} { + p := DeriveClientProfile(id) + if strings.HasPrefix(p.systemVersion, "linux") { + t.Fatalf("accountID=%q derived Linux systemVersion %q — must never happen", id, p.systemVersion) + } + } +} + +// TestDeriveClientProfileStable xác nhận hàm thuần xác định: cùng accountID → cùng bộ ba dấu vân +// tay, qua nhiều lần gọi (một tài khoản trông như một desktop cố định trên các yêu cầu/khởi động lại). +func TestDeriveClientProfileStable(t *testing.T) { + a := DeriveClientProfile("acct-stable") + b := DeriveClientProfile("acct-stable") + if a != b { + t.Fatalf("same accountID must yield identical profile:\n a=%+v\n b=%+v", a, b) + } +} + +// TestDeriveClientProfileDistributionMatchesWeights xác nhận việc phân bổ theo trọng số đạt được: +// tính toán các lựa chọn cho một loạt các accountID và kỳ vọng tỷ lệ thô khớp với trọng số pool +// (mac được ưu tiên, một số ít dùng win). Sử dụng các ID được băm, không phải ngẫu nhiên, để +// không bị sai lệch (flaky). +func TestDeriveClientProfileDistributionMatchesWeights(t *testing.T) { + mac, win := 0, 0 + n := 5000 + for i := 0; i < n; i++ { + p := DeriveClientProfile("acct-" + string(rune('a'+i%26)) + string(rune('0'+i%10)) + "x") + if strings.HasPrefix(p.systemVersion, "darwin") { + mac++ + } else if strings.HasPrefix(p.systemVersion, "win32") { + win++ + } + } + // Trọng số pool: mac=85, win=15 → mac nên chiếm phần lớn. Chấp nhận sai số ±5%. + if mac < int(float64(n)*0.80) { + t.Fatalf("mac share too low: %d/%d (expected ~85%%)", mac, n) + } + if win < int(float64(n)*0.10) { + t.Fatalf("win share too low: %d/%d (expected ~15%%)", win, n) + } + if mac+win != n { + t.Fatalf("all profiles must be mac or win, got other: %d", n-mac-win) + } +} + +// TestDeriveClientProfileUsesVerifiedVersions xác nhận mỗi cấu hình pool mang theo đúng bộ Kiro +// 0.11.107 / node 22.22.0 đã được xác minh (chỉ đa dạng hóa nền tảng; đồng nhất phiên bản là +// được kỳ vọng cho người dùng thực tế tự động cập nhật). +func TestDeriveClientProfileUsesVerifiedVersions(t *testing.T) { + for _, id := range []string{"a", "b", "c"} { + p := DeriveClientProfile(id) + if p.kiroVersion != "0.11.107" { + t.Fatalf("kiroVersion must be the verified 0.11.107, got %q", p.kiroVersion) + } + if p.nodeVersion != "22.22.0" { + t.Fatalf("nodeVersion must be the verified 22.22.0, got %q", p.nodeVersion) + } + } +} + +// TestGetKiroClientConfigOperatorOverrideWinsPerField xác nhận các ghi đè từ config +// (KiroVersion/SystemVersion/NodeVersion) có mức ưu tiên cao hơn cho mỗi trường, giữ nguyên đường +// dẫn ghi đè thủ công. Các trường không được đặt sẽ lấy từ cấu hình được dẫn xuất. +func TestGetKiroClientConfigOperatorOverrideWinsPerField(t *testing.T) { + cfgLock.Lock() + old := cfg + cfg = &Config{SystemVersion: "win32#10.0.19045"} // chỉ ghi đè system + cfgLock.Unlock() + t.Cleanup(func() { + cfgLock.Lock() + cfg = old + cfgLock.Unlock() + }) + + got := GetKiroClientConfig("any-account") + if got.SystemVersion != "win32#10.0.19045" { + t.Fatalf("operator SystemVersion override must win, got %q", got.SystemVersion) + } + // Các trường chưa được đặt vẫn dùng các phiên bản đã được xác minh. + if got.KiroVersion != "0.11.107" || got.NodeVersion != "22.22.0" { + t.Fatalf("unset fields must fall back to verified versions, got %+v", got) + } +} diff --git a/config/config.go b/config/config.go index cc90aae5..34f9471c 100644 --- a/config/config.go +++ b/config/config.go @@ -12,10 +12,12 @@ package config import ( "crypto/rand" + "crypto/sha256" + "encoding/hex" "encoding/json" "fmt" "os" - "runtime" + "strings" "sync" "time" ) @@ -32,6 +34,101 @@ func GenerateMachineId() string { bytes[0:4], bytes[4:6], bytes[6:8], bytes[8:10], bytes[10:16]) } +// DeriveMachineId deterministically derives a stable 64-hex device id from an +// account ID via sha256("kiro-device-" + accountID). It is a pure function: +// the same account ID always yields the same value, so an account looks like a +// single fixed device across requests and across restarts. Ported from +// kiro-tutu (zero-dep: stdlib crypto/sha256 + encoding/hex). +// +// Used as the empty-MachineId fallback in kiro_headers.go so every account's +// User-Agent always carries a unique, stable device suffix instead of all +// empty-id accounts sharing an identical UA (the strongest cross-account +// association signal upstream can correlate on). +func DeriveMachineId(accountID string) string { + sum := sha256.Sum256([]byte("kiro-device-" + accountID)) + return hex.EncodeToString(sum[:]) +} + +// clientProfile is one entry in the desktop fingerprint pool: a real +// {OS#kernel, KiroIDE version, node version} triple that a genuine Kiro IDE +// install would report. Every value here MUST be a combination that actually +// shipped — fabricating version numbers would create a fresh "non-genuine +// client" signal, which is exactly what we are trying to avoid. Ported from +// kiro-tutu (zero-dep). +type clientProfile struct { + systemVersion string + kiroVersion string + nodeVersion string + weight int +} + +// clientProfilePool is the weighted distribution of desktop fingerprints used +// to derive a per-account {platform, version} triple. +// +// Design constraints: +// - macOS-weighted, with a small Windows minority, matching the real account +// source distribution. A wrong distribution manufactures a new +// "unrealistic population" signal. +// - NEVER contains Linux. A Linux server kernel in a Kiro IDE UA is the +// single strongest "not a real user" tell; this pool guarantees the output +// platform is always mac/win regardless of runtime.GOOS — so a deploy on a +// Linux box no longer advertises linux#6.6.87 in every request's UA. +// - Every systemVersion is a verified real macOS↔Darwin / Windows build: +// darwin#23.6.0 = macOS 14.6 (Sonoma), darwin#24.5.0 = macOS 15.5, +// darwin#24.6.0 = macOS 15.6 (Sequoia); win32#10.0.22631 = Win11 23H2, +// win32#10.0.19045 = Win10 22H2. +// - KiroVersion/NodeVersion stay on the one combination confirmed against a +// genuine client (0.11.107 / 22.22.0). Version uniformity is expected +// (real users auto-update), so only the platform dimension is diversified. +var clientProfilePool = []clientProfile{ + {systemVersion: "darwin#24.6.0", kiroVersion: "0.11.107", nodeVersion: "22.22.0", weight: 35}, // macOS 15.6 Sequoia + {systemVersion: "darwin#24.5.0", kiroVersion: "0.11.107", nodeVersion: "22.22.0", weight: 25}, // macOS 15.5 Sequoia + {systemVersion: "darwin#23.6.0", kiroVersion: "0.11.107", nodeVersion: "22.22.0", weight: 25}, // macOS 14.6 Sonoma + {systemVersion: "win32#10.0.22631", kiroVersion: "0.11.107", nodeVersion: "22.22.0", weight: 10}, // Win11 23H2 + {systemVersion: "win32#10.0.19045", kiroVersion: "0.11.107", nodeVersion: "22.22.0", weight: 5}, // Win10 22H2 +} + +// DeriveClientProfile deterministically derives a stable desktop fingerprint +// for an account by hashing sha256("kiro-profile-" + accountID) and selecting +// from clientProfilePool by cumulative weight. It is a pure function: the same +// account ID always yields the same {platform, version} triple, across +// requests and across restarts, so an account looks like a single fixed +// desktop install — and it never returns a Linux platform regardless of the +// host OS. +// +// The "kiro-profile-" prefix differs from DeriveMachineId's "kiro-device-" so +// the platform selection is independent of the device-id hash, while both +// remain stably bound to the same account. +func DeriveClientProfile(accountID string) clientProfile { + totalWeight := 0 + for _, p := range clientProfilePool { + totalWeight += p.weight + } + if totalWeight <= 0 || len(clientProfilePool) == 0 { + // Defensive: a misconfigured pool must still never emit Linux. Fall back + // to the first verified mac entry rather than runtime.GOOS. + return clientProfile{systemVersion: "darwin#24.6.0", kiroVersion: "0.11.107", nodeVersion: "22.22.0"} + } + + sum := sha256.Sum256([]byte("kiro-profile-" + accountID)) + // Use the first 8 bytes as an unsigned 64-bit selector, then mod by total + // weight. Deterministic and uniform enough for distribution purposes. + var selector uint64 + for i := 0; i < 8; i++ { + selector = selector<<8 | uint64(sum[i]) + } + target := int(selector % uint64(totalWeight)) + + cumulative := 0 + for _, p := range clientProfilePool { + cumulative += p.weight + if target < cumulative { + return p + } + } + return clientProfilePool[len(clientProfilePool)-1] // unreachable, defensive +} + // Account represents a Kiro API account with authentication credentials and usage statistics. type Account struct { // Basic identification @@ -45,14 +142,25 @@ type Account struct { RefreshToken string `json:"refreshToken"` // OAuth refresh token for token renewal ClientID string `json:"clientId,omitempty"` // OIDC client ID (for IdC auth) ClientSecret string `json:"clientSecret,omitempty"` // OIDC client secret (for IdC auth) - AuthMethod string `json:"authMethod"` // Authentication method: "idc" (AWS IdC) or "social" (GitHub/Google) - Provider string `json:"provider,omitempty"` // Identity provider name (e.g., "BuilderId", "GitHub") + AuthMethod string `json:"authMethod"` // Authentication method: "idc" (AWS IdC), "social" (GitHub/Google), "external_idp" (enterprise SSO, e.g. Azure AD), or "api_key" (Kiro API key used directly as bearer) + KiroApiKey string `json:"kiroApiKey,omitempty"` // API key credential, used directly as the bearer token when AuthMethod == "api_key" + AuthRegion string `json:"authRegion,omitempty"` // Region for token-refresh endpoints; falls back to Region + ApiRegion string `json:"apiRegion,omitempty"` // Region for API request hosts; falls back to Region + Provider string `json:"provider,omitempty"` // Identity provider name (e.g., "BuilderId", "GitHub", "AzureAD") Region string `json:"region"` // AWS region for OIDC endpoints StartUrl string `json:"startUrl,omitempty"` // AWS SSO start URL ExpiresAt int64 `json:"expiresAt,omitempty"` // Token expiration timestamp (Unix seconds) MachineId string `json:"machineId,omitempty"` // UUID machine identifier for request tracking ProfileArn string `json:"profileArn,omitempty"` // CodeWhisperer/Kiro profile ARN for generation requests + // External IdP (enterprise SSO, e.g. Microsoft 365 / Entra ID / Azure AD) refresh material. + // When AuthMethod == "external_idp" the credential is an IdP-issued OAuth token refreshed + // against TokenEndpoint using ClientID and Scopes (refresh_token grant), NOT the AWS SSO + // OIDC endpoint. IssuerURL is the OIDC issuer the endpoints were discovered from. + TokenEndpoint string `json:"tokenEndpoint,omitempty"` // External IdP OAuth2 token endpoint (refresh) + IssuerURL string `json:"issuerUrl,omitempty"` // External IdP OIDC issuer URL + Scopes string `json:"scopes,omitempty"` // Space-separated scopes granted by the external IdP + // Per-account outbound proxy (falls back to global ProxyURL if empty) ProxyURL string `json:"proxyURL,omitempty"` @@ -109,6 +217,62 @@ type Account struct { TotalCredits float64 `json:"totalCredits,omitempty"` // Cumulative credits consumed } +// IsApiKeyCredential reports whether the account authenticates with a Kiro API +// key used directly as the bearer token (tokentype: API_KEY), bypassing OAuth +// refresh and profile-ARN resolution. True when KiroApiKey is set OR the +// authMethod is "api_key"/"apikey" (case-insensitive). +func (a *Account) IsApiKeyCredential() bool { + if a == nil { + return false + } + if a.KiroApiKey != "" { + return true + } + m := strings.ToLower(a.AuthMethod) + return m == "api_key" || m == "apikey" +} + +// EffectiveAuthRegion returns the region used for token-refresh endpoints, with +// the fallback chain: account.AuthRegion > account.Region > global AuthRegion +// (if set and not us-east-1) > global Region (if set and not us-east-1) > us-east-1. +func (a *Account) EffectiveAuthRegion() string { + if a != nil { + if r := strings.TrimSpace(a.AuthRegion); r != "" { + return r + } + if r := strings.TrimSpace(a.Region); r != "" { + return r + } + } + if r := GetGlobalAuthRegion(); r != "" && r != "us-east-1" { + return r + } + if r := GetGlobalRegion(); r != "" && r != "us-east-1" { + return r + } + return "us-east-1" +} + +// EffectiveApiRegion returns the region used for API request hosts, with the +// same fallback chain as EffectiveAuthRegion but using ApiRegion. +func (a *Account) EffectiveApiRegion() string { + if a != nil { + if r := strings.TrimSpace(a.ApiRegion); r != "" { + return r + } + if r := strings.TrimSpace(a.Region); r != "" { + return r + } + } + if r := GetGlobalApiRegion(); r != "" && r != "us-east-1" { + return r + } + if r := GetGlobalRegion(); r != "" && r != "us-east-1" { + return r + } + return "us-east-1" +} + // PromptFilterRule defines a single custom prompt sanitization rule. // Type can be: "regex" (regexp find/replace within prompt) or // "lines-containing" (remove lines containing the match substring). @@ -174,6 +338,15 @@ type Config struct { // solely because usageCurrent >= usageLimit. AllowOverUsage bool `json:"allowOverUsage,omitempty"` + // Region defaults for accounts that omit per-account region/authRegion/apiRegion. + // Defaults to "us-east-1" when empty (see GetGlobalRegion* / Account.Effective*Region). + Region string `json:"region,omitempty"` + AuthRegion string `json:"authRegion,omitempty"` + ApiRegion string `json:"apiRegion,omitempty"` + // MaxPayloadBytes caps the serialized Kiro request body before upstream rejects + // it as oversized. <=0 means use DefaultMaxPayloadBytes. + MaxPayloadBytes int `json:"maxPayloadBytes,omitempty"` + // Proxy configuration: optional outbound proxy for Kiro API requests // Format: "socks5://host:port", "socks5://user:pass@host:port", // "http://host:port", "http://user:pass@host:port" @@ -203,6 +376,18 @@ type Config struct { // Can be overridden by the LOG_LEVEL environment variable. LogLevel string `json:"logLevel,omitempty"` + // PromptCacheMaxRatio caps the fraction of input tokens reported as cache_read + // in a single turn. Default 0.85. Raise to 0.95 for "continue"-heavy workloads + // where the newest content is minimal and >85% of input is genuinely from cache. + PromptCacheMaxRatio float64 `json:"promptCacheMaxRatio,omitempty"` + + // PromptCacheMaxEntries bounds the in-memory prompt-cache map; once exceeded, + // the least-recently-used entries are evicted (LRU). Default 131072. Sized so + // the prefix write-rate × TTL does not evict multi-turn history prefixes + // before the next turn reuses them (mirrors kiro-rs's 131072 default). The + // tracker clamps explicit small values up to 256. + PromptCacheMaxEntries int `json:"promptCacheMaxEntries,omitempty"` + // Global statistics (persisted across restarts) TotalRequests int `json:"totalRequests,omitempty"` // Total API requests received SuccessRequests int `json:"successRequests,omitempty"` // Successful requests count @@ -231,9 +416,19 @@ type AccountInfo struct { TrialExpiresAt int64 } +// DefaultMaxPayloadBytes is the runtime-configurable cap (2 MB) for the Kiro +// request body, below the observed ~2.15 MB AWS upstream rejection threshold. +const DefaultMaxPayloadBytes = 2_000_000 + // Version current version const Version = "1.1.2" +const ( + autoQuarantineSuspicious429Reason = "auto-quarantine: suspicious 429 pattern" + operatorDisabledReason = "operator-disabled" + autoQuarantineDuration = 30 * time.Minute +) + var ( cfg *Config cfgLock sync.RWMutex @@ -333,12 +528,27 @@ func newUUID() string { // Save persists the current configuration to the JSON file. // Uses indented formatting for human readability. +// +// Atomic write: marshalled JSON is written to a sibling temp file, then renamed +// over config.json. A crash mid-write can no longer truncate/corrupt the file — +// which would lose every account token, API key, and the admin password, since +// this is the single source of truth for all credentials. rename is atomic on +// the same filesystem (same dir), so readers see either the old or the new file +// in full, never a partial write. func Save() error { data, err := json.MarshalIndent(cfg, "", " ") if err != nil { return err } - return os.WriteFile(cfgPath, data, 0600) + tmp := cfgPath + ".tmp" + if err := os.WriteFile(tmp, data, 0600); err != nil { + return err + } + if err := os.Rename(tmp, cfgPath); err != nil { + _ = os.Remove(tmp) + return err + } + return nil } // SetPassword updates the admin password. @@ -398,16 +608,35 @@ func GetHost() string { } func GetAccounts() []Account { - cfgLock.RLock() - defer cfgLock.RUnlock() + cfgLock.Lock() + defer cfgLock.Unlock() + applyAutoRestoreLocked() accounts := make([]Account, len(cfg.Accounts)) copy(accounts, cfg.Accounts) return accounts } -func GetEnabledAccounts() []Account { +// AccountIDExists reports whether an account with the given ID is already stored. +// Used by the credential-import path to reuse a pasted record's id when it does +// not collide, so re-importing a backup never creates a duplicate entry. +func AccountIDExists(id string) bool { + if id == "" { + return false + } cfgLock.RLock() defer cfgLock.RUnlock() + for _, a := range cfg.Accounts { + if a.ID == id { + return true + } + } + return false +} + +func GetEnabledAccounts() []Account { + cfgLock.Lock() + defer cfgLock.Unlock() + applyAutoRestoreLocked() var accounts []Account for _, a := range cfg.Accounts { if a.Enabled { @@ -420,6 +649,17 @@ func GetEnabledAccounts() []Account { func AddAccount(account Account) error { cfgLock.Lock() defer cfgLock.Unlock() + // Reject a duplicate id under the write lock. The import path pre-checks with + // AccountIDExists (RLock) and mints a fresh id on collision, but that check and this + // append are not atomic; two concurrent imports of the same pasted id could both + // pass the pre-check. This makes "add if id absent" the atomic invariant. + if account.ID != "" { + for _, a := range cfg.Accounts { + if a.ID == account.ID { + return fmt.Errorf("account with id %s already exists", account.ID) + } + } + } cfg.Accounts = append(cfg.Accounts, account) return Save() } @@ -470,7 +710,11 @@ func SetAccountEnabled(id string, enabled bool) error { for i, a := range cfg.Accounts { if a.ID == id { cfg.Accounts[i].Enabled = enabled - if !enabled { + if enabled { + cfg.Accounts[i].BanStatus = "ACTIVE" + cfg.Accounts[i].BanReason = "" + cfg.Accounts[i].BanTime = 0 + } else { cfg.Accounts[i].BanStatus = "DISABLED" cfg.Accounts[i].BanTime = time.Now().Unix() } @@ -499,6 +743,169 @@ func SetAccountBanStatus(id, status, reason string) error { return nil } +func AutoQuarantineSuspicious429Reason() string { + return autoQuarantineSuspicious429Reason +} + +// OperatorDisabledReason is the BanReason stamped when a human operator disables +// an account. It marks the account as DISABLED (not SUSPENDED), which keeps the +// auto-restore sweep from ever re-enabling it. +func OperatorDisabledReason() string { + return operatorDisabledReason +} + +func shouldAutoRestoreSuspendedAccount(a Account, now time.Time) bool { + return a.BanStatus == "SUSPENDED" && a.BanReason == autoQuarantineSuspicious429Reason && a.BanTime > 0 && now.Unix()-a.BanTime >= int64(autoQuarantineDuration/time.Second) +} + +func applyAutoRestoreLocked() bool { + if cfg == nil { + return false + } + now := time.Now() + changed := false + for i := range cfg.Accounts { + if shouldAutoRestoreSuspendedAccount(cfg.Accounts[i], now) { + cfg.Accounts[i].Enabled = true + cfg.Accounts[i].BanStatus = "ACTIVE" + cfg.Accounts[i].BanReason = "" + cfg.Accounts[i].BanTime = 0 + changed = true + } + } + if changed { + _ = Save() + } + return changed +} + +// AddAccounts appends multiple accounts in a single locked pass and persists +// with exactly one Save(), avoiding the O(n²) write amplification that calling +// AddAccount in a loop would cause (each AddAccount re-serializes the entire +// config.json). Accounts whose RefreshToken already exists (against the current +// config or earlier entries in the same batch) are skipped to keep bulk imports +// idempotent across retries/re-pastes. Entries with an empty RefreshToken are +// also skipped — there is no stable identity to dedup on and they cannot be +// activated later. Returns how many were added and how many were skipped. +// +// Save() is only invoked when at least one account is actually added, so a +// fully-duplicate batch does not churn the config file. +func AddAccounts(accounts []Account) (added int, skipped int, err error) { + cfgLock.Lock() + defer cfgLock.Unlock() + + // Seed the seen-set with refresh tokens already persisted so the batch + // dedups against existing accounts, not just within itself. + seen := make(map[string]struct{}, len(cfg.Accounts)+len(accounts)) + for i := range cfg.Accounts { + if rt := cfg.Accounts[i].RefreshToken; rt != "" { + seen[rt] = struct{}{} + } + } + + for _, a := range accounts { + if a.RefreshToken == "" { + skipped++ + continue + } + if _, dup := seen[a.RefreshToken]; dup { + skipped++ + continue + } + seen[a.RefreshToken] = struct{}{} + cfg.Accounts = append(cfg.Accounts, a) + added++ + } + + if added == 0 { + return 0, skipped, nil + } + if err := Save(); err != nil { + // Roll back the in-memory appends so a failed persist does not leave + // the running pool out of sync with what is on disk. + cfg.Accounts = cfg.Accounts[:len(cfg.Accounts)-added] + return 0, skipped, err + } + return added, skipped, nil +} + +// RefreshTokenExists reports whether any account already holds the given refresh +// token. Used by bulk import to dedup candidates before spending an upstream +// token-exchange round-trip on a duplicate. +func RefreshTokenExists(refreshToken string) bool { + if refreshToken == "" { + return false + } + cfgLock.RLock() + defer cfgLock.RUnlock() + for i := range cfg.Accounts { + if cfg.Accounts[i].RefreshToken == refreshToken { + return true + } + } + return false +} + +// FindAccountIDByRefreshToken returns the id of the account that already holds +// the given refresh token, or "" if none. Used by the single credential-import +// path to update an existing entry in place when a backup is re-imported, rather +// than minting a fresh id and leaving two live accounts sharing the same token. +// Mirrors the refresh-token dedup the bulk path applies (AddAccounts). +func FindAccountIDByRefreshToken(refreshToken string) string { + if refreshToken == "" { + return "" + } + cfgLock.RLock() + defer cfgLock.RUnlock() + for i := range cfg.Accounts { + if cfg.Accounts[i].RefreshToken == refreshToken { + return cfg.Accounts[i].ID + } + } + return "" +} + +func SuspendAccountTemporarily(id, reason string) error { + cfgLock.Lock() + defer cfgLock.Unlock() + now := time.Now().Unix() + for i, a := range cfg.Accounts { + if a.ID == id { + cfg.Accounts[i].Enabled = false + cfg.Accounts[i].BanStatus = "SUSPENDED" + cfg.Accounts[i].BanReason = reason + cfg.Accounts[i].BanTime = now + return Save() + } + } + return nil +} + +// ClearAccountCurrentOverages zeroes the cached CurrentOverages for an account +// while preserving the OverageStatus switch and the cap/rate billing config. +// Called when upstream usage has fallen back within the subscription quota +// (e.g. after a billing-period reset): overage points are zero by definition +// when usage is within quota, so stale points from a previous period must not +// linger in the UI/scheduler. Returns without writing if already zero, so the +// periodic refresh loop does not churn the config file every cycle. +func ClearAccountCurrentOverages(id string, checkedAt int64) error { + cfgLock.Lock() + defer cfgLock.Unlock() + for i, a := range cfg.Accounts { + if a.ID == id { + if cfg.Accounts[i].CurrentOverages == 0 { + return nil + } + cfg.Accounts[i].CurrentOverages = 0 + if checkedAt > 0 { + cfg.Accounts[i].OverageCheckedAt = checkedAt + } + return Save() + } + } + return nil +} + func UpdateAccountProfileArn(id, profileArn string) error { cfgLock.Lock() defer cfgLock.Unlock() @@ -837,6 +1244,55 @@ func UpdateAllowOverUsage(allow bool) error { return Save() } +// GetMaxPayloadBytes returns the configured payload cap, falling back to +// DefaultMaxPayloadBytes when unset (<=0). +func GetMaxPayloadBytes() int { + cfgLock.RLock() + defer cfgLock.RUnlock() + if cfg == nil || cfg.MaxPayloadBytes <= 0 { + return DefaultMaxPayloadBytes + } + return cfg.MaxPayloadBytes +} + +// UpdateMaxPayloadBytes sets the payload cap and persists the change. +func UpdateMaxPayloadBytes(n int) error { + cfgLock.Lock() + defer cfgLock.Unlock() + cfg.MaxPayloadBytes = n + return Save() +} + +// GetGlobalRegion returns the configured default region (empty → "us-east-1"). +func GetGlobalRegion() string { + cfgLock.RLock() + defer cfgLock.RUnlock() + if cfg == nil || cfg.Region == "" { + return "us-east-1" + } + return cfg.Region +} + +// GetGlobalAuthRegion returns the configured default token-refresh region. +func GetGlobalAuthRegion() string { + cfgLock.RLock() + defer cfgLock.RUnlock() + if cfg == nil || cfg.AuthRegion == "" { + return "us-east-1" + } + return cfg.AuthRegion +} + +// GetGlobalApiRegion returns the configured default API-request host region. +func GetGlobalApiRegion() string { + cfgLock.RLock() + defer cfgLock.RUnlock() + if cfg == nil || cfg.ApiRegion == "" { + return "us-east-1" + } + return cfg.ApiRegion +} + // GetLogLevel returns the configured log level (debug/info/warn/error). Defaults to "info". func GetLogLevel() string { cfgLock.RLock() @@ -847,6 +1303,54 @@ func GetLogLevel() string { return cfg.LogLevel } +// GetPromptCacheMaxRatio returns the cache-read cap ratio (0.0-1.0). Defaults to 0.85. +func GetPromptCacheMaxRatio() float64 { + cfgLock.RLock() + defer cfgLock.RUnlock() + if cfg == nil || cfg.PromptCacheMaxRatio <= 0 || cfg.PromptCacheMaxRatio > 1 { + return 0.85 + } + return cfg.PromptCacheMaxRatio +} + +// UpdatePromptCacheMaxRatio sets the cache-read cap ratio and persists the change. +func UpdatePromptCacheMaxRatio(ratio float64) error { + cfgLock.Lock() + defer cfgLock.Unlock() + cfg.PromptCacheMaxRatio = ratio + return Save() +} + +const defaultPromptCacheMaxEntries = 131072 +const minPromptCacheEntries = 256 + +// GetPromptCacheMaxEntries returns the prompt-cache LRU bound. Defaults to +// 131072 when unset (≤ 0); an explicit small value is clamped up to +// minPromptCacheEntries (256) so a misconfigured tiny value cannot make the +// cache useless. This is the production safety floor — the tracker constructor +// trusts its caller (tests may use any capacity). +func GetPromptCacheMaxEntries() int { + cfgLock.RLock() + defer cfgLock.RUnlock() + if cfg == nil || cfg.PromptCacheMaxEntries <= 0 { + return defaultPromptCacheMaxEntries + } + if cfg.PromptCacheMaxEntries < minPromptCacheEntries { + return minPromptCacheEntries + } + return cfg.PromptCacheMaxEntries +} + +// UpdatePromptCacheMaxEntries sets the prompt-cache LRU bound and persists it. +// Applies on the next tracker construction (restart); it does not resize a +// live tracker. +func UpdatePromptCacheMaxEntries(n int) error { + cfgLock.Lock() + defer cfgLock.Unlock() + cfg.PromptCacheMaxEntries = n + return Save() +} + // UpdateLogLevel updates the log level setting and persists the change. func UpdateLogLevel(level string) error { cfgLock.Lock() @@ -855,32 +1359,43 @@ func UpdateLogLevel(level string) error { return Save() } +// KiroClientConfig is the {KiroVersion, SystemVersion, NodeVersion} triple +// embedded into the upstream User-Agent. Populated by GetKiroClientConfig. type KiroClientConfig struct { KiroVersion string SystemVersion string NodeVersion string } -func GetKiroClientConfig() KiroClientConfig { +// GetKiroClientConfig resolves the {KiroVersion, SystemVersion, NodeVersion} +// triple for accountID. When no operator override is set on a field, it derives +// a stable desktop fingerprint from the account via DeriveClientProfile — which +// selects from clientProfilePool (mac/win-weighted) and NEVER returns Linux, so +// a deploy on a Linux box no longer advertises linux#6.6.87 in every request's +// User-Agent. An empty accountID still yields a stable mac-default profile +// (never Linux). Ported from kiro-tutu (zero-dep). +// +// Operator overrides from config (KiroVersion/SystemVersion/NodeVersion) take +// precedence per-field when explicitly set, preserving the manual-override path. +func GetKiroClientConfig(accountID string) KiroClientConfig { cfgLock.RLock() defer cfgLock.RUnlock() - kiroVersion := "0.11.107" - if cfg != nil && cfg.KiroVersion != "" { - kiroVersion = cfg.KiroVersion - } + profile := DeriveClientProfile(accountID) + kiroVersion := profile.kiroVersion + systemVersion := profile.systemVersion + nodeVersion := profile.nodeVersion - systemVersion := "" if cfg != nil { - systemVersion = cfg.SystemVersion - } - if systemVersion == "" { - systemVersion = defaultSystemVersion() - } - - nodeVersion := "22.22.0" - if cfg != nil && cfg.NodeVersion != "" { - nodeVersion = cfg.NodeVersion + if cfg.KiroVersion != "" { + kiroVersion = cfg.KiroVersion + } + if cfg.SystemVersion != "" { + systemVersion = cfg.SystemVersion + } + if cfg.NodeVersion != "" { + nodeVersion = cfg.NodeVersion + } } return KiroClientConfig{ @@ -889,14 +1404,3 @@ func GetKiroClientConfig() KiroClientConfig { NodeVersion: nodeVersion, } } - -func defaultSystemVersion() string { - switch runtime.GOOS { - case "windows": - return "win32#10.0.22631" - case "darwin": - return "darwin#24.6.0" - default: - return "linux#6.6.87" - } -} diff --git a/config/config_test.go b/config/config_test.go index d9d63b5e..ed93c828 100644 --- a/config/config_test.go +++ b/config/config_test.go @@ -127,3 +127,75 @@ func TestAccountAllowOverageMigration(t *testing.T) { } } } + +func TestIsApiKeyCredential(t *testing.T) { + cases := []struct { + name string + a Account + want bool + }{ + {"key present", Account{KiroApiKey: "k"}, true}, + {"api_key lower", Account{AuthMethod: "api_key"}, true}, + {"apikey lower", Account{AuthMethod: "apikey"}, true}, + {"API_KEY upper", Account{AuthMethod: "API_KEY"}, true}, + {"key wins over idc", Account{KiroApiKey: "k", AuthMethod: "idc"}, true}, + {"idc not api key", Account{AuthMethod: "idc"}, false}, + {"external_idp not api key", Account{AuthMethod: "external_idp"}, false}, + {"empty", Account{}, false}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + if got := tc.a.IsApiKeyCredential(); got != tc.want { + t.Fatalf("IsApiKeyCredential() = %v, want %v", got, tc.want) + } + }) + } +} + +func TestEffectiveRegionsFallbackChain(t *testing.T) { + if err := Init(filepath.Join(t.TempDir(), "config.json")); err != nil { + t.Fatalf("init: %v", err) + } + + // account-level wins + a := &Account{AuthRegion: "eu-west-1", ApiRegion: "ap-southeast-1", Region: "us-west-2"} + if got := a.EffectiveAuthRegion(); got != "eu-west-1" { + t.Fatalf("auth: want eu-west-1, got %q", got) + } + if got := a.EffectiveApiRegion(); got != "ap-southeast-1" { + t.Fatalf("api: want ap-southeast-1, got %q", got) + } + + // falls back to Region + b := &Account{Region: "us-west-2"} + if got := b.EffectiveAuthRegion(); got != "us-west-2" { + t.Fatalf("auth fallback: want us-west-2, got %q", got) + } + if got := b.EffectiveApiRegion(); got != "us-west-2" { + t.Fatalf("api fallback: want us-west-2, got %q", got) + } + + // empty account → default us-east-1 + c := &Account{} + if got := c.EffectiveAuthRegion(); got != "us-east-1" { + t.Fatalf("auth default: want us-east-1, got %q", got) + } + if got := c.EffectiveApiRegion(); got != "us-east-1" { + t.Fatalf("api default: want us-east-1, got %q", got) + } +} + +func TestMaxPayloadBytesDefaultFallbackAndPersist(t *testing.T) { + if err := Init(filepath.Join(t.TempDir(), "config.json")); err != nil { + t.Fatalf("init: %v", err) + } + if got := GetMaxPayloadBytes(); got != DefaultMaxPayloadBytes { + t.Fatalf("default: want %d, got %d", DefaultMaxPayloadBytes, got) + } + if err := UpdateMaxPayloadBytes(2100000); err != nil { + t.Fatalf("update: %v", err) + } + if got := GetMaxPayloadBytes(); got != 2100000 { + t.Fatalf("after set: want 2100000, got %d", got) + } +} diff --git a/docker-compose.yml b/docker-compose.yml index 78ebab89..661182f0 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -5,8 +5,23 @@ services: build: . ports: - "8080:8080" + # Enterprise SSO (Microsoft 365) callback. The hosted-portal sign-in + # redirects the browser to localhost:3128, so the operator must run that + # browser ON THE DOCKER HOST. Publishing host-side to 127.0.0.1 only keeps + # the callback off the host's external interfaces while still letting the + # host browser's localhost:3128 reach the container listener below. + - "127.0.0.1:3128:3128" volumes: - ./data:/app/data environment: - CONFIG_PATH=/app/data/config.json + # Bind host for the Enterprise SSO callback listener INSIDE the container. + # Default (unset) is loopback-only (127.0.0.1/::1), which the published port + # above cannot reach; 0.0.0.0 lets the forwarded connection land. The callback + # is transient (~10-minute login window) and anti-CSRF state-matched, and the + # host-side publish is pinned to loopback above, so it is not reachable off the + # host. Note 0.0.0.0 does expose 3128 to peer containers on the same Docker + # network during that window; on a shared/untrusted Docker network, pin + # KIRO_SSO_CALLBACK_BIND to a specific interface or isolate the network. + - KIRO_SSO_CALLBACK_BIND=0.0.0.0 restart: unless-stopped diff --git a/docs/superpowers/plans/2026-06-27-cache-improvements.md b/docs/superpowers/plans/2026-06-27-cache-improvements.md new file mode 100644 index 00000000..1a38748d --- /dev/null +++ b/docs/superpowers/plans/2026-06-27-cache-improvements.md @@ -0,0 +1,547 @@ +# Cache Improvements (C1/C2/C3) Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. + +**Goal:** Improve the prompt-cache accounting layer (`proxy/cache_tracker.go`) — cross-account sharing (C1), configurable cache cap (C2), and disk persistence (C3) — so hit-rate accuracy and restart-survival improve for multi-account pools. + +**Architecture:** The cache tracker is an in-memory SHA-256 fingerprint store that mirrors Anthropic's upstream prompt caching to report `cache_creation/cache_read` tokens. C1 removes the per-account isolation (biggest hit-rate win); C2 makes the 85% cap configurable; C3 persists entries to disk so restart doesn't lose them. + +**Tech Stack:** Go 1.21 (module `kiro-go`, stdlib), `encoding/json` for disk persistence. + +## Global Constraints + +- The cache tracker is ACCOUNTING ONLY — it reports cache tokens; it does NOT cache responses. No upstream behavior changes. +- Canonical fingerprint = SHA-256 running hash of canonicalized blocks (system + tools + messages). This is correct; do NOT change fingerprinting. +- TTL: 5min default, 1h max. Min cacheable: 1024 tokens (4096 opus). These match Anthropic spec — do NOT change. +- Module path `kiro-go`. Run `go test ./...` from `C:\Users\Admin\Kiro-Go`. +- Disk file: `data/prompt_cache.json` (already exists, currently unused). Load on startup, write debounced. + +--- + +### Task 1: C1 — Cross-account cache sharing (remove per-account isolation) + +**Files:** +- Modify: `proxy/cache_tracker.go` (struct `promptCacheTracker`, `Compute`, `Update`, `pruneExpiredLocked`) +- Modify: `proxy/cache_tracker_test.go` (add cross-account hit test) + +**Why:** Currently `entriesByAccount[accountID]` means N accounts in a pool each build their own cache → hit rate drops to ~1/N. Since all accounts are in the same Anthropic org (tenant codezdevbatman), upstream cache IS shared. Making the fingerprint store global (drop accountID from the key) corrects the accounting → hit rate jumps from ~(1/N)×70% to ~70%. + +- [ ] **Step 1: Write the failing test** + +Append to `proxy/cache_tracker_test.go`: + +```go +// TestPromptCacheCrossAccountSharing verifies C1: two different accountIDs with +// the SAME prompt fingerprint share cache entries. Account B's request should +// HIT on the fingerprint Account A stored — no per-account isolation. +func TestPromptCacheCrossAccountSharing(t *testing.T) { + tracker := newPromptCacheTracker(5 * time.Minute) + + // Build a profile with one explicit cache_control breakpoint above the + // min-token threshold. + block := cacheablePromptBlock{ + Value: map[string]interface{}{"kind": "system", "block": map[string]interface{}{ + "type": "text", "text": strings.Repeat("x ", 600), // ~600 tokens > 1024? use more + "cache_control": map[string]interface{}{"type": "ephemeral"}, + }}, + Tokens: 1200, + TTL: 5 * time.Minute, + } + hasher := sha256.New() + writeHashChunk(hasher, canonicalizeCacheValue(block.Value)) + var fp [32]byte + copy(fp[:], hasher.Sum(nil)) + + profile := &promptCacheProfile{ + Breakpoints: []promptCacheBreakpoint{{Fingerprint: fp, CumulativeTokens: 1200, TTL: 5 * time.Minute}}, + TotalInputTokens: 1200, + Model: "claude-sonnet-4-5", + } + + // Account A: first request → cache_creation. + usageA := tracker.Compute("account-A", profile) + if usageA.CacheCreationInputTokens == 0 { + t.Fatalf("account A: expected cache_creation > 0, got %d", usageA.CacheCreationInputTokens) + } + tracker.Update("account-A", profile) + + // Account B: SAME prompt, DIFFERENT account → should be cache_read (C1 fix). + // Before C1: account B had its own empty store → cache_creation. + // After C1: account B shares the global store → cache_read. + usageB := tracker.Compute("account-B", profile) + if usageB.CacheReadInputTokens == 0 { + t.Fatalf("account B: expected cache_read > 0 (cross-account sharing), got 0. usage=%+v", usageB) + } +} +``` + +Ensure `crypto/sha256`, `strings`, `time` are imported in the test file. + +- [ ] **Step 2: Run test to verify it fails** + +Run: `go test ./proxy/ -run TestPromptCacheCrossAccountSharing -v` +Expected: FAIL — `account B: expected cache_read > 0` (current per-account isolation means account B's store is empty → cache_creation, not cache_read). + +- [ ] **Step 3: Change the struct to use a global entries map** + +In `proxy/cache_tracker.go`, replace the struct definition: + +```go +type promptCacheTracker struct { + mu sync.Mutex + entries map[[32]byte]promptCacheEntry + maxSupportedTTL time.Duration +} +``` + +And update `newPromptCacheTracker`: + +```go +func newPromptCacheTracker(maxTTL time.Duration) *promptCacheTracker { + if maxTTL <= 0 { + maxTTL = defaultPromptCacheTTL + } + return &promptCacheTracker{ + entries: make(map[[32]byte]promptCacheEntry), + maxSupportedTTL: maxTTL, + } +} +``` + +- [ ] **Step 4: Update `Compute` to use the global map** + +Replace the body of `Compute` — change `entries := t.entriesByAccount[accountID]` and all `entries[...]` references to use `t.entries` directly. The `accountID` parameter stays in the signature (backward compat for callers) but is now IGNORED (the store is global). + +In `Compute`, replace: +```go + entries := t.entriesByAccount[accountID] + if len(entries) == 0 { +``` +With: +```go + if len(t.entries) == 0 { +``` +And replace all `entries[breakpoint.Fingerprint]` with `t.entries[breakpoint.Fingerprint]`, and `entry, ok := entries[...]` with `entry, ok := t.entries[...]`. + +- [ ] **Step 5: Update `Update` to use the global map** + +In `Update`, replace: +```go + entries := t.entriesByAccount[accountID] + if entries == nil { + entries = make(map[[32]byte]promptCacheEntry) + t.entriesByAccount[accountID] = entries + } +``` +With: +```go + // entries is the global map now (C1: cross-account sharing). +``` +And replace `entries[breakpoint.Fingerprint] = ...` with `t.entries[breakpoint.Fingerprint] = ...`. + +- [ ] **Step 6: Update `pruneExpiredLocked` to iterate the single map** + +Replace: +```go +func (t *promptCacheTracker) pruneExpiredLocked(now time.Time) { + for accountID, entries := range t.entriesByAccount { + for fingerprint, entry := range entries { + if !entry.ExpiresAt.After(now) { + delete(entries, fingerprint) + } + } + if len(entries) == 0 { + delete(t.entriesByAccount, accountID) + } + } +} +``` +With: +```go +func (t *promptCacheTracker) pruneExpiredLocked(now time.Time) { + for fingerprint, entry := range t.entries { + if !entry.ExpiresAt.After(now) { + delete(t.entries, fingerprint) + } + } +} +``` + +- [ ] **Step 7: Run the test to verify it passes** + +Run: `go test ./proxy/ -run TestPromptCacheCrossAccountSharing -v` +Expected: PASS — account B now hits the global store → cache_read > 0. +Then run the full cache test suite: `go test ./proxy/ -run TestPromptCache -v` +Expected: all PASS (existing tests use accountID but it's now ignored — they still pass because the store works regardless). + +- [ ] **Step 8: Commit** + +```bash +git add proxy/cache_tracker.go proxy/cache_tracker_test.go +git commit -m "fix(cache): cross-account prompt-cache sharing (C1) + +Remove per-account isolation from the fingerprint store so N accounts in a +pool share cache entries (they're in the same Anthropic org). Hit rate jumps +from ~(1/N)*70% to ~70%." +``` + +--- + +### Task 2: C2 — Configurable cache cap (replace hardcoded 0.85) + +**Files:** +- Modify: `proxy/cache_tracker.go` (`promptCacheTracker` struct, `Compute`) +- Modify: `proxy/cache_tracker_test.go` (cap test) +- Modify: `config/config.go` (`Config` struct + getter) + +**Why:** The 85% cap (`maxCacheable = total * 0.85`) is a heuristic that under-reports cache hits on "continue" requests (tiny new turn, 90%+ from cache). Making it configurable lets operators raise it to 0.95 for workloads where the newest content is minimal. + +- [ ] **Step 1: Add the config field** + +In `config/config.go`, add to the `Config` struct (after `LogLevel` or near the cache-related fields): + +```go + // PromptCacheMaxRatio caps the fraction of input tokens reported as cache_read + // in a single turn. Default 0.85. Raise to 0.95 for "continue"-heavy workloads + // where the newest content is minimal and >85% of input is genuinely from cache. + PromptCacheMaxRatio float64 `json:"promptCacheMaxRatio,omitempty"` +``` + +Add a getter after `GetLogLevel`: + +```go +// GetPromptCacheMaxRatio returns the cache-read cap ratio (0.0-1.0). Defaults to 0.85. +func GetPromptCacheMaxRatio() float64 { + cfgLock.RLock() + defer cfgLock.RUnlock() + if cfg == nil || cfg.PromptCacheMaxRatio <= 0 || cfg.PromptCacheMaxRatio > 1 { + return 0.85 + } + return cfg.PromptCacheMaxRatio +} +``` + +- [ ] **Step 2: Use the config value in Compute** + +In `proxy/cache_tracker.go` `Compute`, replace: +```go + maxCacheable := int(float64(profile.TotalInputTokens) * 0.85) +``` +With: +```go + maxCacheable := int(float64(profile.TotalInputTokens) * config.GetPromptCacheMaxRatio()) +``` + +Ensure `"kiro-go/config"` is imported in `cache_tracker.go`. + +- [ ] **Step 3: Write the test** + +Append to `proxy/cache_tracker_test.go`: + +```go +// TestPromptCacheCapConfigurable verifies C2: the cache-read cap can be set +// above the default 0.85 via config, so a request where 90% of input is from +// cache reports the full 90% (not clamped to 85%). +func TestPromptCacheCapConfigurable(t *testing.T) { + cfgFile := t.TempDir() + "/config.json" + if err := config.Init(cfgFile); err != nil { + t.Fatalf("config.Init: %v", err) + } + + // Build a profile: 1000 tokens cached, 100 new (total 1100). Default cap + // would clamp to 0.85*1100=935. With cap=0.95, allows up to 1045. + hasher := sha256.New() + writeHashChunk(hasher, canonicalizeCacheValue(map[string]interface{}{"k": strings.Repeat("v ", 500)})) + var fp [32]byte + copy(fp[:], hasher.Sum(nil)) + profile := &promptCacheProfile{ + Breakpoints: []promptCacheBreakpoint{{Fingerprint: fp, CumulativeTokens: 1000, TTL: 5 * time.Minute}}, + TotalInputTokens: 1100, + Model: "claude-sonnet-4-5", + } + tracker := newPromptCacheTracker(5 * time.Minute) + tracker.Update("acc", profile) // store it + + // Default cap 0.85: cache_read clamped to min(1000, 0.85*1100=935) = 935. + usage85 := tracker.Compute("acc", profile) + if usage85.CacheReadInputTokens > 940 { + t.Fatalf("default cap: expected cache_read ~935, got %d", usage85.CacheReadInputTokens) + } + + // Raise cap to 0.95: cache_read should be min(1000, 0.95*1100=1045) = 1000. + config.UpdatePromptCacheMaxRatio(0.95) + defer config.UpdatePromptCacheMaxRatio(0.85) + usage95 := tracker.Compute("acc", profile) + if usage95.CacheReadInputTokens < 990 { + t.Fatalf("cap 0.95: expected cache_read ~1000, got %d", usage95.CacheReadInputTokens) + } +} +``` + +Add `UpdatePromptCacheMaxRatio` to `config/config.go`: +```go +func UpdatePromptCacheMaxRatio(ratio float64) error { + cfgLock.Lock() + defer cfgLock.Unlock() + cfg.PromptCacheMaxRatio = ratio + return Save() +} +``` + +- [ ] **Step 4: Run tests** + +Run: `go test ./proxy/ -run 'TestPromptCacheCap|TestPromptCacheCrossAccount' -v` +Expected: PASS. +Then: `go test ./proxy/ ./config/ -v` +Expected: all PASS. + +- [ ] **Step 5: Commit** + +```bash +git add proxy/cache_tracker.go proxy/cache_tracker_test.go config/config.go +git commit -m "feat(cache): configurable cache-read cap (C2) + +Replace the hardcoded 0.85 cache-read cap with a config value +(promptCacheMaxRatio, default 0.85). Operators can raise it to 0.95 for +continue-heavy workloads where >85% of input is genuinely from cache." +``` + +--- + +### Task 3: C3 — Disk persistence (load on startup, debounce-write) + +**Files:** +- Modify: `proxy/cache_tracker.go` (`promptCacheTracker`, add `Load`, `saveLoop`, `dirty` flag) +- Modify: `proxy/cache_tracker_test.go` (disk round-trip test) +- Modify: `proxy/handler.go` (call `Load` on tracker init + start `saveLoop`) + +**Why:** Currently the cache is purely in-memory → restart loses all entries → every request after restart is `cache_creation` until fingerprints rebuild. Persisting to `data/prompt_cache.json` (which already exists but is unused) lets entries survive restart. + +- [ ] **Step 1: Add persistence fields to the struct** + +In `proxy/cache_tracker.go`, update the struct: + +```go +type promptCacheTracker struct { + mu sync.Mutex + entries map[[32]byte]promptCacheEntry + maxSupportedTTL time.Duration + dirty bool + stopChan chan struct{} +} +``` + +- [ ] **Step 2: Add Load (called on startup)** + +Add after `newPromptCacheTracker`: + +```go +// on-disk format for prompt-cache persistence (C3). +type promptCacheEntryOnDisk struct { + Fingerprint [32]byte + ExpiresAt int64 // unix seconds + TTLSeconds int64 +} + +// Load reads persisted cache entries from path. Entries already expired (by the +// time load finishes) are dropped. Best-effort: a corrupt/missing file is not +// fatal — the tracker starts empty, same as the pre-C3 behavior. +func (t *promptCacheTracker) Load(path string) { + data, err := os.ReadFile(path) + if err != nil { + return // missing file = fresh start (normal on first run) + } + var disk struct { + Version int `json:"version"` + Entries []promptCacheEntryOnDisk `json:"entries"` + } + if json.Unmarshal(data, &disk) != nil { + return // corrupt = fresh start + } + now := time.Now() + t.mu.Lock() + defer t.mu.Unlock() + for _, e := range disk.Entries { + exp := time.Unix(e.ExpiresAt, 0) + if !exp.After(now) { + continue // already expired + } + t.entries[e.Fingerprint] = promptCacheEntry{ + ExpiresAt: exp, + TTL: time.Duration(e.TTLSeconds) * time.Second, + } + } +} +``` + +Ensure `"os"` and `"encoding/json"` are imported in `cache_tracker.go`. + +- [ ] **Step 3: Add saveLoop (debounced background writer)** + +```go +// startSaveLoop launches a background goroutine that flushes the cache to path +// every flushInterval (if dirty). Call once after Load. The goroutine exits when +// stopChan is closed. +func (t *promptCacheTracker) startSaveLoop(path string, flushInterval time.Duration) { + t.stopChan = make(chan struct{}) + go func() { + ticker := time.NewTicker(flushInterval) + defer ticker.Stop() + for { + select { + case <-ticker.C: + t.flush(path) + case <-t.stopChan: + t.flush(path) // final flush + return + } + } + }() +} + +func (t *promptCacheTracker) Stop() { + if t.stopChan != nil { + close(t.stopChan) + } +} + +func (t *promptCacheTracker) flush(path string) { + t.mu.Lock() + if !t.dirty { + t.mu.Unlock() + return + } + now := time.Now() + entries := make([]promptCacheEntryOnDisk, 0, len(t.entries)) + for fp, e := range t.entries { + if !e.ExpiresAt.After(now) { + continue + } + entries = append(entries, promptCacheEntryOnDisk{ + Fingerprint: fp, + ExpiresAt: e.ExpiresAt.Unix(), + TTLSeconds: int64(e.TTL.Seconds()), + }) + } + t.dirty = false + t.mu.Unlock() + + data, _ := json.MarshalIndent(map[string]interface{}{ + "version": 1, + "entries": entries, + }, "", " ") + _ = os.WriteFile(path, data, 0600) +} +``` + +- [ ] **Step 4: Mark dirty on Update** + +In `Update`, after storing entries, add: +```go + t.dirty = true +``` +(right before `t.mu.Unlock()` — or after the loop that writes entries, inside the lock). + +- [ ] **Step 5: Wire Load + startSaveLoop in the handler** + +Find where `newPromptCacheTracker` is called (in `proxy/handler.go` `NewHandler` or similar). After creating the tracker, add: + +```go + cachePath := filepath.Join(config.GetConfigDir(), "prompt_cache.json") + promptCache.Load(cachePath) + promptCache.startSaveLoop(cachePath, 30*time.Second) +``` + +(Adjust the variable name to match the actual tracker variable in handler.go. Ensure `"path/filepath"` and `"time"` are imported.) + +- [ ] **Step 6: Write the disk round-trip test** + +Append to `proxy/cache_tracker_test.go`: + +```go +// TestPromptCacheDiskPersistence verifies C3: entries saved to disk are +// reloaded on startup, surviving a "restart" (new tracker instance). +func TestPromptCacheDiskPersistence(t *testing.T) { + path := t.TempDir() + "/prompt_cache.json" + + // Tracker 1: store an entry, flush to disk. + t1 := newPromptCacheTracker(5 * time.Minute) + hasher := sha256.New() + writeHashChunk(hasher, "test-cache-value-disk") + var fp [32]byte + copy(fp[:], hasher.Sum(nil)) + t1.mu.Lock() + t1.entries[fp] = promptCacheEntry{ + ExpiresAt: time.Now().Add(3 * time.Minute), + TTL: 5 * time.Minute, + } + t1.dirty = true + t1.mu.Unlock() + t1.flush(path) + + // Tracker 2: load from disk → should have the entry. + t2 := newPromptCacheTracker(5 * time.Minute) + t2.Load(path) + t2.mu.Lock() + _, ok := t2.entries[fp] + t2.mu.Unlock() + if !ok { + t.Fatalf("C3: entry not reloaded from disk after 'restart'") + } + + // Expired entry should NOT reload. + path2 := t.TempDir() + "/expired.json" + t1b := newPromptCacheTracker(5 * time.Minute) + t1b.mu.Lock() + var fpExpired [32]byte + copy(fpExpired[:], sha256.Sum256([]byte("expired"))) + t1b.entries[fpExpired] = promptCacheEntry{ + ExpiresAt: time.Now().Add(-1 * time.Minute), // already expired + TTL: 5 * time.Minute, + } + t1b.dirty = true + t1b.mu.Unlock() + t1b.flush(path2) + + t3 := newPromptCacheTracker(5 * time.Minute) + t3.Load(path2) + t3.mu.Lock() + _, okExpired := t3.entries[fpExpired] + t3.mu.Unlock() + if okExpired { + t.Fatalf("C3: expired entry should not be reloaded") + } +} +``` + +Ensure `"crypto/sha256"` is imported. + +- [ ] **Step 7: Run tests** + +Run: `go test ./proxy/ -run 'TestPromptCacheDisk|TestPromptCacheCrossAccount|TestPromptCacheCap' -v` +Expected: all PASS. +Then: `go test ./... ` +Expected: all PASS. + +- [ ] **Step 8: Commit** + +```bash +git add proxy/cache_tracker.go proxy/cache_tracker_test.go proxy/handler.go +git commit -m "feat(cache): persist prompt-cache entries to disk (C3) + +Load data/prompt_cache.json on startup and debounce-write every 30s so +cache entries survive restart. Expired entries are dropped on load. The +file already existed but was unused — now it's read and written." +``` + +--- + +## Final Verification + +- [ ] `go test ./...` passes +- [ ] `go vet ./...` clean +- [ ] `go build -o kiro-go.exe .` succeeds + +## Out of scope (deferred to dispatch plan) + +- Part A (dispatch): health-aware scoring, circuit breaker, session affinity, auto-recovery, weight-as-probability. See `docs/superpowers/specs/2026-06-27-dispatch-cache-improvements-design.md` Part A. Separate plan TBD. diff --git a/docs/superpowers/plans/2026-06-27-credential-json-external-idp-adapter.md b/docs/superpowers/plans/2026-06-27-credential-json-external-idp-adapter.md new file mode 100644 index 00000000..9c1a77d3 --- /dev/null +++ b/docs/superpowers/plans/2026-06-27-credential-json-external-idp-adapter.md @@ -0,0 +1,1036 @@ +# external_idp Credential JSON Adapter — Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Make the "Add credential JSON" import path (`apiImportCredentials` handler + `importCredentials` frontend) accept and persist `external_idp` (Azure AD / Microsoft 365) accounts from any JSON shape without erroring. + +**Architecture:** Extend the backend handler and two frontend functions to recognize `external_idp`, read its refresh material (`tokenEndpoint`/`issuerUrl`/`scopes`), validate the user-supplied IdP endpoint against the existing allow-list (SSRF guard), and route refresh through the already-working `refreshExternalIdpToken`. No changes to the refresh/dispatch logic in `auth/oidc.go`. + +**Tech Stack:** Go 1.21 (module `kiro-go`, stdlib + `github.com/google/uuid v1.6.0`); vanilla JS statically served by the Go server. + +## Global Constraints + +- Canonical AuthMethod value is exactly `external_idp` (snake_case); all aliases normalize to it. +- External IdP endpoints MUST be https, non-IP-literal, and host must `HasSuffix` one of `allowedExternalIdpIssuerSuffixes` (`.microsoftonline.com`, `.microsoftonline.us`, `.microsoftonline.cn`); reject with HTTP 400 otherwise. +- Import MUST succeed at one live token refresh before persisting (regression: `TestApiImportCredentialsRejectsWhenRefreshFails`); refresh failure → HTTP 400, no account persisted. +- External IdP token responses use snake_case JSON (`access_token`/`refresh_token`/`expires_in`); IdC/social use camelCase. +- Provider defaults when empty: external_idp→`AzureAD`, social→`Google`, idc→`BuilderId`. +- Alias `enterprise` maps to `idc` (Kiro Account Manager contract) — do NOT remap to external_idp. +- `external_idp` accounts carry `clientId` but NOT `clientSecret`; detection must precede the clientId+clientSecret→idc default. +- No new third-party Go dependencies. +- Module path is `kiro-go`; run `go test ./...` from the repo root `C:\Users\Admin\Kiro-Go`. + +--- + +## File Structure + +| File | Responsibility | Change | +|---|---|---| +| `auth/kiro_sso.go` | Exported external-IdP endpoint validator | add `externalIdpEndpointValidator` var + `ValidateExternalIdpEndpoint` | +| `auth/testhooks.go` | Test seams | add `SetExternalIdpValidatorForTest` | +| `auth/kiro_sso_test.go` | Validator + seam tests | add 3 tests | +| `proxy/handler.go` | `apiImportCredentials` + authMethod normalization helper | modify handler; add `normalizeImportAuthMethod` + `externalIdpAuthMethodAliases` | +| `proxy/import_credentials_test.go` | import + normalization tests | add happy/SSRF/refresh-fail/identity tests + normalization table test | +| `config/config.go` | account store | add `AccountIDExists` | +| `web/app.js` | current admin UI `importCredentials` | map fields + normalize + payload | +| `web/index-legacy.html` | legacy admin UI `importCredentials` | parity edits | + +--- + +### Task 1: Export `ValidateExternalIdpEndpoint` + add `SetExternalIdpValidatorForTest` seam + +**Files:** +- Modify: `auth/kiro_sso.go` (after the closing brace of `validateExternalIdpEndpoint`, line ~534) +- Modify: `auth/testhooks.go` (after `SetGlobalAuthClientForTest`, line ~30) +- Test: `auth/kiro_sso_test.go` (append; file is `package auth` — call functions with no package prefix) + +**Interfaces:** +- Produces: `auth.ValidateExternalIdpEndpoint(rawURL string) error` (used by Task 3); `auth.SetExternalIdpValidatorForTest(fn func(string) error) func(string) error` (used by Task 3 tests to bypass the allow-list for `http://127.0.0.1` httptest servers). + +- [ ] **Step 1: Write the failing tests** + +Append to `auth/kiro_sso_test.go`: + +```go +// TestValidateExternalIdpEndpointAcceptsAllowListed verifies the exported validator +// accepts real Azure / Microsoft 365 token endpoints (com/global, us-gov, china). +func TestValidateExternalIdpEndpointAcceptsAllowListed(t *testing.T) { + for _, raw := range []string{ + "https://login.microsoftonline.com/tenant/oauth2/v2.0/token", + "https://login.microsoftonline.us/tenant/v2.0", + "https://login.partner.microsoftonline.cn/tenant/oauth2/v2.0/token", + } { + if err := ValidateExternalIdpEndpoint(raw); err != nil { + t.Errorf("expected %q accepted, got %v", raw, err) + } + } +} + +// TestValidateExternalIdpEndpointRejectsUnsafe verifies the validator rejects the +// SSRF shapes a pasted credential JSON could carry: cleartext http, IP literals, +// and non-allow-listed hosts. +func TestValidateExternalIdpEndpointRejectsUnsafe(t *testing.T) { + for _, raw := range []string{ + "http://login.microsoftonline.com/x", // not https + "https://127.0.0.1/oauth/token", // IP literal + "https://evil.example.com/oauth/token", // not allow-listed + } { + if err := ValidateExternalIdpEndpoint(raw); err == nil { + t.Errorf("expected %q rejected, got nil", raw) + } + } +} + +// TestSetExternalIdpValidatorForTestSwapsAndRestores verifies the test seam lets a +// test override (and restore) the validator so happy-path import tests can POST +// against an httptest server (http + 127.0.0.1) that the real allow-list rejects. +func TestSetExternalIdpValidatorForTestSwapsAndRestores(t *testing.T) { + restore := SetExternalIdpValidatorForTest(func(string) error { return nil }) + defer SetExternalIdpValidatorForTest(restore) + if err := ValidateExternalIdpEndpoint("https://evil.example.com/x"); err != nil { + t.Fatalf("expected swapped no-op validator to accept, got %v", err) + } +} +``` + +- [ ] **Step 2: Run tests to verify they fail** + +Run: `go test ./auth/ -run 'TestValidateExternalIdpEndpoint|TestSetExternalIdpValidatorForTest' -v` +Expected: FAIL — `undefined: ValidateExternalIdpEndpoint` and `undefined: SetExternalIdpValidatorForTest` (compile error). + +- [ ] **Step 3: Add the exported validator (kiro_sso.go)** + +In `auth/kiro_sso.go`, immediately after the closing `}` of `validateExternalIdpEndpoint` (after line 534), add: + +```go +// externalIdpEndpointValidator is the function ValidateExternalIdpEndpoint delegates +// to. Tests override it via SetExternalIdpValidatorForTest so a happy-path import +// test can POST against an httptest server (http + 127.0.0.1) that the real +// allow-list would reject. +var externalIdpEndpointValidator = validateExternalIdpEndpoint + +// ValidateExternalIdpEndpoint is the exported entry point for validating a user- or +// discovery-supplied external IdP endpoint URL. The credential-import path +// (package proxy) uses this to guard against SSRF / refresh-token exfiltration: a +// pasted tokenEndpoint pointing at an internal or attacker-controlled host would +// otherwise cause the server to POST the account's refresh token there. +func ValidateExternalIdpEndpoint(rawURL string) error { + return externalIdpEndpointValidator(rawURL) +} +``` + +- [ ] **Step 4: Add the test seam (testhooks.go)** + +In `auth/testhooks.go`, after `SetGlobalAuthClientForTest` (after line 30), add: + +```go +// SetExternalIdpValidatorForTest swaps the validator behind ValidateExternalIdpEndpoint +// and returns the previous one so callers can restore it. Tests POST against httptest +// servers (http + 127.0.0.1), which the real allow-list validator rejects, so tests +// install a no-op validator here. Mirrors SetGlobalAuthClientForTest's swap-and-restore +// shape. Test-only. +func SetExternalIdpValidatorForTest(fn func(string) error) func(string) error { + old := externalIdpEndpointValidator + if fn != nil { + externalIdpEndpointValidator = fn + } + return old +} +``` + +- [ ] **Step 5: Run tests to verify they pass** + +Run: `go test ./auth/ -run 'TestValidateExternalIdpEndpoint|TestSetExternalIdpValidatorForTest' -v` +Expected: PASS (3 tests). If a real-network call is suspected, note these tests hit no network — they are pure string validation. + +- [ ] **Step 6: Commit** + +```bash +git add auth/kiro_sso.go auth/testhooks.go auth/kiro_sso_test.go +git commit -m "feat(auth): export ValidateExternalIdpEndpoint + test seam" +``` + +--- + +### Task 2: `normalizeImportAuthMethod` helper + +**Files:** +- Modify: `proxy/handler.go` (append after `apiImportCredentials`, after line 3109) +- Test: `proxy/import_credentials_test.go` (append) + +**Interfaces:** +- Produces: `normalizeImportAuthMethod(authMethod, clientID, clientSecret, tokenEndpoint string) string` (used by Task 3). +- Consumes: package `strings` (already imported in handler.go). + +- [ ] **Step 1: Write the failing test** + +Append to `proxy/import_credentials_test.go`: + +```go +// TestNormalizeImportAuthMethod pins the auth-method normalization for import, +// including the key regression: external_idp accounts carry clientId but NO +// clientSecret, so the old default branch misclassified them as "social". +func TestNormalizeImportAuthMethod(t *testing.T) { + cases := []struct { + name string + authMethod string + clientID string + clientSecret string + tokenEndpoint string + want string + }{ + {"explicit external_idp", "external_idp", "c", "", "https://login.microsoftonline.com/t/oauth2/v2.0/token", "external_idp"}, + {"azure alias", "AzureAD", "c", "", "https://login.microsoftonline.com/t/oauth2/v2.0/token", "external_idp"}, + {"microsoft alias", "microsoft", "c", "", "https://login.microsoftonline.com/t/oauth2/v2.0/token", "external_idp"}, + {"inferred from tokenEndpoint", "", "c", "", "https://login.microsoftonline.com/t/oauth2/v2.0/token", "external_idp"}, + {"external_idp even with clientSecret", "external_idp", "c", "s", "https://login.microsoftonline.com/t/oauth2/v2.0/token", "external_idp"}, + {"enterprise stays idc", "enterprise", "c", "s", "", "idc"}, + {"idc with clientid+secret", "idc", "c", "s", "", "idc"}, + {"empty + clientid (no secret) -> idc", "", "c", "", "", "idc"}, + {"empty no clientid -> social", "", "", "", "", "social"}, + {"social explicit", "social", "", "", "", "social"}, + {"google alias", "google", "", "", "", "social"}, + {"unrecognized with clientid+secret -> idc", "weird", "c", "s", "", "idc"}, + {"unrecognized without secret -> social", "weird", "c", "", "", "social"}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + if got := normalizeImportAuthMethod(tc.authMethod, tc.clientID, tc.clientSecret, tc.tokenEndpoint); got != tc.want { + t.Fatalf("normalizeImportAuthMethod(%q,%q,%q,%q) = %q, want %q", + tc.authMethod, tc.clientID, tc.clientSecret, tc.tokenEndpoint, got, tc.want) + } + }) + } +} +``` + +- [ ] **Step 2: Run test to verify it fails** + +Run: `go test ./proxy/ -run TestNormalizeImportAuthMethod -v` +Expected: FAIL — `undefined: normalizeImportAuthMethod` (compile error). + +- [ ] **Step 3: Write the helper** + +In `proxy/handler.go`, append after `apiImportCredentials` (after line 3109): + +```go +// externalIdpAuthMethodAliases are lower-cased authMethod values (or Kiro Account +// Manager provider labels) that mean "external IdP / enterprise SSO" and must +// normalize to "external_idp". +var externalIdpAuthMethodAliases = map[string]bool{ + "external_idp": true, + "azuread": true, + "azure": true, + "entra": true, + "entra-id": true, + "entra_id": true, + "microsoft": true, + "m365": true, + "office365": true, + "external": true, +} + +// normalizeImportAuthMethod maps a pasted credential JSON's authMethod (plus its +// clientId/clientSecret/tokenEndpoint) onto one of the three canonical methods +// ("external_idp" | "idc" | "social"). external_idp MUST be detected before the +// clientId+clientSecret→idc inference, because external_idp accounts carry clientId +// but NO clientSecret, so the old default branch misclassified them as "social" and +// refresh hit the wrong endpoint. +// +// It preserves the pre-existing idc/social heuristics: +// - empty authMethod + clientId present -> idc +// - empty authMethod, no clientId -> social +// - "enterprise" (Kiro Account Manager IdC label) -> idc +// - unrecognized non-empty + clientId+clientSecret -> idc, else social +func normalizeImportAuthMethod(authMethod, clientID, clientSecret, tokenEndpoint string) string { + am := strings.ToLower(strings.TrimSpace(authMethod)) + switch { + case externalIdpAuthMethodAliases[am]: + return "external_idp" + case tokenEndpoint != "": // infer when not declared explicitly + return "external_idp" + case am == "social" || am == "google" || am == "github": + return "social" + case am == "idc" || am == "builderid" || am == "enterprise": + return "idc" + } + if am == "" { + if clientID != "" { + return "idc" + } + return "social" + } + if clientID != "" && clientSecret != "" { + return "idc" + } + return "social" +} +``` + +- [ ] **Step 4: Run test to verify it passes** + +Run: `go test ./proxy/ -run TestNormalizeImportAuthMethod -v` +Expected: PASS (13 subtests). + +- [ ] **Step 5: Commit** + +```bash +git add proxy/handler.go proxy/import_credentials_test.go +git commit -m "feat(proxy): add normalizeImportAuthMethod helper" +``` + +--- + +### Task 3: Wire `external_idp` into `apiImportCredentials` + +**Files:** +- Modify: `proxy/handler.go` `apiImportCredentials` (lines 3008–3109) +- Test: `proxy/import_credentials_test.go` (append) + +**Interfaces:** +- Consumes: `normalizeImportAuthMethod` (Task 2), `auth.ValidateExternalIdpEndpoint` + `auth.SetExternalIdpValidatorForTest` (Task 1). +- Produces: a backend that accepts, validates, refreshes, and persists `external_idp` accounts. + +- [ ] **Step 1: Write the failing tests** + +Append to `proxy/import_credentials_test.go`: + +```go +// TestApiImportCredentialsExternalIdpHappyPath verifies an external_idp credential +// imports successfully: authMethod normalizes to external_idp, refresh hits the +// (fake) IdP token endpoint, and the account is persisted with all refresh material. +func TestApiImportCredentialsExternalIdpHappyPath(t *testing.T) { + cfgFile := t.TempDir() + "/config.json" + if err := config.Init(cfgFile); err != nil { + t.Fatalf("config.Init: %v", err) + } + defer installCleanAuthClient(t)() + + const upstreamExpiresIn = 3600 + fake := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if err := r.ParseForm(); err != nil { + http.Error(w, "bad form", http.StatusBadRequest) + return + } + if got := r.PostForm.Get("grant_type"); got != "refresh_token" { + t.Errorf("expected grant_type=refresh_token, got %q", got) + } + w.Header().Set("Content-Type", "application/json") + // external IdP token responses are snake_case. + fmt.Fprintf(w, `{"access_token":"at-ext","refresh_token":"rt-rotated","expires_in":%d}`, upstreamExpiresIn) + })) + defer fake.Close() + + // fake.URL is http + 127.0.0.1; bypass the allow-list validator for this test. + restore := auth.SetExternalIdpValidatorForTest(func(string) error { return nil }) + defer auth.SetExternalIdpValidatorForTest(restore) + + h := &Handler{pool: accountpool.GetPool()} + + body := fmt.Sprintf(`{"authMethod":"external_idp","refreshToken":"rt-ext","clientId":"ext-client","tokenEndpoint":%q,"issuerUrl":"https://login.microsoftonline.com/t/v2.0","scopes":"api://x/codewhisperer:conversations offline_access","region":"eu-central-1"}`, fake.URL) + req := httptest.NewRequest("POST", "/auth/credentials", strings.NewReader(body)) + rec := httptest.NewRecorder() + before := time.Now().Unix() + h.apiImportCredentials(rec, req) + after := time.Now().Unix() + + if rec.Code != http.StatusOK { + t.Fatalf("expected 200, got %d body=%s", rec.Code, rec.Body.String()) + } + + accs := config.GetAccounts() + if len(accs) != 1 { + t.Fatalf("expected 1 account persisted, got %d", len(accs)) + } + got := accs[0] + if got.AuthMethod != "external_idp" { + t.Fatalf("AuthMethod: want external_idp, got %q", got.AuthMethod) + } + if got.AccessToken != "at-ext" { + t.Fatalf("AccessToken: want at-ext, got %q", got.AccessToken) + } + if got.RefreshToken != "rt-rotated" { + t.Fatalf("RefreshToken: want rt-rotated (rotated), got %q", got.RefreshToken) + } + if got.TokenEndpoint != fake.URL { + t.Fatalf("TokenEndpoint not persisted: got %q", got.TokenEndpoint) + } + if got.ClientID != "ext-client" { + t.Fatalf("ClientID not persisted: got %q", got.ClientID) + } + if got.Scopes == "" { + t.Fatalf("Scopes not persisted: got %q", got.Scopes) + } + if got.Provider != "AzureAD" { + t.Fatalf("Provider default: want AzureAD, got %q", got.Provider) + } + if got.Region != "eu-central-1" { + t.Fatalf("Region: want eu-central-1, got %q", got.Region) + } + if got.ExpiresAt < before+upstreamExpiresIn-5 || got.ExpiresAt > after+upstreamExpiresIn+5 { + t.Fatalf("ExpiresAt not from upstream expiresIn: got %d (want ~now+%d)", got.ExpiresAt, upstreamExpiresIn) + } +} + +// TestApiImportCredentialsExternalIdpRejectsNonAllowListedEndpoint verifies the SSRF +// guard: a tokenEndpoint outside the IdP allow-list is rejected with 400 before any +// refresh POST, and nothing is persisted. (Validator is NOT bypassed here.) +func TestApiImportCredentialsExternalIdpRejectsNonAllowListedEndpoint(t *testing.T) { + cfgFile := t.TempDir() + "/config.json" + if err := config.Init(cfgFile); err != nil { + t.Fatalf("config.Init: %v", err) + } + defer installCleanAuthClient(t)() + + h := &Handler{pool: accountpool.GetPool()} + + body := `{"authMethod":"external_idp","refreshToken":"rt","clientId":"c","tokenEndpoint":"https://evil.example.com/oauth/token","region":"us-east-1"}` + req := httptest.NewRequest("POST", "/auth/credentials", strings.NewReader(body)) + rec := httptest.NewRecorder() + h.apiImportCredentials(rec, req) + + if rec.Code != http.StatusBadRequest { + t.Fatalf("expected 400, got %d body=%s", rec.Code, rec.Body.String()) + } + var resp map[string]string + if err := json.Unmarshal(rec.Body.Bytes(), &resp); err != nil { + t.Fatalf("decode response: %v", err) + } + if !strings.Contains(resp["error"], "endpoint rejected") { + t.Fatalf("expected endpoint-rejected error, got %q", resp["error"]) + } + if accs := config.GetAccounts(); len(accs) != 0 { + t.Fatalf("expected no account persisted, got %d", len(accs)) + } +} + +// TestApiImportCredentialsExternalIdpRejectsWhenRefreshFails verifies the refresh +// gate holds for external_idp: a refresh that 400s (invalid_grant) must reject the +// import and persist nothing. +func TestApiImportCredentialsExternalIdpRejectsWhenRefreshFails(t *testing.T) { + cfgFile := t.TempDir() + "/config.json" + if err := config.Init(cfgFile); err != nil { + t.Fatalf("config.Init: %v", err) + } + defer installCleanAuthClient(t)() + + fake := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + http.Error(w, `{"error":"invalid_grant"}`, http.StatusBadRequest) + })) + defer fake.Close() + + restore := auth.SetExternalIdpValidatorForTest(func(string) error { return nil }) + defer auth.SetExternalIdpValidatorForTest(restore) + + h := &Handler{pool: accountpool.GetPool()} + + body := fmt.Sprintf(`{"authMethod":"external_idp","refreshToken":"rt-broken","clientId":"c","tokenEndpoint":%q,"region":"us-east-1"}`, fake.URL) + req := httptest.NewRequest("POST", "/auth/credentials", strings.NewReader(body)) + rec := httptest.NewRecorder() + h.apiImportCredentials(rec, req) + + if rec.Code != http.StatusBadRequest { + t.Fatalf("expected 400, got %d body=%s", rec.Code, rec.Body.String()) + } + var resp map[string]string + if err := json.Unmarshal(rec.Body.Bytes(), &resp); err != nil { + t.Fatalf("decode response: %v", err) + } + if !strings.Contains(resp["error"], "Token refresh failed") { + t.Fatalf("expected refresh-failed error, got %q", resp["error"]) + } + if accs := config.GetAccounts(); len(accs) != 0 { + t.Fatalf("expected no account persisted, got %d", len(accs)) + } +} +``` + +- [ ] **Step 2: Run tests to verify they fail** + +Run: `go test ./proxy/ -run 'TestApiImportCredentialsExternalIdp' -v` +Expected: `TestApiImportCredentialsExternalIdpHappyPath` FAIL (current code normalizes external_idp→social→refresh hits the real AWS social endpoint, not `fake.URL`, so returns 400). The reject-SSRF test FAIL too (no validation yet → it would attempt refresh / return a different error). + +- [ ] **Step 3: Modify the request struct (handler.go 3009–3017)** + +Replace: +```go + var req struct { + AccessToken string `json:"accessToken"` + RefreshToken string `json:"refreshToken"` + ClientID string `json:"clientId"` + ClientSecret string `json:"clientSecret"` + AuthMethod string `json:"authMethod"` + Provider string `json:"provider"` + Region string `json:"region"` + } +``` +With: +```go + var req struct { + AccessToken string `json:"accessToken"` + RefreshToken string `json:"refreshToken"` + ClientID string `json:"clientId"` + ClientSecret string `json:"clientSecret"` + AuthMethod string `json:"authMethod"` + Provider string `json:"provider"` + Region string `json:"region"` + // external_idp (enterprise SSO / Azure AD) refresh material. + TokenEndpoint string `json:"tokenEndpoint"` + IssuerURL string `json:"issuerUrl"` + Scopes string `json:"scopes"` + // Optional identity preservation when pasting a full account record. + ID string `json:"id"` + Email string `json:"email"` + ProfileArn string `json:"profileArn"` + } +``` + +- [ ] **Step 4: Replace the authMethod normalization block (handler.go 3034–3053)** + +Replace (the `if req.AuthMethod == "" { ... }` block AND the `switch strings.ToLower(req.AuthMethod) { ... }` block): +```go + if req.AuthMethod == "" { + if req.ClientID != "" { + req.AuthMethod = "idc" + } else { + req.AuthMethod = "social" + } + } + // 标准化 authMethod + switch strings.ToLower(req.AuthMethod) { + case "idc", "builderid", "enterprise": + req.AuthMethod = "idc" + case "social", "google", "github": + req.AuthMethod = "social" + default: + if req.ClientID != "" && req.ClientSecret != "" { + req.AuthMethod = "idc" + } else { + req.AuthMethod = "social" + } + } +``` +With: +```go + // 标准化 authMethod。external_idp 必须先于 clientId+clientSecret→idc 的推断被识别 + //（external_idp 带 clientId 但没有 clientSecret），否则会被误判成 social 而 refresh 到错误端点。 + req.AuthMethod = normalizeImportAuthMethod(req.AuthMethod, req.ClientID, req.ClientSecret, req.TokenEndpoint) + + // external_idp 的 tokenEndpoint 是用户可填的新信任边界：必须经 allow-list 校验， + // 否则一份不信任的 credential JSON 可指向内网/攻击者主机，导致 refresh token 被外泄。 + if req.AuthMethod == "external_idp" { + if req.ClientID == "" || req.TokenEndpoint == "" { + w.WriteHeader(400) + json.NewEncoder(w).Encode(map[string]string{"error": "external_idp requires clientId and tokenEndpoint"}) + return + } + if err := auth.ValidateExternalIdpEndpoint(req.TokenEndpoint); err != nil { + w.WriteHeader(400) + json.NewEncoder(w).Encode(map[string]string{"error": "external IdP endpoint rejected: " + err.Error()}) + return + } + if req.IssuerURL != "" { + if err := auth.ValidateExternalIdpEndpoint(req.IssuerURL); err != nil { + w.WriteHeader(400) + json.NewEncoder(w).Encode(map[string]string{"error": "external IdP issuer rejected: " + err.Error()}) + return + } + } + } +``` + +- [ ] **Step 5: Carry refresh material in tempAccount (handler.go 3058–3064)** + +Replace: +```go + tempAccount := &config.Account{ + RefreshToken: req.RefreshToken, + ClientID: req.ClientID, + ClientSecret: req.ClientSecret, + AuthMethod: req.AuthMethod, + Region: req.Region, + } +``` +With: +```go + tempAccount := &config.Account{ + RefreshToken: req.RefreshToken, + ClientID: req.ClientID, + ClientSecret: req.ClientSecret, + AuthMethod: req.AuthMethod, + Region: req.Region, + TokenEndpoint: req.TokenEndpoint, + Scopes: req.Scopes, + } +``` + +- [ ] **Step 6: Persist external_idp fields + provider default in created account (handler.go 3079–3093)** + +Replace: +```go + // 创建账号 + account := config.Account{ + ID: auth.GenerateAccountID(), + Email: email, + AccessToken: accessToken, + RefreshToken: req.RefreshToken, + ClientID: req.ClientID, + ClientSecret: req.ClientSecret, + AuthMethod: req.AuthMethod, + Provider: req.Provider, + Region: req.Region, + ExpiresAt: expiresAt, + Enabled: true, + MachineId: config.GenerateMachineId(), + ProfileArn: newProfileArn, + } +``` +With: +```go + // 创建账号 + provider := req.Provider + if provider == "" && req.AuthMethod == "external_idp" { + provider = "AzureAD" + } + account := config.Account{ + ID: auth.GenerateAccountID(), + Email: email, + AccessToken: accessToken, + RefreshToken: req.RefreshToken, + ClientID: req.ClientID, + ClientSecret: req.ClientSecret, + AuthMethod: req.AuthMethod, + Provider: provider, + Region: req.Region, + ExpiresAt: expiresAt, + Enabled: true, + MachineId: config.GenerateMachineId(), + ProfileArn: newProfileArn, + TokenEndpoint: req.TokenEndpoint, + IssuerURL: req.IssuerURL, + Scopes: req.Scopes, + } +``` + +- [ ] **Step 7: Run tests to verify they pass** + +Run: `go test ./proxy/ -run 'TestApiImportCredentialsExternalIdp' -v` +Expected: PASS (3 tests). Also run the full handler suite to confirm no regressions: +Run: `go test ./proxy/ ./auth/ ./config/ ./pool/ -v` +Expected: all PASS (including pre-existing `TestApiImportCredentialsRejectsWhenRefreshFails` and `TestApiImportCredentialsUsesUpstreamExpiresAt`). + +- [ ] **Step 8: Commit** + +```bash +git add proxy/handler.go proxy/import_credentials_test.go +git commit -m "feat(auth): import external_idp credentials via Add credential JSON" +``` + +--- + +### Task 4: Identity preservation for full-record paste + +**Files:** +- Modify: `config/config.go` (after `GetAccounts`, line ~414) +- Modify: `proxy/handler.go` `apiImportCredentials` (email line ~3076, created-account block) +- Test: `proxy/import_credentials_test.go` (append) + +**Interfaces:** +- Consumes: `req.ID`, `req.Email`, `req.ProfileArn` (added to struct in Task 3). +- Produces: `config.AccountIDExists(id string) bool` (used by the handler to avoid duplicate IDs when re-importing a backup). + +- [ ] **Step 1: Write the failing test** + +Append to `proxy/import_credentials_test.go`: + +```go +// TestApiImportCredentialsExternalIdpPreservesFullRecordIdentity verifies that when +// a full account record (with id/email/profileArn) is pasted, those are preserved +// rather than regenerated, so re-importing a backup does not duplicate accounts. +func TestApiImportCredentialsExternalIdpPreservesFullRecordIdentity(t *testing.T) { + cfgFile := t.TempDir() + "/config.json" + if err := config.Init(cfgFile); err != nil { + t.Fatalf("config.Init: %v", err) + } + defer installCleanAuthClient(t)() + + fake := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + fmt.Fprint(w, `{"access_token":"at-ext","refresh_token":"rt-rotated","expires_in":3600}`) + })) + defer fake.Close() + + restore := auth.SetExternalIdpValidatorForTest(func(string) error { return nil }) + defer auth.SetExternalIdpValidatorForTest(restore) + + h := &Handler{pool: accountpool.GetPool()} + + const providedID = "11111111-2222-3333-4444-555555555555" + body := fmt.Sprintf(`{"id":%q,"email":"ada@example.com","profileArn":"arn:aws:codewhisperer:eu-central-1:1:profile/PRESERVED","authMethod":"external_idp","refreshToken":"rt","clientId":"c","tokenEndpoint":%q,"region":"eu-central-1"}`, providedID, fake.URL) + req := httptest.NewRequest("POST", "/auth/credentials", strings.NewReader(body)) + rec := httptest.NewRecorder() + h.apiImportCredentials(rec, req) + + if rec.Code != http.StatusOK { + t.Fatalf("expected 200, got %d body=%s", rec.Code, rec.Body.String()) + } + got := config.GetAccounts()[0] + if got.ID != providedID { + t.Fatalf("ID: want reused %q, got %q", providedID, got.ID) + } + if got.Email != "ada@example.com" { + t.Fatalf("Email: want ada@example.com (GetUserInfo empty in test → fallback), got %q", got.Email) + } + if got.ProfileArn != "arn:aws:codewhisperer:eu-central-1:1:profile/PRESERVED" { + t.Fatalf("ProfileArn: want preserved, got %q", got.ProfileArn) + } +} +``` + +- [ ] **Step 2: Run test to verify it fails** + +Run: `go test ./proxy/ -run TestApiImportCredentialsExternalIdpPreservesFullRecordIdentity -v` +Expected: FAIL — the ID is a fresh `GenerateAccountID()` (not `providedID`), email is whatever `GetUserInfo` returned (likely empty, not `ada@example.com`), and ProfileArn is `""` (external_idp refresh returns no profileArn), not the preserved value. + +- [ ] **Step 3: Add `config.AccountIDExists` (config.go)** + +In `config/config.go`, immediately after `GetAccounts` (after line 414, before `GetEnabledAccounts`), add: + +```go +// AccountIDExists reports whether an account with the given ID is already stored. +// Used by the credential-import path to reuse a pasted record's id when it does +// not collide, so re-importing a backup never creates a duplicate entry. +func AccountIDExists(id string) bool { + if id == "" { + return false + } + cfgLock.RLock() + defer cfgLock.RUnlock() + for _, a := range cfg.Accounts { + if a.ID == id { + return true + } + } + return false +} +``` + +- [ ] **Step 4: Add email fallback (handler.go, the `GetUserInfo` line ~3076)** + +Replace: +```go + // 获取用户信息 + email, _, _ := auth.GetUserInfo(accessToken) +``` +With: +```go + // 获取用户信息 + email, _, _ := auth.GetUserInfo(accessToken) + if email == "" { + email = req.Email // fall back to a pasted full record's email + } +``` + +- [ ] **Step 5: Add id reuse + profileArn fallback (handler.go created-account block)** + +In Task 3's created-account block, replace: +```go + // 创建账号 + provider := req.Provider + if provider == "" && req.AuthMethod == "external_idp" { + provider = "AzureAD" + } + account := config.Account{ + ID: auth.GenerateAccountID(), + Email: email, +``` +With: +```go + // 创建账号 + provider := req.Provider + if provider == "" && req.AuthMethod == "external_idp" { + provider = "AzureAD" + } + // Reuse a pasted record's id when it does not collide; otherwise mint a fresh one + // so re-importing a backup never creates a duplicate entry. + id := req.ID + if id == "" || config.AccountIDExists(id) { + id = auth.GenerateAccountID() + } + profileArn := newProfileArn + if profileArn == "" { + profileArn = req.ProfileArn // external_idp refresh returns no profileArn + } + account := config.Account{ + ID: id, + Email: email, +``` + +And in the same struct literal, replace `ProfileArn: newProfileArn,` with `ProfileArn: profileArn,`: +```go + ProfileArn: profileArn, +``` + +- [ ] **Step 6: Run test to verify it passes** + +Run: `go test ./proxy/ -run TestApiImportCredentialsExternalIdpPreservesFullRecordIdentity -v` +Expected: PASS. +Then the full suite: `go test ./... ` (from repo root) +Expected: all PASS. + +- [ ] **Step 7: Commit** + +```bash +git add config/config.go proxy/handler.go proxy/import_credentials_test.go +git commit -m "feat(auth): preserve id/email/profileArn when importing full records" +``` + +--- + +### Task 5: Frontend `importCredentials` (current UI, `web/app.js`) + +**Files:** +- Modify: `web/app.js` `importCredentials` (lines 2289–2299, 2319–2333) +- Verification: manual (no JS test harness in this repo) + +**Interfaces:** +- Consumes: backend `/auth/credentials` (now accepts `tokenEndpoint`/`issuerUrl`/`scopes`/`id`/`email`/`profileArn` and recognizes `external_idp`). +- Produces: a UI that pastes all three JSON shapes (single object, array/Kiro export, full record) into a correct payload. + +- [ ] **Step 1: Extend the `json.accounts` map (app.js 2289–2299)** + +Replace: +```js + items = json.accounts.map(a => { + const c = a.credentials || {}; + return { + refreshToken: c.refreshToken || a.refreshToken, + clientId: c.clientId || a.clientId, + clientSecret: c.clientSecret || a.clientSecret, + region: c.region || a.region, + authMethod: c.authMethod || a.authMethod, + provider: c.provider || a.provider || a.idp + }; + }); +``` +With: +```js + items = json.accounts.map(a => { + const c = a.credentials || {}; + return { + refreshToken: c.refreshToken || a.refreshToken, + clientId: c.clientId || a.clientId, + clientSecret: c.clientSecret || a.clientSecret, + region: c.region || a.region, + authMethod: c.authMethod || a.authMethod, + provider: c.provider || a.provider || a.idp, + tokenEndpoint: c.tokenEndpoint || a.tokenEndpoint, + issuerUrl: c.issuerUrl || a.issuerUrl, + scopes: c.scopes || a.scopes, + id: a.id, + email: c.email || a.email, + profileArn: c.profileArn || a.profileArn + }; + }); +``` + +- [ ] **Step 2: Recognize `external_idp` + extend payload (app.js 2319–2333)** + +Replace: +```js + let authMethod = item.authMethod || ''; + if (item.clientId && item.clientSecret) authMethod = 'idc'; + else if (!authMethod || authMethod === 'social') authMethod = 'social'; + else authMethod = authMethod.toLowerCase() === 'idc' ? 'idc' : 'social'; + let provider = item.provider || ''; + if (!provider && authMethod === 'social') provider = 'Google'; + if (!provider && authMethod === 'idc') provider = 'BuilderId'; + const payload = { + refreshToken: item.refreshToken, + accessToken: item.accessToken || '', + clientId: item.clientId || '', + clientSecret: item.clientSecret || '', + authMethod, provider, + region: item.region || 'us-east-1' + }; +``` +With: +```js + const EXTERNAL_IDP = ['external_idp','azuread','azure','entra','entra-id','microsoft','m365','office365','external']; + let authMethod = (item.authMethod || '').toLowerCase(); + if (EXTERNAL_IDP.includes(authMethod) || item.tokenEndpoint) { + authMethod = 'external_idp'; + } else if (item.clientId && item.clientSecret) { + authMethod = 'idc'; + } else if (!authMethod || authMethod === 'social') { + authMethod = 'social'; + } else { + authMethod = authMethod === 'idc' ? 'idc' : 'social'; + } + let provider = item.provider || ''; + if (!provider && authMethod === 'external_idp') provider = 'AzureAD'; + if (!provider && authMethod === 'social') provider = 'Google'; + if (!provider && authMethod === 'idc') provider = 'BuilderId'; + const payload = { + refreshToken: item.refreshToken, + accessToken: item.accessToken || '', + clientId: item.clientId || '', + clientSecret: item.clientSecret || '', + authMethod, provider, + region: item.region || 'us-east-1', + tokenEndpoint: item.tokenEndpoint || '', + issuerUrl: item.issuerUrl || '', + scopes: item.scopes || '', + ...(item.id ? { id: item.id } : {}), + ...(item.email ? { email: item.email } : {}), + ...(item.profileArn ? { profileArn: item.profileArn } : {}) + }; +``` + +- [ ] **Step 3: Verify (manual)** + +Rebuild & restart the server (so any embedded assets pick up the change): +Run: `go build -o kiro-go.exe .` +Launch: `./kiro-go.exe` (in a separate terminal) +Then in the admin panel (log in), go to **Accounts → Add → "Credential JSON"**: + +a. **Single object:** paste +```json +{"authMethod":"external_idp","refreshToken":"","clientId":"","tokenEndpoint":"https://login.microsoftonline.com//oauth2/v2.0/token","issuerUrl":"https://login.microsoftonline.com//v2.0","scopes":"","region":"eu-central-1"} +``` +Expected: toast "Imported: 1", account appears with provider "AzureAD", region eu-central-1, the real email. + +b. **Full record:** paste a record copied from `data/config.json` (the `external_idp` object with `id`/`email`/`profileArn`). +Expected: account imported with the SAME `id` (no duplicate) and the preserved `profileArn`/`email`. + +c. **Kiro export:** paste `{"version":2,"accounts":[{"credentials":{...external_idp fields...}}]}`. +Expected: "Imported: 1". + +d. **Negative (SSRF):** paste the single object but change `tokenEndpoint` to `https://evil.example.com/x`. +Expected: "Failed: 1" in the toast (exact reason visible in the browser devtools Network tab on the `/auth/credentials` response: `external IdP endpoint rejected: ...`). + +If a step fails, re-read the relevant diff above; do not edit by guess. + +- [ ] **Step 4: Commit** + +```bash +git add web/app.js +git commit -m "feat(ui): import external_idp credentials from Add credential JSON" +``` + +--- + +### Task 6: Frontend parity (`web/index-legacy.html`) + +**Files:** +- Modify: `web/index-legacy.html` `importCredentials` (lines 2850–2861, 2869–2881) +- Verification: manual + +**Interfaces:** same as Task 5 — keep the two UIs in sync so the legacy page also imports `external_idp`. + +- [ ] **Step 1: Extend the `json.accounts` map (index-legacy.html 2850–2861)** + +Replace: +```js + items = json.accounts.map(a => { + const c = a.credentials || {}; + return { + refreshToken: c.refreshToken || a.refreshToken, + clientId: c.clientId || a.clientId, + clientSecret: c.clientSecret || a.clientSecret, + region: c.region || a.region, + // 不传 accessToken，强制后端用 refreshToken 刷新获取新 token + authMethod: c.authMethod || a.authMethod, + provider: c.provider || a.provider || a.idp + }; + }); +``` +With: +```js + items = json.accounts.map(a => { + const c = a.credentials || {}; + return { + refreshToken: c.refreshToken || a.refreshToken, + clientId: c.clientId || a.clientId, + clientSecret: c.clientSecret || a.clientSecret, + region: c.region || a.region, + // 不传 accessToken，强制后端用 refreshToken 刷新获取新 token + authMethod: c.authMethod || a.authMethod, + provider: c.provider || a.provider || a.idp, + tokenEndpoint: c.tokenEndpoint || a.tokenEndpoint, + issuerUrl: c.issuerUrl || a.issuerUrl, + scopes: c.scopes || a.scopes, + id: a.id, + email: c.email || a.email, + profileArn: c.profileArn || a.profileArn + }; + }); +``` + +- [ ] **Step 2: Recognize `external_idp` + extend payload (index-legacy.html 2869–2881)** + +Replace: +```js + // 映射 authMethod: IdC/idc -> idc, social -> social + let authMethod = item.authMethod || ''; + if (item.clientId && item.clientSecret) { + authMethod = 'idc'; + } else if (!authMethod || authMethod === 'social') { + authMethod = 'social'; + } else { + authMethod = authMethod.toLowerCase() === 'idc' ? 'idc' : 'social'; + } + // 映射 provider + let provider = item.provider || ''; + if (!provider && authMethod === 'social') provider = 'Google'; + if (!provider && authMethod === 'idc') provider = 'BuilderId'; + const payload = { refreshToken: item.refreshToken, accessToken: item.accessToken || '', clientId: item.clientId || '', clientSecret: item.clientSecret || '', authMethod: authMethod, provider: provider, region: item.region || 'us-east-1' }; +``` +With: +```js + // 映射 authMethod: external_idp -> external_idp, IdC/idc -> idc, social -> social + const EXTERNAL_IDP = ['external_idp','azuread','azure','entra','entra-id','microsoft','m365','office365','external']; + let authMethod = (item.authMethod || '').toLowerCase(); + if (EXTERNAL_IDP.includes(authMethod) || item.tokenEndpoint) { + authMethod = 'external_idp'; + } else if (item.clientId && item.clientSecret) { + authMethod = 'idc'; + } else if (!authMethod || authMethod === 'social') { + authMethod = 'social'; + } else { + authMethod = authMethod.toLowerCase() === 'idc' ? 'idc' : 'social'; + } + // 映射 provider + let provider = item.provider || ''; + if (!provider && authMethod === 'external_idp') provider = 'AzureAD'; + if (!provider && authMethod === 'social') provider = 'Google'; + if (!provider && authMethod === 'idc') provider = 'BuilderId'; + const payload = { refreshToken: item.refreshToken, accessToken: item.accessToken || '', clientId: item.clientId || '', clientSecret: item.clientSecret || '', authMethod: authMethod, provider: provider, region: item.region || 'us-east-1', tokenEndpoint: item.tokenEndpoint || '', issuerUrl: item.issuerUrl || '', scopes: item.scopes || '', ...(item.id ? { id: item.id } : {}), ...(item.email ? { email: item.email } : {}), ...(item.profileArn ? { profileArn: item.profileArn } : {}) }; +``` + +- [ ] **Step 3: Verify (manual)** + +Rebuild (`go build -o kiro-go.exe .`) and restart. Open the legacy page and repeat the four paste scenarios from Task 5 Step 3. Expected outcomes are identical (alert shows "Imported: 1" / "Failed: 1"). + +- [ ] **Step 4: Commit** + +```bash +git add web/index-legacy.html +git commit -m "feat(ui): legacy parity for external_idp credential import" +``` + +--- + +## Final Verification + +- [ ] Full test suite passes: `go test ./...` (from `C:\Users\Admin\Kiro-Go`) +- [ ] `go vet ./...` is clean +- [ ] Server builds: `go build -o kiro-go.exe .` +- [ ] Manual end-to-end: paste each of the three JSON shapes (single object, full record, Kiro export) for a real `external_idp` account → all import successfully and the account refreshes + serves a request; pasting an off-allow-list `tokenEndpoint` is rejected. + +## Out of scope (explicitly deferred) + +- No JWT `exp` trust-on-import (refresh-on-import remains mandatory). +- No changes to `parseLineCredentials` (line format does not carry `tokenEndpoint`/`scopes`). +- No new external IdPs beyond the Microsoft allow-list; add new suffixes to `allowedExternalIdpIssuerSuffixes` separately if needed. +- No upsert-by-email; duplicate IDs are avoided via `config.AccountIDExists`, but re-importing a record whose `email` already exists under a different `id` creates a second entry (acceptable; out of scope). diff --git a/docs/superpowers/plans/2026-07-02-cache-capacity-lru-metrics.md b/docs/superpowers/plans/2026-07-02-cache-capacity-lru-metrics.md new file mode 100644 index 00000000..fe3bac3a --- /dev/null +++ b/docs/superpowers/plans/2026-07-02-cache-capacity-lru-metrics.md @@ -0,0 +1,945 @@ +# Cache Capacity, O(1) LRU, and Metrics Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Stop prefix-churn under load (bump the prompt-cache LRU from 4096 → 131072, configurable), replace O(n log n) sort-under-lock eviction with O(1) `container/list`-based LRU, and add hit/miss/eviction/expiration counters exposed via `/v1/stats`. + +**Architecture:** The prompt-cache tracker (`proxy/cache_tracker.go`) is an *accounting-only* fingerprint store. Its entries map is converted from value-typed to pointer-typed and paired with a doubly-linked list (`container/list`) that records recency; eviction pops the list back in O(1). Capacity becomes configurable via `config.PromptCacheMaxEntries` (default 131072, clamp ≥ 256). Four atomic counters are added to the tracker and surfaced through a new `Stats()` method folded into the existing `/v1/stats` response. + +**Tech Stack:** Go 1.x, standard library (`container/list`, `sync`, `sync/atomic`), package `proxy` (tracker + handler), package `config`. + +## Global Constraints + +- **Module:** `kiro-go`. Tracker + handler live in package `proxy`; config in package `config`. +- **TDD:** every task is RED → GREEN (write the failing test first, run it, implement, run it green, commit). +- **Disk format:** version 1 unchanged — `flush`/`Load` serialize the same `promptCacheEntryOnDisk{Fingerprint,ExpiresAt,TTLSeconds}`. The list is NOT persisted. +- **Invariants preserved (do not regress):** + - Global cross-account sharing — the `accountID` parameter on `Compute`/`Update` stays ignored (comment `entries is the global map now (C1: cross-account sharing)` stays). + - `Update` runs only on the request-success path (`handler.go:1270`). + - `dirty`-on-hit + idempotent `Stop` (commit `c35e792`). + - `clampCacheBreakdownToCreation` + `splitAgainstTotal` (commits `c01dfb9`, `501443b`) unchanged. + - Single mutex; ALL list/map mutation under `t.mu`. +- **Capacity:** default 131072, configurable floor 256, min-cacheable threshold 1024 (4096 opus) unchanged. +- **Commits:** frequent, conventional-commit messages, one per task. + +## Spec deviation (note for the implementer) + +The spec (`docs/superpowers/specs/2026-07-02-cache-capacity-lru-metrics-design.md`) describes a single constructor `newPromptCacheTracker(maxTTL, maxEntries)`. The plan refines this to **two constructors** to avoid churning ~15 call sites: keep `newPromptCacheTracker(maxTTL)` delegating to `config.GetPromptCacheMaxEntries()`, and add `newPromptCacheTrackerWithCapacity(maxTTL, maxEntries)` for tests. Behavior is identical; only the call-site footprint differs. + +## File Structure + +- `config/config.go` — add `PromptCacheMaxEntries` field + `GetPromptCacheMaxEntries`/`UpdatePromptCacheMaxEntries` + `defaultPromptCacheMaxEntries` const (Task 1). +- `proxy/cache_tracker.go` — pointer entries + `container/list` LRU + `putLocked`/`evictOverflowLocked` + two constructors + drop `LastHit`/`maxPromptCacheEntries`/`evictLRULocked` (Task 2); atomic counters + `PromptCacheStats`/`Stats()` (Task 4). +- `proxy/cache_tracker_test.go` — config test, churn regression tests, metrics test, rewrites of the 3 direct-insert tests. +- `proxy/cache_tracker_hardening_test.go` — rewrite the LRU eviction test + the dirty-on-hit test. +- `proxy/handler.go` — wire `cache` into `/v1/stats` (Task 5). +- `proxy/handler_test.go` — handler-level stats test (Task 5). + +--- + +### Task 1: Config — `PromptCacheMaxEntries` (field + getter + setter + default) + +**Files:** +- Modify: `config/config.go` (field after `:226`; const + getter + setter after `:949`) +- Test: `proxy/cache_tracker_test.go` (append a black-box config test, same pattern as `TestPromptCacheCapConfigurable`) + +**Interfaces:** +- Produces: `config.GetPromptCacheMaxEntries() int` (default 131072, ≤0 → default), `config.UpdatePromptCacheMaxEntries(n int) error`, `config.defaultPromptCacheMaxEntries` const. Consumed by Task 2's delegating constructor. + +- [ ] **Step 1: Write the failing test** + +Append to `proxy/cache_tracker_test.go`: + +```go +// TestPromptCacheMaxEntriesConfigurable verifies the cache LRU bound is +// configurable via config and defaults to 131072. +func TestPromptCacheMaxEntriesConfigurable(t *testing.T) { + cfgFile := t.TempDir() + "/config.json" + if err := config.Init(cfgFile); err != nil { + t.Fatalf("config.Init: %v", err) + } + + if got := config.GetPromptCacheMaxEntries(); got != 131072 { + t.Fatalf("default cap: expected 131072, got %d", got) + } + + if err := config.UpdatePromptCacheMaxEntries(50000); err != nil { + t.Fatalf("update: %v", err) + } + if got := config.GetPromptCacheMaxEntries(); got != 50000 { + t.Fatalf("after update: expected 50000, got %d", got) + } + + // ≤ 0 falls back to the default. + if err := config.UpdatePromptCacheMaxEntries(0); err != nil { + t.Fatalf("reset: %v", err) + } + if got := config.GetPromptCacheMaxEntries(); got != 131072 { + t.Fatalf("zero should fall back to default 131072, got %d", got) + } +} +``` + +- [ ] **Step 2: Run test to verify it fails** + +Run: `go test ./proxy/ -run TestPromptCacheMaxEntriesConfigurable -v` +Expected: FAIL — `config.GetPromptCacheMaxEntries undefined`. + +- [ ] **Step 3: Add the config field** + +In `config/config.go`, immediately after the `PromptCacheMaxRatio` field (currently `:226`), add: + +```go + // PromptCacheMaxEntries bounds the in-memory prompt-cache map; once exceeded, + // the least-recently-used entries are evicted (LRU). Default 131072. Sized so + // the prefix write-rate × TTL does not evict multi-turn history prefixes + // before the next turn reuses them (mirrors kiro-rs's 131072 default). The + // tracker clamps explicit small values up to 256. + PromptCacheMaxEntries int `json:"promptCacheMaxEntries,omitempty"` +``` + +- [ ] **Step 4: Add the const + getter + setter** + +In `config/config.go`, immediately after `UpdatePromptCacheMaxRatio` (currently ends at `:949`), add: + +```go +const defaultPromptCacheMaxEntries = 131072 + +// GetPromptCacheMaxEntries returns the prompt-cache LRU bound. Defaults to +// 131072 when unset (≤ 0). Explicit small values are clamped up to 256 by the +// tracker constructor. +func GetPromptCacheMaxEntries() int { + cfgLock.RLock() + defer cfgLock.RUnlock() + if cfg == nil || cfg.PromptCacheMaxEntries <= 0 { + return defaultPromptCacheMaxEntries + } + return cfg.PromptCacheMaxEntries +} + +// UpdatePromptCacheMaxEntries sets the prompt-cache LRU bound and persists it. +// Applies on the next tracker construction (restart); it does not resize a +// live tracker. +func UpdatePromptCacheMaxEntries(n int) error { + cfgLock.Lock() + defer cfgLock.Unlock() + cfg.PromptCacheMaxEntries = n + return Save() +} +``` + +- [ ] **Step 5: Run test to verify it passes** + +Run: `go test ./proxy/ -run TestPromptCacheMaxEntriesConfigurable -v` +Expected: PASS. + +- [ ] **Step 6: Build the whole module to confirm no breakage** + +Run: `go build ./...` +Expected: succeeds. + +- [ ] **Step 7: Commit** + +```bash +git add config/config.go proxy/cache_tracker_test.go +git commit -m "feat(config): add PromptCacheMaxEntries (default 131072)" +``` + +--- + +### Task 2: O(1) LRU via `container/list` + capacity field + two constructors + +This is the structural change. It converts `entries` to pointer-typed, adds a recency list, drops `LastHit` (list position is authoritative), replaces `evictLRULocked` (sort) with `evictOverflowLocked` (O(1) pop-back), adds `putLocked`, and adds `maxEntries` + the two constructors. Five existing tests are rewritten to compile against the new shape. + +**Files:** +- Modify: `proxy/cache_tracker.go` (imports, consts, entry type, struct, constructors, `Load`, `Compute`, `Update`, `pruneExpiredLocked`, replace `evictLRULocked`) +- Modify: `proxy/cache_tracker_hardening_test.go` (rewrite `TestPromptCacheEvictsLRUWhenOverCapacity`, rewrite `TestComputeSetsDirtyOnCacheHit`) +- Modify: `proxy/cache_tracker_test.go` (rewrite the 3 direct-insert sites in `TestPromptCacheDiskPersistence` and `TestComputeBreakdownClampedToCreation`) +- Test: `proxy/cache_tracker_hardening_test.go` (new `TestPromptCacheLRUEvictsOldestUnused` + `TestPromptCacheTrackerClampsSmallCapacity`) + +**Interfaces:** +- Consumes: `config.GetPromptCacheMaxEntries()` (Task 1). +- Produces: `newPromptCacheTrackerWithCapacity(maxTTL, maxEntries)`, `(*promptCacheTracker).putLocked(fp, expiresAt, ttl)`, `(*promptCacheTracker).evictOverflowLocked()`, `promptCacheTracker.maxEntries`, pointer-typed `promptCacheEntry` (with `lruElem *list.Element`). Consumed by Tasks 3–5. + +- [ ] **Step 1: Write the failing tests** + +Replace the entire `TestPromptCacheEvictsLRUWhenOverCapacity` function (currently `proxy/cache_tracker_hardening_test.go:8-45`) with: + +```go +// TestPromptCacheLRUEvictsOldestUnused verifies the list-based LRU: with cap 3, +// after inserting a,b,c, touching a, then inserting d, the evicted entry is b +// (the least-recently-used), not a. +func TestPromptCacheLRUEvictsOldestUnused(t *testing.T) { + tr := newPromptCacheTrackerWithCapacity(time.Hour, 3) + now := time.Now() + + tr.mu.Lock() + tr.putLocked([32]byte{1}, now.Add(time.Hour), time.Hour) // a + tr.putLocked([32]byte{2}, now.Add(time.Hour), time.Hour) // b + tr.putLocked([32]byte{3}, now.Add(time.Hour), time.Hour) // c + tr.putLocked([32]byte{1}, now.Add(time.Hour), time.Hour) // touch a → front + tr.putLocked([32]byte{4}, now.Add(time.Hour), time.Hour) // d + tr.evictOverflowLocked() // cap 3 → evict back (b) + tr.mu.Unlock() + + if _, ok := tr.entries[[32]byte{2}]; ok { + t.Fatalf("expected least-recently-used entry (b) to be evicted") + } + for _, want := range [][32]byte{{1}, {3}, {4}} { + if _, ok := tr.entries[want]; !ok { + t.Fatalf("expected entry %v to survive", want) + } + } + if got := len(tr.entries); got != 3 { + t.Fatalf("expected cap=3 after eviction, got %d", got) + } +} + +// TestPromptCacheTrackerClampsSmallCapacity verifies maxEntries < 256 is clamped +// up to 256 so a misconfigured (or tiny test) value cannot make the cache useless. +func TestPromptCacheTrackerClampsSmallCapacity(t *testing.T) { + tr := newPromptCacheTrackerWithCapacity(time.Hour, 1) + if tr.maxEntries != 256 { + t.Fatalf("expected capacity clamped to 256, got %d", tr.maxEntries) + } +} +``` + +- [ ] **Step 2: Run tests to verify they fail** + +Run: `go test ./proxy/ -run 'TestPromptCacheLRUEvictsOldestUnused|TestPromptCacheTrackerClampsSmallCapacity' -v` +Expected: FAIL — does not compile (`newPromptCacheTrackerWithCapacity undefined`, `tr.putLocked undefined`, `tr.evictOverflowLocked undefined`, `tr.maxEntries undefined`). + +- [ ] **Step 3: Update the import block** + +In `proxy/cache_tracker.go`, add `container/list`: + +```go +import ( + "bytes" + "container/list" + "crypto/sha256" + "encoding/json" + "kiro-go/config" + "os" + "sort" + "strconv" + "strings" + "sync" + "time" +) +``` + +- [ ] **Step 4: Drop the `maxPromptCacheEntries` const** + +In `proxy/cache_tracker.go`, delete these three lines (currently `:25-27`): + +```go +// maxPromptCacheEntries bounds the in-memory cache map; once exceeded, the +// least-recently-hit entries are evicted (LRU), mirroring kiro-rs's cap. +const maxPromptCacheEntries = 4096 +``` + +- [ ] **Step 5: Replace the entry type and tracker struct** + +Replace the current `promptCacheEntry` struct (currently `:56-60`) and `promptCacheTracker` struct (currently `:62-69`) with: + +```go +type promptCacheEntry struct { + ExpiresAt time.Time + TTL time.Duration + lruElem *list.Element // back-ref into t.order; Value = fingerprint [32]byte +} + +// minPromptCacheEntries is the floor for maxEntries; an explicit smaller value +// (or a test value) is clamped up so the cache is never useless. +const minPromptCacheEntries = 256 + +type promptCacheTracker struct { + mu sync.Mutex + entries map[[32]byte]*promptCacheEntry + order *list.List // front = most-recently-used; Element.Value = [32]byte fingerprint + maxEntries int + maxSupportedTTL time.Duration + dirty bool + stopChan chan struct{} + stopOnce sync.Once +} +``` + +- [ ] **Step 6: Replace the constructor** + +Replace the current `newPromptCacheTracker` (currently `:71-79`) with the two-constructor form: + +```go +func newPromptCacheTracker(maxTTL time.Duration) *promptCacheTracker { + return newPromptCacheTrackerWithCapacity(maxTTL, config.GetPromptCacheMaxEntries()) +} + +func newPromptCacheTrackerWithCapacity(maxTTL time.Duration, maxEntries int) *promptCacheTracker { + if maxTTL <= 0 { + maxTTL = defaultPromptCacheTTL + } + if maxEntries < minPromptCacheEntries { + maxEntries = minPromptCacheEntries + } + return &promptCacheTracker{ + entries: make(map[[32]byte]*promptCacheEntry), + order: list.New(), + maxEntries: maxEntries, + maxSupportedTTL: maxTTL, + } +} +``` + +- [ ] **Step 7: Update `Load` to use `putLocked`** + +In `proxy/cache_tracker.go`, replace the body of `Load`'s locked loop (currently `:106-116`): + +```go + for _, e := range disk.Entries { + exp := time.Unix(e.ExpiresAt, 0) + if !exp.After(now) { + continue // already expired + } + t.putLocked(e.Fingerprint, exp, time.Duration(e.TTLSeconds)*time.Second) + } +``` + +- [ ] **Step 8: Update `Compute` hit path** + +In `Compute` (currently `:279-291`), replace the matched-entry block — remove the `LastHit` write and re-insert, add `MoveToFront`: + +```go + entry, ok := t.entries[breakpoint.Fingerprint] + if !ok || entry.ExpiresAt.Before(now) { + continue + } + entry.ExpiresAt = now.Add(entry.TTL) + t.order.MoveToFront(entry.lruElem) + t.dirty = true // hit extends TTL — persist so a flush before the next Update doesn't lose it + matchedTokens = minInt(breakpoint.CumulativeTokens, profile.TotalInputTokens) + if matchedTokens > lastTokens { + matchedTokens = lastTokens + } + break +``` + +- [ ] **Step 9: Update `Update` to use `putLocked` + `evictOverflowLocked`** + +In `Update` (currently `:317-329`), replace the storage loop and eviction: + +```go + // entries is the global map now (C1: cross-account sharing). + for _, breakpoint := range profile.Breakpoints { + // Skip breakpoints below the minimum cacheable token threshold. + if breakpoint.CumulativeTokens < minTokens { + continue + } + t.putLocked(breakpoint.Fingerprint, now.Add(breakpoint.TTL), breakpoint.TTL) + } + t.dirty = true + t.evictOverflowLocked() +``` + +- [ ] **Step 10: Update `pruneExpiredLocked` to keep the list consistent** + +Replace `pruneExpiredLocked` (currently `:332-338`): + +```go +func (t *promptCacheTracker) pruneExpiredLocked(now time.Time) { + for fingerprint, entry := range t.entries { + if !entry.ExpiresAt.After(now) { + t.order.Remove(entry.lruElem) + delete(t.entries, fingerprint) + } + } +} +``` + +- [ ] **Step 11: Replace `evictLRULocked` with `putLocked` + `evictOverflowLocked`** + +Delete the entire `evictLRULocked` function (currently `:340-361`) and its doc comment, and add in its place: + +```go +// putLocked inserts a fingerprint or refreshes its existing entry, marking it +// most-recently-used. Caller holds t.mu. +func (t *promptCacheTracker) putLocked(fp [32]byte, expiresAt time.Time, ttl time.Duration) { + if e, ok := t.entries[fp]; ok { + e.ExpiresAt = expiresAt + e.TTL = ttl + t.order.MoveToFront(e.lruElem) + return + } + elem := t.order.PushFront(fp) + t.entries[fp] = &promptCacheEntry{ExpiresAt: expiresAt, TTL: ttl, lruElem: elem} +} + +// evictOverflowLocked bounds the entries map to maxEntries by evicting the +// least-recently-used entries (the back of the order list). O(1) per eviction. +// Caller holds t.mu. +func (t *promptCacheTracker) evictOverflowLocked() { + for len(t.entries) > t.maxEntries { + back := t.order.Back() + if back == nil { + return + } + fp := back.Value.([32]byte) + t.order.Remove(back) + delete(t.entries, fp) + } +} +``` + +- [ ] **Step 12: Rewrite the dirty-on-hit test** + +In `proxy/cache_tracker_hardening_test.go`, replace the direct-insert line in `TestComputeSetsDirtyOnCacheHit` (currently `:109`): + +```go + tr.entries[[32]byte{1}] = promptCacheEntry{ExpiresAt: now.Add(time.Hour), TTL: time.Hour, LastHit: now} +``` + +with: + +```go + tr.mu.Lock() + tr.putLocked([32]byte{1}, now.Add(time.Hour), time.Hour) + tr.mu.Unlock() +``` + +- [ ] **Step 13: Rewrite the breakdown-clamp test's direct insert** + +In `proxy/cache_tracker_test.go`, replace the direct-insert line in `TestComputeBreakdownClampedToCreation` (currently `:420`): + +```go + tr.entries[[32]byte{1}] = promptCacheEntry{ExpiresAt: now.Add(time.Hour), TTL: time.Hour, LastHit: now} +``` + +with: + +```go + tr.mu.Lock() + tr.putLocked([32]byte{1}, now.Add(time.Hour), time.Hour) + tr.mu.Unlock() +``` + +- [ ] **Step 14: Rewrite the two disk-persistence direct inserts** + +In `proxy/cache_tracker_test.go`, `TestPromptCacheDiskPersistence`: + +Replace the t1 insert block (currently `:361-367`): + +```go + t1.mu.Lock() + t1.entries[fp] = promptCacheEntry{ + ExpiresAt: time.Now().Add(3 * time.Minute), + TTL: 5 * time.Minute, + } + t1.dirty = true + t1.mu.Unlock() +``` + +with: + +```go + t1.mu.Lock() + t1.putLocked(fp, time.Now().Add(3*time.Minute), 5*time.Minute) + t1.dirty = true + t1.mu.Unlock() +``` + +Replace the t1b insert block (currently `:383-390`): + +```go + t1b.mu.Lock() + fpExpired := sha256.Sum256([]byte("expired")) + t1b.entries[fpExpired] = promptCacheEntry{ + ExpiresAt: time.Now().Add(-1 * time.Minute), // already expired + TTL: 5 * time.Minute, + } + t1b.dirty = true + t1b.mu.Unlock() +``` + +with: + +```go + t1b.mu.Lock() + fpExpired := sha256.Sum256([]byte("expired")) + t1b.putLocked(fpExpired, time.Now().Add(-1*time.Minute), 5*time.Minute) // already expired + t1b.dirty = true + t1b.mu.Unlock() +``` + +- [ ] **Step 15: Run the two new tests to verify they pass** + +Run: `go test ./proxy/ -run 'TestPromptCacheLRUEvictsOldestUnused|TestPromptCacheTrackerClampsSmallCapacity' -v` +Expected: PASS. + +- [ ] **Step 16: Run the full proxy package test suite** + +Run: `go test ./proxy/ -v` +Expected: PASS — all existing tests (cross-account sharing, billing-header drift, breakdown clamp, dirty-on-hit, idempotent Stop, implicit breakpoint, disk persistence, configurable cap) still green. + +- [ ] **Step 17: Commit** + +```bash +git add proxy/cache_tracker.go proxy/cache_tracker_test.go proxy/cache_tracker_hardening_test.go +git commit -m "refactor(cache): O(1) LRU via container/list + 131072 default capacity" +``` + +--- + +### Task 3: Churn regression tests + +Locks in the capacity fix. Two characterization tests: at cap 131072 a 5000-prefix workload does NOT churn (the oldest prefix survives → `cache_read > 0`); at cap 4096 the same workload DOES churn (the oldest, never-re-touched prefix is evicted → `cache_read == 0`). These pass immediately after Task 2 and guard against a future capacity regression; the cap value (131072) is separately guarded by `TestPromptCacheMaxEntriesConfigurable`. + +**Files:** +- Test: `proxy/cache_tracker_test.go` (append two tests) + +**Interfaces:** +- Consumes: `newPromptCacheTrackerWithCapacity`, `putLocked`, `evictOverflowLocked` (Task 2). + +- [ ] **Step 1: Write the tests** + +Append to `proxy/cache_tracker_test.go`: + +```go +// TestCacheDoesNotChurnAtHighCapacity is the regression guard for the prefix +// churn bug: at the production default capacity (131072), a 5000-prefix +// workload (far above the old 4096 cap) does not evict the oldest seeded prefix +// before it is replayed, so the replay reads from cache. +func TestCacheDoesNotChurnAtHighCapacity(t *testing.T) { + tr := newPromptCacheTrackerWithCapacity(time.Hour, 131072) + now := time.Now() + tr.mu.Lock() + for i := 0; i < 5000; i++ { + var fp [32]byte + fp[0] = byte(i) + fp[1] = byte(i >> 8) + fp[2] = byte(i >> 16) + tr.putLocked(fp, now.Add(time.Hour), time.Hour) + } + tr.evictOverflowLocked() + tr.mu.Unlock() + + // The oldest seeded prefix (i=0, all-zero bytes) must survive at cap 131072. + var fp0 [32]byte + profile := &promptCacheProfile{ + Model: "claude-sonnet-4-6", + TotalInputTokens: 2000, + Breakpoints: []promptCacheBreakpoint{{Fingerprint: fp0, CumulativeTokens: 2000, TTL: time.Hour}}, + } + if usage := tr.Compute("acct", profile); usage.CacheReadInputTokens == 0 { + t.Fatalf("expected old prefix to survive at cap=131072 (no churn); got cache_read=0") + } +} + +// TestCacheChurnsAtLowCapacity proves the churn mechanism: at the old cap 4096 +// the same 5000-prefix workload evicts the oldest prefix (i=0), so its replay +// misses. This is the sensitivity companion to TestCacheDoesNotChurnAtHighCapacity. +func TestCacheChurnsAtLowCapacity(t *testing.T) { + tr := newPromptCacheTrackerWithCapacity(time.Hour, 4096) + now := time.Now() + tr.mu.Lock() + for i := 0; i < 5000; i++ { + var fp [32]byte + fp[0] = byte(i) + fp[1] = byte(i >> 8) + fp[2] = byte(i >> 16) + tr.putLocked(fp, now.Add(time.Hour), time.Hour) + } + tr.evictOverflowLocked() + tr.mu.Unlock() + + // 5000 > 4096 → LRU evicts the 904 oldest; i=0 (oldest, never re-touched) is gone. + var fp0 [32]byte + profile := &promptCacheProfile{ + Model: "claude-sonnet-4-6", + TotalInputTokens: 2000, + Breakpoints: []promptCacheBreakpoint{{Fingerprint: fp0, CumulativeTokens: 2000, TTL: time.Hour}}, + } + if usage := tr.Compute("acct", profile); usage.CacheReadInputTokens != 0 { + t.Fatalf("expected oldest prefix churned at cap=4096; got cache_read=%d", usage.CacheReadInputTokens) + } +} +``` + +- [ ] **Step 2: Run the tests** + +Run: `go test ./proxy/ -run 'TestCacheDoesNotChurnAtHighCapacity|TestCacheChurnsAtLowCapacity' -v` +Expected: PASS (characterization tests — the fix from Task 2 is already in). + +- [ ] **Step 3: Commit** + +```bash +git add proxy/cache_tracker_test.go +git commit -m "test(cache): regression guard for prefix churn at low capacity" +``` + +--- + +### Task 4: Metrics — atomic counters + `Stats()` + +Adds four atomic counters (`hits`, `misses`, `evictions`, `expirations`) to the tracker, instruments `Compute`/`evictOverflowLocked`/`pruneExpiredLocked`, and exposes them via a new `Stats()` method returning a JSON-tagged `PromptCacheStats`. + +**Files:** +- Modify: `proxy/cache_tracker.go` (imports, struct fields, `Compute` signature + defer, `evictOverflowLocked`, `pruneExpiredLocked`, new `PromptCacheStats` type + `Stats()`) +- Test: `proxy/cache_tracker_test.go` (append `TestCacheStats`) + +**Interfaces:** +- Consumes: the Task 2 tracker shape. +- Produces: `(*promptCacheTracker).Stats() PromptCacheStats`. Consumed by Task 5. + +- [ ] **Step 1: Write the failing test** + +Append to `proxy/cache_tracker_test.go`: + +```go +// TestCacheStats verifies the atomic counters and Stats() snapshot: one miss, +// one hit, one expiration, and one LRU eviction are all counted, and capacity +// reflects the configured bound. +func TestCacheStats(t *testing.T) { + tr := newPromptCacheTrackerWithCapacity(time.Hour, 3) + hit := &promptCacheProfile{ + Model: "claude-sonnet-4-6", + TotalInputTokens: 2000, + Breakpoints: []promptCacheBreakpoint{{Fingerprint: [32]byte{7}, CumulativeTokens: 2000, TTL: time.Hour}}, + } + now := time.Now() + + // Seed an already-expired entry, then Compute: pruneExpiredLocked drops it + // (expirations=1) and the empty cache yields a miss (misses=1). + tr.mu.Lock() + tr.putLocked([32]byte{99}, now.Add(-time.Minute), time.Hour) + tr.mu.Unlock() + tr.Compute("acct", hit) + + // Store the hit profile, then Compute → hit (hits=1). + tr.Update("acct", hit) + tr.Compute("acct", hit) + + // Overflow: entries were {99-evicted, 7}; add 3 more → {7,1,2,3} len=4 → + // evictOverflowLocked pops the LRU back (7), leaving 3 (evictions=1). + tr.mu.Lock() + for i := 1; i <= 3; i++ { + tr.putLocked([32]byte{byte(i)}, now.Add(time.Hour), time.Hour) + } + tr.evictOverflowLocked() + tr.mu.Unlock() + + stats := tr.Stats() + if stats.Hits != 1 { + t.Errorf("hits = %d, want 1", stats.Hits) + } + if stats.Misses != 1 { + t.Errorf("misses = %d, want 1", stats.Misses) + } + if stats.Evictions != 1 { + t.Errorf("evictions = %d, want 1", stats.Evictions) + } + if stats.Expirations != 1 { + t.Errorf("expirations = %d, want 1", stats.Expirations) + } + if stats.Capacity != 3 { + t.Errorf("capacity = %d, want 3", stats.Capacity) + } + if stats.Entries != 3 { + t.Errorf("entries = %d, want 3", stats.Entries) + } +} +``` + +- [ ] **Step 2: Run test to verify it fails** + +Run: `go test ./proxy/ -run TestCacheStats -v` +Expected: FAIL — `tr.Stats undefined`. + +- [ ] **Step 3: Add `sync/atomic` import** + +In `proxy/cache_tracker.go`, add `"sync/atomic"` to the import block: + +```go +import ( + "bytes" + "container/list" + "crypto/sha256" + "encoding/json" + "kiro-go/config" + "os" + "sort" + "strconv" + "strings" + "sync" + "sync/atomic" + "time" +) +``` + +- [ ] **Step 4: Add the counter fields to the struct** + +Replace the `promptCacheTracker` struct (the version from Task 2 Step 5) with: + +```go +type promptCacheTracker struct { + mu sync.Mutex + entries map[[32]byte]*promptCacheEntry + order *list.List // front = most-recently-used; Element.Value = [32]byte fingerprint + maxEntries int + maxSupportedTTL time.Duration + hits int64 // atomic — Compute calls returning CacheReadInputTokens > 0 + misses int64 // atomic — Compute calls returning CacheReadInputTokens == 0 + evictions int64 // atomic — LRU pop-backs in evictOverflowLocked + expirations int64 // atomic — TTL removals in pruneExpiredLocked + dirty bool + stopChan chan struct{} + stopOnce sync.Once +} +``` + +- [ ] **Step 5: Instrument `Compute` (named return + defer)** + +Change the `Compute` signature to a named return and add a counting defer immediately after the guard. The function header becomes: + +```go +func (t *promptCacheTracker) Compute(accountID string, profile *promptCacheProfile) (u promptCacheUsage) { + if t == nil || profile == nil || len(profile.Breakpoints) == 0 || accountID == "" { + return promptCacheUsage{} + } + defer func() { + if u.CacheReadInputTokens > 0 { + atomic.AddInt64(&t.hits, 1) + } else { + atomic.AddInt64(&t.misses, 1) + } + }() +``` + +Leave the rest of `Compute` unchanged — both later `return promptCacheUsage{...}` statements assign to the named return `u`, which the defer reads. The guard return (before the defer) is not counted. + +- [ ] **Step 6: Instrument `evictOverflowLocked`** + +In `evictOverflowLocked` (from Task 2 Step 11), add the eviction counter inside the loop: + +```go +func (t *promptCacheTracker) evictOverflowLocked() { + for len(t.entries) > t.maxEntries { + back := t.order.Back() + if back == nil { + return + } + fp := back.Value.([32]byte) + t.order.Remove(back) + delete(t.entries, fp) + atomic.AddInt64(&t.evictions, 1) + } +} +``` + +- [ ] **Step 7: Instrument `pruneExpiredLocked`** + +In `pruneExpiredLocked` (from Task 2 Step 10), add the expiration counter per delete: + +```go +func (t *promptCacheTracker) pruneExpiredLocked(now time.Time) { + for fingerprint, entry := range t.entries { + if !entry.ExpiresAt.After(now) { + t.order.Remove(entry.lruElem) + delete(t.entries, fingerprint) + atomic.AddInt64(&t.expirations, 1) + } + } +} +``` + +- [ ] **Step 8: Add `PromptCacheStats` + `Stats()`** + +Add immediately after `evictOverflowLocked`: + +```go +// PromptCacheStats is a point-in-time snapshot of cache counters, surfaced via +// /v1/stats. All counters are cumulative since tracker construction. +type PromptCacheStats struct { + Entries int `json:"entries"` + Capacity int `json:"capacity"` + Hits int64 `json:"hits"` + Misses int64 `json:"misses"` + Evictions int64 `json:"evictions"` + Expirations int64 `json:"expirations"` +} + +func (t *promptCacheTracker) Stats() PromptCacheStats { + if t == nil { + return PromptCacheStats{} + } + t.mu.Lock() + entries := len(t.entries) + capacity := t.maxEntries + t.mu.Unlock() + return PromptCacheStats{ + Entries: entries, + Capacity: capacity, + Hits: atomic.LoadInt64(&t.hits), + Misses: atomic.LoadInt64(&t.misses), + Evictions: atomic.LoadInt64(&t.evictions), + Expirations: atomic.LoadInt64(&t.expirations), + } +} +``` + +- [ ] **Step 9: Run the test to verify it passes** + +Run: `go test ./proxy/ -run TestCacheStats -v` +Expected: PASS. + +- [ ] **Step 10: Run the full proxy package test suite** + +Run: `go test ./proxy/ -v` +Expected: PASS. + +- [ ] **Step 11: Commit** + +```bash +git add proxy/cache_tracker.go proxy/cache_tracker_test.go +git commit -m "feat(cache): hit/miss/eviction metrics + Stats()" +``` + +--- + +### Task 5: Expose cache stats in `/v1/stats` + +Folds `h.promptCache.Stats()` into the existing `/v1/stats` JSON response (reuses `validateApiKey` auth at the route dispatch; zero new routes). + +**Files:** +- Modify: `proxy/handler.go` (the stats response map in `handleStats`, currently `:453-464`) +- Test: `proxy/handler_test.go` (append `TestStatsIncludesCacheMetrics`) + +**Interfaces:** +- Consumes: `(*promptCacheTracker).Stats()` (Task 4). + +- [ ] **Step 1: Write the failing test** + +Append to `proxy/handler_test.go`: + +```go +// TestStatsIncludesCacheMetrics verifies /v1/stats surfaces a "cache" object +// populated from the prompt-cache tracker's counters. +func TestStatsIncludesCacheMetrics(t *testing.T) { + cfgFile := t.TempDir() + "/config.json" + if err := config.Init(cfgFile); err != nil { + t.Fatalf("config.Init: %v", err) + } + p := accountpool.GetPool() + p.Reload() + + tr := newPromptCacheTracker(defaultPromptCacheTTL) + profile := &promptCacheProfile{ + Model: "claude-sonnet-4-6", + TotalInputTokens: 2000, + Breakpoints: []promptCacheBreakpoint{{Fingerprint: [32]byte{9}, CumulativeTokens: 2000, TTL: time.Hour}}, + } + tr.Compute("acct", profile) // miss + tr.Update("acct", profile) + tr.Compute("acct", profile) // hit + + h := &Handler{ + pool: p, + promptCache: tr, + startTime: time.Now().Unix(), + } + + rec := httptest.NewRecorder() + h.handleStats(rec, httptest.NewRequest(http.MethodGet, "/v1/stats", nil)) + + var got map[string]interface{} + if err := json.NewDecoder(rec.Body).Decode(&got); err != nil { + t.Fatalf("decode stats: %v", err) + } + cache, ok := got["cache"].(map[string]interface{}) + if !ok { + t.Fatalf("expected a cache object in stats, got %#v", got) + } + if cache["hits"].(float64) != 1 { + t.Fatalf("expected 1 hit, got %v", cache["hits"]) + } + if cache["misses"].(float64) != 1 { + t.Fatalf("expected 1 miss, got %v", cache["misses"]) + } +} +``` + +> Note: `handleStats` also reads `h.pool.Count()`, `h.pool.AvailableCount()`, `h.getCredits()`, and `h.startTime`. The literal above provides a real pool + startTime; `totalTokens`/`totalRequests` default to 0. If `getCredits()` panics with this minimal `Handler` (it should not — it reads config/h.totalCredits), inspect `getCredits` and add the needed field to the literal. + +- [ ] **Step 2: Run test to verify it fails** + +Run: `go test ./proxy/ -run TestStatsIncludesCacheMetrics -v` +Expected: FAIL — `got["cache"]` is nil (no `cache` key in the response). + +- [ ] **Step 3: Wire `cache` into the stats response** + +In `proxy/handler.go`, in `handleStats` (the `json.NewEncoder(w).Encode(map[string]interface{}{...})` literal, currently `:453-464`), add the `cache` key: + +```go + json.NewEncoder(w).Encode(map[string]interface{}{ + "status": "ok", + "version": config.Version, + "accounts": h.pool.Count(), + "available": h.pool.AvailableCount(), + "totalRequests": atomic.LoadInt64(&h.totalRequests), + "successRequests": atomic.LoadInt64(&h.successRequests), + "failedRequests": atomic.LoadInt64(&h.failedRequests), + "totalTokens": atomic.LoadInt64(&h.totalTokens), + "totalCredits": h.getCredits(), + "cache": h.promptCache.Stats(), + "uptime": time.Now().Unix() - h.startTime, + }) +``` + +- [ ] **Step 4: Run the test to verify it passes** + +Run: `go test ./proxy/ -run TestStatsIncludesCacheMetrics -v` +Expected: PASS. + +- [ ] **Step 5: Commit** + +```bash +git add proxy/handler.go proxy/handler_test.go +git commit -m "feat(handler): expose prompt-cache stats in /v1/stats" +``` + +--- + +### Task 6: Full verification + +**Files:** none (verification only). + +- [ ] **Step 1: Build the whole module** + +Run: `go build ./...` +Expected: succeeds. + +- [ ] **Step 2: Vet the whole module** + +Run: `go vet ./...` +Expected: no issues. + +- [ ] **Step 3: Run the entire test suite** + +Run: `go test ./...` +Expected: PASS. (Pre-existing unrelated failures, if any, should match the known baseline — only the two pre-existing translator image-test failures noted in the branch history are acceptable.) + +- [ ] **Step 4: Manual smoke (optional)** + +If a running instance is available: `curl -s /v1/stats -H 'x-api-key: ' | jq .cache` should show `hits`, `misses`, `evictions`, `expirations`, `entries`, `capacity`. Under multi-turn load, `hits/(hits+misses)` should rise toward 1 and `evictions` should stay near 0 — the visible proof the churn bug is fixed. + +--- + +## Self-Review (completed) + +**Spec coverage:** capacity configurable + default 131072 (Task 1 + Task 2) ✓; O(1) LRU via `container/list` (Task 2) ✓; drop `LastHit` (Task 2) ✓; min clamp 256 (Task 2) ✓; atomic counters hits/misses/evictions/expirations (Task 4) ✓; `Stats()` + `/v1/stats` exposure (Tasks 4–5) ✓; churn regression (Task 3) ✓; invariants preserved (Global Constraints; verified by the unchanged existing tests in Task 2 Step 16) ✓. + +**Placeholder scan:** none — every code step shows complete code. + +**Type/signature consistency:** `putLocked(fp [32]byte, expiresAt time.Time, ttl time.Duration)`, `evictOverflowLocked()`, `newPromptCacheTrackerWithCapacity(maxTTL time.Duration, maxEntries int)`, `Stats() PromptCacheStats` — identical across all tasks that reference them ✓. `entries` is `map[[32]byte]*promptCacheEntry` everywhere after Task 2 ✓. + +**Known GREEN-on-arrival tests:** Task 3's churn guards pass immediately after Task 2 (they are regression guards, not new behavior); the cap value 131072 is independently guarded by Task 1's `TestPromptCacheMaxEntriesConfigurable`. This is documented in Task 3's intro. diff --git a/docs/superpowers/specs/2026-06-27-credential-json-external-idp-adapter-design.md b/docs/superpowers/specs/2026-06-27-credential-json-external-idp-adapter-design.md new file mode 100644 index 00000000..d5ff650c --- /dev/null +++ b/docs/superpowers/specs/2026-06-27-credential-json-external-idp-adapter-design.md @@ -0,0 +1,377 @@ +# Adapter "Add credential JSON" cho account external_idp + +- **Ngày:** 2026-06-27 +- **Trạng thái:** Chờ duyệt (design) +- **Branch:** `feat/azure-tenant-sso` +- **Tác giả:** brainstorming session + +## 1. Bối cảnh & vấn đề + +kiro-go hỗ trợ 3 loại auth method cho account: `idc` (AWS Builder ID / IdC), `social` +(GitHub/Google), và `external_idp` (Enterprise SSO — Microsoft 365 / Entra ID / Azure AD). +Loại `external_idp` đã được thêm gần đây cho luồng **đăng nhập SSO tương tác** +(`auth/kiro_sso.go`) và **chạy bình thường** khi đăng nhập qua trình duyệt — record thật +đang hoạt động trong `data/config.json`. + +Tuy nhiên **luồng "Add credential JSON"** (paste JSON credential vào admin panel để thêm +account) **chưa được cập nhật** cho `external_idp`. Khi paste một credential JSON +`external_idp`, request bị từ chối với lỗi `"Token refresh failed"`. + +## 2. Root cause + +**Backend `proxy/handler.go` `apiImportCredentials` (line ~3008):** +- Request struct (3009–3017) chỉ có `accessToken/refreshToken/clientId/clientSecret/ + authMethod/provider/region` — **thiếu `tokenEndpoint`, `issuerUrl`, `scopes`** (và + không đọc `id/email/profileArn` cho trường hợp paste full record). +- Khối chuẩn hóa `authMethod` (3042–3053) chỉ nhận `idc`/`social`. `external_idp` rơi + vào nhánh `default`: có `clientId` nhưng **không có `clientSecret`** → bị ép thành + `social` → `auth.RefreshToken` dispatch sai sang `refreshSocialToken` → POST refresh + token lên endpoint AWS social → upstream trả lỗi → `"Token refresh failed"`. +- `tempAccount` (3058–3064) không mang `TokenEndpoint`/`Scopes`, nên ngay cả khi + `authMethod` đúng thì `refreshExternalIdpToken` cũng nhận `tokenEndpoint=""` và fail. +- Account được tạo (3079–3093) không set `TokenEndpoint/IssuerURL/Scopes/Provider` → + ngay cả khi import được thì account không có vật liệu refresh → chết sau khi access + token hết hạn. + +**Frontend `web/app.js` `importCredentials` (line ~2281):** +- Map `json.accounts` (2289–2299) không đọc `tokenEndpoint/issuerUrl/scopes` → các + field này bị drop ngay từ đầu (kể cả khi user paste Kiro export có sẵn chúng). +- Chuẩn hóa `authMethod` (2319–2322): `external_idp` có `clientId` nhưng không có + `clientSecret` → rơi nhánh cuối → `toLowerCase()!=='idc'` → `"social"`. +- Payload (2326–2333) không gửi `tokenEndpoint/issuerUrl/scopes`. + +→ Kết quả: credential `external_idp` không thể thêm qua UI, dù đã hoạt động qua SSO flow. + +## 3. Mục tiêu & non-goals + +**Mục tiêu:** +- "Add credential JSON" nhận và import được account `external_idp` ở **mọi dạng JSON** + (xem mục 4), persist đủ vật liệu refresh (`tokenEndpoint/issuerUrl/scopes/clientId`) + để account sống sót qua các lần refresh tiếp theo. +- Giữ nguyên contract hiện tại: import phải **refresh thành công** trước khi persist + (phù hợp regression test `TestApiImportCredentialsRejectsWhenRefreshFails`). + +**Non-goals (YAGNI):** +- Không thêm trust-on-import (parse JWT `exp` để bỏ qua refresh). +- Không thay đổi `auth.RefreshToken` / `refreshExternalIdpToken`. +- Không mở rộng `parseLineCredentials` (external_idp không vừa định dạng dòng cột). +- Không thêm IdP mới ngoài allow-list hiện tại. + +## 4. Các dạng JSON đầu vào phải hỗ trợ + +User chọn "tất cả các dạng" → adapter phải linh hoạt: + +**A. Object credential đơn lẻ** (chỉ vật liệu refresh — sẽ refresh để lấy accessToken): +```json +{ "authMethod": "external_idp", "refreshToken": "...", + "clientId": "fa6d79bf-...", "tokenEndpoint": "https://login.microsoftonline.com//oauth2/v2.0/token", + "issuerUrl": "https://login.microsoftonline.com//v2.0", "scopes": "...", "region": "eu-central-1" } +``` + +**B. Nguyên record account** (như trong `config.json`: có sẵn `id/email/accessToken/ +profileArn/expiresAt`): +```json +{ "id": "...", "email": "...", "accessToken": "eyJ...(JWT)", "refreshToken": "...", + "clientId": "...", "authMethod": "external_idp", "provider": "AzureAD", "region": "eu-central-1", + "profileArn": "arn:aws:codewhisperer:eu-central-1:...:profile/...", + "tokenEndpoint": "...", "issuerUrl": "...", "scopes": "..." } +``` + +**C. Kiro Account Manager export** (`{version, accounts:[{credentials:{...}}]}`) — frontend +đã xử lý `a.credentials` ở dòng 2288; chỉ cần bổ sung các field mới vào cùng map đó. + +Dạng mảng `[{...},{...}]` và object đơn đều đi qua nhánh `Array.isArray(json) ? json : [json]`. + +## 5. Thiết kế + +### 5.1 Backend — `proxy/handler.go` `apiImportCredentials` + +**(a) Mở rộng request struct** thêm field external_idp + optional identity: +```go +var req struct { + AccessToken string `json:"accessToken"` + RefreshToken string `json:"refreshToken"` + ClientID string `json:"clientId"` + ClientSecret string `json:"clientSecret"` + AuthMethod string `json:"authMethod"` + Provider string `json:"provider"` + Region string `json:"region"` + // external_idp refresh material + TokenEndpoint string `json:"tokenEndpoint"` + IssuerURL string `json:"issuerUrl"` + Scopes string `json:"scopes"` + // full-record preservation (optional, only used when provided) + ID string `json:"id"` + Email string `json:"email"` + ProfileArn string `json:"profileArn"` +} +``` + +**(b) Thay khối chuẩn hóa authMethod** bằng helper có thể unit-test, đặt detection +`external_idp` **trước** logic `clientId+clientSecret→idc`: +```go +req.AuthMethod = normalizeImportAuthMethod(req.AuthMethod, req.ClientID, req.ClientSecret, req.TokenEndpoint) +``` +```go +// normalizeImportAuthMethod chuẩn hóa auth method cho import. Quan trọng: external_idp +// phải được phát hiện TRƯỚC nhánh clientId+clientSecret→idc, vì external_idp có +// clientId nhưng không có clientSecret. +func normalizeImportAuthMethod(authMethod, clientID, clientSecret, tokenEndpoint string) string { + am := strings.ToLower(strings.TrimSpace(authMethod)) + switch { + case am == "external_idp" || am == "azuread" || am == "azure" || am == "entra" || + am == "entra-id" || am == "microsoft" || am == "m365" || am == "office365" || + am == "external": + return "external_idp" + case tokenEndpoint != "": // inference khi không khai báo rõ + return "external_idp" + case am == "social" || am == "google" || am == "github": + return "social" + case am == "idc" || am == "builderid" || am == "enterprise": + return "idc" + default: // chưa khai báo + if clientID != "" && clientSecret != "" { + return "idc" + } + if clientID != "" { + return "idc" // có clientId nhưng không clientSecret → vẫn là idc (public client IdC) + } + return "social" + } +} +``` +> **Lưu ý naming collision:** alias `"enterprise"` hiện ánh xạ sang `idc` (Kiro Account +> Manager gán provider `"Enterprise"` cho account IdC có `clientId+clientSecret`). Giữ +> nguyên contract đó để không phá import idc hiện có. Phân biệt `external_idp` bằng alias +> tường minh + inference `tokenEndpoint`. + +**(c) Validate (mới — bảo mật, xem mục 7):** sau khi chuẩn hóa, nếu `external_idp`: +- Yêu cầu `clientID != ""` và `tokenEndpoint != ""` (đúng yêu cầu của + `refreshExternalIdpToken`) — thiếu → 400. +- Validate `tokenEndpoint` (và `issuerUrl` nếu có) bằng allow-list + `auth.ValidateExternalIdpEndpoint(...)` — ngoài allow-list → 400. + +**(d) tempAccount** mang đủ vật liệu để `auth.RefreshToken` dispatch đúng: +```go +tempAccount := &config.Account{ + RefreshToken: req.RefreshToken, + ClientID: req.ClientID, + ClientSecret: req.ClientSecret, + AuthMethod: req.AuthMethod, + Region: req.Region, + TokenEndpoint: req.TokenEndpoint, + Scopes: req.Scopes, +} +``` + +**(e) Account được tạo:** set các field external_idp, default provider, bảo lưu identity +khi được cung cấp: +```go +provider := req.Provider +if provider == "" && req.AuthMethod == "external_idp" { + provider = "AzureAD" +} +email := email // từ auth.GetUserInfo(accessToken) +if email == "" { + email = req.Email // fallback từ full record +} +// Reuse id nếu được cung cấp và chưa tồn tại (tránh duplicate khi import lại backup). +id := req.ID +if id == "" || idTaken(id) { + id = auth.GenerateAccountID() +} +profileArn := newProfileArn +if profileArn == "" { + profileArn = req.ProfileArn // external_idp refresh không trả profileArn +} +account := config.Account{ + ID: id, + Email: email, + AccessToken: accessToken, + RefreshToken: req.RefreshToken, + ClientID: req.ClientID, + ClientSecret: req.ClientSecret, + AuthMethod: req.AuthMethod, + Provider: provider, + Region: req.Region, + ExpiresAt: expiresAt, + Enabled: true, + MachineId: config.GenerateMachineId(), + ProfileArn: profileArn, + TokenEndpoint: req.TokenEndpoint, + IssuerURL: req.IssuerURL, + Scopes: req.Scopes, +} +``` +> `idTaken(id)` duyệt `config.GetAccounts()` xem id đã tồn tại chưa. Helper nhỏ, +> cùng file hoặc trong package `config`. + +### 5.2 Frontend — `web/app.js` `importCredentials` + +**(a) Map `json.accounts`** (2289–2299) — bổ sung field mới: +```js +items = json.accounts.map(a => { + const c = a.credentials || {}; + return { + refreshToken: c.refreshToken || a.refreshToken, + accessToken: c.accessToken || a.accessToken, + clientId: c.clientId || a.clientId, + clientSecret: c.clientSecret || a.clientSecret, + region: c.region || a.region, + authMethod: c.authMethod || a.authMethod, + provider: c.provider || a.provider || a.idp, + tokenEndpoint: c.tokenEndpoint || a.tokenEndpoint, + issuerUrl: c.issuerUrl || a.issuerUrl, + scopes: c.scopes || a.scopes, + id: a.id, + email: c.email || a.email, + profileArn: c.profileArn || a.profileArn, + }; +}); +``` + +**(b) Chuẩn hóa authMethod** (2319–2322) — nhận `external_idp` trước: +```js +const EXTERNAL_IDP = ['external_idp','azuread','azure','entra','entra-id','microsoft','m365','office365','external']; +let authMethod = (item.authMethod || '').toLowerCase(); +if (EXTERNAL_IDP.includes(authMethod) || item.tokenEndpoint) { + authMethod = 'external_idp'; +} else if (item.clientId && item.clientSecret) { + authMethod = 'idc'; +} else if (!authMethod || authMethod === 'social') { + authMethod = 'social'; +} else { + authMethod = authMethod === 'idc' ? 'idc' : 'social'; +} +``` + +**(c) Default provider** (2323–2325) — thêm nhánh external_idp: +```js +let provider = item.provider || ''; +if (!provider && authMethod === 'external_idp') provider = 'AzureAD'; +if (!provider && authMethod === 'social') provider = 'Google'; +if (!provider && authMethod === 'idc') provider = 'BuilderId'; +``` + +**(d) Payload** (2326–2333) — thêm field mới + optional identity: +```js +const payload = { + refreshToken: item.refreshToken, + accessToken: item.accessToken || '', + clientId: item.clientId || '', + clientSecret: item.clientSecret || '', + authMethod, provider, + region: item.region || 'us-east-1', + tokenEndpoint: item.tokenEndpoint || '', + issuerUrl: item.issuerUrl || '', + scopes: item.scopes || '', + ...(item.id ? { id: item.id } : {}), + ...(item.email ? { email: item.email } : {}), + ...(item.profileArn ? { profileArn: item.profileArn } : {}), +}; +``` + +> **Parity (thứ cấp):** `web/index-legacy.html` cũng có `importCredentials` (line ~2843) +> với logic tương đương. Cập nhật tương tự để giữ nhất quán, nhưng ưu tiên thấp hơn +> (UI hiện hành dùng `app.js`). + +## 6. Luồng dữ liệu + +``` +paste JSON → app.js: gom field (gồm tokenEndpoint/issuerUrl/scopes + optional id) + → chuẩn hóa authMethod → POST /auth/credentials + → handler: decode + validate refreshToken + → normalizeImportAuthMethod → "external_idp" + → validate TokenEndpoint/IssuerURL qua allow-list + → tempAccount(đủ vật liệu) → auth.RefreshToken + → dispatch theo AuthMethod == "external_idp" + → refreshExternalIdpToken → POST login.microsoftonline.com (refresh_token grant) + → persist account (TokenEndpoint/IssuerURL/Scopes/Provider=AzureAD) + → pool.Reload() → 200 +``` + +## 7. Bảo mật (SSRF / credential exfiltration) + +Đường import là **trust boundary mới cho `tokenEndpoint`**: user paste endpoint tùy ý, +`postExternalIdpToken` (auth/oidc.go) POST thẳng `client_id + refresh_token + scope` tới +endpoint đó mà **không validate** (luôn dựa vào caller đã validate — đúng cho luồng SSO, +vì endpoint đến từ OIDC discovery đã qua `validateExternalIdpEndpoint`). Với import, một +JSON credential không tin cậy (vd. gói account chia sẻ) có thể chỉ `tokenEndpoint` vào host +nội bộ / do kẻ tấn công kiểm soát → **rò rỉ refresh token**. + +→ **Bắt buộc:** import path phải validate `tokenEndpoint` (+ `issuerUrl` nếu có) bằng cùng +allow-list mà SSO discovery dùng: `auth.validateExternalIdpEndpoint` (kiro_sso.go:512). Hàm +này kiểm tra: parse OK, **bắt buộc https**, **không chấp nhận host dạng IP literal** +(`net.ParseIP`), và host phải khớp một suffix trong `allowedExternalIdpIssuerSuffixes` +(allow-list Microsoft: `login.microsoftonline.com` …) — đủ cho mọi record thật. + +**Export + test seam:** hàm chưa export và package `proxy` cần gọi chéo, nên thêm: +- `auth.ValidateExternalIdpEndpoint(u string) error` — exported, gọi qua package var + `var externalIdpEndpointValidator = validateExternalIdpEndpoint` (mặc định = hàm thật). + 2 call site nội bộ trong `kiro_sso.go` giữ nguyên gọi `validateExternalIdpEndpoint` trực + tiếp (không đổi → giảm churn); handler import gọi `auth.ValidateExternalIdpEndpoint`. +- `auth.SetExternalIdpValidatorForTest(fn func(string) error) func(string) error` trong + `auth/testhooks.go` (thay `externalIdpEndpointValidator`, trả về giá trị cũ để restore). + Khớp pattern `SetOIDCTokenURLForTest` / `SetGlobalAuthClientForTest` sẵn có. + +Test happy-path POST lên `http://127.0.0.1:` (httptest) sẽ bị validator thật chặn +(http + IP literal + ngoài allow-list), nên test override validator bằng no-op rồi restore. + +## 8. Xử lý lỗi + +| Tình huống | Mã | Thông báo | +|---|---|---| +| JSON sai | 400 | `"Invalid JSON"` | +| Thiếu `refreshToken` | 400 | `"refreshToken is required"` | +| `external_idp` thiếu `clientId`/`tokenEndpoint` | 400 | `"external_idp requires clientId and tokenEndpoint"` | +| `tokenEndpoint`/`issuerUrl` ngoài allow-list | 400 | `"external IdP endpoint rejected: ..."` | +| Refresh fail | 400 | `"Token refresh failed: ..."` + **không persist** | +| `config.AddAccount` fail | 500 | lỗi từ config | + +Refresh-fail không persist là cố ý — giữ regression test +`TestApiImportCredentialsRejectsWhenRefreshFails`. + +## 9. Testing + +**Backend (Go, thêm vào `proxy/import_credentials_test.go` hoặc file mới):** +1. **Happy path external_idp:** dựng `httptest` fake token endpoint trả + `{accessToken, refreshToken, expiresIn}`. Body: + `{"authMethod":"external_idp","refreshToken":"rt","clientId":"c", + "tokenEndpoint":"","scopes":"s","region":"eu-central-1"}`. Assert 200 + + account persist có `TokenEndpoint/Scopes/Provider="AzureAD"/AuthMethod="external_idp"` + + ExpiresAt ≈ now+expiresIn. + - *Test seam (đã xác định):* fakeURL là `http://127.0.0.1:...`, bị validator thật chặn + (http + IP literal + ngoài allow-list). Dùng `restore := auth.SetExternalIdpValidatorForTest( + func(string) error { return nil }); defer auth.SetExternalIdpValidatorForTest(restore)` + để bypass rồi restore — không cần seam mới ngoài chính hàm đó. +2. **Reject khi tokenEndpoint ngoài allow-list:** body với `tokenEndpoint:"http://evil/"` + → 400, error chứa "endpoint rejected", không persist. +3. **Reject khi refresh fail** cho external_idp (fake endpoint trả 400 `invalid_grant`) + → 400 "Token refresh failed", không persist. +4. **Unit test `normalizeImportAuthMethod`** bảng sự thật: alias → method, + inference theo `tokenEndpoint`, fallback `clientId`/`clientId+clientSecret`. + +**Frontend:** verify thủ công qua UI (không có test JS tự động hiện có) — paste cả 3 dạng +A/B/C, kiểm toast success và account xuất hiện trong list với đầy đủ field. + +## 10. File ảnh hưởng + +| File | Thay đổi | +|---|---| +| `proxy/handler.go` | `apiImportCredentials`: struct + normalization + validate + tempAccount + created account; thêm `normalizeImportAuthMethod` (+ `idTaken` helper) | +| `proxy/import_credentials_test.go` (hoặc file mới) | test happy/reject external_idp + unit test normalization | +| `auth/oidc.go` | không đổi (refresh logic đã đúng) | +| `auth/kiro_sso.go` | thêm `var externalIdpEndpointValidator = validateExternalIdpEndpoint` + `func ValidateExternalIdpEndpoint(u string) error` (exported, gọi qua var). 2 call site nội bộ giữ nguyên. | +| `auth/testhooks.go` | thêm `SetExternalIdpValidatorForTest(fn) func(string) error` (thay + restore `externalIdpEndpointValidator`) | +| `web/app.js` | `importCredentials`: map field + normalization + provider default + payload | +| `web/index-legacy.html` | parity update (thứ cấp) | + +## 11. Rủi ro & cân nhắc + +- **Naming collision `"enterprise"`:** giữ `enterprise → idc` (contract cũ); document rõ. +- **Refresh-on-import cần egress mạng ra Microsoft:** nếu triển khai sau firewall chặn + `login.microsoftonline.com`, import sẽ fail (giống idc/social — chấp nhận được). +- **id reuse / duplicate:** chỉ reuse `id` khi chưa tồn tại; nếu trùng → sinh id mới (không + upsert) để giữ thay đổi tối thiểu. Upsert theo email ra khỏi scope (YAGNI). +- **Allow-list hẹp:** chỉ cho phép host Microsoft. Nếu sau này cần IdP khác, mở rộng + allow-list (riêng biệt, có thể tái sử dụng cho cả SSO discovery). diff --git a/docs/superpowers/specs/2026-07-02-cache-capacity-lru-metrics-design.md b/docs/superpowers/specs/2026-07-02-cache-capacity-lru-metrics-design.md new file mode 100644 index 00000000..c9106cc4 --- /dev/null +++ b/docs/superpowers/specs/2026-07-02-cache-capacity-lru-metrics-design.md @@ -0,0 +1,147 @@ +# Cache Capacity, O(1) LRU, and Metrics — Design Spec + +- **Date:** 2026-07-02 +- **Status:** Approved +- **Branch:** `feat/dispatch-cache-hardening` +- **Scope:** `proxy/cache_tracker.go`, `config/config.go`, `proxy/handler.go`, plus tests + +## Context + +`proxy/cache_tracker.go` is a **prompt-prefix fingerprint accounting layer** (NOT a response cache). It mirrors +Anthropic's upstream prompt caching to populate `cache_creation_input_tokens` / `cache_read_input_tokens` / +5m+1h breakdown in the usage map returned to clients — the Kiro/CodeWhisperer upstream does not emit those +fields, so the proxy fabricates them. Correctness is measured against matching Anthropic's *actual* upstream +caching behavior. + +Two problems: + +1. **Prefix churn.** The LRU bound is `maxPromptCacheEntries = 4096` (`cache_tracker.go:27`), and eviction is an + O(n log n) **sort-under-lock** (`evictLRULocked:342-361`). The sibling Rust project `kiro.rs` + (`src/anthropic/cache_metering.rs`) is the **same product with the same Claude Code workload**, and it hit and + fixed a real production churn bug at 4096: at ~90 req/min the table rotated in ~60s, evicting history prefixes + before the next turn (often >60s apart) reused them → "cache_read always 0, cache_creation always high, full + rebuild every turn". Rust bumped capacity to **131072** (formula: capacity ≥ write-rate × TTL). kiro-go is the + same single-global-map topology with the same load, so it is very likely affected — just undiagnosed, and with + no metrics to detect it. +2. **No observability.** `Compute` returns a usage value silently; there are no hit/miss/eviction counters + (shared gap with kiro.rs). The cache's only job is accuracy, yet its behavior is unmeasurable. + +## Goals + +1. Stop prefix churn under load → restore accurate `cache_read` reporting. +2. Remove the O(n log n)-under-lock eviction hot spot (which Rust itself flags at 131072). +3. Add observability (hit/miss/eviction counters) so cache behavior is measurable and the fix is verifiable. + +## Non-Goals + +- No change to the **global cross-account sharing** philosophy (C1; `accountID` ignored). +- No change to **disk format** (version 1). +- No new response caching, no sharding, no per-session isolation. +- No live-resize of capacity (config change requires restart — YAGNI). + +## Approaches Considered (LRU — the core change) + +- **A (chosen): `container/list` O(1) LRU.** Doubly-linked list (front = most-recently-used) + map of + `key → *entry`, where each entry holds a back-ref to its `*list.Element`. Hit/update move the element to the + front in O(1); overflow pops the back in O(1). Removes the sort-under-lock Rust itself flags at 131072. +- **B: keep sort-based, bump capacity only.** One-constant change, but inherits the O(n log n)-under-lock cost + Rust flagged (~2.2M comparisons/overflow at 131072, on every `Update` once near capacity). kiro-go is global + (higher write-rate) → worse. +- **C: sharded LRU.** Overkill — the bottleneck is the sort, not lock contention; A resolves it. Rejected (YAGNI). + +## Design + +### 1. Capacity: configurable + bumped default + +- **`config/config.go`**: add field `PromptCacheMaxEntries int` (`json:"promptCacheMaxEntries,omitempty"`) + immediately after `PromptCacheMaxRatio` (`:226`); add `GetPromptCacheMaxEntries()` (default **131072**, + clamp ≥ 256) and `UpdatePromptCacheMaxEntries(int) error`, mirroring the `PromptCacheMaxRatio` pattern + (`:933-949`). +- **`proxy/cache_tracker.go`**: remove `const maxPromptCacheEntries = 4096` (`:27`); the tracker gains a + `maxEntries int` field; `newPromptCacheTracker(maxTTL, maxEntries)` clamps ≥ 256. +- **`proxy/handler.go`**: pass `config.GetPromptCacheMaxEntries()` into the constructor (`:248`). + +Rationale for 131072: it is the empirically measured value for the same workload (Rust). Same product, same +single global map, same Claude Code load → sizing transfers directly. Not a guess. + +### 2. O(1) LRU via `container/list` + +- `import "container/list"`. The tracker gains `order *list.List` (front = MRU) alongside `entries`. +- The entry becomes pointer-typed with a back-ref; `LastHit` is dropped (list position is authoritative for LRU + order, and `LastHit` was never persisted — reset to `now` on `Load`, so post-restart order is already flat; + behavior unchanged): + +```go +entries map[[32]byte]*promptCacheEntry +order *list.List // front = most-recent; Element.Value = fingerprint [32]byte + +type promptCacheEntry struct { + ExpiresAt time.Time + TTL time.Duration + lruElem *list.Element // back-ref into t.order; Value = fingerprint [32]byte +} +``` + +- `Update` (`:317-329`): for each breakpoint — if the key exists, `MoveToFront` + update expiry/TTL; else + `PushFront(fp)` + insert entry. Then bound the map: + `for len(t.entries) > t.maxEntries { back := t.order.Back(); fp := back.Value.([32]byte); t.order.Remove(back); delete(t.entries, fp); t.evictions++ }`. +- `Compute` hit (`:279-291`): `t.order.MoveToFront(entry.lruElem)` + update expiry. +- `pruneExpiredLocked` (`:332-338`): on delete, also `t.order.Remove(entry.lruElem)`. +- `Load` (`:106-116`): rebuild map + list (`PushFront`/`PushBack` each loaded entry; post-restart order is + arbitrary — acceptable). +- `flush` (`:149-175`): unchanged (serializes the map; the list is not persisted). Disk format version 1 + unchanged. +- `newPromptCacheTracker`: initialize `order: list.New()`. + +### 3. Metrics (atomic counters) + +- The tracker gains `hits, misses, evictions, expirations int64`, updated via `atomic.AddInt64` and read via + `atomic.LoadInt64` (matching `handler.go:44` `totalRequests` style). + - `hits++` in `Compute` when the returned `CacheReadInputTokens > 0`; `misses++` when `== 0` (covers both the + empty-cache early return at `:248-262` and the main return at `:297-302`; the nil-profile/empty-account guard + at `:235` does NOT count — no cache decision is made there). + - `evictions++` per pop-back in `Update`. + - `expirations++` per deleted entry in `pruneExpiredLocked`. +- New method `Stats() PromptCacheStats { Entries int; Capacity int; Hits, Misses, Evictions, Expirations int64 }`. +- Exposure: fold into the `/v1/stats` JSON response (`handler.go:456-460`) as a `"cache"` object. Reuses + `validateApiKey` auth (`:427`); zero new routes. +- `hitRate = hits / (hits + misses)` is computable client-side → directly surfaces the "cache_read always 0" + failure mode (high misses + high evictions) versus healthy (high hitRate, low evictions). + +## Invariants Preserved (do not change) + +- Global cross-account sharing — `accountID` ignored (`:316`). +- `Update` runs only on the request-success path (`handler.go:1270`). +- `dirty`-on-hit + idempotent `Stop` (commit `c35e792`). +- `clampCacheBreakdownToCreation` + `splitAgainstTotal` (commits `c01dfb9`, `501443b`). +- Disk format version 1. +- Single mutex; all list ops under `t.mu`. + +## Testing (TDD, RED → GREEN) + +- **Churn reproduction (regression):** tracker at `cap = 131072`; feed ~5000 distinct fingerprints within the TTL; + replay an old one → still a hit (`CacheReadInputTokens > 0`). The same scenario at `cap = 4096` → miss. Mirrors + the Rust regression. +- **O(1) LRU correctness:** insert a, b, c (cap 3); touch a; insert d → evicted is b (not a); hit moves to front; + eviction counter increments by exactly the overflow count. +- **Capacity config:** `GetPromptCacheMaxEntries` default 131072; clamp ≥ 256 (0 / 1 / negative → 256); + `Update` persists. +- **Metrics:** after N hit + M miss `Compute` calls, `Stats()` reports correct counts; `expirations` increments + when TTL elapses. +- **Load rebuild:** after `Load`, `len(order) == len(entries)`; subsequent eviction works without panic or leaked + elements. +- **`LastHit` removal:** grep tests for `LastHit` and `maxPromptCacheEntries`; update any that reference them. The + existing `TestPromptCacheEvictsLRUWhenOverCapacity` must be rewritten for list-order semantics (insertion + + access order, no explicit `LastHit`). +- All existing tests stay green: cross-account sharing, billing-header drift, breakdown clamp, dirty-on-hit, + idempotent `Stop`, implicit breakpoint. + +## Risks + +- **Removing `LastHit`:** any test or constructor literal referencing it breaks compilation → fix in the RED + phase. If a test pins eviction order via `LastHit`, rewrite it for list-order semantics. +- **Constructor signature change** (`+maxEntries`): update every caller — `handler.go:248` and the test call + sites that construct a tracker directly. +- **RAM:** ~131072 × ~56 B ≈ 7–10 MiB. Negligible (Rust measured ~10 MiB). +- **Runtime resize:** `maxEntries` is set once at construction; changing config at runtime requires a restart. A + live-resize setter is out of scope (YAGNI). diff --git a/main.go b/main.go index 6d84df44..ae96326a 100644 --- a/main.go +++ b/main.go @@ -14,6 +14,8 @@ package main import ( + "context" + "errors" "fmt" "kiro-go/config" "kiro-go/logger" @@ -22,7 +24,9 @@ import ( "log" "net/http" "os" + "os/signal" "path/filepath" + "syscall" "time" ) @@ -51,9 +55,23 @@ func main() { config.SetPassword(envPassword) } + // 生产安全态势告警（不改变行为，仅提示）：默认密码 / 公网无认证裸奔。 + // 这两个组合是真实事故的高频根因，启动时一次性提示，方便运维尽早发现。 + if config.GetPassword() == "changeme" { + logger.Warnf("[SECURITY] 管理密码仍为默认 'changeme'，请通过 ADMIN_PASSWORD 环境变量或管理面板尽快修改") + } + if config.GetHost() == "0.0.0.0" && !config.IsApiKeyRequired() { + logger.Warnf("[SECURITY] 监听 0.0.0.0 且未强制 API Key 鉴权：服务对外无认证开放；生产环境请设置 requireApiKey=true 或限制绑定地址") + } + // 初始化账号池 pool.GetPool() + // 初始化请求审计持久化（默认 JSONL，失败自动禁用；KIRO_AUDIT_DISABLE 可关闭）。 + // 审计写入异步非阻塞，缓冲满即丢弃并计数（kiro_audit_dropped_total），绝不拖垮主流程。 + proxy.InitAuditStore(filepath.Dir(configPath)) + defer proxy.CloseAuditStore() + // 创建 HTTP 处理器（包含后台刷新任务） handler := proxy.NewHandler() @@ -64,18 +82,83 @@ func main() { logger.Infof("Claude API: http://%s/v1/messages", addr) logger.Infof("OpenAI API: http://%s/v1/chat/completions", addr) + // WrapHardening adds the outer zero-dep hardening chain (panic recovery, + // request-id, security headers, 64 MiB body cap) around the existing handler + // without altering its routing. + rootHandler := proxy.WrapHardening(handler) + // WriteTimeout intentionally 0: SSE streams can run for minutes while the // upstream model produces tokens. ReadHeaderTimeout + ReadTimeout still // guard against slowloris-style header/body stalls. srv := &http.Server{ Addr: addr, - Handler: handler, + Handler: rootHandler, ReadHeaderTimeout: 30 * time.Second, ReadTimeout: 60 * time.Second, IdleTimeout: 120 * time.Second, } - if err := srv.ListenAndServe(); err != nil { + // Graceful shutdown: SIGINT/SIGTERM drains in-flight requests (including + // long SSE streams) for up to shutdownTimeout before forcing exit. Without + // this a rolling deploy / Ctrl+C killed the process mid-stream, dropping + // every in-flight request and leaving clients hanging on a half-open stream. + // + // The serve/shutdown loop is factored into runServerWithDrain so the + // signal->drain->shutdown path is unit-testable (Windows taskkill cannot + // deliver SIGINT to a detached console process, so the only reliable way to + // exercise this path here is via a test harness). + const shutdownTimeout = 55 * time.Second + signals := make(chan os.Signal, 1) + signal.Notify(signals, syscall.SIGINT, syscall.SIGTERM) + defer signal.Stop(signals) + + if err := runServerWithDrain(srv.ListenAndServe, srv.Shutdown, nil, signals, shutdownTimeout); err != nil { logger.Fatalf("Server failed: %v", err) } } + +// runServerWithDrain runs serve() until either a signal arrives or serve errors. +// On signal it calls drain (if non-nil), then shutdown(ctx) bounded by timeout, +// and finally waits for serve to return (or the timeout to elapse) before +// returning. A nil drain is a no-op — this repo has no cluster worker to drain. +func runServerWithDrain(serve func() error, shutdown func(context.Context) error, drain func(), signals <-chan os.Signal, timeout time.Duration) error { + errCh := make(chan error, 1) + go func() { errCh <- serve() }() + + select { + case sig := <-signals: + logger.Infof("Received %s, shutting down (draining in-flight requests for up to %s)", sig, timeout) + if drain != nil { + drain() + } + if timeout <= 0 { + timeout = 55 * time.Second + } + ctx, cancel := context.WithTimeout(context.Background(), timeout) + if err := shutdown(ctx); err != nil { + cancel() + return err + } + // shutdown returned (all in-flight drained) — wait for ListenAndServe + // to actually unblock, or give up when the deadline passes. + select { + case err := <-errCh: + cancel() + return normalizeServerShutdownError(err) + case <-ctx.Done(): + cancel() + return ctx.Err() + } + case err := <-errCh: + return normalizeServerShutdownError(err) + } +} + +// normalizeServerShutdownError treats http.ErrServerClosed as success — that is +// the error ListenAndServe returns after a graceful Shutdown, not a real failure. +func normalizeServerShutdownError(err error) error { + if err == nil || errors.Is(err, http.ErrServerClosed) { + return nil + } + return err +} diff --git a/main_test.go b/main_test.go new file mode 100644 index 00000000..84dfd2e2 --- /dev/null +++ b/main_test.go @@ -0,0 +1,148 @@ +package main + +import ( + "context" + "errors" + "net/http" + "os" + "sync/atomic" + "syscall" + "testing" + "time" +) + +// TestRunServerWithDrainDrainsBeforeShutdown proves the signal->drain->shutdown +// ordering on Windows without needing to deliver a real SIGINT (which +// Windows taskkill cannot do to a detached console process). drain must run +// before shutdown(ctx) is called. +func TestRunServerWithDrainDrainsBeforeShutdown(t *testing.T) { + signals := make(chan os.Signal, 1) + serveStarted := make(chan struct{}) + serveDone := make(chan struct{}) + var drained atomic.Bool + var shutdownSawDrain atomic.Bool + + serve := func() error { + close(serveStarted) + <-serveDone + return http.ErrServerClosed + } + shutdown := func(ctx context.Context) error { + shutdownSawDrain.Store(drained.Load()) + close(serveDone) + return nil + } + drain := func() { + drained.Store(true) + } + + errCh := make(chan error, 1) + go func() { + errCh <- runServerWithDrain(serve, shutdown, drain, signals, time.Second) + }() + + <-serveStarted + signals <- syscall.SIGTERM + + select { + case err := <-errCh: + if err != nil { + t.Fatalf("runServerWithDrain returned error: %v", err) + } + case <-time.After(time.Second): + t.Fatal("runServerWithDrain did not return") + } + if !drained.Load() { + t.Fatal("expected drain to be called") + } + if !shutdownSawDrain.Load() { + t.Fatal("expected shutdown to run after drain") + } +} + +// TestRunServerWithDrainNilDrainIsNoop confirms a nil drain (this repo has no +// cluster worker) does not panic and the path still shuts down cleanly. +func TestRunServerWithDrainNilDrainIsNoop(t *testing.T) { + signals := make(chan os.Signal, 1) + serveStarted := make(chan struct{}) + serveDone := make(chan struct{}) + + serve := func() error { + close(serveStarted) + <-serveDone + return http.ErrServerClosed + } + shutdown := func(ctx context.Context) error { + close(serveDone) + return nil + } + + errCh := make(chan error, 1) + go func() { + errCh <- runServerWithDrain(serve, shutdown, nil, signals, time.Second) + }() + + <-serveStarted + signals <- syscall.SIGINT + + select { + case err := <-errCh: + if err != nil { + t.Fatalf("runServerWithDrain returned error: %v", err) + } + case <-time.After(time.Second): + t.Fatal("runServerWithDrain did not return") + } +} + +// TestRunServerWithDrainReturnsWhenShutdownContextExpires confirms that if the +// in-flight serve() doesn't unblock before the shutdown deadline, the function +// returns context.DeadlineExceeded rather than hanging forever. +func TestRunServerWithDrainReturnsWhenShutdownContextExpires(t *testing.T) { + signals := make(chan os.Signal, 1) + serveStarted := make(chan struct{}) + serveDone := make(chan struct{}) + + serve := func() error { + close(serveStarted) + <-serveDone + return http.ErrServerClosed + } + shutdown := func(ctx context.Context) error { + return nil + } + + errCh := make(chan error, 1) + go func() { + errCh <- runServerWithDrain(serve, shutdown, nil, signals, 20*time.Millisecond) + }() + + <-serveStarted + signals <- syscall.SIGTERM + + select { + case err := <-errCh: + close(serveDone) + if !errors.Is(err, context.DeadlineExceeded) { + t.Fatalf("expected shutdown wait to end with context deadline, got %v", err) + } + case <-time.After(250 * time.Millisecond): + close(serveDone) + t.Fatal("runServerWithDrain kept waiting for serve after shutdown context expired") + } +} + +// TestNormalizeServerShutdownError confirms http.ErrServerClosed (the normal +// result of a graceful Shutdown) is treated as success, while real errors pass through. +func TestNormalizeServerShutdownError(t *testing.T) { + if err := normalizeServerShutdownError(nil); err != nil { + t.Fatalf("nil should normalize to nil, got %v", err) + } + if err := normalizeServerShutdownError(http.ErrServerClosed); err != nil { + t.Fatalf("ErrServerClosed should normalize to nil, got %v", err) + } + real := errors.New("bind: address already in use") + if err := normalizeServerShutdownError(real); !errors.Is(err, real) { + t.Fatalf("real error should pass through, got %v", err) + } +} diff --git a/pool/account.go b/pool/account.go index 5505a391..40079730 100644 --- a/pool/account.go +++ b/pool/account.go @@ -4,14 +4,49 @@ package pool import ( "kiro-go/config" + "math/rand" "strings" "sync" - "sync/atomic" "time" ) const tokenRefreshSkewSeconds int64 = 120 +// Error-cooldown tuning. When an account hits errorCooldownThreshold consecutive +// errors it is cooled down with exponential backoff (base, 2×base, 4×base …) capped +// at errorCooldownMax, plus ±errorCooldownJitterFraction jitter so co-failing +// accounts recover at staggered times instead of stampeding upstream together. +const ( + errorCooldownBase = 1 * time.Minute + errorCooldownMax = 8 * time.Minute + errorCooldownThreshold = 3 + errorCooldownJitterFraction = 0.1 +) + +// errorCooldownDuration returns the cooldown for a given consecutive-error count: +// exponential from errorCooldownBase (starting at the threshold) up to +// errorCooldownMax, with ±10% jitter to desynchronize simultaneous recoveries. +func errorCooldownDuration(consecutiveErrors int) time.Duration { + steps := consecutiveErrors - errorCooldownThreshold + if steps < 0 { + steps = 0 + } + // Bound the shift to prevent overflow and to match the cap semantics. + if steps > 16 { + steps = 16 + } + backoff := errorCooldownBase << uint(steps) + if backoff > errorCooldownMax || backoff <= 0 { + backoff = errorCooldownMax + } + jitter := time.Duration(float64(backoff) * errorCooldownJitterFraction * (2*rand.Float64() - 1)) + d := backoff + jitter + if d <= 0 { + d = backoff + } + return d +} + // AccountPool 账号池 type AccountPool struct { mu sync.RWMutex @@ -21,6 +56,7 @@ type AccountPool struct { cooldowns map[string]time.Time // 账号冷却时间 errorCounts map[string]int // 连续错误计数 modelLists map[string]map[string]bool // accountID → set of modelIDs (from ListAvailableModels) + runtimeStats map[string]*accountRuntimeStats // accountID → dispatch health/load signals (health-scoring) } var ( @@ -32,9 +68,10 @@ var ( func GetPool() *AccountPool { poolOnce.Do(func() { pool = &AccountPool{ - cooldowns: make(map[string]time.Time), - errorCounts: make(map[string]int), - modelLists: make(map[string]map[string]bool), + cooldowns: make(map[string]time.Time), + errorCounts: make(map[string]int), + modelLists: make(map[string]map[string]bool), + runtimeStats: make(map[string]*accountRuntimeStats), } pool.Reload() }) @@ -71,74 +108,16 @@ func (p *AccountPool) GetNext() *config.Account { } // GetNextExcluding 获取下一个可用账号（加权轮询），并跳过指定账号。 +// +// Selection is delegated to selectAccountLocked (pool/dispatch.go), which among +// the accounts passing the hard filters (not cooling down, token valid, quota +// available) picks the healthiest/least-loaded one via the smart scheduler and +// stamps it as in-flight. The least-cooldown fallback is preserved inside +// selectAccountLocked so availability behaviour is unchanged. func (p *AccountPool) GetNextExcluding(excluded map[string]bool) *config.Account { - p.mu.RLock() - defer p.mu.RUnlock() - - if len(p.accounts) == 0 { - return nil - } - - allowOverUsage := config.GetAllowOverUsage() - now := time.Now() - n := len(p.accounts) - seen := make(map[string]bool) - - // 加权轮询查找可用账号 - for i := 0; i < n; i++ { - idx := atomic.AddUint64(&p.currentIndex, 1) % uint64(n) - acc := &p.accounts[idx] - - if excluded != nil && excluded[acc.ID] { - seen[acc.ID] = true - continue - } - if seen[acc.ID] { - continue - } - - // 跳过冷却中的账号 - if cooldown, ok := p.cooldowns[acc.ID]; ok && now.Before(cooldown) { - seen[acc.ID] = true - continue - } - - // 跳过即将过期的 Token - if acc.ExpiresAt > 0 && time.Now().Unix() > acc.ExpiresAt-tokenRefreshSkewSeconds { - seen[acc.ID] = true - continue - } - - // Skip accounts whose quota is exhausted, unless overrides apply. - if isQuotaBlocked(*acc, allowOverUsage) { - seen[acc.ID] = true - continue - } - - return acc - } - - // 无可用账号，返回冷却时间最短的（排除额度用尽的，除非允许超额） - var best *config.Account - var earliest time.Time - for i := range p.accounts { - acc := &p.accounts[i] - if excluded != nil && excluded[acc.ID] { - continue - } - if isQuotaBlocked(*acc, allowOverUsage) { - continue - } - if cooldown, ok := p.cooldowns[acc.ID]; ok { - if best == nil || cooldown.Before(earliest) { - best = acc - earliest = cooldown - } - } else { - return acc - } - } - return best + p.mu.Lock() + defer p.mu.Unlock() + return p.selectAccountLocked("", excluded, false) } // SetModelList 缓存账号支持的模型集合（由 handler 在刷新后调用） @@ -186,73 +165,11 @@ func (p *AccountPool) GetNextForModel(model string) *config.Account { } // GetNextForModelExcluding 获取下一个支持指定模型的可用账号，并跳过指定账号。 +// Delegates to selectAccountLocked (requireModel=true). See GetNextExcluding. func (p *AccountPool) GetNextForModelExcluding(model string, excluded map[string]bool) *config.Account { - p.mu.RLock() - defer p.mu.RUnlock() - - if len(p.accounts) == 0 { - return nil - } - - allowOverUsage := config.GetAllowOverUsage() - now := time.Now() - n := len(p.accounts) - seen := make(map[string]bool) - - for i := 0; i < n; i++ { - idx := atomic.AddUint64(&p.currentIndex, 1) % uint64(n) - acc := &p.accounts[idx] - - if excluded != nil && excluded[acc.ID] { - seen[acc.ID] = true - continue - } - if seen[acc.ID] { - continue - } - if !p.accountHasModel(acc.ID, model) { - seen[acc.ID] = true - continue - } - if cooldown, ok := p.cooldowns[acc.ID]; ok && now.Before(cooldown) { - seen[acc.ID] = true - continue - } - if acc.ExpiresAt > 0 && time.Now().Unix() > acc.ExpiresAt-tokenRefreshSkewSeconds { - seen[acc.ID] = true - continue - } - if isQuotaBlocked(*acc, allowOverUsage) { - seen[acc.ID] = true - continue - } - return acc - } - - // fallback：找冷却时间最短且支持该模型的账号 - var best *config.Account - var earliest time.Time - for i := range p.accounts { - acc := &p.accounts[i] - if excluded != nil && excluded[acc.ID] { - continue - } - if !p.accountHasModel(acc.ID, model) { - continue - } - if isQuotaBlocked(*acc, allowOverUsage) { - continue - } - if cooldown, ok := p.cooldowns[acc.ID]; ok { - if best == nil || cooldown.Before(earliest) { - best = acc - earliest = cooldown - } - } else { - return acc - } - } - return best + p.mu.Lock() + defer p.mu.Unlock() + return p.selectAccountLocked(model, excluded, true) } // GetByID 根据 ID 获取账号 @@ -261,7 +178,7 @@ func (p *AccountPool) GetByID(id string) *config.Account { defer p.mu.RUnlock() for i := range p.accounts { if p.accounts[i].ID == id { - return &p.accounts[i] + return copyAccount(&p.accounts[i]) } } return nil @@ -271,6 +188,8 @@ func (p *AccountPool) GetByID(id string) *config.Account { func (p *AccountPool) RecordSuccess(id string) { p.mu.Lock() defer p.mu.Unlock() + p.ensurePoolMapsLocked() + p.finishAccountUseLocked(id, true, false, time.Now()) delete(p.cooldowns, id) p.errorCounts[id] = 0 } @@ -280,14 +199,18 @@ func (p *AccountPool) RecordError(id string, isQuotaError bool) { p.mu.Lock() defer p.mu.Unlock() + p.ensurePoolMapsLocked() + p.finishAccountUseLocked(id, false, isQuotaError, time.Now()) p.errorCounts[id]++ if isQuotaError { // 配额错误，冷却 1 小时 p.cooldowns[id] = time.Now().Add(time.Hour) - } else if p.errorCounts[id] >= 3 { - // 连续 3 次错误，冷却 1 分钟 - p.cooldowns[id] = time.Now().Add(time.Minute) + } else if p.errorCounts[id] >= errorCooldownThreshold { + // 连续错误达阈值：指数退避（1m→2m→4m→…封顶 8m）+ ±10% 抖动。 + // 越错越久给上游喘息；抖动错开各账号解禁时刻，避免被罚账号一窝蜂同时 + // 涌回再一起撞墙（thundering herd）。 + p.cooldowns[id] = time.Now().Add(errorCooldownDuration(p.errorCounts[id])) } } @@ -364,6 +287,10 @@ func (p *AccountPool) DisableAccount(id, reason string) { _ = err } p.mu.Lock() + p.ensurePoolMapsLocked() + // Release any in-flight slot the account held before evicting it, so the + // smart scheduler's in-flight counts stay accurate across disable/re-enable. + p.finishAccountUseLocked(id, false, false, time.Now()) // Long cooldown as a safety net in case Reload races p.cooldowns[id] = time.Now().Add(24 * time.Hour) p.mu.Unlock() @@ -376,6 +303,8 @@ func (p *AccountPool) DisableAccount(id, reason string) { // the next attempt picks a different account, then reload. func (p *AccountPool) MarkOverLimit(id string) { p.mu.Lock() + p.ensurePoolMapsLocked() + p.finishAccountUseLocked(id, false, false, time.Now()) p.cooldowns[id] = time.Now().Add(time.Hour) p.mu.Unlock() p.Reload() @@ -499,3 +428,10 @@ func effectiveWeight(weight int) int { } return weight } + +// copyAccount returns a shallow copy of the account, safe to use after +// the pool lock is released. +func copyAccount(acc *config.Account) *config.Account { + c := *acc + return &c +} diff --git a/pool/cooldown_test.go b/pool/cooldown_test.go new file mode 100644 index 00000000..13b993e1 --- /dev/null +++ b/pool/cooldown_test.go @@ -0,0 +1,77 @@ +package pool + +import ( + "testing" + "time" +) + +// jitterBand returns the inclusive [lo, hi] bounds for a ±10% jitter window +// around base, computed with integer Duration arithmetic (exact for the minute +// multiples used here, and avoids float→Duration constant conversions that Go +// rejects at compile time). +func jitterBand(base time.Duration) (time.Duration, time.Duration) { + return base * 9 / 10, base * 11 / 10 +} + +// TestErrorCooldownDurationBelowThresholdIsBase confirms that below the threshold +// the shift is clamped to 0, so the cooldown equals the base duration (±jitter). +func TestErrorCooldownDurationBelowThresholdIsBase(t *testing.T) { + lo, hi := jitterBand(errorCooldownBase) + for n := 0; n < errorCooldownThreshold; n++ { + for i := 0; i < 50; i++ { + if d := errorCooldownDuration(n); d < lo || d > hi { + t.Fatalf("errorCooldownDuration(%d)=%v outside [%v,%v]", n, d, lo, hi) + } + } + } +} + +// TestErrorCooldownDurationRespectsCap confirms that once the exponential shift +// would exceed errorCooldownMax, the cooldown is clamped at the cap (±jitter) and +// never grows unbounded no matter how high the error count climbs. +func TestErrorCooldownDurationRespectsCap(t *testing.T) { + lo, hi := jitterBand(errorCooldownMax) + // n=6 -> steps=3 -> base<<3 = 8m == cap; n>=7 -> capped. Both land at the cap. + for n := 6; n <= 100; n++ { + for i := 0; i < 20; i++ { + if d := errorCooldownDuration(n); d < lo || d > hi { + t.Fatalf("errorCooldownDuration(%d)=%v outside cap band [%v,%v]", n, d, lo, hi) + } + } + } +} + +// TestErrorCooldownDurationStaysPositive confirms jitter can never drive the +// duration to zero or negative (defensive guard d<=0 → backoff). +func TestErrorCooldownDurationStaysPositive(t *testing.T) { + for n := 0; n < 1000; n++ { + for i := 0; i < 5; i++ { + if d := errorCooldownDuration(n); d <= 0 { + t.Fatalf("errorCooldownDuration(%d) returned non-positive %v", n, d) + } + } + } +} + +// TestErrorCooldownDurationGrowsBeforeCap confirms the nominal (jitter-centre) +// backoff grows exponentially with the error count until it hits the cap: the +// minimum observed duration at a higher step should not be below the maximum +// observed at a much lower step. +func TestErrorCooldownDurationGrowsBeforeCap(t *testing.T) { + var minHigh time.Duration = 1 << 62 + var maxLow time.Duration + for i := 0; i < 200; i++ { + // step 1: backoff = 2×base (just above threshold) + if d := errorCooldownDuration(errorCooldownThreshold + 1); d < minHigh { + minHigh = d + } + // step 0: backoff = base (at threshold) + if d := errorCooldownDuration(errorCooldownThreshold); d > maxLow { + maxLow = d + } + } + // 2×base minus 10% jitter (=1.8×base) must exceed 1×base plus 10% jitter (=1.1×base). + if minHigh <= maxLow { + t.Fatalf("expected higher-step cooldown to exceed lower-step: minHigh=%v maxLow=%v", minHigh, maxLow) + } +} diff --git a/pool/dispatch.go b/pool/dispatch.go new file mode 100644 index 00000000..53ed936c --- /dev/null +++ b/pool/dispatch.go @@ -0,0 +1,332 @@ +package pool + +// Health-aware account selection (production dispatch quality). +// +// Replaces pure weighted round-robin with a smart scheduler that, among the +// accounts passing the hard filters (model support, not cooling down, token +// valid, quota available), picks the one least likely to fail or stall: +// +// lowest in-flight -> fewest recent failures -> fewest recent selections +// -> fewest recent successes +// +// in-flight is incremented when an account is selected and decremented when the +// dispatch finishes (RecordSuccess / RecordError / RecordPermanentRejection / +// DisableAccount / MarkOverLimit), so under concurrency the scheduler spreads +// load instead of piling onto the same account. The per-account concurrency cap +// (pool/tuning.go, default 0 = unlimited) and the selection strategy are opt-in +// via env and default to the smart scheduler. +// +// Ported from kiro-tutu. Differences from tutu, deliberately kept: +// - This repo's least-cooldown fallback is retained: when no account passes +// all hard filters, the account with the earliest cooldown is returned +// (rather than nil -> 503), preserving existing availability behaviour. +// - endpoint429 / affinity machinery is NOT ported (no cache-affinity layer +// in this repo; endpoint-429 signals only fed affinity circuit-breaking). +// +// Safety: every finish helper clamps in-flight at 0, so a missed or double +// finish can never drive the counter negative. With the default concurrency cap +// of 0 (unlimited) a leaked in-flight only skews the smart tie-break, never +// starves an account. +import ( + "kiro-go/config" + "time" +) + +// accountFairnessWindow is the rolling window over which "recent" selection / +// success / failure counts are tallied. Beyond it the counters rotate (reset), +// so a transient burst does not permanently skew a healthy account's score. +const accountFairnessWindow = 5 * time.Minute + +// AccountRuntimeStats is an in-memory snapshot of dispatch-level account usage, +// exposed for observability (e.g. the kiro_account_inflight metric). +type AccountRuntimeStats struct { + SelectedCount int64 `json:"selectedCount"` + SuccessCount int64 `json:"successCount"` + FailureCount int64 `json:"failureCount"` + InFlight int64 `json:"inFlight"` +} + +// accountRuntimeStats tracks per-account dispatch signals used by the smart +// scheduler. The recent* fields rotate every accountFairnessWindow. +type accountRuntimeStats struct { + selectedCount int64 + successCount int64 + failureCount int64 + inFlight int64 + + windowStartedAt time.Time + recentSelectedCount int64 + recentSuccessCount int64 + recentFailureCount int64 +} + +// accountCandidate is a selectable account paired with the health/load signals +// the smart scheduler compares (accountCandidateLess). +type accountCandidate struct { + account *config.Account + inFlight int64 + recentFailureCount int64 + recentSelectedCount int64 + recentSuccessCount int64 +} + +// ensurePoolMapsLocked lazily initialises the pool's maps so callers (and tests +// that construct an AccountPool directly) cannot nil-deref. Caller holds p.mu. +func (p *AccountPool) ensurePoolMapsLocked() { + if p.cooldowns == nil { + p.cooldowns = make(map[string]time.Time) + } + if p.errorCounts == nil { + p.errorCounts = make(map[string]int) + } + if p.modelLists == nil { + p.modelLists = make(map[string]map[string]bool) + } + if p.runtimeStats == nil { + p.runtimeStats = make(map[string]*accountRuntimeStats) + } +} + +func (p *AccountPool) statsForAccountLocked(id string, now time.Time) *accountRuntimeStats { + p.ensurePoolMapsLocked() + stats := p.runtimeStats[id] + if stats == nil { + stats = &accountRuntimeStats{windowStartedAt: now} + p.runtimeStats[id] = stats + } + rotateAccountStatsWindow(stats, now) + return stats +} + +func rotateAccountStatsWindow(stats *accountRuntimeStats, now time.Time) { + if stats.windowStartedAt.IsZero() || now.Before(stats.windowStartedAt) { + stats.windowStartedAt = now + return + } + if now.Sub(stats.windowStartedAt) < accountFairnessWindow { + return + } + stats.windowStartedAt = now + stats.recentSelectedCount = 0 + stats.recentSuccessCount = 0 + stats.recentFailureCount = 0 +} + +func (p *AccountPool) accountCandidateLocked(acc *config.Account, now time.Time) accountCandidate { + stats := p.statsForAccountLocked(acc.ID, now) + return accountCandidate{ + account: acc, + inFlight: stats.inFlight, + recentFailureCount: stats.recentFailureCount, + recentSelectedCount: stats.recentSelectedCount, + recentSuccessCount: stats.recentSuccessCount, + } +} + +// accountCandidateLess orders by production safety first: avoid busy accounts, +// then avoid recently unhealthy accounts, then spread selections and successes. +func accountCandidateLess(a, b accountCandidate) bool { + switch { + case a.inFlight != b.inFlight: + return a.inFlight < b.inFlight + case a.recentFailureCount != b.recentFailureCount: + return a.recentFailureCount < b.recentFailureCount + case a.recentSelectedCount != b.recentSelectedCount: + return a.recentSelectedCount < b.recentSelectedCount + default: + return a.recentSuccessCount < b.recentSuccessCount + } +} + +func (p *AccountPool) markAccountSelectedLocked(id string, now time.Time) { + stats := p.statsForAccountLocked(id, now) + stats.selectedCount++ + stats.recentSelectedCount++ + stats.inFlight++ +} + +// finishAccountUseLocked returns the in-flight slot and records the outcome. +// success -> recentSuccess; otherwise -> recentFailure (and recent429 when +// quotaError). in-flight is clamped at 0 so a missed/double finish is harmless. +func (p *AccountPool) finishAccountUseLocked(id string, success, quotaError bool, now time.Time) { + stats := p.statsForAccountLocked(id, now) + if stats.inFlight > 0 { + stats.inFlight-- + } + if success { + stats.successCount++ + stats.recentSuccessCount++ + return + } + stats.failureCount++ + stats.recentFailureCount++ +} + +// finishAccountUseNeutralLocked returns the in-flight slot without recording a +// success or failure. Used when the dispatch outcome reflects the request +// payload (a permanent upstream rejection), not the account's health. +func (p *AccountPool) finishAccountUseNeutralLocked(id string, now time.Time) { + stats := p.statsForAccountLocked(id, now) + if stats.inFlight > 0 { + stats.inFlight-- + } +} + +// nextStartIndexLocked advances the round-robin cursor and returns the starting +// index for a selection sweep. Caller holds p.mu. +func (p *AccountPool) nextStartIndexLocked(n int) int { + p.currentIndex++ + return int(p.currentIndex % uint64(n)) +} + +// selectAccountLocked is the single selection routine shared by GetNextExcluding +// (requireModel=false) and GetNextForModelExcluding (requireModel=true). It +// gathers eligible candidates, applies the optional concurrency cap, and picks +// one via the configured strategy (default: smart health/load-aware). If no +// candidate passes the hard filters, it falls back to the least-cooldown +// eligible account (never nil when capacity exists) — preserving this repo's +// availability behaviour, which tutu lacks. +func (p *AccountPool) selectAccountLocked(model string, excluded map[string]bool, requireModel bool) *config.Account { + if len(p.accounts) == 0 { + return nil + } + + allowOverUsage := config.GetAllowOverUsage() + now := time.Now() + nowUnix := now.Unix() + n := len(p.accounts) + start := p.nextStartIndexLocked(n) + seen := make(map[string]bool) + candidates := make([]accountCandidate, 0, n) + + for i := 0; i < n; i++ { + idx := (start + i) % n + acc := &p.accounts[idx] + if excluded != nil && excluded[acc.ID] { + seen[acc.ID] = true + continue + } + if seen[acc.ID] { + continue + } + seen[acc.ID] = true + if requireModel && !p.accountHasModel(acc.ID, model) { + continue + } + if acc.ExpiresAt > 0 && nowUnix > acc.ExpiresAt-tokenRefreshSkewSeconds { + continue + } + if isQuotaBlocked(*acc, allowOverUsage) { + continue + } + if cooldown, ok := p.cooldowns[acc.ID]; ok && now.Before(cooldown) { + continue + } + candidates = append(candidates, p.accountCandidateLocked(acc, now)) + } + + if selected := p.bestCandidateLocked(candidates, now); selected != nil { + return selected + } + + // Fallback: no account passed all hard filters (all cooling down). Return the + // eligible account (model + quota OK, not excluded) with the earliest cooldown + // so the caller still gets capacity rather than a hard 503. + var best *config.Account + var earliest time.Time + for i := range p.accounts { + acc := &p.accounts[i] + if excluded != nil && excluded[acc.ID] { + continue + } + if requireModel && !p.accountHasModel(acc.ID, model) { + continue + } + if isQuotaBlocked(*acc, allowOverUsage) { + continue + } + if cooldown, ok := p.cooldowns[acc.ID]; ok { + if best == nil || cooldown.Before(earliest) { + best = acc + earliest = cooldown + } + } else { + // No cooldown at all — prefer over any cooling-down account. + p.markAccountSelectedLocked(acc.ID, now) + return copyAccount(acc) + } + } + if best != nil { + p.markAccountSelectedLocked(best.ID, now) + return copyAccount(best) + } + return nil +} + +// bestCandidateLocked applies the concurrency cap, picks one candidate via the +// configured strategy, marks it selected (in-flight++), and returns a value +// copy. Returns nil when candidates is empty. +func (p *AccountPool) bestCandidateLocked(candidates []accountCandidate, now time.Time) *config.Account { + if len(candidates) == 0 { + return nil + } + candidates = applyConcurrencyCap(candidates) + best := selectCandidate(candidates) + p.markAccountSelectedLocked(best.account.ID, now) + // Return a value copy, never an interior pointer into p.accounts: callers + // (ensureValidToken etc.) mutate AccessToken/ExpiresAt without holding p.mu, + // which would data-race with UpdateToken/Reload on the shared slice. + return copyAccount(best.account) +} + +// RecordPermanentRejection finishes an account dispatch that failed because the +// upstream rejected the request itself (e.g. "improperly formed request"). Such +// a rejection is inherent to the request payload, not the account, so it returns +// the in-flight slot WITHOUT counting any success or failure: the account's +// health must not be penalised for a bad request it merely relayed. +func (p *AccountPool) RecordPermanentRejection(id string) { + p.mu.Lock() + defer p.mu.Unlock() + p.ensurePoolMapsLocked() + p.finishAccountUseNeutralLocked(id, time.Now()) +} + +// CooldownAccount keeps an account out of routing for the given duration without +// recording a dispatch outcome. Used by failover paths that already accounted for +// the attempt via RecordError but need a fresh short cooldown. +func (p *AccountPool) CooldownAccount(id string, duration time.Duration) { + p.mu.Lock() + defer p.mu.Unlock() + p.ensurePoolMapsLocked() + p.cooldowns[id] = time.Now().Add(duration) +} + +// ResetTransientState clears per-account cooldowns, error counters and runtime +// stats and rewinds the round-robin cursor. Intended for test isolation: the +// pool is a process-wide singleton, so a prior test's failures/cooldowns would +// otherwise leak into later tests that reuse the same account IDs. +func (p *AccountPool) ResetTransientState() { + p.mu.Lock() + defer p.mu.Unlock() + p.cooldowns = make(map[string]time.Time) + p.errorCounts = make(map[string]int) + p.runtimeStats = make(map[string]*accountRuntimeStats) + p.currentIndex = 0 +} + +// GetRuntimeStatsSnapshot returns a copy of the per-account dispatch stats for +// observability (e.g. the kiro_account_inflight Prometheus gauge). +func (p *AccountPool) GetRuntimeStatsSnapshot() map[string]AccountRuntimeStats { + p.mu.RLock() + defer p.mu.RUnlock() + out := make(map[string]AccountRuntimeStats, len(p.runtimeStats)) + for id, s := range p.runtimeStats { + out[id] = AccountRuntimeStats{ + SelectedCount: s.selectedCount, + SuccessCount: s.successCount, + FailureCount: s.failureCount, + InFlight: s.inFlight, + } + } + return out +} diff --git a/pool/dispatch_test.go b/pool/dispatch_test.go new file mode 100644 index 00000000..e00c9c7d --- /dev/null +++ b/pool/dispatch_test.go @@ -0,0 +1,164 @@ +package pool + +import ( + "kiro-go/config" + "path/filepath" + "testing" + "time" +) + +// newHealthTestPool builds a pool with accounts whose IDs are given, so tests +// can assert dispatch ordering without relying on config persistence. +func newHealthTestPool(ids ...string) *AccountPool { + accounts := make([]config.Account, 0, len(ids)) + for _, id := range ids { + accounts = append(accounts, config.Account{ID: id}) + } + return &AccountPool{ + accounts: accounts, + cooldowns: make(map[string]time.Time), + errorCounts: make(map[string]int), + modelLists: make(map[string]map[string]bool), + runtimeStats: make(map[string]*accountRuntimeStats), + } +} + +// TestSmartSchedulerPrefersLowerInFlight proves the health-aware tie-break: +// among otherwise-equal accounts, the one with fewer in-flight dispatches is +// selected. Two equal candidates: selecting the first (RecordSuccess finishes +// it), then selecting again must pick a different account because the first +// still has in-flight=0 but the second has been picked once (recentSelected). +func TestSmartSchedulerPrefersLowerInFlight(t *testing.T) { + p := newHealthTestPool("a", "b") + now := time.Now() + + // Pretend "a" already has one in-flight dispatch. + p.markAccountSelectedLocked("a", now) + + acc := p.GetNextExcluding(nil) + if acc == nil { + t.Fatal("expected an account") + } + if acc.ID != "b" { + t.Fatalf("smart scheduler should prefer lower-in-flight b over a, got %q", acc.ID) + } +} + +// TestRecordSuccessDecrementsInFlight proves finishing a dispatch releases the +// in-flight slot, so the account becomes selectable again on equal footing. +func TestRecordSuccessDecrementsInFlight(t *testing.T) { + p := newHealthTestPool("only") + p.markAccountSelectedLocked("only", time.Now()) + + snap := p.GetRuntimeStatsSnapshot() + if snap["only"].InFlight != 1 { + t.Fatalf("expected inFlight=1 after selection, got %d", snap["only"].InFlight) + } + + p.RecordSuccess("only") + + snap = p.GetRuntimeStatsSnapshot() + if snap["only"].InFlight != 0 { + t.Fatalf("expected inFlight=0 after RecordSuccess, got %d", snap["only"].InFlight) + } + if snap["only"].SuccessCount != 1 { + t.Fatalf("expected successCount=1, got %d", snap["only"].SuccessCount) + } +} + +// TestRecordPermanentRejectionIsNeutral proves a permanent upstream rejection +// releases the in-flight slot WITHOUT recording a success or failure: the +// account's health must not be penalised for a bad request it merely relayed. +func TestRecordPermanentRejectionIsNeutral(t *testing.T) { + p := newHealthTestPool("acc") + p.markAccountSelectedLocked("acc", time.Now()) + + p.RecordPermanentRejection("acc") + + snap := p.GetRuntimeStatsSnapshot() + if snap["acc"].InFlight != 0 { + t.Fatalf("expected inFlight=0 after permanent rejection, got %d", snap["acc"].InFlight) + } + if snap["acc"].SuccessCount != 0 { + t.Fatalf("permanent rejection must not count as success, got %d", snap["acc"].SuccessCount) + } + if snap["acc"].FailureCount != 0 { + t.Fatalf("permanent rejection must not count as failure, got %d", snap["acc"].FailureCount) + } +} + +// TestInFlightClampedAtZero proves a double-finish (a code path bug, or +// RecordError called without a matching selection) can never drive in-flight +// negative. This is the safety guarantee that prevents a leaked finish from +// corrupting the scheduler. +func TestInFlightClampedAtZero(t *testing.T) { + p := newHealthTestPool("acc") + // Finish twice without any selection — must not panic and must clamp at 0. + p.RecordError("acc", false) + p.RecordError("acc", false) + + snap := p.GetRuntimeStatsSnapshot() + if snap["acc"].InFlight != 0 { + t.Fatalf("expected inFlight clamped at 0, got %d", snap["acc"].InFlight) + } +} + +// TestDisableAccountReleasesInFlight proves disabling an account that holds an +// in-flight slot releases it, so the runtimeStats stays accurate across a +// disable/re-enable lifecycle. config must be initialised first because +// DisableAccount persists via config.SetAccountBanStatus. +func TestDisableAccountReleasesInFlight(t *testing.T) { + cfgFile := filepath.Join(t.TempDir(), "config.json") + if err := config.Init(cfgFile); err != nil { + t.Fatalf("config.Init: %v", err) + } + + p := newHealthTestPool("acc") + p.markAccountSelectedLocked("acc", time.Now()) + if snap := p.GetRuntimeStatsSnapshot()["acc"]; snap.InFlight != 1 { + t.Fatalf("expected inFlight=1 before disable, got %d", snap.InFlight) + } + + p.DisableAccount("acc", "test") + + if snap, ok := p.GetRuntimeStatsSnapshot()["acc"]; ok && snap.InFlight != 0 { + t.Fatalf("expected inFlight=0 after disable, got %d", snap.InFlight) + } +} + +// TestLeastCooldownFallbackPreserved proves the repo's availability behaviour +// (return the earliest-cooldown account when all are cooling down, rather than +// nil/503) survived the selection rewrite. +func TestLeastCooldownFallbackPreserved(t *testing.T) { + p := newHealthTestPool("soon", "later") + now := time.Now() + p.cooldowns["soon"] = now.Add(1 * time.Second) + p.cooldowns["later"] = now.Add(1 * time.Hour) + + acc := p.GetNextExcluding(nil) + if acc == nil { + t.Fatal("expected least-cooldown fallback, got nil") + } + if acc.ID != "soon" { + t.Fatalf("expected earliest-cooldown account 'soon', got %q", acc.ID) + } +} + +// TestSelectionSpansBothAccounts proves that across many selections without +// finishing, the scheduler does not starve one account entirely (round-robin +// fairness via recentSelectedCount tie-break once in-flight ties). +func TestSelectionSpansBothAccounts(t *testing.T) { + p := newHealthTestPool("x", "y") + seen := map[string]bool{} + for i := 0; i < 50; i++ { + acc := p.GetNextExcluding(nil) + if acc != nil { + seen[acc.ID] = true + } + } + // Without finishing, both accounts should eventually be selected as the + // recentSelectedCount tie-break rotates the choice. + if !seen["x"] || !seen["y"] { + t.Fatalf("expected both accounts selected across 50 picks, got %v", seen) + } +} diff --git a/pool/tuning.go b/pool/tuning.go new file mode 100644 index 00000000..e0fcceb1 --- /dev/null +++ b/pool/tuning.go @@ -0,0 +1,125 @@ +package pool + +// Pool scheduling tuning: a hard per-account concurrency cap and opt-in account +// selection strategies. Both are configured via environment and default to the +// existing behaviour (no cap, "smart" soft scheduler) so enabling them never +// regresses current deployments. +// +// KIRO_MAX_CONCURRENT_PER_ACCOUNT int hard in-flight cap per account (0 = unlimited) +// KIRO_SELECTION_STRATEGY enum smart|round_robin|random|least_used|most_used +// +// "least_used"/"most_used" rank by remaining quota (UsageLimit-UsageCurrent); +// accounts with no usage limit count as effectively unlimited remaining. +// +// Ported verbatim from kiro-tutu (zero-dep). +import ( + "kiro-go/config" + "math" + "math/rand" + "os" + "strconv" + "strings" +) + +type selectionStrategy int + +const ( + strategySmart selectionStrategy = iota + strategyRoundRobin + strategyRandom + strategyLeastUsed + strategyMostUsed +) + +var ( + poolMaxConcurrent = envIntTuning("KIRO_MAX_CONCURRENT_PER_ACCOUNT", 0) + poolStrategy = parseStrategy(os.Getenv("KIRO_SELECTION_STRATEGY")) +) + +func envIntTuning(name string, def int) int { + if v := strings.TrimSpace(os.Getenv(name)); v != "" { + if n, err := strconv.Atoi(v); err == nil && n >= 0 { + return n + } + } + return def +} + +func parseStrategy(s string) selectionStrategy { + switch strings.ToLower(strings.TrimSpace(s)) { + case "round_robin", "roundrobin", "ordered": + return strategyRoundRobin + case "random": + return strategyRandom + case "least_used", "leastused", "least-used": + return strategyLeastUsed + case "most_used", "mostused", "most-used": + return strategyMostUsed + default: + return strategySmart + } +} + +// applyConcurrencyCap drops candidates whose in-flight count has reached the +// per-account cap. If every candidate is saturated it returns the original set +// unchanged, so a burst degrades to "least-busy" (via the scheduler) rather than +// a hard 503 while capacity technically exists. +func applyConcurrencyCap(candidates []accountCandidate) []accountCandidate { + if poolMaxConcurrent <= 0 || len(candidates) == 0 { + return candidates + } + filtered := make([]accountCandidate, 0, len(candidates)) + for _, c := range candidates { + if c.inFlight < int64(poolMaxConcurrent) { + filtered = append(filtered, c) + } + } + if len(filtered) == 0 { + return candidates + } + return filtered +} + +// selectCandidate chooses one candidate according to the configured strategy. +// candidates must be non-empty. +func selectCandidate(candidates []accountCandidate) accountCandidate { + switch poolStrategy { + case strategyRandom: + return candidates[rand.Intn(len(candidates))] + case strategyRoundRobin: + // Candidates are gathered in rotation order (selectAccountLocked starts + // at an advancing index), so taking the first entry yields round-robin. + return candidates[0] + case strategyLeastUsed: + best := candidates[0] + for _, c := range candidates[1:] { + if remainingQuota(c.account) > remainingQuota(best.account) { + best = c + } + } + return best + case strategyMostUsed: + best := candidates[0] + for _, c := range candidates[1:] { + if remainingQuota(c.account) < remainingQuota(best.account) { + best = c + } + } + return best + default: // strategySmart — the health/load-aware soft scheduler + best := candidates[0] + for _, c := range candidates[1:] { + if accountCandidateLess(c, best) { + best = c + } + } + return best + } +} + +func remainingQuota(acc *config.Account) float64 { + if acc == nil || acc.UsageLimit <= 0 { + return math.MaxFloat64 + } + return acc.UsageLimit - acc.UsageCurrent +} diff --git a/proxy/account_failover.go b/proxy/account_failover.go index 4f2da150..8936222a 100644 --- a/proxy/account_failover.go +++ b/proxy/account_failover.go @@ -3,11 +3,56 @@ package proxy import ( "kiro-go/config" "kiro-go/logger" + "math/rand" "strings" "time" ) -const maxAccountRetryAttempts = 3 +// absoluteMaxAccountRetryAttempts is a defensive ceiling on per-request account +// retries. The budget otherwise means "iterate every selectable account" (see +// resolveAccountRetryBudget): each retry excludes already-tried accounts and +// must pick a different one, so the real retry count converges to the number of +// accounts and never runs away. This only guards pathological cases. +const absoluteMaxAccountRetryAttempts = 64 + +// resolveAccountRetryBudget returns how many account attempts one request may +// make. Previously this was a fixed 3, so with >3 accounts a request only tried +// the first 3 and gave up while the rest sat idle — under a 429 storm that meant +// an outright failure. Now it scales with the account count so every account +// can be tried. totalAccounts<=0 falls back to 1 (try at least once). +func resolveAccountRetryBudget(totalAccounts int) int { + if totalAccounts <= 0 { + return 1 + } + if totalAccounts > absoluteMaxAccountRetryAttempts { + return absoluteMaxAccountRetryAttempts + } + return totalAccounts +} + +// accountRetryBackoff returns the wait before the next retry attempt: +// exponential 200ms→2s cap + jitter, so a transient upstream blip isn't +// amplified by back-to-back retries against the next account. attempt is 0-based. +func accountRetryBackoff(attempt int) time.Duration { + const baseMS = 200 + const maxMS = 2000 + if attempt < 0 { + attempt = 0 + } + shift := attempt + if shift > 4 { + shift = 4 + } + backoff := baseMS << shift + if backoff > maxMS { + backoff = maxMS + } + jitter := 0 + if j := backoff / 4; j > 0 { + jitter = rand.Intn(j) + } + return time.Duration(backoff+jitter) * time.Millisecond +} func isQuotaErrorMessage(msg string) bool { msg = strings.ToLower(msg) @@ -95,6 +140,15 @@ func (h *Handler) handleAccountFailure(account *config.Account, err error) { errMsg := err.Error() switch { + case isUpstreamPermanentError(err): + // The request itself is malformed/rejected by the upstream regardless of + // which account relays it. Penalising the account's health, cooling it + // down, or rotating to another account would be wrong (the account is + // healthy) and harmful (scatters cache affinity, wastes upstream hits). + // Neutral: release the in-flight slot the selection took (RecordPermanentRejection) + // without recording any success or failure, then return. + h.pool.RecordPermanentRejection(account.ID) + return case isOverageErrorMessage(errMsg): h.disableAccountOverage(account) h.pool.RecordError(account.ID, false) diff --git a/proxy/apikey_batch.go b/proxy/apikey_batch.go new file mode 100644 index 00000000..22cf93b7 --- /dev/null +++ b/proxy/apikey_batch.go @@ -0,0 +1,160 @@ +package proxy + +import ( + "encoding/json" + "net/http" + "strings" + + "kiro-go/auth" + "kiro-go/config" + "kiro-go/logger" +) + +// ApiKeyImportResult is the per-key outcome of a bulk API-key import. +type ApiKeyImportResult struct { + MaskedKey string `json:"maskedKey"` + AccountID string `json:"accountId,omitempty"` + Email string `json:"email,omitempty"` + Subscription string `json:"subscription,omitempty"` + UsageCurrent float64 `json:"usageCurrent,omitempty"` + UsageLimit float64 `json:"usageLimit,omitempty"` + Imported bool `json:"imported"` + Skipped bool `json:"skipped"` + InfoOK bool `json:"infoOk,omitempty"` + Error string `json:"error,omitempty"` +} + +// ApiKeyImportSummary aggregates a bulk import run. +type ApiKeyImportSummary struct { + Total int `json:"total"` + Imported int `json:"imported"` + Skipped int `json:"skipped"` + InfoFailed int `json:"infoFailed"` + Results []ApiKeyImportResult `json:"results"` +} + +// maskKey returns head-6 … tail-4 of a key (or "****" when the key is too short +// to safely mask). +func maskKey(key string) string { + if len(key) <= 10 { + return "****" + } + return key[:6] + "…" + key[len(key)-4:] +} + +// ImportApiKeys parses newline-separated Kiro API keys and imports each as an +// enabled api_key account. Dedups against the existing snapshot (by KiroApiKey) +// and within the batch itself. For each new key it best-effort fetches usage +// info (email/subscription/credit) via RefreshAccountInfo. Returns a per-key +// result summary; never aborts the whole batch on a single key's failure. +func (h *Handler) ImportApiKeys(rawText, region, authRegion, apiRegion string) ApiKeyImportSummary { + summary := ApiKeyImportSummary{Results: []ApiKeyImportResult{}} + + // Snapshot existing keys to dedup against. + existing := make(map[string]bool) + for _, a := range config.GetAccounts() { + if a.KiroApiKey != "" { + existing[a.KiroApiKey] = true + } + } + seenInBatch := make(map[string]bool) + + if region == "" { + region = "us-east-1" + } + + for _, raw := range strings.Split(rawText, "\n") { + key := strings.TrimSpace(raw) + if key == "" { + continue + } + summary.Total++ + masked := maskKey(key) + + if existing[key] || seenInBatch[key] { + summary.Skipped++ + summary.Results = append(summary.Results, ApiKeyImportResult{MaskedKey: masked, Skipped: true}) + continue + } + seenInBatch[key] = true + + account := config.Account{ + ID: auth.GenerateAccountID(), + AuthMethod: "api_key", + KiroApiKey: key, + AccessToken: key, // pool/dispatch/metrics compatibility + Region: region, + AuthRegion: authRegion, + ApiRegion: apiRegion, + ExpiresAt: 0, // api_key: never refresh + Enabled: true, + MachineId: config.GenerateMachineId(), + } + if err := config.AddAccount(account); err != nil { + summary.Results = append(summary.Results, ApiKeyImportResult{MaskedKey: masked, Error: err.Error()}) + continue + } + summary.Imported++ + + result := ApiKeyImportResult{MaskedKey: masked, Imported: true, AccountID: account.ID} + // Best-effort usage info. RefreshAccountInfo uses KiroApiKey (mirrored into + // AccessToken) as the bearer; profile-ARN resolution is skipped for api_key. + if info, err := RefreshAccountInfo(&account); err != nil { + summary.InfoFailed++ + logger.Warnf("[ApiKeyBatch] RefreshAccountInfo failed for %s: %v", masked, err) + } else if info != nil { + result.InfoOK = true + result.Email = info.Email + result.Subscription = info.SubscriptionTitle + if result.Subscription == "" { + result.Subscription = info.SubscriptionType + } + result.UsageCurrent = info.UsageCurrent + result.UsageLimit = info.UsageLimit + } + // Async model-cache refresh (non-blocking). + go func(acc config.Account) { + if err := h.fetchAndCacheAccountModels(&acc); err != nil { + logger.Warnf("[ModelsCache] Auto-refresh failed for api_key account %s: %v", accountEmailForLog(&acc), err) + } + }(account) + summary.Results = append(summary.Results, result) + } + + if summary.Imported > 0 { + h.pool.Reload() + } + return summary +} + +// apiImportApiKeys handles POST /auth/apikeys-batch: bulk import of newline- +// separated Kiro API keys. +func (h *Handler) apiImportApiKeys(w http.ResponseWriter, r *http.Request) { + r.Body = http.MaxBytesReader(w, r.Body, 2<<20) // 2 MB cap on the keys payload + var req struct { + Keys string `json:"keys"` + Region string `json:"region"` + AuthRegion string `json:"authRegion"` + ApiRegion string `json:"apiRegion"` + } + if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + w.WriteHeader(400) + json.NewEncoder(w).Encode(map[string]string{"error": "Invalid JSON"}) + return + } + if strings.TrimSpace(req.Keys) == "" { + w.WriteHeader(400) + json.NewEncoder(w).Encode(map[string]string{"error": "keys is required"}) + return + } + + summary := h.ImportApiKeys(req.Keys, req.Region, req.AuthRegion, req.ApiRegion) + json.NewEncoder(w).Encode(map[string]interface{}{ + "success": true, + "total": summary.Total, + "imported": summary.Imported, + "skipped": summary.Skipped, + "infoFailed": summary.InfoFailed, + "results": summary.Results, + }) +} diff --git a/proxy/apikey_batch_test.go b/proxy/apikey_batch_test.go new file mode 100644 index 00000000..850ef3ea --- /dev/null +++ b/proxy/apikey_batch_test.go @@ -0,0 +1,134 @@ +package proxy + +import ( + "encoding/json" + "io" + "kiro-go/config" + accountpool "kiro-go/pool" + "net/http" + "net/http/httptest" + "strings" + "testing" +) + +// TestMaskKey verifies head-6 … tail-4 masking and the short-key fallback. +func TestMaskKey(t *testing.T) { + if got := maskKey("ABCDEFGHIJKLMNOP"); got != "ABCDEF…MNOP" { + t.Fatalf("mask long: want ABCDEF…MNOP, got %q", got) + } + if got := maskKey("short"); got != "****" { + t.Fatalf("mask short: want ****, got %q", got) + } +} + +// TestImportApiKeysDedupAndPerKeyResult verifies dedup against existing accounts +// and within the batch, plus per-key result + maskedKey in the summary. The +// best-effort RefreshAccountInfo is exercised against a stubbed REST store that +// 404s (best-effort → InfoFailed++, never a fatal/ban path). +func TestImportApiKeysDedupAndPerKeyResult(t *testing.T) { + if err := config.Init(t.TempDir() + "/config.json"); err != nil { + t.Fatalf("config.Init: %v", err) + } + // Seed an existing api_key account to dedup against. + if err := config.AddAccount(config.Account{ID: "existing-1", AuthMethod: "api_key", KiroApiKey: "key-already-present", AccessToken: "key-already-present", Enabled: true}); err != nil { + t.Fatalf("seed: %v", err) + } + + // Stub the REST store so the best-effort RefreshAccountInfo / model-cache + // fetches don't hit real network. A 404 is NOT an auth error, so it never + // trips the ban path — it just records InfoFailed. + kiroRestHttpStore.Store(&http.Client{ + Transport: roundTripFunc(func(*http.Request) (*http.Response, error) { + return &http.Response{ + StatusCode: http.StatusNotFound, + Body: io.NopCloser(strings.NewReader("not found")), + Header: make(http.Header), + }, nil + }), + }) + t.Cleanup(func() { InitKiroHttpClient("") }) + + h := &Handler{pool: accountpool.GetPool()} + // One duplicate of the seeded key, one within-batch duplicate, two fresh keys. + keys := "key-already-present\nKEY-FRESH-1234567890\nKEY-FRESH-1234567890\nKEY-FRESH-2234567890" + summary := h.ImportApiKeys(keys, "eu-central-1", "", "") + + if summary.Total != 4 { + t.Fatalf("total: want 4, got %d", summary.Total) + } + // 1 existing-skip + 1 within-batch-skip = 2 skipped; 2 imported. + if summary.Imported != 2 { + t.Fatalf("imported: want 2, got %d", summary.Imported) + } + if summary.Skipped != 2 { + t.Fatalf("skipped: want 2, got %d", summary.Skipped) + } + if len(summary.Results) != 4 { + t.Fatalf("results: want 4, got %d", len(summary.Results)) + } + + // Persisted accounts: the 1 seed + 2 fresh = 3. + accs := config.GetAccounts() + if len(accs) != 3 { + t.Fatalf("persisted accounts: want 3, got %d", len(accs)) + } + for _, a := range accs { + if a.AuthMethod != "api_key" { + t.Fatalf("account %s AuthMethod: want api_key, got %q", a.ID, a.AuthMethod) + } + if a.ExpiresAt != 0 { + t.Fatalf("account %s ExpiresAt: want 0, got %d", a.ID, a.ExpiresAt) + } + } +} + +// TestApiImportApiKeysHandler verifies the POST /auth/apikeys-batch handler +// decodes the body and returns the summary shape. +func TestApiImportApiKeysHandler(t *testing.T) { + if err := config.Init(t.TempDir() + "/config.json"); err != nil { + t.Fatalf("config.Init: %v", err) + } + kiroRestHttpStore.Store(&http.Client{ + Transport: roundTripFunc(func(*http.Request) (*http.Response, error) { + return &http.Response{ + StatusCode: http.StatusNotFound, + Body: io.NopCloser(strings.NewReader("not found")), + Header: make(http.Header), + }, nil + }), + }) + t.Cleanup(func() { InitKiroHttpClient("") }) + + h := &Handler{pool: accountpool.GetPool()} + body := `{"keys":"KEY-A-1234567890\nKEY-B-1234567890","region":"us-east-1"}` + req := httptest.NewRequest("POST", "/auth/apikeys-batch", strings.NewReader(body)) + rec := httptest.NewRecorder() + h.apiImportApiKeys(rec, req) + + if rec.Code != 200 { + t.Fatalf("expected 200, got %d body=%s", rec.Code, rec.Body.String()) + } + var resp map[string]json.RawMessage + if err := json.Unmarshal(rec.Body.Bytes(), &resp); err != nil { + t.Fatalf("decode: %v", err) + } + for _, k := range []string{"success", "total", "imported", "skipped", "infoFailed", "results"} { + if string(resp[k]) == "" { + t.Fatalf("response missing key %q: %v", k, resp) + } + } +} + +// TestApiImportApiKeysRejectsEmpty verifies the 400 guard. +func TestApiImportApiKeysRejectsEmpty(t *testing.T) { + if err := config.Init(t.TempDir() + "/config.json"); err != nil { + t.Fatalf("config.Init: %v", err) + } + h := &Handler{pool: accountpool.GetPool()} + req := httptest.NewRequest("POST", "/auth/apikeys-batch", strings.NewReader(`{"keys":""}`)) + rec := httptest.NewRecorder() + h.apiImportApiKeys(rec, req) + if rec.Code != 400 { + t.Fatalf("expected 400 for empty keys, got %d", rec.Code) + } +} diff --git a/proxy/apikey_characterization_test.go b/proxy/apikey_characterization_test.go new file mode 100644 index 00000000..0511b89d --- /dev/null +++ b/proxy/apikey_characterization_test.go @@ -0,0 +1,43 @@ +package proxy + +import ( + "kiro-go/auth" + "kiro-go/config" + accountpool "kiro-go/pool" + "testing" +) + +// TestApiKeyAccountDoesNotRefreshOnHandlerPath pins spec §6: an api_key account +// (ExpiresAt=0, RefreshToken="") must not trigger auth.RefreshToken on the +// handler refresh path. refreshAccountToken's not-near-expiry guard returns +// early (force=false, ExpiresAt==0), so the OIDC URL builder is never consulted. +func TestApiKeyAccountDoesNotRefreshOnHandlerPath(t *testing.T) { + if err := config.Init(t.TempDir() + "/config.json"); err != nil { + t.Fatalf("config.Init: %v", err) + } + // Sentinel: if a refresh fires, the OIDC URL builder is consulted (to derive + // the token endpoint), which would mean a round-trip was about to happen. + prev := auth.GetOIDCTokenURLForTest() + auth.SetOIDCTokenURLForTest(func(string) string { + t.Fatal("refresh must not be called for an api_key account") + return "" + }) + defer auth.SetOIDCTokenURLForTest(prev) + + account := &config.Account{ + ID: "ak-1", + AuthMethod: "api_key", + KiroApiKey: "k", + AccessToken: "k", + ExpiresAt: 0, + RefreshToken: "", + Enabled: true, + } + if !account.IsApiKeyCredential() { + t.Fatalf("expected IsApiKeyCredential() == true") + } + h := &Handler{pool: accountpool.GetPool()} + if err := h.refreshAccountToken(account, false); err != nil { + t.Fatalf("refreshAccountToken returned error: %v", err) + } +} diff --git a/proxy/auth.go b/proxy/auth.go index 922ea220..b54d14ac 100644 --- a/proxy/auth.go +++ b/proxy/auth.go @@ -2,6 +2,7 @@ package proxy import ( "context" + "crypto/subtle" "kiro-go/config" "net/http" "strings" @@ -85,7 +86,11 @@ func (h *Handler) authenticate(r *http.Request) (*config.ApiKeyEntry, error) { // Auth required but nothing configured → fail closed. return nil, newAuthError(http.StatusUnauthorized, "authentication_error", "API key authentication is required but no keys are configured") } - if provided == "" || provided != expected { + // Constant-time compare so a network observer cannot time-attack the key + // (byte-by-byte early-exit of `provided != expected` would leak the length of + // the correct prefix). Matches the constant-time check already used for the + // admin password (handler.go) and audit-log password (prod_middleware.go). + if provided == "" || subtle.ConstantTimeCompare([]byte(provided), []byte(expected)) != 1 { return nil, newAuthError(http.StatusUnauthorized, "authentication_error", "Invalid or missing API key") } return nil, nil diff --git a/proxy/cache_hitrate_sim_test.go b/proxy/cache_hitrate_sim_test.go new file mode 100644 index 00000000..2f21947e --- /dev/null +++ b/proxy/cache_hitrate_sim_test.go @@ -0,0 +1,112 @@ +package proxy + +import ( + "fmt" + "strings" + "testing" + "time" +) + +// TestCacheHitRateSimulation drives the prompt-cache tracker with realistic +// multi-session / multi-turn Claude Code workloads and reports the observed hit +// rate from the Stats() counters. It does NOT touch the network — it replays the +// same BuildClaudeProfile → Compute → Update flow the real handler runs, which is +// exactly what the cache metrics measure. +// +// Run: go test ./proxy/ -run TestCacheHitRateSimulation -v +func TestCacheHitRateSimulation(t *testing.T) { + // Shared system prompt, large enough (~1500 tokens) to clear the 1024-token + // min-cacheable threshold. Represents the common cached system prompt. + sharedSystem := strings.Repeat( + "You are a meticulous senior engineer who writes careful, idiomatic Go and reviews diffs critically. ", 95) + + // buildReq assembles a multi-turn request that grows turn-by-turn, the way a + // real conversation does (turn N's message history is a superset of turn N-1's). + buildReq := func(userText string, turns int) *ClaudeRequest { + messages := make([]ClaudeMessage, 0, turns*2) + for i := 0; i < turns; i++ { + messages = append(messages, ClaudeMessage{ + Role: "user", + Content: fmt.Sprintf("%s (turn %d)", userText, i), + }) + messages = append(messages, ClaudeMessage{ + Role: "assistant", + Content: strings.Repeat("Here is a careful, idiomatic answer with rationale. ", 25), + }) + } + return &ClaudeRequest{ + Model: "claude-sonnet-4-6", + System: []interface{}{ + map[string]interface{}{ + "type": "text", + "text": sharedSystem, + "cache_control": map[string]interface{}{"type": "ephemeral", "ttl": "5m"}, + }, + }, + Messages: messages, + } + } + + run := func(tr *promptCacheTracker, req *ClaudeRequest) { + prof := tr.BuildClaudeProfile(req, estimateClaudeRequestInputTokens(req)) + if prof == nil { + return + } + tr.Compute("acct", prof) // estimate-domain lookup → counts a hit or miss + tr.Update("acct", prof) // success path: store refreshed breakpoints + } + + report := func(label string, tr *promptCacheTracker) { + s := tr.Stats() + total := s.Hits + s.Misses + rate := 0.0 + if total > 0 { + rate = float64(s.Hits) / float64(total) + } + t.Logf("") + t.Logf("=== %s ===", label) + t.Logf(" requests : %d", total) + t.Logf(" hits : %d", s.Hits) + t.Logf(" misses : %d", s.Misses) + t.Logf(" HIT RATE : %.1f%%", rate*100) + t.Logf(" evictions : %d (LRU pop-backs)", s.Evictions) + t.Logf(" expirations : %d (TTL)", s.Expirations) + t.Logf(" entries now : %d / %d capacity", s.Entries, s.Capacity) + } + + // Scenario A — realistic Claude Code load: 30 sessions share one system + // prompt; each grows 1→6 turns (180 requests total). Cross-session sharing + // (C1) means even a brand-new session hits on the shared system prefix. + trA := newPromptCacheTracker(time.Hour) + for s := 0; s < 30; s++ { + tag := fmt.Sprintf("session-%d", s) + for n := 1; n <= 6; n++ { + run(trA, buildReq(tag, n)) + } + } + report("A: realistic (30 sessions x 6 turns, shared system prompt)", trA) + + // Scenario B — cold baseline: every request uses a unique system prompt, so + // nothing is reusable. Expect ~0% hit rate. + trB := newPromptCacheTracker(time.Hour) + for i := 0; i < 60; i++ { + req := buildReq(fmt.Sprintf("unique-%d", i), 1) + req.System = []interface{}{ + map[string]interface{}{ + "type": "text", + "text": strings.Repeat(fmt.Sprintf("unique-system-prompt-%d ", i), 95), + "cache_control": map[string]interface{}{"type": "ephemeral", "ttl": "5m"}, + }, + } + run(trB, req) + } + report("B: cold baseline (60 unique system prompts, no sharing)", trB) + + // Scenario C — single long conversation: one session, 12 turns. Only turn 1 + // is a cold miss; turns 2..12 hit at ever-deeper prefixes. Expect ~92%. + trC := newPromptCacheTracker(time.Hour) + for n := 1; n <= 12; n++ { + run(trC, buildReq("one-session", n)) + } + report("C: single conversation (1 session x 12 turns)", trC) +} diff --git a/proxy/cache_metrics.go b/proxy/cache_metrics.go new file mode 100644 index 00000000..d9a3b222 --- /dev/null +++ b/proxy/cache_metrics.go @@ -0,0 +1,367 @@ +package proxy + +import ( + "sync" + "time" + + "kiro-go/config" + "kiro-go/logger" +) + +// claudeCacheDispatchDiagnostics captures one cache dispatch's outcome and +// token usage for aggregation by cacheDispatchMetrics. This is a minimal shape +// of the type kiro-tutu defines in cache_affinity.go (which also carries +// affinity-key / endpoint / request-id fields used for affinity logging). This +// repo deliberately has no cache affinity (global fingerprints + simulated +// cache → session→account stickiness yields no real hit-rate gain), so only the +// fields the metrics consume are retained. The fields kept match what tutu's +// cache_metrics.go reads, so the aggregation logic ports verbatim. +type claudeCacheDispatchDiagnostics struct { + Outcome string // "success" | "failure" + Model string + AffinitySource string + Account *config.Account + Usage promptCacheUsage + InputTokens int + OutputTokens int +} + +// cacheDispatchMetrics 在内存中按时间窗口聚合 [claudeCacheDispatchDiagnostics] 诊断数据， +// 使缓存命中率指标无需开启 info 日志即可审计（见 issue #5）。聚合在诊断汇聚点 +// recordClaudeCacheDispatch 中无条件进行，因此与日志级别完全解耦。 +// +// 数据按分钟桶保留，Snapshot 时合并落在查询窗口内的桶；过期桶在写入时惰性清理， +// 内存占用上界为 retain 分钟 × (1 + 维度基数) 个计数器。 +const ( + cacheMetricsBucketSpan = time.Minute + // 比最大查询窗口（1h）多保留几分钟，避免边界请求被提前裁剪。 + cacheMetricsRetain = 65 * time.Minute +) + +// 缓存命中率安全阀：本网关返回给下游的 cache_read 是本地模拟值，命中率越高下游 +// 付费越少、由我方补贴。业务约定聚合命中率上限 ~88%（见 ComputeScoped 的 cap 注释）。 +// 三个旋钮（前缀断点 / 1h 本地存活 / 0.99 单请求 cap）叠加后真实聚合可能冲破该线， +// 而 cap 只约束单请求、不强制聚合。这里在聚合侧加观测告警：滚动窗口聚合命中率超阈值 +// 时打 WARN，使补贴失控可被发现，而不是只靠注释。仅告警、不改数，不影响计费。 +const ( + // cacheHitRatioWarnThreshold 是触发 WARN 的聚合命中率（read/总输入，正确口径）。 + cacheHitRatioWarnThreshold = 0.88 + // cacheHitRatioWarnWindow 是计算聚合命中率的滚动窗口。 + cacheHitRatioWarnWindow = 5 * time.Minute + // cacheHitRatioWarnInterval 是两次告警的最小间隔，避免高频请求刷屏。 + cacheHitRatioWarnInterval = time.Minute + // cacheHitRatioWarnMinDispatches 是触发告警所需的最小窗口样本量，避免小样本 + // （首轮/冷启动）噪声造成的偶发高比率误报。 + cacheHitRatioWarnMinDispatches = 50 +) + +// cacheMetricCounters 是一组可累加的计数器。token 字段累加的是每条 dispatch 的 +// 原始值：InputTokens 为输入总量（含 cache read+creation），成功路径才非零。 +type cacheMetricCounters struct { + Dispatches int64 `json:"dispatches"` + Success int64 `json:"success"` + Failure int64 `json:"failure"` + CacheReadInputTokens int64 `json:"cacheReadInputTokens"` + CacheCreationInputTokens int64 `json:"cacheCreationInputTokens"` + InputTokens int64 `json:"inputTokens"` + OutputTokens int64 `json:"outputTokens"` +} + +func (c *cacheMetricCounters) add(o cacheMetricCounters) { + c.Dispatches += o.Dispatches + c.Success += o.Success + c.Failure += o.Failure + c.CacheReadInputTokens += o.CacheReadInputTokens + c.CacheCreationInputTokens += o.CacheCreationInputTokens + c.InputTokens += o.InputTokens + c.OutputTokens += o.OutputTokens +} + +type cacheMetricBucket struct { + minute int64 + total cacheMetricCounters + bySource map[string]*cacheMetricCounters + byModel map[string]*cacheMetricCounters + byAccount map[string]*cacheMetricCounters +} + +func newCacheMetricBucket(minute int64) *cacheMetricBucket { + return &cacheMetricBucket{ + minute: minute, + bySource: make(map[string]*cacheMetricCounters), + byModel: make(map[string]*cacheMetricCounters), + byAccount: make(map[string]*cacheMetricCounters), + } +} + +func addDimension(dim map[string]*cacheMetricCounters, key string, delta cacheMetricCounters) { + if key == "" { + key = "unknown" + } + counter := dim[key] + if counter == nil { + counter = &cacheMetricCounters{} + dim[key] = counter + } + counter.add(delta) +} + +type cacheDispatchMetrics struct { + mu sync.Mutex + buckets map[int64]*cacheMetricBucket + retain time.Duration + // lastHitRatioWarn 节流命中率告警，避免高频 dispatch 下刷屏。 + lastHitRatioWarn time.Time +} + +func newCacheDispatchMetrics(retain time.Duration) *cacheDispatchMetrics { + if retain <= 0 { + retain = cacheMetricsRetain + } + return &cacheDispatchMetrics{ + buckets: make(map[int64]*cacheMetricBucket), + retain: retain, + } +} + +// Record 累加一条 dispatch 诊断。now 由调用方注入以便测试。 +func (m *cacheDispatchMetrics) Record(diag claudeCacheDispatchDiagnostics, now time.Time) { + minute := now.Unix() / 60 + delta := cacheMetricCounters{ + Dispatches: 1, + CacheReadInputTokens: int64(diag.Usage.CacheReadInputTokens), + CacheCreationInputTokens: int64(diag.Usage.CacheCreationInputTokens), + InputTokens: int64(diag.InputTokens), + OutputTokens: int64(diag.OutputTokens), + } + if diag.Outcome == "success" { + delta.Success = 1 + } else { + delta.Failure = 1 + } + + source := diag.AffinitySource + model := diag.Model + accountID := "" + if diag.Account != nil { + accountID = diag.Account.ID + } + + m.mu.Lock() + m.pruneLocked(minute) + + bucket := m.buckets[minute] + if bucket == nil { + bucket = newCacheMetricBucket(minute) + m.buckets[minute] = bucket + } + bucket.total.add(delta) + addDimension(bucket.bySource, source, delta) + addDimension(bucket.byModel, model, delta) + addDimension(bucket.byAccount, accountID, delta) + + // 命中率安全阀：在锁内聚合近窗并判定，把日志 I/O 留到释放锁之后，避免持锁做 I/O。 + warn, ratio, dispatches := m.shouldWarnHitRatioLocked(now) + m.mu.Unlock() + + if warn { + logger.Warnf( + "[CacheHitRatioGuard] 聚合缓存命中率 %.1f%% 超过约定上限 %.0f%%（近 %s，样本 %d 条）：下游补贴可能失控，请核对 /v1/stats 与 ComputeScoped cap/TTL 参数", + ratio*100, cacheHitRatioWarnThreshold*100, cacheHitRatioWarnWindow, dispatches, + ) + } +} + +// shouldWarnHitRatioLocked 在持锁状态下聚合近 cacheHitRatioWarnWindow 窗口，判断 +// 聚合命中率是否超阈值且满足节流与最小样本量。命中口径复用 computeCacheMetricRatios +// （read/(read+create+input) 取 max 分母），与 admin 接口一致，规避净输入做分母的 +// 假数 bug。返回是否告警、命中率、窗口样本量；返回 true 时已就地更新 lastHitRatioWarn。 +func (m *cacheDispatchMetrics) shouldWarnHitRatioLocked(now time.Time) (bool, float64, int64) { + if now.Sub(m.lastHitRatioWarn) < cacheHitRatioWarnInterval { + return false, 0, 0 + } + currentMinute := now.Unix() / 60 + windowMinutes := int64((cacheHitRatioWarnWindow + cacheMetricsBucketSpan - 1) / cacheMetricsBucketSpan) + cutoff := currentMinute - windowMinutes + 1 + + var agg cacheMetricCounters + for minute, bucket := range m.buckets { + if minute < cutoff || minute > currentMinute { + continue + } + agg.add(bucket.total) + } + if agg.Dispatches < cacheHitRatioWarnMinDispatches { + return false, 0, 0 + } + ratio := computeCacheMetricRatios(agg).CacheHitRatio + if ratio <= cacheHitRatioWarnThreshold { + return false, 0, 0 + } + m.lastHitRatioWarn = now + return true, ratio, agg.Dispatches +} + +func (m *cacheDispatchMetrics) pruneLocked(currentMinute int64) { + cutoff := currentMinute - int64(m.retain/time.Minute) + for minute := range m.buckets { + if minute < cutoff { + delete(m.buckets, minute) + } + } +} + +// cacheMetricsWindowSnapshot 是单个时间窗口的聚合结果。 +type cacheMetricsWindowSnapshot struct { + WindowSeconds int64 `json:"windowSeconds"` + Total cacheMetricCounters `json:"total"` + Ratios cacheMetricRatios `json:"ratios"` + BySource map[string]cacheMetricGroup `json:"bySource"` + ByModel map[string]cacheMetricGroup `json:"byModel"` + ByAccount map[string]cacheMetricGroup `json:"byAccount"` +} + +// cacheMetricGroup 是某个分组维度下的计数器加该分组自身的命中率。 +type cacheMetricGroup struct { + cacheMetricCounters + Ratios cacheMetricRatios `json:"ratios"` +} + +// cacheMetricRatios 给出多个常用缓存率口径，避免与客户口径对不上。 +// - CacheReadRatio: cache_read / total_input，真正从缓存读取的输入占比。 +// - CacheHitRatio: (cache_read + cache_creation) / total_input，广义缓存覆盖占比。 +// - BilledInputTokens: total_input - cache_read - cache_creation，实际计费的未缓存输入。 +// - SuccessRate: success / dispatches，分发成功率。 +// +// 分母 total_input 取 max(InputTokens, cache_read + cache_creation)：缓存读写的 +// token 本身就是输入的一部分，总输入不可能小于它们之和。InputTokens 累加的是上游 +// 实报/估算的总输入，与 cache_read/creation（本地 prompt cache 模拟器产出）口径不同， +// 长会话高命中时常出现 InputTokens < cache_read，若直接以 InputTokens 作分母会让 +// 比率突破 1（线上实测 cacheHitRatio 达 1.4）。以两者较大值为分母即可保证比率 ∈[0,1] +// 且与 Anthropic 官方口径 read/(read+create+uncached) 一致。分母为 0 时返回 0。 +type cacheMetricRatios struct { + CacheReadRatio float64 `json:"cacheReadRatio"` + CacheHitRatio float64 `json:"cacheHitRatio"` + BilledInputTokens int64 `json:"billedInputTokens"` + SuccessRate float64 `json:"successRate"` +} + +func computeCacheMetricRatios(c cacheMetricCounters) cacheMetricRatios { + ratios := cacheMetricRatios{} + cached := c.CacheReadInputTokens + c.CacheCreationInputTokens + totalInput := max(c.InputTokens, cached) + ratios.BilledInputTokens = totalInput - cached + if totalInput > 0 { + ratios.CacheReadRatio = float64(c.CacheReadInputTokens) / float64(totalInput) + ratios.CacheHitRatio = float64(cached) / float64(totalInput) + } + if c.Dispatches > 0 { + ratios.SuccessRate = float64(c.Success) / float64(c.Dispatches) + } + return ratios +} + +func dimensionSnapshot(dim map[string]*cacheMetricCounters) map[string]cacheMetricGroup { + out := make(map[string]cacheMetricGroup, len(dim)) + for key, counter := range dim { + out[key] = cacheMetricGroup{ + cacheMetricCounters: *counter, + Ratios: computeCacheMetricRatios(*counter), + } + } + return out +} + +// Snapshot 合并落在 [now-window, now] 内的桶。window 超过 retain 时按 retain 截断。 +func (m *cacheDispatchMetrics) Snapshot(window time.Duration, now time.Time) cacheMetricsWindowSnapshot { + if window <= 0 { + window = cacheMetricsBucketSpan + } + if window > m.retain { + window = m.retain + } + currentMinute := now.Unix() / 60 + // 向上取整覆盖整个窗口：window 内涉及的最早分钟桶。 + windowMinutes := int64((window + cacheMetricsBucketSpan - 1) / cacheMetricsBucketSpan) + cutoff := currentMinute - windowMinutes + 1 + + snapshot := cacheMetricsWindowSnapshot{ + WindowSeconds: int64(window / time.Second), + } + merged := newCacheMetricBucket(0) + + m.mu.Lock() + for minute, bucket := range m.buckets { + if minute < cutoff || minute > currentMinute { + continue + } + merged.total.add(bucket.total) + for k, v := range bucket.bySource { + addDimension(merged.bySource, k, *v) + } + for k, v := range bucket.byModel { + addDimension(merged.byModel, k, *v) + } + for k, v := range bucket.byAccount { + addDimension(merged.byAccount, k, *v) + } + } + m.mu.Unlock() + + snapshot.Total = merged.total + snapshot.Ratios = computeCacheMetricRatios(merged.total) + snapshot.BySource = dimensionSnapshot(merged.bySource) + snapshot.ByModel = dimensionSnapshot(merged.byModel) + snapshot.ByAccount = dimensionSnapshot(merged.byAccount) + return snapshot +} + +// Reset 清空所有桶，供测试使用。 +func (m *cacheDispatchMetrics) Reset() { + m.mu.Lock() + m.buckets = make(map[int64]*cacheMetricBucket) + m.mu.Unlock() +} + +// 包级单例：诊断汇聚点为包函数（非 Handler 方法，且被测试直接调用），用单例可零侵入 +// 接入而无需改各调用方签名。 +var defaultCacheDispatchMetrics = newCacheDispatchMetrics(cacheMetricsRetain) + +func recordCacheDispatchMetric(diag claudeCacheDispatchDiagnostics) { + defaultCacheDispatchMetrics.Record(diag, time.Now()) +} + +// recordClaudeCacheDispatch 是分发完成处的汇聚钩子，把一次 Claude 分发（成功/失败、 +// 流式/非流式）喂给窗口化指标。失败路径 token 全部传 0：调用未真正完成，没有 token +// 被读取/生成，喂入模拟的 cache 值会让命中率虚高。source 用 "stream"/"nonstream"， +// 给运维一个真实维度（流式与非流式的缓存覆盖是否不同）；本仓库无 cache affinity， +// 故不复用 tutu 的 affinity source 分类。 +func recordClaudeCacheDispatch(outcome, model, source string, account *config.Account, usage promptCacheUsage, inputTokens, outputTokens int) { + recordCacheDispatchMetric(claudeCacheDispatchDiagnostics{ + Outcome: outcome, + Model: model, + AffinitySource: source, + Account: account, + Usage: usage, + InputTokens: inputTokens, + OutputTokens: outputTokens, + }) +} + +// cacheMetricsWindowDurations 是 /v1/stats 默认输出的窗口口径。 +var cacheMetricsWindowDurations = []struct { + Label string + Duration time.Duration +}{ + {"5m", 5 * time.Minute}, + {"15m", 15 * time.Minute}, + {"1h", time.Hour}, +} + +func cacheMetricsSnapshotByWindows(now time.Time) map[string]cacheMetricsWindowSnapshot { + out := make(map[string]cacheMetricsWindowSnapshot, len(cacheMetricsWindowDurations)) + for _, w := range cacheMetricsWindowDurations { + out[w.Label] = defaultCacheDispatchMetrics.Snapshot(w.Duration, now) + } + return out +} diff --git a/proxy/cache_metrics_test.go b/proxy/cache_metrics_test.go new file mode 100644 index 00000000..0005a873 --- /dev/null +++ b/proxy/cache_metrics_test.go @@ -0,0 +1,239 @@ +package proxy + +import ( + "testing" + "time" + + "kiro-go/config" +) + +func dispatch(source, model, accountID, outcome string, read, creation, input, output int) claudeCacheDispatchDiagnostics { + diag := claudeCacheDispatchDiagnostics{ + Outcome: outcome, + Model: model, + AffinitySource: source, + Usage: promptCacheUsage{ + CacheReadInputTokens: read, + CacheCreationInputTokens: creation, + }, + InputTokens: input, + OutputTokens: output, + } + if accountID != "" { + diag.Account = &config.Account{ID: accountID} + } + return diag +} + +func TestCacheDispatchMetricsAggregatesTotalsAndRatios(t *testing.T) { + m := newCacheDispatchMetrics(cacheMetricsRetain) + now := time.Unix(1_000_000, 0) + + // 两条成功 dispatch，输入总量合计 1000，其中 cache_read 600、creation 200。 + m.Record(dispatch("conversation", "claude-sonnet-4.5", "acct-1", "success", 300, 100, 500, 50), now) + m.Record(dispatch("conversation", "claude-sonnet-4.5", "acct-1", "success", 300, 100, 500, 50), now) + // 一条失败 dispatch（失败路径 token 为 0）。 + m.Record(dispatch("conversation", "claude-sonnet-4.5", "acct-1", "failure", 0, 0, 0, 0), now) + + snap := m.Snapshot(5*time.Minute, now) + if snap.Total.Dispatches != 3 { + t.Fatalf("dispatches = %d, want 3", snap.Total.Dispatches) + } + if snap.Total.Success != 2 || snap.Total.Failure != 1 { + t.Fatalf("success/failure = %d/%d, want 2/1", snap.Total.Success, snap.Total.Failure) + } + if snap.Total.InputTokens != 1000 || snap.Total.CacheReadInputTokens != 600 || snap.Total.CacheCreationInputTokens != 200 { + t.Fatalf("token totals = in:%d read:%d creation:%d, want 1000/600/200", + snap.Total.InputTokens, snap.Total.CacheReadInputTokens, snap.Total.CacheCreationInputTokens) + } + // cache_read / total_input = 600/1000 = 0.6 + if got := snap.Ratios.CacheReadRatio; got < 0.5999 || got > 0.6001 { + t.Fatalf("cacheReadRatio = %v, want 0.6", got) + } + // (read + creation) / total = 800/1000 = 0.8 + if got := snap.Ratios.CacheHitRatio; got < 0.7999 || got > 0.8001 { + t.Fatalf("cacheHitRatio = %v, want 0.8", got) + } + // billed = 1000 - 600 - 200 = 200 + if snap.Ratios.BilledInputTokens != 200 { + t.Fatalf("billedInputTokens = %d, want 200", snap.Ratios.BilledInputTokens) + } + // success rate = 2/3 + if got := snap.Ratios.SuccessRate; got < 0.6665 || got > 0.6668 { + t.Fatalf("successRate = %v, want ~0.6667", got) + } +} + +func TestCacheDispatchMetricsGroupsByDimension(t *testing.T) { + m := newCacheDispatchMetrics(cacheMetricsRetain) + now := time.Unix(2_000_000, 0) + + m.Record(dispatch("metadata-conversation", "claude-opus", "acct-A", "success", 100, 0, 200, 10), now) + m.Record(dispatch("cache-prefix", "claude-sonnet-4.5", "acct-B", "success", 50, 0, 100, 5), now) + + snap := m.Snapshot(time.Hour, now) + if len(snap.BySource) != 2 { + t.Fatalf("bySource keys = %d, want 2 (%v)", len(snap.BySource), snap.BySource) + } + if g, ok := snap.BySource["metadata-conversation"]; !ok || g.InputTokens != 200 { + t.Fatalf("metadata-conversation group missing or wrong: %+v ok=%v", g, ok) + } + if g, ok := snap.ByModel["claude-opus"]; !ok || g.CacheReadInputTokens != 100 { + t.Fatalf("claude-opus group missing or wrong: %+v ok=%v", g, ok) + } + if g, ok := snap.ByAccount["acct-B"]; !ok || g.InputTokens != 100 { + t.Fatalf("acct-B group missing or wrong: %+v ok=%v", g, ok) + } +} + +func TestCacheDispatchMetricsWindowExcludesOldBuckets(t *testing.T) { + m := newCacheDispatchMetrics(cacheMetricsRetain) + now := time.Unix(3_000_000, 0) + + // 10 分钟前的一条记录，落在 5m 窗口外、1h 窗口内。 + m.Record(dispatch("conversation", "m", "a", "success", 10, 0, 20, 1), now.Add(-10*time.Minute)) + // 当前一条记录。 + m.Record(dispatch("conversation", "m", "a", "success", 30, 0, 60, 2), now) + + short := m.Snapshot(5*time.Minute, now) + if short.Total.Dispatches != 1 || short.Total.InputTokens != 60 { + t.Fatalf("5m window = dispatches:%d input:%d, want 1/60", short.Total.Dispatches, short.Total.InputTokens) + } + + long := m.Snapshot(time.Hour, now) + if long.Total.Dispatches != 2 || long.Total.InputTokens != 80 { + t.Fatalf("1h window = dispatches:%d input:%d, want 2/80", long.Total.Dispatches, long.Total.InputTokens) + } +} + +func TestCacheDispatchMetricsPrunesBeyondRetain(t *testing.T) { + m := newCacheDispatchMetrics(cacheMetricsRetain) + now := time.Unix(4_000_000, 0) + + m.Record(dispatch("conversation", "m", "a", "success", 1, 0, 2, 1), now.Add(-2*time.Hour)) + // 一条新记录触发对超过 retain 的旧桶的惰性清理。 + m.Record(dispatch("conversation", "m", "a", "success", 1, 0, 2, 1), now) + + m.mu.Lock() + buckets := len(m.buckets) + m.mu.Unlock() + if buckets != 1 { + t.Fatalf("buckets after prune = %d, want 1", buckets) + } +} + +func TestCacheDispatchMetricsReset(t *testing.T) { + m := newCacheDispatchMetrics(cacheMetricsRetain) + now := time.Unix(5_000_000, 0) + m.Record(dispatch("conversation", "m", "a", "success", 1, 0, 2, 1), now) + m.Reset() + snap := m.Snapshot(time.Hour, now) + if snap.Total.Dispatches != 0 { + t.Fatalf("dispatches after reset = %d, want 0", snap.Total.Dispatches) + } +} + +func TestCacheDispatchMetricsEmptyRatios(t *testing.T) { + m := newCacheDispatchMetrics(cacheMetricsRetain) + now := time.Unix(6_000_000, 0) + snap := m.Snapshot(time.Hour, now) + if snap.Ratios.CacheReadRatio != 0 || snap.Ratios.CacheHitRatio != 0 || snap.Ratios.SuccessRate != 0 { + t.Fatalf("empty ratios should be 0, got %+v", snap.Ratios) + } + if snap.Ratios.BilledInputTokens != 0 { + t.Fatalf("empty billed should be 0, got %d", snap.Ratios.BilledInputTokens) + } +} + +// cache_read/creation 来自本地 prompt cache 模拟器，InputTokens 来自上游实报/估算， +// 两者口径不同：长会话高命中时 cache_read 常超过 InputTokens。此时命中率必须仍 +// ≤ 1（分母取 max(input, read+creation)），否则面板会出现 >100% 的假命中率。 +// 复现线上实测 cacheHitRatio=1.4 的回归。 +func TestCacheDispatchMetricsRatiosNeverExceedOneWhenCachedExceedsInput(t *testing.T) { + m := newCacheDispatchMetrics(cacheMetricsRetain) + now := time.Unix(7_000_000, 0) + + // read=205000、creation=49000、input(净/估算)=181000 —— read 已超过 input。 + m.Record(dispatch("agent-prefix", "claude-opus-4.6", "acct-1", "success", 205000, 49000, 181000, 100), now) + + snap := m.Snapshot(5*time.Minute, now) + if r := snap.Ratios.CacheHitRatio; r > 1.0 { + t.Fatalf("cacheHitRatio = %v, must be <= 1", r) + } + if r := snap.Ratios.CacheReadRatio; r > 1.0 { + t.Fatalf("cacheReadRatio = %v, must be <= 1", r) + } + // 分母 = max(181000, 254000) = 254000；hit = 254000/254000 = 1.0；read = 205000/254000 ≈ 0.807。 + if r := snap.Ratios.CacheHitRatio; r < 0.9999 || r > 1.0001 { + t.Fatalf("cacheHitRatio = %v, want ~1.0", r) + } + if r := snap.Ratios.CacheReadRatio; r < 0.8070 || r > 0.8072 { + t.Fatalf("cacheReadRatio = %v, want ~0.8071", r) + } + // billed = 254000 - 254000 = 0（净输入被缓存量覆盖，不可为负）。 + if snap.Ratios.BilledInputTokens != 0 { + t.Fatalf("billedInputTokens = %d, want 0", snap.Ratios.BilledInputTokens) + } +} + +// TestCacheHitRatioGuardWarnsAboveThreshold 验证命中率安全阀：滚动窗口聚合命中率 +// 超过 cacheHitRatioWarnThreshold 且样本量达标时触发告警，并受节流与最小样本量保护。 +// 直接测纯判定函数 shouldWarnHitRatioLocked（无日志 I/O），覆盖阈值/样本量/节流三条。 +func TestCacheHitRatioGuardWarnsAboveThreshold(t *testing.T) { + m := newCacheDispatchMetrics(cacheMetricsRetain) + now := time.Unix(8_000_000, 0) + + // 喂入命中率 ~95% 的高命中样本，数量越过最小样本量门槛。每条 read=950、 + // creation=0、input=50 → 分母 max(50, 950)=950，hit=950/950=1.0（单条）， + // 聚合命中率远超 0.88。 + for i := 0; i < cacheHitRatioWarnMinDispatches+5; i++ { + m.Record(dispatch("agent-prefix", "claude-opus-4.6", "acct-1", "success", 950, 0, 50, 100), now) + } + + // 安全阀应已在某条 Record 中触发并更新 lastHitRatioWarn（节流时间戳被设为 now）。 + if m.lastHitRatioWarn.IsZero() { + t.Fatalf("高命中率 + 足量样本应已触发命中率告警，但 lastHitRatioWarn 未设置") + } + + // 节流：紧接着同一时刻再判定应被抑制（距上次告警 < cacheHitRatioWarnInterval）。 + m.mu.Lock() + warnAgain, _, _ := m.shouldWarnHitRatioLocked(now) + m.mu.Unlock() + if warnAgain { + t.Fatalf("节流失效：同一时刻不应重复告警") + } + + // 过了节流间隔后应能再次告警（命中率仍高）。 + later := now.Add(cacheHitRatioWarnInterval + time.Second) + m.mu.Lock() + warnLater, ratioLater, _ := m.shouldWarnHitRatioLocked(later) + m.mu.Unlock() + if !warnLater { + t.Fatalf("过节流间隔后命中率仍高，应再次告警，got ratio=%.3f", ratioLater) + } +} + +// TestCacheHitRatioGuardSilentBelowThresholdAndSmallSample 验证安全阀不误报： +// 命中率低于阈值、或样本量不足时都不告警。 +func TestCacheHitRatioGuardSilentBelowThresholdAndSmallSample(t *testing.T) { + // 场景一：足量样本但命中率低（~40%），不应告警。 + low := newCacheDispatchMetrics(cacheMetricsRetain) + now := time.Unix(8_100_000, 0) + for i := 0; i < cacheHitRatioWarnMinDispatches+5; i++ { + // read=300、creation=0、input=700 → 分母 max(700,300)=700，hit≈0.43，低于 0.88。 + low.Record(dispatch("agent-prefix", "claude-opus-4.6", "acct-1", "success", 300, 0, 700, 100), now) + } + if !low.lastHitRatioWarn.IsZero() { + t.Fatalf("命中率低于阈值不应告警，但 lastHitRatioWarn 被设置") + } + + // 场景二：命中率高但样本量不足（< 最小样本量），不应告警。 + small := newCacheDispatchMetrics(cacheMetricsRetain) + now2 := time.Unix(8_200_000, 0) + for i := 0; i < cacheHitRatioWarnMinDispatches-1; i++ { + small.Record(dispatch("agent-prefix", "claude-opus-4.6", "acct-1", "success", 950, 0, 50, 100), now2) + } + if !small.lastHitRatioWarn.IsZero() { + t.Fatalf("样本量不足不应告警，但 lastHitRatioWarn 被设置") + } +} diff --git a/proxy/cache_tracker.go b/proxy/cache_tracker.go index 024e4d39..f8b4cc50 100644 --- a/proxy/cache_tracker.go +++ b/proxy/cache_tracker.go @@ -2,12 +2,17 @@ package proxy import ( "bytes" + "container/list" "crypto/sha256" "encoding/json" + "kiro-go/config" + "os" + "path/filepath" "sort" "strconv" "strings" "sync" + "sync/atomic" "time" ) @@ -18,7 +23,13 @@ const defaultPromptCacheTTL = 5 * time.Minute // matching and storage to avoid reporting unrealistic 100% cache hits on // short requests. const defaultMinCacheableTokens = 1024 -const opusMinCacheableTokens = 4096 + +// opusMinCacheableTokens is the minimum cacheable prefix size for Opus models. +// Anthropic's real prompt-cache threshold is ~1024 tokens (Claude 4 family); +// the previous 4096 excluded many valid Opus breakpoints, so an Opus +// conversation would report far fewer cache hits than it actually earned. +// Ported from kiro-tutu (opusMinCacheableTokens = 1024). +const opusMinCacheableTokens = 1024 type promptCacheUsage struct { CacheCreationInputTokens int @@ -50,22 +61,249 @@ func minCacheableTokensForModel(model string) int { type promptCacheEntry struct { ExpiresAt time.Time TTL time.Duration + lruElem *list.Element // back-ref into t.order; Value = fingerprint [32]byte } type promptCacheTracker struct { - mu sync.Mutex - entriesByAccount map[string]map[[32]byte]promptCacheEntry - maxSupportedTTL time.Duration + mu sync.Mutex + entries map[[32]byte]*promptCacheEntry + order *list.List // front = most-recently-used; Element.Value = [32]byte fingerprint + maxEntries int + maxSupportedTTL time.Duration + hits int64 // atomic — Compute calls returning CacheReadInputTokens > 0 + misses int64 // atomic — Compute calls returning CacheReadInputTokens == 0 + evictions int64 // atomic — LRU pop-backs in evictOverflowLocked + expirations int64 // atomic — TTL removals in pruneExpiredLocked + dirty bool + stopChan chan struct{} + stopOnce sync.Once } func newPromptCacheTracker(maxTTL time.Duration) *promptCacheTracker { + return newPromptCacheTrackerWithCapacity(maxTTL, config.GetPromptCacheMaxEntries()) +} + +func newPromptCacheTrackerWithCapacity(maxTTL time.Duration, maxEntries int) *promptCacheTracker { if maxTTL <= 0 { maxTTL = defaultPromptCacheTTL } return &promptCacheTracker{ - entriesByAccount: make(map[string]map[[32]byte]promptCacheEntry), - maxSupportedTTL: maxTTL, + entries: make(map[[32]byte]*promptCacheEntry), + order: list.New(), + maxEntries: maxEntries, + maxSupportedTTL: maxTTL, + } +} + +// on-disk format for prompt-cache persistence (C3). +type promptCacheEntryOnDisk struct { + Fingerprint [32]byte + ExpiresAt int64 // unix seconds + TTLSeconds int64 +} + +// Load reads persisted cache entries from path. Entries already expired (by the +// time load finishes) are dropped. Best-effort: a corrupt/missing file is not +// fatal — the tracker starts empty, same as the pre-C3 behavior. +func (t *promptCacheTracker) Load(path string) { + data, err := os.ReadFile(path) + if err != nil { + return // missing file = fresh start (normal on first run) + } + var disk struct { + Version int `json:"version"` + Entries []promptCacheEntryOnDisk `json:"entries"` + } + if json.Unmarshal(data, &disk) != nil { + return // corrupt = fresh start + } + now := time.Now() + t.mu.Lock() + defer t.mu.Unlock() + for _, e := range disk.Entries { + exp := time.Unix(e.ExpiresAt, 0) + if !exp.After(now) { + continue // already expired + } + t.putLocked(e.Fingerprint, exp, time.Duration(e.TTLSeconds)*time.Second) + } +} + +// startSaveLoop launches a background goroutine that flushes the cache to path +// every flushInterval (if dirty). Call once after Load. The goroutine exits when +// stopChan is closed. +func (t *promptCacheTracker) startSaveLoop(path string, flushInterval time.Duration) { + t.stopChan = make(chan struct{}) + go func() { + ticker := time.NewTicker(flushInterval) + defer ticker.Stop() + for { + select { + case <-ticker.C: + t.flush(path) + case <-t.stopChan: + t.flush(path) // final flush + return + } + } + }() +} + +func (t *promptCacheTracker) Stop() { + // Idempotent: a second Stop (e.g. test cleanup + main shutdown) would + // otherwise close an already-closed channel and panic. + t.stopOnce.Do(func() { + if t.stopChan != nil { + close(t.stopChan) + } + }) +} + +func (t *promptCacheTracker) flush(path string) { + t.mu.Lock() + if !t.dirty { + t.mu.Unlock() + return + } + now := time.Now() + entries := make([]promptCacheEntryOnDisk, 0, len(t.entries)) + for fp, e := range t.entries { + if !e.ExpiresAt.After(now) { + continue + } + entries = append(entries, promptCacheEntryOnDisk{ + Fingerprint: fp, + ExpiresAt: e.ExpiresAt.Unix(), + TTLSeconds: int64(e.TTL.Seconds()), + }) + } + t.dirty = false + t.mu.Unlock() + + data, _ := json.MarshalIndent(map[string]interface{}{ + "version": 1, + "entries": entries, + }, "", " ") + _ = writeFileAtomic(path, data) +} + +// writeFileAtomic writes data to path via a temp file in the same directory +// then renames it into place, so a crash mid-flush cannot leave a truncated +// state file (which would silently reset the cache on the next Load). Mirrors +// the Python shim's tmp+os.replace path. os.Rename is atomic on the same +// filesystem; the same-dir temp guarantees same-fs. Best-effort: a rename +// failure falls back to a direct write (matching pre-existing behavior). +func writeFileAtomic(path string, data []byte) error { + if dir := filepath.Dir(path); dir != "" && dir != "." { + _ = os.MkdirAll(dir, 0o700) + } + tmp, err := os.CreateTemp(dirOf(path), ".prompt-cache-*.tmp") + if err != nil { + return err + } + tmpName := tmp.Name() + closed := false + defer func() { + if !closed { + tmp.Close() + } + _ = os.Remove(tmpName) + }() + if _, err := tmp.Write(data); err != nil { + return err + } + if err := tmp.Sync(); err != nil { + return err + } + if err := tmp.Close(); err != nil { + return err + } + closed = true + _ = os.Chmod(tmpName, 0o600) + if err := os.Rename(tmpName, path); err != nil { + return os.WriteFile(path, data, 0o600) + } + return nil +} + +func dirOf(path string) string { + if dir := filepath.Dir(path); dir != "" { + return dir + } + return "." +} + +// interfaceHasCacheControl reports whether the value tree contains any +// "cache_control" key on a dict node. Mirrors the Python shim's has_cache_control: +// cache_control can nest inside system blocks or message content blocks. +func interfaceHasCacheControl(v interface{}) bool { + switch val := v.(type) { + case map[string]interface{}: + if _, ok := val["cache_control"]; ok { + return true + } + for _, x := range val { + if interfaceHasCacheControl(x) { + return true + } + } + case []interface{}: + for _, x := range val { + if interfaceHasCacheControl(x) { + return true + } + } } + return false +} + +// claudeRequestHasCacheControl reports whether any cache_control marker appears +// in the request's system or message content. ClaudeTool is a typed struct with +// no cache_control field, so only system/messages can carry one. Used as the +// BuildClaudeProfile fast-path guard: no marker → no breakpoints possible → nil. +func claudeRequestHasCacheControl(req *ClaudeRequest) bool { + if req == nil { + return false + } + if interfaceHasCacheControl(req.System) { + return true + } + for _, m := range req.Messages { + if interfaceHasCacheControl(m.Content) { + return true + } + } + return false +} + +// detectMaxTTL returns the largest cache_control TTL present anywhere in the +// request (1h wins over 5m), defaulting to defaultPromptCacheTTL when no +// cache_control marker exists. Mirrors Rust detect_max_ttl: it seeds the +// auto-prefix activeTTL so a request WITHOUT explicit cache_control still +// produces message-end breakpoints at the default TTL — reproducing Anthropic's +// automatic prompt caching (a stable history prefix hits across turns even +// though the client never set cache_control). ClaudeTool is a typed struct with +// no cache_control field, so only system/message blocks can contribute. +func detectMaxTTL(req *ClaudeRequest) time.Duration { + max := defaultPromptCacheTTL + bump := func(value interface{}) { + if normalized := normalizePromptCacheTTL(extractPromptCacheTTL(value)); normalized > max { + max = normalized + } + } + if arr, ok := req.System.([]interface{}); ok { + for _, block := range arr { + bump(block) + } + } + for _, msg := range req.Messages { + if arr, ok := msg.Content.([]interface{}); ok { + for _, block := range arr { + bump(block) + } + } + } + return max } func (t *promptCacheTracker) BuildClaudeProfile(req *ClaudeRequest, totalInputTokens int) *promptCacheProfile { @@ -77,7 +315,9 @@ func (t *promptCacheTracker) BuildClaudeProfile(req *ClaudeRequest, totalInputTo hasher := sha256.New() breakpoints := make([]promptCacheBreakpoint, 0) cumulativeTokens := 0 - var activeTTL time.Duration + // Auto-prefix: seed activeTTL with detectMaxTTL(req) so message-end + // breakpoints fire even with no explicit cache_control (see detectMaxTTL). + activeTTL := detectMaxTTL(req) for _, block := range blocks { canonical := canonicalizeCacheValue(block.Value) @@ -125,10 +365,17 @@ func (t *promptCacheTracker) BuildClaudeProfile(req *ClaudeRequest, totalInputTo } } -func (t *promptCacheTracker) Compute(accountID string, profile *promptCacheProfile) promptCacheUsage { +func (t *promptCacheTracker) Compute(accountID string, profile *promptCacheProfile) (u promptCacheUsage) { if t == nil || profile == nil || len(profile.Breakpoints) == 0 || accountID == "" { return promptCacheUsage{} } + defer func() { + if u.CacheReadInputTokens > 0 { + atomic.AddInt64(&t.hits, 1) + } else { + atomic.AddInt64(&t.misses, 1) + } + }() minTokens := minCacheableTokensForModel(profile.Model) last := profile.Breakpoints[len(profile.Breakpoints)-1] @@ -139,14 +386,14 @@ func (t *promptCacheTracker) Compute(accountID string, profile *promptCacheProfi defer t.mu.Unlock() t.pruneExpiredLocked(now) - entries := t.entriesByAccount[accountID] - if len(entries) == 0 { + if len(t.entries) == 0 { // First request for this account: report creation only if above threshold. effectiveCreation := lastTokens if effectiveCreation < minTokens { effectiveCreation = 0 } cache5m, cache1h := computePromptCacheTTLBreakdown(profile, 0) + cache5m, cache1h = clampCacheBreakdownToCreation(cache5m, cache1h, effectiveCreation) return promptCacheUsage{ CacheCreationInputTokens: effectiveCreation, CacheReadInputTokens: 0, @@ -158,7 +405,7 @@ func (t *promptCacheTracker) Compute(accountID string, profile *promptCacheProfi // Cap cacheable tokens at 85% of total input to ensure a realistic // uncached portion. The newest content in a request is never fully // served from cache on the current turn. - maxCacheable := int(float64(profile.TotalInputTokens) * 0.85) + maxCacheable := int(float64(profile.TotalInputTokens) * config.GetPromptCacheMaxRatio()) if lastTokens > maxCacheable { lastTokens = maxCacheable } @@ -170,12 +417,13 @@ func (t *promptCacheTracker) Compute(accountID string, profile *promptCacheProfi if breakpoint.CumulativeTokens < minTokens { continue } - entry, ok := entries[breakpoint.Fingerprint] + entry, ok := t.entries[breakpoint.Fingerprint] if !ok || entry.ExpiresAt.Before(now) { continue } entry.ExpiresAt = now.Add(entry.TTL) - entries[breakpoint.Fingerprint] = entry + t.order.MoveToFront(entry.lruElem) + t.dirty = true // hit extends TTL — persist so a flush before the next Update doesn't lose it matchedTokens = minInt(breakpoint.CumulativeTokens, profile.TotalInputTokens) if matchedTokens > lastTokens { matchedTokens = lastTokens @@ -185,6 +433,7 @@ func (t *promptCacheTracker) Compute(accountID string, profile *promptCacheProfi creation := maxInt(lastTokens-matchedTokens, 0) cache5m, cache1h := computePromptCacheTTLBreakdown(profile, matchedTokens) + cache5m, cache1h = clampCacheBreakdownToCreation(cache5m, cache1h, creation) return promptCacheUsage{ CacheCreationInputTokens: creation, CacheReadInputTokens: matchedTokens, @@ -204,34 +453,128 @@ func (t *promptCacheTracker) Update(accountID string, profile *promptCacheProfil defer t.mu.Unlock() t.pruneExpiredLocked(now) - entries := t.entriesByAccount[accountID] - if entries == nil { - entries = make(map[[32]byte]promptCacheEntry) - t.entriesByAccount[accountID] = entries - } - + // entries is the global map now (C1: cross-account sharing). for _, breakpoint := range profile.Breakpoints { // Skip breakpoints below the minimum cacheable token threshold. if breakpoint.CumulativeTokens < minTokens { continue } - entries[breakpoint.Fingerprint] = promptCacheEntry{ - ExpiresAt: now.Add(breakpoint.TTL), - TTL: breakpoint.TTL, - } + t.putLocked(breakpoint.Fingerprint, now.Add(breakpoint.TTL), breakpoint.TTL) } + t.dirty = true + t.evictOverflowLocked() } func (t *promptCacheTracker) pruneExpiredLocked(now time.Time) { - for accountID, entries := range t.entriesByAccount { - for fingerprint, entry := range entries { - if !entry.ExpiresAt.After(now) { - delete(entries, fingerprint) - } + for fingerprint, entry := range t.entries { + if !entry.ExpiresAt.After(now) { + t.order.Remove(entry.lruElem) + delete(t.entries, fingerprint) + atomic.AddInt64(&t.expirations, 1) + } + } +} + +// putLocked inserts a fingerprint or refreshes its existing entry, marking it +// most-recently-used. Caller holds t.mu. +func (t *promptCacheTracker) putLocked(fp [32]byte, expiresAt time.Time, ttl time.Duration) { + if e, ok := t.entries[fp]; ok { + e.ExpiresAt = expiresAt + e.TTL = ttl + t.order.MoveToFront(e.lruElem) + return + } + elem := t.order.PushFront(fp) + t.entries[fp] = &promptCacheEntry{ExpiresAt: expiresAt, TTL: ttl, lruElem: elem} +} + +// evictOverflowLocked bounds the entries map to maxEntries by evicting the +// least-recently-used entries (the back of the order list). O(1) per eviction. +// Caller holds t.mu. +func (t *promptCacheTracker) evictOverflowLocked() { + for len(t.entries) > t.maxEntries { + back := t.order.Back() + if back == nil { + return } - if len(entries) == 0 { - delete(t.entriesByAccount, accountID) + fp := back.Value.([32]byte) + t.order.Remove(back) + delete(t.entries, fp) + atomic.AddInt64(&t.evictions, 1) + } +} + +// PromptCacheStats is a point-in-time snapshot of cache counters, surfaced via +// /v1/stats. All counters are cumulative since tracker construction. +type PromptCacheStats struct { + Entries int `json:"entries"` + Capacity int `json:"capacity"` + Hits int64 `json:"hits"` + Misses int64 `json:"misses"` + Evictions int64 `json:"evictions"` + Expirations int64 `json:"expirations"` +} + +func (t *promptCacheTracker) Stats() PromptCacheStats { + if t == nil { + return PromptCacheStats{} + } + t.mu.Lock() + entries := len(t.entries) + capacity := t.maxEntries + t.mu.Unlock() + return PromptCacheStats{ + Entries: entries, + Capacity: capacity, + Hits: atomic.LoadInt64(&t.hits), + Misses: atomic.LoadInt64(&t.misses), + Evictions: atomic.LoadInt64(&t.evictions), + Expirations: atomic.LoadInt64(&t.expirations), + } +} + +// splitAgainstTotal rescales an estimate-domain cache split onto a real input +// total, preserving input + creation + read == realTotal. The cache-covered +// fraction of the prompt (in estimate units) is applied to realTotal, then +// divided into read vs creation by their estimate-domain ratio; the 5m/1h +// breakdown is scaled to the new creation total. +func (u promptCacheUsage) splitAgainstTotal(estTotal, realTotal int) promptCacheUsage { + if realTotal <= 0 { + return u + } + coveredEst := u.CacheCreationInputTokens + u.CacheReadInputTokens + if coveredEst <= 0 || estTotal <= 0 { + return promptCacheUsage{} // no cache coverage → all fresh input + } + ratio := float64(coveredEst) / float64(estTotal) + if ratio > 1 { + ratio = 1 + } + cacheTotal := int(float64(realTotal)*ratio + 0.5) + if cacheTotal > realTotal { + cacheTotal = realTotal + } + read := int(float64(cacheTotal)*float64(u.CacheReadInputTokens)/float64(coveredEst) + 0.5) + if read > cacheTotal { + read = cacheTotal + } + if read < 0 { + read = 0 + } + creation := cacheTotal - read + cache5m, cache1h := creation, 0 + if u.CacheCreationInputTokens > 0 { + cache1h = int(float64(u.CacheCreation1hInputTokens)*float64(creation)/float64(u.CacheCreationInputTokens) + 0.5) + if cache1h > creation { + cache1h = creation } + cache5m = creation - cache1h + } + return promptCacheUsage{ + CacheCreationInputTokens: creation, + CacheReadInputTokens: read, + CacheCreation5mInputTokens: cache5m, + CacheCreation1hInputTokens: cache1h, } } @@ -295,11 +638,32 @@ func appendSystemCacheBlocks(blocks *[]cacheablePromptBlock, system interface{}) }, }, false) case []interface{}: + // Structural skip (mirrors Rust extract_segments): Claude Code injects a + // per-turn DYNAMIC block at system[0] WITHOUT cache_control, followed by + // the stable cached block(s). Fingerprinting from that volatile head would + // drift every downstream prefix fingerprint across turns, collapsing + // cross-turn cache hits to zero. So skip every leading system block that + // lacks cache_control — accumulate from the first block carrying one. If + // no block has cache_control, skipUntil stays 0 (include all, matching + // Rust's unwrap_or(0)). system_index is a stripped position key, so the + // shift in index values does not affect fingerprints. The + // isAnthropicBillingHeaderBlock drop inside appendPromptBlock stays as a + // backstop for a non-leading billing-header block. + skipUntil := 0 + found := false for i, block := range v { + if interfaceHasCacheControl(block) { + skipUntil = i + found = true + break + } + } + _ = found + for i := skipUntil; i < len(v); i++ { appendPromptBlock(blocks, map[string]interface{}{ "kind": "system", "system_index": i, - "block": block, + "block": v[i], }, false) } case []string: @@ -506,6 +870,26 @@ func computePromptCacheTTLBreakdown(profile *promptCacheProfile, matchedTokens i return cache5m, cache1h } +// clampCacheBreakdownToCreation scales the 5m/1h cache-creation split down to +// creation when the raw breakpoint deltas (uncapped by the 0.85 cacheable ratio) +// exceed it, preserving the 1h:5m ratio. Guarantees the Anthropic invariant +// cache_creation_input_tokens == ephemeral_5m + ephemeral_1h. +func clampCacheBreakdownToCreation(cache5m, cache1h, creation int) (int, int) { + total := cache5m + cache1h + if total <= creation || total <= 0 { + return cache5m, cache1h + } + scale := float64(creation) / float64(total) + one := int(float64(cache1h)*scale + 0.5) + if one > creation { + one = creation + } + if one < 0 { + one = 0 + } + return creation - one, one +} + func billedClaudeInputTokens(inputTokens int, usage promptCacheUsage) int { return maxInt(inputTokens-usage.CacheCreationInputTokens-usage.CacheReadInputTokens, 0) } diff --git a/proxy/cache_tracker_hardening_test.go b/proxy/cache_tracker_hardening_test.go new file mode 100644 index 00000000..a8eaabf9 --- /dev/null +++ b/proxy/cache_tracker_hardening_test.go @@ -0,0 +1,235 @@ +package proxy + +import ( + "crypto/sha256" + "encoding/json" + "os" + "strings" + "testing" + "time" +) + +// TestPromptCacheLRUEvictsOldestUnused verifies the list-based LRU: with cap 3, +// after inserting a,b,c, touching a, then inserting d, the evicted entry is b +// (the least-recently-used), not a. +func TestPromptCacheLRUEvictsOldestUnused(t *testing.T) { + tr := newPromptCacheTrackerWithCapacity(time.Hour, 3) + now := time.Now() + + tr.mu.Lock() + tr.putLocked([32]byte{1}, now.Add(time.Hour), time.Hour) // a + tr.putLocked([32]byte{2}, now.Add(time.Hour), time.Hour) // b + tr.putLocked([32]byte{3}, now.Add(time.Hour), time.Hour) // c + tr.putLocked([32]byte{1}, now.Add(time.Hour), time.Hour) // touch a → front + tr.putLocked([32]byte{4}, now.Add(time.Hour), time.Hour) // d + tr.evictOverflowLocked() // cap 3 → evict back (b) + tr.mu.Unlock() + + if _, ok := tr.entries[[32]byte{2}]; ok { + t.Fatalf("expected least-recently-used entry (b) to be evicted") + } + for _, want := range [][32]byte{{1}, {3}, {4}} { + if _, ok := tr.entries[want]; !ok { + t.Fatalf("expected entry %v to survive", want) + } + } + if got := len(tr.entries); got != 3 { + t.Fatalf("expected cap=3 after eviction, got %d", got) + } +} + +// TestSplitAgainstTotalAnchorsToRealTotal verifies the cache split is rescaled +// onto the real upstream input total while preserving the accounting identity +// input + creation + read == realTotal. +func TestSplitAgainstTotalAnchorsToRealTotal(t *testing.T) { + // Estimate domain: covered = 800 (creation 500 + read 300), estTotal = 1000. + u := promptCacheUsage{ + CacheCreationInputTokens: 500, + CacheReadInputTokens: 300, + CacheCreation5mInputTokens: 500, + } + got := u.splitAgainstTotal(1000, 2000) + + // ratio = 800/1000 = 0.8 → cacheTotal = 1600; read = 1600*300/800 = 600; creation = 1000. + if got.CacheReadInputTokens != 600 { + t.Fatalf("cache_read = %d, want 600", got.CacheReadInputTokens) + } + if got.CacheCreationInputTokens != 1000 { + t.Fatalf("cache_creation = %d, want 1000", got.CacheCreationInputTokens) + } + billed := billedClaudeInputTokens(2000, got) + if billed != 400 { + t.Fatalf("billed input = %d, want 400", billed) + } + if sum := billed + got.CacheCreationInputTokens + got.CacheReadInputTokens; sum != 2000 { + t.Fatalf("identity broken: input+creation+read = %d, want 2000", sum) + } +} + +// TestSplitAgainstTotalNoCoverageIsAllInput verifies that with no cache +// coverage the whole real total is billed as fresh input. +func TestSplitAgainstTotalNoCoverageIsAllInput(t *testing.T) { + got := promptCacheUsage{}.splitAgainstTotal(1000, 2000) + if got.CacheReadInputTokens != 0 || got.CacheCreationInputTokens != 0 { + t.Fatalf("expected no cache tokens, got read=%d creation=%d", got.CacheReadInputTokens, got.CacheCreationInputTokens) + } + if billed := billedClaudeInputTokens(2000, got); billed != 2000 { + t.Fatalf("expected all 2000 billed as input, got %d", billed) + } +} + +// TestStopIsIdempotent verifies Stop() can be called twice (e.g. test cleanup +// + main shutdown) without panicking on a double close of stopChan. +func TestStopIsIdempotent(t *testing.T) { + tr := newPromptCacheTracker(time.Hour) + tr.stopChan = make(chan struct{}) + tr.Stop() + tr.Stop() // before fix: close of closed channel → panic +} + +// TestComputeSetsDirtyOnCacheHit verifies a cache hit marks the tracker dirty +// so the extended TTL/LastHit is persisted even if the save loop flushes +// between Compute and the following Update. +func TestComputeSetsDirtyOnCacheHit(t *testing.T) { + tr := newPromptCacheTracker(time.Hour) + profile := &promptCacheProfile{ + Model: "claude-sonnet-4-6", + TotalInputTokens: 10000, + Breakpoints: []promptCacheBreakpoint{ + {Fingerprint: [32]byte{1}, CumulativeTokens: 2000, TTL: time.Hour}, + }, + } + now := time.Now() + tr.mu.Lock() + tr.putLocked([32]byte{1}, now.Add(time.Hour), time.Hour) + tr.mu.Unlock() + tr.dirty = false + + tr.Compute("acct", profile) + + if !tr.dirty { + t.Fatal("Compute must set dirty on a cache hit so the refreshed TTL/LastHit is persisted") + } +} + +// TestWriteFileAtomicPersistsAndLeavesNoTemp verifies the atomic-flush path: +// flush writes a complete, reloadable file via tmp+rename, with no stray +// .prompt-cache-*.tmp left in the directory (which would indicate a torn write). +func TestWriteFileAtomicPersistsAndLeavesNoTemp(t *testing.T) { + dir := t.TempDir() + path := dir + "/prompt_cache.json" + + tr := newPromptCacheTracker(5 * time.Minute) + hasher := sha256.New() + writeHashChunk(hasher, "atomic-flush-marker") + var fp [32]byte + copy(fp[:], hasher.Sum(nil)) + tr.mu.Lock() + tr.putLocked(fp, time.Now().Add(10*time.Minute), 5*time.Minute) + tr.dirty = true + tr.mu.Unlock() + tr.flush(path) + + // File must be present and valid JSON (reloadable by Load). + raw, err := os.ReadFile(path) + if err != nil { + t.Fatalf("expected flushed state file, got error: %v", err) + } + var disk struct { + Version int `json:"version"` + Entries []promptCacheEntryOnDisk `json:"entries"` + } + if json.Unmarshal(raw, &disk) != nil { + t.Fatalf("flushed file is not valid JSON: %s", string(raw)) + } + if len(disk.Entries) != 1 { + t.Fatalf("expected 1 persisted entry, got %d", len(disk.Entries)) + } + + // No torn-write temp files should remain after a successful flush. + entries, _ := os.ReadDir(dir) + for _, f := range entries { + if strings.HasSuffix(f.Name(), ".tmp") { + t.Fatalf("leftover temp file after flush: %s", f.Name()) + } + } + + // Reload via Load confirms the round-trip. + t2 := newPromptCacheTracker(5 * time.Minute) + t2.Load(path) + t2.mu.Lock() + _, ok := t2.entries[fp] + t2.mu.Unlock() + if !ok { + t.Fatal("entry not reloaded after atomic flush") + } +} + +// TestBuildClaudeProfileAutoPrefixWithoutCacheControl locks the auto-prefix +// behavior (mirrors Rust prefix_chain_works_without_any_cache_control): a +// request with NO cache_control marker still builds a profile and — across two +// turns sharing a stable history prefix — the second turn reads cache +// (cache_read > 0). This reproduces Anthropic's automatic prompt caching, where +// stable prefixes are reused across turns without the client marking +// cache_control. A request carrying cache_control still builds a profile. +func TestBuildClaudeProfileAutoPrefixWithoutCacheControl(t *testing.T) { + tr := newPromptCacheTracker(time.Hour) + body := strings.Repeat("lorem ipsum dolor sit amet consectetur adipiscing elit sed do eiusmod tempor incididunt ut labore ", 80) + + mk := func(final string) *ClaudeRequest { + return &ClaudeRequest{ + Model: "claude-sonnet-4.5", + System: strings.Repeat("system prompt without any cache marker ", 200), + Messages: []ClaudeMessage{ + {Role: "user", Content: body}, + {Role: "assistant", Content: body}, + {Role: "user", Content: final}, + }, + } + } + + // No cache_control anywhere — predicate confirms it, yet a profile is built. + plain := mk("question one") + if claudeRequestHasCacheControl(plain) { + t.Fatal("claudeRequestHasCacheControl should be false for a request with no cache_control") + } + p1 := tr.BuildClaudeProfile(plain, 5000) + if p1 == nil { + t.Fatal("auto-prefix: expected a profile even when no cache_control is present") + } + first := tr.Compute("acct-1", p1) + if first.CacheReadInputTokens != 0 { + t.Fatalf("first turn has no prior cache to read, got %+v", first) + } + tr.Update("acct-1", p1) + + // Turn 2: identical history prefix, new final message, still no cache_control. + // The shared prefix must hit — auto-prefix cache, no explicit marker needed. + p2 := tr.BuildClaudeProfile(mk("question two"), 5000) + if p2 == nil { + t.Fatal("auto-prefix: turn-2 profile should be built") + } + second := tr.Compute("acct-1", p2) + if second.CacheReadInputTokens == 0 { + t.Fatalf("auto-prefix: expected cross-turn cache_read without cache_control, got %+v", second) + } + + // A request WITH cache_control still builds a profile (unchanged path). + withCache := &ClaudeRequest{ + Model: "claude-sonnet-4.5", + System: []interface{}{ + map[string]interface{}{ + "type": "text", + "text": strings.Repeat("cached system prompt ", 200), + "cache_control": map[string]interface{}{"type": "ephemeral"}, + }, + }, + Messages: []ClaudeMessage{{Role: "user", Content: "follow up"}}, + } + if !claudeRequestHasCacheControl(withCache) { + t.Fatal("claudeRequestHasCacheControl should be true when cache_control is present") + } + if prof := tr.BuildClaudeProfile(withCache, 2048); prof == nil { + t.Fatal("expected a profile when cache_control is present above threshold") + } +} diff --git a/proxy/cache_tracker_test.go b/proxy/cache_tracker_test.go index 6b0262c6..1ba0ec53 100644 --- a/proxy/cache_tracker_test.go +++ b/proxy/cache_tracker_test.go @@ -1,6 +1,8 @@ package proxy import ( + "crypto/sha256" + "kiro-go/config" "strings" "testing" "time" @@ -166,6 +168,71 @@ func TestPromptCacheStableWhenBillingHeaderAppearsOrDisappears(t *testing.T) { } } +// TestPromptCacheStructuralSkipOfLeadingDynamicSystemBlock generalizes the two +// billing-header guards above: Claude Code injects a per-turn DYNAMIC block at +// system[0] WITHOUT cache_control, followed by the stable cached block(s). +// Fingerprinting from that volatile head would drift every downstream prefix +// fingerprint across turns, collapsing cross-turn cache hits to zero. The +// tracker must STRUCTURALLY skip every leading system block that lacks +// cache_control (mirrors Rust extract_segments), not just the literal +// x-anthropic-billing-header string the legacy matcher recognizes. So this must +// hold even when the dynamic header is an arbitrary changing string. +func TestPromptCacheStructuralSkipOfLeadingDynamicSystemBlock(t *testing.T) { + tracker := newPromptCacheTracker(time.Hour) + stableSys := strings.Repeat("You are a careful senior engineer who reviews diffs critically and writes idiomatic Go. ", 80) + body := strings.Repeat("implement the feature step by step with tests and clear rationale. ", 60) + + build := func(dynamicHeader string, msgs []ClaudeMessage) *ClaudeRequest { + return &ClaudeRequest{ + Model: "claude-sonnet-4.5", + System: []interface{}{ + // sys[0]: per-turn dynamic marker, NO cache_control (not a billing string). + map[string]interface{}{"type": "text", "text": dynamicHeader}, + // sys[1]: stable cached block. + map[string]interface{}{ + "type": "text", + "text": stableSys, + "cache_control": map[string]interface{}{"type": "ephemeral"}, + }, + }, + Messages: msgs, + } + } + + turn1 := build("now=1001", []ClaudeMessage{ + {Role: "user", Content: body}, + {Role: "assistant", Content: body}, + {Role: "user", Content: body}, + }) + profile1 := tracker.BuildClaudeProfile(turn1, 8192) + if profile1 == nil { + t.Fatalf("profile1 should be built") + } + first := tracker.Compute("acct-1", profile1) + if first.CacheReadInputTokens != 0 { + t.Fatalf("turn1 has no prior cache to read, got %+v", first) + } + tracker.Update("acct-1", profile1) + + // Turn 2: dynamic header changed to "now=2002"; one more a/u pair appended. + // After skipping the volatile head, sys[1]+history prefix is byte-identical → must hit. + turn2 := build("now=2002", []ClaudeMessage{ + {Role: "user", Content: body}, + {Role: "assistant", Content: body}, + {Role: "user", Content: body}, + {Role: "assistant", Content: body}, + {Role: "user", Content: body}, + }) + profile2 := tracker.BuildClaudeProfile(turn2, 8192) + if profile2 == nil { + t.Fatalf("profile2 should be built") + } + second := tracker.Compute("acct-1", profile2) + if second.CacheReadInputTokens == 0 { + t.Fatalf("structural skip: dynamic system header change must not break cross-turn cache hit, got %+v", second) + } +} + func TestCanonicalCacheValueIgnoresPositionKeys(t *testing.T) { first := canonicalizeCacheValue(stripCachePositionKeys(map[string]interface{}{ "kind": "system", @@ -262,3 +329,343 @@ func TestPromptCacheImplicitBreakpointAtMessageEnd(t *testing.T) { t.Fatalf("expected cache read via implicit message-end breakpoint, got %+v", result) } } + +// TestPromptCacheCrossAccountSharing verifies C1: two different accountIDs with +// the SAME prompt fingerprint share cache entries. Account B's request should +// HIT on the fingerprint Account A stored — no per-account isolation. +func TestPromptCacheCrossAccountSharing(t *testing.T) { + tracker := newPromptCacheTracker(5 * time.Minute) + + // Build a profile with one explicit cache_control breakpoint above the + // min-token threshold. + block := cacheablePromptBlock{ + Value: map[string]interface{}{"kind": "system", "block": map[string]interface{}{ + "type": "text", "text": strings.Repeat("x ", 600), // ~600 tokens > 1024? use more + "cache_control": map[string]interface{}{"type": "ephemeral"}, + }}, + Tokens: 1200, + TTL: 5 * time.Minute, + } + hasher := sha256.New() + writeHashChunk(hasher, canonicalizeCacheValue(block.Value)) + var fp [32]byte + copy(fp[:], hasher.Sum(nil)) + + profile := &promptCacheProfile{ + Breakpoints: []promptCacheBreakpoint{{Fingerprint: fp, CumulativeTokens: 1200, TTL: 5 * time.Minute}}, + TotalInputTokens: 1200, + Model: "claude-sonnet-4-5", + } + + // Account A: first request → cache_creation. + usageA := tracker.Compute("account-A", profile) + if usageA.CacheCreationInputTokens == 0 { + t.Fatalf("account A: expected cache_creation > 0, got %d", usageA.CacheCreationInputTokens) + } + tracker.Update("account-A", profile) + + // Account B: SAME prompt, DIFFERENT account → should be cache_read (C1 fix). + // Before C1: account B had its own empty store → cache_creation. + // After C1: account B shares the global store → cache_read. + usageB := tracker.Compute("account-B", profile) + if usageB.CacheReadInputTokens == 0 { + t.Fatalf("account B: expected cache_read > 0 (cross-account sharing), got 0. usage=%+v", usageB) + } +} + +// TestPromptCacheCapConfigurable verifies C2: the cache-read cap can be set +// above the default 0.85 via config, so a request where 90% of input is from +// cache reports the full 90% (not clamped to 85%). +func TestPromptCacheCapConfigurable(t *testing.T) { + cfgFile := t.TempDir() + "/config.json" + if err := config.Init(cfgFile); err != nil { + t.Fatalf("config.Init: %v", err) + } + + // Build a profile: 1024 tokens cached (>= defaultMinCacheableTokens), total + // 1100. Default cap clamps cache_read to 0.85*1100=935; cap=0.95 allows up + // to 0.95*1100=1045 >= 1024, so the full breakpoint is reported. + hasher := sha256.New() + writeHashChunk(hasher, canonicalizeCacheValue(map[string]interface{}{"k": strings.Repeat("v ", 500)})) + var fp [32]byte + copy(fp[:], hasher.Sum(nil)) + profile := &promptCacheProfile{ + Breakpoints: []promptCacheBreakpoint{{Fingerprint: fp, CumulativeTokens: 1024, TTL: 5 * time.Minute}}, + TotalInputTokens: 1100, + Model: "claude-sonnet-4-5", + } + tracker := newPromptCacheTracker(5 * time.Minute) + tracker.Update("acc", profile) // store it + + // Default cap 0.85: cache_read clamped to min(1000, 0.85*1100=935) = 935. + usage85 := tracker.Compute("acc", profile) + if usage85.CacheReadInputTokens > 940 { + t.Fatalf("default cap: expected cache_read ~935, got %d", usage85.CacheReadInputTokens) + } + + // Raise cap to 0.95: cache_read should be min(1000, 0.95*1100=1045) = 1000. + config.UpdatePromptCacheMaxRatio(0.95) + defer config.UpdatePromptCacheMaxRatio(0.85) + usage95 := tracker.Compute("acc", profile) + if usage95.CacheReadInputTokens < 990 { + t.Fatalf("cap 0.95: expected cache_read ~1000, got %d", usage95.CacheReadInputTokens) + } +} + +// TestPromptCacheDiskPersistence verifies C3: entries saved to disk are +// reloaded on startup, surviving a "restart" (new tracker instance). +func TestPromptCacheDiskPersistence(t *testing.T) { + path := t.TempDir() + "/prompt_cache.json" + + // Tracker 1: store an entry, flush to disk. + t1 := newPromptCacheTracker(5 * time.Minute) + hasher := sha256.New() + writeHashChunk(hasher, "test-cache-value-disk") + var fp [32]byte + copy(fp[:], hasher.Sum(nil)) + t1.mu.Lock() + t1.putLocked(fp, time.Now().Add(3*time.Minute), 5*time.Minute) + t1.dirty = true + t1.mu.Unlock() + t1.flush(path) + + // Tracker 2: load from disk → should have the entry. + t2 := newPromptCacheTracker(5 * time.Minute) + t2.Load(path) + t2.mu.Lock() + _, ok := t2.entries[fp] + t2.mu.Unlock() + if !ok { + t.Fatalf("C3: entry not reloaded from disk after 'restart'") + } + + // Expired entry should NOT reload. + path2 := t.TempDir() + "/expired.json" + t1b := newPromptCacheTracker(5 * time.Minute) + t1b.mu.Lock() + fpExpired := sha256.Sum256([]byte("expired")) + t1b.putLocked(fpExpired, time.Now().Add(-1*time.Minute), 5*time.Minute) // already expired + t1b.dirty = true + t1b.mu.Unlock() + t1b.flush(path2) + + t3 := newPromptCacheTracker(5 * time.Minute) + t3.Load(path2) + t3.mu.Lock() + _, okExpired := t3.entries[fpExpired] + t3.mu.Unlock() + if okExpired { + t.Fatalf("C3: expired entry should not be reloaded") + } +} + +// TestComputeBreakdownClampedToCreation verifies the Anthropic usage invariant +// cache_creation_input_tokens == ephemeral_5m + ephemeral_1h holds even when the +// 0.85 cacheable cap reduces creation below the raw breakpoint deltas. +// Profile: TotalInputTokens=10000 → maxCacheable=8500. A matched breakpoint at +// 2000 leaves creation=6500, but the uncapped breakdown from matched=2000 to the +// 10000 breakpoint is 8000 (1h). Without clamping, 5m+1h (8000) != creation (6500). +func TestComputeBreakdownClampedToCreation(t *testing.T) { + tr := newPromptCacheTracker(time.Hour) + profile := &promptCacheProfile{ + Model: "claude-sonnet-4-6", + TotalInputTokens: 10000, + Breakpoints: []promptCacheBreakpoint{ + {Fingerprint: [32]byte{1}, CumulativeTokens: 2000, TTL: 5 * time.Minute}, + {Fingerprint: [32]byte{2}, CumulativeTokens: 10000, TTL: time.Hour}, + }, + } + now := time.Now() + tr.mu.Lock() + tr.putLocked([32]byte{1}, now.Add(time.Hour), time.Hour) + tr.mu.Unlock() + + usage := tr.Compute("acct", profile) + + if usage.CacheCreation5mInputTokens+usage.CacheCreation1hInputTokens != usage.CacheCreationInputTokens { + t.Fatalf("breakdown must sum to creation: 5m=%d 1h=%d creation=%d", + usage.CacheCreation5mInputTokens, usage.CacheCreation1hInputTokens, usage.CacheCreationInputTokens) + } + if usage.CacheCreation1hInputTokens > usage.CacheCreationInputTokens { + t.Fatalf("1h breakdown must not exceed creation: 1h=%d creation=%d", + usage.CacheCreation1hInputTokens, usage.CacheCreationInputTokens) + } +} + +// TestComputeBreakdownClampedToCreationEmptyCacheBelowMin verifies the empty-cache +// (first-request) path still preserves the Anthropic invariant +// cache_creation_input_tokens == ephemeral_5m + ephemeral_1h when the input is +// below the minimum-cacheable threshold. effectiveCreation is zeroed +// (lastTokens < minTokens), but computePromptCacheTTLBreakdown(profile, 0) still +// returns lastTokens > 0 — the empty-path clamp (cache_tracker.go:255) must force +// 5m/1h to 0 so message_start doesn't emit an invariant-violating usage. Guards +// the empty-cache clamp site, which the matched-path TestComputeBreakdownClampedToCreation +// does not exercise (it always has lastTokens >= minTokens). +func TestComputeBreakdownClampedToCreationEmptyCacheBelowMin(t *testing.T) { + tr := newPromptCacheTracker(time.Hour) // empty → empty-cache (first-request) path + profile := &promptCacheProfile{ + Model: "claude-sonnet-4-6", // minTokens = 1024 (non-opus default) + TotalInputTokens: 200, + Breakpoints: []promptCacheBreakpoint{ + {Fingerprint: [32]byte{1}, CumulativeTokens: 100, TTL: time.Hour}, // 100 < 1024 → effectiveCreation = 0 + }, + } + + usage := tr.Compute("acct", profile) + + if usage.CacheCreationInputTokens != 0 { + t.Fatalf("below-minTokens first request should yield zero creation; got %d", usage.CacheCreationInputTokens) + } + if sum := usage.CacheCreation5mInputTokens + usage.CacheCreation1hInputTokens; sum != usage.CacheCreationInputTokens { + t.Fatalf("breakdown must sum to creation (empty path, below min): 5m=%d 1h=%d creation=%d", + usage.CacheCreation5mInputTokens, usage.CacheCreation1hInputTokens, usage.CacheCreationInputTokens) + } +} + +// TestPromptCacheMaxEntriesConfigurable verifies the cache LRU bound is +// configurable via config and defaults to 131072. +func TestPromptCacheMaxEntriesConfigurable(t *testing.T) { + cfgFile := t.TempDir() + "/config.json" + if err := config.Init(cfgFile); err != nil { + t.Fatalf("config.Init: %v", err) + } + + if got := config.GetPromptCacheMaxEntries(); got != 131072 { + t.Fatalf("default cap: expected 131072, got %d", got) + } + + if err := config.UpdatePromptCacheMaxEntries(50000); err != nil { + t.Fatalf("update: %v", err) + } + if got := config.GetPromptCacheMaxEntries(); got != 50000 { + t.Fatalf("after update: expected 50000, got %d", got) + } + + // ≤ 0 falls back to the default. + if err := config.UpdatePromptCacheMaxEntries(0); err != nil { + t.Fatalf("reset: %v", err) + } + if got := config.GetPromptCacheMaxEntries(); got != 131072 { + t.Fatalf("zero should fall back to default 131072, got %d", got) + } + + // An explicit small value is clamped to the 256 floor. + if err := config.UpdatePromptCacheMaxEntries(10); err != nil { + t.Fatalf("set small: %v", err) + } + if got := config.GetPromptCacheMaxEntries(); got != 256 { + t.Fatalf("small value 10 should clamp to 256, got %d", got) + } +} + +// TestCacheDoesNotChurnAtHighCapacity is the regression guard for the prefix +// churn bug: at the production default capacity (131072), a 5000-prefix +// workload (far above the old 4096 cap) does not evict the oldest seeded prefix +// before it is replayed, so the replay reads from cache. +func TestCacheDoesNotChurnAtHighCapacity(t *testing.T) { + tr := newPromptCacheTrackerWithCapacity(time.Hour, 131072) + now := time.Now() + tr.mu.Lock() + for i := 0; i < 5000; i++ { + var fp [32]byte + fp[0] = byte(i) + fp[1] = byte(i >> 8) + fp[2] = byte(i >> 16) + tr.putLocked(fp, now.Add(time.Hour), time.Hour) + } + tr.evictOverflowLocked() + tr.mu.Unlock() + + // The oldest seeded prefix (i=0, all-zero bytes) must survive at cap 131072. + var fp0 [32]byte + profile := &promptCacheProfile{ + Model: "claude-sonnet-4-6", + TotalInputTokens: 2000, + Breakpoints: []promptCacheBreakpoint{{Fingerprint: fp0, CumulativeTokens: 2000, TTL: time.Hour}}, + } + if usage := tr.Compute("acct", profile); usage.CacheReadInputTokens == 0 { + t.Fatalf("expected old prefix to survive at cap=131072 (no churn); got cache_read=0") + } +} + +// TestCacheChurnsAtLowCapacity proves the churn mechanism: at the old cap 4096 +// the same 5000-prefix workload evicts the oldest prefix (i=0), so its replay +// misses. This is the sensitivity companion to TestCacheDoesNotChurnAtHighCapacity. +func TestCacheChurnsAtLowCapacity(t *testing.T) { + tr := newPromptCacheTrackerWithCapacity(time.Hour, 4096) + now := time.Now() + tr.mu.Lock() + for i := 0; i < 5000; i++ { + var fp [32]byte + fp[0] = byte(i) + fp[1] = byte(i >> 8) + fp[2] = byte(i >> 16) + tr.putLocked(fp, now.Add(time.Hour), time.Hour) + } + tr.evictOverflowLocked() + tr.mu.Unlock() + + // 5000 > 4096 → LRU evicts the 904 oldest; i=0 (oldest, never re-touched) is gone. + var fp0 [32]byte + profile := &promptCacheProfile{ + Model: "claude-sonnet-4-6", + TotalInputTokens: 2000, + Breakpoints: []promptCacheBreakpoint{{Fingerprint: fp0, CumulativeTokens: 2000, TTL: time.Hour}}, + } + if usage := tr.Compute("acct", profile); usage.CacheReadInputTokens != 0 { + t.Fatalf("expected oldest prefix churned at cap=4096; got cache_read=%d", usage.CacheReadInputTokens) + } +} + +// TestCacheStats verifies the atomic counters and Stats() snapshot: one miss, +// one hit, one expiration, and one LRU eviction are all counted, and capacity +// reflects the configured bound. +func TestCacheStats(t *testing.T) { + tr := newPromptCacheTrackerWithCapacity(time.Hour, 3) + hit := &promptCacheProfile{ + Model: "claude-sonnet-4-6", + TotalInputTokens: 2000, + Breakpoints: []promptCacheBreakpoint{{Fingerprint: [32]byte{7}, CumulativeTokens: 2000, TTL: time.Hour}}, + } + now := time.Now() + + // Seed an already-expired entry, then Compute: pruneExpiredLocked drops it + // (expirations=1) and the empty cache yields a miss (misses=1). + tr.mu.Lock() + tr.putLocked([32]byte{99}, now.Add(-time.Minute), time.Hour) + tr.mu.Unlock() + tr.Compute("acct", hit) + + // Store the hit profile, then Compute → hit (hits=1). + tr.Update("acct", hit) + tr.Compute("acct", hit) + + // Overflow: entries were {7}; add 3 more → {7,1,2,3} len=4 → + // evictOverflowLocked pops the LRU back (7), leaving 3 (evictions=1). + tr.mu.Lock() + for i := 1; i <= 3; i++ { + tr.putLocked([32]byte{byte(i)}, now.Add(time.Hour), time.Hour) + } + tr.evictOverflowLocked() + tr.mu.Unlock() + + stats := tr.Stats() + if stats.Hits != 1 { + t.Errorf("hits = %d, want 1", stats.Hits) + } + if stats.Misses != 1 { + t.Errorf("misses = %d, want 1", stats.Misses) + } + if stats.Evictions != 1 { + t.Errorf("evictions = %d, want 1", stats.Evictions) + } + if stats.Expirations != 1 { + t.Errorf("expirations = %d, want 1", stats.Expirations) + } + if stats.Capacity != 3 { + t.Errorf("capacity = %d, want 3", stats.Capacity) + } + if stats.Entries != 3 { + t.Errorf("entries = %d, want 3", stats.Entries) + } +} diff --git a/proxy/handler.go b/proxy/handler.go index a14f71fd..4162b5e8 100644 --- a/proxy/handler.go +++ b/proxy/handler.go @@ -1,6 +1,7 @@ package proxy import ( + "crypto/subtle" "encoding/json" "fmt" "io" @@ -9,6 +10,7 @@ import ( "kiro-go/logger" "kiro-go/pool" "net/http" + "path/filepath" "strings" "sync" "sync/atomic" @@ -21,16 +23,16 @@ const tokenRefreshSkewSeconds int64 = 120 // RequestLog stores details about a single API request (success or failure). type RequestLog struct { - Time int64 `json:"time"` // Unix timestamp - Endpoint string `json:"endpoint"` // claude/openai/responses - Model string `json:"model"` // Requested model - AccountID string `json:"accountId"` // Account used - Status string `json:"status"` // "success" or "error" - Error string `json:"error"` // Error message (empty on success) - ErrorType string `json:"errorType"` // Error category (empty on success) - Tokens int `json:"tokens"` // Total tokens (input+output, 0 on failure) - Credits float64 `json:"credits"` // Credits consumed (0 on failure) - Duration int64 `json:"duration"` // Request duration in ms + Time int64 `json:"time"` // Unix timestamp + Endpoint string `json:"endpoint"` // claude/openai/responses + Model string `json:"model"` // Requested model + AccountID string `json:"accountId"` // Account used + Status string `json:"status"` // "success" or "error" + Error string `json:"error"` // Error message (empty on success) + ErrorType string `json:"errorType"` // Error category (empty on success) + Tokens int `json:"tokens"` // Total tokens (input+output, 0 on failure) + Credits float64 `json:"credits"` // Credits consumed (0 on failure) + Duration int64 `json:"duration"` // Request duration in ms } const requestLogsMaxSize = 500 @@ -245,8 +247,14 @@ func NewHandler() *Handler { stopStatsSaver: make(chan struct{}), promptCache: newPromptCacheTracker(defaultPromptCacheTTL), } + cachePath := filepath.Join(config.GetConfigDir(), "prompt_cache.json") + h.promptCache.Load(cachePath) + h.promptCache.startSaveLoop(cachePath, 30*time.Second) // 启动后台刷新 go h.backgroundRefresh() + // 启动 token 预热 worker（2 分钟一轮，提前 15 分钟刷新即将过期的 token， + // 避免请求路径上同步刷新造成的延迟尖峰） + go h.backgroundWarmup() // 启动后台统计保存 (每30秒保存一次) go h.backgroundStatsSaver() // 清理过期的 stored responses（>30 天） @@ -284,25 +292,11 @@ func (h *Handler) refreshAllAccounts() { continue } - // 检查 token 是否需要刷新 - if account.ExpiresAt > 0 && time.Now().Unix() > account.ExpiresAt-tokenRefreshSkewSeconds { - newAccessToken, newRefreshToken, newExpiresAt, profileArn, err := auth.RefreshToken(account) - if err != nil { - logger.Warnf("[BackgroundRefresh] Token refresh failed for %s: %v", account.Email, err) - h.handleAccountFailure(account, err) - continue - } - account.AccessToken = newAccessToken - if newRefreshToken != "" { - account.RefreshToken = newRefreshToken - } - account.ExpiresAt = newExpiresAt - config.UpdateAccountToken(account.ID, newAccessToken, newRefreshToken, newExpiresAt) - h.pool.UpdateToken(account.ID, newAccessToken, newRefreshToken, newExpiresAt) - if profileArn != "" { - account.ProfileArn = profileArn - config.UpdateAccountProfileArn(account.ID, profileArn) - } + // 刷新 token（去重 + 持久化由 refreshAccountToken 统一处理） + if err := h.refreshAccountToken(account, false); err != nil { + logger.Warnf("[BackgroundRefresh] Token refresh failed for %s: %v", account.Email, err) + h.handleAccountFailure(account, err) + continue } // 刷新账户信息 @@ -455,6 +449,8 @@ func (h *Handler) handleStats(w http.ResponseWriter, r *http.Request) { "failedRequests": atomic.LoadInt64(&h.failedRequests), "totalTokens": atomic.LoadInt64(&h.totalTokens), "totalCredits": h.getCredits(), + "cache": h.promptCache.Stats(), + "cacheWindows": cacheMetricsSnapshotByWindows(time.Now()), "uptime": time.Now().Unix() - h.startTime, }) } @@ -513,6 +509,8 @@ func buildAnthropicModelsResponse(cached []ModelInfo, thinkingSuffix string) []m func fallbackAnthropicModels(thinkingSuffix string) []map[string]interface{} { return []map[string]interface{}{ + buildModelInfo("claude-opus-4.8", "anthropic", true), + buildModelInfo("claude-opus-4.8"+thinkingSuffix, "anthropic", true), buildModelInfo("claude-sonnet-4.6", "anthropic", true), buildModelInfo("claude-sonnet-4.6"+thinkingSuffix, "anthropic", true), buildModelInfo("claude-opus-4.6", "anthropic", true), @@ -843,6 +841,12 @@ func (h *Handler) handleClaudeMessagesInternal(w http.ResponseWriter, r *http.Re } } +// streamKeepaliveInterval 是流式响应期间向客户端发送 SSE 保活注释行的间隔。远小于 +// Cloudflare ~120s 的 Proxy Read Timeout，确保上游长静默（opus 高强度 thinking 两个数据 +// 块之间、或去掉总超时后的超长推理）期间客户端/CDN 连接不被中途判超时断开。为 var 以便 +// 测试注入更小的间隔来覆盖"流中途静默触发 ping 与数据并发写"路径。Ported from kiro-tutu。 +var streamKeepaliveInterval = 10 * time.Second + // handleClaudeStream Claude 流式响应 func (h *Handler) handleClaudeStream(w http.ResponseWriter, payload *KiroPayload, model string, thinking bool, thinkingOpts claudeThinkingResponseOptions, estimatedInputTokens int, cacheProfile *promptCacheProfile, apiKeyID string) { w.Header().Set("Content-Type", "text/event-stream; charset=utf-8") @@ -855,6 +859,58 @@ func (h *Handler) handleClaudeStream(w http.ResponseWriter, payload *KiroPayload return } + // SSE 保活：整个流式期间（不止首字节前）周期性发送注释行，使 200 响应头立即抵达 + // 客户端/CDN 并持续刷新其 Proxy Read Timeout——既覆盖上游 prefill 迟迟不返回首字节， + // 也覆盖上游输出中途长静默（opus 高强度 thinking 两个数据块之间可静默数十秒，且本服务 + // 已去掉总超时），避免客户端/CDN 在 ~120s 处判超时断开（用户侧表现为"说一半就断"）。 + // ':' 开头的注释行被 Anthropic SDK / Claude Code 等客户端忽略，不污染响应内容。 + // + // 并发：保活 goroutine 与主 goroutine 都会写 w，故所有真实事件写出统一走 emit()，与 + // 保活共用 hbMu 串行化；lastWriteNano 记录最近一次写出时刻，保活仅在距上次写出已满 + // streamKeepaliveInterval（确有静默）时才补发，数据正常流动时不插入多余注释行。 + var hbMu sync.Mutex + hbStopped := false + committed := false // 200 响应头是否已 flush 到客户端（由保活触发） + lastWriteNano := time.Now().UnixNano() + hbDone := make(chan struct{}) + var hbStopOnce sync.Once + stopHeartbeat := func() { + hbStopOnce.Do(func() { + hbMu.Lock() + hbStopped = true + hbMu.Unlock() + close(hbDone) + }) + } + defer stopHeartbeat() + + emit := func(event string, data interface{}) { + hbMu.Lock() + h.sendSSE(w, flusher, event, data) + lastWriteNano = time.Now().UnixNano() + hbMu.Unlock() + } + + go func() { + ticker := time.NewTicker(streamKeepaliveInterval) + defer ticker.Stop() + for { + select { + case <-hbDone: + return + case <-ticker.C: + hbMu.Lock() + if !hbStopped && time.Since(time.Unix(0, lastWriteNano)) >= streamKeepaliveInterval { + fmt.Fprint(w, ": keepalive\n\n") + flusher.Flush() + committed = true + lastWriteNano = time.Now().UnixNano() + } + hbMu.Unlock() + } + } + }() + // 获取 thinking 输出格式配置 thinkingFormat := thinkingOpts.Format @@ -870,7 +926,7 @@ func (h *Handler) handleClaudeStream(w http.ResponseWriter, payload *KiroPayload if messageStarted { return } - h.sendSSE(w, flusher, "message_start", map[string]interface{}{ + emit("message_start", map[string]interface{}{ "type": "message_start", "message": map[string]interface{}{ "id": msgID, @@ -886,7 +942,11 @@ func (h *Handler) handleClaudeStream(w http.ResponseWriter, payload *KiroPayload messageStarted = true } - for attempt := 0; attempt < maxAccountRetryAttempts; attempt++ { + retryBudget := resolveAccountRetryBudget(h.pool.Count()) + for attempt := 0; attempt < retryBudget; attempt++ { + if attempt > 0 { + time.Sleep(accountRetryBackoff(attempt - 1)) + } account := h.pool.GetNextForModelExcluding(model, excluded) if account == nil { break @@ -914,7 +974,7 @@ func (h *Handler) handleClaudeStream(w http.ResponseWriter, payload *KiroPayload if activeBlockIndex < 0 { return } - h.sendSSE(w, flusher, "content_block_stop", map[string]interface{}{ + emit("content_block_stop", map[string]interface{}{ "type": "content_block_stop", "index": activeBlockIndex, }) @@ -933,7 +993,7 @@ func (h *Handler) handleClaudeStream(w http.ResponseWriter, payload *KiroPayload nextContentIndex++ if blockType == "thinking" { - h.sendSSE(w, flusher, "content_block_start", map[string]interface{}{ + emit("content_block_start", map[string]interface{}{ "type": "content_block_start", "index": idx, "content_block": map[string]string{ @@ -942,7 +1002,7 @@ func (h *Handler) handleClaudeStream(w http.ResponseWriter, payload *KiroPayload }, }) } else { - h.sendSSE(w, flusher, "content_block_start", map[string]interface{}{ + emit("content_block_start", map[string]interface{}{ "type": "content_block_start", "index": idx, "content_block": map[string]string{ @@ -969,7 +1029,7 @@ func (h *Handler) handleClaudeStream(w http.ResponseWriter, payload *KiroPayload return } startContentBlock("text") - h.sendSSE(w, flusher, "content_block_delta", map[string]interface{}{ + emit("content_block_delta", map[string]interface{}{ "type": "content_block_delta", "index": activeBlockIndex, "delta": map[string]string{"type": "text_delta", "text": text}, @@ -996,7 +1056,7 @@ func (h *Handler) handleClaudeStream(w http.ResponseWriter, payload *KiroPayload return } startContentBlock("text") - h.sendSSE(w, flusher, "content_block_delta", map[string]interface{}{ + emit("content_block_delta", map[string]interface{}{ "type": "content_block_delta", "index": activeBlockIndex, "delta": map[string]string{"type": "text_delta", "text": outputText}, @@ -1006,7 +1066,7 @@ func (h *Handler) handleClaudeStream(w http.ResponseWriter, payload *KiroPayload return } startContentBlock("text") - h.sendSSE(w, flusher, "content_block_delta", map[string]interface{}{ + emit("content_block_delta", map[string]interface{}{ "type": "content_block_delta", "index": activeBlockIndex, "delta": map[string]string{"type": "text_delta", "text": text}, @@ -1033,7 +1093,7 @@ func (h *Handler) handleClaudeStream(w http.ResponseWriter, payload *KiroPayload } if text != "" { startContentBlock("thinking") - h.sendSSE(w, flusher, "content_block_delta", map[string]interface{}{ + emit("content_block_delta", map[string]interface{}{ "type": "content_block_delta", "index": activeBlockIndex, "delta": map[string]string{"type": "thinking_delta", "thinking": text}, @@ -1177,7 +1237,7 @@ func (h *Handler) handleClaudeStream(w http.ResponseWriter, payload *KiroPayload idx := nextContentIndex nextContentIndex++ - h.sendSSE(w, flusher, "content_block_start", map[string]interface{}{ + emit("content_block_start", map[string]interface{}{ "type": "content_block_start", "index": idx, "content_block": map[string]interface{}{ @@ -1189,7 +1249,7 @@ func (h *Handler) handleClaudeStream(w http.ResponseWriter, payload *KiroPayload }) inputJSON, _ := json.Marshal(tu.Input) - h.sendSSE(w, flusher, "content_block_delta", map[string]interface{}{ + emit("content_block_delta", map[string]interface{}{ "type": "content_block_delta", "index": idx, "delta": map[string]interface{}{ @@ -1198,7 +1258,7 @@ func (h *Handler) handleClaudeStream(w http.ResponseWriter, payload *KiroPayload }, }) - h.sendSSE(w, flusher, "content_block_stop", map[string]interface{}{ + emit("content_block_stop", map[string]interface{}{ "type": "content_block_stop", "index": idx, }) @@ -1220,13 +1280,22 @@ func (h *Handler) handleClaudeStream(w http.ResponseWriter, payload *KiroPayload lastErr = err excluded[account.ID] = true h.handleAccountFailure(account, err) + // Cache-dispatch observability: failed dispatch contributes 0 tokens so + // it cannot inflate the simulated cache_read subsidy (see cache_metrics.go). + recordClaudeCacheDispatch("failure", model, "stream", account, promptCacheUsage{}, 0, 0) + if isUpstreamPermanentError(err) { + // Malformed request — no other account can succeed. Stop retrying; + // lastErr flows to the client. The relaying account was left neutral + // by handleAccountFailure (not penalised). + break + } if !messageStarted { continue } h.recordFailureWithDetails("claude", model, account.ID, err) - h.sendSSE(w, flusher, "error", map[string]interface{}{ + emit("error", map[string]interface{}{ "type": "error", - "error": map[string]string{"type": "api_error", "message": err.Error()}, + "error": map[string]string{"type": "api_error", "message": improperlyFormedClientMessage(err)}, }) return } @@ -1242,6 +1311,12 @@ func (h *Handler) handleClaudeStream(w http.ResponseWriter, payload *KiroPayload } else if inputTokens <= 0 { inputTokens = estimatedInputTokens } + // Re-anchor the cache split to the real upstream input total so + // input+creation+read stays consistent (cacheUsage was computed against + // the pre-call token estimate). + if cacheProfile != nil && realInputTokens > 0 { + cacheUsage = cacheUsage.splitAgainstTotal(cacheProfile.TotalInputTokens, inputTokens) + } outputContent, extractedReasoning := extractThinkingFromContent(rawContentBuilder.String()) thinkingOutput := rawThinkingBuilder.String() if thinking && thinkingOutput == "" && extractedReasoning != "" { @@ -1254,9 +1329,14 @@ func (h *Handler) handleClaudeStream(w http.ResponseWriter, payload *KiroPayload h.recordSuccessForApiKey(apiKeyID, inputTokens, outputTokens, credits) h.pool.RecordSuccess(account.ID) + // h.pool.RecordLatency(account.ID, float64(time.Since(reqStart).Milliseconds())) + // ↑ dispatch-only latency EWMA (commit fbc99b1), absent on this cache-only + // branch off main. No-op here; stripped so handler compiles against main's + // pool. Re-added verbatim when dispatch lands. h.pool.UpdateStats(account.ID, inputTokens+outputTokens, credits) h.promptCache.Update(account.ID, cacheProfile) h.recordSuccessLog("claude", model, account.ID, inputTokens+outputTokens, credits, time.Since(reqStart).Milliseconds()) + recordClaudeCacheDispatch("success", model, "stream", account, cacheUsage, inputTokens, outputTokens) stopReason := "end_turn" if len(toolUses) > 0 { @@ -1264,7 +1344,7 @@ func (h *Handler) handleClaudeStream(w http.ResponseWriter, payload *KiroPayload } ensureMessageStart() - h.sendSSE(w, flusher, "message_delta", map[string]interface{}{ + emit("message_delta", map[string]interface{}{ "type": "message_delta", "delta": map[string]interface{}{ "stop_reason": stopReason, @@ -1272,7 +1352,41 @@ func (h *Handler) handleClaudeStream(w http.ResponseWriter, payload *KiroPayload "usage": buildClaudeUsageMap(inputTokens, outputTokens, cacheUsage, cacheProfile != nil), }) - h.sendSSE(w, flusher, "message_stop", map[string]interface{}{ + emit("message_stop", map[string]interface{}{ + "type": "message_stop", + }) + return + } + + // 循环外的错误/无账号路径。保活 goroutine 可能仍在写 w，先停掉再决定响应形态， + // 避免 sendClaudeError 的 WriteHeader+JSON 与心跳的 ': keepalive' 并发写同一个 w。 + stopHeartbeat() + if lastErr != nil { + h.recordFailureWithDetails("claude", model, "", lastErr) + } + hbMu.Lock() + didCommit := committed + hbMu.Unlock() + + if didCommit { + // 心跳已落地 200 响应头：HTTP 状态码已无法更改，客户端正按 SSE 解析——必须以 + // SSE error 序列收尾，而非 sendClaudeError 的 WriteHeader+JSON（后者会触发 + // superfluous WriteHeader 告警且破坏流式协议）。 + msg := "No available accounts" + if lastErr != nil { + msg = improperlyFormedClientMessage(lastErr) + } + ensureMessageStart() + emit("error", map[string]interface{}{ + "type": "error", + "error": map[string]string{"type": "api_error", "message": msg}, + }) + emit("message_delta", map[string]interface{}{ + "type": "message_delta", + "delta": map[string]interface{}{"stop_reason": "error"}, + "usage": map[string]interface{}{"output_tokens": 0}, + }) + emit("message_stop", map[string]interface{}{ "type": "message_stop", }) return @@ -1283,8 +1397,7 @@ func (h *Handler) handleClaudeStream(w http.ResponseWriter, payload *KiroPayload return } - h.recordFailureWithDetails("claude", model, "", lastErr) - h.sendClaudeError(w, 500, "api_error", lastErr.Error()) + h.sendClaudeError(w, 500, "api_error", improperlyFormedClientMessage(lastErr)) } func (h *Handler) sendSSE(w http.ResponseWriter, flusher http.Flusher, event string, data interface{}) { @@ -1446,7 +1559,11 @@ func (h *Handler) handleClaudeNonStream(w http.ResponseWriter, payload *KiroPayl var lastErr error reqStart := time.Now() - for attempt := 0; attempt < maxAccountRetryAttempts; attempt++ { + retryBudget := resolveAccountRetryBudget(h.pool.Count()) + for attempt := 0; attempt < retryBudget; attempt++ { + if attempt > 0 { + time.Sleep(accountRetryBackoff(attempt - 1)) + } account := h.pool.GetNextForModelExcluding(model, excluded) if account == nil { break @@ -1494,6 +1611,12 @@ func (h *Handler) handleClaudeNonStream(w http.ResponseWriter, payload *KiroPayl lastErr = err excluded[account.ID] = true h.handleAccountFailure(account, err) + // Cache-dispatch observability: failed dispatch contributes 0 tokens so + // it cannot inflate the simulated cache_read subsidy (see cache_metrics.go). + recordClaudeCacheDispatch("failure", model, "nonstream", account, promptCacheUsage{}, 0, 0) + if isUpstreamPermanentError(err) { + break + } continue } @@ -1512,13 +1635,24 @@ func (h *Handler) handleClaudeNonStream(w http.ResponseWriter, payload *KiroPayl } else if inputTokens <= 0 { inputTokens = estimatedInputTokens } + // Re-anchor the cache split to the real upstream input total so + // input+creation+read stays consistent (cacheUsage was computed against + // the pre-call token estimate). + if cacheProfile != nil && realInputTokens > 0 { + cacheUsage = cacheUsage.splitAgainstTotal(cacheProfile.TotalInputTokens, inputTokens) + } outputTokens = estimateClaudeOutputTokens(finalContent, rawThinkingContent, toolUses) h.recordSuccessForApiKey(apiKeyID, inputTokens, outputTokens, credits) h.pool.RecordSuccess(account.ID) + // h.pool.RecordLatency(account.ID, float64(time.Since(reqStart).Milliseconds())) + // ↑ dispatch-only latency EWMA (commit fbc99b1), absent on this cache-only + // branch off main. No-op here; stripped so handler compiles against main's + // pool. Re-added verbatim when dispatch lands. h.pool.UpdateStats(account.ID, inputTokens+outputTokens, credits) h.promptCache.Update(account.ID, cacheProfile) h.recordSuccessLog("claude", model, account.ID, inputTokens+outputTokens, credits, time.Since(reqStart).Milliseconds()) + recordClaudeCacheDispatch("success", model, "nonstream", account, cacheUsage, inputTokens, outputTokens) responseThinkingContent := rawThinkingContent includeEmptyThinkingBlock := thinking && thinkingOpts.OmitDisplay && rawThinkingContent != "" @@ -1559,7 +1693,7 @@ func (h *Handler) handleClaudeNonStream(w http.ResponseWriter, payload *KiroPayl } h.recordFailureWithDetails("claude", model, "", lastErr) - h.sendClaudeError(w, 500, "api_error", lastErr.Error()) + h.sendClaudeError(w, 500, "api_error", improperlyFormedClientMessage(lastErr)) } func (h *Handler) sendClaudeError(w http.ResponseWriter, status int, errType, message string) { @@ -1625,6 +1759,59 @@ func (h *Handler) handleOpenAIStream(w http.ResponseWriter, payload *KiroPayload return } + // 流式保活：与 handleClaudeStream 同构。整个流式期间周期性发送 SSE 注释行（OpenAI 流 + // 亦以 ':' 注释行合法、客户端忽略），落地 200 头并持续刷新 CF/客户端的 Proxy Read + // Timeout，既覆盖上游 prefill 迟迟无首字节、也覆盖输出中途长静默（opus 高强度 thinking + // 两个数据块之间、去掉 5min 总超时后的超长推理），避免连接在 ~120s 处被判超时断开 + // （用户侧表现为“说一半就断”）。 + // + // 并发：保活 goroutine 与主 goroutine 都会写 w，故所有真实 chunk 写出统一走 emitRaw()， + // 与保活共用 hbMu 串行化；lastWriteNano 记录最近写出时刻，保活仅在确有静默时才补发。 + var hbMu sync.Mutex + hbStopped := false + committed := false // 200 响应头是否已 flush 到客户端（由保活或首个真实 chunk 触发） + lastWriteNano := time.Now().UnixNano() + hbDone := make(chan struct{}) + var hbStopOnce sync.Once + stopHeartbeat := func() { + hbStopOnce.Do(func() { + hbMu.Lock() + hbStopped = true + hbMu.Unlock() + close(hbDone) + }) + } + defer stopHeartbeat() + + // emitRaw 串行化所有真实 chunk 写出，与保活 goroutine 互斥（共用 hbMu），并记录写出时刻。 + emitRaw := func(s string) { + hbMu.Lock() + fmt.Fprint(w, s) + flusher.Flush() + lastWriteNano = time.Now().UnixNano() + hbMu.Unlock() + } + + go func() { + ticker := time.NewTicker(streamKeepaliveInterval) + defer ticker.Stop() + for { + select { + case <-hbDone: + return + case <-ticker.C: + hbMu.Lock() + if !hbStopped && time.Since(time.Unix(0, lastWriteNano)) >= streamKeepaliveInterval { + fmt.Fprint(w, ": keepalive\n\n") + flusher.Flush() + committed = true + lastWriteNano = time.Now().UnixNano() + } + hbMu.Unlock() + } + } + }() + // 获取 thinking 输出格式配置 thinkingFormat := config.GetThinkingConfig().OpenAIFormat @@ -1633,7 +1820,11 @@ func (h *Handler) handleOpenAIStream(w http.ResponseWriter, payload *KiroPayload var lastErr error reqStart := time.Now() - for attempt := 0; attempt < maxAccountRetryAttempts; attempt++ { + retryBudget := resolveAccountRetryBudget(h.pool.Count()) + for attempt := 0; attempt < retryBudget; attempt++ { + if attempt > 0 { + time.Sleep(accountRetryBackoff(attempt - 1)) + } account := h.pool.GetNextForModelExcluding(model, excluded) if account == nil { break @@ -1753,8 +1944,7 @@ func (h *Handler) handleOpenAIStream(w http.ResponseWriter, payload *KiroPayload } } data, _ := json.Marshal(chunk) - fmt.Fprintf(w, "data: %s\n\n", string(data)) - flusher.Flush() + emitRaw(fmt.Sprintf("data: %s\n\n", string(data))) responseStarted = true } @@ -1910,8 +2100,7 @@ func (h *Handler) handleOpenAIStream(w http.ResponseWriter, payload *KiroPayload } toolCallIndex++ data, _ := json.Marshal(chunk) - fmt.Fprintf(w, "data: %s\n\n", string(data)) - flusher.Flush() + emitRaw(fmt.Sprintf("data: %s\n\n", string(data))) responseStarted = true }, OnComplete: func(inTok, outTok int) { @@ -1931,6 +2120,9 @@ func (h *Handler) handleOpenAIStream(w http.ResponseWriter, payload *KiroPayload lastErr = err excluded[account.ID] = true h.handleAccountFailure(account, err) + if isUpstreamPermanentError(err) { + break + } if !responseStarted { continue } @@ -1964,6 +2156,10 @@ func (h *Handler) handleOpenAIStream(w http.ResponseWriter, payload *KiroPayload h.recordSuccessForApiKey(apiKeyID, inputTokens, outputTokens, credits) h.pool.RecordSuccess(account.ID) + // h.pool.RecordLatency(account.ID, float64(time.Since(reqStart).Milliseconds())) + // ↑ dispatch-only latency EWMA (commit fbc99b1), absent on this cache-only + // branch off main. No-op here; stripped so handler compiles against main's + // pool. Re-added verbatim when dispatch lands. h.pool.UpdateStats(account.ID, inputTokens+outputTokens, credits) h.recordSuccessLog("openai", model, account.ID, inputTokens+outputTokens, credits, time.Since(reqStart).Milliseconds()) @@ -1989,9 +2185,44 @@ func (h *Handler) handleOpenAIStream(w http.ResponseWriter, payload *KiroPayload }, } data, _ := json.Marshal(chunk) - fmt.Fprintf(w, "data: %s\n\n", string(data)) - fmt.Fprintf(w, "data: [DONE]\n\n") - flusher.Flush() + emitRaw(fmt.Sprintf("data: %s\n\n", string(data))) + emitRaw("data: [DONE]\n\n") + return + } + + // 循环外的错误/无账号路径。保活 goroutine 可能仍在写 w，先停掉再决定响应形态， + // 避免 sendOpenAIError 的 WriteHeader+JSON 与心跳的 ': keepalive' 并发写同一个 w。 + stopHeartbeat() + if lastErr != nil { + h.recordFailureWithDetails("openai", model, "", lastErr) + } + hbMu.Lock() + didCommit := committed + hbMu.Unlock() + + if didCommit { + // 心跳已落地 200 响应头：HTTP 状态码已无法更改，客户端正按 SSE 解析——必须以 + // SSE error chunk + [DONE] 收尾，而非 sendOpenAIError 的 WriteHeader+JSON（后者 + // 会触发 superfluous WriteHeader 且破坏流式协议）。 + msg := "No available accounts" + if lastErr != nil { + msg = improperlyFormedClientMessage(lastErr) + } + errChunk := map[string]interface{}{ + "id": chatID, + "object": "chat.completion.chunk", + "created": time.Now().Unix(), + "model": model, + "choices": []map[string]interface{}{{ + "index": 0, + "delta": map[string]interface{}{}, + "finish_reason": "error", + }}, + "error": map[string]string{"message": msg}, + } + errData, _ := json.Marshal(errChunk) + emitRaw(fmt.Sprintf("data: %s\n\n", string(errData))) + emitRaw("data: [DONE]\n\n") return } @@ -2000,8 +2231,7 @@ func (h *Handler) handleOpenAIStream(w http.ResponseWriter, payload *KiroPayload return } - h.recordFailureWithDetails("openai", model, "", lastErr) - h.sendOpenAIError(w, 500, "server_error", lastErr.Error()) + h.sendOpenAIError(w, 500, "server_error", improperlyFormedClientMessage(lastErr)) } // handleOpenAINonStream OpenAI 非流式响应 @@ -2010,7 +2240,11 @@ func (h *Handler) handleOpenAINonStream(w http.ResponseWriter, payload *KiroPayl var lastErr error reqStart := time.Now() - for attempt := 0; attempt < maxAccountRetryAttempts; attempt++ { + retryBudget := resolveAccountRetryBudget(h.pool.Count()) + for attempt := 0; attempt < retryBudget; attempt++ { + if attempt > 0 { + time.Sleep(accountRetryBackoff(attempt - 1)) + } account := h.pool.GetNextForModelExcluding(model, excluded) if account == nil { break @@ -2050,6 +2284,9 @@ func (h *Handler) handleOpenAINonStream(w http.ResponseWriter, payload *KiroPayl lastErr = err excluded[account.ID] = true h.handleAccountFailure(account, err) + if isUpstreamPermanentError(err) { + break + } continue } @@ -2069,6 +2306,10 @@ func (h *Handler) handleOpenAINonStream(w http.ResponseWriter, payload *KiroPayl h.recordSuccessForApiKey(apiKeyID, inputTokens, outputTokens, credits) h.pool.RecordSuccess(account.ID) + // h.pool.RecordLatency(account.ID, float64(time.Since(reqStart).Milliseconds())) + // ↑ dispatch-only latency EWMA (commit fbc99b1), absent on this cache-only + // branch off main. No-op here; stripped so handler compiles against main's + // pool. Re-added verbatim when dispatch lands. h.pool.UpdateStats(account.ID, inputTokens+outputTokens, credits) h.recordSuccessLog("openai", model, account.ID, inputTokens+outputTokens, credits, time.Since(reqStart).Milliseconds()) @@ -2085,7 +2326,7 @@ func (h *Handler) handleOpenAINonStream(w http.ResponseWriter, payload *KiroPayl } h.recordFailureWithDetails("openai", model, "", lastErr) - h.sendOpenAIError(w, 500, "server_error", lastErr.Error()) + h.sendOpenAIError(w, 500, "server_error", improperlyFormedClientMessage(lastErr)) } func (h *Handler) sendOpenAIError(w http.ResponseWriter, status int, errType, message string) { @@ -2099,32 +2340,43 @@ func (h *Handler) sendOpenAIError(w http.ResponseWriter, status int, errType, me }) } -// ensureValidToken 确保 token 有效 -func (h *Handler) ensureValidToken(account *config.Account) error { - if account.ExpiresAt == 0 || time.Now().Unix() < account.ExpiresAt-tokenRefreshSkewSeconds { +// refreshAccountToken is the single locked entry point for refreshing an +// account's access token. Every caller — the background refresh sweep, the +// per-request ensureValidToken, and the admin refresh endpoints — must funnel +// through here so concurrent refreshes of the same account deduplicate rather +// than racing the IdP. IdPs that rotate refresh tokens (one-time-use — e.g. +// external_idp / Azure AD; see the comment in auth/oidc.go refreshExternalIdpToken) +// reject a refresh that reuses an already-consumed token with invalid_grant, and +// that flows into handleAccountFailure → disableAccount("BANNED"). +// +// force bypasses the not-near-expiry fast path for explicit admin "refresh now" +// actions and the 401/403 retry path. +func (h *Handler) refreshAccountToken(account *config.Account, force bool) error { + if !force && (account.ExpiresAt == 0 || time.Now().Unix() < account.ExpiresAt-tokenRefreshSkewSeconds) { return nil } h.tokenRefreshMu.Lock() defer h.tokenRefreshMu.Unlock() - // Another concurrent request may have refreshed this account while we waited. + // Another goroutine may have refreshed this account while we waited on the + // lock. Adopt the freshest persisted token before refreshing so we never + // reuse an already-rotated (one-time-use) refresh token. if latest := h.pool.GetByID(account.ID); latest != nil { account.AccessToken = latest.AccessToken account.RefreshToken = latest.RefreshToken account.ExpiresAt = latest.ExpiresAt account.ProfileArn = latest.ProfileArn - if account.ExpiresAt == 0 || time.Now().Unix() < account.ExpiresAt-tokenRefreshSkewSeconds { + if !force && (account.ExpiresAt == 0 || time.Now().Unix() < account.ExpiresAt-tokenRefreshSkewSeconds) { return nil } } - accessToken, refreshToken, expiresAt, profileArn, err := auth.RefreshToken(account) + accessToken, refreshToken, expiresAt, profileArn, err := authRefreshToken(account) if err != nil { return err } - // 更新内存 h.pool.UpdateToken(account.ID, accessToken, refreshToken, expiresAt) account.AccessToken = accessToken if refreshToken != "" { @@ -2135,13 +2387,15 @@ func (h *Handler) ensureValidToken(account *config.Account) error { account.ProfileArn = profileArn config.UpdateAccountProfileArn(account.ID, profileArn) } - - // 持久化 config.UpdateAccountToken(account.ID, accessToken, refreshToken, expiresAt) - return nil } +// ensureValidToken 确保 token 有效 +func (h *Handler) ensureValidToken(account *config.Account) error { + return h.refreshAccountToken(account, false) +} + // ==================== 管理 API ==================== func (h *Handler) handleAdminAPI(w http.ResponseWriter, r *http.Request) { @@ -2154,7 +2408,8 @@ func (h *Handler) handleAdminAPI(w http.ResponseWriter, r *http.Request) { } } - if password != config.GetPassword() { + expected := config.GetPassword() + if expected != "" && subtle.ConstantTimeCompare([]byte(password), []byte(expected)) != 1 { w.WriteHeader(401) json.NewEncoder(w).Encode(map[string]string{"error": "Unauthorized"}) return @@ -2211,10 +2466,20 @@ func (h *Handler) handleAdminAPI(w http.ResponseWriter, r *http.Request) { h.apiStartBuilderIdLogin(w, r) case path == "/auth/builderid/poll" && r.Method == "POST": h.apiPollBuilderIdAuth(w, r) + case path == "/auth/kiro-sso/start" && r.Method == "POST": + h.apiStartKiroSso(w, r) + case path == "/auth/kiro-sso/poll" && r.Method == "POST": + h.apiPollKiroSso(w, r) + case path == "/auth/kiro-sso/cancel" && r.Method == "POST": + h.apiCancelKiroSso(w, r) + case path == "/auth/kiro-sso/callback" && r.Method == "POST": + h.apiCompleteKiroSso(w, r) case path == "/auth/sso-token" && r.Method == "POST": h.apiImportSsoToken(w, r) case path == "/auth/credentials" && r.Method == "POST": h.apiImportCredentials(w, r) + case path == "/auth/apikeys-batch" && r.Method == "POST": + h.apiImportApiKeys(w, r) case path == "/status" && r.Method == "GET": h.apiGetStatus(w, r) case path == "/settings" && r.Method == "GET": @@ -2347,6 +2612,21 @@ func (h *Handler) apiAddAccount(w http.ResponseWriter, r *http.Request) { account.Region = "us-east-1" } + // api_key: use the Kiro API key directly as the bearer (tokentype: API_KEY). + // Mirror it into AccessToken for pool/dispatch/metrics compatibility, zero + // ExpiresAt + RefreshToken so no refresh path ever fires. + if account.IsApiKeyCredential() { + if account.KiroApiKey == "" { + w.WriteHeader(400) + json.NewEncoder(w).Encode(map[string]string{"error": "kiroApiKey is required for api_key accounts"}) + return + } + account.AuthMethod = "api_key" + account.AccessToken = account.KiroApiKey + account.RefreshToken = "" + account.ExpiresAt = 0 + } + if err := config.AddAccount(account); err != nil { w.WriteHeader(500) json.NewEncoder(w).Encode(map[string]string{"error": err.Error()}) @@ -2355,10 +2635,10 @@ func (h *Handler) apiAddAccount(w http.ResponseWriter, r *http.Request) { h.pool.Reload() // 新账号若已启用且有 token，立即拉取并缓存模型列表 - if account.Enabled && account.AccessToken != "" { + if account.Enabled && (account.AccessToken != "" || account.IsApiKeyCredential()) { go func(acc config.Account) { if err := h.fetchAndCacheAccountModels(&acc); err != nil { - logger.Warnf("[ModelsCache] Auto-refresh failed for new account %s: %v", acc.Email, err) + logger.Warnf("[ModelsCache] Auto-refresh failed for new account %s: %v", accountEmailForLog(&acc), err) } }(account) } @@ -2427,7 +2707,7 @@ func (h *Handler) apiUpdateAccount(w http.ResponseWriter, r *http.Request, id st if !oldEnabled && existing.Enabled && existing.AccessToken != "" { go func(acc config.Account) { if err := h.fetchAndCacheAccountModels(&acc); err != nil { - logger.Warnf("[ModelsCache] Auto-refresh failed for re-enabled account %s: %v", acc.Email, err) + logger.Warnf("[ModelsCache] Auto-refresh failed for re-enabled account %s: %v", accountEmailForLog(&acc), err) } }(*existing) } @@ -2592,20 +2872,10 @@ func (h *Handler) apiBatchAccounts(w http.ResponseWriter, r *http.Request) { failCount++ continue } - // 刷新 token + // 刷新 token（去重 + 持久化由 refreshAccountToken 统一处理） if account.RefreshToken != "" { - if newAccess, newRefresh, newExpires, profileArn, err := auth.RefreshToken(account); err == nil { - account.AccessToken = newAccess - if newRefresh != "" { - account.RefreshToken = newRefresh - } - account.ExpiresAt = newExpires - config.UpdateAccountToken(id, newAccess, newRefresh, newExpires) - if profileArn != "" { - account.ProfileArn = profileArn - config.UpdateAccountProfileArn(id, profileArn) - } - h.pool.UpdateToken(id, newAccess, newRefresh, newExpires) + if err := h.refreshAccountToken(account, true); err != nil { + logger.Warnf("[ApiBatch] Token refresh failed for %s: %v", account.Email, err) } } // 刷新账户信息 @@ -2805,6 +3075,159 @@ func (h *Handler) apiPollBuilderIdAuth(w http.ResponseWriter, r *http.Request) { }) } +// apiStartKiroSso starts the Kiro hosted-portal sign-in (Enterprise SSO — Microsoft 365 / +// Entra ID, plus Google/GitHub). It binds the loopback callback listener and returns the +// sign-in URL the operator opens in a browser ON THE SAME HOST as the proxy (the OAuth +// redirect targets 127.0.0.1:3128). The browser is driven through the enterprise external-IdP +// leg automatically; the front end polls /auth/kiro-sso/poll until completion. +func (h *Handler) apiStartKiroSso(w http.ResponseWriter, r *http.Request) { + var req struct { + Region string `json:"region"` + } + // Region is optional (defaults to us-east-1 in StartKiroSsoLogin), so a decode + // error (including an empty body) is intentionally tolerated — mirrors + // apiStartBuilderIdLogin. + json.NewDecoder(r.Body).Decode(&req) + + session, signInURL, err := auth.StartKiroSsoLogin(req.Region) + if err != nil { + w.WriteHeader(500) + json.NewEncoder(w).Encode(map[string]string{"error": err.Error()}) + return + } + + json.NewEncoder(w).Encode(map[string]interface{}{ + "sessionId": session.ID, + "signInUrl": signInURL, + "interval": 2, + }) +} + +// apiCancelKiroSso tears down an in-flight hosted-portal sign-in (operator closed or +// cancelled the modal), freeing the loopback callback port immediately instead of +// waiting for the deadline. +func (h *Handler) apiCancelKiroSso(w http.ResponseWriter, r *http.Request) { + var req struct { + SessionID string `json:"sessionId"` + } + json.NewDecoder(r.Body).Decode(&req) + if req.SessionID != "" { + auth.CancelKiroSsoLogin(req.SessionID) + } + json.NewEncoder(w).Encode(map[string]interface{}{"success": true}) +} + +// apiCompleteKiroSso lets an operator paste the OAuth callback URL for remote +// deployments where localhost isn't the proxy host. The callback URL is fed +// through the session's state machine; if a redirect is returned (enterprise +// SSO leg-1) the front end should open it in a browser. +func (h *Handler) apiCompleteKiroSso(w http.ResponseWriter, r *http.Request) { + var req struct { + SessionID string `json:"session_id"` + CallbackURL string `json:"callback_url"` + } + if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + w.WriteHeader(400) + json.NewEncoder(w).Encode(map[string]string{"error": "Invalid JSON"}) + return + } + if req.SessionID == "" || req.CallbackURL == "" { + w.WriteHeader(400) + json.NewEncoder(w).Encode(map[string]string{"error": "session_id and callback_url are required"}) + return + } + + redirectURL, err := auth.FeedCallbackURL(req.SessionID, req.CallbackURL) + if err != nil { + w.WriteHeader(400) + json.NewEncoder(w).Encode(map[string]string{"error": err.Error()}) + return + } + + if redirectURL != "" { + json.NewEncoder(w).Encode(map[string]interface{}{ + "status": "redirect", + "authorize_url": redirectURL, + }) + return + } + + json.NewEncoder(w).Encode(map[string]interface{}{ + "status": "accepted", + }) +} + +// apiPollKiroSso reports the hosted-portal sign-in status. While the user is signing in it +// returns completed=false; once the listener captures the authorization code it exchanges it, +// persists the account (AuthMethod "external_idp" for an Azure tenant, "social" otherwise), and +// returns completed=true. The profileArn is resolved lazily on first use (the EXTERNAL_IDP +// token type header is now sent on CodeWhisperer calls), so it is not required here. +func (h *Handler) apiPollKiroSso(w http.ResponseWriter, r *http.Request) { + var req struct { + SessionID string `json:"sessionId"` + } + if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + w.WriteHeader(400) + json.NewEncoder(w).Encode(map[string]string{"error": "Invalid JSON"}) + return + } + + result, status, err := auth.PollKiroSsoAuth(req.SessionID) + if err != nil { + w.WriteHeader(400) + json.NewEncoder(w).Encode(map[string]interface{}{ + "success": false, + "error": err.Error(), + }) + return + } + + if status == "pending" { + json.NewEncoder(w).Encode(map[string]interface{}{ + "success": true, + "completed": false, + "status": "pending", + }) + return + } + + // 授权完成，创建账号 + account := config.Account{ + ID: auth.GenerateAccountID(), + Email: result.Email, + AccessToken: result.AccessToken, + RefreshToken: result.RefreshToken, + ClientID: result.ClientID, + AuthMethod: result.AuthMethod, + Provider: result.Provider, + Region: result.Region, + ProfileArn: result.ProfileArn, + TokenEndpoint: result.TokenEndpoint, + IssuerURL: result.IssuerURL, + Scopes: result.Scopes, + ExpiresAt: time.Now().Unix() + int64(result.ExpiresIn), + Enabled: true, + MachineId: config.GenerateMachineId(), + } + + if err := config.AddAccount(account); err != nil { + w.WriteHeader(500) + json.NewEncoder(w).Encode(map[string]string{"error": err.Error()}) + return + } + + h.pool.Reload() + json.NewEncoder(w).Encode(map[string]interface{}{ + "success": true, + "completed": true, + "account": map[string]interface{}{ + "id": account.ID, + "email": account.Email, + "authMethod": account.AuthMethod, + }, + }) +} + func (h *Handler) apiImportSsoToken(w http.ResponseWriter, r *http.Request) { var req struct { BearerToken string `json:"bearerToken"` @@ -2887,6 +3310,11 @@ func (h *Handler) apiImportSsoToken(w http.ResponseWriter, r *http.Request) { } func (h *Handler) apiImportCredentials(w http.ResponseWriter, r *http.Request) { + // Cap the body: accessToken becomes attacker-influenced input that is base64- and + // JSON-decoded twice (issuerFromAccessTokenJWT / ExpFromAccessTokenJWT). Without a + // limit an oversized token is a memory-amplification DoS. Mirrors the io.LimitReader + // guard on outbound IdP responses in auth/kiro_sso.go's oidcDiscover. + r.Body = http.MaxBytesReader(w, r.Body, 1<<20) var req struct { AccessToken string `json:"accessToken"` RefreshToken string `json:"refreshToken"` @@ -2895,6 +3323,22 @@ func (h *Handler) apiImportCredentials(w http.ResponseWriter, r *http.Request) { AuthMethod string `json:"authMethod"` Provider string `json:"provider"` Region string `json:"region"` + // external_idp (enterprise SSO / Azure AD) refresh material. + TokenEndpoint string `json:"tokenEndpoint"` + IssuerURL string `json:"issuerUrl"` + Scopes string `json:"scopes"` + // api_key credential: Kiro API key used directly as bearer (no OAuth refresh). + KiroApiKey string `json:"kiroApiKey"` + AuthRegion string `json:"authRegion"` + ApiRegion string `json:"apiRegion"` + // Optional identity preservation when pasting a full account record. + ID string `json:"id"` + Email string `json:"email"` + Nickname string `json:"nickname"` + ProfileArn string `json:"profileArn"` + // userId (account-level in Kiro Account Manager exports) embeds the Azure + // tenant, from which tokenEndpoint/issuerUrl/scopes are derived when missing. + UserID string `json:"userId"` } if err := json.NewDecoder(r.Body).Decode(&req); err != nil { w.WriteHeader(400) @@ -2902,6 +3346,62 @@ func (h *Handler) apiImportCredentials(w http.ResponseWriter, r *http.Request) { return } + // api_key: import a Kiro API key used directly as bearer (no OAuth refresh, + // no profile ARN). Mirror the key into AccessToken for pool/dispatch/metrics + // compatibility. Best-effort RefreshAccountInfo populates email/usage/credit. + if req.KiroApiKey != "" || strings.EqualFold(req.AuthMethod, "api_key") || strings.EqualFold(req.AuthMethod, "apikey") { + if req.KiroApiKey == "" { + w.WriteHeader(400) + json.NewEncoder(w).Encode(map[string]string{"error": "kiroApiKey is required for api_key accounts"}) + return + } + if req.Region == "" { + req.Region = "us-east-1" + } + id := req.ID + if id == "" || config.AccountIDExists(id) { + id = auth.GenerateAccountID() + } + account := config.Account{ + ID: id, + Email: req.Email, + Nickname: req.Nickname, + AuthMethod: "api_key", + KiroApiKey: req.KiroApiKey, + AccessToken: req.KiroApiKey, + Region: req.Region, + AuthRegion: req.AuthRegion, + ApiRegion: req.ApiRegion, + ExpiresAt: 0, + Enabled: true, + MachineId: config.GenerateMachineId(), + } + if err := config.AddAccount(account); err != nil { + w.WriteHeader(500) + json.NewEncoder(w).Encode(map[string]string{"error": err.Error()}) + return + } + h.pool.Reload() + // Best-effort: populate email/usage/credit from upstream. RefreshAccountInfo + // uses KiroApiKey (mirrored into AccessToken) as the bearer via + // applyKiroBaseHeaders; profile-ARN resolution is skipped for api_key. + go func(acc config.Account) { + if _, infoErr := RefreshAccountInfo(&acc); infoErr != nil { + logger.Warnf("[Import] RefreshAccountInfo failed for api_key account %s: %v", accountEmailForLog(&acc), infoErr) + } + }(account) + go func(acc config.Account) { + if err := h.fetchAndCacheAccountModels(&acc); err != nil { + logger.Warnf("[ModelsCache] Auto-refresh failed for api_key account %s: %v", accountEmailForLog(&acc), err) + } + }(account) + json.NewEncoder(w).Encode(map[string]interface{}{ + "success": true, + "account": map[string]interface{}{"id": account.ID, "email": account.Email}, + }) + return + } + if req.RefreshToken == "" { w.WriteHeader(400) json.NewEncoder(w).Encode(map[string]string{"error": "refreshToken is required"}) @@ -2912,71 +3412,158 @@ func (h *Handler) apiImportCredentials(w http.ResponseWriter, r *http.Request) { if req.Region == "" { req.Region = "us-east-1" } - if req.AuthMethod == "" { - if req.ClientID != "" { - req.AuthMethod = "idc" - } else { - req.AuthMethod = "social" + // 标准化 authMethod。external_idp 必须先于 clientId+clientSecret→idc 的推断被识别 + //（external_idp 带 clientId 但没有 clientSecret），否则会被误判成 social 而 refresh 到错误端点。 + req.AuthMethod = normalizeImportAuthMethod(req.AuthMethod, req.ClientID, req.ClientSecret, req.TokenEndpoint) + + // Resolve Azure endpoints from userId (Kiro export, account level) or the + // accessToken JWT issuer (bare blobs: clientId + token only). A derivation that + // also clears the allow-list is itself proof the credential is external_idp — + // IdC/social access tokens are not microsoftonline JWTs, so a bare IdC blob (its + // iss is an AWS host) won't clear the list and won't be misclassified. + derivedTE, derivedIss, derivedSc := auth.DeriveExternalIdpEndpoints(req.UserID, req.ClientID, req.AccessToken) + if derivedTE != "" && auth.ValidateExternalIdpEndpoint(derivedTE) == nil && req.AuthMethod != "external_idp" { + req.AuthMethod = "external_idp" + } + + // external_idp 的 tokenEndpoint 是用户可填的新信任边界：必须经 allow-list 校验， + // 否则一份不信任的 credential JSON 可指向内网/攻击者主机，导致 refresh token 被外泄。 + if req.AuthMethod == "external_idp" { + // Kiro Account Manager exports and bare blobs omit tokenEndpoint/issuerUrl/ + // scopes; fill them from the derived (userId or accessToken-JWT) tenant. + if req.TokenEndpoint == "" { + req.TokenEndpoint = derivedTE } - } - // 标准化 authMethod - switch strings.ToLower(req.AuthMethod) { - case "idc", "builderid", "enterprise": - req.AuthMethod = "idc" - case "social", "google", "github": - req.AuthMethod = "social" - default: - if req.ClientID != "" && req.ClientSecret != "" { - req.AuthMethod = "idc" - } else { - req.AuthMethod = "social" + if req.IssuerURL == "" { + req.IssuerURL = derivedIss + } + if req.Scopes == "" { + req.Scopes = derivedSc + } + if req.ClientID == "" || req.TokenEndpoint == "" { + w.WriteHeader(400) + json.NewEncoder(w).Encode(map[string]string{"error": "external_idp requires clientId and tokenEndpoint (or userId/accessToken to derive it)"}) + return + } + if err := auth.ValidateExternalIdpEndpoint(req.TokenEndpoint); err != nil { + w.WriteHeader(400) + json.NewEncoder(w).Encode(map[string]string{"error": "external IdP endpoint rejected: " + err.Error()}) + return + } + if req.IssuerURL != "" { + if err := auth.ValidateExternalIdpEndpoint(req.IssuerURL); err != nil { + w.WriteHeader(400) + json.NewEncoder(w).Encode(map[string]string{"error": "external IdP issuer rejected: " + err.Error()}) + return + } } } - // 用 refreshToken 刷新获取新的 accessToken。导入必须以一次成功的刷新为前提： - // 本地缓存里的 accessToken 不携带可信的过期时间，盲猜短 TTL 会让账号在选号时 - // 永远被跳过，导致后台/按需刷新都无法触发（详见 ensureValidToken 与 Pick 的过期判定）。 - tempAccount := &config.Account{ - RefreshToken: req.RefreshToken, - ClientID: req.ClientID, - ClientSecret: req.ClientSecret, - AuthMethod: req.AuthMethod, - Region: req.Region, + // Resolve the access token to persist. For external_idp we prefer TRUST-ON-IMPORT: + // when the pasted JSON carries an Azure AD access token (a JWT with a real exp), + // persist it directly WITHOUT a live refresh round-trip. The JSON can then be + // imported repeatedly / into multiple instances without each import consuming + // (rotating) the refresh token, and without requiring egress to Microsoft at + // import time. The runtime background refresh (backgroundRefresh / + // ensureValidToken) renews it later when the account is actually used. Falls + // back to refresh-at-import for idc/social and for external_idp credentials + // carrying only a refreshToken (so the regression gate — reject when refresh + // fails — still holds there). + var ( + accessToken string + expiresAt int64 + profileArn string + ) + email := req.Email + if req.AuthMethod == "external_idp" && req.AccessToken != "" { + if exp := auth.ExpFromAccessTokenJWT(req.AccessToken); exp > 0 { + accessToken = req.AccessToken + expiresAt = exp + profileArn = req.ProfileArn + } } - accessToken, newRefreshToken, expiresAt, newProfileArn, err := auth.RefreshToken(tempAccount) - if err != nil { - w.WriteHeader(400) - json.NewEncoder(w).Encode(map[string]string{"error": "Token refresh failed: " + err.Error()}) - return + if accessToken == "" { + tempAccount := &config.Account{ + RefreshToken: req.RefreshToken, + ClientID: req.ClientID, + ClientSecret: req.ClientSecret, + AuthMethod: req.AuthMethod, + Region: req.Region, + TokenEndpoint: req.TokenEndpoint, + Scopes: req.Scopes, + } + a, newRT, ea, newPA, err := authRefreshToken(tempAccount) + if err != nil { + w.WriteHeader(400) + json.NewEncoder(w).Encode(map[string]string{"error": "Token refresh failed: " + err.Error()}) + return + } + accessToken = a + expiresAt = ea + profileArn = newPA + if newRT != "" { + req.RefreshToken = newRT + } + if fetchedEmail, _, _ := auth.GetUserInfo(accessToken); fetchedEmail != "" { + email = fetchedEmail + } } - if newRefreshToken != "" { - req.RefreshToken = newRefreshToken + if profileArn == "" { + profileArn = req.ProfileArn // external_idp refresh returns no profileArn } - // 获取用户信息 - email, _, _ := auth.GetUserInfo(accessToken) - // 创建账号 - account := config.Account{ - ID: auth.GenerateAccountID(), - Email: email, - AccessToken: accessToken, - RefreshToken: req.RefreshToken, - ClientID: req.ClientID, - ClientSecret: req.ClientSecret, - AuthMethod: req.AuthMethod, - Provider: req.Provider, - Region: req.Region, - ExpiresAt: expiresAt, - Enabled: true, - MachineId: config.GenerateMachineId(), - ProfileArn: newProfileArn, + provider := req.Provider + if provider == "" && req.AuthMethod == "external_idp" { + provider = "AzureAD" + } + // Re-importing a backup must never create a duplicate entry. If an account + // with the same (resolved) refresh token already exists, update it in place + // (reusing its id) instead of minting a fresh id and leaving two live + // accounts sharing the token — which would interleave refresh/rotation + // conflicts. Mirrors the bulk import path's dedup (config.AddAccounts). + existingID := "" + if req.RefreshToken != "" { + existingID = config.FindAccountIDByRefreshToken(req.RefreshToken) + } + id := req.ID + if existingID != "" { + id = existingID + } else if id == "" || config.AccountIDExists(id) { + id = auth.GenerateAccountID() } - - if err := config.AddAccount(account); err != nil { - w.WriteHeader(500) - json.NewEncoder(w).Encode(map[string]string{"error": err.Error()}) - return + account := config.Account{ + ID: id, + Email: email, + AccessToken: accessToken, + RefreshToken: req.RefreshToken, + ClientID: req.ClientID, + ClientSecret: req.ClientSecret, + AuthMethod: req.AuthMethod, + Provider: provider, + Region: req.Region, + ExpiresAt: expiresAt, + Enabled: true, + MachineId: config.GenerateMachineId(), + ProfileArn: profileArn, + TokenEndpoint: req.TokenEndpoint, + IssuerURL: req.IssuerURL, + Scopes: req.Scopes, + } + + if existingID != "" { + // Update the existing entry in place rather than creating a duplicate. + if err := config.UpdateAccount(id, account); err != nil { + w.WriteHeader(500) + json.NewEncoder(w).Encode(map[string]string{"error": err.Error()}) + return + } + } else { + if err := config.AddAccount(account); err != nil { + w.WriteHeader(500) + json.NewEncoder(w).Encode(map[string]string{"error": err.Error()}) + return + } } h.pool.Reload() @@ -2989,27 +3576,80 @@ func (h *Handler) apiImportCredentials(w http.ResponseWriter, r *http.Request) { }) } +// externalIdpAuthMethodAliases are lower-cased authMethod values (or Kiro Account +// Manager provider labels) that mean "external IdP / enterprise SSO" and must +// normalize to "external_idp". +var externalIdpAuthMethodAliases = map[string]bool{ + "external_idp": true, + "azuread": true, + "azure": true, + "entra": true, + "entra-id": true, + "entra_id": true, + "microsoft": true, + "m365": true, + "office365": true, + "external": true, +} + +// normalizeImportAuthMethod maps a pasted credential JSON's authMethod (plus its +// clientId/clientSecret/tokenEndpoint) onto one of the three canonical methods +// ("external_idp" | "idc" | "social"). external_idp MUST be detected before the +// clientId+clientSecret→idc inference, because external_idp accounts carry clientId +// but NO clientSecret, so the old default branch misclassified them as "social" and +// refresh hit the wrong endpoint. +// +// It preserves the pre-existing idc/social heuristics: +// - empty authMethod + clientId present -> idc +// - empty authMethod, no clientId -> social +// - "enterprise" (Kiro Account Manager IdC label) -> idc +// - unrecognized non-empty + clientId+clientSecret -> idc, else social +func normalizeImportAuthMethod(authMethod, clientID, clientSecret, tokenEndpoint string) string { + am := strings.ToLower(strings.TrimSpace(authMethod)) + switch { + case externalIdpAuthMethodAliases[am]: + return "external_idp" + case tokenEndpoint != "": // infer when not declared explicitly + return "external_idp" + case am == "social" || am == "google" || am == "github": + return "social" + case am == "idc" || am == "builderid" || am == "enterprise": + return "idc" + } + if am == "" { + if clientID != "" { + return "idc" + } + return "social" + } + if clientID != "" && clientSecret != "" { + return "idc" + } + return "social" +} + func (h *Handler) apiGetStatus(w http.ResponseWriter, r *http.Request) { json.NewEncoder(w).Encode(map[string]interface{}{ "version": config.Version, "accounts": h.pool.Count(), "available": h.pool.AvailableCount(), - "totalRequests": h.totalRequests, - "successRequests": h.successRequests, - "failedRequests": h.failedRequests, - "totalTokens": h.totalTokens, - "totalCredits": h.totalCredits, + "totalRequests": atomic.LoadInt64(&h.totalRequests), + "successRequests": atomic.LoadInt64(&h.successRequests), + "failedRequests": atomic.LoadInt64(&h.failedRequests), + "totalTokens": atomic.LoadInt64(&h.totalTokens), + "totalCredits": h.getCredits(), "uptime": time.Now().Unix() - h.startTime, }) } func (h *Handler) apiGetSettings(w http.ResponseWriter, r *http.Request) { json.NewEncoder(w).Encode(map[string]interface{}{ - "apiKey": config.GetApiKey(), - "requireApiKey": config.IsApiKeyRequired(), - "port": config.GetPort(), - "host": config.GetHost(), - "allowOverUsage": config.GetAllowOverUsage(), + "apiKey": config.GetApiKey(), + "requireApiKey": config.IsApiKeyRequired(), + "port": config.GetPort(), + "host": config.GetHost(), + "allowOverUsage": config.GetAllowOverUsage(), + "maxPayloadBytes": config.GetMaxPayloadBytes(), }) } @@ -3058,10 +3698,11 @@ func (h *Handler) apiUpdatePromptFilter(w http.ResponseWriter, r *http.Request) func (h *Handler) apiUpdateSettings(w http.ResponseWriter, r *http.Request) { var req struct { - ApiKey *string `json:"apiKey,omitempty"` - RequireApiKey *bool `json:"requireApiKey,omitempty"` - Password string `json:"password,omitempty"` - AllowOverUsage *bool `json:"allowOverUsage,omitempty"` + ApiKey *string `json:"apiKey,omitempty"` + RequireApiKey *bool `json:"requireApiKey,omitempty"` + Password string `json:"password,omitempty"` + AllowOverUsage *bool `json:"allowOverUsage,omitempty"` + MaxPayloadBytes *int `json:"maxPayloadBytes,omitempty"` } if err := json.NewDecoder(r.Body).Decode(&req); err != nil { w.WriteHeader(400) @@ -3086,6 +3727,14 @@ func (h *Handler) apiUpdateSettings(w http.ResponseWriter, r *http.Request) { h.pool.Reload() } + if req.MaxPayloadBytes != nil { + if err := config.UpdateMaxPayloadBytes(*req.MaxPayloadBytes); err != nil { + w.WriteHeader(500) + json.NewEncoder(w).Encode(map[string]string{"error": err.Error()}) + return + } + } + json.NewEncoder(w).Encode(map[string]bool{"success": true}) } @@ -3220,22 +3869,7 @@ func (h *Handler) apiRefreshAccount(w http.ResponseWriter, r *http.Request, id s if account.RefreshToken == "" { return nil } - newAccessToken, newRefreshToken, newExpiresAt, profileArn, err := auth.RefreshToken(account) - if err != nil { - return err - } - account.AccessToken = newAccessToken - if newRefreshToken != "" { - account.RefreshToken = newRefreshToken - } - account.ExpiresAt = newExpiresAt - config.UpdateAccountToken(id, newAccessToken, newRefreshToken, newExpiresAt) - h.pool.UpdateToken(id, newAccessToken, newRefreshToken, newExpiresAt) - if profileArn != "" { - account.ProfileArn = profileArn - config.UpdateAccountProfileArn(id, profileArn) - } - return nil + return h.refreshAccountToken(account, true) } // 检查 token 是否快过期，先刷新 @@ -3262,7 +3896,7 @@ func (h *Handler) apiRefreshAccount(w http.ResponseWriter, r *http.Request, id s } // 如果是 403/401，说明 token 无效，尝试刷新后重试 - if strings.Contains(errMsg, "403") || strings.Contains(errMsg, "401") || strings.Contains(errMsg, "invalid") || strings.Contains(errMsg, "expired") { + if isAuthErrorMessage(errMsg) { if refreshErr := refreshTokenIfNeeded(); refreshErr == nil { // 重试 info, err = RefreshAccountInfo(account) diff --git a/proxy/handler_keepalive_test.go b/proxy/handler_keepalive_test.go new file mode 100644 index 00000000..211c2120 --- /dev/null +++ b/proxy/handler_keepalive_test.go @@ -0,0 +1,200 @@ +package proxy + +import ( + "net/http" + "net/http/httptest" + "strings" + "sync" + "testing" + "time" + + accountpool "kiro-go/pool" + "kiro-go/config" +) + +// lockedFlushRecorder 是一个并发安全的 ResponseWriter+Flusher。流式 handler 的保活 +// goroutine 与主 goroutine 会并发写 w（上游中途静默时，保活补发 ': keepalive' 而主 +// goroutine 仍可能在静默结束后立刻写数据）。生产代码以 hbMu 串行化这两条写路径；本 +// recorder 自带锁，仅为让单测在并发写下不因 recorder 自身非线程安全而误报——真正要由 +// -race 验证的是 handler 内部 emit/emitRaw 与 ping 是否正确互斥。若 hbMu 失效，race +// detector 会先抓到 handler 内部对 lastWriteNano / w 的竞争。 +type lockedFlushRecorder struct { + mu sync.Mutex + buf strings.Builder +} + +func newLockedFlushRecorder() *lockedFlushRecorder { return &lockedFlushRecorder{} } + +func (r *lockedFlushRecorder) Header() http.Header { return make(http.Header) } + +func (r *lockedFlushRecorder) Write(p []byte) (int, error) { + r.mu.Lock() + defer r.mu.Unlock() + return r.buf.Write(p) +} + +func (r *lockedFlushRecorder) WriteHeader(int) {} + +func (r *lockedFlushRecorder) Flush() {} + +func (r *lockedFlushRecorder) body() string { + r.mu.Lock() + defer r.mu.Unlock() + return r.buf.String() +} + +// newStalledEventStreamServer 返回一个上游 server：先写一帧 assistant 文本，flush 后 +// 静默 stallFor（模拟 opus 两个数据块之间的长 thinking 静默），再写第二帧。用于触发 +// 流式 handler 在“数据流中途”发出保活 ping 的路径。 +func newStalledEventStreamServer(t *testing.T, stallFor time.Duration) *httptest.Server { + t.Helper() + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + flusher, ok := w.(http.Flusher) + if !ok { + t.Errorf("upstream test server ResponseWriter is not a Flusher") + return + } + w.WriteHeader(http.StatusOK) + _, _ = w.Write(awsEventStreamFrame(t, "assistantResponseEvent", map[string]interface{}{ + "content": "first ", + })) + flusher.Flush() + time.Sleep(stallFor) + _, _ = w.Write(awsEventStreamFrame(t, "assistantResponseEvent", map[string]interface{}{ + "content": "second", + })) + flusher.Flush() + })) +} + +// withFastKeepalive 临时把保活间隔调小，并返回恢复函数。 +func withFastKeepalive(t *testing.T, d time.Duration) func() { + t.Helper() + old := streamKeepaliveInterval + streamKeepaliveInterval = d + return func() { streamKeepaliveInterval = old } +} + +// withNoTotalTimeoutStreamClient 临时把流式 client 换成无总超时的，避免我们故意制造的 +// 静默期间 client 把流砍断、掩盖保活路径。 +func withNoTotalTimeoutStreamClient(t *testing.T) func() { + t.Helper() + old := kiroHttpStore.Load() + kiroHttpStore.Store(&http.Client{Timeout: 0, Transport: &http.Transport{}}) + return func() { kiroHttpStore.Store(old) } +} + +func setupKeepaliveTestAccount(t *testing.T, id string) { + t.Helper() + cfgFile := t.TempDir() + "/config.json" + if err := config.Init(cfgFile); err != nil { + t.Fatalf("config.Init: %v", err) + } + if err := config.AddAccount(config.Account{ + ID: id, + Enabled: true, + AccessToken: "token-" + id, + ProfileArn: "arn:aws:codewhisperer:profile/" + id, + }); err != nil { + t.Fatalf("add account: %v", err) + } + if err := config.UpdatePreferredEndpoint("kiro"); err != nil { + t.Fatalf("set preferred endpoint: %v", err) + } + if err := config.UpdateEndpointFallback(false); err != nil { + t.Fatalf("disable endpoint fallback: %v", err) + } +} + +func keepaliveTestPayload(model string) *KiroPayload { + payload := &KiroPayload{} + payload.ConversationState.CurrentMessage.UserInputMessage = KiroUserInputMessage{ + Content: "hello", + ModelID: model, + Origin: "AI_EDITOR", + } + return payload +} + +// TestClaudeStreamEmitsKeepaliveDuringUpstreamStall 验证：上游在两帧数据之间长静默时， +// handleClaudeStream 会在数据流中途向客户端发出 ': keepalive'，且两帧真实文本仍完整下发、 +// 协议收尾正确。保活 goroutine 与主 goroutine 经 hbMu 串行化、无并发写冲突。 +func TestClaudeStreamEmitsKeepaliveDuringUpstreamStall(t *testing.T) { + setupKeepaliveTestAccount(t, "claude-keepalive") + + // 间隔 20ms，静默 200ms（≥ 多个保活周期），确保至少补发一次 ping。 + defer withFastKeepalive(t, 20*time.Millisecond)() + + server := newStalledEventStreamServer(t, 200*time.Millisecond) + defer server.Close() + + oldEndpoints := kiroEndpoints + kiroEndpoints = []kiroEndpoint{{URL: server.URL, Origin: "AI_EDITOR", Name: "test"}} + defer func() { kiroEndpoints = oldEndpoints }() + defer withNoTotalTimeoutStreamClient(t)() + + p := accountpool.GetPool() + p.Reload() + h := &Handler{ + pool: p, + promptCache: newPromptCacheTracker(defaultPromptCacheTTL), + } + + model := "claude-opus-4-8" + rec := newLockedFlushRecorder() + h.handleClaudeStream(rec, keepaliveTestPayload(model), model, false, claudeThinkingResponseOptions{}, 1000, nil, "") + + body := rec.body() + if !strings.Contains(body, ": keepalive") { + t.Fatalf("expected a keepalive comment emitted during upstream stall, got body=%s", body) + } + // 两帧真实文本必须都送达，保活不能吞掉或打断数据。 + if !strings.Contains(body, "first ") || !strings.Contains(body, "second") { + t.Fatalf("expected both upstream text frames in output, got body=%s", body) + } + // 协议收尾必须存在（message_start 在前、message_stop 在后）。 + if !strings.Contains(body, "message_start") { + t.Fatalf("expected message_start, got body=%s", body) + } + if !strings.Contains(body, "message_stop") { + t.Fatalf("expected message_stop, got body=%s", body) + } +} + +// TestOpenAIStreamEmitsKeepaliveDuringUpstreamStall 与上等价，覆盖 handleOpenAIStream +// 的 emitRaw/ping 并发路径。 +func TestOpenAIStreamEmitsKeepaliveDuringUpstreamStall(t *testing.T) { + setupKeepaliveTestAccount(t, "openai-keepalive") + + defer withFastKeepalive(t, 20*time.Millisecond)() + + server := newStalledEventStreamServer(t, 200*time.Millisecond) + defer server.Close() + + oldEndpoints := kiroEndpoints + kiroEndpoints = []kiroEndpoint{{URL: server.URL, Origin: "AI_EDITOR", Name: "test"}} + defer func() { kiroEndpoints = oldEndpoints }() + defer withNoTotalTimeoutStreamClient(t)() + + p := accountpool.GetPool() + p.Reload() + h := &Handler{ + pool: p, + promptCache: newPromptCacheTracker(defaultPromptCacheTTL), + } + + model := "claude-opus-4-8" + rec := newLockedFlushRecorder() + h.handleOpenAIStream(rec, keepaliveTestPayload(model), model, false, 1000, "") + + body := rec.body() + if !strings.Contains(body, ": keepalive") { + t.Fatalf("expected a keepalive comment emitted during upstream stall, got body=%s", body) + } + if !strings.Contains(body, "first ") || !strings.Contains(body, "second") { + t.Fatalf("expected both upstream text frames in output, got body=%s", body) + } + if !strings.Contains(body, "[DONE]") { + t.Fatalf("expected [DONE] terminator, got body=%s", body) + } +} diff --git a/proxy/handler_refresh_test.go b/proxy/handler_refresh_test.go new file mode 100644 index 00000000..5e7842db --- /dev/null +++ b/proxy/handler_refresh_test.go @@ -0,0 +1,86 @@ +package proxy + +import ( + "io" + "kiro-go/auth" + "kiro-go/config" + accountpool "kiro-go/pool" + "net/http" + "net/http/httptest" + "sync" + "sync/atomic" + "testing" + "time" +) + +// TestRefreshAccountTokenDedupConcurrentCalls verifies the central token-refresh +// path deduplicates concurrent refreshes of the SAME account: N goroutines all +// needing a refresh must collapse into a SINGLE upstream auth.RefreshToken call. +// +// Without the lock + pool re-check, each goroutine refreshes independently. For +// an IdP that rotates refresh tokens (one-time-use — e.g. external_idp / Azure +// AD, see oidc.go comment), the loser of the race sends an already-consumed +// refresh token and gets invalid_grant → handleAccountFailure bans the account. +// This is Bug A. +func TestRefreshAccountTokenDedupConcurrentCalls(t *testing.T) { + cfgFile := t.TempDir() + "/config.json" + if err := config.Init(cfgFile); err != nil { + t.Fatalf("config.Init: %v", err) + } + + var refreshCount int32 + authServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + atomic.AddInt32(&refreshCount, 1) + // Hold the request open briefly so concurrent callers queue on the + // token-refresh lock, widening the race window the dedup must close. + time.Sleep(75 * time.Millisecond) + w.Header().Set("Content-Type", "application/json") + _, _ = io.WriteString(w, `{"accessToken":"new-access","refreshToken":"new-refresh","expiresIn":3600,"profileArn":"arn:aws:codewhisperer:profile/dedup"}`) + })) + t.Cleanup(authServer.Close) + + oldTokenURL := auth.GetOIDCTokenURLForTest() + auth.SetOIDCTokenURLForTest(func(string) string { return authServer.URL }) + t.Cleanup(func() { auth.SetOIDCTokenURLForTest(oldTokenURL) }) + oldAuthClient := auth.SetGlobalAuthClientForTest(authServer.Client()) + t.Cleanup(func() { auth.SetGlobalAuthClientForTest(oldAuthClient) }) + + account := config.Account{ + ID: "acct-dedup", + Email: "dedup@example.com", + AccessToken: "stale-access", + RefreshToken: "stale-refresh", + ClientID: "client-id", + ClientSecret: "client-secret", + AuthMethod: "idc", + Region: "us-east-1", + Enabled: true, + ExpiresAt: time.Now().Unix() - 1, // already expired → refresh forced + } + if err := config.AddAccount(account); err != nil { + t.Fatalf("add account: %v", err) + } + + p := accountpool.GetPool() + p.Reload() + h := &Handler{pool: p} + + const N = 8 + var wg sync.WaitGroup + start := make(chan struct{}) + for i := 0; i < N; i++ { + wg.Add(1) + go func() { + defer wg.Done() + <-start // release all goroutines together to maximize concurrency + local := account // each goroutine simulates an independent per-request copy + _ = h.refreshAccountToken(&local, false) + }() + } + close(start) + wg.Wait() + + if got := atomic.LoadInt32(&refreshCount); got != 1 { + t.Fatalf("expected exactly 1 upstream refresh (dedup collapsed concurrent callers), got %d", got) + } +} diff --git a/proxy/handler_test.go b/proxy/handler_test.go index 728f88f7..7dd4828a 100644 --- a/proxy/handler_test.go +++ b/proxy/handler_test.go @@ -2,6 +2,7 @@ package proxy import ( "encoding/json" + "fmt" "kiro-go/config" accountpool "kiro-go/pool" "net/http" @@ -480,3 +481,142 @@ func TestBuildAnthropicModelsResponseGeneratesThinkingVariants(t *testing.T) { t.Fatalf("expected image capability to be preserved, got %#v", models[0]["supports_image"]) } } + +// TestStatsIncludesCacheMetrics verifies /v1/stats surfaces a "cache" object +// populated from the prompt-cache tracker's atomic counters (hits/misses). +func TestStatsIncludesCacheMetrics(t *testing.T) { + cfgFile := t.TempDir() + "/config.json" + if err := config.Init(cfgFile); err != nil { + t.Fatalf("config.Init: %v", err) + } + p := accountpool.GetPool() + p.Reload() + + tr := newPromptCacheTracker(defaultPromptCacheTTL) + profile := &promptCacheProfile{ + Model: "claude-sonnet-4-6", + TotalInputTokens: 2000, + Breakpoints: []promptCacheBreakpoint{{Fingerprint: [32]byte{9}, CumulativeTokens: 2000, TTL: time.Hour}}, + } + tr.Compute("acct", profile) // miss (empty cache) + tr.Update("acct", profile) + tr.Compute("acct", profile) // hit + + h := &Handler{ + pool: p, + promptCache: tr, + startTime: time.Now().Unix(), + } + + rec := httptest.NewRecorder() + h.handleStats(rec, httptest.NewRequest(http.MethodGet, "/v1/stats", nil)) + + var got map[string]interface{} + if err := json.NewDecoder(rec.Body).Decode(&got); err != nil { + t.Fatalf("decode stats: %v", err) + } + cache, ok := got["cache"].(map[string]interface{}) + if !ok { + t.Fatalf("expected a cache object in stats, got %#v", got) + } + if cache["hits"].(float64) != 1 { + t.Fatalf("expected 1 hit, got %v", cache["hits"]) + } + if cache["misses"].(float64) != 1 { + t.Fatalf("expected 1 miss, got %v", cache["misses"]) + } +} + +// TestApiAddAccountApiKeyBranch verifies POST /admin/api/accounts with an +// api_key account normalizes AuthMethod, zeroes ExpiresAt, mirrors AccessToken, +// and skips the OAuth refresh requirement. +func TestApiAddAccountApiKeyBranch(t *testing.T) { + if err := config.Init(t.TempDir() + "/config.json"); err != nil { + t.Fatalf("config.Init: %v", err) + } + // Stub the REST store so the background model-cache fetch fails fast without + // real network egress. + kiroRestHttpStore.Store(&http.Client{ + Transport: roundTripFunc(func(*http.Request) (*http.Response, error) { + return nil, fmt.Errorf("test: no network") + }), + }) + t.Cleanup(func() { InitKiroHttpClient("") }) + + h := &Handler{pool: accountpool.GetPool()} + body := `{"kiroApiKey":"key-123","authMethod":"api_key","region":"eu-central-1","enabled":true}` + req := httptest.NewRequest("POST", "/admin/api/accounts", strings.NewReader(body)) + rec := httptest.NewRecorder() + h.apiAddAccount(rec, req) + + if rec.Code != 200 { + t.Fatalf("expected 200, got %d body=%s", rec.Code, rec.Body.String()) + } + accs := config.GetAccounts() + if len(accs) != 1 { + t.Fatalf("expected 1 account, got %d", len(accs)) + } + got := accs[0] + if got.AuthMethod != "api_key" { + t.Fatalf("AuthMethod: want api_key, got %q", got.AuthMethod) + } + if got.AccessToken != "key-123" { + t.Fatalf("AccessToken mirror: want key-123, got %q", got.AccessToken) + } + if got.ExpiresAt != 0 { + t.Fatalf("ExpiresAt: want 0, got %d", got.ExpiresAt) + } + if got.KiroApiKey != "key-123" { + t.Fatalf("KiroApiKey: want key-123, got %q", got.KiroApiKey) + } + if got.RefreshToken != "" { + t.Fatalf("RefreshToken: want empty, got %q", got.RefreshToken) + } +} + +// TestApiAddAccountApiKeyRequiresKey verifies a missing key on an api_key +// request is rejected with 400. +func TestApiAddAccountApiKeyRequiresKey(t *testing.T) { + if err := config.Init(t.TempDir() + "/config.json"); err != nil { + t.Fatalf("config.Init: %v", err) + } + h := &Handler{pool: accountpool.GetPool()} + body := `{"authMethod":"api_key"}` + req := httptest.NewRequest("POST", "/admin/api/accounts", strings.NewReader(body)) + rec := httptest.NewRecorder() + h.apiAddAccount(rec, req) + if rec.Code != 400 { + t.Fatalf("expected 400 for api_key without key, got %d body=%s", rec.Code, rec.Body.String()) + } +} + +// TestSettingsRoundTripMaxPayloadBytes verifies GET /settings reports +// maxPayloadBytes and POST /settings persists it. +func TestSettingsRoundTripMaxPayloadBytes(t *testing.T) { + if err := config.Init(t.TempDir() + "/config.json"); err != nil { + t.Fatalf("config.Init: %v", err) + } + h := &Handler{pool: accountpool.GetPool()} + + // GET reports the default. + getRec := httptest.NewRecorder() + h.apiGetSettings(getRec, httptest.NewRequest("GET", "/settings", nil)) + var got map[string]json.RawMessage + if err := json.Unmarshal(getRec.Body.Bytes(), &got); err != nil { + t.Fatalf("decode GET: %v", err) + } + if _, ok := got["maxPayloadBytes"]; !ok { + t.Fatalf("expected maxPayloadBytes in GET /settings, got %v", got) + } + + // POST sets it. + postBody := `{"maxPayloadBytes":2100000}` + postRec := httptest.NewRecorder() + h.apiUpdateSettings(postRec, httptest.NewRequest("POST", "/settings", strings.NewReader(postBody))) + if postRec.Code != 200 { + t.Fatalf("POST /settings: expected 200, got %d body=%s", postRec.Code, postRec.Body.String()) + } + if config.GetMaxPayloadBytes() != 2100000 { + t.Fatalf("maxPayloadBytes not persisted: want 2100000, got %d", config.GetMaxPayloadBytes()) + } +} diff --git a/proxy/import_credentials_test.go b/proxy/import_credentials_test.go index 4aae8d11..cf64b606 100644 --- a/proxy/import_credentials_test.go +++ b/proxy/import_credentials_test.go @@ -1,8 +1,10 @@ package proxy import ( + "encoding/base64" "encoding/json" "fmt" + "io" "kiro-go/auth" "kiro-go/config" accountpool "kiro-go/pool" @@ -80,6 +82,49 @@ func TestApiImportCredentialsRejectsWhenRefreshFails(t *testing.T) { // TestApiImportCredentialsUsesUpstreamExpiresAt verifies the happy path: when // refresh succeeds, the persisted ExpiresAt reflects the upstream expiresIn, // not a hard-coded 300s. +// TestApiImportCredentialsDedupsByRefreshToken verifies the invariant stated at +// config.go:437 / handler.go:3234 ("re-importing a backup never creates a +// duplicate entry"): re-importing the same credential JSON must update the +// existing account in place rather than mint a fresh id and leave two live +// accounts sharing the same refresh token (which causes interleaved rotation +// conflicts). The bulk import path already dedups on refresh token; the single +// import path did not. +func TestApiImportCredentialsDedupsByRefreshToken(t *testing.T) { + cfgFile := t.TempDir() + "/config.json" + if err := config.Init(cfgFile); err != nil { + t.Fatalf("config.Init: %v", err) + } + defer installCleanAuthClient(t)() + + // Fake OIDC issues a valid access token and echoes a fixed refresh token, so + // both imports resolve the same persisted RefreshToken and reach the dedup. + fake := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + fmt.Fprintf(w, `{"accessToken":"at-new","refreshToken":"rt-shared","expiresIn":3600,"profileArn":"arn:aws:codewhisperer:profile/test"}`) + })) + defer fake.Close() + + oldOIDC := authOidcURL() + auth.SetOIDCTokenURLForTest(func(string) string { return fake.URL }) + defer auth.SetOIDCTokenURLForTest(oldOIDC) + + h := &Handler{pool: accountpool.GetPool()} + body := `{"refreshToken":"rt-good","clientId":"c","clientSecret":"s","authMethod":"idc","region":"us-east-1"}` + + for i := 0; i < 2; i++ { + rec := httptest.NewRecorder() + h.apiImportCredentials(rec, httptest.NewRequest("POST", "/auth/credentials", strings.NewReader(body))) + if rec.Code != http.StatusOK { + t.Fatalf("import #%d failed: status=%d body=%s", i+1, rec.Code, rec.Body.String()) + } + } + + accs := config.GetAccounts() + if len(accs) != 1 { + t.Fatalf("re-importing the same credential must update in place, not create a duplicate sharing the refresh token; got %d accounts", len(accs)) + } +} + func TestApiImportCredentialsUsesUpstreamExpiresAt(t *testing.T) { cfgFile := t.TempDir() + "/config.json" if err := config.Init(cfgFile); err != nil { @@ -137,3 +182,418 @@ func TestApiImportCredentialsUsesUpstreamExpiresAt(t *testing.T) { // authOidcURL captures the current oidc URL builder so the test can restore it. func authOidcURL() func(string) string { return auth.GetOIDCTokenURLForTest() } + +// TestNormalizeImportAuthMethod pins the auth-method normalization for import, +// including the key regression: external_idp accounts carry clientId but NO +// clientSecret, so the old default branch misclassified them as "social". +func TestNormalizeImportAuthMethod(t *testing.T) { + cases := []struct { + name string + authMethod string + clientID string + clientSecret string + tokenEndpoint string + want string + }{ + {"explicit external_idp", "external_idp", "c", "", "https://login.microsoftonline.com/t/oauth2/v2.0/token", "external_idp"}, + {"azure alias", "AzureAD", "c", "", "https://login.microsoftonline.com/t/oauth2/v2.0/token", "external_idp"}, + {"microsoft alias", "microsoft", "c", "", "https://login.microsoftonline.com/t/oauth2/v2.0/token", "external_idp"}, + {"inferred from tokenEndpoint", "", "c", "", "https://login.microsoftonline.com/t/oauth2/v2.0/token", "external_idp"}, + {"external_idp even with clientSecret", "external_idp", "c", "s", "https://login.microsoftonline.com/t/oauth2/v2.0/token", "external_idp"}, + {"enterprise stays idc", "enterprise", "c", "s", "", "idc"}, + {"idc with clientid+secret", "idc", "c", "s", "", "idc"}, + {"empty + clientid (no secret) -> idc", "", "c", "", "", "idc"}, + {"empty no clientid -> social", "", "", "", "", "social"}, + {"social explicit", "social", "", "", "", "social"}, + {"google alias", "google", "", "", "", "social"}, + {"unrecognized with clientid+secret -> idc", "weird", "c", "s", "", "idc"}, + {"unrecognized without secret -> social", "weird", "c", "", "", "social"}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + if got := normalizeImportAuthMethod(tc.authMethod, tc.clientID, tc.clientSecret, tc.tokenEndpoint); got != tc.want { + t.Fatalf("normalizeImportAuthMethod(%q,%q,%q,%q) = %q, want %q", + tc.authMethod, tc.clientID, tc.clientSecret, tc.tokenEndpoint, got, tc.want) + } + }) + } +} + +// TestApiImportCredentialsExternalIdpHappyPath verifies an external_idp credential +// imports successfully: authMethod normalizes to external_idp, refresh hits the +// (fake) IdP token endpoint, and the account is persisted with all refresh material. +func TestApiImportCredentialsExternalIdpHappyPath(t *testing.T) { + cfgFile := t.TempDir() + "/config.json" + if err := config.Init(cfgFile); err != nil { + t.Fatalf("config.Init: %v", err) + } + defer installCleanAuthClient(t)() + + const upstreamExpiresIn = 3600 + fake := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if err := r.ParseForm(); err != nil { + http.Error(w, "bad form", http.StatusBadRequest) + return + } + if got := r.PostForm.Get("grant_type"); got != "refresh_token" { + t.Errorf("expected grant_type=refresh_token, got %q", got) + } + w.Header().Set("Content-Type", "application/json") + // external IdP token responses are snake_case. + fmt.Fprintf(w, `{"access_token":"at-ext","refresh_token":"rt-rotated","expires_in":%d}`, upstreamExpiresIn) + })) + defer fake.Close() + + // fake.URL is http + 127.0.0.1; bypass the allow-list validator for this test. + restore := auth.SetExternalIdpValidatorForTest(func(string) error { return nil }) + defer auth.SetExternalIdpValidatorForTest(restore) + + h := &Handler{pool: accountpool.GetPool()} + + body := fmt.Sprintf(`{"authMethod":"external_idp","refreshToken":"rt-ext","clientId":"ext-client","tokenEndpoint":%q,"issuerUrl":"https://login.microsoftonline.com/t/v2.0","scopes":"api://x/codewhisperer:conversations offline_access","region":"eu-central-1"}`, fake.URL) + req := httptest.NewRequest("POST", "/auth/credentials", strings.NewReader(body)) + rec := httptest.NewRecorder() + before := time.Now().Unix() + h.apiImportCredentials(rec, req) + after := time.Now().Unix() + + if rec.Code != http.StatusOK { + t.Fatalf("expected 200, got %d body=%s", rec.Code, rec.Body.String()) + } + + accs := config.GetAccounts() + if len(accs) != 1 { + t.Fatalf("expected 1 account persisted, got %d", len(accs)) + } + got := accs[0] + if got.AuthMethod != "external_idp" { + t.Fatalf("AuthMethod: want external_idp, got %q", got.AuthMethod) + } + if got.AccessToken != "at-ext" { + t.Fatalf("AccessToken: want at-ext, got %q", got.AccessToken) + } + if got.RefreshToken != "rt-rotated" { + t.Fatalf("RefreshToken: want rt-rotated (rotated), got %q", got.RefreshToken) + } + if got.TokenEndpoint != fake.URL { + t.Fatalf("TokenEndpoint not persisted: got %q", got.TokenEndpoint) + } + if got.ClientID != "ext-client" { + t.Fatalf("ClientID not persisted: got %q", got.ClientID) + } + if got.Scopes == "" { + t.Fatalf("Scopes not persisted: got %q", got.Scopes) + } + if got.Provider != "AzureAD" { + t.Fatalf("Provider default: want AzureAD, got %q", got.Provider) + } + if got.Region != "eu-central-1" { + t.Fatalf("Region: want eu-central-1, got %q", got.Region) + } + if got.ExpiresAt < before+upstreamExpiresIn-5 || got.ExpiresAt > after+upstreamExpiresIn+5 { + t.Fatalf("ExpiresAt not from upstream expiresIn: got %d (want ~now+%d)", got.ExpiresAt, upstreamExpiresIn) + } +} + +// TestApiImportCredentialsExternalIdpRejectsNonAllowListedEndpoint verifies the SSRF +// guard: a tokenEndpoint outside the IdP allow-list is rejected with 400 before any +// refresh POST, and nothing is persisted. (Validator is NOT bypassed here.) +func TestApiImportCredentialsExternalIdpRejectsNonAllowListedEndpoint(t *testing.T) { + cfgFile := t.TempDir() + "/config.json" + if err := config.Init(cfgFile); err != nil { + t.Fatalf("config.Init: %v", err) + } + defer installCleanAuthClient(t)() + + h := &Handler{pool: accountpool.GetPool()} + + body := `{"authMethod":"external_idp","refreshToken":"rt","clientId":"c","tokenEndpoint":"https://evil.example.com/oauth/token","region":"us-east-1"}` + req := httptest.NewRequest("POST", "/auth/credentials", strings.NewReader(body)) + rec := httptest.NewRecorder() + h.apiImportCredentials(rec, req) + + if rec.Code != http.StatusBadRequest { + t.Fatalf("expected 400, got %d body=%s", rec.Code, rec.Body.String()) + } + var resp map[string]string + if err := json.Unmarshal(rec.Body.Bytes(), &resp); err != nil { + t.Fatalf("decode response: %v", err) + } + if !strings.Contains(resp["error"], "endpoint rejected") { + t.Fatalf("expected endpoint-rejected error, got %q", resp["error"]) + } + if accs := config.GetAccounts(); len(accs) != 0 { + t.Fatalf("expected no account persisted, got %d", len(accs)) + } +} + +// TestApiImportCredentialsExternalIdpRejectsWhenRefreshFails verifies the refresh +// gate holds for external_idp: a refresh that 400s (invalid_grant) must reject the +// import and persist nothing. +func TestApiImportCredentialsExternalIdpRejectsWhenRefreshFails(t *testing.T) { + cfgFile := t.TempDir() + "/config.json" + if err := config.Init(cfgFile); err != nil { + t.Fatalf("config.Init: %v", err) + } + defer installCleanAuthClient(t)() + + fake := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + http.Error(w, `{"error":"invalid_grant"}`, http.StatusBadRequest) + })) + defer fake.Close() + + restore := auth.SetExternalIdpValidatorForTest(func(string) error { return nil }) + defer auth.SetExternalIdpValidatorForTest(restore) + + h := &Handler{pool: accountpool.GetPool()} + + body := fmt.Sprintf(`{"authMethod":"external_idp","refreshToken":"rt-broken","clientId":"c","tokenEndpoint":%q,"region":"us-east-1"}`, fake.URL) + req := httptest.NewRequest("POST", "/auth/credentials", strings.NewReader(body)) + rec := httptest.NewRecorder() + h.apiImportCredentials(rec, req) + + if rec.Code != http.StatusBadRequest { + t.Fatalf("expected 400, got %d body=%s", rec.Code, rec.Body.String()) + } + var resp map[string]string + if err := json.Unmarshal(rec.Body.Bytes(), &resp); err != nil { + t.Fatalf("decode response: %v", err) + } + if !strings.Contains(resp["error"], "Token refresh failed") { + t.Fatalf("expected refresh-failed error, got %q", resp["error"]) + } + if accs := config.GetAccounts(); len(accs) != 0 { + t.Fatalf("expected no account persisted, got %d", len(accs)) + } +} + +// TestApiImportCredentialsExternalIdpPreservesFullRecordIdentity verifies that when +// a full account record (with id/email/profileArn) is pasted, those are preserved +// rather than regenerated, so re-importing a backup does not duplicate accounts. +func TestApiImportCredentialsExternalIdpPreservesFullRecordIdentity(t *testing.T) { + cfgFile := t.TempDir() + "/config.json" + if err := config.Init(cfgFile); err != nil { + t.Fatalf("config.Init: %v", err) + } + defer installCleanAuthClient(t)() + + fake := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + fmt.Fprint(w, `{"access_token":"at-ext","refresh_token":"rt-rotated","expires_in":3600}`) + })) + defer fake.Close() + + restore := auth.SetExternalIdpValidatorForTest(func(string) error { return nil }) + defer auth.SetExternalIdpValidatorForTest(restore) + + h := &Handler{pool: accountpool.GetPool()} + + const providedID = "11111111-2222-3333-4444-555555555555" + body := fmt.Sprintf(`{"id":%q,"email":"ada@example.com","profileArn":"arn:aws:codewhisperer:eu-central-1:1:profile/PRESERVED","authMethod":"external_idp","refreshToken":"rt","clientId":"c","tokenEndpoint":%q,"region":"eu-central-1"}`, providedID, fake.URL) + req := httptest.NewRequest("POST", "/auth/credentials", strings.NewReader(body)) + rec := httptest.NewRecorder() + h.apiImportCredentials(rec, req) + + if rec.Code != http.StatusOK { + t.Fatalf("expected 200, got %d body=%s", rec.Code, rec.Body.String()) + } + got := config.GetAccounts()[0] + if got.ID != providedID { + t.Fatalf("ID: want reused %q, got %q", providedID, got.ID) + } + if got.Email != "ada@example.com" { + t.Fatalf("Email: want ada@example.com (GetUserInfo empty in test → fallback), got %q", got.Email) + } + if got.ProfileArn != "arn:aws:codewhisperer:eu-central-1:1:profile/PRESERVED" { + t.Fatalf("ProfileArn: want preserved, got %q", got.ProfileArn) + } +} + +// TestApiImportCredentialsExternalIdpDerivesEndpointsFromUserId verifies the +// Kiro Account Manager export shape: a credential carrying only refreshToken + +// clientId + userId (NO tokenEndpoint/issuerUrl/scopes) is accepted, with the +// endpoints+scopes derived from userId's embedded Azure tenant. +func TestApiImportCredentialsExternalIdpDerivesEndpointsFromUserId(t *testing.T) { + cfgFile := t.TempDir() + "/config.json" + if err := config.Init(cfgFile); err != nil { + t.Fatalf("config.Init: %v", err) + } + defer installCleanAuthClient(t)() + + fake := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + fmt.Fprint(w, `{"access_token":"at-derived","refresh_token":"rt-d2","expires_in":3600}`) + })) + defer fake.Close() + + restore := auth.SetExternalIdpValidatorForTest(func(string) error { return nil }) + defer auth.SetExternalIdpValidatorForTest(restore) + + h := &Handler{pool: accountpool.GetPool()} + + // userId points at the fake so the derived tokenEndpoint hits it. + userID := fake.URL + "/5fbc183e-3d09-4043-b36f-0c49d3665977/v2.0.8db0e2eb-d491-4a1a-98f1-cbdc12bb60a0" + body := fmt.Sprintf(`{"authMethod":"external_idp","refreshToken":"rt","clientId":"fa6d79bf-cdaa-495e-8359-78aab7c7cd9b","userId":%q,"region":"eu-central-1"}`, userID) + req := httptest.NewRequest("POST", "/auth/credentials", strings.NewReader(body)) + rec := httptest.NewRecorder() + h.apiImportCredentials(rec, req) + + if rec.Code != http.StatusOK { + t.Fatalf("expected 200, got %d body=%s", rec.Code, rec.Body.String()) + } + got := config.GetAccounts()[0] + wantTE := fake.URL + "/5fbc183e-3d09-4043-b36f-0c49d3665977/oauth2/v2.0/token" + if got.TokenEndpoint != wantTE { + t.Fatalf("derived TokenEndpoint: want %q, got %q", wantTE, got.TokenEndpoint) + } + if got.IssuerURL != fake.URL+"/5fbc183e-3d09-4043-b36f-0c49d3665977/v2.0" { + t.Fatalf("derived IssuerURL: got %q", got.IssuerURL) + } + if !strings.Contains(got.Scopes, "codewhisperer:conversations") || !strings.Contains(got.Scopes, "offline_access") { + t.Fatalf("derived Scopes: got %q", got.Scopes) + } + if got.AccessToken != "at-derived" { + t.Fatalf("AccessToken: want at-derived, got %q", got.AccessToken) + } +} + +// TestApiImportCredentialsExternalIdpDerivesFromAccessTokenJWT verifies a bare +// credential blob (clientId + accessToken + refreshToken, NO authMethod/userId/ +// tokenEndpoint) imports via TRUST-ON-IMPORT: the accessToken's JWT issuer +// classifies it as external_idp + yields the tenant to derive the endpoint from, +// and — because the token is an Azure AD JWT with a real exp — the credential is +// persisted verbatim WITHOUT a live refresh round-trip (so the same JSON can be +// re-imported without the refresh token getting consumed/rotated). +func TestApiImportCredentialsExternalIdpDerivesFromAccessTokenJWT(t *testing.T) { + cfgFile := t.TempDir() + "/config.json" + if err := config.Init(cfgFile); err != nil { + t.Fatalf("config.Init: %v", err) + } + defer installCleanAuthClient(t)() + + // Sentinel: if trust-on-import were broken, a refresh would hit this server and + // the persisted AccessToken would be "at-jwt" instead of the pasted JWT. + fake := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + fmt.Fprint(w, `{"access_token":"at-jwt","refresh_token":"rt-j2","expires_in":3600}`) + })) + defer fake.Close() + + restore := auth.SetExternalIdpValidatorForTest(func(string) error { return nil }) + defer auth.SetExternalIdpValidatorForTest(restore) + + // Bare blob: only clientId + accessToken + refreshToken. The access token is an + // Azure AD JWT (iss + future exp) → external_idp + derived tenant + trust-on-import + // (NO live refresh; pasted tokens persist verbatim, ExpiresAt from JWT exp). + tenant := "5fbc183e-3d09-4043-b36f-0c49d3665977" + const exp int64 = 2000000000 + jwt := "eyJhbGciOiJub25lIn0." + + base64.RawURLEncoding.EncodeToString([]byte(fmt.Sprintf(`{"iss":%q,"exp":%d}`, fake.URL+"/"+tenant+"/v2.0", exp))) + "." + + h := &Handler{pool: accountpool.GetPool()} + + body := fmt.Sprintf(`{"clientId":"fa6d79bf-cdaa-495e-8359-78aab7c7cd9b","accessToken":%q,"refreshToken":"rt","region":"eu-central-1"}`, jwt) + req := httptest.NewRequest("POST", "/auth/credentials", strings.NewReader(body)) + rec := httptest.NewRecorder() + h.apiImportCredentials(rec, req) + + if rec.Code != http.StatusOK { + t.Fatalf("expected 200 (trust-on-import), got %d body=%s", rec.Code, rec.Body.String()) + } + got := config.GetAccounts()[0] + if got.AuthMethod != "external_idp" { + t.Fatalf("AuthMethod: want external_idp (derived from accessToken JWT), got %q", got.AuthMethod) + } + wantTE := fake.URL + "/" + tenant + "/oauth2/v2.0/token" + if got.TokenEndpoint != wantTE { + t.Fatalf("derived TokenEndpoint: want %q, got %q", wantTE, got.TokenEndpoint) + } + if got.AccessToken != jwt { + t.Fatalf("AccessToken: want the pasted JWT persisted verbatim (trust-on-import), got %q", got.AccessToken) + } + if got.ExpiresAt != exp { + t.Fatalf("ExpiresAt: want %d (from JWT exp, trust-on-import), got %d", exp, got.ExpiresAt) + } + if got.RefreshToken != "rt" { + t.Fatalf("RefreshToken: want the pasted token (not rotated), got %q", got.RefreshToken) + } +} + +// TestApiImportCredentialsApiKeyBranch verifies POST /auth/credentials with a +// kiroApiKey imports an api_key account (AuthMethod normalized, ExpiresAt=0, +// AccessToken mirrored) without any OAuth refresh round-trip. The background +// best-effort RefreshAccountInfo is exercised against a stubbed REST store. +func TestApiImportCredentialsApiKeyBranch(t *testing.T) { + cfgFile := t.TempDir() + "/config.json" + if err := config.Init(cfgFile); err != nil { + t.Fatalf("config.Init: %v", err) + } + defer installCleanAuthClient(t)() + + // Stub the REST store so the background RefreshAccountInfo / model-cache + // fetches don't hit real network. A 200 with empty JSON decodes cleanly and + // never trips the auth-error ban path. + kiroRestHttpStore.Store(&http.Client{ + Transport: roundTripFunc(func(*http.Request) (*http.Response, error) { + return &http.Response{ + StatusCode: 200, + Body: io.NopCloser(strings.NewReader(`{}`)), + Header: make(http.Header), + }, nil + }), + }) + t.Cleanup(func() { InitKiroHttpClient("") }) + + h := &Handler{pool: accountpool.GetPool()} + body := `{"kiroApiKey":"my-key","authMethod":"api_key","region":"eu-central-1"}` + req := httptest.NewRequest("POST", "/auth/credentials", strings.NewReader(body)) + rec := httptest.NewRecorder() + h.apiImportCredentials(rec, req) + + if rec.Code != 200 { + t.Fatalf("expected 200, got %d body=%s", rec.Code, rec.Body.String()) + } + accs := config.GetAccounts() + if len(accs) != 1 { + t.Fatalf("expected 1 account, got %d", len(accs)) + } + got := accs[0] + if got.AuthMethod != "api_key" { + t.Fatalf("AuthMethod: want api_key, got %q", got.AuthMethod) + } + if got.AccessToken != "my-key" { + t.Fatalf("AccessToken mirror: want my-key, got %q", got.AccessToken) + } + if got.KiroApiKey != "my-key" { + t.Fatalf("KiroApiKey: want my-key, got %q", got.KiroApiKey) + } + if got.ExpiresAt != 0 { + t.Fatalf("ExpiresAt: want 0, got %d", got.ExpiresAt) + } + if got.RefreshToken != "" { + t.Fatalf("RefreshToken: want empty, got %q", got.RefreshToken) + } + if got.Region != "eu-central-1" { + t.Fatalf("Region: want eu-central-1, got %q", got.Region) + } +} + +// TestApiImportCredentialsApiKeyRequiresKey verifies a missing key on an +// api_key import request is rejected with 400 and nothing is persisted. +func TestApiImportCredentialsApiKeyRequiresKey(t *testing.T) { + if err := config.Init(t.TempDir() + "/config.json"); err != nil { + t.Fatalf("config.Init: %v", err) + } + defer installCleanAuthClient(t)() + h := &Handler{pool: accountpool.GetPool()} + body := `{"authMethod":"api_key"}` + req := httptest.NewRequest("POST", "/auth/credentials", strings.NewReader(body)) + rec := httptest.NewRecorder() + h.apiImportCredentials(rec, req) + if rec.Code != 400 { + t.Fatalf("expected 400 for api_key without key, got %d body=%s", rec.Code, rec.Body.String()) + } + if accs := config.GetAccounts(); len(accs) != 0 { + t.Fatalf("expected no account persisted, got %d", len(accs)) + } +} diff --git a/proxy/inference_test.go b/proxy/inference_test.go new file mode 100644 index 00000000..704f41c9 --- /dev/null +++ b/proxy/inference_test.go @@ -0,0 +1,75 @@ +package proxy + +import ( + "encoding/json" + "strings" + "testing" +) + +// Guards the temperature:0 fix: an explicitly-provided temperature of 0 must be +// transmitted to Kiro (not dropped by omitempty), while an unset temperature is +// still omitted. Ported from kiro-tutu. +func TestInferenceConfigTransmitsExplicitZeroTemperature(t *testing.T) { + zero := 0.0 + withZero, _ := json.Marshal(&InferenceConfig{Temperature: &zero}) + if !strings.Contains(string(withZero), `"temperature":0`) { + t.Fatalf("explicit temperature 0 must be transmitted, got %s", withZero) + } + + unset, _ := json.Marshal(&InferenceConfig{}) + if strings.Contains(string(unset), "temperature") { + t.Fatalf("unset temperature must be omitted, got %s", unset) + } +} + +// TestClaudeRequestTransmitsExplicitZeroTemperature locks the full user-facing +// fix: a Claude request with temperature:0 (deterministic decoding — a real, +// deliberate setting) must survive JSON marshalling and reach the translator +// path, instead of being silently dropped by float64+omitempty (which treated +// 0 == "unset" and stripped the field, silently switching the client to the +// default temperature). +func TestClaudeRequestTransmitsExplicitZeroTemperature(t *testing.T) { + zero := 0.0 + raw, _ := json.Marshal(&ClaudeRequest{Model: "claude-sonnet-4.5", Temperature: &zero}) + if !strings.Contains(string(raw), `"temperature":0`) { + t.Fatalf("ClaudeRequest must transmit explicit temperature 0, got %s", raw) + } +} + +// TestOpenAIRequestTransmitsExplicitZeroTemperature is the OpenAI-side guard +// for the same temperature:0 drop bug. +func TestOpenAIRequestTransmitsExplicitZeroTemperature(t *testing.T) { + zero := 0.0 + raw, _ := json.Marshal(&OpenAIRequest{Model: "gpt-4o", Temperature: &zero}) + if !strings.Contains(string(raw), `"temperature":0`) { + t.Fatalf("OpenAIRequest must transmit explicit temperature 0, got %s", raw) + } +} + +// TestClaudeToKiroPropagatesExplicitZeroTemperature verifies the end-to-end +// translation path: temperature:0 on an inbound Claude request must land on the +// Kiro InferenceConfig (also as a non-nil 0), not be dropped — so deterministic +// decoding reaches the upstream model rather than silently reverting to default. +func TestClaudeToKiroPropagatesExplicitZeroTemperature(t *testing.T) { + zero := 0.0 + req := &ClaudeRequest{ + Model: "claude-sonnet-4.5", + MaxTokens: 1024, + Messages: []ClaudeMessage{{Role: "user", Content: "hi"}}, + Temperature: &zero, + } + payload := ClaudeToKiro(req, false) + if payload.InferenceConfig == nil { + t.Fatal("InferenceConfig must be set when temperature is explicitly 0") + } + if payload.InferenceConfig.Temperature == nil { + t.Fatal("temperature must be propagated as a non-nil pointer when explicitly 0") + } + if *payload.InferenceConfig.Temperature != 0 { + t.Fatalf("temperature must be 0, got %v", *payload.InferenceConfig.Temperature) + } + kiroRaw, _ := json.Marshal(payload.InferenceConfig) + if !strings.Contains(string(kiroRaw), `"temperature":0`) { + t.Fatalf("Kiro InferenceConfig must transmit temperature:0, got %s", kiroRaw) + } +} diff --git a/proxy/kiro.go b/proxy/kiro.go index b2fef49d..b204ff14 100644 --- a/proxy/kiro.go +++ b/proxy/kiro.go @@ -224,9 +224,9 @@ type KiroToolUse struct { } type InferenceConfig struct { - MaxTokens int `json:"maxTokens,omitempty"` - Temperature float64 `json:"temperature,omitempty"` - TopP float64 `json:"topP,omitempty"` + MaxTokens int `json:"maxTokens,omitempty"` + Temperature *float64 `json:"temperature,omitempty"` + TopP float64 `json:"topP,omitempty"` } // ==================== Stream Callbacks ==================== @@ -340,7 +340,13 @@ func CallKiroAPI(account *config.Account, payload *KiroPayload, callback *KiroSt // Update the origin field for the selected endpoint. payload.ConversationState.CurrentMessage.UserInputMessage.Origin = ep.Origin - // Target the profile's data-plane region; endpoint URLs are declared for us-east-1. + // Target the data-plane region via the unified regionalizeURLForProfile. + // kiroRegionForProfile branches on IsApiKeyCredential: api_key accounts + // resolve from EffectiveApiRegion (no profile ARN), OAuth accounts derive + // from the profile ARN. regionalizeURLForRegion collapses both us-east-1 + // hosts (q. and codewhisperer.) onto q. for non-us-east-1 — there + // is no regional codewhisperer host — so this also routes the CodeWhisperer + // endpoint correctly for regional api_key accounts. epURL := regionalizeURLForProfile(ep.URL, account, payload.ProfileArn) reqBody, _ := json.Marshal(payload) @@ -384,6 +390,16 @@ func CallKiroAPI(account *config.Account, payload *KiroPayload, callback *KiroSt if resp.StatusCode != 200 { errBody, _ := io.ReadAll(resp.Body) resp.Body.Close() + // A 400 "Improperly formed request" is inherent to the request payload — + // no endpoint or account can fix it. Wrap it as a permanent rejection so + // CallKiroAPI's caller short-circuits (no other endpoint/account tried) + // and the relaying account is not penalised for a bad request it merely + // forwarded. Detected before any event-stream bytes are read, so streaming + // handlers have not yet sent anything to the client. + if resp.StatusCode == 400 && isImproperlyFormedRejection(string(errBody)) { + lastErr = newUpstreamPermanentError(resp.StatusCode, string(errBody)) + return lastErr + } lastErr = fmt.Errorf("HTTP %d from %s: %s", resp.StatusCode, ep.Name, string(errBody)) // Authentication errors and payment errors are not retried across endpoints. if resp.StatusCode == 401 || resp.StatusCode == 403 || resp.StatusCode == 402 { @@ -404,10 +420,22 @@ func CallKiroAPI(account *config.Account, payload *KiroPayload, callback *KiroSt return fmt.Errorf("all endpoints failed") } +// accountEmailForLog returns the most informative stable identifier for an +// account in log lines. OAuth accounts log their email. api_key accounts have +// no email until their first successful refresh (RefreshAccountInfo populates +// it from getUsageLimits), so without a fallback every refresh/ ban/ model- +// cache line for them was blind ("for : ..."). Fall back to the masked key, +// which is stable and safe to log. func accountEmailForLog(account *config.Account) string { if account == nil { return "" } + if email := strings.TrimSpace(account.Email); email != "" { + return email + } + if account.IsApiKeyCredential() && account.KiroApiKey != "" { + return maskKey(account.KiroApiKey) + } return account.Email } diff --git a/proxy/kiro_api.go b/proxy/kiro_api.go index 5571f52f..4c141e45 100644 --- a/proxy/kiro_api.go +++ b/proxy/kiro_api.go @@ -4,11 +4,11 @@ import ( "encoding/json" "fmt" "io" - "kiro-go/auth" "kiro-go/config" "kiro-go/logger" "net/http" neturl "net/url" + "os" "strings" "sync" "time" @@ -37,6 +37,16 @@ func kiroRegion(account *config.Account) string { } func kiroRegionForProfile(account *config.Account, profileArn string) string { + // api_key accounts have no profile ARN, so the profile-derived region is + // meaningless. Resolve from the EffectiveApiRegion fallback chain + // (ApiRegion > Region > global > us-east-1) — the same chain the chat path + // (CallKiroAPI) uses — so REST calls (getUsageLimits, ListAvailableModels, + // overage) target the same region as chat. Previously the REST path saw only + // account.Region, so ApiRegion and the global region were silently ignored, + // diverging from the chat path and triggering spurious 403s. + if account != nil && account.IsApiKeyCredential() { + return account.EffectiveApiRegion() + } if r := regionFromProfileArn(profileArn); r != "" { return r } @@ -52,17 +62,30 @@ func kiroRegionForProfile(account *config.Account, profileArn string) string { } // regionalizeURL points a hardcoded us-east-1 Kiro endpoint at the profile's -// data-plane region. Amazon Q is regional (q.{region}.amazonaws.com), but the CodeWhisperer -// REST host only exists in us-east-1 — non-us-east-1 accounts are served by the -// regional Amazon Q host instead. Both us-east-1 hosts map -// to q.{region}. It is a no-op for us-east-1 profiles. +// data-plane region (see regionalizeURLForProfile). It is a no-op for us-east-1. func regionalizeURL(rawURL string, account *config.Account) string { return regionalizeURLForProfile(rawURL, account, "") } +// regionalizeURLForProfile points a hardcoded us-east-1 Kiro endpoint at the +// data-plane region derived from the profile (payload ARN first, then the account's +// cached ARN, then account.Region). account.Region is the auth/OIDC region and can +// differ from the profile's region, so the profile ARN is preferred. func regionalizeURLForProfile(rawURL string, account *config.Account, profileArn string) string { - region := kiroRegionForProfile(account, profileArn) - if region == "us-east-1" { + return regionalizeURLForRegion(rawURL, kiroRegionForProfile(account, profileArn)) +} + +// regionalizeURLForRegion rewrites a hardcoded us-east-1 Kiro endpoint to target +// the given region. Amazon Q is regional (q.{region}.amazonaws.com), but the +// CodeWhisperer REST host only exists in us-east-1 — every other region is served +// by the regional Amazon Q host instead. So for a non-us-east-1 region BOTH +// us-east-1 hosts (q.us-east-1.* and codewhisperer.us-east-1.*) collapse onto +// q.{region}.amazonaws.com; there is deliberately no codewhisperer.{region} host. +// It is a no-op for us-east-1 or an empty region. This region-targeted primitive +// also backs cross-region profile probing (listAvailableProfilesInRegion). +func regionalizeURLForRegion(rawURL, region string) string { + region = strings.TrimSpace(region) + if region == "" || region == "us-east-1" { return rawURL } regionalHost := "q." + region + ".amazonaws.com" @@ -72,6 +95,71 @@ func regionalizeURLForProfile(rawURL string, account *config.Account, profileArn ).Replace(rawURL) } +// defaultKiroProfileRegions is the ordered set of regions probed when an account's +// home region is unknown. us-east-1 is the historical default every login falls +// back to; eu-central-1 is where EU-provisioned Azure-tenant profiles +// (e.g. KiroProfile-eu-central-1) live. Override or extend with the +// KIRO_PROFILE_REGIONS env var (comma-separated) to onboard further regions +// without a code change. +var defaultKiroProfileRegions = []string{"us-east-1", "eu-central-1"} + +// kiroProfileRegionCandidates returns the ordered, de-duplicated list of regions +// to probe for an account's Kiro profile. The account's currently-configured region +// is always tried first. Cross-region fallbacks are added for auth methods whose +// login region does not guarantee the profile region: external_idp (Azure-tenant, +// defaults to us-east-1) and idc (IAM Identity Center / enterprise SSO, where a +// us-east-1 portal can provision profiles in eu-central-1). social and Builder ID +// accounts carry their authoritative region and are probed against that single +// region only. KIRO_PROFILE_REGIONS, when set, replaces the built-in fallback set +// (the account region is still tried first). +func kiroProfileRegionCandidates(account *config.Account) []string { + seen := make(map[string]bool) + var out []string + add := func(region string) { + region = strings.TrimSpace(region) + if region == "" || seen[region] { + return + } + seen[region] = true + out = append(out, region) + } + + if account != nil { + add(account.Region) + } + if !shouldProbeFallbackRegions(account) { + return out + } + if env := strings.TrimSpace(os.Getenv("KIRO_PROFILE_REGIONS")); env != "" { + for _, r := range strings.Split(env, ",") { + add(r) + } + return out + } + for _, r := range defaultKiroProfileRegions { + add(r) + } + return out +} + +// shouldProbeFallbackRegions reports whether an account should probe fallback +// regions when its home region returns no profile. external_idp accounts always +// qualify (region defaulted to us-east-1 at login). idc (IAM Identity Center / +// enterprise SSO) accounts also qualify: an enterprise SSO portal in us-east-1 +// can provision profiles in eu-central-1, so the auth region does not imply the +// profile region. social and Builder ID accounts carry their authoritative +// region and are not probed further. +func shouldProbeFallbackRegions(account *config.Account) bool { + if account == nil { + return true + } + if strings.TrimSpace(account.Region) == "" { + return true + } + method := strings.TrimSpace(account.AuthMethod) + return strings.EqualFold(method, "external_idp") || strings.EqualFold(method, "idc") +} + // GetUsageLimits 获取账户使用量和订阅信息 func GetUsageLimits(account *config.Account) (*UsageLimitsResponse, error) { if err := ensureRestProfileArn(account); err != nil { @@ -182,6 +270,11 @@ func ResolveProfileArn(account *config.Account) (string, error) { if account == nil { return "", fmt.Errorf("account is nil") } + // api_key accounts authenticate directly (tokentype: API_KEY) and have no + // profile ARN. Skip the ListAvailableProfiles probe entirely. + if account.IsApiKeyCredential() { + return "", nil + } if profileArn := strings.TrimSpace(account.ProfileArn); profileArn != "" { return profileArn, nil } @@ -191,8 +284,14 @@ func ResolveProfileArn(account *config.Account) (string, error) { var profileUnsupported bool if !profileLookupSuppressed { - // Try ListAvailableProfiles first, retrying on transient failures. - profileArn, err := listAvailableProfilesWithRetry(account) + // Probe ListAvailableProfiles across candidate regions, retrying transient + // failures. The home region is unknown at login for Azure-tenant + // (external_idp) accounts (they default to us-east-1), so the probe is what + // discovers a profile that lives outside the account's configured region. The + // cached ARN then drives the data-plane region via kiroRegionForProfile — no + // separate region persistence is needed (and account.Region stays the auth + // region, which can legitimately differ from the profile's region). + profileArn, err := resolveProfileArnAcrossRegions(account) if err == nil && profileArn != "" { if updateErr := config.UpdateAccountProfileArn(account.ID, profileArn); updateErr != nil { logger.Warnf("[ProfileArn] Failed to cache profile ARN for %s: %v", account.Email, updateErr) @@ -206,7 +305,7 @@ func ResolveProfileArn(account *config.Account) (string, error) { // Fallback: refresh token to get profileArn from auth response if account.RefreshToken != "" { - _, _, _, refreshedArn, refreshErr := auth.RefreshToken(account) + _, _, _, refreshedArn, refreshErr := authRefreshToken(account) if refreshErr == nil && refreshedArn != "" { if updateErr := config.UpdateAccountProfileArn(account.ID, refreshedArn); updateErr != nil { logger.Warnf("[ProfileArn] Failed to cache profile ARN for %s: %v", account.Email, updateErr) @@ -305,16 +404,43 @@ func ensureRestProfileArn(account *config.Account) error { return nil } -func listAvailableProfilesWithRetry(account *config.Account) (string, error) { - // Retry transient failures (network errors, 5xx, 429) with short backoff. - // An empty profile list or 4xx (other than 429) is treated as authoritative - // and not retried — they reflect account state, not upstream flakiness. +// resolveProfileArnAcrossRegions probes ListAvailableProfiles against each +// candidate region (the account's configured region first, then the fallbacks) and +// returns the first profile ARN found. This is what lets an account whose profile +// lives outside its configured region — every Azure-tenant (external_idp) login +// defaults to us-east-1 — discover that profile (e.g. in eu-central-1) on first use. +// The returned ARN carries its own region, which kiroRegionForProfile then uses for +// data-plane calls. A correctly-regioned account resolves on the first probe. A +// Builder ID "unsupported" 403 is authoritative across all regions, so it +// short-circuits the probe rather than repeating per region. +func resolveProfileArnAcrossRegions(account *config.Account) (string, error) { + var lastErr error + for _, region := range kiroProfileRegionCandidates(account) { + arn, probeErr := listAvailableProfilesWithRetryInRegion(account, region) + if probeErr == nil && strings.TrimSpace(arn) != "" { + return arn, nil + } + if probeErr != nil { + lastErr = probeErr + if isBuilderIDProfileUnsupportedError(account, probeErr) { + return "", probeErr + } + } + } + return "", lastErr +} + +// listAvailableProfilesWithRetryInRegion calls ListAvailableProfiles against a +// specific region, retrying transient failures (network errors, 5xx, 429) with +// short backoff. An empty profile list or 4xx (other than 429) is treated as +// authoritative and not retried — they reflect account state, not upstream flakiness. +func listAvailableProfilesWithRetryInRegion(account *config.Account, region string) (string, error) { const maxAttempts = 3 backoff := 200 * time.Millisecond var lastErr error for attempt := 1; attempt <= maxAttempts; attempt++ { - profileArn, err := listAvailableProfiles(account) + profileArn, err := listAvailableProfilesInRegion(account, region) if err == nil { return profileArn, nil } @@ -322,8 +448,8 @@ func listAvailableProfilesWithRetry(account *config.Account) (string, error) { if !isTransientProfileFetchError(err) || attempt == maxAttempts { return "", err } - logger.Debugf("[ProfileArn] ListAvailableProfiles transient failure for %s (attempt %d/%d): %v", - account.Email, attempt, maxAttempts, err) + logger.Debugf("[ProfileArn] ListAvailableProfiles transient failure for %s in %s (attempt %d/%d): %v", + account.Email, region, attempt, maxAttempts, err) time.Sleep(backoff) backoff *= 2 } @@ -348,8 +474,14 @@ func isTransientProfileFetchError(err error) bool { return true } -func listAvailableProfiles(account *config.Account) (string, error) { - req, err := http.NewRequest("POST", regionalizeURL(fmt.Sprintf("%s/ListAvailableProfiles", kiroRestAPIBase), account), strings.NewReader(`{"maxResults":10}`)) +// listAvailableProfilesInRegion calls ListAvailableProfiles with the request host +// pointed at a specific region (q.{region} for non-us-east-1, the CodeWhisperer +// REST host for us-east-1). Targeting an explicit region — rather than the account's +// stored one — is what makes cross-region detection possible: the same credential is +// probed against each candidate region until one returns a profile. +func listAvailableProfilesInRegion(account *config.Account, region string) (string, error) { + endpoint := regionalizeURLForRegion(fmt.Sprintf("%s/ListAvailableProfiles", kiroRestAPIBase), region) + req, err := http.NewRequest("POST", endpoint, strings.NewReader(`{"maxResults":10}`)) if err != nil { return "", err } @@ -418,7 +550,7 @@ func RefreshAccountInfo(account *config.Account) (*config.AccountInfo, error) { errMsg := err.Error() if strings.Contains(errMsg, "TEMPORARILY_SUSPENDED") { // 账户被暂时封禁，自动禁用并标记封禁状态 - logger.Warnf("[RefreshAccountInfo] Account %s is temporarily suspended: %v", account.Email, err) + logger.Warnf("[RefreshAccountInfo] Account %s is temporarily suspended: %v", accountEmailForLog(account), err) // 更新账户封禁状态并自动禁用 updatedAccount := *account @@ -433,10 +565,12 @@ func RefreshAccountInfo(account *config.Account) (*config.AccountInfo, error) { } return nil, fmt.Errorf("Account suspended: %w", err) - } else if strings.Contains(errMsg, "403") || strings.Contains(errMsg, "401") || - strings.Contains(errMsg, "invalid") || strings.Contains(errMsg, "expired") { - // Token 相关错误，可能需要重新认证 - logger.Warnf("[RefreshAccountInfo] Authentication error for %s: %v", account.Email, err) + } else if isAuthErrorMessage(errMsg) && !account.IsApiKeyCredential() { + // Token 相关错误，可能需要重新认证。api_key 账户不在 refresh 路径封禁： + // usage-limit 403 并非权威信号（可能是 region/配置问题），且 api_key 没有 + // profile-ARN 跨区探测来交叉验证。chat 路径 (handleAccountFailure) 才是 + // api_key 有效性的权威来源——真实请求的 403 会在那里封禁。 + logger.Warnf("[RefreshAccountInfo] Authentication error for %s: %v", accountEmailForLog(account), err) // 更新账户封禁状态为认证失败并自动禁用 updatedAccount := *account @@ -456,7 +590,7 @@ func RefreshAccountInfo(account *config.Account) (*config.AccountInfo, error) { // 如果成功获取信息，清除封禁状态（如果之前被标记） if account.BanStatus != "" && account.BanStatus != "ACTIVE" { - logger.Infof("[RefreshAccountInfo] Account %s is now active, clearing ban status", account.Email) + logger.Infof("[RefreshAccountInfo] Account %s is now active, clearing ban status", accountEmailForLog(account)) updatedAccount := *account updatedAccount.BanStatus = "ACTIVE" diff --git a/proxy/kiro_api_test.go b/proxy/kiro_api_test.go index a4c12dfc..efb30a91 100644 --- a/proxy/kiro_api_test.go +++ b/proxy/kiro_api_test.go @@ -57,6 +57,26 @@ func TestRegionalizeURLForProfileUsesPayloadProfileArnRegion(t *testing.T) { } } +// TestRegionalizeURLApiKeyUsesEffectiveApiRegion guards the REST/chat region +// inconsistency: an api_key account's REST calls (getUsageLimits, +// ListAvailableModels, overage) must resolve the same region as the chat path +// (CallKiroAPI), i.e. use the EffectiveApiRegion fallback chain (ApiRegion > +// Region > global > us-east-1), not just account.Region. Previously the REST +// path saw only account.Region, so ApiRegion (and global) was ignored. +func TestRegionalizeURLApiKeyUsesEffectiveApiRegion(t *testing.T) { + account := &config.Account{ + AuthMethod: "api_key", + KiroApiKey: "k", + Region: "us-east-1", + ApiRegion: "eu-central-1", + } + got := regionalizeURL("https://q.us-east-1.amazonaws.com/getUsageLimits", account) + want := "https://q.eu-central-1.amazonaws.com/getUsageLimits" + if got != want { + t.Fatalf("api_key REST region: want %q, got %q", want, got) + } +} + func TestResolveProfileArnFetchesAndCachesProfile(t *testing.T) { configPath := filepath.Join(t.TempDir(), "config.json") if err := config.Init(configPath); err != nil { @@ -118,6 +138,28 @@ func TestResolveProfileArnFetchesAndCachesProfile(t *testing.T) { } } +// TestResolveProfileArnSkipsApiKeyAccount asserts an api_key account returns +// an empty profile ARN without any HTTP probe (the key authenticates directly). +func TestResolveProfileArnSkipsApiKeyAccount(t *testing.T) { + // If a request were made, this round-tripper fails the test. + kiroRestHttpStore.Store(&http.Client{ + Transport: roundTripFunc(func(*http.Request) (*http.Response, error) { + t.Fatal("unexpected HTTP request for api_key account profile ARN") + return nil, nil + }), + }) + t.Cleanup(func() { InitKiroHttpClient("") }) + + account := &config.Account{AuthMethod: "api_key", KiroApiKey: "key-abc", AccessToken: "key-abc"} + got, err := ResolveProfileArn(account) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if got != "" { + t.Fatalf("expected empty profile ARN for api_key account, got %q", got) + } +} + func TestResolveProfileArnSuppressesBuilderIDUnsupportedLookup(t *testing.T) { clearProfileArnResolutionCooldowns() t.Cleanup(clearProfileArnResolutionCooldowns) @@ -282,6 +324,120 @@ func TestRefreshAccountInfoDoesNotDisableBuilderIDWhenProfileLookupUnsupported(t } } +// TestRefreshAccountInfoDoesNotBanOnTransientUpstreamError guards the false-ban: +// a transient NON-auth upstream failure (HTTP 5xx) whose body happens to contain +// "expired"/"invalid" must NOT be misclassified as an authentication failure and +// must NOT disable + BAN an otherwise-healthy account. Only genuine auth errors +// (HTTP 401/403, invalid_grant, expired-token variants) warrant a ban. +func TestRefreshAccountInfoDoesNotBanOnTransientUpstreamError(t *testing.T) { + clearProfileArnResolutionCooldowns() + t.Cleanup(clearProfileArnResolutionCooldowns) + + configPath := filepath.Join(t.TempDir(), "config.json") + if err := config.Init(configPath); err != nil { + t.Fatalf("init config: %v", err) + } + account := config.Account{ + ID: "acct-transient-500", + Email: "user@example.com", + AccessToken: "access-token", + Provider: "BuilderId", + Region: "us-east-1", + Enabled: true, + } + if err := config.AddAccount(account); err != nil { + t.Fatalf("add account: %v", err) + } + + kiroRestHttpStore.Store(&http.Client{ + Transport: roundTripFunc(func(req *http.Request) (*http.Response, error) { + switch req.URL.Path { + case "/ListAvailableProfiles": + return &http.Response{ + StatusCode: http.StatusForbidden, + Body: io.NopCloser(strings.NewReader(`{"message":"AWS Builder ID is not supported for this operation.","reason":null}`)), + Header: make(http.Header), + }, nil + case "/getUsageLimits": + // Transient 5xx whose body text contains "expired" — must not be + // read as an auth/token failure. + return &http.Response{ + StatusCode: http.StatusInternalServerError, + Body: io.NopCloser(strings.NewReader(`{"message":"The cached configuration is expired, please retry"}`)), + Header: make(http.Header), + }, nil + default: + t.Fatalf("unexpected request path %s", req.URL.Path) + return nil, nil + } + }), + }) + t.Cleanup(func() { InitKiroHttpClient("") }) + + requestAccount := account + if _, err := RefreshAccountInfo(&requestAccount); err == nil { + t.Fatalf("expected GetUsageLimits 500 error to propagate, got nil") + } + + accounts := config.GetAccounts() + if len(accounts) != 1 { + t.Fatalf("expected one account, got %d", len(accounts)) + } + if !accounts[0].Enabled || accounts[0].BanStatus == "BANNED" { + t.Fatalf("transient 5xx with body containing \"expired\" must NOT ban the account, got enabled=%v banStatus=%q banReason=%q", + accounts[0].Enabled, accounts[0].BanStatus, accounts[0].BanReason) + } +} + +// TestRefreshAccountInfoDoesNotBanApiKeyOnUsageLimit403 pins the invariant that an +// api_key account is NOT auto-banned when a usage-limit refresh returns 403. api_key +// accounts authenticate directly and have no profile-ARN cross-region probe, so a +// usage-limit 403 is not authoritative (could be region/config). The chat path +// (handleAccountFailure) is the source of truth and bans a genuinely-invalid key on a +// real request. Banning here would permanently disable a key over a spurious 403. +func TestRefreshAccountInfoDoesNotBanApiKeyOnUsageLimit403(t *testing.T) { + configPath := filepath.Join(t.TempDir(), "config.json") + if err := config.Init(configPath); err != nil { + t.Fatalf("init config: %v", err) + } + account := config.Account{ + ID: "apikey-refresh-1", + AuthMethod: "api_key", + KiroApiKey: "ksk_test_key_value", + AccessToken: "ksk_test_key_value", + Region: "us-east-1", + Enabled: true, + } + if err := config.AddAccount(account); err != nil { + t.Fatalf("add account: %v", err) + } + + kiroRestHttpStore.Store(&http.Client{ + Transport: roundTripFunc(func(*http.Request) (*http.Response, error) { + return &http.Response{ + StatusCode: http.StatusForbidden, + Body: io.NopCloser(strings.NewReader(`{"message":"The bearer token included in the request is invalid.","reason":null}`)), + Header: make(http.Header), + }, nil + }), + }) + t.Cleanup(func() { InitKiroHttpClient("") }) + + requestAccount := account + if _, err := RefreshAccountInfo(&requestAccount); err == nil { + t.Fatalf("expected GetUsageLimits 403 to propagate, got nil") + } + + accounts := config.GetAccounts() + if len(accounts) != 1 { + t.Fatalf("expected one account, got %d", len(accounts)) + } + if !accounts[0].Enabled || accounts[0].BanStatus == "BANNED" { + t.Fatalf("api_key account must NOT be banned on usage-limit 403 (chat path is source of truth), got enabled=%v banStatus=%q banReason=%q", + accounts[0].Enabled, accounts[0].BanStatus, accounts[0].BanReason) + } +} + func clearProfileArnResolutionCooldowns() { profileArnResolutionCooldowns.Range(func(key, _ interface{}) bool { profileArnResolutionCooldowns.Delete(key) diff --git a/proxy/kiro_external_idp_test.go b/proxy/kiro_external_idp_test.go new file mode 100644 index 00000000..8420ec44 --- /dev/null +++ b/proxy/kiro_external_idp_test.go @@ -0,0 +1,49 @@ +package proxy + +import ( + "kiro-go/config" + "net/http" + "testing" +) + +// TestSetKiroHeadersExternalIdpSendsTokenType verifies that external-IdP +// (enterprise SSO, e.g. Azure AD) accounts get the TokenType: EXTERNAL_IDP +// header. Without it, CodeWhisperer silently returns an empty profile list and +// rejects data-plane calls, which is exactly what broke Azure-tenant SSO. +func TestSetKiroHeadersExternalIdpSendsTokenType(t *testing.T) { + account := &config.Account{ + AccessToken: "azure-access-token", + AuthMethod: "external_idp", + MachineId: "machine-xyz", + } + req, err := http.NewRequest("POST", "https://codewhisperer.us-east-1.amazonaws.com/ListAvailableProfiles", nil) + if err != nil { + t.Fatalf("new request: %v", err) + } + + setKiroHeaders(req, account) + + if got := req.Header.Get("TokenType"); got != "EXTERNAL_IDP" { + t.Fatalf("expected TokenType=EXTERNAL_IDP for external_idp account, got %q", got) + } + if got := req.Header.Get("Authorization"); got != "Bearer azure-access-token" { + t.Fatalf("expected bearer authorization, got %q", got) + } +} + +// TestSetKiroHeadersNonExternalIdpOmitsTokenType verifies the header is NOT sent +// for the other auth methods, so existing Builder ID / IDC / social accounts are +// unaffected. +func TestSetKiroHeadersNonExternalIdpOmitsTokenType(t *testing.T) { + for _, method := range []string{"idc", "social", ""} { + account := &config.Account{AccessToken: "tok", AuthMethod: method, MachineId: "m"} + req, err := http.NewRequest("POST", "https://codewhisperer.us-east-1.amazonaws.com/ListAvailableProfiles", nil) + if err != nil { + t.Fatalf("new request: %v", err) + } + setKiroHeaders(req, account) + if got := req.Header.Get("TokenType"); got != "" { + t.Fatalf("auth method %q should not set TokenType, got %q", method, got) + } + } +} diff --git a/proxy/kiro_headers.go b/proxy/kiro_headers.go index baf3fc64..63bf29a6 100644 --- a/proxy/kiro_headers.go +++ b/proxy/kiro_headers.go @@ -26,10 +26,27 @@ func buildRuntimeHeaderValues(account *config.Account, host string) kiroHeaderVa } func buildKiroHeaderValues(account *config.Account, host, apiName, sdkVersion, mode string) kiroHeaderValues { - clientCfg := config.GetKiroClientConfig() + accountID := "" + if account != nil { + accountID = account.ID + } + // Derive the desktop fingerprint (platform/version) from the same accountID + // used for the device id below, so an account's platform and device stay + // mutually consistent — one stable mac/win install, never a Linux server. + clientCfg := config.GetKiroClientConfig(accountID) machineID := "" if account != nil { machineID = account.MachineId + // Never let machineId be empty: an empty value drops the UA suffix and + // makes this account's User-Agent identical to every other empty-id + // account, which is the strongest cross-account association signal + // upstream can correlate on. Fall back to a stable, account-derived + // 64-hex id (sha256) so the UA always carries a unique, fixed device + // fingerprint — same value per account across requests and restarts. + // Ported from kiro-tutu (zero-dep: config.DeriveMachineId is stdlib). + if machineID == "" && account.ID != "" { + machineID = config.DeriveMachineId(account.ID) + } } userAgent := fmt.Sprintf( @@ -56,12 +73,26 @@ func buildKiroHeaderValues(account *config.Account, host, apiName, sdkVersion, m } func applyKiroBaseHeaders(req *http.Request, account *config.Account, values kiroHeaderValues) { - if account != nil && account.AccessToken != "" { - req.Header.Set("Authorization", "Bearer "+account.AccessToken) + if account != nil { + if account.IsApiKeyCredential() { + if account.KiroApiKey != "" { + req.Header.Set("Authorization", "Bearer "+account.KiroApiKey) + req.Header.Set("tokentype", "API_KEY") + } + } else if account.AccessToken != "" { + req.Header.Set("Authorization", "Bearer "+account.AccessToken) + } } req.Header.Set("User-Agent", values.UserAgent) req.Header.Set("x-amz-user-agent", values.AmzUserAgent) req.Header.Set("x-amzn-codewhisperer-optout", "true") + // External IdP (enterprise SSO, e.g. Azure AD) tokens MUST carry this header or + // CodeWhisperer does not recognize the token type and silently returns an empty + // profile list (and rejects data-plane calls). With it, a provisioned account + // resolves its profile; an unprovisioned one gets a clear 403. + if account != nil && account.AuthMethod == "external_idp" { + req.Header.Set("TokenType", "EXTERNAL_IDP") + } if values.Host != "" { req.Host = values.Host } diff --git a/proxy/kiro_headers_test.go b/proxy/kiro_headers_test.go index a4b08055..d01ee535 100644 --- a/proxy/kiro_headers_test.go +++ b/proxy/kiro_headers_test.go @@ -2,6 +2,8 @@ package proxy import ( "kiro-go/config" + "net/http" + "net/http/httptest" "strings" "testing" ) @@ -41,3 +43,88 @@ func TestBuildRuntimeHeaderValuesUsesRuntimeAPIFormat(t *testing.T) { t.Fatalf("expected runtime mode marker in user agent, got %q", values.UserAgent) } } + +// TestBuildKiroHeaderValuesFallsBackToDerivedMachineId asserts an account with an +// empty MachineId still gets a unique, stable device suffix in the User-Agent. +// Without the fallback every empty-id account shares an identical UA, which is the +// strongest cross-account association signal upstream can correlate on. The fallback +// must be deterministic so the same account always looks like the same device. +func TestBuildKiroHeaderValuesFallsBackToDerivedMachineId(t *testing.T) { + acc := &config.Account{ID: "acct-empty-id", MachineId: ""} + a := buildKiroHeaderValues(acc, "q.us-east-1.amazonaws.com", "codewhispererstreaming", "1.0.34", "m/E") + b := buildKiroHeaderValues(acc, "q.us-east-1.amazonaws.com", "codewhispererstreaming", "1.0.34", "m/E") + + // Deterministic: same account → same UA. + if a.UserAgent != b.UserAgent { + t.Fatalf("empty-id fallback must be stable across calls:\n a=%q\n b=%q", a.UserAgent, b.UserAgent) + } + // UA must carry a device suffix (the derived id) rather than ending at the version. + suffix := "-" + config.DeriveMachineId("acct-empty-id") + if !strings.Contains(a.UserAgent, suffix) { + t.Fatalf("expected UA to carry derived machine id suffix %q, got %q", suffix, a.UserAgent) + } + + // Distinct across accounts: two empty-id accounts must NOT share a UA. + other := &config.Account{ID: "acct-other-id", MachineId: ""} + c := buildKiroHeaderValues(other, "q.us-east-1.amazonaws.com", "codewhispererstreaming", "1.0.34", "m/E") + if a.UserAgent == c.UserAgent { + t.Fatalf("two distinct empty-id accounts must get distinct device suffixes, both got %q", a.UserAgent) + } + + // An explicitly-set MachineId still wins over the fallback. + explicit := &config.Account{ID: "acct-empty-id", MachineId: "explicit-machine"} + e := buildKiroHeaderValues(explicit, "q.us-east-1.amazonaws.com", "codewhispererstreaming", "1.0.34", "m/E") + if !strings.Contains(e.UserAgent, "-explicit-machine") { + t.Fatalf("explicit MachineId must win over fallback, got %q", e.UserAgent) + } + + // A nil account (no ID to derive from) keeps the suffix-less UA — nothing to derive from. + none := buildKiroHeaderValues(nil, "q.us-east-1.amazonaws.com", "codewhispererstreaming", "1.0.34", "m/E") + if strings.Contains(none.UserAgent, suffix) { + t.Fatalf("nil account must not carry a derived suffix, got %q", none.UserAgent) + } +} + +// TestApplyKiroBaseHeadersApiKeyBranch asserts an api_key account sends the key +// as the bearer with tokentype: API_KEY, while an OAuth account keeps the plain +// bearer and the external_idp TokenType branch is preserved. +func TestApplyKiroBaseHeadersApiKeyBranch(t *testing.T) { + mk := func(account *config.Account) http.Header { + req := httptest.NewRequest("POST", "https://q.us-east-1.amazonaws.com/x", nil) + applyKiroBaseHeaders(req, account, kiroHeaderValues{UserAgent: "ua", AmzUserAgent: "aua"}) + return req.Header + } + + // api_key: bearer is the key, tokentype is API_KEY. + h := mk(&config.Account{AuthMethod: "api_key", KiroApiKey: "key-xyz", AccessToken: "key-xyz"}) + if got := h.Get("Authorization"); got != "Bearer key-xyz" { + t.Fatalf("api_key bearer: want %q, got %q", "Bearer key-xyz", got) + } + if got := h.Get("tokentype"); got != "API_KEY" { + t.Fatalf("api_key tokentype: want API_KEY, got %q", got) + } + + // OAuth (idc): plain bearer, no tokentype header. + h = mk(&config.Account{AuthMethod: "idc", AccessToken: "at-oauth"}) + if got := h.Get("Authorization"); got != "Bearer at-oauth" { + t.Fatalf("oauth bearer: want %q, got %q", "Bearer at-oauth", got) + } + if got := h.Get("tokentype"); got != "" { + t.Fatalf("oauth must NOT set tokentype, got %q", got) + } + + // external_idp preserved: the (canonical) Tokentype header carries EXTERNAL_IDP. + // "tokentype" and "TokenType" canonicalize to the same key in net/http, but + // api_key and external_idp are mutually exclusive account types, so the two + // values never collide at runtime. + h = mk(&config.Account{AuthMethod: "external_idp", AccessToken: "at-ext"}) + if got := h.Get("TokenType"); got != "EXTERNAL_IDP" { + t.Fatalf("external_idp TokenType: want EXTERNAL_IDP, got %q", got) + } + + // empty account: no Authorization header. + h = mk(&config.Account{}) + if got := h.Get("Authorization"); got != "" { + t.Fatalf("empty account must not set Authorization, got %q", got) + } +} diff --git a/proxy/kiro_region_test.go b/proxy/kiro_region_test.go new file mode 100644 index 00000000..d5040cac --- /dev/null +++ b/proxy/kiro_region_test.go @@ -0,0 +1,138 @@ +package proxy + +import ( + "kiro-go/config" + "testing" +) + +// TestRegionalizeURLForRegion asserts that a non-us-east-1 region collapses BOTH +// hardcoded us-east-1 hosts (q.* and codewhisperer.*) onto q.{region} — there is no +// codewhisperer.{region} host — and that us-east-1/empty are no-ops. +func TestRegionalizeURLForRegion(t *testing.T) { + cases := []struct { + name string + rawURL string + region string + want string + }{ + { + name: "codewhisperer host to q.eu-central-1", + rawURL: "https://codewhisperer.us-east-1.amazonaws.com/ListAvailableProfiles", + region: "eu-central-1", + want: "https://q.eu-central-1.amazonaws.com/ListAvailableProfiles", + }, + { + name: "q host to q.eu-central-1", + rawURL: "https://q.us-east-1.amazonaws.com/generateAssistantResponse", + region: "eu-central-1", + want: "https://q.eu-central-1.amazonaws.com/generateAssistantResponse", + }, + { + name: "us-east-1 is a no-op (codewhisperer host kept)", + rawURL: "https://codewhisperer.us-east-1.amazonaws.com/ListAvailableProfiles", + region: "us-east-1", + want: "https://codewhisperer.us-east-1.amazonaws.com/ListAvailableProfiles", + }, + { + name: "empty region is a no-op", + rawURL: "https://q.us-east-1.amazonaws.com/x", + region: "", + want: "https://q.us-east-1.amazonaws.com/x", + }, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + got := regionalizeURLForRegion(tc.rawURL, tc.region) + if got != tc.want { + t.Fatalf("regionalizeURLForRegion(%q, %q) = %q, want %q", tc.rawURL, tc.region, got, tc.want) + } + }) + } +} + +// TestRegionalizeURLForRegionNoCodewhispererRegionalHost guards the user-stated +// invariant directly: a regionalized URL must never produce codewhisperer.{region}. +func TestRegionalizeURLForRegionNoCodewhispererRegionalHost(t *testing.T) { + got := regionalizeURLForRegion("https://codewhisperer.us-east-1.amazonaws.com/GetUserInfo", "eu-central-1") + if want := "https://q.eu-central-1.amazonaws.com/GetUserInfo"; got != want { + t.Fatalf("got %q, want %q (must not be codewhisperer.eu-central-1)", got, want) + } +} + +// TestKiroProfileRegionCandidatesExternalIdp checks that an external_idp account — +// whose home region is unknown and defaults to us-east-1 — probes the account region +// first and then the built-in fallbacks, de-duplicated. +func TestKiroProfileRegionCandidatesExternalIdp(t *testing.T) { + // Default us-east-1 external_idp login: fallbacks follow. + got := kiroProfileRegionCandidates(&config.Account{AuthMethod: "external_idp", Region: "us-east-1"}) + assertOrder(t, got, []string{"us-east-1", "eu-central-1"}) + + // Already-detected eu-central-1 leads; us-east-1 fallback follows. + got = kiroProfileRegionCandidates(&config.Account{AuthMethod: "external_idp", Region: "eu-central-1"}) + assertOrder(t, got, []string{"eu-central-1", "us-east-1"}) + + // A non-default region leads, both defaults follow. + got = kiroProfileRegionCandidates(&config.Account{AuthMethod: "external_idp", Region: "ap-southeast-2"}) + assertOrder(t, got, []string{"ap-southeast-2", "us-east-1", "eu-central-1"}) +} + +// TestKiroProfileRegionCandidatesNoRegion checks an account with no region set falls +// back across the defaults regardless of auth method. +func TestKiroProfileRegionCandidatesNoRegion(t *testing.T) { + got := kiroProfileRegionCandidates(&config.Account{}) + assertOrder(t, got, []string{"us-east-1", "eu-central-1"}) +} + +// TestKiroProfileRegionCandidatesSingleRegionAuthMethods checks that social/ +// Builder ID accounts — which already carry their authoritative region — are probed +// against that single region only, with no fallback probing. +func TestKiroProfileRegionCandidatesSingleRegionAuthMethods(t *testing.T) { + for _, method := range []string{"social", "builderId", ""} { + got := kiroProfileRegionCandidates(&config.Account{AuthMethod: method, Region: "eu-central-1"}) + if len(got) != 1 || got[0] != "eu-central-1" { + t.Fatalf("authMethod %q: candidate regions = %v, want [eu-central-1] only", method, got) + } + } +} + +// TestKiroProfileRegionCandidatesIdcFallback checks that idc (IAM Identity Center / +// enterprise SSO) accounts — whose SSO portal region (e.g. us-east-1) does not +// guarantee the profile region — probe fallback regions just like external_idp. +func TestKiroProfileRegionCandidatesIdcFallback(t *testing.T) { + // Default us-east-1 idc login: fallbacks follow. + got := kiroProfileRegionCandidates(&config.Account{AuthMethod: "idc", Region: "us-east-1"}) + assertOrder(t, got, []string{"us-east-1", "eu-central-1"}) + + // Already eu-central-1 leads; us-east-1 fallback follows. + got = kiroProfileRegionCandidates(&config.Account{AuthMethod: "idc", Region: "eu-central-1"}) + assertOrder(t, got, []string{"eu-central-1", "us-east-1"}) +} + +// TestKiroProfileRegionCandidatesEnvOverride checks KIRO_PROFILE_REGIONS replaces +// the built-in fallbacks (external_idp only) while the account region is tried first. +func TestKiroProfileRegionCandidatesEnvOverride(t *testing.T) { + t.Setenv("KIRO_PROFILE_REGIONS", "eu-west-1, ap-south-1 ,eu-west-1") + got := kiroProfileRegionCandidates(&config.Account{AuthMethod: "external_idp", Region: "us-east-1"}) + // us-east-1 (account) first; env values de-duplicated and trimmed; no built-in defaults. + assertOrder(t, got, []string{"us-east-1", "eu-west-1", "ap-south-1"}) + + // An idc account also uses env fallbacks (enterprise SSO can have cross-region profiles). + got = kiroProfileRegionCandidates(&config.Account{AuthMethod: "idc", Region: "us-east-1"}) + assertOrder(t, got, []string{"us-east-1", "eu-west-1", "ap-south-1"}) + + // A social account ignores the env fallbacks entirely. + got = kiroProfileRegionCandidates(&config.Account{AuthMethod: "social", Region: "us-east-1"}) + assertOrder(t, got, []string{"us-east-1"}) +} + +func assertOrder(t *testing.T, got, want []string) { + t.Helper() + if len(got) != len(want) { + t.Fatalf("candidate regions = %v, want %v", got, want) + } + for i := range want { + if got[i] != want[i] { + t.Fatalf("candidate regions = %v, want %v", got, want) + } + } +} diff --git a/proxy/kiro_test.go b/proxy/kiro_test.go index 20d64366..8f4bab21 100644 --- a/proxy/kiro_test.go +++ b/proxy/kiro_test.go @@ -4,9 +4,11 @@ import ( "bytes" "encoding/binary" "encoding/json" + "io" "kiro-go/config" "net/http" "net/url" + "strings" "testing" "time" ) @@ -269,3 +271,54 @@ func awsEventStreamFrame(t *testing.T, eventType string, payload map[string]inte frame = append(frame, 0, 0, 0, 0) return frame } + +// TestCallKiroAPIRebuildsRegionForApiKeyAccount asserts an api_key account's +// endpoint host is rebuilt from EffectiveApiRegion (q. for non-us-east-1 +// endpoints), since it has no profile ARN to derive the region from. +func TestCallKiroAPIRebuildsRegionForApiKeyAccount(t *testing.T) { + if err := config.Init(t.TempDir() + "/config.json"); err != nil { + t.Fatalf("config.Init: %v", err) + } + var gotHost string + // CallKiroAPI uses the streaming store (kiroHttpStore), not the REST store. + kiroHttpStore.Store(&http.Client{ + Transport: roundTripFunc(func(req *http.Request) (*http.Response, error) { + gotHost = req.URL.Host + return &http.Response{ + StatusCode: 200, + Body: io.NopCloser(strings.NewReader("")), + Header: make(http.Header), + }, nil + }), + }) + t.Cleanup(func() { InitKiroHttpClient("") }) + + apiKeyAcct := &config.Account{ + AuthMethod: "api_key", + KiroApiKey: "k", + AccessToken: "k", + ApiRegion: "eu-central-1", + } + payload := &KiroPayload{} + _ = CallKiroAPI(apiKeyAcct, payload, nil) + if gotHost != "q.eu-central-1.amazonaws.com" { + t.Fatalf("api_key host: want q.eu-central-1.amazonaws.com, got %q", gotHost) + } +} + +// TestAccountEmailForLogMasksApiKey ensures api_key accounts (which have no email +// until their first successful refresh) are not blind in logs: fall back to the masked +// key. When an email is present it wins. +func TestAccountEmailForLogMasksApiKey(t *testing.T) { + acct := &config.Account{AuthMethod: "api_key", KiroApiKey: "ksk_abcdefghijklmnop"} + if got := accountEmailForLog(acct); got != "ksk_ab…mnop" { + t.Fatalf("api_key no-email: want masked key %q, got %q", "ksk_ab…mnop", got) + } + acct.Email = "user@example.com" + if got := accountEmailForLog(acct); got != "user@example.com" { + t.Fatalf("email present: want %q, got %q", "user@example.com", got) + } + if got := accountEmailForLog(nil); got != "" { + t.Fatalf("nil account: want %q, got %q", "", got) + } +} diff --git a/proxy/persist.go b/proxy/persist.go new file mode 100644 index 00000000..5eaf6f04 --- /dev/null +++ b/proxy/persist.go @@ -0,0 +1,269 @@ +package proxy + +// Durable request audit / access log (production persistence + observability). +// +// Records one structured entry per HTTP request: request id, hashed API key, +// method, path, status, latency and response size. Writes are asynchronous and +// non-blocking — if the buffer is full an entry is dropped (and counted) rather +// than stalling the proxied request. +// +// Backend: append-only JSON-lines file (zero dependency). The tutu reference +// defaults to SQLite (modernc.org/sqlite) with a JSONL fallback; this repo keeps +// a single dependency (github.com/google/uuid), so only the JSONL backend is +// compiled in. KIRO_AUDIT_DISABLE turns audit off entirely. +// +// Each entry is flushed immediately (bufio.Flush after every write) so a crash +// loses at most the single in-flight record — no WAL or batched fsync needed. +import ( + "bufio" + "crypto/sha256" + "encoding/hex" + "encoding/json" + "fmt" + "kiro-go/logger" + "os" + "path/filepath" + "strings" + "sync" +) + +// AuditEntry is one persisted request record. +type AuditEntry struct { + RequestID string `json:"request_id"` + TimeUnixMs int64 `json:"ts_ms"` + APIKeyHash string `json:"api_key_hash,omitempty"` + Method string `json:"method"` + Path string `json:"path"` + Status int `json:"status"` + LatencyMs int64 `json:"latency_ms"` + Bytes int `json:"bytes"` +} + +type auditBackend interface { + write(AuditEntry) error + queryRecent(limit int) ([]AuditEntry, error) + Close() error +} + +// ---- JSON-lines backend ---- + +// jsonlAuditMaxTailBytes bounds how much of the file queryRecent reads when +// looking for the most recent entries. 4 MiB is far more than limit*entrySize +// for any sane limit, while keeping a tail read cheap on a large audit file. +const jsonlAuditMaxTailBytes int64 = 4 << 20 + +type jsonlAudit struct { + mu sync.Mutex + f *os.File + w *bufio.Writer +} + +func newJSONLAudit(path string) (*jsonlAudit, error) { + f, err := os.OpenFile(path, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0o600) + if err != nil { + return nil, err + } + return &jsonlAudit{f: f, w: bufio.NewWriter(f)}, nil +} + +func (j *jsonlAudit) write(e AuditEntry) error { + b, err := json.Marshal(e) + if err != nil { + return err + } + j.mu.Lock() + defer j.mu.Unlock() + if _, err := j.w.Write(b); err != nil { + return err + } + if err := j.w.WriteByte('\n'); err != nil { + return err + } + return j.w.Flush() // flush each entry so a crash loses at most the in-flight write +} + +// queryRecent returns the most recent up-to-limit entries by reading the file +// tail. The file is append-only, so the newest entries are at the end. +func (j *jsonlAudit) queryRecent(limit int) ([]AuditEntry, error) { + if limit <= 0 || limit > 1000 { + limit = 100 + } + j.mu.Lock() + path := j.f.Name() + j.mu.Unlock() + + // Open a separate read handle so the writer's appends keep flowing. + rf, err := os.Open(path) + if err != nil { + return nil, err + } + defer rf.Close() + + st, err := rf.Stat() + if err != nil { + return nil, err + } + size := st.Size() + off := int64(0) + if size > jsonlAuditMaxTailBytes { + off = size - jsonlAuditMaxTailBytes + } + if _, err := rf.Seek(off, 0); err != nil { + return nil, err + } + + // Sliding window of the last `limit` decoded entries. + var ring []AuditEntry + sc := bufio.NewScanner(rf) + sc.Buffer(make([]byte, 64*1024), 4*1024*1024) + for sc.Scan() { + line := strings.TrimSpace(sc.Text()) + if line == "" { + continue + } + var e AuditEntry + if err := json.Unmarshal([]byte(line), &e); err != nil { + continue // skip malformed/partial line (e.g. mid-write tail) + } + ring = append(ring, e) + if len(ring) > limit { + ring = ring[len(ring)-limit:] + } + } + if err := sc.Err(); err != nil { + return nil, err + } + return ring, nil +} + +func (j *jsonlAudit) Close() error { + j.mu.Lock() + defer j.mu.Unlock() + if err := j.w.Flush(); err != nil { + _ = j.f.Close() + return err + } + return j.f.Close() +} + +// ---- Async store ---- + +// AuditStore drains entries from a buffered channel to its backend on a single +// writer goroutine. Record never blocks the caller. +type AuditStore struct { + ch chan AuditEntry + backend auditBackend + wg sync.WaitGroup + closeOnce sync.Once +} + +func NewAuditStore(backend auditBackend, bufSize int) *AuditStore { + if bufSize <= 0 { + bufSize = 4096 + } + s := &AuditStore{ch: make(chan AuditEntry, bufSize), backend: backend} + s.wg.Add(1) + go s.loop() + return s +} + +func (s *AuditStore) loop() { + defer s.wg.Done() + for e := range s.ch { + if err := s.backend.write(e); err != nil { + logger.Warnf("[Audit] write failed: %v", err) + } + } +} + +// Record enqueues an entry without blocking; a full buffer drops the entry. +func (s *AuditStore) Record(e AuditEntry) { + select { + case s.ch <- e: + default: + metricsAuditDropped() + } +} + +// Close stops the writer, flushing buffered entries, and closes the backend. +func (s *AuditStore) Close() { + s.closeOnce.Do(func() { + close(s.ch) + s.wg.Wait() + if err := s.backend.Close(); err != nil { + logger.Warnf("[Audit] close failed: %v", err) + } + }) +} + +// ---- Package-level wiring ---- + +var globalAudit *AuditStore + +// InitAuditStore initializes the process-wide audit store under dir. It is a +// no-op when KIRO_AUDIT_DISABLE is set. +func InitAuditStore(dir string) { + if strings.TrimSpace(os.Getenv("KIRO_AUDIT_DISABLE")) != "" { + logger.Infof("[Audit] disabled via KIRO_AUDIT_DISABLE") + return + } + if dir == "" { + dir = "." + } + if err := os.MkdirAll(dir, 0o755); err != nil { + logger.Warnf("[Audit] cannot create dir %s: %v", dir, err) + return + } + + // KIRO_AUDIT_BACKEND is accepted for compatibility with kiro-tutu, but only + // the JSONL backend is compiled in here (zero-dep). Any non-"jsonl" value + // falls back to JSONL with a one-time info note so operators aren't confused. + backendName := strings.ToLower(strings.TrimSpace(os.Getenv("KIRO_AUDIT_BACKEND"))) + if backendName != "" && backendName != "jsonl" { + logger.Infof("[Audit] backend %q requested but only JSONL is compiled in; using JSONL", backendName) + } + backendName = "jsonl" + + jb, err := newJSONLAudit(filepath.Join(dir, "kiro-audit.jsonl")) + if err != nil { + logger.Errorf("[Audit] disabled: cannot open JSONL backend: %v", err) + return + } + + globalAudit = NewAuditStore(jb, 4096) + logger.Infof("[Audit] enabled (backend=%s, dir=%s)", backendName, dir) +} + +// CloseAuditStore flushes and closes the audit store, if any. +func CloseAuditStore() { + if globalAudit != nil { + globalAudit.Close() + } +} + +func recordAudit(e AuditEntry) { + if globalAudit != nil { + globalAudit.Record(e) + } +} + +func hashAPIKey(key string) string { + if key == "" { + return "" + } + sum := sha256.Sum256([]byte(key)) + return hex.EncodeToString(sum[:]) +} + +// QueryRecent returns the most recent audit entries when the backend supports it. +func (s *AuditStore) QueryRecent(limit int) ([]AuditEntry, error) { + return s.backend.queryRecent(limit) +} + +// AuditRecent reads recent audit entries from the process-wide store. +func AuditRecent(limit int) ([]AuditEntry, error) { + if globalAudit == nil { + return nil, fmt.Errorf("audit store not initialized") + } + return globalAudit.QueryRecent(limit) +} diff --git a/proxy/prod_gauges.go b/proxy/prod_gauges.go new file mode 100644 index 00000000..03a3fec1 --- /dev/null +++ b/proxy/prod_gauges.go @@ -0,0 +1,51 @@ +package proxy + +// Pool-backed Prometheus gauges. Registered as the metrics gaugeProvider so a +// /metrics scrape reflects live pool state: account counts, per-account +// in-flight dispatch load, and remaining credential quota. Collection is +// guarded so a scrape can never panic the metrics endpoint. +// +// Ported from kiro-tutu. kiro_account_inflight is emitted from +// pool.GetRuntimeStatsSnapshot, populated by the health-scoring + inFlight +// dispatch slice. +import "kiro-go/pool" + +func init() { + gaugeProvider = collectPoolGauges +} + +func collectPoolGauges() (out []gaugeSample) { + defer func() { + if r := recover(); r != nil { + out = nil + } + }() + + p := pool.GetPool() + out = append(out, + gaugeSample{name: "kiro_accounts_total", value: float64(p.Count())}, + gaugeSample{name: "kiro_accounts_available", value: float64(p.AvailableCount())}, + ) + + // In-flight per account, exposed now that the pool tracks dispatch-level + // in-flight load (health-scoring slice). A non-zero value means the smart + // scheduler is actively routing that account. + for id, s := range p.GetRuntimeStatsSnapshot() { + out = append(out, gaugeSample{ + name: "kiro_account_inflight", + labels: [][2]string{{"account", id}}, + value: float64(s.InFlight), + }) + } + + for _, acc := range p.GetAllAccounts() { + if acc.UsageLimit > 0 { + out = append(out, gaugeSample{ + name: "kiro_credential_quota_remaining", + labels: [][2]string{{"account", acc.ID}}, + value: acc.UsageLimit - acc.UsageCurrent, + }) + } + } + return out +} diff --git a/proxy/prod_metrics.go b/proxy/prod_metrics.go new file mode 100644 index 00000000..060f9f31 --- /dev/null +++ b/proxy/prod_metrics.go @@ -0,0 +1,212 @@ +package proxy + +// Dependency-free Prometheus metrics exposition (production observability). +// +// Rather than pull in a metrics client library, this exposes a tiny, correct +// subset of the Prometheus text format (v0.0.4). Collection is lock-guarded and +// off the request hot path only by a single map insert + a few atomics, so it +// never blocks the proxied stream. Ported from kiro-tutu. The audit counter +// kiro_audit_dropped_total is incremented by the async audit store +// (persist.go) when its buffer overflows. +import ( + "fmt" + "sort" + "strings" + "sync" + "time" +) + +var globalMetrics = newProdMetrics() + +type prodMetrics struct { + mu sync.Mutex + counters map[string]int64 // fully-rendered series (name or name{labels}) -> value + + hBuckets []float64 + hCounts []uint64 // cumulative counts aligned with hBuckets + hSum float64 + hCount uint64 +} + +func newProdMetrics() *prodMetrics { + buckets := []float64{0.05, 0.1, 0.25, 0.5, 1, 2.5, 5, 10, 30, 60} + return &prodMetrics{ + counters: make(map[string]int64), + hBuckets: buckets, + hCounts: make([]uint64, len(buckets)), + } +} + +var counterHelp = map[string]string{ + "kiro_http_requests_total": "Total HTTP requests handled, labelled by method, code and path class.", + "kiro_ratelimit_rejected_total": "Total inbound requests rejected by the rate limiter.", + "kiro_panics_recovered_total": "Total panics recovered by the recovery middleware.", + "kiro_audit_dropped_total": "Audit entries dropped because the async buffer was full.", + "kiro_token_refresh_total": "Total OAuth token refresh attempts, by result.", +} + +const durationMetric = "kiro_http_request_duration_seconds" + +// gaugeSample is one gauge reading emitted at scrape time by gaugeProvider. +type gaugeSample struct { + name string + labels [][2]string + value float64 +} + +// gaugeProvider, when set, supplies dynamic gauges (e.g. pool state) at render time. +var gaugeProvider func() []gaugeSample + +var gaugeHelp = map[string]string{ + "kiro_accounts_total": "Total configured accounts in the pool.", + "kiro_accounts_available": "Accounts currently routable (not cooling down).", + "kiro_account_inflight": "In-flight dispatches currently assigned to the account.", + "kiro_credential_quota_remaining": "Remaining usage quota per account (UsageLimit-UsageCurrent).", +} + +func escapeLabel(s string) string { + s = strings.ReplaceAll(s, `\`, `\\`) + s = strings.ReplaceAll(s, "\n", `\n`) + s = strings.ReplaceAll(s, `"`, `\"`) + return s +} + +func metricSeries(name string, labels [][2]string) string { + if len(labels) == 0 { + return name + } + var b strings.Builder + b.WriteString(name) + b.WriteByte('{') + for i, kv := range labels { + if i > 0 { + b.WriteByte(',') + } + b.WriteString(kv[0]) + b.WriteString(`="`) + b.WriteString(escapeLabel(kv[1])) + b.WriteByte('"') + } + b.WriteByte('}') + return b.String() +} + +func (m *prodMetrics) incCounter(series string) { + m.mu.Lock() + m.counters[series]++ + m.mu.Unlock() +} + +func (m *prodMetrics) observe(seconds float64) { + m.mu.Lock() + m.hCount++ + m.hSum += seconds + for i, ub := range m.hBuckets { + if seconds <= ub { + m.hCounts[i]++ + } + } + m.mu.Unlock() +} + +// Render returns the current metrics in Prometheus text exposition format. +func (m *prodMetrics) Render() string { + m.mu.Lock() + families := make(map[string][]string) + for series, val := range m.counters { + name := series + if i := strings.IndexByte(series, '{'); i >= 0 { + name = series[:i] + } + families[name] = append(families[name], fmt.Sprintf("%s %d", series, val)) + } + hBuckets := append([]float64(nil), m.hBuckets...) + hCounts := append([]uint64(nil), m.hCounts...) + hSum := m.hSum + hCount := m.hCount + m.mu.Unlock() + + var b strings.Builder + names := make([]string, 0, len(families)) + for name := range families { + names = append(names, name) + } + sort.Strings(names) + for _, name := range names { + help := counterHelp[name] + if help == "" { + help = "Counter." + } + fmt.Fprintf(&b, "# HELP %s %s\n", name, help) + fmt.Fprintf(&b, "# TYPE %s counter\n", name) + lines := families[name] + sort.Strings(lines) + for _, l := range lines { + b.WriteString(l) + b.WriteByte('\n') + } + } + + fmt.Fprintf(&b, "# HELP %s HTTP request duration in seconds.\n", durationMetric) + fmt.Fprintf(&b, "# TYPE %s histogram\n", durationMetric) + for i, ub := range hBuckets { + fmt.Fprintf(&b, "%s_bucket{le=\"%g\"} %d\n", durationMetric, ub, hCounts[i]) + } + fmt.Fprintf(&b, "%s_bucket{le=\"+Inf\"} %d\n", durationMetric, hCount) + fmt.Fprintf(&b, "%s_sum %g\n", durationMetric, hSum) + fmt.Fprintf(&b, "%s_count %d\n", durationMetric, hCount) + + if gaugeProvider != nil { + gm := make(map[string][]string) + for _, s := range gaugeProvider() { + gm[s.name] = append(gm[s.name], fmt.Sprintf("%s %g", metricSeries(s.name, s.labels), s.value)) + } + gnames := make([]string, 0, len(gm)) + for name := range gm { + gnames = append(gnames, name) + } + sort.Strings(gnames) + for _, name := range gnames { + help := gaugeHelp[name] + if help == "" { + help = "Gauge." + } + fmt.Fprintf(&b, "# HELP %s %s\n", name, help) + fmt.Fprintf(&b, "# TYPE %s gauge\n", name) + lines := gm[name] + sort.Strings(lines) + for _, l := range lines { + b.WriteString(l) + b.WriteByte('\n') + } + } + } + + return b.String() +} + +func metricsRequest(method, code, pathClass string) { + globalMetrics.incCounter(metricSeries("kiro_http_requests_total", [][2]string{ + {"method", method}, {"code", code}, {"path", pathClass}, + })) +} +func metricsRateLimited() { globalMetrics.incCounter("kiro_ratelimit_rejected_total") } +func metricsPanic() { globalMetrics.incCounter("kiro_panics_recovered_total") } +func metricsAuditDropped() { globalMetrics.incCounter("kiro_audit_dropped_total") } +func metricsTokenRefresh(result string) { + globalMetrics.incCounter(metricSeries("kiro_token_refresh_total", [][2]string{{"result", result}})) +} + +func recordTokenRefreshResult(err error) { + if err != nil { + metricsTokenRefresh("error") + return + } + metricsTokenRefresh("success") +} +func metricsDuration(d time.Duration) { + if d < 0 { + d = 0 + } + globalMetrics.observe(d.Seconds()) +} diff --git a/proxy/prod_middleware.go b/proxy/prod_middleware.go new file mode 100644 index 00000000..c9f39ad8 --- /dev/null +++ b/proxy/prod_middleware.go @@ -0,0 +1,334 @@ +package proxy + +// Production hardening middleware chain (zero-dependency). +// +// WrapHardening composes, from outermost to innermost: +// +// recover -> request-id -> /metrics router -> security headers +// -> instrument (status + latency + counters + 64 MiB body cap) +// -> rate limit -> handler +// +// The chain wraps the existing *Handler without touching its large ServeHTTP, +// preserving all current routing/behaviour. Streaming (SSE) is preserved because +// statusRecorder implements http.Flusher and Unwrap. Rate limiting and metrics +// are opt-in / always-on respectively and add no dependency. +// +// The async JSONL audit store (persist.go) is wired from main and fed by +// instrumentMiddleware; kiro_audit_dropped_total reflects its buffer overflow. +// Deliberately NOT included (needs deps/infrastructure absent from this repo): +// cluster drain. +import ( + "context" + "crypto/subtle" + "encoding/json" + "io" + "kiro-go/config" + "kiro-go/logger" + "net/http" + "runtime/debug" + "strconv" + "strings" + "time" + + "github.com/google/uuid" +) + +// maxInboundRequestBytes caps every inbound request body so a hostile/oversized +// upload cannot exhaust memory via io.ReadAll on the API endpoints. 64 MiB is +// well above any real Anthropic/OpenAI request (including base64 images). +const maxInboundRequestBytes = 64 << 20 + +type ctxKeyRequestID struct{} + +func requestIDFromContext(ctx context.Context) string { + v, _ := ctx.Value(ctxKeyRequestID{}).(string) + return v +} + +// statusRecorder captures the response status/size while transparently +// forwarding streaming capabilities used by SSE handlers. +type statusRecorder struct { + http.ResponseWriter + status int + wrote bool + bytes int +} + +func (s *statusRecorder) WriteHeader(code int) { + if !s.wrote { + s.status = code + s.wrote = true + } + s.ResponseWriter.WriteHeader(code) +} + +func (s *statusRecorder) Write(b []byte) (int, error) { + if !s.wrote { + s.status = http.StatusOK + s.wrote = true + } + n, err := s.ResponseWriter.Write(b) + s.bytes += n + return n, err +} + +// Flush forwards to the underlying writer so SSE streaming keeps working. +func (s *statusRecorder) Flush() { + if f, ok := s.ResponseWriter.(http.Flusher); ok { + f.Flush() + } +} + +// Unwrap lets http.ResponseController reach the underlying writer (deadlines etc.). +func (s *statusRecorder) Unwrap() http.ResponseWriter { return s.ResponseWriter } + +// WrapHardening wraps inner with the full production middleware chain, reading +// rate limits from the environment. +func WrapHardening(inner http.Handler) http.Handler { + return wrapHardeningWith(inner, newRateLimiterFromEnv()) +} + +func wrapHardeningWith(inner http.Handler, rl *rateLimiter) http.Handler { + h := inner + h = rateLimitMiddleware(h, rl) + h = instrumentMiddleware(h) + h = securityHeadersMiddleware(h) + h = metricsRouter(h) + h = requestIDMiddleware(h) + h = recoverMiddleware(h) + return h +} + +// recoverMiddleware turns an unexpected panic in the handler chain into a clean +// 500 instead of crashing the connection (which leaves an SSE client hanging +// forever). http.ErrAbortHandler is re-panicked to preserve net/http's +// intentional abort semantics. +func recoverMiddleware(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + defer func() { + rec := recover() + if rec == nil { + return + } + if rec == http.ErrAbortHandler { + panic(rec) + } + metricsPanic() + logger.Errorf("[Recover] panic on %s %s (req=%s): %v\n%s", + r.Method, r.URL.Path, requestIDFromContext(r.Context()), rec, debug.Stack()) + func() { + defer func() { _ = recover() }() // guard against write-after-headers panics + w.Header().Set("Content-Type", "application/json; charset=utf-8") + w.WriteHeader(http.StatusInternalServerError) + _, _ = io.WriteString(w, `{"type":"error","error":{"type":"api_error","message":"internal server error"}}`) + }() + }() + next.ServeHTTP(w, r) + }) +} + +// requestIDMiddleware propagates an inbound X-Request-Id or mints a fresh UUID, +// echoes it back on the response and carries it through context for log lines. +func requestIDMiddleware(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + rid := strings.TrimSpace(r.Header.Get("X-Request-Id")) + if rid == "" { + rid = uuid.NewString() + } + w.Header().Set("x-request-id", rid) + ctx := context.WithValue(r.Context(), ctxKeyRequestID{}, rid) + next.ServeHTTP(w, r.WithContext(ctx)) + }) +} + +// metricsRouter serves /metrics (Prometheus text exposition) and /debug/audit-log +// without forwarding to the inner handler. Everything else passes through. +func metricsRouter(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/metrics": + w.Header().Set("Content-Type", "text/plain; version=0.0.4; charset=utf-8") + _, _ = io.WriteString(w, globalMetrics.Render()) + return + case "/debug/audit-log": + serveAuditLog(w, r) + return + } + next.ServeHTTP(w, r) + }) +} + +// adminPasswordOK guards the audit-log endpoint with the admin password +// (header X-Admin-Password or ?password=), compared in constant time so a +// network observer cannot time-attack the password. Empty configured password +// fails closed. +func adminPasswordOK(r *http.Request) bool { + pw := config.GetPassword() + if pw == "" { + return false + } + got := r.Header.Get("X-Admin-Password") + if got == "" { + got = r.URL.Query().Get("password") + } + return subtle.ConstantTimeCompare([]byte(got), []byte(pw)) == 1 +} + +// serveAuditLog exposes recent structured audit entries as JSON for the admin +// UI's error/access-log view. Read-only; requires the admin password. +func serveAuditLog(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json; charset=utf-8") + if !adminPasswordOK(r) { + w.WriteHeader(http.StatusUnauthorized) + _, _ = io.WriteString(w, `{"type":"error","error":{"type":"authentication_error","message":"admin password required"}}`) + return + } + limit := 100 + if v := strings.TrimSpace(r.URL.Query().Get("limit")); v != "" { + if n, err := strconv.Atoi(v); err == nil && n > 0 { + limit = n + } + } + entries, err := AuditRecent(limit) + if err != nil { + w.WriteHeader(http.StatusServiceUnavailable) + _, _ = io.WriteString(w, `{"type":"error","error":{"type":"api_error","message":"audit log unavailable"}}`) + return + } + _ = json.NewEncoder(w).Encode(map[string]interface{}{"count": len(entries), "entries": entries}) +} + +// securityHeadersMiddleware sets baseline browser-security headers (admin panel +// clickjacking / MIME-sniffing protection). +func securityHeadersMiddleware(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + h := w.Header() + h.Set("X-Content-Type-Options", "nosniff") + h.Set("X-Frame-Options", "DENY") + h.Set("Referrer-Policy", "no-referrer") + next.ServeHTTP(w, r) + }) +} + +// instrumentMiddleware records request count, latency histogram, and bounds the +// request body to 64 MiB. The body cap means a subsequent io.ReadAll inside a +// handler returns an error when the limit is exceeded, which each handler +// already maps to a 400 — the goal (memory bound) is met without changing +// per-endpoint error handling. +func instrumentMiddleware(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + rec := &statusRecorder{ResponseWriter: w, status: http.StatusOK} + if r.Body != nil { + r.Body = http.MaxBytesReader(rec, r.Body, maxInboundRequestBytes) + } + start := time.Now() + next.ServeHTTP(rec, r) + lat := time.Since(start) + metricsDuration(lat) + metricsRequest(r.Method, strconv.Itoa(rec.status), pathClass(r.URL.Path)) + // Audit: one structured entry per request (request id, SHA-256 of the + // provided API key, method, path, status, latency, response size). Async + // + non-blocking (recordAudit drops on a full buffer and counts it). The + // API key is hashed, never stored, so the audit log is safe to expose via + // /debug/audit-log to an authenticated admin. + recordAudit(AuditEntry{ + RequestID: requestIDFromContext(r.Context()), + TimeUnixMs: start.UnixMilli(), + APIKeyHash: hashAPIKey(extractProvidedKey(r)), + Method: r.Method, + Path: r.URL.Path, + Status: rec.status, + LatencyMs: lat.Milliseconds(), + Bytes: rec.bytes, + }) + }) +} + +// rateLimitMiddleware enforces the global + per-key token buckets when the +// limiter is enabled. Health/admin/static paths are exempt so the admin UI and +// liveness checks stay reachable even while API traffic is being throttled. +func rateLimitMiddleware(next http.Handler, rl *rateLimiter) http.Handler { + if !rl.enabled() { + return next + } + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method == http.MethodOptions || rateLimitExempt(r.URL.Path) { + next.ServeHTTP(w, r) + return + } + key := extractProvidedKey(r) + if key == "" { + key = "ip:" + clientIP(r) + } + if ok, retry := rl.Allow(key); !ok { + metricsRateLimited() + secs := int(retry.Seconds()) + if secs < 1 { + secs = 1 + } + h := w.Header() + h.Set("Access-Control-Allow-Origin", "*") + h.Set("Retry-After", strconv.Itoa(secs)) + h.Set("Content-Type", "application/json; charset=utf-8") + w.WriteHeader(http.StatusTooManyRequests) + _, _ = io.WriteString(w, `{"type":"error","error":{"type":"rate_limit_error","message":"rate limit exceeded, retry later"}}`) + return + } + next.ServeHTTP(w, r) + }) +} + +// rateLimitExempt reports paths that bypass rate limiting: health/metrics +// endpoints, and the admin UI / static assets. +func rateLimitExempt(path string) bool { + switch { + case path == "/health" || path == "/" || path == "/metrics" || path == "/favicon.ico": + return true + case strings.HasPrefix(path, "/admin"), + strings.HasPrefix(path, "/web"), + strings.HasPrefix(path, "/static"): + return true + } + return false +} + +// clientIP extracts the client address, preferring the first X-Forwarded-For +// entry (set by proxies) and falling back to r.RemoteAddr. +func clientIP(r *http.Request) string { + if xff := r.Header.Get("X-Forwarded-For"); xff != "" { + if i := strings.IndexByte(xff, ','); i >= 0 { + return strings.TrimSpace(xff[:i]) + } + return strings.TrimSpace(xff) + } + host := r.RemoteAddr + if i := strings.LastIndexByte(host, ':'); i >= 0 { + host = host[:i] + } + return host +} + +// pathClass collapses a request path to a stable cardinality for the +// kiro_http_requests_total{path} label, so the metrics series stay bounded. +func pathClass(p string) string { + switch { + case p == "/v1/messages" || p == "/messages" || p == "/anthropic/v1/messages": + return "/v1/messages" + case strings.HasSuffix(p, "/count_tokens"): + return "/v1/messages/count_tokens" + case p == "/v1/chat/completions" || p == "/chat/completions": + return "/v1/chat/completions" + case strings.HasPrefix(p, "/v1/responses"): + return "/v1/responses" + case p == "/v1/models" || p == "/models": + return "/v1/models" + case p == "/health" || p == "/": + return "/health" + case p == "/metrics": + return "/metrics" + case strings.HasPrefix(p, "/admin"): + return "/admin" + default: + return "other" + } +} diff --git a/proxy/ratelimit.go b/proxy/ratelimit.go new file mode 100644 index 00000000..283338a5 --- /dev/null +++ b/proxy/ratelimit.go @@ -0,0 +1,182 @@ +package proxy + +// Inbound request rate limiting (production hardening). +// +// Provides a dependency-free token-bucket limiter with two tiers: +// - a global bucket (all inbound API traffic), and +// - a per-API-key (tenant) bucket. +// +// Both tiers are OFF by default (rate 0 == unlimited) so enabling the limiter +// never silently regresses existing deployments; operators opt in via env: +// +// KIRO_RATE_LIMIT_RPM global requests/minute (0 = unlimited) +// KIRO_RATE_LIMIT_PER_KEY_RPM per-API-key requests/minute (0 = unlimited) +// KIRO_RATE_LIMIT_BURST_SECONDS burst window in seconds (default 10) +// +// Per-credential (single-account) concurrency limiting is enforced at the pool +// layer, not here, because the credential is only chosen deep inside the +// request handler. +// +// Ported verbatim from kiro-tutu (zero-dep). +import ( + "os" + "strconv" + "strings" + "sync" + "time" +) + +// tokenBucket is a mutex-guarded token bucket. A rate <= 0 means "unlimited". +type tokenBucket struct { + mu sync.Mutex + rate float64 // tokens per second + burst float64 // maximum tokens held + tokens float64 + last time.Time +} + +func newTokenBucket(rate, burst float64, now time.Time) *tokenBucket { + if burst < 1 { + burst = 1 + } + return &tokenBucket{rate: rate, burst: burst, tokens: burst, last: now} +} + +// allow reports whether a token is available at now. When denied it also returns +// the estimated duration until the next token becomes available. +func (b *tokenBucket) allow(now time.Time) (bool, time.Duration) { + if b.rate <= 0 { + return true, 0 + } + b.mu.Lock() + defer b.mu.Unlock() + if elapsed := now.Sub(b.last).Seconds(); elapsed > 0 { + b.tokens += elapsed * b.rate + if b.tokens > b.burst { + b.tokens = b.burst + } + b.last = now + } + if b.tokens >= 1 { + b.tokens-- + return true, 0 + } + wait := (1 - b.tokens) / b.rate + return false, time.Duration(wait * float64(time.Second)) +} + +type tokenBucketEntry struct { + bucket *tokenBucket + lastAt time.Time +} + +// rateLimiter enforces a global bucket plus lazily-created per-key buckets. +type rateLimiter struct { + globalRate float64 + perKeyRate float64 + burstSeconds float64 + + global *tokenBucket + + mu sync.Mutex + perKey map[string]*tokenBucketEntry + sweepN int + nowFn func() time.Time +} + +func bucketBurst(rate, seconds float64) float64 { + b := rate * seconds + if b < 1 { + b = 1 + } + return b +} + +// newRateLimiter builds a limiter from explicit requests-per-minute values. +func newRateLimiter(globalRPM, perKeyRPM, burstSeconds float64) *rateLimiter { + if burstSeconds <= 0 { + burstSeconds = 10 + } + rl := &rateLimiter{ + globalRate: globalRPM / 60, + perKeyRate: perKeyRPM / 60, + burstSeconds: burstSeconds, + perKey: make(map[string]*tokenBucketEntry), + nowFn: time.Now, + } + if rl.globalRate > 0 { + rl.global = newTokenBucket(rl.globalRate, bucketBurst(rl.globalRate, burstSeconds), rl.nowFn()) + } + return rl +} + +// newRateLimiterFromEnv reads limits from the KIRO_RATE_LIMIT_* environment. +func newRateLimiterFromEnv() *rateLimiter { + return newRateLimiter( + envFloat("KIRO_RATE_LIMIT_RPM", 0), + envFloat("KIRO_RATE_LIMIT_PER_KEY_RPM", 0), + envFloat("KIRO_RATE_LIMIT_BURST_SECONDS", 10), + ) +} + +func envFloat(name string, def float64) float64 { + if v := strings.TrimSpace(os.Getenv(name)); v != "" { + if f, err := strconv.ParseFloat(v, 64); err == nil && f >= 0 { + return f + } + } + return def +} + +func (rl *rateLimiter) enabled() bool { + return rl != nil && (rl.globalRate > 0 || rl.perKeyRate > 0) +} + +func (rl *rateLimiter) bucketFor(key string, now time.Time) *tokenBucket { + rl.mu.Lock() + defer rl.mu.Unlock() + rl.sweepLocked(now) + e := rl.perKey[key] + if e == nil { + e = &tokenBucketEntry{bucket: newTokenBucket(rl.perKeyRate, bucketBurst(rl.perKeyRate, rl.burstSeconds), now)} + rl.perKey[key] = e + } + e.lastAt = now + return e.bucket +} + +// sweepLocked evicts per-key buckets idle for more than 10 minutes. Called under rl.mu. +func (rl *rateLimiter) sweepLocked(now time.Time) { + rl.sweepN++ + if rl.sweepN < 1024 { + return + } + rl.sweepN = 0 + cutoff := now.Add(-10 * time.Minute) + for k, e := range rl.perKey { + if e.lastAt.Before(cutoff) { + delete(rl.perKey, k) + } + } +} + +// Allow reports whether a request identified by key may proceed. When denied it +// returns the suggested Retry-After duration. key is typically the API key, or a +// client-IP fallback for unauthenticated paths. +func (rl *rateLimiter) Allow(key string) (bool, time.Duration) { + if !rl.enabled() { + return true, 0 + } + now := rl.nowFn() + if rl.perKeyRate > 0 && key != "" { + if ok, retry := rl.bucketFor(key, now).allow(now); !ok { + return false, retry + } + } + if rl.global != nil { + if ok, retry := rl.global.allow(now); !ok { + return false, retry + } + } + return true, 0 +} diff --git a/proxy/responses_handler.go b/proxy/responses_handler.go index a58106f6..d25de8b0 100644 --- a/proxy/responses_handler.go +++ b/proxy/responses_handler.go @@ -5,6 +5,7 @@ import ( "fmt" "io" "kiro-go/config" + "kiro-go/logger" "net/http" "strings" "time" @@ -42,11 +43,16 @@ func (h *Handler) handleOpenAIResponses(w http.ResponseWriter, r *http.Request) } var historyMessages []OpenAIMessage + apiKeyID := apiKeyIDFromContext(r.Context()) if req.PreviousResponseID != "" { - prev, loadErr := loadResponse(req.PreviousResponseID) + prev, loadErr := loadResponseForOwner(req.PreviousResponseID, apiKeyID) if loadErr != nil { - h.sendOpenAIError(w, 404, "invalid_request_error", - fmt.Sprintf("previous_response_id not found: %v", loadErr)) + // Log the real reason server-side for ops, but return a single fixed + // generic message to the client so a bad id never leaks the server's + // filesystem path (the raw os.PathError) nor distinguishes missing vs + // expired vs not-owner. + logger.Warnf("[Responses] previous_response_id %q load failed: %v", req.PreviousResponseID, loadErr) + h.sendOpenAIError(w, 404, "invalid_request_error", "previous_response_id not found") return } historyMessages = expandPreviousResponseHistory(prev) @@ -97,7 +103,7 @@ func (h *Handler) handleOpenAIResponses(w http.ResponseWriter, r *http.Request) Tools: req.Tools, } if req.Temperature != nil { - openaiReq.Temperature = *req.Temperature + openaiReq.Temperature = req.Temperature } if req.MaxOutputTokens != nil { openaiReq.MaxTokens = *req.MaxOutputTokens @@ -110,7 +116,6 @@ func (h *Handler) handleOpenAIResponses(w http.ResponseWriter, r *http.Request) estimatedInputTokens := estimateOpenAIRequestInputTokens(openaiReq) kiroPayload := OpenAIToKiro(openaiReq, thinking) - apiKeyID := apiKeyIDFromContext(r.Context()) respID := generateResponseID() if req.Stream { @@ -132,7 +137,11 @@ func (h *Handler) handleResponsesNonStream( var lastErr error reqStart := time.Now() - for attempt := 0; attempt < maxAccountRetryAttempts; attempt++ { + retryBudget := resolveAccountRetryBudget(h.pool.Count()) + for attempt := 0; attempt < retryBudget; attempt++ { + if attempt > 0 { + time.Sleep(accountRetryBackoff(attempt - 1)) + } account := h.pool.GetNextForModelExcluding(model, excluded) if account == nil { break @@ -171,6 +180,9 @@ func (h *Handler) handleResponsesNonStream( lastErr = err excluded[account.ID] = true h.handleAccountFailure(account, err) + if isUpstreamPermanentError(err) { + break + } continue } @@ -194,6 +206,7 @@ func (h *Handler) handleResponsesNonStream( respObj := buildResponsesObject(respID, model, finalContent, toolUses, inputTokens, outputTokens, req) respObj.StoredInput = storedInput respObj.Instructions = req.Instructions + respObj.OwnerKeyID = apiKeyID if storeResponse { if saveErr := saveResponse(respObj); saveErr != nil { @@ -211,7 +224,7 @@ func (h *Handler) handleResponsesNonStream( return } h.recordFailureWithDetails("responses", model, "", lastErr) - h.sendOpenAIError(w, 500, "server_error", lastErr.Error()) + h.sendOpenAIError(w, 500, "server_error", improperlyFormedClientMessage(lastErr)) } func buildResponsesObject( @@ -317,7 +330,11 @@ func (h *Handler) handleResponsesStream( responseStarted := false reqStart := time.Now() - for attempt := 0; attempt < maxAccountRetryAttempts; attempt++ { + retryBudget := resolveAccountRetryBudget(h.pool.Count()) + for attempt := 0; attempt < retryBudget; attempt++ { + if attempt > 0 { + time.Sleep(accountRetryBackoff(attempt - 1)) + } account := h.pool.GetNextForModelExcluding(model, excluded) if account == nil { break @@ -472,10 +489,13 @@ func (h *Handler) handleResponsesStream( err := CallKiroAPI(account, payload, callback) if err != nil { + lastErr = err + excluded[account.ID] = true + h.handleAccountFailure(account, err) + if isUpstreamPermanentError(err) { + break + } if !responseStarted { - lastErr = err - excluded[account.ID] = true - h.handleAccountFailure(account, err) continue } send("response.failed", map[string]interface{}{ @@ -485,7 +505,7 @@ func (h *Handler) handleResponsesStream( "status": "failed", "error": map[string]string{ "type": "server_error", - "message": err.Error(), + "message": improperlyFormedClientMessage(err), }, }, }) @@ -542,6 +562,7 @@ func (h *Handler) handleResponsesStream( respObj.CreatedAt = createdAt respObj.StoredInput = storedInput respObj.Instructions = req.Instructions + respObj.OwnerKeyID = apiKeyID if storeResponse { if saveErr := saveResponse(respObj); saveErr != nil { @@ -580,7 +601,7 @@ func (h *Handler) handleResponsesStream( "status": "failed", "error": map[string]string{ "type": "server_error", - "message": lastErr.Error(), + "message": improperlyFormedClientMessage(lastErr), }, }, }) diff --git a/proxy/responses_handler_test.go b/proxy/responses_handler_test.go index 72003eb9..bba16935 100644 --- a/proxy/responses_handler_test.go +++ b/proxy/responses_handler_test.go @@ -257,6 +257,44 @@ func TestResponsesContinuationKeepsNewInstructions(t *testing.T) { } } +// TestResponsesPreviousResponseIDNotFoundIsGeneric verifies that a bad +// previous_response_id does NOT leak the server's filesystem path (the raw +// os.PathError from os.ReadFile, e.g. "open C:\...\responses\resp_xxx.json: ...") +// and does not distinguish missing vs expired vs not-owner — all must surface +// the same fixed generic message to the client. +func TestResponsesPreviousResponseIDNotFoundIsGeneric(t *testing.T) { + h, cleanup := setupResponsesTestHandler(t) + defer cleanup() + + body := strings.NewReader(`{"model":"claude-sonnet-4.5","input":"hi","previous_response_id":"resp_does_not_exist_xyz","store":false}`) + req := httptest.NewRequest(http.MethodPost, "/v1/responses", body) + rec := httptest.NewRecorder() + h.handleOpenAIResponses(rec, req) + + if rec.Code != http.StatusNotFound { + t.Fatalf("expected 404 for a bad previous_response_id, got %d body=%s", rec.Code, rec.Body.String()) + } + + var errResp struct { + Error struct { + Message string `json:"message"` + } `json:"error"` + } + if err := json.Unmarshal(rec.Body.Bytes(), &errResp); err != nil { + t.Fatalf("decode error body: %v body=%s", err, rec.Body.String()) + } + msg := errResp.Error.Message + + // The client-facing message must be a fixed generic string — no filesystem + // path, no file extension, no OS error text, no per-case detail. ".json" is + // the cross-platform marker that the raw os.PathError was surfaced verbatim. + for _, leak := range []string{".json", "no such file", `responses\`, "responses/", "stored response expired", "stored response not found"} { + if strings.Contains(msg, leak) { + t.Fatalf("previous_response_id error leaks server detail (%q): %q", leak, msg) + } + } +} + func setupResponsesTestHandler(t *testing.T) (*Handler, func()) { t.Helper() cfgFile := filepath.Join(t.TempDir(), "config.json") diff --git a/proxy/responses_store.go b/proxy/responses_store.go index 94cee166..bb111eec 100644 --- a/proxy/responses_store.go +++ b/proxy/responses_store.go @@ -62,6 +62,7 @@ func saveResponse(resp *ResponsesObject) error { Metadata: resp.Metadata, Instructions: resp.Instructions, StoredInput: resp.StoredInput, + OwnerKeyID: resp.OwnerKeyID, StoredAt: resp.StoredAt, } @@ -110,10 +111,28 @@ func loadResponse(id string) (*ResponsesObject, error) { Metadata: doc.Metadata, Instructions: doc.Instructions, StoredInput: doc.StoredInput, + OwnerKeyID: doc.OwnerKeyID, StoredAt: doc.StoredAt, }, nil } +// loadResponseForOwner loads a stored response and verifies ownership. +// If the stored response has an OwnerKeyID set and it doesn't match the +// given apiKeyID, a generic "not found" error is returned (identical to +// a missing ID) to prevent response-ID enumeration. Responses created +// before the ownership field was added (OwnerKeyID=="") are accessible +// by anyone for backward compatibility. +func loadResponseForOwner(id, apiKeyID string) (*ResponsesObject, error) { + resp, err := loadResponse(id) + if err != nil { + return nil, err + } + if resp.OwnerKeyID != "" && apiKeyID != "" && resp.OwnerKeyID != apiKeyID { + return nil, fmt.Errorf("stored response not found") + } + return resp, nil +} + func purgeExpiredResponses(ttl time.Duration) { if ttl <= 0 { ttl = responsesDefaultTTL @@ -178,5 +197,6 @@ type storedResponseDoc struct { Metadata map[string]string `json:"metadata,omitempty"` Instructions string `json:"instructions,omitempty"` StoredInput json.RawMessage `json:"stored_input,omitempty"` + OwnerKeyID string `json:"owner_key_id,omitempty"` StoredAt int64 `json:"stored_at"` } diff --git a/proxy/responses_store_test.go b/proxy/responses_store_test.go new file mode 100644 index 00000000..084f559a --- /dev/null +++ b/proxy/responses_store_test.go @@ -0,0 +1,106 @@ +package proxy + +import ( + "encoding/json" + "kiro-go/config" + "path/filepath" + "strings" + "testing" + "time" +) + +func TestLoadResponseForOwnerIsolation(t *testing.T) { + // Set up a temp config dir so responses are written to a temp directory. + cfgFile := filepath.Join(t.TempDir(), "config.json") + if err := config.Init(cfgFile); err != nil { + t.Fatalf("config.Init: %v", err) + } + + // Helper to create and save a minimal response with the given ID and owner. + makeResponse := func(id, ownerKeyID string) *ResponsesObject { + return &ResponsesObject{ + ID: id, + Object: "response", + CreatedAt: time.Now().Unix(), + Status: "completed", + Model: "claude-sonnet-4.5", + Output: []ResponseOutputItem{{ + ID: "msg_test", + Type: "message", + Role: "assistant", + Content: []ResponseContentPart{{ + Type: "output_text", + Text: "test output", + }}, + }}, + StoredInput: json.RawMessage(`"test input"`), + OwnerKeyID: ownerKeyID, + StoredAt: time.Now().Unix(), + } + } + + t.Run("access allowed same key", func(t *testing.T) { + resp := makeResponse("resp_owner_allowed", "key-A") + if err := saveResponse(resp); err != nil { + t.Fatalf("save: %v", err) + } + + loaded, err := loadResponseForOwner("resp_owner_allowed", "key-A") + if err != nil { + t.Fatalf("expected access allowed for same key, got error: %v", err) + } + if loaded.ID != "resp_owner_allowed" { + t.Fatalf("loaded ID mismatch: got %q", loaded.ID) + } + if loaded.OwnerKeyID != "key-A" { + t.Fatalf("loaded OwnerKeyID mismatch: got %q", loaded.OwnerKeyID) + } + }) + + t.Run("access denied different key", func(t *testing.T) { + resp := makeResponse("resp_owner_denied", "key-A") + if err := saveResponse(resp); err != nil { + t.Fatalf("save: %v", err) + } + + _, err := loadResponseForOwner("resp_owner_denied", "key-B") + if err == nil { + t.Fatalf("expected error for different key, got nil") + } + // Error message is intentionally generic ("not found") to prevent + // response-ID enumeration — do not assert specific wording. + if !strings.Contains(err.Error(), "not found") { + t.Fatalf("expected 'not found' in error, got: %v", err) + } + }) + + t.Run("legacy compatibility empty owner", func(t *testing.T) { + resp := makeResponse("resp_legacy_no_owner", "") + if err := saveResponse(resp); err != nil { + t.Fatalf("save: %v", err) + } + + loaded, err := loadResponseForOwner("resp_legacy_no_owner", "key-B") + if err != nil { + t.Fatalf("expected legacy response accessible by any key, got error: %v", err) + } + if loaded.ID != "resp_legacy_no_owner" { + t.Fatalf("loaded ID mismatch: got %q", loaded.ID) + } + }) + + t.Run("empty caller key allowed", func(t *testing.T) { + resp := makeResponse("resp_empty_caller", "key-A") + if err := saveResponse(resp); err != nil { + t.Fatalf("save: %v", err) + } + + loaded, err := loadResponseForOwner("resp_empty_caller", "") + if err != nil { + t.Fatalf("expected unauthenticated access allowed, got error: %v", err) + } + if loaded.ID != "resp_empty_caller" { + t.Fatalf("loaded ID mismatch: got %q", loaded.ID) + } + }) +} diff --git a/proxy/responses_types.go b/proxy/responses_types.go index c36413a0..70a7f4fa 100644 --- a/proxy/responses_types.go +++ b/proxy/responses_types.go @@ -30,6 +30,7 @@ type ResponsesObject struct { Instructions string `json:"instructions,omitempty"` StoredInput json.RawMessage `json:"-"` StoredInstr string `json:"-"` + OwnerKeyID string `json:"-"` // API key that created this response (not exposed in API) StoredAt int64 `json:"stored_at,omitempty"` } diff --git a/proxy/token_refresh.go b/proxy/token_refresh.go new file mode 100644 index 00000000..f70f74b1 --- /dev/null +++ b/proxy/token_refresh.go @@ -0,0 +1,18 @@ +package proxy + +import ( + "kiro-go/auth" + "kiro-go/config" +) + +// authRefreshToken wraps auth.RefreshToken so every refresh attempt — on the +// request path, in the background refresher, and in the warm-up worker — is +// counted in the kiro_token_refresh_total metric from a single choke point. +// +// Ported from kiro-tutu. Behavior is identical to auth.RefreshToken; only the +// metric counter is added. +func authRefreshToken(account *config.Account) (string, string, int64, string, error) { + accessToken, refreshToken, expiresAt, profileArn, err := auth.RefreshToken(account) + recordTokenRefreshResult(err) + return accessToken, refreshToken, expiresAt, profileArn, err +} diff --git a/proxy/tool_compression.go b/proxy/tool_compression.go new file mode 100644 index 00000000..bd80e2e9 --- /dev/null +++ b/proxy/tool_compression.go @@ -0,0 +1,214 @@ +package proxy + +import ( + "encoding/json" + "os" + "strconv" + "strings" + "unicode/utf8" + + "kiro-go/logger" +) + +// 工具定义体积压缩 +// +// 当一次请求携带的工具定义（name + description + inputSchema）序列化后的总字节数 +// 超过阈值时，对齐 kiro-rs 的 compress_tools_if_needed 做两步渐进压缩，避免把超大的 +// tools 数组原样上送、被 Kiro 上游以 "Improperly formed request." 拒绝： +// 1. 递归简化每个工具的 inputSchema：仅保留结构骨架（type / required / properties 的 +// key 与其 type），剥掉 description / examples / enum 等纯说明性字段（这些往往是体 +// 积大头，但对模型选参的结构理解非必需）。 +// 2. 若简化 schema 后仍超阈值，按超出比例截断每个工具的 description（UTF-8 安全，至少 +// 保留 minToolDescChars 个字符，以免描述被砍到无法辨识工具用途）。 +// +// 注意：本压缩是“总量兜底”，与既有的 per-tool 上限（maxToolDescLen）和非法字段清理 +// （cleanSchema）互补，不替代它们——前两者在单个工具维度先行处理，本函数只在所有工具 +// 加起来仍然过大时才介入。 +// +// 与 upstream_error.go 的 isImproperlyFormedRejection 互补：那边在被上游拒绝后做检测、 +// 永久短路本请求；这里在发送前做预防、尽量不让请求被拒。 + +// defaultToolsSizeThreshold 是触发工具压缩的总字节阈值。20KB 对齐 kiro-rs 的 +// TOOL_SIZE_THRESHOLD 经验值；Kiro 上游真实红线未公开，可经 env 按实测调整。 +const defaultToolsSizeThreshold = 20 * 1024 + +// minToolDescChars 是 description 截断后至少保留的字符数（rune 计，UTF-8 安全）， +// 保证压缩后描述仍可辨识工具用途。对齐 kiro-rs 的 MIN_DESCRIPTION_CHARS。 +const minToolDescChars = 50 + +// resolveToolsSizeThreshold 读取 env 覆盖（KIRO_TOOLS_COMPRESS_THRESHOLD_BYTES）。 +// 非法值回落默认；0 表示禁用压缩（退回原样上送，由上游与既有 per-tool 上限兜底）。 +func resolveToolsSizeThreshold() int { + raw := strings.TrimSpace(os.Getenv("KIRO_TOOLS_COMPRESS_THRESHOLD_BYTES")) + if raw == "" { + return defaultToolsSizeThreshold + } + n, err := strconv.Atoi(raw) + if err != nil || n < 0 { + return defaultToolsSizeThreshold + } + return n +} + +// compressToolsIfNeeded 在工具定义总体积超过阈值时执行两步压缩，否则原样返回。 +// 入参 tools 来自 convertClaudeTools / convertOpenAITools 的产出（已做过 per-tool +// 截断与 schema 清理），本函数只负责总量收敛。 +func compressToolsIfNeeded(tools []KiroToolWrapper) []KiroToolWrapper { + threshold := resolveToolsSizeThreshold() + if threshold <= 0 || len(tools) == 0 { + return tools + } + + total := estimateToolsBytes(tools) + if total <= threshold { + return tools + } + + // 第一步：递归简化每个工具的 inputSchema。 + for i := range tools { + tools[i].ToolSpecification.InputSchema.JSON = simplifyToolSchema(tools[i].ToolSpecification.InputSchema.JSON) + } + + afterSchema := estimateToolsBytes(tools) + if afterSchema <= threshold { + logger.Infof("[ToolCompress] %d tools compressed by schema simplification: %d -> %d bytes (threshold %d)", + len(tools), total, afterSchema, threshold) + return tools + } + + // 第二步：按比例截断 description（基于字节比例，UTF-8 安全，保底 minToolDescChars）。 + ratio := float64(threshold) / float64(afterSchema) + for i := range tools { + desc := tools[i].ToolSpecification.Description + tools[i].ToolSpecification.Description = truncateDescByRatio(desc, ratio) + } + + final := estimateToolsBytes(tools) + logger.Infof("[ToolCompress] %d tools compressed (schema + description): %d -> %d -> %d bytes (threshold %d)", + len(tools), total, afterSchema, final, threshold) + return tools +} + +// estimateToolsBytes 估算工具列表序列化后的总字节数。直接对整个 slice 做 json.Marshal +// 取长度，口径与上送前 json.Marshal(payload) 内 tools 部分一致——把 toolSpecification / +// name / description / inputSchema 等 JSON key、引号、括号、逗号等结构开销也算进去， +// 避免只累加字段长度导致的系统性低估（低估会让临界请求“该压不压”仍被上游拒）。 +// Marshal 失败（理论上不会，KiroToolWrapper 全是可序列化字段）时回退到字段长度累加。 +func estimateToolsBytes(tools []KiroToolWrapper) int { + if raw, err := json.Marshal(tools); err == nil { + return len(raw) + } + total := 0 + for i := range tools { + spec := tools[i].ToolSpecification + total += len(spec.Name) + len(spec.Description) + if raw, err := json.Marshal(spec.InputSchema.JSON); err == nil { + total += len(raw) + } + } + return total +} + +// truncateDescByRatio 按 ratio（字节比例）截断 desc，UTF-8 安全，至少保留 +// minToolDescChars 个字符。ratio>=1 时不截断。 +// +// ratio = threshold / afterSchemaBytes 是按字节算出的，故 target 也按字节算（而非乘 +// rune 数），口径自洽——否则对 CJK 这类多字节文本会把字节比误当字符比，导致过度截断。 +// 取到目标字节数后回退到不超过它的最近 UTF-8 字符边界，保证不切坏多字节字符。 +func truncateDescByRatio(desc string, ratio float64) string { + if ratio >= 1 { + return desc + } + if len([]rune(desc)) <= minToolDescChars { + return desc + } + + targetBytes := int(float64(len(desc)) * ratio) + + // 保底：至少保留 minToolDescChars 个字符对应的字节数。 + minBytes := len(desc) + if r := []rune(desc); len(r) > minToolDescChars { + minBytes = len(string(r[:minToolDescChars])) + } + if targetBytes < minBytes { + targetBytes = minBytes + } + if targetBytes >= len(desc) { + return desc + } + + // 回退到 <= targetBytes 的最近字符边界（避免切坏多字节 rune）。 + cut := targetBytes + for cut > 0 && !utf8.RuneStart(desc[cut]) { + cut-- + } + return desc[:cut] +} + +// simplifyToolSchema 递归简化 JSON Schema：保留结构与约束骨架，剥掉纯说明性字段。 +// - 保留：type、required、enum、$ref、anyOf/oneOf/allOf、properties、items（结构与 +// 选参约束相关，去掉会改变上游/模型对参数的理解）。 +// - 移除：description、examples、default、title、$comment 等纯说明性字段（体积大头， +// 对模型选参非必需）。 +// +// 仅处理 map[string]interface{} 形态的 schema（ensureObjectSchema 已保证顶层为该形态）； +// 其他形态原样返回。注意：本函数返回新 map，不修改入参。 +func simplifyToolSchema(schema interface{}) interface{} { + m, ok := schema.(map[string]interface{}) + if !ok { + return schema + } + + result := make(map[string]interface{}) + + // 保留顶层结构 / 约束字段（剔除 additionalProperties——cleanSchema 已要求移除它）。 + // $ref / anyOf / oneOf / allOf 在缺 type 时是该节点唯一的语义来源，必须保留，否则 + // 节点会塌成 {} 让模型与上游都无法理解。 + copySchemaKeptKeys(m, result) + + // properties：递归简化每个属性。 + if props, ok := m["properties"].(map[string]interface{}); ok { + simplified := make(map[string]interface{}, len(props)) + for name, prop := range props { + simplified[name] = simplifyToolSchema(prop) + } + result["properties"] = simplified + } + + // items（数组元素 schema）：递归简化（可能是单个 schema 或 schema 数组）。 + if items, exists := m["items"]; exists { + result["items"] = simplifyToolSchema(items) + } + + return result +} + +// schemaKeptKeys 是简化时保留的非递归字段：结构（type/required）、选参约束（enum）、 +// 组合/引用（$ref/anyOf/oneOf/allOf）、以及 $schema。description/examples/default/title +// 等说明性字段不在此列，会被剥除。 +var schemaKeptKeys = []string{"$schema", "type", "required", "enum", "$ref", "anyOf", "oneOf", "allOf"} + +// copySchemaKeptKeys 把 src 中 schemaKeptKeys 列出的字段拷到 dst。对 anyOf/oneOf/allOf +// 这类「子 schema 数组」会递归简化其中每个元素。 +func copySchemaKeptKeys(src, dst map[string]interface{}) { + for _, key := range schemaKeptKeys { + v, exists := src[key] + if !exists { + continue + } + switch key { + case "anyOf", "oneOf", "allOf": + if arr, ok := v.([]interface{}); ok { + simplifiedArr := make([]interface{}, len(arr)) + for i, sub := range arr { + simplifiedArr[i] = simplifyToolSchema(sub) + } + dst[key] = simplifiedArr + continue + } + dst[key] = v + default: + dst[key] = v + } + } +} diff --git a/proxy/tool_compression_test.go b/proxy/tool_compression_test.go new file mode 100644 index 00000000..29e1ccc8 --- /dev/null +++ b/proxy/tool_compression_test.go @@ -0,0 +1,301 @@ +package proxy + +import ( + "encoding/json" + "strings" + "testing" +) + +// makeKiroTool 构造一个工具 wrapper，desc 为给定描述，schema 为给定 JSON 对象。 +func makeKiroTool(name, desc string, schema map[string]interface{}) KiroToolWrapper { + w := KiroToolWrapper{} + w.ToolSpecification.Name = name + w.ToolSpecification.Description = desc + w.ToolSpecification.InputSchema = InputSchema{JSON: schema} + return w +} + +// schemaWithDescriptions 造一个靠 description/examples 撑大的 schema，用于验证简化能 +// 剥掉这些说明性字段、显著减小体积。刻意不带 enum——enum 属于“保留”字段（见 +// TestSimplifyToolSchemaPreservesConstraints），混进来会让本类“简化即达标”的体积断言 +// 受 enum 保留量干扰。 +func schemaWithDescriptions(propCount int, descPerProp string) map[string]interface{} { + props := make(map[string]interface{}, propCount) + for i := 0; i < propCount; i++ { + key := "field" + strings.Repeat("x", 3) + string(rune('a'+i%26)) + string(rune('0'+i%10)) + props[key] = map[string]interface{}{ + "type": "string", + "description": descPerProp, + "examples": []interface{}{descPerProp, descPerProp}, + "title": descPerProp, + } + } + return map[string]interface{}{ + "type": "object", + "required": []interface{}{"fieldxxxa0"}, + "properties": props, + } +} + +func TestCompressToolsNoOpUnderThreshold(t *testing.T) { + t.Setenv("KIRO_TOOLS_COMPRESS_THRESHOLD_BYTES", "") + tools := []KiroToolWrapper{ + makeKiroTool("smallTool", "a short description", map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{"x": map[string]interface{}{"type": "string"}}, + }), + } + before, _ := json.Marshal(tools) + out := compressToolsIfNeeded(tools) + after, _ := json.Marshal(out) + if string(before) != string(after) { + t.Fatalf("expected no-op under threshold\nbefore=%s\nafter=%s", before, after) + } +} + +func TestCompressToolsSchemaSimplificationBringsUnderThreshold(t *testing.T) { + // 阈值设小，让“简化 schema”这一步就足以达标，验证不会进入 description 截断。 + t.Setenv("KIRO_TOOLS_COMPRESS_THRESHOLD_BYTES", "2048") + desc := "keep me readable" // 短描述，简化 schema 后总量应已达标 + tools := []KiroToolWrapper{ + makeKiroTool("bigSchemaTool", desc, schemaWithDescriptions(40, strings.Repeat("verbose ", 20))), + } + if estimateToolsBytes(tools) <= 2048 { + t.Fatalf("precondition: tools should exceed threshold before compression") + } + + out := compressToolsIfNeeded(tools) + + if got := estimateToolsBytes(out); got > 2048 { + t.Fatalf("expected size <= threshold after schema simplification, got %d", got) + } + // description 未被截断（schema 简化已足够）。 + if out[0].ToolSpecification.Description != desc { + t.Fatalf("description should be intact when schema simplification suffices, got %q", out[0].ToolSpecification.Description) + } + // 简化后：结构骨架保留，说明性字段剥离。 + schema := out[0].ToolSpecification.InputSchema.JSON.(map[string]interface{}) + if schema["type"] != "object" { + t.Fatalf("top-level type must be preserved") + } + if _, ok := schema["required"]; !ok { + t.Fatalf("required must be preserved") + } + props := schema["properties"].(map[string]interface{}) + for name, p := range props { + pm := p.(map[string]interface{}) + if pm["type"] != "string" { + t.Fatalf("prop %s type must be preserved", name) + } + if _, ok := pm["description"]; ok { + t.Fatalf("prop %s description must be stripped", name) + } + if _, ok := pm["enum"]; ok { + t.Fatalf("prop %s enum must be stripped", name) + } + if _, ok := pm["examples"]; ok { + t.Fatalf("prop %s examples must be stripped", name) + } + } +} + +func TestCompressToolsTruncatesDescriptionWhenSchemaNotEnough(t *testing.T) { + // 阈值极小，简化 schema 仍不够 → 必须截断 description；验证保底字符数与 UTF-8 安全。 + t.Setenv("KIRO_TOOLS_COMPRESS_THRESHOLD_BYTES", "200") + longDesc := strings.Repeat("用途说明", 200) // 多字节字符，验证不切坏 rune + tools := []KiroToolWrapper{ + makeKiroTool("toolWithLongDesc", longDesc, schemaWithDescriptions(10, "x")), + } + + out := compressToolsIfNeeded(tools) + + gotDesc := out[0].ToolSpecification.Description + if len([]rune(gotDesc)) >= len([]rune(longDesc)) { + t.Fatalf("description should be truncated") + } + if len([]rune(gotDesc)) < minToolDescChars { + t.Fatalf("description must keep at least %d chars, got %d", minToolDescChars, len([]rune(gotDesc))) + } + // UTF-8 安全：截断结果仍是合法字符串（rune 往返一致）。 + if string([]rune(gotDesc)) != gotDesc { + t.Fatalf("truncation broke UTF-8 boundary") + } +} + +func TestCompressToolsDisabledByZeroThreshold(t *testing.T) { + t.Setenv("KIRO_TOOLS_COMPRESS_THRESHOLD_BYTES", "0") + tools := []KiroToolWrapper{ + makeKiroTool("bigTool", "d", schemaWithDescriptions(40, strings.Repeat("verbose ", 20))), + } + before, _ := json.Marshal(tools) + out := compressToolsIfNeeded(tools) + after, _ := json.Marshal(out) + if string(before) != string(after) { + t.Fatalf("threshold=0 must disable compression (pass-through)") + } +} + +func TestResolveToolsSizeThreshold(t *testing.T) { + cases := []struct { + name string + env string + want int + }{ + {"empty falls back to default", "", defaultToolsSizeThreshold}, + {"explicit zero disables", "0", 0}, + {"valid override", "4096", 4096}, + {"negative falls back", "-1", defaultToolsSizeThreshold}, + {"garbage falls back", "abc", defaultToolsSizeThreshold}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + t.Setenv("KIRO_TOOLS_COMPRESS_THRESHOLD_BYTES", tc.env) + if got := resolveToolsSizeThreshold(); got != tc.want { + t.Fatalf("env=%q want %d got %d", tc.env, tc.want, got) + } + }) + } +} + +// TestCompressToolsEndToEndViaConvertClaude 验证压缩已接到 Claude 工具转换路径： +// 构造少量但 schema 臃肿的工具，经 convertClaudeTools 后体积应被显著压缩。 +// 注意：对齐 kiro-rs，压缩是“尽力而为”——简化 schema + 截断 description 两步后即返回， +// 不保证一定降到阈值内（工具数量本身极多时可能仍超），故这里断言“显著变小”而非“必达阈值”。 +func TestCompressToolsEndToEndViaConvertClaude(t *testing.T) { + t.Setenv("KIRO_TOOLS_COMPRESS_THRESHOLD_BYTES", "2048") + claudeTools := make([]ClaudeTool, 0, 4) + for i := 0; i < 4; i++ { + claudeTools = append(claudeTools, ClaudeTool{ + Name: "tool" + string(rune('A'+i%26)) + string(rune('0'+i%10)), + Description: "a tool", + InputSchema: schemaWithDescriptions(8, strings.Repeat("verbose ", 16)), + }) + } + + // 不经压缩时的体积（关掉压缩量一次基线）。 + t.Setenv("KIRO_TOOLS_COMPRESS_THRESHOLD_BYTES", "0") + rawWrappers, _ := convertClaudeTools(claudeTools) + rawBytes := estimateToolsBytes(rawWrappers) + + // 开启压缩。 + t.Setenv("KIRO_TOOLS_COMPRESS_THRESHOLD_BYTES", "2048") + wrappers, _ := convertClaudeTools(claudeTools) + got := estimateToolsBytes(wrappers) + + if got > 2048 { + t.Fatalf("少量臃肿工具应被压到阈值内，got %d bytes", got) + } + if got >= rawBytes { + t.Fatalf("压缩后应显著小于原始体积：raw=%d compressed=%d", rawBytes, got) + } +} + +// TestSimplifyToolSchemaPreservesConstraints 锁定 code review 发现的高危回归：简化 +// schema 时必须保留 type/required/enum 及 $ref/anyOf/oneOf/allOf，且嵌套 object 的 +// required 不能丢——否则“压缩”反而把原本能过的请求改成上游拒收。 +func TestSimplifyToolSchemaPreservesConstraints(t *testing.T) { + schema := map[string]interface{}{ + "type": "object", + "required": []interface{}{"mode", "addr"}, + "properties": map[string]interface{}{ + // 枚举约束必须保留（模型据此选值）。 + "mode": map[string]interface{}{ + "type": "string", + "enum": []interface{}{"r", "w"}, + "description": "should be stripped", + }, + // 嵌套 object：其 required 与子属性 type 必须保留。 + "addr": map[string]interface{}{ + "type": "object", + "required": []interface{}{"host"}, + "properties": map[string]interface{}{ + "host": map[string]interface{}{"type": "string", "description": "drop me"}, + "port": map[string]interface{}{"type": "integer"}, + }, + }, + // 只有 $ref、无 type：不能被压成 {}，$ref 必须保留。 + "ref": map[string]interface{}{ + "$ref": "#/$defs/Foo", + "description": "drop me", + }, + // anyOf 子 schema 需递归简化但保留结构。 + "choice": map[string]interface{}{ + "anyOf": []interface{}{ + map[string]interface{}{"type": "string", "description": "drop"}, + map[string]interface{}{"type": "null"}, + }, + }, + }, + } + + out := simplifyToolSchema(schema).(map[string]interface{}) + + // 顶层 required 保留。 + if req, _ := out["required"].([]interface{}); len(req) != 2 { + t.Fatalf("top-level required must be preserved, got %v", out["required"]) + } + props := out["properties"].(map[string]interface{}) + + // enum 保留、description 剥离。 + mode := props["mode"].(map[string]interface{}) + if _, ok := mode["enum"]; !ok { + t.Fatal("enum must be preserved (model needs it to pick valid values)") + } + if _, ok := mode["description"]; ok { + t.Fatal("description must be stripped") + } + + // 嵌套 object 的 required 保留。 + addr := props["addr"].(map[string]interface{}) + if req, _ := addr["required"].([]interface{}); len(req) != 1 { + t.Fatalf("nested object required must be preserved, got %v", addr["required"]) + } + addrProps := addr["properties"].(map[string]interface{}) + if addrProps["host"].(map[string]interface{})["type"] != "string" { + t.Fatal("nested property type must be preserved") + } + if _, ok := addrProps["host"].(map[string]interface{})["description"]; ok { + t.Fatal("nested property description must be stripped") + } + + // 只有 $ref 的节点不能塌成空对象。 + ref := props["ref"].(map[string]interface{}) + if ref["$ref"] != "#/$defs/Foo" { + t.Fatalf("$ref must be preserved to avoid empty {} node, got %v", ref) + } + if len(ref) == 0 { + t.Fatal("ref node collapsed to empty object — model/upstream cannot interpret it") + } + + // anyOf 保留且子 schema 简化。 + choice := props["choice"].(map[string]interface{}) + anyOf, ok := choice["anyOf"].([]interface{}) + if !ok || len(anyOf) != 2 { + t.Fatalf("anyOf must be preserved with its subschemas, got %v", choice) + } + if anyOf[0].(map[string]interface{})["type"] != "string" { + t.Fatal("anyOf subschema type must survive simplification") + } + if _, ok := anyOf[0].(map[string]interface{})["description"]; ok { + t.Fatal("anyOf subschema description must be stripped") + } +} + +// TestTruncateDescByRatioCJKSafe 验证 CJK 描述按字节比例截断、不切坏 rune、保底字符数。 +func TestTruncateDescByRatioCJKSafe(t *testing.T) { + desc := strings.Repeat("中文说明", 100) // 1200 bytes (每字 3B), 400 runes + got := truncateDescByRatio(desc, 0.25) + + // 结果合法 UTF-8、未切坏 rune。 + if string([]rune(got)) != got { + t.Fatal("truncation broke a multi-byte rune boundary") + } + // 比原文短。 + if len(got) >= len(desc) { + t.Fatalf("expected truncation, got %d >= %d bytes", len(got), len(desc)) + } + // 保底字符数。 + if len([]rune(got)) < minToolDescChars { + t.Fatalf("must keep at least %d chars, got %d", minToolDescChars, len([]rune(got))) + } +} diff --git a/proxy/translator.go b/proxy/translator.go index f29e2d13..f3b33b7c 100644 --- a/proxy/translator.go +++ b/proxy/translator.go @@ -46,15 +46,15 @@ const minimalFallbackUserContent = "." const toolResultsContinuationPrefix = "Tool results:" const toolResultImagePlaceholder = "[Tool returned an image; the image is attached to this message.]" -// maxPayloadBytes is the upper bound for the serialized Kiro request body. -// Kiro's upstream rejects oversized requests with HTTP 400 -// "Input is too long." (CONTENT_LENGTH_EXCEEDS_THRESHOLD). When a converted -// payload exceeds this size we drop the oldest history turns (keeping the -// system priming, the most recent turns, the active tool turn, and the current -// message) and insert a placeholder note so the model knows context was elided. -// The limit is kept conservatively below the observed upstream threshold to -// leave room for headers and minor serialization overhead. -const maxPayloadBytes = 900 * 1024 +// The payload size cap is runtime-configurable via config.GetMaxPayloadBytes() +// (default 2 MB, see config.DefaultMaxPayloadBytes). Kiro's upstream rejects +// oversized requests with HTTP 400 "Input is too long." +// (CONTENT_LENGTH_EXCEEDS_THRESHOLD), observed near ~2.15 MB. When a converted +// payload exceeds the cap we drop the oldest history turns (keeping the system +// priming, the most recent turns, the active tool turn, and the current message) +// and insert a placeholder note so the model knows context was elided. The cap +// is kept conservatively below the upstream threshold to leave room for headers +// and serialization overhead. // truncationPlaceholder is inserted in history where older turns were dropped to // fit within maxPayloadBytes. @@ -123,7 +123,7 @@ type ClaudeRequest struct { Model string `json:"model"` Messages []ClaudeMessage `json:"messages"` MaxTokens int `json:"max_tokens"` - Temperature float64 `json:"temperature,omitempty"` + Temperature *float64 `json:"temperature,omitempty"` TopP float64 `json:"top_p,omitempty"` Stream bool `json:"stream,omitempty"` System interface{} `json:"system,omitempty"` // string or []SystemBlock @@ -287,16 +287,26 @@ func ClaudeToKiro(req *ClaudeRequest, thinking bool) *KiroPayload { // 构建最终内容 finalContent := "" + toolResultsFolded := false if currentContent != "" { finalContent = currentContent } else if len(currentImages) > 0 { finalContent = normalizeUserContent("", true) } else if len(currentToolResults) > 0 { finalContent = buildToolResultsContinuation(currentToolResults) + toolResultsFolded = true } else { finalContent = minimalFallbackUserContent } + // Orphaned current tool results (no matching assistant tool turn in history) + // cannot stay structured — upstream accepts only a single active tool turn and + // these have none to answer — so fold their text into the message instead of + // losing it (e.g. a tool's textual output alongside an extracted image). + if !keepCurrentToolResults && len(currentToolResults) > 0 && !toolResultsFolded { + finalContent = joinHistoryText(finalContent, buildToolResultsContinuation(currentToolResults)) + } + // 转换工具 kiroTools, toolNameMap := convertClaudeTools(req.Tools) @@ -331,7 +341,7 @@ func ClaudeToKiro(req *ClaudeRequest, thinking bool) *KiroPayload { payload.ConversationState.History = history } - if req.MaxTokens > 0 || req.Temperature > 0 || req.TopP > 0 { + if req.MaxTokens > 0 || req.Temperature != nil || req.TopP > 0 { payload.InferenceConfig = &InferenceConfig{ MaxTokens: req.MaxTokens, Temperature: req.Temperature, @@ -797,6 +807,7 @@ func convertClaudeTools(tools []ClaudeTool) ([]KiroToolWrapper, map[string]strin w.ToolSpecification.InputSchema = InputSchema{JSON: ensureObjectSchema(tool.InputSchema)} result = append(result, w) } + result = compressToolsIfNeeded(result) return result, nameMap } @@ -980,7 +991,7 @@ type OpenAIRequest struct { Model string `json:"model"` Messages []OpenAIMessage `json:"messages"` MaxTokens int `json:"max_tokens,omitempty"` - Temperature float64 `json:"temperature,omitempty"` + Temperature *float64 `json:"temperature,omitempty"` TopP float64 `json:"top_p,omitempty"` Stream bool `json:"stream,omitempty"` Tools []OpenAITool `json:"tools,omitempty"` @@ -1232,15 +1243,22 @@ func OpenAIToKiro(req *OpenAIRequest, thinking bool) *KiroPayload { // 构建最终内容 finalContent := currentContent + toolResultsFolded := false if finalContent == "" { if len(currentImages) > 0 { finalContent = normalizeUserContent("", true) } else if len(currentToolResults) > 0 { finalContent = buildToolResultsContinuation(currentToolResults) + toolResultsFolded = true } else { finalContent = minimalFallbackUserContent } } + // Orphaned current tool results: see ClaudeToKiro — fold them into the message + // text instead of dropping them when an image placeholder or user text was chosen. + if !keepCurrentToolResults && len(currentToolResults) > 0 && !toolResultsFolded { + finalContent = joinHistoryText(finalContent, buildToolResultsContinuation(currentToolResults)) + } // 转换工具 kiroTools := convertOpenAITools(req.Tools) @@ -1271,7 +1289,7 @@ func OpenAIToKiro(req *OpenAIRequest, thinking bool) *KiroPayload { payload.ConversationState.History = history } - if req.MaxTokens > 0 || req.Temperature > 0 || req.TopP > 0 { + if req.MaxTokens > 0 || req.Temperature != nil || req.TopP > 0 { payload.InferenceConfig = &InferenceConfig{ MaxTokens: req.MaxTokens, Temperature: req.Temperature, @@ -1625,7 +1643,8 @@ func truncatePayloadToLimit(payload *KiroPayload, hasPriming bool) { if payload == nil { return } - if payloadByteSize(payload) <= maxPayloadBytes { + limit := config.GetMaxPayloadBytes() + if payloadByteSize(payload) <= limit { return } @@ -1667,7 +1686,7 @@ func truncatePayloadToLimit(payload *KiroPayload, hasPriming bool) { for i := len(conversation) - 1; i >= 0; i-- { running += entrySizes[i] kept := len(conversation) - i - if running > maxPayloadBytes && kept > minRecentHistoryTurns { + if running > limit && kept > minRecentHistoryTurns { break } keepFrom = i @@ -1686,7 +1705,7 @@ func truncatePayloadToLimit(payload *KiroPayload, hasPriming bool) { // If still too large (current message or retained tail alone exceeds the // limit), shrink the current message content as a last resort. - if payloadByteSize(payload) > maxPayloadBytes { + if payloadByteSize(payload) > limit { truncateCurrentMessage(payload) } } @@ -1729,7 +1748,7 @@ func currentMessageModelID(payload *KiroPayload) string { func truncateCurrentMessage(payload *KiroPayload) { cur := &payload.ConversationState.CurrentMessage.UserInputMessage overhead := payloadByteSize(payload) - len(cur.Content) - budget := maxPayloadBytes - overhead + budget := config.GetMaxPayloadBytes() - overhead if budget < 0 { budget = 0 } @@ -2016,6 +2035,7 @@ func convertOpenAITools(tools []OpenAITool) []KiroToolWrapper { wrapper.ToolSpecification.InputSchema = InputSchema{JSON: ensureObjectSchema(tool.Function.Parameters)} result = append(result, wrapper) } + result = compressToolsIfNeeded(result) return result } diff --git a/proxy/translator_test.go b/proxy/translator_test.go index ad6ef174..1cc44306 100644 --- a/proxy/translator_test.go +++ b/proxy/translator_test.go @@ -554,12 +554,42 @@ func TestClaudeToolResultMixedTextAndImage(t *testing.T) { if len(cur.Images) != 1 { t.Fatalf("expected one image extracted, got %d", len(cur.Images)) } - if cur.UserInputMessageContext == nil || len(cur.UserInputMessageContext.ToolResults) != 1 { - t.Fatalf("expected one tool result") + // The tool_result is orphaned: it is the only message, with no preceding + // assistant tool_use to answer. It therefore cannot stay structured (upstream + // rejects a dangling tool_result), so its text must be folded into the message + // content rather than dropped — alongside the extracted image. + if !strings.Contains(cur.Content, "here is the screenshot") { + t.Fatalf("expected orphaned tool_result text folded into content, got %q", cur.Content) + } +} + +func TestOpenAIToolResultMixedTextAndImageOrphaned(t *testing.T) { + const dataURL = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" + req := &OpenAIRequest{ + Model: "claude-sonnet-4.5", + Messages: []OpenAIMessage{ + {Role: "user", Content: "look at the file"}, + { + Role: "tool", + ToolCallID: "call_orph", + Content: []interface{}{ + map[string]interface{}{"type": "text", "text": "here is the screenshot"}, + map[string]interface{}{"type": "image_url", "image_url": map[string]interface{}{"url": dataURL}}, + }, + }, + }, + } + + payload := OpenAIToKiro(req, false) + cur := payload.ConversationState.CurrentMessage.UserInputMessage + if len(cur.Images) != 1 { + t.Fatalf("expected one image extracted, got %d", len(cur.Images)) } - gotText := cur.UserInputMessageContext.ToolResults[0].Content[0].Text - if gotText != "here is the screenshot" { - t.Fatalf("expected original tool text preserved, got %q", gotText) + // The trailing tool message is orphaned (no matching assistant tool_call in + // history), so its text must be folded into the message content alongside the + // extracted image rather than dropped. Mirrors TestClaudeToolResultMixedTextAndImage. + if !strings.Contains(cur.Content, "here is the screenshot") { + t.Fatalf("expected orphaned tool_result text folded into content, got %q", cur.Content) } } @@ -624,15 +654,18 @@ func TestOpenAIToolResultImageCarriedWhenFollowedByUser(t *testing.T) { payload := OpenAIToKiro(req, false) - var toolHistImages int + // The flushed tool-result entry keeps its image: sanitizeKiroHistory narrates + // the structured ToolResults into text and clears UserInputMessageContext, but + // leaves the attached Images in place. So count the image on the history + // entry directly (the "look at the file" user turn carries none). + var histImages int for _, h := range payload.ConversationState.History { - if h.UserInputMessage != nil && h.UserInputMessage.UserInputMessageContext != nil && - len(h.UserInputMessage.UserInputMessageContext.ToolResults) > 0 { - toolHistImages += len(h.UserInputMessage.Images) + if h.UserInputMessage != nil { + histImages += len(h.UserInputMessage.Images) } } - if toolHistImages != 1 { - t.Fatalf("expected tool image carried on the flushed tool-result history entry, got %d", toolHistImages) + if histImages != 1 { + t.Fatalf("expected tool image carried on the flushed tool-result history entry, got %d", histImages) } cur := payload.ConversationState.CurrentMessage.UserInputMessage diff --git a/proxy/translator_truncate_test.go b/proxy/translator_truncate_test.go index 34f33a02..849c26fe 100644 --- a/proxy/translator_truncate_test.go +++ b/proxy/translator_truncate_test.go @@ -2,22 +2,29 @@ package proxy import ( "encoding/json" + "kiro-go/config" + "path/filepath" "strings" "testing" ) // TestClaudeToKiroTruncatesOversizedHistory builds a conversation whose history // far exceeds the upstream input limit and verifies the converted payload is -// trimmed below maxPayloadBytes, that a truncation placeholder is inserted, and +// trimmed below the runtime payload cap, that a truncation placeholder is inserted, and // that the current message is preserved. func TestClaudeToKiroTruncatesOversizedHistory(t *testing.T) { + if err := config.Init(filepath.Join(t.TempDir(), "config.json")); err != nil { + t.Fatalf("config.Init: %v", err) + } + limit := config.GetMaxPayloadBytes() + // ~2KB chunk repeated across many turns to blow past the byte limit. big := strings.Repeat("lorem ipsum dolor sit amet ", 80) // ~2.1KB msgs := []ClaudeMessage{ {Role: "user", Content: "start the long task"}, } - for i := 0; i < 800; i++ { + for i := 0; i < 1200; i++ { msgs = append(msgs, ClaudeMessage{Role: "assistant", Content: "step result: " + big}, ClaudeMessage{Role: "user", Content: "next: " + big}, @@ -37,8 +44,8 @@ func TestClaudeToKiroTruncatesOversizedHistory(t *testing.T) { if err != nil { t.Fatalf("marshal failed: %v", err) } - if len(raw) > maxPayloadBytes { - t.Fatalf("payload size %d exceeds limit %d after truncation", len(raw), maxPayloadBytes) + if len(raw) > limit { + t.Fatalf("payload size %d exceeds limit %d after truncation", len(raw), limit) } // The current message must be preserved. diff --git a/proxy/upstream_error.go b/proxy/upstream_error.go new file mode 100644 index 00000000..73f54aa1 --- /dev/null +++ b/proxy/upstream_error.go @@ -0,0 +1,82 @@ +package proxy + +import ( + "errors" + "strings" +) + +// improperlyFormedMarker is the substring AWS returns in the 400 response body +// when it rejects the request payload itself ("Improperly formed request."). +// Matched case-insensitively so a minor upstream wording/casing change does not +// silently bypass the permanent-error short-circuit. Ported from kiro-tutu. +const improperlyFormedMarker = "improperly formed request" + +// isImproperlyFormedRejection reports whether an upstream error body marks the +// request as improperly formed. Both the endpoint short-circuit (kiro.go) and +// handleAccountFailure key off this single predicate so they can never disagree. +func isImproperlyFormedRejection(errBody string) bool { + return strings.Contains(strings.ToLower(errBody), improperlyFormedMarker) +} + +// upstreamPermanentErrorCode marks a request that the upstream rejects for a +// reason inherent to the request itself (e.g. "Improperly formed request"). +// Retrying it — across endpoints or accounts — cannot succeed and only +// amplifies the damage (wasted upstream hits, scattered cache affinity, +// unwarranted account health penalties), so it must short-circuit both layers. +const upstreamPermanentErrorCode = "upstream_permanent_rejection" + +// upstreamPermanentError signals an upstream rejection that is inherent to the +// request and therefore not retryable. Both the endpoint loop (kiro.go) and the +// account loop short-circuit on it: no other endpoint/account is tried, account +// health is not penalised (handleAccountFailure treats it as neutral). +type upstreamPermanentError struct { + status int + code string + message string +} + +func newUpstreamPermanentError(status int, message string) error { + return &upstreamPermanentError{ + status: status, + code: upstreamPermanentErrorCode, + message: message, + } +} + +func (e *upstreamPermanentError) Error() string { + return e.message +} + +func asUpstreamPermanentError(err error) (*upstreamPermanentError, bool) { + var target *upstreamPermanentError + if errors.As(err, &target) { + return target, true + } + return nil, false +} + +func isUpstreamPermanentError(err error) bool { + _, ok := asUpstreamPermanentError(err) + return ok +} + +// improperlyFormedClientMessage translates the upstream's opaque "Improperly +// formed request" rejection into a readable client hint. In practice the vast +// majority of these rejections occur when the request structure is perfectly +// valid but the tool definitions are too numerous or the payload too large +// (upstream limits on large requests); the raw wording makes users think the +// gateway is buggy. Other errors pass through unchanged. +// +// Complements tool_compression.go (compressToolsIfNeeded): compression prevents +// the rejection when it can; when it cannot, this surfaces a actionable message +// instead of the raw upstream text. Ported from kiro-tutu (zero-dep). +func improperlyFormedClientMessage(err error) string { + if err == nil { + return "" + } + msg := err.Error() + if isImproperlyFormedRejection(msg) { + return "upstream rejected the request (commonly caused by too many tool definitions or an oversized payload); try reducing the number of tools or shortening the context" + } + return msg +} diff --git a/proxy/upstream_error_test.go b/proxy/upstream_error_test.go new file mode 100644 index 00000000..c29f96fc --- /dev/null +++ b/proxy/upstream_error_test.go @@ -0,0 +1,53 @@ +package proxy + +import ( + "errors" + "strings" + "testing" +) + +// TestImproperlyFormedClientMessageTranslatesOpaqueRejection verifies the upstream's +// opaque "Improperly formed request" wording is translated into a readable, actionable +// hint. Case-insensitive so a minor upstream wording/casing change does not bypass it. +func TestImproperlyFormedClientMessageTranslatesOpaqueRejection(t *testing.T) { + cases := []string{ + "Improperly formed request.", + "improperly formed request", + "HTTP 400: IMPROPERLY FORMED REQUEST body", + } + for _, raw := range cases { + got := improperlyFormedClientMessage(errors.New(raw)) + if got == raw { + t.Fatalf("expected translation for %q, got raw message back", raw) + } + if !strings.Contains(got, "tool definitions") || !strings.Contains(got, "reducing") { + t.Fatalf("expected actionable hint about tools/payload, got %q", got) + } + } +} + +// TestImproperlyFormedClientMessagePassesThroughOtherErrors verifies non-matching +// errors pass through unchanged — the helper must only rewrite the one known opaque +// rejection, not munge unrelated upstream/client messages. +func TestImproperlyFormedClientMessagePassesThroughOtherErrors(t *testing.T) { + cases := []string{ + "upstream returned empty assistant response", + "HTTP 500 from q: internal error", + "quota exhausted on codewhisperer", + "context length exceeded", + } + for _, raw := range cases { + got := improperlyFormedClientMessage(errors.New(raw)) + if got != raw { + t.Fatalf("expected pass-through for %q, got %q", raw, got) + } + } +} + +// TestImproperlyFormedClientMessageHandlesNil locks the nil-safety contract so a +// nil error (a missed guard at a call site) returns "" rather than panicking. +func TestImproperlyFormedClientMessageHandlesNil(t *testing.T) { + if got := improperlyFormedClientMessage(nil); got != "" { + t.Fatalf("nil error must yield empty string, got %q", got) + } +} diff --git a/proxy/warmup.go b/proxy/warmup.go new file mode 100644 index 00000000..d9cbd40e --- /dev/null +++ b/proxy/warmup.go @@ -0,0 +1,67 @@ +package proxy + +// Token warm-up worker (production resilience). +// +// Complements the coarse 30-minute backgroundRefresh with a tighter 2-minute +// cadence and a 15-minute look-ahead: tokens nearing expiry are refreshed off +// the request path, so requests rarely block on a synchronous refresh and a +// transient refresh failure has time to recover before the token actually +// expires. It reuses tokenRefreshMu (shared with the request path) and +// re-checks under lock, so it never races or double-refreshes. +// +// Ported from kiro-tutu (zero-dep). +import ( + "kiro-go/config" + "kiro-go/logger" + "time" +) + +const ( + warmupTokenWindowSeconds int64 = 15 * 60 + warmupInterval = 2 * time.Minute +) + +func (h *Handler) backgroundWarmup() { + ticker := time.NewTicker(warmupInterval) + defer ticker.Stop() + for range ticker.C { + h.warmupExpiringTokens() + } +} + +func (h *Handler) warmupExpiringTokens() { + now := time.Now().Unix() + for _, acc := range h.pool.GetAllAccounts() { + if acc.ExpiresAt <= 0 || now < acc.ExpiresAt-warmupTokenWindowSeconds { + continue // not near expiry yet + } + account := acc // work on a copy; refresh results are written back via pool/config + + h.tokenRefreshMu.Lock() + // Re-read latest state under lock in case the request path just refreshed it. + if latest := h.pool.GetByID(account.ID); latest != nil { + account.AccessToken = latest.AccessToken + account.RefreshToken = latest.RefreshToken + account.ExpiresAt = latest.ExpiresAt + account.ProfileArn = latest.ProfileArn + } + if account.ExpiresAt <= 0 || time.Now().Unix() < account.ExpiresAt-warmupTokenWindowSeconds { + h.tokenRefreshMu.Unlock() + continue // already refreshed by another path + } + + accessToken, refreshToken, expiresAt, profileArn, err := authRefreshToken(&account) + if err != nil { + h.tokenRefreshMu.Unlock() + logger.Warnf("[Warmup] refresh failed for %s: %v", account.ID, err) + continue + } + h.pool.UpdateToken(account.ID, accessToken, refreshToken, expiresAt) + if profileArn != "" { + config.UpdateAccountProfileArn(account.ID, profileArn) + } + config.UpdateAccountToken(account.ID, accessToken, refreshToken, expiresAt) + h.tokenRefreshMu.Unlock() + logger.Infof("[Warmup] proactively refreshed token for %s", account.ID) + } +} diff --git a/web/app.js b/web/app.js index 085183c0..c5970abe 100644 --- a/web/app.js +++ b/web/app.js @@ -21,6 +21,8 @@ let promptRules = []; let builderIdSession = ''; let builderIdPollTimer = null; + let kiroSsoSession = ''; + let kiroSsoPollTimer = null; let iamSession = ''; let exportSelectedIds = new Set(); let currentVersion = ''; @@ -1535,6 +1537,8 @@ const d = await res.json(); $('requireApiKey').checked = d.requireApiKey; $('allowOverUsage').checked = d.allowOverUsage || false; + const maxPayloadEl = document.getElementById('maxPayloadBytes'); + if (maxPayloadEl) maxPayloadEl.value = String(d.maxPayloadBytes || 2000000); await Promise.all([loadThinkingConfig(), loadEndpointConfig(), loadProxyConfig(), loadPromptFilter(), loadApiKeys()]); refreshCustomSelects(); } @@ -1640,7 +1644,9 @@ } async function saveOverUsageConfig() { const allowOverUsage = $('allowOverUsage').checked; - await api('/settings', { method: 'POST', body: JSON.stringify({ allowOverUsage }) }); + const maxPayloadEl = document.getElementById('maxPayloadBytes'); + const maxPayloadBytes = maxPayloadEl ? parseInt(maxPayloadEl.value || '0', 10) : 0; + await api('/settings', { method: 'POST', body: JSON.stringify({ allowOverUsage, maxPayloadBytes }) }); toast(t('settings.overUsageSaved'), 'success'); } async function changePassword() { @@ -2019,10 +2025,13 @@ var METHOD_ICONS = { builderid: 'fa-solid fa-id-card', iam: 'fa-solid fa-key', + enterprisesso: 'fa-brands fa-microsoft', sso: 'fa-solid fa-shield-halved', local: 'fa-solid fa-folder-open', credentials: 'fa-solid fa-code', - cookie: 'fa-solid fa-cookie-bite' + cookie: 'fa-solid fa-cookie-bite', + apikey: 'fa-solid fa-key', + apikeybatch: 'fa-solid fa-layer-group' }; function methodCard(type, title, desc) { var icon = METHOD_ICONS[type] || 'fa-solid fa-circle-plus'; @@ -2042,10 +2051,13 @@ if (type === 'add') modalAdd(title, body); else if (type === 'builderid') modalBuilderId(title, body); else if (type === 'iam') modalIam(title, body); + else if (type === 'enterprisesso') modalEnterpriseSso(title, body); else if (type === 'sso') modalSso(title, body); else if (type === 'local') modalLocal(title, body); else if (type === 'credentials') modalCredentials(title, body); else if (type === 'cookie') modalCookie(title, body); + else if (type === 'apikey') modalApiKey(title, body); + else if (type === 'apikeybatch') modalApiKeyBatch(title, body); if (!modal.classList.contains('active')) openDialog('addModal'); enhanceCustomSelects(body); } @@ -2054,6 +2066,14 @@ iamSession = ''; if (builderIdPollTimer) { clearTimeout(builderIdPollTimer); builderIdPollTimer = null; } builderIdSession = ''; + if (kiroSsoPollTimer) { clearTimeout(kiroSsoPollTimer); kiroSsoPollTimer = null; } + // If a hosted-portal sign-in is still in flight (modal closed via X/backdrop + // before completion), release the loopback port now. On successful completion + // the poller clears kiroSsoSession first, so this no-ops then. + if (kiroSsoSession) { + api('/auth/kiro-sso/cancel', { method: 'POST', body: JSON.stringify({ sessionId: kiroSsoSession }) }).catch(() => {}); + } + kiroSsoSession = ''; } function modalAdd(title, body) { title.textContent = t('modal.addAccount'); @@ -2061,10 +2081,13 @@ '
' + methodCard('builderid', t('modal.builderIdTitle'), t('modal.builderIdDesc')) + methodCard('iam', t('modal.iamTitle'), t('modal.iamDesc')) + + methodCard('enterprisesso', t('modal.enterpriseSsoTitle'), t('modal.enterpriseSsoDesc')) + methodCard('sso', t('modal.ssoTitle'), t('modal.ssoDesc')) + methodCard('local', t('modal.localTitle'), t('modal.localDesc')) + methodCard('credentials', t('modal.credentialsTitle'), t('modal.credentialsDesc')) + methodCard('cookie', t('modal.cookieTitle'), t('modal.cookieDesc')) + + methodCard('apikey', t('modal.apikeyTitle'), t('modal.apikeyDesc')) + + methodCard('apikeybatch', t('modal.apikeyBatchTitle'), t('modal.apikeyBatchDesc')) + '
' + '
'; } @@ -2221,6 +2244,35 @@ ''; $('importCookieBtn').addEventListener('click', importFromCookie); } + function modalApiKey(title, body) { + title.textContent = t('modal.apikeyTitle'); + body.innerHTML = + '

' + escapeHtml(t('apikey.hint')) + '

' + + '
' + + '' + + '
' + + '
' + + '
' + + '' + + '' + + '
'; + $('importApiKeyBtn').addEventListener('click', importApiKey); + } + function modalApiKeyBatch(title, body) { + title.textContent = t('modal.apikeyBatchTitle'); + body.innerHTML = + '

' + escapeHtml(t('apikeyBatch.hint')) + '

' + + '
' + + '' + + '' + escapeHtml(t('apikeyBatch.onePerLine')) + '' + + '
' + + '
' + + '
' + + '' + + '' + + '
'; + $('importApiKeyBatchBtn').addEventListener('click', importApiKeysBatch); + } function updateLocalFields() { const p = $('localProvider').value; $('localClientGroup').classList.toggle('hidden', p === 'Google' || p === 'Github'); @@ -2277,11 +2329,19 @@ const c = a.credentials || {}; return { refreshToken: c.refreshToken || a.refreshToken, + accessToken: c.accessToken || a.accessToken, clientId: c.clientId || a.clientId, clientSecret: c.clientSecret || a.clientSecret, region: c.region || a.region, authMethod: c.authMethod || a.authMethod, - provider: c.provider || a.provider || a.idp + provider: c.provider || a.provider || a.idp, + tokenEndpoint: c.tokenEndpoint || a.tokenEndpoint, + issuerUrl: c.issuerUrl || a.issuerUrl, + scopes: c.scopes || a.scopes, + id: a.id, + email: c.email || a.email, + profileArn: c.profileArn || a.profileArn, + userId: a.userId }; }); } else { @@ -2303,11 +2363,19 @@ let ok = 0, fail = 0, newIds = []; for (const item of items) { if (!item.refreshToken) { fail++; continue; } - let authMethod = item.authMethod || ''; - if (item.clientId && item.clientSecret) authMethod = 'idc'; - else if (!authMethod || authMethod === 'social') authMethod = 'social'; - else authMethod = authMethod.toLowerCase() === 'idc' ? 'idc' : 'social'; + const EXTERNAL_IDP = ['external_idp','azuread','azure','entra','entra-id','microsoft','m365','office365','external']; + let authMethod = (item.authMethod || '').toLowerCase(); + if (EXTERNAL_IDP.includes(authMethod) || item.tokenEndpoint) { + authMethod = 'external_idp'; + } else if (item.clientId && item.clientSecret) { + authMethod = 'idc'; + } else if (!authMethod || authMethod === 'social') { + authMethod = 'social'; + } else { + authMethod = authMethod === 'idc' ? 'idc' : 'social'; + } let provider = item.provider || ''; + if (!provider && authMethod === 'external_idp') provider = 'AzureAD'; if (!provider && authMethod === 'social') provider = 'Google'; if (!provider && authMethod === 'idc') provider = 'BuilderId'; const payload = { @@ -2316,7 +2384,13 @@ clientId: item.clientId || '', clientSecret: item.clientSecret || '', authMethod, provider, - region: item.region || 'us-east-1' + region: item.region || 'us-east-1', + tokenEndpoint: item.tokenEndpoint || '', + issuerUrl: item.issuerUrl || '', + scopes: item.scopes || '', + ...(item.id ? { id: item.id } : {}), + ...(item.email ? { email: item.email } : {}), + ...(item.profileArn ? { profileArn: item.profileArn } : {}) }; try { const res = await api('/auth/credentials', { method: 'POST', body: JSON.stringify(payload) }); @@ -2370,6 +2444,39 @@ autoRefreshNewAccount(d.account?.id); } else toastError(t('common.failed') + ': ' + (d.error || '')); } + async function importApiKey() { + const kiroApiKey = $('apiKeyValue').value.trim(); + if (!kiroApiKey) return toastWarning(t('apikey.keyMissing')); + const region = $('apiKeyRegion').value || 'us-east-1'; + const res = await api('/auth/credentials', { method: 'POST', body: JSON.stringify({ kiroApiKey, authMethod: 'api_key', region }) }); + const d = await res.json(); + if (d.success) { + closeModal(); loadAccounts(); loadStats(); + toastPrimary(t('apikey.importSuccess') + ': ' + (d.account?.email || d.account?.id)); + autoRefreshNewAccount(d.account?.id); + } else toastError(t('common.failed') + ': ' + (d.error || '')); + } + async function importApiKeysBatch() { + const keys = $('apiKeyBatchValue').value.trim(); + if (!keys) return toastWarning(t('apikeyBatch.keysMissing')); + const region = $('apiKeyBatchRegion').value || 'us-east-1'; + const res = await api('/auth/apikeys-batch', { method: 'POST', body: JSON.stringify({ keys, region }) }); + const d = await res.json(); + if (d.success) { + closeModal(); loadAccounts(); loadStats(); + let msg = t('apikeyBatch.summary', d.imported || 0, d.total || 0, d.skipped || 0); + if (d.infoFailed > 0) msg += t('apikeyBatch.infoFailed', d.infoFailed); + toastPrimary(msg, { duration: 6000 }); + renderApiKeyBatchResults(d.results || []); + } else toastError(t('common.failed') + ': ' + (d.error || '')); + } + function renderApiKeyBatchResults(results) { + // Per-key detail surfaces via the toast summary above; log masked failures + // for the operator so a full per-key panel can be built later if needed. + results.forEach(r => { + if (r.error) console.warn('[ApiKeyBatch]', r.maskedKey, r.error); + }); + } async function importSsoToken() { const res = await api('/auth/sso-token', { method: 'POST', body: JSON.stringify({ @@ -2429,6 +2536,146 @@ builderIdSession = ''; showModal('add'); } + // Enterprise SSO — Microsoft 365 / Entra ID (Azure AD), via the Kiro hosted sign-in portal. + // The backend binds a loopback listener and returns the sign-in URL; the browser is driven + // through the external-IdP leg automatically, and we poll until the account is created. + function modalEnterpriseSso(title, body) { + title.textContent = t('modal.enterpriseSsoTitle'); + body.innerHTML = + '

' + escapeHtml(t('modal.enterpriseSsoDesc')) + '

' + + '
' + + '

' + escapeHtml(t('kirosso.hostNote')) + '

' + + '
' + + '' + + '' + + '
' + + '
' + + ''; + $('startKiroSsoBtn').addEventListener('click', startKiroSsoLogin); + $('kiroSsoSubmitCallbackBtn').addEventListener('click', submitKiroSsoCallback); + $('kiroSsoCallbackUrl').addEventListener('keydown', (e) => { if (e.key === 'Enter') submitKiroSsoCallback(); }); + } + async function startKiroSsoLogin() { + // No region prompt: the data-plane region is derived from the profile ARN + // returned by SSO (social) or discovered via the cross-region profile probe + // (external_idp / Azure), so the operator never has to know it up front. + const res = await api('/auth/kiro-sso/start', { method: 'POST', body: JSON.stringify({}) }); + const d = await res.json(); + if (d.sessionId && d.signInUrl) { + kiroSsoSession = d.sessionId; + $('kiroSsoSignInUrl').textContent = d.signInUrl; + $('kiroSsoStep1').classList.add('hidden'); + $('kiroSsoStep2').classList.remove('hidden'); + $('kiroSsoOpenBtn').addEventListener('click', () => window.open($('kiroSsoSignInUrl').textContent, '_blank')); + $('kiroSsoCopyBtn').addEventListener('click', async () => { + await copyText($('kiroSsoSignInUrl').textContent); + toast(t('common.copied'), 'primary'); + }); + $('kiroSsoCancelBtn').addEventListener('click', cancelKiroSsoLogin); + // Open the sign-in tab immediately (works when the admin panel is viewed on the proxy host). + window.open(d.signInUrl, '_blank'); + pollKiroSso(d.interval || 2); + } else toastError(t('common.failed') + ': ' + (d.error || '')); + } + async function submitKiroSsoCallback() { + const url = $('kiroSsoCallbackUrl').value.trim(); + if (!url) return; + if (!kiroSsoSession) { toastError('No active SSO session'); return; } + $('kiroSsoSubmitCallbackBtn').disabled = true; + $('kiroSsoStatus').textContent = 'Processing callback...'; + $('kiroSsoStatus').style.color = 'var(--muted)'; + try { + const res = await api('/auth/kiro-sso/callback', { + method: 'POST', body: JSON.stringify({ sessionId: kiroSsoSession, callbackUrl: url }) + }); + const d = await res.json(); + if (d.success && d.redirectUrl) { + // Enterprise SSO leg-1: show Microsoft login, advance to step 2 + $('kiroSsoStepBadge').textContent = 'Step 2'; + $('kiroSsoStepLabel').textContent = 'Paste the final redirect URL after Microsoft login'; + $('kiroSsoStepHint').innerHTML = 'After Microsoft 365 authentication, your browser will again redirect to localhost:3128. Copy and paste that final URL.'; + $('kiroSsoRedirectLink').href = d.redirectUrl; + $('kiroSsoRedirectArea').classList.remove('hidden'); + $('kiroSsoCallbackUrl').placeholder = 'http://localhost:3128/oauth/callback?code=...'; + $('kiroSsoStatus').textContent = '↑ Open the Microsoft login link above, then paste the next redirect URL'; + $('kiroSsoStatus').style.color = 'var(--warning)'; + $('kiroSsoCallbackUrl').value = ''; + // Auto-open the Microsoft login link + window.open(d.redirectUrl, '_blank'); + } else if (d.success) { + // Leg-2 or social: polling will pick up the result + $('kiroSsoStatus').textContent = 'Callback accepted — completing login...'; + $('kiroSsoStatus').style.color = 'var(--success)'; + $('kiroSsoCallbackUrl').value = ''; + } else { + toastError(t('common.failed') + ': ' + (d.error || '')); + } + } catch (e) { + toastError('Failed to submit callback: ' + e.message); + } finally { + $('kiroSsoSubmitCallbackBtn').disabled = false; + } + } + function pollKiroSso(interval) { + kiroSsoPollTimer = setTimeout(async () => { + const res = await api('/auth/kiro-sso/poll', { method: 'POST', body: JSON.stringify({ sessionId: kiroSsoSession }) }); + const d = await res.json(); + if (d.completed) { + // Session is already consumed server-side; clear it so closeModal() does + // not fire a redundant cancel for an account that succeeded. + kiroSsoSession = ''; + closeModal(); loadAccounts(); loadStats(); + toastPrimary(t('builderid.success') + ': ' + (d.account?.email || d.account?.id)); + autoRefreshNewAccount(d.account?.id); + } else if (d.success && !d.completed) { + // Don't overwrite a manual-instruction status message + if (!$('kiroSsoStatus').textContent.includes('↑') && !$('kiroSsoStatus').textContent.includes('Step')) { + $('kiroSsoStatus').textContent = 'Waiting for callback URL — paste the redirect from your browser above'; + $('kiroSsoStatus').style.color = 'var(--warning)'; + } + pollKiroSso(interval); + } else { + toastError(t('common.failed') + ': ' + (d.error || '')); + cancelKiroSsoLogin(); + } + }, interval * 1000); + } + function cancelKiroSsoLogin() { + if (kiroSsoPollTimer) { clearTimeout(kiroSsoPollTimer); kiroSsoPollTimer = null; } + // Tell the backend to release the loopback callback port now instead of waiting + // for the deadline (fire-and-forget; ignore the result). + if (kiroSsoSession) { + api('/auth/kiro-sso/cancel', { method: 'POST', body: JSON.stringify({ sessionId: kiroSsoSession }) }).catch(() => {}); + } + kiroSsoSession = ''; + showModal('add'); + } async function startIamSso() { if (iamSession) { const res = await api('/auth/iam-sso/complete', { diff --git a/web/index-legacy.html b/web/index-legacy.html index bde09215..596b0d90 100644 --- a/web/index-legacy.html +++ b/web/index-legacy.html @@ -2851,12 +2851,21 @@

+
+ + + +
diff --git a/web/locales/en.json b/web/locales/en.json index ffac0bfc..affb3e8a 100644 --- a/web/locales/en.json +++ b/web/locales/en.json @@ -227,6 +227,10 @@ "modal.credentialsDesc": "Add account via Kiro Account Manager exported credentials", "modal.cookieTitle": "Kiro Web Cookie", "modal.cookieDesc": "Add account via RefreshToken obtained from Kiro Web Cookie", + "modal.enterpriseSsoTitle": "Enterprise SSO - Microsoft 365", + "modal.enterpriseSsoDesc": "Add a Microsoft 365 / Entra ID (Azure AD) tenant account via Kiro hosted SSO", + "kirosso.hostNote": "Open the sign-in link on this machine (the same host as the proxy). The browser is redirected through your Microsoft 365 login automatically.", + "kirosso.openInstruction": "A browser tab was opened. Sign in with your Microsoft 365 work/school account, then return here. Tip: use a guest/incognito window to avoid a cached session.", "cookie.howToGet": "How to get the RefreshToken?", "cookie.step1": "Open your browser and sign in at", "cookie.step2": "Right-click on blank area → Inspect → Application → Cookies → https://app.kiro.dev", @@ -239,6 +243,24 @@ "cookie.refreshTokenMissing": "RefreshToken is required", "cookie.importSuccess": "Account added successfully", "cookie.link": "https://app.kiro.dev/account/usage", + "modal.apikeyTitle": "API Key", + "modal.apikeyDesc": "Add an account using a Kiro API key directly as the bearer token (no OAuth refresh)", + "modal.apikeyBatchTitle": "Bulk API Key Import", + "modal.apikeyBatchDesc": "Import many Kiro API keys at once (one per line)", + "apikey.hint": "Paste a Kiro API key. It is used directly as the bearer token (tokentype: API_KEY); no refresh token or profile ARN is needed.", + "apikey.key": "API Key", + "apikey.keyPlaceholder": "Paste the Kiro API key here", + "apikey.keyMissing": "API key is required", + "apikey.importSuccess": "Account added", + "apikeyBatch.hint": "One key per line. Duplicates (already imported or repeated in this batch) are skipped.", + "apikeyBatch.keys": "API Keys", + "apikeyBatch.keysPlaceholder": "one-key-per-line...", + "apikeyBatch.onePerLine": "One key per line", + "apikeyBatch.keysMissing": "At least one API key is required", + "apikeyBatch.summary": "Imported {0} of {1} keys ({2} skipped)", + "apikeyBatch.infoFailed": ", {0} usage-info lookups failed", + "settings.maxPayloadBytes": "Max Payload Size", + "settings.maxPayloadBytesHint": "Cap on the request body size before the upstream rejects it as too large. Higher = more context per request.", "builderid.startLogin": "Start Login", "builderid.verifyCode": "Enter the code above in your browser", "builderid.verifyUrl": "Verification URL", diff --git a/web/locales/zh.json b/web/locales/zh.json index f02a71d3..145d0f00 100644 --- a/web/locales/zh.json +++ b/web/locales/zh.json @@ -227,6 +227,10 @@ "modal.credentialsDesc": "通过 Kiro Account Manager 导出的凭证添加账号", "modal.cookieTitle": "Kiro 网页 Cookie", "modal.cookieDesc": "通过从 Kiro 网页 Cookie 中获取的 RefreshToken 添加账号", + "modal.enterpriseSsoTitle": "企业 SSO - Microsoft 365", + "modal.enterpriseSsoDesc": "通过 Kiro 托管 SSO 添加 Microsoft 365 / Entra ID（Azure AD）租户账号", + "kirosso.hostNote": "请在运行代理的本机打开登录链接（与代理同一主机）。浏览器会自动跳转到 Microsoft 365 登录流程。", + "kirosso.openInstruction": "已打开浏览器标签页。请使用 Microsoft 365 工作/学校账号登录，然后返回此处。建议使用访客/无痕窗口以避免缓存会话干扰。", "cookie.howToGet": "如何获取 RefreshToken？", "cookie.step1": "打开浏览器，访问并登录", "cookie.step2": "在页面空白处右键 → 检查 → Application → Cookies → https://app.kiro.dev", @@ -239,6 +243,24 @@ "cookie.refreshTokenMissing": "请填写 RefreshToken", "cookie.importSuccess": "账号添加成功", "cookie.link": "https://app.kiro.dev/account/usage", + "modal.apikeyTitle": "API 密钥", + "modal.apikeyDesc": "使用 Kiro API 密钥直接作为 Bearer 令牌添加账号（无需 OAuth 刷新）", + "modal.apikeyBatchTitle": "批量导入 API 密钥", + "modal.apikeyBatchDesc": "一次导入多个 Kiro API 密钥（每行一个）", + "apikey.hint": "粘贴 Kiro API 密钥。它将直接作为 Bearer 令牌使用（tokentype: API_KEY），无需 refresh token 或 profile ARN。", + "apikey.key": "API 密钥", + "apikey.keyPlaceholder": "在此粘贴 Kiro API 密钥", + "apikey.keyMissing": "API 密钥不能为空", + "apikey.importSuccess": "账号添加成功", + "apikeyBatch.hint": "每行一个密钥。重复项（已导入或本批次内重复）将被跳过。", + "apikeyBatch.keys": "API 密钥", + "apikeyBatch.keysPlaceholder": "每行一个密钥...", + "apikeyBatch.onePerLine": "每行一个密钥", + "apikeyBatch.keysMissing": "至少需要一个 API 密钥", + "apikeyBatch.summary": "已导入 {1} 个中的 {0} 个密钥（{2} 个跳过）", + "apikeyBatch.infoFailed": "，{0} 个用量信息查询失败", + "settings.maxPayloadBytes": "最大请求体大小", + "settings.maxPayloadBytesHint": "上游拒绝过大请求前的请求体大小上限。数值越大，单次请求可携带的上下文越多。", "builderid.startLogin": "开始登录", "builderid.verifyCode": "请在浏览器中输入上方验证码", "builderid.verifyUrl": "验证链接",