From 78b8dea66f3dfdb33e52c2160244e4bc93cac5b8 Mon Sep 17 00:00:00 2001 From: Vasanthdev2004 Date: Tue, 8 Sep 2026 00:28:06 +0530 Subject: [PATCH 1/8] feat(tui): offer mid-run model escalation behind --allow-escalation Mid-run escalation existed but only exec offered it. The TUI already handled every consequence of a switch, re-deriving the compaction threshold "after a mid-run escalate_model switch" and resolving the summarizer against the active profile, while nothing on that surface could cause one: escalate_model was registered by exec alone. This closes that, opt-in, the same conservative way exec answers the question. Opt-in rather than default because escalation moves a run onto a different model, which changes what it costs and which provider sees the conversation. That is the operator's call, and the exec flag already answers it. The switchers are built per turn from m.providerProfile rather than once at launch. A TUI session can change models with /model, so a closure capturing the startup profile would escalate from whatever the session began with instead of what is in force now, and would carry that stale profile's base URL and credential with it. exec has no such problem because its profile cannot change mid-run, which is why the wiring lives in different places on the two surfaces. The nil-contract logic is now shared rather than copied. providers. EscalationSwitchers is the single implementation both surfaces call: the loop swaps only on a non-nil provider, so a (nil, nil) return has to leave everything untouched including the caller's usage attribution, and the session switcher is installed only when the run STARTED optimized so escalation cannot change the transport underneath a session. A second copy of that would only ever have been exercised on one surface. exec's behaviour is unchanged: its six existing escalation tests pass against the shared builder. The tool and the switchers ride on one flag and a test asserts they cannot be separated. Registering escalate_model without wiring a switcher ships a tool the loop silently ignores, so the model reports an escalation that never happened; that is the specific failure worth a guard rather than a comment. Note this does not answer all of #554. Assigning models per phase, plan with one and execute with another, is still not built. This makes the mechanism reachable from the interactive surface, which is the part that was one flag away. Refs #554 --- internal/cli/app.go | 62 +++++++++++- internal/cli/completions.go | 2 +- internal/cli/exec.go | 63 +++---------- internal/cli/setup.go | 2 +- internal/cli/tui_escalation_test.go | 125 +++++++++++++++++++++++++ internal/providers/escalation.go | 89 ++++++++++++++++++ internal/providers/escalation_test.go | 130 ++++++++++++++++++++++++++ internal/tui/model.go | 51 +++++++--- internal/tui/options.go | 8 +- 9 files changed, 460 insertions(+), 72 deletions(-) create mode 100644 internal/cli/tui_escalation_test.go create mode 100644 internal/providers/escalation.go create mode 100644 internal/providers/escalation_test.go diff --git a/internal/cli/app.go b/internal/cli/app.go index 44fa370ff..39189ec26 100644 --- a/internal/cli/app.go +++ b/internal/cli/app.go @@ -327,9 +327,15 @@ func runWithDeps(args []string, stdout io.Writer, stderr io.Writer, deps appDeps return writeAppError(stderr, err.Error(), 1) } addDirs = append(addDirs, moreDirs...) + // --allow-escalation opts the interactive session into mid-run model + // escalation, mirroring the exec flag of the same name. + allowEscalation, args, err := splitLeadingAllowEscalationFlag(args) + if err != nil { + return writeAppError(stderr, err.Error(), 1) + } 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 @@ -374,6 +380,13 @@ func runWithDeps(args []string, stdout io.Writer, stderr io.Writer, deps appDeps return writeAppError(stderr, err.Error(), 1) } moreDirs = append(moreDirs, evenMoreDirs...) + // --allow-escalation may sit on either side of --skip-permissions-unsafe, + // like --theme and --add-dir, so re-split it here and OR the two results + // rather than letting the side it was written on decide. + skipAllowEscalation, rest, err := splitLeadingAllowEscalationFlag(rest) + if err != nil { + return writeAppError(stderr, err.Error(), 1) + } // 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 +403,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 @@ -694,11 +707,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 +803,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 +1054,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 +1400,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 @@ -1434,6 +1457,35 @@ func splitLeadingAddDirFlags(args []string) ([]string, []string, error) { return addDirs, 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/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..4f0a81c53 --- /dev/null +++ b/internal/cli/tui_escalation_test.go @@ -0,0 +1,125 @@ +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() + 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/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/model.go b/internal/tui/model.go index 9473de06a..a35dde7f1 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 @@ -999,6 +1002,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, @@ -5496,6 +5500,23 @@ 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, nil, + ) + } // 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, 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 From a339c7064ca86b3d65995def6377e03d68e607bc Mon Sep 17 00:00:00 2001 From: Vasanthdev2004 Date: Tue, 8 Sep 2026 10:20:13 +0530 Subject: [PATCH 2/8] fix(tui): attribute usage after an escalation to the model in force exec reassigns currentModel when the escalation switcher fires; the TUI captured usageModelID once per run and handed the switcher no callback, so every usage event after a mid-run escalation was billed to the model the run started on. The switcher now updates usageModelID, and the final response carries the model per usage event so the batch fallback does not swing the other way and bill the events before the switch to the escalated model. Covered by a run through a scripted escalation that checks both the live and the batch attribution. --- internal/tui/escalation_usage_test.go | 119 ++++++++++++++++++++++++++ internal/tui/model.go | 44 +++++++--- 2 files changed, 153 insertions(+), 10 deletions(-) create mode 100644 internal/tui/escalation_usage_test.go diff --git a/internal/tui/escalation_usage_test.go b/internal/tui/escalation_usage_test.go new file mode 100644 index 000000000..cded19bd7 --- /dev/null +++ b/internal/tui/escalation_usage_test.go @@ -0,0 +1,119 @@ +package tui + +import ( + "context" + "reflect" + "testing" + + tea "charm.land/bubbletea/v2" + + "github.com/Gitlawb/zero/internal/config" + "github.com/Gitlawb/zero/internal/sessions" + "github.com/Gitlawb/zero/internal/tools" + "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) + } +} diff --git a/internal/tui/model.go b/internal/tui/model.go index a35dde7f1..fd6b0eae6 100644 --- a/internal/tui/model.go +++ b/internal/tui/model.go @@ -661,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 @@ -677,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 @@ -2550,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) } @@ -2637,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) } @@ -5431,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) @@ -5514,7 +5533,11 @@ func (m model) runAgentWithOptions(runID int, runCtx context.Context, prompt str // 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, nil, + 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 }, ) } @@ -5882,6 +5905,7 @@ 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) sessionEvents = append(sessionEvents, pendingSessionEvent{ Type: sessions.EventUsage, Payload: usage.EventUsagePayload(event), @@ -5899,7 +5923,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 == "" { @@ -5909,10 +5933,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) @@ -5933,7 +5957,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} } } From 49a9a632e564d7d4007e037f83374d1933338ad3 Mon Sep 17 00:00:00 2001 From: Vasanthdev2004 Date: Tue, 8 Sep 2026 10:40:03 +0530 Subject: [PATCH 3/8] fix(cli): accept the root flags 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 load-bearing: `zero --allow-escalation --theme auto` stranded the theme as an unknown command and exited with an argument error instead of launching, and the same happened after --skip-permissions-unsafe. The three splitters now run until they make no progress, at the root and after the unsafe flag, so every ordering reaches the TUI with every flag applied. A --theme written before --skip-permissions-unsafe was dropped on that path; it is kept now, with a later one winning as the last occurrence does. --- internal/cli/app.go | 108 ++++++++++++++++----------- internal/cli/root_flag_order_test.go | 90 ++++++++++++++++++++++ 2 files changed, 156 insertions(+), 42 deletions(-) create mode 100644 internal/cli/root_flag_order_test.go diff --git a/internal/cli/app.go b/internal/cli/app.go index 39189ec26..86059a81d 100644 --- a/internal/cli/app.go +++ b/internal/cli/app.go @@ -311,28 +311,13 @@ 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) - 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...) - // --allow-escalation opts the interactive session into mid-run model - // escalation, mirroring the exec flag of the same name. - allowEscalation, args, err := splitLeadingAllowEscalationFlag(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) } + addDirs, theme, allowEscalation := root.addDirs, root.theme, root.allowEscalation if len(args) == 0 { return runInteractiveTUI(stderr, deps, agent.PermissionModeAsk, addDirs, theme, allowEscalation) @@ -360,33 +345,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 = append(moreDirs, evenMoreDirs...) - // --allow-escalation may sit on either side of --skip-permissions-unsafe, - // like --theme and --add-dir, so re-split it here and OR the two results - // rather than letting the side it was written on decide. - skipAllowEscalation, rest, err := splitLeadingAllowEscalationFlag(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 } + 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 { @@ -1457,6 +1433,54 @@ 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. diff --git a/internal/cli/root_flag_order_test.go b/internal/cli/root_flag_order_test.go new file mode 100644 index 000000000..7af3bed5f --- /dev/null +++ b/internal/cli/root_flag_order_test.go @@ -0,0 +1,90 @@ +package cli + +import ( + "path/filepath" + "strings" + "testing" + + "github.com/Gitlawb/zero/internal/agent" +) + +// 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, on the ask path and the unsafe path, 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) + } + for _, args := range [][]string{ + {"--allow-escalation", "--theme", "auto", "--add-dir", extra}, + {"--allow-escalation", "--add-dir", extra, "--theme", "auto"}, + {"--theme", "auto", "--allow-escalation", "--add-dir", extra}, + {"--add-dir", extra, "--allow-escalation", "--theme", "auto"}, + {"--add-dir", extra, "--theme", "auto", "--allow-escalation"}, + {"--allow-escalation", "--add-dir", extra, "--theme", "auto", "--allow-escalation"}, + {"--skip-permissions-unsafe", "--allow-escalation", "--theme", "auto", "--add-dir", extra}, + {"--allow-escalation", "--skip-permissions-unsafe", "--add-dir", extra, "--theme", "auto"}, + {"--theme", "auto", "--skip-permissions-unsafe", "--allow-escalation", "--add-dir", extra}, + {"--add-dir", extra, "--theme", "auto", "--allow-escalation", "--skip-permissions-unsafe"}, + } { + 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) + } + }) + } +} From bdba4049840f8f7221336ac5539ec9352c86dd0e Mon Sep 17 00:00:00 2001 From: Vasanthdev2004 Date: Tue, 8 Sep 2026 10:55:39 +0530 Subject: [PATCH 4/8] fix(cli): forward a root --allow-escalation into exec, reject it elsewhere The root flag was stripped before dispatch, so `zero --allow-escalation exec ...` and `zero --allow-escalation -p ...` ran without the opt-in the help text promises, and `zero --allow-escalation version` accepted a flag it could only discard. It now follows --add-dir: re-synthesised into the exec argument list for the exec and -p shapes, so exec parses it as its own flag, and rejected loudly for every other command. The ordering test now generates every permutation of the three root flags, with --skip-permissions-unsafe absent or at any position. --- internal/cli/app.go | 49 ++++++--- internal/cli/root_flag_order_test.go | 149 ++++++++++++++++++++++++--- 2 files changed, 167 insertions(+), 31 deletions(-) diff --git a/internal/cli/app.go b/internal/cli/app.go index 86059a81d..160eabc5b 100644 --- a/internal/cli/app.go +++ b/internal/cli/app.go @@ -323,19 +323,23 @@ func runWithDeps(args []string, stdout io.Writer, stderr io.Writer, deps appDeps 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] { @@ -412,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": @@ -1393,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. diff --git a/internal/cli/root_flag_order_test.go b/internal/cli/root_flag_order_test.go index 7af3bed5f..84357f736 100644 --- a/internal/cli/root_flag_order_test.go +++ b/internal/cli/root_flag_order_test.go @@ -1,38 +1,55 @@ 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, on the ask path and the unsafe path, has to reach the TUI with all -// three applied. +// 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) } - for _, args := range [][]string{ - {"--allow-escalation", "--theme", "auto", "--add-dir", extra}, - {"--allow-escalation", "--add-dir", extra, "--theme", "auto"}, - {"--theme", "auto", "--allow-escalation", "--add-dir", extra}, - {"--add-dir", extra, "--allow-escalation", "--theme", "auto"}, - {"--add-dir", extra, "--theme", "auto", "--allow-escalation"}, - {"--allow-escalation", "--add-dir", extra, "--theme", "auto", "--allow-escalation"}, - {"--skip-permissions-unsafe", "--allow-escalation", "--theme", "auto", "--add-dir", extra}, - {"--allow-escalation", "--skip-permissions-unsafe", "--add-dir", extra, "--theme", "auto"}, - {"--theme", "auto", "--skip-permissions-unsafe", "--allow-escalation", "--add-dir", extra}, - {"--add-dir", extra, "--theme", "auto", "--allow-escalation", "--skip-permissions-unsafe"}, - } { + 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 @@ -88,3 +105,107 @@ func TestRootThemeLastOccurrenceWinsAcrossOtherFlags(t *testing.T) { }) } } + +// 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() + t.Setenv("XDG_DATA_HOME", t.TempDir()) + 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 +} From 36f2eafefbad031b14a2bbf7e498c5df38dab008 Mon Sep 17 00:00:00 2001 From: Vasanthdev2004 Date: Wed, 9 Sep 2026 11:06:36 +0530 Subject: [PATCH 5/8] fix(tui): record the model on an escalated run persisted usage events The live usage record carried the model in force; the persisted session payload did not. The usage report rebuilds cost from that payload and falls back to the session-wide model when an event does not name one, so an escalated interactive run was priced end to end at the model it started on. Written only under escalation, matching what exec records under the same flag, so an ordinary run persists the same compact payload as before. --- internal/tui/escalation_usage_test.go | 184 ++++++++++++++++++++++++++ internal/tui/model.go | 12 +- 2 files changed, 195 insertions(+), 1 deletion(-) diff --git a/internal/tui/escalation_usage_test.go b/internal/tui/escalation_usage_test.go index cded19bd7..4903bb078 100644 --- a/internal/tui/escalation_usage_test.go +++ b/internal/tui/escalation_usage_test.go @@ -2,14 +2,17 @@ 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" ) @@ -117,3 +120,184 @@ func TestUsageModelIDAtFallsBackToTheRunModel(t *testing.T) { 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 fd6b0eae6..c23935980 100644 --- a/internal/tui/model.go +++ b/internal/tui/model.go @@ -5906,9 +5906,19 @@ func (m model) runAgentWithOptions(runID int, runCtx context.Context, prompt str 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 { From a548791b09c76af31ef82551114aca1263530c0c Mon Sep 17 00:00:00 2001 From: Vasanthdev2004 Date: Wed, 9 Sep 2026 11:09:16 +0530 Subject: [PATCH 6/8] test(cli): isolate the new launch helpers from real user state A temporary working directory isolates project files and nothing else. runWithDeps fills the dependencies a test leaves out with production ones, and the interactive launch path reads user config, opens stores, refreshes the models.dev cache and migrates any inline plaintext API key into the credential store before it reaches an injected runTUI callback. Running these tests on a developer machine therefore rewrote that developer's config.json and wrote a credential file beside it, and the results depended on whatever providers, MCP servers and plugins the machine had. Both new helpers now point every per-user base directory at a throwaway root, pin the models.dev cache path, disable its background fetch, and force the file credential backend so nothing reaches the host keyring. A new test seeds a config with an inline key outside those fixture roots and fails if either helper rewrites it or leaves a file beside it; without the isolation it reports the rewritten config and the credentials.enc files. --- internal/cli/root_flag_order_test.go | 2 +- internal/cli/tui_escalation_test.go | 1 + internal/cli/user_state_isolation_test.go | 96 +++++++++++++++++++++++ 3 files changed, 98 insertions(+), 1 deletion(-) create mode 100644 internal/cli/user_state_isolation_test.go diff --git a/internal/cli/root_flag_order_test.go b/internal/cli/root_flag_order_test.go index 84357f736..864c6d55a 100644 --- a/internal/cli/root_flag_order_test.go +++ b/internal/cli/root_flag_order_test.go @@ -160,7 +160,7 @@ func TestRootAllowEscalationIsRejectedWhereNothingConsumesIt(t *testing.T) { // whether escalate_model was among them. func execAdvertisesEscalateModel(t *testing.T, args []string) bool { t.Helper() - t.Setenv("XDG_DATA_HOME", t.TempDir()) + isolateCLIUserState(t) provider := &toolListingProvider{} var stdout, stderr bytes.Buffer exitCode := runWithDeps(args, &stdout, &stderr, appDeps{ diff --git a/internal/cli/tui_escalation_test.go b/internal/cli/tui_escalation_test.go index 4f0a81c53..d8e01eab8 100644 --- a/internal/cli/tui_escalation_test.go +++ b/internal/cli/tui_escalation_test.go @@ -13,6 +13,7 @@ import ( // 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() diff --git a/internal/cli/user_state_isolation_test.go b/internal/cli/user_state_isolation_test.go new file mode 100644 index 000000000..aaf7e0c20 --- /dev/null +++ b/internal/cli/user_state_isolation_test.go @@ -0,0 +1,96 @@ +package cli + +import ( + "os" + "path/filepath" + "testing" + + "github.com/Gitlawb/zero/internal/config" +) + +// 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 interactive +// launch path reads user config, opens stores, refreshes the models.dev cache +// and migrates any inline plaintext API key into the credential store before it +// ever reaches an injected runTUI callback. Without this, running these tests on +// a developer machine rewrites that developer's config.json and moves their key +// into their keychain, and the results depend on whatever providers, MCP servers +// and plugins that machine has configured. +func isolateCLIUserState(t *testing.T) { + t.Helper() + root := t.TempDir() + // One root behind every per-user base directory, so the platform-specific + // resolution lands inside it whichever branch it takes. + for _, name := range []string{ + "HOME", "USERPROFILE", "APPDATA", "LOCALAPPDATA", + "XDG_CONFIG_HOME", "XDG_DATA_HOME", "XDG_CACHE_HOME", "XDG_STATE_HOME", + } { + 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. A seeded config outside the helper's own +// fixture root stands in for a developer's real one: the launch helpers must +// leave it exactly as they found it, inline API key included, and must not +// write a credential file beside it. +func TestCLILaunchHelpersLeaveUserStateAlone(t *testing.T) { + seedRoot := t.TempDir() + for _, name := range []string{ + "HOME", "USERPROFILE", "APPDATA", "LOCALAPPDATA", + "XDG_CONFIG_HOME", "XDG_DATA_HOME", "XDG_CACHE_HOME", "XDG_STATE_HOME", + } { + t.Setenv(name, seedRoot) + } + configPath, err := config.DefaultUserConfigPath() + if err != nil { + t.Fatal(err) + } + if !filepath.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) + } + + // Both helpers added by this PR, each of which reaches the real launch path. + _ = captureTUIOptions(t, "--allow-escalation") + _ = execAdvertisesEscalateModel(t, []string{"--allow-escalation", "exec", "say hi"}) + + 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) + } +} From aff1967d2accfa6ee09777d33b47b96c85336fe8 Mon Sep 17 00:00:00 2001 From: Vasanthdev2004 Date: Wed, 9 Sep 2026 11:10:17 +0530 Subject: [PATCH 7/8] test(cli): check each launch helper isolation on its own t.Setenv restores at the end of the test, not when the helper returns, so one test exercising both helpers let the first isolation cover for the second. Each helper now seeds its own emulated user config in a subtest, and removing the isolation from either one fails that subtest. --- internal/cli/user_state_isolation_test.go | 101 ++++++++++++++-------- 1 file changed, 63 insertions(+), 38 deletions(-) diff --git a/internal/cli/user_state_isolation_test.go b/internal/cli/user_state_isolation_test.go index aaf7e0c20..8bc2b62ad 100644 --- a/internal/cli/user_state_isolation_test.go +++ b/internal/cli/user_state_isolation_test.go @@ -3,11 +3,20 @@ 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. // @@ -22,12 +31,7 @@ import ( func isolateCLIUserState(t *testing.T) { t.Helper() root := t.TempDir() - // One root behind every per-user base directory, so the platform-specific - // resolution lands inside it whichever branch it takes. - for _, name := range []string{ - "HOME", "USERPROFILE", "APPDATA", "LOCALAPPDATA", - "XDG_CONFIG_HOME", "XDG_DATA_HOME", "XDG_CACHE_HOME", "XDG_STATE_HOME", - } { + for _, name := range userStateEnvNames { t.Setenv(name, root) } // The models.dev overlay cache, and the background refresh that writes it: @@ -40,23 +44,66 @@ func isolateCLIUserState(t *testing.T) { t.Setenv("ZERO_CRED_STORAGE", "encrypted-file") } -// AND THE ISOLATION IS ITSELF PINNED. A seeded config outside the helper's own -// fixture root stands in for a developer's real one: the launch helpers must -// leave it exactly as they found it, inline API key included, and must not -// write a credential file beside it. +// 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, +// and must not write a credential file beside it. +// +// 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) { - seedRoot := t.TempDir() - for _, name := range []string{ - "HOME", "USERPROFILE", "APPDATA", "LOCALAPPDATA", - "XDG_CONFIG_HOME", "XDG_DATA_HOME", "XDG_CACHE_HOME", "XDG_STATE_HOME", + 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) { + 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) + } + }) + } +} + +// 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 the config path, its bytes, +// and the directory listing to compare against. +func seedEmulatedUserConfig(t *testing.T) (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 !filepath.HasPrefix(configPath, seedRoot) { + 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 { @@ -70,27 +117,5 @@ func TestCLILaunchHelpersLeaveUserStateAlone(t *testing.T) { if err != nil { t.Fatal(err) } - - // Both helpers added by this PR, each of which reaches the real launch path. - _ = captureTUIOptions(t, "--allow-escalation") - _ = execAdvertisesEscalateModel(t, []string{"--allow-escalation", "exec", "say hi"}) - - 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) - } + return configPath, seeded, before } From 7e7407beb5f9b6667c1ef919cf57badedcb84455 Mon Sep 17 00:00:00 2001 From: Vasanthdev2004 Date: Wed, 9 Sep 2026 11:11:57 +0530 Subject: [PATCH 8/8] test(cli): pin the isolation itself, not only what this platform writes The exec helper writes its models.dev cache from a background goroutine, so a footprint check could not prove its isolation deterministically. Each subtest now also asserts the helper moved the per-user paths off the seeded root and disabled the fetch and the keyring backend, which fails for either helper when the isolation call is removed. --- internal/cli/user_state_isolation_test.go | 44 ++++++++++++++++------- 1 file changed, 31 insertions(+), 13 deletions(-) diff --git a/internal/cli/user_state_isolation_test.go b/internal/cli/user_state_isolation_test.go index 8bc2b62ad..007c96b8d 100644 --- a/internal/cli/user_state_isolation_test.go +++ b/internal/cli/user_state_isolation_test.go @@ -21,13 +21,13 @@ var userStateEnvNames = []string{ // 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 interactive -// launch path reads user config, opens stores, refreshes the models.dev cache -// and migrates any inline plaintext API key into the credential store before it -// ever reaches an injected runTUI callback. Without this, running these tests on -// a developer machine rewrites that developer's config.json and moves their key -// into their keychain, and the results depend on whatever providers, MCP servers -// and plugins that machine has configured. +// 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() @@ -47,7 +47,8 @@ func isolateCLIUserState(t *testing.T) { // 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, -// and must not write a credential file beside it. +// 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 @@ -63,7 +64,7 @@ func TestCLILaunchHelpersLeaveUserStateAlone(t *testing.T) { }}, } { t.Run(testCase.name, func(t *testing.T) { - configPath, seeded, before := seedEmulatedUserConfig(t) + seedRoot, configPath, seeded, before := seedEmulatedUserConfig(t) testCase.exercise(t) @@ -85,15 +86,32 @@ func TestCLILaunchHelpersLeaveUserStateAlone(t *testing.T) { } 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 the config path, its bytes, -// and the directory listing to compare against. -func seedEmulatedUserConfig(t *testing.T) (string, []byte, []os.DirEntry) { +// 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 { @@ -117,5 +135,5 @@ func seedEmulatedUserConfig(t *testing.T) (string, []byte, []os.DirEntry) { if err != nil { t.Fatal(err) } - return configPath, seeded, before + return seedRoot, configPath, seeded, before }