diff --git a/README.md b/README.md index 8830674..9249039 100644 --- a/README.md +++ b/README.md @@ -19,6 +19,7 @@ replaced by default. - [First run](#first-run) - [Supported harnesses](#supported-harnesses) - [Usage](#usage) +- [Cost attribution](#cost-attribution) - [Merge or overwrite](#merge-or-overwrite) - [Backups and how to revert](#backups-and-how-to-revert) - [Keys](#keys) @@ -88,11 +89,12 @@ cannot be configured. The footer always shows the keys available on the current ↑/↓ move · space configure · r refresh · q/esc quit ``` -**Choose a model, then choose how to write the config** +**Choose a model, then how requests are attributed, then how to write the config** The model list is fetched from your account with your key, so it reflects the models and policies -you actually have access to. Then choose **merge** (recommended) or **overwrite**. The CLI writes -the files, the row flips to `[✓] active`, and the header count goes up. +you actually have access to. Then decide whether requests should say where they came from, see +[Cost attribution](#cost-attribution). Finally choose **merge** (recommended) or **overwrite**. The +CLI writes the files, the row flips to `[✓] active`, and the header count goes up. **Restart the harness** @@ -114,8 +116,9 @@ Pi and Hermes are configured against the native Anthropic Messages format, which Requesty apply [automatic prompt caching](https://docs.requesty.ai/features/auto-caching) to those harnesses. -Where a harness supports custom headers, the CLI also sets an `X-Title` header naming the tool, so -the [Requesty dashboard](https://app.requesty.ai/analytics) can break spend down per harness. +Where a harness supports custom headers, the CLI sets an `X-Title` header naming the tool, so the +[Requesty dashboard](https://app.requesty.ai/analytics) can break spend down per harness. The same +headers carry the repository and branch when you turn on [cost attribution](#cost-attribution). DeepSeek Harness needs the models of a custom provider listed in its own settings, so the CLI writes the model you selected. You can add more Requesty models later inside the harness, under @@ -128,6 +131,48 @@ Above the harness list, the CLI shows spend, requests and tokens for the last 30 you onboarded with, and refreshes them on demand. Full breakdowns by model, user and tool live in the [Requesty dashboard](https://app.requesty.ai/analytics). +## Cost attribution + +By default a request tells Requesty which harness it came from and nothing else. The wizard can +also attribute each request to the repository and branch you are working in, and to you, so a bill +can be read per project rather than per key: + +| Header | Value | +| --- | --- | +| `X-Requesty-Repo` | The `origin` remote of the current directory, without the scheme or the `.git` suffix, for example `github.com/requestyai/cli` | +| `X-Requesty-Branch` | The branch checked out in the current directory | +| `X-Requesty-User` | Your username on this machine | + +Requesty turns each of these into a dimension you can group by in the +[dashboard](https://app.requesty.ai/analytics) and the management API, as `extra.Repo`, +`extra.Branch` and `extra.User`, and strips the headers before forwarding the request to a model +provider. Outside a repository, `Repo` and `Branch` are sent as `none`. The harness itself is +already named by `X-Title`, so there is no separate header for it. + +**This is off unless you ask for it.** Repository and branch names are often private, so the +attribution step starts on `Keep requests unattributed` and you have to select the other row. + +**How the repository and branch are read.** They depend on where a harness runs, not on where the +CLI ran, so they cannot be written into a config file once and be right afterwards. Instead: + +- **Pi** runs a `git` command itself on every launch, from its own configuration. Nothing else is + needed. +- **Claude Code, Codex and OpenCode** read them from their environment. Choosing attribution writes + `~/.requesty/shell/attribution.sh` (or `attribution.fish`) and adds a marked block to your + `.zshrc`, `.bashrc` or `config.fish` that loads it. The hook refreshes the values before each + prompt, so a `git checkout` is picked up without opening a new shell. Open a new shell once after + the change. +- **DeepSeek Harness** can only be given fixed header values, so it carries `X-Requesty-User` and + leaves the repository and branch out rather than reporting a stale one. +- **Hermes** ignores the headers its configuration asks for when it speaks the Anthropic Messages + format, which is the format Requesty uses for it, so it cannot be attributed at all. + +To stop sending them, run the wizard again and pick `Keep requests unattributed`, which takes the +headers back out of that harness's configuration. Claude Code has none of its own to take out, as it +reads them from the shell hook, so it also needs the hook removed: delete the block marked +`# --- Requesty attribution ---` from your shell startup file, along with `~/.requesty/shell`. The +hook is otherwise left in place, because one hook serves every harness. + ## Merge or overwrite The last step of the wizard asks how to write the configuration. diff --git a/internal/attribution/attribution.go b/internal/attribution/attribution.go new file mode 100644 index 0000000..c29fc0a --- /dev/null +++ b/internal/attribution/attribution.go @@ -0,0 +1,211 @@ +// Package attribution builds the request attribution headers a harness sends to +// Requesty, so spend can be grouped by the repository and branch a request came +// from and by the person who made it. +// +// Requesty turns any X-Requesty- request header into an extra. +// dimension on the request, which the usage APIs can group by, and strips the +// header before forwarding the request to a provider. The values reach Requesty +// and nothing else. +// +// Attribution is off unless the user asks for it: the zero Set writes no +// headers, so a harness configured without it sends exactly what it sent before. +package attribution + +import ( + "fmt" + "os/user" + "strings" +) + +const ( + // The headers below arrive at Requesty as extra.Repo, extra.Branch and + // extra.User. + headerRepo = "X-Requesty-Repo" + headerBranch = "X-Requesty-Branch" + headerUser = "X-Requesty-User" + + // The variables below are the ones the shell hook exports for harnesses + // that read header values from the environment. + envRepo = "REQUESTY_REPO" + envBranch = "REQUESTY_BRANCH" + + // unknown stands in for a value the shell could not read, so a harness + // never has to send, or refuse to send, an empty header. + unknown = "none" +) + +// The commands below are written in POSIX shell, and always print a value on a +// successful exit, so the shell hook, the fish hook and Pi can share one +// definition of how a dimension is read. They must stay free of single quotes: +// the fish hook passes them to `sh -c '…'`. +const ( + // repoShell prints owner/name for the repository holding the working + // directory, dropping the transport, any credentials, the host and the .git + // suffix from the origin remote. + repoShell = `url=$(git remote get-url origin 2>/dev/null); url=${url%.git}; url=${url#*://}; url=${url#*@}; url=${url#*[:/]}; echo "${url:-` + unknown + `}"` + + // branchShell prints the checked out branch, and falls back when HEAD is + // detached or the directory is not a repository at all. + branchShell = `branch=$(git symbolic-ref --quiet --short HEAD 2>/dev/null); echo "${branch:-` + unknown + `}"` +) + +// Dimension is one thing a request can be attributed to, carried in one header. +// +// A dimension is either static, and holds the Value the CLI resolved while +// configuring, or dynamic, and holds the Env and Shell a harness needs to read +// the value itself on every launch. The repository and branch have to be +// dynamic: the CLI is configured once, from one directory, while harnesses run +// later from anywhere, so a value baked in now would be confidently wrong. +type Dimension struct { + Header string + Value string + Env string + Shell string +} + +// Dynamic reports whether the value depends on where a harness runs, and so +// cannot be resolved while configuring. +func (d Dimension) Dynamic() bool { + return d.Env != "" +} + +// spell returns the value to write for the dimension, asking reference for the +// dynamic ones. +func (d Dimension) spell(reference Reference) string { + if d.Dynamic() { + return reference(d) + } + + return d.Value +} + +// Reference spells a dynamic dimension the way one harness reads a value it +// cannot be given upfront. Harnesses differ: Codex names an environment +// variable, OpenCode expands an {env:…} placeholder, and Pi runs a command. +type Reference func(Dimension) string + +// EnvName names the environment variable, for Codex, whose env_http_headers +// table holds variable names rather than values. +func EnvName(dimension Dimension) string { + return dimension.Env +} + +// EnvPlaceholder wraps the environment variable in the {env:…} placeholder +// OpenCode expands wherever it appears in its config file. +func EnvPlaceholder(dimension Dimension) string { + return fmt.Sprintf("{env:%s}", dimension.Env) +} + +// ShellCommand asks Pi to run the command itself on every launch. Pi treats a +// header it cannot resolve as a configuration error and refuses to start, so it +// reads the repository and branch directly instead of relying on a shell that +// may never have loaded the hook. +func ShellCommand(dimension Dimension) string { + return "!" + dimension.Shell +} + +// Set is the dimensions to attribute requests by. A nil Set means the user did +// not opt in, which is the default. +type Set []Dimension + +// New builds the dimensions for this machine. +func New() (Set, error) { + name, err := currentUser() + if err != nil { + return nil, err + } + + return Set{ + {Header: headerRepo, Env: envRepo, Shell: repoShell}, + {Header: headerBranch, Env: envBranch, Shell: branchShell}, + {Header: headerUser, Value: name}, + }, nil +} + +// Headers names every header attribution can write, whether or not the current +// set carries it. A config written again is stripped of the ones it is no longer +// asked for, so turning attribution off actually stops them being sent. +func Headers() []string { + return []string{headerRepo, headerBranch, headerUser} +} + +// Static adds the dimensions the CLI already resolved to headers, for harnesses +// that can only be given fixed values. +func (s Set) Static(headers map[string]string) map[string]string { + for _, dimension := range s { + if !dimension.Dynamic() { + headers = put(headers, dimension.Header, dimension.Value) + } + } + + return headers +} + +// Dynamic returns the dimensions a harness has to read on every launch, spelled +// with reference. Codex keeps those apart from its fixed headers, so they come +// back on their own. +func (s Set) Dynamic(reference Reference) map[string]string { + var headers map[string]string + for _, dimension := range s { + if dimension.Dynamic() { + headers = put(headers, dimension.Header, reference(dimension)) + } + } + + return headers +} + +// All adds every dimension to headers, spelling the dynamic ones with +// reference, for harnesses that hold resolved and unresolved values together. +func (s Set) All(headers map[string]string, reference Reference) map[string]string { + for _, dimension := range s { + headers = put(headers, dimension.Header, dimension.spell(reference)) + } + + return headers +} + +// currentUser names the person a harness runs as. It is resolved while +// configuring because, unlike the repository and branch, it does not change with +// the directory a harness runs from. +func currentUser() (string, error) { + current, err := user.Current() + if err != nil { + return "", fmt.Errorf("failed to resolve current user: %w", err) + } + + return sanitize(current.Username), nil +} + +// sanitize keeps only the characters a header value may safely carry, so a user +// name can never break the header list the shell hook writes. +func sanitize(value string) string { + var sanitized strings.Builder + for _, char := range value { + switch { + case char >= 'a' && char <= 'z', + char >= 'A' && char <= 'Z', + char >= '0' && char <= '9', + char == '.', char == '_', char == '-', char == '@': + sanitized.WriteRune(char) + default: + sanitized.WriteRune('-') + } + } + + if sanitized.Len() == 0 { + return unknown + } + + return sanitized.String() +} + +// put adds a header, allocating the map when the caller had none to add to. +func put(headers map[string]string, name, value string) map[string]string { + if headers == nil { + headers = make(map[string]string, 1) + } + headers[name] = value + + return headers +} diff --git a/internal/attribution/attribution_test.go b/internal/attribution/attribution_test.go new file mode 100644 index 0000000..b0a681e --- /dev/null +++ b/internal/attribution/attribution_test.go @@ -0,0 +1,166 @@ +package attribution + +import ( + "os/exec" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestSetSpellsDimensionsPerHarness(t *testing.T) { + set := Set{ + {Header: headerRepo, Env: envRepo, Shell: repoShell}, + {Header: headerUser, Value: "ada"}, + } + + assert.Equal(t, map[string]string{ + "X-Title": "Pi", + "X-Requesty-User": "ada", + "X-Requesty-Repo": "!" + repoShell, + }, set.All(map[string]string{"X-Title": "Pi"}, ShellCommand)) + + assert.Equal(t, map[string]string{ + "X-Requesty-Repo": "{env:REQUESTY_REPO}", + "X-Requesty-User": "ada", + }, set.All(nil, EnvPlaceholder)) + + assert.Equal(t, map[string]string{ + "X-Requesty-User": "ada", + }, set.Static(nil)) + + assert.Equal(t, map[string]string{ + "X-Requesty-Repo": "REQUESTY_REPO", + }, set.Dynamic(EnvName)) +} + +func TestEmptySetWritesNoHeaders(t *testing.T) { + var set Set + + assert.Nil(t, set.Static(nil)) + assert.Nil(t, set.Dynamic(EnvName)) + assert.Nil(t, set.All(nil, ShellCommand)) + assert.Equal(t, map[string]string{"X-Title": "Pi"}, + set.All(map[string]string{"X-Title": "Pi"}, ShellCommand)) +} + +func TestNewResolvesTheUserAndDefersTheRest(t *testing.T) { + set, err := New() + require.NoError(t, err) + + byHeader := make(map[string]Dimension, len(set)) + for _, dimension := range set { + byHeader[dimension.Header] = dimension + } + + assert.False(t, byHeader[headerUser].Dynamic()) + assert.NotEmpty(t, byHeader[headerUser].Value) + assert.True(t, byHeader[headerRepo].Dynamic()) + assert.True(t, byHeader[headerBranch].Dynamic()) +} + +func TestSanitizeKeepsHeaderValuesSafe(t *testing.T) { + assert.Equal(t, "ada.lovelace_1", sanitize("ada.lovelace_1")) + assert.Equal(t, "CORP-ada", sanitize(`CORP\ada`)) + assert.Equal(t, "ada-branch--x", sanitize("ada\nbranch: x")) + assert.Equal(t, unknown, sanitize("")) +} + +// The fish hook passes the shell commands to `sh -c '…'`, which a single quote +// in one of them would cut short. +func TestShellCommandsCarryNoSingleQuotes(t *testing.T) { + for _, command := range []string{repoShell, branchShell} { + assert.NotContains(t, command, "'") + } +} + +// Pi refuses to start when a header value comes back empty, so both commands +// have to print something and exit cleanly wherever they run. +func TestShellCommandsAlwaysPrintAValue(t *testing.T) { + for name, dir := range map[string]string{ + "outside a repository": t.TempDir(), + "inside a repository": gitRepo(t), + } { + t.Run(name, func(t *testing.T) { + for _, command := range []string{repoShell, branchShell} { + output, err := run(t, dir, "sh", "-c", command) + require.NoError(t, err) + assert.NotEmpty(t, output) + } + }) + } +} + +func TestShellCommandsReadTheRepositoryAndBranch(t *testing.T) { + dir := gitRepo(t) + + repo, err := run(t, dir, "sh", "-c", repoShell) + require.NoError(t, err) + assert.Equal(t, "requestyai/cli", repo) + + branch, err := run(t, dir, "sh", "-c", branchShell) + require.NoError(t, err) + assert.Equal(t, "main", branch) +} + +func TestShellCommandsFallBackOnADetachedHead(t *testing.T) { + dir := gitRepo(t) + _, err := run(t, dir, "git", "checkout", "--detach") + require.NoError(t, err) + + branch, err := run(t, dir, "sh", "-c", branchShell) + require.NoError(t, err) + assert.Equal(t, unknown, branch) +} + +func TestRepoShellTrimsEveryRemoteForm(t *testing.T) { + remotes := map[string]string{ + "https://github.com/requestyai/cli.git": "requestyai/cli", + "https://token@github.com/requestyai/cli.git": "requestyai/cli", + "git@github.com:requestyai/cli.git": "requestyai/cli", + "ssh://git@github.com/requestyai/cli.git": "requestyai/cli", + "https://gitlab.com/requestyai/team/cli": "requestyai/team/cli", + } + + for remote, expected := range remotes { + t.Run(remote, func(t *testing.T) { + dir := gitRepo(t) + _, err := run(t, dir, "git", "remote", "set-url", "origin", remote) + require.NoError(t, err) + + repo, err := run(t, dir, "sh", "-c", repoShell) + require.NoError(t, err) + assert.Equal(t, expected, repo) + }) + } +} + +// gitRepo creates a repository on branch main with an origin remote and one +// commit, so HEAD points at a branch the shell commands can read. +func gitRepo(t *testing.T) string { + t.Helper() + + dir := t.TempDir() + for _, args := range [][]string{ + {"git", "init", "--initial-branch", "main"}, + {"git", "remote", "add", "origin", "https://github.com/requestyai/cli.git"}, + {"git", "-c", "user.email=cli@requesty.ai", "-c", "user.name=CLI", + "commit", "--allow-empty", "--message", "init"}, + } { + _, err := run(t, dir, args[0], args[1:]...) + require.NoError(t, err) + } + + return dir +} + +func run(t *testing.T, dir, name string, args ...string) (string, error) { + t.Helper() + + command := exec.Command(name, args...) + command.Dir = dir + output, err := command.Output() + + return strings.TrimSpace(string(output)), err +} diff --git a/internal/attribution/shell.go b/internal/attribution/shell.go new file mode 100644 index 0000000..a59aeb3 --- /dev/null +++ b/internal/attribution/shell.go @@ -0,0 +1,333 @@ +package attribution + +import ( + "fmt" + "os" + "path/filepath" + "slices" + "strings" + + "github.com/requestyai/cli/internal/config" + "github.com/requestyai/cli/internal/fileio" +) + +const ( + // markerStart and markerEnd delimit the block the CLI owns in a startup + // file, so it can update or remove its own lines and leave the rest of the + // file alone. + markerStart = "# --- Requesty attribution ---" + markerEnd = "# --- End Requesty attribution ---" + + // hookDir holds the hook itself, which the startup file only sources. The + // hook can then grow without the startup file changing again. + hookDir = "shell" + + hookPerm = 0o600 + // startupFilePerm applies to a startup file the user does not have yet. + startupFilePerm = 0o644 +) + +// shell is an interactive shell the hook supports. The two flavours differ only +// in syntax: both export the same variables from the same commands. +type shell struct { + // name is the executable the user's SHELL points at. + name string + // startupFile is read for every interactive session, relative to the home + // directory. + startupFile string + // hookFile is where the CLI writes the hook, relative to its own directory. + hookFile string + // hook renders the hook body, and source renders the line loading it. + hook func(Set) string + source func(hookPath string) string +} + +var shells = []shell{ + { + name: "zsh", + startupFile: ".zshrc", + hookFile: "attribution.sh", + hook: posixHook, + source: posixSource, + }, + { + name: "bash", + startupFile: ".bashrc", + hookFile: "attribution.sh", + hook: posixHook, + source: posixSource, + }, + { + name: "fish", + startupFile: filepath.Join(".config", "fish", "config.fish"), + hookFile: "attribution.fish", + hook: fishHook, + source: fishSource, + }, +} + +// ShellHook is where the CLI put the hook and which startup file loads it, so a +// caller can tell the user what changed. +type ShellHook struct { + HookPath string + StartupFilePath string +} + +// InstallShellHook writes the hook that exports the dynamic dimensions and makes +// the user's shell load it. Harnesses that read header values from the +// environment need it; Pi does not, and neither does a harness configured +// without attribution. +// +// It is idempotent: installing again rewrites the hook and leaves one block +// behind in the startup file. +func InstallShellHook(set Set) (ShellHook, error) { + shell, paths, err := resolveShellHook() + if err != nil { + return ShellHook{}, err + } + + if err := fileio.Write(paths.HookPath, []byte(shell.hook(set)), hookPerm); err != nil { + return ShellHook{}, fmt.Errorf("failed to write hook: %w", err) + } + + block := fmt.Sprintf("%s\n%s\n%s\n", markerStart, shell.source(paths.HookPath), markerEnd) + if err := rewriteStartupFile(paths.StartupFilePath, func(content string) string { + return withBlock(content, block) + }); err != nil { + return ShellHook{}, err + } + + return paths, nil +} + +// RemoveShellHook takes the startup file back to how it read before the hook was +// installed and deletes the hook itself. +func RemoveShellHook() error { + _, paths, err := resolveShellHook() + if err != nil { + return err + } + + if err := rewriteStartupFile(paths.StartupFilePath, withoutBlock); err != nil { + return err + } + + if err := os.Remove(paths.HookPath); err != nil && !os.IsNotExist(err) { + return fmt.Errorf("failed to remove hook: %w", err) + } + + return nil +} + +// ShellHookInstalled reports whether the startup file loads the hook, which is +// how the CLI knows the user opted in on an earlier run. +func ShellHookInstalled() (bool, error) { + _, paths, err := resolveShellHook() + if err != nil { + return false, err + } + + content, err := os.ReadFile(paths.StartupFilePath) + if os.IsNotExist(err) { + return false, nil + } + if err != nil { + return false, fmt.Errorf("failed to read startup file: %w", err) + } + + return strings.Contains(string(content), markerStart), nil +} + +// resolveShellHook picks the shell from the environment and works out where its +// files live. +func resolveShellHook() (shell, ShellHook, error) { + name := filepath.Base(os.Getenv("SHELL")) + index := slices.IndexFunc(shells, func(candidate shell) bool { + return candidate.name == name + }) + if index < 0 { + return shell{}, ShellHook{}, fmt.Errorf( + "unsupported shell %q: attribution needs zsh, bash or fish", name) + } + current := shells[index] + + home, err := os.UserHomeDir() + if err != nil { + return shell{}, ShellHook{}, fmt.Errorf("failed to find home directory: %w", err) + } + + dir, err := config.Dir() + if err != nil { + return shell{}, ShellHook{}, err + } + + return current, ShellHook{ + HookPath: filepath.Join(dir, hookDir, current.hookFile), + StartupFilePath: filepath.Join(home, current.startupFile), + }, nil +} + +// rewriteStartupFile applies rewrite to a file the user owns, keeping the mode it +// already has and taking one backup before the first change. +func rewriteStartupFile(path string, rewrite func(content string) string) error { + content, err := os.ReadFile(path) + if err != nil && !os.IsNotExist(err) { + return fmt.Errorf("failed to read startup file: %w", err) + } + + perm, err := fileio.Mode(path, startupFilePerm) + if err != nil { + return err + } + + rewritten := rewrite(string(content)) + if rewritten == string(content) { + return nil + } + + if err := fileio.BackupAndWrite(path, []byte(rewritten), perm); err != nil { + return fmt.Errorf("failed to write startup file: %w", err) + } + + return nil +} + +// withBlock replaces the CLI's block, or appends it when there is none yet. +func withBlock(content, block string) string { + content = strings.TrimRight(withoutBlock(content), "\n") + if content != "" { + content += "\n\n" + } + + return content + block +} + +// withoutBlock drops the CLI's block, and the blank line ahead of it, so +// installing and removing repeatedly cannot grow the file. +func withoutBlock(content string) string { + lines := strings.Split(content, "\n") + kept := make([]string, 0, len(lines)) + + inside := false + for _, line := range lines { + trimmed := strings.TrimSpace(line) + switch { + case trimmed == markerStart: + inside = true + kept = trimTrailingBlankLines(kept) + case inside: + inside = trimmed != markerEnd + default: + kept = append(kept, line) + } + } + + return strings.Join(kept, "\n") +} + +func trimTrailingBlankLines(lines []string) []string { + for len(lines) > 0 && strings.TrimSpace(lines[len(lines)-1]) == "" { + lines = lines[:len(lines)-1] + } + + return lines +} + +func posixSource(hookPath string) string { + return fmt.Sprintf("[ -r %q ] && . %q", hookPath, hookPath) +} + +func fishSource(hookPath string) string { + return fmt.Sprintf("test -r %q && source %q", hookPath, hookPath) +} + +// posixHook renders the hook for zsh and bash. It runs before every prompt +// rather than once per session because the branch changes under a shell that +// never leaves the directory. +func posixHook(set Set) string { + var hook strings.Builder + + hook.WriteString(hookHeader("#")) + hook.WriteString("\n__requesty_attribution() {\n") + for _, dimension := range set.dynamic() { + fmt.Fprintf(&hook, "\texport %s=$(%s)\n", dimension.Env, dimension.Shell) + } + // The header list has to carry real newlines, which a shell keeps inside + // double quotes, and reach the shell unexpanded, which is why it is not + // quoted as a Go string. + fmt.Fprintf(&hook, "\texport ANTHROPIC_CUSTOM_HEADERS=\"%s\"\n}\n", set.customHeaders(shellVar)) + + hook.WriteString(` +# zsh runs precmd hooks before each prompt, bash runs PROMPT_COMMAND. +if [ -n "$ZSH_VERSION" ]; then + autoload -Uz add-zsh-hook + add-zsh-hook precmd __requesty_attribution +elif [ -n "$BASH_VERSION" ]; then + case ";$PROMPT_COMMAND;" in + *";__requesty_attribution;"*) ;; + *) PROMPT_COMMAND="__requesty_attribution${PROMPT_COMMAND:+;$PROMPT_COMMAND}" ;; + esac +fi + +__requesty_attribution +`) + + return hook.String() +} + +// fishHook renders the hook for fish, which reaches the shared commands through +// sh so there is one definition of how a dimension is read. +func fishHook(set Set) string { + var hook strings.Builder + + hook.WriteString(hookHeader("#")) + hook.WriteString("\n# fish emits fish_prompt before each prompt, and the branch can change\n" + + "# without the directory changing.\n" + + "function __requesty_attribution --on-event fish_prompt\n") + for _, dimension := range set.dynamic() { + fmt.Fprintf(&hook, "\tset -gx %s (sh -c '%s')\n", dimension.Env, dimension.Shell) + } + fmt.Fprintf(&hook, "\tset -gx ANTHROPIC_CUSTOM_HEADERS \"%s\"\nend\n\n__requesty_attribution\n", + set.customHeaders(shellVar)) + + return hook.String() +} + +func hookHeader(comment string) string { + return fmt.Sprintf(`%[1]s Requesty request attribution. Written by the Requesty CLI: change it with +%[1]s the CLI rather than by hand, since installing again replaces this file. +%[1]s +%[1]s It exports the repository and branch of the current directory so the +%[1]s harnesses routing through Requesty can send them as X-Requesty-* headers. +%[1]s Requesty strips those headers before forwarding a request to a provider. +`, comment) +} + +// shellVar reads a dynamic dimension from the variable the hook exports, which +// is how the hook builds the header list Claude Code reads. +func shellVar(dimension Dimension) string { + return "$" + dimension.Env +} + +// customHeaders renders the header list Claude Code takes from +// ANTHROPIC_CUSTOM_HEADERS, one "Name: value" pair per line. +func (s Set) customHeaders(reference Reference) string { + lines := make([]string, 0, len(s)) + for _, dimension := range s { + lines = append(lines, dimension.Header+": "+dimension.spell(reference)) + } + + return strings.Join(lines, "\n") +} + +// dynamic returns the dimensions the hook has to read on every prompt. +func (s Set) dynamic() Set { + dimensions := make(Set, 0, len(s)) + for _, dimension := range s { + if dimension.Dynamic() { + dimensions = append(dimensions, dimension) + } + } + + return dimensions +} diff --git a/internal/attribution/shell_test.go b/internal/attribution/shell_test.go new file mode 100644 index 0000000..6a725e4 --- /dev/null +++ b/internal/attribution/shell_test.go @@ -0,0 +1,179 @@ +package attribution + +import ( + "os" + "os/exec" + "path/filepath" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestInstallShellHookIsIdempotent(t *testing.T) { + home := shellHome(t, "zsh") + startupPath := filepath.Join(home, ".zshrc") + require.NoError(t, os.WriteFile(startupPath, []byte("alias ll='ls -l'\n"), 0o644)) + + set, err := New() + require.NoError(t, err) + + hook, err := InstallShellHook(set) + require.NoError(t, err) + assert.Equal(t, startupPath, hook.StartupFilePath) + assert.Equal(t, filepath.Join(home, ".requesty", "shell", "attribution.sh"), hook.HookPath) + + installed, err := ShellHookInstalled() + require.NoError(t, err) + assert.True(t, installed) + + after := readFile(t, startupPath) + assert.Contains(t, after, "alias ll='ls -l'") + assert.Contains(t, after, hook.HookPath) + assert.Equal(t, 1, strings.Count(after, markerStart)) + + // The startup file the user had is kept once, before the first change. + assert.Equal(t, "alias ll='ls -l'\n", readFile(t, startupPath+".requesty.bak")) + + _, err = InstallShellHook(set) + require.NoError(t, err) + assert.Equal(t, after, readFile(t, startupPath)) +} + +func TestRemoveShellHookRestoresTheStartupFile(t *testing.T) { + home := shellHome(t, "bash") + startupPath := filepath.Join(home, ".bashrc") + original := "export EDITOR=vim\n" + require.NoError(t, os.WriteFile(startupPath, []byte(original), 0o644)) + + set, err := New() + require.NoError(t, err) + hook, err := InstallShellHook(set) + require.NoError(t, err) + + require.NoError(t, RemoveShellHook()) + + assert.Equal(t, original, readFile(t, startupPath)) + _, err = os.Stat(hook.HookPath) + assert.True(t, os.IsNotExist(err)) + + installed, err := ShellHookInstalled() + require.NoError(t, err) + assert.False(t, installed) + + // Removing again is not an error, and leaves the file alone. + require.NoError(t, RemoveShellHook()) + assert.Equal(t, original, readFile(t, startupPath)) +} + +func TestInstallShellHookCreatesAMissingStartupFile(t *testing.T) { + home := shellHome(t, "fish") + + set, err := New() + require.NoError(t, err) + hook, err := InstallShellHook(set) + require.NoError(t, err) + + assert.Equal(t, filepath.Join(home, ".config", "fish", "config.fish"), hook.StartupFilePath) + assert.Contains(t, readFile(t, hook.StartupFilePath), "source") + assert.Contains(t, readFile(t, hook.HookPath), "--on-event fish_prompt") +} + +func TestInstallShellHookRejectsAnUnsupportedShell(t *testing.T) { + shellHome(t, "tcsh") + + _, err := InstallShellHook(Set{}) + require.Error(t, err) + assert.Contains(t, err.Error(), "unsupported shell") +} + +// The hook is only useful if a shell that sources it really ends up with the +// dimensions exported, so run it. +func TestPosixHookExportsTheDimensions(t *testing.T) { + home := shellHome(t, "bash") + set, err := New() + require.NoError(t, err) + hook, err := InstallShellHook(set) + require.NoError(t, err) + + repo := gitRepo(t) + output, err := run(t, repo, "sh", "-c", + ". "+hook.HookPath+`; printf '%s|%s|%s' "$REQUESTY_REPO" "$REQUESTY_BRANCH" "$ANTHROPIC_CUSTOM_HEADERS"`) + require.NoError(t, err) + + fields := strings.SplitN(output, "|", 3) + require.Len(t, fields, 3) + assert.Equal(t, "requestyai/cli", fields[0]) + assert.Equal(t, "main", fields[1]) + assert.Equal(t, strings.Join([]string{ + "X-Requesty-Repo: requestyai/cli", + "X-Requesty-Branch: main", + "X-Requesty-User: " + userDimension(t, set).Value, + }, "\n"), fields[2]) + + // Nothing in the hook depends on the home directory it was written from. + assert.NotContains(t, readFile(t, hook.HookPath), filepath.Join(home, ".zshrc")) +} + +func TestFishHookExportsTheDimensions(t *testing.T) { + if _, err := exec.LookPath("fish"); err != nil { + t.Skip("fish is not installed") + } + + shellHome(t, "fish") + set, err := New() + require.NoError(t, err) + hook, err := InstallShellHook(set) + require.NoError(t, err) + + output, err := run(t, gitRepo(t), "fish", "--no-config", "--command", + "source "+hook.HookPath+ + `; printf '%s|%s' "$REQUESTY_REPO" "$ANTHROPIC_CUSTOM_HEADERS"`) + require.NoError(t, err) + + fields := strings.SplitN(output, "|", 2) + require.Len(t, fields, 2) + assert.Equal(t, "requestyai/cli", fields[0]) + assert.Contains(t, fields[1], "X-Requesty-Branch: main") +} + +func TestWithoutBlockLeavesTheRestOfTheFileAlone(t *testing.T) { + content := "export EDITOR=vim\n\n" + markerStart + "\nsource hook\n" + markerEnd + "\n\nalias g=git\n" + + assert.Equal(t, "export EDITOR=vim\n\nalias g=git\n", withoutBlock(content)) +} + +// shellHome points the CLI at a throwaway home directory and shell, so +// installing the hook cannot touch the startup file of whoever runs the tests. +func shellHome(t *testing.T, shell string) string { + t.Helper() + + home := t.TempDir() + t.Setenv("HOME", home) + t.Setenv("SHELL", filepath.Join("/opt/homebrew/bin", shell)) + + return home +} + +func userDimension(t *testing.T, set Set) Dimension { + t.Helper() + + for _, dimension := range set { + if dimension.Header == headerUser { + return dimension + } + } + t.Fatal("the set has no user dimension") + + return Dimension{} +} + +func readFile(t *testing.T, path string) string { + t.Helper() + + content, err := os.ReadFile(path) + require.NoError(t, err) + + return string(content) +} diff --git a/internal/config/config.go b/internal/config/config.go index bf4c04d..32f7a57 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -71,11 +71,21 @@ func Save(config Config) error { return nil } -func configPath() (string, error) { +// Dir is where the CLI keeps everything it owns, settings included. +func Dir() (string, error) { home, err := os.UserHomeDir() if err != nil { return "", fmt.Errorf("failed to find home directory: %w", err) } - return filepath.Join(home, dirName, fileName), nil + return filepath.Join(home, dirName), nil +} + +func configPath() (string, error) { + dir, err := Dir() + if err != nil { + return "", err + } + + return filepath.Join(dir, fileName), nil } diff --git a/internal/fileio/fileio.go b/internal/fileio/fileio.go new file mode 100644 index 0000000..18f8a58 --- /dev/null +++ b/internal/fileio/fileio.go @@ -0,0 +1,131 @@ +// Package fileio updates files the CLI does not own without losing what was +// there before: a write lands atomically, so an interrupted run cannot leave a +// half written config behind, and the first write keeps a copy of the original +// next to it. +package fileio + +import ( + "fmt" + "io/fs" + "os" + "path/filepath" +) + +// BackupSuffix marks the copy taken of a file before the CLI first writes it. +const BackupSuffix = ".requesty.bak" + +// BackupAndWrite copies path once, then replaces it with data. +func BackupAndWrite(path string, data []byte, perm fs.FileMode) error { + if err := Backup(path); err != nil { + return fmt.Errorf("failed to backup file: %w", err) + } + + if err := Write(path, data, perm); err != nil { + return fmt.Errorf("failed to write file: %w", err) + } + + return nil +} + +// Write replaces path with data, creating any missing parent directories. The +// write goes to a temporary file in the same directory and is renamed over the +// original, so a reader sees either the old file or the new one. +func Write(path string, data []byte, perm fs.FileMode) (err error) { + dir := filepath.Dir(path) + if err := os.MkdirAll(dir, 0o700); err != nil { + return fmt.Errorf("failed to create directory: %w", err) + } + + tmp, err := os.CreateTemp(dir, "."+filepath.Base(path)+".tmp-*") + if err != nil { + return fmt.Errorf("failed to create temporary file: %w", err) + } + defer tmp.Close() + defer os.Remove(tmp.Name()) + + if err := tmp.Chmod(perm); err != nil { + return fmt.Errorf("failed to chmod temporary file: %w", err) + } + if _, err := tmp.Write(data); err != nil { + return fmt.Errorf("failed to write temporary file: %w", err) + } + if err := tmp.Sync(); err != nil { + return fmt.Errorf("failed to sync temporary file: %w", err) + } + if err := tmp.Close(); err != nil { + return fmt.Errorf("failed to close temporary file: %w", err) + } + + if err := os.Rename(tmp.Name(), path); err != nil { + return fmt.Errorf("failed to rename temporary file: %w", err) + } + + return nil +} + +// Backup copies path next to itself, keeping its mode. It does nothing when +// path is missing, or when a backup was already taken, so the copy always holds +// the file as it was before the CLI first touched it. +func Backup(path string) error { + backupPath := path + BackupSuffix + + exists, err := Exists(backupPath) + if err != nil { + return fmt.Errorf("failed to check backup exists: %w", err) + } + if exists { + return nil + } + + srcExists, err := Exists(path) + if err != nil { + return fmt.Errorf("failed to check source exists: %w", err) + } + if !srcExists { + return nil + } + + data, err := os.ReadFile(path) + if err != nil { + return fmt.Errorf("failed to read source file: %w", err) + } + + info, err := os.Stat(path) + if err != nil { + return fmt.Errorf("failed to stat source file: %w", err) + } + + if err := os.WriteFile(backupPath, data, info.Mode().Perm()); err != nil { + return fmt.Errorf("failed to write backup file: %w", err) + } + + return nil +} + +// Exists reports whether path is there. A missing file is not an error. +func Exists(path string) (bool, error) { + _, err := os.Stat(path) + switch { + case err == nil: + return true, nil + case os.IsNotExist(err): + return false, nil + default: + return false, err + } +} + +// Mode returns the permissions of path, falling back to fallback when path does +// not exist yet. It lets an update keep the mode a file already has instead of +// imposing one. +func Mode(path string, fallback fs.FileMode) (fs.FileMode, error) { + info, err := os.Stat(path) + if os.IsNotExist(err) { + return fallback, nil + } + if err != nil { + return 0, fmt.Errorf("failed to stat file: %w", err) + } + + return info.Mode().Perm(), nil +} diff --git a/internal/harnesses/attribution_test.go b/internal/harnesses/attribution_test.go new file mode 100644 index 0000000..10764b1 --- /dev/null +++ b/internal/harnesses/attribution_test.go @@ -0,0 +1,281 @@ +package harnesses + +import ( + "encoding/json" + "os" + "path/filepath" + "testing" + + "github.com/pelletier/go-toml/v2" + "github.com/requestyai/cli/internal/attribution" + "github.com/requestyai/cli/internal/config" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "gopkg.in/yaml.v3" +) + +const attributionTestModel = "anthropic/claude-sonnet-4-5" + +// attributionSet holds one dimension of each kind: two that depend on where the +// harness runs, so every harness has to spell them its own way, and one that is +// known when the CLI runs and reaches every harness as a plain value. +func attributionSet() attribution.Set { + return attribution.Set{ + {Header: "X-Requesty-Repo", Env: "REQUESTY_REPO", Shell: "echo repo"}, + {Header: "X-Requesty-Branch", Env: "REQUESTY_BRANCH", Shell: "echo branch"}, + {Header: "X-Requesty-User", Value: "ada"}, + } +} + +func TestPiCarriesAttributionAsShellCommands(t *testing.T) { + eachConfigurePath(t, func(t *testing.T, overwrite bool) { + configDir := configureWithAttribution(t, func(config config.Config, dir string) Harness { + return NewPiHarness(config, dir) + }, overwrite) + + var models piModels + readJSONFile(t, filepath.Join(configDir, "models.json"), &models) + + assert.Equal(t, map[string]string{ + "HTTP-Referer": "https://pi.dev", + "X-Title": "Pi", + "X-Requesty-Repo": "!echo repo", + "X-Requesty-Branch": "!echo branch", + "X-Requesty-User": "ada", + }, models.Providers[piProvider].Headers) + }) +} + +func TestOpenCodeCarriesAttributionAsEnvironmentPlaceholders(t *testing.T) { + eachConfigurePath(t, func(t *testing.T, overwrite bool) { + configDir := configureWithAttribution(t, func(config config.Config, dir string) Harness { + return NewOpenCodeHarness(config, dir) + }, overwrite) + + var settings openCodeConfig + readJSONFile(t, filepath.Join(configDir, "opencode.json"), &settings) + + assert.Equal(t, map[string]string{ + "X-Title": "OpenCode", + "X-Requesty-Repo": "{env:REQUESTY_REPO}", + "X-Requesty-Branch": "{env:REQUESTY_BRANCH}", + "X-Requesty-User": "ada", + }, settings.Providers[openCodeProvider].Options.Headers) + }) +} + +func TestCodexCarriesAttributionAsEnvironmentHeaders(t *testing.T) { + eachConfigurePath(t, func(t *testing.T, overwrite bool) { + configDir := configureWithAttribution(t, func(config config.Config, dir string) Harness { + return NewCodexHarness(config, dir) + }, overwrite) + + var parsed codexConfig + readTOMLFile(t, filepath.Join(configDir, "config.toml"), &parsed) + provider := parsed.ModelProviders[codexModelProvider] + + assert.Equal(t, map[string]string{ + "X-Title": "OpenAI Codex", + "X-Requesty-User": "ada", + }, provider.HTTPHeaders) + + // Codex resolves these on every launch, from what the shell hook exported. + assert.Equal(t, map[string]string{ + "X-Requesty-Repo": "REQUESTY_REPO", + "X-Requesty-Branch": "REQUESTY_BRANCH", + }, provider.EnvHTTPHeaders) + }) +} + +// TestDeepSeekCarriesOnlyStaticAttribution pins that the repository and branch +// are left out rather than frozen to wherever the CLI ran, as the harness has no +// way to read a header value from the environment or a command. +func TestDeepSeekCarriesOnlyStaticAttribution(t *testing.T) { + eachConfigurePath(t, func(t *testing.T, overwrite bool) { + configDir := configureWithAttribution(t, func(config config.Config, dir string) Harness { + return NewDeepSeekHarness(config, dir) + }, overwrite) + + var settings deepseekSettings + readYAMLFile(t, filepath.Join(configDir, "settings.yaml"), &settings) + + assert.Equal(t, map[string]string{ + "HTTP-Referer": "https://requesty.ai", + "X-Title": "DeepSeek Harness", + "X-Requesty-User": "ada", + }, settings.PiAI.Providers[deepseekProvider].Headers) + }) +} + +// TestWithoutAttributionHeadersAreUnchanged pins what a user who did not opt in +// gets: the headers that name the harness to Requesty, and nothing else. +func TestWithoutAttributionHeadersAreUnchanged(t *testing.T) { + eachConfigurePath(t, func(t *testing.T, overwrite bool) { + configDir := t.TempDir() + harness := NewCodexHarness(attributionTestConfig(), configDir) + + require.NoError(t, harness.Configure(ConfigureOptions{ + Model: attributionTestModel, + Overwrite: overwrite, + })) + + configPath := filepath.Join(configDir, "config.toml") + + var parsed codexConfig + readTOMLFile(t, configPath, &parsed) + provider := parsed.ModelProviders[codexModelProvider] + assert.Equal(t, map[string]string{"X-Title": "OpenAI Codex"}, provider.HTTPHeaders) + assert.Empty(t, provider.EnvHTTPHeaders) + + // An empty table would be noise in a file the user reads and edits. + contents, err := os.ReadFile(configPath) + require.NoError(t, err) + assert.NotContains(t, string(contents), "env_http_headers") + }) +} + +// TestAttributionKeepsHeadersAddedByHand covers the merge path only, as the +// overwrite path is the user asking for the file to be replaced. +func TestAttributionKeepsHeadersAddedByHand(t *testing.T) { + configDir := t.TempDir() + modelsPath := filepath.Join(configDir, "models.json") + require.NoError(t, os.WriteFile(modelsPath, []byte(`{ + "providers": { + "requesty": { + "headers": {"X-Team": "platform"} + } + } + }`), 0o600)) + + harness := NewPiHarness(attributionTestConfig(), configDir) + require.NoError(t, harness.Configure(ConfigureOptions{ + Model: attributionTestModel, + Attribution: attributionSet(), + })) + + var models piModels + readJSONFile(t, modelsPath, &models) + assert.Equal(t, "platform", models.Providers[piProvider].Headers["X-Team"]) + assert.Equal(t, "!echo repo", models.Providers[piProvider].Headers["X-Requesty-Repo"]) +} + +// TestOptingOutRemovesAttributionHeaders is what makes attribution something a +// user can take back: configuring again without it has to leave none of the +// headers an earlier run wrote, in a merge as well as an overwrite. +func TestOptingOutRemovesAttributionHeaders(t *testing.T) { + eachConfigurePath(t, func(t *testing.T, overwrite bool) { + configDir := configureWithAttribution(t, func(config config.Config, dir string) Harness { + return NewCodexHarness(config, dir) + }, overwrite) + + harness := NewCodexHarness(attributionTestConfig(), configDir) + require.NoError(t, harness.Configure(ConfigureOptions{ + Model: attributionTestModel, + Overwrite: overwrite, + })) + + configPath := filepath.Join(configDir, "config.toml") + + var parsed codexConfig + readTOMLFile(t, configPath, &parsed) + provider := parsed.ModelProviders[codexModelProvider] + assert.Equal(t, map[string]string{"X-Title": "OpenAI Codex"}, provider.HTTPHeaders) + assert.Empty(t, provider.EnvHTTPHeaders) + + contents, err := os.ReadFile(configPath) + require.NoError(t, err) + assert.NotContains(t, string(contents), "X-Requesty") + assert.NotContains(t, string(contents), "env_http_headers") + }) +} + +// TestOptingOutKeepsHeadersAddedByHand pins that taking the attribution headers +// out of a merged config takes only those, and not a header a user put there. +func TestOptingOutKeepsHeadersAddedByHand(t *testing.T) { + configDir := t.TempDir() + modelsPath := filepath.Join(configDir, "models.json") + require.NoError(t, os.WriteFile(modelsPath, []byte(`{ + "providers": { + "requesty": { + "headers": {"X-Team": "platform"} + } + } + }`), 0o600)) + + harness := NewPiHarness(attributionTestConfig(), configDir) + require.NoError(t, harness.Configure(ConfigureOptions{ + Model: attributionTestModel, + Attribution: attributionSet(), + })) + require.NoError(t, harness.Configure(ConfigureOptions{Model: attributionTestModel})) + + var models piModels + readJSONFile(t, modelsPath, &models) + assert.Equal(t, map[string]string{ + "HTTP-Referer": "https://pi.dev", + "X-Title": "Pi", + "X-Team": "platform", + }, models.Providers[piProvider].Headers) +} + +// eachConfigurePath runs the test against both ways a harness writes its config, +// as the headers have to land the same way whether the file was merged into or +// replaced. +func eachConfigurePath(t *testing.T, test func(t *testing.T, overwrite bool)) { + t.Helper() + + t.Run("merge", func(t *testing.T) { test(t, false) }) + t.Run("overwrite", func(t *testing.T) { test(t, true) }) +} + +// configureWithAttribution configures a harness in a directory of its own and +// returns that directory to read the written files from. +func configureWithAttribution( + t *testing.T, + newHarness func(config.Config, string) Harness, + overwrite bool, +) string { + t.Helper() + + configDir := t.TempDir() + harness := newHarness(attributionTestConfig(), configDir) + + require.NoError(t, harness.Configure(ConfigureOptions{ + Model: attributionTestModel, + Overwrite: overwrite, + Attribution: attributionSet(), + })) + + return configDir +} + +func attributionTestConfig() config.Config { + return config.Config{ + RouterBaseURL: "https://router.requesty.ai", + APIKey: "my-api-key", + } +} + +func readJSONFile(t *testing.T, path string, into any) { + t.Helper() + + contents, err := os.ReadFile(path) + require.NoError(t, err) + require.NoError(t, json.Unmarshal(contents, into)) +} + +func readTOMLFile(t *testing.T, path string, into any) { + t.Helper() + + contents, err := os.ReadFile(path) + require.NoError(t, err) + require.NoError(t, toml.Unmarshal(contents, into)) +} + +func readYAMLFile(t *testing.T, path string, into any) { + t.Helper() + + contents, err := os.ReadFile(path) + require.NoError(t, err) + require.NoError(t, yaml.Unmarshal(contents, into)) +} diff --git a/internal/harnesses/codex.go b/internal/harnesses/codex.go index 40b652c..fe7021d 100644 --- a/internal/harnesses/codex.go +++ b/internal/harnesses/codex.go @@ -8,6 +8,7 @@ import ( "path/filepath" "github.com/pelletier/go-toml/v2" + "github.com/requestyai/cli/internal/attribution" "github.com/requestyai/cli/internal/config" ) @@ -27,9 +28,13 @@ type codexConfig struct { } type codexProvider struct { - Name string `toml:"name"` - BaseURL string `toml:"base_url"` - HTTPHeaders map[string]string `toml:"http_headers"` + Name string `toml:"name"` + BaseURL string `toml:"base_url"` + + // HTTPHeaders holds fixed values, EnvHTTPHeaders the names of environment + // variables Codex reads on every launch. + HTTPHeaders map[string]string `toml:"http_headers"` + EnvHTTPHeaders map[string]string `toml:"env_http_headers,omitempty"` } type codexAuth struct { @@ -133,13 +138,7 @@ func (c *CodexHarness) configureMerge(opts ConfigureOptions) error { "web_search": "live", "personality": "pragmatic", "model_providers": map[string]any{ - codexModelProvider: map[string]any{ - "name": "Requesty", - "base_url": fmt.Sprintf("%s/v1", c.config.RouterBaseURL), - "http_headers": map[string]any{ - "X-Title": "OpenAI Codex", - }, - }, + codexModelProvider: c.provider(opts), }, }) if err != nil { @@ -172,11 +171,10 @@ func (c *CodexHarness) configureOverwrite(opts ConfigureOptions) error { ModelProvider: codexModelProvider, ModelProviders: map[string]codexProvider{ codexModelProvider: { - Name: "Requesty", - BaseURL: fmt.Sprintf("%s/v1", c.config.RouterBaseURL), - HTTPHeaders: map[string]string{ - "X-Title": "OpenAI Codex", - }, + Name: "Requesty", + BaseURL: fmt.Sprintf("%s/v1", c.config.RouterBaseURL), + HTTPHeaders: c.headers(opts), + EnvHTTPHeaders: c.envHeaders(opts), }, }, ModelReasoningEffort: "high", @@ -201,6 +199,33 @@ func (c *CodexHarness) configureOverwrite(opts ConfigureOptions) error { return nil } +// provider describes the Requesty entry Codex routes through. Both header tables +// are patched even when empty, so attribution headers written by an earlier run +// are taken back out; a table left with nothing in it is dropped by the merge. +func (c *CodexHarness) provider(opts ConfigureOptions) map[string]any { + return map[string]any{ + "name": "Requesty", + "base_url": fmt.Sprintf("%s/v1", c.config.RouterBaseURL), + "http_headers": headerPatch(c.headers(opts)), + "env_http_headers": headerPatch(c.envHeaders(opts)), + } +} + +// headers name Codex to Requesty and carry the attribution dimensions it can be +// given upfront. +func (c *CodexHarness) headers(opts ConfigureOptions) map[string]string { + return opts.Attribution.Static(map[string]string{ + "X-Title": "OpenAI Codex", + }) +} + +// envHeaders carry the dimensions that depend on where Codex runs. Codex reads +// them from the environment on every launch, from the variables the shell hook +// exports, which is why they are named here instead of resolved. +func (c *CodexHarness) envHeaders(opts ConfigureOptions) map[string]string { + return opts.Attribution.Dynamic(attribution.EnvName) +} + func (c *CodexHarness) configPath() string { return filepath.Join(c.configDir, "config.toml") } diff --git a/internal/harnesses/deepseek.go b/internal/harnesses/deepseek.go index 0519bc5..7603b0f 100644 --- a/internal/harnesses/deepseek.go +++ b/internal/harnesses/deepseek.go @@ -184,11 +184,8 @@ func (d *DeepSeekHarness) configureMerge(opts ConfigureOptions) error { "baseURL": d.routerBaseURL(), "defaultContextWindow": deepseekContextWindow, "defaultMaxTokens": deepseekMaxTokens, - "headers": map[string]any{ - "HTTP-Referer": "https://requesty.ai", - "X-Title": "DeepSeek Harness", - }, - "models": models, + "headers": headerPatch(d.headers(opts)), + "models": models, }, }, }, @@ -223,10 +220,7 @@ func (d *DeepSeekHarness) configureOverwrite(opts ConfigureOptions) error { BaseURL: d.routerBaseURL(), DefaultContextWindow: deepseekContextWindow, DefaultMaxTokens: deepseekMaxTokens, - Headers: map[string]string{ - "HTTP-Referer": "https://requesty.ai", - "X-Title": "DeepSeek Harness", - }, + Headers: d.headers(opts), Models: []deepseekModelConfig{ {ID: opts.Model}, }, @@ -307,6 +301,17 @@ func (d *DeepSeekHarness) writeCredentials() error { return nil } +// headers name the harness to Requesty and carry the attribution dimensions it +// can be given upfront. The harness has no way to read a header value from the +// environment or a command, so the repository and branch are left out rather +// than frozen to wherever the CLI happened to run. +func (d *DeepSeekHarness) headers(opts ConfigureOptions) map[string]string { + return opts.Attribution.Static(map[string]string{ + "HTTP-Referer": "https://requesty.ai", + "X-Title": "DeepSeek Harness", + }) +} + // routerBaseURL keeps the /v1 suffix the OpenAI Chat Completions format needs. func (d *DeepSeekHarness) routerBaseURL() string { return fmt.Sprintf("%s/v1", d.config.RouterBaseURL) diff --git a/internal/harnesses/files.go b/internal/harnesses/files.go index 6aee69f..2340df1 100644 --- a/internal/harnesses/files.go +++ b/internal/harnesses/files.go @@ -3,14 +3,17 @@ package harnesses import ( "encoding/json" "fmt" - "io/fs" - "os" - "path/filepath" "github.com/pelletier/go-toml/v2" + "github.com/requestyai/cli/internal/attribution" + "github.com/requestyai/cli/internal/fileio" "gopkg.in/yaml.v3" ) +// configFilePerm keeps harness config files readable by their owner only, since +// most of them hold the Requesty API key. +const configFilePerm = 0o600 + func backupAndWriteConfigFileAsJSON(path string, data any) error { dataBytes, err := json.MarshalIndent(data, "", " ") if err != nil { @@ -39,94 +42,27 @@ func backupAndWriteConfigFileAsYAML(path string, data any) error { } func backupAndWriteFile(path string, data []byte) error { - if err := backupFile(path); err != nil { - return fmt.Errorf("failed to backup file: %w", err) - } - - if err := writeFile(path, data, 0o600); err != nil { - return fmt.Errorf("failed to write file: %w", err) - } - - return nil + return fileio.BackupAndWrite(path, data, configFilePerm) } -func writeFile(path string, data []byte, perm fs.FileMode) (err error) { - dir := filepath.Dir(path) - if err := os.MkdirAll(dir, 0o700); err != nil { - return fmt.Errorf("failed to create directory: %w", err) - } - - tmp, err := os.CreateTemp(dir, "."+filepath.Base(path)+".tmp-*") - if err != nil { - return fmt.Errorf("failed to create temporary file: %w", err) - } - defer tmp.Close() - defer os.Remove(tmp.Name()) - - if err := tmp.Chmod(perm); err != nil { - return fmt.Errorf("failed to chmod temporary file: %w", err) - } - if _, err := tmp.Write(data); err != nil { - return fmt.Errorf("failed to write temporary file: %w", err) - } - if err := tmp.Sync(); err != nil { - return fmt.Errorf("failed to sync temporary file: %w", err) - } - if err := tmp.Close(); err != nil { - return fmt.Errorf("failed to close temporary file: %w", err) - } - - if err := os.Rename(tmp.Name(), path); err != nil { - return fmt.Errorf("failed to rename temporary file: %w", err) - } - - return nil +func pathExists(path string) (bool, error) { + return fileio.Exists(path) } -func backupFile(path string) error { - backupPath := path + ".requesty.bak" - - exists, err := pathExists(backupPath) - if err != nil { - return fmt.Errorf("failed to check backup exists: %w", err) - } - if exists { - return nil - } - - srcExists, err := pathExists(path) - if err != nil { - return fmt.Errorf("failed to check source exists: %w", err) - } - if !srcExists { - return nil - } - - data, err := os.ReadFile(path) - if err != nil { - return fmt.Errorf("failed to read source file: %w", err) +// headerPatch turns a header map into the shape a config patch needs: set what is +// asked for, and remove every attribution header that is not, so a user who turns +// attribution back off stops sending the headers an earlier run wrote. +// +// A patch merges recursively only into a map[string]any, so anything else here +// would replace the headers a user added by hand instead of joining them. +func headerPatch(headers map[string]string) map[string]any { + patch := make(map[string]any, len(headers)+len(attribution.Headers())) + for _, name := range attribution.Headers() { + patch[name] = removal{} } - - info, err := os.Stat(path) - if err != nil { - return fmt.Errorf("failed to stat source file: %w", err) - } - - if err := os.WriteFile(backupPath, data, info.Mode().Perm()); err != nil { - return fmt.Errorf("failed to write backup file: %w", err) + for name, value := range headers { + patch[name] = value } - return nil -} - -func pathExists(path string) (bool, error) { - _, err := os.Stat(path) - switch { - case err == nil: - return true, nil - case os.IsNotExist(err): - return false, nil - default: - return false, err - } + return patch } diff --git a/internal/harnesses/harnesses.go b/internal/harnesses/harnesses.go index 0b0ac63..b9db6c0 100644 --- a/internal/harnesses/harnesses.go +++ b/internal/harnesses/harnesses.go @@ -3,6 +3,7 @@ package harnesses import ( "fmt" + "github.com/requestyai/cli/internal/attribution" "github.com/requestyai/cli/internal/config" ) @@ -15,6 +16,11 @@ type Status struct { type ConfigureOptions struct { Model string Overwrite bool + + // Attribution is the set of dimensions to attribute requests by. The zero + // value writes no attribution headers, which is what a user who did not opt + // in gets. + Attribution attribution.Set } type Harness interface { diff --git a/internal/harnesses/hermes.go b/internal/harnesses/hermes.go index 6e1adfb..2d3b1f8 100644 --- a/internal/harnesses/hermes.go +++ b/internal/harnesses/hermes.go @@ -15,7 +15,9 @@ const ( hermesProvider = "requesty" // hermesAPIMode is the native Anthropic Messages format, which lets Requesty - // apply automatic prompt caching to Hermes' large system prompt. + // apply automatic prompt caching to Hermes' large system prompt. In this mode + // Hermes ignores the headers its config file asks for, which makes it the one + // harness that cannot carry request attribution. hermesAPIMode = "anthropic_messages" ) @@ -25,9 +27,8 @@ type hermesConfig struct { } type hermesModel struct { - Default string `yaml:"default"` - Provider string `yaml:"provider"` - DefaultHeaders map[string]string `yaml:"default_headers,omitempty"` + Default string `yaml:"default"` + Provider string `yaml:"provider"` } type hermesCustomProvider struct { @@ -127,10 +128,6 @@ func (h *HermesHarness) configureMerge(opts ConfigureOptions) error { "model": map[string]any{ "default": opts.Model, "provider": hermesProvider, - "default_headers": map[string]any{ - "HTTP-Referer": "https://hermes-agent.nousresearch.com", - "X-Origin-Title": "Hermes", - }, }, }) if err != nil { @@ -157,10 +154,6 @@ func (h *HermesHarness) configureOverwrite(opts ConfigureOptions) error { Model: hermesModel{ Default: opts.Model, Provider: hermesProvider, - DefaultHeaders: map[string]string{ - "HTTP-Referer": "https://hermes-agent.nousresearch.com", - "X-Origin-Title": "Hermes", - }, }, CustomProviders: []hermesCustomProvider{ { diff --git a/internal/harnesses/hermes_test.go b/internal/harnesses/hermes_test.go index 7d51f24..364df57 100644 --- a/internal/harnesses/hermes_test.go +++ b/internal/harnesses/hermes_test.go @@ -56,10 +56,6 @@ custom_providers: "model": map[string]any{ "default": "anthropic/claude-fable-5", "provider": "requesty", - "default_headers": map[string]any{ - "HTTP-Referer": "https://hermes-agent.nousresearch.com", - "X-Origin-Title": "Hermes", - }, }, "custom_providers": []any{ map[string]any{ @@ -139,10 +135,6 @@ func TestHermesHarnessConfigureCreatesMissingConfig(t *testing.T) { "model": map[string]any{ "default": "anthropic/claude-fable-5", "provider": "requesty", - "default_headers": map[string]any{ - "HTTP-Referer": "https://hermes-agent.nousresearch.com", - "X-Origin-Title": "Hermes", - }, }, "custom_providers": []any{ map[string]any{ diff --git a/internal/harnesses/merge.go b/internal/harnesses/merge.go index 88a4450..ff27a80 100644 --- a/internal/harnesses/merge.go +++ b/internal/harnesses/merge.go @@ -111,10 +111,19 @@ func mergeYAMLConfigFileWithOptions(path string, patch map[string]any, options m return data, nil } +// removal marks a patch key the merge should delete rather than set, which is how +// a value the CLI wrote on an earlier run is taken back out of a config. +type removal struct{} + // mergePatch recursively applies patch values while preserving unrelated fields. // It rejects nested patches when the existing value is not an object or table. func mergePatch(destination, patch map[string]any, path string) error { for key, patchValue := range patch { + if _, isRemoval := patchValue.(removal); isRemoval { + delete(destination, key) + continue + } + patchMap, isMap := patchValue.(map[string]any) if !isMap { destination[key] = patchValue @@ -128,8 +137,7 @@ func mergePatch(destination, patch map[string]any, path string) error { destinationValue, exists := destination[key] if !exists || destinationValue == nil { - destination[key] = patchMap - continue + destinationValue = make(map[string]any) } destinationMap, ok := destinationValue.(map[string]any) @@ -139,6 +147,15 @@ func mergePatch(destination, patch map[string]any, path string) error { if err := mergePatch(destinationMap, patchMap, fieldPath); err != nil { return err } + + // A patch of nothing but removals leaves an object with nothing in it, + // which says no more than having no object at all. + if len(destinationMap) == 0 { + delete(destination, key) + continue + } + + destination[key] = destinationMap } return nil diff --git a/internal/harnesses/opencode.go b/internal/harnesses/opencode.go index fd456bd..25045a4 100644 --- a/internal/harnesses/opencode.go +++ b/internal/harnesses/opencode.go @@ -8,6 +8,7 @@ import ( "os/exec" "path/filepath" + "github.com/requestyai/cli/internal/attribution" "github.com/requestyai/cli/internal/config" ) @@ -116,9 +117,7 @@ func (o *OpenCodeHarness) configureMerge(opts ConfigureOptions) error { "options": map[string]any{ "baseURL": o.baseURL(), "apiKey": o.config.APIKey, - "headers": map[string]any{ - "X-Title": "OpenCode", - }, + "headers": headerPatch(o.headers(opts)), }, }, }, @@ -143,9 +142,7 @@ func (o *OpenCodeHarness) configureOverwrite(opts ConfigureOptions) error { Options: openCodeProviderOptions{ BaseURL: o.baseURL(), APIKey: o.config.APIKey, - Headers: map[string]string{ - "X-Title": "OpenCode", - }, + Headers: o.headers(opts), }, }, }, @@ -158,6 +155,15 @@ func (o *OpenCodeHarness) configureOverwrite(opts ConfigureOptions) error { return nil } +// headers name OpenCode to Requesty and carry the attribution dimensions. +// OpenCode expands an {env:…} placeholder anywhere in its config file, so the +// dynamic dimensions are read from what the shell hook exported. +func (o *OpenCodeHarness) headers(opts ConfigureOptions) map[string]string { + return opts.Attribution.All(map[string]string{ + "X-Title": "OpenCode", + }, attribution.EnvPlaceholder) +} + // modelID qualifies the model with the provider, as OpenCode addresses models // as "/". func (o *OpenCodeHarness) modelID(model string) string { diff --git a/internal/harnesses/pi.go b/internal/harnesses/pi.go index 964caf9..916612f 100644 --- a/internal/harnesses/pi.go +++ b/internal/harnesses/pi.go @@ -8,6 +8,7 @@ import ( "os/exec" "path/filepath" + "github.com/requestyai/cli/internal/attribution" "github.com/requestyai/cli/internal/config" ) @@ -143,10 +144,7 @@ func (p *PiHarness) configureMerge(opts ConfigureOptions) error { "baseUrl": p.config.RouterBaseURL, "api": piAPI, "apiKey": p.config.APIKey, - "headers": map[string]any{ - "HTTP-Referer": "https://pi.dev", - "X-Title": "Pi", - }, + "headers": headerPatch(p.headers(opts)), "models": []any{ map[string]any{"id": opts.Model}, }, @@ -186,10 +184,7 @@ func (p *PiHarness) configureOverwrite(opts ConfigureOptions) error { BaseURL: p.config.RouterBaseURL, API: piAPI, APIKey: p.config.APIKey, - Headers: map[string]string{ - "HTTP-Referer": "https://pi.dev", - "X-Title": "Pi", - }, + Headers: p.headers(opts), Models: []piModelConfig{ {ID: opts.Model}, }, @@ -214,6 +209,17 @@ func (p *PiHarness) configureOverwrite(opts ConfigureOptions) error { return nil } +// headers name Pi to Requesty and carry the attribution dimensions. Pi resolves +// each one by running the shell command itself on every launch, rather than +// reading what the shell hook exported, because it refuses to start when a +// header value cannot be resolved and the hook may never have been loaded. +func (p *PiHarness) headers(opts ConfigureOptions) map[string]string { + return opts.Attribution.All(map[string]string{ + "HTTP-Referer": "https://pi.dev", + "X-Title": "Pi", + }, attribution.ShellCommand) +} + func (p *PiHarness) modelsPath() string { return filepath.Join(p.configDir, "models.json") } diff --git a/internal/tui/pages/requesty/dashboard/integrations.go b/internal/tui/pages/requesty/dashboard/integrations.go index 80b17c8..1c8bbe7 100644 --- a/internal/tui/pages/requesty/dashboard/integrations.go +++ b/internal/tui/pages/requesty/dashboard/integrations.go @@ -9,6 +9,7 @@ import ( "charm.land/bubbles/v2/textinput" tea "charm.land/bubbletea/v2" "charm.land/lipgloss/v2" + "github.com/requestyai/cli/internal/attribution" "github.com/requestyai/cli/internal/client" "github.com/requestyai/cli/internal/config" "github.com/requestyai/cli/internal/harnesses" @@ -27,6 +28,7 @@ type integrationWizardStep uint8 const ( integrationModelWizardStep integrationWizardStep = iota + integrationAttributionWizardStep integrationModeWizardStep ) @@ -35,11 +37,12 @@ type integrationWizardState struct { open bool step integrationWizardStep - options harnesses.ConfigureOptions - models []client.Model - modelCursor int - modelSearch textinput.Model - modeCursor int + options harnesses.ConfigureOptions + models []client.Model + modelCursor int + modelSearch textinput.Model + attributionCursor int + modeCursor int modelsErr error configureErr error @@ -75,7 +78,8 @@ type integrationModelsLoadedMsg struct { // integrationConfiguredMsg carries the result of configuring the selected harness type integrationConfiguredMsg struct { - err error + hook attribution.ShellHook + err error } // integrationState holds the integrations list and configuration wizard state. @@ -88,6 +92,10 @@ type integrationState struct { refreshing bool loadErr error + // notice reports what the last configuration changed outside the harness + // itself, which is the startup file the shell hook is sourced from. + notice string + wizard integrationWizardState } @@ -157,11 +165,38 @@ func (m integrationState) loadModels() tea.Msg { return integrationModelsLoadedMsg{models: models, err: err} } +// configure writes the harness config, and the shell hook alongside it when the +// user asked for attribution, as most harnesses read the repository and branch +// from the environment the hook exports. func (m integrationState) configure(harness harnesses.Harness, options harnesses.ConfigureOptions) tea.Cmd { return func() tea.Msg { - err := harness.Configure(options) - return integrationConfiguredMsg{err: err} + hook, err := installAttributionShellHook(options.Attribution) + if err != nil { + return integrationConfiguredMsg{err: err} + } + + if err := harness.Configure(options); err != nil { + return integrationConfiguredMsg{err: err} + } + + return integrationConfiguredMsg{hook: hook} + } +} + +// installAttributionShellHook installs the hook when attribution is on. When it +// is off the hook is left where it is rather than removed, as one hook serves +// every harness and another may still be attributed through it. +func installAttributionShellHook(set attribution.Set) (attribution.ShellHook, error) { + if len(set) == 0 { + return attribution.ShellHook{}, nil + } + + hook, err := attribution.InstallShellHook(set) + if err != nil { + return attribution.ShellHook{}, fmt.Errorf("could not set up attribution: %w", err) } + + return hook, nil } func (m integrationState) update(msg tea.Msg) (integrationState, tea.Cmd) { @@ -192,6 +227,7 @@ func (m integrationState) update(msg tea.Msg) (integrationState, tea.Cmd) { m.wizard.configureErr = fmt.Errorf("could not configure harness: %w", typedMsg.err) return m, nil } + m.notice = shellHookNotice(typedMsg.hook) m.wizard = integrationWizardState{} m.refreshing = true return m, m.load @@ -212,6 +248,7 @@ func (m integrationState) update(msg tea.Msg) (integrationState, tea.Cmd) { case "enter": if m.canConfigure() { var focus tea.Cmd + m.notice = "" m.wizard, focus = newIntegrationWizardState() return m, tea.Batch(m.loadModels, focus) } @@ -229,6 +266,8 @@ func (m integrationState) updateWizard(msg tea.KeyPressMsg) (integrationState, t switch m.wizard.step { case integrationModelWizardStep: return m.updateModelStep(msg) + case integrationAttributionWizardStep: + return m.updateAttributionStep(msg) case integrationModeWizardStep: return m.updateModeStep(msg) default: @@ -255,8 +294,8 @@ func (m integrationState) updateModelStep(msg tea.KeyPressMsg) (integrationState len(models) > 0 && m.cursor < len(m.items) { m.wizard.options.Model = models[m.wizard.modelCursor].ID - m.wizard.step = integrationModeWizardStep - m.wizard.modeCursor = 0 + m.wizard.step = integrationAttributionWizardStep + m.wizard.attributionCursor = 0 m.wizard.configureErr = nil } default: @@ -287,11 +326,48 @@ func (m integrationWizardState) filteredModels() []client.Model { return models } -func (m integrationState) updateModeStep(msg tea.KeyPressMsg) (integrationState, tea.Cmd) { +// updateAttributionStep asks whether requests should carry where they came from. +// The answer is resolved into a set of dimensions here, so a machine that cannot +// name its user is reported before anything is written. +func (m integrationState) updateAttributionStep(msg tea.KeyPressMsg) (integrationState, tea.Cmd) { switch msg.String() { case "esc": m.wizard.step = integrationModelWizardStep m.wizard.options.Model = "" + m.wizard.attributionCursor = 0 + m.wizard.configureErr = nil + case "up", "k": + if m.wizard.attributionCursor > 0 { + m.wizard.attributionCursor-- + } + case "down", "j": + if m.wizard.attributionCursor < 1 { + m.wizard.attributionCursor++ + } + case "enter": + m.wizard.options.Attribution = nil + if m.wizard.attributionCursor == 1 { + set, err := attribution.New() + if err != nil { + m.wizard.configureErr = fmt.Errorf("could not read attribution: %w", err) + return m, nil + } + m.wizard.options.Attribution = set + } + + m.wizard.step = integrationModeWizardStep + m.wizard.modeCursor = 0 + m.wizard.configureErr = nil + } + + return m, nil +} + +func (m integrationState) updateModeStep(msg tea.KeyPressMsg) (integrationState, tea.Cmd) { + switch msg.String() { + case "esc": + m.wizard.step = integrationAttributionWizardStep + m.wizard.options.Attribution = nil m.wizard.modeCursor = 0 m.wizard.configureErr = nil case "up", "k": @@ -314,6 +390,19 @@ func (m integrationState) updateModeStep(msg tea.KeyPressMsg) (integrationState, return m, nil } +// shellHookNotice reports the startup file the hook was added to, as the +// variables it exports only reach a shell started after the change. +func shellHookNotice(hook attribution.ShellHook) string { + if hook.StartupFilePath == "" { + return "" + } + + return fmt.Sprintf( + "Attribution added to %s. Open a new shell for it to take effect.", + hook.StartupFilePath, + ) +} + func (m integrationState) canConfigure() bool { if m.refreshing || m.cursor < 0 || m.cursor >= len(m.items) { return false @@ -442,6 +531,9 @@ func (m integrationState) detail(width int) string { if m.canConfigure() { lines = append(lines, text.LineSeparator, theme.Key.Render("enter")+" "+theme.Footer.Render("to configure")) } + if m.notice != "" { + lines = append(lines, text.LineSeparator, wrap.Render(theme.Good.Render(m.notice))) + } if m.refreshing { lines = append(lines, theme.Muted.Render("Refreshing…")) } @@ -488,6 +580,8 @@ func (m integrationState) wizardPage(inner int) integrationWizardPage { switch m.wizard.step { case integrationModelWizardStep: return m.modelStepPage(inner) + case integrationAttributionWizardStep: + return m.attributionStepPage(inner) case integrationModeWizardStep: return m.modeStepPage(inner) default: @@ -534,6 +628,45 @@ func (m integrationState) modelStepPage(inner int) integrationWizardPage { return page } +func (m integrationState) attributionStepPage(inner int) integrationWizardPage { + body := table.Table{ + Cols: []table.Column{ + {Title: "ATTRIBUTION", Width: inner - 2, Align: table.Left}, + }, + Rows: [][]string{ + {"Keep requests unattributed"}, + {"Attribute requests by repository, branch and user"}, + }, + Cursor: m.wizard.attributionCursor, + Height: 2, + Style: table.CellStyle(m.wizard.attributionCursor), + }.Render() + + description := []string{"Requests carry the name of the harness and nothing more."} + if m.wizard.attributionCursor == 1 { + description = []string{ + "Adds the repository, branch and your username to each request, as " + + "headers Requesty groups spend by and removes before calling a model " + + "provider.", + "The repository and branch are read when the harness starts, from a hook " + + "this adds to your shell startup file.", + } + } + + return integrationWizardPage{ + title: "Choose what requests are attributed to", + body: lipgloss.JoinVertical(lipgloss.Left, + body, + text.LineSeparator, + lipgloss.NewStyle().Width(inner-2).Render( + theme.Muted.Render(strings.Join(description, "\n\n")), + ), + ), + enterHint: "continue", + escapeHint: "back", + } +} + func (m integrationState) modeStepPage(inner int) integrationWizardPage { body := table.Table{ Cols: []table.Column{ diff --git a/internal/tui/pages/requesty/dashboard/integrations_test.go b/internal/tui/pages/requesty/dashboard/integrations_test.go new file mode 100644 index 0000000..da925e3 --- /dev/null +++ b/internal/tui/pages/requesty/dashboard/integrations_test.go @@ -0,0 +1,161 @@ +package dashboard + +import ( + "testing" + + tea "charm.land/bubbletea/v2" + "github.com/requestyai/cli/internal/client" + "github.com/requestyai/cli/internal/harnesses" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// stubHarness records what the wizard asked for instead of writing any files. +type stubHarness struct { + options harnesses.ConfigureOptions +} + +func (s *stubHarness) Name() string { return "Stub" } +func (s *stubHarness) Description() []string { return nil } + +func (s *stubHarness) Status() (harnesses.Status, error) { + return harnesses.Status{Executable: true}, nil +} + +func (s *stubHarness) Configure(options harnesses.ConfigureOptions) error { + s.options = options + return nil +} + +// TestIntegrationWizardLeavesAttributionOff pins that walking the wizard without +// choosing anything configures the harness the way it always was: a user has to +// ask for their repository and branch to be sent. +func TestIntegrationWizardLeavesAttributionOff(t *testing.T) { + harness := &stubHarness{} + state := openWizard(t, harness) + + state = press(t, state, "enter") // pick the model + require.Equal(t, integrationAttributionWizardStep, state.wizard.step) + + state = press(t, state, "enter") // keep requests unattributed + require.Equal(t, integrationModeWizardStep, state.wizard.step) + assert.Empty(t, state.wizard.options.Attribution) + + state = configureFrom(t, state) + assert.Empty(t, harness.options.Attribution) + assert.Equal(t, "anthropic/claude-sonnet-4-5", harness.options.Model) +} + +// TestIntegrationWizardOptsIntoAttribution walks the same steps, choosing the +// second row of the attribution step. +func TestIntegrationWizardOptsIntoAttribution(t *testing.T) { + harness := &stubHarness{} + state := openWizard(t, harness) + + state = press(t, state, "enter") + state = press(t, state, "down") + state = press(t, state, "enter") + require.Equal(t, integrationModeWizardStep, state.wizard.step) + + headers := make([]string, 0, len(state.wizard.options.Attribution)) + for _, dimension := range state.wizard.options.Attribution { + headers = append(headers, dimension.Header) + } + assert.Equal(t, []string{"X-Requesty-Repo", "X-Requesty-Branch", "X-Requesty-User"}, headers) +} + +// TestIntegrationWizardForgetsAttributionOnTheWayBack covers stepping back out of +// the last step, which has to leave nothing behind from the answer given before. +func TestIntegrationWizardForgetsAttributionOnTheWayBack(t *testing.T) { + state := openWizard(t, &stubHarness{}) + + state = press(t, state, "enter") + state = press(t, state, "down") + state = press(t, state, "enter") + require.NotEmpty(t, state.wizard.options.Attribution) + + state = press(t, state, "esc") + assert.Equal(t, integrationAttributionWizardStep, state.wizard.step) + assert.Empty(t, state.wizard.options.Attribution) + + state = press(t, state, "esc") + assert.Equal(t, integrationModelWizardStep, state.wizard.step) + assert.Empty(t, state.wizard.options.Model) +} + +// TestIntegrationWizardRendersEveryStep is a smoke test over the wizard pages, +// as a step that renders nothing is invisible in the flow tests above. +func TestIntegrationWizardRendersEveryStep(t *testing.T) { + state := openWizard(t, &stubHarness{}) + + for _, step := range []integrationWizardStep{ + integrationModelWizardStep, + integrationAttributionWizardStep, + integrationModeWizardStep, + } { + state.wizard.step = step + assert.NotEmpty(t, state.wizardView(96, 40), "step %d rendered nothing", step) + } +} + +// openWizard opens the wizard on a harness with the model catalogue already +// loaded, which is where a user starts once they press enter on the list. +func openWizard(t *testing.T, harness harnesses.Harness) integrationState { + t.Helper() + + state := integrationState{ + items: []integrationItem{{ + harness: harness, + status: harnesses.Status{Executable: true}, + }}, + } + + state.wizard, _ = newIntegrationWizardState() + state.wizard.models = []client.Model{{ID: "anthropic/claude-sonnet-4-5"}} + + return state +} + +// configureFrom presses enter on the last step and applies the message the +// resulting command produces, as the wizard writes files off the update loop. +func configureFrom(t *testing.T, state integrationState) integrationState { + t.Helper() + + state, cmd := state.update(keyPress(t, "enter")) + require.True(t, state.wizard.configuring) + require.NotNil(t, cmd) + + state, _ = state.update(cmd()) + require.False(t, state.wizard.configuring) + require.NoError(t, state.wizard.configureErr) + + return state +} + +func press(t *testing.T, state integrationState, key string) integrationState { + t.Helper() + + state, _ = state.update(keyPress(t, key)) + + return state +} + +func keyPress(t *testing.T, key string) tea.KeyPressMsg { + t.Helper() + + msg := tea.KeyPressMsg{} + switch key { + case "enter": + msg.Code = tea.KeyEnter + case "esc": + msg.Code = tea.KeyEscape + case "down": + msg.Code = tea.KeyDown + default: + t.Fatalf("unsupported key %q", key) + } + + require.Equal(t, key, msg.String()) + + return msg +}