diff --git a/internal/cli/app.go b/internal/cli/app.go index 44fa370ff..160eabc5b 100644 --- a/internal/cli/app.go +++ b/internal/cli/app.go @@ -311,40 +311,35 @@ func runWithDeps(args []string, stdout io.Writer, stderr io.Writer, deps appDeps // cache file on the machine. The refresh itself is fired in exec/TUI startup. modelregistry.EnableModelsDevOverlay() - addDirs, args, err := splitLeadingAddDirFlags(args) + // --add-dir, --theme and --allow-escalation may be written in any order; + // see splitLeadingRootFlags for why they are split together. + root, args, err := splitLeadingRootFlags(args) if err != nil { return writeAppError(stderr, err.Error(), 1) } - // --theme selects the TUI palette non-interactively (auto or any registered - // theme; populates tui.Options.Theme, which resolveThemeMode prefers over - // ZERO_THEME). Re-split --add-dir afterward so it may appear on either side of --theme. - theme, args, err := splitLeadingThemeFlag(args) - if err != nil { - return writeAppError(stderr, err.Error(), 1) - } - moreDirs, args, err := splitLeadingAddDirFlags(args) - if err != nil { - return writeAppError(stderr, err.Error(), 1) - } - addDirs = append(addDirs, moreDirs...) + addDirs, theme, allowEscalation := root.addDirs, root.theme, root.allowEscalation if len(args) == 0 { - return runInteractiveTUI(stderr, deps, agent.PermissionModeAsk, addDirs, theme) + return runInteractiveTUI(stderr, deps, agent.PermissionModeAsk, addDirs, theme, allowEscalation) } - // --add-dir grants an extra write root, and only the interactive TUI and - // exec dispatch paths consume one. Fail loud everywhere else rather than - // silently discarding an explicit grant — including help/version, which - // run no agent and could only ignore it. The allowlist names exactly the - // cases below that forward addDirs; a future subcommand is rejected by - // default until it opts in here. - if len(addDirs) > 0 { - switch args[0] { - case "--skip-permissions-unsafe", "-p", "--prompt", "exec": - // Forwarded by the matching case below. - default: - return writeAppError(stderr, "--add-dir is only supported for the interactive TUI and exec", 1) - } + // --add-dir grants an extra write root and --allow-escalation opts a run + // into mid-run model escalation; only the interactive TUI and exec consume + // either. Both are forwarded to exec below and rejected loudly everywhere + // else rather than silently discarded, including help/version, which run + // no agent and could only ignore them. The allowlist names exactly the + // cases below that forward; a future subcommand is rejected by default + // until it opts in here. + forwardsRootFlags := false + switch args[0] { + case "--skip-permissions-unsafe", "-p", "--prompt", "exec": + forwardsRootFlags = true + } + if len(addDirs) > 0 && !forwardsRootFlags { + return writeAppError(stderr, "--add-dir is only supported for the interactive TUI and exec", 1) + } + if allowEscalation && !forwardsRootFlags { + return writeAppError(stderr, "--allow-escalation is only supported for the interactive TUI and exec", 1) } switch args[0] { @@ -354,26 +349,24 @@ func runWithDeps(args []string, stdout io.Writer, stderr io.Writer, deps appDeps // reach unsafe mode in the shell — and the "!" shell escape (which is // gated behind unsafe) was therefore unreachable. // - // --add-dir may legally appear on either side of the flag, so re-split - // the remaining args and merge with the dirs already collected. Any - // trailing non-flag args were ignored on this path before --add-dir - // existed and still are — but an --add-dir hidden BEHIND one would be - // silently dropped with them, so reject that misplacement loudly. - moreDirs, rest, err := splitLeadingAddDirFlags(args[1:]) - if err != nil { - return writeAppError(stderr, err.Error(), 1) - } - // --theme may appear here too; extract it before the stray-arg checks so it is - // not rejected as an unexpected positional, then re-split --add-dir after it. - skipTheme, rest, err := splitLeadingThemeFlag(rest) + // The root flags may legally appear on either side of this flag, so + // split them again from what follows it and merge with what the root + // already took. Any trailing non-flag args were ignored on this path + // before --add-dir existed and still are, but an --add-dir hidden + // BEHIND one would be silently dropped with them, so reject that + // misplacement loudly below. + more, rest, err := splitLeadingRootFlags(args[1:]) if err != nil { return writeAppError(stderr, err.Error(), 1) } - evenMoreDirs, rest, err := splitLeadingAddDirFlags(rest) - if err != nil { - return writeAppError(stderr, err.Error(), 1) + moreDirs := more.addDirs + // A --theme written before the flag was taken at the root and used to be + // dropped here; one written after it wins, as the last occurrence does. + skipTheme := theme + if more.theme != "" { + skipTheme = more.theme } - moreDirs = append(moreDirs, evenMoreDirs...) + skipAllowEscalation := more.allowEscalation // A misplaced --add-dir anywhere in the remainder is the more specific error, // so check for it across all of rest before rejecting stray args. for _, arg := range rest { @@ -390,7 +383,7 @@ func runWithDeps(args []string, stdout io.Writer, stderr io.Writer, deps appDeps return writeAppError(stderr, "--skip-permissions-unsafe launches the interactive TUI and takes no prompt or subcommand; for a one-shot unsafe run use `zero exec --skip-permissions-unsafe -p \"...\"`", 1) } } - return runInteractiveTUI(stderr, deps, agent.PermissionModeUnsafe, append(append([]string{}, addDirs...), moreDirs...), skipTheme) + return runInteractiveTUI(stderr, deps, agent.PermissionModeUnsafe, append(append([]string{}, addDirs...), moreDirs...), skipTheme, allowEscalation || skipAllowEscalation) case "-h", "--help", "help": if err := writeHelp(stdout); err != nil { return 1 @@ -423,16 +416,16 @@ func runWithDeps(args []string, stdout io.Writer, stderr io.Writer, deps appDeps if len(args) < 2 { return writePromptRequired(stderr) } - // Forward leading --add-dir occurrences so exec's own parser collects them. + // Forward the root flags exec consumes so its own parser collects them. // Use the inline --prompt= form so a prompt whose first character is a // dash (e.g. `zero -p "-foo"`) is taken verbatim instead of being mistaken for // a flag and rejected with "--prompt requires a value" (matches the cron path). - execArgs := append(addDirFlagArgs(addDirs), "--prompt="+args[1]) + execArgs := append(rootFlagArgs(addDirs, allowEscalation), "--prompt="+args[1]) execArgs = append(execArgs, args[2:]...) return runExec(execArgs, stdout, stderr, deps) case "exec": - // Forward leading --add-dir occurrences so exec's own parser collects them. - return runExec(append(addDirFlagArgs(addDirs), args[1:]...), stdout, stderr, deps) + // Forward the root flags exec consumes so its own parser collects them. + return runExec(append(rootFlagArgs(addDirs, allowEscalation), args[1:]...), stdout, stderr, deps) case "completions": return runCompletions(args[1:], stdout, stderr) case "daemon": @@ -694,11 +687,11 @@ func fillAppDeps(deps appDeps) appDeps { return deps } -func runInteractiveTUI(stderr io.Writer, deps appDeps, permissionMode agent.PermissionMode, addDirs []string, theme string) int { - return runInteractiveTUIWithSetup(stderr, deps, permissionMode, addDirs, theme, false) +func runInteractiveTUI(stderr io.Writer, deps appDeps, permissionMode agent.PermissionMode, addDirs []string, theme string, allowEscalation bool) int { + return runInteractiveTUIWithSetup(stderr, deps, permissionMode, addDirs, theme, false, allowEscalation) } -func runInteractiveTUIWithSetup(stderr io.Writer, deps appDeps, permissionMode agent.PermissionMode, addDirs []string, theme string, forceSetup bool) int { +func runInteractiveTUIWithSetup(stderr io.Writer, deps appDeps, permissionMode agent.PermissionMode, addDirs []string, theme string, forceSetup bool, allowEscalation bool) int { // Refresh the models.dev pricing/limits cache in the background when stale; // the overlay is read at registry construction from the cache file, so this // benefits the next run and never blocks or fails this one. @@ -790,6 +783,14 @@ func runInteractiveTUIWithSetup(stderr io.Writer, deps appDeps, permissionMode a registry := newCoreRegistryScoped(workspaceRoot, scope) registerLocalControlTools(registry, workspaceRoot, resolved.LocalControl) + // Mid-run model escalation is opt-in on this surface too. The tool is present + // only when the operator asked for it with --allow-escalation, and the + // switchers that make it do anything ride on the same flag through + // Options.AllowEscalation below. Registering one without the other ships a + // tool the loop will never act on, which looks like a feature and is not. + if allowEscalation { + registry.Register(tools.NewEscalateModelTool()) + } executionRunner := execution.NewRunner(nil) sandboxStore, err := deps.newSandboxStore() if err != nil { @@ -1033,6 +1034,7 @@ func runInteractiveTUIWithSetup(stderr io.Writer, deps appDeps, permissionMode a Specialists: specialistRuntime.specialists, Skills: pluginActivation.skillInfos(deps.skillsDir()), }, + AllowEscalation: allowEscalation, // LoadSkills backs /skills and direct / invocation in the TUI. // It resolves against the same merged set (default dir + plugin skill // roots) as the skill tool and the system-prompt list, re-read per use so @@ -1378,6 +1380,7 @@ Flags: -v, --version Print version -p, --prompt Run a one-shot prompt --add-dir Allow writes in an extra directory (repeatable) + --allow-escalation Let the agent escalate to a stronger model mid-run via escalate_model --skip-permissions-unsafe Launch the interactive shell in unsafe mode (enables the ! shell escape) `) return err @@ -1394,6 +1397,17 @@ func addDirFlagArgs(addDirs []string) []string { return flags } +// rootFlagArgs re-synthesises the root flags exec consumes, in the spelling +// its own parser accepts, so a flag written before the subcommand reaches +// the run exactly as one written after it would. +func rootFlagArgs(addDirs []string, allowEscalation bool) []string { + flags := addDirFlagArgs(addDirs) + if allowEscalation { + flags = append(flags, "--allow-escalation") + } + return flags +} + // splitLeadingAddDirFlags strips leading --add-dir flags from the root // argument list (zero --add-dir [--add-dir ] [subcommand …]). // Subcommands like exec parse their own --add-dir occurrences. @@ -1434,6 +1448,83 @@ func splitLeadingAddDirFlags(args []string) ([]string, []string, error) { return addDirs, args, nil } +// rootFlags is what the leading root flags amount to once every one of them +// has been stripped from the front of the argument list. +type rootFlags struct { + addDirs []string + theme string + allowEscalation bool +} + +// splitLeadingRootFlags strips --add-dir, --theme and --allow-escalation from +// the front of args in whatever order they were written, stopping at the +// first token none of them claims. +// +// EACH SPLITTER STOPS AT THE FIRST TOKEN IT DOES NOT OWN, so running them once +// in a fixed sequence made the order the operator wrote them in load-bearing: +// a flag handled late in the sequence stranded every flag written after it as +// an unknown command, and `zero --allow-escalation --theme auto` exited with +// an argument error instead of launching. Running the sequence until it makes +// no progress accepts every ordering, including a flag repeated on both sides +// of another. +func splitLeadingRootFlags(args []string) (rootFlags, []string, error) { + var flags rootFlags + for { + before := len(args) + addDirs, rest, err := splitLeadingAddDirFlags(args) + if err != nil { + return rootFlags{}, nil, err + } + flags.addDirs = append(flags.addDirs, addDirs...) + theme, rest, err := splitLeadingThemeFlag(rest) + if err != nil { + return rootFlags{}, nil, err + } + if theme != "" { + // The last occurrence wins across passes, as it does within one. + flags.theme = theme + } + allowEscalation, rest, err := splitLeadingAllowEscalationFlag(rest) + if err != nil { + return rootFlags{}, nil, err + } + flags.allowEscalation = flags.allowEscalation || allowEscalation + args = rest + if len(args) == before { + return flags, args, nil + } + } +} + +// splitLeadingAllowEscalationFlag strips a leading --allow-escalation from the +// root argument list, opting the interactive session into mid-run model +// escalation. +// +// OPT-IN, THE SAME WAY exec IS. Escalation moves a run onto a different model, +// which changes what the run costs and which provider sees the conversation, so +// it is a decision the operator makes rather than a default. The exec flag +// already answers this conservatively and the interactive surface should not +// answer it differently. +// +// Bare flag only: repeating it is harmless, and an =value form is rejected so a +// mistyped --allow-escalation=false is a loud error instead of silently enabling +// the thing it was trying to turn off. +func splitLeadingAllowEscalationFlag(args []string) (bool, []string, error) { + allow := false + for len(args) > 0 { + switch { + case args[0] == "--allow-escalation": + allow = true + args = args[1:] + case strings.HasPrefix(args[0], "--allow-escalation="): + return false, nil, errors.New("--allow-escalation takes no value; pass it bare to enable mid-run model escalation, or omit it") + default: + return allow, args, nil + } + } + return allow, args, nil +} + // splitLeadingThemeFlag strips a leading --theme (space or =form) // from the root argument list and validates it against the registered themes. The // last occurrence wins. A value outside the allowed set is a loud error rather than diff --git a/internal/cli/completions.go b/internal/cli/completions.go index 5f9058946..dfae2379c 100644 --- a/internal/cli/completions.go +++ b/internal/cli/completions.go @@ -20,7 +20,7 @@ type completionContext struct { } var completionRoot = completionNode{ - flags: []string{"-h", "--help", "-v", "--version", "-p", "--prompt", "--add-dir", "--theme", "--skip-permissions-unsafe"}, + flags: []string{"-h", "--help", "-v", "--version", "-p", "--prompt", "--add-dir", "--allow-escalation", "--theme", "--skip-permissions-unsafe"}, children: []completionNode{ {names: []string{"exec"}, flags: []string{ "-h", "--help", "-f", "--file", "--image", "--add-dir", "--mode", "-m", "--model", diff --git a/internal/cli/exec.go b/internal/cli/exec.go index 63d22f8bf..751c0b6da 100644 --- a/internal/cli/exec.go +++ b/internal/cli/exec.go @@ -442,57 +442,22 @@ func runExec(args []string, stdout io.Writer, stderr io.Writer, deps appDeps) in // the resolved model and is reassigned by the model switcher on a mid-run // escalation so post-switch turns are attributed to the escalated model. currentModel := resolved.Provider.Model - var modelSwitcher func(context.Context, string) (agent.Provider, error) - if options.allowEscalation { - modelSwitcher = func(_ context.Context, modelID string) (agent.Provider, error) { - // deps.newProvider is wrapped (fillAppDeps) to apply the stored key, so - // the escalated provider is authenticated even though resolved.Provider - // is the pure profile — no per-site key handling here. - switchedProfile := resolved.Provider - switchedProfile.Model = modelID - switchedProvider, err := deps.newProvider(switchedProfile) - if err != nil { - return nil, err - } - // Mirror the agent loop's switch guard (it only reassigns the provider - // when newProvider != nil). Updating currentModel only on a non-nil - // provider keeps usage attribution consistent with whether the loop - // actually switched — a (nil, nil) return leaves both untouched. - if switchedProvider != nil { - currentModel = modelID - } - return switchedProvider, nil - } - } - - // Optimized OpenAI turn sessions (ZERO_OPENAI_TURN_SESSION, default off). - // nil when gated off or the profile is ineligible: agent.Run then wraps the - // provider in its default adapter — the exact code path of today. The - // session switcher is installed only when the run START is optimized, so - // the legacy ModelSwitcher path above stays untouched otherwise. + // Optimized OpenAI turn sessions (ZERO_OPENAI_TURN_SESSION, default off). nil + // when gated off or the profile is ineligible: agent.Run then wraps the + // provider in its default adapter. This is the run's STARTING session provider + // and is used whether or not escalation is enabled. turnSessions, _ := providers.OptimizedTurnSessions(resolved.Provider, provider, providers.Options{}) + // Both switchers come from one shared builder so exec and the interactive TUI + // cannot drift on the nil contracts the agent loop depends on. The session + // switcher is nil unless this run STARTED optimized, which is what keeps a + // default-adapter run on the default adapter. + var modelSwitcher func(context.Context, string) (agent.Provider, error) var modelSessionSwitcher func(context.Context, string) (zeroruntime.TurnSessionProvider, error) - if options.allowEscalation && turnSessions != nil { - modelSessionSwitcher = func(_ context.Context, modelID string) (zeroruntime.TurnSessionProvider, error) { - switchedProfile := resolved.Provider - switchedProfile.Model = modelID - switchedProvider, err := deps.newProvider(switchedProfile) - if err != nil { - return nil, err - } - if switchedProvider == nil { - // The loop treats a nil session source as "no swap" — mirror the - // legacy closure's (nil, nil) contract. - return nil, nil - } - currentModel = modelID - if optimized, ok := providers.OptimizedTurnSessions(switchedProfile, switchedProvider, providers.Options{}); ok { - return optimized, nil - } - // Ineligible switch target: default adapter, but with the switched - // model's resolved capability projection preserved. - return providers.DefaultTurnSessions(switchedProfile, switchedProvider, providers.Options{}), nil - } + if options.allowEscalation { + modelSwitcher, modelSessionSwitcher = providers.EscalationSwitchers( + resolved.Provider, provider, deps.newProvider, + func(modelID string) { currentModel = modelID }, + ) } runMetadata, err := resolveExecRunMetadata(resolved.Provider) diff --git a/internal/cli/root_flag_order_test.go b/internal/cli/root_flag_order_test.go new file mode 100644 index 000000000..864c6d55a --- /dev/null +++ b/internal/cli/root_flag_order_test.go @@ -0,0 +1,211 @@ +package cli + +import ( + "bytes" + "context" + "path/filepath" + "strings" + "testing" + + "github.com/Gitlawb/zero/internal/agent" + "github.com/Gitlawb/zero/internal/config" + "github.com/Gitlawb/zero/internal/sandbox" + "github.com/Gitlawb/zero/internal/tui" + "github.com/Gitlawb/zero/internal/zeroruntime" +) + +// rootFlagOrderings is every ordering of the given flag groups, each group +// kept intact so a flag stays next to its value. +func rootFlagOrderings(groups [][]string) [][]string { + if len(groups) == 0 { + return [][]string{{}} + } + var orderings [][]string + for index, group := range groups { + rest := append(append([][]string{}, groups[:index]...), groups[index+1:]...) + for _, tail := range rootFlagOrderings(rest) { + orderings = append(orderings, append(append([]string{}, group...), tail...)) + } + } + return orderings +} + +// ROOT FLAGS COMPOSE IN ANY ORDER. Each leading-flag splitter stops at the first +// token it does not own, so running them once in a fixed sequence made the +// order the operator wrote them in load-bearing: `zero --allow-escalation +// --theme auto` stranded `--theme auto` as an unknown command and exited with +// an argument error instead of launching. Every ordering of the three root +// flags, with --skip-permissions-unsafe absent or at any position, has to reach +// the TUI with all three applied. +func TestRootFlagsComposeInAnyOrder(t *testing.T) { + extra := t.TempDir() + resolvedExtra, err := filepath.EvalSymlinks(extra) + if err != nil { + t.Fatal(err) + } + groups := [][]string{{"--add-dir", extra}, {"--theme", "auto"}, {"--allow-escalation"}} + cases := rootFlagOrderings(groups) + cases = append(cases, rootFlagOrderings(append(groups, []string{"--skip-permissions-unsafe"}))...) + if len(cases) != 6+24 { + t.Fatalf("SETUP INVALID: %d orderings, want 30", len(cases)) + } + for _, args := range cases { + t.Run(strings.Join(args, " "), func(t *testing.T) { + options := captureTUIOptions(t, args...) + wantMode := agent.PermissionModeAsk + for _, arg := range args { + if arg == "--skip-permissions-unsafe" { + wantMode = agent.PermissionModeUnsafe + } + } + if options.PermissionMode != wantMode { + t.Errorf("PermissionMode = %q, want %q", options.PermissionMode, wantMode) + } + if !options.AllowEscalation { + t.Error("escalation not enabled") + } + if options.Theme != "auto" { + t.Errorf("Theme = %q, want auto", options.Theme) + } + if options.AgentOptions.Sandbox == nil { + t.Fatal("no sandbox engine on the launched options") + } + roots := options.AgentOptions.Sandbox.Scope().Roots() + found := false + for _, root := range roots { + if root == resolvedExtra { + found = true + } + } + if !found { + t.Errorf("scope roots = %v, want the --add-dir root %q", roots, resolvedExtra) + } + }) + } +} + +// The last --theme wins even with another root flag between two of them, and a +// --theme written before --skip-permissions-unsafe reaches the TUI: the unsafe +// path used to launch with only the theme written after the flag, dropping one +// written before it. +func TestRootThemeLastOccurrenceWinsAcrossOtherFlags(t *testing.T) { + for _, testCase := range []struct { + args []string + want string + }{ + {[]string{"--theme", "light", "--allow-escalation", "--theme", "auto"}, "auto"}, + {[]string{"--theme", "auto", "--skip-permissions-unsafe"}, "auto"}, + {[]string{"--theme", "light", "--skip-permissions-unsafe", "--theme", "auto"}, "auto"}, + } { + t.Run(strings.Join(testCase.args, " "), func(t *testing.T) { + options := captureTUIOptions(t, testCase.args...) + if options.Theme != testCase.want { + t.Fatalf("Theme = %q, want %q", options.Theme, testCase.want) + } + }) + } +} + +// THE ROOT FLAG REACHES EXEC. A --allow-escalation written before the +// subcommand is forwarded the way --add-dir is, so `zero --allow-escalation -p +// "..."` runs with the opt-in the help text promises instead of silently +// without it. The negative control is the same run without the flag, which +// must not advertise escalate_model. +func TestRootAllowEscalationForwardsIntoExec(t *testing.T) { + for _, testCase := range []struct { + name string + args []string + }{ + {"exec subcommand", []string{"--allow-escalation", "exec", "say hi"}}, + {"prompt flag", []string{"--allow-escalation", "-p", "say hi"}}, + } { + t.Run(testCase.name, func(t *testing.T) { + if !execAdvertisesEscalateModel(t, testCase.args) { + t.Fatalf("%v ran without escalate_model advertised: the root flag was dropped before exec", testCase.args) + } + if execAdvertisesEscalateModel(t, testCase.args[1:]) { + t.Fatalf("%v advertised escalate_model without the flag", testCase.args[1:]) + } + }) + } +} + +// Everywhere else the flag is rejected loudly rather than discarded, the way +// --add-dir is: help and version run no agent and could only ignore it. +func TestRootAllowEscalationIsRejectedWhereNothingConsumesIt(t *testing.T) { + for _, args := range [][]string{ + {"--allow-escalation", "version"}, + {"--allow-escalation", "help"}, + {"--allow-escalation", "completions", "bash"}, + } { + t.Run(strings.Join(args, " "), func(t *testing.T) { + var stdout, stderr bytes.Buffer + launched := false + exitCode := runWithDeps(args, &stdout, &stderr, appDeps{ + getwd: func() (string, error) { return t.TempDir(), nil }, + runTUI: func(context.Context, tui.Options) int { launched = true; return exitSuccess }, + }) + if exitCode == exitSuccess || launched { + t.Fatalf("exit = %d launched = %v: the flag was discarded instead of rejected", exitCode, launched) + } + if !strings.Contains(stderr.String(), "--allow-escalation is only supported for the interactive TUI and exec") { + t.Fatalf("stderr does not name the flag: %s", stderr.String()) + } + }) + } +} + +// execAdvertisesEscalateModel dispatches args through the root command with a +// provider that records the tools advertised on the first request, and reports +// whether escalate_model was among them. +func execAdvertisesEscalateModel(t *testing.T, args []string) bool { + t.Helper() + isolateCLIUserState(t) + provider := &toolListingProvider{} + var stdout, stderr bytes.Buffer + exitCode := runWithDeps(args, &stdout, &stderr, appDeps{ + getwd: func() (string, error) { return t.TempDir(), nil }, + resolveConfig: func(string, config.Overrides) (config.ResolvedConfig, error) { + return execResolvedConfig(), nil + }, + newProvider: func(config.ProviderProfile) (zeroruntime.Provider, error) { + return provider, nil + }, + newSandboxStore: func() (*sandbox.GrantStore, error) { + return sandbox.NewGrantStore(sandbox.StoreOptions{FilePath: filepath.Join(t.TempDir(), "sandbox-grants.json")}) + }, + }) + if exitCode != exitSuccess { + t.Fatalf("%v: exit = %d, stderr = %s", args, exitCode, stderr.String()) + } + if provider.requests == 0 { + t.Fatalf("SETUP INVALID: %v never reached the provider", args) + } + for _, name := range provider.toolNames { + if name == "escalate_model" { + return true + } + } + return false +} + +// toolListingProvider answers every request with a one-line text reply and +// keeps the tool names advertised on the first one. +type toolListingProvider struct { + requests int + toolNames []string +} + +func (provider *toolListingProvider) StreamCompletion(_ context.Context, request zeroruntime.CompletionRequest) (<-chan zeroruntime.StreamEvent, error) { + provider.requests++ + if provider.requests == 1 { + for _, tool := range request.Tools { + provider.toolNames = append(provider.toolNames, tool.Name) + } + } + ch := make(chan zeroruntime.StreamEvent, 2) + ch <- zeroruntime.StreamEvent{Type: zeroruntime.StreamEventText, Content: "hi"} + ch <- zeroruntime.StreamEvent{Type: zeroruntime.StreamEventDone} + close(ch) + return ch, nil +} diff --git a/internal/cli/setup.go b/internal/cli/setup.go index 766cea69b..41ccbffa2 100644 --- a/internal/cli/setup.go +++ b/internal/cli/setup.go @@ -42,7 +42,7 @@ func runSetup(args []string, stdout io.Writer, stderr io.Writer, deps appDeps) i return exitSuccess } if strings.TrimSpace(options.catalogID) == "" { - return runInteractiveTUIWithSetup(stderr, deps, "", nil, "", true) + return runInteractiveTUIWithSetup(stderr, deps, "", nil, "", true, false) } result, err := saveSetupProvider(deps, tui.SetupSelection{ diff --git a/internal/cli/tui_escalation_test.go b/internal/cli/tui_escalation_test.go new file mode 100644 index 000000000..d8e01eab8 --- /dev/null +++ b/internal/cli/tui_escalation_test.go @@ -0,0 +1,126 @@ +package cli + +import ( + "bytes" + "context" + "strings" + "testing" + + "github.com/Gitlawb/zero/internal/tui" +) + +// captureTUIOptions runs the root command with the TUI launch intercepted, and +// returns the options the interactive session would have started with. +func captureTUIOptions(t *testing.T, args ...string) tui.Options { + t.Helper() + isolateCLIUserState(t) + var captured tui.Options + var stdout, stderr bytes.Buffer + workspace := t.TempDir() + exitCode := runWithDeps(args, &stdout, &stderr, appDeps{ + getwd: func() (string, error) { return workspace, nil }, + runTUI: func(_ context.Context, options tui.Options) int { + captured = options + return exitSuccess + }, + }) + if exitCode != exitSuccess { + t.Fatalf("exitCode = %d stdout=%s stderr=%s", exitCode, stdout.String(), stderr.String()) + } + return captured +} + +func hasEscalateModel(options tui.Options) bool { + if options.AgentOptions.Registry == nil { + return false + } + _, registered := options.AgentOptions.Registry.Get("escalate_model") + return registered +} + +// ESCALATION IS OFF UNLESS THE OPERATOR ASKS FOR IT. +// +// Escalation moves a run onto a different model, changing what it costs and which +// provider sees the conversation, so the interactive surface answers this the same +// conservative way `zero exec --allow-escalation` already does. +func TestInteractiveTUIRegistersEscalateModelOnlyWithTheFlag(t *testing.T) { + if hasEscalateModel(captureTUIOptions(t)) { + t.Error("escalate_model was registered without --allow-escalation") + } + if !hasEscalateModel(captureTUIOptions(t, "--allow-escalation")) { + t.Error("--allow-escalation did not register escalate_model") + } +} + +// THE TOOL AND THE SWITCHERS ARE ONE FEATURE, NOT TWO. +// +// This is the failure the whole change exists to avoid. The agent loop performs a +// switch only when a switcher is wired, so registering escalate_model without +// wiring one ships a tool the model can call and the loop will silently ignore: +// the run reports an escalation that never happened. The reverse is merely dead +// code. Both halves ride on the same flag and this asserts they cannot be +// separated, whichever half a future change touches. +func TestInteractiveTUIEscalationToolAndSwitchersAreWiredTogether(t *testing.T) { + for _, testCase := range []struct { + name string + args []string + want bool + }{ + {name: "default", args: nil, want: false}, + {name: "flagged", args: []string{"--allow-escalation"}, want: true}, + } { + t.Run(testCase.name, func(t *testing.T) { + options := captureTUIOptions(t, testCase.args...) + if got := hasEscalateModel(options); got != testCase.want { + t.Fatalf("escalate_model registered = %v, want %v", got, testCase.want) + } + if options.AllowEscalation != testCase.want { + t.Errorf("AllowEscalation = %v, want %v: the tool is registered on one gate and the switchers on another, so they can drift apart", + options.AllowEscalation, testCase.want) + } + }) + } +} + +// The flag is accepted on either side of --skip-permissions-unsafe, like the +// other root flags that may appear there. +func TestInteractiveTUIAcceptsAllowEscalationAroundUnsafe(t *testing.T) { + for _, args := range [][]string{ + {"--allow-escalation", "--skip-permissions-unsafe"}, + {"--skip-permissions-unsafe", "--allow-escalation"}, + } { + options := captureTUIOptions(t, args...) + if !options.AllowEscalation { + t.Errorf("%v did not enable escalation", args) + } + if !hasEscalateModel(options) { + t.Errorf("%v did not register escalate_model", args) + } + } +} + +// An =value form is a loud error rather than a silent enable, so a mistyped +// --allow-escalation=false cannot turn the feature ON. +func TestAllowEscalationRejectsAValue(t *testing.T) { + var stdout, stderr bytes.Buffer + exitCode := runWithDeps([]string{"--allow-escalation=false"}, &stdout, &stderr, appDeps{ + getwd: func() (string, error) { return t.TempDir(), nil }, + runTUI: func(context.Context, tui.Options) int { return exitSuccess }, + }) + if exitCode == exitSuccess { + t.Fatal("--allow-escalation=false was accepted; a mistyped disable must not silently enable escalation") + } + if !strings.Contains(stderr.String(), "--allow-escalation takes no value") { + t.Errorf("stderr does not explain the flag: %s", stderr.String()) + } +} + +func TestRootHelpDocumentsAllowEscalation(t *testing.T) { + var stdout, stderr bytes.Buffer + if exitCode := runWithDeps([]string{"--help"}, &stdout, &stderr, appDeps{}); exitCode != exitSuccess { + t.Fatalf("exitCode = %d stderr=%s", exitCode, stderr.String()) + } + if !strings.Contains(stdout.String(), "--allow-escalation") { + t.Error("the root help does not mention --allow-escalation, so the flag is undiscoverable") + } +} diff --git a/internal/cli/user_state_isolation_test.go b/internal/cli/user_state_isolation_test.go new file mode 100644 index 000000000..007c96b8d --- /dev/null +++ b/internal/cli/user_state_isolation_test.go @@ -0,0 +1,139 @@ +package cli + +import ( + "os" + "path/filepath" + "strings" + "testing" + + "github.com/Gitlawb/zero/internal/config" +) + +// userStateEnvNames is every per-user base directory the root command can +// resolve against, across platforms. One throwaway root behind all of them +// lands the resolution inside it whichever branch it takes. +var userStateEnvNames = []string{ + "HOME", "USERPROFILE", "APPDATA", "LOCALAPPDATA", + "XDG_CONFIG_HOME", "XDG_DATA_HOME", "XDG_CACHE_HOME", "XDG_STATE_HOME", +} + +// isolateCLIUserState points every user-scoped path the root command touches at +// throwaway directories, and keeps credential work off the host keyring. +// +// A TEMPORARY WORKING DIRECTORY IS NOT ISOLATION. runWithDeps fills the +// dependencies a test leaves out with the production ones, and the launch paths +// read user config, open stores, refresh the models.dev cache and, for an +// interactive run, migrate any inline plaintext API key into the credential +// store before they ever reach an injected callback. Without this, running these +// tests on a developer machine rewrites that developer's config.json and moves +// their key into their keychain, makes a real network call, and reports results +// that depend on whatever providers, MCP servers and plugins that machine has. +func isolateCLIUserState(t *testing.T) { + t.Helper() + root := t.TempDir() + for _, name := range userStateEnvNames { + t.Setenv(name, root) + } + // The models.dev overlay cache, and the background refresh that writes it: + // a wiring test has no business making a network call or leaving a cache + // file behind. + t.Setenv("ZERO_MODELS_CACHE_PATH", filepath.Join(root, "modelsdev.json")) + t.Setenv("ZERO_DISABLE_MODELS_FETCH", "1") + // Credentials resolve keyring-first. This keeps a migrated key in a file + // under the fixture root rather than in the developer's OS keychain. + t.Setenv("ZERO_CRED_STORAGE", "encrypted-file") +} + +// AND THE ISOLATION IS ITSELF PINNED, ONE HELPER AT A TIME. A seeded config +// outside the helper fixture roots stands in for a developer's real one: the +// launch helper must leave it exactly as it found it, inline API key included, +// must not write a credential file beside it, and must have moved the per-user +// paths off that root before running anything. +// +// EACH HELPER GETS ITS OWN SUBTEST because t.Setenv restores at the end of the +// test, not when the helper returns: exercising both in one test would let the +// first helper isolation cover for a second helper that had none. +func TestCLILaunchHelpersLeaveUserStateAlone(t *testing.T) { + for _, testCase := range []struct { + name string + exercise func(*testing.T) + }{ + {"captureTUIOptions", func(t *testing.T) { _ = captureTUIOptions(t, "--allow-escalation") }}, + {"execAdvertisesEscalateModel", func(t *testing.T) { + _ = execAdvertisesEscalateModel(t, []string{"--allow-escalation", "exec", "say hi"}) + }}, + } { + t.Run(testCase.name, func(t *testing.T) { + seedRoot, configPath, seeded, before := seedEmulatedUserConfig(t) + + testCase.exercise(t) + + after, err := os.ReadFile(configPath) + if err != nil { + t.Fatal(err) + } + if string(after) != string(seeded) { + t.Errorf("the seeded user config was rewritten:\n before: %s\n after: %s", seeded, after) + } + entries, err := os.ReadDir(filepath.Dir(configPath)) + if err != nil { + t.Fatal(err) + } + if len(entries) != len(before) { + names := []string{} + for _, entry := range entries { + names = append(names, entry.Name()) + } + t.Errorf("the user config directory gained entries: %v", names) + } + + // The checks above catch what this platform happens to write. This + // one catches the missing isolation itself, so a startup path that + // writes somewhere else, or only sometimes, is covered too. + resolved, err := config.DefaultUserConfigPath() + if err != nil { + t.Fatal(err) + } + if strings.HasPrefix(resolved, seedRoot) { + t.Errorf("user config still resolves inside the seeded root (%s): the helper installed no isolation", resolved) + } + if os.Getenv("ZERO_DISABLE_MODELS_FETCH") == "" { + t.Error("the helper left the background models.dev refresh enabled, so this test can make a network call and write a cache file") + } + if os.Getenv("ZERO_CRED_STORAGE") == "" { + t.Error("the helper left credential storage resolving keyring-first, so a migrated key can reach the host keychain") + } + }) + } +} + +// seedEmulatedUserConfig points the per-user base directories at a throwaway +// root and writes a config carrying an inline plaintext API key there, which is +// what the startup migration rewrites. It returns that root, the config path, +// its bytes, and the directory listing to compare against. +func seedEmulatedUserConfig(t *testing.T) (string, string, []byte, []os.DirEntry) { + t.Helper() + seedRoot := t.TempDir() + for _, name := range userStateEnvNames { + t.Setenv(name, seedRoot) + } + configPath, err := config.DefaultUserConfigPath() + if err != nil { + t.Fatal(err) + } + if !strings.HasPrefix(configPath, seedRoot) { + t.Fatalf("SETUP INVALID: user config resolved to %q, outside the seeded root %q", configPath, seedRoot) + } + if err := os.MkdirAll(filepath.Dir(configPath), 0o755); err != nil { + t.Fatal(err) + } + seeded := []byte(`{"providers":[{"name":"anthropic","providerKind":"anthropic","model":"claude-haiku-4.5","apiKey":"sk-seeded-plaintext-key"}]}`) + if err := os.WriteFile(configPath, seeded, 0o600); err != nil { + t.Fatal(err) + } + before, err := os.ReadDir(filepath.Dir(configPath)) + if err != nil { + t.Fatal(err) + } + return seedRoot, configPath, seeded, before +} diff --git a/internal/providers/escalation.go b/internal/providers/escalation.go new file mode 100644 index 000000000..317cdba13 --- /dev/null +++ b/internal/providers/escalation.go @@ -0,0 +1,89 @@ +package providers + +import ( + "context" + + "github.com/Gitlawb/zero/internal/config" + "github.com/Gitlawb/zero/internal/zeroruntime" +) + +// EscalationSwitchers builds the two mid-run model switchers a run installs when +// it has opted into escalation. +// +// ONE IMPLEMENTATION, TWO SURFACES. exec grew this first, and the interactive TUI +// needs the same behaviour rather than a similar one. The part that is easy to get +// subtly wrong is not the switch, it is the nil handling around it, and a second +// copy of that would only ever be exercised on one surface. +// +// THE NIL CONTRACTS ARE THE CONTRACT. The agent loop reassigns the provider only +// when a switcher returns a non-nil one, so a (nil, nil) return means "no swap" +// and has to leave everything as it was, including whatever the caller tracks +// through onSwitch. onSwitch therefore fires only on a real swap, and may be nil. +// An error is reported to the loop, which records a note and continues on the +// current model. +// +// The session switcher comes back nil unless the run STARTED optimized. A run +// that began on the default adapter stays on it, so escalation cannot quietly +// change the transport underneath a session. +func EscalationSwitchers( + profile config.ProviderProfile, + provider zeroruntime.Provider, + newProvider func(config.ProviderProfile) (zeroruntime.Provider, error), + onSwitch func(modelID string), +) ( + func(context.Context, string) (zeroruntime.Provider, error), + func(context.Context, string) (zeroruntime.TurnSessionProvider, error), +) { + if newProvider == nil { + return nil, nil + } + // The escalated profile is the run's profile with the model replaced, so the + // credential, base URL and headers travel with it. Callers pass a newProvider + // that already applies the stored key, which is why there is no per-site key + // handling here. + switchTo := func(modelID string) (config.ProviderProfile, zeroruntime.Provider, error) { + switched := profile + switched.Model = modelID + provider, err := newProvider(switched) + return switched, provider, err + } + + modelSwitcher := func(_ context.Context, modelID string) (zeroruntime.Provider, error) { + _, switchedProvider, err := switchTo(modelID) + if err != nil { + return nil, err + } + if switchedProvider == nil { + return nil, nil + } + if onSwitch != nil { + onSwitch(modelID) + } + return switchedProvider, nil + } + + turnSessions, _ := OptimizedTurnSessions(profile, provider, Options{}) + if turnSessions == nil { + return modelSwitcher, nil + } + + sessionSwitcher := func(_ context.Context, modelID string) (zeroruntime.TurnSessionProvider, error) { + switchedProfile, switchedProvider, err := switchTo(modelID) + if err != nil { + return nil, err + } + if switchedProvider == nil { + return nil, nil + } + if onSwitch != nil { + onSwitch(modelID) + } + if optimized, ok := OptimizedTurnSessions(switchedProfile, switchedProvider, Options{}); ok { + return optimized, nil + } + // Ineligible target: the default adapter, but carrying the switched + // model's own capability projection rather than the original's. + return DefaultTurnSessions(switchedProfile, switchedProvider, Options{}), nil + } + return modelSwitcher, sessionSwitcher +} diff --git a/internal/providers/escalation_test.go b/internal/providers/escalation_test.go new file mode 100644 index 000000000..8507fb0fe --- /dev/null +++ b/internal/providers/escalation_test.go @@ -0,0 +1,130 @@ +package providers + +import ( + "context" + "errors" + "testing" + + "github.com/Gitlawb/zero/internal/config" + "github.com/Gitlawb/zero/internal/zeroruntime" +) + +type escalationStubProvider struct{ model string } + +func (escalationStubProvider) Name() string { return "stub" } + +func (escalationStubProvider) StreamCompletion(context.Context, zeroruntime.CompletionRequest) (<-chan zeroruntime.StreamEvent, error) { + ch := make(chan zeroruntime.StreamEvent) + close(ch) + return ch, nil +} + +// THE ESCALATED PROFILE IS THE RUN'S PROFILE WITH THE MODEL REPLACED. +// +// Everything else has to travel: the base URL, the credential and the headers. +// Building a bare profile with only the model set would escalate onto an +// unauthenticated endpoint, and the failure would look like a provider outage. +func TestEscalationSwitchersKeepTheProfileAndReplaceOnlyTheModel(t *testing.T) { + profile := config.ProviderProfile{Name: "acme", Model: "small", BaseURL: "https://acme.test", APIKey: "sk-live"} + var asked config.ProviderProfile + switcher, _ := EscalationSwitchers(profile, escalationStubProvider{}, func(p config.ProviderProfile) (zeroruntime.Provider, error) { + asked = p + return escalationStubProvider{model: p.Model}, nil + }, nil) + if switcher == nil { + t.Fatal("no model switcher was built for a run that opted into escalation") + } + if _, err := switcher(context.Background(), "large"); err != nil { + t.Fatalf("switch: %v", err) + } + if asked.Model != "large" { + t.Errorf("escalated to model %q, want large", asked.Model) + } + if asked.BaseURL != profile.BaseURL || asked.APIKey != profile.APIKey || asked.Name != profile.Name { + t.Errorf("the escalated profile lost the run's provider identity: %+v", asked) + } +} + +// A (nil, nil) RETURN MEANS NO SWAP, AND MUST CHANGE NOTHING. +// +// The agent loop reassigns the provider only when it gets a non-nil one, so a nil +// provider leaves the run on its current model. Anything the caller tracks +// through onSwitch has to stay in step with that, or usage gets attributed to a +// model the run never moved to. +func TestEscalationSwitchersDoNotReportASwitchThatDidNotHappen(t *testing.T) { + switched := []string{} + switcher, _ := EscalationSwitchers( + config.ProviderProfile{Model: "small"}, escalationStubProvider{}, + func(config.ProviderProfile) (zeroruntime.Provider, error) { return nil, nil }, + func(modelID string) { switched = append(switched, modelID) }, + ) + provider, err := switcher(context.Background(), "large") + if err != nil { + t.Fatalf("a nil provider is not an error: %v", err) + } + if provider != nil { + t.Errorf("provider = %v, want nil", provider) + } + if len(switched) != 0 { + t.Errorf("onSwitch fired for a swap that never happened: %v", switched) + } +} + +// An error reaches the loop, which records it and stays on the current model. +func TestEscalationSwitchersReportProviderErrors(t *testing.T) { + want := errors.New("no credential for that model") + switched := []string{} + switcher, _ := EscalationSwitchers( + config.ProviderProfile{Model: "small"}, escalationStubProvider{}, + func(config.ProviderProfile) (zeroruntime.Provider, error) { return nil, want }, + func(modelID string) { switched = append(switched, modelID) }, + ) + if _, err := switcher(context.Background(), "large"); !errors.Is(err, want) { + t.Errorf("err = %v, want %v", err, want) + } + if len(switched) != 0 { + t.Errorf("onSwitch fired for a failed switch: %v", switched) + } +} + +// And a real swap does report itself, or the caller's attribution never moves. +func TestEscalationSwitchersReportARealSwitch(t *testing.T) { + switched := []string{} + switcher, _ := EscalationSwitchers( + config.ProviderProfile{Model: "small"}, escalationStubProvider{}, + func(p config.ProviderProfile) (zeroruntime.Provider, error) { + return escalationStubProvider{model: p.Model}, nil + }, + func(modelID string) { switched = append(switched, modelID) }, + ) + if _, err := switcher(context.Background(), "large"); err != nil { + t.Fatalf("switch: %v", err) + } + if len(switched) != 1 || switched[0] != "large" { + t.Errorf("onSwitch recorded %v, want one switch to large", switched) + } +} + +// A RUN THAT STARTED ON THE DEFAULT ADAPTER STAYS ON IT. +// +// The session switcher is installed only when the run START is optimized, so +// escalation cannot quietly change the transport underneath a session that never +// had one. +func TestEscalationSwitchersOmitTheSessionSwitcherForAnUnoptimizedStart(t *testing.T) { + _, sessionSwitcher := EscalationSwitchers( + config.ProviderProfile{Name: "acme", Model: "small"}, escalationStubProvider{}, + func(p config.ProviderProfile) (zeroruntime.Provider, error) { return escalationStubProvider{}, nil }, + nil, + ) + if sessionSwitcher != nil { + t.Error("a run that did not start with optimized turn sessions was given a session switcher") + } +} + +// No provider factory means no escalation rather than a switcher that panics. +func TestEscalationSwitchersRequireAProviderFactory(t *testing.T) { + switcher, sessionSwitcher := EscalationSwitchers(config.ProviderProfile{}, escalationStubProvider{}, nil, nil) + if switcher != nil || sessionSwitcher != nil { + t.Error("switchers were built without a provider factory to build providers with") + } +} diff --git a/internal/tui/escalation_usage_test.go b/internal/tui/escalation_usage_test.go new file mode 100644 index 000000000..4903bb078 --- /dev/null +++ b/internal/tui/escalation_usage_test.go @@ -0,0 +1,303 @@ +package tui + +import ( + "context" + "encoding/json" + "reflect" + "testing" + + tea "charm.land/bubbletea/v2" + + "github.com/Gitlawb/zero/internal/config" + "github.com/Gitlawb/zero/internal/modelregistry" + "github.com/Gitlawb/zero/internal/sessions" + "github.com/Gitlawb/zero/internal/tools" + "github.com/Gitlawb/zero/internal/usage" + "github.com/Gitlawb/zero/internal/zeroruntime" +) + +// USAGE FOLLOWS THE SWITCH. exec reassigns its currentModel when the escalation +// switcher fires; the TUI captured usageModelID once per run and handed the +// switcher no callback, so every usage event after an escalation was billed to +// the model the run started on. Both attribution paths are pinned here: the +// live per-event message, and the per-event record the final response carries +// for the batch fallback, which must not swing the other way and bill the +// events before the switch to the escalated model. +func TestEscalatedRunAttributesUsageToTheModelInForce(t *testing.T) { + store := testSessionStore(t) + session, err := store.Create(sessions.CreateInput{SessionID: "escalation_usage"}) + if err != nil { + t.Fatal(err) + } + + // A catalog model with an upgrade target; the target is whatever the + // catalog says, read back from the provider factory rather than assumed. + const startModel = "claude-haiku-4.5" + starting := &scriptedProvider{scripts: [][]zeroruntime.StreamEvent{{ + {Type: zeroruntime.StreamEventUsage, Usage: zeroruntime.Usage{InputTokens: 10, OutputTokens: 1}}, + {Type: zeroruntime.StreamEventToolCallStart, ToolCallID: "escalate", ToolName: "escalate_model"}, + {Type: zeroruntime.StreamEventToolCallDelta, ToolCallID: "escalate", ArgumentsFragment: `{"reason":"harder than it looked"}`}, + {Type: zeroruntime.StreamEventToolCallEnd, ToolCallID: "escalate"}, + {Type: zeroruntime.StreamEventDone}, + }}} + escalated := &scriptedProvider{scripts: [][]zeroruntime.StreamEvent{{ + {Type: zeroruntime.StreamEventText, Content: "Done on the stronger model."}, + {Type: zeroruntime.StreamEventUsage, Usage: zeroruntime.Usage{InputTokens: 20, OutputTokens: 2}}, + {Type: zeroruntime.StreamEventDone}, + }}} + registry := tools.NewRegistry() + registry.Register(tools.NewEscalateModelTool()) + + switchedTo := "" + m := newModel(context.Background(), Options{ + ProviderName: "anthropic", + ModelName: startModel, + Provider: starting, + Registry: registry, + SessionStore: store, + AllowEscalation: true, + ProviderProfile: config.ProviderProfile{Model: startModel}, + NewProvider: func(profile config.ProviderProfile) (zeroruntime.Provider, error) { + switchedTo = profile.Model + return escalated, nil + }, + }) + m.activeSession = session + m.agentOptions.Model = startModel + var live []string + m.runtimeMessageSink = func(msg tea.Msg) { + if usage, ok := msg.(agentUsageMsg); ok { + live = append(live, usage.modelID) + } + } + + msg := execCmd(m.runAgentWithOptions(1, context.Background(), "do the hard thing", nil, tuiAgentRunOptions{})) + response, ok := msg.(agentResponseMsg) + if !ok { + t.Fatalf("run returned %T, want agentResponseMsg", msg) + } + if response.err != nil { + t.Fatalf("run failed: %v", response.err) + } + if switchedTo == "" || switchedTo == startModel { + t.Fatalf("SETUP INVALID: the escalation never switched providers (switched to %q), so nothing here exercises attribution", switchedTo) + } + if len(escalated.requests) == 0 { + t.Fatal("SETUP INVALID: the escalated provider was never asked for a completion") + } + if len(response.usageEvents) != 2 { + t.Fatalf("usage events = %d, want one before the switch and one after", len(response.usageEvents)) + } + + want := []string{startModel, switchedTo} + if !reflect.DeepEqual(live, want) { + t.Fatalf("live usage attribution = %v, want %v", live, want) + } + if !reflect.DeepEqual(response.usageModelIDs, want) { + t.Fatalf("per-event usage attribution on the response = %v, want %v", response.usageModelIDs, want) + } + for index := range response.usageEvents { + if got := response.usageModelIDAt(index); got != want[index] { + t.Fatalf("usageModelIDAt(%d) = %q, want %q", index, got, want[index]) + } + } +} + +// A response built without the per-event record, which is every constructor +// that predates escalation, still attributes through the run-level model. +func TestUsageModelIDAtFallsBackToTheRunModel(t *testing.T) { + msg := agentResponseMsg{usageModelID: "gpt-4.1", usageEvents: make([]zeroruntime.Usage, 2)} + for index := range msg.usageEvents { + if got := msg.usageModelIDAt(index); got != "gpt-4.1" { + t.Fatalf("usageModelIDAt(%d) = %q, want the run-level model", index, got) + } + } + msg.usageModelIDs = []string{"gpt-4.1-mini"} + if got := msg.usageModelIDAt(0); got != "gpt-4.1-mini" { + t.Fatalf("usageModelIDAt(0) = %q, want the per-event model", got) + } + if got := msg.usageModelIDAt(1); got != "gpt-4.1" { + t.Fatalf("usageModelIDAt(1) = %q, want the run-level fallback past the record", got) + } +} + +// AND THE PRICE THAT COMES BACK OUT OF THE SESSION LOG. The live record and the +// persisted event are two representations of one usage event, and only the +// first carried the model. `zero usage report` rebuilds cost from the persisted +// payload, falling back to the session's own model when the event does not name +// one, so an escalated run was priced end to end at the model it started on. +// This drives the real persistence path and reconstructs the report from it. +func TestEscalatedRunPricesPersistedUsageAtEachModel(t *testing.T) { + registry, err := modelregistry.DefaultRegistry() + if err != nil { + t.Fatal(err) + } + store := testSessionStore(t) + session, err := store.Create(sessions.CreateInput{SessionID: "escalation_cost"}) + if err != nil { + t.Fatal(err) + } + + const startModel = "claude-haiku-4.5" + before := zeroruntime.Usage{InputTokens: 10, OutputTokens: 1} + after := zeroruntime.Usage{InputTokens: 20, OutputTokens: 2} + starting := &scriptedProvider{scripts: [][]zeroruntime.StreamEvent{{ + {Type: zeroruntime.StreamEventUsage, Usage: before}, + {Type: zeroruntime.StreamEventToolCallStart, ToolCallID: "escalate", ToolName: "escalate_model"}, + {Type: zeroruntime.StreamEventToolCallDelta, ToolCallID: "escalate", ArgumentsFragment: `{"reason":"harder than it looked"}`}, + {Type: zeroruntime.StreamEventToolCallEnd, ToolCallID: "escalate"}, + {Type: zeroruntime.StreamEventDone}, + }}} + escalated := &scriptedProvider{scripts: [][]zeroruntime.StreamEvent{{ + {Type: zeroruntime.StreamEventText, Content: "Done on the stronger model."}, + {Type: zeroruntime.StreamEventUsage, Usage: after}, + {Type: zeroruntime.StreamEventDone}, + }}} + toolRegistry := tools.NewRegistry() + toolRegistry.Register(tools.NewEscalateModelTool()) + + switchedTo := "" + m := newModel(context.Background(), Options{ + ProviderName: "anthropic", + ModelName: startModel, + Provider: starting, + Registry: toolRegistry, + SessionStore: store, + AllowEscalation: true, + ProviderProfile: config.ProviderProfile{Model: startModel}, + NewProvider: func(profile config.ProviderProfile) (zeroruntime.Provider, error) { + switchedTo = profile.Model + return escalated, nil + }, + }) + m.activeSession = session + m.agentOptions.Model = startModel + + msg := execCmd(m.runAgentWithOptions(1, context.Background(), "do the hard thing", nil, tuiAgentRunOptions{})) + response, ok := msg.(agentResponseMsg) + if !ok || response.err != nil { + t.Fatalf("run returned %T err=%v", msg, response.err) + } + if switchedTo == "" || switchedTo == startModel { + t.Fatalf("SETUP INVALID: no escalation happened (switched to %q)", switchedTo) + } + + // Persist through the same path a real run uses, then read the events back. + m, rows := m.appendSessionEvents(response.sessionEvents) + for _, row := range rows { + if row.kind == rowError { + t.Fatalf("session record error: %s", row.text) + } + } + events, err := store.ReadEvents(session.SessionID) + if err != nil { + t.Fatal(err) + } + persisted := []string{} + for _, event := range events { + if event.Type != sessions.EventUsage { + continue + } + var payload struct { + Model string `json:"model"` + } + if err := json.Unmarshal(event.Payload, &payload); err != nil { + t.Fatal(err) + } + persisted = append(persisted, payload.Model) + } + if want := []string{startModel, switchedTo}; !reflect.DeepEqual(persisted, want) { + t.Fatalf("persisted usage models = %v, want %v", persisted, want) + } + + // The report prices from those payloads. Its session metadata names only the + // starting model, which is the fallback that used to price both events. + metadata := []sessions.Metadata{{SessionID: session.SessionID, ModelID: startModel}} + report, err := usage.BuildReport(events, metadata, ®istry, 0) + if err != nil { + t.Fatal(err) + } + + // Expected cost comes from the catalog, one completion at each model. + expected := 0.0 + for _, priced := range []struct { + modelID string + usage zeroruntime.Usage + }{{startModel, before}, {switchedTo, after}} { + model, err := registry.Require(priced.modelID) + if err != nil { + t.Fatal(err) + } + cost, err := modelregistry.CalculateCost(model, priced.usage) + if err != nil { + t.Fatal(err) + } + expected += cost.TotalCost + } + // And the wrong answer, for contrast: both events at the starting model. + startingOnly := 0.0 + startingEntry, err := registry.Require(startModel) + if err != nil { + t.Fatal(err) + } + for _, event := range []zeroruntime.Usage{before, after} { + cost, err := modelregistry.CalculateCost(startingEntry, event) + if err != nil { + t.Fatal(err) + } + startingOnly += cost.TotalCost + } + if expected == startingOnly { + t.Fatal("SETUP INVALID: the two models price this usage identically, so the assertion below proves nothing") + } + if difference := report.Total.TotalCost - expected; difference > 1e-12 || difference < -1e-12 { + t.Fatalf("report cost = %v, want %v (pricing both events at %s gives %v)", report.Total.TotalCost, expected, startModel, startingOnly) + } +} + +// A run without escalation persists no model on its usage events: the model +// cannot change mid-run, so the session-wide identity is the whole story and +// the payload stays compact, exactly as exec does it. +func TestUnescalatedRunLeavesTheUsagePayloadAlone(t *testing.T) { + store := testSessionStore(t) + session, err := store.Create(sessions.CreateInput{SessionID: "no_escalation_usage"}) + if err != nil { + t.Fatal(err) + } + provider := &scriptedProvider{scripts: [][]zeroruntime.StreamEvent{{ + {Type: zeroruntime.StreamEventText, Content: "done"}, + {Type: zeroruntime.StreamEventUsage, Usage: zeroruntime.Usage{InputTokens: 10, OutputTokens: 1}}, + {Type: zeroruntime.StreamEventDone}, + }}} + m := newModel(context.Background(), Options{ + ProviderName: "anthropic", + ModelName: "claude-haiku-4.5", + Provider: provider, + Registry: tools.NewRegistry(), + SessionStore: store, + }) + m.activeSession = session + + msg := execCmd(m.runAgentWithOptions(1, context.Background(), "hello", nil, tuiAgentRunOptions{})) + response, ok := msg.(agentResponseMsg) + if !ok || response.err != nil { + t.Fatalf("run returned %T err=%v", msg, response.err) + } + usageEvents := 0 + for _, event := range response.sessionEvents { + if event.Type != sessions.EventUsage { + continue + } + usageEvents++ + payload, ok := event.Payload.(map[string]any) + if !ok { + t.Fatalf("usage payload = %T, want map", event.Payload) + } + if _, present := payload["model"]; present { + t.Errorf("a non-escalation run recorded a model on its usage payload: %v", payload) + } + } + if usageEvents == 0 { + t.Fatal("SETUP INVALID: the run recorded no usage event") + } +} diff --git a/internal/tui/model.go b/internal/tui/model.go index 9473de06a..c23935980 100644 --- a/internal/tui/model.go +++ b/internal/tui/model.go @@ -70,21 +70,24 @@ const dragEdgeScrollInterval = 70 * time.Millisecond const dragEdgeScrollStep = 1 type model struct { - ctx context.Context - cwd string - appVersion string - userCommands []usercommands.Command // file-sourced /commands (.zero/commands) - loadSkills func() []skills.Skill // lazy installed-skills loader for /skills + / - userConfigPath string - doctorUserConfigPath string - projectConfigPath string - gitBranch string - providerName string - modelName string - modelCatalog modelregistry.Registry - providerProfile config.ProviderProfile - savedProviders []config.ProviderProfile - provider zeroruntime.Provider + ctx context.Context + cwd string + appVersion string + userCommands []usercommands.Command // file-sourced /commands (.zero/commands) + loadSkills func() []skills.Skill // lazy installed-skills loader for /skills + / + userConfigPath string + doctorUserConfigPath string + projectConfigPath string + gitBranch string + providerName string + modelName string + modelCatalog modelregistry.Registry + providerProfile config.ProviderProfile + savedProviders []config.ProviderProfile + provider zeroruntime.Provider + // allowEscalation mirrors Options.AllowEscalation: it gates the per-run model + // switchers, and the caller gates the escalate_model tool on the same flag. + allowEscalation bool newProvider func(config.ProviderProfile) (zeroruntime.Provider, error) newTurnSessionProvider func(config.ProviderProfile, zeroruntime.Provider) zeroruntime.TurnSessionProvider probeProviderHealth func(context.Context, providerhealth.Options) providerhealth.Result @@ -658,10 +661,16 @@ type agentUsageMsg struct { } type agentResponseMsg struct { - runID int - rows []transcriptRow - usageEvents []zeroruntime.Usage + runID int + rows []transcriptRow + usageEvents []zeroruntime.Usage + // usageModelID is the model in force when the run ended. usageModelIDs is + // the model in force when each usageEvents entry fired: a mid-run + // escalation changes it partway through the run, and billing the events + // before the switch to the escalated model would be as wrong as billing + // the ones after it to the starting model. Read through usageModelIDAt. usageModelID string + usageModelIDs []string sessionEvents []pendingSessionEvent specReview *pendingSpecReviewPrompt err error @@ -674,6 +683,16 @@ type agentResponseMsg struct { ttft time.Duration } +// usageModelIDAt is the model in force when usageEvents[index] fired. The +// per-event record wins; usageModelID is the fallback for a message built +// without one, which is what every constructor before escalation produced. +func (msg agentResponseMsg) usageModelIDAt(index int) string { + if index < len(msg.usageModelIDs) && msg.usageModelIDs[index] != "" { + return msg.usageModelIDs[index] + } + return msg.usageModelID +} + type peerMessageMsg struct { message peermsg.InboundMessage admit chan<- bool @@ -999,6 +1018,7 @@ func newModel(ctx context.Context, options Options) model { mcpCommand: options.MCPCommand, sandboxSetupCommand: options.SandboxSetupCommand, agentOptions: options.AgentOptions, + allowEscalation: options.AllowEscalation, sessionCompactor: options.SessionCompactor, runtimeMessageSink: options.RuntimeMessageSink, permissionMode: permissionMode, @@ -2546,7 +2566,7 @@ func (m model) updateModel(msg tea.Msg) (tea.Model, tea.Cmd) { continue } var usageRows []transcriptRow - m, usageRows = m.recordUsageEvent(msg.usageModelID, event) + m, usageRows = m.recordUsageEvent(msg.usageModelIDAt(index), event) for _, row := range usageRows { m.transcript = appendTranscriptRow(m.transcript, row) } @@ -2633,7 +2653,7 @@ func (m model) updateModel(msg tea.Msg) (tea.Model, tea.Cmd) { continue } var usageRows []transcriptRow - m, usageRows = m.recordUsageEvent(msg.usageModelID, event) + m, usageRows = m.recordUsageEvent(msg.usageModelIDAt(index), event) for _, row := range usageRows { m.transcript = appendTranscriptRow(m.transcript, row) } @@ -5427,6 +5447,9 @@ func (m model) runAgentWithOptions(runID int, runCtx context.Context, prompt str usageEvents := []zeroruntime.Usage{} sessionEvents := []pendingSessionEvent{} usageModelID := m.modelName + // usageModelIDs records, per usage event, the model in force when it + // fired; the escalation switcher reassigns usageModelID mid-run. + usageModelIDs := []string{} var specReview *pendingSpecReviewPrompt if m.awaitToolReadiness != nil { m.awaitToolReadiness(runCtx) @@ -5496,6 +5519,27 @@ func (m model) runAgentWithOptions(runID int, runCtx context.Context, prompt str options.ContextWindowFor = func(modelID string) int { return modelregistry.AgentContextWindow(m.modelContextWindow(modelID)) } + // And make that switch reachable, when the operator asked for it. The + // consequences of an escalation were already handled here (the window + // above, and the summarizer resolved against the active profile) while + // nothing on this surface could cause one: escalate_model was registered + // only by exec. + // + // BUILT FROM THE ACTIVE PROFILE, NOT THE STARTUP ONE. A TUI session can + // change models with /model, so escalating from the profile captured at + // launch would switch from whatever the session began with rather than + // from what is in force now, and would carry that stale profile's base URL + // and credential with it. m.providerProfile tracks the switches, which is + // why this is built per turn rather than once in the caller. + if m.allowEscalation { + options.ModelSwitcher, options.ModelSessionSwitcher = providers.EscalationSwitchers( + m.providerProfile, m.provider, m.newProvider, + // Usage attribution follows the switch, as exec reassigns its + // currentModel: every usage event after a real escalation is billed + // to the escalated model, not the one the run started on. + func(modelID string) { usageModelID = modelID }, + ) + } // Post-edit self-correction is on by default in the TUI but kept FAST: it // runs LSP diagnostics over the changed files only — cheap, change-scoped, @@ -5861,9 +5905,20 @@ func (m model) runAgentWithOptions(runID int, runCtx context.Context, prompt str onUsage := options.OnUsage options.OnUsage = func(event zeroruntime.Usage) { usageEvents = append(usageEvents, event) + usageModelIDs = append(usageModelIDs, usageModelID) + payload := usage.EventUsagePayload(event) + // AND ON THE PERSISTED EVENT TOO, not only the in-memory record: the + // report reconstructs cost from the payload, falling back to the + // session-wide model, so an escalated run would be priced entirely at + // the model it started on. Written only under escalation, which is the + // only way the model in force can change mid-run, matching what exec + // records under the same flag. + if m.allowEscalation { + payload["model"] = usageModelID + } sessionEvents = append(sessionEvents, pendingSessionEvent{ Type: sessions.EventUsage, - Payload: usage.EventUsagePayload(event), + Payload: payload, }) m.sendAgentUsage(runID, usageModelID, event) if onUsage != nil { @@ -5878,7 +5933,7 @@ func (m model) runAgentWithOptions(runID int, runCtx context.Context, prompt str Type: sessions.EventError, Payload: map[string]any{"message": err.Error()}, }) - return agentResponseMsg{runID: runID, rows: rows, usageEvents: usageEvents, usageModelID: usageModelID, sessionEvents: sessionEvents, err: err, goalAware: goalAwareRun, turnTools: toolCalls, turnElapsed: m.activeTurnElapsed(started)} + return agentResponseMsg{runID: runID, rows: rows, usageEvents: usageEvents, usageModelID: usageModelID, usageModelIDs: usageModelIDs, sessionEvents: sessionEvents, err: err, goalAware: goalAwareRun, turnTools: toolCalls, turnElapsed: m.activeTurnElapsed(started)} } if runOptions.specDraft { if result.StopReason != agent.StopReasonSpecReviewRequired || specReview == nil || specReview.SpecID == "" || specReview.SpecFilePath == "" { @@ -5888,10 +5943,10 @@ func (m model) runAgentWithOptions(runID int, runCtx context.Context, prompt str Type: sessions.EventError, Payload: map[string]any{"message": err.Error()}, }) - return agentResponseMsg{runID: runID, rows: rows, usageEvents: usageEvents, usageModelID: usageModelID, sessionEvents: sessionEvents, err: err, goalAware: goalAwareRun, turnTools: toolCalls, turnElapsed: m.activeTurnElapsed(started)} + return agentResponseMsg{runID: runID, rows: rows, usageEvents: usageEvents, usageModelID: usageModelID, usageModelIDs: usageModelIDs, sessionEvents: sessionEvents, err: err, goalAware: goalAwareRun, turnTools: toolCalls, turnElapsed: m.activeTurnElapsed(started)} } flushReasoning(m.now()) - return agentResponseMsg{runID: runID, rows: rows, usageEvents: usageEvents, usageModelID: usageModelID, sessionEvents: sessionEvents, specReview: specReview, goalAware: goalAwareRun, turnTools: toolCalls, turnElapsed: m.activeTurnElapsed(started)} + return agentResponseMsg{runID: runID, rows: rows, usageEvents: usageEvents, usageModelID: usageModelID, usageModelIDs: usageModelIDs, sessionEvents: sessionEvents, specReview: specReview, goalAware: goalAwareRun, turnTools: toolCalls, turnElapsed: m.activeTurnElapsed(started)} } flushReasoning(m.now()) elapsed := m.activeTurnElapsed(started) @@ -5912,7 +5967,7 @@ func (m model) runAgentWithOptions(runID int, runCtx context.Context, prompt str "content": result.FinalAnswer, }, }) - return agentResponseMsg{runID: runID, rows: rows, usageEvents: usageEvents, usageModelID: usageModelID, sessionEvents: sessionEvents, goalAware: goalAwareRun, turnTools: toolCalls, turnElapsed: elapsed, ttft: firstTokenElapsed} + return agentResponseMsg{runID: runID, rows: rows, usageEvents: usageEvents, usageModelID: usageModelID, usageModelIDs: usageModelIDs, sessionEvents: sessionEvents, goalAware: goalAwareRun, turnTools: toolCalls, turnElapsed: elapsed, ttft: firstTokenElapsed} } } diff --git a/internal/tui/options.go b/internal/tui/options.go index e73c7eaa5..638e82085 100644 --- a/internal/tui/options.go +++ b/internal/tui/options.go @@ -70,7 +70,13 @@ type Options struct { // invocation. Called lazily per use so newly installed skills are picked up // without a restart. Nil means the session has no skills wiring (skills stay // model-pulled via the skill tool only). - LoadSkills func() []skills.Skill + LoadSkills func() []skills.Skill + // AllowEscalation opts this session into mid-run model escalation, set from + // --allow-escalation. It wires the model switchers onto every turn's options; + // the escalate_model tool itself is registered by the caller on the same flag. + // Both halves are required: the tool without the switchers is inert, and the + // switchers without the tool are unreachable. + AllowEscalation bool PermissionMode agent.PermissionMode ReasoningEffort modelregistry.ReasoningEffort ResponseStyle string