From d8c170e5117b8681ca40b03d28443c33519d9607 Mon Sep 17 00:00:00 2001 From: LuneZ99 Date: Wed, 21 Oct 2026 01:50:54 +0800 Subject: [PATCH 01/10] test(workbuddy): record private UUID buffer invariant --- internal/util/nocopy_invariant_test.go | 1 + 1 file changed, 1 insertion(+) diff --git a/internal/util/nocopy_invariant_test.go b/internal/util/nocopy_invariant_test.go index 9dbfb5ceec2..ad2315b0e4c 100644 --- a/internal/util/nocopy_invariant_test.go +++ b/internal/util/nocopy_invariant_test.go @@ -112,6 +112,7 @@ var reviewedInPlaceByteWrites = map[string]reviewedInPlaceByteWrite{ "internal/client/codex/live/tcp_proxy.go": {1, "copies header and payload into a freshly allocated frame"}, "internal/home/client.go": {1, "zeroes a secret buffer after json.Unmarshal has copied every value out"}, "internal/pluginstore/auth.go": {1, "zeroes a locally built credential buffer after base64 encoding copied it out"}, + "internal/workbuddy/identity.go": {1, "copies a hash digest into a new local UUID array; no JSON payload or GJSON result can alias that array"}, } // TestInPlaceByteWritesAreReviewed keeps the set of in-place byte writes small From 9622e784386f0082b2a003f4b7dffee3d8af3985 Mon Sep 17 00:00:00 2001 From: LuneZ99 Date: Wed, 21 Oct 2026 01:52:53 +0800 Subject: [PATCH 02/10] feat(workbuddy): classify failures and persist scoped risk state --- internal/workbuddy/fault.go | 207 ++++++++++++ internal/workbuddy/fault_test.go | 78 +++++ internal/workbuddy/governor.go | 481 ++++++++++++++++++++++++++++ internal/workbuddy/governor_test.go | 244 ++++++++++++++ 4 files changed, 1010 insertions(+) create mode 100644 internal/workbuddy/fault.go create mode 100644 internal/workbuddy/fault_test.go create mode 100644 internal/workbuddy/governor.go create mode 100644 internal/workbuddy/governor_test.go diff --git a/internal/workbuddy/fault.go b/internal/workbuddy/fault.go new file mode 100644 index 00000000000..bdf082a7f7b --- /dev/null +++ b/internal/workbuddy/fault.go @@ -0,0 +1,207 @@ +package workbuddy + +import ( + "context" + "crypto/tls" + "crypto/x509" + "encoding/json" + "errors" + "io" + "net" + "regexp" + "strconv" + "strings" + "syscall" + "time" +) + +type FaultKind string + +const ( + FaultAuth FaultKind = "authentication" + FaultBalance FaultKind = "balance" + FaultRestricted FaultKind = "account_restricted" + FaultContent FaultKind = "content_rejected" + FaultAccountRate FaultKind = "account_rate" + FaultModelRate FaultKind = "model_rate" + FaultModelCap FaultKind = "model_cap" + FaultModelRetired FaultKind = "model_retired" + FaultProviderUnavailable FaultKind = "provider_unavailable" + FaultEntitlement FaultKind = "model_entitlement" + FaultGateway FaultKind = "gateway" + FaultHard FaultKind = "hard_failure" + FaultUnknown FaultKind = "unknown_failure" + FaultRequest FaultKind = "request_rejected" + FaultLocalNetwork FaultKind = "local_network" + FaultTLS FaultKind = "tls_verification" + FaultCanceled FaultKind = "canceled" + FaultConnection FaultKind = "connection" +) + +// Fault contains only classification evidence. Upstream bodies can contain +// credentials or user content and must not enter persisted risk state. +type Fault struct { + Kind FaultKind + Status int + Code int + ResetAt time.Time + Soft bool +} + +var ( + codePattern = regexp.MustCompile(`"code"\s*:\s*"?([0-9]+)"?`) + resetPhrasePattern = regexp.MustCompile(`(?:reset at|将在)\s*(\d{4}-\d{2}-\d{2}[ T]\d{2}:\d{2}:\d{2})`) + resetClockPattern = regexp.MustCompile(`(\d{4}-\d{2}-\d{2}[ T]\d{2}:\d{2}:\d{2})\s*UTC[+-]\d{1,2}`) + offsetPattern = regexp.MustCompile(`UTC([+-])(\d{1,2})(?::?(\d{2}))?`) + accountScopePattern = regexp.MustCompile(`(?i)\b(?:account-level|account level|account-wide|account scope|account quota|account rate limit|account limit|per account|per-account|this account|your account|the account|an account|credential-level|credential level|credential scope|credential limit|per credential|this credential|the credential|your credential)\b`) +) + +// faultText also exposes JSON-escaped localized display messages to the +// classifier. It never substitutes a default meaning for unknown responses. +func faultText(raw []byte) string { + text := string(raw) + var payload any + if json.Unmarshal(raw, &payload) != nil { + return text + } + var values []string + var collect func(any, int) + collect = func(value any, depth int) { + if depth > 4 { + return + } + switch v := value.(type) { + case string: + values = append(values, v) + case map[string]any: + for _, child := range v { + collect(child, depth+1) + } + case []any: + for _, child := range v { + collect(child, depth+1) + } + } + } + collect(payload, 0) + return text + "\n" + strings.Join(values, "\n") +} + +func ParseResetTime(raw []byte) time.Time { + text := faultText(raw) + match := resetPhrasePattern.FindStringSubmatch(text) + if len(match) == 0 { + match = resetClockPattern.FindStringSubmatch(text) + } + if len(match) == 0 { + return time.Time{} + } + offset := 0 + if zone := offsetPattern.FindStringSubmatch(text); len(zone) > 0 { + hours, errHours := strconv.Atoi(zone[2]) + minutes := 0 + if zone[3] != "" { + var errMinutes error + minutes, errMinutes = strconv.Atoi(zone[3]) + if errMinutes != nil { + return time.Time{} + } + } + if errHours != nil || hours > 23 || minutes > 59 { + return time.Time{} + } + offset = hours*3600 + minutes*60 + if zone[1] == "-" { + offset = -offset + } + } + stamp, errStamp := time.ParseInLocation("2006-01-02 15:04:05", strings.ReplaceAll(match[1], "T", " "), time.FixedZone("upstream", offset)) + if errStamp != nil { + return time.Time{} + } + return stamp.UTC() +} + +func ClassifyHTTP(status int, raw []byte) Fault { + fault := Fault{Status: status} + text := strings.ToLower(faultText(raw)) + // Capture the explicit code field, never a coincidental request ID. + if match := codePattern.FindSubmatch(raw); len(match) > 0 { + fault.Code, _ = strconv.Atoi(string(match[1])) + } + switch status { + case 401: + fault.Kind = FaultAuth + case 402: + fault.Kind = FaultBalance + case 403: + fault.Kind = FaultContent + if fault.Code == 11140 || strings.Contains(text, "unusual activity") || strings.Contains(text, "账号存在异常行为") || strings.Contains(text, "已被限制") { + fault.Kind = FaultRestricted + } + case 429: + fault.Kind = FaultModelRate + fault.ResetAt = ParseResetTime(raw) + if !fault.ResetAt.IsZero() { + if fault.Code == 6004 { + fault.Kind = FaultModelCap + } + return fault + } + accountScoped := accountScopePattern.MatchString(text) + for _, phrase := range []string{"账号级", "账户级", "该账号", "此账号", "该账户", "此账户", "账号限流", "账户限流", "凭证级", "该凭证"} { + accountScoped = accountScoped || strings.Contains(text, phrase) + } + if accountScoped { + fault.Kind = FaultAccountRate + return fault + } + fault.Soft = fault.Code == 14003 || strings.Contains(text, "usage exceeds frequency limit") || + strings.Contains(text, "too many requests") || strings.Contains(text, "请求过于频繁") + default: + if status == 400 || status >= 500 { + switch fault.Code { + case 11102: + fault.Kind = FaultModelRetired + return fault + case 11134: + fault.Kind = FaultProviderUnavailable + return fault + case 11158: + fault.Kind = FaultEntitlement + return fault + } + } + switch { + case status == 502 || status == 503 || status == 504: + fault.Kind = FaultGateway + case status >= 500: + fault.Kind = FaultHard + default: + fault.Kind = FaultRequest + } + } + return fault +} + +func ClassifyTransport(err error) Fault { + if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) { + return Fault{Kind: FaultCanceled} + } + var unknownAuthority x509.UnknownAuthorityError + var hostname x509.HostnameError + var invalidCertificate x509.CertificateInvalidError + var verification *tls.CertificateVerificationError + if errors.As(err, &unknownAuthority) || errors.As(err, &hostname) || errors.As(err, &invalidCertificate) || errors.As(err, &verification) { + return Fault{Kind: FaultTLS} + } + var dns *net.DNSError + if errors.As(err, &dns) || errors.Is(err, syscall.ENETDOWN) || errors.Is(err, syscall.ENETUNREACH) || errors.Is(err, syscall.EHOSTUNREACH) { + return Fault{Kind: FaultLocalNetwork} + } + var operation *net.OpError + if errors.As(err, &operation) || errors.Is(err, io.EOF) || errors.Is(err, io.ErrUnexpectedEOF) { + return Fault{Kind: FaultConnection} + } + return Fault{Kind: FaultUnknown} +} diff --git a/internal/workbuddy/fault_test.go b/internal/workbuddy/fault_test.go new file mode 100644 index 00000000000..44a66121471 --- /dev/null +++ b/internal/workbuddy/fault_test.go @@ -0,0 +1,78 @@ +package workbuddy + +import ( + "context" + "crypto/x509" + "fmt" + "io" + "net" + "syscall" + "testing" + "time" +) + +func TestFaultScopeAndResetClassification(t *testing.T) { + for _, tc := range []struct { + name string + status int + body string + kind FaultKind + soft bool + reset string + }{ + {"free_cap_en", 429, `{"code":6004,"msg":"your usage will reset at 2026-10-02 16:00:34 UTC+8"}`, FaultModelCap, false, "2026-10-02T08:00:34Z"}, + {"free_cap_cn", 429, `{"code":"6004","msg":"将在 2026-10-02T16:00:34 UTC+8 重置"}`, FaultModelCap, false, "2026-10-02T08:00:34Z"}, + {"escaped_cn", 429, `{"code":6004,"msg":"\u5c06\u5728 2026-10-02 16:00:34 UTC+05:30"}`, FaultModelCap, false, "2026-10-02T10:30:34Z"}, + {"negative_half_hour", 429, `{"code":6004,"msg":"reset at 2026-10-02 16:00:34 UTC-00:30"}`, FaultModelCap, false, "2026-10-02T16:30:34Z"}, + {"soft_clock_is_not_cap", 429, `{"code":14003,"msg":"your account usage will reset at 2026-10-02 16:00:34 UTC+8"}`, FaultModelRate, false, "2026-10-02T08:00:34Z"}, + {"cap_without_reset", 429, `{"code":6004}`, FaultModelRate, false, ""}, + {"account_phrase", 429, `{"msg":"your account is rate limited"}`, FaultAccountRate, false, ""}, + {"account_metadata_is_not_scope", 429, `{"accountId":"u","credentialId":"x","msg":"too many requests"}`, FaultModelRate, true, ""}, + {"account_cn", 429, `{"msg":"\u8d26\u53f7\u9650\u6d41"}`, FaultAccountRate, false, ""}, + {"soft_text", 429, "usage exceeds frequency limit", FaultModelRate, true, ""}, + {"unknown_shape", 429, "busy", FaultModelRate, false, ""}, + {"restricted_code", 403, `{"code":11140,"msg":"request illegal"}`, FaultRestricted, false, ""}, + {"restricted_phrase", 403, `{"msg":"Your account has been restricted due to unusual activity"}`, FaultRestricted, false, ""}, + {"request_id_is_not_code", 403, `{"requestId":"0011140abc","msg":"content rejected"}`, FaultContent, false, ""}, + {"auth", 401, "", FaultAuth, false, ""}, + {"credits", 402, "", FaultBalance, false, ""}, + {"retired", 500, `{"code":11102}`, FaultModelRetired, false, ""}, + {"provider_unavailable", 503, `{"code":11134}`, FaultProviderUnavailable, false, ""}, + {"entitlement", 400, `{"code":11158}`, FaultEntitlement, false, ""}, + {"gateway", 502, "", FaultGateway, false, ""}, + {"hard", 550, "", FaultHard, false, ""}, + {"request", 400, "unknown", FaultRequest, false, ""}, + } { + t.Run(tc.name, func(t *testing.T) { + got := ClassifyHTTP(tc.status, []byte(tc.body)) + if got.Kind != tc.kind || got.Soft != tc.soft { + t.Fatalf("got kind=%s soft=%t", got.Kind, got.Soft) + } + if tc.reset == "" { + if !got.ResetAt.IsZero() { + t.Fatal("invented reset time") + } + } else if got.ResetAt.Format(time.RFC3339) != tc.reset { + t.Fatalf("reset=%s want %s", got.ResetAt, tc.reset) + } + }) + } +} + +func TestTransportFaultDoesNotBlameCredentialForLocalFailures(t *testing.T) { + for _, tc := range []struct { + err error + kind FaultKind + }{ + {fmt.Errorf("wrapped: %w", context.Canceled), FaultCanceled}, + {fmt.Errorf("wrapped: %w", &net.DNSError{Err: "no such host", Name: "fixture.invalid"}), FaultLocalNetwork}, + {&net.OpError{Op: "dial", Net: "tcp", Err: syscall.ENETUNREACH}, FaultLocalNetwork}, + {x509.UnknownAuthorityError{}, FaultTLS}, + {io.ErrUnexpectedEOF, FaultConnection}, + {fmt.Errorf("fixture unknown error"), FaultUnknown}, + } { + if got := ClassifyTransport(tc.err); got.Kind != tc.kind { + t.Fatalf("kind=%s want %s", got.Kind, tc.kind) + } + } +} diff --git a/internal/workbuddy/governor.go b/internal/workbuddy/governor.go new file mode 100644 index 00000000000..cf1b9a40e0b --- /dev/null +++ b/internal/workbuddy/governor.go @@ -0,0 +1,481 @@ +package workbuddy + +import ( + "crypto/sha256" + "encoding/hex" + "encoding/json" + "errors" + "fmt" + "math" + "os" + "sync" + "time" +) + +// Limits are enforced on observed usage. Concurrent requests can finish beyond +// a configured budget; no estimated token or credit amount is stored as fact. +type Limits struct { + ReserveCredits float64 + DailyTokens int64 + DailyCredits float64 + ModelTokens map[string]int64 +} + +type GovernorOptions struct { + Now func() time.Time + Location *time.Location + Limits Limits + Realms map[Realm]Limits +} + +type modelRisk struct { + CooldownUntil time.Time `json:"-"` + CapUntil time.Time `json:"cap_until,omitempty"` +} + +type accountRisk struct { + Version uint64 `json:"-"` + RestrictedUntil time.Time `json:"restricted_until,omitempty"` + BalanceUntil time.Time `json:"balance_until,omitempty"` + CooldownUntil time.Time `json:"-"` + BreakerUntil time.Time `json:"-"` + DegradeUntil time.Time `json:"-"` + LastFailure time.Time `json:"-"` + LastSoft time.Time `json:"-"` + Failures int `json:"-"` + UnknownFailures int `json:"-"` + SoftStreak int `json:"-"` + UnscopedStreak int `json:"-"` + Models map[string]*modelRisk `json:"models"` + Day string `json:"day"` + Tokens int64 `json:"tokens"` + CreditsSpent float64 `json:"credits_spent"` + ModelTokens map[string]int64 `json:"model_tokens"` + CreditsRemain *float64 `json:"credits_remain,omitempty"` + BalanceStartedAt time.Time `json:"balance_started_at,omitempty"` +} + +type realmRisk struct { + HoldUntil time.Time `json:"hold_until,omitempty"` + Events map[string]time.Time `json:"events"` +} + +type governorState struct { + Version int `json:"version"` + Accounts map[string]*accountRisk `json:"accounts"` + Realms map[Realm]*realmRisk `json:"realms"` +} + +type Governor struct { + mu sync.Mutex + path string + options GovernorOptions + state governorState + saveErr error +} + +type Gate struct { + Reason string + Until time.Time + Stop bool + Version uint64 +} + +func (g Gate) Allowed() bool { return g.Reason == "" } + +// Attempt records one credential's state at dispatch. A later success may not +// clear restrictions or cooldowns raised by a different in-flight request. +type Attempt struct { + Account string + Model string + Version uint64 + StartedAt time.Time +} + +func OpenGovernor(path string, options GovernorOptions) (*Governor, error) { + if path == "" { + return nil, fmt.Errorf("workbuddy: risk state path is required") + } + if options.Now == nil { + options.Now = time.Now + } + if options.Location == nil { + location, errLocation := time.LoadLocation("Asia/Shanghai") + if errLocation != nil { + return nil, errLocation + } + options.Location = location + } + validateLimits := func(limits Limits) error { + if !validCredit(limits.ReserveCredits) || limits.DailyTokens < 0 || !validCredit(limits.DailyCredits) { + return fmt.Errorf("workbuddy: limits must be finite and nonnegative") + } + for _, limit := range limits.ModelTokens { + if limit < 0 { + return fmt.Errorf("workbuddy: model limits cannot be negative") + } + } + return nil + } + if err := validateLimits(options.Limits); err != nil { + return nil, err + } + for realm, limits := range options.Realms { + if realm != China && realm != International { + return nil, fmt.Errorf("workbuddy: invalid limit region") + } + if err := validateLimits(limits); err != nil { + return nil, err + } + } + g := &Governor{path: path, options: options, state: governorState{Version: 1, Accounts: make(map[string]*accountRisk), Realms: make(map[Realm]*realmRisk)}} + raw, errRead := os.ReadFile(path) + if errRead == nil { + if errDecode := json.Unmarshal(raw, &g.state); errDecode != nil || g.state.Version != 1 || g.state.Accounts == nil || g.state.Realms == nil { + return nil, fmt.Errorf("workbuddy: invalid risk state; dispatch remains disabled") + } + for _, state := range g.state.Accounts { + if state == nil || state.Models == nil || state.ModelTokens == nil { + return nil, fmt.Errorf("workbuddy: incomplete account risk state") + } + if state.Tokens < 0 || !validCredit(state.CreditsSpent) || (state.CreditsRemain != nil && !validCredit(*state.CreditsRemain)) { + return nil, fmt.Errorf("workbuddy: invalid persisted usage") + } + for _, tokens := range state.ModelTokens { + if tokens < 0 { + return nil, fmt.Errorf("workbuddy: invalid persisted model usage") + } + } + for _, model := range state.Models { + if model == nil { + return nil, fmt.Errorf("workbuddy: incomplete model risk state") + } + } + } + for realm, state := range g.state.Realms { + if (realm != China && realm != International) || state == nil || state.Events == nil { + return nil, fmt.Errorf("workbuddy: incomplete region risk state") + } + } + } else if !errors.Is(errRead, os.ErrNotExist) { + return nil, fmt.Errorf("workbuddy: read risk state: %w", errRead) + } + if err := g.persistLocked(); err != nil { + return nil, err + } + return g, nil +} + +func RiskAccountKey(c Credential) string { + sum := sha256.Sum256([]byte(string(c.Realm) + "\x00" + c.UID)) + return hex.EncodeToString(sum[:]) +} + +func (g *Governor) accountLocked(c Credential) *accountRisk { + key := RiskAccountKey(c) + state := g.state.Accounts[key] + if state == nil { + state = &accountRisk{Models: make(map[string]*modelRisk), ModelTokens: make(map[string]int64)} + g.state.Accounts[key] = state + } + day := g.options.Now().In(g.options.Location).Format("2006-01-02") + if state.Day != day { + state.Day = day + state.Tokens, state.CreditsSpent = 0, 0 + clear(state.ModelTokens) + } + return state +} + +func (g *Governor) persistLocked() error { + if err := WriteJSON(g.path, g.state); err != nil { + g.saveErr = err + return fmt.Errorf("workbuddy: persist risk state; dispatch suspended: %w", err) + } + g.saveErr = nil + return nil +} + +func (g *Governor) Flush() error { + g.mu.Lock() + defer g.mu.Unlock() + return g.persistLocked() +} + +// Check is shared by chat, probes and activity jobs. A balance-only request may +// bypass balance/budget gates to observe recovery, but never account or region +// restrictions, ordinary cooldowns or circuit breakers. +func (g *Governor) Check(c Credential, model string, paid, balanceOnly bool) (Gate, Attempt) { + g.mu.Lock() + defer g.mu.Unlock() + now := g.options.Now() + if err := c.Validate(); err != nil { + return Gate{Reason: "invalid_credential", Stop: true}, Attempt{} + } + state := g.accountLocked(c) + attempt := Attempt{Account: RiskAccountKey(c), Model: model, Version: state.Version, StartedAt: now} + gate := Gate{Version: state.Version} + if g.saveErr != nil { + gate.Reason, gate.Stop = "state_persistence_failed", true + return gate, attempt + } + if c.Disabled || c.AccessToken == "" { + gate.Reason = "credential_disabled" + return gate, attempt + } + if c.ExpiresAt > 0 && c.ExpiresAt <= now.Unix() { + gate.Reason = "credential_expired" + return gate, attempt + } + if region := g.state.Realms[c.Realm]; region != nil && now.Before(region.HoldUntil) { + gate.Reason, gate.Until, gate.Stop = "region_restricted", region.HoldUntil, true + return gate, attempt + } + for _, blocked := range []struct { + reason string + until time.Time + }{ + {"account_restricted", state.RestrictedUntil}, {"account_cooldown", state.CooldownUntil}, + {"account_breaker", state.BreakerUntil}, {"account_degraded", state.DegradeUntil}, + } { + if now.Before(blocked.until) && blocked.until.After(gate.Until) { + gate.Reason, gate.Until = blocked.reason, blocked.until + } + } + if gate.Reason != "" { + return gate, attempt + } + if balanceOnly { + return gate, attempt + } + if now.Before(state.BalanceUntil) { + gate.Reason, gate.Until = "balance_empty", state.BalanceUntil + return gate, attempt + } + if modelState := state.Models[model]; modelState != nil { + until := later(modelState.CapUntil, modelState.CooldownUntil) + if now.Before(until) { + gate.Reason, gate.Until = "model_cooldown", until + return gate, attempt + } + } + limits := g.options.Limits + if override, exists := g.options.Realms[c.Realm]; exists { + limits = override + } + if limits.ReserveCredits > 0 { + if state.CreditsRemain == nil { + gate.Reason = "balance_unknown" + return gate, attempt + } + if *state.CreditsRemain <= limits.ReserveCredits { + gate.Reason = "credit_reserve" + return gate, attempt + } + } + switch { + case limits.DailyTokens > 0 && state.Tokens >= limits.DailyTokens: + gate.Reason = "daily_tokens" + case paid && limits.DailyCredits > 0 && state.CreditsSpent >= limits.DailyCredits: + gate.Reason = "daily_credits" + case limits.ModelTokens[model] > 0 && state.ModelTokens[model] >= limits.ModelTokens[model]: + gate.Reason = "model_daily_tokens" + } + if gate.Reason != "" { + gate.Until = nextLocalHour(now, g.options.Location, 0) + } + return gate, attempt +} + +func nextLocalHour(now time.Time, location *time.Location, hour int) time.Time { + local := now.In(location) + next := time.Date(local.Year(), local.Month(), local.Day(), hour, 0, 0, 0, location) + if !next.After(now) { + next = next.AddDate(0, 0, 1) + } + return next +} + +func later(a, b time.Time) time.Time { + if b.After(a) { + return b + } + return a +} + +func backoff(count, threshold int, base, maximum time.Duration) time.Duration { + steps := min(max(count-threshold, 0), 20) + return min(base*time.Duration(1< 15*time.Minute { + delete(region.Events, uid) + } + } + region.Events[RiskAccountKey(c)] = now + if len(region.Events) >= 2 { + region.HoldUntil = later(region.HoldUntil, now.Add(30*time.Minute)) + } + case FaultBalance: + state.BalanceUntil = later(state.BalanceUntil, nextLocalHour(now, g.options.Location, 4)) + gate.Until = state.BalanceUntil + case FaultAccountRate: + g.decaySoftLocked(state, now) + state.SoftStreak++ + state.CooldownUntil = later(state.CooldownUntil, now.Add(backoff(state.SoftStreak, 1, 10*time.Minute, 2*time.Hour))) + gate.Until = state.CooldownUntil + case FaultModelRate, FaultModelCap: + wait := time.Minute + if fault.ResetAt.IsZero() && fault.Soft { + g.decaySoftLocked(state, now) + state.UnscopedStreak++ + wait = backoff(state.UnscopedStreak, 1, time.Minute, 10*time.Minute) + } + until := now.Add(wait) + if !fault.ResetAt.IsZero() { + until = later(now.Add(time.Second), fault.ResetAt) + } + modelState := state.Models[model] + if modelState == nil { + modelState = &modelRisk{} + state.Models[model] = modelState + } + modelState.CooldownUntil = later(modelState.CooldownUntil, until) + if fault.Kind == FaultModelCap && options.FreeModel { + modelState.CapUntil = later(modelState.CapUntil, until) + } + gate.Until = until + case FaultHard, FaultGateway, FaultUnknown, FaultConnection: + if !state.LastFailure.IsZero() && now.Sub(state.LastFailure) > 10*time.Minute { + state.Failures, state.UnknownFailures = 0, 0 + } + state.LastFailure = now + state.Failures++ + if fault.Kind == FaultUnknown || fault.Kind == FaultConnection { + state.UnknownFailures++ + } + if state.Failures >= 3 { + state.BreakerUntil = later(state.BreakerUntil, now.Add(backoff(state.Failures, 3, 30*time.Minute, 6*time.Hour))) + } + if state.UnknownFailures >= 5 { + state.DegradeUntil = later(state.DegradeUntil, now.Add(backoff(state.UnknownFailures, 5, 10*time.Minute, 2*time.Hour))) + } + gate.Until = later(state.BreakerUntil, state.DegradeUntil) + if fault.Kind == FaultUnknown { + wait := time.Minute + if options.SingleAccount { + wait = 3 * time.Second + } + state.CooldownUntil = later(state.CooldownUntil, now.Add(wait)) + gate.Until = later(gate.Until, state.CooldownUntil) + } + case FaultAuth: + // Called only after the request's one refresh attempt has failed. + wait := time.Minute + if options.SingleAccount { + wait = 3 * time.Second + } + state.CooldownUntil = later(state.CooldownUntil, now.Add(wait)) + gate.Until = state.CooldownUntil + default: + changed = false + gate.Stop = fault.Kind != FaultLocalNetwork + } + if changed { + state.Version++ + if err := g.persistLocked(); err != nil { + return Gate{Reason: "state_persistence_failed", Stop: true}, err + } + } + gate.Version = state.Version + return gate, nil +} + +func (g *Governor) decaySoftLocked(state *accountRisk, now time.Time) { + if !state.LastSoft.IsZero() && now.Sub(state.LastSoft) > 10*time.Minute { + state.SoftStreak, state.UnscopedStreak = 0, 0 + } + state.LastSoft = now +} + +// Complete records supplied observed usage even if another request raised a +// newer penalty. A success clears penalties only when its dispatch version +// still matches. A cap survives success until its own reset time. +func (g *Governor) Complete(c Credential, attempt Attempt, tokens int64, credits *float64, success bool) error { + g.mu.Lock() + defer g.mu.Unlock() + if attempt.Account != RiskAccountKey(c) || tokens < 0 || (credits != nil && !validCredit(*credits)) { + return fmt.Errorf("workbuddy: invalid completion accounting") + } + state := g.accountLocked(c) + if tokens > math.MaxInt64-state.Tokens || tokens > math.MaxInt64-state.ModelTokens[attempt.Model] || + (credits != nil && !validCredit(state.CreditsSpent+*credits)) { + return fmt.Errorf("workbuddy: usage counter overflow") + } + state.Tokens += tokens + state.ModelTokens[attempt.Model] += tokens + if credits != nil { + state.CreditsSpent += *credits + } + if success && attempt.Version == state.Version { + state.RestrictedUntil, state.CooldownUntil, state.BreakerUntil, state.DegradeUntil = time.Time{}, time.Time{}, time.Time{}, time.Time{} + state.LastFailure, state.LastSoft = time.Time{}, time.Time{} + state.Failures, state.UnknownFailures, state.SoftStreak, state.UnscopedStreak = 0, 0, 0, 0 + if model := state.Models[attempt.Model]; model != nil { + model.CooldownUntil = time.Time{} + } + state.Version++ + } + return g.persistLocked() +} + +func (g *Governor) ObserveBalance(c Credential, attempt Attempt, remain float64) error { + g.mu.Lock() + defer g.mu.Unlock() + if attempt.Account != RiskAccountKey(c) || !validCredit(remain) { + return fmt.Errorf("workbuddy: invalid balance observation") + } + state := g.accountLocked(c) + if attempt.Version != state.Version || attempt.StartedAt.Before(state.BalanceStartedAt) { + return nil + } + state.CreditsRemain = &remain + state.BalanceStartedAt = attempt.StartedAt + if remain > 0 && attempt.Version == state.Version { + state.BalanceUntil = time.Time{} + } + return g.persistLocked() +} + +func validCredit(value float64) bool { + return value >= 0 && !math.IsNaN(value) && !math.IsInf(value, 0) +} diff --git a/internal/workbuddy/governor_test.go b/internal/workbuddy/governor_test.go new file mode 100644 index 00000000000..8943316d20b --- /dev/null +++ b/internal/workbuddy/governor_test.go @@ -0,0 +1,244 @@ +package workbuddy + +import ( + "os" + "path/filepath" + "strings" + "testing" + "time" +) + +func governorFixture(t *testing.T, limits Limits) (*Governor, *time.Time, string) { + t.Helper() + now := time.Date(2026, 10, 21, 1, 0, 0, 0, time.UTC) + path := filepath.Join(t.TempDir(), "risk.json") + g, err := OpenGovernor(path, GovernorOptions{Now: func() time.Time { return now }, Location: time.FixedZone("CST", 8*3600), Limits: limits}) + if err != nil { + t.Fatal(err) + } + return g, &now, path +} + +func TestGovernorRestrictionWaveAndStaleSuccessSurviveRestart(t *testing.T) { + g, now, path := governorFixture(t, Limits{}) + first := testCredential() + second := first + second.UID = "second-user" + _, oldAttempt := g.Check(first, "model", false, false) + fault := Fault{Kind: FaultRestricted, Status: 403, Code: 11140} + result, err := g.RecordFailure(first, "model", fault, FailureOptions{}) + if err != nil || !result.Stop || result.Until.Sub(*now) != 2*time.Hour { + t.Fatalf("restriction result=%+v err=%v", result, err) + } + if err := g.Complete(first, oldAttempt, 10, nil, true); err != nil { + t.Fatal(err) + } + if gate, _ := g.Check(first, "model", false, false); gate.Reason != "account_restricted" { + t.Fatal("stale success cleared restriction") + } + // The same UID in another product cannot create a second regional event. + first.Product = VSCode + if _, err := g.RecordFailure(first, "model", fault, FailureOptions{}); err != nil { + t.Fatal(err) + } + if gate, _ := g.Check(second, "model", false, false); !gate.Allowed() { + t.Fatal("same account was counted twice") + } + if _, err := g.RecordFailure(second, "model", fault, FailureOptions{}); err != nil { + t.Fatal(err) + } + third := first + third.UID = "third-user" + if gate, _ := g.Check(third, "model", false, true); gate.Reason != "region_restricted" || !gate.Stop { + t.Fatal("background balance request bypassed region hold") + } + restarted, errOpen := OpenGovernor(path, g.options) + if errOpen != nil { + t.Fatal(errOpen) + } + third.AccessToken = "fresh-token" + if gate, _ := restarted.Check(third, "model", false, false); gate.Reason != "region_restricted" { + t.Fatal("restart cleared region hold") + } + *now = now.Add(31 * time.Minute) + if gate, _ := restarted.Check(third, "model", false, false); !gate.Allowed() { + t.Fatal("regional hold did not expire") + } + if gate, _ := restarted.Check(first, "model", false, false); gate.Reason != "account_restricted" { + t.Fatal("identity switch cleared account restriction") + } + *now = now.Add(90 * time.Minute) + if gate, _ := restarted.Check(first, "model", false, false); !gate.Allowed() { + t.Fatal("account restriction did not expire") + } + raw, errRead := os.ReadFile(path) + if errRead != nil { + t.Fatal(errRead) + } + if strings.Contains(string(raw), "fixture-user") || strings.Contains(string(raw), "old-refresh") { + t.Fatal("risk state contains raw identity or credentials") + } +} + +func TestGovernor429ScopesCountersDecayAndPersistentCaps(t *testing.T) { + g, now, path := governorFixture(t, Limits{}) + c := testCredential() + for index := 0; index < 3; index++ { + gate, err := g.RecordFailure(c, "limited", Fault{Kind: FaultModelRate, Soft: true}, FailureOptions{}) + if err != nil { + t.Fatal(err) + } + if want := time.Minute * time.Duration(1< Date: Wed, 21 Oct 2026 02:20:38 +0800 Subject: [PATCH 03/10] feat(workbuddy): integrate governed native chat execution --- .github/workflows/workbuddy-checks.yml | 4 + docs/workbuddy-native.md | 56 + internal/config/config.go | 2 + internal/config/config_v8.go | 1 + internal/config/workbuddy.go | 20 + internal/config/workbuddy_test.go | 38 + .../executor/helps/workbuddy_payload.go | 184 +++ .../executor/helps/workbuddy_payload_test.go | 42 + .../executor/helps/workbuddy_runtime.go | 213 ++++ .../runtime/executor/workbuddy_executor.go | 431 +++++++ .../executor/workbuddy_executor_test.go | 464 ++++++++ internal/workbuddy/catalog.go | 60 + internal/workbuddy/catalog.json | 1036 +++++++++++++++++ internal/workbuddy/dispatch.go | 176 +++ internal/workbuddy/dispatch_test.go | 100 ++ internal/workbuddy/governor.go | 38 + internal/workbuddy/stream.go | 335 ++++++ internal/workbuddy/stream_test.go | 80 ++ sdk/cliproxy/auth/conductor_execution.go | 33 + sdk/cliproxy/auth/conductor_refresh.go | 3 + .../auth/conductor_request_scoped_errors.go | 16 + sdk/cliproxy/auth/conductor_workbuddy_test.go | 103 ++ sdk/cliproxy/service_executors.go | 13 + sdk/cliproxy/service_models.go | 32 + sdk/cliproxy/service_workbuddy_test.go | 46 + 25 files changed, 3526 insertions(+) create mode 100644 internal/config/workbuddy.go create mode 100644 internal/config/workbuddy_test.go create mode 100644 internal/runtime/executor/helps/workbuddy_payload.go create mode 100644 internal/runtime/executor/helps/workbuddy_payload_test.go create mode 100644 internal/runtime/executor/helps/workbuddy_runtime.go create mode 100644 internal/runtime/executor/workbuddy_executor.go create mode 100644 internal/runtime/executor/workbuddy_executor_test.go create mode 100644 internal/workbuddy/catalog.go create mode 100644 internal/workbuddy/catalog.json create mode 100644 internal/workbuddy/dispatch.go create mode 100644 internal/workbuddy/dispatch_test.go create mode 100644 internal/workbuddy/stream.go create mode 100644 internal/workbuddy/stream_test.go create mode 100644 sdk/cliproxy/auth/conductor_workbuddy_test.go create mode 100644 sdk/cliproxy/service_workbuddy_test.go diff --git a/.github/workflows/workbuddy-checks.yml b/.github/workflows/workbuddy-checks.yml index f4b3fd7712b..058fa1f8e44 100644 --- a/.github/workflows/workbuddy-checks.yml +++ b/.github/workflows/workbuddy-checks.yml @@ -29,3 +29,7 @@ jobs: env: CGO_ENABLED: '1' run: go test -race -timeout 120s ./internal/workbuddy ./sdk/auth ./internal/watcher/synthesizer + - name: Test native executor and SDK retry integration under the race detector + env: + CGO_ENABLED: '1' + run: go test -race -timeout 180s ./internal/runtime/executor ./internal/runtime/executor/helps ./sdk/cliproxy/auth ./sdk/cliproxy -run TestWorkBuddy diff --git a/docs/workbuddy-native.md b/docs/workbuddy-native.md index 41a03870f76..150afe88285 100644 --- a/docs/workbuddy-native.md +++ b/docs/workbuddy-native.md @@ -44,6 +44,60 @@ upstream acceptance, current account entitlement, award amounts or actual anti-abuse behavior. Those checks remain pending until a test account is provided. +## Implementation status + +The credential/transport foundation is merged. Core execution is being developed +in `feat/workbuddy-proxy`; deployment is still pending. Offline coverage now +includes the native executor, OpenAI Chat/Responses, Claude and Gemini response +translation, SSE/nonstream collection, parallel tool calls, usage accounting, +401 refresh with persistence before replay, and SDK-scoped retry decisions. +HTTP errors and error envelopes inside HTTP 200 streams share the risk classifier. +A truncated stream is an error; once a chunk was emitted it cannot be replayed +on another credential. Only a pre-output connection failure gets one retry. + +The static model catalog is an offline snapshot. It is not evidence of live +model availability. Native credential `models` can explicitly restrict the +published models, including synthetic models used by offline tests. + +The governor and transport pools are shared across executor replacements. +Configuration reloads preserve cooldowns, restrictions and usage. Changing the +state directory requires a restart; the old directory must be migrated explicitly. +Only one process may own a state directory. The isolated deployment will use a +single replica and its own PVC. + +The remaining stage-2 work includes live model discovery, conversation paging, +affinity policy, content-rejection degradation, region/model entitlement holds, +optional exit backoff, identity switching and paid fallback. Stage-3 quota and +activity jobs are not implemented or enabled yet. These gaps must be closed +before claiming parity with the reference project's risk controls. + +An explicit v8 configuration enables the provider: + +```yaml +oauth: + auth-dir: /data/workbuddy-auth + providers: + workbuddy: + enabled: true + state-dir: /data/workbuddy-state + max-concurrency: 32 + per-caller-concurrency: 2 + queue-seconds: 30 + limits: + daily-tokens: 0 + daily-credits: 0 + reserve-credits: 0 +``` + +Zero budgets disable that local limit. Regional limit blocks replace the global +budget block for that region. A positive credit reserve requires an observed +balance, so requests remain gated while balance is unknown. Tokens and credits +are recorded from upstream observations only; credits are never inferred from +token counts. Concurrent requests can complete beyond a configured budget. +Queue deadlines end at admission and never become upstream network deadlines. +The general management HTTP tunnel is unsupported for this provider; arbitrary +requests must not bypass its risk governor or export credentials to another host. + ## Credential setup Use a separate configuration and auth directory for this test service. Never @@ -87,6 +141,8 @@ contact WorkBuddy or inspect any existing account files. ```sh go test ./internal/workbuddy ./sdk/auth ./internal/watcher/synthesizer \ ./internal/cmd ./cmd/server +go test ./internal/runtime/executor ./internal/runtime/executor/helps \ + ./sdk/cliproxy/auth ./sdk/cliproxy ./internal/config -run TestWorkBuddy go build -o test-output ./cmd/server rm test-output ``` diff --git a/internal/config/config.go b/internal/config/config.go index cbe344933f3..de474defd74 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -125,6 +125,8 @@ type Config struct { // Antigravity configures provider-wide Antigravity request behavior. Antigravity AntigravityConfig `yaml:"antigravity" json:"antigravity"` + WorkBuddy WorkBuddyConfig `yaml:"workbuddy" json:"workbuddy"` + // Devin configures provider-wide Devin request behavior. Devin DevinConfig `yaml:"devin" json:"devin"` diff --git a/internal/config/config_v8.go b/internal/config/config_v8.go index 68bd7d98d19..cd10df183fb 100644 --- a/internal/config/config_v8.go +++ b/internal/config/config_v8.go @@ -99,6 +99,7 @@ func buildV8Paths() []configPath { {"antigravity-signature-cache-enabled", "oauth.providers.antigravity.signature-cache-enabled"}, {"antigravity-signature-bypass-strict", "oauth.providers.antigravity.signature-bypass-strict"}, {"quota-exceeded.antigravity-credits", "oauth.providers.antigravity.antigravity-credits"}, + {"workbuddy", "oauth.providers.workbuddy"}, {"xai", "upstream.xai"}, {"devin", "oauth.providers.devin"}, {"disable-image-generation", "multimedia.disable-image-generation"}, {"gpt-image-2-base-model", "multimedia.gpt-image-2-base-model"}, {"video-result-auth-cache-ttl", "multimedia.video-result-auth-cache-ttl"}, diff --git a/internal/config/workbuddy.go b/internal/config/workbuddy.go new file mode 100644 index 00000000000..cb6b9013960 --- /dev/null +++ b/internal/config/workbuddy.go @@ -0,0 +1,20 @@ +package config + +// WorkBuddyConfig is opt-in. State belongs to a dedicated provider volume and +// must never be shared between independent server processes. +type WorkBuddyConfig struct { + Enabled bool `yaml:"enabled" json:"enabled"` + StateDir string `yaml:"state-dir" json:"state-dir"` + MaxConcurrency int `yaml:"max-concurrency" json:"max-concurrency"` + PerCallerConcurrency int `yaml:"per-caller-concurrency" json:"per-caller-concurrency"` + QueueSeconds int `yaml:"queue-seconds" json:"queue-seconds"` + Limits WorkBuddyLimits `yaml:"limits" json:"limits"` + RealmLimits map[string]WorkBuddyLimits `yaml:"realm-limits" json:"realm-limits"` +} + +type WorkBuddyLimits struct { + ReserveCredits float64 `yaml:"reserve-credits" json:"reserve-credits"` + DailyTokens int64 `yaml:"daily-tokens" json:"daily-tokens"` + DailyCredits float64 `yaml:"daily-credits" json:"daily-credits"` + ModelTokens map[string]int64 `yaml:"model-tokens" json:"model-tokens"` +} diff --git a/internal/config/workbuddy_test.go b/internal/config/workbuddy_test.go new file mode 100644 index 00000000000..0a6c4844a2a --- /dev/null +++ b/internal/config/workbuddy_test.go @@ -0,0 +1,38 @@ +package config + +import "testing" + +func TestWorkBuddyV8Configuration(t *testing.T) { + cfg, err := ParseConfigBytes([]byte(`version: 8 +oauth: + auth-dir: /private/workbuddy-auth + providers: + workbuddy: + enabled: true + state-dir: /private/workbuddy-state + max-concurrency: 8 + per-caller-concurrency: 2 + queue-seconds: 15 + limits: + reserve-credits: 10 + daily-tokens: 10000 + model-tokens: + fixture-model: 5000 + realm-limits: + cn: + daily-tokens: 2000 +`)) + if err != nil { + t.Fatal(err) + } + if !cfg.WorkBuddy.Enabled || cfg.WorkBuddy.StateDir != "/private/workbuddy-state" || cfg.WorkBuddy.PerCallerConcurrency != 2 || cfg.WorkBuddy.Limits.ModelTokens["fixture-model"] != 5000 || cfg.WorkBuddy.RealmLimits["cn"].DailyTokens != 2000 { + t.Fatalf("configuration not mapped: %+v", cfg.WorkBuddy) + } + defaultCfg, err := ParseConfigBytes([]byte("version: 8\n")) + if err != nil { + t.Fatal(err) + } + if defaultCfg.WorkBuddy.Enabled { + t.Fatal("provider enabled implicitly") + } +} diff --git a/internal/runtime/executor/helps/workbuddy_payload.go b/internal/runtime/executor/helps/workbuddy_payload.go new file mode 100644 index 00000000000..d392c92840c --- /dev/null +++ b/internal/runtime/executor/helps/workbuddy_payload.go @@ -0,0 +1,184 @@ +package helps + +import ( + "bytes" + "encoding/json" + "fmt" + "strings" + + "github.com/router-for-me/CLIProxyAPI/v8/internal/workbuddy" +) + +const workBuddySystemPrompt = "You are a helpful coding assistant." + +// NormalizeWorkBuddyPayload runs before the payload finalizer. It intentionally +// preserves caller text, tool schemas and unknown public fields. +func NormalizeWorkBuddyPayload(raw []byte, c workbuddy.Credential, model string) ([]byte, error) { + if len(raw) > 16<<20 { + return nil, &workbuddy.DispatchError{Reason: "request_body_too_large", Status: 413, Stop: true} + } + var body map[string]any + decoder := json.NewDecoder(bytes.NewReader(raw)) + decoder.UseNumber() + if err := decoder.Decode(&body); err != nil || body == nil { + return nil, fmt.Errorf("workbuddy: invalid chat payload") + } + messages, ok := body["messages"].([]any) + if !ok { + return nil, fmt.Errorf("workbuddy: messages must be an array") + } + deepseek := strings.HasPrefix(strings.ToLower(model), "deepseek") + thinking, _ := body["thinking"].(map[string]any) + effort, _ := body["reasoning_effort"].(string) + if effort == "" { + effort, _ = body["reasoningEffort"].(string) + } + enabled := deepseek && thinking["type"] != "disabled" && effort != "none" + for _, item := range messages { + message, ok := item.(map[string]any) + if !ok { + return nil, fmt.Errorf("workbuddy: invalid message") + } + if message["role"] == "developer" { + message["role"] = "system" + } + if enabled && message["role"] == "assistant" { + if _, ok := message["reasoning_content"]; !ok { + message["reasoning_content"] = "" + } + } + } + if len(messages) == 0 || messages[0].(map[string]any)["role"] != "system" { + messages = append([]any{map[string]any{"role": "system", "content": workBuddySystemPrompt}}, messages...) + } + body["messages"] = repairWorkBuddyTools(messages) + body["model"] = model + for key := range body { + if strings.HasPrefix(key, "_") { + delete(body, key) + } + } + if value, ok := body["max_completion_tokens"]; ok { + if _, exists := body["max_tokens"]; !exists { + body["max_tokens"] = value + } + delete(body, "max_completion_tokens") + } + if value, ok := body["max_tokens"].(json.Number); ok && strings.HasPrefix(model, "gpt-") { + if n, err := value.Int64(); err == nil && n < 16 { + body["max_tokens"] = 16 + } + } + if _, ok := body["max_tokens"]; !ok { + if meta, found := workbuddy.ModelFor(c.Realm, model); found && meta.Output > 0 { + body["max_tokens"] = meta.Output + } + } + if enabled { + if _, ok := body["thinking"]; !ok { + body["thinking"] = map[string]any{"type": "enabled"} + } + if effort == "" { + body["reasoning_effort"] = "high" + } + } + if stop, ok := body["stop"].(string); ok { + body["stop"] = []string{stop} + } + body["stream"] = true + if _, ok := body["stream_options"]; !ok { + body["stream_options"] = map[string]any{"include_usage": true} + } + if c.Product == workbuddy.CLI { + for _, key := range []string{"metadata", "conversation_id", "previous_response_id"} { + delete(body, key) + } + body["cache_retention"] = "long" + } + var compact bytes.Buffer + encoder := json.NewEncoder(&compact) + encoder.SetEscapeHTML(false) + if err := encoder.Encode(body); err != nil { + return nil, err + } + return bytes.TrimSuffix(compact.Bytes(), []byte("\n")), nil +} + +// repairWorkBuddyTools joins split adjacent call declarations, moves available +// results next to their declarations, and removes incomplete call/result pairs. +// Intervening ordinary messages retain their relative order and text. +func repairWorkBuddyTools(messages []any) []any { + merged := make([]map[string]any, 0, len(messages)) + for _, item := range messages { + m := item.(map[string]any) + calls, _ := m["tool_calls"].([]any) + if len(merged) > 0 && m["role"] == "assistant" && len(calls) > 0 { + prev := merged[len(merged)-1] + previous, _ := prev["tool_calls"].([]any) + // Content-bearing assistant turns remain independent. + if prev["role"] == "assistant" && len(previous) > 0 && (m["content"] == nil || m["content"] == "") { + prev["tool_calls"] = append(previous, calls...) + continue + } + } + merged = append(merged, m) + } + consumed := make(map[int]bool) + out := make([]any, 0, len(merged)) + for i, m := range merged { + if consumed[i] || m["role"] == "tool" { + continue + } + calls, _ := m["tool_calls"].([]any) + if len(calls) == 0 { + out = append(out, m) + continue + } + kept := make([]any, 0, len(calls)) + results := make([]any, 0, len(calls)) + usedIDs := make(map[string]bool) + for _, raw := range calls { + call, ok := raw.(map[string]any) + if !ok { + continue + } + id, _ := call["id"].(string) + if id == "" || usedIDs[id] { + continue + } + usedIDs[id] = true + for j := i + 1; j < len(merged); j++ { + candidate := merged[j] + // A new declaration for the same ID bounds this call's result scope. + if nextCalls, ok := candidate["tool_calls"].([]any); ok { + repeated := false + for _, next := range nextCalls { + if nextMap, ok := next.(map[string]any); ok && nextMap["id"] == id { + repeated = true + break + } + } + if repeated { + break + } + } + if !consumed[j] && candidate["role"] == "tool" && candidate["tool_call_id"] == id { + kept = append(kept, call) + results = append(results, candidate) + consumed[j] = true + break + } + } + } + if len(kept) > 0 { + m["tool_calls"] = kept + } else { + delete(m, "tool_calls") + } + if len(kept) > 0 || m["content"] != nil && m["content"] != "" { + out = append(out, m) + out = append(out, results...) + } + } + return out +} diff --git a/internal/runtime/executor/helps/workbuddy_payload_test.go b/internal/runtime/executor/helps/workbuddy_payload_test.go new file mode 100644 index 00000000000..d1643439d78 --- /dev/null +++ b/internal/runtime/executor/helps/workbuddy_payload_test.go @@ -0,0 +1,42 @@ +package helps + +import ( + "testing" + + "github.com/router-for-me/CLIProxyAPI/v8/internal/workbuddy" + "github.com/tidwall/gjson" +) + +func TestWorkBuddyToolRepairAndThinking(t *testing.T) { + raw := []byte(`{"messages":[{"role":"developer","content":"keep client instructions"},{"role":"assistant","tool_calls":[{"id":"a","type":"function","function":{"name":"one","arguments":"{}"}}]},{"role":"assistant","tool_calls":[{"id":"b","type":"function","function":{"name":"two","arguments":"{}"}},{"id":"missing","type":"function","function":{"name":"three","arguments":"{}"}}]},{"role":"user","content":"intervening"},{"role":"tool","tool_call_id":"b","content":"B"},{"role":"tool","tool_call_id":"a","content":"A"},{"role":"tool","tool_call_id":"orphan","content":"drop"}],"_private_marker":true,"max_completion_tokens":50}`) + body, err := NormalizeWorkBuddyPayload(raw, workbuddy.Credential{Realm: workbuddy.International, Product: workbuddy.CLI}, "deepseek-v4.1-flash") + if err != nil { + t.Fatal(err) + } + for path, want := range map[string]string{"messages.#": "5", "messages.0.role": "system", "messages.0.content": "keep client instructions", "messages.1.tool_calls.#": "2", "messages.2.tool_call_id": "a", "messages.3.tool_call_id": "b", "messages.4.content": "intervening", "thinking.type": "enabled", "reasoning_effort": "high", "max_tokens": "50", "cache_retention": "long"} { + if got := gjson.GetBytes(body, path).String(); got != want { + t.Errorf("%s=%q want %q; body=%s", path, got, want, body) + } + } + if !gjson.GetBytes(body, "messages.1.reasoning_content").Exists() { + t.Error("missing reasoning backfill") + } + for _, path := range []string{"_private_marker", "max_completion_tokens"} { + if gjson.GetBytes(body, path).Exists() { + t.Errorf("unexpected %s", path) + } + } +} +func TestWorkBuddyPayloadKeepsExplicitReasoningOptOut(t *testing.T) { + raw := []byte(`{"messages":[{"role":"assistant","content":"prior"},{"role":"user","content":"next"}],"thinking":{"type":"disabled"},"tools":[{"type":"function","function":{"name":"tool","parameters":{"type":"object","properties":{"n":{"type":"integer"}}}}}]}`) + body, err := NormalizeWorkBuddyPayload(raw, workbuddy.Credential{Realm: workbuddy.International, Product: workbuddy.Desktop}, "deepseek-v4.1-flash") + if err != nil { + t.Fatal(err) + } + if gjson.GetBytes(body, "reasoning_effort").Exists() || gjson.GetBytes(body, "messages.1.reasoning_content").Exists() { + t.Fatal("thinking opt out overwritten") + } + if gjson.GetBytes(body, "tools.0.function.parameters.properties.n.type").String() != "integer" { + t.Fatal("tool schema modified") + } +} diff --git a/internal/runtime/executor/helps/workbuddy_runtime.go b/internal/runtime/executor/helps/workbuddy_runtime.go new file mode 100644 index 00000000000..d301adefa5d --- /dev/null +++ b/internal/runtime/executor/helps/workbuddy_runtime.go @@ -0,0 +1,213 @@ +package helps + +import ( + "context" + "encoding/json" + "fmt" + "path/filepath" + "sync" + "time" + + "github.com/router-for-me/CLIProxyAPI/v8/internal/config" + "github.com/router-for-me/CLIProxyAPI/v8/internal/workbuddy" + coreauth "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/auth" +) + +// WorkBuddyRuntime is shared by executor replacements and all products/accounts. +// There must be one owner per state directory in a server process. +type WorkBuddyRuntime struct { + Governor *workbuddy.Governor + Transports *workbuddy.Transports + Limiter *workbuddy.Limiter + // BaseURL is an injection seam for offline tests, never populated from auth files. + BaseURL string + StateDir string + RefreshGate workbuddy.RefreshGate + mu sync.Mutex + auths map[string]*coreauth.Auth + manager *coreauth.Manager + Save func(context.Context, *coreauth.Auth, *coreauth.Auth) error +} + +func NewWorkBuddyRuntime(cfg config.WorkBuddyConfig) (*WorkBuddyRuntime, error) { + if cfg.StateDir == "" { + return nil, fmt.Errorf("workbuddy: state-dir is required") + } + r := &WorkBuddyRuntime{Transports: workbuddy.NewTransports(256, nil), StateDir: filepath.Clean(cfg.StateDir), auths: make(map[string]*coreauth.Auth)} + var err error + r.Governor, err = workbuddy.OpenGovernor(filepath.Join(r.StateDir, "risk.json"), workbuddy.GovernorOptions{}) + if err != nil { + return nil, err + } + r.Limiter, err = workbuddy.NewLimiter(cfg.MaxConcurrency, cfg.PerCallerConcurrency) + if err != nil { + return nil, err + } + r.Save = func(context.Context, *coreauth.Auth, *coreauth.Auth) error { + return fmt.Errorf("workbuddy: credential persistence is not configured") + } + if err = r.Configure(cfg); err != nil { + return nil, err + } + return r, nil +} + +// BindManager uses the configured manager store and its concurrency fences. +// It never falls back to a process-global token directory. +func (r *WorkBuddyRuntime) BindManager(manager *coreauth.Manager) { + r.mu.Lock() + defer r.mu.Unlock() + r.manager = manager + r.Save = func(ctx context.Context, original, updated *coreauth.Auth) error { + base, ok := manager.GetByID(updated.ID) + if !ok || base.Disabled { + return fmt.Errorf("workbuddy: credential removed or disabled during refresh") + } + // The manager validates registration/credential versions before persistence. + committed, err := manager.UpdateRefreshedAuth(ctx, original, updated) + if err != nil { + return err + } + if committed == nil || committed.Metadata["access_token"] != updated.Metadata["access_token"] { + return fmt.Errorf("workbuddy: credential changed during refresh") + } + r.Track(committed) + return nil + } +} + +func (r *WorkBuddyRuntime) Configure(cfg config.WorkBuddyConfig) error { + if filepath.Clean(cfg.StateDir) != r.StateDir { + return fmt.Errorf("workbuddy: changing state-dir requires a process restart") + } + if cfg.QueueSeconds < 0 { + return fmt.Errorf("workbuddy: queue-seconds cannot be negative") + } + convert := func(l config.WorkBuddyLimits) workbuddy.Limits { + return workbuddy.Limits{ReserveCredits: l.ReserveCredits, DailyTokens: l.DailyTokens, DailyCredits: l.DailyCredits, ModelTokens: l.ModelTokens} + } + realms := make(map[workbuddy.Realm]workbuddy.Limits) + for k, v := range cfg.RealmLimits { + realms[workbuddy.Realm(k)] = convert(v) + } + if err := r.Governor.UpdateLimits(convert(cfg.Limits), realms); err != nil { + return err + } + return r.Limiter.Configure(cfg.MaxConcurrency, cfg.PerCallerConcurrency) +} +func WorkBuddyCredential(a *coreauth.Auth) (workbuddy.Credential, error) { + if a == nil || a.Provider != "workbuddy" { + return workbuddy.Credential{}, fmt.Errorf("workbuddy: missing credential") + } + raw, err := json.Marshal(a.Metadata) + if err != nil { + return workbuddy.Credential{}, fmt.Errorf("workbuddy: invalid credential metadata") + } + c, err := workbuddy.DecodeCredential(raw) + if err != nil { + return c, err + } + c.Disabled = c.Disabled || a.Disabled + if a.ProxyURL != "" { + c.ProxyURL = a.ProxyURL + } + return c, nil +} +func (r *WorkBuddyRuntime) Track(a *coreauth.Auth) { + if a == nil { + return + } + r.mu.Lock() + defer r.mu.Unlock() + if prev := r.auths[a.ID]; prev != nil { + if prev.RegistrationEpoch > a.RegistrationEpoch || + (prev.RegistrationEpoch == a.RegistrationEpoch && prev.Generation > a.Generation) || + (prev.RegistrationEpoch == a.RegistrationEpoch && prev.Generation == a.Generation && prev.UpdatedAt.After(a.UpdatedAt)) { + return + } + } + r.auths[a.ID] = a.Clone() +} +func (r *WorkBuddyRuntime) Current(a *coreauth.Auth) *coreauth.Auth { + r.mu.Lock() + manager := r.manager + r.mu.Unlock() + if manager != nil && a != nil { + current, ok := manager.GetByID(a.ID) + if !ok { + return nil + } + r.Track(current) + } + r.Track(a) + r.mu.Lock() + defer r.mu.Unlock() + if a != nil && r.auths[a.ID] != nil { + return r.auths[a.ID].Clone() + } + return a +} +func (r *WorkBuddyRuntime) Accounts() int { + r.mu.Lock() + defer r.mu.Unlock() + unique := make(map[string]bool) + for _, a := range r.auths { + c, err := WorkBuddyCredential(a) + if err == nil && !c.Disabled { + unique[workbuddy.RiskAccountKey(c)] = true + } + } + return len(unique) +} +func (r *WorkBuddyRuntime) Refresh(ctx context.Context, a *coreauth.Auth, proxy string) (*coreauth.Auth, error) { + a = r.Current(a) + c, err := WorkBuddyCredential(a) + if err != nil { + return nil, err + } + client := &workbuddy.AuthClient{Transports: r.Transports, Versions: workbuddy.DefaultVersions(), ProxyURL: proxy, BaseURL: r.BaseURL} + var saved, loaded *coreauth.Auth + refreshed, err := r.RefreshGate.Run(ctx, workbuddy.RiskAccountKey(c), c.AccessToken, + func() (workbuddy.Credential, error) { + loaded = r.Current(a) + return WorkBuddyCredential(loaded) + }, + func(next workbuddy.Credential) error { + original := loaded.Clone() + updated := original.Clone() + raw, errMarshal := json.Marshal(next) + if errMarshal != nil { + return errMarshal + } + if errDecode := json.Unmarshal(raw, &updated.Metadata); errDecode != nil { + return errDecode + } + // Zero means no reported expiration; do not retain an older value. + updated.Metadata["expires_at"] = next.ExpiresAt + updated.LastRefreshedAt = time.Now() + updated.UpdatedAt = updated.LastRefreshedAt + r.mu.Lock() + save := r.Save + r.mu.Unlock() + if errSave := save(ctx, original, updated); errSave != nil { + return errSave + } + r.Track(updated) + saved = r.Current(updated) + return nil + }, client) + if err != nil { + return nil, &workbuddy.DispatchError{Reason: "credential_refresh_failed", Status: 401, Stop: true} + } + if saved != nil { + return saved, nil + } + current := r.Current(a) + if current == nil { + return nil, fmt.Errorf("workbuddy: credential removed during refresh") + } + if current.Metadata["access_token"] != refreshed.AccessToken { + return nil, fmt.Errorf("workbuddy: refreshed credential not installed") + } + return current, nil +} diff --git a/internal/runtime/executor/workbuddy_executor.go b/internal/runtime/executor/workbuddy_executor.go new file mode 100644 index 00000000000..f284c12e367 --- /dev/null +++ b/internal/runtime/executor/workbuddy_executor.go @@ -0,0 +1,431 @@ +package executor + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "strings" + "time" + + "github.com/router-for-me/CLIProxyAPI/v8/internal/config" + "github.com/router-for-me/CLIProxyAPI/v8/internal/runtime/executor/helps" + "github.com/router-for-me/CLIProxyAPI/v8/internal/thinking" + "github.com/router-for-me/CLIProxyAPI/v8/internal/workbuddy" + coreauth "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/auth" + execution "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/executor" + sdktranslator "github.com/router-for-me/CLIProxyAPI/v8/sdk/translator" + log "github.com/sirupsen/logrus" + "github.com/tidwall/gjson" +) + +var _ coreauth.ProviderExecutor = (*WorkBuddyExecutor)(nil) +var _ coreauth.RequestAuthPreparer = (*WorkBuddyExecutor)(nil) + +// WorkBuddyExecutor owns no account-specific transport or risk state. Those +// belong to the shared runtime so rotations and reloads cannot bypass gates. +type WorkBuddyExecutor struct { + cfg *config.Config + runtime *helps.WorkBuddyRuntime + initErr error +} + +func NewWorkBuddyExecutor(cfg *config.Config, previous *WorkBuddyExecutor) *WorkBuddyExecutor { + e := &WorkBuddyExecutor{cfg: cfg} + if previous != nil { + e.runtime = previous.runtime + } + if cfg == nil || !cfg.WorkBuddy.Enabled { + e.initErr = fmt.Errorf("workbuddy: provider is disabled") + return e + } + if e.runtime == nil { + e.runtime, e.initErr = helps.NewWorkBuddyRuntime(cfg.WorkBuddy) + } else { + e.initErr = e.runtime.Configure(cfg.WorkBuddy) + } + return e +} +func (e *WorkBuddyExecutor) Identifier() string { return "workbuddy" } +func (e *WorkBuddyExecutor) UsesConfig(cfg *config.Config) bool { return e.cfg == cfg } +func (e *WorkBuddyExecutor) BindCredentialManager(manager *coreauth.Manager) { + if e.runtime != nil { + e.runtime.BindManager(manager) + } +} +func (e *WorkBuddyExecutor) TrackAuth(a *coreauth.Auth) { + if e.runtime != nil { + e.runtime.Track(a) + } +} +func (e *WorkBuddyExecutor) PrepareExecutionContext(ctx context.Context, _ execution.Request, _ execution.Options) (context.Context, error) { + if e.initErr != nil { + return ctx, &workbuddy.DispatchError{Reason: "provider_unavailable", Status: 503, Stop: true} + } + return workbuddy.WithRequestState(ctx, e.runtime.Accounts()) +} +func (e *WorkBuddyExecutor) ShouldPrepareRequestAuth(a *coreauth.Auth) bool { + c, err := helps.WorkBuddyCredential(a) + return err == nil && c.RefreshToken != "" && c.Expiring(time.Now()) +} +func (e *WorkBuddyExecutor) PrepareRequestAuth(ctx context.Context, a *coreauth.Auth) (*coreauth.Auth, error) { + return e.Refresh(ctx, a) +} +func (e *WorkBuddyExecutor) Refresh(ctx context.Context, a *coreauth.Auth) (*coreauth.Auth, error) { + if e.initErr != nil { + return nil, e.initErr + } + c, err := helps.WorkBuddyCredential(a) + if err != nil { + return nil, err + } + proxy := c.ProxyURL + if proxy == "" && e.cfg != nil { + proxy = e.cfg.ProxyURL + } + return e.runtime.Refresh(ctx, a, proxy) +} + +// Arbitrary management HTTP requests cannot bypass the provider governor or +// send WorkBuddy credentials to a caller-specified destination. +func (e *WorkBuddyExecutor) HttpRequest(context.Context, *coreauth.Auth, *http.Request) (*http.Response, error) { + return nil, &workbuddy.DispatchError{Reason: "arbitrary_http_requests_are_not_supported", Status: 400, Stop: true} +} + +func (e *WorkBuddyExecutor) build(ctx context.Context, c workbuddy.Credential, req execution.Request, opts execution.Options) ([]byte, error) { + if opts.Alt != "" { + return nil, &workbuddy.DispatchError{Reason: "unsupported_endpoint", Status: 400, Stop: true} + } + baseModel := thinking.ParseSuffix(req.Model).ModelName + from, to := opts.SourceFormat, sdktranslator.FromString("openai") + original := req.Payload + if len(opts.OriginalRequest) > 0 { + original = opts.OriginalRequest + } + originalTranslated, body, changed, err := helps.TranslateRequestPairWithAPIKeyModelCompatibilityAndUpdateIntent(ctx, opts.Headers, e.cfg, from, to, baseModel, original, req.Payload, opts.Stream, helps.APIKeyModelIsCompat(req)) + if err != nil { + return nil, err + } + body = helps.SetStringIfDifferent(body, "model", baseModel) + body, err = helps.ApplyRequestThinking(body, req, opts, from.String(), to.String(), e.Identifier(), changed) + if err != nil { + return nil, err + } + body, err = helps.NormalizeWorkBuddyPayload(body, c, baseModel) + if err != nil { + return nil, err + } + return helps.NewPayloadFinalizer(e.cfg, e.Identifier(), baseModel, to.String(), "", originalTranslated, req, opts)(body), nil +} + +func (e *WorkBuddyExecutor) execute(ctx context.Context, a *coreauth.Auth, req execution.Request, opts execution.Options, consume func([]byte, []byte) error) (errResult error) { + if e.initErr != nil { + return &workbuddy.DispatchError{Reason: "provider_unavailable", Status: 503, Stop: true} + } + ctx, errResult = e.PrepareExecutionContext(ctx, req, opts) + if errResult != nil { + return errResult + } + state := workbuddy.State(ctx) + a = e.runtime.Current(a) + c, err := helps.WorkBuddyCredential(a) + if err != nil { + return &workbuddy.DispatchError{Reason: "invalid_credential", Status: 401, Stop: true} + } + if c.Expiring(time.Now()) && c.RefreshToken != "" { + a, err = e.Refresh(ctx, a) + if err != nil { + return err + } + c, err = helps.WorkBuddyCredential(a) + if err != nil { + return err + } + } + caller, _ := opts.Metadata[execution.CallerScopeMetadataKey].(string) + queue := time.Duration(e.cfg.WorkBuddy.QueueSeconds) * time.Second + if queue == 0 { + queue = 30 * time.Second + } + release, err := e.runtime.Limiter.Acquire(ctx, caller, queue) + if err != nil { + return err + } + defer release() + reporter := helps.NewExecutorUsageReporter(ctx, e, req.Model, a) + var reportedUsage json.RawMessage + defer func() { + detail := helps.ParseOpenAIUsage(append(append([]byte(`{"usage":`), reportedUsage...), '}')) + if errResult != nil { + reporter.PublishFailureWithDetail(ctx, detail, errResult) + } else { + reporter.Publish(ctx, detail) + } + }() + preOutputRetry := false + for { + body, errBuild := e.build(ctx, c, req, opts) + if errBuild != nil { + return errBuild + } + // Validation after the final barrier is read-only. Explicit conflicting user + // rules are rejected rather than silently overwritten by provider defaults. + if len(body) > 16<<20 { + return &workbuddy.DispatchError{Reason: "request_body_too_large", Status: 413, Stop: true} + } + model := gjson.GetBytes(body, "model").String() + if !gjson.ValidBytes(body) || model == "" || !gjson.GetBytes(body, "stream").Bool() { + return &workbuddy.DispatchError{Reason: "final_payload_requires_model_and_stream", Status: 400, Stop: true} + } + gate, attempt := e.runtime.Governor.Check(c, model, workbuddy.IsPaidModel(c.Realm, model), false) + if !gate.Allowed() { + return &workbuddy.DispatchError{Reason: gate.Reason, Status: 429, Stop: gate.Stop, Until: gate.Until} + } + if errBegin := state.Begin(); errBegin != nil { + return errBegin + } + session := helps.DerivedSessionID(opts.Metadata, req.Metadata) + if c.Product == workbuddy.VSCode { + session = state.Turn + } + headers, order, errHeaders := workbuddy.IdentityHeaders(c, workbuddy.DefaultVersions(), session, state.Turn, model, time.Now()) + if errHeaders != nil { + return errHeaders + } + framed, encoding, errFrame := workbuddy.Frame(c.Product, body) + if errFrame != nil { + return errFrame + } + if encoding != "" { + headers.Set("Content-Encoding", encoding) + } + base, _, errBase := workbuddy.Endpoint(c.Realm, c.Product) + if errBase != nil { + return errBase + } + if e.runtime.BaseURL != "" { + base = e.runtime.BaseURL + } + proxy := c.ProxyURL + if proxy == "" { + proxy = e.cfg.ProxyURL + } + if override := execution.RequestProxyURL(ctx); override != "" { + proxy = override + } + client, errClient := e.runtime.Transports.Client(c, base, proxy, order) + if errClient != nil { + return errClient + } + httpReq, errRequest := http.NewRequestWithContext(ctx, http.MethodPost, strings.TrimSuffix(base, "/")+"/v2/chat/completions", bytes.NewReader(framed)) + if errRequest != nil { + return errRequest + } + httpReq.Header = headers + client = reporter.TrackHTTPClient(client) + reporter.SetTranslatedReasoningEffort(body, "openai") + execution.MarkUpstreamAttempt(ctx) + response, errDo := client.Do(httpReq) + if errDo != nil { + return e.recordFailure(c, model, state, workbuddy.ClassifyTransport(errDo)) + } + reporter.ObserveResponse(response) + if response.StatusCode < 200 || response.StatusCode >= 300 { + raw, errRead := io.ReadAll(io.LimitReader(response.Body, 1<<20)) + closeWorkBuddyBody(response.Body) + if errRead != nil { + return e.recordFailure(c, model, state, workbuddy.ClassifyTransport(errRead)) + } + fault := workbuddy.ClassifyHTTP(response.StatusCode, raw) + if fault.Kind == workbuddy.FaultAuth && c.RefreshToken != "" && state.Refresh(c) { + updated, errRefresh := e.runtime.Refresh(ctx, a, proxy) + if errRefresh == nil { + a = updated + c, err = helps.WorkBuddyCredential(a) + if err != nil { + return err + } + continue + } + } + return e.recordFailure(c, model, state, fault) + } + emitted := false + var usage json.RawMessage + var tokens int64 + finished := make(map[int]bool) + errStream := workbuddy.ReadEvents(response.Body, func(raw []byte) error { + event, errEvent := workbuddy.DecodeEvent(raw) + if errEvent != nil { + return errEvent + } + if len(event.Usage) > 0 && string(event.Usage) != "null" { + usage = append(usage[:0], event.Usage...) + if observed, ok := workbuddy.ObservedTokens(usage); ok { + tokens = observed + } + } + reporter.ObserveResponseModel(raw) + reporter.ObserveTokenEvent(gjson.GetBytes(raw, "choices.0.delta.content").String() != "" || gjson.GetBytes(raw, "choices.0.delta.reasoning_content").String() != "" || gjson.GetBytes(raw, "choices.0.delta.tool_calls.#").Int() > 0) + for _, choice := range event.Choices { + if choice.FinishReason != nil { + finished[choice.Index] = true + } else if _, ok := finished[choice.Index]; !ok { + finished[choice.Index] = false + } + } + if errConsume := consume(raw, body); errConsume != nil { + return errConsume + } + emitted = true + return nil + }) + closeWorkBuddyBody(response.Body) + if errStream == nil { + if len(finished) == 0 { + errStream = &workbuddy.StreamError{Reason: "upstream completed without a choice"} + } + for _, done := range finished { + if !done { + errStream = &workbuddy.StreamError{Reason: "upstream completed without a finish reason"} + break + } + } + } + if len(usage) > 0 { + reportedUsage = usage + } + if errComplete := e.runtime.Governor.Complete(c, attempt, tokens, nil, errStream == nil); errComplete != nil { + return &workbuddy.DispatchError{Reason: "usage_persistence_failed", Status: 503, Stop: true} + } + if errStream == nil { + return consume([]byte("[DONE]"), body) + } + if ctx.Err() != nil { + return ctx.Err() + } + var upstreamFault *workbuddy.StreamFault + if errors.As(errStream, &upstreamFault) { + failure := e.recordFailure(c, model, state, upstreamFault.Fault) + var dispatch *workbuddy.DispatchError + if emitted && errors.As(failure, &dispatch) { + dispatch.Stop = true + } + return failure + } + if !emitted && !preOutputRetry { + var streamError *workbuddy.StreamError + fault := workbuddy.ClassifyTransport(errStream) + if errors.As(errStream, &streamError) { + fault.Kind = workbuddy.FaultConnection + } + if fault.Kind == workbuddy.FaultConnection { + preOutputRetry = true + continue + } + } + failure := e.recordFailure(c, model, state, workbuddy.Fault{Kind: workbuddy.FaultUnknown, Status: 502}) + // Once output escaped, no SDK credential rotation may replay the prompt. + var dispatch *workbuddy.DispatchError + if errors.As(failure, &dispatch) { + dispatch.Stop = true + } + return failure + } +} +func (e *WorkBuddyExecutor) recordFailure(c workbuddy.Credential, model string, state *workbuddy.RequestState, fault workbuddy.Fault) error { + // Hard/unknown failures count once per UID per logical request. Explicit + // scopes and caps always apply, including new evidence after refresh. + charge := true + if fault.Kind == workbuddy.FaultHard || fault.Kind == workbuddy.FaultUnknown { + charge = state.Charge(c) + } + if charge { + if _, err := e.runtime.Governor.RecordFailure(c, model, fault, workbuddy.FailureOptions{FreeModel: workbuddy.IsFreeCapModel(c.Realm, model), SingleAccount: e.runtime.Accounts() <= 1}); err != nil { + return &workbuddy.DispatchError{Reason: "risk_persistence_failed", Status: 503, Stop: true} + } + } + return workbuddy.ErrorForFault(fault) +} +func closeWorkBuddyBody(body io.Closer) { + if err := body.Close(); err != nil { + log.WithField("provider", "workbuddy").Debug("upstream body close failed") + } +} + +func (e *WorkBuddyExecutor) Execute(ctx context.Context, a *coreauth.Auth, req execution.Request, opts execution.Options) (execution.Response, error) { + var collector workbuddy.Collector + var finalBody []byte + err := e.execute(ctx, a, req, opts, func(raw, body []byte) error { + finalBody = body + if bytes.Equal(raw, []byte("[DONE]")) { + return nil + } + return collector.Add(raw) + }) + if err != nil { + return execution.Response{}, err + } + body, err := collector.Response() + if err != nil { + return execution.Response{}, err + } + var param any + out := sdktranslator.TranslateNonStream(ctx, sdktranslator.FromString("openai"), execution.ResponseFormatOrSource(opts), req.Model, helps.ApplyPatchOriginalRequest(req, opts), finalBody, body, ¶m) + return execution.Response{Payload: out}, nil +} +func (e *WorkBuddyExecutor) ExecuteStream(ctx context.Context, a *coreauth.Auth, req execution.Request, opts execution.Options) (*execution.StreamResult, error) { + if e.initErr != nil { + return nil, &workbuddy.DispatchError{Reason: "provider_unavailable", Status: 503, Stop: true} + } + out := make(chan execution.StreamChunk) + go func() { + defer close(out) + var param any + err := e.execute(ctx, a, req, opts, func(raw, body []byte) error { + line := append([]byte("data: "), raw...) + chunks := sdktranslator.TranslateStream(ctx, sdktranslator.FromString("openai"), execution.ResponseFormatOrSource(opts), req.Model, helps.ApplyPatchOriginalRequest(req, opts), body, line, ¶m) + for _, chunk := range chunks { + select { + case out <- execution.StreamChunk{Payload: chunk}: + case <-ctx.Done(): + return ctx.Err() + } + } + return nil + }) + if err != nil { + select { + case out <- execution.StreamChunk{Err: err}: + case <-ctx.Done(): + } + } + }() + return &execution.StreamResult{Chunks: out}, nil +} +func (e *WorkBuddyExecutor) CountTokens(ctx context.Context, a *coreauth.Auth, req execution.Request, opts execution.Options) (execution.Response, error) { + if e.initErr != nil { + return execution.Response{}, e.initErr + } + c, err := helps.WorkBuddyCredential(a) + if err != nil { + return execution.Response{}, err + } + body, err := e.build(ctx, c, req, opts) + if err != nil { + return execution.Response{}, err + } + enc, err := helps.TokenizerForModel(thinking.ParseSuffix(req.Model).ModelName) + if err != nil { + return execution.Response{}, err + } + count, err := helps.CountOpenAIChatTokens(enc, body) + if err != nil { + return execution.Response{}, err + } + result := sdktranslator.TranslateTokenCount(ctx, sdktranslator.FromString("openai"), execution.ResponseFormatOrSource(opts), count, helps.BuildOpenAIUsageJSON(count)) + return execution.Response{Payload: result}, nil +} diff --git a/internal/runtime/executor/workbuddy_executor_test.go b/internal/runtime/executor/workbuddy_executor_test.go new file mode 100644 index 00000000000..f4e6261d269 --- /dev/null +++ b/internal/runtime/executor/workbuddy_executor_test.go @@ -0,0 +1,464 @@ +package executor + +import ( + "compress/gzip" + "context" + "encoding/json" + "errors" + "io" + "net/http" + "net/http/httptest" + "strings" + "sync" + "sync/atomic" + "testing" + + "github.com/router-for-me/CLIProxyAPI/v8/internal/config" + "github.com/router-for-me/CLIProxyAPI/v8/internal/runtime/executor/helps" + "github.com/router-for-me/CLIProxyAPI/v8/internal/workbuddy" + sdkauth "github.com/router-for-me/CLIProxyAPI/v8/sdk/auth" + coreauth "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/auth" + execution "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/executor" + sdktranslator "github.com/router-for-me/CLIProxyAPI/v8/sdk/translator" + "github.com/tidwall/gjson" + "github.com/tidwall/sjson" +) + +const wbSuccess = "data: {\"id\":\"chat-test\",\"model\":\"fixture-model\",\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\",\"content\":\"hello\"}}]}\n\ndata: {\"choices\":[{\"index\":0,\"delta\":{},\"finish_reason\":\"stop\"}]}\n\ndata: {\"choices\":[],\"usage\":{\"prompt_tokens\":8,\"completion_tokens\":2,\"total_tokens\":10}}\n\ndata: [DONE]\n\n" + +func wbFixture(t *testing.T, handler http.HandlerFunc) (*WorkBuddyExecutor, *coreauth.Auth) { + t.Helper() + server := httptest.NewServer(handler) + t.Cleanup(server.Close) + cfg := &config.Config{WorkBuddy: config.WorkBuddyConfig{Enabled: true, StateDir: t.TempDir()}} + e := NewWorkBuddyExecutor(cfg, nil) + if e.initErr != nil { + t.Fatal(e.initErr) + } + e.runtime.BaseURL = server.URL + t.Cleanup(e.runtime.Transports.CloseIdleConnections) + a, err := sdkauth.WorkBuddyAuthRecord(workbuddy.Credential{Type: "workbuddy", UID: "offline-user", Realm: workbuddy.International, Product: workbuddy.Desktop, AccessToken: "fixture-only-token", Models: []string{"fixture-model"}}) + if err != nil { + t.Fatal(err) + } + e.TrackAuth(a) + return e, a +} +func wbRequest() (execution.Request, execution.Options) { + return execution.Request{Model: "fixture-model", Payload: []byte(`{"model":"fixture-model","messages":[{"role":"user","content":"test"}]}`)}, execution.Options{SourceFormat: sdktranslator.FromString("openai")} +} +func TestWorkBuddyExecutorNonStreamAndStream(t *testing.T) { + for _, stream := range []bool{false, true} { + t.Run(map[bool]string{false: "aggregate", true: "stream"}[stream], func(t *testing.T) { + var requests atomic.Int32 + e, a := wbFixture(t, func(w http.ResponseWriter, r *http.Request) { + requests.Add(1) + if r.URL.Path != "/v2/chat/completions" { + t.Errorf("path = %s", r.URL.Path) + } + raw, err := io.ReadAll(r.Body) + if err != nil { + t.Error(err) + } + if !gjson.GetBytes(raw, "stream").Bool() { + t.Error("upstream was not streaming") + } + if r.Header.Get("X-IDE-Type") != "WorkBuddy" { + t.Error("missing identity") + } + w.Header().Set("Content-Type", "text/event-stream") + _, _ = io.WriteString(w, wbSuccess) + }) + req, opts := wbRequest() + if stream { + opts.Stream = true + result, err := e.ExecuteStream(context.Background(), a, req, opts) + if err != nil { + t.Fatal(err) + } + var output strings.Builder + for chunk := range result.Chunks { + if chunk.Err != nil { + t.Fatal(chunk.Err) + } + output.Write(chunk.Payload) + } + if !strings.Contains(output.String(), "hello") || !strings.Contains(output.String(), `"finish_reason":"stop"`) { + t.Fatalf("output = %s", output.String()) + } + } else { + result, err := e.Execute(context.Background(), a, req, opts) + if err != nil { + t.Fatal(err) + } + if gjson.GetBytes(result.Payload, "choices.0.message.content").String() != "hello" || gjson.GetBytes(result.Payload, "usage.total_tokens").Int() != 10 { + t.Fatalf("response = %s", result.Payload) + } + } + if requests.Load() != 1 { + t.Fatalf("requests = %d", requests.Load()) + } + raw, err := json.Marshal(a.Metadata) + if err != nil { + t.Fatal(err) + } + c, err := workbuddy.DecodeCredential(raw) + if err != nil { + t.Fatal(err) + } + if err := e.runtime.Governor.UpdateLimits(workbuddy.Limits{DailyTokens: 10}, nil); err != nil { + t.Fatal(err) + } + if gate, _ := e.runtime.Governor.Check(c, "fixture-model", true, false); gate.Allowed() { + t.Fatal("observed usage was not charged") + } + }) + } +} +func TestWorkBuddyExecutorPayloadBarrierOnBothPaths(t *testing.T) { + for _, stream := range []bool{false, true} { + t.Run(map[bool]string{false: "aggregate", true: "stream"}[stream], func(t *testing.T) { + e, a := wbFixture(t, func(w http.ResponseWriter, r *http.Request) { + reader, err := gzip.NewReader(r.Body) + if err != nil { + t.Error(err) + w.WriteHeader(500) + return + } + defer func() { _ = reader.Close() }() + body, err := io.ReadAll(reader) + if err != nil { + t.Error(err) + } + for path, want := range map[string]string{"max_tokens": "7", "cache_retention": "short", "metadata.user": "configured", "barrier": "first"} { + if got := gjson.GetBytes(body, path).String(); got != want { + t.Errorf("%s=%q want %q", path, got, want) + } + } + if gjson.GetBytes(body, "stream_options").Exists() { + t.Error("filtered field restored") + } + _, _ = io.WriteString(w, wbSuccess) + }) + a.Metadata["product"] = "cli" + e.cfg.Payload = config.PayloadConfig{ + Override: []config.PayloadRule{{Models: []config.PayloadModelRule{{Name: "*"}}, Params: map[string]any{"max_tokens": 7, "cache_retention": "short", "metadata.user": "configured", "barrier": "first"}}}, + Filter: []config.PayloadFilterRule{{Models: []config.PayloadModelRule{{Name: "*"}}, Params: []string{"stream_options"}}}, + } + req, opts := wbRequest() + if stream { + opts.Stream = true + result, err := e.ExecuteStream(context.Background(), a, req, opts) + if err != nil { + t.Fatal(err) + } + for chunk := range result.Chunks { + if chunk.Err != nil { + t.Fatal(chunk.Err) + } + } + } else { + if _, err := e.Execute(context.Background(), a, req, opts); err != nil { + t.Fatal(err) + } + } + }) + } +} +func TestWorkBuddyRestrictionSurvivesExecutorReload(t *testing.T) { + var requests atomic.Int32 + e, a := wbFixture(t, func(w http.ResponseWriter, r *http.Request) { + requests.Add(1) + w.WriteHeader(403) + _, _ = io.WriteString(w, `{"code":11140,"message":"sensitive-content"}`) + }) + req, opts := wbRequest() + _, err := e.Execute(context.Background(), a, req, opts) + var failure *workbuddy.DispatchError + if !errors.As(err, &failure) || !failure.Stop || strings.Contains(err.Error(), "sensitive-content") { + t.Fatalf("restriction = %v", err) + } + cfg := *e.cfg + cfg.WorkBuddy.MaxConcurrency = 2 + next := NewWorkBuddyExecutor(&cfg, e) + if next.runtime != e.runtime { + t.Fatal("reload lost runtime") + } + _, err = next.Execute(context.Background(), a, req, opts) + if !errors.As(err, &failure) || !strings.Contains(failure.Reason, "restricted") { + t.Fatalf("gate = %v", err) + } + if requests.Load() != 1 { + t.Fatal("restriction allowed new upstream request") + } +} +func TestWorkBuddyMidStreamFailureDoesNotReplayOrFakeDone(t *testing.T) { + var requests atomic.Int32 + e, a := wbFixture(t, func(w http.ResponseWriter, r *http.Request) { + requests.Add(1) + _, _ = io.WriteString(w, "data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"partial\"}}]}\n\n") + }) + req, opts := wbRequest() + opts.Stream = true + result, err := e.ExecuteStream(context.Background(), a, req, opts) + if err != nil { + t.Fatal(err) + } + var text strings.Builder + var terminal error + for chunk := range result.Chunks { + text.Write(chunk.Payload) + if chunk.Err != nil { + terminal = chunk.Err + } + } + var failure *workbuddy.DispatchError + if !errors.As(terminal, &failure) || !failure.Stop { + t.Fatalf("terminal = %v", terminal) + } + if !strings.Contains(text.String(), "partial") || strings.Contains(text.String(), "[DONE]") { + t.Fatalf("output = %s", text.String()) + } + if requests.Load() != 1 { + t.Fatal("replayed visible output") + } +} +func TestWorkBuddyEmptyStreamRetriesOnlyOnce(t *testing.T) { + var requests atomic.Int32 + e, a := wbFixture(t, func(w http.ResponseWriter, r *http.Request) { + if requests.Add(1) == 1 { + return + } + _, _ = io.WriteString(w, wbSuccess) + }) + req, opts := wbRequest() + if _, err := e.Execute(context.Background(), a, req, opts); err != nil { + t.Fatal(err) + } + if requests.Load() != 2 { + t.Fatalf("attempts = %d", requests.Load()) + } +} +func TestWorkBuddyRefreshBeforeRetryPersistsPair(t *testing.T) { + var chats atomic.Int32 + var refreshes atomic.Int32 + var saved atomic.Int32 + e, a := wbFixture(t, func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/v2/plugin/auth/token/refresh" { + refreshes.Add(1) + _, _ = io.WriteString(w, `{"code":0,"data":{"accessToken":"eyJhbGciOiJub25lIn0.eyJzdWIiOiJvZmZsaW5lLXVzZXIifQ.sig","refreshToken":"rotated-refresh","expiresAt":4102444800000}}`) + return + } + if chats.Add(1) == 1 { + w.WriteHeader(401) + return + } + if saved.Load() != 1 { + t.Error("chat resumed before durable credential save") + } + _, _ = io.WriteString(w, wbSuccess) + }) + a.Metadata["refresh_token"] = "fixture-refresh" + var mu sync.Mutex + var persisted *coreauth.Auth + e.runtime.Save = func(_ context.Context, _ *coreauth.Auth, next *coreauth.Auth) error { + mu.Lock() + defer mu.Unlock() + persisted = next.Clone() + saved.Add(1) + return nil + } + req, opts := wbRequest() + if _, err := e.Execute(context.Background(), a, req, opts); err != nil { + t.Fatal(err) + } + if refreshes.Load() != 1 || chats.Load() != 2 { + t.Fatalf("refreshes=%d chats=%d", refreshes.Load(), chats.Load()) + } + mu.Lock() + defer mu.Unlock() + if persisted == nil || persisted.Metadata["refresh_token"] != "rotated-refresh" { + t.Fatal("rotated refresh token missing") + } +} +func TestWorkBuddyDisabledExecutorAndArbitraryHTTP(t *testing.T) { + e := NewWorkBuddyExecutor(&config.Config{}, nil) + req, opts := wbRequest() + if _, err := e.Execute(context.Background(), nil, req, opts); err == nil { + t.Fatal("disabled provider executed") + } + if _, err := e.HttpRequest(context.Background(), nil, nil); err == nil { + t.Fatal("ungoverned HTTP supported") + } +} +func TestWorkBuddyCredentialUsesLatestRotation(t *testing.T) { + e, a := wbFixture(t, func(http.ResponseWriter, *http.Request) {}) + old := a.Clone() + newer := a.Clone() + newer.UpdatedAt = newer.UpdatedAt.Add(1) + newer.Metadata["access_token"] = "new-token" + e.TrackAuth(newer) + c, err := helps.WorkBuddyCredential(e.runtime.Current(old)) + if err != nil { + t.Fatal(err) + } + if c.AccessToken != "new-token" { + t.Fatal("stale SDK credential restored consumed token") + } +} + +func TestWorkBuddyDownstreamProtocols(t *testing.T) { + cases := []struct{ format, body, contentPath, streamMarker string }{ + {"claude", `{"model":"fixture-model","max_tokens":100,"messages":[{"role":"user","content":"test"}]}`, "content.0.text", "content_block_delta"}, + {"openai-response", `{"model":"fixture-model","input":"test"}`, "output.0.content.0.text", "response.output_text.delta"}, + {"gemini", `{"contents":[{"role":"user","parts":[{"text":"test"}]}]}`, "candidates.0.content.parts.0.text", "hello"}, + } + for _, tc := range cases { + t.Run(tc.format, func(t *testing.T) { + e, a := wbFixture(t, func(w http.ResponseWriter, r *http.Request) { + body, _ := io.ReadAll(r.Body) + if gjson.GetBytes(body, "messages.#").Int() < 2 { + t.Errorf("untranslated request: %s", body) + } + _, _ = io.WriteString(w, wbSuccess) + }) + req := execution.Request{Model: "fixture-model", Payload: []byte(tc.body)} + opts := execution.Options{SourceFormat: sdktranslator.FromString(tc.format)} + response, err := e.Execute(context.Background(), a, req, opts) + if err != nil { + t.Fatal(err) + } + if gjson.GetBytes(response.Payload, tc.contentPath).String() != "hello" { + t.Fatalf("nonstream response = %s", response.Payload) + } + opts.Stream = true + req.Payload, _ = sjson.SetBytes(req.Payload, "stream", true) + result, err := e.ExecuteStream(context.Background(), a, req, opts) + if err != nil { + t.Fatal(err) + } + var text strings.Builder + for chunk := range result.Chunks { + if chunk.Err != nil { + t.Fatal(chunk.Err) + } + text.Write(chunk.Payload) + } + if !strings.Contains(text.String(), tc.streamMarker) { + t.Fatalf("untranslated stream = %s", text.String()) + } + }) + } +} + +func TestWorkBuddyStreamEnvelopeRestrictionStopsAndPersists(t *testing.T) { + var requests atomic.Int32 + e, a := wbFixture(t, func(w http.ResponseWriter, r *http.Request) { + requests.Add(1) + _, _ = io.WriteString(w, "data: {\"error\":{\"code\":11140,\"message\":\"private upstream detail\"}}\n\n") + }) + req, opts := wbRequest() + _, err := e.Execute(context.Background(), a, req, opts) + var failure *workbuddy.DispatchError + if !errors.As(err, &failure) || !failure.Stop || failure.Reason != string(workbuddy.FaultRestricted) { + t.Fatalf("SSE fault = %v", err) + } + if strings.Contains(err.Error(), "private") { + t.Fatal("error body exposed") + } + if _, err = e.Execute(context.Background(), a, req, opts); err == nil { + t.Fatal("restriction lost") + } + if requests.Load() != 1 { + t.Fatal("stream error bypassed restriction gate") + } +} + +func TestWorkBuddyRefreshUsesConfiguredManagerStore(t *testing.T) { + var chats atomic.Int32 + e, a := wbFixture(t, func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/v2/plugin/auth/token/refresh" { + _, _ = io.WriteString(w, `{"code":0,"data":{"accessToken":"eyJhbGciOiJub25lIn0.eyJzdWIiOiJvZmZsaW5lLXVzZXIifQ.sig","refreshToken":"manager-rotated","expiresAt":4102444800000}}`) + return + } + if chats.Add(1) == 1 { + w.WriteHeader(401) + return + } + _, _ = io.WriteString(w, wbSuccess) + }) + a.Metadata["refresh_token"] = "fixture-refresh" + store := sdkauth.NewFileTokenStore() + dir := t.TempDir() + store.SetBaseDir(dir) + manager := coreauth.NewManager(store, nil, nil) + registered, err := manager.Register(context.Background(), a) + if err != nil { + t.Fatal(err) + } + e.BindCredentialManager(manager) + e.TrackAuth(registered) + manager.RegisterExecutor(e) + req, opts := wbRequest() + if _, err := e.Execute(context.Background(), registered, req, opts); err != nil { + t.Fatal(err) + } + current, ok := manager.GetByID(registered.ID) + if !ok || current.Metadata["refresh_token"] != "manager-rotated" { + t.Fatal("manager not updated") + } + loaded, err := store.List(context.Background()) + if err != nil { + t.Fatal(err) + } + if len(loaded) != 1 || loaded[0].Metadata["refresh_token"] != "manager-rotated" { + t.Fatal("configured store did not receive rotation") + } + // A later SDK refresh takes the same provider path and preserves manager fences. + if _, err := manager.ForceRefreshAuth(context.Background(), registered.ID); err != nil { + t.Fatal(err) + } +} + +func TestWorkBuddyInflightRefreshCannotOverwriteExternalRotation(t *testing.T) { + started, release := make(chan struct{}), make(chan struct{}) + e, a := wbFixture(t, func(w http.ResponseWriter, r *http.Request) { + close(started) + <-release + _, _ = io.WriteString(w, `{"code":0,"data":{"accessToken":"eyJhbGciOiJub25lIn0.eyJzdWIiOiJvZmZsaW5lLXVzZXIifQ.sig","refreshToken":"stale-rotation","expiresAt":4102444800000}}`) + }) + a.Metadata["refresh_token"] = "fixture-refresh" + store := sdkauth.NewFileTokenStore() + store.SetBaseDir(t.TempDir()) + manager := coreauth.NewManager(store, nil, nil) + registered, err := manager.Register(context.Background(), a) + if err != nil { + t.Fatal(err) + } + e.BindCredentialManager(manager) + e.TrackAuth(registered) + done := make(chan error, 1) + go func() { _, errRefresh := e.Refresh(context.Background(), registered); done <- errRefresh }() + <-started + replacement := registered.Clone() + replacement.Metadata["access_token"] = "newer-external-token" + replacement.Metadata["refresh_token"] = "newer-external-refresh" + if _, err := manager.Update(context.Background(), replacement); err != nil { + t.Fatal(err) + } + close(release) + if err := <-done; err == nil { + t.Fatal("stale refresh was accepted") + } + current, ok := manager.GetByID(registered.ID) + if !ok || current.Metadata["access_token"] != "newer-external-token" { + t.Fatal("external rotation overwritten") + } + loaded, err := store.List(context.Background()) + if err != nil { + t.Fatal(err) + } + if len(loaded) != 1 || loaded[0].Metadata["access_token"] != "newer-external-token" { + t.Fatal("external persisted rotation overwritten") + } +} diff --git a/internal/workbuddy/catalog.go b/internal/workbuddy/catalog.go new file mode 100644 index 00000000000..29090a2339d --- /dev/null +++ b/internal/workbuddy/catalog.go @@ -0,0 +1,60 @@ +package workbuddy + +import ( + _ "embed" + "encoding/json" + "strings" +) + +// catalogSnapshot is a versioned offline fallback, not a claim of live access. +// Source: workbuddy2api-hub af706fc31e4a22cb2dea7442361b2b15185c3cee. +// +//go:embed catalog.json +var catalogSnapshot []byte + +type Model struct { + ID string `json:"id"` + Name string `json:"name"` + Credits string `json:"credits"` + Input int `json:"maxInputTokens"` + Output int `json:"maxOutputTokens"` + Context int `json:"maxAllowedSize"` + Images bool `json:"supportsImages"` + Tools bool `json:"supportsToolCall"` + Reasoning bool `json:"supportsReasoning"` + Advertise bool `json:"advertise"` + Thinking ModelThinking `json:"reasoning"` +} + +type ModelThinking struct { + Effort string `json:"effort"` + Default string `json:"defaultEffort"` + Levels []string `json:"supportedEfforts"` + CanDisable *bool `json:"canDisableThinking"` +} + +func Catalog(realm Realm) []Model { + var catalogs map[Realm][]Model + if json.Unmarshal(catalogSnapshot, &catalogs) != nil { + return nil + } + return catalogs[realm] +} +func ModelFor(realm Realm, id string) (Model, bool) { + for _, m := range Catalog(realm) { + if m.ID == id { + return m, true + } + } + return Model{ID: id}, false +} +func IsFreeCapModel(realm Realm, model string) bool { + return (realm == China || realm == International) && (model == "hy3" || model == "hy4-preview-f" || realm == International && model == "deepseek-v4.1-flash") +} +func IsPaidModel(realm Realm, model string) bool { + m, ok := ModelFor(realm, model) + if IsFreeCapModel(realm, model) { + return false + } + return !ok || !strings.Contains(m.Credits, "x0.00") +} diff --git a/internal/workbuddy/catalog.json b/internal/workbuddy/catalog.json new file mode 100644 index 00000000000..4691440707a --- /dev/null +++ b/internal/workbuddy/catalog.json @@ -0,0 +1,1036 @@ +{ + "intl": [ + { + "id": "deepseek-v4.1-flash", + "name": "Deepseek-V4.1-Flash", + "credits": "x0.00", + "maxInputTokens": 1000000, + "maxOutputTokens": 128000, + "maxAllowedSize": 1000000, + "supportsImages": true, + "supportsToolCall": true, + "supportsReasoning": true, + "reasoning": { + "defaultEffort": "high", + "supportedEfforts": [ + "low", + "high", + "max" + ], + "summary": "auto" + }, + "advertise": true + }, + { + "id": "gpt-6-astra", + "name": "GPT-6-Astra", + "credits": "x6.67", + "maxInputTokens": 1000000, + "maxOutputTokens": 128000, + "maxAllowedSize": 1000000, + "supportsImages": true, + "supportsToolCall": true, + "supportsReasoning": true, + "reasoning": { + "canDisableThinking": true, + "defaultEffort": "high", + "summary": "auto", + "supportedEfforts": [ + "low", + "medium", + "high", + "xhigh", + "max" + ] + }, + "advertise": true + }, + { + "id": "hy4-preview-f", + "name": "Hy4 preview", + "credits": "x0.00", + "maxInputTokens": 1000000, + "maxOutputTokens": 64000, + "maxAllowedSize": 1000000, + "supportsImages": true, + "supportsToolCall": true, + "supportsReasoning": true, + "reasoning": { + "canDisableThinking": false, + "defaultEffort": "high", + "summary": "auto", + "supportedEfforts": [ + "high" + ] + }, + "advertise": true + }, + { + "id": "hy4-preview", + "name": "Hy4 preview", + "credits": "x0.29", + "maxInputTokens": 1000000, + "maxOutputTokens": 64000, + "maxAllowedSize": 1000000, + "supportsImages": true, + "supportsToolCall": true, + "supportsReasoning": true, + "reasoning": { + "canDisableThinking": false, + "defaultEffort": "high", + "summary": "auto", + "supportedEfforts": [ + "high" + ] + }, + "advertise": false + }, + { + "id": "hy3", + "name": "Hy3", + "credits": "x0.00", + "maxInputTokens": 192000, + "maxOutputTokens": 64000, + "maxAllowedSize": 192000, + "supportsImages": true, + "supportsToolCall": true, + "supportsReasoning": true, + "reasoning": { + "canDisableThinking": false, + "defaultEffort": "high", + "summary": "auto", + "supportedEfforts": [ + "low", + "high" + ] + }, + "advertise": true + }, + { + "id": "gpt-5.6-sol", + "name": "GPT-5.6-Sol", + "credits": "x3.47", + "maxInputTokens": 1000000, + "maxOutputTokens": 128000, + "maxAllowedSize": 1000000, + "supportsImages": true, + "supportsToolCall": true, + "supportsReasoning": true, + "reasoning": { + "canDisableThinking": true, + "defaultEffort": "high", + "summary": "auto", + "supportedEfforts": [ + "low", + "medium", + "high", + "xhigh", + "max" + ] + }, + "advertise": true + }, + { + "id": "gpt-5.6-terra", + "name": "GPT-5.6-Terra", + "credits": "x1.39", + "maxInputTokens": 1000000, + "maxOutputTokens": 128000, + "maxAllowedSize": 1000000, + "supportsImages": true, + "supportsToolCall": true, + "supportsReasoning": true, + "reasoning": { + "canDisableThinking": true, + "defaultEffort": "high", + "summary": "auto", + "supportedEfforts": [ + "low", + "medium", + "high", + "xhigh", + "max" + ] + }, + "advertise": true + }, + { + "id": "gpt-5.6-luna", + "name": "GPT-5.6-Luna", + "credits": "x0.14", + "maxInputTokens": 1000000, + "maxOutputTokens": 128000, + "maxAllowedSize": 1000000, + "supportsImages": true, + "supportsToolCall": true, + "supportsReasoning": true, + "reasoning": { + "canDisableThinking": true, + "defaultEffort": "high", + "summary": "auto", + "supportedEfforts": [ + "low", + "medium", + "high", + "xhigh", + "max" + ] + }, + "advertise": true + }, + { + "id": "gpt-5.5", + "name": "GPT-5.5", + "credits": "x3.31", + "maxInputTokens": 1000000, + "maxOutputTokens": 128000, + "maxAllowedSize": 1000000, + "supportsImages": true, + "supportsToolCall": true, + "supportsReasoning": true, + "reasoning": { + "canDisableThinking": false, + "defaultEffort": "high", + "summary": "auto", + "supportedEfforts": [ + "low", + "medium", + "high", + "xhigh" + ] + }, + "advertise": true + }, + { + "id": "gpt-5.4", + "name": "GPT-5.4", + "credits": "x1.65", + "maxInputTokens": 272000, + "maxOutputTokens": 72000, + "maxAllowedSize": 272000, + "supportsImages": true, + "supportsToolCall": true, + "supportsReasoning": true, + "reasoning": { + "canDisableThinking": false, + "defaultEffort": "high", + "summary": "auto", + "supportedEfforts": [ + "low", + "medium", + "high", + "xhigh" + ] + }, + "advertise": true + }, + { + "id": "grok-4.7", + "name": "Grok-4.7", + "credits": "x1.90", + "maxInputTokens": 500000, + "maxOutputTokens": 128000, + "maxAllowedSize": 500000, + "supportsImages": true, + "supportsToolCall": true, + "supportsReasoning": true, + "reasoning": { + "canDisableThinking": false, + "defaultEffort": "high", + "supportedEfforts": [ + "low", + "medium", + "high", + "xhigh" + ] + }, + "advertise": true + }, + { + "id": "gpt-5.3-codex", + "name": "GPT-5.3-Codex", + "credits": "x1.25", + "maxInputTokens": 272000, + "maxOutputTokens": 72000, + "maxAllowedSize": 272000, + "supportsImages": true, + "supportsToolCall": true, + "supportsReasoning": true, + "reasoning": { + "effort": "medium", + "summary": "auto" + }, + "advertise": false + }, + { + "id": "gemini-3.5-flash", + "name": "Gemini-3.5-Flash", + "credits": "x0.99", + "maxInputTokens": 1000000, + "maxOutputTokens": 65536, + "maxAllowedSize": 1000000, + "supportsImages": true, + "supportsToolCall": true, + "supportsReasoning": true, + "reasoning": { + "effort": "medium", + "summary": "auto" + }, + "advertise": true + }, + { + "id": "glm-5.3", + "name": "GLM-5.3", + "credits": "x0.79", + "maxInputTokens": 1000000, + "maxOutputTokens": 48000, + "maxAllowedSize": 1000000, + "supportsImages": true, + "supportsToolCall": true, + "supportsReasoning": true, + "reasoning": { + "canDisableThinking": true, + "defaultEffort": "high", + "summary": "auto", + "supportedEfforts": [ + "low", + "high", + "max" + ] + }, + "advertise": true + }, + { + "id": "glm-5.3-flash", + "name": "GLM-5.3-Flash", + "credits": "x0.06", + "maxInputTokens": 1000000, + "maxOutputTokens": 32000, + "maxAllowedSize": 1000000, + "supportsImages": true, + "supportsToolCall": true, + "supportsReasoning": true, + "reasoning": { + "canDisableThinking": true, + "defaultEffort": "high", + "summary": "auto", + "supportedEfforts": [ + "low", + "high", + "max" + ] + }, + "advertise": true + }, + { + "id": "glm-5.2", + "name": "GLM-5.2", + "credits": "x0.79", + "maxInputTokens": 1000000, + "maxOutputTokens": 48000, + "maxAllowedSize": 1000000, + "supportsImages": true, + "supportsToolCall": true, + "supportsReasoning": true, + "reasoning": { + "canDisableThinking": true, + "defaultEffort": "high", + "summary": "auto", + "supportedEfforts": [ + "high", + "xhigh" + ] + }, + "advertise": true + }, + { + "id": "kimi-k3", + "name": "Kimi-K3", + "credits": "x1.62", + "maxInputTokens": 1000000, + "maxOutputTokens": 32000, + "maxAllowedSize": 1000000, + "supportsImages": true, + "supportsToolCall": true, + "supportsReasoning": true, + "reasoning": { + "effort": "medium", + "summary": "auto" + }, + "advertise": true + }, + { + "id": "kimi-k2.8-preview", + "name": "Kimi-K2.8-Preview", + "credits": "x0.77", + "maxInputTokens": 1000000, + "maxOutputTokens": 32000, + "maxAllowedSize": 1000000, + "supportsImages": true, + "supportsToolCall": true, + "supportsReasoning": true, + "reasoning": { + "canDisableThinking": true, + "defaultEffort": "high", + "summary": "auto", + "supportedEfforts": [ + "low", + "high", + "max" + ] + }, + "advertise": true + }, + { + "id": "kimi-k2.6", + "name": "Kimi-K2.6", + "credits": "x0.52", + "maxInputTokens": 256000, + "maxOutputTokens": 32000, + "maxAllowedSize": 256000, + "supportsImages": true, + "supportsToolCall": true, + "supportsReasoning": true, + "reasoning": { + "effort": "medium", + "summary": "auto" + }, + "advertise": true + } + ], + "cn": [ + { + "id": "hy3", + "name": "Hy3", + "credits": "x0.00", + "maxInputTokens": 192000, + "maxOutputTokens": 64000, + "maxAllowedSize": 192000, + "supportsImages": true, + "supportsToolCall": true, + "supportsReasoning": true, + "reasoning": { + "canDisableThinking": false, + "defaultEffort": "high", + "summary": "auto", + "supportedEfforts": [ + "low", + "high" + ] + }, + "advertise": true + }, + { + "id": "hy3-x", + "name": "Hy3", + "credits": "x0.05", + "maxInputTokens": 192000, + "maxOutputTokens": 64000, + "maxAllowedSize": 192000, + "supportsImages": true, + "supportsToolCall": true, + "supportsReasoning": true, + "reasoning": { + "canDisableThinking": false, + "defaultEffort": "high", + "summary": "auto", + "supportedEfforts": [ + "low", + "high" + ] + }, + "advertise": false + }, + { + "id": "hy4-preview", + "name": "Hy4 preview", + "credits": "x0.29", + "maxInputTokens": 1000000, + "maxOutputTokens": 64000, + "maxAllowedSize": 1000000, + "supportsImages": true, + "supportsToolCall": true, + "supportsReasoning": true, + "reasoning": { + "canDisableThinking": false, + "defaultEffort": "high", + "summary": "auto", + "supportedEfforts": [ + "high" + ] + }, + "advertise": false + }, + { + "id": "hy4-preview-dev", + "name": "Hy4 preview", + "credits": "x0.29", + "maxInputTokens": 1000000, + "maxOutputTokens": 64000, + "maxAllowedSize": 1000000, + "supportsImages": true, + "supportsToolCall": true, + "supportsReasoning": true, + "reasoning": { + "canDisableThinking": false, + "defaultEffort": "high", + "summary": "auto", + "supportedEfforts": [ + "high" + ] + }, + "advertise": false + }, + { + "id": "hy4-preview-f", + "name": "Hy4 preview", + "credits": "x0.00", + "maxInputTokens": 1000000, + "maxOutputTokens": 64000, + "maxAllowedSize": 1000000, + "supportsImages": true, + "supportsToolCall": true, + "supportsReasoning": true, + "reasoning": { + "canDisableThinking": false, + "defaultEffort": "high", + "summary": "auto", + "supportedEfforts": [ + "high" + ] + }, + "advertise": true + }, + { + "id": "hy4-preview-x", + "name": "Hy4 preview", + "credits": "x0.29", + "maxInputTokens": 1000000, + "maxOutputTokens": 64000, + "maxAllowedSize": 1000000, + "supportsImages": true, + "supportsToolCall": true, + "supportsReasoning": true, + "reasoning": { + "canDisableThinking": false, + "defaultEffort": "high", + "summary": "auto", + "supportedEfforts": [ + "high" + ] + }, + "advertise": false + }, + { + "id": "minimax-m2.5", + "name": "MiniMax-M2.5", + "credits": "x0.18", + "maxInputTokens": 200000, + "maxOutputTokens": 48000, + "maxAllowedSize": 200000, + "supportsImages": true, + "supportsToolCall": true, + "supportsReasoning": true, + "reasoning": { + "effort": "medium", + "summary": "auto" + }, + "advertise": false + }, + { + "id": "glm-5v-turbo", + "name": "GLM-5v-Turbo", + "credits": "x0.71", + "maxInputTokens": 200000, + "maxOutputTokens": 64000, + "maxAllowedSize": 200000, + "supportsImages": true, + "supportsToolCall": true, + "supportsReasoning": true, + "reasoning": { + "effort": "medium", + "summary": "auto" + }, + "advertise": true + }, + { + "id": "glm-5.3", + "name": "GLM-5.3", + "credits": "x0.79", + "maxInputTokens": 1000000, + "maxOutputTokens": 48000, + "maxAllowedSize": 1000000, + "supportsImages": true, + "supportsToolCall": true, + "supportsReasoning": true, + "reasoning": { + "canDisableThinking": true, + "defaultEffort": "high", + "summary": "auto", + "supportedEfforts": [ + "low", + "high", + "max" + ] + }, + "advertise": true + }, + { + "id": "glm-5.3-flash", + "name": "GLM-5.3-Flash", + "credits": "x0.06", + "maxInputTokens": 1000000, + "maxOutputTokens": 32000, + "maxAllowedSize": 1000000, + "supportsImages": true, + "supportsToolCall": true, + "supportsReasoning": true, + "reasoning": { + "canDisableThinking": true, + "defaultEffort": "high", + "summary": "auto", + "supportedEfforts": [ + "low", + "high", + "max" + ] + }, + "advertise": true + }, + { + "id": "glm-5.2", + "name": "GLM-5.2", + "credits": "x0.79", + "maxInputTokens": 1000000, + "maxOutputTokens": 48000, + "maxAllowedSize": 1000000, + "supportsImages": true, + "supportsToolCall": true, + "supportsReasoning": true, + "reasoning": { + "canDisableThinking": true, + "defaultEffort": "high", + "summary": "auto", + "supportedEfforts": [ + "high", + "xhigh" + ] + }, + "advertise": true + }, + { + "id": "glm-5.1", + "name": "GLM-5.1", + "credits": "x0.79", + "maxInputTokens": 200000, + "maxOutputTokens": 48000, + "maxAllowedSize": 200000, + "supportsImages": true, + "supportsToolCall": true, + "supportsReasoning": true, + "reasoning": { + "effort": "medium", + "summary": "auto" + }, + "advertise": true + }, + { + "id": "glm-5.0-turbo", + "name": "GLM-5.0-Turbo", + "credits": "x0.95", + "maxInputTokens": 200000, + "maxOutputTokens": 48000, + "maxAllowedSize": 200000, + "supportsImages": true, + "supportsToolCall": true, + "supportsReasoning": true, + "reasoning": { + "effort": "medium", + "summary": "auto" + }, + "advertise": false + }, + { + "id": "glm-4.6v", + "name": "GLM-4.6V", + "credits": "x0.11", + "maxInputTokens": 128000, + "maxOutputTokens": 32000, + "maxAllowedSize": 128000, + "supportsImages": true, + "supportsToolCall": true, + "supportsReasoning": true, + "reasoning": { + "effort": "medium", + "summary": "auto" + }, + "advertise": false + }, + { + "id": "kimi-k3-1", + "name": "Kimi-K3", + "credits": "x1.62", + "maxInputTokens": 1000000, + "maxOutputTokens": 32000, + "maxAllowedSize": 1000000, + "supportsImages": true, + "supportsToolCall": true, + "supportsReasoning": true, + "reasoning": { + "canDisableThinking": true, + "defaultEffort": "high", + "summary": "auto", + "supportedEfforts": [ + "low", + "high", + "xhigh" + ] + }, + "advertise": true + }, + { + "id": "kimi-k2.8-preview", + "name": "Kimi-K2.8-Preview", + "credits": "x0.77", + "maxInputTokens": 1000000, + "maxOutputTokens": 32000, + "maxAllowedSize": 1000000, + "supportsImages": true, + "supportsToolCall": true, + "supportsReasoning": true, + "reasoning": { + "canDisableThinking": true, + "defaultEffort": "high", + "summary": "auto", + "supportedEfforts": [ + "low", + "high", + "max" + ] + }, + "advertise": true + }, + { + "id": "kimi-k2.7", + "name": "Kimi-K2.7-Code", + "credits": "x0.57", + "maxInputTokens": 256000, + "maxOutputTokens": 32000, + "maxAllowedSize": 256000, + "supportsImages": true, + "supportsToolCall": true, + "supportsReasoning": true, + "reasoning": { + "effort": "medium", + "summary": "auto" + }, + "advertise": true + }, + { + "id": "kimi-k2.6", + "name": "Kimi-K2.6", + "credits": "x0.52", + "maxInputTokens": 256000, + "maxOutputTokens": 32000, + "maxAllowedSize": 256000, + "supportsImages": true, + "supportsToolCall": true, + "supportsReasoning": true, + "reasoning": { + "effort": "medium", + "summary": "auto" + }, + "advertise": true + }, + { + "id": "kimi-k2.5", + "name": "Kimi-K2.5", + "credits": "x0.45", + "maxInputTokens": 256000, + "maxOutputTokens": 32000, + "maxAllowedSize": 256000, + "supportsImages": true, + "supportsToolCall": true, + "supportsReasoning": true, + "reasoning": { + "effort": "medium", + "summary": "auto" + }, + "advertise": false + }, + { + "id": "kimi-k2-thinking", + "name": "Kimi-K2-Thinking", + "credits": "x0.54", + "maxInputTokens": 256000, + "maxOutputTokens": 32000, + "maxAllowedSize": 256000, + "supportsImages": true, + "supportsToolCall": true, + "supportsReasoning": true, + "reasoning": { + "effort": "medium", + "summary": "auto" + }, + "advertise": false + }, + { + "id": "minimax-m3", + "name": "MiniMax-M3", + "credits": "x0.25", + "maxInputTokens": 512000, + "maxOutputTokens": 128000, + "maxAllowedSize": 512000, + "supportsImages": true, + "supportsToolCall": true, + "supportsReasoning": true, + "reasoning": { + "effort": "medium", + "summary": "auto" + }, + "advertise": true + }, + { + "id": "minimax-m2.7", + "name": "MiniMax-M2.7", + "credits": "x0.26", + "maxInputTokens": 200000, + "maxOutputTokens": 48000, + "maxAllowedSize": 200000, + "supportsImages": true, + "supportsToolCall": true, + "supportsReasoning": true, + "reasoning": { + "effort": "medium", + "summary": "auto" + }, + "advertise": false + }, + { + "id": "glm-4.6", + "name": "GLM-4.6", + "credits": "x0.23", + "maxInputTokens": 168000, + "maxOutputTokens": 32000, + "maxAllowedSize": 168000, + "supportsImages": true, + "supportsToolCall": true, + "supportsReasoning": true, + "advertise": false + }, + { + "id": "deepseek-v4-flash", + "name": "Deepseek-V4-Flash", + "credits": "x0.17", + "maxInputTokens": 1000000, + "maxOutputTokens": 50000, + "maxAllowedSize": 1000000, + "supportsImages": true, + "supportsToolCall": true, + "supportsReasoning": true, + "reasoning": { + "canDisableThinking": true, + "defaultEffort": "high", + "summary": "auto", + "supportedEfforts": [ + "high", + "xhigh" + ] + }, + "advertise": false + }, + { + "id": "deepseek-v4.1-flash", + "name": "Deepseek-V4.1-Flash", + "credits": "x0.11", + "maxInputTokens": 1000000, + "maxOutputTokens": 128000, + "maxAllowedSize": 1000000, + "supportsImages": true, + "supportsToolCall": true, + "supportsReasoning": true, + "reasoning": { + "defaultEffort": "high", + "supportedEfforts": [ + "low", + "high", + "max" + ], + "summary": "auto" + }, + "advertise": true + }, + { + "id": "deepseek-v4-pro", + "name": "Deepseek-V4-Pro", + "credits": "x0.51", + "maxInputTokens": 1000000, + "maxOutputTokens": 50000, + "maxAllowedSize": 1000000, + "supportsImages": true, + "supportsToolCall": true, + "supportsReasoning": true, + "reasoning": { + "canDisableThinking": true, + "defaultEffort": "high", + "summary": "auto", + "supportedEfforts": [ + "high", + "xhigh" + ] + }, + "advertise": true + }, + { + "id": "deepseek-v3-2-volc", + "name": "DeepSeek-V3.2", + "credits": "x0.29", + "maxInputTokens": 96000, + "maxOutputTokens": 32000, + "maxAllowedSize": 96000, + "supportsImages": true, + "supportsToolCall": true, + "supportsReasoning": true, + "reasoning": { + "effort": "medium", + "summary": "auto" + }, + "advertise": false + }, + { + "id": "deepseek-v3-1-volc", + "name": "DeepSeek-V3-1-Terminus", + "credits": "x0.52", + "maxInputTokens": 96000, + "maxOutputTokens": 32000, + "maxAllowedSize": 128000, + "supportsImages": true, + "supportsToolCall": true, + "advertise": false + }, + { + "id": "deepseek-v3-1-lkeap", + "name": "DeepSeek-V3-1", + "credits": "x0.52", + "maxInputTokens": 96000, + "maxOutputTokens": 32000, + "maxAllowedSize": 128000, + "supportsImages": true, + "supportsToolCall": true, + "advertise": false + }, + { + "id": "deepseek-v3-1", + "name": "DeepSeek-V3.1", + "credits": "x0.52", + "maxInputTokens": 96000, + "maxOutputTokens": 32000, + "maxAllowedSize": 128000, + "supportsImages": true, + "supportsToolCall": true, + "advertise": false + }, + { + "id": "deepseek-v3-0324-lkeap", + "name": "DeepSeek-V3-0324", + "credits": "x0.52", + "maxInputTokens": 112000, + "maxOutputTokens": 16000, + "maxAllowedSize": 128000, + "supportsImages": true, + "supportsToolCall": true, + "advertise": false + }, + { + "id": "deepseek-r1-0528-lkeap", + "name": "DeepSeek-R1-0528", + "credits": "", + "maxInputTokens": 96000, + "maxOutputTokens": 16000, + "maxAllowedSize": 112000, + "supportsImages": true, + "supportsToolCall": true, + "advertise": false + }, + { + "id": "kimi-k2-instruct-taiji", + "name": "Kimi-K2", + "credits": "", + "maxInputTokens": 31000, + "maxOutputTokens": 8192, + "maxAllowedSize": 31000, + "supportsImages": true, + "supportsToolCall": true, + "advertise": false + }, + { + "id": "default-1.1", + "name": "Claude-3.7-Sonnet", + "credits": "", + "maxInputTokens": 200000, + "maxOutputTokens": 8192, + "maxAllowedSize": 200000, + "supportsImages": true, + "supportsToolCall": true, + "supportsReasoning": true, + "advertise": false + }, + { + "id": "deepseek-r1-0528", + "name": "deepseek-r1", + "credits": "", + "maxInputTokens": 96000, + "maxOutputTokens": 8192, + "maxAllowedSize": 96000, + "advertise": false + }, + { + "id": "deepseek-v3-0324", + "name": "deepseek-v3", + "credits": "", + "maxInputTokens": 96000, + "maxOutputTokens": 8192, + "maxAllowedSize": 96000, + "advertise": false + }, + { + "id": "default-1.2", + "name": "Claude-4.0-Sonnet", + "credits": "", + "maxInputTokens": 200000, + "maxOutputTokens": 24000, + "maxAllowedSize": 200000, + "supportsImages": true, + "supportsToolCall": true, + "supportsReasoning": true, + "advertise": false + }, + { + "id": "hunyuan-2.0-instruct", + "name": "Hunyuan-2.0-Instruct", + "credits": "", + "maxInputTokens": 128000, + "maxOutputTokens": 16000, + "maxAllowedSize": 128000, + "supportsImages": true, + "supportsToolCall": true, + "supportsReasoning": true, + "reasoning": { + "effort": "medium", + "summary": "auto" + }, + "advertise": false + }, + { + "id": "hunyuan-chat", + "name": "Hunyuan-Turbos", + "credits": "", + "maxInputTokens": 128000, + "maxOutputTokens": 8192, + "maxAllowedSize": 128000, + "supportsImages": true, + "supportsToolCall": true, + "advertise": false + } + ] +} diff --git a/internal/workbuddy/dispatch.go b/internal/workbuddy/dispatch.go new file mode 100644 index 00000000000..2710b7a3e9b --- /dev/null +++ b/internal/workbuddy/dispatch.go @@ -0,0 +1,176 @@ +package workbuddy + +import ( + "context" + "fmt" + "net/http" + "sync" + "time" +) + +// DispatchError declares provider-owned retry policy. The SDK must not convert +// model-only caps into generic credential cooldowns or rotate after restriction. +type DispatchError struct { + Reason string + Status int + Stop bool + Until time.Time +} + +func (e *DispatchError) Error() string { return "workbuddy: " + e.Reason } +func (e *DispatchError) StatusCode() int { + if e.Status > 0 { + return e.Status + } + return http.StatusServiceUnavailable +} +func (e *DispatchError) NativeExecutionErrorAction() string { + if e.Stop { + return "stop" + } + return "continue" +} + +func ErrorForFault(f Fault) *DispatchError { + stop := false + switch f.Kind { + case FaultRestricted, FaultContent, FaultModelRetired, FaultProviderUnavailable, FaultEntitlement, FaultTLS, FaultCanceled, FaultRequest: + stop = true + } + status := f.Status + if status == 0 { + status = http.StatusBadGateway + } + return &DispatchError{Reason: string(f.Kind), Status: status, Stop: stop, Until: f.ResetAt} +} + +type requestStateKey struct{} +type RequestState struct { + mu sync.Mutex + Turn string + attempts, max int + charged map[string]bool + refreshed map[string]bool +} + +func WithRequestState(ctx context.Context, accounts int) (context.Context, error) { + if _, ok := ctx.Value(requestStateKey{}).(*RequestState); ok { + return ctx, nil + } + turn, err := NewTurnID() + if err != nil { + return ctx, err + } + if accounts < 2 { + accounts = 2 + } + state := &RequestState{Turn: turn, max: accounts + 1, charged: make(map[string]bool), refreshed: make(map[string]bool)} + return context.WithValue(ctx, requestStateKey{}, state), nil +} +func State(ctx context.Context) *RequestState { + state, _ := ctx.Value(requestStateKey{}).(*RequestState) + return state +} +func (s *RequestState) Begin() error { + s.mu.Lock() + defer s.mu.Unlock() + if s.attempts >= s.max { + return &DispatchError{Reason: "request_attempt_limit", Status: 503, Stop: true} + } + s.attempts++ + return nil +} +func (s *RequestState) Charge(c Credential) bool { + s.mu.Lock() + defer s.mu.Unlock() + key := RiskAccountKey(c) + if s.charged[key] { + return false + } + s.charged[key] = true + return true +} +func (s *RequestState) Refresh(c Credential) bool { + s.mu.Lock() + defer s.mu.Unlock() + key := RiskAccountKey(c) + if s.refreshed[key] { + return false + } + s.refreshed[key] = true + return true +} + +// Limiter bounds both total upstream concurrency and individual caller usage. +// Queue deadlines only govern admission; they never become network deadlines. +type Limiter struct { + mu sync.Mutex + global, perKey int + active int + callers map[string]int + changed chan struct{} +} + +func NewLimiter(global, perKey int) (*Limiter, error) { + if global == 0 { + global = 32 + } + if global < 1 || perKey < 0 { + return nil, fmt.Errorf("workbuddy: invalid concurrency limits") + } + return &Limiter{global: global, perKey: perKey, callers: make(map[string]int), changed: make(chan struct{})}, nil +} +func (l *Limiter) Acquire(ctx context.Context, key string, wait time.Duration) (func(), error) { + queueCtx, cancel := context.WithTimeout(ctx, wait) + defer cancel() + for { + if err := ctx.Err(); err != nil { + return nil, err + } + l.mu.Lock() + if l.active < l.global && (l.perKey == 0 || l.callers[key] < l.perKey) { + l.active++ + l.callers[key]++ + l.mu.Unlock() + var once sync.Once + return func() { + once.Do(func() { + l.mu.Lock() + defer l.mu.Unlock() + l.active-- + l.callers[key]-- + if l.callers[key] == 0 { + delete(l.callers, key) + } + close(l.changed) + l.changed = make(chan struct{}) + }) + }, nil + } + changed := l.changed + l.mu.Unlock() + select { + case <-changed: + case <-queueCtx.Done(): + if err := ctx.Err(); err != nil { + return nil, err + } + return nil, &DispatchError{Reason: "concurrency_queue_full", Status: 503, Stop: true} + } + } +} + +func (l *Limiter) Configure(global, perKey int) error { + if global == 0 { + global = 32 + } + if global < 1 || perKey < 0 { + return fmt.Errorf("workbuddy: invalid concurrency limits") + } + l.mu.Lock() + defer l.mu.Unlock() + l.global, l.perKey = global, perKey + close(l.changed) + l.changed = make(chan struct{}) + return nil +} diff --git a/internal/workbuddy/dispatch_test.go b/internal/workbuddy/dispatch_test.go new file mode 100644 index 00000000000..42799a914fc --- /dev/null +++ b/internal/workbuddy/dispatch_test.go @@ -0,0 +1,100 @@ +package workbuddy + +import ( + "context" + "errors" + "testing" + "time" +) + +func TestWorkBuddyLogicalRequestBudgetAndIdentity(t *testing.T) { + ctx, err := WithRequestState(context.Background(), 2) + if err != nil { + t.Fatal(err) + } + second, err := WithRequestState(ctx, 200) + if err != nil { + t.Fatal(err) + } + s := State(second) + if s != State(ctx) { + t.Fatal("retry replaced request state") + } + for i := 0; i < 3; i++ { + if err := s.Begin(); err != nil { + t.Fatal(err) + } + } + var dispatch *DispatchError + if err := s.Begin(); !errors.As(err, &dispatch) || !dispatch.Stop { + t.Fatalf("budget error = %v", err) + } + c := testCredential() + if !s.Charge(c) || s.Charge(c) { + t.Fatal("failure count repeated") + } + c.Product = CLI + if s.Charge(c) { + t.Fatal("product bypassed UID failure scope") + } + if !s.Refresh(c) || s.Refresh(c) { + t.Fatal("refresh not bounded") + } +} +func TestWorkBuddyConcurrencyAdmissionAndCancellation(t *testing.T) { + limiter, err := NewLimiter(2, 1) + if err != nil { + t.Fatal(err) + } + ctx := context.Background() + releaseA, err := limiter.Acquire(ctx, "a", time.Second) + if err != nil { + t.Fatal(err) + } + releaseB, err := limiter.Acquire(ctx, "b", time.Second) + if err != nil { + t.Fatal(err) + } + releaseB() + releaseB() + if _, err := limiter.Acquire(ctx, "a", 0); err == nil { + t.Fatal("caller cap bypassed") + } + releaseA() + release, err := limiter.Acquire(ctx, "a", 0) + if err != nil { + t.Fatal(err) + } + release() + canceled, cancel := context.WithCancel(ctx) + cancel() + if _, err := limiter.Acquire(canceled, "c", time.Second); !errors.Is(err, context.Canceled) { + t.Fatalf("cancel = %v", err) + } + if limiter.active != 0 || len(limiter.callers) != 0 { + t.Fatal("admission leaked") + } +} +func TestWorkBuddyConcurrencyWakeup(t *testing.T) { + limiter, _ := NewLimiter(1, 0) + release, _ := limiter.Acquire(context.Background(), "a", time.Second) + admitted := make(chan func(), 1) + go func() { + r, err := limiter.Acquire(context.Background(), "b", time.Minute) + if err != nil { + admitted <- nil + return + } + admitted <- r + }() + release() + select { + case r := <-admitted: + if r == nil { + t.Fatal("wait failed") + } + r() + case <-time.After(5 * time.Second): + t.Fatal("waiter not released") + } +} diff --git a/internal/workbuddy/governor.go b/internal/workbuddy/governor.go index cf1b9a40e0b..bb761c6c49f 100644 --- a/internal/workbuddy/governor.go +++ b/internal/workbuddy/governor.go @@ -479,3 +479,41 @@ func (g *Governor) ObserveBalance(c Credential, attempt Attempt, remain float64) func validCredit(value float64) bool { return value >= 0 && !math.IsNaN(value) && !math.IsInf(value, 0) } + +// UpdateLimits applies hot-reloaded budgets without rebuilding cooldown state. +func (g *Governor) UpdateLimits(limits Limits, realms map[Realm]Limits) error { + validate := func(l Limits) bool { + if !validCredit(l.ReserveCredits) || !validCredit(l.DailyCredits) || l.DailyTokens < 0 { + return false + } + for _, n := range l.ModelTokens { + if n < 0 { + return false + } + } + return true + } + clone := func(l Limits) Limits { + m := make(map[string]int64, len(l.ModelTokens)) + for k, v := range l.ModelTokens { + m[k] = v + } + l.ModelTokens = m + return l + } + if !validate(limits) { + return fmt.Errorf("workbuddy: invalid budgets") + } + copyRealms := make(map[Realm]Limits, len(realms)) + for realm, l := range realms { + if (realm != International && realm != China) || !validate(l) { + return fmt.Errorf("workbuddy: invalid regional budgets") + } + copyRealms[realm] = clone(l) + } + g.mu.Lock() + defer g.mu.Unlock() + g.options.Limits = clone(limits) + g.options.Realms = copyRealms + return nil +} diff --git a/internal/workbuddy/stream.go b/internal/workbuddy/stream.go new file mode 100644 index 00000000000..ec6ce9aaec2 --- /dev/null +++ b/internal/workbuddy/stream.go @@ -0,0 +1,335 @@ +package workbuddy + +import ( + "bufio" + "bytes" + "encoding/json" + "fmt" + "io" + "sort" + "strconv" + "strings" +) + +const maxStreamEvent = 16 << 20 +const maxCollectedResponse = 64 << 20 + +// StreamError intentionally excludes upstream payloads and partial user text. +type StreamError struct{ Reason string } + +func (e *StreamError) Error() string { return "workbuddy: " + e.Reason } + +// ReadEvents implements SSE framing, including CRLF, comments and multiline +// data. An EOF without the upstream terminator is never a successful completion. +func ReadEvents(r io.Reader, consume func([]byte) error) error { + scanner := bufio.NewScanner(r) + scanner.Buffer(make([]byte, 4096), maxStreamEvent) + var data []byte + dispatch := func() (bool, error) { + if len(data) == 0 { + return false, nil + } + event := bytes.TrimSuffix(data, []byte("\n")) + if bytes.Equal(bytes.TrimSpace(event), []byte("[DONE]")) { + return true, nil + } + if !json.Valid(event) { + return false, &StreamError{"invalid upstream event"} + } + if err := consume(event); err != nil { + return false, err + } + data = data[:0] + return false, nil + } + for scanner.Scan() { + line := scanner.Bytes() + if len(line) == 0 { + done, err := dispatch() + if err != nil { + return err + } + if done { + return nil + } + continue + } + if bytes.HasPrefix(line, []byte("data:")) { + value := line[5:] + if len(value) > 0 && value[0] == ' ' { + value = value[1:] + } + if len(data)+len(value)+1 > maxStreamEvent { + return &StreamError{"upstream event exceeds limit"} + } + data = append(data, value...) + data = append(data, '\n') + } + } + if err := scanner.Err(); err != nil { + return err + } + // A final data line need not be followed by an empty line. + done, err := dispatch() + if err != nil { + return err + } + if done { + return nil + } + return &StreamError{"upstream stream ended before completion"} +} + +// StreamFault retains risk classification without retaining the error body. +type StreamFault struct{ Fault Fault } + +func (e *StreamFault) Error() string { return "workbuddy: upstream stream " + string(e.Fault.Kind) } + +type chatDelta struct { + Role string `json:"role,omitempty"` + Content *string `json:"content,omitempty"` + Reasoning *string `json:"reasoning_content,omitempty"` + Refusal *string `json:"refusal,omitempty"` + ToolCalls []struct { + Index int `json:"index"` + ID string `json:"id,omitempty"` + Type string `json:"type,omitempty"` + Function struct { + Name string `json:"name,omitempty"` + Arguments string `json:"arguments,omitempty"` + } `json:"function"` + } `json:"tool_calls,omitempty"` +} + +type ChatEvent struct { + ID string `json:"id"` + Model string `json:"model"` + Created int64 `json:"created"` + Choices []struct { + Index int `json:"index"` + Delta chatDelta `json:"delta"` + FinishReason *string `json:"finish_reason"` + Logprobs json.RawMessage `json:"logprobs,omitempty"` + } `json:"choices"` + Usage json.RawMessage `json:"usage,omitempty"` + Error json.RawMessage `json:"error,omitempty"` +} + +// DecodeEvent refuses SSE error envelopes even when the HTTP status was 200. +func DecodeEvent(raw []byte) (ChatEvent, error) { + var event ChatEvent + if err := json.Unmarshal(raw, &event); err != nil { + return event, &StreamError{"invalid upstream chat event"} + } + if len(event.Error) > 0 && string(event.Error) != "null" { + var detail map[string]any + _ = json.Unmarshal(event.Error, &detail) + status := 502 + for _, key := range []string{"status", "status_code", "http_status"} { + if value, ok := detail[key].(float64); ok && value >= 400 && value <= 599 { + status = int(value) + break + } + } + if match := codePattern.FindSubmatch(raw); len(match) == 2 { + code, _ := strconv.Atoi(string(match[1])) + switch code { + case 11140: + status = 403 + case 6004, 14003: + status = 429 + } + } + return event, &StreamFault{Fault: ClassifyHTTP(status, raw)} + } + if event.Choices == nil && len(event.Usage) == 0 { + return event, &StreamError{"unrecognized upstream chat event"} + } + for _, choice := range event.Choices { + if choice.Index < 0 || choice.Index > 127 { + return event, &StreamError{"invalid upstream choice index"} + } + for _, tool := range choice.Delta.ToolCalls { + if tool.Index < 0 || tool.Index > 1023 { + return event, &StreamError{"invalid upstream tool index"} + } + } + } + return event, nil +} + +// ObservedTokens returns only vendor-reported usage. Missing usage is distinct +// from an explicit zero and must not be replaced by an estimate. +func ObservedTokens(raw json.RawMessage) (int64, bool) { + var u struct { + Total *int64 `json:"total_tokens"` + Prompt *int64 `json:"prompt_tokens"` + Completion *int64 `json:"completion_tokens"` + } + if json.Unmarshal(raw, &u) != nil { + return 0, false + } + if u.Total != nil && *u.Total >= 0 { + return *u.Total, true + } + if u.Prompt != nil && u.Completion != nil && *u.Prompt >= 0 && *u.Completion >= 0 && *u.Prompt <= int64(^uint64(0)>>1)-*u.Completion { + return *u.Prompt + *u.Completion, true + } + return 0, false +} + +type collectedTool struct { + ID string `json:"id"` + Type string `json:"type"` + Function struct { + Name string `json:"name"` + Arguments string `json:"arguments"` + } `json:"function"` +} +type collectedChoice struct { + role string + content strings.Builder + reasoning strings.Builder + refusal strings.Builder + tools map[int]*collectedTool + finish *string + logprobs []json.RawMessage +} + +// Collector retains independent choice/tool indices instead of concatenating +// parallel tool calls into one invalid function invocation. +type Collector struct { + id, model string + created int64 + choices map[int]*collectedChoice + usage json.RawMessage + size int +} + +func (c *Collector) Add(raw []byte) error { + if len(raw) > maxCollectedResponse-c.size { + return &StreamError{"collected response exceeds limit"} + } + c.size += len(raw) + event, err := DecodeEvent(raw) + if err != nil { + return err + } + if event.ID != "" { + c.id = event.ID + } + if event.Model != "" { + c.model = event.Model + } + if event.Created != 0 { + c.created = event.Created + } + if len(event.Usage) > 0 && string(event.Usage) != "null" { + c.usage = append(c.usage[:0], event.Usage...) + } + if c.choices == nil { + c.choices = make(map[int]*collectedChoice) + } + for _, delta := range event.Choices { + choice := c.choices[delta.Index] + if choice == nil { + choice = &collectedChoice{role: "assistant", tools: make(map[int]*collectedTool)} + c.choices[delta.Index] = choice + } + if delta.Delta.Role != "" { + choice.role = delta.Delta.Role + } + if delta.Delta.Content != nil { + choice.content.WriteString(*delta.Delta.Content) + } + if delta.Delta.Reasoning != nil { + choice.reasoning.WriteString(*delta.Delta.Reasoning) + } + if delta.Delta.Refusal != nil { + choice.refusal.WriteString(*delta.Delta.Refusal) + } + if delta.FinishReason != nil { + finish := *delta.FinishReason + choice.finish = &finish + } + for _, td := range delta.Delta.ToolCalls { + tool := choice.tools[td.Index] + if tool == nil { + tool = &collectedTool{Type: "function"} + choice.tools[td.Index] = tool + } + if td.ID != "" { + tool.ID = td.ID + } + if td.Type != "" { + tool.Type = td.Type + } + tool.Function.Name += td.Function.Name + tool.Function.Arguments += td.Function.Arguments + } + if len(delta.Logprobs) > 0 && string(delta.Logprobs) != "null" { + choice.logprobs = append(choice.logprobs, append(json.RawMessage(nil), delta.Logprobs...)) + } + } + return nil +} +func (c *Collector) Usage() json.RawMessage { return c.usage } +func (c *Collector) Response() ([]byte, error) { + if len(c.choices) == 0 { + return nil, &StreamError{"upstream completed without a choice"} + } + indexes := make([]int, 0, len(c.choices)) + for i := range c.choices { + indexes = append(indexes, i) + } + sort.Ints(indexes) + choices := make([]any, 0, len(indexes)) + for _, i := range indexes { + choice := c.choices[i] + if choice.finish == nil { + return nil, &StreamError{"upstream completed without a finish reason"} + } + message := map[string]any{"role": choice.role, "content": choice.content.String()} + if choice.reasoning.Len() > 0 { + message["reasoning_content"] = choice.reasoning.String() + } + if choice.refusal.Len() > 0 { + message["refusal"] = choice.refusal.String() + } + if len(choice.tools) > 0 { + ids := make([]int, 0, len(choice.tools)) + for id := range choice.tools { + ids = append(ids, id) + } + sort.Ints(ids) + tools := make([]*collectedTool, 0, len(ids)) + for _, id := range ids { + tools = append(tools, choice.tools[id]) + } + message["tool_calls"] = tools + if choice.content.Len() == 0 { + message["content"] = nil + } + } + out := map[string]any{"index": i, "message": message, "finish_reason": choice.finish} + if len(choice.logprobs) > 0 { + merged := map[string][]json.RawMessage{} + for _, raw := range choice.logprobs { + var part map[string][]json.RawMessage + if err := json.Unmarshal(raw, &part); err != nil { + return nil, fmt.Errorf("workbuddy: invalid logprobs") + } + for key, items := range part { + merged[key] = append(merged[key], items...) + } + } + out["logprobs"] = merged + } + choices = append(choices, out) + } + result := map[string]any{"id": c.id, "object": "chat.completion", "created": c.created, "model": c.model, "choices": choices} + if len(c.usage) > 0 { + result["usage"] = c.usage + } + return json.Marshal(result) +} diff --git a/internal/workbuddy/stream_test.go b/internal/workbuddy/stream_test.go new file mode 100644 index 00000000000..ffd0993014d --- /dev/null +++ b/internal/workbuddy/stream_test.go @@ -0,0 +1,80 @@ +package workbuddy + +import ( + "errors" + "io" + "strings" + "testing" + + "github.com/tidwall/gjson" +) + +func TestWorkBuddySSEFramingAndCollection(t *testing.T) { + input := ": heartbeat\r\ndata: {\r\ndata: \"id\":\"chat-1\",\"model\":\"test\",\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\",\"reasoning_content\":\"think\",\"tool_calls\":[{\"index\":1,\"id\":\"b\",\"function\":{\"name\":\"second\",\"arguments\":\"{\"}},{\"index\":0,\"id\":\"a\",\"function\":{\"name\":\"first\",\"arguments\":\"{}\"}}]}}]}\r\n\r\n" + + `data: {"choices":[{"index":0,"delta":{"tool_calls":[{"index":1,"function":{"arguments":"}"}}]},"finish_reason":"tool_calls"}]}` + "\n\n" + + `data: {"choices":[],"usage":{"prompt_tokens":12,"completion_tokens":3,"total_tokens":15}}` + "\n\n" + "data: [DONE]" + var collector Collector + if err := ReadEvents(strings.NewReader(input), collector.Add); err != nil { + t.Fatal(err) + } + raw, err := collector.Response() + if err != nil { + t.Fatal(err) + } + for path, want := range map[string]string{"choices.0.message.tool_calls.0.id": "a", "choices.0.message.tool_calls.1.id": "b", "choices.0.message.tool_calls.1.function.arguments": "{}", "choices.0.message.reasoning_content": "think", "usage.total_tokens": "15", "choices.0.finish_reason": "tool_calls"} { + if got := gjson.GetBytes(raw, path).String(); got != want { + t.Errorf("%s = %q, want %q", path, got, want) + } + } + if n, ok := ObservedTokens(collector.Usage()); !ok || n != 15 { + t.Fatalf("usage = %d %v", n, ok) + } +} +func TestWorkBuddyStreamFailuresAreNotCompleted(t *testing.T) { + for _, input := range []string{ + "", `data: {"choices":[{"index":0,"delta":{"content":"partial"}}]}` + "\n\n", + "data: not-json\n\n", `data: {"error":{"message":"secret-token"}}` + "\n\ndata: [DONE]\n\n", + `data: {"choices":[{"index":-1,"delta":{}}]}` + "\n\ndata: [DONE]\n\n", + } { + var c Collector + err := ReadEvents(strings.NewReader(input), c.Add) + if err == nil { + t.Fatalf("accepted malformed stream %q", input) + } + if strings.Contains(err.Error(), "secret-token") { + t.Fatal("error leaked upstream body") + } + } + var c Collector + if err := ReadEvents(strings.NewReader(`data: {"choices":[{"index":0,"delta":{"content":"partial"}}]}`+"\n\ndata: [DONE]\n\n"), c.Add); err != nil { + t.Fatal(err) + } + if _, err := c.Response(); err == nil { + t.Fatal("accepted missing finish reason") + } +} +func TestWorkBuddySSEConsumerErrorStopsImmediately(t *testing.T) { + want := errors.New("downstream canceled") + calls := 0 + err := ReadEvents(strings.NewReader("data: {}\n\ndata: {}\n\n"), func([]byte) error { calls++; return want }) + if !errors.Is(err, want) || calls != 1 { + t.Fatalf("err=%v calls=%d", err, calls) + } +} +func TestWorkBuddySSEReadErrorPreserved(t *testing.T) { + err := ReadEvents(errorReader{}, func([]byte) error { return nil }) + if !errors.Is(err, io.ErrUnexpectedEOF) { + t.Fatalf("error = %v", err) + } +} + +type errorReader struct{} + +func (errorReader) Read([]byte) (int, error) { return 0, io.ErrUnexpectedEOF } +func TestWorkBuddyObservedTokensMissingAndOverflow(t *testing.T) { + for _, raw := range []string{`{}`, `{"total_tokens":-1}`, `{"prompt_tokens":9223372036854775807,"completion_tokens":1}`, `null`} { + if n, ok := ObservedTokens([]byte(raw)); ok { + t.Fatalf("invented usage %d for %s", n, raw) + } + } +} diff --git a/sdk/cliproxy/auth/conductor_execution.go b/sdk/cliproxy/auth/conductor_execution.go index 17a1b458ca8..64f21b10958 100644 --- a/sdk/cliproxy/auth/conductor_execution.go +++ b/sdk/cliproxy/auth/conductor_execution.go @@ -131,6 +131,12 @@ func (m *Manager) Execute(ctx context.Context, providers []string, req cliproxye return resp, unwrapExecutionBoundaryError(errHome) } + var errPrepare error + ctx, errPrepare = m.prepareExecutionContext(ctx, normalized, req, opts) + if errPrepare != nil { + return cliproxyexecutor.Response{}, errPrepare + } + defaultRequestRetry, maxRetryCredentials, maxWait := m.retrySettings() var lastErr error @@ -245,6 +251,12 @@ func (m *Manager) ExecuteStream(ctx context.Context, providers []string, req cli return nil, &Error{Code: "provider_not_found", Message: "no provider supplied"} } + var errPrepare error + ctx, errPrepare = m.prepareExecutionContext(ctx, normalized, req, opts) + if errPrepare != nil { + return nil, errPrepare + } + defaultRequestRetry, maxRetryCredentials, maxWait := m.retrySettings() var lastErr error @@ -2175,3 +2187,24 @@ func syncMetadataSessionToContext(ctx context.Context, metadata map[string]any) ctx = logging.WithClientRequestMetadata(ctx, clientMeta) return util.WithSessionID(ctx, clientMeta.SessionID) } + +// prepareExecutionContext initializes optional provider state once, before the +// credential and request retry loops. It does not alter nonparticipating providers. +func (m *Manager) prepareExecutionContext(ctx context.Context, providers []string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (context.Context, error) { + for _, provider := range providers { + executor, ok := m.Executor(provider) + if !ok { + continue + } + if preparer, ok := executor.(interface { + PrepareExecutionContext(context.Context, cliproxyexecutor.Request, cliproxyexecutor.Options) (context.Context, error) + }); ok { + next, err := preparer.PrepareExecutionContext(ctx, req, opts) + if err != nil { + return ctx, err + } + ctx = next + } + } + return ctx, nil +} diff --git a/sdk/cliproxy/auth/conductor_refresh.go b/sdk/cliproxy/auth/conductor_refresh.go index caed3ebc100..68ba50b4dfd 100644 --- a/sdk/cliproxy/auth/conductor_refresh.go +++ b/sdk/cliproxy/auth/conductor_refresh.go @@ -714,6 +714,9 @@ func (m *Manager) tryRefreshAfterUnauthorized(ctx context.Context, auth *Auth, e if m == nil || auth == nil || alreadyTried || execErr == nil { return auth, false } + if _, handled := nativeExecutionErrorAction(execErr); handled { + return auth, false + } // Request-scoped failures describe this request, not stale credentials. // Refreshing would turn a direct error response into an implicit retry. if isRequestScopedError(execErr) { diff --git a/sdk/cliproxy/auth/conductor_request_scoped_errors.go b/sdk/cliproxy/auth/conductor_request_scoped_errors.go index 15801a12812..5bf204011fa 100644 --- a/sdk/cliproxy/auth/conductor_request_scoped_errors.go +++ b/sdk/cliproxy/auth/conductor_request_scoped_errors.go @@ -187,6 +187,9 @@ func matchRequestScopedErrorAction(auth *Auth, err error, cfg *internalconfig.Co if err == nil { return "", false } + if action, ok := nativeExecutionErrorAction(err); ok { + return action, true + } rules := extractRequestScopedErrorRules(auth, cfg) if len(rules) == 0 { return "", false @@ -253,3 +256,16 @@ func applyRequestScopedActionToResult(action string, okAction bool, result *Resu func isRequestScopedStop(action string, okAction bool) bool { return okAction && (action == RequestScopedActionStop || action == RequestScopedActionStopAndCooldown) } + +// nativeExecutionErrorAction is an opt-in policy for executors that already +// own auth refresh, retries and scope-specific cooldowns. +func nativeExecutionErrorAction(err error) (string, bool) { + var native interface{ NativeExecutionErrorAction() string } + if errors.As(err, &native) { + switch action := native.NativeExecutionErrorAction(); action { + case RequestScopedActionStop, RequestScopedActionContinue: + return action, true + } + } + return "", false +} diff --git a/sdk/cliproxy/auth/conductor_workbuddy_test.go b/sdk/cliproxy/auth/conductor_workbuddy_test.go new file mode 100644 index 00000000000..3c5ab4f4125 --- /dev/null +++ b/sdk/cliproxy/auth/conductor_workbuddy_test.go @@ -0,0 +1,103 @@ +package auth + +import ( + "context" + "errors" + "testing" + + "github.com/router-for-me/CLIProxyAPI/v8/internal/registry" + "github.com/router-for-me/CLIProxyAPI/v8/internal/workbuddy" + execution "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/executor" +) + +type workBuddyContextExecutor struct { + mockCustomErrorExecutor + preparations int +} + +func (e *workBuddyContextExecutor) PrepareExecutionContext(ctx context.Context, _ execution.Request, _ execution.Options) (context.Context, error) { + e.preparations++ + return workbuddy.WithRequestState(ctx, 2) +} +func TestWorkBuddyNativeFailureScopeInManager(t *testing.T) { + for _, stop := range []bool{true, false} { + t.Run(map[bool]string{true: "restriction_stops_rotation", false: "model_cap_rotates_without_generic_penalty"}[stop], func(t *testing.T) { + manager := NewManager(nil, nil, nil) + ids := []string{t.Name() + "-a", t.Name() + "-b"} + for _, id := range ids { + a := &Auth{ID: id, Provider: "workbuddy", Status: StatusActive} + if _, err := manager.Register(context.Background(), a); err != nil { + t.Fatal(err) + } + registry.GetGlobalRegistry().RegisterClient(id, "workbuddy", []*registry.ModelInfo{{ID: "fixture-model"}}) + t.Cleanup(func() { registry.GetGlobalRegistry().UnregisterClient(id) }) + } + calls := 0 + var firstState *workbuddy.RequestState + e := &workBuddyContextExecutor{mockCustomErrorExecutor: mockCustomErrorExecutor{identifier: "workbuddy"}} + e.executeFn = func(ctx context.Context, a *Auth, _ execution.Request, _ execution.Options) (execution.Response, error) { + calls++ + execution.MarkUpstreamAttempt(ctx) + state := workbuddy.State(ctx) + if state == nil { + t.Error("request context not prepared") + } + if firstState == nil { + firstState = state + } else if state != firstState { + t.Error("rotation replaced logical request") + } + if calls == 1 { + return execution.Response{}, &workbuddy.DispatchError{Reason: "fixture", Status: 429, Stop: stop} + } + return execution.Response{Payload: []byte(`{"ok":true}`)}, nil + } + manager.RegisterExecutor(e) + _, err := manager.Execute(context.Background(), []string{"workbuddy"}, execution.Request{Model: "fixture-model"}, execution.Options{}) + if stop { + var native *workbuddy.DispatchError + if !errors.As(err, &native) || calls != 1 { + t.Fatalf("stop: calls=%d err=%v", calls, err) + } + } else if err != nil || calls != 2 { + t.Fatalf("continue: calls=%d err=%v", calls, err) + } + if e.preparations != 1 { + t.Fatalf("context preparations = %d", e.preparations) + } + for _, id := range ids { + a, ok := manager.GetByID(id) + if !ok { + t.Fatal("auth disappeared") + } + if a.Unavailable || !a.NextRetryAfter.IsZero() { + t.Fatal("generic account cooldown replaced native scope") + } + } + }) + } +} + +func TestWorkBuddyNativeAuthFailureDoesNotRefreshAgain(t *testing.T) { + manager := NewManager(nil, nil, nil) + a := &Auth{ID: "wb-refresh-once", Provider: "workbuddy", Status: StatusActive, Metadata: map[string]any{"access_token": "fixture", "refresh_token": "fixture-refresh"}} + if _, err := manager.Register(context.Background(), a); err != nil { + t.Fatal(err) + } + executor := &workBuddyRefreshTrap{mockCustomErrorExecutor: mockCustomErrorExecutor{identifier: "workbuddy"}} + manager.RegisterExecutor(executor) + _, refreshed := manager.tryRefreshAfterUnauthorized(context.Background(), a, &workbuddy.DispatchError{Reason: "authentication", Status: 401}, false) + if refreshed || executor.calls != 0 { + t.Fatal("SDK repeated provider-owned refresh") + } +} + +type workBuddyRefreshTrap struct { + mockCustomErrorExecutor + calls int +} + +func (e *workBuddyRefreshTrap) Refresh(context.Context, *Auth) (*Auth, error) { + e.calls++ + return nil, errors.New("unexpected refresh") +} diff --git a/sdk/cliproxy/service_executors.go b/sdk/cliproxy/service_executors.go index 2c2083d959c..872cd9fabdd 100644 --- a/sdk/cliproxy/service_executors.go +++ b/sdk/cliproxy/service_executors.go @@ -323,6 +323,19 @@ func (s *Service) registerExecutorForAuth(a *coreauth.Auth, forceReplace bool) { } } s.coreManager.RegisterExecutor(executor.NewXAIAutoExecutor(cfg)) + case "workbuddy": + var previous *executor.WorkBuddyExecutor + if current, ok := s.coreManager.Executor("workbuddy"); ok { + previous, _ = current.(*executor.WorkBuddyExecutor) + if previous != nil && previous.UsesConfig(cfg) { + previous.TrackAuth(a) + return + } + } + next := executor.NewWorkBuddyExecutor(cfg, previous) + next.BindCredentialManager(s.coreManager) + next.TrackAuth(a) + s.coreManager.RegisterExecutor(next) case "devin": s.coreManager.RegisterExecutor(executor.NewDevinExecutor(cfg)) case "meta": diff --git a/sdk/cliproxy/service_models.go b/sdk/cliproxy/service_models.go index 7e478af9420..91919057dc6 100644 --- a/sdk/cliproxy/service_models.go +++ b/sdk/cliproxy/service_models.go @@ -9,6 +9,8 @@ import ( "github.com/router-for-me/CLIProxyAPI/v8/internal/constant" "github.com/router-for-me/CLIProxyAPI/v8/internal/modelconfig" "github.com/router-for-me/CLIProxyAPI/v8/internal/registry" + "github.com/router-for-me/CLIProxyAPI/v8/internal/runtime/executor/helps" + "github.com/router-for-me/CLIProxyAPI/v8/internal/workbuddy" coreauth "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/auth" "github.com/router-for-me/CLIProxyAPI/v8/sdk/config" ) @@ -196,6 +198,36 @@ func (s *Service) registerModelsForAuthWithConfigAndCache(ctx context.Context, c } } models = applyExcludedModels(models, excluded) + case "workbuddy": + if cfg != nil && cfg.WorkBuddy.Enabled { + if credential, err := helps.WorkBuddyCredential(a); err == nil { + if len(credential.Models) > 0 { + for _, id := range credential.Models { + models = append(models, ®istry.ModelInfo{ID: id, Object: "model", OwnedBy: "workbuddy", Type: "workbuddy", UserDefined: true}) + } + } else { + for _, model := range workbuddy.Catalog(credential.Realm) { + if !model.Advertise { + continue + } + modalities := []string{"TEXT"} + if model.Images { + modalities = append(modalities, "IMAGE") + } + info := ®istry.ModelInfo{ID: model.ID, Object: "model", OwnedBy: "workbuddy", Type: "workbuddy", DisplayName: model.Name, ContextLength: model.Context, MaxCompletionTokens: model.Output, InputTokenLimit: model.Input, OutputTokenLimit: model.Output, SupportedInputModalities: modalities} + if model.Reasoning { + levels := append([]string(nil), model.Thinking.Levels...) + if len(levels) == 0 && model.Thinking.Effort != "" { + levels = []string{model.Thinking.Effort} + } + info.Thinking = ®istry.ThinkingSupport{Levels: levels, ZeroAllowed: model.Thinking.CanDisable == nil || *model.Thinking.CanDisable} + } + models = append(models, info) + } + } + } + } + models = applyExcludedModels(models, excluded) case "devin": models = registry.GetDevinModels() models = applyExcludedModels(models, excluded) diff --git a/sdk/cliproxy/service_workbuddy_test.go b/sdk/cliproxy/service_workbuddy_test.go new file mode 100644 index 00000000000..ca3938395fa --- /dev/null +++ b/sdk/cliproxy/service_workbuddy_test.go @@ -0,0 +1,46 @@ +package cliproxy + +import ( + "context" + "testing" + + internalconfig "github.com/router-for-me/CLIProxyAPI/v8/internal/config" + "github.com/router-for-me/CLIProxyAPI/v8/internal/workbuddy" + sdkauth "github.com/router-for-me/CLIProxyAPI/v8/sdk/auth" +) + +func TestWorkBuddyModelsRequireOptInAndRespectRealm(t *testing.T) { + c := workbuddy.Credential{Type: "workbuddy", UID: "fixture-model-user", AccessToken: "fixture-only", Realm: workbuddy.China, Product: workbuddy.Desktop} + a, err := sdkauth.WorkBuddyAuthRecord(c) + if err != nil { + t.Fatal(err) + } + registry := GlobalModelRegistry() + t.Cleanup(func() { registry.UnregisterClient(a.ID) }) + cfg := &internalconfig.Config{} + service := &Service{cfg: cfg} + service.registerModelsForAuth(context.Background(), a) + if len(registry.GetModelsForClient(a.ID)) != 0 { + t.Fatal("disabled provider advertised models") + } + cfg.WorkBuddy.Enabled = true + service.registerModelsForAuth(context.Background(), a) + models := registry.GetModelsForClient(a.ID) + if len(models) == 0 { + t.Fatal("no CN fallback catalog") + } + for _, m := range models { + if m.ID == "hy3" && (m.Thinking == nil || len(m.Thinking.Levels) == 0) { + t.Fatal("catalog reasoning capabilities were lost") + } + if m.ID == "gpt-6-astra" { + t.Fatal("CN catalog leaked intl model") + } + } + a.Metadata["models"] = []string{"fixture-model"} + service.registerModelsForAuth(context.Background(), a) + models = registry.GetModelsForClient(a.ID) + if len(models) != 1 || models[0].ID != "fixture-model" { + t.Fatalf("explicit models = %+v", models) + } +} From 7ce419c5507da409054b3983937e68ef967868af Mon Sep 17 00:00:00 2001 From: LuneZ99 Date: Wed, 21 Oct 2026 02:26:12 +0800 Subject: [PATCH 04/10] fix(workbuddy): preserve model entitlement evidence without account penalties --- docs/workbuddy-native.md | 10 ++- .../runtime/executor/workbuddy_executor.go | 8 ++- internal/workbuddy/dispatch.go | 13 +++- internal/workbuddy/fault.go | 27 ++++++-- internal/workbuddy/governor.go | 50 +++++++++++++- internal/workbuddy/governor_test.go | 66 +++++++++++++++++++ 6 files changed, 161 insertions(+), 13 deletions(-) diff --git a/docs/workbuddy-native.md b/docs/workbuddy-native.md index 150afe88285..a0945905235 100644 --- a/docs/workbuddy-native.md +++ b/docs/workbuddy-native.md @@ -65,9 +65,15 @@ state directory requires a restart; the old directory must be migrated explicitl Only one process may own a state directory. The isolated deployment will use a single replica and its own PVC. +Model entitlement error `11158` is retained as region/model evidence for 24 hours, +including the reported user tier when known. As in the reference, this is not a +blanket dispatch ban: a free account's refusal cannot establish whether a paid +account may use the model. A later successful check clears the mark; older +in-flight completions cannot erase newer evidence. Model configuration errors +`11129` and `11150` also stop rotation without penalizing the account. + The remaining stage-2 work includes live model discovery, conversation paging, -affinity policy, content-rejection degradation, region/model entitlement holds, -optional exit backoff, identity switching and paid fallback. Stage-3 quota and +affinity policy, content-rejection degradation, optional exit backoff, identity switching and paid fallback. Stage-3 quota and activity jobs are not implemented or enabled yet. These gaps must be closed before claiming parity with the reference project's risk controls. diff --git a/internal/runtime/executor/workbuddy_executor.go b/internal/runtime/executor/workbuddy_executor.go index f284c12e367..af3f05d2ffc 100644 --- a/internal/runtime/executor/workbuddy_executor.go +++ b/internal/runtime/executor/workbuddy_executor.go @@ -343,12 +343,16 @@ func (e *WorkBuddyExecutor) recordFailure(c workbuddy.Credential, model string, if fault.Kind == workbuddy.FaultHard || fault.Kind == workbuddy.FaultUnknown { charge = state.Charge(c) } + failure := workbuddy.ErrorForFault(fault) if charge { - if _, err := e.runtime.Governor.RecordFailure(c, model, fault, workbuddy.FailureOptions{FreeModel: workbuddy.IsFreeCapModel(c.Realm, model), SingleAccount: e.runtime.Accounts() <= 1}); err != nil { + gate, err := e.runtime.Governor.RecordFailure(c, model, fault, workbuddy.FailureOptions{FreeModel: workbuddy.IsFreeCapModel(c.Realm, model), SingleAccount: e.runtime.Accounts() <= 1}) + if err != nil { return &workbuddy.DispatchError{Reason: "risk_persistence_failed", Status: 503, Stop: true} } + failure.Until = gate.Until + failure.Stop = failure.Stop || gate.Stop } - return workbuddy.ErrorForFault(fault) + return failure } func closeWorkBuddyBody(body io.Closer) { if err := body.Close(); err != nil { diff --git a/internal/workbuddy/dispatch.go b/internal/workbuddy/dispatch.go index 2710b7a3e9b..e981470a7a2 100644 --- a/internal/workbuddy/dispatch.go +++ b/internal/workbuddy/dispatch.go @@ -24,6 +24,17 @@ func (e *DispatchError) StatusCode() int { } return http.StatusServiceUnavailable } +func (e *DispatchError) RetryAfter() *time.Duration { + if e.Until.IsZero() { + return nil + } + wait := time.Until(e.Until) + if wait < 0 { + wait = 0 + } + return &wait +} + func (e *DispatchError) NativeExecutionErrorAction() string { if e.Stop { return "stop" @@ -34,7 +45,7 @@ func (e *DispatchError) NativeExecutionErrorAction() string { func ErrorForFault(f Fault) *DispatchError { stop := false switch f.Kind { - case FaultRestricted, FaultContent, FaultModelRetired, FaultProviderUnavailable, FaultEntitlement, FaultTLS, FaultCanceled, FaultRequest: + case FaultRestricted, FaultContent, FaultModelRetired, FaultProviderUnavailable, FaultEntitlement, FaultModelParameter, FaultTLS, FaultCanceled, FaultRequest: stop = true } status := f.Status diff --git a/internal/workbuddy/fault.go b/internal/workbuddy/fault.go index bdf082a7f7b..dd3009543d1 100644 --- a/internal/workbuddy/fault.go +++ b/internal/workbuddy/fault.go @@ -28,6 +28,7 @@ const ( FaultModelRetired FaultKind = "model_retired" FaultProviderUnavailable FaultKind = "provider_unavailable" FaultEntitlement FaultKind = "model_entitlement" + FaultModelParameter FaultKind = "model_parameter" FaultGateway FaultKind = "gateway" FaultHard FaultKind = "hard_failure" FaultUnknown FaultKind = "unknown_failure" @@ -41,11 +42,12 @@ const ( // Fault contains only classification evidence. Upstream bodies can contain // credentials or user content and must not enter persisted risk state. type Fault struct { - Kind FaultKind - Status int - Code int - ResetAt time.Time - Soft bool + Kind FaultKind + Status int + Code int + ResetAt time.Time + Soft bool + UserType string } var ( @@ -167,8 +169,12 @@ func ClassifyHTTP(status int, raw []byte) Fault { case 11134: fault.Kind = FaultProviderUnavailable return fault + case 11129, 11150: + fault.Kind = FaultModelParameter + return fault case 11158: fault.Kind = FaultEntitlement + fault.UserType = entitlementUserType(raw) return fault } } @@ -205,3 +211,14 @@ func ClassifyTransport(err error) Fault { } return Fault{Kind: FaultUnknown} } + +// entitlementUserType retains only known tier labels, never arbitrary messages. +func entitlementUserType(raw []byte) string { + text := faultText(raw) + for _, tier := range []string{"personal_free_user", "personal_paid_user", "enterprise_user"} { + if strings.Contains(text, tier) { + return tier + } + } + return "" +} diff --git a/internal/workbuddy/governor.go b/internal/workbuddy/governor.go index bb761c6c49f..8889a1d32e1 100644 --- a/internal/workbuddy/governor.go +++ b/internal/workbuddy/governor.go @@ -60,10 +60,20 @@ type realmRisk struct { Events map[string]time.Time `json:"events"` } +// EntitlementMark records a refusal observed for a region/model and reported +// user tier. It is evidence, not a region-wide permission decision: one free +// account's refusal must not block paid accounts in the same region. +type EntitlementMark struct { + CheckedAt time.Time `json:"checked_at"` + Until time.Time `json:"until"` + UserType string `json:"user_type,omitempty"` +} + type governorState struct { - Version int `json:"version"` - Accounts map[string]*accountRisk `json:"accounts"` - Realms map[Realm]*realmRisk `json:"realms"` + Version int `json:"version"` + Entitlements map[Realm]map[string]EntitlementMark `json:"entitlements,omitempty"` + Accounts map[string]*accountRisk `json:"accounts"` + Realms map[Realm]*realmRisk `json:"realms"` } type Governor struct { @@ -160,6 +170,19 @@ func OpenGovernor(path string, options GovernorOptions) (*Governor, error) { } else if !errors.Is(errRead, os.ErrNotExist) { return nil, fmt.Errorf("workbuddy: read risk state: %w", errRead) } + if g.state.Entitlements == nil { + g.state.Entitlements = make(map[Realm]map[string]EntitlementMark) + } + for realm, models := range g.state.Entitlements { + if realm != China && realm != International { + return nil, fmt.Errorf("workbuddy: invalid entitlement region") + } + for model, mark := range models { + if model == "" || mark.CheckedAt.IsZero() || mark.Until.Before(mark.CheckedAt) { + return nil, fmt.Errorf("workbuddy: invalid entitlement evidence") + } + } + } if err := g.persistLocked(); err != nil { return nil, err } @@ -398,6 +421,14 @@ func (g *Governor) RecordFailure(c Credential, model string, fault Fault, option state.CooldownUntil = later(state.CooldownUntil, now.Add(wait)) gate.Until = later(gate.Until, state.CooldownUntil) } + case FaultEntitlement: + models := g.state.Entitlements[c.Realm] + if models == nil { + models = make(map[string]EntitlementMark) + g.state.Entitlements[c.Realm] = models + } + models[model] = EntitlementMark{CheckedAt: now, Until: now.Add(24 * time.Hour), UserType: fault.UserType} + gate.Stop = true case FaultAuth: // Called only after the request's one refresh attempt has failed. wait := time.Minute @@ -446,6 +477,11 @@ func (g *Governor) Complete(c Credential, attempt Attempt, tokens int64, credits if credits != nil { state.CreditsSpent += *credits } + if success { + if mark, ok := g.state.Entitlements[c.Realm][attempt.Model]; ok && !attempt.StartedAt.Before(mark.CheckedAt) { + delete(g.state.Entitlements[c.Realm], attempt.Model) + } + } if success && attempt.Version == state.Version { state.RestrictedUntil, state.CooldownUntil, state.BreakerUntil, state.DegradeUntil = time.Time{}, time.Time{}, time.Time{}, time.Time{} state.LastFailure, state.LastSoft = time.Time{}, time.Time{} @@ -517,3 +553,11 @@ func (g *Governor) UpdateLimits(limits Limits, realms map[Realm]Limits) error { g.options.Realms = copyRealms return nil } + +// Entitlement returns fresh refusal evidence without changing dispatch policy. +func (g *Governor) Entitlement(realm Realm, model string) (EntitlementMark, bool) { + g.mu.Lock() + defer g.mu.Unlock() + mark, ok := g.state.Entitlements[realm][model] + return mark, ok && g.options.Now().Before(mark.Until) +} diff --git a/internal/workbuddy/governor_test.go b/internal/workbuddy/governor_test.go index 8943316d20b..6c6e04832a0 100644 --- a/internal/workbuddy/governor_test.go +++ b/internal/workbuddy/governor_test.go @@ -242,3 +242,69 @@ func TestGovernorInvalidStateAndPersistenceFailureStopDispatch(t *testing.T) { t.Fatal("unknown risk schema accepted") } } + +func TestGovernorEntitlementEvidenceIsScopedAndExpires(t *testing.T) { + g, now, path := governorFixture(t, Limits{}) + free := testCredential() + paid := free + paid.UID = "paid-fixture-user" + paid.EnterpriseID = "fixture-enterprise" + _, older := g.Check(paid, "model", true, false) + *now = now.Add(time.Second) + fault := ClassifyHTTP(500, []byte(`{"code":11158,"message":"model is not allowed for user type [personal_free_user]"}`)) + if gate, err := g.RecordFailure(free, "model", fault, FailureOptions{}); err != nil || !gate.Stop { + t.Fatalf("record: %+v %v", gate, err) + } + restarted, err := OpenGovernor(path, g.options) + if err != nil { + t.Fatal(err) + } + mark, ok := restarted.Entitlement(International, "model") + if !ok || mark.UserType != "personal_free_user" { + t.Fatalf("mark = %+v %v", mark, ok) + } + if gate, _ := restarted.Check(paid, "model", true, false); !gate.Allowed() { + t.Fatal("free refusal blocked paid account") + } + if gate, _ := restarted.Check(free, "sibling", true, false); !gate.Allowed() { + t.Fatal("model refusal penalized sibling model") + } + if err := restarted.Complete(paid, older, 0, nil, true); err != nil { + t.Fatal(err) + } + if _, ok := restarted.Entitlement(International, "model"); !ok { + t.Fatal("stale completion erased newer refusal evidence") + } + _, fresh := restarted.Check(paid, "model", true, false) + if err := restarted.Complete(paid, fresh, 0, nil, true); err != nil { + t.Fatal(err) + } + if _, ok := restarted.Entitlement(International, "model"); ok { + t.Fatal("successful recheck did not clear mark") + } + if _, err := restarted.RecordFailure(free, "model", fault, FailureOptions{}); err != nil { + t.Fatal(err) + } + *now = now.Add(24 * time.Hour) + if _, ok := restarted.Entitlement(International, "model"); ok { + t.Fatal("entitlement mark did not expire") + } +} +func TestGovernorModelConfigurationFailuresDoNotPoisonAccount(t *testing.T) { + g, _, _ := governorFixture(t, Limits{}) + c := testCredential() + for _, code := range []string{"11129", "11150"} { + fault := ClassifyHTTP(550, []byte(`{"code":`+code+`}`)) + if fault.Kind != FaultModelParameter || !ErrorForFault(fault).Stop { + t.Fatalf("code %s classified %+v", code, fault) + } + for i := 0; i < 8; i++ { + if _, err := g.RecordFailure(c, "model", fault, FailureOptions{}); err != nil { + t.Fatal(err) + } + } + } + if gate, _ := g.Check(c, "sibling", true, false); !gate.Allowed() { + t.Fatalf("configuration failure poisoned account: %+v", gate) + } +} From c81e2ae812295e888b7924096f3f69f2506c59d1 Mon Sep 17 00:00:00 2001 From: LuneZ99 Date: Wed, 21 Oct 2026 02:44:59 +0800 Subject: [PATCH 05/10] feat(workbuddy): add bounded recovery policies and session affinity --- docs/workbuddy-native.md | 22 +- internal/config/workbuddy.go | 14 ++ .../executor/helps/workbuddy_payload.go | 73 +++++- .../executor/helps/workbuddy_runtime.go | 88 ++++++- .../runtime/executor/workbuddy_executor.go | 224 +++++++++++++++++- .../executor/workbuddy_executor_test.go | 165 +++++++++++++ internal/workbuddy/affinity.go | 216 +++++++++++++++++ internal/workbuddy/affinity_test.go | 76 ++++++ internal/workbuddy/dispatch.go | 63 ++++- internal/workbuddy/governor.go | 49 +++- internal/workbuddy/governor_test.go | 14 ++ internal/workbuddy/policy.go | 186 +++++++++++++++ internal/workbuddy/policy_test.go | 125 ++++++++++ sdk/cliproxy/auth/conductor.go | 6 + sdk/cliproxy/auth/conductor_selection.go | 51 +++- sdk/cliproxy/auth/conductor_workbuddy_test.go | 45 ++++ 16 files changed, 1382 insertions(+), 35 deletions(-) create mode 100644 internal/workbuddy/affinity.go create mode 100644 internal/workbuddy/affinity_test.go create mode 100644 internal/workbuddy/policy.go create mode 100644 internal/workbuddy/policy_test.go diff --git a/docs/workbuddy-native.md b/docs/workbuddy-native.md index a0945905235..f318a7653a2 100644 --- a/docs/workbuddy-native.md +++ b/docs/workbuddy-native.md @@ -72,10 +72,24 @@ account may use the model. A later successful check clears the mark; older in-flight completions cannot erase newer evidence. Model configuration errors `11129` and `11150` also stop rotation without penalizing the account. -The remaining stage-2 work includes live model discovery, conversation paging, -affinity policy, content-rejection degradation, optional exit backoff, identity switching and paid fallback. Stage-3 quota and -activity jobs are not implemented or enabled yet. These gaps must be closed -before claiming parity with the reference project's risk controls. +Conversation affinity keeps three recent eligible accounts per caller/session +for two hours. Conversations beyond 400 messages use stable account pages +(default size three); cooldowns do not change page membership. Native selection +still honors SDK pins, exclusions, priority and the final governor check. + +Content rejection can replace the built-in prompt with a neutral prompt until +CST midnight, with one retry per request. Custom prompts are excluded. Optional +exit backoff shares failures by origin/proxy and caps waiting at five seconds +per logical request. Optional product switching persists the new identity before +retrying and permits at most four switches per account/model per minute. +Restrictions and durable free-model caps survive identity changes. Optional paid +fallback applies only to international DeepSeek Flash when every enabled account +in that region has an active free-model cap; paid budgets are checked again. +All rebuilt requests pass through the final payload-rule barrier exactly once. + +The remaining stage-2 work is live model discovery. Stage-3 quota and activity +jobs are not implemented or enabled yet. These gaps must be closed before +claiming parity with the reference project's risk controls. An explicit v8 configuration enables the provider: diff --git a/internal/config/workbuddy.go b/internal/config/workbuddy.go index cb6b9013960..b451d77f5f5 100644 --- a/internal/config/workbuddy.go +++ b/internal/config/workbuddy.go @@ -3,6 +3,14 @@ package config // WorkBuddyConfig is opt-in. State belongs to a dedicated provider volume and // must never be shared between independent server processes. type WorkBuddyConfig struct { + AffinityByPrefix *bool `yaml:"affinity-by-prefix" json:"affinity-by-prefix"` + AffinityMaxMessages *int `yaml:"affinity-max-messages" json:"affinity-max-messages"` + AffinityPageSize *int `yaml:"affinity-page-size" json:"affinity-page-size"` + EgressBackoff bool `yaml:"egress-backoff" json:"egress-backoff"` + AutoSwitchProduct bool `yaml:"auto-switch-product" json:"auto-switch-product"` + AutoSwitchProducts []string `yaml:"auto-switch-products" json:"auto-switch-products"` + PaidFallback bool `yaml:"paid-fallback" json:"paid-fallback"` + Prompt WorkBuddyPrompt `yaml:"prompt" json:"prompt"` Enabled bool `yaml:"enabled" json:"enabled"` StateDir string `yaml:"state-dir" json:"state-dir"` MaxConcurrency int `yaml:"max-concurrency" json:"max-concurrency"` @@ -18,3 +26,9 @@ type WorkBuddyLimits struct { DailyCredits float64 `yaml:"daily-credits" json:"daily-credits"` ModelTokens map[string]int64 `yaml:"model-tokens" json:"model-tokens"` } + +type WorkBuddyPrompt struct { + Mode string `yaml:"mode" json:"mode"` + Text string `yaml:"text" json:"text"` + File string `yaml:"file" json:"file"` +} diff --git a/internal/runtime/executor/helps/workbuddy_payload.go b/internal/runtime/executor/helps/workbuddy_payload.go index d392c92840c..62ba470befd 100644 --- a/internal/runtime/executor/helps/workbuddy_payload.go +++ b/internal/runtime/executor/helps/workbuddy_payload.go @@ -4,6 +4,9 @@ import ( "bytes" "encoding/json" "fmt" + "github.com/router-for-me/CLIProxyAPI/v8/internal/config" + "io" + "os" "strings" "github.com/router-for-me/CLIProxyAPI/v8/internal/workbuddy" @@ -13,7 +16,7 @@ const workBuddySystemPrompt = "You are a helpful coding assistant." // NormalizeWorkBuddyPayload runs before the payload finalizer. It intentionally // preserves caller text, tool schemas and unknown public fields. -func NormalizeWorkBuddyPayload(raw []byte, c workbuddy.Credential, model string) ([]byte, error) { +func NormalizeWorkBuddyPayload(raw []byte, c workbuddy.Credential, model string, prompt ...WorkBuddyPromptPolicy) ([]byte, error) { if len(raw) > 16<<20 { return nil, &workbuddy.DispatchError{Reason: "request_body_too_large", Status: 413, Stop: true} } @@ -48,6 +51,9 @@ func NormalizeWorkBuddyPayload(raw []byte, c workbuddy.Credential, model string) } } } + if len(prompt) > 0 { + messages = applyWorkBuddyPrompt(messages, prompt[0]) + } if len(messages) == 0 || messages[0].(map[string]any)["role"] != "system" { messages = append([]any{map[string]any{"role": "system", "content": workBuddySystemPrompt}}, messages...) } @@ -182,3 +188,68 @@ func repairWorkBuddyTools(messages []any) []any { } return out } + +type WorkBuddyPromptPolicy struct { + Mode, Text string + Degraded bool +} + +// WorkBuddyPrompt resolves an explicitly configured local prompt file. A failed +// read preserves the caller's prompt, matching the reference's fail-open mode. +func WorkBuddyPrompt(cfg config.WorkBuddyPrompt, degraded bool) WorkBuddyPromptPolicy { + mode := cfg.Mode + if mode == "" { + mode = "passthrough" + } + policy := WorkBuddyPromptPolicy{Mode: mode, Text: cfg.Text, Degraded: degraded} + if mode == "custom" || mode == "append" { + if cfg.File != "" { + file, err := os.Open(cfg.File) + if err != nil { + return WorkBuddyPromptPolicy{Mode: "passthrough"} + } + raw, errRead := io.ReadAll(io.LimitReader(file, (1<<20)+1)) + errClose := file.Close() + if errRead != nil || errClose != nil || len(raw) > 1<<20 || len(raw) == 0 { + return WorkBuddyPromptPolicy{Mode: "passthrough"} + } + policy.Text = string(raw) + } + if policy.Text == "" { + policy.Text = "You are an engineering assistant. Understand the code and context, follow existing conventions, make focused changes, verify results, and report uncertainty honestly." + } + } + return policy +} +func applyWorkBuddyPrompt(messages []any, policy WorkBuddyPromptPolicy) []any { + text := policy.Text + rewrite := policy.Mode == "custom" + if policy.Mode != "custom" && policy.Degraded { + rewrite = true + text = workbuddy.DegradedPrompt + } + if rewrite && text != "" { + out := []any{map[string]any{"role": "system", "content": text}} + for _, item := range messages { + m := item.(map[string]any) + if m["role"] != "system" && m["role"] != "developer" { + out = append(out, item) + } + } + return out + } + if policy.Mode == "append" && text != "" { + i := 0 + for i < len(messages) { + role := messages[i].(map[string]any)["role"] + if role != "system" && role != "developer" { + break + } + i++ + } + out := append([]any(nil), messages[:i]...) + out = append(out, map[string]any{"role": "system", "content": text}) + return append(out, messages[i:]...) + } + return messages +} diff --git a/internal/runtime/executor/helps/workbuddy_runtime.go b/internal/runtime/executor/helps/workbuddy_runtime.go index d301adefa5d..b9041110f64 100644 --- a/internal/runtime/executor/helps/workbuddy_runtime.go +++ b/internal/runtime/executor/helps/workbuddy_runtime.go @@ -19,6 +19,8 @@ type WorkBuddyRuntime struct { Governor *workbuddy.Governor Transports *workbuddy.Transports Limiter *workbuddy.Limiter + Policy *workbuddy.PolicyState + Affinity *workbuddy.Affinity // BaseURL is an injection seam for offline tests, never populated from auth files. BaseURL string StateDir string @@ -33,7 +35,7 @@ func NewWorkBuddyRuntime(cfg config.WorkBuddyConfig) (*WorkBuddyRuntime, error) if cfg.StateDir == "" { return nil, fmt.Errorf("workbuddy: state-dir is required") } - r := &WorkBuddyRuntime{Transports: workbuddy.NewTransports(256, nil), StateDir: filepath.Clean(cfg.StateDir), auths: make(map[string]*coreauth.Auth)} + r := &WorkBuddyRuntime{Affinity: workbuddy.NewAffinity(nil), Policy: workbuddy.NewPolicyState(nil), Transports: workbuddy.NewTransports(256, nil), StateDir: filepath.Clean(cfg.StateDir), auths: make(map[string]*coreauth.Auth)} var err error r.Governor, err = workbuddy.OpenGovernor(filepath.Join(r.StateDir, "risk.json"), workbuddy.GovernorOptions{}) if err != nil { @@ -80,6 +82,19 @@ func (r *WorkBuddyRuntime) Configure(cfg config.WorkBuddyConfig) error { if filepath.Clean(cfg.StateDir) != r.StateDir { return fmt.Errorf("workbuddy: changing state-dir requires a process restart") } + for _, product := range cfg.AutoSwitchProducts { + if product != string(workbuddy.Desktop) && product != string(workbuddy.VSCode) && product != string(workbuddy.CLI) { + return fmt.Errorf("workbuddy: invalid automatic product selection") + } + } + if cfg.AffinityMaxMessages != nil && *cfg.AffinityMaxMessages < 0 || cfg.AffinityPageSize != nil && *cfg.AffinityPageSize < 0 { + return fmt.Errorf("workbuddy: affinity limits cannot be negative") + } + switch cfg.Prompt.Mode { + case "", "passthrough", "custom", "append": + default: + return fmt.Errorf("workbuddy: invalid prompt mode") + } if cfg.QueueSeconds < 0 { return fmt.Errorf("workbuddy: queue-seconds cannot be negative") } @@ -148,17 +163,13 @@ func (r *WorkBuddyRuntime) Current(a *coreauth.Auth) *coreauth.Auth { return a } func (r *WorkBuddyRuntime) Accounts() int { - r.mu.Lock() - defer r.mu.Unlock() unique := make(map[string]bool) - for _, a := range r.auths { - c, err := WorkBuddyCredential(a) - if err == nil && !c.Disabled { - unique[workbuddy.RiskAccountKey(c)] = true - } + for _, c := range r.Credentials() { + unique[workbuddy.RiskAccountKey(c)] = true } return len(unique) } + func (r *WorkBuddyRuntime) Refresh(ctx context.Context, a *coreauth.Auth, proxy string) (*coreauth.Auth, error) { a = r.Current(a) c, err := WorkBuddyCredential(a) @@ -211,3 +222,64 @@ func (r *WorkBuddyRuntime) Refresh(ctx context.Context, a *coreauth.Auth, proxy } return current, nil } + +func (r *WorkBuddyRuntime) Credentials() []workbuddy.Credential { + r.mu.Lock() + manager := r.manager + var records []*coreauth.Auth + if manager == nil { + for _, a := range r.auths { + records = append(records, a.Clone()) + } + } + r.mu.Unlock() + if manager != nil { + records = manager.List() + } + out := make([]workbuddy.Credential, 0, len(records)) + for _, a := range records { + if c, err := WorkBuddyCredential(a); err == nil && !c.Disabled { + out = append(out, c) + } + } + return out +} +func (r *WorkBuddyRuntime) AllFreeCapped(realm workbuddy.Realm, model string) bool { + found := false + for _, c := range r.Credentials() { + if c.Realm != realm { + continue + } + found = true + if !r.Governor.Capped(c, model) { + return false + } + } + return found +} +func (r *WorkBuddyRuntime) SwitchProduct(ctx context.Context, a *coreauth.Auth, product workbuddy.Product) (*coreauth.Auth, error) { + original := r.Current(a) + if original == nil || original.Disabled { + return nil, fmt.Errorf("workbuddy: credential removed or disabled") + } + updated := original.Clone() + updated.Metadata["product"] = string(product) + updated.UpdatedAt = time.Now() + r.mu.Lock() + manager, save := r.manager, r.Save + r.mu.Unlock() + if manager != nil { + committed, err := manager.UpdatePreparedAuth(ctx, original, updated) + if err != nil { + return nil, err + } + if committed == nil || committed.Disabled || committed.Metadata["product"] != string(product) { + return nil, fmt.Errorf("workbuddy: product change not committed") + } + updated = committed + } else if err := save(ctx, original, updated); err != nil { + return nil, err + } + r.Track(updated) + return r.Current(updated), nil +} diff --git a/internal/runtime/executor/workbuddy_executor.go b/internal/runtime/executor/workbuddy_executor.go index af3f05d2ffc..27bd8693640 100644 --- a/internal/runtime/executor/workbuddy_executor.go +++ b/internal/runtime/executor/workbuddy_executor.go @@ -8,6 +8,7 @@ import ( "fmt" "io" "net/http" + "sort" "strings" "time" @@ -61,12 +62,54 @@ func (e *WorkBuddyExecutor) TrackAuth(a *coreauth.Auth) { e.runtime.Track(a) } } -func (e *WorkBuddyExecutor) PrepareExecutionContext(ctx context.Context, _ execution.Request, _ execution.Options) (context.Context, error) { +func (e *WorkBuddyExecutor) PrepareExecutionContext(ctx context.Context, req execution.Request, opts execution.Options) (context.Context, error) { if e.initErr != nil { return ctx, &workbuddy.DispatchError{Reason: "provider_unavailable", Status: 503, Stop: true} } - return workbuddy.WithRequestState(ctx, e.runtime.Accounts()) + if workbuddy.State(ctx) != nil { + return ctx, nil + } + extra := 0 + if e.cfg.WorkBuddy.AutoSwitchProduct { + extra = 4 + } + next, err := workbuddy.WithRequestState(ctx, e.runtime.Accounts(), extra) + if err != nil { + return ctx, err + } + prefix, maxMessages, pageSize := true, 400, 3 + if e.cfg.WorkBuddy.AffinityByPrefix != nil { + prefix = *e.cfg.WorkBuddy.AffinityByPrefix + } + if e.cfg.WorkBuddy.AffinityMaxMessages != nil { + maxMessages = *e.cfg.WorkBuddy.AffinityMaxMessages + } + if e.cfg.WorkBuddy.AffinityPageSize != nil { + pageSize = *e.cfg.WorkBuddy.AffinityPageSize + } + caller, _ := opts.Metadata[execution.CallerScopeMetadataKey].(string) + explicit, _ := opts.Metadata[execution.ExecutionSessionMetadataKey].(string) + if explicit == "" { + for _, key := range []string{"X-Conversation-Id", "X-Session-Id", "Session_id", "X-Codex-Session-Id"} { + if value := opts.Headers.Get(key); value != "" { + explicit = value + break + } + } + } + if explicit == "" { + for _, key := range []string{"conversation_id", "session_id", "metadata.session_id"} { + if value := gjson.GetBytes(req.Payload, key).String(); value != "" { + explicit = value + break + } + } + } + payload := sdktranslator.TranslateRequest(opts.SourceFormat, sdktranslator.FromString("openai"), req.Model, req.Payload, opts.Stream) + workbuddy.State(next).Route = workbuddy.SessionRouting(caller, explicit, payload, prefix, maxMessages, pageSize) + return next, nil } + func (e *WorkBuddyExecutor) ShouldPrepareRequestAuth(a *coreauth.Auth) bool { c, err := helps.WorkBuddyCredential(a) return err == nil && c.RefreshToken != "" && c.Expiring(time.Now()) @@ -95,11 +138,14 @@ func (e *WorkBuddyExecutor) HttpRequest(context.Context, *coreauth.Auth, *http.R return nil, &workbuddy.DispatchError{Reason: "arbitrary_http_requests_are_not_supported", Status: 400, Stop: true} } -func (e *WorkBuddyExecutor) build(ctx context.Context, c workbuddy.Credential, req execution.Request, opts execution.Options) ([]byte, error) { +func (e *WorkBuddyExecutor) build(ctx context.Context, c workbuddy.Credential, req execution.Request, opts execution.Options, override ...string) ([]byte, error) { if opts.Alt != "" { return nil, &workbuddy.DispatchError{Reason: "unsupported_endpoint", Status: 400, Stop: true} } baseModel := thinking.ParseSuffix(req.Model).ModelName + if len(override) > 0 && override[0] != "" { + baseModel = override[0] + } from, to := opts.SourceFormat, sdktranslator.FromString("openai") original := req.Payload if len(opts.OriginalRequest) > 0 { @@ -114,7 +160,7 @@ func (e *WorkBuddyExecutor) build(ctx context.Context, c workbuddy.Credential, r if err != nil { return nil, err } - body, err = helps.NormalizeWorkBuddyPayload(body, c, baseModel) + body, err = helps.NormalizeWorkBuddyPayload(body, c, baseModel, helps.WorkBuddyPrompt(e.cfg.WorkBuddy.Prompt, e.runtime.Policy.Degraded())) if err != nil { return nil, err } @@ -166,8 +212,9 @@ func (e *WorkBuddyExecutor) execute(ctx context.Context, a *coreauth.Auth, req e } }() preOutputRetry := false + paidModel := "" for { - body, errBuild := e.build(ctx, c, req, opts) + body, errBuild := e.build(ctx, c, req, opts, paidModel) if errBuild != nil { return errBuild } @@ -182,12 +229,19 @@ func (e *WorkBuddyExecutor) execute(ctx context.Context, a *coreauth.Auth, req e } gate, attempt := e.runtime.Governor.Check(c, model, workbuddy.IsPaidModel(c.Realm, model), false) if !gate.Allowed() { + if !gate.Stop && gate.Reason == "model_cooldown" && e.paidFallback(c, model, state, paidModel) { + paidModel = "deepseek-v4.1-flash-sg" + continue + } return &workbuddy.DispatchError{Reason: gate.Reason, Status: 429, Stop: gate.Stop, Until: gate.Until} } if errBegin := state.Begin(); errBegin != nil { return errBegin } session := helps.DerivedSessionID(opts.Metadata, req.Metadata) + if state.Route.Key != "" { + session = state.Route.Key + } if c.Product == workbuddy.VSCode { session = state.Turn } @@ -230,7 +284,7 @@ func (e *WorkBuddyExecutor) execute(ctx context.Context, a *coreauth.Auth, req e execution.MarkUpstreamAttempt(ctx) response, errDo := client.Do(httpReq) if errDo != nil { - return e.recordFailure(c, model, state, workbuddy.ClassifyTransport(errDo)) + return e.connectionFailure(ctx, c, model, state, workbuddy.ClassifyTransport(errDo), base, proxy) } reporter.ObserveResponse(response) if response.StatusCode < 200 || response.StatusCode >= 300 { @@ -251,7 +305,22 @@ func (e *WorkBuddyExecutor) execute(ctx context.Context, a *coreauth.Auth, req e continue } } - return e.recordFailure(c, model, state, fault) + updated, failure, retry := e.recoverFailure(ctx, a, c, model, state, fault, base, proxy, paidModel == "") + if retry { + if updated != nil { + a = updated + c, err = helps.WorkBuddyCredential(a) + if err != nil { + return err + } + } + continue + } + if fault.Kind == workbuddy.FaultModelCap && e.paidFallback(c, model, state, paidModel) { + paidModel = "deepseek-v4.1-flash-sg" + continue + } + return failure } emitted := false var usage json.RawMessage @@ -302,6 +371,10 @@ func (e *WorkBuddyExecutor) execute(ctx context.Context, a *coreauth.Auth, req e return &workbuddy.DispatchError{Reason: "usage_persistence_failed", Status: 503, Stop: true} } if errStream == nil { + e.runtime.Policy.ProductSuccess(c, model) + if e.cfg.WorkBuddy.EgressBackoff { + e.runtime.Policy.ExitSuccess(base, proxy) + } return consume([]byte("[DONE]"), body) } if ctx.Err() != nil { @@ -309,6 +382,24 @@ func (e *WorkBuddyExecutor) execute(ctx context.Context, a *coreauth.Auth, req e } var upstreamFault *workbuddy.StreamFault if errors.As(errStream, &upstreamFault) { + if !emitted { + updated, failure, retry := e.recoverFailure(ctx, a, c, model, state, upstreamFault.Fault, base, proxy, paidModel == "") + if retry { + if updated != nil { + a = updated + c, err = helps.WorkBuddyCredential(a) + if err != nil { + return err + } + } + continue + } + if upstreamFault.Fault.Kind == workbuddy.FaultModelCap && e.paidFallback(c, model, state, paidModel) { + paidModel = "deepseek-v4.1-flash-sg" + continue + } + return failure + } failure := e.recordFailure(c, model, state, upstreamFault.Fault) var dispatch *workbuddy.DispatchError if emitted && errors.As(failure, &dispatch) { @@ -324,10 +415,20 @@ func (e *WorkBuddyExecutor) execute(ctx context.Context, a *coreauth.Auth, req e } if fault.Kind == workbuddy.FaultConnection { preOutputRetry = true + if e.cfg.WorkBuddy.EgressBackoff { + if errWait := e.runtime.Policy.Backoff(ctx, state, base, proxy); errWait != nil { + return errWait + } + } continue } } - failure := e.recordFailure(c, model, state, workbuddy.Fault{Kind: workbuddy.FaultUnknown, Status: 502}) + var failure error + if e.cfg.WorkBuddy.EgressBackoff { + failure = e.connectionFailure(ctx, c, model, state, workbuddy.Fault{Kind: workbuddy.FaultConnection, Status: 502}, base, proxy) + } else { + failure = e.recordFailure(c, model, state, workbuddy.Fault{Kind: workbuddy.FaultUnknown, Status: 502}) + } // Once output escaped, no SDK credential rotation may replay the prompt. var dispatch *workbuddy.DispatchError if errors.As(failure, &dispatch) { @@ -343,6 +444,7 @@ func (e *WorkBuddyExecutor) recordFailure(c workbuddy.Credential, model string, if fault.Kind == workbuddy.FaultHard || fault.Kind == workbuddy.FaultUnknown { charge = state.Charge(c) } + e.runtime.Affinity.Demote(c.Realm, state.Route, c.UID) failure := workbuddy.ErrorForFault(fault) if charge { gate, err := e.runtime.Governor.RecordFailure(c, model, fault, workbuddy.FailureOptions{FreeModel: workbuddy.IsFreeCapModel(c.Realm, model), SingleAccount: e.runtime.Accounts() <= 1}) @@ -433,3 +535,109 @@ func (e *WorkBuddyExecutor) CountTokens(ctx context.Context, a *coreauth.Auth, r result := sdktranslator.TranslateTokenCount(ctx, sdktranslator.FromString("openai"), execution.ResponseFormatOrSource(opts), count, helps.BuildOpenAIUsageJSON(count)) return execution.Response{Payload: result}, nil } + +// recoverFailure runs only before downstream output. Each retry rebuilds from +// the original request and crosses the payload finalizer once again. +func (e *WorkBuddyExecutor) recoverFailure(ctx context.Context, a *coreauth.Auth, c workbuddy.Credential, model string, state *workbuddy.RequestState, fault workbuddy.Fault, base, proxy string, allowSwitch bool) (*coreauth.Auth, error, bool) { + if fault.Kind == workbuddy.FaultContent && e.cfg.WorkBuddy.Prompt.Mode != "custom" && state.TryDegrade() { + e.runtime.Policy.Degrade() + return nil, nil, true + } + if fault.Kind == workbuddy.FaultGateway && e.cfg.WorkBuddy.EgressBackoff { + return nil, e.connectionFailure(ctx, c, model, state, fault, base, proxy), false + } + failure := e.recordFailure(c, model, state, fault) + var dispatch *workbuddy.DispatchError + if !errors.As(failure, &dispatch) || dispatch.Stop { + return nil, failure, false + } + if allowSwitch && e.cfg.WorkBuddy.AutoSwitchProduct && (fault.Kind == workbuddy.FaultModelRate || fault.Kind == workbuddy.FaultModelCap) && !e.runtime.Governor.Capped(c, model) { + allowed := make([]workbuddy.Product, 0, len(e.cfg.WorkBuddy.AutoSwitchProducts)) + for _, p := range e.cfg.WorkBuddy.AutoSwitchProducts { + allowed = append(allowed, workbuddy.Product(p)) + } + if next, ok := e.runtime.Policy.NextProduct(c, model, allowed); ok { + updated, err := e.runtime.SwitchProduct(ctx, a, next) + if err == nil && e.runtime.Governor.ClearIdentityModelCooldown(c, model) { + return updated, nil, true + } + } + } + return nil, failure, false +} +func (e *WorkBuddyExecutor) connectionFailure(ctx context.Context, c workbuddy.Credential, model string, state *workbuddy.RequestState, fault workbuddy.Fault, base, proxy string) error { + if e.cfg.WorkBuddy.EgressBackoff && (fault.Kind == workbuddy.FaultGateway || fault.Kind == workbuddy.FaultConnection) { + if err := e.runtime.Policy.Backoff(ctx, state, base, proxy); err != nil { + return err + } + return workbuddy.ErrorForFault(fault) + } + return e.recordFailure(c, model, state, fault) +} +func (e *WorkBuddyExecutor) paidFallback(c workbuddy.Credential, model string, state *workbuddy.RequestState, current string) bool { + return current == "" && e.cfg.WorkBuddy.PaidFallback && c.Realm == workbuddy.International && model == "deepseek-v4.1-flash" && e.runtime.AllFreeCapped(c.Realm, model) && state.TryPaidFallback() +} + +// SelectAuth runs after SDK model, pin, caller eligibility and priority checks. +// It never introduces a credential outside the candidate set supplied by SDK. +func (e *WorkBuddyExecutor) SelectAuth(ctx context.Context, model string, opts execution.Options, candidates []*coreauth.Auth) (*coreauth.Auth, error) { + if len(candidates) == 0 { + return nil, &workbuddy.DispatchError{Reason: "no_eligible_credentials", Status: 503, Stop: true} + } + state := workbuddy.State(ctx) + if state == nil { + return candidates[0], nil + } + baseModel := thinking.ParseSuffix(model).ModelName + byRealm := make(map[workbuddy.Realm]map[string]*coreauth.Auth) + for _, candidate := range candidates { + c, err := helps.WorkBuddyCredential(candidate) + if err != nil { + continue + } + gate, _ := e.runtime.Governor.Check(c, baseModel, workbuddy.IsPaidModel(c.Realm, baseModel), false) + eligible := gate.Allowed() || gate.Reason == "credential_expired" && c.RefreshToken != "" + if !eligible && gate.Reason == "model_cooldown" && e.cfg.WorkBuddy.PaidFallback && baseModel == "deepseek-v4.1-flash" && c.Realm == workbuddy.International && e.runtime.AllFreeCapped(c.Realm, baseModel) { + paidGate, _ := e.runtime.Governor.Check(c, "deepseek-v4.1-flash-sg", true, false) + eligible = paidGate.Allowed() + } + if !eligible { + continue + } + if byRealm[c.Realm] == nil { + byRealm[c.Realm] = make(map[string]*coreauth.Auth) + } + previous := byRealm[c.Realm][c.UID] + if previous == nil || candidate.ID < previous.ID { + byRealm[c.Realm][c.UID] = candidate + } + } + // Shared model names prefer the international realm. Explicit auth prefixes + // and pins are already narrowed by SDK and preserve China-only routes. + for _, realm := range []workbuddy.Realm{workbuddy.International, workbuddy.China} { + group := byRealm[realm] + if len(group) == 0 { + continue + } + eligible := make([]string, 0, len(group)) + for uid := range group { + eligible = append(eligible, uid) + } + allSet := make(map[string]bool) + for _, c := range e.runtime.Credentials() { + if c.Realm == realm { + allSet[c.UID] = true + } + } + all := make([]string, 0, len(allSet)) + for uid := range allSet { + all = append(all, uid) + } + uid := e.runtime.Affinity.Pick(realm, state.Route, eligible, all) + return group[uid], nil + } + // Let the executor report the precise gate and the SDK's normal retry loop + // handle exhausted tiers. Selection itself must not bypass a restriction. + sort.Slice(candidates, func(i, j int) bool { return candidates[i].ID < candidates[j].ID }) + return candidates[0], nil +} diff --git a/internal/runtime/executor/workbuddy_executor_test.go b/internal/runtime/executor/workbuddy_executor_test.go index f4e6261d269..c24ed656d2c 100644 --- a/internal/runtime/executor/workbuddy_executor_test.go +++ b/internal/runtime/executor/workbuddy_executor_test.go @@ -462,3 +462,168 @@ func TestWorkBuddyInflightRefreshCannotOverwriteExternalRotation(t *testing.T) { t.Fatal("external persisted rotation overwritten") } } + +func TestWorkBuddyContentDegradationRebuildsBeforeFinalBarrier(t *testing.T) { + var calls atomic.Int32 + e, a := wbFixture(t, func(w http.ResponseWriter, r *http.Request) { + body, _ := io.ReadAll(r.Body) + if gjson.GetBytes(body, "configured").String() != "yes" { + t.Error("payload rules missing on rebuilt attempt") + } + if calls.Add(1) == 1 { + if gjson.GetBytes(body, "messages.0.content").String() != "caller-system" { + t.Error("initial prompt rewritten") + } + w.WriteHeader(403) + _, _ = io.WriteString(w, `{"message":"content review"}`) + return + } + if gjson.GetBytes(body, "messages.0.content").String() != workbuddy.DegradedPrompt { + t.Error("retry did not degrade system prompt") + } + if gjson.GetBytes(body, "messages.1.content").String() != "user-content" { + t.Error("degradation altered user content") + } + _, _ = io.WriteString(w, wbSuccess) + }) + e.cfg.Payload = config.PayloadConfig{Override: []config.PayloadRule{{Models: []config.PayloadModelRule{{Name: "*"}}, Params: map[string]any{"configured": "yes"}}}} + req, opts := wbRequest() + req.Payload = []byte(`{"model":"fixture-model","messages":[{"role":"system","content":"caller-system"},{"role":"user","content":"user-content"}]}`) + if _, err := e.Execute(context.Background(), a, req, opts); err != nil { + t.Fatal(err) + } + if calls.Load() != 2 { + t.Fatalf("attempts=%d", calls.Load()) + } +} +func TestWorkBuddyCustomPromptNeverDegradesOrRetries(t *testing.T) { + var calls atomic.Int32 + e, a := wbFixture(t, func(w http.ResponseWriter, r *http.Request) { + calls.Add(1) + body, _ := io.ReadAll(r.Body) + if gjson.GetBytes(body, "messages.0.content").String() != "custom-prompt" { + t.Error("custom prompt not applied") + } + w.WriteHeader(403) + }) + e.cfg.WorkBuddy.Prompt = config.WorkBuddyPrompt{Mode: "custom", Text: "custom-prompt"} + req, opts := wbRequest() + if _, err := e.Execute(context.Background(), a, req, opts); err == nil { + t.Fatal("content rejection lost") + } + if calls.Load() != 1 || e.runtime.Policy.Degraded() { + t.Fatal("custom prompt opted into degradation") + } +} +func TestWorkBuddyProductSwitchPersistsBeforeRetry(t *testing.T) { + var calls atomic.Int32 + var saved atomic.Int32 + e, a := wbFixture(t, func(w http.ResponseWriter, r *http.Request) { + if calls.Add(1) == 1 { + w.WriteHeader(429) + _, _ = io.WriteString(w, `{"code":14003,"message":"too many requests"}`) + return + } + if r.Header.Get("X-IDE-Type") != "CLI" || saved.Load() != 1 { + t.Error("identity did not persist before retry") + } + _, _ = io.WriteString(w, wbSuccess) + }) + e.cfg.WorkBuddy.AutoSwitchProduct = true + e.cfg.WorkBuddy.AutoSwitchProducts = []string{"workbuddy", "cli"} + e.runtime.Save = func(_ context.Context, _ *coreauth.Auth, next *coreauth.Auth) error { + if next.Metadata["product"] != "cli" { + t.Error("unexpected saved product") + } + saved.Add(1) + return nil + } + req, opts := wbRequest() + if _, err := e.Execute(context.Background(), a, req, opts); err != nil { + t.Fatal(err) + } + if calls.Load() != 2 { + t.Fatalf("requests=%d", calls.Load()) + } +} +func TestWorkBuddyPaidFallbackRequiresOptInAndEveryFreeAccountCapped(t *testing.T) { + for _, tc := range []struct { + name string + enabled, other bool + want int + }{{"disabled", false, false, 1}, {"other-free-account", true, true, 1}, {"all-capped", true, false, 2}} { + t.Run(tc.name, func(t *testing.T) { + var calls atomic.Int32 + e, a := wbFixture(t, func(w http.ResponseWriter, r *http.Request) { + body, _ := io.ReadAll(r.Body) + model := gjson.GetBytes(body, "model").String() + calls.Add(1) + if model == "deepseek-v4.1-flash" { + w.WriteHeader(429) + _, _ = io.WriteString(w, `{"code":6004,"message":"reset at 2099-10-22 00:00:00 UTC+8"}`) + return + } + if model != "deepseek-v4.1-flash-sg" { + t.Errorf("unexpected billed model %s", model) + } + _, _ = io.WriteString(w, wbSuccess) + }) + e.cfg.WorkBuddy.PaidFallback = tc.enabled + if tc.other { + second := a.Clone() + second.ID = "second-fixture" + second.Metadata["uid"] = "second-fixture-user" + e.TrackAuth(second) + } + req, opts := wbRequest() + req.Model = "deepseek-v4.1-flash" + _, err := e.Execute(context.Background(), a, req, opts) + if tc.want == 2 && err != nil { + t.Fatal(err) + } + if tc.want == 1 && err == nil { + t.Fatal("cap ignored") + } + if calls.Load() != int32(tc.want) { + t.Fatalf("requests=%d want %d", calls.Load(), tc.want) + } + c, _ := helps.WorkBuddyCredential(a) + if !e.runtime.Governor.Capped(c, "deepseek-v4.1-flash") { + t.Fatal("paid fallback cleared free cap") + } + }) + } +} + +func TestWorkBuddyNativeSelectionUsesWarmHistoryAndRiskGates(t *testing.T) { + e, a := wbFixture(t, func(w http.ResponseWriter, r *http.Request) { _, _ = io.WriteString(w, wbSuccess) }) + b := a.Clone() + b.ID = "workbuddy-fixture-b" + b.Metadata["uid"] = "fixture-b" + e.TrackAuth(b) + req, opts := wbRequest() + ctx, err := e.PrepareExecutionContext(context.Background(), req, opts) + if err != nil { + t.Fatal(err) + } + first, err := e.SelectAuth(ctx, "fixture-model", opts, []*coreauth.Auth{a, b}) + if err != nil { + t.Fatal(err) + } + again, err := e.SelectAuth(ctx, "fixture-model", opts, []*coreauth.Auth{b, a}) + if err != nil || again.ID != first.ID { + t.Fatal("affinity followed candidate ordering") + } + c, _ := helps.WorkBuddyCredential(first) + if _, err := e.runtime.Governor.RecordFailure(c, "fixture-model", workbuddy.Fault{Kind: workbuddy.FaultModelRate}, workbuddy.FailureOptions{}); err != nil { + t.Fatal(err) + } + switched, err := e.SelectAuth(ctx, "fixture-model", opts, []*coreauth.Auth{a, b}) + if err != nil || switched.ID == first.ID { + t.Fatal("affinity bypassed model cooldown") + } + sibling, err := e.SelectAuth(ctx, "another-model", opts, []*coreauth.Auth{first}) + if err != nil || sibling.ID != first.ID { + t.Fatal("model cooldown poisoned account") + } +} diff --git a/internal/workbuddy/affinity.go b/internal/workbuddy/affinity.go new file mode 100644 index 00000000000..2ff62573f08 --- /dev/null +++ b/internal/workbuddy/affinity.go @@ -0,0 +1,216 @@ +package workbuddy + +import ( + "crypto/sha256" + "encoding/binary" + "encoding/hex" + "encoding/json" + "sort" + "strconv" + "sync" + "time" +) + +type SessionRoute struct { + Key string + Page int + PageSize int +} +type affinityEntry struct { + account string + until time.Time +} +type affinityBinding struct { + entries []affinityEntry + touched time.Time +} +type Affinity struct { + mu sync.Mutex + now func() time.Time + bindings map[string]affinityBinding + cursors map[string]int +} + +func NewAffinity(now func() time.Time) *Affinity { + if now == nil { + now = time.Now + } + return &Affinity{now: now, bindings: make(map[string]affinityBinding), cursors: make(map[string]int)} +} + +// SessionRouting hashes only a stable prefix and caller scope; no prompt text +// is retained. Explicit sessions keep ordinary affinity regardless of length. +func SessionRouting(caller, explicit string, payload []byte, prefix bool, maxMessages, pageSize int) SessionRoute { + var source []byte + count := 0 + if explicit != "" { + source = []byte(explicit) + } else if prefix { + var body struct { + Messages []json.RawMessage `json:"messages"` + } + if json.Unmarshal(payload, &body) != nil || len(body.Messages) == 0 { + return SessionRoute{} + } + count = len(body.Messages) + head := body.Messages + if len(head) > 2 { + head = head[:2] + } + // Canonical object key ordering makes equivalent JSON share a session. + var value any + raw, _ := json.Marshal(head) + if json.Unmarshal(raw, &value) != nil { + return SessionRoute{} + } + source, _ = json.Marshal(value) + } else { + return SessionRoute{} + } + digest := sha256.Sum256(append(append([]byte(caller), 0), source...)) + route := SessionRoute{Key: hex.EncodeToString(digest[:])} + if explicit == "" && maxMessages > 0 && count > maxMessages { + if pageSize == 0 { + return SessionRoute{} + } + if pageSize > 1 { + route.PageSize = pageSize + route.Page = int(binary.BigEndian.Uint32(digest[:4])) % pageSize + } + } + return route +} +func affinityKey(realm Realm, route SessionRoute) string { + if route.Key == "" { + return "" + } + return string(realm) + "\x00" + route.Key +} +func (a *Affinity) live(key string) affinityBinding { + binding := a.bindings[key] + now := a.now() + live := binding.entries[:0] + for _, entry := range binding.entries { + if now.Before(entry.until) { + live = append(live, entry) + } + } + binding.entries = live + if len(live) == 0 { + delete(a.bindings, key) + } else { + a.bindings[key] = binding + } + return binding +} + +// Pick accepts already eligible account keys. Page ranks use all enabled realm +// accounts so a temporary cooldown does not reshuffle the conversation's page. +func (a *Affinity) Pick(realm Realm, route SessionRoute, eligible, all []string) string { + if len(eligible) == 0 { + return "" + } + a.mu.Lock() + defer a.mu.Unlock() + key := affinityKey(realm, route) + set := make(map[string]bool, len(eligible)) + for _, id := range eligible { + set[id] = true + } + group := eligible + if route.PageSize > 1 { + ranks := append([]string(nil), all...) + sort.Strings(ranks) + page := make(map[string]bool) + if route.PageSize < len(ranks) { + for i, id := range ranks { + if i%route.PageSize == route.Page%route.PageSize { + page[id] = true + } + } + } else { + for _, id := range ranks { + page[id] = true + } + } + narrowed := make([]string, 0, len(eligible)) + for _, id := range eligible { + if page[id] { + narrowed = append(narrowed, id) + } + } + if len(narrowed) > 0 { + group = narrowed + } + } else if key != "" { + for _, entry := range a.live(key).entries { + if set[entry.account] { + a.bindLocked(key, entry.account) + return entry.account + } + } + } + sort.Strings(group) + cursorKey := string(realm) + if route.PageSize > 1 { + cursorKey = string(realm) + ":" + strconv.Itoa(route.Page) + "/" + strconv.Itoa(route.PageSize) + } + if len(a.cursors) >= 5000 { + delete(a.cursors, cursorKey) + if len(a.cursors) >= 5000 { + for k := range a.cursors { + delete(a.cursors, k) + break + } + } + } + index := a.cursors[cursorKey] % len(group) + selected := group[index] + a.cursors[cursorKey] = (index + 1) % len(group) + if key != "" { + a.bindLocked(key, selected) + } + return selected +} +func (a *Affinity) bindLocked(key, id string) { + now := a.now() + binding := a.live(key) + entries := []affinityEntry{{account: id, until: now.Add(2 * time.Hour)}} + for _, entry := range binding.entries { + if entry.account != id && len(entries) < 3 { + entries = append(entries, entry) + } + } + if _, exists := a.bindings[key]; !exists && len(a.bindings) >= 5000 { + oldest := "" + var stamp time.Time + for k, b := range a.bindings { + if oldest == "" || b.touched.Before(stamp) { + oldest, stamp = k, b.touched + } + } + delete(a.bindings, oldest) + } + a.bindings[key] = affinityBinding{entries: entries, touched: now} +} +func (a *Affinity) Demote(realm Realm, route SessionRoute, id string) { + key := affinityKey(realm, route) + if key == "" { + return + } + a.mu.Lock() + defer a.mu.Unlock() + binding := a.live(key) + kept := binding.entries[:0] + for _, entry := range binding.entries { + if entry.account != id { + kept = append(kept, entry) + } + } + binding.entries = kept + if len(kept) == 0 { + delete(a.bindings, key) + } else { + a.bindings[key] = binding + } +} diff --git a/internal/workbuddy/affinity_test.go b/internal/workbuddy/affinity_test.go new file mode 100644 index 00000000000..3c2feb94604 --- /dev/null +++ b/internal/workbuddy/affinity_test.go @@ -0,0 +1,76 @@ +package workbuddy + +import ( + "testing" + "time" +) + +func TestWorkBuddyPrefixRoutingScopeAndPaging(t *testing.T) { + first := []byte(`{"messages":[{"role":"system","content":"rules"},{"role":"user","content":"first"}]}`) + follow := []byte(`{"messages":[{"content":"rules","role":"system"},{"content":"first","role":"user"},{"role":"assistant","content":"reply"},{"role":"user","content":"next"}]}`) + a := SessionRouting("caller-a", "", first, true, 2, 3) + b := SessionRouting("caller-a", "", follow, true, 2, 3) + if a.Key == "" || a.Key != b.Key || a.PageSize != 0 || b.PageSize != 3 { + t.Fatalf("routes = %+v %+v", a, b) + } + if other := SessionRouting("caller-b", "", first, true, 2, 3); other.Key == a.Key { + t.Fatal("caller scopes share affinity") + } + if explicit := SessionRouting("caller-a", "explicit", follow, true, 2, 3); explicit.PageSize != 0 || explicit.Key == a.Key { + t.Fatal("explicit session paged or collided") + } + if off := SessionRouting("caller-a", "", follow, true, 2, 0); off.Key != "" { + t.Fatal("legacy full rotation mode did not disable long affinity") + } +} +func TestWorkBuddyAffinityHistoryAndExpiry(t *testing.T) { + now := time.Date(2026, 10, 21, 0, 0, 0, 0, time.UTC) + a := NewAffinity(func() time.Time { return now }) + route := SessionRoute{Key: "conversation"} + all := []string{"a", "b", "c"} + if got := a.Pick(International, route, []string{"a"}, all); got != "a" { + t.Fatal(got) + } + now = now.Add(time.Minute) + if got := a.Pick(International, route, []string{"b"}, all); got != "b" { + t.Fatal(got) + } + if got := a.Pick(International, route, []string{"a", "c"}, all); got != "a" { + t.Fatal("warm history not preferred") + } + a.Demote(International, route, "a") + if got := a.Pick(International, route, []string{"b", "c"}, all); got != "b" { + t.Fatal("demotion discarded other warm accounts") + } + if got := a.Pick(China, route, []string{"c"}, all); got != "c" { + t.Fatal("realm route corrupted") + } + now = now.Add(2*time.Hour + time.Second) + a.mu.Lock() + live := a.live(affinityKey(International, route)) + a.mu.Unlock() + if len(live.entries) != 0 { + t.Fatal("expired affinity survived") + } +} +func TestWorkBuddyConversationPageRanksDoNotShiftWithCooldown(t *testing.T) { + a := NewAffinity(nil) + route := SessionRoute{Key: "long-conversation", Page: 1, PageSize: 3} + all := []string{"a", "b", "c", "d", "e", "f", "g", "h"} + seen := make(map[string]bool) + for i := 0; i < 12; i++ { + seen[a.Pick(International, route, append([]string(nil), all...), all)] = true + } + if len(seen) != 3 || !seen["b"] || !seen["e"] || !seen["h"] { + t.Fatalf("page = %+v", seen) + } + for i := 0; i < 6; i++ { + got := a.Pick(International, route, []string{"a", "c", "d", "e", "f", "g", "h"}, all) + if got != "e" && got != "h" { + t.Fatalf("cooldown shifted page: %s", got) + } + } + if got := a.Pick(International, route, []string{"a", "c", "d"}, all); got != "a" && got != "c" && got != "d" { + t.Fatal("empty page did not fall back") + } +} diff --git a/internal/workbuddy/dispatch.go b/internal/workbuddy/dispatch.go index e981470a7a2..7299f6984a6 100644 --- a/internal/workbuddy/dispatch.go +++ b/internal/workbuddy/dispatch.go @@ -57,14 +57,18 @@ func ErrorForFault(f Fault) *DispatchError { type requestStateKey struct{} type RequestState struct { - mu sync.Mutex - Turn string - attempts, max int - charged map[string]bool - refreshed map[string]bool + Route SessionRoute + mu sync.Mutex + Turn string + attempts, max int + charged map[string]bool + refreshed map[string]bool + exits map[string]bool + degraded, paid bool + delay time.Duration } -func WithRequestState(ctx context.Context, accounts int) (context.Context, error) { +func WithRequestState(ctx context.Context, accounts int, extra ...int) (context.Context, error) { if _, ok := ctx.Value(requestStateKey{}).(*RequestState); ok { return ctx, nil } @@ -75,7 +79,11 @@ func WithRequestState(ctx context.Context, accounts int) (context.Context, error if accounts < 2 { accounts = 2 } - state := &RequestState{Turn: turn, max: accounts + 1, charged: make(map[string]bool), refreshed: make(map[string]bool)} + additional := 0 + if len(extra) > 0 && extra[0] > 0 { + additional = extra[0] + } + state := &RequestState{Turn: turn, max: accounts + 1 + additional, exits: make(map[string]bool), charged: make(map[string]bool), refreshed: make(map[string]bool)} return context.WithValue(ctx, requestStateKey{}, state), nil } func State(ctx context.Context) *RequestState { @@ -185,3 +193,44 @@ func (l *Limiter) Configure(global, perKey int) error { l.changed = make(chan struct{}) return nil } + +func (s *RequestState) TryDegrade() bool { + s.mu.Lock() + defer s.mu.Unlock() + if s.degraded { + return false + } + s.degraded = true + return true +} +func (s *RequestState) TryPaidFallback() bool { + s.mu.Lock() + defer s.mu.Unlock() + if s.paid { + return false + } + s.paid = true + return true +} +func (s *RequestState) ChargeExit(key string) bool { + s.mu.Lock() + defer s.mu.Unlock() + if s.exits[key] { + return false + } + s.exits[key] = true + return true +} +func (s *RequestState) ReserveDelay(requested time.Duration) time.Duration { + s.mu.Lock() + defer s.mu.Unlock() + remaining := 5*time.Second - s.delay + if requested > remaining { + requested = remaining + } + if requested < 0 { + requested = 0 + } + s.delay += requested + return requested +} diff --git a/internal/workbuddy/governor.go b/internal/workbuddy/governor.go index 8889a1d32e1..dd22a78b6b3 100644 --- a/internal/workbuddy/governor.go +++ b/internal/workbuddy/governor.go @@ -268,10 +268,7 @@ func (g *Governor) Check(c Credential, model string, paid, balanceOnly bool) (Ga if gate.Reason != "" { return gate, attempt } - if balanceOnly { - return gate, attempt - } - if now.Before(state.BalanceUntil) { + if !balanceOnly && now.Before(state.BalanceUntil) { gate.Reason, gate.Until = "balance_empty", state.BalanceUntil return gate, attempt } @@ -282,6 +279,9 @@ func (g *Governor) Check(c Credential, model string, paid, balanceOnly bool) (Ga return gate, attempt } } + if balanceOnly { + return gate, attempt + } limits := g.options.Limits if override, exists := g.options.Realms[c.Realm]; exists { limits = override @@ -561,3 +561,44 @@ func (g *Governor) Entitlement(realm Realm, model string) (EntitlementMark, bool mark, ok := g.state.Entitlements[realm][model] return mark, ok && g.options.Now().Before(mark.Until) } + +func (g *Governor) Capped(c Credential, model string) bool { + g.mu.Lock() + defer g.mu.Unlock() + state := g.state.Accounts[RiskAccountKey(c)] + if state == nil { + return false + } + cap := state.Models[model] + return cap != nil && g.options.Now().Before(cap.CapUntil) +} + +// ClearIdentityModelCooldown clears only the ordinary model cooldown after a +// successful product change. Durable caps and account/region restrictions stay. +func (g *Governor) ClearIdentityModelCooldown(c Credential, model string) bool { + g.mu.Lock() + defer g.mu.Unlock() + now := g.options.Now() + state := g.accountLocked(c) + if g.saveErr != nil || c.Disabled { + return false + } + if realm := g.state.Realms[c.Realm]; realm != nil && now.Before(realm.HoldUntil) { + return false + } + for _, until := range []time.Time{state.RestrictedUntil, state.BalanceUntil, state.CooldownUntil, state.BreakerUntil, state.DegradeUntil} { + if now.Before(until) { + return false + } + } + current := state.Models[model] + if current == nil { + return true + } + if now.Before(current.CapUntil) { + return false + } + current.CooldownUntil = time.Time{} + state.Version++ + return true +} diff --git a/internal/workbuddy/governor_test.go b/internal/workbuddy/governor_test.go index 6c6e04832a0..79aa90a8c4d 100644 --- a/internal/workbuddy/governor_test.go +++ b/internal/workbuddy/governor_test.go @@ -19,6 +19,20 @@ func governorFixture(t *testing.T, limits Limits) (*Governor, *time.Time, string return g, &now, path } +func TestWorkBuddyBalanceReadRespectsControlCooldown(t *testing.T) { + g, _, _ := governorFixture(t, Limits{ReserveCredits: 10}) + c := testCredential() + if gate, _ := g.Check(c, "balance", false, true); !gate.Allowed() { + t.Fatal("unknown balance prevented a balance read") + } + if _, err := g.RecordFailure(c, "balance", Fault{Kind: FaultModelRate, Soft: true}, FailureOptions{}); err != nil { + t.Fatal(err) + } + if gate, _ := g.Check(c, "balance", false, true); gate.Reason != "model_cooldown" { + t.Fatal("balance read bypassed endpoint cooldown") + } +} + func TestGovernorRestrictionWaveAndStaleSuccessSurviveRestart(t *testing.T) { g, now, path := governorFixture(t, Limits{}) first := testCredential() diff --git a/internal/workbuddy/policy.go b/internal/workbuddy/policy.go new file mode 100644 index 00000000000..31e3127df3c --- /dev/null +++ b/internal/workbuddy/policy.go @@ -0,0 +1,186 @@ +package workbuddy + +import ( + "context" + "crypto/sha256" + "encoding/hex" + "math/rand/v2" + "net/url" + "sync" + "time" +) + +const DegradedPrompt = "You are a helpful assistant. Respond in the user's language, follow the user's instructions, and be direct and concise." + +type exitFailure struct { + count int + last time.Time + wait time.Duration +} +type productSwitch struct { + count int + last time.Time +} + +// PolicyState contains deliberately process-local recovery state. Restrictions, +// model caps and usage remain in Governor and cannot be cleared by this object. +type PolicyState struct { + mu sync.Mutex + now func() time.Time + jitter func() float64 + sleep func(context.Context, time.Duration) error + degradedUntil time.Time + exits map[string]exitFailure + switches map[string]productSwitch +} + +func NewPolicyState(now func() time.Time) *PolicyState { + if now == nil { + now = time.Now + } + return &PolicyState{now: now, jitter: func() float64 { return .75 + rand.Float64()*.5 }, sleep: waitBeforeConnect, exits: make(map[string]exitFailure), switches: make(map[string]productSwitch)} +} +func waitBeforeConnect(ctx context.Context, d time.Duration) error { + if d <= 0 { + return ctx.Err() + } + timer := time.NewTimer(d) + defer timer.Stop() + select { + case <-ctx.Done(): + return ctx.Err() + case <-timer.C: + return nil + } +} +func (p *PolicyState) Degraded() bool { + p.mu.Lock() + defer p.mu.Unlock() + return p.now().Before(p.degradedUntil) +} +func (p *PolicyState) Degrade() { + p.mu.Lock() + defer p.mu.Unlock() + now := p.now() + if now.Before(p.degradedUntil) { + return + } + cst := time.FixedZone("CST", 8*3600) + local := now.In(cst) + p.degradedUntil = time.Date(local.Year(), local.Month(), local.Day()+1, 0, 0, 0, 0, cst) +} +func ExitKey(base, proxy string) string { + u, err := url.Parse(base) + target := base + if err == nil { + target = u.Scheme + "://" + u.Host + } + sum := sha256.Sum256([]byte(target + "\x00" + proxy)) + return hex.EncodeToString(sum[:]) +} + +// Backoff charges each exit once per logical request and caps total waiting at +// five seconds. Call only after closing the failed connection/response body. +func (p *PolicyState) Backoff(ctx context.Context, state *RequestState, base, proxy string) error { + key := ExitKey(base, proxy) + charge := state.ChargeExit(key) + p.mu.Lock() + now := p.now() + entry := p.exits[key] + if now.Sub(entry.last) > time.Minute { + entry = exitFailure{} + } + if charge { + entry.count++ + entry.last = now + step := entry.count - 1 + if step > 4 { + step = 4 + } + seconds := 1 << step + if seconds > 15 { + seconds = 15 + } + entry.wait = time.Duration(float64(time.Duration(seconds)*time.Second) * p.jitter()) + if len(p.exits) >= 1024 { + for k, v := range p.exits { + if now.Sub(v.last) > time.Minute { + delete(p.exits, k) + } + } + } + if _, exists := p.exits[key]; exists || len(p.exits) < 1024 { + p.exits[key] = entry + } + } + p.mu.Unlock() + return p.sleep(ctx, state.ReserveDelay(entry.wait)) +} +func (p *PolicyState) ExitSuccess(base, proxy string) { + p.mu.Lock() + defer p.mu.Unlock() + delete(p.exits, ExitKey(base, proxy)) +} + +func (p *PolicyState) NextProduct(c Credential, model string, allowed []Product) (Product, bool) { + p.mu.Lock() + defer p.mu.Unlock() + now := p.now() + key := RiskAccountKey(c) + "\x00" + model + entry := p.switches[key] + if now.Sub(entry.last) > time.Minute { + entry = productSwitch{} + } + if entry.count >= 4 { + return "", false + } + sequence := []Product{Desktop, VSCode, CLI} + start := 0 + for i, product := range sequence { + if product == c.Product { + start = i + break + } + } + permitted := func(product Product) bool { + if len(allowed) == 0 { + return true + } + for _, item := range allowed { + if item == product { + return true + } + } + return false + } + for i := 1; i < len(sequence); i++ { + next := sequence[(start+i)%len(sequence)] + if !permitted(next) { + continue + } + // CN VSCode personal accounts require a real device token. Do not invent one. + if next == VSCode && c.Realm == China && c.EnterpriseID == "" && c.DeviceToken == "" { + continue + } + if len(p.switches) >= 500 { + for k, v := range p.switches { + if now.Sub(v.last) > time.Minute { + delete(p.switches, k) + } + } + } + if _, exists := p.switches[key]; !exists && len(p.switches) >= 500 { + return "", false + } + entry.count++ + entry.last = now + p.switches[key] = entry + return next, true + } + return "", false +} +func (p *PolicyState) ProductSuccess(c Credential, model string) { + p.mu.Lock() + defer p.mu.Unlock() + delete(p.switches, RiskAccountKey(c)+"\x00"+model) +} diff --git a/internal/workbuddy/policy_test.go b/internal/workbuddy/policy_test.go new file mode 100644 index 00000000000..506ebdbaea8 --- /dev/null +++ b/internal/workbuddy/policy_test.go @@ -0,0 +1,125 @@ +package workbuddy + +import ( + "context" + "errors" + "testing" + "time" +) + +func TestWorkBuddyPromptDegradeEndsAtCSTMidnight(t *testing.T) { + now := time.Date(2026, 10, 21, 15, 59, 0, 0, time.UTC) + p := NewPolicyState(func() time.Time { return now }) + p.Degrade() + if !p.Degraded() { + t.Fatal("missing degradation") + } + now = now.Add(30 * time.Second) + p.Degrade() + now = now.Add(30 * time.Second) + if p.Degraded() { + t.Fatal("repeated rejection extended midnight window") + } +} +func TestWorkBuddyExitBackoffScopeDecayAndRequestWaitBudget(t *testing.T) { + now := time.Date(2026, 10, 21, 0, 0, 0, 0, time.UTC) + p := NewPolicyState(func() time.Time { return now }) + p.jitter = func() float64 { return 1 } + var waits []time.Duration + p.sleep = func(_ context.Context, d time.Duration) error { waits = append(waits, d); return nil } + ctx, _ := WithRequestState(context.Background(), 2) + s := State(ctx) + for i := 0; i < 8; i++ { + if err := p.Backoff(ctx, s, "https://upstream.test/path", "http://user:secret@proxy.test"); err != nil { + t.Fatal(err) + } + } + var total time.Duration + for _, d := range waits { + total += d + } + if total != 5*time.Second { + t.Fatalf("logical request wait = %v", total) + } + key := ExitKey("https://upstream.test/another", "http://user:secret@proxy.test") + if p.exits[key].count != 1 { + t.Fatal("same request charged exit repeatedly") + } + second, _ := WithRequestState(context.Background(), 2) + if err := p.Backoff(second, State(second), "https://upstream.test", "http://user:secret@proxy.test"); err != nil { + t.Fatal(err) + } + if waits[len(waits)-1] != 2*time.Second { + t.Fatal("new request did not advance exit streak") + } + if key == ExitKey("https://different.test", "http://user:secret@proxy.test") { + t.Fatal("different origins shared failure scope") + } + now = now.Add(61 * time.Second) + third, _ := WithRequestState(context.Background(), 2) + _ = p.Backoff(third, State(third), "https://upstream.test", "http://user:secret@proxy.test") + if waits[len(waits)-1] != time.Second { + t.Fatal("exit streak did not decay") + } + p.ExitSuccess("https://upstream.test", "http://user:secret@proxy.test") + if _, ok := p.exits[key]; ok { + t.Fatal("success did not clear exit streak") + } + canceled, cancel := context.WithCancel(context.Background()) + cancel() + if err := waitBeforeConnect(canceled, time.Second); !errors.Is(err, context.Canceled) { + t.Fatalf("cancel = %v", err) + } +} +func TestWorkBuddyProductSwitchBudgetAndCNDeviceGate(t *testing.T) { + now := time.Date(2026, 10, 21, 0, 0, 0, 0, time.UTC) + p := NewPolicyState(func() time.Time { return now }) + c := testCredential() + expected := []Product{VSCode, CLI, Desktop, VSCode} + for _, want := range expected { + got, ok := p.NextProduct(c, "model", nil) + if !ok || got != want { + t.Fatalf("next=%s/%v want %s", got, ok, want) + } + c.Product = got + } + if _, ok := p.NextProduct(c, "model", nil); ok { + t.Fatal("fifth switch allowed") + } + now = now.Add(61 * time.Second) + if _, ok := p.NextProduct(c, "model", nil); !ok { + t.Fatal("switch window did not expire") + } + c.Realm = China + c.Product = Desktop + c.DeviceToken = "" + c.EnterpriseID = "" + if got, ok := p.NextProduct(c, "model", nil); !ok || got != CLI { + t.Fatal("CN personal VSCode identity fabricated device token") + } + if _, ok := p.NextProduct(c, "only-current", []Product{Desktop}); ok { + t.Fatal("single allowed identity switched") + } +} +func TestWorkBuddyIdentityChangeCannotClearCapOrRestriction(t *testing.T) { + g, now, _ := governorFixture(t, Limits{}) + c := testCredential() + if _, err := g.RecordFailure(c, "free", Fault{Kind: FaultModelCap, ResetAt: now.Add(time.Hour)}, FailureOptions{FreeModel: true}); err != nil { + t.Fatal(err) + } + if g.ClearIdentityModelCooldown(c, "free") { + t.Fatal("identity change cleared durable cap") + } + if _, err := g.RecordFailure(c, "ordinary", Fault{Kind: FaultModelRate}, FailureOptions{}); err != nil { + t.Fatal(err) + } + if !g.ClearIdentityModelCooldown(c, "ordinary") { + t.Fatal("ordinary model cooldown not cleared") + } + if _, err := g.RecordFailure(c, "ordinary", Fault{Kind: FaultRestricted}, FailureOptions{}); err != nil { + t.Fatal(err) + } + if g.ClearIdentityModelCooldown(c, "ordinary") { + t.Fatal("identity change cleared restriction") + } +} diff --git a/sdk/cliproxy/auth/conductor.go b/sdk/cliproxy/auth/conductor.go index 6f16017bfef..1da9577cda5 100644 --- a/sdk/cliproxy/auth/conductor.go +++ b/sdk/cliproxy/auth/conductor.go @@ -300,3 +300,9 @@ func (m *Manager) ResultPolicy() ResultPolicy { } return holder.policy } + +// ProviderAuthSelector optionally refines selection among already eligible, +// authorized credentials. Implementations must return one supplied candidate. +type ProviderAuthSelector interface { + SelectAuth(context.Context, string, cliproxyexecutor.Options, []*Auth) (*Auth, error) +} diff --git a/sdk/cliproxy/auth/conductor_selection.go b/sdk/cliproxy/auth/conductor_selection.go index 7cdacb7185f..5ddb4b214c9 100644 --- a/sdk/cliproxy/auth/conductor_selection.go +++ b/sdk/cliproxy/auth/conductor_selection.go @@ -1896,7 +1896,14 @@ func (m *Manager) pickNextLegacy(ctx context.Context, provider, model string, op if !handled { selectorCtx := selectorContextForAvailableAuths(ctx, selector, model) selectorCtx = withSessionAffinityTransientAuthIDs(selectorCtx, transientAuthIDs) - selected, errPick = selector.Pick(selectorCtx, provider, selectionArgForSelector(selector, model), opts, selectorAuths) + if native, ok := executor.(ProviderAuthSelector); ok { + selected, errPick = native.SelectAuth(selectorCtx, model, opts, available) + if errPick == nil { + selected, errPick = validateNativeSelection(selected, available) + } + } else { + selected, errPick = selector.Pick(selectorCtx, provider, selectionArgForSelector(selector, model), opts, selectorAuths) + } if errPick != nil { if isBuiltInSelector(selector) { errPick = restoreModelCooldownErrorModel(errPick, model) @@ -2094,7 +2101,7 @@ func (m *Manager) pickNext(ctx context.Context, provider, model string, opts cli opts.Metadata[cliproxyexecutor.SessionAffinityProviderMetadataKey] = provider opts.Metadata[cliproxyexecutor.SessionAffinityModelMetadataKey] = model - if m.hasPluginScheduler() || !m.useSchedulerFastPath() { + if m.hasNativeAuthSelector([]string{provider}) || m.hasPluginScheduler() || !m.useSchedulerFastPath() { return m.pickNextLegacy(ctx, provider, model, opts, tried) } eligibility := authSelectionEligibilityForRequest(ctx, opts) @@ -2259,6 +2266,23 @@ func (m *Manager) pickNextMixedLegacy(ctx context.Context, providers []string, m if !okExecutor { return nil, nil, "", &Error{Code: "executor_not_found", Message: "executor not registered"} } + if !handled { + if native, ok := executor.(ProviderAuthSelector); ok { + group := make([]*Auth, 0, len(available)) + for _, candidate := range available { + if executorKeyFromAuth(candidate) == providerKey { + group = append(group, candidate) + } + } + selected, errPick = native.SelectAuth(ctx, model, opts, group) + if errPick == nil { + selected, errPick = validateNativeSelection(selected, group) + } + if errPick != nil { + return nil, nil, "", errPick + } + } + } authCopy := selected.Clone() if !selected.indexAssigned { m.mu.Lock() @@ -2279,7 +2303,7 @@ func (m *Manager) pickNextMixed(ctx context.Context, providers []string, model s opts.Metadata[cliproxyexecutor.SessionAffinityProviderMetadataKey] = "mixed" opts.Metadata[cliproxyexecutor.SessionAffinityModelMetadataKey] = model - if m.hasPluginScheduler() || !m.useSchedulerFastPath() { + if m.hasNativeAuthSelector(providers) || m.hasPluginScheduler() || !m.useSchedulerFastPath() { return m.pickNextMixedLegacy(ctx, providers, model, opts, tried) } @@ -2464,3 +2488,24 @@ func (m *Manager) warnLogAuthUnavailable(ctx context.Context, providers []string } } } + +func (m *Manager) hasNativeAuthSelector(providers []string) bool { + for _, provider := range providers { + if executor, ok := m.Executor(provider); ok { + if _, native := executor.(ProviderAuthSelector); native { + return true + } + } + } + return false +} +func validateNativeSelection(selected *Auth, candidates []*Auth) (*Auth, error) { + if selected != nil { + for _, candidate := range candidates { + if candidate.ID == selected.ID { + return candidate, nil + } + } + } + return nil, &Error{Code: "auth_not_found", Message: "provider selector returned an ineligible auth"} +} diff --git a/sdk/cliproxy/auth/conductor_workbuddy_test.go b/sdk/cliproxy/auth/conductor_workbuddy_test.go index 3c5ab4f4125..6ff4f8381dc 100644 --- a/sdk/cliproxy/auth/conductor_workbuddy_test.go +++ b/sdk/cliproxy/auth/conductor_workbuddy_test.go @@ -101,3 +101,48 @@ func (e *workBuddyRefreshTrap) Refresh(context.Context, *Auth) (*Auth, error) { e.calls++ return nil, errors.New("unexpected refresh") } + +type workBuddyNativeSelectorExecutor struct { + mockCustomErrorExecutor + want string + calls int +} + +func (e *workBuddyNativeSelectorExecutor) SelectAuth(_ context.Context, _ string, _ execution.Options, candidates []*Auth) (*Auth, error) { + e.calls++ + for _, a := range candidates { + if a.ID == e.want { + return a, nil + } + } + return &Auth{ID: e.want, Provider: "workbuddy"}, nil +} +func TestWorkBuddyNativeSelectorRespectsSDKCandidateBoundary(t *testing.T) { + manager := NewManager(nil, &RoundRobinSelector{}, nil) + for _, id := range []string{"wb-select-a", "wb-select-b"} { + a := &Auth{ID: id, Provider: "workbuddy", Status: StatusActive} + if _, err := manager.Register(context.Background(), a); err != nil { + t.Fatal(err) + } + registry.GetGlobalRegistry().RegisterClient(id, "workbuddy", []*registry.ModelInfo{{ID: "fixture-model"}}) + t.Cleanup(func() { registry.GetGlobalRegistry().UnregisterClient(id) }) + } + executor := &workBuddyNativeSelectorExecutor{mockCustomErrorExecutor: mockCustomErrorExecutor{identifier: "workbuddy"}, want: "wb-select-b"} + manager.RegisterExecutor(executor) + selected, _, err := manager.pickNext(context.Background(), "workbuddy", "fixture-model", execution.Options{}, nil) + if err != nil || selected.ID != "wb-select-b" { + t.Fatalf("native selection = %v %v", selected, err) + } + selected, _, _, err = manager.pickNextMixed(context.Background(), []string{"workbuddy"}, "fixture-model", execution.Options{}, nil) + if err != nil || selected.ID != "wb-select-b" { + t.Fatalf("mixed native selection = %v %v", selected, err) + } + _, _, err = manager.pickNext(context.Background(), "workbuddy", "fixture-model", execution.Options{Metadata: map[string]any{execution.PinnedAuthMetadataKey: "wb-select-a"}}, nil) + if err == nil { + t.Fatal("provider selector bypassed explicit pinned auth") + } + _, _, err = manager.pickNext(context.Background(), "workbuddy", "fixture-model", execution.Options{}, map[string]struct{}{"wb-select-b": {}}) + if err == nil { + t.Fatal("provider selector reintroduced excluded auth") + } +} From 3b68d35114417cf3c93ff549860ab1c37869cc5d Mon Sep 17 00:00:00 2001 From: LuneZ99 Date: Wed, 21 Oct 2026 03:08:54 +0800 Subject: [PATCH 06/10] feat(workbuddy): add governed maintenance and isolated image delivery --- .github/workflows/workbuddy-image.yml | 65 ++++ Dockerfile | 1 + docs/workbuddy-native.md | 53 ++- internal/config/workbuddy.go | 11 + .../runtime/executor/helps/workbuddy_jobs.go | 215 ++++++++++++ .../executor/helps/workbuddy_runtime.go | 29 ++ .../runtime/executor/workbuddy_executor.go | 57 +++- .../executor/workbuddy_executor_test.go | 144 ++++++++ internal/workbuddy/balance.go | 228 +++++++++++++ internal/workbuddy/balance_test.go | 91 +++++ internal/workbuddy/control.go | 323 ++++++++++++++++++ internal/workbuddy/control_test.go | 175 ++++++++++ internal/workbuddy/discovery.go | 201 +++++++++++ internal/workbuddy/discovery_test.go | 54 +++ internal/workbuddy/governor.go | 42 ++- internal/workbuddy/governor_test.go | 28 ++ internal/workbuddy/growth.go | 142 ++++++++ internal/workbuddy/jobs.go | 184 ++++++++++ internal/workbuddy/jobs_test.go | 207 +++++++++++ internal/workbuddy/probes.go | 37 ++ internal/workbuddy/stream.go | 12 + sdk/cliproxy/service_lifecycle.go | 1 + sdk/cliproxy/service_models.go | 16 +- sdk/cliproxy/service_workbuddy.go | 50 +++ 24 files changed, 2348 insertions(+), 18 deletions(-) create mode 100644 .github/workflows/workbuddy-image.yml create mode 100644 internal/runtime/executor/helps/workbuddy_jobs.go create mode 100644 internal/workbuddy/balance.go create mode 100644 internal/workbuddy/balance_test.go create mode 100644 internal/workbuddy/control.go create mode 100644 internal/workbuddy/control_test.go create mode 100644 internal/workbuddy/discovery.go create mode 100644 internal/workbuddy/discovery_test.go create mode 100644 internal/workbuddy/growth.go create mode 100644 internal/workbuddy/jobs.go create mode 100644 internal/workbuddy/jobs_test.go create mode 100644 internal/workbuddy/probes.go create mode 100644 sdk/cliproxy/service_workbuddy.go diff --git a/.github/workflows/workbuddy-image.yml b/.github/workflows/workbuddy-image.yml new file mode 100644 index 00000000000..f6cc71d3280 --- /dev/null +++ b/.github/workflows/workbuddy-image.yml @@ -0,0 +1,65 @@ +name: WorkBuddy isolated image + +on: + workflow_dispatch: + push: + branches: [main] + paths: + - 'internal/workbuddy/**' + - 'internal/runtime/executor/workbuddy*' + - 'internal/runtime/executor/helps/workbuddy*' + - 'sdk/cliproxy/**' + - 'internal/config/**' + - 'go.mod' + - 'go.sum' + - 'Dockerfile' + - '.github/workflows/workbuddy-image.yml' + +permissions: + contents: read + packages: write + +jobs: + build: + if: github.repository == 'CorgiCortex/CLIProxyAPI' + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-go@v5 + with: + go-version-file: go.mod + cache: true + - name: Verify native WorkBuddy contracts + env: + CGO_ENABLED: '1' + run: | + go test -race -timeout 180s ./internal/workbuddy ./sdk/auth ./internal/watcher/synthesizer + go test -race -timeout 180s ./internal/runtime/executor ./internal/runtime/executor/helps ./sdk/cliproxy/auth ./sdk/cliproxy -run TestWorkBuddy + go build -o test-output ./cmd/server + rm test-output + - uses: docker/setup-buildx-action@v3 + - uses: docker/login-action@v3 + with: + registry: ghcr.io + username: ${{ github.actor }} + password: ${{ secrets.GITHUB_TOKEN }} + - name: Build immutable test-service image + id: image + uses: docker/build-push-action@v6 + with: + context: . + platforms: linux/amd64 + push: true + tags: ghcr.io/corgicortex/cliproxyapi-wb:sha-${{ github.sha }} + build-args: | + VERSION=workbuddy-${{ github.sha }} + COMMIT=${{ github.sha }} + cache-from: type=gha,scope=workbuddy + cache-to: type=gha,mode=max,scope=workbuddy + - name: Record image receipt + env: + IMAGE_DIGEST: ${{ steps.image.outputs.digest }} + run: | + echo "Image: ghcr.io/corgicortex/cliproxyapi-wb:sha-${GITHUB_SHA}" >> "$GITHUB_STEP_SUMMARY" + echo "Digest: ${IMAGE_DIGEST}" >> "$GITHUB_STEP_SUMMARY" + echo "Deployment requires a separate GitOps image pin." >> "$GITHUB_STEP_SUMMARY" diff --git a/Dockerfile b/Dockerfile index a24a8d6156f..d3cda51bbd6 100644 --- a/Dockerfile +++ b/Dockerfile @@ -25,6 +25,7 @@ RUN mkdir /CLIProxyAPI COPY --from=builder ./app/CLIProxyAPI /CLIProxyAPI/CLIProxyAPI COPY config.example.yaml /CLIProxyAPI/config.example.yaml +COPY internal/workbuddy/NOTICE /usr/share/doc/cliproxyapi/WORKBUDDY-NOTICE WORKDIR /CLIProxyAPI diff --git a/docs/workbuddy-native.md b/docs/workbuddy-native.md index f318a7653a2..056f68fc1fe 100644 --- a/docs/workbuddy-native.md +++ b/docs/workbuddy-native.md @@ -55,9 +55,12 @@ HTTP errors and error envelopes inside HTTP 200 streams share the risk classifie A truncated stream is an error; once a chunk was emitted it cannot be replayed on another credential. Only a pre-output connection failure gets one retry. -The static model catalog is an offline snapshot. It is not evidence of live -model availability. Native credential `models` can explicitly restrict the -published models, including synthetic models used by offline tests. +The static model catalog is an offline snapshot. Opt-in discovery reads +`/v3/config`, prefers the CLI agent's model list and merges observed metadata +with the snapshot. Empty responses cannot erase the catalog. Discovery persists +per-account results and republishes model registrations with the existing SDK +publication fences. Neither source establishes entitlement to run a model. +Native credential `models` explicitly restrict the published models. The governor and transport pools are shared across executor replacements. Configuration reloads preserve cooldowns, restrictions and usage. Changing the @@ -87,9 +90,34 @@ fallback applies only to international DeepSeek Flash when every enabled account in that region has an active free-model cap; paid budgets are checked again. All rebuilt requests pass through the final payload-rule barrier exactly once. -The remaining stage-2 work is live model discovery. Stage-3 quota and activity -jobs are not implemented or enabled yet. These gaps must be closed before -claiming parity with the reference project's risk controls. +Background operations now share chat's governor, connection pools and limiter: + +- Model discovery and balance observations run in six-hour slots. Resource + pagination must complete before publishing totals. Missing numeric fields + remain unknown, including the enterprise endpoint's unreported remainder. +- CN personal sign-in runs after 09:00 CST. An explicit already-checked-in + response is completion evidence, with no invented reward amount. +- Growth reconciliation runs after 09:00 and 21:00 CST. It verifies individual + acceptance results and claims only vendor-observed completion. Desktop, + subscription, donation and other real actions are not fabricated as telemetry. +- Optional quota scans use the three known free-cap models as applicable to the + realm. Reset probes wait for the recorded reset time; both execute through the + native chat path with paid fallback and product switching disabled. Each cycle + admits at most 16 probes, with 400 ms spacing. +- Atomic job receipts precede actions. Completed slots survive restart; ordinary + failures retry after 15 minutes, at most six times per slot. Reset failures + retry after at least 15 seconds, at most 20 times per reset window, subject to + any longer governor cooldown. Reload and shutdown cancel active maintenance. + +Background jobs default to disabled. They record only observed credits and +rewards. Inference `usage.credit` feeds the paid daily budget and local debits +against the last observed balance, protecting the reserve between refreshes. +The persisted balance snapshot retains the vendor's original observation. + +Automated international web activity and actual task-specific growth actions +remain follow-up work. The isolated deployment and real-account validation are +still pending; this implementation must not be described as fully validated +upstream parity. An explicit v8 configuration enables the provider: @@ -103,6 +131,14 @@ oauth: max-concurrency: 32 per-caller-concurrency: 2 queue-seconds: 30 + jobs: + enabled: false + model-discovery: false + balance-refresh: false + checkin: false + growth: false + quota-scan: false + reset-probes: false limits: daily-tokens: 0 daily-credits: 0 @@ -118,6 +154,11 @@ Queue deadlines end at admission and never become upstream network deadlines. The general management HTTP tunnel is unsupported for this provider; arbitrary requests must not bypass its risk governor or export credentials to another host. +The `WorkBuddy isolated image` workflow publishes only +`ghcr.io/corgicortex/cliproxyapi-wb:sha-` after offline race checks. +It never updates the existing DockerHub image or deploys a workload. Deployment +requires a separately reviewed GitOps pin. + ## Credential setup Use a separate configuration and auth directory for this test service. Never diff --git a/internal/config/workbuddy.go b/internal/config/workbuddy.go index b451d77f5f5..e7aaad80a84 100644 --- a/internal/config/workbuddy.go +++ b/internal/config/workbuddy.go @@ -3,6 +3,7 @@ package config // WorkBuddyConfig is opt-in. State belongs to a dedicated provider volume and // must never be shared between independent server processes. type WorkBuddyConfig struct { + Jobs WorkBuddyJobs `yaml:"jobs" json:"jobs"` AffinityByPrefix *bool `yaml:"affinity-by-prefix" json:"affinity-by-prefix"` AffinityMaxMessages *int `yaml:"affinity-max-messages" json:"affinity-max-messages"` AffinityPageSize *int `yaml:"affinity-page-size" json:"affinity-page-size"` @@ -20,6 +21,16 @@ type WorkBuddyConfig struct { RealmLimits map[string]WorkBuddyLimits `yaml:"realm-limits" json:"realm-limits"` } +type WorkBuddyJobs struct { + Enabled bool `yaml:"enabled" json:"enabled"` + ModelDiscovery bool `yaml:"model-discovery" json:"model-discovery"` + BalanceRefresh bool `yaml:"balance-refresh" json:"balance-refresh"` + Checkin bool `yaml:"checkin" json:"checkin"` + Growth bool `yaml:"growth" json:"growth"` + QuotaScan bool `yaml:"quota-scan" json:"quota-scan"` + ResetProbes bool `yaml:"reset-probes" json:"reset-probes"` +} + type WorkBuddyLimits struct { ReserveCredits float64 `yaml:"reserve-credits" json:"reserve-credits"` DailyTokens int64 `yaml:"daily-tokens" json:"daily-tokens"` diff --git a/internal/runtime/executor/helps/workbuddy_jobs.go b/internal/runtime/executor/helps/workbuddy_jobs.go new file mode 100644 index 00000000000..de73e5e9753 --- /dev/null +++ b/internal/runtime/executor/helps/workbuddy_jobs.go @@ -0,0 +1,215 @@ +package helps + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "os" + "path/filepath" + "sort" + "time" + + "github.com/router-for-me/CLIProxyAPI/v8/internal/config" + "github.com/router-for-me/CLIProxyAPI/v8/internal/workbuddy" + coreauth "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/auth" + log "github.com/sirupsen/logrus" +) + +type workBuddyCatalogCache struct { + ObservedAt time.Time `json:"observed_at"` + Models []workbuddy.Model `json:"models"` +} + +func (r *WorkBuddyRuntime) Models(c workbuddy.Credential) ([]workbuddy.Model, error) { + raw, err := os.ReadFile(filepath.Join(r.StateDir, "catalog-"+workbuddy.RiskAccountKey(c)+".json")) + if errors.Is(err, os.ErrNotExist) { + return workbuddy.Catalog(c.Realm), nil + } + if err != nil { + return nil, fmt.Errorf("workbuddy: read model catalog: %w", err) + } + var cached workBuddyCatalogCache + if err = json.Unmarshal(raw, &cached); err != nil || len(cached.Models) == 0 { + return nil, fmt.Errorf("workbuddy: invalid cached model catalog") + } + return cached.Models, nil +} + +// Maintenance is a single process-wide lane. Reload cancels an active cycle; +// every control request also reloads the authoritative credential before dispatch. +func (r *WorkBuddyRuntime) Maintenance(ctx context.Context, cfg *config.Config, epoch uint64, probe ...func(context.Context, *coreauth.Auth, string) error) ([]string, error) { + if cfg == nil || !cfg.WorkBuddy.Enabled || !cfg.WorkBuddy.Jobs.Enabled { + return nil, nil + } + if !r.jobMu.TryLock() { + return nil, nil + } + defer r.jobMu.Unlock() + ctx, cancel := context.WithCancel(ctx) + defer cancel() + r.mu.Lock() + if epoch != r.jobEpoch { + r.mu.Unlock() + return nil, nil + } + r.jobCancel = cancel + manager := r.manager + var records []*coreauth.Auth + if manager == nil { + for _, a := range r.auths { + records = append(records, a.Clone()) + } + } + r.mu.Unlock() + defer func() { + r.mu.Lock() + if r.jobEpoch == epoch { + r.jobCancel = nil + } + r.mu.Unlock() + }() + if manager != nil { + records = manager.List() + } + sort.Slice(records, func(i, j int) bool { return records[i].ID < records[j].ID }) + seen := make(map[string]bool) + var changed []string + var failures []error + policy := cfg.WorkBuddy.Jobs + probeCount := 0 + for _, auth := range records { + if err := ctx.Err(); err != nil { + return changed, err + } + c, err := WorkBuddyCredential(r.Current(auth)) + if err != nil || c.Disabled { + continue + } + key := workbuddy.RiskAccountKey(c) + if seen[key] { + continue + } + seen[key] = true + load := func() (workbuddy.Credential, error) { + r.mu.Lock() + current := r.jobEpoch == epoch + r.mu.Unlock() + if !current { + return workbuddy.Credential{}, context.Canceled + } + credential, errLoad := WorkBuddyCredential(r.Current(auth)) + if errLoad == nil && workbuddy.RiskAccountKey(credential) != key { + return credential, fmt.Errorf("workbuddy: account changed during background work") + } + return credential, errLoad + } + client := &workbuddy.ControlClient{Governor: r.Governor, Transports: r.Transports, Limiter: r.Limiter, Versions: workbuddy.DefaultVersions(), ProxyURL: cfg.ProxyURL, BaseURL: r.BaseURL, QueueWait: time.Duration(cfg.WorkBuddy.QueueSeconds) * time.Second, MinimumInterval: time.Second, Load: load} + client.Refresh = func(ctx context.Context) (workbuddy.Credential, error) { + if _, errLoad := load(); errLoad != nil { + return workbuddy.Credential{}, errLoad + } + updated, errRefresh := r.Refresh(ctx, auth, cfg.ProxyURL) + if errRefresh != nil { + return workbuddy.Credential{}, errRefresh + } + return WorkBuddyCredential(updated) + } + jobs := []struct { + name string + enabled bool + }{{"models", policy.ModelDiscovery}, {"balance", policy.BalanceRefresh}, {"checkin", policy.Checkin && c.Realm == workbuddy.China && c.EnterpriseID == ""}, {"growth", policy.Growth && c.Realm == workbuddy.China && c.EnterpriseID == ""}} + for _, job := range jobs { + if !job.enabled { + continue + } + slot, due := workbuddy.JobSlot(time.Now(), job.name) + if !due { + continue + } + _, ran, errRun := r.Jobs.Run(ctx, c, job.name, slot, func(ctx context.Context) (workbuddy.JobReceipt, error) { + switch job.name { + case "models": + models, errModels := client.DiscoverModels(ctx) + if errModels != nil { + return workbuddy.JobReceipt{}, errModels + } + errWrite := workbuddy.WriteJSON(filepath.Join(r.StateDir, "catalog-"+key+".json"), workBuddyCatalogCache{ObservedAt: time.Now(), Models: models}) + if errWrite == nil { + changed = append(changed, auth.ID) + } + return workbuddy.JobReceipt{}, errWrite + case "balance": + balance, errBalance := client.FetchBalance(ctx) + if errBalance != nil { + return workbuddy.JobReceipt{}, errBalance + } + return workbuddy.JobReceipt{}, workbuddy.WriteJSON(filepath.Join(r.StateDir, "balance-"+key+".json"), balance) + case "checkin": + result, errCheckin := client.Call(ctx, workbuddy.ControlRequest{Operation: workbuddy.ControlCheckin}) + if errCheckin != nil { + return workbuddy.JobReceipt{}, errCheckin + } + return workbuddy.RewardReceipt(result) + case "growth": + pending, errGrowth := client.ReconcileGrowth(ctx, r.Jobs) + return workbuddy.JobReceipt{PendingTasks: pending}, errGrowth + default: + return workbuddy.JobReceipt{}, fmt.Errorf("workbuddy: unsupported maintenance task") + } + }) + if errRun != nil { + reason := "operation_failed" + var dispatch *workbuddy.DispatchError + if errors.As(errRun, &dispatch) { + reason = dispatch.Reason + } + log.WithFields(log.Fields{"provider": "workbuddy", "account_key": key, "job": job.name, "reason": reason}).Warn("background task incomplete") + failures = append(failures, fmt.Errorf("workbuddy: %s: %w", job.name, errRun)) + } else if ran { + log.WithFields(log.Fields{"provider": "workbuddy", "account_key": key, "job": job.name}).Debug("background task recorded") + } + } + if len(probe) > 0 && (policy.QuotaScan || policy.ResetProbes) { + plans := r.Governor.DueResetProbes(c) + if !policy.ResetProbes { + plans = nil + } + if policy.QuotaScan { + for _, model := range workbuddy.FreeProbeModels(c.Realm) { + plans = append(plans, workbuddy.ResetProbe{Model: model}) + } + } + for _, plan := range plans { + if probeCount >= 16 { + return changed, errors.Join(failures...) + } + job, slot := "reset_probe:"+plan.Model, plan.ResetAt.UTC().Format(time.RFC3339Nano) + if plan.ResetAt.IsZero() { + job = "quota_scan:" + plan.Model + slot, _ = workbuddy.JobSlot(time.Now(), "quota") + } + _, _, errProbe := r.Jobs.Run(ctx, c, job, slot, func(ctx context.Context) (workbuddy.JobReceipt, error) { + if _, errLoad := load(); errLoad != nil { + return workbuddy.JobReceipt{}, errLoad + } + if probeCount > 0 { + timer := time.NewTimer(400 * time.Millisecond) + select { + case <-ctx.Done(): + timer.Stop() + return workbuddy.JobReceipt{}, ctx.Err() + case <-timer.C: + } + } + probeCount++ + return workbuddy.JobReceipt{}, probe[0](ctx, r.Current(auth), plan.Model) + }) + if errProbe != nil { + failures = append(failures, fmt.Errorf("workbuddy: quota probe: %w", errProbe)) + } + } + } + } + return changed, errors.Join(failures...) +} diff --git a/internal/runtime/executor/helps/workbuddy_runtime.go b/internal/runtime/executor/helps/workbuddy_runtime.go index b9041110f64..8dd8114226d 100644 --- a/internal/runtime/executor/helps/workbuddy_runtime.go +++ b/internal/runtime/executor/helps/workbuddy_runtime.go @@ -21,6 +21,10 @@ type WorkBuddyRuntime struct { Limiter *workbuddy.Limiter Policy *workbuddy.PolicyState Affinity *workbuddy.Affinity + Jobs *workbuddy.JobLedger + jobMu sync.Mutex + jobCancel context.CancelFunc + jobEpoch uint64 // BaseURL is an injection seam for offline tests, never populated from auth files. BaseURL string StateDir string @@ -41,6 +45,10 @@ func NewWorkBuddyRuntime(cfg config.WorkBuddyConfig) (*WorkBuddyRuntime, error) if err != nil { return nil, err } + r.Jobs, err = workbuddy.OpenJobLedger(filepath.Join(r.StateDir, "jobs.json"), nil) + if err != nil { + return nil, err + } r.Limiter, err = workbuddy.NewLimiter(cfg.MaxConcurrency, cfg.PerCallerConcurrency) if err != nil { return nil, err @@ -79,6 +87,7 @@ func (r *WorkBuddyRuntime) BindManager(manager *coreauth.Manager) { } func (r *WorkBuddyRuntime) Configure(cfg config.WorkBuddyConfig) error { + r.CancelMaintenance() if filepath.Clean(cfg.StateDir) != r.StateDir { return fmt.Errorf("workbuddy: changing state-dir requires a process restart") } @@ -110,6 +119,22 @@ func (r *WorkBuddyRuntime) Configure(cfg config.WorkBuddyConfig) error { } return r.Limiter.Configure(cfg.MaxConcurrency, cfg.PerCallerConcurrency) } + +func (r *WorkBuddyRuntime) CancelMaintenance() { + r.mu.Lock() + defer r.mu.Unlock() + r.jobEpoch++ + if r.jobCancel != nil { + r.jobCancel() + r.jobCancel = nil + } +} + +func (r *WorkBuddyRuntime) MaintenanceEpoch() uint64 { + r.mu.Lock() + defer r.mu.Unlock() + return r.jobEpoch +} func WorkBuddyCredential(a *coreauth.Auth) (workbuddy.Credential, error) { if a == nil || a.Provider != "workbuddy" { return workbuddy.Credential{}, fmt.Errorf("workbuddy: missing credential") @@ -176,6 +201,10 @@ func (r *WorkBuddyRuntime) Refresh(ctx context.Context, a *coreauth.Auth, proxy if err != nil { return nil, err } + gate, _ := r.Governor.Check(c, "credential_refresh", false, true) + if !gate.Allowed() && gate.Reason != "credential_expired" { + return nil, &workbuddy.DispatchError{Reason: gate.Reason, Status: 429, Stop: gate.Stop, Until: gate.Until} + } client := &workbuddy.AuthClient{Transports: r.Transports, Versions: workbuddy.DefaultVersions(), ProxyURL: proxy, BaseURL: r.BaseURL} var saved, loaded *coreauth.Auth refreshed, err := r.RefreshGate.Run(ctx, workbuddy.RiskAccountKey(c), c.AccessToken, diff --git a/internal/runtime/executor/workbuddy_executor.go b/internal/runtime/executor/workbuddy_executor.go index 27bd8693640..62c25dc0203 100644 --- a/internal/runtime/executor/workbuddy_executor.go +++ b/internal/runtime/executor/workbuddy_executor.go @@ -29,9 +29,10 @@ var _ coreauth.RequestAuthPreparer = (*WorkBuddyExecutor)(nil) // WorkBuddyExecutor owns no account-specific transport or risk state. Those // belong to the shared runtime so rotations and reloads cannot bypass gates. type WorkBuddyExecutor struct { - cfg *config.Config - runtime *helps.WorkBuddyRuntime - initErr error + cfg *config.Config + runtime *helps.WorkBuddyRuntime + initErr error + runtimeEpoch uint64 } func NewWorkBuddyExecutor(cfg *config.Config, previous *WorkBuddyExecutor) *WorkBuddyExecutor { @@ -40,6 +41,9 @@ func NewWorkBuddyExecutor(cfg *config.Config, previous *WorkBuddyExecutor) *Work e.runtime = previous.runtime } if cfg == nil || !cfg.WorkBuddy.Enabled { + if e.runtime != nil { + e.runtime.CancelMaintenance() + } e.initErr = fmt.Errorf("workbuddy: provider is disabled") return e } @@ -48,6 +52,9 @@ func NewWorkBuddyExecutor(cfg *config.Config, previous *WorkBuddyExecutor) *Work } else { e.initErr = e.runtime.Configure(cfg.WorkBuddy) } + if e.runtime != nil { + e.runtimeEpoch = e.runtime.MaintenanceEpoch() + } return e } func (e *WorkBuddyExecutor) Identifier() string { return "workbuddy" } @@ -62,6 +69,44 @@ func (e *WorkBuddyExecutor) TrackAuth(a *coreauth.Auth) { e.runtime.Track(a) } } + +func (e *WorkBuddyExecutor) ModelCatalog(a *coreauth.Auth) ([]workbuddy.Model, error) { + if e.initErr != nil { + return nil, e.initErr + } + c, err := helps.WorkBuddyCredential(a) + if err != nil { + return nil, err + } + return e.runtime.Models(c) +} + +func (e *WorkBuddyExecutor) RunMaintenance(ctx context.Context) ([]string, error) { + if e.initErr != nil { + return nil, e.initErr + } + return e.runtime.Maintenance(ctx, e.cfg, e.runtimeEpoch, func(ctx context.Context, auth *coreauth.Auth, model string) error { + probeExecutor := *e + probeConfig := *e.cfg + probeConfig.WorkBuddy.PaidFallback = false + probeConfig.WorkBuddy.AutoSwitchProduct = false + probeConfig.WorkBuddy.Prompt = config.WorkBuddyPrompt{Mode: "custom", Text: workbuddy.DegradedPrompt} + probeExecutor.cfg = &probeConfig + payload, err := json.Marshal(map[string]any{ + "model": model, + "messages": []map[string]string{{"role": "user", "content": "Reply with OK."}}, + "max_tokens": 16, + }) + if err != nil { + return err + } + opts := execution.Options{SourceFormat: sdktranslator.FromString("openai"), Metadata: map[string]any{ + execution.ExecutionSessionMetadataKey: "workbuddy-quota-probe:" + model, + }} + _, err = probeExecutor.Execute(ctx, auth, execution.Request{Model: model, Payload: payload}, opts) + return err + }) +} func (e *WorkBuddyExecutor) PrepareExecutionContext(ctx context.Context, req execution.Request, opts execution.Options) (context.Context, error) { if e.initErr != nil { return ctx, &workbuddy.DispatchError{Reason: "provider_unavailable", Status: 503, Stop: true} @@ -325,6 +370,7 @@ func (e *WorkBuddyExecutor) execute(ctx context.Context, a *coreauth.Auth, req e emitted := false var usage json.RawMessage var tokens int64 + var credits *float64 finished := make(map[int]bool) errStream := workbuddy.ReadEvents(response.Body, func(raw []byte) error { event, errEvent := workbuddy.DecodeEvent(raw) @@ -336,6 +382,9 @@ func (e *WorkBuddyExecutor) execute(ctx context.Context, a *coreauth.Auth, req e if observed, ok := workbuddy.ObservedTokens(usage); ok { tokens = observed } + if observed := workbuddy.ObservedCredits(usage); observed != nil { + credits = observed + } } reporter.ObserveResponseModel(raw) reporter.ObserveTokenEvent(gjson.GetBytes(raw, "choices.0.delta.content").String() != "" || gjson.GetBytes(raw, "choices.0.delta.reasoning_content").String() != "" || gjson.GetBytes(raw, "choices.0.delta.tool_calls.#").Int() > 0) @@ -367,7 +416,7 @@ func (e *WorkBuddyExecutor) execute(ctx context.Context, a *coreauth.Auth, req e if len(usage) > 0 { reportedUsage = usage } - if errComplete := e.runtime.Governor.Complete(c, attempt, tokens, nil, errStream == nil); errComplete != nil { + if errComplete := e.runtime.Governor.Complete(c, attempt, tokens, credits, errStream == nil); errComplete != nil { return &workbuddy.DispatchError{Reason: "usage_persistence_failed", Status: 503, Stop: true} } if errStream == nil { diff --git a/internal/runtime/executor/workbuddy_executor_test.go b/internal/runtime/executor/workbuddy_executor_test.go index c24ed656d2c..3d77017039c 100644 --- a/internal/runtime/executor/workbuddy_executor_test.go +++ b/internal/runtime/executor/workbuddy_executor_test.go @@ -8,10 +8,12 @@ import ( "io" "net/http" "net/http/httptest" + "path/filepath" "strings" "sync" "sync/atomic" "testing" + "time" "github.com/router-for-me/CLIProxyAPI/v8/internal/config" "github.com/router-for-me/CLIProxyAPI/v8/internal/runtime/executor/helps" @@ -627,3 +629,145 @@ func TestWorkBuddyNativeSelectionUsesWarmHistoryAndRiskGates(t *testing.T) { t.Fatal("model cooldown poisoned account") } } + +func TestWorkBuddyMaintenanceDefaultsOffAndCachesDiscoveredModels(t *testing.T) { + var calls atomic.Int32 + e, a := wbFixture(t, func(w http.ResponseWriter, r *http.Request) { + calls.Add(1) + if r.URL.Path != "/v3/config" { + t.Error("unexpected background endpoint") + } + _, _ = io.WriteString(w, `{"code":0,"data":{"agents":{"cli":{"models":["fixture-discovered"]}},"models":[{"id":"fixture-discovered","maxInputTokens":2048}]}}`) + }) + if _, err := e.RunMaintenance(context.Background()); err != nil || calls.Load() != 0 { + t.Fatal("disabled background jobs ran") + } + e.cfg.WorkBuddy.Jobs = config.WorkBuddyJobs{Enabled: true, ModelDiscovery: true} + changed, err := e.RunMaintenance(context.Background()) + if err != nil || len(changed) != 1 || changed[0] != a.ID || calls.Load() != 1 { + t.Fatalf("changed=%v calls=%d err=%v", changed, calls.Load(), err) + } + if _, err = e.RunMaintenance(context.Background()); err != nil || calls.Load() != 1 { + t.Fatal("completed slot repeated") + } + restarted := NewWorkBuddyExecutor(e.cfg, nil) + restarted.runtime.BaseURL = e.runtime.BaseURL + restarted.TrackAuth(a) + t.Cleanup(restarted.runtime.Transports.CloseIdleConnections) + if _, err = restarted.RunMaintenance(context.Background()); err != nil || calls.Load() != 1 { + t.Fatal("restart repeated completed discovery") + } + models, err := restarted.ModelCatalog(a) + if err != nil { + t.Fatal(err) + } + found := false + for _, model := range models { + if model.ID == "fixture-discovered" && model.Input == 2048 { + found = true + } + } + if !found { + t.Fatal("discovered catalog did not survive restart") + } +} + +func TestWorkBuddyDisableReloadCancelsBackgroundConnection(t *testing.T) { + started := make(chan struct{}) + e, _ := wbFixture(t, func(w http.ResponseWriter, r *http.Request) { close(started); <-r.Context().Done() }) + e.cfg.WorkBuddy.Jobs = config.WorkBuddyJobs{Enabled: true, ModelDiscovery: true} + done := make(chan error, 1) + go func() { _, err := e.RunMaintenance(context.Background()); done <- err }() + <-started + nextCfg := *e.cfg + nextCfg.WorkBuddy.Enabled = false + disabled := NewWorkBuddyExecutor(&nextCfg, e) + if disabled.initErr == nil { + t.Fatal("disabled provider remained available") + } + if err := <-done; err == nil { + t.Fatal("canceled background connection reported success") + } +} + +func TestWorkBuddyUpstreamCreditUsageEnforcesPaidDailyBudget(t *testing.T) { + for _, stream := range []bool{false, true} { + t.Run(map[bool]string{false: "aggregate", true: "stream"}[stream], func(t *testing.T) { + var calls atomic.Int32 + e, a := wbFixture(t, func(w http.ResponseWriter, r *http.Request) { + calls.Add(1) + _, _ = io.WriteString(w, strings.Replace(wbSuccess, `"total_tokens":10`, `"total_tokens":10,"credit":2.5`, 1)) + }) + e.cfg.WorkBuddy.Limits.DailyCredits = 2 + if err := e.runtime.Configure(e.cfg.WorkBuddy); err != nil { + t.Fatal(err) + } + req, opts := wbRequest() + if stream { + opts.Stream = true + result, err := e.ExecuteStream(context.Background(), a, req, opts) + if err != nil { + t.Fatal(err) + } + for chunk := range result.Chunks { + if chunk.Err != nil { + t.Fatal(chunk.Err) + } + } + } else if _, err := e.Execute(context.Background(), a, req, opts); err != nil { + t.Fatal(err) + } + _, err := e.Execute(context.Background(), a, req, opts) + if err == nil || !strings.Contains(err.Error(), "daily_credits") || calls.Load() != 1 { + t.Fatalf("calls=%d err=%v", calls.Load(), err) + } + }) + } +} + +func TestWorkBuddyResetProbeReusesGovernorAndStopsAfterObservedRecovery(t *testing.T) { + var calls atomic.Int32 + e, a := wbFixture(t, func(w http.ResponseWriter, r *http.Request) { calls.Add(1); _, _ = io.WriteString(w, wbSuccess) }) + now := time.Now().Add(-2 * time.Minute) + governor, err := workbuddy.OpenGovernor(filepath.Join(t.TempDir(), "risk.json"), workbuddy.GovernorOptions{Now: func() time.Time { return now }}) + if err != nil { + t.Fatal(err) + } + e.runtime.Governor = governor + e.cfg.WorkBuddy.Jobs = config.WorkBuddyJobs{Enabled: true, ResetProbes: true} + c, err := helps.WorkBuddyCredential(a) + if err != nil { + t.Fatal(err) + } + if _, err = governor.RecordFailure(c, "hy3", workbuddy.Fault{Kind: workbuddy.FaultModelCap, ResetAt: now.Add(time.Minute)}, workbuddy.FailureOptions{FreeModel: true}); err != nil { + t.Fatal(err) + } + if _, err = e.RunMaintenance(context.Background()); err != nil || calls.Load() != 0 { + t.Fatal("probe ran before reset") + } + now = now.Add(2 * time.Minute) + if _, err = e.RunMaintenance(context.Background()); err != nil || calls.Load() != 1 { + t.Fatalf("calls=%d err=%v", calls.Load(), err) + } + if len(governor.DueResetProbes(c)) != 0 { + t.Fatal("confirmed recovery retained pending reset") + } + if _, err = e.RunMaintenance(context.Background()); err != nil || calls.Load() != 1 { + t.Fatal("recovered model was probed again") + } +} + +func TestWorkBuddyStaleExecutorCannotStartMaintenanceAfterDisable(t *testing.T) { + var calls atomic.Int32 + e, _ := wbFixture(t, func(w http.ResponseWriter, r *http.Request) { + calls.Add(1) + _, _ = io.WriteString(w, `{"code":0,"data":{}}`) + }) + e.cfg.WorkBuddy.Jobs = config.WorkBuddyJobs{Enabled: true, ModelDiscovery: true} + disabledConfig := *e.cfg + disabledConfig.WorkBuddy.Enabled = false + _ = NewWorkBuddyExecutor(&disabledConfig, e) + if changed, err := e.RunMaintenance(context.Background()); err != nil || len(changed) != 0 || calls.Load() != 0 { + t.Fatal("old executor started maintenance after disable") + } +} diff --git a/internal/workbuddy/balance.go b/internal/workbuddy/balance.go new file mode 100644 index 00000000000..bb0fd8dff73 --- /dev/null +++ b/internal/workbuddy/balance.go @@ -0,0 +1,228 @@ +package workbuddy + +import ( + "context" + "encoding/json" + "fmt" + "strconv" + "time" +) + +type CreditPackage struct { + ID string `json:"id,omitempty"` + Code string `json:"code,omitempty"` + Remaining *float64 `json:"remaining,omitempty"` + Used *float64 `json:"used,omitempty"` + Capacity *float64 `json:"capacity,omitempty"` + CycleStart string `json:"cycle_start,omitempty"` + CycleEnd string `json:"cycle_end,omitempty"` +} +type BalanceSnapshot struct { + ObservedAt time.Time `json:"observed_at"` + Remaining *float64 `json:"remaining,omitempty"` + Used *float64 `json:"used,omitempty"` + Capacity *float64 `json:"capacity,omitempty"` + Packages []CreditPackage `json:"packages"` + Enterprise bool `json:"enterprise"` +} + +// observedNumber preserves absence and explicit zero. It never derives a +// missing field from another field and never treats malformed values as zero. +func observedNumber(fields map[string]json.RawMessage, keys ...string) (*float64, error) { + for _, key := range keys { + raw, exists := fields[key] + if !exists || string(raw) == "null" || string(raw) == `""` { + continue + } + var value float64 + if len(raw) > 0 && raw[0] == '"' { + var text string + if err := json.Unmarshal(raw, &text); err != nil { + return nil, fmt.Errorf("workbuddy: malformed numeric field %s", key) + } + var err error + value, err = strconv.ParseFloat(text, 64) + if err != nil { + return nil, fmt.Errorf("workbuddy: malformed numeric field %s", key) + } + } else if err := json.Unmarshal(raw, &value); err != nil { + return nil, fmt.Errorf("workbuddy: malformed numeric field %s", key) + } + if !validCredit(value) { + return nil, fmt.Errorf("workbuddy: invalid credit field %s", key) + } + return &value, nil + } + return nil, nil +} +func decodeCreditPackage(raw json.RawMessage) (CreditPackage, error) { + var fields map[string]json.RawMessage + var p CreditPackage + if err := json.Unmarshal(raw, &fields); err != nil || fields == nil { + return p, fmt.Errorf("workbuddy: invalid credit package") + } + var err error + p.Remaining, err = observedNumber(fields, "CycleCapacityRemainPrecise", "CycleCapacityRemain", "CapacityRemainPrecise", "CapacityRemain") + if err != nil { + return p, err + } + p.Used, err = observedNumber(fields, "CycleCapacityUsedPrecise", "CycleCapacityUsed", "CapacityUsedPrecise", "CapacityUsed") + if err != nil { + return p, err + } + p.Capacity, err = observedNumber(fields, "CycleCapacitySizePrecise", "CycleCapacitySize", "CapacitySizePrecise", "CapacitySize") + if err != nil { + return p, err + } + for key, target := range map[string]*string{"AccountId": &p.ID, "PackageCode": &p.Code, "CycleStartTime": &p.CycleStart, "CycleEndTime": &p.CycleEnd} { + if value, ok := fields[key]; ok && string(value) != "null" { + if err = json.Unmarshal(value, target); err != nil { + return p, fmt.Errorf("workbuddy: invalid package metadata %s", key) + } + } + } + return p, nil +} + +type resourcePage struct { + Accounts *[]json.RawMessage `json:"Accounts"` + TotalCount *int `json:"TotalCount"` +} + +func decodeResourcePage(raw json.RawMessage, nested bool) (resourcePage, error) { + var result resourcePage + if nested { + var data struct { + Response struct { + Data json.RawMessage `json:"Data"` + } `json:"Response"` + } + if err := json.Unmarshal(raw, &data); err != nil { + return result, fmt.Errorf("workbuddy: invalid resource response") + } + raw = data.Response.Data + } + if err := json.Unmarshal(raw, &result); err != nil || result.Accounts == nil || result.TotalCount != nil && *result.TotalCount < 0 { + return result, fmt.Errorf("workbuddy: missing or invalid resource accounts") + } + return result, nil +} + +func sumObserved(packages []CreditPackage, selectValue func(CreditPackage) *float64) *float64 { + total := 0.0 + for _, p := range packages { + v := selectValue(p) + if v == nil { + return nil + } + total += *v + if !validCredit(total) { + return nil + } + } + return &total +} + +// FetchBalance follows every reported page before publishing a total. Partial +// reads and duplicate package IDs are errors, not a lower balance observation. +func (a *ControlClient) FetchBalance(ctx context.Context) (BalanceSnapshot, error) { + var snapshot BalanceSnapshot + c, err := a.Load() + if err != nil { + return snapshot, err + } + now := time.Now + if a.Now != nil { + now = a.Now + } + var last ControlResult + if c.EnterpriseID != "" { + last, err = a.Call(ctx, ControlRequest{Operation: ControlEnterpriseUsage}) + if err != nil { + return snapshot, err + } + var fields map[string]json.RawMessage + if err = json.Unmarshal(last.Data, &fields); err != nil || fields == nil { + return snapshot, fmt.Errorf("workbuddy: invalid enterprise balance") + } + p := CreditPackage{Code: "enterprise"} + p.Used, err = observedNumber(fields, "credit") + if err != nil { + return snapshot, err + } + p.Capacity, err = observedNumber(fields, "limitNum") + if err != nil { + return snapshot, err + } + // The source exposes usage and a cycle limit, not an observed remainder. + // Leave Remaining unknown instead of fabricating a reported balance. + for key, target := range map[string]*string{"cycleStartTime": &p.CycleStart, "cycleEndTime": &p.CycleEnd} { + if raw, ok := fields[key]; ok && string(raw) != "null" { + if err = json.Unmarshal(raw, target); err != nil { + return snapshot, fmt.Errorf("workbuddy: invalid enterprise cycle") + } + } + } + snapshot.Packages = []CreditPackage{p} + snapshot.Enterprise = true + } else { + seen := make(map[string]bool) + total := -1 + complete := false + for page := 1; page <= 100; page++ { + last, err = a.Call(ctx, ControlRequest{Operation: ControlResources, Page: page}) + if err != nil { + return snapshot, err + } + data, errPage := decodeResourcePage(last.Data, true) + if errPage != nil { + return snapshot, errPage + } + if data.TotalCount != nil { + if total >= 0 && total != *data.TotalCount { + return snapshot, fmt.Errorf("workbuddy: balance pagination changed during observation") + } + total = *data.TotalCount + } + for _, raw := range *data.Accounts { + p, errPackage := decodeCreditPackage(raw) + if errPackage != nil { + return snapshot, errPackage + } + if p.ID != "" { + if seen[p.ID] { + return snapshot, fmt.Errorf("workbuddy: duplicate package in balance pages") + } + seen[p.ID] = true + } + snapshot.Packages = append(snapshot.Packages, p) + } + count := len(snapshot.Packages) + if total >= 0 && count > total { + return snapshot, fmt.Errorf("workbuddy: balance exceeds reported package count") + } + if total >= 0 && count == total || total < 0 && len(*data.Accounts) < 100 { + complete = true + break + } + if len(*data.Accounts) == 0 { + return snapshot, fmt.Errorf("workbuddy: incomplete balance pages") + } + } + if !complete { + return snapshot, fmt.Errorf("workbuddy: balance page limit exceeded") + } + } + snapshot.ObservedAt = now() + snapshot.Remaining = sumObserved(snapshot.Packages, func(p CreditPackage) *float64 { return p.Remaining }) + snapshot.Used = sumObserved(snapshot.Packages, func(p CreditPackage) *float64 { return p.Used }) + snapshot.Capacity = sumObserved(snapshot.Packages, func(p CreditPackage) *float64 { return p.Capacity }) + if snapshot.Remaining != nil { + if err = a.Governor.ObserveBalance(last.Credential, last.Attempt, *snapshot.Remaining); err != nil { + return snapshot, err + } + } else if err = a.Governor.ForgetBalance(last.Credential, last.Attempt); err != nil { + return snapshot, err + } + return snapshot, nil +} diff --git a/internal/workbuddy/balance_test.go b/internal/workbuddy/balance_test.go new file mode 100644 index 00000000000..09c1b4ee358 --- /dev/null +++ b/internal/workbuddy/balance_test.go @@ -0,0 +1,91 @@ +package workbuddy + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "sync/atomic" + "testing" +) + +func TestWorkBuddyBalancePreservesUnknownZeroAndPrecision(t *testing.T) { + p, err := decodeCreditPackage(json.RawMessage(`{"CapacitySize":"100","CapacityUsedPrecise":"0","CapacityRemainPrecise":"0.125","CapacityRemain":0}`)) + if err != nil || p.Capacity == nil || *p.Capacity != 100 || p.Used == nil || *p.Used != 0 || p.Remaining == nil || *p.Remaining != .125 { + t.Fatalf("package=%+v err=%v", p, err) + } + p, err = decodeCreditPackage(json.RawMessage(`{"CapacitySize":100,"CapacityUsed":20}`)) + if err != nil || p.Remaining != nil { + t.Fatal("missing remainder was derived") + } + for _, raw := range []string{`{"CapacityRemain":"NaN"}`, `{"CapacityRemain":-1}`, `{"CapacityRemain":true}`} { + if _, err = decodeCreditPackage(json.RawMessage(raw)); err == nil { + t.Fatal("invalid number accepted") + } + } +} +func TestWorkBuddyBalanceAggregatesAllPagesBeforeCreditRecovery(t *testing.T) { + var calls atomic.Int32 + a, c, _ := controlFixture(t, func(w http.ResponseWriter, r *http.Request) { + calls.Add(1) + var request struct{ PageNumber int } + if err := json.NewDecoder(r.Body).Decode(&request); err != nil { + t.Error(err) + } + fmt.Fprintf(w, `{"code":0,"data":{"Response":{"Data":{"TotalCount":2,"Accounts":[{"AccountId":"p%d","CapacityRemainPrecise":"12.5","CapacityUsed":2,"CapacitySize":14.5}]}}}}`, request.PageNumber) + }) + if _, err := a.Governor.RecordFailure(*c, "chat", Fault{Kind: FaultBalance, Status: 402}, FailureOptions{}); err != nil { + t.Fatal(err) + } + snapshot, err := a.FetchBalance(context.Background()) + if err != nil || snapshot.Remaining == nil || *snapshot.Remaining != 25 || len(snapshot.Packages) != 2 || calls.Load() != 2 { + t.Fatalf("snapshot=%+v err=%v", snapshot, err) + } + if gate, _ := a.Governor.Check(*c, "chat", false, false); !gate.Allowed() { + t.Fatal("positive balance did not clear credit parking") + } +} +func TestWorkBuddyBalanceMissingFieldClearsOldObservedRemainder(t *testing.T) { + a, c, _ := controlFixture(t, func(w http.ResponseWriter, r *http.Request) { + fmt.Fprint(w, `{"code":0,"data":{"Response":{"Data":{"Accounts":[{"CapacityUsed":2,"CapacitySize":10}]}}}}`) + }) + _, attempt := a.Governor.Check(*c, "balance", false, true) + if err := a.Governor.ObserveBalance(*c, attempt, 50); err != nil { + t.Fatal(err) + } + if err := a.Governor.UpdateLimits(Limits{ReserveCredits: 5}, nil); err != nil { + t.Fatal(err) + } + snapshot, err := a.FetchBalance(context.Background()) + if err != nil || snapshot.Remaining != nil { + t.Fatalf("snapshot=%+v err=%v", snapshot, err) + } + if gate, _ := a.Governor.Check(*c, "chat", false, false); gate.Reason != "balance_unknown" { + t.Fatal("old balance carried forward") + } +} +func TestWorkBuddyBalanceRejectsIncompleteAndDuplicatePages(t *testing.T) { + for _, accounts := range []string{`[]`, `[{"AccountId":"duplicate","CapacityRemain":10}]`} { + t.Run(accounts, func(t *testing.T) { + a, _, _ := controlFixture(t, func(w http.ResponseWriter, r *http.Request) { + fmt.Fprintf(w, `{"code":0,"data":{"Response":{"Data":{"TotalCount":2,"Accounts":%s}}}}`, accounts) + }) + if _, err := a.FetchBalance(context.Background()); err == nil { + t.Fatal("partial total accepted") + } + }) + } +} +func TestWorkBuddyEnterpriseBalanceDoesNotInventRemainder(t *testing.T) { + a, c, _ := controlFixture(t, func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/billing/meter/get-enterprise-user-usage" { + t.Error("personal endpoint used for enterprise") + } + fmt.Fprint(w, `{"code":0,"data":{"credit":"30.5","limitNum":100,"cycleStartTime":"2026-10-01"}}`) + }) + c.EnterpriseID = "fixture-enterprise" + result, err := a.FetchBalance(context.Background()) + if err != nil || !result.Enterprise || result.Remaining != nil || result.Used == nil || *result.Used != 30.5 { + t.Fatalf("result=%+v err=%v", result, err) + } +} diff --git a/internal/workbuddy/control.go b/internal/workbuddy/control.go new file mode 100644 index 00000000000..2c78ae7ad98 --- /dev/null +++ b/internal/workbuddy/control.go @@ -0,0 +1,323 @@ +package workbuddy + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "regexp" + "sync" + "time" +) + +// ControlOperation is a closed set: background work cannot use this client as +// an arbitrary authenticated HTTP tunnel. +type ControlOperation string + +const ( + ControlModels ControlOperation = "models" + ControlResources ControlOperation = "resources" + ControlResourceSummary ControlOperation = "resource_summary" + ControlFreePackages ControlOperation = "free_packages" + ControlPaidPackages ControlOperation = "paid_packages" + ControlEnterpriseUsage ControlOperation = "enterprise_usage" + ControlCheckin ControlOperation = "checkin" + ControlCheckinStatus ControlOperation = "checkin_status" + ControlGrowthTasks ControlOperation = "growth_tasks" + ControlGrowthAccept ControlOperation = "growth_accept" + ControlGrowthClaim ControlOperation = "growth_claim" +) + +type ControlRequest struct { + Operation ControlOperation + Page int + PackageCodes []string + TaskCodes []string + TaskCode string +} +type ControlResult struct { + Data json.RawMessage + AlreadyDone bool + Attempt Attempt + Credential Credential +} +type controlSpec struct { + method, path string + billing, observeBalance bool + body any +} + +var taskCodePattern = regexp.MustCompile(`^[A-Za-z0-9_-]{1,128}$`) + +func (r ControlRequest) spec(c Credential, now time.Time) (controlSpec, error) { + s := controlSpec{method: http.MethodPost, body: struct{}{}} + page := r.Page + if page == 0 { + page = 1 + } + if page < 1 || page > 10000 { + return s, fmt.Errorf("workbuddy: invalid resource page") + } + switch r.Operation { + case ControlModels: + s.method, s.path, s.body = http.MethodGet, "/v3/config", nil + case ControlResources: + s.billing, s.observeBalance, s.path = true, true, "/v2/billing/meter/get-user-resource" + s.body = map[string]any{"PageNumber": page, "PageSize": 100, "ProductCode": "p_tcaca", "Status": []int{0, 3}, "PackageEndTimeRangeBegin": now.In(time.FixedZone("CST", 8*3600)).Format("2006-01-02 15:04:05"), "PackageEndTimeRangeEnd": "2036-01-01 00:00:00"} + case ControlResourceSummary: + s.billing, s.observeBalance, s.path = true, true, "/billing/meter/get-user-resource-summary" + case ControlFreePackages, ControlPaidPackages: + s.billing, s.observeBalance, s.path = true, true, "/billing/meter/get-user-resource-free-packages" + if len(r.PackageCodes) == 0 || len(r.PackageCodes) > 200 { + return s, fmt.Errorf("workbuddy: invalid package selection") + } + s.body = map[string]any{"PageNumber": page, "PageSize": 200, "PackageCodes": r.PackageCodes, "Status": []int{0}} + if r.Operation == ControlPaidPackages { + s.path = "/billing/meter/get-user-resource-paid-packages" + s.body = map[string]any{"PageNumber": page, "PageSize": 200, "PackageCodes": r.PackageCodes, "Status": []int{0, 3}, "NeedRenewInfo": true} + } + case ControlEnterpriseUsage: + if c.EnterpriseID == "" { + return s, fmt.Errorf("workbuddy: enterprise usage requires an enterprise credential") + } + s.billing, s.observeBalance, s.path = true, true, "/billing/meter/get-enterprise-user-usage" + case ControlCheckin, ControlCheckinStatus, ControlGrowthTasks, ControlGrowthAccept, ControlGrowthClaim: + if c.Realm != China || c.EnterpriseID != "" { + return s, fmt.Errorf("workbuddy: activity requires a personal CN credential") + } + switch r.Operation { + case ControlCheckin: + s.billing, s.path = true, "/v2/billing/meter/daily-checkin" + // Acquiring credits must work while the account is parked for empty credit. + s.observeBalance = true + case ControlCheckinStatus: + s.path, s.observeBalance = "/v2/billing/meter/checkin-activity-status", true + case ControlGrowthTasks: + s.method, s.path, s.body = http.MethodGet, "/v2/activity/growth/tasks", nil + case ControlGrowthAccept: + if len(r.TaskCodes) == 0 || len(r.TaskCodes) > 20 { + return s, fmt.Errorf("workbuddy: growth accept requires 1 to 20 task codes") + } + for _, code := range r.TaskCodes { + if !taskCodePattern.MatchString(code) { + return s, fmt.Errorf("workbuddy: invalid task code") + } + } + s.path, s.body = "/v2/activity/growth/tasks/accept", map[string]any{"task_codes": r.TaskCodes} + case ControlGrowthClaim: + if !taskCodePattern.MatchString(r.TaskCode) { + return s, fmt.Errorf("workbuddy: invalid task code") + } + s.path = "/activity/growth/tasks/" + r.TaskCode + "/claim" + } + default: + return s, fmt.Errorf("workbuddy: unsupported control operation") + } + return s, nil +} + +// ControlClient shares chat's transports, limiter and governor. Load must read +// the authoritative credential; Refresh must persist rotation before returning. +// No upstream deadline is introduced and no mutating request is replayed after +// a connection failure with an uncertain result. +type ControlClient struct { + Governor *Governor + Transports *Transports + Limiter *Limiter + Versions IdentityVersions + ProxyURL string + BaseURL string // Offline fixture seam; never loaded from credential metadata. + QueueWait time.Duration + Now func() time.Time + Load func() (Credential, error) + Refresh func(context.Context) (Credential, error) + MinimumInterval time.Duration + Wait func(context.Context, time.Duration) error + paceMu sync.Mutex + nextCall time.Time +} + +func (a *ControlClient) Call(ctx context.Context, request ControlRequest) (ControlResult, error) { + var result ControlResult + if a.Governor == nil || a.Transports == nil || a.Limiter == nil || a.Load == nil { + return result, fmt.Errorf("workbuddy: incomplete control client") + } + now := time.Now + if a.Now != nil { + now = a.Now + } + if a.MinimumInterval > 0 { + a.paceMu.Lock() + start := now() + if a.nextCall.After(start) { + start = a.nextCall + } + delay := start.Sub(now()) + a.nextCall = start.Add(a.MinimumInterval) + a.paceMu.Unlock() + wait := a.Wait + if wait == nil { + wait = waitBeforeConnect + } + if err := wait(ctx, delay); err != nil { + return result, err + } + } + c, err := a.Load() + if err != nil { + return result, err + } + if _, err = request.spec(c, now()); err != nil { + return result, err + } + wait := a.QueueWait + if wait == 0 { + wait = 30 * time.Second + } + release, err := a.Limiter.Acquire(ctx, "background:"+RiskAccountKey(c), wait) + if err != nil { + return result, err + } + defer release() + refreshed := false + for { + c, err = a.Load() + if err != nil { + return result, err + } + spec, errSpec := request.spec(c, now()) + if errSpec != nil { + return result, errSpec + } + scope := "control:" + string(request.Operation) + gate, attempt := a.Governor.Check(c, scope, false, spec.observeBalance) + if gate.Reason == "credential_expired" && !refreshed && a.Refresh != nil { + refreshed = true + if _, err = a.Refresh(ctx); err != nil { + return result, err + } + continue + } + if !gate.Allowed() { + return result, &DispatchError{Reason: gate.Reason, Stop: gate.Stop, Until: gate.Until, Status: 429} + } + // Proactive acquisition is allowed only after restriction/cooldown checks. + if c.Expiring(now()) && !refreshed && a.Refresh != nil { + refreshed = true + if _, err = a.Refresh(ctx); err != nil { + return result, err + } + continue + } + base, _, errBase := Endpoint(c.Realm, Desktop) + if errBase != nil { + return result, errBase + } + headers, errHeaders := BillingHeaders(c, a.Versions) + var order []string + if !spec.billing { + turn, errTurn := NewTurnID() + if errTurn != nil { + return result, errTurn + } + // These control endpoints belong to the desktop API regardless of the + // product selected for inference. + desktop := c + desktop.Product = Desktop + headers, order, errHeaders = IdentityHeaders(desktop, a.Versions, turn, turn, "", now()) + } else { + base = "https://www.workbuddy.ai" + if c.Realm == China { + base = "https://www.codebuddy.cn" + } + headers.Set("X-Client-Platform", "web") + } + if errHeaders != nil { + return result, errHeaders + } + if a.BaseURL != "" { + if err = validateBaseURL(a.BaseURL); err != nil { + return result, err + } + base = a.BaseURL + } + proxy := a.ProxyURL + if c.ProxyURL != "" { + proxy = c.ProxyURL + } + client, errClient := a.Transports.Client(c, base, proxy, order) + if errClient != nil { + return result, errClient + } + var body []byte + if spec.body != nil { + body, err = json.Marshal(spec.body) + if err != nil { + return result, err + } + } + req, errReq := http.NewRequestWithContext(ctx, spec.method, base+spec.path, bytes.NewReader(body)) + if errReq != nil { + return result, errReq + } + req.Header = headers + response, errDo := client.Do(req) + if errDo != nil { + return result, a.failure(c, scope, ClassifyTransport(errDo)) + } + raw, errRead := io.ReadAll(io.LimitReader(response.Body, 4<<20+1)) + errClose := response.Body.Close() + if errRead != nil { + return result, a.failure(c, scope, ClassifyTransport(errRead)) + } + if errClose != nil { + return result, a.failure(c, scope, ClassifyTransport(errClose)) + } + if len(raw) > 4<<20 { + return result, a.failure(c, scope, Fault{Kind: FaultUnknown, Status: 502}) + } + var envelope authEnvelope + errJSON := json.Unmarshal(raw, &envelope) + already := request.Operation == ControlCheckin && (response.StatusCode == 200 || response.StatusCode == 400) && errJSON == nil && envelope.Code != nil && *envelope.Code == 10001 + success := response.StatusCode >= 200 && response.StatusCode < 300 && errJSON == nil && envelope.Code != nil && *envelope.Code == 0 + if success || already { + return ControlResult{Data: envelope.Data, AlreadyDone: already, Attempt: attempt, Credential: c}, nil + } + status := response.StatusCode + if status >= 200 && status < 300 { + status = 400 + if envelope.Code != nil { + switch *envelope.Code { + case 401, 402, 403, 429: + status = *envelope.Code + case 11140: + status = 403 + case 6004, 14003: + status = 429 + } + } + if errJSON != nil || envelope.Code == nil { + return result, a.failure(c, scope, Fault{Kind: FaultUnknown, Status: 502}) + } + } + fault := ClassifyHTTP(status, raw) + if fault.Kind == FaultAuth && !refreshed && a.Refresh != nil { + refreshed = true + if _, err = a.Refresh(ctx); err == nil { + continue + } + } + return result, a.failure(c, scope, fault) + } +} +func (a *ControlClient) failure(c Credential, scope string, fault Fault) error { + gate, err := a.Governor.RecordFailure(c, scope, fault, FailureOptions{}) + if err != nil { + return err + } + failure := ErrorForFault(fault) + failure.Stop, failure.Until = gate.Stop, gate.Until + return failure +} diff --git a/internal/workbuddy/control_test.go b/internal/workbuddy/control_test.go new file mode 100644 index 00000000000..9921bcb1aa7 --- /dev/null +++ b/internal/workbuddy/control_test.go @@ -0,0 +1,175 @@ +package workbuddy + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "net/http" + "net/http/httptest" + "strings" + "sync/atomic" + "testing" + "time" +) + +func controlFixture(t *testing.T, handler http.HandlerFunc) (*ControlClient, *Credential, *time.Time) { + t.Helper() + server := httptest.NewServer(handler) + t.Cleanup(server.Close) + g, now, _ := governorFixture(t, Limits{}) + c := testCredential() + c.Realm = China + limiter, err := NewLimiter(4, 1) + if err != nil { + t.Fatal(err) + } + a := &ControlClient{Governor: g, Transports: NewTransports(32, nil), Limiter: limiter, Versions: DefaultVersions(), BaseURL: server.URL, Now: func() time.Time { return *now }, Load: func() (Credential, error) { return c, nil }} + return a, &c, now +} + +func TestWorkBuddyControlRestrictionBlocksRefreshAndAllOperations(t *testing.T) { + var calls, refreshes atomic.Int32 + a, c, now := controlFixture(t, func(w http.ResponseWriter, r *http.Request) { calls.Add(1); fmt.Fprint(w, `{"code":0,"data":{}}`) }) + a.Refresh = func(context.Context) (Credential, error) { refreshes.Add(1); return *c, nil } + c.ExpiresAt = now.Add(-time.Minute).Unix() + if _, err := a.Governor.RecordFailure(*c, "chat", Fault{Kind: FaultRestricted, Status: 403, Code: 11140}, FailureOptions{}); err != nil { + t.Fatal(err) + } + for _, op := range []ControlOperation{ControlModels, ControlResources, ControlCheckin, ControlGrowthTasks} { + _, err := a.Call(context.Background(), ControlRequest{Operation: op}) + var dispatch *DispatchError + if !errors.As(err, &dispatch) || dispatch.Reason != "account_restricted" { + t.Fatalf("%s: %v", op, err) + } + } + if calls.Load() != 0 || refreshes.Load() != 0 { + t.Fatal("restricted credential contacted upstream") + } +} + +func TestWorkBuddyControl429IsSharedAndBalanceCannotBypassIt(t *testing.T) { + var calls atomic.Int32 + a, _, _ := controlFixture(t, func(w http.ResponseWriter, r *http.Request) { + calls.Add(1) + w.WriteHeader(429) + fmt.Fprint(w, `{"code":14003,"msg":"too many requests"}`) + }) + for i := 0; i < 2; i++ { + if _, err := a.Call(context.Background(), ControlRequest{Operation: ControlResources}); err == nil { + t.Fatal("429 accepted") + } + } + if calls.Load() != 1 { + t.Fatal("balance bypassed cooldown") + } +} + +func TestWorkBuddyControlRefreshOnceAndPersistBeforeReplay(t *testing.T) { + var calls atomic.Int32 + var saved atomic.Bool + a, c, _ := controlFixture(t, func(w http.ResponseWriter, r *http.Request) { + calls.Add(1) + if r.Header.Get("Authorization") != "Bearer rotated" { + w.WriteHeader(401) + fmt.Fprint(w, `{"code":401}`) + return + } + if !saved.Load() { + t.Error("rotation replayed before persistence") + } + if r.URL.Path != "/v2/billing/meter/daily-checkin" { + t.Error("wrong path") + } + fmt.Fprint(w, `{"code":0,"data":{"credit":25}}`) + }) + a.Refresh = func(context.Context) (Credential, error) { + c.AccessToken = "rotated" + saved.Store(true) + return *c, nil + } + result, err := a.Call(context.Background(), ControlRequest{Operation: ControlCheckin}) + if err != nil || calls.Load() != 2 || result.AlreadyDone || string(result.Data) != `{"credit":25}` { + t.Fatalf("result=%+v calls=%d err=%v", result, calls.Load(), err) + } +} + +func TestWorkBuddyControlCheckinAlreadyDoneHasNoInventedReward(t *testing.T) { + a, _, _ := controlFixture(t, func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(400) + fmt.Fprint(w, `{"code":10001,"data":null}`) + }) + result, err := a.Call(context.Background(), ControlRequest{Operation: ControlCheckin}) + if err != nil || !result.AlreadyDone || string(result.Data) != "null" { + t.Fatalf("result=%+v err=%v", result, err) + } +} + +func TestWorkBuddyControlUncertainMutationIsNeverReplayed(t *testing.T) { + var calls atomic.Int32 + a, _, _ := controlFixture(t, func(w http.ResponseWriter, r *http.Request) { + calls.Add(1) + conn, _, err := w.(http.Hijacker).Hijack() + if err != nil { + t.Error(err) + return + } + if err = conn.Close(); err != nil { + t.Error(err) + } + }) + if _, err := a.Call(context.Background(), ControlRequest{Operation: ControlGrowthClaim, TaskCode: "chat_5"}); err == nil { + t.Fatal("uncertain claim reported success") + } + if calls.Load() != 1 { + t.Fatal("uncertain mutation replayed") + } +} + +func TestWorkBuddyControlRejectsUnknownEnvelopesAndInjectedPaths(t *testing.T) { + for _, body := range []string{`{}`, `{"data":{}}`, `{"code":0`, strings.Repeat("x", 4<<20+1)} { + t.Run(fmt.Sprint(len(body)), func(t *testing.T) { + a, _, _ := controlFixture(t, func(w http.ResponseWriter, r *http.Request) { fmt.Fprint(w, body) }) + if _, err := a.Call(context.Background(), ControlRequest{Operation: ControlResources}); err == nil { + t.Fatal("malformed envelope accepted") + } + }) + } + var calls atomic.Int32 + a, _, _ := controlFixture(t, func(w http.ResponseWriter, r *http.Request) { calls.Add(1) }) + for _, r := range []ControlRequest{{Operation: "https://other.test"}, {Operation: ControlGrowthClaim, TaskCode: "../../auth"}, {Operation: ControlGrowthAccept, TaskCodes: []string{"bad/code"}}} { + if _, err := a.Call(context.Background(), r); err == nil { + t.Fatal("invalid request accepted") + } + } + if calls.Load() != 0 { + t.Fatal("invalid request reached network") + } +} + +func TestWorkBuddyControlBodyAndContext(t *testing.T) { + a, _, _ := controlFixture(t, func(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodPost || r.Header.Get("X-Client-Platform") != "web" { + t.Error("wrong billing request") + } + var body struct { + PageNumber int + PageSize int + } + if err := json.NewDecoder(r.Body).Decode(&body); err != nil { + t.Error(err) + } + if body.PageNumber != 2 || body.PageSize != 100 { + t.Error("pagination lost") + } + fmt.Fprint(w, `{"code":0,"data":{"Response":{"Data":{"Accounts":[]}}}}`) + }) + if _, err := a.Call(context.Background(), ControlRequest{Operation: ControlResources, Page: 2}); err != nil { + t.Fatal(err) + } + c, _ := a.Load() + client, err := a.Transports.Client(c, a.BaseURL, "", nil) + if err != nil || client.Timeout != 0 { + t.Fatal("background upstream has a deadline") + } +} diff --git a/internal/workbuddy/discovery.go b/internal/workbuddy/discovery.go new file mode 100644 index 00000000000..8a9817d4050 --- /dev/null +++ b/internal/workbuddy/discovery.go @@ -0,0 +1,201 @@ +package workbuddy + +import ( + "context" + "encoding/json" + "fmt" + "sort" + "strings" +) + +// ParseRemoteCatalog accepts the endpoint data or the equivalent desktop cache +// object. An empty agent list is not evidence that an account has no models. +func ParseRemoteCatalog(raw json.RawMessage, realm Realm) ([]Model, error) { + var root map[string]json.RawMessage + if err := json.Unmarshal(raw, &root); err != nil || root == nil { + return nil, fmt.Errorf("workbuddy: invalid model catalog") + } + roots := []map[string]json.RawMessage{root} + if data, ok := root["data"]; ok { + var nested map[string]json.RawMessage + if err := json.Unmarshal(data, &nested); err != nil { + return nil, fmt.Errorf("workbuddy: invalid catalog data") + } + if nested != nil { + roots = append(roots, nested) + } + } + var cli, other [][]string + for _, r := range roots { + agentsRaw := r["agents"] + if len(agentsRaw) == 0 { + continue + } + var agents []struct { + Name string `json:"name"` + Models []string `json:"models"` + } + if agentsRaw[0] == '[' { + // Decode entries independently: non-chat agents may have a different + // model configuration shape and cannot invalidate a valid chat list. + var entries []json.RawMessage + if err := json.Unmarshal(agentsRaw, &entries); err != nil { + return nil, fmt.Errorf("workbuddy: invalid catalog agents") + } + for _, entry := range entries { + var item struct { + Name string `json:"name"` + Models []string `json:"models"` + } + if json.Unmarshal(entry, &item) == nil { + agents = append(agents, item) + } + } + } else { + var entries map[string]json.RawMessage + if err := json.Unmarshal(agentsRaw, &entries); err != nil { + return nil, fmt.Errorf("workbuddy: invalid catalog agents") + } + names := make([]string, 0, len(entries)) + for name := range entries { + names = append(names, name) + } + sort.Strings(names) + for _, name := range names { + var item struct { + Name string `json:"name"` + Models []string `json:"models"` + } + if json.Unmarshal(entries[name], &item) == nil { + item.Name = name + agents = append(agents, item) + } + } + } + for _, agent := range agents { + if len(agent.Models) == 0 { + continue + } + if strings.EqualFold(strings.TrimSpace(agent.Name), "cli") { + cli = append(cli, agent.Models) + } else { + other = append(other, agent.Models) + } + } + } + pool := cli + if len(pool) == 0 { + pool = other + } + var ids []string + for _, list := range pool { + if len(list) > len(ids) { + ids = list + } + } + if len(ids) == 0 { + return nil, fmt.Errorf("workbuddy: empty remote model catalog") + } + metadata := make(map[string]json.RawMessage) + var walk func(json.RawMessage, int) + walk = func(node json.RawMessage, depth int) { + if depth > 16 || len(node) == 0 { + return + } + switch node[0] { + case '{': + var object map[string]json.RawMessage + if json.Unmarshal(node, &object) != nil { + return + } + if models, ok := object["models"]; ok { + var entries []json.RawMessage + if json.Unmarshal(models, &entries) == nil { + for _, entry := range entries { + var item struct { + ID string `json:"id"` + } + if json.Unmarshal(entry, &item) == nil && validModelID(item.ID) { + if _, exists := metadata[item.ID]; !exists { + metadata[item.ID] = entry + } + } + } + } + } + names := make([]string, 0, len(object)) + for name := range object { + names = append(names, name) + } + sort.Strings(names) + for _, name := range names { + walk(object[name], depth+1) + } + case '[': + var entries []json.RawMessage + if json.Unmarshal(node, &entries) == nil { + for _, entry := range entries { + walk(entry, depth+1) + } + } + } + } + walk(raw, 0) + models := Catalog(realm) + byID := make(map[string]int, len(models)) + for i, m := range models { + byID[m.ID] = i + } + seen := make(map[string]bool) + valid := 0 + for _, id := range ids { + id = strings.TrimSpace(id) + if !validModelID(id) || seen[id] { + continue + } + seen[id] = true + valid++ + m := Model{ID: id, Advertise: true} + index, exists := byID[id] + if exists { + m = models[index] + } + if entry, ok := metadata[id]; ok { + // Decode onto the snapshot so absent metadata never deletes known limits. + if err := json.Unmarshal(entry, &m); err != nil { + return nil, fmt.Errorf("workbuddy: malformed model metadata") + } + m.ID = id + } + if strings.HasSuffix(id, "-sg") || strings.HasSuffix(id, "-x") { + sibling := strings.TrimSuffix(strings.TrimSuffix(id, "-sg"), "-x") + if _, ok := byID[sibling]; ok { + m.Advertise = false + } + } + if exists { + models[index] = m + } else { + byID[id] = len(models) + models = append(models, m) + } + } + if valid == 0 { + return nil, fmt.Errorf("workbuddy: no usable models in remote catalog") + } + return models, nil +} +func validModelID(id string) bool { + if len(id) == 0 || len(id) > 200 || strings.ContainsAny(id, "\r\n\t /\\") { + return false + } + lower := strings.ToLower(id) + return lower != "auto" && lower != "default" && !strings.HasPrefix(lower, "default-model") && !strings.Contains(lower, "completion") +} +func (a *ControlClient) DiscoverModels(ctx context.Context) ([]Model, error) { + result, err := a.Call(ctx, ControlRequest{Operation: ControlModels}) + if err != nil { + return nil, err + } + return ParseRemoteCatalog(result.Data, result.Credential.Realm) +} diff --git a/internal/workbuddy/discovery_test.go b/internal/workbuddy/discovery_test.go new file mode 100644 index 00000000000..327e239c8a4 --- /dev/null +++ b/internal/workbuddy/discovery_test.go @@ -0,0 +1,54 @@ +package workbuddy + +import ( + "encoding/json" + "testing" +) + +func TestWorkBuddyRemoteCatalogPrefersCLIAndKeepsSnapshot(t *testing.T) { + raw := json.RawMessage(`{"data":{"agents":{"desktop":{"models":["other-a","other-b"]},"cli":{"models":["fixture-new","auto","fixture-new"]}},"models":[{"id":"fixture-new","maxInputTokens":1234,"supportsToolCall":true}]}}`) + models, err := ParseRemoteCatalog(raw, International) + if err != nil { + t.Fatal(err) + } + found, static := false, false + for _, m := range models { + switch m.ID { + case "fixture-new": + if found || m.Input != 1234 || !m.Tools || !m.Advertise { + t.Fatal("bad remote model") + } + found = true + case "deepseek-v4.1-flash": + static = true + case "other-a", "other-b", "auto": + t.Fatal("non-CLI or virtual model advertised") + } + } + if !found || !static { + t.Fatal("remote merge lost models") + } +} +func TestWorkBuddyRemoteCatalogEmptyCannotDeleteFallback(t *testing.T) { + for _, raw := range []string{`{}`, `{"agents":[]}`, `{"agents":{"cli":{"models":[]}}}`, `{"agents":[{"name":"cli","models":["auto"]}]}`} { + if _, err := ParseRemoteCatalog(json.RawMessage(raw), China); err == nil { + t.Fatal("empty catalog accepted") + } + } + models, err := ParseRemoteCatalog(json.RawMessage(`{"agents":[{"name":"other","models":["first"]},{"name":"long","models":["second","third"]}]}`), China) + if err != nil { + t.Fatal(err) + } + found := 0 + for _, m := range models { + if m.ID == "first" { + t.Fatal("shorter list chosen") + } + if m.ID == "second" || m.ID == "third" { + found++ + } + } + if found != 2 { + t.Fatal("longest fallback list missing") + } +} diff --git a/internal/workbuddy/governor.go b/internal/workbuddy/governor.go index dd22a78b6b3..1a2552e6245 100644 --- a/internal/workbuddy/governor.go +++ b/internal/workbuddy/governor.go @@ -52,6 +52,7 @@ type accountRisk struct { CreditsSpent float64 `json:"credits_spent"` ModelTokens map[string]int64 `json:"model_tokens"` CreditsRemain *float64 `json:"credits_remain,omitempty"` + CreditsDebited float64 `json:"credits_debited,omitempty"` BalanceStartedAt time.Time `json:"balance_started_at,omitempty"` } @@ -148,7 +149,7 @@ func OpenGovernor(path string, options GovernorOptions) (*Governor, error) { if state == nil || state.Models == nil || state.ModelTokens == nil { return nil, fmt.Errorf("workbuddy: incomplete account risk state") } - if state.Tokens < 0 || !validCredit(state.CreditsSpent) || (state.CreditsRemain != nil && !validCredit(*state.CreditsRemain)) { + if state.Tokens < 0 || !validCredit(state.CreditsSpent) || !validCredit(state.CreditsDebited) || (state.CreditsRemain != nil && !validCredit(*state.CreditsRemain)) { return nil, fmt.Errorf("workbuddy: invalid persisted usage") } for _, tokens := range state.ModelTokens { @@ -246,10 +247,6 @@ func (g *Governor) Check(c Credential, model string, paid, balanceOnly bool) (Ga gate.Reason = "credential_disabled" return gate, attempt } - if c.ExpiresAt > 0 && c.ExpiresAt <= now.Unix() { - gate.Reason = "credential_expired" - return gate, attempt - } if region := g.state.Realms[c.Realm]; region != nil && now.Before(region.HoldUntil) { gate.Reason, gate.Until, gate.Stop = "region_restricted", region.HoldUntil, true return gate, attempt @@ -279,6 +276,10 @@ func (g *Governor) Check(c Credential, model string, paid, balanceOnly bool) (Ga return gate, attempt } } + if c.ExpiresAt > 0 && c.ExpiresAt <= now.Unix() { + gate.Reason = "credential_expired" + return gate, attempt + } if balanceOnly { return gate, attempt } @@ -291,7 +292,7 @@ func (g *Governor) Check(c Credential, model string, paid, balanceOnly bool) (Ga gate.Reason = "balance_unknown" return gate, attempt } - if *state.CreditsRemain <= limits.ReserveCredits { + if *state.CreditsRemain-state.CreditsDebited <= limits.ReserveCredits { gate.Reason = "credit_reserve" return gate, attempt } @@ -469,13 +470,14 @@ func (g *Governor) Complete(c Credential, attempt Attempt, tokens int64, credits } state := g.accountLocked(c) if tokens > math.MaxInt64-state.Tokens || tokens > math.MaxInt64-state.ModelTokens[attempt.Model] || - (credits != nil && !validCredit(state.CreditsSpent+*credits)) { + (credits != nil && (!validCredit(state.CreditsSpent+*credits) || !validCredit(state.CreditsDebited+*credits))) { return fmt.Errorf("workbuddy: usage counter overflow") } state.Tokens += tokens state.ModelTokens[attempt.Model] += tokens if credits != nil { state.CreditsSpent += *credits + state.CreditsDebited += *credits } if success { if mark, ok := g.state.Entitlements[c.Realm][attempt.Model]; ok && !attempt.StartedAt.Before(mark.CheckedAt) { @@ -488,8 +490,15 @@ func (g *Governor) Complete(c Credential, attempt Attempt, tokens int64, credits state.Failures, state.UnknownFailures, state.SoftStreak, state.UnscopedStreak = 0, 0, 0, 0 if model := state.Models[attempt.Model]; model != nil { model.CooldownUntil = time.Time{} + if !g.options.Now().Before(model.CapUntil) { + model.CapUntil = time.Time{} + } } state.Version++ + } else if credits != nil && *credits > 0 { + // An older balance request must not overwrite a subsequently observed + // charge, including one from a truncated response. + state.Version++ } return g.persistLocked() } @@ -505,6 +514,7 @@ func (g *Governor) ObserveBalance(c Credential, attempt Attempt, remain float64) return nil } state.CreditsRemain = &remain + state.CreditsDebited = 0 state.BalanceStartedAt = attempt.StartedAt if remain > 0 && attempt.Version == state.Version { state.BalanceUntil = time.Time{} @@ -512,6 +522,24 @@ func (g *Governor) ObserveBalance(c Credential, attempt Attempt, remain float64) return g.persistLocked() } +// ForgetBalance records an incomplete observation without carrying a previous +// balance forward. It does not clear credit parking or any risk restriction. +func (g *Governor) ForgetBalance(c Credential, attempt Attempt) error { + g.mu.Lock() + defer g.mu.Unlock() + if attempt.Account != RiskAccountKey(c) { + return fmt.Errorf("workbuddy: invalid balance observation") + } + state := g.accountLocked(c) + if attempt.Version != state.Version || attempt.StartedAt.Before(state.BalanceStartedAt) { + return nil + } + state.CreditsRemain = nil + state.CreditsDebited = 0 + state.BalanceStartedAt = attempt.StartedAt + return g.persistLocked() +} + func validCredit(value float64) bool { return value >= 0 && !math.IsNaN(value) && !math.IsInf(value, 0) } diff --git a/internal/workbuddy/governor_test.go b/internal/workbuddy/governor_test.go index 79aa90a8c4d..20ea514166f 100644 --- a/internal/workbuddy/governor_test.go +++ b/internal/workbuddy/governor_test.go @@ -322,3 +322,31 @@ func TestGovernorModelConfigurationFailuresDoNotPoisonAccount(t *testing.T) { t.Fatalf("configuration failure poisoned account: %+v", gate) } } + +func TestWorkBuddyObservedCreditChargesProtectReserveBetweenBalanceReads(t *testing.T) { + g, _, _ := governorFixture(t, Limits{ReserveCredits: 5}) + c := testCredential() + _, balance := g.Check(c, "balance", false, true) + if err := g.ObserveBalance(c, balance, 20); err != nil { + t.Fatal(err) + } + _, attempt := g.Check(c, "paid", true, false) + _, staleBalance := g.Check(c, "balance", false, true) + credit := 16.0 + if err := g.Complete(c, attempt, 1, &credit, false); err != nil { + t.Fatal(err) + } + if err := g.ObserveBalance(c, staleBalance, 20); err != nil { + t.Fatal(err) + } + if gate, _ := g.Check(c, "paid", true, false); gate.Reason != "credit_reserve" { + t.Fatal("observed spend did not protect reserve") + } + _, fresh := g.Check(c, "balance", false, true) + if err := g.ObserveBalance(c, fresh, 30); err != nil { + t.Fatal(err) + } + if gate, _ := g.Check(c, "paid", true, false); !gate.Allowed() { + t.Fatal("fresh balance did not replace local debits") + } +} diff --git a/internal/workbuddy/growth.go b/internal/workbuddy/growth.go new file mode 100644 index 00000000000..218def9392a --- /dev/null +++ b/internal/workbuddy/growth.go @@ -0,0 +1,142 @@ +package workbuddy + +import ( + "context" + "encoding/json" + "fmt" + "time" +) + +type GrowthTask struct { + Code string `json:"task_code"` + Status string `json:"accept_status"` + Type string `json:"task_type"` + Progress struct { + Current *int64 `json:"current"` + Target *int64 `json:"target"` + } `json:"progress"` +} + +func (a *ControlClient) GrowthTasks(ctx context.Context) ([]GrowthTask, error) { + result, err := a.Call(ctx, ControlRequest{Operation: ControlGrowthTasks}) + if err != nil { + return nil, err + } + var data struct { + Tasks *[]GrowthTask `json:"tasks"` + } + if err = json.Unmarshal(result.Data, &data); err != nil || data.Tasks == nil { + return nil, fmt.Errorf("workbuddy: missing growth task list") + } + seen := make(map[string]bool) + for _, task := range *data.Tasks { + if !taskCodePattern.MatchString(task.Code) || seen[task.Code] { + return nil, fmt.Errorf("workbuddy: invalid or duplicate growth task") + } + seen[task.Code] = true + switch task.Status { + case "not_accepted", "accepted", "completed", "claimed": + default: + return nil, fmt.Errorf("workbuddy: unknown growth task status") + } + if task.Progress.Current != nil && *task.Progress.Current < 0 || task.Progress.Target != nil && *task.Progress.Target < 0 { + return nil, fmt.Errorf("workbuddy: invalid growth progress") + } + } + return *data.Tasks, nil +} +func (a *ControlClient) AcceptGrowthTasks(ctx context.Context, codes []string) error { + for len(codes) > 0 { + size := min(20, len(codes)) + part := codes[:size] + result, err := a.Call(ctx, ControlRequest{Operation: ControlGrowthAccept, TaskCodes: part}) + if err != nil { + return err + } + var data struct { + Results []struct { + Code string `json:"task_code"` + Status string `json:"status"` + } `json:"results"` + } + if err = json.Unmarshal(result.Data, &data); err != nil { + return fmt.Errorf("workbuddy: invalid growth acceptance") + } + accepted := make(map[string]bool) + for _, item := range data.Results { + if item.Status == "accepted" || item.Status == "already_accepted" { + accepted[item.Code] = true + } + } + for _, code := range part { + if !accepted[code] { + return fmt.Errorf("workbuddy: growth acceptance incomplete") + } + } + codes = codes[size:] + } + return nil +} + +// ReconcileGrowth accepts eligible tasks and claims only completion observed +// from the vendor. It never fabricates report events or pretends a desktop, +// subscription, donation or other real-world action has taken place. +func (a *ControlClient) ReconcileGrowth(ctx context.Context, ledger *JobLedger) (int, error) { + tasks, err := a.GrowthTasks(ctx) + if err != nil { + return 0, err + } + var accept []string + for _, task := range tasks { + switch task.Code { + case "wb_wechat_oa_subscribe_task", "Expert_Philanthropy", "RichMeow_Chat", "Library_read", "Buddy_App", "Buddy_App_QQ": + continue + } + if task.Status == "not_accepted" && task.Type != "auto" { + accept = append(accept, task.Code) + } + } + if len(accept) > 0 { + if err = a.AcceptGrowthTasks(ctx, accept); err != nil { + return 0, err + } + tasks, err = a.GrowthTasks(ctx) + if err != nil { + return 0, err + } + } + now := time.Now() + if a.Now != nil { + now = a.Now() + } + slot := now.In(time.FixedZone("CST", 8*3600)).Format("2006-01-02") + c, err := a.Load() + if err != nil { + return 0, err + } + pending := 0 + for _, task := range tasks { + if task.Status == "claimed" { + continue + } + complete := task.Status == "completed" || (task.Status == "accepted" || task.Type == "auto") && task.Progress.Current != nil && task.Progress.Target != nil && *task.Progress.Target > 0 && *task.Progress.Current >= *task.Progress.Target + if !complete { + pending++ + continue + } + receipt, _, errClaim := ledger.Run(ctx, c, "growth_claim:"+task.Code, slot, func(ctx context.Context) (JobReceipt, error) { + result, errCall := a.Call(ctx, ControlRequest{Operation: ControlGrowthClaim, TaskCode: task.Code}) + if errCall != nil { + return JobReceipt{}, errCall + } + return RewardReceipt(result) + }) + if errClaim != nil { + return pending, errClaim + } + if receipt.Status != "success" { + pending++ + } + } + return pending, nil +} diff --git a/internal/workbuddy/jobs.go b/internal/workbuddy/jobs.go new file mode 100644 index 00000000000..4b9986c17c9 --- /dev/null +++ b/internal/workbuddy/jobs.go @@ -0,0 +1,184 @@ +package workbuddy + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "os" + "strings" + "sync" + "time" +) + +type JobReceipt struct { + Account string `json:"account"` + Job string `json:"job"` + Slot string `json:"slot"` + Status string `json:"status"` + Attempts int `json:"attempts"` + StartedAt time.Time `json:"started_at"` + FinishedAt time.Time `json:"finished_at,omitempty"` + NextAttempt time.Time `json:"next_attempt"` + Reason string `json:"reason,omitempty"` + AlreadyDone bool `json:"already_done,omitempty"` + Credit *float64 `json:"credit,omitempty"` + Energy *float64 `json:"energy,omitempty"` + PendingTasks int `json:"pending_tasks,omitempty"` +} +type jobState struct { + Version int `json:"version"` + Receipts map[string]JobReceipt `json:"receipts"` +} +type JobLedger struct { + mu sync.Mutex + path string + now func() time.Time + state jobState + active map[string]bool + saveErr error +} + +func OpenJobLedger(path string, now func() time.Time) (*JobLedger, error) { + if now == nil { + now = time.Now + } + l := &JobLedger{path: path, now: now, state: jobState{Version: 1, Receipts: make(map[string]JobReceipt)}, active: make(map[string]bool)} + raw, err := os.ReadFile(path) + if errors.Is(err, os.ErrNotExist) { + return l, nil + } + if err != nil { + return nil, fmt.Errorf("workbuddy: read job receipts: %w", err) + } + if err = json.Unmarshal(raw, &l.state); err != nil { + return nil, fmt.Errorf("workbuddy: decode job receipts: %w", err) + } + if l.state.Version != 1 || l.state.Receipts == nil { + return nil, fmt.Errorf("workbuddy: unsupported job receipt version") + } + return l, nil +} +func (l *JobLedger) persistLocked() error { + l.saveErr = WriteJSON(l.path, l.state) + if l.saveErr != nil { + return fmt.Errorf("workbuddy: persist job receipt: %w", l.saveErr) + } + return nil +} + +// Run persists intent before the action. A process restart retains retry limits +// and success receipts; an uncertain outcome is never recorded as completion. +// run must reconcile upstream state before retrying non-idempotent actions. +func (l *JobLedger) Run(ctx context.Context, c Credential, job, slot string, run func(context.Context) (JobReceipt, error)) (JobReceipt, bool, error) { + account := RiskAccountKey(c) + key := account + ":" + job + ":" + slot + lane := account + ":" + job + l.mu.Lock() + receipt := l.state.Receipts[key] + if l.saveErr != nil { + err := l.saveErr + l.mu.Unlock() + return receipt, false, err + } + now := l.now() + maxAttempts, retryDelay := 6, 15*time.Minute + if strings.HasPrefix(job, "reset_probe:") { + maxAttempts, retryDelay = 20, 15*time.Second + } + if receipt.Status == "success" || receipt.Attempts >= maxAttempts || now.Before(receipt.NextAttempt) || l.active[lane] { + l.mu.Unlock() + return receipt, false, nil + } + // Keep seven days of receipts. Pending/uncertain outcomes remain available + // for reconciliation instead of silently becoming new work. + for k, r := range l.state.Receipts { + if r.Status == "success" && now.Sub(r.FinishedAt) > 7*24*time.Hour { + delete(l.state.Receipts, k) + } + } + if _, exists := l.state.Receipts[key]; !exists && len(l.state.Receipts) >= 10000 { + l.mu.Unlock() + return receipt, false, fmt.Errorf("workbuddy: job receipt limit reached") + } + receipt.Account, receipt.Job, receipt.Slot = account, job, slot + receipt.Status, receipt.Reason = "pending", "" + receipt.StartedAt, receipt.FinishedAt = now, time.Time{} + receipt.Attempts++ + receipt.NextAttempt = now.Add(retryDelay) + l.state.Receipts[key] = receipt + if err := l.persistLocked(); err != nil { + l.mu.Unlock() + return receipt, false, err + } + l.active[lane] = true + l.mu.Unlock() + outcome, err := run(ctx) + l.mu.Lock() + defer l.mu.Unlock() + delete(l.active, lane) + receipt.FinishedAt = l.now() + receipt.Status = "success" + receipt.AlreadyDone, receipt.Credit, receipt.Energy = outcome.AlreadyDone, outcome.Credit, outcome.Energy + receipt.PendingTasks = outcome.PendingTasks + if err != nil { + receipt.Status, receipt.Reason = "failed", "operation_failed" + if errors.Is(err, context.Canceled) { + receipt.Reason = "canceled" + } + var dispatch *DispatchError + if errors.As(err, &dispatch) { + receipt.Reason = dispatch.Reason + if dispatch.Until.After(receipt.NextAttempt) { + receipt.NextAttempt = dispatch.Until + } + } + } + l.state.Receipts[key] = receipt + if errSave := l.persistLocked(); errSave != nil { + return receipt, true, errSave + } + return receipt, true, err +} + +func RewardReceipt(result ControlResult) (JobReceipt, error) { + receipt := JobReceipt{AlreadyDone: result.AlreadyDone} + if len(result.Data) == 0 || string(result.Data) == "null" { + return receipt, nil + } + var fields map[string]json.RawMessage + if err := json.Unmarshal(result.Data, &fields); err != nil { + return receipt, fmt.Errorf("workbuddy: invalid activity receipt") + } + var err error + receipt.Credit, err = observedNumber(fields, "credit") + if err != nil { + return receipt, err + } + receipt.Energy, err = observedNumber(fields, "energy") + return receipt, err +} + +// JobSlot separates recurring observations from daily benefits. No timer or +// scheduler touches an account until the corresponding policy is enabled. +func JobSlot(now time.Time, job string) (string, bool) { + local := now.In(time.FixedZone("CST", 8*3600)) + switch job { + case "models", "balance", "quota": + return fmt.Sprintf("%s/%d", local.Format("2006-01-02"), local.Hour()/6), true + case "checkin", "growth": + if local.Hour() < 9 { + return "", false + } + if job == "growth" { + hour := 9 + if local.Hour() >= 21 { + hour = 21 + } + return fmt.Sprintf("%s/%d", local.Format("2006-01-02"), hour), true + } + return local.Format("2006-01-02"), true + default: + return "", false + } +} diff --git a/internal/workbuddy/jobs_test.go b/internal/workbuddy/jobs_test.go new file mode 100644 index 00000000000..1ce593cf75c --- /dev/null +++ b/internal/workbuddy/jobs_test.go @@ -0,0 +1,207 @@ +package workbuddy + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "net/http" + "os" + "path/filepath" + "strings" + "testing" + "time" +) + +func TestWorkBuddyJobIntentPersistsBeforeActionAndSurvivesRestart(t *testing.T) { + now := time.Date(2026, 10, 21, 1, 0, 0, 0, time.UTC) + path := filepath.Join(t.TempDir(), "jobs.json") + ledger, err := OpenJobLedger(path, func() time.Time { return now }) + if err != nil { + t.Fatal(err) + } + c := testCredential() + calls := 0 + run := func(context.Context) (JobReceipt, error) { + calls++ + raw, errRead := os.ReadFile(path) + if errRead != nil || !strings.Contains(string(raw), `"pending"`) { + t.Fatal("action ran without durable intent") + } + if strings.Contains(string(raw), c.UID) || strings.Contains(string(raw), c.AccessToken) { + t.Fatal("receipt leaked credential") + } + return JobReceipt{AlreadyDone: true}, nil + } + receipt, ran, err := ledger.Run(context.Background(), c, "checkin", "2026-10-21", run) + if err != nil || !ran || receipt.Status != "success" { + t.Fatalf("receipt=%+v err=%v", receipt, err) + } + restarted, err := OpenJobLedger(path, func() time.Time { return now }) + if err != nil { + t.Fatal(err) + } + if _, ran, err = restarted.Run(context.Background(), c, "checkin", "2026-10-21", run); err != nil || ran || calls != 1 { + t.Fatal("restart repeated completed work") + } +} +func TestWorkBuddyJobCooldownAndRetryBudget(t *testing.T) { + now := time.Date(2026, 10, 21, 1, 0, 0, 0, time.UTC) + ledger, err := OpenJobLedger(filepath.Join(t.TempDir(), "jobs.json"), func() time.Time { return now }) + if err != nil { + t.Fatal(err) + } + c := testCredential() + calls := 0 + run := func(context.Context) (JobReceipt, error) { + calls++ + return JobReceipt{}, &DispatchError{Reason: "region_restricted", Until: now.Add(time.Hour)} + } + for i := 0; i < 6; i++ { + receipt, ran, errRun := ledger.Run(context.Background(), c, "balance", "slot", run) + if errRun == nil || !ran || receipt.Status != "failed" || receipt.Reason != "region_restricted" { + t.Fatal("failed job reported success") + } + if _, ran, _ = ledger.Run(context.Background(), c, "balance", "slot", run); ran { + t.Fatal("job ignored cooldown") + } + now = now.Add(time.Hour) + } + if _, ran, _ := ledger.Run(context.Background(), c, "balance", "slot", run); ran || calls != 6 { + t.Fatal("retry budget exceeded") + } +} +func TestWorkBuddyJobCannotRunWhenIntentPersistenceFails(t *testing.T) { + path := filepath.Join(t.TempDir(), "receipts") + ledger, err := OpenJobLedger(path, nil) + if err != nil { + t.Fatal(err) + } + if err = os.Mkdir(path, 0700); err != nil { + t.Fatal(err) + } + ranAction := false + _, _, err = ledger.Run(context.Background(), testCredential(), "checkin", "day", func(context.Context) (JobReceipt, error) { ranAction = true; return JobReceipt{}, nil }) + if err == nil || ranAction { + t.Fatal("failed persistence allowed action") + } +} +func TestWorkBuddyJobSlotAndUnknownReward(t *testing.T) { + before := time.Date(2026, 10, 21, 0, 59, 0, 0, time.UTC) + if _, due := JobSlot(before, "checkin"); due { + t.Fatal("checkin before 9 CST") + } + morning, _ := JobSlot(before.Add(time.Minute), "growth") + night, _ := JobSlot(before.Add(13*time.Hour), "growth") + if morning == night { + t.Fatal("evening growth reconciliation missing") + } + receipt, err := RewardReceipt(ControlResult{Data: json.RawMessage(`{}`)}) + if err != nil || receipt.Credit != nil || receipt.Energy != nil { + t.Fatal("missing reward replaced with zero") + } +} +func TestWorkBuddyGrowthClaimsOnlyObservedCompletion(t *testing.T) { + var paths []string + a, c, now := controlFixture(t, func(w http.ResponseWriter, r *http.Request) { + paths = append(paths, r.URL.Path) + switch r.URL.Path { + case "/v2/activity/growth/tasks": + fmt.Fprint(w, `{"code":0,"data":{"tasks":[{"task_code":"chat_5","accept_status":"completed","task_type":"single"},{"task_code":"needs_real_action","accept_status":"accepted","task_type":"single","progress":{"current":0,"target":1}},{"task_code":"missing_progress","accept_status":"accepted","task_type":"single"}]}}`) + case "/activity/growth/tasks/chat_5/claim": + fmt.Fprint(w, `{"code":0,"data":{"credit":100}}`) + default: + t.Error("unexpected operation", r.URL.Path) + w.WriteHeader(500) + } + }) + path := filepath.Join(t.TempDir(), "jobs.json") + ledger, err := OpenJobLedger(path, func() time.Time { return *now }) + if err != nil { + t.Fatal(err) + } + pending, err := a.ReconcileGrowth(context.Background(), ledger) + if err != nil || pending != 2 || len(paths) != 2 { + t.Fatalf("pending=%d paths=%v err=%v", pending, paths, err) + } + // The same stale vendor list must not cause a second payout request. + pending, err = a.ReconcileGrowth(context.Background(), ledger) + if err != nil || pending != 2 || len(paths) != 3 { + t.Fatal("reconciliation replayed an already recorded claim") + } + raw, err := os.ReadFile(path) + if err != nil { + t.Fatal(err) + } + var persisted jobState + if err = json.Unmarshal(raw, &persisted); err != nil { + t.Fatal(err) + } + rewardObserved := false + for _, receipt := range persisted.Receipts { + if receipt.Job == "growth_claim:chat_5" && receipt.Credit != nil && *receipt.Credit == 100 { + rewardObserved = true + } + } + if !rewardObserved || strings.Contains(string(raw), c.AccessToken) { + t.Fatal("invalid reward receipt") + } +} +func TestWorkBuddyGrowthAcceptanceRequiresPerTaskAck(t *testing.T) { + a, _, _ := controlFixture(t, func(w http.ResponseWriter, r *http.Request) { + fmt.Fprint(w, `{"code":0,"data":{"results":[{"task_code":"one","status":"accepted"}]}}`) + }) + if err := a.AcceptGrowthTasks(context.Background(), []string{"one", "two"}); err == nil { + t.Fatal("missing acceptance acknowledged") + } + a, _, _ = controlFixture(t, func(w http.ResponseWriter, r *http.Request) { + fmt.Fprint(w, `{"code":0,"data":{"tasks":[{"task_code":"one"}]}}`) + }) + if _, err := a.GrowthTasks(context.Background()); err == nil { + t.Fatal("missing status inferred") + } +} +func TestWorkBuddyJobSameLaneDoesNotOverlap(t *testing.T) { + ledger, err := OpenJobLedger(filepath.Join(t.TempDir(), "jobs.json"), nil) + if err != nil { + t.Fatal(err) + } + started, finish := make(chan struct{}), make(chan struct{}) + result := make(chan error, 1) + go func() { + _, _, errRun := ledger.Run(context.Background(), testCredential(), "balance", "slot1", func(context.Context) (JobReceipt, error) { + close(started) + <-finish + return JobReceipt{}, errors.New("fixture error") + }) + result <- errRun + }() + <-started + _, ran, err := ledger.Run(context.Background(), testCredential(), "balance", "slot2", func(context.Context) (JobReceipt, error) { t.Error("overlapping lane"); return JobReceipt{}, nil }) + close(finish) + if err != nil || ran || <-result == nil { + t.Fatal("lane or failure result incorrect") + } +} + +func TestWorkBuddyResetProbeUsesShortRetryAndTwentyAttemptBudget(t *testing.T) { + now := time.Date(2026, 10, 21, 1, 0, 0, 0, time.UTC) + ledger, err := OpenJobLedger(filepath.Join(t.TempDir(), "jobs.json"), func() time.Time { return now }) + if err != nil { + t.Fatal(err) + } + calls := 0 + run := func(context.Context) (JobReceipt, error) { + calls++ + return JobReceipt{}, errors.New("fixture connection error") + } + for i := 0; i < 20; i++ { + if _, ran, errRun := ledger.Run(context.Background(), testCredential(), "reset_probe:hy3", "reset-window", run); !ran || errRun == nil { + t.Fatal("reset attempt skipped") + } + now = now.Add(15 * time.Second) + } + if _, ran, _ := ledger.Run(context.Background(), testCredential(), "reset_probe:hy3", "reset-window", run); ran || calls != 20 { + t.Fatal("reset retry budget exceeded") + } +} diff --git a/internal/workbuddy/probes.go b/internal/workbuddy/probes.go new file mode 100644 index 00000000000..ed17fa6b78b --- /dev/null +++ b/internal/workbuddy/probes.go @@ -0,0 +1,37 @@ +package workbuddy + +import ( + "sort" + "time" +) + +type ResetProbe struct { + Model string + ResetAt time.Time +} + +func (g *Governor) DueResetProbes(c Credential) []ResetProbe { + g.mu.Lock() + defer g.mu.Unlock() + now := g.options.Now() + var probes []ResetProbe + for model, state := range g.accountLocked(c).Models { + if IsFreeCapModel(c.Realm, model) && !state.CapUntil.IsZero() && !now.Before(state.CapUntil) { + probes = append(probes, ResetProbe{Model: model, ResetAt: state.CapUntil}) + } + } + sort.Slice(probes, func(i, j int) bool { + if probes[i].ResetAt.Equal(probes[j].ResetAt) { + return probes[i].Model < probes[j].Model + } + return probes[i].ResetAt.Before(probes[j].ResetAt) + }) + return probes +} +func FreeProbeModels(realm Realm) []string { + models := []string{"hy3", "hy4-preview-f"} + if realm == International { + models = append(models, "deepseek-v4.1-flash") + } + return models +} diff --git a/internal/workbuddy/stream.go b/internal/workbuddy/stream.go index ec6ce9aaec2..d9718705cbf 100644 --- a/internal/workbuddy/stream.go +++ b/internal/workbuddy/stream.go @@ -178,6 +178,18 @@ func ObservedTokens(raw json.RawMessage) (int64, bool) { return 0, false } +func ObservedCredits(raw json.RawMessage) *float64 { + var fields map[string]json.RawMessage + if json.Unmarshal(raw, &fields) != nil { + return nil + } + value, err := observedNumber(fields, "credit") + if err != nil { + return nil + } + return value +} + type collectedTool struct { ID string `json:"id"` Type string `json:"type"` diff --git a/sdk/cliproxy/service_lifecycle.go b/sdk/cliproxy/service_lifecycle.go index f08eb1b4cf3..3763806f35e 100644 --- a/sdk/cliproxy/service_lifecycle.go +++ b/sdk/cliproxy/service_lifecycle.go @@ -233,6 +233,7 @@ func (s *Service) Run(ctx context.Context) error { s.registerModelRefreshCallback() if !homeEnabled { go s.runAntigravityModelRefresh(ctx) + go s.runWorkBuddyMaintenance(ctx) } select { diff --git a/sdk/cliproxy/service_models.go b/sdk/cliproxy/service_models.go index 91919057dc6..e144a03788b 100644 --- a/sdk/cliproxy/service_models.go +++ b/sdk/cliproxy/service_models.go @@ -9,10 +9,12 @@ import ( "github.com/router-for-me/CLIProxyAPI/v8/internal/constant" "github.com/router-for-me/CLIProxyAPI/v8/internal/modelconfig" "github.com/router-for-me/CLIProxyAPI/v8/internal/registry" + "github.com/router-for-me/CLIProxyAPI/v8/internal/runtime/executor" "github.com/router-for-me/CLIProxyAPI/v8/internal/runtime/executor/helps" "github.com/router-for-me/CLIProxyAPI/v8/internal/workbuddy" coreauth "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/auth" "github.com/router-for-me/CLIProxyAPI/v8/sdk/config" + log "github.com/sirupsen/logrus" ) // registerModelsForAuth (re)binds provider models in the global registry using the core auth ID as client identifier. @@ -206,7 +208,19 @@ func (s *Service) registerModelsForAuthWithConfigAndCache(ctx context.Context, c models = append(models, ®istry.ModelInfo{ID: id, Object: "model", OwnedBy: "workbuddy", Type: "workbuddy", UserDefined: true}) } } else { - for _, model := range workbuddy.Catalog(credential.Realm) { + catalog := workbuddy.Catalog(credential.Realm) + if s.coreManager != nil { + if provider, ok := s.coreManager.Executor("workbuddy"); ok { + if native, okNative := provider.(*executor.WorkBuddyExecutor); okNative && native.UsesConfig(cfg) { + if observed, errCatalog := native.ModelCatalog(a); errCatalog == nil { + catalog = observed + } else { + log.WithField("provider", "workbuddy").Warn("cached model catalog unavailable; using static snapshot") + } + } + } + } + for _, model := range catalog { if !model.Advertise { continue } diff --git a/sdk/cliproxy/service_workbuddy.go b/sdk/cliproxy/service_workbuddy.go new file mode 100644 index 00000000000..d9d9c87d79f --- /dev/null +++ b/sdk/cliproxy/service_workbuddy.go @@ -0,0 +1,50 @@ +package cliproxy + +import ( + "context" + "time" + + "github.com/router-for-me/CLIProxyAPI/v8/internal/runtime/executor" + log "github.com/sirupsen/logrus" +) + +func (s *Service) runWorkBuddyMaintenance(ctx context.Context) { + timer := time.NewTimer(5 * time.Second) + defer timer.Stop() + for { + select { + case <-ctx.Done(): + return + case <-timer.C: + } + s.workBuddyMaintenanceCycle(ctx) + timer.Reset(15 * time.Second) + } +} + +func (s *Service) workBuddyMaintenanceCycle(ctx context.Context) { + cfg := s.currentConfig() + if cfg == nil || !cfg.WorkBuddy.Enabled || !cfg.WorkBuddy.Jobs.Enabled || s.coreManager == nil { + return + } + provider, ok := s.coreManager.Executor("workbuddy") + if !ok { + return + } + native, ok := provider.(*executor.WorkBuddyExecutor) + if !ok || !native.UsesConfig(cfg) { + return + } + changed, err := native.RunMaintenance(ctx) + if err != nil && ctx.Err() == nil { + log.WithField("provider", "workbuddy").Warn("maintenance cycle has incomplete tasks; inspect durable receipts") + } + if ctx.Err() != nil || !native.UsesConfig(s.currentConfig()) { + return + } + for _, id := range changed { + if current, exists := s.coreManager.GetByID(id); exists && !current.Disabled { + s.registerModelsForAuth(ctx, current) + } + } +} From a8e5212af06faa1961a35a96363aeb5207d804d0 Mon Sep 17 00:00:00 2001 From: LuneZ99 Date: Wed, 21 Oct 2026 03:09:32 +0800 Subject: [PATCH 07/10] ci(workbuddy): publish isolated candidate images for test rollout --- .github/workflows/workbuddy-image.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/workbuddy-image.yml b/.github/workflows/workbuddy-image.yml index f6cc71d3280..b169017412c 100644 --- a/.github/workflows/workbuddy-image.yml +++ b/.github/workflows/workbuddy-image.yml @@ -3,7 +3,7 @@ name: WorkBuddy isolated image on: workflow_dispatch: push: - branches: [main] + branches: [main, feat/workbuddy-proxy] paths: - 'internal/workbuddy/**' - 'internal/runtime/executor/workbuddy*' From 71767a80441fda1ebd1485945175cd99057ae452 Mon Sep 17 00:00:00 2001 From: LuneZ99 Date: Wed, 21 Oct 2026 03:20:16 +0800 Subject: [PATCH 08/10] fix(workbuddy): persist quota scan pacing across schedule boundaries --- docs/workbuddy-native.md | 6 +- .../runtime/executor/helps/workbuddy_jobs.go | 18 +++- internal/workbuddy/jobs.go | 74 ++++++++++++++++ internal/workbuddy/jobs_test.go | 86 +++++++++++++++++++ 4 files changed, 182 insertions(+), 2 deletions(-) diff --git a/docs/workbuddy-native.md b/docs/workbuddy-native.md index 056f68fc1fe..ff90e382bcc 100644 --- a/docs/workbuddy-native.md +++ b/docs/workbuddy-native.md @@ -103,7 +103,11 @@ Background operations now share chat's governor, connection pools and limiter: - Optional quota scans use the three known free-cap models as applicable to the realm. Reset probes wait for the recorded reset time; both execute through the native chat path with paid fallback and product switching disabled. Each cycle - admits at most 16 probes, with 400 ms spacing. + admits at most 16 probes, including at most six quota scans, with 400 ms spacing. + New quota pairs start after five seconds and are staggered by 15 seconds up to + one hour; the schedule survives restart. Successful scans remain at least six + hours apart even across a scheduling-slot boundary. Registered cap windows are + handled only by reset probes until recovery is observed. - Atomic job receipts precede actions. Completed slots survive restart; ordinary failures retry after 15 minutes, at most six times per slot. Reset failures retry after at least 15 seconds, at most 20 times per reset window, subject to diff --git a/internal/runtime/executor/helps/workbuddy_jobs.go b/internal/runtime/executor/helps/workbuddy_jobs.go index de73e5e9753..37e0c1269c1 100644 --- a/internal/runtime/executor/helps/workbuddy_jobs.go +++ b/internal/runtime/executor/helps/workbuddy_jobs.go @@ -78,6 +78,7 @@ func (r *WorkBuddyRuntime) Maintenance(ctx context.Context, cfg *config.Config, var failures []error policy := cfg.WorkBuddy.Jobs probeCount := 0 + scanCount := 0 for _, auth := range records { if err := ctx.Err(); err != nil { return changed, err @@ -172,12 +173,18 @@ func (r *WorkBuddyRuntime) Maintenance(ctx context.Context, cfg *config.Config, } if len(probe) > 0 && (policy.QuotaScan || policy.ResetProbes) { plans := r.Governor.DueResetProbes(c) + pendingReset := make(map[string]bool, len(plans)) + for _, plan := range plans { + pendingReset[plan.Model] = true + } if !policy.ResetProbes { plans = nil } if policy.QuotaScan { for _, model := range workbuddy.FreeProbeModels(c.Realm) { - plans = append(plans, workbuddy.ResetProbe{Model: model}) + if !pendingReset[model] && !r.Governor.Capped(c, model) { + plans = append(plans, workbuddy.ResetProbe{Model: model}) + } } } for _, plan := range plans { @@ -188,6 +195,12 @@ func (r *WorkBuddyRuntime) Maintenance(ctx context.Context, cfg *config.Config, if plan.ResetAt.IsZero() { job = "quota_scan:" + plan.Model slot, _ = workbuddy.JobSlot(time.Now(), "quota") + if errSchedule := r.Jobs.PrepareQuota(c, plan.Model, slot); errSchedule != nil { + return changed, errSchedule + } + if scanCount >= 6 { + continue + } } _, _, errProbe := r.Jobs.Run(ctx, c, job, slot, func(ctx context.Context) (workbuddy.JobReceipt, error) { if _, errLoad := load(); errLoad != nil { @@ -203,6 +216,9 @@ func (r *WorkBuddyRuntime) Maintenance(ctx context.Context, cfg *config.Config, } } probeCount++ + if plan.ResetAt.IsZero() { + scanCount++ + } return workbuddy.JobReceipt{}, probe[0](ctx, r.Current(auth), plan.Model) }) if errProbe != nil { diff --git a/internal/workbuddy/jobs.go b/internal/workbuddy/jobs.go index 4b9986c17c9..a8d6733bbcd 100644 --- a/internal/workbuddy/jobs.go +++ b/internal/workbuddy/jobs.go @@ -82,6 +82,44 @@ func (l *JobLedger) Run(ctx context.Context, c Credential, job, slot string, run return receipt, false, err } now := l.now() + if strings.HasPrefix(job, "quota_scan:") { + var latest JobReceipt + for previousKey, previous := range l.state.Receipts { + if previous.Account == account && previous.Job == job && previous.StartedAt.After(latest.StartedAt) { + latest = previous + } + if previous.Account == account && previous.Job == job && previous.Status == "scheduled" { + if now.Before(previous.NextAttempt) { + l.mu.Unlock() + return previous, false, nil + } + if previousKey != key { + delete(l.state.Receipts, previousKey) + } + } + if previous.Account == account && previous.Job == job && previous.Status == "success" && now.Before(previous.FinishedAt.Add(6*time.Hour)) { + l.mu.Unlock() + return previous, false, nil + } + } + if latest.Slot != slot && (latest.Status == "failed" || latest.Status == "pending") { + until := latest.NextAttempt + lastAttempt := latest.FinishedAt + if lastAttempt.IsZero() { + lastAttempt = latest.StartedAt + } + if latest.Attempts >= 6 && lastAttempt.Add(6*time.Hour).After(until) { + until = lastAttempt.Add(6 * time.Hour) + } + if now.Before(until) { + l.mu.Unlock() + return latest, false, nil + } + if latest.Attempts < 6 { + receipt.Attempts = latest.Attempts + } + } + } maxAttempts, retryDelay := 6, 15*time.Minute if strings.HasPrefix(job, "reset_probe:") { maxAttempts, retryDelay = 20, 15*time.Second @@ -141,6 +179,42 @@ func (l *JobLedger) Run(ctx context.Context, c Credential, job, slot string, run return receipt, true, err } +// PrepareQuota spreads newly observed account/model pairs over the initial +// hour. Persisting the schedule prevents a restart from bursting the whole pool. +func (l *JobLedger) PrepareQuota(c Credential, model, slot string) error { + l.mu.Lock() + defer l.mu.Unlock() + if l.saveErr != nil { + return l.saveErr + } + account, job := RiskAccountKey(c), "quota_scan:"+model + for _, receipt := range l.state.Receipts { + if receipt.Account == account && receipt.Job == job { + return nil + } + } + if len(l.state.Receipts) >= 10000 { + return fmt.Errorf("workbuddy: job receipt limit reached") + } + now := l.now() + next := now.Add(5 * time.Second) + var latest time.Time + for _, receipt := range l.state.Receipts { + if strings.HasPrefix(receipt.Job, "quota_scan:") && receipt.Status == "scheduled" && receipt.NextAttempt.After(latest) { + latest = receipt.NextAttempt + } + } + if spaced := latest.Add(15 * time.Second); spaced.After(next) { + next = spaced + } + if limit := now.Add(time.Hour); next.After(limit) { + next = limit + } + key := account + ":" + job + ":" + slot + l.state.Receipts[key] = JobReceipt{Account: account, Job: job, Slot: slot, Status: "scheduled", NextAttempt: next} + return l.persistLocked() +} + func RewardReceipt(result ControlResult) (JobReceipt, error) { receipt := JobReceipt{AlreadyDone: result.AlreadyDone} if len(result.Data) == 0 || string(result.Data) == "null" { diff --git a/internal/workbuddy/jobs_test.go b/internal/workbuddy/jobs_test.go index 1ce593cf75c..d747a01966d 100644 --- a/internal/workbuddy/jobs_test.go +++ b/internal/workbuddy/jobs_test.go @@ -205,3 +205,89 @@ func TestWorkBuddyResetProbeUsesShortRetryAndTwentyAttemptBudget(t *testing.T) { t.Fatal("reset retry budget exceeded") } } + +func TestWorkBuddyQuotaScheduleSurvivesRestartAndSlotBoundary(t *testing.T) { + now := time.Date(2026, 10, 21, 21, 59, 59, 0, time.UTC) + path := filepath.Join(t.TempDir(), "jobs.json") + ledger, err := OpenJobLedger(path, func() time.Time { return now }) + if err != nil { + t.Fatal(err) + } + first := testCredential() + second := first + second.UID = "second-quota-fixture" + slot, _ := JobSlot(now, "quota") + if err = ledger.PrepareQuota(first, "hy3", slot); err != nil { + t.Fatal(err) + } + if err = ledger.PrepareQuota(second, "hy3", slot); err != nil { + t.Fatal(err) + } + firstKey := RiskAccountKey(first) + ":quota_scan:hy3:" + slot + secondKey := RiskAccountKey(second) + ":quota_scan:hy3:" + slot + if ledger.state.Receipts[firstKey].NextAttempt.Sub(now) != 5*time.Second || ledger.state.Receipts[secondKey].NextAttempt.Sub(now) != 20*time.Second { + t.Fatal("first scans were not staggered") + } + ledger, err = OpenJobLedger(path, func() time.Time { return now }) + if err != nil { + t.Fatal(err) + } + calls := 0 + run := func(context.Context) (JobReceipt, error) { calls++; return JobReceipt{}, nil } + now = now.Add(4 * time.Second) + nextSlot, _ := JobSlot(now, "quota") + if nextSlot == slot { + t.Fatal("fixture did not cross a six-hour boundary") + } + if _, ran, _ := ledger.Run(context.Background(), first, "quota_scan:hy3", nextSlot, run); ran { + t.Fatal("slot rollover bypassed persisted initial delay") + } + now = now.Add(time.Second) + if _, ran, errRun := ledger.Run(context.Background(), first, "quota_scan:hy3", nextSlot, run); !ran || errRun != nil { + t.Fatal("due scan did not run") + } + now = now.Add(6*time.Hour - time.Second) + thirdSlot, _ := JobSlot(now, "quota") + if _, ran, _ := ledger.Run(context.Background(), first, "quota_scan:hy3", thirdSlot, run); ran { + t.Fatal("slot rollover shortened the six-hour scan interval") + } + now = now.Add(time.Second) + if _, ran, errRun := ledger.Run(context.Background(), first, "quota_scan:hy3", thirdSlot, run); !ran || errRun != nil || calls != 2 { + t.Fatal("six-hour scan did not resume") + } +} + +func TestWorkBuddyQuotaRetryBudgetDoesNotResetAtSlotBoundary(t *testing.T) { + now := time.Date(2026, 10, 21, 21, 59, 0, 0, time.UTC) + ledger, err := OpenJobLedger(filepath.Join(t.TempDir(), "jobs.json"), func() time.Time { return now }) + if err != nil { + t.Fatal(err) + } + c := testCredential() + calls := 0 + run := func(context.Context) (JobReceipt, error) { + calls++ + return JobReceipt{}, errors.New("fixture quota failure") + } + var last time.Time + for i := 1; i <= 6; i++ { + slot, _ := JobSlot(now, "quota") + receipt, ran, errRun := ledger.Run(context.Background(), c, "quota_scan:hy3", slot, run) + if !ran || errRun == nil || receipt.Attempts != i { + t.Fatalf("attempt %d: receipt=%+v ran=%v err=%v", i, receipt, ran, errRun) + } + last = now + now = now.Add(15 * time.Minute) + } + now = last.Add(6*time.Hour - time.Second) + slot, _ := JobSlot(now, "quota") + if _, ran, _ := ledger.Run(context.Background(), c, "quota_scan:hy3", slot, run); ran { + t.Fatal("slot rollover restarted exhausted retry budget early") + } + now = now.Add(time.Second) + slot, _ = JobSlot(now, "quota") + receipt, ran, errRun := ledger.Run(context.Background(), c, "quota_scan:hy3", slot, run) + if !ran || errRun == nil || receipt.Attempts != 1 || calls != 7 { + t.Fatal("new quota period did not reset attempt budget") + } +} From f285826b89e865b6a5fa266df332a36c8b8bdd22 Mon Sep 17 00:00:00 2001 From: LuneZ99 Date: Wed, 21 Oct 2026 03:24:20 +0800 Subject: [PATCH 09/10] fix(workbuddy): honor model exclusions in background probes --- .../runtime/executor/helps/workbuddy_jobs.go | 26 ++++++++++- .../executor/workbuddy_executor_test.go | 46 +++++++++++++++++++ 2 files changed, 71 insertions(+), 1 deletion(-) diff --git a/internal/runtime/executor/helps/workbuddy_jobs.go b/internal/runtime/executor/helps/workbuddy_jobs.go index 37e0c1269c1..9658a0b1955 100644 --- a/internal/runtime/executor/helps/workbuddy_jobs.go +++ b/internal/runtime/executor/helps/workbuddy_jobs.go @@ -7,10 +7,12 @@ import ( "fmt" "os" "path/filepath" + "slices" "sort" "time" "github.com/router-for-me/CLIProxyAPI/v8/internal/config" + "github.com/router-for-me/CLIProxyAPI/v8/internal/registry" "github.com/router-for-me/CLIProxyAPI/v8/internal/workbuddy" coreauth "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/auth" log "github.com/sirupsen/logrus" @@ -172,6 +174,17 @@ func (r *WorkBuddyRuntime) Maintenance(ctx context.Context, cfg *config.Config, } } if len(probe) > 0 && (policy.QuotaScan || policy.ResetProbes) { + allowed := func(a *coreauth.Auth, model string) bool { + credential, errCredential := WorkBuddyCredential(a) + if errCredential != nil || credential.Disabled || len(credential.Models) > 0 && !slices.Contains(credential.Models, model) { + return false + } + if manager == nil { + return true + } + models := registry.GetGlobalRegistry() + return models.ClientSupportsModel(a.ID, model) || a.Prefix != "" && models.ClientSupportsModel(a.ID, a.Prefix+"/"+model) + } plans := r.Governor.DueResetProbes(c) pendingReset := make(map[string]bool, len(plans)) for _, plan := range plans { @@ -188,6 +201,9 @@ func (r *WorkBuddyRuntime) Maintenance(ctx context.Context, cfg *config.Config, } } for _, plan := range plans { + if !allowed(r.Current(auth), plan.Model) { + continue + } if probeCount >= 16 { return changed, errors.Join(failures...) } @@ -206,6 +222,10 @@ func (r *WorkBuddyRuntime) Maintenance(ctx context.Context, cfg *config.Config, if _, errLoad := load(); errLoad != nil { return workbuddy.JobReceipt{}, errLoad } + current := r.Current(auth) + if !allowed(current, plan.Model) { + return workbuddy.JobReceipt{}, &workbuddy.DispatchError{Reason: "model_not_allowed_for_probe", Status: 400, Stop: true} + } if probeCount > 0 { timer := time.NewTimer(400 * time.Millisecond) select { @@ -219,7 +239,11 @@ func (r *WorkBuddyRuntime) Maintenance(ctx context.Context, cfg *config.Config, if plan.ResetAt.IsZero() { scanCount++ } - return workbuddy.JobReceipt{}, probe[0](ctx, r.Current(auth), plan.Model) + current = r.Current(auth) + if !allowed(current, plan.Model) { + return workbuddy.JobReceipt{}, &workbuddy.DispatchError{Reason: "model_not_allowed_for_probe", Status: 400, Stop: true} + } + return workbuddy.JobReceipt{}, probe[0](ctx, current, plan.Model) }) if errProbe != nil { failures = append(failures, fmt.Errorf("workbuddy: quota probe: %w", errProbe)) diff --git a/internal/runtime/executor/workbuddy_executor_test.go b/internal/runtime/executor/workbuddy_executor_test.go index 3d77017039c..656e09cfcb2 100644 --- a/internal/runtime/executor/workbuddy_executor_test.go +++ b/internal/runtime/executor/workbuddy_executor_test.go @@ -728,6 +728,8 @@ func TestWorkBuddyUpstreamCreditUsageEnforcesPaidDailyBudget(t *testing.T) { func TestWorkBuddyResetProbeReusesGovernorAndStopsAfterObservedRecovery(t *testing.T) { var calls atomic.Int32 e, a := wbFixture(t, func(w http.ResponseWriter, r *http.Request) { calls.Add(1); _, _ = io.WriteString(w, wbSuccess) }) + a.Metadata["models"] = []string{"hy3"} + e.TrackAuth(a) now := time.Now().Add(-2 * time.Minute) governor, err := workbuddy.OpenGovernor(filepath.Join(t.TempDir(), "risk.json"), workbuddy.GovernorOptions{Now: func() time.Time { return now }}) if err != nil { @@ -771,3 +773,47 @@ func TestWorkBuddyStaleExecutorCannotStartMaintenanceAfterDisable(t *testing.T) t.Fatal("old executor started maintenance after disable") } } + +func TestWorkBuddyResetProbeRespectsCredentialAndPublishedModelRestrictions(t *testing.T) { + for _, withManager := range []bool{false, true} { + t.Run(map[bool]string{false: "credential-models", true: "published-models"}[withManager], func(t *testing.T) { + var calls atomic.Int32 + e, a := wbFixture(t, func(w http.ResponseWriter, r *http.Request) { calls.Add(1); _, _ = io.WriteString(w, wbSuccess) }) + if withManager { + a.Metadata["models"] = []string{"hy3"} + store := sdkauth.NewFileTokenStore() + store.SetBaseDir(t.TempDir()) + manager := coreauth.NewManager(store, nil, nil) + registered, err := manager.Register(context.Background(), a) + if err != nil { + t.Fatal(err) + } + e.BindCredentialManager(manager) + e.TrackAuth(registered) + a = registered + // No hy3 publication: a model excluded by SDK policy must stay unprobed. + } + now := time.Now().Add(-2 * time.Minute) + governor, err := workbuddy.OpenGovernor(filepath.Join(t.TempDir(), "risk.json"), workbuddy.GovernorOptions{Now: func() time.Time { return now }}) + if err != nil { + t.Fatal(err) + } + e.runtime.Governor = governor + e.cfg.WorkBuddy.Jobs = config.WorkBuddyJobs{Enabled: true, ResetProbes: true} + c, err := helps.WorkBuddyCredential(a) + if err != nil { + t.Fatal(err) + } + if _, err = governor.RecordFailure(c, "hy3", workbuddy.Fault{Kind: workbuddy.FaultModelCap, ResetAt: now.Add(time.Minute)}, workbuddy.FailureOptions{FreeModel: true}); err != nil { + t.Fatal(err) + } + now = now.Add(2 * time.Minute) + if _, err = e.RunMaintenance(context.Background()); err != nil || calls.Load() != 0 { + t.Fatalf("excluded model was probed: calls=%d err=%v", calls.Load(), err) + } + if len(governor.DueResetProbes(c)) != 1 { + t.Fatal("excluded probe falsely recorded recovery") + } + }) + } +} From 462179c9d4e910064a6a9cf21a393b3be7d45e15 Mon Sep 17 00:00:00 2001 From: LuneZ99 Date: Wed, 21 Oct 2026 03:26:55 +0800 Subject: [PATCH 10/10] fix(workbuddy): allow non-billable controls during credit parking --- internal/workbuddy/control.go | 5 ++++- internal/workbuddy/control_test.go | 28 ++++++++++++++++++++++++++++ 2 files changed, 32 insertions(+), 1 deletion(-) diff --git a/internal/workbuddy/control.go b/internal/workbuddy/control.go index 2c78ae7ad98..2a996e637a4 100644 --- a/internal/workbuddy/control.go +++ b/internal/workbuddy/control.go @@ -52,7 +52,10 @@ type controlSpec struct { var taskCodePattern = regexp.MustCompile(`^[A-Za-z0-9_-]{1,128}$`) func (r ControlRequest) spec(c Credential, now time.Time) (controlSpec, error) { - s := controlSpec{method: http.MethodPost, body: struct{}{}} + // These closed-set endpoints inspect metadata/balances or acquire rewards; + // none performs paid inference. Credit parking must not prevent recovery. + // The governor still checks every account, region and endpoint cooldown. + s := controlSpec{method: http.MethodPost, body: struct{}{}, observeBalance: true} page := r.Page if page == 0 { page = 1 diff --git a/internal/workbuddy/control_test.go b/internal/workbuddy/control_test.go index 9921bcb1aa7..1644a672961 100644 --- a/internal/workbuddy/control_test.go +++ b/internal/workbuddy/control_test.go @@ -173,3 +173,31 @@ func TestWorkBuddyControlBodyAndContext(t *testing.T) { t.Fatal("background upstream has a deadline") } } + +func TestWorkBuddyRewardAndMetadataReadsCanRecoverCreditParking(t *testing.T) { + var calls atomic.Int32 + a, c, _ := controlFixture(t, func(w http.ResponseWriter, r *http.Request) { calls.Add(1); fmt.Fprint(w, `{"code":0,"data":{}}`) }) + if err := a.Governor.UpdateLimits(Limits{ReserveCredits: 10}, nil); err != nil { + t.Fatal(err) + } + if _, err := a.Governor.RecordFailure(*c, "chat", Fault{Kind: FaultBalance, Status: 402}, FailureOptions{}); err != nil { + t.Fatal(err) + } + for _, request := range []ControlRequest{{Operation: ControlModels}, {Operation: ControlGrowthTasks}, {Operation: ControlGrowthAccept, TaskCodes: []string{"chat_5"}}, {Operation: ControlGrowthClaim, TaskCode: "chat_5"}} { + if _, err := a.Call(context.Background(), request); err != nil { + t.Fatalf("non-billable recovery operation %s blocked: %v", request.Operation, err) + } + } + if gate, _ := a.Governor.Check(*c, "chat", true, false); gate.Reason != "balance_empty" { + t.Fatal("metadata/reward response cleared credit parking without a balance observation") + } + if calls.Load() != 4 { + t.Fatal("recovery operations missing") + } + if _, err := a.Governor.RecordFailure(*c, "chat", Fault{Kind: FaultRestricted, Status: 403, Code: 11140}, FailureOptions{}); err != nil { + t.Fatal(err) + } + if _, err := a.Call(context.Background(), ControlRequest{Operation: ControlGrowthClaim, TaskCode: "chat_5"}); err == nil || calls.Load() != 4 { + t.Fatal("reward operation bypassed account restriction") + } +}