From b44b5d45e45b617d91e0460399ae3d379a1f9467 Mon Sep 17 00:00:00 2001 From: Bruno Bornsztein Date: Sat, 15 Aug 2026 17:08:37 -0500 Subject: [PATCH 1/3] feat(plugins): route tasks to the Claude profile with the most headroom MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two Claude logins (two CLAUDE_CONFIG_DIRs) means a choice on every task — which account should this one spend? Until now that choice was manual and uninformed, so one account got hammered to a 429 while the other sat idle. Three pieces, each useful on its own: - internal/claudeusage reads a profile's stored OAuth token (macOS Keychain, namespaced per config dir by a sha256 prefix, or .credentials.json) and calls the same /api/oauth/usage endpoint Claude Code's own /usage command uses. It is strictly read-only: no token is refreshed, rewritten, or printed. The endpoint rate-limits and routing probes it on every spawn, so snapshots are cached on disk for a minute, and a cached snapshot up to 30 minutes old rescues a failed probe rather than leaving routing blind. - `ty usage` surfaces those numbers per profile. --percent prints one bare number, because the caller that matters is a shell script with no JSON parser. - task.route is a new plugin hook, and the first one TaskYou *waits* on: it fires before a task spawns and reads the script's stdout back as a decision (CLAUDE_CONFIG_DIR=…, HOLD=1, REASON=…). Every other hook can only react to a task having run; a router has to answer before the command is built. The decision is carried by Task.ClaudeConfigDir, which both command builders have always honored — so there is no second spawn path and no way for the daemon and the TUI to disagree about which profile is in play. It also gives session affinity for free: a routed task is stamped once and stays there, which is required, not merely tidy, since a Claude session lives inside one config dir and resuming elsewhere would silently start a fresh conversation. Routing never becomes a reason work doesn't happen: no router, a failing script, a timeout, or an unreadable profile all mean the task spawns exactly as it would have. An explicit config dir (set by hand or by a workflow step) is never overruled. HOLD leaves a task queued, never blocked, so it starts by itself once limits reset — and is ignored on a manual run. examples/plugins/claude-profile-router is the working plugin: list your profile dirs, and each task goes to the account with the most headroom. Co-Authored-By: Claude Opus 5 --- README.md | 4 +- cmd/task/main.go | 1 + cmd/task/usage.go | 203 +++++++++++ docs/plugin-ideas.md | 17 +- docs/plugins.md | 60 ++++ .../plugins/claude-profile-router/README.md | 95 +++++ .../claude-profile-router/config.example.env | 21 ++ .../plugins/claude-profile-router/plugin.yaml | 18 + .../plugins/claude-profile-router/route.sh | 102 ++++++ .../plugins/claude-profile-router/status.sh | 31 ++ internal/claudeusage/cache.go | 109 ++++++ internal/claudeusage/cache_test.go | 195 +++++++++++ internal/claudeusage/credentials.go | 157 +++++++++ internal/claudeusage/usage.go | 327 ++++++++++++++++++ internal/claudeusage/usage_test.go | 261 ++++++++++++++ internal/db/tasks.go | 17 + internal/executor/executor.go | 14 + internal/executor/routing.go | 113 ++++++ internal/executor/routing_test.go | 277 +++++++++++++++ internal/hooks/route.go | 184 ++++++++++ internal/hooks/route_test.go | 205 +++++++++++ 21 files changed, 2408 insertions(+), 3 deletions(-) create mode 100644 cmd/task/usage.go create mode 100644 examples/plugins/claude-profile-router/README.md create mode 100644 examples/plugins/claude-profile-router/config.example.env create mode 100644 examples/plugins/claude-profile-router/plugin.yaml create mode 100755 examples/plugins/claude-profile-router/route.sh create mode 100755 examples/plugins/claude-profile-router/status.sh create mode 100644 internal/claudeusage/cache.go create mode 100644 internal/claudeusage/cache_test.go create mode 100644 internal/claudeusage/credentials.go create mode 100644 internal/claudeusage/usage.go create mode 100644 internal/claudeusage/usage_test.go create mode 100644 internal/executor/routing.go create mode 100644 internal/executor/routing_test.go create mode 100644 internal/hooks/route.go create mode 100644 internal/hooks/route_test.go diff --git a/README.md b/README.md index 4698cd63..78cc827a 100644 --- a/README.md +++ b/README.md @@ -688,7 +688,9 @@ A **plugin** is a self-contained directory under `~/.config/task/plugins/` with - **workflows** (`workflows/*.yaml`) — new `ty pipeline -d ` definitions - **hooks** — scripts that fire on task events. Unlike the one-script-per-event hooks - dir above, any number of plugins can handle the same event and **all of them run** + dir above, any number of plugins can handle the same event and **all of them run**. + One of them, `task.route`, fires *before* a task spawns and lets the plugin pick + which Claude account it runs under — see [Routing](docs/plugins.md#routing-pre-spawn) - **actions** — user-invoked commands (`ty plugins run `) Install one — or a whole collection, since a single git repo can hold many plugins — diff --git a/cmd/task/main.go b/cmd/task/main.go index 55dc1b36..387d4d86 100644 --- a/cmd/task/main.go +++ b/cmd/task/main.go @@ -531,6 +531,7 @@ Examples: // Plugins subcommand - inspect installed task plugins rootCmd.AddCommand(newPluginsCmd()) + rootCmd.AddCommand(newUsageCmd()) // Workflow artifact store over the CLI — the transport-independent twin of the // taskyou_get_artifact/taskyou_set_artifact MCP tools, so a workflow phase can diff --git a/cmd/task/usage.go b/cmd/task/usage.go new file mode 100644 index 00000000..968fb78e --- /dev/null +++ b/cmd/task/usage.go @@ -0,0 +1,203 @@ +package main + +import ( + "context" + "encoding/json" + "fmt" + "os" + "sort" + "strings" + "time" + + "github.com/spf13/cobra" + + "github.com/bborn/workflow/internal/claudeusage" + "github.com/bborn/workflow/internal/db" + "github.com/bborn/workflow/internal/executor" +) + +// `ty usage` exposes, per Claude profile, how much of that account's rate +// limits are already spent. It exists so a decision that used to be guesswork — +// "which of my logins should this task run under?" — can be made from real +// numbers, by a person at the terminal or by a plugin's routing script. +// +// The plugin path is why the output is shaped the way it is. A hook script has +// no JSON parser to lean on, so --percent prints one bare number and nothing +// else; comparing profiles is then a numeric sort in shell. --json is for +// anything richer. + +// profileResult pairs a probed config dir with its outcome. Errors are carried +// rather than returned so one unreadable profile (an expired login, say) still +// lets the others report. +type profileResult struct { + ConfigDir string `json:"config_dir"` + Snapshot *claudeusage.Snapshot `json:"snapshot,omitempty"` + Error string `json:"error,omitempty"` +} + +func newUsageCmd() *cobra.Command { + cmd := &cobra.Command{ + Use: "usage", + Short: "Show how much of each Claude profile's rate limits are used", + Long: `Report Claude subscription usage for one or more profiles. + +A profile is a CLAUDE_CONFIG_DIR — one logged-in Claude account. With no +--config-dir flags, every config dir ty knows about is probed: the default +(~/.claude) plus each distinct dir configured on a project. + +Usage is read from the same endpoint Claude Code's own /usage command uses, +authenticated with the credentials already stored for that profile. Nothing is +written: no token is refreshed, rewritten, or printed. + +Results are cached for a minute, because that endpoint rate-limits and routing +calls it on every task spawn. Use --refresh to force a live read. + +Examples: + ty usage # every profile ty knows about + ty usage --config-dir ~/.claude-work # one profile, with account email + ty usage --config-dir ~/.claude-work --percent # just "42" — for scripts + ty usage --json # full detail + ty usage --refresh # ignore the cache`, + SilenceUsage: true, + RunE: func(cmd *cobra.Command, args []string) error { + dirs, _ := cmd.Flags().GetStringArray("config-dir") + asJSON, _ := cmd.Flags().GetBool("json") + percentOnly, _ := cmd.Flags().GetBool("percent") + refresh, _ := cmd.Flags().GetBool("refresh") + + if len(dirs) == 0 { + dirs = knownConfigDirs() + } + if len(dirs) == 0 { + return fmt.Errorf("no Claude config dirs to check") + } + if percentOnly && len(dirs) != 1 { + return fmt.Errorf("--percent needs exactly one --config-dir (got %d)", len(dirs)) + } + + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + + client := claudeusage.NewClient() + client.NoCache = refresh + results := probeProfiles(ctx, client, dirs, !percentOnly) + + switch { + case percentOnly: + if results[0].Error != "" { + return fmt.Errorf("%s", results[0].Error) + } + // One bare number, no styling, no trailing prose: routing + // scripts read this with $(...) and compare it numerically. + fmt.Printf("%.0f\n", results[0].Snapshot.UsedPercent()) + case asJSON: + enc := json.NewEncoder(os.Stdout) + enc.SetIndent("", " ") + return enc.Encode(results) + default: + printUsageTable(results) + } + return nil + }, + } + + cmd.Flags().StringArray("config-dir", nil, "Claude config dir to check (repeatable; defaults to every dir ty knows about)") + cmd.Flags().Bool("json", false, "Emit JSON") + cmd.Flags().Bool("percent", false, "Print only the binding limit's used percent (requires one --config-dir)") + cmd.Flags().Bool("refresh", false, "Bypass the cache and read live") + return cmd +} + +// probeProfiles fetches usage for each dir. withAccount adds the account email, +// which costs a second request per profile — worth it for a human reading a +// table, wasted for a script that only wants a number. +func probeProfiles(ctx context.Context, client *claudeusage.Client, dirs []string, withAccount bool) []profileResult { + results := make([]profileResult, 0, len(dirs)) + for _, dir := range dirs { + res := profileResult{ConfigDir: dir} + var snap *claudeusage.Snapshot + var err error + if withAccount { + snap, err = client.FetchWithAccount(ctx, dir) + } else { + snap, err = client.Fetch(ctx, dir) + } + if err != nil { + res.Error = err.Error() + } else { + res.Snapshot = snap + res.ConfigDir = snap.ConfigDir + } + results = append(results, res) + } + return results +} + +// knownConfigDirs collects every Claude config dir ty is already aware of: the +// default, plus whatever projects have been pointed at. Opening the DB is best +// effort — `ty usage` must still work from a machine with no board. +func knownConfigDirs() []string { + seen := map[string]bool{} + var dirs []string + add := func(d string) { + d = executor.ResolveClaudeConfigDir(d) + if d == "" || seen[d] { + return + } + seen[d] = true + dirs = append(dirs, d) + } + + add("") // the default dir + + if database, err := openTaskDB(db.DefaultPath()); err == nil { + defer database.Close() //nolint:errcheck // read-only lookup + if projects, err := database.ListProjects(); err == nil { + for _, p := range projects { + if strings.TrimSpace(p.ClaudeConfigDir) != "" { + add(p.ClaudeConfigDir) + } + } + } + } + + sort.Strings(dirs[1:]) // keep the default first, order the rest stably + return dirs +} + +func printUsageTable(results []profileResult) { + for _, r := range results { + fmt.Println(boldStyle.Render(r.ConfigDir)) + if r.Error != "" { + fmt.Println(" " + errorStyle.Render(r.Error)) + fmt.Println() + continue + } + if r.Snapshot.Email != "" { + fmt.Println(" " + dimStyle.Render(r.Snapshot.Email)) + } + style := successStyle + switch used := r.Snapshot.UsedPercent(); { + case used >= 90: + style = errorStyle + case used >= 70: + style = warnStyle + } + line := r.Snapshot.Describe() + if r.Snapshot.Stale { + line += fmt.Sprintf(" (cached %s ago; the usage API is unreachable)", r.Snapshot.Age().Round(time.Minute)) + } + fmt.Println(" " + style.Render(line)) + for _, l := range r.Snapshot.Limits { + line := fmt.Sprintf(" %-16s %3.0f%%", l.Kind, l.Percent) + if l.Scope != "" { + line += " " + l.Scope + } + if l.ResetsAt != nil { + line += " resets " + l.ResetsAt.Local().Format("Mon 15:04") + } + fmt.Println(dimStyle.Render(line)) + } + fmt.Println() + } +} diff --git a/docs/plugin-ideas.md b/docs/plugin-ideas.md index caa138a3..0a7501f5 100644 --- a/docs/plugin-ideas.md +++ b/docs/plugin-ideas.md @@ -28,6 +28,19 @@ can be any language and can bundle its own config/binaries. - **status-file** — maintain a tiny JSON of live counts for a tmux statusline or menubar widget. +## Routing (the `task.route` hook) + +Fires before a task spawns and its stdout is read back as a decision — the one +hook that changes how a task runs rather than reporting on it. See +[Routing](plugins.md#routing-pre-spawn). + +- ✅ **claude-profile-router** — send each task to whichever Claude account has + the most rate-limit headroom; hold the task when both are spent. +- **quiet-hours** — `HOLD=1` outside working hours, so overnight queueing doesn't + spend your weekly limit while you sleep. +- **cheap-account-first** — route routine task types (docs, chores) to a Pro + account and keep the Max one for the heavy work. + ## Worktree & quality (actions) - ✅ **worktree** — show the task's diff; run its tests. @@ -51,8 +64,8 @@ can be any language and can bundle its own config/binaries. ## Where should plugins live? (in-repo vs. own repo) - **In-repo `examples/plugins/`** — small, canonical, copy-paste starting points - that ship with TaskYou and are covered by the loader's tests. The three above - live here. Best for anything short enough to read in one sitting. + that ship with TaskYou and are covered by the loader's tests. The ✅ entries + above live here. Best for anything short enough to read in one sitting. - **Its own repo** — when a plugin grows an independent release cadence, ships a compiled binary or heavier dependencies, or has a real surface of its own (config, docs, versioning). Install by dropping (or symlinking) its directory diff --git a/docs/plugins.md b/docs/plugins.md index 2660502f..37fa8f25 100644 --- a/docs/plugins.md +++ b/docs/plugins.md @@ -146,10 +146,68 @@ system. The ones dispatched today: | `task.blocked` | Task needs input | | `task.failed` | Agent execution failed | | `task.auth_required` | Executor session needs re-authentication | +| `task.route` | **Before** a task spawns — the one event you *answer*. See [Routing](#routing-pre-spawn) | A plugin may declare any event string; it only runs for events TaskYou actually emits, so unknown events are harmless. +## Routing (pre-spawn) + +Every hook above is a notification: it fires after the fact, runs detached, and +nothing waits for it. `task.route` is the exception. It fires *before* a task is +spawned, TaskYou waits for it, and it reads your script's **stdout back as a +decision** — which is the only way a plugin can influence how a task runs rather +than just react to it having run. + +The decision format is KEY=VALUE, one per line: + +```sh +#!/bin/sh +echo "CLAUDE_CONFIG_DIR=$HOME/.claude-work" # run this task under that Claude profile +echo "REASON=7% of its limits used" # optional note for the task log +``` + +| Key | Effect | +|-----|--------| +| `CLAUDE_CONFIG_DIR` | Run the task under that Claude profile (config dir) | +| `HOLD=1` | Don't start this task yet; leave it queued and reconsider next tick | +| `REASON=…` | Free text, written to the task log alongside the decision | + +Everything else on stdout is ignored, so unknown keys and stray output are +harmless — but keep diagnostics on **stderr** (which goes to the daemon log), +since stdout is the decision channel. + +**Guarantees.** Printing nothing is always safe, and so is failing: + +- **Silence means carry on.** No router installed, a script that errors, exceeds + the 15s timeout, or prints nothing usable — the task spawns exactly as it would + have. Routing is an optimization; failing to optimize never blocks work. +- **An explicit choice wins.** A task that already names a config dir (set by + hand, or by a workflow step's `config_dir:`) is left alone. +- **A routed task stays put.** The decision is written to the task, so a task + resumed later runs under the same profile it started on. This is required, not + merely tidy: a Claude session lives inside one config dir, and resuming under a + different one would find no session and quietly start a fresh conversation. The + cost is that a long-lived task pinned to a profile waits for *that* profile's + limits to reset — it cannot be migrated mid-conversation. +- **First answer stands.** Plugins are consulted in name order and the first + non-empty decision is used, so two installed routers give a deterministic + result. +- **`HOLD` keeps a task queued**, never blocked — it starts by itself once a + later tick gets a different answer. A hold is honored for queued tasks; a task + you started by hand (`ty run`, "start now") runs regardless. +- **Claude only.** `CLAUDE_CONFIG_DIR` means nothing to the codex or gemini + executors, so tasks using them are not routed. + +Extra environment on a routing hook, beyond the standard `TASK_*` and +`TASK_PLUGIN_*` variables: `TASK_EXECUTOR` and `TASK_CLAUDE_CONFIG_DIR` (the +task's current config dir, empty when unset). + +The worked example is +[`examples/plugins/claude-profile-router/`](../examples/plugins/claude-profile-router/), +which routes each task to whichever of your Claude accounts has the most +rate-limit headroom — see also [`ty usage`](#inspecting). + ## Environment Every hook receives the standard task variables: @@ -247,6 +305,7 @@ only for a process that must stay *up* (a socket connection, a listening port). ```bash ty plugins list # what's installed and which events each handles ty plugins dir # the plugins directory path +ty usage # rate-limit usage per Claude profile — what a router routes on ``` Set `TY_PLUGINS_DIR` to use a directory other than `~/.config/task/plugins/`. @@ -261,6 +320,7 @@ Complete, copy-pasteable plugins live in [`examples/plugins/`](../examples/plugi | [`slack`](../examples/plugins/slack/) | hooks | webhook integration; bundled `config.env` | | [`worktree`](../examples/plugins/worktree/) | actions | task-scoped `diff` / `test` using `WORKTREE_PATH` | | [`heartbeat`](../examples/plugins/heartbeat/) | service | a daemon-supervised long-running process | +| [`claude-profile-router`](../examples/plugins/claude-profile-router/) | route hook + action | picking a Claude account per task from live usage | ```bash cp -R examples/plugins/desktop-notify ~/.config/task/plugins/ diff --git a/examples/plugins/claude-profile-router/README.md b/examples/plugins/claude-profile-router/README.md new file mode 100644 index 00000000..1bd959a8 --- /dev/null +++ b/examples/plugins/claude-profile-router/README.md @@ -0,0 +1,95 @@ +# claude-profile-router + +Route each task to whichever of your Claude accounts has the most rate-limit +headroom left. + +If you have two logins — a personal one and a work one, say — you already have +two Claude config dirs. This plugin checks how much of each account's 5-hour and +weekly limits are spent and points every task ty spawns at the one with room. If +both are spent, it holds the task in the queue instead of burning a session on a +429. + +## Setup + +1. **Have two profiles.** A profile is a `CLAUDE_CONFIG_DIR` with its own login: + + ```bash + CLAUDE_CONFIG_DIR=~/.claude-work claude # then /login as the second account + ``` + +2. **Install the plugin:** + + ```bash + cp -R examples/plugins/claude-profile-router ~/.config/task/plugins/ + cd ~/.config/task/plugins/claude-profile-router + cp config.example.env config.env + $EDITOR config.env # list your profile dirs in TY_CLAUDE_PROFILES + ``` + +3. **Check it sees both accounts:** + + ```bash + ty plugins run claude-profile-router status + ``` + + ``` + Routing threshold: skip a profile at or above 90% used + + /Users/me/.claude + me@personal.example + 3% used (5-hour session, resets Sat 14:00) — 97% headroom + + /Users/me/.claude-work + me@work.example + 71% used (weekly, resets Thu 20:00) — 29% headroom + ``` + +That's it — the next task ty spawns is routed. `ty logs` and the task's own log +record which profile it landed on and why. + +## How it decides + +- Each profile's **binding limit** is the worst of its reported windows (5-hour + session, weekly, per-model weekly). A session window at 98% blocks the next + task even when the weekly one is untouched, so the worst window is the one + that matters. +- The profile with the **lowest** binding percent wins. +- Profiles at or above `TY_CLAUDE_MAX_PERCENT` (default 90) are skipped. The + margin exists because usage is sampled at spawn, not metered continuously — + a long task started at 89% can still cross the line mid-run. +- If every profile is over the threshold, the task is **held**: it stays queued + and is reconsidered on the next daemon tick, with one log line saying why. +- A task that already names a config dir (set by hand, or by a workflow step) is + left alone. Routing fills a vacuum; it doesn't overrule you. +- **A task is routed once and stays there.** Its Claude session lives inside that + config dir, so a resume has to happen under the same profile or it would start + a fresh conversation. A task already running on a profile therefore waits for + *that* profile to reset rather than hopping to the other one. +- Anything that goes wrong — no credentials, an expired login, `ty` not on the + daemon's `PATH` — means the plugin says nothing and the task spawns exactly as + it would have without it. + +## Configuration + +See [`config.example.env`](config.example.env). The knobs: + +| Variable | Default | Meaning | +| --- | --- | --- | +| `TY_CLAUDE_PROFILES` | *(required)* | Space-separated config dirs to route between | +| `TY_CLAUDE_MAX_PERCENT` | `90` | Skip a profile at or above this percent used | +| `TY_CLAUDE_PROJECTS` | *(all)* | Only route tasks in these projects | +| `TY_BIN` | `ty` | Path to the ty binary, if the daemon's `PATH` lacks it | + +## Caveats + +- **A config dir is more than an account.** It also carries that profile's + plugins, MCP servers, and trusted-worktree state. Set both profiles up the + same way, or a task routed to the quieter one may find tools missing. If you + want to swap only credentials, use a per-task `env` override instead — see + [docs/plugins.md](../../../docs/plugins.md). +- **Usage is read, never written.** The plugin reads each profile's stored OAuth + token to call the same endpoint Claude Code's `/usage` uses. It never + refreshes or rewrites a credential. A profile whose token has gone stale + reports as unavailable until you run a `claude` session under it. +- **Two probes per spawn.** Each is a single HTTPS GET; ty caps the whole hook + at 15s and spawns normally if it overruns. diff --git a/examples/plugins/claude-profile-router/config.example.env b/examples/plugins/claude-profile-router/config.example.env new file mode 100644 index 00000000..19550d8a --- /dev/null +++ b/examples/plugins/claude-profile-router/config.example.env @@ -0,0 +1,21 @@ +# Copy to config.env (next to this file) and edit. + +# The Claude profiles to route between: a space-separated list of +# CLAUDE_CONFIG_DIR paths. Each is one logged-in account. +# +# To create a second profile, log in with a different config dir: +# CLAUDE_CONFIG_DIR=~/.claude-work claude # then /login +# +# Check what each one is with: ty usage +TY_CLAUDE_PROFILES="$HOME/.claude $HOME/.claude-work" + +# Skip a profile once its binding limit (5-hour session or weekly, whichever is +# worse) is at or above this percent. When every profile is above it, tasks are +# held in the queue instead of spawning into a 429. +TY_CLAUDE_MAX_PERCENT=90 + +# Only route tasks in these projects (space-separated). Empty = all projects. +TY_CLAUDE_PROJECTS="" + +# Path to the ty binary, if the daemon's PATH doesn't include it. +# TY_BIN=/usr/local/bin/ty diff --git a/examples/plugins/claude-profile-router/plugin.yaml b/examples/plugins/claude-profile-router/plugin.yaml new file mode 100644 index 00000000..17df1a60 --- /dev/null +++ b/examples/plugins/claude-profile-router/plugin.yaml @@ -0,0 +1,18 @@ +# Route each task to whichever Claude account has the most rate-limit headroom. +# +# Copy this directory to ~/.config/task/plugins/claude-profile-router/, copy +# config.example.env to config.env, list your profiles in it, and every task ty +# spawns picks its account automatically. See README.md. +name: claude-profile-router +version: 0.1.0 +description: Send each task to the Claude profile with the most usage headroom. + +hooks: + # task.route is the one hook ty waits on: it fires just before a task spawns + # and reads the script's stdout back as a decision. See docs/plugins.md. + task.route: route.sh + +actions: + - id: status + label: Show usage for each Claude profile + command: status.sh diff --git a/examples/plugins/claude-profile-router/route.sh b/examples/plugins/claude-profile-router/route.sh new file mode 100755 index 00000000..bd72ef74 --- /dev/null +++ b/examples/plugins/claude-profile-router/route.sh @@ -0,0 +1,102 @@ +#!/bin/bash +# task.route hook: pick the Claude profile with the most rate-limit headroom. +# +# ty runs this synchronously just before it spawns a task and reads stdout back +# as the decision, so stdout carries KEY=VALUE lines and nothing else — every +# diagnostic goes to stderr (which lands in the daemon log). +# +# CLAUDE_CONFIG_DIR= run this task under that profile +# HOLD=1 / REASON= every profile is spent; keep the task queued +# (no output) no opinion; ty spawns as already configured +# +# Printing nothing is always safe, so every failure path here does exactly that. +set -uo pipefail + +say() { echo "claude-profile-router: $*" >&2; } + +# config.env holds the settings, but anything already in the environment wins — +# that is what makes a one-off `TY_CLAUDE_MAX_PERCENT=50 ./route.sh` a usable way +# to try a threshold without editing the file. +_pre_profiles="${TY_CLAUDE_PROFILES:-}" +_pre_max="${TY_CLAUDE_MAX_PERCENT:-}" +_pre_projects="${TY_CLAUDE_PROJECTS:-}" +_pre_bin="${TY_BIN:-}" +if [[ -n "${TASK_PLUGIN_DIR:-}" && -f "$TASK_PLUGIN_DIR/config.env" ]]; then + # shellcheck disable=SC1091 + source "$TASK_PLUGIN_DIR/config.env" +fi +[[ -n "$_pre_profiles" ]] && TY_CLAUDE_PROFILES="$_pre_profiles" +[[ -n "$_pre_max" ]] && TY_CLAUDE_MAX_PERCENT="$_pre_max" +[[ -n "$_pre_projects" ]] && TY_CLAUDE_PROJECTS="$_pre_projects" +[[ -n "$_pre_bin" ]] && TY_BIN="$_pre_bin" + +TY="${TY_BIN:-ty}" +MAX_PERCENT="${TY_CLAUDE_MAX_PERCENT:-90}" + +if [[ -z "${TY_CLAUDE_PROFILES:-}" ]]; then + say "TY_CLAUDE_PROFILES not set (see config.example.env)" + exit 0 +fi + +if ! command -v "$TY" >/dev/null 2>&1; then + say "ty not found on PATH (set TY_BIN in config.env)" + exit 0 +fi + +# Optional project allowlist, so routing can be tried on one project first. +if [[ -n "${TY_CLAUDE_PROJECTS:-}" ]]; then + match="" + for p in $TY_CLAUDE_PROJECTS; do + [[ "$p" == "${TASK_PROJECT:-}" ]] && match=1 && break + done + [[ -z "$match" ]] && exit 0 +fi + +best_dir="" +best_pct="" +exhausted_low="" # lowest usage among profiles that were over the threshold + +for raw in $TY_CLAUDE_PROFILES; do + dir="${raw/#\~/$HOME}" + + # --percent prints one bare number: the binding limit's used percent. + if ! pct=$("$TY" usage --config-dir "$dir" --percent 2>/dev/null); then + say "skipping $dir (usage unavailable — expired login?)" + continue + fi + if [[ ! "$pct" =~ ^[0-9]+$ ]]; then + say "skipping $dir (unparseable usage: '$pct')" + continue + fi + + if (( pct >= MAX_PERCENT )); then + say "$dir at ${pct}% (>= ${MAX_PERCENT}%), skipping" + if [[ -z "$exhausted_low" ]] || (( pct < exhausted_low )); then + exhausted_low="$pct" + fi + continue + fi + + if [[ -z "$best_pct" ]] || (( pct < best_pct )); then + best_pct="$pct" + best_dir="$dir" + fi +done + +if [[ -n "$best_dir" ]]; then + say "routing to $best_dir (${best_pct}% used)" + echo "CLAUDE_CONFIG_DIR=$best_dir" + echo "REASON=${best_pct}% of its binding limit used" + exit 0 +fi + +# Nothing usable. Only hold the task if we actually saw an exhausted profile — +# if every probe merely failed, stay out of the way and let ty spawn normally +# rather than parking the whole board behind a broken credential lookup. +if [[ -n "$exhausted_low" ]]; then + echo "HOLD=1" + echo "REASON=every Claude profile is at or above ${MAX_PERCENT}% (best is ${exhausted_low}%)" + exit 0 +fi + +say "no profile could be evaluated; leaving this task alone" diff --git a/examples/plugins/claude-profile-router/status.sh b/examples/plugins/claude-profile-router/status.sh new file mode 100755 index 00000000..bd630a92 --- /dev/null +++ b/examples/plugins/claude-profile-router/status.sh @@ -0,0 +1,31 @@ +#!/bin/bash +# `ty plugins run claude-profile-router status` — show what the router sees. +# Same numbers route.sh decides on, so a surprising routing choice can be +# checked against reality without reading the daemon log. +set -uo pipefail + +_pre_profiles="${TY_CLAUDE_PROFILES:-}" +_pre_max="${TY_CLAUDE_MAX_PERCENT:-}" +_pre_bin="${TY_BIN:-}" +if [[ -n "${TASK_PLUGIN_DIR:-}" && -f "$TASK_PLUGIN_DIR/config.env" ]]; then + # shellcheck disable=SC1091 + source "$TASK_PLUGIN_DIR/config.env" +fi +[[ -n "$_pre_profiles" ]] && TY_CLAUDE_PROFILES="$_pre_profiles" +[[ -n "$_pre_max" ]] && TY_CLAUDE_MAX_PERCENT="$_pre_max" +[[ -n "$_pre_bin" ]] && TY_BIN="$_pre_bin" + +TY="${TY_BIN:-ty}" +MAX_PERCENT="${TY_CLAUDE_MAX_PERCENT:-90}" + +if [[ -z "${TY_CLAUDE_PROFILES:-}" ]]; then + echo "No profiles configured. Copy config.example.env to config.env and set TY_CLAUDE_PROFILES." + exit 0 +fi + +echo "Routing threshold: skip a profile at or above ${MAX_PERCENT}% used" +echo +for raw in $TY_CLAUDE_PROFILES; do + dir="${raw/#\~/$HOME}" + "$TY" usage --config-dir "$dir" || echo " (unavailable)" +done diff --git a/internal/claudeusage/cache.go b/internal/claudeusage/cache.go new file mode 100644 index 00000000..86ba14af --- /dev/null +++ b/internal/claudeusage/cache.go @@ -0,0 +1,109 @@ +package claudeusage + +import ( + "crypto/sha256" + "encoding/hex" + "encoding/json" + "os" + "path/filepath" + "time" +) + +// The usage endpoint rate-limits, and routing calls it on every spawn — a busy +// board would otherwise walk straight into a 429 and lose the very numbers it +// spawns on. So snapshots are cached on disk, briefly. +// +// On disk rather than in memory because the caller that matters is a shell +// script: a routing plugin shells out to `ty usage`, so every probe is a fresh +// process and an in-process cache would never be read. +const ( + // CacheTTL is how long a snapshot is served without asking the API again. + // Rate-limit utilization moves in percentage points over minutes; a minute + // of staleness cannot flip a routing decision that wasn't already marginal. + CacheTTL = 60 * time.Second + + // StaleTTL is how old a cached snapshot may be and still be used when a + // live fetch fails. Routing on a 10-minute-old number beats routing on + // nothing, which is what a transient 429 would otherwise leave us with. + StaleTTL = 30 * time.Minute +) + +// DefaultCacheDir is where snapshots are stored. +func DefaultCacheDir() string { + base, err := os.UserCacheDir() + if err != nil { + return "" + } + return filepath.Join(base, "ty", "claude-usage") +} + +func (c *Client) cacheDir() string { + if c.CacheDir != "" { + return c.CacheDir + } + return DefaultCacheDir() +} + +// cachePath names a profile's cache file by a hash of its config dir, so the +// path is flat and safe regardless of what the dir itself is called. +func (c *Client) cachePath(configDir string) string { + dir := c.cacheDir() + if dir == "" { + return "" + } + sum := sha256.Sum256([]byte(normalizeDir(configDir))) + return filepath.Join(dir, hex.EncodeToString(sum[:])[:16]+".json") +} + +// readCache returns a cached snapshot if one exists and is younger than maxAge. +// Any problem — no file, unreadable, corrupt — reads as a miss; a cache is never +// a reason to fail. +func (c *Client) readCache(configDir string, maxAge time.Duration) (*Snapshot, bool) { + if c.NoCache { + return nil, false + } + path := c.cachePath(configDir) + if path == "" { + return nil, false + } + data, err := os.ReadFile(path) //nolint:gosec // path is a hash under our own cache dir + if err != nil { + return nil, false + } + var snap Snapshot + if err := json.Unmarshal(data, &snap); err != nil { + return nil, false + } + if snap.FetchedAt.IsZero() || time.Since(snap.FetchedAt) > maxAge { + return nil, false + } + return &snap, true +} + +// writeCache stores a snapshot. Failures are ignored: a cache that can't be +// written costs an extra API call, nothing more. +func (c *Client) writeCache(configDir string, snap *Snapshot) { + if c.NoCache { + return + } + path := c.cachePath(configDir) + if path == "" { + return + } + if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil { + return + } + data, err := json.Marshal(snap) + if err != nil { + return + } + // Write-then-rename so a concurrent reader never sees a half-written file: + // several spawns can probe the same profile at once. + tmp := path + ".tmp" + if err := os.WriteFile(tmp, data, 0o600); err != nil { + return + } + if err := os.Rename(tmp, path); err != nil { + _ = os.Remove(tmp) + } +} diff --git a/internal/claudeusage/cache_test.go b/internal/claudeusage/cache_test.go new file mode 100644 index 00000000..47bd3dd4 --- /dev/null +++ b/internal/claudeusage/cache_test.go @@ -0,0 +1,195 @@ +package claudeusage + +import ( + "context" + "net/http" + "net/http/httptest" + "os" + "strings" + "sync/atomic" + "testing" + "time" +) + +// countingServer serves the usage fixture and reports how many times it was hit. +func countingServer(t *testing.T, hits *atomic.Int32) *httptest.Server { + t.Helper() + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + hits.Add(1) + _, _ = w.Write([]byte(usageBody)) + })) + t.Cleanup(srv.Close) + return srv +} + +func TestFetch_SecondCallIsServedFromCache(t *testing.T) { + // This is the whole reason the cache exists: routing probes every profile on + // every spawn, and the usage endpoint rate-limits. Without this, a busy + // board 429s itself out of the numbers it routes on. + var hits atomic.Int32 + srv := countingServer(t, &hits) + client := testClient(t, srv) + dir := writeCredsDir(t, "tok", time.Now().Add(time.Hour)) + + for i := 0; i < 3; i++ { + if _, err := client.Fetch(context.Background(), dir); err != nil { + t.Fatalf("Fetch %d: %v", i, err) + } + } + if got := hits.Load(); got != 1 { + t.Errorf("hit the API %d times across 3 fetches, want 1", got) + } +} + +func TestFetch_NoCacheAlwaysReadsLive(t *testing.T) { + var hits atomic.Int32 + srv := countingServer(t, &hits) + client := testClient(t, srv) + client.NoCache = true + dir := writeCredsDir(t, "tok", time.Now().Add(time.Hour)) + + for i := 0; i < 2; i++ { + if _, err := client.Fetch(context.Background(), dir); err != nil { + t.Fatalf("Fetch %d: %v", i, err) + } + } + if got := hits.Load(); got != 2 { + t.Errorf("hit the API %d times with NoCache, want 2", got) + } +} + +func TestFetch_CacheIsPerProfile(t *testing.T) { + // One profile's snapshot must never be served for another — that would route + // tasks to an account based on a different account's headroom. + var hits atomic.Int32 + srv := countingServer(t, &hits) + client := testClient(t, srv) + a := writeCredsDir(t, "tok-a", time.Now().Add(time.Hour)) + b := writeCredsDir(t, "tok-b", time.Now().Add(time.Hour)) + + snapA, err := client.Fetch(context.Background(), a) + if err != nil { + t.Fatal(err) + } + snapB, err := client.Fetch(context.Background(), b) + if err != nil { + t.Fatal(err) + } + if hits.Load() != 2 { + t.Errorf("hit the API %d times for 2 profiles, want 2", hits.Load()) + } + if snapA.ConfigDir == snapB.ConfigDir { + t.Errorf("both snapshots claim config dir %q", snapA.ConfigDir) + } +} + +func TestFetch_StaleCacheRescuesAFailedRequest(t *testing.T) { + // A 429 from the usage endpoint should not leave routing blind: a snapshot + // from a few minutes ago is a far better basis for a decision than none. + var fail atomic.Bool + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if fail.Load() { + w.WriteHeader(http.StatusTooManyRequests) + return + } + _, _ = w.Write([]byte(usageBody)) + })) + defer srv.Close() + + client := testClient(t, srv) + dir := writeCredsDir(t, "tok", time.Now().Add(time.Hour)) + + first, err := client.Fetch(context.Background(), dir) + if err != nil { + t.Fatalf("priming fetch: %v", err) + } + // Age the cached entry past CacheTTL but well inside StaleTTL. + first.FetchedAt = time.Now().Add(-5 * time.Minute) + client.writeCache(dir, first) + + fail.Store(true) + snap, err := client.Fetch(context.Background(), dir) + if err != nil { + t.Fatalf("Fetch should have fallen back to cache: %v", err) + } + if !snap.Stale { + t.Error("a cache-rescued snapshot must be marked stale") + } + if got := snap.UsedPercent(); got != 71 { + t.Errorf("UsedPercent = %v, want the cached 71", got) + } +} + +func TestFetch_CacheTooOldToRescue(t *testing.T) { + var fail atomic.Bool + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if fail.Load() { + w.WriteHeader(http.StatusTooManyRequests) + return + } + _, _ = w.Write([]byte(usageBody)) + })) + defer srv.Close() + + client := testClient(t, srv) + dir := writeCredsDir(t, "tok", time.Now().Add(time.Hour)) + + first, err := client.Fetch(context.Background(), dir) + if err != nil { + t.Fatal(err) + } + first.FetchedAt = time.Now().Add(-2 * StaleTTL) + client.writeCache(dir, first) + + fail.Store(true) + _, err = client.Fetch(context.Background(), dir) + if err == nil { + t.Fatal("a snapshot older than StaleTTL should not rescue a failed fetch") + } + if !strings.Contains(err.Error(), "429") { + t.Errorf("error should surface the 429, got %v", err) + } +} + +func TestFetch_ExpiredCredentialsIgnoreTheCache(t *testing.T) { + // A stale *credential* is a configuration problem, not a transient one. + // Serving a cached number would keep sending tasks to a profile that can no + // longer run them. + var hits atomic.Int32 + srv := countingServer(t, &hits) + client := testClient(t, srv) + dir := writeCredsDir(t, "tok", time.Now().Add(time.Hour)) + + if _, err := client.Fetch(context.Background(), dir); err != nil { + t.Fatal(err) + } + // Same profile, same warm cache entry — but its token has since expired. + writeCredsInto(t, dir, "tok", time.Now().Add(-time.Hour)) + if _, err := client.Fetch(context.Background(), dir); err == nil { + t.Error("expired credentials should error even with a warm cache") + } +} + +func TestCorruptCacheEntryIsAMiss(t *testing.T) { + var hits atomic.Int32 + srv := countingServer(t, &hits) + client := testClient(t, srv) + dir := writeCredsDir(t, "tok", time.Now().Add(time.Hour)) + + if _, err := client.Fetch(context.Background(), dir); err != nil { + t.Fatal(err) + } + if err := writeFile(client.cachePath(dir), "{not json"); err != nil { + t.Fatal(err) + } + if _, err := client.Fetch(context.Background(), dir); err != nil { + t.Fatalf("a corrupt cache entry must not fail the fetch: %v", err) + } + if hits.Load() != 2 { + t.Errorf("API hits = %d, want 2 (the corrupt entry should have been refetched)", hits.Load()) + } +} + +func writeFile(path, body string) error { + return os.WriteFile(path, []byte(body), 0o600) +} diff --git a/internal/claudeusage/credentials.go b/internal/claudeusage/credentials.go new file mode 100644 index 00000000..2434e46c --- /dev/null +++ b/internal/claudeusage/credentials.go @@ -0,0 +1,157 @@ +package claudeusage + +import ( + "crypto/sha256" + "encoding/hex" + "encoding/json" + "fmt" + "os" + "os/exec" + "path/filepath" + "runtime" + "strings" + "time" +) + +// keychainPrefix is the macOS Keychain service name Claude Code stores its +// OAuth credentials under. For a non-default CLAUDE_CONFIG_DIR the service is +// suffixed with a short hash of the config dir, so every profile gets its own +// entry (see KeychainService). +const keychainPrefix = "Claude Code-credentials" + +// credentialsFile is the on-disk credential store Claude Code uses where no OS +// keychain is available (Linux, containers). It lives inside the config dir. +const credentialsFile = ".credentials.json" + +// Credentials is the subset of a Claude Code credential blob we need: the OAuth +// access token used for api.anthropic.com/api/oauth/* calls, plus enough +// metadata to tell a stale token from a missing one. +type Credentials struct { + AccessToken string + ExpiresAt time.Time // zero when the store didn't record one + Subscription string +} + +// Expired reports whether the access token's recorded expiry has passed. A +// zero ExpiresAt is treated as not expired — we'd rather attempt the request +// and let the API decide than refuse on missing metadata. +func (c Credentials) Expired() bool { + return !c.ExpiresAt.IsZero() && time.Now().After(c.ExpiresAt) +} + +// credentialsBlob mirrors the JSON shape of a Claude Code credential store. We +// only decode the claudeAiOauth object; the sibling mcpOAuth map holds +// per-MCP-server tokens that are none of our business. +type credentialsBlob struct { + ClaudeAIOAuth struct { + AccessToken string `json:"accessToken"` + ExpiresAt int64 `json:"expiresAt"` // epoch milliseconds + SubscriptionType string `json:"subscriptionType"` + } `json:"claudeAiOauth"` +} + +// KeychainService returns the macOS Keychain service name holding the +// credentials for a given CLAUDE_CONFIG_DIR. +// +// Claude Code namespaces each config dir's credentials by appending the first 8 +// hex digits of the SHA-256 of the *absolute, cleaned* config-dir path to the +// base service name — that is what makes two logged-in profiles able to coexist +// on one machine. Reproducing the derivation here is what lets ty read a +// profile's usage without shelling out to `claude` (which has no usage command) +// or making the user paste a token anywhere. +func KeychainService(configDir string) string { + sum := sha256.Sum256([]byte(normalizeDir(configDir))) + return keychainPrefix + "-" + hex.EncodeToString(sum[:])[:8] +} + +// normalizeDir expands a leading ~ and cleans the path, so the hash we compute +// matches the one Claude Code computed from its own resolved config dir. +func normalizeDir(dir string) string { + dir = strings.TrimSpace(dir) + if strings.HasPrefix(dir, "~") { + if home, err := os.UserHomeDir(); err == nil { + dir = filepath.Join(home, dir[1:]) + } + } + return filepath.Clean(dir) +} + +// defaultConfigDir is ~/.claude — the config dir Claude Code uses when +// CLAUDE_CONFIG_DIR is unset. Credentials for it may still live under the +// unsuffixed keychain service from before per-profile namespacing existed. +func defaultConfigDir() string { + home, err := os.UserHomeDir() + if err != nil { + return "" + } + return filepath.Join(home, ".claude") +} + +// LoadCredentials finds the OAuth credentials for one Claude profile. +// +// Lookup order, first hit wins: +// 1. the per-config-dir macOS Keychain entry (the normal case on a Mac); +// 2. /.credentials.json (Linux, containers, and anywhere the +// keychain isn't used); +// 3. for the default ~/.claude only, the legacy unsuffixed keychain entry. +// +// A store that exists but holds an empty access token counts as a miss, so a +// half-migrated profile falls through to the next candidate instead of failing +// with a confusing 401 later. +func LoadCredentials(configDir string) (Credentials, error) { + dir := normalizeDir(configDir) + + if runtime.GOOS == "darwin" { + if c, ok := readKeychain(KeychainService(dir)); ok { + return c, nil + } + } + + if c, ok := readCredentialsFile(filepath.Join(dir, credentialsFile)); ok { + return c, nil + } + + if runtime.GOOS == "darwin" && dir == defaultConfigDir() { + if c, ok := readKeychain(keychainPrefix); ok { + return c, nil + } + } + + return Credentials{}, fmt.Errorf("no Claude credentials found for %s (log in once with CLAUDE_CONFIG_DIR=%s claude)", dir, dir) +} + +// readKeychain pulls a credential blob out of the macOS Keychain. A missing +// entry, a locked keychain, or an entry without an access token all report +// "not found" rather than an error: every one of them means "try the next +// candidate", and none is worth failing the whole lookup over. +func readKeychain(service string) (Credentials, bool) { + out, err := exec.Command("security", "find-generic-password", "-s", service, "-w").Output() + if err != nil { + return Credentials{}, false + } + return parseCredentials(out) +} + +func readCredentialsFile(path string) (Credentials, bool) { + data, err := os.ReadFile(path) //nolint:gosec // path is derived from the caller's own config dir + if err != nil { + return Credentials{}, false + } + return parseCredentials(data) +} + +func parseCredentials(data []byte) (Credentials, bool) { + var blob credentialsBlob + if err := json.Unmarshal(data, &blob); err != nil { + return Credentials{}, false + } + token := strings.TrimSpace(blob.ClaudeAIOAuth.AccessToken) + if token == "" { + return Credentials{}, false + } + c := Credentials{AccessToken: token, Subscription: blob.ClaudeAIOAuth.SubscriptionType} + if ms := blob.ClaudeAIOAuth.ExpiresAt; ms > 0 { + c.ExpiresAt = time.UnixMilli(ms) + } + return c, true +} diff --git a/internal/claudeusage/usage.go b/internal/claudeusage/usage.go new file mode 100644 index 00000000..c6aa4623 --- /dev/null +++ b/internal/claudeusage/usage.go @@ -0,0 +1,327 @@ +// Package claudeusage reads how much of a Claude subscription's rate limits a +// given Claude Code profile has burned through. +// +// A "profile" here is a CLAUDE_CONFIG_DIR: one logged-in Claude account with +// its own credentials. Someone with two accounts (say a personal one and a work +// one) keeps them in two config dirs and points ty at whichever they want a task +// to run under. This package answers the question that makes *choosing* between +// them possible — "how much headroom does each one have left?" — by reading the +// profile's stored OAuth token and calling the same usage endpoint Claude Code's +// own /usage command uses. +// +// It is strictly read-only: it never refreshes, rewrites, or otherwise touches +// stored credentials. A profile whose access token has gone stale simply reports +// an error, and the caller decides what to do about it. +package claudeusage + +import ( + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "strings" + "time" +) + +// DefaultBaseURL is the Anthropic API host serving the OAuth usage/profile +// endpoints. Overridable on Client for tests. +const DefaultBaseURL = "https://api.anthropic.com" + +// DefaultTimeout bounds a single usage probe. Routing a task waits on this, so +// it is deliberately short: a slow or unreachable API should fall back to the +// existing config rather than stall a spawn. +const DefaultTimeout = 10 * time.Second + +// Client fetches usage for Claude profiles. +type Client struct { + BaseURL string + HTTP *http.Client + // CacheDir overrides where snapshots are cached ("" = DefaultCacheDir). + CacheDir string + // NoCache bypasses the cache entirely, in both directions. + NoCache bool +} + +// NewClient returns a Client with the default endpoint and timeout. +func NewClient() *Client { + return &Client{BaseURL: DefaultBaseURL, HTTP: &http.Client{Timeout: DefaultTimeout}} +} + +// Limit is one rate-limit window the API reports for an account. +type Limit struct { + Kind string `json:"kind"` // "session", "weekly_all", "weekly_scoped", … + Group string `json:"group"` // "session" or "weekly" + Percent float64 `json:"percent"` // 0-100 of this window consumed + Severity string `json:"severity"` // provider's own label, e.g. "normal" + ResetsAt *time.Time `json:"resets_at,omitempty"` + Scope string `json:"scope,omitempty"` // model name for per-model windows +} + +// Snapshot is one profile's usage at a point in time. +type Snapshot struct { + ConfigDir string `json:"config_dir"` + Email string `json:"email,omitempty"` // filled in only when asked for; a second API call + Limits []Limit `json:"limits"` + FetchedAt time.Time `json:"fetched_at"` + // Stale marks a snapshot served from cache after a live fetch failed. The + // numbers are still worth routing on, but a caller showing them to a person + // should say so. + Stale bool `json:"stale,omitempty"` +} + +// Age is how long ago this snapshot was read from the API. +func (s Snapshot) Age() time.Duration { return time.Since(s.FetchedAt) } + +// UsedPercent is the profile's *binding* constraint: the highest utilization +// across every window the API reports. Routing cares about the window that will +// stop work first, not the average — a 5-hour session window at 98% blocks the +// next task even when the weekly window is nearly untouched. +func (s Snapshot) UsedPercent() float64 { + worst := 0.0 + for _, l := range s.Limits { + if l.Percent > worst { + worst = l.Percent + } + } + return worst +} + +// Headroom is how much of the binding window is still available, in percent. +// This is the number to route on: higher wins. +func (s Snapshot) Headroom() float64 { return 100 - s.UsedPercent() } + +// BindingLimit returns the window behind UsedPercent, for display and for +// telling the user *when* an exhausted profile frees up. +func (s Snapshot) BindingLimit() (Limit, bool) { + var worst Limit + found := false + for _, l := range s.Limits { + if !found || l.Percent > worst.Percent { + worst, found = l, true + } + } + return worst, found +} + +// apiUsage mirrors the /api/oauth/usage response. The endpoint carries a good +// deal more (dollar spend, extra-usage credits, unreleased window names); we +// decode only the parts that bear on "can this account do more work right now". +type apiUsage struct { + Limits []apiLimit `json:"limits"` + FiveHour *apiWindow `json:"five_hour"` + SevenDay *apiWindow `json:"seven_day"` +} + +type apiWindow struct { + Utilization float64 `json:"utilization"` + ResetsAt *time.Time `json:"resets_at"` +} + +type apiLimit struct { + Kind string `json:"kind"` + Group string `json:"group"` + Percent float64 `json:"percent"` + Severity string `json:"severity"` + ResetsAt *time.Time `json:"resets_at"` + Scope *struct { + Model *struct { + DisplayName string `json:"display_name"` + } `json:"model"` + } `json:"scope"` +} + +type apiProfile struct { + Account struct { + Email string `json:"email"` + } `json:"account"` +} + +// Fetch reads the usage for one profile (a CLAUDE_CONFIG_DIR), serving a recent +// cached snapshot when there is one. +// +// A missing or expired *credential* fails immediately — that is a configuration +// problem, and papering over it with a cached number would let a profile keep +// receiving tasks it can no longer run. A failed *request* is different: the API +// rate-limits, and a cached snapshot up to StaleTTL old is a far better basis +// for routing than nothing at all. +func (c *Client) Fetch(ctx context.Context, configDir string) (*Snapshot, error) { + creds, err := LoadCredentials(configDir) + if err != nil { + return nil, err + } + if creds.Expired() { + return nil, fmt.Errorf("credentials for %s expired at %s (run a claude session with CLAUDE_CONFIG_DIR=%s to refresh)", + normalizeDir(configDir), creds.ExpiresAt.Format(time.RFC3339), normalizeDir(configDir)) + } + + if snap, ok := c.readCache(configDir, CacheTTL); ok { + return snap, nil + } + + body, err := c.get(ctx, "/api/oauth/usage", creds.AccessToken) + if err != nil { + if stale, ok := c.readCache(configDir, StaleTTL); ok { + stale.Stale = true + return stale, nil + } + return nil, err + } + var raw apiUsage + if err := json.Unmarshal(body, &raw); err != nil { + return nil, fmt.Errorf("decode usage response: %w", err) + } + + snap := &Snapshot{ + ConfigDir: normalizeDir(configDir), + Limits: limitsFrom(raw), + FetchedAt: time.Now(), + } + c.writeCache(configDir, snap) + return snap, nil +} + +// FetchWithAccount is Fetch plus the account email, for surfaces that show the +// user *which* login a profile is. It costs a second round trip, so routing — +// which only needs the numbers — uses plain Fetch. A failure to resolve the +// email is not fatal: the usage numbers are the point. +func (c *Client) FetchWithAccount(ctx context.Context, configDir string) (*Snapshot, error) { + snap, err := c.Fetch(ctx, configDir) + if err != nil { + return nil, err + } + if snap.Email != "" { + return snap, nil // came back from cache with the email already on it + } + if email, err := c.Account(ctx, configDir); err == nil && email != "" { + snap.Email = email + // Re-cache so the next reader gets the email without a second request. + if !snap.Stale { + c.writeCache(configDir, snap) + } + } + return snap, nil +} + +// Account returns the email address a profile is logged in as. +func (c *Client) Account(ctx context.Context, configDir string) (string, error) { + creds, err := LoadCredentials(configDir) + if err != nil { + return "", err + } + body, err := c.get(ctx, "/api/oauth/profile", creds.AccessToken) + if err != nil { + return "", err + } + var p apiProfile + if err := json.Unmarshal(body, &p); err != nil { + return "", fmt.Errorf("decode profile response: %w", err) + } + return p.Account.Email, nil +} + +func (c *Client) get(ctx context.Context, path, token string) ([]byte, error) { + base := c.BaseURL + if base == "" { + base = DefaultBaseURL + } + req, err := http.NewRequestWithContext(ctx, http.MethodGet, strings.TrimRight(base, "/")+path, nil) + if err != nil { + return nil, err + } + req.Header.Set("Authorization", "Bearer "+token) + + httpClient := c.HTTP + if httpClient == nil { + httpClient = &http.Client{Timeout: DefaultTimeout} + } + resp, err := httpClient.Do(req) + if err != nil { + return nil, fmt.Errorf("%s: %w", path, err) + } + defer resp.Body.Close() //nolint:errcheck // read-only GET + body, err := io.ReadAll(io.LimitReader(resp.Body, 1<<20)) + if err != nil { + return nil, fmt.Errorf("%s: read body: %w", path, err) + } + if resp.StatusCode == http.StatusUnauthorized { + return nil, fmt.Errorf("%s: credentials rejected (401) — this profile needs a fresh login", path) + } + if resp.StatusCode == http.StatusTooManyRequests { + // The usage endpoint has its own rate limit, separate from the + // subscription limits it reports. Name it, so nobody reads this as the + // account being out of quota. + return nil, fmt.Errorf("%s: the usage API itself is rate-limiting (429) — this is not the account's quota; retry shortly", path) + } + if resp.StatusCode != http.StatusOK { + return nil, fmt.Errorf("%s: unexpected status %d: %s", path, resp.StatusCode, strings.TrimSpace(string(body))) + } + return body, nil +} + +// limitsFrom flattens the API's usage payload into our Limit list. The modern +// response carries a `limits` array; when it is absent or empty we synthesize +// the two headline windows from the older `five_hour`/`seven_day` objects, so a +// rollback on the API side doesn't leave routing blind. +func limitsFrom(raw apiUsage) []Limit { + if len(raw.Limits) > 0 { + out := make([]Limit, 0, len(raw.Limits)) + for _, l := range raw.Limits { + lim := Limit{ + Kind: l.Kind, + Group: l.Group, + Percent: l.Percent, + Severity: l.Severity, + ResetsAt: l.ResetsAt, + } + if l.Scope != nil && l.Scope.Model != nil { + lim.Scope = l.Scope.Model.DisplayName + } + out = append(out, lim) + } + return out + } + + var out []Limit + if raw.FiveHour != nil { + out = append(out, Limit{Kind: "session", Group: "session", Percent: raw.FiveHour.Utilization, ResetsAt: raw.FiveHour.ResetsAt}) + } + if raw.SevenDay != nil { + out = append(out, Limit{Kind: "weekly_all", Group: "weekly", Percent: raw.SevenDay.Utilization, ResetsAt: raw.SevenDay.ResetsAt}) + } + return out +} + +// Describe renders a one-line human summary of a snapshot, e.g. +// "94% used (session, resets 14:00) — 6% headroom". +func (s Snapshot) Describe() string { + var b strings.Builder + fmt.Fprintf(&b, "%.0f%% used", s.UsedPercent()) + if l, ok := s.BindingLimit(); ok { + b.WriteString(" (") + b.WriteString(limitLabel(l)) + if l.ResetsAt != nil { + fmt.Fprintf(&b, ", resets %s", l.ResetsAt.Local().Format("Mon 15:04")) + } + b.WriteString(")") + } + fmt.Fprintf(&b, " — %.0f%% headroom", s.Headroom()) + return b.String() +} + +// limitLabel turns an API window kind into something readable. +func limitLabel(l Limit) string { + label := l.Kind + switch l.Kind { + case "session": + label = "5-hour session" + case "weekly_all": + label = "weekly" + case "weekly_scoped": + label = "weekly" + if l.Scope != "" { + label = "weekly " + l.Scope + } + } + return label +} diff --git a/internal/claudeusage/usage_test.go b/internal/claudeusage/usage_test.go new file mode 100644 index 00000000..fcada876 --- /dev/null +++ b/internal/claudeusage/usage_test.go @@ -0,0 +1,261 @@ +package claudeusage + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strings" + "testing" + "time" +) + +// usageBody is a trimmed but structurally faithful /api/oauth/usage response. +const usageBody = `{ + "five_hour": {"utilization": 3.0, "resets_at": "2026-08-15T13:59:59.981513+00:00"}, + "seven_day": {"utilization": 7.0, "resets_at": "2026-08-20T20:59:59.981535+00:00"}, + "limits": [ + {"kind": "session", "group": "session", "percent": 3, "severity": "normal", + "resets_at": "2026-08-15T13:59:59.981513+00:00", "scope": null, "is_active": false}, + {"kind": "weekly_all", "group": "weekly", "percent": 71, "severity": "normal", + "resets_at": "2026-08-20T20:59:59.981535+00:00", "scope": null, "is_active": true}, + {"kind": "weekly_scoped", "group": "weekly", "percent": 12, "severity": "normal", + "resets_at": null, "scope": {"model": {"id": null, "display_name": "Opus"}}, "is_active": false} + ] +}` + +// writeCredsDir makes a config dir holding a .credentials.json, which is the +// non-macOS credential path and the one a test can exercise hermetically. +func writeCredsDir(t *testing.T, token string, expiresAt time.Time) string { + t.Helper() + dir := t.TempDir() + writeCredsInto(t, dir, token, expiresAt) + return dir +} + +// writeCredsInto (re)writes a credential blob into an existing config dir, so a +// test can age a profile's token without changing its identity. +func writeCredsInto(t *testing.T, dir, token string, expiresAt time.Time) { + t.Helper() + blob := map[string]any{ + "claudeAiOauth": map[string]any{ + "accessToken": token, + "subscriptionType": "max", + }, + } + if !expiresAt.IsZero() { + blob["claudeAiOauth"].(map[string]any)["expiresAt"] = expiresAt.UnixMilli() + } + data, err := json.Marshal(blob) + if err != nil { + t.Fatal(err) + } + if err := os.WriteFile(filepath.Join(dir, credentialsFile), data, 0o600); err != nil { + t.Fatal(err) + } +} + +// testClient points the client at a test server and an isolated cache dir, so +// no test can be helped (or hurt) by another test's cached snapshot, nor leave +// anything behind in the real user cache. +func testClient(t *testing.T, srv *httptest.Server) *Client { + t.Helper() + return &Client{BaseURL: srv.URL, HTTP: srv.Client(), CacheDir: t.TempDir()} +} + +func TestFetch_ParsesLimitsAndAuthenticates(t *testing.T) { + var gotAuth, gotPath string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotAuth, gotPath = r.Header.Get("Authorization"), r.URL.Path + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(usageBody)) + })) + defer srv.Close() + + dir := writeCredsDir(t, "tok-123", time.Now().Add(time.Hour)) + client := testClient(t, srv) + + snap, err := client.Fetch(context.Background(), dir) + if err != nil { + t.Fatalf("Fetch: %v", err) + } + if gotAuth != "Bearer tok-123" { + t.Errorf("Authorization = %q", gotAuth) + } + if gotPath != "/api/oauth/usage" { + t.Errorf("path = %q", gotPath) + } + if len(snap.Limits) != 3 { + t.Fatalf("got %d limits, want 3", len(snap.Limits)) + } + if snap.Limits[2].Scope != "Opus" { + t.Errorf("scoped limit scope = %q, want Opus", snap.Limits[2].Scope) + } + if snap.Limits[0].ResetsAt == nil { + t.Error("session limit lost its resets_at") + } +} + +func TestUsedPercentIsTheWorstWindow(t *testing.T) { + // The point of routing is to avoid the window that stops work first. An + // account 3% into its session but 71% into its week has 29% of headroom, + // not 97% — averaging or taking the session window alone would send tasks + // to an account about to run out. + snap := Snapshot{Limits: []Limit{ + {Kind: "session", Percent: 3}, + {Kind: "weekly_all", Percent: 71}, + {Kind: "weekly_scoped", Percent: 12}, + }} + if got := snap.UsedPercent(); got != 71 { + t.Errorf("UsedPercent = %v, want 71", got) + } + if got := snap.Headroom(); got != 29 { + t.Errorf("Headroom = %v, want 29", got) + } + binding, ok := snap.BindingLimit() + if !ok || binding.Kind != "weekly_all" { + t.Errorf("BindingLimit = %+v, ok=%v, want weekly_all", binding, ok) + } +} + +func TestUsedPercentOnEmptySnapshot(t *testing.T) { + var snap Snapshot + if got := snap.UsedPercent(); got != 0 { + t.Errorf("UsedPercent = %v, want 0", got) + } + if _, ok := snap.BindingLimit(); ok { + t.Error("BindingLimit should report not-found with no limits") + } +} + +func TestLimitsFallBackToLegacyWindows(t *testing.T) { + // If the API ever drops back to the older shape (no `limits` array), routing + // must keep working off five_hour/seven_day rather than seeing 0% used and + // happily piling work onto an exhausted account. + body := `{"five_hour": {"utilization": 88.0, "resets_at": null}, + "seven_day": {"utilization": 40.0, "resets_at": null}}` + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _, _ = w.Write([]byte(body)) + })) + defer srv.Close() + + dir := writeCredsDir(t, "tok", time.Time{}) + client := testClient(t, srv) + snap, err := client.Fetch(context.Background(), dir) + if err != nil { + t.Fatalf("Fetch: %v", err) + } + if len(snap.Limits) != 2 { + t.Fatalf("got %d limits, want 2 synthesized", len(snap.Limits)) + } + if got := snap.UsedPercent(); got != 88 { + t.Errorf("UsedPercent = %v, want 88", got) + } +} + +func TestFetch_ExpiredTokenIsNotSentToTheAPI(t *testing.T) { + called := false + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + called = true + })) + defer srv.Close() + + dir := writeCredsDir(t, "stale", time.Now().Add(-time.Hour)) + client := testClient(t, srv) + _, err := client.Fetch(context.Background(), dir) + if err == nil { + t.Fatal("expected an error for an expired token") + } + if called { + t.Error("expired token should not reach the API") + } + if !strings.Contains(err.Error(), "expired") { + t.Errorf("error should say the credentials expired: %v", err) + } +} + +func TestFetch_UnauthorizedIsExplainedNotDumped(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusUnauthorized) + _, _ = w.Write([]byte(`{"type":"error","error":{"type":"authentication_error"}}`)) + })) + defer srv.Close() + + dir := writeCredsDir(t, "tok", time.Time{}) + client := testClient(t, srv) + _, err := client.Fetch(context.Background(), dir) + if err == nil || !strings.Contains(err.Error(), "fresh login") { + t.Errorf("401 should point at re-login, got %v", err) + } +} + +func TestFetch_MissingCredentialsErrors(t *testing.T) { + client := &Client{BaseURL: "http://127.0.0.1:1", HTTP: &http.Client{Timeout: time.Second}, CacheDir: t.TempDir()} + _, err := client.Fetch(context.Background(), t.TempDir()) + if err == nil || !strings.Contains(err.Error(), "no Claude credentials") { + t.Errorf("want a missing-credentials error, got %v", err) + } +} + +func TestLoadCredentials_EmptyTokenCountsAsMissing(t *testing.T) { + // A half-migrated profile leaves a credential blob with an empty token. + // Treating it as found would surface as a baffling 401 later. + dir := writeCredsDir(t, "", time.Time{}) + if _, err := LoadCredentials(dir); err == nil { + t.Error("an empty access token should not count as credentials") + } +} + +func TestKeychainServiceDerivation(t *testing.T) { + // Claude Code namespaces each config dir's keychain entry by the first 8 hex + // digits of the SHA-256 of the absolute path. Pinning a known value here is + // what catches the derivation drifting: get it wrong and every profile + // silently reports "no credentials" on a Mac. + if got, want := KeychainService("/Users/bruno/.claude-ik"), "Claude Code-credentials-eaf7266a"; got != want { + t.Errorf("KeychainService = %q, want %q", got, want) + } + if got, want := KeychainService("/Users/bruno/.claude"), "Claude Code-credentials-5561fe67"; got != want { + t.Errorf("KeychainService = %q, want %q", got, want) + } + // A trailing slash or a redundant segment is the same dir, so it must hash + // the same — Claude Code hashed its own cleaned path. + if KeychainService("/Users/bruno/.claude-ik/") != KeychainService("/Users/bruno/.claude-ik") { + t.Error("trailing slash changed the keychain service name") + } + if KeychainService("/Users/bruno/foo/../.claude-ik") != KeychainService("/Users/bruno/.claude-ik") { + t.Error("uncleaned path changed the keychain service name") + } +} + +func TestKeychainServiceExpandsTilde(t *testing.T) { + home, err := os.UserHomeDir() + if err != nil { + t.Skip("no home dir") + } + if KeychainService("~/.claude-x") != KeychainService(filepath.Join(home, ".claude-x")) { + t.Error("~ was not expanded before hashing") + } +} + +func TestDescribe(t *testing.T) { + reset := time.Date(2026, 8, 20, 20, 0, 0, 0, time.UTC) + snap := Snapshot{Limits: []Limit{ + {Kind: "session", Percent: 3}, + {Kind: "weekly_all", Percent: 71, ResetsAt: &reset}, + }} + got := snap.Describe() + for _, want := range []string{"71% used", "weekly", "29% headroom", "resets"} { + if !strings.Contains(got, want) { + t.Errorf("Describe() = %q, missing %q", got, want) + } + } +} + +func TestDescribeNamesTheScopedModel(t *testing.T) { + snap := Snapshot{Limits: []Limit{{Kind: "weekly_scoped", Percent: 95, Scope: "Opus"}}} + if got := snap.Describe(); !strings.Contains(got, "weekly Opus") { + t.Errorf("Describe() = %q, want the model named", got) + } +} diff --git a/internal/db/tasks.go b/internal/db/tasks.go index 74c23c0e..a312e9a4 100644 --- a/internal/db/tasks.go +++ b/internal/db/tasks.go @@ -937,6 +937,23 @@ func (db *DB) UpdateTaskPermissionMode(taskID int64, mode string) error { return nil } +// UpdateTaskClaudeConfigDir sets the per-task CLAUDE_CONFIG_DIR override, +// which is how a task is pinned to one Claude profile (account). Writing it as +// its own column update — rather than through UpdateTask — matters at spawn +// time: the routing decision is made from a task struct the daemon has been +// holding, and a full-row write would stomp any field another surface (the TUI, +// a hook) changed in the meantime. +func (db *DB) UpdateTaskClaudeConfigDir(taskID int64, configDir string) error { + _, err := db.Exec(` + UPDATE tasks SET claude_config_dir = ?, updated_at = CURRENT_TIMESTAMP + WHERE id = ? + `, configDir, taskID) + if err != nil { + return fmt.Errorf("update task claude config dir: %w", err) + } + return nil +} + // UpdateTaskPinned updates only the pinned flag for a task. func (db *DB) UpdateTaskPinned(taskID int64, pinned bool) error { _, err := db.Exec(` diff --git a/internal/executor/executor.go b/internal/executor/executor.go index 2cb0138d..d19e17b5 100644 --- a/internal/executor/executor.go +++ b/internal/executor/executor.go @@ -1755,6 +1755,15 @@ func (e *Executor) processNextTask(ctx context.Context) { continue } + // Last decision before the spawn: which Claude profile does this run + // under? A routing plugin may pick one (stamping task.ClaudeConfigDir, + // which both command builders already honor) or ask to hold the task + // when every account is out of headroom. With no router installed this + // is a no-op. See routing.go. + if !e.routeTask(ctx, task, true) { + continue + } + // Atomically check-and-set to prevent race where two ticks // both see the task as not-running and spawn duplicate goroutines e.mu.Lock() @@ -1818,6 +1827,11 @@ func (e *Executor) ExecuteNow(ctx context.Context, taskID int64) error { e.runningTasks[taskID] = true e.mu.Unlock() + // Route this run to a Claude profile too, so a manually started task lands + // on the same account the queue would have chosen. A hold is not honored + // here: the user asked for this task to run now. + e.routeTask(ctx, task, false) + e.executeTask(ctx, task) return nil } diff --git a/internal/executor/routing.go b/internal/executor/routing.go new file mode 100644 index 00000000..33947d0e --- /dev/null +++ b/internal/executor/routing.go @@ -0,0 +1,113 @@ +package executor + +import ( + "context" + "fmt" + "strings" + "sync" + + "github.com/bborn/workflow/internal/db" + "github.com/bborn/workflow/internal/hooks" +) + +// Profile routing gives a plugin the last word on which Claude account a task +// runs under, at the only moment where that word is still worth anything: after +// the task is cleared to run, before its command is built. +// +// Everything downstream already supports this — Task.ClaudeConfigDir has always +// been the per-task profile lever, honored identically by the daemon's command +// builder and the TUI's. What was missing was anyone to set it automatically. +// Routing fills that in: it stamps the column and lets the existing machinery +// carry the decision the rest of the way, so there is no second code path for a +// routed task and no chance of the two builders disagreeing about which profile +// is in play. +// +// Two rules keep it from getting in the way: +// +// - An explicit choice always wins. A task that already names a config dir — +// set by hand, by a workflow step, or by an earlier routing pass — is left +// alone. Routing fills a vacuum; it does not overrule a person. +// - Silence means "carry on". No router installed, a script that fails, times +// out, or prints nothing: the task spawns exactly as it would have before +// any of this existed. + +// routeHoldLog remembers the last hold reason logged per task, so a task parked +// behind exhausted profiles writes one log line rather than one per daemon tick. +var routeHoldLog sync.Map // taskID -> last reason written + +// routeTask consults the task.route plugin hook and applies its decision. +// +// It returns false only when a router asked to hold the task — every other +// outcome, including every kind of failure, returns true and lets the spawn +// proceed. allowHold is false on the manual "run this now" path: a person who +// explicitly started a task has already made the call, and silently refusing +// would look like the button was broken. +func (e *Executor) routeTask(ctx context.Context, task *db.Task, allowHold bool) bool { + if task == nil || e.hooks == nil { + return true + } + // CLAUDE_CONFIG_DIR is a Claude concept; a codex or gemini task has no + // profile to route between. + if task.Executor != "" && task.Executor != db.ExecutorClaude { + return true + } + if strings.TrimSpace(task.ClaudeConfigDir) != "" { + return true + } + if !e.hooks.HandlesRoute() { + return true + } + + decision := e.hooks.Route(ctx, task) + if decision.Empty() { + routeHoldLog.Delete(task.ID) + return true + } + + if decision.Hold && allowHold { + e.noteRouteHold(task, decision) + return false + } + routeHoldLog.Delete(task.ID) + + dir := strings.TrimSpace(decision.ClaudeConfigDir) + if dir == "" { + return true + } + resolved := ResolveClaudeConfigDir(dir) + if err := e.db.UpdateTaskClaudeConfigDir(task.ID, resolved); err != nil { + // The write is what makes the decision visible to the TUI and to a + // later resume. If it fails, don't apply the route in memory either — + // a task whose spawned profile disagrees with its recorded one is the + // exact confusion this feature is supposed to remove. + e.logger.Error("Failed to record routed Claude profile", "id", task.ID, "dir", resolved, "error", err) + return true + } + task.ClaudeConfigDir = resolved + + msg := fmt.Sprintf("Routed to Claude profile %s (by plugin %q)", resolved, decision.Plugin) + if decision.Reason != "" { + msg += ": " + decision.Reason + } + e.logger.Info("Routed task to Claude profile", "id", task.ID, "dir", resolved, "plugin", decision.Plugin) + e.logLine(task.ID, "system", msg) + return true +} + +// noteRouteHold records a hold, writing to the task log only when the reason +// changes. The daemon reconsiders a queued task every tick, so an unconditional +// log line would bury the task's real history under thousands of repeats of +// "waiting for headroom". +func (e *Executor) noteRouteHold(task *db.Task, decision hooks.RouteDecision) { + reason := decision.Reason + if reason == "" { + reason = "no Claude profile has headroom right now" + } + e.logger.Info("Holding task: no Claude profile available", "id", task.ID, "plugin", decision.Plugin, "reason", reason) + + if prev, ok := routeHoldLog.Load(task.ID); ok && prev == reason { + return + } + routeHoldLog.Store(task.ID, reason) + e.logLine(task.ID, "system", fmt.Sprintf("Waiting to start — %s (plugin %q). Will retry automatically.", reason, decision.Plugin)) +} diff --git a/internal/executor/routing_test.go b/internal/executor/routing_test.go new file mode 100644 index 00000000..eb0ab290 --- /dev/null +++ b/internal/executor/routing_test.go @@ -0,0 +1,277 @@ +package executor + +import ( + "context" + "os" + "path/filepath" + "strings" + "testing" + + "github.com/bborn/workflow/internal/config" + "github.com/bborn/workflow/internal/db" +) + +// newRoutingExecutor builds an Executor whose plugins come from a temp dir, so a +// routing test exercises the real hook path (subprocess, stdout parsing, DB +// write) without depending on what is installed on the machine. +func newRoutingExecutor(t *testing.T, routeScript string) (*Executor, *db.DB) { + t.Helper() + + pluginsDir := t.TempDir() + if routeScript != "" { + dir := filepath.Join(pluginsDir, "router") + if err := os.MkdirAll(dir, 0o755); err != nil { + t.Fatal(err) + } + manifest := "name: router\nhooks:\n task.route: route.sh\n" + if err := os.WriteFile(filepath.Join(dir, "plugin.yaml"), []byte(manifest), 0o644); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(filepath.Join(dir, "route.sh"), []byte(routeScript), 0o755); err != nil { + t.Fatal(err) + } + } + t.Setenv("TY_PLUGINS_DIR", pluginsDir) + + tmpFile, err := os.CreateTemp("", "test-routing-*.db") + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { os.Remove(tmpFile.Name()) }) + tmpFile.Close() + + database, err := db.Open(tmpFile.Name()) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { database.Close() }) + if err := database.CreateProject(&db.Project{Name: "test", Path: "/tmp/test"}); err != nil { + t.Fatal(err) + } + + // The hold-log memo is package state; keep tests from leaking into each other. + t.Cleanup(func() { routeHoldLog.Range(func(k, _ any) bool { routeHoldLog.Delete(k); return true }) }) + + return New(database, &config.Config{}), database +} + +func newRoutingTask(t *testing.T, database *db.DB, executorName string) *db.Task { + t.Helper() + task := &db.Task{Title: "route me", Type: "task", Project: "test", Executor: executorName} + if err := database.CreateTask(task); err != nil { + t.Fatal(err) + } + return task +} + +func TestRouteTask_AppliesAndPersistsConfigDir(t *testing.T) { + e, database := newRoutingExecutor(t, "#!/bin/sh\necho CLAUDE_CONFIG_DIR=/tmp/claude-work\necho REASON=12% used\n") + task := newRoutingTask(t, database, db.ExecutorClaude) + + if ok := e.routeTask(context.Background(), task, true); !ok { + t.Fatal("routeTask returned false, want the task cleared to run") + } + if task.ClaudeConfigDir != "/tmp/claude-work" { + t.Errorf("in-memory ClaudeConfigDir = %q", task.ClaudeConfigDir) + } + + // The write matters as much as the in-memory value: the TUI and any later + // resume read the column, and a disagreement there is exactly the confusion + // routing is meant to remove. + reloaded, err := database.GetTask(task.ID) + if err != nil { + t.Fatal(err) + } + if reloaded.ClaudeConfigDir != "/tmp/claude-work" { + t.Errorf("persisted ClaudeConfigDir = %q", reloaded.ClaudeConfigDir) + } + + logs, err := database.GetTaskLogs(task.ID, 10) + if err != nil { + t.Fatal(err) + } + found := false + for _, l := range logs { + if strings.Contains(l.Content, "/tmp/claude-work") && strings.Contains(l.Content, "12% used") { + found = true + } + } + if !found { + t.Errorf("routing decision was not written to the task log: %+v", logs) + } +} + +func TestRouteTask_ExplicitConfigDirIsNotOverridden(t *testing.T) { + // A dir chosen by a person or a workflow step is a decision, not a default. + e, database := newRoutingExecutor(t, "#!/bin/sh\necho CLAUDE_CONFIG_DIR=/tmp/router-choice\n") + task := newRoutingTask(t, database, db.ExecutorClaude) + task.ClaudeConfigDir = "/tmp/chosen-by-hand" + + if ok := e.routeTask(context.Background(), task, true); !ok { + t.Fatal("routeTask returned false") + } + if task.ClaudeConfigDir != "/tmp/chosen-by-hand" { + t.Errorf("router overrode an explicit config dir: %q", task.ClaudeConfigDir) + } +} + +func TestRouteTask_ResumedTaskStaysOnItsProfile(t *testing.T) { + // Session affinity, and it is not optional: a Claude session lives inside + // one config dir, so a task resumed under a different profile would find no + // session to resume and silently start a fresh conversation. Once a task has + // been routed, the stamped dir must pin it for the rest of its life — even + // when the router would now prefer somewhere else. + e, database := newRoutingExecutor(t, "#!/bin/sh\necho CLAUDE_CONFIG_DIR=/tmp/now-emptier\n") + task := newRoutingTask(t, database, db.ExecutorClaude) + + // First spawn: the router picks a profile and it is recorded. + if ok := e.routeTask(context.Background(), task, true); !ok { + t.Fatal("routeTask returned false") + } + first := task.ClaudeConfigDir + if first == "" { + t.Fatal("first spawn was not routed") + } + task.ClaudeSessionID = "sess-abc" + + // Second pass (a resume after the task was blocked, say) must not move it. + if ok := e.routeTask(context.Background(), task, true); !ok { + t.Fatal("routeTask returned false on resume") + } + if task.ClaudeConfigDir != first { + t.Errorf("resumed task moved profiles: %q -> %q", first, task.ClaudeConfigDir) + } +} + +func TestRouteTask_NonClaudeExecutorIsUntouched(t *testing.T) { + // CLAUDE_CONFIG_DIR means nothing to codex; setting it would be noise at + // best and a misleading task log at worst. + e, database := newRoutingExecutor(t, "#!/bin/sh\necho CLAUDE_CONFIG_DIR=/tmp/nope\n") + task := newRoutingTask(t, database, db.ExecutorCodex) + + if ok := e.routeTask(context.Background(), task, true); !ok { + t.Fatal("routeTask returned false") + } + if task.ClaudeConfigDir != "" { + t.Errorf("codex task was routed: %q", task.ClaudeConfigDir) + } +} + +func TestRouteTask_HoldKeepsTaskQueued(t *testing.T) { + e, database := newRoutingExecutor(t, "#!/bin/sh\necho HOLD=1\necho 'REASON=every profile above 90%'\n") + task := newRoutingTask(t, database, db.ExecutorClaude) + if err := database.UpdateTaskStatus(task.ID, db.StatusQueued); err != nil { + t.Fatal(err) + } + + if ok := e.routeTask(context.Background(), task, true); ok { + t.Fatal("routeTask returned true, want the spawn held") + } + + // A held task must stay queued: parking it as blocked would take a human to + // undo, when the whole point is that it starts by itself once limits reset. + reloaded, err := database.GetTask(task.ID) + if err != nil { + t.Fatal(err) + } + if reloaded.Status != db.StatusQueued { + t.Errorf("status = %q, want it left queued", reloaded.Status) + } +} + +func TestRouteTask_RepeatedHoldLogsOnce(t *testing.T) { + // The daemon reconsiders a queued task every tick. Logging each refusal + // would bury the task's real history under thousands of identical lines. + e, database := newRoutingExecutor(t, "#!/bin/sh\necho HOLD=1\necho 'REASON=every profile above 90%'\n") + task := newRoutingTask(t, database, db.ExecutorClaude) + + for i := 0; i < 3; i++ { + if ok := e.routeTask(context.Background(), task, true); ok { + t.Fatal("routeTask returned true, want held") + } + } + + logs, err := database.GetTaskLogs(task.ID, 50) + if err != nil { + t.Fatal(err) + } + holds := 0 + for _, l := range logs { + if strings.Contains(l.Content, "Waiting to start") { + holds++ + } + } + if holds != 1 { + t.Errorf("wrote %d hold log lines across 3 ticks, want 1", holds) + } +} + +func TestRouteTask_HoldIgnoredOnManualRun(t *testing.T) { + // `ty run` / "start now" is an explicit instruction. Silently refusing it + // would read as a broken button. + e, database := newRoutingExecutor(t, "#!/bin/sh\necho HOLD=1\n") + task := newRoutingTask(t, database, db.ExecutorClaude) + + if ok := e.routeTask(context.Background(), task, false); !ok { + t.Error("a manual run should not be held") + } +} + +func TestRouteTask_NoRouterIsANoOp(t *testing.T) { + e, database := newRoutingExecutor(t, "") + task := newRoutingTask(t, database, db.ExecutorClaude) + + if ok := e.routeTask(context.Background(), task, true); !ok { + t.Fatal("routeTask returned false with no router installed") + } + if task.ClaudeConfigDir != "" { + t.Errorf("ClaudeConfigDir = %q, want untouched", task.ClaudeConfigDir) + } +} + +func TestRouteTask_FailingRouterStillSpawns(t *testing.T) { + // Routing is an optimization. Failing to optimize must never be why a task + // doesn't run. + e, database := newRoutingExecutor(t, "#!/bin/sh\necho boom >&2\nexit 1\n") + task := newRoutingTask(t, database, db.ExecutorClaude) + + if ok := e.routeTask(context.Background(), task, true); !ok { + t.Error("a failing router should not block the spawn") + } + if task.ClaudeConfigDir != "" { + t.Errorf("ClaudeConfigDir = %q, want untouched", task.ClaudeConfigDir) + } +} + +func TestRouteTask_ExpandsTildeInRoutedDir(t *testing.T) { + // A router written in shell may well emit a literal ~; the stored value has + // to be the resolved path, since it is spliced straight into the spawn + // command as CLAUDE_CONFIG_DIR="…". + e, database := newRoutingExecutor(t, "#!/bin/sh\necho 'CLAUDE_CONFIG_DIR=~/.claude-work'\n") + task := newRoutingTask(t, database, db.ExecutorClaude) + + if ok := e.routeTask(context.Background(), task, true); !ok { + t.Fatal("routeTask returned false") + } + home, err := os.UserHomeDir() + if err != nil { + t.Skip("no home dir") + } + want := filepath.Join(home, ".claude-work") + if task.ClaudeConfigDir != want { + t.Errorf("ClaudeConfigDir = %q, want %q", task.ClaudeConfigDir, want) + } +} + +func TestRouteTask_EmptyExecutorIsTreatedAsClaude(t *testing.T) { + // Older tasks carry no executor; claude is the default, so they should route. + e, database := newRoutingExecutor(t, "#!/bin/sh\necho CLAUDE_CONFIG_DIR=/tmp/x\n") + task := newRoutingTask(t, database, "") + + if ok := e.routeTask(context.Background(), task, true); !ok { + t.Fatal("routeTask returned false") + } + if task.ClaudeConfigDir != "/tmp/x" { + t.Errorf("ClaudeConfigDir = %q, want /tmp/x", task.ClaudeConfigDir) + } +} diff --git a/internal/hooks/route.go b/internal/hooks/route.go new file mode 100644 index 00000000..2f40661b --- /dev/null +++ b/internal/hooks/route.go @@ -0,0 +1,184 @@ +package hooks + +import ( + "bufio" + "bytes" + "context" + "fmt" + "os/exec" + "strings" + "time" + + "github.com/bborn/workflow/internal/db" +) + +// EventTaskRoute is fired immediately before a task is spawned, and is the one +// event a plugin can *answer* rather than merely observe. +// +// Every other hook is a notification: it fires after the fact, runs detached, +// and nothing waits for it. A router has to be the opposite — the decision it +// makes (which Claude profile this task runs under) is only useful before the +// command is built, so this hook runs synchronously, in the foreground of the +// spawn, and its stdout is read back. +// +// That inversion is deliberate, and bounded: RouteTimeout caps the wait, a +// failing or silent script yields no decision and the task spawns exactly as it +// would have, and the first plugin to answer wins so a slow one can't be made to +// re-litigate a settled choice. +const EventTaskRoute = "task.route" + +// RouteTimeout bounds a routing hook. A task spawn blocks on this, so it is +// tight: a router that needs longer than this to pick a profile is a router that +// should be caching, and the safe answer while it does is "spawn as configured". +const RouteTimeout = 15 * time.Second + +// RouteDecision is what a routing hook answers with. The zero value means "no +// opinion" — the caller proceeds with the task's existing configuration. +type RouteDecision struct { + // Plugin is the name of the plugin that answered, for logging. + Plugin string + // ClaudeConfigDir routes the task to a particular Claude profile. Empty + // leaves the task's existing (project or per-task) config dir alone. + ClaudeConfigDir string + // Hold asks the caller not to start this task yet — every candidate profile + // is out of headroom, and running now would only burn a session on a 429. + // The task stays queued and is reconsidered on the next tick. + Hold bool + // Reason explains a Hold (or annotates a routing choice) for the task log. + Reason string +} + +// Empty reports whether the decision carries no instruction at all. +func (d RouteDecision) Empty() bool { + return !d.Hold && strings.TrimSpace(d.ClaudeConfigDir) == "" +} + +// HandlesRoute reports whether any loaded plugin declares a task.route hook. +// Spawn checks this first so the overwhelmingly common case — nobody has +// installed a router — costs a slice scan instead of a subprocess. +func (r *Runner) HandlesRoute() bool { + for _, p := range r.plugins { + if _, ok := p.ScriptFor(EventTaskRoute); ok { + return true + } + } + return false +} + +// Route asks every plugin that handles task.route what to do with this task, +// in plugin-name order, and returns the first non-empty decision. +// +// Plugins are consulted in order rather than in parallel and the first answer +// stands, which keeps the outcome deterministic when more than one router is +// installed — the alternative (merging or last-write-wins) makes the effective +// policy depend on which script happened to finish first. +// +// A hook that errors, times out, or prints nothing usable is skipped: routing +// is an optimization, and failing to optimize must never be the reason a task +// doesn't run. +func (r *Runner) Route(ctx context.Context, task *db.Task) RouteDecision { + if task == nil { + return RouteDecision{} + } + for _, p := range r.plugins { + script, ok := p.ScriptFor(EventTaskRoute) + if !ok { + continue + } + env := append(taskEnv(EventTaskRoute, task, ""), + "TASK_PLUGIN_NAME="+p.Name, + "TASK_PLUGIN_DIR="+p.Dir, + "TASK_EXECUTOR="+task.Executor, + "TASK_CLAUDE_CONFIG_DIR="+task.ClaudeConfigDir, + ) + + decision, err := runRouteScript(ctx, script, p.Dir, env) + if err != nil { + r.logger.Warn("route hook failed", "plugin", p.Name, "task", task.ID, "error", err) + continue + } + if decision.Empty() { + continue + } + decision.Plugin = p.Name + return decision + } + return RouteDecision{} +} + +// runRouteScript executes one routing script and parses its verdict. +func runRouteScript(ctx context.Context, script, workDir string, env []string) (RouteDecision, error) { + ctx, cancel := context.WithTimeout(ctx, RouteTimeout) + defer cancel() + + cmd := exec.CommandContext(ctx, script) + cmd.Dir = workDir + cmd.Env = env + + // stdout is the decision channel and stderr is free for the script to log + // on, so they are captured separately — otherwise an `echo "checking..." >&2` + // in a router would be parsed as part of its answer. + var stdout, stderr bytes.Buffer + cmd.Stdout = &stdout + cmd.Stderr = &stderr + + if err := cmd.Run(); err != nil { + return RouteDecision{}, fmt.Errorf("%w: %s", err, strings.TrimSpace(stderr.String())) + } + return ParseRouteOutput(stdout.String()), nil +} + +// ParseRouteOutput reads a routing script's stdout. +// +// The format is deliberately the dullest thing that works — KEY=VALUE, one per +// line, unknown keys ignored — because the scripts writing it are shell. A +// router shouldn't need a JSON encoder to say "use this directory". +// +// CLAUDE_CONFIG_DIR=/Users/me/.claude-work +// HOLD=1 +// REASON=both profiles above 90%, next reset 14:00 +// +// Values are taken literally after the first '='; surrounding quotes are +// stripped so `CLAUDE_CONFIG_DIR="$dir"` from a script that quoted its output +// still parses. +func ParseRouteOutput(out string) RouteDecision { + var d RouteDecision + scanner := bufio.NewScanner(strings.NewReader(out)) + for scanner.Scan() { + line := strings.TrimSpace(scanner.Text()) + if line == "" || strings.HasPrefix(line, "#") { + continue + } + key, value, found := strings.Cut(line, "=") + if !found { + continue + } + value = unquote(strings.TrimSpace(value)) + switch strings.ToUpper(strings.TrimSpace(key)) { + case "CLAUDE_CONFIG_DIR": + d.ClaudeConfigDir = value + case "HOLD", "DEFER": + d.Hold = isTruthy(value) + case "REASON": + d.Reason = value + } + } + return d +} + +func unquote(s string) string { + if len(s) >= 2 { + if (s[0] == '"' && s[len(s)-1] == '"') || (s[0] == '\'' && s[len(s)-1] == '\'') { + return s[1 : len(s)-1] + } + } + return s +} + +func isTruthy(s string) bool { + switch strings.ToLower(strings.TrimSpace(s)) { + case "1", "true", "yes", "y", "on": + return true + } + return false +} diff --git a/internal/hooks/route_test.go b/internal/hooks/route_test.go new file mode 100644 index 00000000..11e64385 --- /dev/null +++ b/internal/hooks/route_test.go @@ -0,0 +1,205 @@ +package hooks + +import ( + "context" + "os" + "testing" + + "github.com/charmbracelet/log" + + "github.com/bborn/workflow/internal/db" +) + +func routeRunner(t *testing.T, root string) *Runner { + t.Helper() + return newRunner("", root, log.NewWithOptions(os.Stderr, log.Options{Level: log.FatalLevel})) +} + +func TestParseRouteOutput(t *testing.T) { + tests := []struct { + name string + out string + want RouteDecision + }{ + { + name: "config dir", + out: "CLAUDE_CONFIG_DIR=/home/me/.claude-work\n", + want: RouteDecision{ClaudeConfigDir: "/home/me/.claude-work"}, + }, + { + name: "quoted value", + out: "CLAUDE_CONFIG_DIR=\"/home/me/my claude\"\n", + want: RouteDecision{ClaudeConfigDir: "/home/me/my claude"}, + }, + { + name: "hold with reason", + out: "HOLD=1\nREASON=all profiles above 90%\n", + want: RouteDecision{Hold: true, Reason: "all profiles above 90%"}, + }, + { + name: "defer is an alias for hold", + out: "DEFER=true\n", + want: RouteDecision{Hold: true}, + }, + { + name: "hold=0 is not a hold", + out: "HOLD=0\nCLAUDE_CONFIG_DIR=/a\n", + want: RouteDecision{ClaudeConfigDir: "/a"}, + }, + { + // A router's stdout is decision-only, but scripts still leak the odd + // line. Anything unrecognized must be inert rather than fatal. + name: "noise, comments and blank lines are ignored", + out: "\n# picking a profile\nchecking usage...\nCLAUDE_CONFIG_DIR=/a\nUNKNOWN_KEY=x\n", + want: RouteDecision{ClaudeConfigDir: "/a"}, + }, + { + name: "value containing = is kept whole", + out: "REASON=used=93%\nHOLD=yes\n", + want: RouteDecision{Hold: true, Reason: "used=93%"}, + }, + { + name: "empty output is no opinion", + out: "", + want: RouteDecision{}, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + got := ParseRouteOutput(tc.out) + if got != tc.want { + t.Errorf("ParseRouteOutput(%q) = %+v, want %+v", tc.out, got, tc.want) + } + }) + } +} + +func TestRouteDecisionEmpty(t *testing.T) { + if !(RouteDecision{}).Empty() { + t.Error("zero decision should be empty") + } + if (RouteDecision{ClaudeConfigDir: "/a"}).Empty() { + t.Error("decision with a config dir is not empty") + } + if (RouteDecision{Hold: true}).Empty() { + t.Error("hold decision is not empty") + } + // A reason on its own carries no instruction, so it must not count as an + // answer — otherwise a script that only logged would silently shadow the + // next router in line. + if !(RouteDecision{Reason: "just saying"}).Empty() { + t.Error("reason-only decision should be empty") + } +} + +func TestRoute_AppliesDecisionAndInjectsEnv(t *testing.T) { + root := t.TempDir() + writePlugin(t, root, "router", + "name: router\nhooks:\n task.route: route.sh\n", + map[string]string{"route.sh": "#!/bin/sh\necho \"CLAUDE_CONFIG_DIR=/dirs/$TASK_PROJECT-$TASK_ID-$TASK_EXECUTOR\"\n"}) + + r := routeRunner(t, root) + if !r.HandlesRoute() { + t.Fatal("HandlesRoute() = false, want true") + } + + task := &db.Task{ID: 7, Title: "t", Project: "proj", Executor: db.ExecutorClaude} + got := r.Route(context.Background(), task) + if got.ClaudeConfigDir != "/dirs/proj-7-claude" { + t.Errorf("ClaudeConfigDir = %q", got.ClaudeConfigDir) + } + if got.Plugin != "router" { + t.Errorf("Plugin = %q, want router", got.Plugin) + } +} + +func TestRoute_FirstNonEmptyDecisionWinsInNameOrder(t *testing.T) { + root := t.TempDir() + // "a-quiet" sorts first but abstains; "b-router" must then be consulted. + writePlugin(t, root, "a-quiet", + "name: a-quiet\nhooks:\n task.route: route.sh\n", + map[string]string{"route.sh": "#!/bin/sh\nexit 0\n"}) + writePlugin(t, root, "b-router", + "name: b-router\nhooks:\n task.route: route.sh\n", + map[string]string{"route.sh": "#!/bin/sh\necho CLAUDE_CONFIG_DIR=/from-b\n"}) + writePlugin(t, root, "c-router", + "name: c-router\nhooks:\n task.route: route.sh\n", + map[string]string{"route.sh": "#!/bin/sh\necho CLAUDE_CONFIG_DIR=/from-c\n"}) + + r := routeRunner(t, root) + got := r.Route(context.Background(), &db.Task{ID: 1, Executor: db.ExecutorClaude}) + if got.ClaudeConfigDir != "/from-b" { + t.Errorf("ClaudeConfigDir = %q, want /from-b (first answering plugin by name)", got.ClaudeConfigDir) + } +} + +func TestRoute_FailingScriptIsSkipped(t *testing.T) { + root := t.TempDir() + // A router that prints a decision *and* exits non-zero must not be trusted: + // a half-finished script's last echo is not a decision. + writePlugin(t, root, "a-broken", + "name: a-broken\nhooks:\n task.route: route.sh\n", + map[string]string{"route.sh": "#!/bin/sh\necho CLAUDE_CONFIG_DIR=/bad\nexit 3\n"}) + writePlugin(t, root, "b-good", + "name: b-good\nhooks:\n task.route: route.sh\n", + map[string]string{"route.sh": "#!/bin/sh\necho CLAUDE_CONFIG_DIR=/good\n"}) + + r := routeRunner(t, root) + got := r.Route(context.Background(), &db.Task{ID: 1, Executor: db.ExecutorClaude}) + if got.ClaudeConfigDir != "/good" { + t.Errorf("ClaudeConfigDir = %q, want /good", got.ClaudeConfigDir) + } +} + +func TestRoute_StderrIsNotParsedAsDecision(t *testing.T) { + root := t.TempDir() + writePlugin(t, root, "router", + "name: router\nhooks:\n task.route: route.sh\n", + map[string]string{"route.sh": "#!/bin/sh\necho 'CLAUDE_CONFIG_DIR=/from-stderr' >&2\necho CLAUDE_CONFIG_DIR=/from-stdout\n"}) + + r := routeRunner(t, root) + got := r.Route(context.Background(), &db.Task{ID: 1, Executor: db.ExecutorClaude}) + if got.ClaudeConfigDir != "/from-stdout" { + t.Errorf("ClaudeConfigDir = %q, want /from-stdout", got.ClaudeConfigDir) + } +} + +func TestRoute_NoRoutePluginsIsNoOpinion(t *testing.T) { + root := t.TempDir() + writePlugin(t, root, "notifier", + "name: notifier\nhooks:\n task.done: done.sh\n", + map[string]string{"done.sh": "#!/bin/sh\n"}) + + r := routeRunner(t, root) + if r.HandlesRoute() { + t.Error("HandlesRoute() = true with no task.route hook") + } + if got := r.Route(context.Background(), &db.Task{ID: 1}); !got.Empty() { + t.Errorf("Route = %+v, want empty", got) + } +} + +func TestRoute_NilTaskIsSafe(t *testing.T) { + root := t.TempDir() + writePlugin(t, root, "router", + "name: router\nhooks:\n task.route: route.sh\n", + map[string]string{"route.sh": "#!/bin/sh\necho CLAUDE_CONFIG_DIR=/x\n"}) + + if got := routeRunner(t, root).Route(context.Background(), nil); !got.Empty() { + t.Errorf("Route(nil) = %+v, want empty", got) + } +} + +func TestRoute_CancelledContextYieldsNoDecision(t *testing.T) { + root := t.TempDir() + writePlugin(t, root, "router", + "name: router\nhooks:\n task.route: route.sh\n", + map[string]string{"route.sh": "#!/bin/sh\necho CLAUDE_CONFIG_DIR=/x\n"}) + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + if got := routeRunner(t, root).Route(ctx, &db.Task{ID: 1}); !got.Empty() { + t.Errorf("Route with cancelled ctx = %+v, want empty (spawn as configured)", got) + } +} From 099a56951d30cf07fad3ff8b21f192ba7c454b9a Mon Sep 17 00:00:00 2001 From: Bruno Bornsztein Date: Sat, 15 Aug 2026 17:24:53 -0500 Subject: [PATCH 2/3] Move claude-profile-router to the community plugin collection MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit It has config, a README, and a surface of its own — the repo's own rule is "start in examples/; graduate to a repo when it earns one", and this one earns it. examples/plugins/ stays what it says on the tin: short, canonical starting points you read in one sitting. The hook stays here, because that is core: task.route is a ty event, and only ty can emit it. What moves is the policy that answers it. Plugin lands in taskyou/plugins; docs here point at the collection instead of a local directory. Co-Authored-By: Claude Opus 5 --- docs/plugin-ideas.md | 7 +- docs/plugins.md | 15 ++- .../plugins/claude-profile-router/README.md | 95 ---------------- .../claude-profile-router/config.example.env | 21 ---- .../plugins/claude-profile-router/plugin.yaml | 18 ---- .../plugins/claude-profile-router/route.sh | 102 ------------------ .../plugins/claude-profile-router/status.sh | 31 ------ 7 files changed, 15 insertions(+), 274 deletions(-) delete mode 100644 examples/plugins/claude-profile-router/README.md delete mode 100644 examples/plugins/claude-profile-router/config.example.env delete mode 100644 examples/plugins/claude-profile-router/plugin.yaml delete mode 100755 examples/plugins/claude-profile-router/route.sh delete mode 100755 examples/plugins/claude-profile-router/status.sh diff --git a/docs/plugin-ideas.md b/docs/plugin-ideas.md index 0a7501f5..d8cc259f 100644 --- a/docs/plugin-ideas.md +++ b/docs/plugin-ideas.md @@ -7,6 +7,8 @@ demand with the task's env). A plugin is just a directory with executables; it can be any language and can bundle its own config/binaries. ✅ = shipped as an example in [`examples/plugins/`](../examples/plugins/). +📦 = shipped in the [community collection](https://github.com/taskyou/plugins) +(`ty plugins add https://github.com/taskyou/plugins`). ## Notifications & awareness (hooks) @@ -34,8 +36,9 @@ Fires before a task spawns and its stdout is read back as a decision — the one hook that changes how a task runs rather than reporting on it. See [Routing](plugins.md#routing-pre-spawn). -- ✅ **claude-profile-router** — send each task to whichever Claude account has - the most rate-limit headroom; hold the task when both are spent. +- 📦 **claude-profile-router** — send each task to whichever Claude account has + the most rate-limit headroom; hold the task when both are spent. Ships in the + [community collection](https://github.com/taskyou/plugins). - **quiet-hours** — `HOLD=1` outside working hours, so overnight queueing doesn't spend your weekly limit while you sleep. - **cheap-account-first** — route routine task types (docs, chores) to a Pro diff --git a/docs/plugins.md b/docs/plugins.md index 37fa8f25..6dc560be 100644 --- a/docs/plugins.md +++ b/docs/plugins.md @@ -203,10 +203,16 @@ Extra environment on a routing hook, beyond the standard `TASK_*` and `TASK_PLUGIN_*` variables: `TASK_EXECUTOR` and `TASK_CLAUDE_CONFIG_DIR` (the task's current config dir, empty when unset). -The worked example is -[`examples/plugins/claude-profile-router/`](../examples/plugins/claude-profile-router/), -which routes each task to whichever of your Claude accounts has the most -rate-limit headroom — see also [`ty usage`](#inspecting). +The worked example is **claude-profile-router** in the +[community collection](https://github.com/taskyou/plugins), which routes each task +to whichever of your Claude accounts has the most rate-limit headroom and holds a +task when every account is spent: + +```bash +ty plugins add https://github.com/taskyou/plugins +``` + +It runs on [`ty usage`](#inspecting), which reports the same numbers directly. ## Environment @@ -320,7 +326,6 @@ Complete, copy-pasteable plugins live in [`examples/plugins/`](../examples/plugi | [`slack`](../examples/plugins/slack/) | hooks | webhook integration; bundled `config.env` | | [`worktree`](../examples/plugins/worktree/) | actions | task-scoped `diff` / `test` using `WORKTREE_PATH` | | [`heartbeat`](../examples/plugins/heartbeat/) | service | a daemon-supervised long-running process | -| [`claude-profile-router`](../examples/plugins/claude-profile-router/) | route hook + action | picking a Claude account per task from live usage | ```bash cp -R examples/plugins/desktop-notify ~/.config/task/plugins/ diff --git a/examples/plugins/claude-profile-router/README.md b/examples/plugins/claude-profile-router/README.md deleted file mode 100644 index 1bd959a8..00000000 --- a/examples/plugins/claude-profile-router/README.md +++ /dev/null @@ -1,95 +0,0 @@ -# claude-profile-router - -Route each task to whichever of your Claude accounts has the most rate-limit -headroom left. - -If you have two logins — a personal one and a work one, say — you already have -two Claude config dirs. This plugin checks how much of each account's 5-hour and -weekly limits are spent and points every task ty spawns at the one with room. If -both are spent, it holds the task in the queue instead of burning a session on a -429. - -## Setup - -1. **Have two profiles.** A profile is a `CLAUDE_CONFIG_DIR` with its own login: - - ```bash - CLAUDE_CONFIG_DIR=~/.claude-work claude # then /login as the second account - ``` - -2. **Install the plugin:** - - ```bash - cp -R examples/plugins/claude-profile-router ~/.config/task/plugins/ - cd ~/.config/task/plugins/claude-profile-router - cp config.example.env config.env - $EDITOR config.env # list your profile dirs in TY_CLAUDE_PROFILES - ``` - -3. **Check it sees both accounts:** - - ```bash - ty plugins run claude-profile-router status - ``` - - ``` - Routing threshold: skip a profile at or above 90% used - - /Users/me/.claude - me@personal.example - 3% used (5-hour session, resets Sat 14:00) — 97% headroom - - /Users/me/.claude-work - me@work.example - 71% used (weekly, resets Thu 20:00) — 29% headroom - ``` - -That's it — the next task ty spawns is routed. `ty logs` and the task's own log -record which profile it landed on and why. - -## How it decides - -- Each profile's **binding limit** is the worst of its reported windows (5-hour - session, weekly, per-model weekly). A session window at 98% blocks the next - task even when the weekly one is untouched, so the worst window is the one - that matters. -- The profile with the **lowest** binding percent wins. -- Profiles at or above `TY_CLAUDE_MAX_PERCENT` (default 90) are skipped. The - margin exists because usage is sampled at spawn, not metered continuously — - a long task started at 89% can still cross the line mid-run. -- If every profile is over the threshold, the task is **held**: it stays queued - and is reconsidered on the next daemon tick, with one log line saying why. -- A task that already names a config dir (set by hand, or by a workflow step) is - left alone. Routing fills a vacuum; it doesn't overrule you. -- **A task is routed once and stays there.** Its Claude session lives inside that - config dir, so a resume has to happen under the same profile or it would start - a fresh conversation. A task already running on a profile therefore waits for - *that* profile to reset rather than hopping to the other one. -- Anything that goes wrong — no credentials, an expired login, `ty` not on the - daemon's `PATH` — means the plugin says nothing and the task spawns exactly as - it would have without it. - -## Configuration - -See [`config.example.env`](config.example.env). The knobs: - -| Variable | Default | Meaning | -| --- | --- | --- | -| `TY_CLAUDE_PROFILES` | *(required)* | Space-separated config dirs to route between | -| `TY_CLAUDE_MAX_PERCENT` | `90` | Skip a profile at or above this percent used | -| `TY_CLAUDE_PROJECTS` | *(all)* | Only route tasks in these projects | -| `TY_BIN` | `ty` | Path to the ty binary, if the daemon's `PATH` lacks it | - -## Caveats - -- **A config dir is more than an account.** It also carries that profile's - plugins, MCP servers, and trusted-worktree state. Set both profiles up the - same way, or a task routed to the quieter one may find tools missing. If you - want to swap only credentials, use a per-task `env` override instead — see - [docs/plugins.md](../../../docs/plugins.md). -- **Usage is read, never written.** The plugin reads each profile's stored OAuth - token to call the same endpoint Claude Code's `/usage` uses. It never - refreshes or rewrites a credential. A profile whose token has gone stale - reports as unavailable until you run a `claude` session under it. -- **Two probes per spawn.** Each is a single HTTPS GET; ty caps the whole hook - at 15s and spawns normally if it overruns. diff --git a/examples/plugins/claude-profile-router/config.example.env b/examples/plugins/claude-profile-router/config.example.env deleted file mode 100644 index 19550d8a..00000000 --- a/examples/plugins/claude-profile-router/config.example.env +++ /dev/null @@ -1,21 +0,0 @@ -# Copy to config.env (next to this file) and edit. - -# The Claude profiles to route between: a space-separated list of -# CLAUDE_CONFIG_DIR paths. Each is one logged-in account. -# -# To create a second profile, log in with a different config dir: -# CLAUDE_CONFIG_DIR=~/.claude-work claude # then /login -# -# Check what each one is with: ty usage -TY_CLAUDE_PROFILES="$HOME/.claude $HOME/.claude-work" - -# Skip a profile once its binding limit (5-hour session or weekly, whichever is -# worse) is at or above this percent. When every profile is above it, tasks are -# held in the queue instead of spawning into a 429. -TY_CLAUDE_MAX_PERCENT=90 - -# Only route tasks in these projects (space-separated). Empty = all projects. -TY_CLAUDE_PROJECTS="" - -# Path to the ty binary, if the daemon's PATH doesn't include it. -# TY_BIN=/usr/local/bin/ty diff --git a/examples/plugins/claude-profile-router/plugin.yaml b/examples/plugins/claude-profile-router/plugin.yaml deleted file mode 100644 index 17df1a60..00000000 --- a/examples/plugins/claude-profile-router/plugin.yaml +++ /dev/null @@ -1,18 +0,0 @@ -# Route each task to whichever Claude account has the most rate-limit headroom. -# -# Copy this directory to ~/.config/task/plugins/claude-profile-router/, copy -# config.example.env to config.env, list your profiles in it, and every task ty -# spawns picks its account automatically. See README.md. -name: claude-profile-router -version: 0.1.0 -description: Send each task to the Claude profile with the most usage headroom. - -hooks: - # task.route is the one hook ty waits on: it fires just before a task spawns - # and reads the script's stdout back as a decision. See docs/plugins.md. - task.route: route.sh - -actions: - - id: status - label: Show usage for each Claude profile - command: status.sh diff --git a/examples/plugins/claude-profile-router/route.sh b/examples/plugins/claude-profile-router/route.sh deleted file mode 100755 index bd72ef74..00000000 --- a/examples/plugins/claude-profile-router/route.sh +++ /dev/null @@ -1,102 +0,0 @@ -#!/bin/bash -# task.route hook: pick the Claude profile with the most rate-limit headroom. -# -# ty runs this synchronously just before it spawns a task and reads stdout back -# as the decision, so stdout carries KEY=VALUE lines and nothing else — every -# diagnostic goes to stderr (which lands in the daemon log). -# -# CLAUDE_CONFIG_DIR= run this task under that profile -# HOLD=1 / REASON= every profile is spent; keep the task queued -# (no output) no opinion; ty spawns as already configured -# -# Printing nothing is always safe, so every failure path here does exactly that. -set -uo pipefail - -say() { echo "claude-profile-router: $*" >&2; } - -# config.env holds the settings, but anything already in the environment wins — -# that is what makes a one-off `TY_CLAUDE_MAX_PERCENT=50 ./route.sh` a usable way -# to try a threshold without editing the file. -_pre_profiles="${TY_CLAUDE_PROFILES:-}" -_pre_max="${TY_CLAUDE_MAX_PERCENT:-}" -_pre_projects="${TY_CLAUDE_PROJECTS:-}" -_pre_bin="${TY_BIN:-}" -if [[ -n "${TASK_PLUGIN_DIR:-}" && -f "$TASK_PLUGIN_DIR/config.env" ]]; then - # shellcheck disable=SC1091 - source "$TASK_PLUGIN_DIR/config.env" -fi -[[ -n "$_pre_profiles" ]] && TY_CLAUDE_PROFILES="$_pre_profiles" -[[ -n "$_pre_max" ]] && TY_CLAUDE_MAX_PERCENT="$_pre_max" -[[ -n "$_pre_projects" ]] && TY_CLAUDE_PROJECTS="$_pre_projects" -[[ -n "$_pre_bin" ]] && TY_BIN="$_pre_bin" - -TY="${TY_BIN:-ty}" -MAX_PERCENT="${TY_CLAUDE_MAX_PERCENT:-90}" - -if [[ -z "${TY_CLAUDE_PROFILES:-}" ]]; then - say "TY_CLAUDE_PROFILES not set (see config.example.env)" - exit 0 -fi - -if ! command -v "$TY" >/dev/null 2>&1; then - say "ty not found on PATH (set TY_BIN in config.env)" - exit 0 -fi - -# Optional project allowlist, so routing can be tried on one project first. -if [[ -n "${TY_CLAUDE_PROJECTS:-}" ]]; then - match="" - for p in $TY_CLAUDE_PROJECTS; do - [[ "$p" == "${TASK_PROJECT:-}" ]] && match=1 && break - done - [[ -z "$match" ]] && exit 0 -fi - -best_dir="" -best_pct="" -exhausted_low="" # lowest usage among profiles that were over the threshold - -for raw in $TY_CLAUDE_PROFILES; do - dir="${raw/#\~/$HOME}" - - # --percent prints one bare number: the binding limit's used percent. - if ! pct=$("$TY" usage --config-dir "$dir" --percent 2>/dev/null); then - say "skipping $dir (usage unavailable — expired login?)" - continue - fi - if [[ ! "$pct" =~ ^[0-9]+$ ]]; then - say "skipping $dir (unparseable usage: '$pct')" - continue - fi - - if (( pct >= MAX_PERCENT )); then - say "$dir at ${pct}% (>= ${MAX_PERCENT}%), skipping" - if [[ -z "$exhausted_low" ]] || (( pct < exhausted_low )); then - exhausted_low="$pct" - fi - continue - fi - - if [[ -z "$best_pct" ]] || (( pct < best_pct )); then - best_pct="$pct" - best_dir="$dir" - fi -done - -if [[ -n "$best_dir" ]]; then - say "routing to $best_dir (${best_pct}% used)" - echo "CLAUDE_CONFIG_DIR=$best_dir" - echo "REASON=${best_pct}% of its binding limit used" - exit 0 -fi - -# Nothing usable. Only hold the task if we actually saw an exhausted profile — -# if every probe merely failed, stay out of the way and let ty spawn normally -# rather than parking the whole board behind a broken credential lookup. -if [[ -n "$exhausted_low" ]]; then - echo "HOLD=1" - echo "REASON=every Claude profile is at or above ${MAX_PERCENT}% (best is ${exhausted_low}%)" - exit 0 -fi - -say "no profile could be evaluated; leaving this task alone" diff --git a/examples/plugins/claude-profile-router/status.sh b/examples/plugins/claude-profile-router/status.sh deleted file mode 100755 index bd630a92..00000000 --- a/examples/plugins/claude-profile-router/status.sh +++ /dev/null @@ -1,31 +0,0 @@ -#!/bin/bash -# `ty plugins run claude-profile-router status` — show what the router sees. -# Same numbers route.sh decides on, so a surprising routing choice can be -# checked against reality without reading the daemon log. -set -uo pipefail - -_pre_profiles="${TY_CLAUDE_PROFILES:-}" -_pre_max="${TY_CLAUDE_MAX_PERCENT:-}" -_pre_bin="${TY_BIN:-}" -if [[ -n "${TASK_PLUGIN_DIR:-}" && -f "$TASK_PLUGIN_DIR/config.env" ]]; then - # shellcheck disable=SC1091 - source "$TASK_PLUGIN_DIR/config.env" -fi -[[ -n "$_pre_profiles" ]] && TY_CLAUDE_PROFILES="$_pre_profiles" -[[ -n "$_pre_max" ]] && TY_CLAUDE_MAX_PERCENT="$_pre_max" -[[ -n "$_pre_bin" ]] && TY_BIN="$_pre_bin" - -TY="${TY_BIN:-ty}" -MAX_PERCENT="${TY_CLAUDE_MAX_PERCENT:-90}" - -if [[ -z "${TY_CLAUDE_PROFILES:-}" ]]; then - echo "No profiles configured. Copy config.example.env to config.env and set TY_CLAUDE_PROFILES." - exit 0 -fi - -echo "Routing threshold: skip a profile at or above ${MAX_PERCENT}% used" -echo -for raw in $TY_CLAUDE_PROFILES; do - dir="${raw/#\~/$HOME}" - "$TY" usage --config-dir "$dir" || echo " (unavailable)" -done From f9ae7ece13286d05618413db8bf382e66e27ec0b Mon Sep 17 00:00:00 2001 From: Bruno Bornsztein Date: Sat, 15 Aug 2026 17:45:53 -0500 Subject: [PATCH 3/3] Move the usage reader into the plugin; keep only the hook in core MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `ty usage` and internal/claudeusage were ~79% of this branch, and none of it had to be here — the plugin can read the keychain and call the endpoint itself, as it now does. What's left is the part a plugin genuinely cannot do. Plugins see four events, all from OnStatusChange, and the earliest of them (task.started) fires after the spawn command is already built. The alternative — a service polling for queued tasks — races a 2s tick and would have to write claude_config_dir straight into SQLite, since the HTTP API exposes that field on projects only. So ty has to offer the moment; it does not have to offer the data. Moving it also puts knowledge of someone else's endpoint where it can be fixed with a git pull instead of a ty release, which matters for something this likely to drift. Core keeps: the task.route hook, its executor wiring, and UpdateTaskClaudeConfigDir. Data and policy both live in taskyou/plugins#1 now. Co-Authored-By: Claude Opus 5 --- cmd/task/main.go | 1 - cmd/task/usage.go | 203 ----------------- docs/plugins.md | 4 +- internal/claudeusage/cache.go | 109 ---------- internal/claudeusage/cache_test.go | 195 ----------------- internal/claudeusage/credentials.go | 157 ------------- internal/claudeusage/usage.go | 327 ---------------------------- internal/claudeusage/usage_test.go | 261 ---------------------- 8 files changed, 2 insertions(+), 1255 deletions(-) delete mode 100644 cmd/task/usage.go delete mode 100644 internal/claudeusage/cache.go delete mode 100644 internal/claudeusage/cache_test.go delete mode 100644 internal/claudeusage/credentials.go delete mode 100644 internal/claudeusage/usage.go delete mode 100644 internal/claudeusage/usage_test.go diff --git a/cmd/task/main.go b/cmd/task/main.go index 387d4d86..55dc1b36 100644 --- a/cmd/task/main.go +++ b/cmd/task/main.go @@ -531,7 +531,6 @@ Examples: // Plugins subcommand - inspect installed task plugins rootCmd.AddCommand(newPluginsCmd()) - rootCmd.AddCommand(newUsageCmd()) // Workflow artifact store over the CLI — the transport-independent twin of the // taskyou_get_artifact/taskyou_set_artifact MCP tools, so a workflow phase can diff --git a/cmd/task/usage.go b/cmd/task/usage.go deleted file mode 100644 index 968fb78e..00000000 --- a/cmd/task/usage.go +++ /dev/null @@ -1,203 +0,0 @@ -package main - -import ( - "context" - "encoding/json" - "fmt" - "os" - "sort" - "strings" - "time" - - "github.com/spf13/cobra" - - "github.com/bborn/workflow/internal/claudeusage" - "github.com/bborn/workflow/internal/db" - "github.com/bborn/workflow/internal/executor" -) - -// `ty usage` exposes, per Claude profile, how much of that account's rate -// limits are already spent. It exists so a decision that used to be guesswork — -// "which of my logins should this task run under?" — can be made from real -// numbers, by a person at the terminal or by a plugin's routing script. -// -// The plugin path is why the output is shaped the way it is. A hook script has -// no JSON parser to lean on, so --percent prints one bare number and nothing -// else; comparing profiles is then a numeric sort in shell. --json is for -// anything richer. - -// profileResult pairs a probed config dir with its outcome. Errors are carried -// rather than returned so one unreadable profile (an expired login, say) still -// lets the others report. -type profileResult struct { - ConfigDir string `json:"config_dir"` - Snapshot *claudeusage.Snapshot `json:"snapshot,omitempty"` - Error string `json:"error,omitempty"` -} - -func newUsageCmd() *cobra.Command { - cmd := &cobra.Command{ - Use: "usage", - Short: "Show how much of each Claude profile's rate limits are used", - Long: `Report Claude subscription usage for one or more profiles. - -A profile is a CLAUDE_CONFIG_DIR — one logged-in Claude account. With no ---config-dir flags, every config dir ty knows about is probed: the default -(~/.claude) plus each distinct dir configured on a project. - -Usage is read from the same endpoint Claude Code's own /usage command uses, -authenticated with the credentials already stored for that profile. Nothing is -written: no token is refreshed, rewritten, or printed. - -Results are cached for a minute, because that endpoint rate-limits and routing -calls it on every task spawn. Use --refresh to force a live read. - -Examples: - ty usage # every profile ty knows about - ty usage --config-dir ~/.claude-work # one profile, with account email - ty usage --config-dir ~/.claude-work --percent # just "42" — for scripts - ty usage --json # full detail - ty usage --refresh # ignore the cache`, - SilenceUsage: true, - RunE: func(cmd *cobra.Command, args []string) error { - dirs, _ := cmd.Flags().GetStringArray("config-dir") - asJSON, _ := cmd.Flags().GetBool("json") - percentOnly, _ := cmd.Flags().GetBool("percent") - refresh, _ := cmd.Flags().GetBool("refresh") - - if len(dirs) == 0 { - dirs = knownConfigDirs() - } - if len(dirs) == 0 { - return fmt.Errorf("no Claude config dirs to check") - } - if percentOnly && len(dirs) != 1 { - return fmt.Errorf("--percent needs exactly one --config-dir (got %d)", len(dirs)) - } - - ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) - defer cancel() - - client := claudeusage.NewClient() - client.NoCache = refresh - results := probeProfiles(ctx, client, dirs, !percentOnly) - - switch { - case percentOnly: - if results[0].Error != "" { - return fmt.Errorf("%s", results[0].Error) - } - // One bare number, no styling, no trailing prose: routing - // scripts read this with $(...) and compare it numerically. - fmt.Printf("%.0f\n", results[0].Snapshot.UsedPercent()) - case asJSON: - enc := json.NewEncoder(os.Stdout) - enc.SetIndent("", " ") - return enc.Encode(results) - default: - printUsageTable(results) - } - return nil - }, - } - - cmd.Flags().StringArray("config-dir", nil, "Claude config dir to check (repeatable; defaults to every dir ty knows about)") - cmd.Flags().Bool("json", false, "Emit JSON") - cmd.Flags().Bool("percent", false, "Print only the binding limit's used percent (requires one --config-dir)") - cmd.Flags().Bool("refresh", false, "Bypass the cache and read live") - return cmd -} - -// probeProfiles fetches usage for each dir. withAccount adds the account email, -// which costs a second request per profile — worth it for a human reading a -// table, wasted for a script that only wants a number. -func probeProfiles(ctx context.Context, client *claudeusage.Client, dirs []string, withAccount bool) []profileResult { - results := make([]profileResult, 0, len(dirs)) - for _, dir := range dirs { - res := profileResult{ConfigDir: dir} - var snap *claudeusage.Snapshot - var err error - if withAccount { - snap, err = client.FetchWithAccount(ctx, dir) - } else { - snap, err = client.Fetch(ctx, dir) - } - if err != nil { - res.Error = err.Error() - } else { - res.Snapshot = snap - res.ConfigDir = snap.ConfigDir - } - results = append(results, res) - } - return results -} - -// knownConfigDirs collects every Claude config dir ty is already aware of: the -// default, plus whatever projects have been pointed at. Opening the DB is best -// effort — `ty usage` must still work from a machine with no board. -func knownConfigDirs() []string { - seen := map[string]bool{} - var dirs []string - add := func(d string) { - d = executor.ResolveClaudeConfigDir(d) - if d == "" || seen[d] { - return - } - seen[d] = true - dirs = append(dirs, d) - } - - add("") // the default dir - - if database, err := openTaskDB(db.DefaultPath()); err == nil { - defer database.Close() //nolint:errcheck // read-only lookup - if projects, err := database.ListProjects(); err == nil { - for _, p := range projects { - if strings.TrimSpace(p.ClaudeConfigDir) != "" { - add(p.ClaudeConfigDir) - } - } - } - } - - sort.Strings(dirs[1:]) // keep the default first, order the rest stably - return dirs -} - -func printUsageTable(results []profileResult) { - for _, r := range results { - fmt.Println(boldStyle.Render(r.ConfigDir)) - if r.Error != "" { - fmt.Println(" " + errorStyle.Render(r.Error)) - fmt.Println() - continue - } - if r.Snapshot.Email != "" { - fmt.Println(" " + dimStyle.Render(r.Snapshot.Email)) - } - style := successStyle - switch used := r.Snapshot.UsedPercent(); { - case used >= 90: - style = errorStyle - case used >= 70: - style = warnStyle - } - line := r.Snapshot.Describe() - if r.Snapshot.Stale { - line += fmt.Sprintf(" (cached %s ago; the usage API is unreachable)", r.Snapshot.Age().Round(time.Minute)) - } - fmt.Println(" " + style.Render(line)) - for _, l := range r.Snapshot.Limits { - line := fmt.Sprintf(" %-16s %3.0f%%", l.Kind, l.Percent) - if l.Scope != "" { - line += " " + l.Scope - } - if l.ResetsAt != nil { - line += " resets " + l.ResetsAt.Local().Format("Mon 15:04") - } - fmt.Println(dimStyle.Render(line)) - } - fmt.Println() - } -} diff --git a/docs/plugins.md b/docs/plugins.md index 6dc560be..7b9756c4 100644 --- a/docs/plugins.md +++ b/docs/plugins.md @@ -212,7 +212,8 @@ task when every account is spent: ty plugins add https://github.com/taskyou/plugins ``` -It runs on [`ty usage`](#inspecting), which reports the same numbers directly. +It reads each account's remaining limits itself — ty supplies the hook, the +plugin supplies the policy and the data it routes on. ## Environment @@ -311,7 +312,6 @@ only for a process that must stay *up* (a socket connection, a listening port). ```bash ty plugins list # what's installed and which events each handles ty plugins dir # the plugins directory path -ty usage # rate-limit usage per Claude profile — what a router routes on ``` Set `TY_PLUGINS_DIR` to use a directory other than `~/.config/task/plugins/`. diff --git a/internal/claudeusage/cache.go b/internal/claudeusage/cache.go deleted file mode 100644 index 86ba14af..00000000 --- a/internal/claudeusage/cache.go +++ /dev/null @@ -1,109 +0,0 @@ -package claudeusage - -import ( - "crypto/sha256" - "encoding/hex" - "encoding/json" - "os" - "path/filepath" - "time" -) - -// The usage endpoint rate-limits, and routing calls it on every spawn — a busy -// board would otherwise walk straight into a 429 and lose the very numbers it -// spawns on. So snapshots are cached on disk, briefly. -// -// On disk rather than in memory because the caller that matters is a shell -// script: a routing plugin shells out to `ty usage`, so every probe is a fresh -// process and an in-process cache would never be read. -const ( - // CacheTTL is how long a snapshot is served without asking the API again. - // Rate-limit utilization moves in percentage points over minutes; a minute - // of staleness cannot flip a routing decision that wasn't already marginal. - CacheTTL = 60 * time.Second - - // StaleTTL is how old a cached snapshot may be and still be used when a - // live fetch fails. Routing on a 10-minute-old number beats routing on - // nothing, which is what a transient 429 would otherwise leave us with. - StaleTTL = 30 * time.Minute -) - -// DefaultCacheDir is where snapshots are stored. -func DefaultCacheDir() string { - base, err := os.UserCacheDir() - if err != nil { - return "" - } - return filepath.Join(base, "ty", "claude-usage") -} - -func (c *Client) cacheDir() string { - if c.CacheDir != "" { - return c.CacheDir - } - return DefaultCacheDir() -} - -// cachePath names a profile's cache file by a hash of its config dir, so the -// path is flat and safe regardless of what the dir itself is called. -func (c *Client) cachePath(configDir string) string { - dir := c.cacheDir() - if dir == "" { - return "" - } - sum := sha256.Sum256([]byte(normalizeDir(configDir))) - return filepath.Join(dir, hex.EncodeToString(sum[:])[:16]+".json") -} - -// readCache returns a cached snapshot if one exists and is younger than maxAge. -// Any problem — no file, unreadable, corrupt — reads as a miss; a cache is never -// a reason to fail. -func (c *Client) readCache(configDir string, maxAge time.Duration) (*Snapshot, bool) { - if c.NoCache { - return nil, false - } - path := c.cachePath(configDir) - if path == "" { - return nil, false - } - data, err := os.ReadFile(path) //nolint:gosec // path is a hash under our own cache dir - if err != nil { - return nil, false - } - var snap Snapshot - if err := json.Unmarshal(data, &snap); err != nil { - return nil, false - } - if snap.FetchedAt.IsZero() || time.Since(snap.FetchedAt) > maxAge { - return nil, false - } - return &snap, true -} - -// writeCache stores a snapshot. Failures are ignored: a cache that can't be -// written costs an extra API call, nothing more. -func (c *Client) writeCache(configDir string, snap *Snapshot) { - if c.NoCache { - return - } - path := c.cachePath(configDir) - if path == "" { - return - } - if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil { - return - } - data, err := json.Marshal(snap) - if err != nil { - return - } - // Write-then-rename so a concurrent reader never sees a half-written file: - // several spawns can probe the same profile at once. - tmp := path + ".tmp" - if err := os.WriteFile(tmp, data, 0o600); err != nil { - return - } - if err := os.Rename(tmp, path); err != nil { - _ = os.Remove(tmp) - } -} diff --git a/internal/claudeusage/cache_test.go b/internal/claudeusage/cache_test.go deleted file mode 100644 index 47bd3dd4..00000000 --- a/internal/claudeusage/cache_test.go +++ /dev/null @@ -1,195 +0,0 @@ -package claudeusage - -import ( - "context" - "net/http" - "net/http/httptest" - "os" - "strings" - "sync/atomic" - "testing" - "time" -) - -// countingServer serves the usage fixture and reports how many times it was hit. -func countingServer(t *testing.T, hits *atomic.Int32) *httptest.Server { - t.Helper() - srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - hits.Add(1) - _, _ = w.Write([]byte(usageBody)) - })) - t.Cleanup(srv.Close) - return srv -} - -func TestFetch_SecondCallIsServedFromCache(t *testing.T) { - // This is the whole reason the cache exists: routing probes every profile on - // every spawn, and the usage endpoint rate-limits. Without this, a busy - // board 429s itself out of the numbers it routes on. - var hits atomic.Int32 - srv := countingServer(t, &hits) - client := testClient(t, srv) - dir := writeCredsDir(t, "tok", time.Now().Add(time.Hour)) - - for i := 0; i < 3; i++ { - if _, err := client.Fetch(context.Background(), dir); err != nil { - t.Fatalf("Fetch %d: %v", i, err) - } - } - if got := hits.Load(); got != 1 { - t.Errorf("hit the API %d times across 3 fetches, want 1", got) - } -} - -func TestFetch_NoCacheAlwaysReadsLive(t *testing.T) { - var hits atomic.Int32 - srv := countingServer(t, &hits) - client := testClient(t, srv) - client.NoCache = true - dir := writeCredsDir(t, "tok", time.Now().Add(time.Hour)) - - for i := 0; i < 2; i++ { - if _, err := client.Fetch(context.Background(), dir); err != nil { - t.Fatalf("Fetch %d: %v", i, err) - } - } - if got := hits.Load(); got != 2 { - t.Errorf("hit the API %d times with NoCache, want 2", got) - } -} - -func TestFetch_CacheIsPerProfile(t *testing.T) { - // One profile's snapshot must never be served for another — that would route - // tasks to an account based on a different account's headroom. - var hits atomic.Int32 - srv := countingServer(t, &hits) - client := testClient(t, srv) - a := writeCredsDir(t, "tok-a", time.Now().Add(time.Hour)) - b := writeCredsDir(t, "tok-b", time.Now().Add(time.Hour)) - - snapA, err := client.Fetch(context.Background(), a) - if err != nil { - t.Fatal(err) - } - snapB, err := client.Fetch(context.Background(), b) - if err != nil { - t.Fatal(err) - } - if hits.Load() != 2 { - t.Errorf("hit the API %d times for 2 profiles, want 2", hits.Load()) - } - if snapA.ConfigDir == snapB.ConfigDir { - t.Errorf("both snapshots claim config dir %q", snapA.ConfigDir) - } -} - -func TestFetch_StaleCacheRescuesAFailedRequest(t *testing.T) { - // A 429 from the usage endpoint should not leave routing blind: a snapshot - // from a few minutes ago is a far better basis for a decision than none. - var fail atomic.Bool - srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - if fail.Load() { - w.WriteHeader(http.StatusTooManyRequests) - return - } - _, _ = w.Write([]byte(usageBody)) - })) - defer srv.Close() - - client := testClient(t, srv) - dir := writeCredsDir(t, "tok", time.Now().Add(time.Hour)) - - first, err := client.Fetch(context.Background(), dir) - if err != nil { - t.Fatalf("priming fetch: %v", err) - } - // Age the cached entry past CacheTTL but well inside StaleTTL. - first.FetchedAt = time.Now().Add(-5 * time.Minute) - client.writeCache(dir, first) - - fail.Store(true) - snap, err := client.Fetch(context.Background(), dir) - if err != nil { - t.Fatalf("Fetch should have fallen back to cache: %v", err) - } - if !snap.Stale { - t.Error("a cache-rescued snapshot must be marked stale") - } - if got := snap.UsedPercent(); got != 71 { - t.Errorf("UsedPercent = %v, want the cached 71", got) - } -} - -func TestFetch_CacheTooOldToRescue(t *testing.T) { - var fail atomic.Bool - srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - if fail.Load() { - w.WriteHeader(http.StatusTooManyRequests) - return - } - _, _ = w.Write([]byte(usageBody)) - })) - defer srv.Close() - - client := testClient(t, srv) - dir := writeCredsDir(t, "tok", time.Now().Add(time.Hour)) - - first, err := client.Fetch(context.Background(), dir) - if err != nil { - t.Fatal(err) - } - first.FetchedAt = time.Now().Add(-2 * StaleTTL) - client.writeCache(dir, first) - - fail.Store(true) - _, err = client.Fetch(context.Background(), dir) - if err == nil { - t.Fatal("a snapshot older than StaleTTL should not rescue a failed fetch") - } - if !strings.Contains(err.Error(), "429") { - t.Errorf("error should surface the 429, got %v", err) - } -} - -func TestFetch_ExpiredCredentialsIgnoreTheCache(t *testing.T) { - // A stale *credential* is a configuration problem, not a transient one. - // Serving a cached number would keep sending tasks to a profile that can no - // longer run them. - var hits atomic.Int32 - srv := countingServer(t, &hits) - client := testClient(t, srv) - dir := writeCredsDir(t, "tok", time.Now().Add(time.Hour)) - - if _, err := client.Fetch(context.Background(), dir); err != nil { - t.Fatal(err) - } - // Same profile, same warm cache entry — but its token has since expired. - writeCredsInto(t, dir, "tok", time.Now().Add(-time.Hour)) - if _, err := client.Fetch(context.Background(), dir); err == nil { - t.Error("expired credentials should error even with a warm cache") - } -} - -func TestCorruptCacheEntryIsAMiss(t *testing.T) { - var hits atomic.Int32 - srv := countingServer(t, &hits) - client := testClient(t, srv) - dir := writeCredsDir(t, "tok", time.Now().Add(time.Hour)) - - if _, err := client.Fetch(context.Background(), dir); err != nil { - t.Fatal(err) - } - if err := writeFile(client.cachePath(dir), "{not json"); err != nil { - t.Fatal(err) - } - if _, err := client.Fetch(context.Background(), dir); err != nil { - t.Fatalf("a corrupt cache entry must not fail the fetch: %v", err) - } - if hits.Load() != 2 { - t.Errorf("API hits = %d, want 2 (the corrupt entry should have been refetched)", hits.Load()) - } -} - -func writeFile(path, body string) error { - return os.WriteFile(path, []byte(body), 0o600) -} diff --git a/internal/claudeusage/credentials.go b/internal/claudeusage/credentials.go deleted file mode 100644 index 2434e46c..00000000 --- a/internal/claudeusage/credentials.go +++ /dev/null @@ -1,157 +0,0 @@ -package claudeusage - -import ( - "crypto/sha256" - "encoding/hex" - "encoding/json" - "fmt" - "os" - "os/exec" - "path/filepath" - "runtime" - "strings" - "time" -) - -// keychainPrefix is the macOS Keychain service name Claude Code stores its -// OAuth credentials under. For a non-default CLAUDE_CONFIG_DIR the service is -// suffixed with a short hash of the config dir, so every profile gets its own -// entry (see KeychainService). -const keychainPrefix = "Claude Code-credentials" - -// credentialsFile is the on-disk credential store Claude Code uses where no OS -// keychain is available (Linux, containers). It lives inside the config dir. -const credentialsFile = ".credentials.json" - -// Credentials is the subset of a Claude Code credential blob we need: the OAuth -// access token used for api.anthropic.com/api/oauth/* calls, plus enough -// metadata to tell a stale token from a missing one. -type Credentials struct { - AccessToken string - ExpiresAt time.Time // zero when the store didn't record one - Subscription string -} - -// Expired reports whether the access token's recorded expiry has passed. A -// zero ExpiresAt is treated as not expired — we'd rather attempt the request -// and let the API decide than refuse on missing metadata. -func (c Credentials) Expired() bool { - return !c.ExpiresAt.IsZero() && time.Now().After(c.ExpiresAt) -} - -// credentialsBlob mirrors the JSON shape of a Claude Code credential store. We -// only decode the claudeAiOauth object; the sibling mcpOAuth map holds -// per-MCP-server tokens that are none of our business. -type credentialsBlob struct { - ClaudeAIOAuth struct { - AccessToken string `json:"accessToken"` - ExpiresAt int64 `json:"expiresAt"` // epoch milliseconds - SubscriptionType string `json:"subscriptionType"` - } `json:"claudeAiOauth"` -} - -// KeychainService returns the macOS Keychain service name holding the -// credentials for a given CLAUDE_CONFIG_DIR. -// -// Claude Code namespaces each config dir's credentials by appending the first 8 -// hex digits of the SHA-256 of the *absolute, cleaned* config-dir path to the -// base service name — that is what makes two logged-in profiles able to coexist -// on one machine. Reproducing the derivation here is what lets ty read a -// profile's usage without shelling out to `claude` (which has no usage command) -// or making the user paste a token anywhere. -func KeychainService(configDir string) string { - sum := sha256.Sum256([]byte(normalizeDir(configDir))) - return keychainPrefix + "-" + hex.EncodeToString(sum[:])[:8] -} - -// normalizeDir expands a leading ~ and cleans the path, so the hash we compute -// matches the one Claude Code computed from its own resolved config dir. -func normalizeDir(dir string) string { - dir = strings.TrimSpace(dir) - if strings.HasPrefix(dir, "~") { - if home, err := os.UserHomeDir(); err == nil { - dir = filepath.Join(home, dir[1:]) - } - } - return filepath.Clean(dir) -} - -// defaultConfigDir is ~/.claude — the config dir Claude Code uses when -// CLAUDE_CONFIG_DIR is unset. Credentials for it may still live under the -// unsuffixed keychain service from before per-profile namespacing existed. -func defaultConfigDir() string { - home, err := os.UserHomeDir() - if err != nil { - return "" - } - return filepath.Join(home, ".claude") -} - -// LoadCredentials finds the OAuth credentials for one Claude profile. -// -// Lookup order, first hit wins: -// 1. the per-config-dir macOS Keychain entry (the normal case on a Mac); -// 2. /.credentials.json (Linux, containers, and anywhere the -// keychain isn't used); -// 3. for the default ~/.claude only, the legacy unsuffixed keychain entry. -// -// A store that exists but holds an empty access token counts as a miss, so a -// half-migrated profile falls through to the next candidate instead of failing -// with a confusing 401 later. -func LoadCredentials(configDir string) (Credentials, error) { - dir := normalizeDir(configDir) - - if runtime.GOOS == "darwin" { - if c, ok := readKeychain(KeychainService(dir)); ok { - return c, nil - } - } - - if c, ok := readCredentialsFile(filepath.Join(dir, credentialsFile)); ok { - return c, nil - } - - if runtime.GOOS == "darwin" && dir == defaultConfigDir() { - if c, ok := readKeychain(keychainPrefix); ok { - return c, nil - } - } - - return Credentials{}, fmt.Errorf("no Claude credentials found for %s (log in once with CLAUDE_CONFIG_DIR=%s claude)", dir, dir) -} - -// readKeychain pulls a credential blob out of the macOS Keychain. A missing -// entry, a locked keychain, or an entry without an access token all report -// "not found" rather than an error: every one of them means "try the next -// candidate", and none is worth failing the whole lookup over. -func readKeychain(service string) (Credentials, bool) { - out, err := exec.Command("security", "find-generic-password", "-s", service, "-w").Output() - if err != nil { - return Credentials{}, false - } - return parseCredentials(out) -} - -func readCredentialsFile(path string) (Credentials, bool) { - data, err := os.ReadFile(path) //nolint:gosec // path is derived from the caller's own config dir - if err != nil { - return Credentials{}, false - } - return parseCredentials(data) -} - -func parseCredentials(data []byte) (Credentials, bool) { - var blob credentialsBlob - if err := json.Unmarshal(data, &blob); err != nil { - return Credentials{}, false - } - token := strings.TrimSpace(blob.ClaudeAIOAuth.AccessToken) - if token == "" { - return Credentials{}, false - } - c := Credentials{AccessToken: token, Subscription: blob.ClaudeAIOAuth.SubscriptionType} - if ms := blob.ClaudeAIOAuth.ExpiresAt; ms > 0 { - c.ExpiresAt = time.UnixMilli(ms) - } - return c, true -} diff --git a/internal/claudeusage/usage.go b/internal/claudeusage/usage.go deleted file mode 100644 index c6aa4623..00000000 --- a/internal/claudeusage/usage.go +++ /dev/null @@ -1,327 +0,0 @@ -// Package claudeusage reads how much of a Claude subscription's rate limits a -// given Claude Code profile has burned through. -// -// A "profile" here is a CLAUDE_CONFIG_DIR: one logged-in Claude account with -// its own credentials. Someone with two accounts (say a personal one and a work -// one) keeps them in two config dirs and points ty at whichever they want a task -// to run under. This package answers the question that makes *choosing* between -// them possible — "how much headroom does each one have left?" — by reading the -// profile's stored OAuth token and calling the same usage endpoint Claude Code's -// own /usage command uses. -// -// It is strictly read-only: it never refreshes, rewrites, or otherwise touches -// stored credentials. A profile whose access token has gone stale simply reports -// an error, and the caller decides what to do about it. -package claudeusage - -import ( - "context" - "encoding/json" - "fmt" - "io" - "net/http" - "strings" - "time" -) - -// DefaultBaseURL is the Anthropic API host serving the OAuth usage/profile -// endpoints. Overridable on Client for tests. -const DefaultBaseURL = "https://api.anthropic.com" - -// DefaultTimeout bounds a single usage probe. Routing a task waits on this, so -// it is deliberately short: a slow or unreachable API should fall back to the -// existing config rather than stall a spawn. -const DefaultTimeout = 10 * time.Second - -// Client fetches usage for Claude profiles. -type Client struct { - BaseURL string - HTTP *http.Client - // CacheDir overrides where snapshots are cached ("" = DefaultCacheDir). - CacheDir string - // NoCache bypasses the cache entirely, in both directions. - NoCache bool -} - -// NewClient returns a Client with the default endpoint and timeout. -func NewClient() *Client { - return &Client{BaseURL: DefaultBaseURL, HTTP: &http.Client{Timeout: DefaultTimeout}} -} - -// Limit is one rate-limit window the API reports for an account. -type Limit struct { - Kind string `json:"kind"` // "session", "weekly_all", "weekly_scoped", … - Group string `json:"group"` // "session" or "weekly" - Percent float64 `json:"percent"` // 0-100 of this window consumed - Severity string `json:"severity"` // provider's own label, e.g. "normal" - ResetsAt *time.Time `json:"resets_at,omitempty"` - Scope string `json:"scope,omitempty"` // model name for per-model windows -} - -// Snapshot is one profile's usage at a point in time. -type Snapshot struct { - ConfigDir string `json:"config_dir"` - Email string `json:"email,omitempty"` // filled in only when asked for; a second API call - Limits []Limit `json:"limits"` - FetchedAt time.Time `json:"fetched_at"` - // Stale marks a snapshot served from cache after a live fetch failed. The - // numbers are still worth routing on, but a caller showing them to a person - // should say so. - Stale bool `json:"stale,omitempty"` -} - -// Age is how long ago this snapshot was read from the API. -func (s Snapshot) Age() time.Duration { return time.Since(s.FetchedAt) } - -// UsedPercent is the profile's *binding* constraint: the highest utilization -// across every window the API reports. Routing cares about the window that will -// stop work first, not the average — a 5-hour session window at 98% blocks the -// next task even when the weekly window is nearly untouched. -func (s Snapshot) UsedPercent() float64 { - worst := 0.0 - for _, l := range s.Limits { - if l.Percent > worst { - worst = l.Percent - } - } - return worst -} - -// Headroom is how much of the binding window is still available, in percent. -// This is the number to route on: higher wins. -func (s Snapshot) Headroom() float64 { return 100 - s.UsedPercent() } - -// BindingLimit returns the window behind UsedPercent, for display and for -// telling the user *when* an exhausted profile frees up. -func (s Snapshot) BindingLimit() (Limit, bool) { - var worst Limit - found := false - for _, l := range s.Limits { - if !found || l.Percent > worst.Percent { - worst, found = l, true - } - } - return worst, found -} - -// apiUsage mirrors the /api/oauth/usage response. The endpoint carries a good -// deal more (dollar spend, extra-usage credits, unreleased window names); we -// decode only the parts that bear on "can this account do more work right now". -type apiUsage struct { - Limits []apiLimit `json:"limits"` - FiveHour *apiWindow `json:"five_hour"` - SevenDay *apiWindow `json:"seven_day"` -} - -type apiWindow struct { - Utilization float64 `json:"utilization"` - ResetsAt *time.Time `json:"resets_at"` -} - -type apiLimit struct { - Kind string `json:"kind"` - Group string `json:"group"` - Percent float64 `json:"percent"` - Severity string `json:"severity"` - ResetsAt *time.Time `json:"resets_at"` - Scope *struct { - Model *struct { - DisplayName string `json:"display_name"` - } `json:"model"` - } `json:"scope"` -} - -type apiProfile struct { - Account struct { - Email string `json:"email"` - } `json:"account"` -} - -// Fetch reads the usage for one profile (a CLAUDE_CONFIG_DIR), serving a recent -// cached snapshot when there is one. -// -// A missing or expired *credential* fails immediately — that is a configuration -// problem, and papering over it with a cached number would let a profile keep -// receiving tasks it can no longer run. A failed *request* is different: the API -// rate-limits, and a cached snapshot up to StaleTTL old is a far better basis -// for routing than nothing at all. -func (c *Client) Fetch(ctx context.Context, configDir string) (*Snapshot, error) { - creds, err := LoadCredentials(configDir) - if err != nil { - return nil, err - } - if creds.Expired() { - return nil, fmt.Errorf("credentials for %s expired at %s (run a claude session with CLAUDE_CONFIG_DIR=%s to refresh)", - normalizeDir(configDir), creds.ExpiresAt.Format(time.RFC3339), normalizeDir(configDir)) - } - - if snap, ok := c.readCache(configDir, CacheTTL); ok { - return snap, nil - } - - body, err := c.get(ctx, "/api/oauth/usage", creds.AccessToken) - if err != nil { - if stale, ok := c.readCache(configDir, StaleTTL); ok { - stale.Stale = true - return stale, nil - } - return nil, err - } - var raw apiUsage - if err := json.Unmarshal(body, &raw); err != nil { - return nil, fmt.Errorf("decode usage response: %w", err) - } - - snap := &Snapshot{ - ConfigDir: normalizeDir(configDir), - Limits: limitsFrom(raw), - FetchedAt: time.Now(), - } - c.writeCache(configDir, snap) - return snap, nil -} - -// FetchWithAccount is Fetch plus the account email, for surfaces that show the -// user *which* login a profile is. It costs a second round trip, so routing — -// which only needs the numbers — uses plain Fetch. A failure to resolve the -// email is not fatal: the usage numbers are the point. -func (c *Client) FetchWithAccount(ctx context.Context, configDir string) (*Snapshot, error) { - snap, err := c.Fetch(ctx, configDir) - if err != nil { - return nil, err - } - if snap.Email != "" { - return snap, nil // came back from cache with the email already on it - } - if email, err := c.Account(ctx, configDir); err == nil && email != "" { - snap.Email = email - // Re-cache so the next reader gets the email without a second request. - if !snap.Stale { - c.writeCache(configDir, snap) - } - } - return snap, nil -} - -// Account returns the email address a profile is logged in as. -func (c *Client) Account(ctx context.Context, configDir string) (string, error) { - creds, err := LoadCredentials(configDir) - if err != nil { - return "", err - } - body, err := c.get(ctx, "/api/oauth/profile", creds.AccessToken) - if err != nil { - return "", err - } - var p apiProfile - if err := json.Unmarshal(body, &p); err != nil { - return "", fmt.Errorf("decode profile response: %w", err) - } - return p.Account.Email, nil -} - -func (c *Client) get(ctx context.Context, path, token string) ([]byte, error) { - base := c.BaseURL - if base == "" { - base = DefaultBaseURL - } - req, err := http.NewRequestWithContext(ctx, http.MethodGet, strings.TrimRight(base, "/")+path, nil) - if err != nil { - return nil, err - } - req.Header.Set("Authorization", "Bearer "+token) - - httpClient := c.HTTP - if httpClient == nil { - httpClient = &http.Client{Timeout: DefaultTimeout} - } - resp, err := httpClient.Do(req) - if err != nil { - return nil, fmt.Errorf("%s: %w", path, err) - } - defer resp.Body.Close() //nolint:errcheck // read-only GET - body, err := io.ReadAll(io.LimitReader(resp.Body, 1<<20)) - if err != nil { - return nil, fmt.Errorf("%s: read body: %w", path, err) - } - if resp.StatusCode == http.StatusUnauthorized { - return nil, fmt.Errorf("%s: credentials rejected (401) — this profile needs a fresh login", path) - } - if resp.StatusCode == http.StatusTooManyRequests { - // The usage endpoint has its own rate limit, separate from the - // subscription limits it reports. Name it, so nobody reads this as the - // account being out of quota. - return nil, fmt.Errorf("%s: the usage API itself is rate-limiting (429) — this is not the account's quota; retry shortly", path) - } - if resp.StatusCode != http.StatusOK { - return nil, fmt.Errorf("%s: unexpected status %d: %s", path, resp.StatusCode, strings.TrimSpace(string(body))) - } - return body, nil -} - -// limitsFrom flattens the API's usage payload into our Limit list. The modern -// response carries a `limits` array; when it is absent or empty we synthesize -// the two headline windows from the older `five_hour`/`seven_day` objects, so a -// rollback on the API side doesn't leave routing blind. -func limitsFrom(raw apiUsage) []Limit { - if len(raw.Limits) > 0 { - out := make([]Limit, 0, len(raw.Limits)) - for _, l := range raw.Limits { - lim := Limit{ - Kind: l.Kind, - Group: l.Group, - Percent: l.Percent, - Severity: l.Severity, - ResetsAt: l.ResetsAt, - } - if l.Scope != nil && l.Scope.Model != nil { - lim.Scope = l.Scope.Model.DisplayName - } - out = append(out, lim) - } - return out - } - - var out []Limit - if raw.FiveHour != nil { - out = append(out, Limit{Kind: "session", Group: "session", Percent: raw.FiveHour.Utilization, ResetsAt: raw.FiveHour.ResetsAt}) - } - if raw.SevenDay != nil { - out = append(out, Limit{Kind: "weekly_all", Group: "weekly", Percent: raw.SevenDay.Utilization, ResetsAt: raw.SevenDay.ResetsAt}) - } - return out -} - -// Describe renders a one-line human summary of a snapshot, e.g. -// "94% used (session, resets 14:00) — 6% headroom". -func (s Snapshot) Describe() string { - var b strings.Builder - fmt.Fprintf(&b, "%.0f%% used", s.UsedPercent()) - if l, ok := s.BindingLimit(); ok { - b.WriteString(" (") - b.WriteString(limitLabel(l)) - if l.ResetsAt != nil { - fmt.Fprintf(&b, ", resets %s", l.ResetsAt.Local().Format("Mon 15:04")) - } - b.WriteString(")") - } - fmt.Fprintf(&b, " — %.0f%% headroom", s.Headroom()) - return b.String() -} - -// limitLabel turns an API window kind into something readable. -func limitLabel(l Limit) string { - label := l.Kind - switch l.Kind { - case "session": - label = "5-hour session" - case "weekly_all": - label = "weekly" - case "weekly_scoped": - label = "weekly" - if l.Scope != "" { - label = "weekly " + l.Scope - } - } - return label -} diff --git a/internal/claudeusage/usage_test.go b/internal/claudeusage/usage_test.go deleted file mode 100644 index fcada876..00000000 --- a/internal/claudeusage/usage_test.go +++ /dev/null @@ -1,261 +0,0 @@ -package claudeusage - -import ( - "context" - "encoding/json" - "net/http" - "net/http/httptest" - "os" - "path/filepath" - "strings" - "testing" - "time" -) - -// usageBody is a trimmed but structurally faithful /api/oauth/usage response. -const usageBody = `{ - "five_hour": {"utilization": 3.0, "resets_at": "2026-08-15T13:59:59.981513+00:00"}, - "seven_day": {"utilization": 7.0, "resets_at": "2026-08-20T20:59:59.981535+00:00"}, - "limits": [ - {"kind": "session", "group": "session", "percent": 3, "severity": "normal", - "resets_at": "2026-08-15T13:59:59.981513+00:00", "scope": null, "is_active": false}, - {"kind": "weekly_all", "group": "weekly", "percent": 71, "severity": "normal", - "resets_at": "2026-08-20T20:59:59.981535+00:00", "scope": null, "is_active": true}, - {"kind": "weekly_scoped", "group": "weekly", "percent": 12, "severity": "normal", - "resets_at": null, "scope": {"model": {"id": null, "display_name": "Opus"}}, "is_active": false} - ] -}` - -// writeCredsDir makes a config dir holding a .credentials.json, which is the -// non-macOS credential path and the one a test can exercise hermetically. -func writeCredsDir(t *testing.T, token string, expiresAt time.Time) string { - t.Helper() - dir := t.TempDir() - writeCredsInto(t, dir, token, expiresAt) - return dir -} - -// writeCredsInto (re)writes a credential blob into an existing config dir, so a -// test can age a profile's token without changing its identity. -func writeCredsInto(t *testing.T, dir, token string, expiresAt time.Time) { - t.Helper() - blob := map[string]any{ - "claudeAiOauth": map[string]any{ - "accessToken": token, - "subscriptionType": "max", - }, - } - if !expiresAt.IsZero() { - blob["claudeAiOauth"].(map[string]any)["expiresAt"] = expiresAt.UnixMilli() - } - data, err := json.Marshal(blob) - if err != nil { - t.Fatal(err) - } - if err := os.WriteFile(filepath.Join(dir, credentialsFile), data, 0o600); err != nil { - t.Fatal(err) - } -} - -// testClient points the client at a test server and an isolated cache dir, so -// no test can be helped (or hurt) by another test's cached snapshot, nor leave -// anything behind in the real user cache. -func testClient(t *testing.T, srv *httptest.Server) *Client { - t.Helper() - return &Client{BaseURL: srv.URL, HTTP: srv.Client(), CacheDir: t.TempDir()} -} - -func TestFetch_ParsesLimitsAndAuthenticates(t *testing.T) { - var gotAuth, gotPath string - srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - gotAuth, gotPath = r.Header.Get("Authorization"), r.URL.Path - w.Header().Set("Content-Type", "application/json") - _, _ = w.Write([]byte(usageBody)) - })) - defer srv.Close() - - dir := writeCredsDir(t, "tok-123", time.Now().Add(time.Hour)) - client := testClient(t, srv) - - snap, err := client.Fetch(context.Background(), dir) - if err != nil { - t.Fatalf("Fetch: %v", err) - } - if gotAuth != "Bearer tok-123" { - t.Errorf("Authorization = %q", gotAuth) - } - if gotPath != "/api/oauth/usage" { - t.Errorf("path = %q", gotPath) - } - if len(snap.Limits) != 3 { - t.Fatalf("got %d limits, want 3", len(snap.Limits)) - } - if snap.Limits[2].Scope != "Opus" { - t.Errorf("scoped limit scope = %q, want Opus", snap.Limits[2].Scope) - } - if snap.Limits[0].ResetsAt == nil { - t.Error("session limit lost its resets_at") - } -} - -func TestUsedPercentIsTheWorstWindow(t *testing.T) { - // The point of routing is to avoid the window that stops work first. An - // account 3% into its session but 71% into its week has 29% of headroom, - // not 97% — averaging or taking the session window alone would send tasks - // to an account about to run out. - snap := Snapshot{Limits: []Limit{ - {Kind: "session", Percent: 3}, - {Kind: "weekly_all", Percent: 71}, - {Kind: "weekly_scoped", Percent: 12}, - }} - if got := snap.UsedPercent(); got != 71 { - t.Errorf("UsedPercent = %v, want 71", got) - } - if got := snap.Headroom(); got != 29 { - t.Errorf("Headroom = %v, want 29", got) - } - binding, ok := snap.BindingLimit() - if !ok || binding.Kind != "weekly_all" { - t.Errorf("BindingLimit = %+v, ok=%v, want weekly_all", binding, ok) - } -} - -func TestUsedPercentOnEmptySnapshot(t *testing.T) { - var snap Snapshot - if got := snap.UsedPercent(); got != 0 { - t.Errorf("UsedPercent = %v, want 0", got) - } - if _, ok := snap.BindingLimit(); ok { - t.Error("BindingLimit should report not-found with no limits") - } -} - -func TestLimitsFallBackToLegacyWindows(t *testing.T) { - // If the API ever drops back to the older shape (no `limits` array), routing - // must keep working off five_hour/seven_day rather than seeing 0% used and - // happily piling work onto an exhausted account. - body := `{"five_hour": {"utilization": 88.0, "resets_at": null}, - "seven_day": {"utilization": 40.0, "resets_at": null}}` - srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - _, _ = w.Write([]byte(body)) - })) - defer srv.Close() - - dir := writeCredsDir(t, "tok", time.Time{}) - client := testClient(t, srv) - snap, err := client.Fetch(context.Background(), dir) - if err != nil { - t.Fatalf("Fetch: %v", err) - } - if len(snap.Limits) != 2 { - t.Fatalf("got %d limits, want 2 synthesized", len(snap.Limits)) - } - if got := snap.UsedPercent(); got != 88 { - t.Errorf("UsedPercent = %v, want 88", got) - } -} - -func TestFetch_ExpiredTokenIsNotSentToTheAPI(t *testing.T) { - called := false - srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - called = true - })) - defer srv.Close() - - dir := writeCredsDir(t, "stale", time.Now().Add(-time.Hour)) - client := testClient(t, srv) - _, err := client.Fetch(context.Background(), dir) - if err == nil { - t.Fatal("expected an error for an expired token") - } - if called { - t.Error("expired token should not reach the API") - } - if !strings.Contains(err.Error(), "expired") { - t.Errorf("error should say the credentials expired: %v", err) - } -} - -func TestFetch_UnauthorizedIsExplainedNotDumped(t *testing.T) { - srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - w.WriteHeader(http.StatusUnauthorized) - _, _ = w.Write([]byte(`{"type":"error","error":{"type":"authentication_error"}}`)) - })) - defer srv.Close() - - dir := writeCredsDir(t, "tok", time.Time{}) - client := testClient(t, srv) - _, err := client.Fetch(context.Background(), dir) - if err == nil || !strings.Contains(err.Error(), "fresh login") { - t.Errorf("401 should point at re-login, got %v", err) - } -} - -func TestFetch_MissingCredentialsErrors(t *testing.T) { - client := &Client{BaseURL: "http://127.0.0.1:1", HTTP: &http.Client{Timeout: time.Second}, CacheDir: t.TempDir()} - _, err := client.Fetch(context.Background(), t.TempDir()) - if err == nil || !strings.Contains(err.Error(), "no Claude credentials") { - t.Errorf("want a missing-credentials error, got %v", err) - } -} - -func TestLoadCredentials_EmptyTokenCountsAsMissing(t *testing.T) { - // A half-migrated profile leaves a credential blob with an empty token. - // Treating it as found would surface as a baffling 401 later. - dir := writeCredsDir(t, "", time.Time{}) - if _, err := LoadCredentials(dir); err == nil { - t.Error("an empty access token should not count as credentials") - } -} - -func TestKeychainServiceDerivation(t *testing.T) { - // Claude Code namespaces each config dir's keychain entry by the first 8 hex - // digits of the SHA-256 of the absolute path. Pinning a known value here is - // what catches the derivation drifting: get it wrong and every profile - // silently reports "no credentials" on a Mac. - if got, want := KeychainService("/Users/bruno/.claude-ik"), "Claude Code-credentials-eaf7266a"; got != want { - t.Errorf("KeychainService = %q, want %q", got, want) - } - if got, want := KeychainService("/Users/bruno/.claude"), "Claude Code-credentials-5561fe67"; got != want { - t.Errorf("KeychainService = %q, want %q", got, want) - } - // A trailing slash or a redundant segment is the same dir, so it must hash - // the same — Claude Code hashed its own cleaned path. - if KeychainService("/Users/bruno/.claude-ik/") != KeychainService("/Users/bruno/.claude-ik") { - t.Error("trailing slash changed the keychain service name") - } - if KeychainService("/Users/bruno/foo/../.claude-ik") != KeychainService("/Users/bruno/.claude-ik") { - t.Error("uncleaned path changed the keychain service name") - } -} - -func TestKeychainServiceExpandsTilde(t *testing.T) { - home, err := os.UserHomeDir() - if err != nil { - t.Skip("no home dir") - } - if KeychainService("~/.claude-x") != KeychainService(filepath.Join(home, ".claude-x")) { - t.Error("~ was not expanded before hashing") - } -} - -func TestDescribe(t *testing.T) { - reset := time.Date(2026, 8, 20, 20, 0, 0, 0, time.UTC) - snap := Snapshot{Limits: []Limit{ - {Kind: "session", Percent: 3}, - {Kind: "weekly_all", Percent: 71, ResetsAt: &reset}, - }} - got := snap.Describe() - for _, want := range []string{"71% used", "weekly", "29% headroom", "resets"} { - if !strings.Contains(got, want) { - t.Errorf("Describe() = %q, missing %q", got, want) - } - } -} - -func TestDescribeNamesTheScopedModel(t *testing.T) { - snap := Snapshot{Limits: []Limit{{Kind: "weekly_scoped", Percent: 95, Scope: "Opus"}}} - if got := snap.Describe(); !strings.Contains(got, "weekly Opus") { - t.Errorf("Describe() = %q, want the model named", got) - } -}