From a2ab3a74fc2385804e37ec1fbd14ae45ca032e48 Mon Sep 17 00:00:00 2001 From: Ian Gulliver Date: Thu, 15 Oct 2026 22:30:32 -0700 Subject: [PATCH] Read Codex rate limits at startup so Status always shows them --- docs/codex-protocol.md | 2 +- docs/operations.md | 2 +- internal/agents/codex.go | 30 ++++- internal/agents/manager.go | 1 + internal/codex/gen/main.go | 10 +- internal/codex/protocol.gen.go | 195 ++++++++++++++++++++------------- internal/codex/session.go | 77 +++++++++---- 7 files changed, 216 insertions(+), 101 deletions(-) diff --git a/docs/codex-protocol.md b/docs/codex-protocol.md index 5f0bb43..c112288 100644 --- a/docs/codex-protocol.md +++ b/docs/codex-protocol.md @@ -43,7 +43,7 @@ Closing stdin ends the process. - `turn/started`, `turn/completed` (`turn.status` `completed`, `interrupted` or `failed` with `turn.error`). A failed request to the model also sends an `error` notification first, with `willRetry`. - `item/started` and `item/completed` with a `ThreadItem`: `userMessage`, `agentMessage` (`text`), `reasoning`, `commandExecution` (`command`, `aggregatedOutput`, `exitCode`), `fileChange`, `mcpToolCall`, `contextCompaction` and others. - `item/agentMessage/delta` (`itemId`, `delta`) streams an agent message; the deltas join to the completed item's `text`. `item/commandExecution/outputDelta` and the reasoning deltas stream the same way. -- `thread/tokenUsage/updated` carries `last` and `total` token counts and `modelContextWindow`; `account/rateLimits/updated` the account's rate-limit windows; `thread/status/changed`, `mcpServer/startupStatus/updated`, `account/updated` and `remoteControl/status/changed` also arrive. +- `thread/tokenUsage/updated` carries `last` and `total` token counts and `modelContextWindow`; `account/rateLimits/updated` the account's rate-limit windows (the same snapshot `account/rateLimits/read` returns on request, as `rateLimits`); `thread/status/changed`, `mcpServer/startupStatus/updated`, `account/updated` and `remoteControl/status/changed` also arrive. - Lines that aren't a JSON object, don't parse, or carry a notification method outside `notificationParams` go to `Options.OnBadLine`; parsed notifications are still delivered as events. `Event.Decode` returns the notification's generated params type. Besides `item/tool/call`, server requests (approvals, user input) shouldn't arrive with approvals off; if one does, command and file-change approvals are accepted and anything else gets a JSON-RPC error. diff --git a/docs/operations.md b/docs/operations.md index a35b509..dbba776 100644 --- a/docs/operations.md +++ b/docs/operations.md @@ -118,7 +118,7 @@ Codex runs with its own home, `codex` under the data directory (`CODEX_HOME`, se ## Status -`GET /api/admin/status` reports `running_turns` (task sessions now `running`, as session, task and vendor; lead turns aren't listed); `rate_limits`, the latest status each vendor sent since byre started (Claude's `rate_limit_info` unchanged; Codex's `account/rateLimits/updated` reshaped by `codexRateLimit` in `internal/agents/codex.go` into the same `status`, `rateLimitType` and `unifiedWindows` fields); `agent_processes`, `zombies`, `stuck_after_seconds` and `protocol_errors` (below); and server counts (start time, agent processes, projects, active tasks, SSE clients, database size including WAL, last event seq). +`GET /api/admin/status` reports `running_turns` (task sessions now `running`, as session, task and vendor; lead turns aren't listed); `rate_limits`, the latest status each vendor sent since byre started (Claude's `rate_limit_info` unchanged; Codex's `account/rateLimits/updated` reshaped by `codexRateLimit` in `internal/agents/codex.go` into the same `status`, `rateLimitType` and `unifiedWindows` fields). Claude only reports during its turns; Codex is also asked once at startup (`readCodexRateLimit`, a short-lived `codex app-server` that calls `account/rateLimits/read` without opening a thread), so its entry is there before any Codex turn runs, and every Codex update broadcasts `admin.health`; `agent_processes`, `zombies`, `stuck_after_seconds` and `protocol_errors` (below); and server counts (start time, agent processes, projects, active tasks, SSE clients, database size including WAL, last event seq). Agent health is `internal/agents/health.go`: diff --git a/internal/agents/codex.go b/internal/agents/codex.go index e579495..8e78d8b 100644 --- a/internal/agents/codex.go +++ b/internal/agents/codex.go @@ -4,8 +4,11 @@ import ( "context" "encoding/json" "log/slog" + "os" "strings" + "time" + "g.fc.run/misc/byre/internal/api" "g.fc.run/misc/byre/internal/codex" ) @@ -66,9 +69,7 @@ func (m *Manager) handleCodex(ctx context.Context, l *live, s *codex.Session, ev if err != nil { return err } - m.mu.Lock() - m.rateLimits[l.vendor] = codexRateLimit(n.RateLimits) - m.mu.Unlock() + m.setCodexRateLimit(n.RateLimits) return nil case codex.MethodTurnCompleted: @@ -194,6 +195,29 @@ func codexToolResult(item codex.ThreadItem) (string, bool) { return status, status != "completed" } +func (m *Manager) readCodexRateLimit() { + err := os.MkdirAll(m.codexHome, 0o700) + if err != nil { + slog.Error("read codex rate limits", "err", err) + return + } + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + r, err := codex.ReadRateLimits(ctx, m.codexHome, []string{"CODEX_HOME=" + m.codexHome}) + if err != nil { + slog.Error("read codex rate limits", "err", err) + return + } + m.setCodexRateLimit(r) +} + +func (m *Manager) setCodexRateLimit(r codex.RateLimitSnapshot) { + m.mu.Lock() + m.rateLimits[api.Codex] = codexRateLimit(r) + m.mu.Unlock() + m.store.Broadcast(api.Event{Kind: "admin.health", Actor: "byre", Payload: map[string]any{}}) +} + func codexRateLimit(r codex.RateLimitSnapshot) map[string]any { status := "allowed" if r.RateLimitReachedType != "" { diff --git a/internal/agents/manager.go b/internal/agents/manager.go index 7e8c5a6..c6b45bd 100644 --- a/internal/agents/manager.go +++ b/internal/agents/manager.go @@ -114,6 +114,7 @@ func New(st *store.Store, worktreeRoot, leadRoot, codexHome string) *Manager { watchStop: make(chan struct{}), } go m.signalLoop() + go m.readCodexRateLimit() return m } diff --git a/internal/codex/gen/main.go b/internal/codex/gen/main.go index 8f0df0f..fc01758 100644 --- a/internal/codex/gen/main.go +++ b/internal/codex/gen/main.go @@ -24,6 +24,7 @@ var clientMethods = []string{ "turn/start", "turn/steer", "turn/interrupt", + "account/rateLimits/read", } type schema struct { @@ -201,8 +202,13 @@ func (g *gen) variants(def string) []method { log.Fatalf("%s method: %v", def, err) } params := "" - if p := v.Properties["params"]; p != nil && p.Ref != "" { - params = refName(p.Ref) + if p := v.Properties["params"]; p != nil { + if alts := nonNull(p.AnyOf); len(alts) == 1 { + p = alts[0] + } + if p.Ref != "" { + params = refName(p.Ref) + } } response := "" if strings.HasSuffix(params, "Params") { diff --git a/internal/codex/protocol.gen.go b/internal/codex/protocol.gen.go index 44b67cd..3458573 100644 --- a/internal/codex/protocol.gen.go +++ b/internal/codex/protocol.gen.go @@ -16,6 +16,7 @@ const ( MethodTurnStart = "turn/start" MethodTurnSteer = "turn/steer" MethodTurnInterrupt = "turn/interrupt" + MethodAccountRateLimitsRead = "account/rateLimits/read" MethodError = "error" MethodThreadStarted = "thread/started" MethodThreadStatusChanged = "thread/status/changed" @@ -113,15 +114,16 @@ const ( ) var clientResults = map[string]func() any{ - MethodInitialize: func() any { return &InitializeResponse{} }, - MethodThreadStart: func() any { return &ThreadStartResponse{} }, - MethodThreadResume: func() any { return &ThreadResumeResponse{} }, - MethodThreadFork: func() any { return &ThreadForkResponse{} }, - MethodThreadCompactStart: func() any { return &ThreadCompactStartResponse{} }, - MethodThreadRead: func() any { return &ThreadReadResponse{} }, - MethodTurnStart: func() any { return &TurnStartResponse{} }, - MethodTurnSteer: func() any { return &TurnSteerResponse{} }, - MethodTurnInterrupt: func() any { return &TurnInterruptResponse{} }, + MethodInitialize: func() any { return &InitializeResponse{} }, + MethodThreadStart: func() any { return &ThreadStartResponse{} }, + MethodThreadResume: func() any { return &ThreadResumeResponse{} }, + MethodThreadFork: func() any { return &ThreadForkResponse{} }, + MethodThreadCompactStart: func() any { return &ThreadCompactStartResponse{} }, + MethodThreadRead: func() any { return &ThreadReadResponse{} }, + MethodTurnStart: func() any { return &TurnStartResponse{} }, + MethodTurnSteer: func() any { return &TurnSteerResponse{} }, + MethodTurnInterrupt: func() any { return &TurnInterruptResponse{} }, + MethodAccountRateLimitsRead: func() any { return &GetAccountRateLimitsResponse{} }, } var notificationParams = map[string]func() any{ @@ -534,6 +536,47 @@ const ( ReasoningSummaryNone ReasoningSummary = "none" ) +const ( + RateLimitResetTypeCodexRateLimits RateLimitResetType = "codexRateLimits" + RateLimitResetTypeUnknown RateLimitResetType = "unknown" +) + +const ( + RateLimitResetCreditStatusAvailable RateLimitResetCreditStatus = "available" + RateLimitResetCreditStatusRedeeming RateLimitResetCreditStatus = "redeeming" + RateLimitResetCreditStatusRedeemed RateLimitResetCreditStatus = "redeemed" + RateLimitResetCreditStatusUnknown RateLimitResetCreditStatus = "unknown" +) + +const ( + PlanTypeFree PlanType = "free" + PlanTypeGo PlanType = "go" + PlanTypePlus PlanType = "plus" + PlanTypePro PlanType = "pro" + PlanTypeProlite PlanType = "prolite" + PlanTypePromax PlanType = "promax" + PlanTypeTeam PlanType = "team" + PlanTypeSelfServeBusinessProlite PlanType = "self_serve_business_prolite" + PlanTypeSelfServeBusinessUsageBased PlanType = "self_serve_business_usage_based" + PlanTypeBusiness PlanType = "business" + PlanTypeEnt26 PlanType = "ent26" + PlanTypeEnterpriseCbpAutomation PlanType = "enterprise_cbp_automation" + PlanTypeEnterpriseCbpUsageBased PlanType = "enterprise_cbp_usage_based" + PlanTypeEnterprise PlanType = "enterprise" + PlanTypeEdu PlanType = "edu" + PlanTypeEduPlus PlanType = "edu_plus" + PlanTypeEduPro PlanType = "edu_pro" + PlanTypeUnknown PlanType = "unknown" +) + +const ( + RateLimitReachedTypeRateLimitReached RateLimitReachedType = "rate_limit_reached" + RateLimitReachedTypeWorkspaceOwnerCreditsDepleted RateLimitReachedType = "workspace_owner_credits_depleted" + RateLimitReachedTypeWorkspaceMemberCreditsDepleted RateLimitReachedType = "workspace_member_credits_depleted" + RateLimitReachedTypeWorkspaceOwnerUsageLimitReached RateLimitReachedType = "workspace_owner_usage_limit_reached" + RateLimitReachedTypeWorkspaceMemberUsageLimitReached RateLimitReachedType = "workspace_member_usage_limit_reached" +) + const ( ThreadAttachmentOperationCreated ThreadAttachmentOperation = "created" ThreadAttachmentOperationDeleted ThreadAttachmentOperation = "deleted" @@ -714,27 +757,6 @@ const ( AuthModeBedrockAccessKeys AuthMode = "bedrockAccessKeys" ) -const ( - PlanTypeFree PlanType = "free" - PlanTypeGo PlanType = "go" - PlanTypePlus PlanType = "plus" - PlanTypePro PlanType = "pro" - PlanTypeProlite PlanType = "prolite" - PlanTypePromax PlanType = "promax" - PlanTypeTeam PlanType = "team" - PlanTypeSelfServeBusinessProlite PlanType = "self_serve_business_prolite" - PlanTypeSelfServeBusinessUsageBased PlanType = "self_serve_business_usage_based" - PlanTypeBusiness PlanType = "business" - PlanTypeEnt26 PlanType = "ent26" - PlanTypeEnterpriseCbpAutomation PlanType = "enterprise_cbp_automation" - PlanTypeEnterpriseCbpUsageBased PlanType = "enterprise_cbp_usage_based" - PlanTypeEnterprise PlanType = "enterprise" - PlanTypeEdu PlanType = "edu" - PlanTypeEduPlus PlanType = "edu_plus" - PlanTypeEduPro PlanType = "edu_pro" - PlanTypeUnknown PlanType = "unknown" -) - const ( GatewayOAuthStatusNotReady GatewayOAuthStatus = "notReady" GatewayOAuthStatusStarted GatewayOAuthStatus = "started" @@ -742,14 +764,6 @@ const ( GatewayOAuthStatusFailed GatewayOAuthStatus = "failed" ) -const ( - RateLimitReachedTypeRateLimitReached RateLimitReachedType = "rate_limit_reached" - RateLimitReachedTypeWorkspaceOwnerCreditsDepleted RateLimitReachedType = "workspace_owner_credits_depleted" - RateLimitReachedTypeWorkspaceMemberCreditsDepleted RateLimitReachedType = "workspace_member_credits_depleted" - RateLimitReachedTypeWorkspaceOwnerUsageLimitReached RateLimitReachedType = "workspace_owner_usage_limit_reached" - RateLimitReachedTypeWorkspaceMemberUsageLimitReached RateLimitReachedType = "workspace_member_usage_limit_reached" -) - const ( RemoteControlConnectionStatusDisabled RemoteControlConnectionStatus = "disabled" RemoteControlConnectionStatusConnecting RemoteControlConnectionStatus = "connecting" @@ -1611,6 +1625,75 @@ type TurnInterruptParams struct { type TurnInterruptResponse = map[string]any +type GetAccountRateLimitsParams struct { + ExcludeResetCreditDetails bool `json:"excludeResetCreditDetails,omitempty"` + SupportsLunaReserve bool `json:"supportsLunaReserve,omitempty"` +} + +type RateLimitResetType string + +type RateLimitResetCreditStatus string + +type RateLimitResetCredit struct { + Description string `json:"description,omitempty"` + ExpiresAt int64 `json:"expiresAt,omitempty"` + GrantedAt int64 `json:"grantedAt"` + ID string `json:"id"` + ResetType RateLimitResetType `json:"resetType"` + Status RateLimitResetCreditStatus `json:"status"` + Title string `json:"title,omitempty"` +} + +type RateLimitResetCreditsSummary struct { + AvailableCount int64 `json:"availableCount"` + Credits []RateLimitResetCredit `json:"credits,omitempty"` +} + +type CreditsSnapshot struct { + Balance string `json:"balance,omitempty"` + HasCredits bool `json:"hasCredits"` + Unlimited bool `json:"unlimited"` +} + +type SpendControlLimitSnapshot struct { + Limit string `json:"limit"` + RemainingPercent int64 `json:"remainingPercent"` + ResetsAt int64 `json:"resetsAt"` + Used string `json:"used"` +} + +type PlanType string + +type RateLimitWindow struct { + ResetsAt int64 `json:"resetsAt,omitempty"` + UsedPercent int64 `json:"usedPercent"` + WindowDurationMins int64 `json:"windowDurationMins,omitempty"` +} + +type RateLimitReachedType string + +type RateLimitSnapshot struct { + Credits *CreditsSnapshot `json:"credits,omitempty"` + IndividualLimit *SpendControlLimitSnapshot `json:"individualLimit,omitempty"` + LimitID string `json:"limitId,omitempty"` + LimitName string `json:"limitName,omitempty"` + NormalModelSlug string `json:"normalModelSlug,omitempty"` + PlanType PlanType `json:"planType,omitempty"` + Primary *RateLimitWindow `json:"primary,omitempty"` + RateLimitReachedType RateLimitReachedType `json:"rateLimitReachedType,omitempty"` + Secondary *RateLimitWindow `json:"secondary,omitempty"` + SpendControlReached bool `json:"spendControlReached,omitempty"` +} + +type GetAccountRateLimitsResponse struct { + AccountID string `json:"accountId,omitempty"` + OrdinaryUsageAllowed bool `json:"ordinaryUsageAllowed,omitempty"` + RateLimitResetCredits *RateLimitResetCreditsSummary `json:"rateLimitResetCredits,omitempty"` + RateLimitUpsell json.RawMessage `json:"rateLimitUpsell,omitempty"` + RateLimits RateLimitSnapshot `json:"rateLimits"` + RateLimitsByLimitID map[string]RateLimitSnapshot `json:"rateLimitsByLimitId,omitempty"` +} + type ErrorNotification struct { Error TurnError `json:"error"` ThreadID string `json:"threadId"` @@ -2065,8 +2148,6 @@ type McpServerEventStreamNotification struct { type AuthMode string -type PlanType string - type AccountUpdatedNotification struct { AuthMode AuthMode `json:"authMode,omitempty"` PlanType PlanType `json:"planType,omitempty"` @@ -2081,40 +2162,6 @@ type GatewayOAuthChangedNotification struct { Status GatewayOAuthStatus `json:"status"` } -type CreditsSnapshot struct { - Balance string `json:"balance,omitempty"` - HasCredits bool `json:"hasCredits"` - Unlimited bool `json:"unlimited"` -} - -type SpendControlLimitSnapshot struct { - Limit string `json:"limit"` - RemainingPercent int64 `json:"remainingPercent"` - ResetsAt int64 `json:"resetsAt"` - Used string `json:"used"` -} - -type RateLimitWindow struct { - ResetsAt int64 `json:"resetsAt,omitempty"` - UsedPercent int64 `json:"usedPercent"` - WindowDurationMins int64 `json:"windowDurationMins,omitempty"` -} - -type RateLimitReachedType string - -type RateLimitSnapshot struct { - Credits *CreditsSnapshot `json:"credits,omitempty"` - IndividualLimit *SpendControlLimitSnapshot `json:"individualLimit,omitempty"` - LimitID string `json:"limitId,omitempty"` - LimitName string `json:"limitName,omitempty"` - NormalModelSlug string `json:"normalModelSlug,omitempty"` - PlanType PlanType `json:"planType,omitempty"` - Primary *RateLimitWindow `json:"primary,omitempty"` - RateLimitReachedType RateLimitReachedType `json:"rateLimitReachedType,omitempty"` - Secondary *RateLimitWindow `json:"secondary,omitempty"` - SpendControlReached bool `json:"spendControlReached,omitempty"` -} - type AccountRateLimitsUpdatedNotification struct { RateLimits RateLimitSnapshot `json:"rateLimits"` } diff --git a/internal/codex/session.go b/internal/codex/session.go index ffc2c47..4add473 100644 --- a/internal/codex/session.go +++ b/internal/codex/session.go @@ -116,6 +116,45 @@ type Session struct { } func Start(ctx context.Context, opts Options) (*Session, error) { + s, err := spawn(opts) + if err != nil { + return nil, err + } + + err = s.open(ctx, opts) + if err != nil { + s.Close() + return nil, err + } + return s, nil +} + +func ReadRateLimits(ctx context.Context, cwd string, env []string) (RateLimitSnapshot, error) { + s, err := spawn(Options{Cwd: cwd, Env: env}) + if err != nil { + return RateLimitSnapshot{}, err + } + defer s.Close() + + go func() { + for range s.events { + } + }() + + err = s.initialize(ctx) + if err != nil { + return RateLimitSnapshot{}, err + } + + resp := GetAccountRateLimitsResponse{} + err = s.Call(ctx, MethodAccountRateLimitsRead, GetAccountRateLimitsParams{ExcludeResetCreditDetails: true}, &resp) + if err != nil { + return RateLimitSnapshot{}, fmt.Errorf("read rate limits: %w", err) + } + return resp.RateLimits, nil +} + +func spawn(opts Options) (*Session, error) { cmd := exec.Command("codex", "app-server", "-c", "check_for_update_on_startup=false") cmd.Dir = opts.Cwd cmd.Stderr = os.Stderr @@ -138,12 +177,6 @@ func Start(ctx context.Context, opts Options) (*Session, error) { s := newSession(opts, stdin, stdout) s.cmd = cmd - - err = s.open(ctx, opts) - if err != nil { - s.Close() - return nil, err - } return s, nil } @@ -167,20 +200,7 @@ func newSession(opts Options, stdin io.WriteCloser, stdout io.Reader) *Session { } func (s *Session) open(ctx context.Context, opts Options) error { - init := InitializeResponse{} - err := s.Call(ctx, MethodInitialize, InitializeParams{ - ClientInfo: ClientInfo{Name: "byre", Title: "byre", Version: "0.0.0"}, - Capabilities: &InitializeCapabilities{ExperimentalApi: true}, - }, &init) - if err != nil { - return fmt.Errorf("initialize: %w", err) - } - - s.stateMu.Lock() - s.userAgent = init.UserAgent - s.stateMu.Unlock() - - err = s.notify("initialized", nil) + err := s.initialize(ctx) if err != nil { return err } @@ -239,6 +259,23 @@ func (s *Session) open(ctx context.Context, opts Options) error { return nil } +func (s *Session) initialize(ctx context.Context) error { + init := InitializeResponse{} + err := s.Call(ctx, MethodInitialize, InitializeParams{ + ClientInfo: ClientInfo{Name: "byre", Title: "byre", Version: "0.0.0"}, + Capabilities: &InitializeCapabilities{ExperimentalApi: true}, + }, &init) + if err != nil { + return fmt.Errorf("initialize: %w", err) + } + + s.stateMu.Lock() + s.userAgent = init.UserAgent + s.stateMu.Unlock() + + return s.notify("initialized", nil) +} + func stdioConfig(servers []mcp.Stdio) map[string]any { if len(servers) == 0 { return nil