mirror of
https://github.com/Twigpine/zero.git
synced 2026-11-08 17:56:14 -10:00
Compare commits
3
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
221fb93041 | ||
|
|
ac3958f860 | ||
|
|
63767f3e42 |
No files matched your search
@@ -273,6 +273,19 @@ zero sandbox grants list
|
||||
Zero includes local file/search/edit/shell tools, `web_fetch` for public URLs,
|
||||
and MCP support for additional tools.
|
||||
|
||||
`web_fetch` refuses loopback, private and other special-use addresses: it checks
|
||||
the URL before asking permission, resolves the host, and dials the address it
|
||||
validated so a name cannot resolve to something else in between.
|
||||
|
||||
If `HTTP_PROXY`/`HTTPS_PROXY` is set, `web_fetch` and the provider connectivity
|
||||
probe use it, and that last step changes: the proxy is dialed and the target
|
||||
hostname is sent to it, so the proxy decides which address the request actually
|
||||
reaches. The URL is still checked and resolved locally first, but a proxy that
|
||||
answers differently can reach a private service. A forward proxy already sees
|
||||
and can rewrite every request through it, so this is the trust you accept by
|
||||
configuring one. Leave the variables unset for the checks to be enforced end to
|
||||
end.
|
||||
|
||||
For local dev servers, use shell commands such as `curl` through `exec_command`
|
||||
so the normal sandbox and permission policy applies. Long-running commands stay
|
||||
attached to a background terminal session and can be listed or stopped from the
|
||||
|
||||
@@ -627,7 +627,12 @@ Core tool groups include:
|
||||
skills, `ask_user`, permission requests.
|
||||
- **Write tools**: `write_file`, `edit_file`, `apply_patch`, `update_plan`.
|
||||
- **Shell tools**: `exec_command`, `write_stdin`, legacy `bash`.
|
||||
- **Network tools**: `web_fetch` and optional web search backend.
|
||||
- **Network tools**: `web_fetch` and optional web search backend. `web_fetch`
|
||||
validates the URL before the permission prompt, refuses loopback/private/
|
||||
special-use addresses, and dials the address it resolved. With
|
||||
`HTTP_PROXY`/`HTTPS_PROXY` set, the proxy is dialed instead and resolves the
|
||||
target itself, so destination enforcement ends at the proxy; see the note in
|
||||
the README.
|
||||
- **Local-control tools**: config-gated wrappers for browser, terminal, desktop,
|
||||
and artifact helpers backed by packages such as `internal/localcontrol` and
|
||||
`internal/browser`.
|
||||
|
||||
+51
-142
@@ -311,35 +311,40 @@ 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()
|
||||
|
||||
// --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)
|
||||
addDirs, args, err := splitLeadingAddDirFlags(args)
|
||||
if err != nil {
|
||||
return writeAppError(stderr, err.Error(), 1)
|
||||
}
|
||||
addDirs, theme, allowEscalation := root.addDirs, root.theme, root.allowEscalation
|
||||
// --theme <name> 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...)
|
||||
|
||||
if len(args) == 0 {
|
||||
return runInteractiveTUI(stderr, deps, agent.PermissionModeAsk, addDirs, theme, allowEscalation)
|
||||
return runInteractiveTUI(stderr, deps, agent.PermissionModeAsk, addDirs, theme)
|
||||
}
|
||||
|
||||
// --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)
|
||||
// --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)
|
||||
}
|
||||
}
|
||||
|
||||
switch args[0] {
|
||||
@@ -349,24 +354,26 @@ 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.
|
||||
//
|
||||
// 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:])
|
||||
// --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)
|
||||
}
|
||||
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
|
||||
// --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)
|
||||
if err != nil {
|
||||
return writeAppError(stderr, err.Error(), 1)
|
||||
}
|
||||
skipAllowEscalation := more.allowEscalation
|
||||
evenMoreDirs, rest, err := splitLeadingAddDirFlags(rest)
|
||||
if err != nil {
|
||||
return writeAppError(stderr, err.Error(), 1)
|
||||
}
|
||||
moreDirs = append(moreDirs, evenMoreDirs...)
|
||||
// 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 {
|
||||
@@ -383,7 +390,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, allowEscalation || skipAllowEscalation)
|
||||
return runInteractiveTUI(stderr, deps, agent.PermissionModeUnsafe, append(append([]string{}, addDirs...), moreDirs...), skipTheme)
|
||||
case "-h", "--help", "help":
|
||||
if err := writeHelp(stdout); err != nil {
|
||||
return 1
|
||||
@@ -416,16 +423,16 @@ func runWithDeps(args []string, stdout io.Writer, stderr io.Writer, deps appDeps
|
||||
if len(args) < 2 {
|
||||
return writePromptRequired(stderr)
|
||||
}
|
||||
// Forward the root flags exec consumes so its own parser collects them.
|
||||
// Forward leading --add-dir occurrences so exec's own parser collects them.
|
||||
// Use the inline --prompt=<value> 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(rootFlagArgs(addDirs, allowEscalation), "--prompt="+args[1])
|
||||
execArgs := append(addDirFlagArgs(addDirs), "--prompt="+args[1])
|
||||
execArgs = append(execArgs, args[2:]...)
|
||||
return runExec(execArgs, stdout, stderr, deps)
|
||||
case "exec":
|
||||
// Forward the root flags exec consumes so its own parser collects them.
|
||||
return runExec(append(rootFlagArgs(addDirs, allowEscalation), args[1:]...), stdout, stderr, deps)
|
||||
// Forward leading --add-dir occurrences so exec's own parser collects them.
|
||||
return runExec(append(addDirFlagArgs(addDirs), args[1:]...), stdout, stderr, deps)
|
||||
case "completions":
|
||||
return runCompletions(args[1:], stdout, stderr)
|
||||
case "daemon":
|
||||
@@ -687,11 +694,11 @@ func fillAppDeps(deps appDeps) appDeps {
|
||||
return deps
|
||||
}
|
||||
|
||||
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 runInteractiveTUI(stderr io.Writer, deps appDeps, permissionMode agent.PermissionMode, addDirs []string, theme string) int {
|
||||
return runInteractiveTUIWithSetup(stderr, deps, permissionMode, addDirs, theme, false)
|
||||
}
|
||||
|
||||
func runInteractiveTUIWithSetup(stderr io.Writer, deps appDeps, permissionMode agent.PermissionMode, addDirs []string, theme string, forceSetup bool, allowEscalation bool) int {
|
||||
func runInteractiveTUIWithSetup(stderr io.Writer, deps appDeps, permissionMode agent.PermissionMode, addDirs []string, theme string, forceSetup 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.
|
||||
@@ -783,14 +790,6 @@ 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 {
|
||||
@@ -1034,7 +1033,6 @@ func runInteractiveTUIWithSetup(stderr io.Writer, deps appDeps, permissionMode a
|
||||
Specialists: specialistRuntime.specialists,
|
||||
Skills: pluginActivation.skillInfos(deps.skillsDir()),
|
||||
},
|
||||
AllowEscalation: allowEscalation,
|
||||
// LoadSkills backs /skills and direct /<skill-name> 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
|
||||
@@ -1380,7 +1378,6 @@ Flags:
|
||||
-v, --version Print version
|
||||
-p, --prompt Run a one-shot prompt
|
||||
--add-dir <path> 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
|
||||
@@ -1397,17 +1394,6 @@ 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 <path> [--add-dir <path>] [subcommand …]).
|
||||
// Subcommands like exec parse their own --add-dir occurrences.
|
||||
@@ -1448,83 +1434,6 @@ func splitLeadingAddDirFlags(args []string) ([]string, []string, error) {
|
||||
return addDirs, args, nil
|
||||
}
|
||||
|
||||
// rootFlags is what the leading root flags amount to once every one of them
|
||||
// has been stripped from the front of the argument list.
|
||||
type rootFlags struct {
|
||||
addDirs []string
|
||||
theme string
|
||||
allowEscalation bool
|
||||
}
|
||||
|
||||
// splitLeadingRootFlags strips --add-dir, --theme and --allow-escalation from
|
||||
// the front of args in whatever order they were written, stopping at the
|
||||
// first token none of them claims.
|
||||
//
|
||||
// EACH SPLITTER STOPS AT THE FIRST TOKEN IT DOES NOT OWN, so running them once
|
||||
// in a fixed sequence made the order the operator wrote them in load-bearing:
|
||||
// a flag handled late in the sequence stranded every flag written after it as
|
||||
// an unknown command, and `zero --allow-escalation --theme auto` exited with
|
||||
// an argument error instead of launching. Running the sequence until it makes
|
||||
// no progress accepts every ordering, including a flag repeated on both sides
|
||||
// of another.
|
||||
func splitLeadingRootFlags(args []string) (rootFlags, []string, error) {
|
||||
var flags rootFlags
|
||||
for {
|
||||
before := len(args)
|
||||
addDirs, rest, err := splitLeadingAddDirFlags(args)
|
||||
if err != nil {
|
||||
return rootFlags{}, nil, err
|
||||
}
|
||||
flags.addDirs = append(flags.addDirs, addDirs...)
|
||||
theme, rest, err := splitLeadingThemeFlag(rest)
|
||||
if err != nil {
|
||||
return rootFlags{}, nil, err
|
||||
}
|
||||
if theme != "" {
|
||||
// The last occurrence wins across passes, as it does within one.
|
||||
flags.theme = theme
|
||||
}
|
||||
allowEscalation, rest, err := splitLeadingAllowEscalationFlag(rest)
|
||||
if err != nil {
|
||||
return rootFlags{}, nil, err
|
||||
}
|
||||
flags.allowEscalation = flags.allowEscalation || allowEscalation
|
||||
args = rest
|
||||
if len(args) == before {
|
||||
return flags, args, nil
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// splitLeadingAllowEscalationFlag strips a leading --allow-escalation from the
|
||||
// root argument list, opting the interactive session into mid-run model
|
||||
// escalation.
|
||||
//
|
||||
// OPT-IN, THE SAME WAY exec IS. Escalation moves a run onto a different model,
|
||||
// which changes what the run costs and which provider sees the conversation, so
|
||||
// it is a decision the operator makes rather than a default. The exec flag
|
||||
// already answers this conservatively and the interactive surface should not
|
||||
// answer it differently.
|
||||
//
|
||||
// Bare flag only: repeating it is harmless, and an =value form is rejected so a
|
||||
// mistyped --allow-escalation=false is a loud error instead of silently enabling
|
||||
// the thing it was trying to turn off.
|
||||
func splitLeadingAllowEscalationFlag(args []string) (bool, []string, error) {
|
||||
allow := false
|
||||
for len(args) > 0 {
|
||||
switch {
|
||||
case args[0] == "--allow-escalation":
|
||||
allow = true
|
||||
args = args[1:]
|
||||
case strings.HasPrefix(args[0], "--allow-escalation="):
|
||||
return false, nil, errors.New("--allow-escalation takes no value; pass it bare to enable mid-run model escalation, or omit it")
|
||||
default:
|
||||
return allow, args, nil
|
||||
}
|
||||
}
|
||||
return allow, args, nil
|
||||
}
|
||||
|
||||
// splitLeadingThemeFlag strips a leading --theme <auto|theme-name> (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
|
||||
|
||||
@@ -20,7 +20,7 @@ type completionContext struct {
|
||||
}
|
||||
|
||||
var completionRoot = completionNode{
|
||||
flags: []string{"-h", "--help", "-v", "--version", "-p", "--prompt", "--add-dir", "--allow-escalation", "--theme", "--skip-permissions-unsafe"},
|
||||
flags: []string{"-h", "--help", "-v", "--version", "-p", "--prompt", "--add-dir", "--theme", "--skip-permissions-unsafe"},
|
||||
children: []completionNode{
|
||||
{names: []string{"exec"}, flags: []string{
|
||||
"-h", "--help", "-f", "--file", "--image", "--add-dir", "--mode", "-m", "--model",
|
||||
|
||||
+49
-14
@@ -442,22 +442,57 @@ 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
|
||||
// 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 {
|
||||
modelSwitcher, modelSessionSwitcher = providers.EscalationSwitchers(
|
||||
resolved.Provider, provider, deps.newProvider,
|
||||
func(modelID string) { currentModel = modelID },
|
||||
)
|
||||
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.
|
||||
turnSessions, _ := providers.OptimizedTurnSessions(resolved.Provider, provider, providers.Options{})
|
||||
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
|
||||
}
|
||||
}
|
||||
|
||||
runMetadata, err := resolveExecRunMetadata(resolved.Provider)
|
||||
|
||||
@@ -1,211 +0,0 @@
|
||||
package cli
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Gitlawb/zero/internal/agent"
|
||||
"github.com/Gitlawb/zero/internal/config"
|
||||
"github.com/Gitlawb/zero/internal/sandbox"
|
||||
"github.com/Gitlawb/zero/internal/tui"
|
||||
"github.com/Gitlawb/zero/internal/zeroruntime"
|
||||
)
|
||||
|
||||
// rootFlagOrderings is every ordering of the given flag groups, each group
|
||||
// kept intact so a flag stays next to its value.
|
||||
func rootFlagOrderings(groups [][]string) [][]string {
|
||||
if len(groups) == 0 {
|
||||
return [][]string{{}}
|
||||
}
|
||||
var orderings [][]string
|
||||
for index, group := range groups {
|
||||
rest := append(append([][]string{}, groups[:index]...), groups[index+1:]...)
|
||||
for _, tail := range rootFlagOrderings(rest) {
|
||||
orderings = append(orderings, append(append([]string{}, group...), tail...))
|
||||
}
|
||||
}
|
||||
return orderings
|
||||
}
|
||||
|
||||
// ROOT FLAGS COMPOSE IN ANY ORDER. Each leading-flag splitter stops at the first
|
||||
// token it does not own, so running them once in a fixed sequence made the
|
||||
// order the operator wrote them in load-bearing: `zero --allow-escalation
|
||||
// --theme auto` stranded `--theme auto` as an unknown command and exited with
|
||||
// an argument error instead of launching. Every ordering of the three root
|
||||
// flags, with --skip-permissions-unsafe absent or at any position, has to reach
|
||||
// the TUI with all three applied.
|
||||
func TestRootFlagsComposeInAnyOrder(t *testing.T) {
|
||||
extra := t.TempDir()
|
||||
resolvedExtra, err := filepath.EvalSymlinks(extra)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
groups := [][]string{{"--add-dir", extra}, {"--theme", "auto"}, {"--allow-escalation"}}
|
||||
cases := rootFlagOrderings(groups)
|
||||
cases = append(cases, rootFlagOrderings(append(groups, []string{"--skip-permissions-unsafe"}))...)
|
||||
if len(cases) != 6+24 {
|
||||
t.Fatalf("SETUP INVALID: %d orderings, want 30", len(cases))
|
||||
}
|
||||
for _, args := range cases {
|
||||
t.Run(strings.Join(args, " "), func(t *testing.T) {
|
||||
options := captureTUIOptions(t, args...)
|
||||
wantMode := agent.PermissionModeAsk
|
||||
for _, arg := range args {
|
||||
if arg == "--skip-permissions-unsafe" {
|
||||
wantMode = agent.PermissionModeUnsafe
|
||||
}
|
||||
}
|
||||
if options.PermissionMode != wantMode {
|
||||
t.Errorf("PermissionMode = %q, want %q", options.PermissionMode, wantMode)
|
||||
}
|
||||
if !options.AllowEscalation {
|
||||
t.Error("escalation not enabled")
|
||||
}
|
||||
if options.Theme != "auto" {
|
||||
t.Errorf("Theme = %q, want auto", options.Theme)
|
||||
}
|
||||
if options.AgentOptions.Sandbox == nil {
|
||||
t.Fatal("no sandbox engine on the launched options")
|
||||
}
|
||||
roots := options.AgentOptions.Sandbox.Scope().Roots()
|
||||
found := false
|
||||
for _, root := range roots {
|
||||
if root == resolvedExtra {
|
||||
found = true
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
t.Errorf("scope roots = %v, want the --add-dir root %q", roots, resolvedExtra)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// The last --theme wins even with another root flag between two of them, and a
|
||||
// --theme written before --skip-permissions-unsafe reaches the TUI: the unsafe
|
||||
// path used to launch with only the theme written after the flag, dropping one
|
||||
// written before it.
|
||||
func TestRootThemeLastOccurrenceWinsAcrossOtherFlags(t *testing.T) {
|
||||
for _, testCase := range []struct {
|
||||
args []string
|
||||
want string
|
||||
}{
|
||||
{[]string{"--theme", "light", "--allow-escalation", "--theme", "auto"}, "auto"},
|
||||
{[]string{"--theme", "auto", "--skip-permissions-unsafe"}, "auto"},
|
||||
{[]string{"--theme", "light", "--skip-permissions-unsafe", "--theme", "auto"}, "auto"},
|
||||
} {
|
||||
t.Run(strings.Join(testCase.args, " "), func(t *testing.T) {
|
||||
options := captureTUIOptions(t, testCase.args...)
|
||||
if options.Theme != testCase.want {
|
||||
t.Fatalf("Theme = %q, want %q", options.Theme, testCase.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// THE ROOT FLAG REACHES EXEC. A --allow-escalation written before the
|
||||
// subcommand is forwarded the way --add-dir is, so `zero --allow-escalation -p
|
||||
// "..."` runs with the opt-in the help text promises instead of silently
|
||||
// without it. The negative control is the same run without the flag, which
|
||||
// must not advertise escalate_model.
|
||||
func TestRootAllowEscalationForwardsIntoExec(t *testing.T) {
|
||||
for _, testCase := range []struct {
|
||||
name string
|
||||
args []string
|
||||
}{
|
||||
{"exec subcommand", []string{"--allow-escalation", "exec", "say hi"}},
|
||||
{"prompt flag", []string{"--allow-escalation", "-p", "say hi"}},
|
||||
} {
|
||||
t.Run(testCase.name, func(t *testing.T) {
|
||||
if !execAdvertisesEscalateModel(t, testCase.args) {
|
||||
t.Fatalf("%v ran without escalate_model advertised: the root flag was dropped before exec", testCase.args)
|
||||
}
|
||||
if execAdvertisesEscalateModel(t, testCase.args[1:]) {
|
||||
t.Fatalf("%v advertised escalate_model without the flag", testCase.args[1:])
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// Everywhere else the flag is rejected loudly rather than discarded, the way
|
||||
// --add-dir is: help and version run no agent and could only ignore it.
|
||||
func TestRootAllowEscalationIsRejectedWhereNothingConsumesIt(t *testing.T) {
|
||||
for _, args := range [][]string{
|
||||
{"--allow-escalation", "version"},
|
||||
{"--allow-escalation", "help"},
|
||||
{"--allow-escalation", "completions", "bash"},
|
||||
} {
|
||||
t.Run(strings.Join(args, " "), func(t *testing.T) {
|
||||
var stdout, stderr bytes.Buffer
|
||||
launched := false
|
||||
exitCode := runWithDeps(args, &stdout, &stderr, appDeps{
|
||||
getwd: func() (string, error) { return t.TempDir(), nil },
|
||||
runTUI: func(context.Context, tui.Options) int { launched = true; return exitSuccess },
|
||||
})
|
||||
if exitCode == exitSuccess || launched {
|
||||
t.Fatalf("exit = %d launched = %v: the flag was discarded instead of rejected", exitCode, launched)
|
||||
}
|
||||
if !strings.Contains(stderr.String(), "--allow-escalation is only supported for the interactive TUI and exec") {
|
||||
t.Fatalf("stderr does not name the flag: %s", stderr.String())
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// execAdvertisesEscalateModel dispatches args through the root command with a
|
||||
// provider that records the tools advertised on the first request, and reports
|
||||
// whether escalate_model was among them.
|
||||
func execAdvertisesEscalateModel(t *testing.T, args []string) bool {
|
||||
t.Helper()
|
||||
isolateCLIUserState(t)
|
||||
provider := &toolListingProvider{}
|
||||
var stdout, stderr bytes.Buffer
|
||||
exitCode := runWithDeps(args, &stdout, &stderr, appDeps{
|
||||
getwd: func() (string, error) { return t.TempDir(), nil },
|
||||
resolveConfig: func(string, config.Overrides) (config.ResolvedConfig, error) {
|
||||
return execResolvedConfig(), nil
|
||||
},
|
||||
newProvider: func(config.ProviderProfile) (zeroruntime.Provider, error) {
|
||||
return provider, nil
|
||||
},
|
||||
newSandboxStore: func() (*sandbox.GrantStore, error) {
|
||||
return sandbox.NewGrantStore(sandbox.StoreOptions{FilePath: filepath.Join(t.TempDir(), "sandbox-grants.json")})
|
||||
},
|
||||
})
|
||||
if exitCode != exitSuccess {
|
||||
t.Fatalf("%v: exit = %d, stderr = %s", args, exitCode, stderr.String())
|
||||
}
|
||||
if provider.requests == 0 {
|
||||
t.Fatalf("SETUP INVALID: %v never reached the provider", args)
|
||||
}
|
||||
for _, name := range provider.toolNames {
|
||||
if name == "escalate_model" {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// toolListingProvider answers every request with a one-line text reply and
|
||||
// keeps the tool names advertised on the first one.
|
||||
type toolListingProvider struct {
|
||||
requests int
|
||||
toolNames []string
|
||||
}
|
||||
|
||||
func (provider *toolListingProvider) StreamCompletion(_ context.Context, request zeroruntime.CompletionRequest) (<-chan zeroruntime.StreamEvent, error) {
|
||||
provider.requests++
|
||||
if provider.requests == 1 {
|
||||
for _, tool := range request.Tools {
|
||||
provider.toolNames = append(provider.toolNames, tool.Name)
|
||||
}
|
||||
}
|
||||
ch := make(chan zeroruntime.StreamEvent, 2)
|
||||
ch <- zeroruntime.StreamEvent{Type: zeroruntime.StreamEventText, Content: "hi"}
|
||||
ch <- zeroruntime.StreamEvent{Type: zeroruntime.StreamEventDone}
|
||||
close(ch)
|
||||
return ch, nil
|
||||
}
|
||||
@@ -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, false)
|
||||
return runInteractiveTUIWithSetup(stderr, deps, "", nil, "", true)
|
||||
}
|
||||
|
||||
result, err := saveSetupProvider(deps, tui.SetupSelection{
|
||||
|
||||
@@ -1,126 +0,0 @@
|
||||
package cli
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Gitlawb/zero/internal/tui"
|
||||
)
|
||||
|
||||
// captureTUIOptions runs the root command with the TUI launch intercepted, and
|
||||
// returns the options the interactive session would have started with.
|
||||
func captureTUIOptions(t *testing.T, args ...string) tui.Options {
|
||||
t.Helper()
|
||||
isolateCLIUserState(t)
|
||||
var captured tui.Options
|
||||
var stdout, stderr bytes.Buffer
|
||||
workspace := t.TempDir()
|
||||
exitCode := runWithDeps(args, &stdout, &stderr, appDeps{
|
||||
getwd: func() (string, error) { return workspace, nil },
|
||||
runTUI: func(_ context.Context, options tui.Options) int {
|
||||
captured = options
|
||||
return exitSuccess
|
||||
},
|
||||
})
|
||||
if exitCode != exitSuccess {
|
||||
t.Fatalf("exitCode = %d stdout=%s stderr=%s", exitCode, stdout.String(), stderr.String())
|
||||
}
|
||||
return captured
|
||||
}
|
||||
|
||||
func hasEscalateModel(options tui.Options) bool {
|
||||
if options.AgentOptions.Registry == nil {
|
||||
return false
|
||||
}
|
||||
_, registered := options.AgentOptions.Registry.Get("escalate_model")
|
||||
return registered
|
||||
}
|
||||
|
||||
// ESCALATION IS OFF UNLESS THE OPERATOR ASKS FOR IT.
|
||||
//
|
||||
// Escalation moves a run onto a different model, changing what it costs and which
|
||||
// provider sees the conversation, so the interactive surface answers this the same
|
||||
// conservative way `zero exec --allow-escalation` already does.
|
||||
func TestInteractiveTUIRegistersEscalateModelOnlyWithTheFlag(t *testing.T) {
|
||||
if hasEscalateModel(captureTUIOptions(t)) {
|
||||
t.Error("escalate_model was registered without --allow-escalation")
|
||||
}
|
||||
if !hasEscalateModel(captureTUIOptions(t, "--allow-escalation")) {
|
||||
t.Error("--allow-escalation did not register escalate_model")
|
||||
}
|
||||
}
|
||||
|
||||
// THE TOOL AND THE SWITCHERS ARE ONE FEATURE, NOT TWO.
|
||||
//
|
||||
// This is the failure the whole change exists to avoid. The agent loop performs a
|
||||
// switch only when a switcher is wired, so registering escalate_model without
|
||||
// wiring one ships a tool the model can call and the loop will silently ignore:
|
||||
// the run reports an escalation that never happened. The reverse is merely dead
|
||||
// code. Both halves ride on the same flag and this asserts they cannot be
|
||||
// separated, whichever half a future change touches.
|
||||
func TestInteractiveTUIEscalationToolAndSwitchersAreWiredTogether(t *testing.T) {
|
||||
for _, testCase := range []struct {
|
||||
name string
|
||||
args []string
|
||||
want bool
|
||||
}{
|
||||
{name: "default", args: nil, want: false},
|
||||
{name: "flagged", args: []string{"--allow-escalation"}, want: true},
|
||||
} {
|
||||
t.Run(testCase.name, func(t *testing.T) {
|
||||
options := captureTUIOptions(t, testCase.args...)
|
||||
if got := hasEscalateModel(options); got != testCase.want {
|
||||
t.Fatalf("escalate_model registered = %v, want %v", got, testCase.want)
|
||||
}
|
||||
if options.AllowEscalation != testCase.want {
|
||||
t.Errorf("AllowEscalation = %v, want %v: the tool is registered on one gate and the switchers on another, so they can drift apart",
|
||||
options.AllowEscalation, testCase.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// The flag is accepted on either side of --skip-permissions-unsafe, like the
|
||||
// other root flags that may appear there.
|
||||
func TestInteractiveTUIAcceptsAllowEscalationAroundUnsafe(t *testing.T) {
|
||||
for _, args := range [][]string{
|
||||
{"--allow-escalation", "--skip-permissions-unsafe"},
|
||||
{"--skip-permissions-unsafe", "--allow-escalation"},
|
||||
} {
|
||||
options := captureTUIOptions(t, args...)
|
||||
if !options.AllowEscalation {
|
||||
t.Errorf("%v did not enable escalation", args)
|
||||
}
|
||||
if !hasEscalateModel(options) {
|
||||
t.Errorf("%v did not register escalate_model", args)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// An =value form is a loud error rather than a silent enable, so a mistyped
|
||||
// --allow-escalation=false cannot turn the feature ON.
|
||||
func TestAllowEscalationRejectsAValue(t *testing.T) {
|
||||
var stdout, stderr bytes.Buffer
|
||||
exitCode := runWithDeps([]string{"--allow-escalation=false"}, &stdout, &stderr, appDeps{
|
||||
getwd: func() (string, error) { return t.TempDir(), nil },
|
||||
runTUI: func(context.Context, tui.Options) int { return exitSuccess },
|
||||
})
|
||||
if exitCode == exitSuccess {
|
||||
t.Fatal("--allow-escalation=false was accepted; a mistyped disable must not silently enable escalation")
|
||||
}
|
||||
if !strings.Contains(stderr.String(), "--allow-escalation takes no value") {
|
||||
t.Errorf("stderr does not explain the flag: %s", stderr.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestRootHelpDocumentsAllowEscalation(t *testing.T) {
|
||||
var stdout, stderr bytes.Buffer
|
||||
if exitCode := runWithDeps([]string{"--help"}, &stdout, &stderr, appDeps{}); exitCode != exitSuccess {
|
||||
t.Fatalf("exitCode = %d stderr=%s", exitCode, stderr.String())
|
||||
}
|
||||
if !strings.Contains(stdout.String(), "--allow-escalation") {
|
||||
t.Error("the root help does not mention --allow-escalation, so the flag is undiscoverable")
|
||||
}
|
||||
}
|
||||
@@ -1,139 +0,0 @@
|
||||
package cli
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Gitlawb/zero/internal/config"
|
||||
)
|
||||
|
||||
// userStateEnvNames is every per-user base directory the root command can
|
||||
// resolve against, across platforms. One throwaway root behind all of them
|
||||
// lands the resolution inside it whichever branch it takes.
|
||||
var userStateEnvNames = []string{
|
||||
"HOME", "USERPROFILE", "APPDATA", "LOCALAPPDATA",
|
||||
"XDG_CONFIG_HOME", "XDG_DATA_HOME", "XDG_CACHE_HOME", "XDG_STATE_HOME",
|
||||
}
|
||||
|
||||
// isolateCLIUserState points every user-scoped path the root command touches at
|
||||
// throwaway directories, and keeps credential work off the host keyring.
|
||||
//
|
||||
// A TEMPORARY WORKING DIRECTORY IS NOT ISOLATION. runWithDeps fills the
|
||||
// dependencies a test leaves out with the production ones, and the launch paths
|
||||
// read user config, open stores, refresh the models.dev cache and, for an
|
||||
// interactive run, migrate any inline plaintext API key into the credential
|
||||
// store before they ever reach an injected callback. Without this, running these
|
||||
// tests on a developer machine rewrites that developer's config.json and moves
|
||||
// their key into their keychain, makes a real network call, and reports results
|
||||
// that depend on whatever providers, MCP servers and plugins that machine has.
|
||||
func isolateCLIUserState(t *testing.T) {
|
||||
t.Helper()
|
||||
root := t.TempDir()
|
||||
for _, name := range userStateEnvNames {
|
||||
t.Setenv(name, root)
|
||||
}
|
||||
// The models.dev overlay cache, and the background refresh that writes it:
|
||||
// a wiring test has no business making a network call or leaving a cache
|
||||
// file behind.
|
||||
t.Setenv("ZERO_MODELS_CACHE_PATH", filepath.Join(root, "modelsdev.json"))
|
||||
t.Setenv("ZERO_DISABLE_MODELS_FETCH", "1")
|
||||
// Credentials resolve keyring-first. This keeps a migrated key in a file
|
||||
// under the fixture root rather than in the developer's OS keychain.
|
||||
t.Setenv("ZERO_CRED_STORAGE", "encrypted-file")
|
||||
}
|
||||
|
||||
// AND THE ISOLATION IS ITSELF PINNED, ONE HELPER AT A TIME. A seeded config
|
||||
// outside the helper fixture roots stands in for a developer's real one: the
|
||||
// launch helper must leave it exactly as it found it, inline API key included,
|
||||
// must not write a credential file beside it, and must have moved the per-user
|
||||
// paths off that root before running anything.
|
||||
//
|
||||
// EACH HELPER GETS ITS OWN SUBTEST because t.Setenv restores at the end of the
|
||||
// test, not when the helper returns: exercising both in one test would let the
|
||||
// first helper isolation cover for a second helper that had none.
|
||||
func TestCLILaunchHelpersLeaveUserStateAlone(t *testing.T) {
|
||||
for _, testCase := range []struct {
|
||||
name string
|
||||
exercise func(*testing.T)
|
||||
}{
|
||||
{"captureTUIOptions", func(t *testing.T) { _ = captureTUIOptions(t, "--allow-escalation") }},
|
||||
{"execAdvertisesEscalateModel", func(t *testing.T) {
|
||||
_ = execAdvertisesEscalateModel(t, []string{"--allow-escalation", "exec", "say hi"})
|
||||
}},
|
||||
} {
|
||||
t.Run(testCase.name, func(t *testing.T) {
|
||||
seedRoot, configPath, seeded, before := seedEmulatedUserConfig(t)
|
||||
|
||||
testCase.exercise(t)
|
||||
|
||||
after, err := os.ReadFile(configPath)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if string(after) != string(seeded) {
|
||||
t.Errorf("the seeded user config was rewritten:\n before: %s\n after: %s", seeded, after)
|
||||
}
|
||||
entries, err := os.ReadDir(filepath.Dir(configPath))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(entries) != len(before) {
|
||||
names := []string{}
|
||||
for _, entry := range entries {
|
||||
names = append(names, entry.Name())
|
||||
}
|
||||
t.Errorf("the user config directory gained entries: %v", names)
|
||||
}
|
||||
|
||||
// The checks above catch what this platform happens to write. This
|
||||
// one catches the missing isolation itself, so a startup path that
|
||||
// writes somewhere else, or only sometimes, is covered too.
|
||||
resolved, err := config.DefaultUserConfigPath()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if strings.HasPrefix(resolved, seedRoot) {
|
||||
t.Errorf("user config still resolves inside the seeded root (%s): the helper installed no isolation", resolved)
|
||||
}
|
||||
if os.Getenv("ZERO_DISABLE_MODELS_FETCH") == "" {
|
||||
t.Error("the helper left the background models.dev refresh enabled, so this test can make a network call and write a cache file")
|
||||
}
|
||||
if os.Getenv("ZERO_CRED_STORAGE") == "" {
|
||||
t.Error("the helper left credential storage resolving keyring-first, so a migrated key can reach the host keychain")
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// seedEmulatedUserConfig points the per-user base directories at a throwaway
|
||||
// root and writes a config carrying an inline plaintext API key there, which is
|
||||
// what the startup migration rewrites. It returns that root, the config path,
|
||||
// its bytes, and the directory listing to compare against.
|
||||
func seedEmulatedUserConfig(t *testing.T) (string, string, []byte, []os.DirEntry) {
|
||||
t.Helper()
|
||||
seedRoot := t.TempDir()
|
||||
for _, name := range userStateEnvNames {
|
||||
t.Setenv(name, seedRoot)
|
||||
}
|
||||
configPath, err := config.DefaultUserConfigPath()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !strings.HasPrefix(configPath, seedRoot) {
|
||||
t.Fatalf("SETUP INVALID: user config resolved to %q, outside the seeded root %q", configPath, seedRoot)
|
||||
}
|
||||
if err := os.MkdirAll(filepath.Dir(configPath), 0o755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
seeded := []byte(`{"providers":[{"name":"anthropic","providerKind":"anthropic","model":"claude-haiku-4.5","apiKey":"sk-seeded-plaintext-key"}]}`)
|
||||
if err := os.WriteFile(configPath, seeded, 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
before, err := os.ReadDir(filepath.Dir(configPath))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return seedRoot, configPath, seeded, before
|
||||
}
|
||||
@@ -19,6 +19,7 @@ import (
|
||||
"github.com/Gitlawb/zero/internal/providercatalog"
|
||||
"github.com/Gitlawb/zero/internal/providers"
|
||||
"github.com/Gitlawb/zero/internal/providers/providerio"
|
||||
"github.com/Gitlawb/zero/internal/proxydial"
|
||||
"github.com/Gitlawb/zero/internal/redaction"
|
||||
)
|
||||
|
||||
@@ -518,9 +519,15 @@ func newConnectivityClient(timeout time.Duration, resolver Resolver, sensitiveHe
|
||||
transport = &http.Transport{}
|
||||
}
|
||||
transport.DialContext = safeDialContext(resolver, allowLoopbackOrPrivate)
|
||||
// WRAPPED, SO THE DIALER LEARNS THE ROUTE THIS TRANSPORT CHOSE. The wrapper
|
||||
// records the proxy selected for each request on that request, and the
|
||||
// dialer reads it back on the dial that request causes; installing the bare
|
||||
// transport instead would leave a proxied dial looking like a loopback
|
||||
// target and fail the probe (#569). The cloned default transport already
|
||||
// reads HTTPS_PROXY and friends; nothing new is enabled.
|
||||
return &http.Client{
|
||||
Timeout: timeout,
|
||||
Transport: transport,
|
||||
Transport: proxydial.Wrap(transport),
|
||||
CheckRedirect: func(req *http.Request, via []*http.Request) error {
|
||||
if len(via) >= maxConnectivityRedirects {
|
||||
return fmt.Errorf("provider connectivity exceeded %d redirects", maxConnectivityRedirects)
|
||||
@@ -583,6 +590,17 @@ func sensitiveAuthHeaderNames(profile config.ProviderProfile, kind config.Provid
|
||||
func safeDialContext(resolver Resolver, allowLoopbackOrPrivate bool) func(context.Context, string, string) (net.Conn, error) {
|
||||
dialer := &net.Dialer{Timeout: 10 * time.Second, KeepAlive: 30 * time.Second}
|
||||
return func(ctx context.Context, network, address string) (net.Conn, error) {
|
||||
// THE PROXY IS NOT THE TARGET. With HTTPS_PROXY set to a local forward
|
||||
// proxy, the transport dials the proxy first and tunnels the request
|
||||
// through it, and this dialer used to refuse that dial as a loopback
|
||||
// address: "proxyconnect tcp: ... loopback hosts are blocked" for a
|
||||
// request whose real target had already been validated (#569). The
|
||||
// target validation is unchanged; only the dial the transport opens to
|
||||
// the proxy it picked for this very request is let past the guard, and a
|
||||
// direct dial is checked as it always was.
|
||||
if proxydial.IsProxyDial(ctx) {
|
||||
return dialer.DialContext(ctx, network, address)
|
||||
}
|
||||
host, port, err := net.SplitHostPort(address)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
|
||||
@@ -0,0 +1,218 @@
|
||||
package providerhealth
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/netip"
|
||||
"net/url"
|
||||
"strconv"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Gitlawb/zero/internal/proxydial"
|
||||
)
|
||||
|
||||
// connectivityClientWithProxy builds the real connectivity client and swaps in
|
||||
// a proxy function, which is how the transport learns its route. The
|
||||
// environment is left alone deliberately: http.ProxyFromEnvironment reads it
|
||||
// once per process, so an environment-driven test would depend on test order.
|
||||
func connectivityClientWithProxy(t *testing.T, resolver Resolver, proxyFor func(*http.Request) (*url.URL, error)) *http.Client {
|
||||
t.Helper()
|
||||
client := newConnectivityClient(3*time.Second, resolver, nil, false)
|
||||
wrapped, ok := client.Transport.(*proxydial.Transport)
|
||||
if !ok {
|
||||
t.Fatalf("transport = %T, want the route-recording wrapper; without it no dial can be recognised as the proxy dial", client.Transport)
|
||||
}
|
||||
wrapped.Proxy = proxyFor
|
||||
return client
|
||||
}
|
||||
|
||||
// recordingProxy is a listener standing in for a local forward proxy. It
|
||||
// records the first line it receives and hangs up.
|
||||
func recordingProxy(t *testing.T) (*url.URL, <-chan string) {
|
||||
t.Helper()
|
||||
listener, err := net.Listen("tcp", "127.0.0.1:0")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = listener.Close() })
|
||||
arrived := make(chan string, 1)
|
||||
go func() {
|
||||
conn, err := listener.Accept()
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
defer conn.Close()
|
||||
buffer := make([]byte, 256)
|
||||
n, _ := conn.Read(buffer)
|
||||
arrived <- string(buffer[:n])
|
||||
}()
|
||||
proxyURL, err := url.Parse("http://" + listener.Addr().String())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return proxyURL, arrived
|
||||
}
|
||||
|
||||
// THE PROXY IS NOT THE TARGET, AND THE GUARD HAS TO KNOW THE DIFFERENCE.
|
||||
//
|
||||
// With HTTPS_PROXY set to a local forward proxy, the transport dials the proxy
|
||||
// first and tunnels the request through it. This dialer refused that dial as a
|
||||
// loopback address, so `zero doctor --connectivity` failed with "proxyconnect
|
||||
// tcp: ... loopback hosts are blocked" for a target that had already been
|
||||
// validated and was public (#569). Driven through the real client, so the route
|
||||
// is the transport's own decision.
|
||||
func TestConnectivityClientTunnelsThroughTheConfiguredProxy(t *testing.T) {
|
||||
proxyURL, arrived := recordingProxy(t)
|
||||
client := connectivityClientWithProxy(t,
|
||||
staticResolver{err: errors.New("resolver must not be called for a proxied request")},
|
||||
func(*http.Request) (*url.URL, error) { return proxyURL, nil })
|
||||
|
||||
request, err := http.NewRequestWithContext(context.Background(), http.MethodGet, "https://api.example.com/v1/models", nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
_, err = client.Do(request)
|
||||
select {
|
||||
case first := <-arrived:
|
||||
if line := strings.SplitN(first, "\r\n", 2)[0]; line != "CONNECT api.example.com:443 HTTP/1.1" {
|
||||
t.Fatalf("proxy received %q, want a CONNECT for the validated target", line)
|
||||
}
|
||||
case <-time.After(5 * time.Second):
|
||||
t.Fatalf("the proxy never received a connection; request error: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// AND A DIRECT REQUEST STAYS GUARDED, EVEN WHEN ITS ORIGIN IS THE PROXY'S OWN
|
||||
// ADDRESS.
|
||||
//
|
||||
// This is the NO_PROXY shape: the proxy function sends one host direct. An
|
||||
// exemption derived from asking that function about a request of our own
|
||||
// invention would answer "proxied" for the excluded host too, and this dial,
|
||||
// whose address equals the configured proxy endpoint, would skip the resolver
|
||||
// and reach whatever the name resolves to at connect time.
|
||||
func TestDirectRequestToTheProxyEndpointIsStillGuarded(t *testing.T) {
|
||||
proxyURL, err := url.Parse("http://rebind.example:80")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
lookups := 0
|
||||
resolver := countingResolver{
|
||||
onLookup: func() { lookups++ },
|
||||
addr: netip.MustParseAddr("10.0.0.5"),
|
||||
}
|
||||
client := connectivityClientWithProxy(t, resolver, func(request *http.Request) (*url.URL, error) {
|
||||
if request.URL.Hostname() == "rebind.example" {
|
||||
return nil, nil
|
||||
}
|
||||
return proxyURL, nil
|
||||
})
|
||||
|
||||
request, err := http.NewRequestWithContext(context.Background(), http.MethodGet, "http://rebind.example/data", nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
_, err = client.Do(request)
|
||||
|
||||
if lookups == 0 {
|
||||
t.Fatal("the direct dial skipped the safety resolver because its address matched the configured proxy endpoint")
|
||||
}
|
||||
var safety endpointSafetyError
|
||||
if !errors.As(err, &safety) {
|
||||
t.Fatalf("err = %v, want the private answer to be refused", err)
|
||||
}
|
||||
}
|
||||
|
||||
// A proxy configured with a Unicode hostname is dialed by its IDNA-encoded
|
||||
// name, and it is still the proxy: the request must reach it rather than be
|
||||
// refused as an ordinary private address.
|
||||
func TestUnicodeProxyHostnameIsStillTheProxy(t *testing.T) {
|
||||
dialed := ""
|
||||
client := connectivityClientWithProxy(t,
|
||||
staticResolver{err: errors.New("resolver must not be called for a proxied request")},
|
||||
func(*http.Request) (*url.URL, error) { return url.Parse("http://büro.example:3067") })
|
||||
wrapped := client.Transport.(*proxydial.Transport)
|
||||
guarded := wrapped.DialContext
|
||||
wrapped.DialContext = func(ctx context.Context, network, address string) (net.Conn, error) {
|
||||
dialed = address
|
||||
return guarded(ctx, network, address)
|
||||
}
|
||||
|
||||
request, err := http.NewRequestWithContext(context.Background(), http.MethodGet, "https://api.example.com/v1/models", nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
_, err = client.Do(request)
|
||||
|
||||
if dialed != "xn--bro-hoa.example:3067" {
|
||||
t.Fatalf("SETUP INVALID: dialed %q, want the IDNA-encoded proxy endpoint", dialed)
|
||||
}
|
||||
var safety endpointSafetyError
|
||||
if errors.As(err, &safety) {
|
||||
t.Fatalf("the configured proxy was refused by the guard: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// Without a proxy the guard is exactly what it was: loopback is refused.
|
||||
func TestSafeDialContextStillRefusesLoopbackWithoutAProxy(t *testing.T) {
|
||||
dial := safeDialContext(staticResolver{err: errors.New("resolver must not be called")}, false)
|
||||
conn, err := dial(context.Background(), "tcp", "127.0.0.1:3067")
|
||||
if conn != nil {
|
||||
_ = conn.Close()
|
||||
t.Fatal("loopback was dialed with no proxy configured")
|
||||
}
|
||||
var safety endpointSafetyError
|
||||
if !errors.As(err, &safety) {
|
||||
t.Fatalf("err = %v, want endpointSafetyError", err)
|
||||
}
|
||||
}
|
||||
|
||||
// A sibling loopback port is not the proxy, and the guard is unchanged for it.
|
||||
func TestSafeDialContextRefusesALoopbackPortThatIsNotTheProxy(t *testing.T) {
|
||||
proxyURL, arrived := recordingProxy(t)
|
||||
client := connectivityClientWithProxy(t,
|
||||
staticResolver{err: errors.New("resolver must not be called")},
|
||||
func(request *http.Request) (*url.URL, error) {
|
||||
// Only the https probe is proxied; the http one below is direct and
|
||||
// happens to point at the neighbouring port.
|
||||
if request.URL.Scheme == "https" {
|
||||
return proxyURL, nil
|
||||
}
|
||||
return nil, nil
|
||||
})
|
||||
|
||||
_, port, _ := net.SplitHostPort(proxyURL.Host)
|
||||
number, _ := strconv.Atoi(port)
|
||||
sibling := "http://" + net.JoinHostPort("127.0.0.1", strconv.Itoa(number+1)) + "/"
|
||||
request, err := http.NewRequestWithContext(context.Background(), http.MethodGet, sibling, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
_, err = client.Do(request)
|
||||
var safety endpointSafetyError
|
||||
if !errors.As(err, &safety) {
|
||||
t.Fatalf("err = %v, want a loopback port beside the proxy to stay blocked", err)
|
||||
}
|
||||
select {
|
||||
case first := <-arrived:
|
||||
t.Fatalf("the proxy received a connection it should not have: %q", first)
|
||||
default:
|
||||
}
|
||||
}
|
||||
|
||||
// countingResolver answers with one address and reports each lookup, so a test
|
||||
// can tell a skipped guard from a guard that ran and allowed the address.
|
||||
type countingResolver struct {
|
||||
onLookup func()
|
||||
addr netip.Addr
|
||||
}
|
||||
|
||||
func (resolver countingResolver) LookupNetIP(context.Context, string, string) ([]netip.Addr, error) {
|
||||
if resolver.onLookup != nil {
|
||||
resolver.onLookup()
|
||||
}
|
||||
return []netip.Addr{resolver.addr}, nil
|
||||
}
|
||||
@@ -1,89 +0,0 @@
|
||||
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
|
||||
}
|
||||
@@ -1,130 +0,0 @@
|
||||
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")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,93 @@
|
||||
// Package proxydial answers one question for an SSRF-guarded dialer: is this
|
||||
// dial the connection dial of a request the transport routed through a forward
|
||||
// proxy?
|
||||
//
|
||||
// A guarded transport validates the TARGET of a request before dialing, and its
|
||||
// dialer refuses loopback, private and link-local addresses so a redirect or a
|
||||
// DNS rebind cannot walk the request onto the local network. When a forward
|
||||
// proxy is configured, the transport dials the PROXY first and tunnels the
|
||||
// request through it, and a local proxy is by far the common shape:
|
||||
// HTTPS_PROXY=http://127.0.0.1:3067. The dialer then sees 127.0.0.1, refuses it,
|
||||
// and the user gets "loopback hosts are blocked" for a request whose real target
|
||||
// was validated and public (#569).
|
||||
//
|
||||
// The proxy address is not the request target, and the guard was never about it.
|
||||
// It is configuration the user set on purpose, and the transport is going to dial
|
||||
// it for every request whether the guard likes it or not. So that dial is let
|
||||
// through, and only that dial.
|
||||
//
|
||||
// THE EXEMPTION FOLLOWS THE TRANSPORT'S OWN ROUTE DECISION, NOT A GUESS AT IT.
|
||||
// Asking the Proxy function a question of one's own invention answers for a
|
||||
// request the transport never made: with NO_PROXY excluding a host, the real
|
||||
// request to that host goes direct while a synthetic probe still comes back
|
||||
// proxied, and a direct dial whose address happens to equal the proxy endpoint
|
||||
// would then skip the guard entirely. Host-dependent proxy callbacks make the
|
||||
// mismatch general, and no amount of care over the probe's spelling closes it,
|
||||
// because the two calls are about different requests.
|
||||
//
|
||||
// So the decision is recorded where it is made, on the request, and read back on
|
||||
// the dial that request causes. Transport.RoundTrip asks the wrapped transport's
|
||||
// Proxy function about the actual request, stores the answer in that request's
|
||||
// context, and hands it on; the transport derives the dial context from the
|
||||
// request context, so the dialer reads back the route this very connection is
|
||||
// being opened for. Per-request, so concurrent requests cannot overwrite each
|
||||
// other, and unforgeable from outside the package: the context key is private.
|
||||
package proxydial
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"net/url"
|
||||
)
|
||||
|
||||
// selectedProxyKey addresses the route decision inside a request context. It is
|
||||
// unexported so nothing outside this package can plant one.
|
||||
type selectedProxyKey struct{}
|
||||
|
||||
// Transport records the proxy its wrapped transport selects for each request in
|
||||
// that request's context, so a guarded DialContext installed on the same
|
||||
// transport can tell the dial to a proxy from a direct dial.
|
||||
//
|
||||
// The embedded transport stays reachable: callers configure Proxy, DialContext
|
||||
// and the rest on it exactly as before, and tests can read them back.
|
||||
type Transport struct {
|
||||
*http.Transport
|
||||
}
|
||||
|
||||
// Wrap returns a Transport around transport. The result is the http.RoundTripper
|
||||
// to install on the client; installing the bare transport instead leaves every
|
||||
// dial looking direct, which is safe (the guard applies) but reopens #569.
|
||||
func Wrap(transport *http.Transport) *Transport {
|
||||
return &Transport{Transport: transport}
|
||||
}
|
||||
|
||||
// RoundTrip resolves the route for this request and passes the decision down to
|
||||
// the dialer through the request context.
|
||||
func (transport *Transport) RoundTrip(request *http.Request) (*http.Response, error) {
|
||||
if transport.Transport.Proxy != nil {
|
||||
// A Proxy error is left for the transport to report: it fails the
|
||||
// request before dialing, so there is no dial to classify. The function
|
||||
// is called once here and once by the transport, which is what a
|
||||
// host-dependent callback already expects from redirects and retries.
|
||||
if proxyURL, err := transport.Transport.Proxy(request); err == nil && proxyURL != nil {
|
||||
request = request.WithContext(context.WithValue(request.Context(), selectedProxyKey{}, proxyURL))
|
||||
}
|
||||
}
|
||||
return transport.Transport.RoundTrip(request)
|
||||
}
|
||||
|
||||
// IsProxyDial reports whether ctx belongs to a request this package's Transport
|
||||
// routed through a proxy, which makes the dial about to happen the dial to that
|
||||
// proxy.
|
||||
//
|
||||
// THE ADDRESS IS NOT COMPARED, AND DOES NOT NEED TO BE. A transport opens
|
||||
// exactly one connection dial per connection, and when a proxy is selected that
|
||||
// dial goes to the proxy; the target is reached through it, never dialed. So the
|
||||
// route decision on the context already identifies the dial, and matching the
|
||||
// address against the proxy URL on top of it could only ever disagree with the
|
||||
// transport, which is what a punycoded Unicode proxy hostname does: Go dials
|
||||
// xn--bro-hoa.example while the URL still spells buero as the user wrote it.
|
||||
func IsProxyDial(ctx context.Context) bool {
|
||||
proxyURL, _ := ctx.Value(selectedProxyKey{}).(*url.URL)
|
||||
return proxyURL != nil
|
||||
}
|
||||
@@ -0,0 +1,224 @@
|
||||
package proxydial
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"sync"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// errDialed ends the request at the dial, after the route decision has been
|
||||
// observed. Nothing here needs a live server.
|
||||
var errDialed = errors.New("dial recorded")
|
||||
|
||||
type dialRecord struct {
|
||||
address string
|
||||
proxied bool
|
||||
}
|
||||
|
||||
// driveRequest sends one request through a wrapped transport whose dialer
|
||||
// records what it was asked for and whether the context says this dial belongs
|
||||
// to a proxied request.
|
||||
func driveRequest(t *testing.T, target string, proxyFor func(*http.Request) (*url.URL, error)) dialRecord {
|
||||
t.Helper()
|
||||
var record dialRecord
|
||||
transport := &http.Transport{Proxy: proxyFor}
|
||||
transport.DialContext = func(ctx context.Context, _, address string) (net.Conn, error) {
|
||||
record = dialRecord{address: address, proxied: IsProxyDial(ctx)}
|
||||
return nil, errDialed
|
||||
}
|
||||
client := &http.Client{Transport: Wrap(transport)}
|
||||
request, err := http.NewRequest(http.MethodGet, target, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := client.Do(request); !errors.Is(err, errDialed) {
|
||||
t.Fatalf("%s: request ended before the dial: %v", target, err)
|
||||
}
|
||||
if record.address == "" {
|
||||
t.Fatalf("%s: the dialer was never called", target)
|
||||
}
|
||||
return record
|
||||
}
|
||||
|
||||
func fixedProxy(raw string) func(*http.Request) (*url.URL, error) {
|
||||
parsed, err := url.Parse(raw)
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
return func(*http.Request) (*url.URL, error) { return parsed, nil }
|
||||
}
|
||||
|
||||
// THE DIAL TO A CONFIGURED PROXY IS THE ONE THAT IS EXEMPT. A proxied request
|
||||
// dials the proxy, never the target, so the route on the context is what
|
||||
// identifies the dial.
|
||||
func TestProxiedRequestDialsTheProxyAndSaysSo(t *testing.T) {
|
||||
record := driveRequest(t, "https://api.example.com/v1/models", fixedProxy("http://127.0.0.1:3067"))
|
||||
if !record.proxied {
|
||||
t.Error("a request routed through a proxy did not reach the dialer as proxied")
|
||||
}
|
||||
if record.address != "127.0.0.1:3067" {
|
||||
t.Errorf("dial address = %q, want the proxy endpoint", record.address)
|
||||
}
|
||||
}
|
||||
|
||||
// AND A DIRECT REQUEST IS NOT, EVEN WHEN ITS ORIGIN IS THE PROXY'S OWN ADDRESS.
|
||||
//
|
||||
// This is the NO_PROXY shape: the proxy function excludes one host, so a request
|
||||
// to that host goes direct. Asking that same function about a request of our own
|
||||
// invention would answer "proxied" (the invented host is not excluded) and the
|
||||
// direct dial, whose address equals the proxy endpoint, would skip the guard
|
||||
// while resolving to whatever the name resolves to at connect time.
|
||||
func TestDirectRequestToTheProxyEndpointIsNotAProxyDial(t *testing.T) {
|
||||
proxyURL, err := url.Parse("http://rebind.example:80")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
noProxyForRebind := func(request *http.Request) (*url.URL, error) {
|
||||
if request.URL.Hostname() == "rebind.example" {
|
||||
return nil, nil
|
||||
}
|
||||
return proxyURL, nil
|
||||
}
|
||||
|
||||
record := driveRequest(t, "http://rebind.example/data", noProxyForRebind)
|
||||
if record.address != "rebind.example:80" {
|
||||
t.Fatalf("SETUP INVALID: dial address = %q, want the direct origin, which is also the proxy endpoint", record.address)
|
||||
}
|
||||
if record.proxied {
|
||||
t.Error("a direct dial was reported as the proxy dial because its address matched the configured proxy")
|
||||
}
|
||||
|
||||
// The same function still routes another host through that proxy, so the
|
||||
// exclusion above is a route difference and not a disabled proxy.
|
||||
if record := driveRequest(t, "http://elsewhere.example/data", noProxyForRebind); !record.proxied {
|
||||
t.Error("the configured proxy was not used for a host outside the exclusion")
|
||||
}
|
||||
}
|
||||
|
||||
// The inverse mismatch: a callback that routes by host must not have its answer
|
||||
// taken from some other host. A proxy selected only for the real target is
|
||||
// honoured.
|
||||
func TestHostDependentProxySelectionIsHonoured(t *testing.T) {
|
||||
proxyURL, err := url.Parse("http://127.0.0.1:3067")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
onlyForOneHost := func(request *http.Request) (*url.URL, error) {
|
||||
if request.URL.Hostname() == "api.example.com" {
|
||||
return proxyURL, nil
|
||||
}
|
||||
return nil, nil
|
||||
}
|
||||
if record := driveRequest(t, "https://api.example.com/v1", onlyForOneHost); !record.proxied || record.address != "127.0.0.1:3067" {
|
||||
t.Errorf("selected proxy not reached: %+v", record)
|
||||
}
|
||||
if record := driveRequest(t, "https://other.example/v1", onlyForOneHost); record.proxied {
|
||||
t.Error("an unproxied host was reported as proxied")
|
||||
}
|
||||
}
|
||||
|
||||
// Scheme separation: an https-only proxy leaves an http request direct, and the
|
||||
// decision recorded is the one made for that request's own scheme.
|
||||
func TestSchemeSeparation(t *testing.T) {
|
||||
proxyURL, err := url.Parse("http://127.0.0.1:3067")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
httpsOnly := func(request *http.Request) (*url.URL, error) {
|
||||
if request.URL.Scheme == "https" {
|
||||
return proxyURL, nil
|
||||
}
|
||||
return nil, nil
|
||||
}
|
||||
if record := driveRequest(t, "https://api.example.com/v1", httpsOnly); !record.proxied {
|
||||
t.Error("the https proxy was not applied to an https request")
|
||||
}
|
||||
if record := driveRequest(t, "http://api.example.com/v1", httpsOnly); record.proxied {
|
||||
t.Error("an http request was treated as proxied by an https-only proxy")
|
||||
}
|
||||
}
|
||||
|
||||
// A UNICODE PROXY HOSTNAME IS STILL THE PROXY. Go dials proxies by their
|
||||
// IDNA-encoded name, so a matcher comparing the URL's own spelling would call
|
||||
// this dial direct and refuse it as an ordinary private address. Nothing is
|
||||
// compared here, so both spellings of one host are one host.
|
||||
func TestUnicodeProxyHostnameIsRecognised(t *testing.T) {
|
||||
unicodeRecord := driveRequest(t, "https://api.example.com/v1", fixedProxy("http://büro.example:3067"))
|
||||
if !unicodeRecord.proxied {
|
||||
t.Error("a proxy configured with a Unicode hostname was not recognised as the proxy")
|
||||
}
|
||||
asciiRecord := driveRequest(t, "https://api.example.com/v1", fixedProxy("http://xn--bro-hoa.example:3067"))
|
||||
if !asciiRecord.proxied {
|
||||
t.Error("a proxy configured with the ASCII spelling was not recognised as the proxy")
|
||||
}
|
||||
// The two spellings name one endpoint, and it is the ASCII one that is
|
||||
// dialed. This is the mismatch a string comparison would have had to close.
|
||||
if unicodeRecord.address != asciiRecord.address {
|
||||
t.Errorf("dial addresses differ across spellings of one host: %q vs %q", unicodeRecord.address, asciiRecord.address)
|
||||
}
|
||||
}
|
||||
|
||||
// A transport with no proxy leaves every dial direct.
|
||||
func TestNoProxyMeansNoExemption(t *testing.T) {
|
||||
if record := driveRequest(t, "http://127.0.0.1:3067/", nil); record.proxied {
|
||||
t.Error("a dial was reported as proxied by a transport with no Proxy function")
|
||||
}
|
||||
if IsProxyDial(context.Background()) {
|
||||
t.Error("a context that never went through the wrapper reported a proxy dial")
|
||||
}
|
||||
}
|
||||
|
||||
// THE DECISION IS PER REQUEST. A shared record of the last selected proxy would
|
||||
// let a proxied request hand its exemption to a direct one running beside it.
|
||||
func TestConcurrentRequestsKeepTheirOwnRoute(t *testing.T) {
|
||||
proxyURL, err := url.Parse("http://127.0.0.1:3067")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
byHost := func(request *http.Request) (*url.URL, error) {
|
||||
if request.URL.Hostname() == "proxied.example" {
|
||||
return proxyURL, nil
|
||||
}
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
var mutex sync.Mutex
|
||||
seen := map[string]bool{}
|
||||
transport := &http.Transport{Proxy: byHost}
|
||||
transport.DialContext = func(ctx context.Context, _, address string) (net.Conn, error) {
|
||||
mutex.Lock()
|
||||
seen[address] = IsProxyDial(ctx)
|
||||
mutex.Unlock()
|
||||
return nil, errDialed
|
||||
}
|
||||
client := &http.Client{Transport: Wrap(transport)}
|
||||
|
||||
var group sync.WaitGroup
|
||||
for range 40 {
|
||||
for _, target := range []string{"http://proxied.example/", "http://direct.example/"} {
|
||||
group.Add(1)
|
||||
go func() {
|
||||
defer group.Done()
|
||||
request, _ := http.NewRequest(http.MethodGet, target, nil)
|
||||
_, _ = client.Do(request)
|
||||
}()
|
||||
}
|
||||
}
|
||||
group.Wait()
|
||||
|
||||
mutex.Lock()
|
||||
defer mutex.Unlock()
|
||||
if !seen["127.0.0.1:3067"] {
|
||||
t.Error("the proxy dial was not recorded as proxied under concurrency")
|
||||
}
|
||||
if seen["direct.example:80"] {
|
||||
t.Error("a direct dial picked up another request's proxy exemption")
|
||||
}
|
||||
if len(seen) != 2 {
|
||||
t.Fatalf("SETUP INVALID: dialed %v, want exactly the proxy and the direct origin", seen)
|
||||
}
|
||||
}
|
||||
@@ -13,6 +13,7 @@ import (
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/Gitlawb/zero/internal/proxydial"
|
||||
"github.com/Gitlawb/zero/internal/redaction"
|
||||
zeroSandbox "github.com/Gitlawb/zero/internal/sandbox"
|
||||
)
|
||||
@@ -308,19 +309,48 @@ func webFetchSafeTransport(roundTripper http.RoundTripper, resolver webFetchReso
|
||||
transport = http.DefaultTransport.(*http.Transport).Clone()
|
||||
case *http.Transport:
|
||||
transport = typed.Clone()
|
||||
case *proxydial.Transport:
|
||||
// Already guarded: rebuild from the inner transport so a second pass
|
||||
// installs one dialer rather than nesting wrappers.
|
||||
transport = typed.Transport.Clone()
|
||||
default:
|
||||
return roundTripper
|
||||
}
|
||||
|
||||
dialer := &net.Dialer{Timeout: webFetchTimeout, KeepAlive: 30 * time.Second}
|
||||
transport.Proxy = nil
|
||||
// THE PROXY IS HONOURED, AND THAT IS A DELIBERATE REVERSAL. The original
|
||||
// guarded fetch set Proxy to nil as part of tightening the SSRF guardrails,
|
||||
// and the reasoning was sound as far as it went: a proxied request is dialed
|
||||
// to the proxy, so the dial-time address pin below never sees the target and
|
||||
// the DNS-rebind protection it provides does not apply. What that traded away
|
||||
// was every fetch on a machine that reaches the network only through a
|
||||
// configured proxy, which fails outright rather than degrades (#569).
|
||||
//
|
||||
// The target is still validated before the request is made, lexically and
|
||||
// through the resolver, with the same block list the dialer applies. What a
|
||||
// forward proxy removes is only the re-check at dial time, and that is
|
||||
// inherent to using any forward proxy; a user who sets HTTPS_PROXY has chosen
|
||||
// it.
|
||||
//
|
||||
// The returned round tripper is wrapped so the dialer learns which route
|
||||
// the transport chose for each request: the one dial it lets past the guard
|
||||
// is the one this transport opens to the proxy it picked for that request.
|
||||
transport.Proxy = http.ProxyFromEnvironment
|
||||
transport.DialContext = webFetchSafeDialContext(resolver, dialer)
|
||||
transport.DialTLSContext = nil
|
||||
return transport
|
||||
return proxydial.Wrap(transport)
|
||||
}
|
||||
|
||||
func webFetchSafeDialContext(resolver webFetchResolver, dialer webFetchDialer) func(context.Context, string, string) (net.Conn, error) {
|
||||
return func(ctx context.Context, network string, address string) (net.Conn, error) {
|
||||
// The proxy is not the target. See webFetchSafeTransport for why the
|
||||
// proxy is honoured at all; here, the dial the transport opens to the
|
||||
// proxy it selected for this request skips the pin, and nothing else
|
||||
// does. A direct dial is pinned as it always was, including one whose
|
||||
// address happens to equal a configured proxy endpoint.
|
||||
if proxydial.IsProxyDial(ctx) {
|
||||
return dialer.DialContext(ctx, network, address)
|
||||
}
|
||||
pinnedAddress, err := webFetchSafeDialAddress(ctx, resolver, address)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
|
||||
@@ -0,0 +1,163 @@
|
||||
package tools
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/netip"
|
||||
"net/url"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Gitlawb/zero/internal/proxydial"
|
||||
)
|
||||
|
||||
// webFetchTransportWithProxy builds the guarded web_fetch round tripper and
|
||||
// swaps in a proxy function, which is how the transport learns its route. The
|
||||
// environment is left alone deliberately: http.ProxyFromEnvironment reads it
|
||||
// once per process, so an environment-driven test would depend on test order.
|
||||
func webFetchTransportWithProxy(t *testing.T, resolver webFetchResolver, proxyFor func(*http.Request) (*url.URL, error)) *proxydial.Transport {
|
||||
t.Helper()
|
||||
roundTripper := webFetchSafeTransport(nil, resolver)
|
||||
wrapped, ok := roundTripper.(*proxydial.Transport)
|
||||
if !ok {
|
||||
t.Fatalf("transport = %T, want the route-recording wrapper; without it no dial can be recognised as the proxy dial", roundTripper)
|
||||
}
|
||||
wrapped.Proxy = proxyFor
|
||||
return wrapped
|
||||
}
|
||||
|
||||
// THE PROXY IS NOT THE TARGET. With HTTPS_PROXY pointing at a local forward
|
||||
// proxy, the transport dials the proxy and tunnels the request through it. The
|
||||
// dialer used to refuse that dial as a loopback address, so every fetch on a
|
||||
// machine behind such a proxy failed with "loopback hosts are blocked" for a
|
||||
// public target that had already been validated (#569).
|
||||
func TestWebFetchTunnelsThroughTheConfiguredProxy(t *testing.T) {
|
||||
listener, err := net.Listen("tcp", "127.0.0.1:0")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer listener.Close()
|
||||
arrived := make(chan string, 1)
|
||||
go func() {
|
||||
conn, err := listener.Accept()
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
defer conn.Close()
|
||||
buffer := make([]byte, 256)
|
||||
n, _ := conn.Read(buffer)
|
||||
arrived <- string(buffer[:n])
|
||||
}()
|
||||
proxyURL, err := url.Parse("http://" + listener.Addr().String())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
// A proxied fetch never resolves the target on this side; the CONNECT
|
||||
// carries the hostname and the proxy resolves it.
|
||||
transport := webFetchTransportWithProxy(t,
|
||||
webFetchResolverFunc(func(context.Context, string, string) ([]netip.Addr, error) {
|
||||
return nil, errors.New("resolver must not be called for a proxied fetch")
|
||||
}),
|
||||
func(*http.Request) (*url.URL, error) { return proxyURL, nil })
|
||||
client := &http.Client{Transport: transport, Timeout: 3 * time.Second}
|
||||
|
||||
_, err = client.Get("https://public.example/resource")
|
||||
select {
|
||||
case first := <-arrived:
|
||||
if line := strings.SplitN(first, "\r\n", 2)[0]; line != "CONNECT public.example:443 HTTP/1.1" {
|
||||
t.Fatalf("proxy received %q, want a CONNECT for the target", line)
|
||||
}
|
||||
case <-time.After(5 * time.Second):
|
||||
t.Fatalf("the proxy never received a connection; request error: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// AND A DIRECT FETCH STAYS PINNED, EVEN WHEN ITS ORIGIN IS THE PROXY'S OWN
|
||||
// ADDRESS.
|
||||
//
|
||||
// This is the NO_PROXY shape: the proxy function sends one host direct. An
|
||||
// exemption derived from asking that function about a request of our own
|
||||
// invention would answer "proxied" for the excluded host too, and this dial
|
||||
// would skip the pin and reach whatever the name resolves to at connect time.
|
||||
func TestWebFetchDirectRequestToTheProxyEndpointIsStillPinned(t *testing.T) {
|
||||
proxyURL, err := url.Parse("http://rebind.example:80")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
lookups := 0
|
||||
transport := webFetchTransportWithProxy(t,
|
||||
webFetchResolverFunc(func(_ context.Context, _ string, host string) ([]netip.Addr, error) {
|
||||
lookups++
|
||||
return []netip.Addr{netip.MustParseAddr("127.0.0.1")}, nil
|
||||
}),
|
||||
func(request *http.Request) (*url.URL, error) {
|
||||
if request.URL.Hostname() == "rebind.example" {
|
||||
return nil, nil
|
||||
}
|
||||
return proxyURL, nil
|
||||
})
|
||||
client := &http.Client{Transport: transport, Timeout: 3 * time.Second}
|
||||
|
||||
_, err = client.Get("http://rebind.example/data")
|
||||
|
||||
if lookups == 0 {
|
||||
t.Fatal("the direct dial skipped the address pin because its address matched the configured proxy endpoint")
|
||||
}
|
||||
if err == nil || !strings.Contains(err.Error(), "loopback hosts are blocked") {
|
||||
t.Fatalf("err = %v, want the private answer to be refused", err)
|
||||
}
|
||||
}
|
||||
|
||||
// A proxy configured with a Unicode hostname is dialed by its IDNA-encoded
|
||||
// name, and it is still the proxy.
|
||||
func TestWebFetchUnicodeProxyHostnameIsStillTheProxy(t *testing.T) {
|
||||
dialed := ""
|
||||
transport := webFetchTransportWithProxy(t,
|
||||
webFetchResolverFunc(func(context.Context, string, string) ([]netip.Addr, error) {
|
||||
return nil, errors.New("resolver must not be called for a proxied fetch")
|
||||
}),
|
||||
func(*http.Request) (*url.URL, error) { return url.Parse("http://büro.example:3067") })
|
||||
guarded := transport.DialContext
|
||||
transport.DialContext = func(ctx context.Context, network, address string) (net.Conn, error) {
|
||||
dialed = address
|
||||
return guarded(ctx, network, address)
|
||||
}
|
||||
client := &http.Client{Transport: transport, Timeout: 3 * time.Second}
|
||||
|
||||
_, err := client.Get("https://public.example/resource")
|
||||
|
||||
if dialed != "xn--bro-hoa.example:3067" {
|
||||
t.Fatalf("SETUP INVALID: dialed %q, want the IDNA-encoded proxy endpoint", dialed)
|
||||
}
|
||||
if err != nil && strings.Contains(err.Error(), "hosts are blocked") {
|
||||
t.Fatalf("the configured proxy was refused by the guard: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// Without a proxy the guard is exactly what it was: loopback is refused and
|
||||
// the dialer never runs.
|
||||
func TestWebFetchSafeDialStillRefusesLoopbackWithoutAProxy(t *testing.T) {
|
||||
dialCalled := false
|
||||
dial := webFetchSafeDialContext(
|
||||
webFetchResolverFunc(func(_ context.Context, network string, host string) ([]netip.Addr, error) {
|
||||
t.Fatalf("resolver consulted for a loopback literal: network=%q host=%q", network, host)
|
||||
return nil, nil
|
||||
}),
|
||||
webFetchDialFunc(func(context.Context, string, string) (net.Conn, error) {
|
||||
dialCalled = true
|
||||
return nil, errors.New("dial should not run")
|
||||
}),
|
||||
)
|
||||
|
||||
_, err := dial(context.Background(), "tcp", "127.0.0.1:3067")
|
||||
if err == nil || !strings.Contains(err.Error(), "loopback hosts are blocked") {
|
||||
t.Fatalf("expected loopback rejection, got %v", err)
|
||||
}
|
||||
if dialCalled {
|
||||
t.Fatal("loopback was dialed with no proxy configured")
|
||||
}
|
||||
}
|
||||
@@ -11,6 +11,7 @@ import (
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Gitlawb/zero/internal/proxydial"
|
||||
zeroSandbox "github.com/Gitlawb/zero/internal/sandbox"
|
||||
)
|
||||
|
||||
@@ -325,15 +326,26 @@ func TestWebFetchToolConfiguresDialTimeSafetyForDefaultTransport(t *testing.T) {
|
||||
}
|
||||
|
||||
client := tool.clientForRun()
|
||||
transport, ok := client.Transport.(*http.Transport)
|
||||
// The round tripper is the route-recording wrapper, not the bare transport:
|
||||
// it is what tells the dialer which dial is the one to the proxy the
|
||||
// transport picked for a given request. The guarded transport it wraps
|
||||
// stays reachable, and is what the rest of this test inspects.
|
||||
transport, ok := client.Transport.(*proxydial.Transport)
|
||||
if !ok {
|
||||
t.Fatalf("client transport = %T, want *http.Transport", client.Transport)
|
||||
t.Fatalf("client transport = %T, want *proxydial.Transport", client.Transport)
|
||||
}
|
||||
if transport.DialContext == nil {
|
||||
t.Fatal("expected web_fetch transport to install a safe DialContext")
|
||||
}
|
||||
if transport.Proxy != nil {
|
||||
t.Fatal("expected web_fetch transport to disable proxy resolution")
|
||||
// REVERSED ON PURPOSE. The guarded fetch used to pin Proxy to nil as part of
|
||||
// the SSRF hardening, and this assertion held it there. #569 is what that
|
||||
// cost: on a machine that reaches the network only through a configured
|
||||
// forward proxy, every fetch failed outright. The proxy is honoured now, and
|
||||
// the wrapper hands the dialer the route chosen for each request so the one
|
||||
// dial it lets past the guard is the one to that proxy. The target is still
|
||||
// validated before the request, which the tests below pin.
|
||||
if transport.Proxy == nil {
|
||||
t.Fatal("expected web_fetch transport to honour the environment proxy; a machine behind a forward proxy cannot fetch at all without it (#569)")
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -1,303 +0,0 @@
|
||||
package tui
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"reflect"
|
||||
"testing"
|
||||
|
||||
tea "charm.land/bubbletea/v2"
|
||||
|
||||
"github.com/Gitlawb/zero/internal/config"
|
||||
"github.com/Gitlawb/zero/internal/modelregistry"
|
||||
"github.com/Gitlawb/zero/internal/sessions"
|
||||
"github.com/Gitlawb/zero/internal/tools"
|
||||
"github.com/Gitlawb/zero/internal/usage"
|
||||
"github.com/Gitlawb/zero/internal/zeroruntime"
|
||||
)
|
||||
|
||||
// USAGE FOLLOWS THE SWITCH. exec reassigns its currentModel when the escalation
|
||||
// switcher fires; the TUI captured usageModelID once per run and handed the
|
||||
// switcher no callback, so every usage event after an escalation was billed to
|
||||
// the model the run started on. Both attribution paths are pinned here: the
|
||||
// live per-event message, and the per-event record the final response carries
|
||||
// for the batch fallback, which must not swing the other way and bill the
|
||||
// events before the switch to the escalated model.
|
||||
func TestEscalatedRunAttributesUsageToTheModelInForce(t *testing.T) {
|
||||
store := testSessionStore(t)
|
||||
session, err := store.Create(sessions.CreateInput{SessionID: "escalation_usage"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
// A catalog model with an upgrade target; the target is whatever the
|
||||
// catalog says, read back from the provider factory rather than assumed.
|
||||
const startModel = "claude-haiku-4.5"
|
||||
starting := &scriptedProvider{scripts: [][]zeroruntime.StreamEvent{{
|
||||
{Type: zeroruntime.StreamEventUsage, Usage: zeroruntime.Usage{InputTokens: 10, OutputTokens: 1}},
|
||||
{Type: zeroruntime.StreamEventToolCallStart, ToolCallID: "escalate", ToolName: "escalate_model"},
|
||||
{Type: zeroruntime.StreamEventToolCallDelta, ToolCallID: "escalate", ArgumentsFragment: `{"reason":"harder than it looked"}`},
|
||||
{Type: zeroruntime.StreamEventToolCallEnd, ToolCallID: "escalate"},
|
||||
{Type: zeroruntime.StreamEventDone},
|
||||
}}}
|
||||
escalated := &scriptedProvider{scripts: [][]zeroruntime.StreamEvent{{
|
||||
{Type: zeroruntime.StreamEventText, Content: "Done on the stronger model."},
|
||||
{Type: zeroruntime.StreamEventUsage, Usage: zeroruntime.Usage{InputTokens: 20, OutputTokens: 2}},
|
||||
{Type: zeroruntime.StreamEventDone},
|
||||
}}}
|
||||
registry := tools.NewRegistry()
|
||||
registry.Register(tools.NewEscalateModelTool())
|
||||
|
||||
switchedTo := ""
|
||||
m := newModel(context.Background(), Options{
|
||||
ProviderName: "anthropic",
|
||||
ModelName: startModel,
|
||||
Provider: starting,
|
||||
Registry: registry,
|
||||
SessionStore: store,
|
||||
AllowEscalation: true,
|
||||
ProviderProfile: config.ProviderProfile{Model: startModel},
|
||||
NewProvider: func(profile config.ProviderProfile) (zeroruntime.Provider, error) {
|
||||
switchedTo = profile.Model
|
||||
return escalated, nil
|
||||
},
|
||||
})
|
||||
m.activeSession = session
|
||||
m.agentOptions.Model = startModel
|
||||
var live []string
|
||||
m.runtimeMessageSink = func(msg tea.Msg) {
|
||||
if usage, ok := msg.(agentUsageMsg); ok {
|
||||
live = append(live, usage.modelID)
|
||||
}
|
||||
}
|
||||
|
||||
msg := execCmd(m.runAgentWithOptions(1, context.Background(), "do the hard thing", nil, tuiAgentRunOptions{}))
|
||||
response, ok := msg.(agentResponseMsg)
|
||||
if !ok {
|
||||
t.Fatalf("run returned %T, want agentResponseMsg", msg)
|
||||
}
|
||||
if response.err != nil {
|
||||
t.Fatalf("run failed: %v", response.err)
|
||||
}
|
||||
if switchedTo == "" || switchedTo == startModel {
|
||||
t.Fatalf("SETUP INVALID: the escalation never switched providers (switched to %q), so nothing here exercises attribution", switchedTo)
|
||||
}
|
||||
if len(escalated.requests) == 0 {
|
||||
t.Fatal("SETUP INVALID: the escalated provider was never asked for a completion")
|
||||
}
|
||||
if len(response.usageEvents) != 2 {
|
||||
t.Fatalf("usage events = %d, want one before the switch and one after", len(response.usageEvents))
|
||||
}
|
||||
|
||||
want := []string{startModel, switchedTo}
|
||||
if !reflect.DeepEqual(live, want) {
|
||||
t.Fatalf("live usage attribution = %v, want %v", live, want)
|
||||
}
|
||||
if !reflect.DeepEqual(response.usageModelIDs, want) {
|
||||
t.Fatalf("per-event usage attribution on the response = %v, want %v", response.usageModelIDs, want)
|
||||
}
|
||||
for index := range response.usageEvents {
|
||||
if got := response.usageModelIDAt(index); got != want[index] {
|
||||
t.Fatalf("usageModelIDAt(%d) = %q, want %q", index, got, want[index])
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// A response built without the per-event record, which is every constructor
|
||||
// that predates escalation, still attributes through the run-level model.
|
||||
func TestUsageModelIDAtFallsBackToTheRunModel(t *testing.T) {
|
||||
msg := agentResponseMsg{usageModelID: "gpt-4.1", usageEvents: make([]zeroruntime.Usage, 2)}
|
||||
for index := range msg.usageEvents {
|
||||
if got := msg.usageModelIDAt(index); got != "gpt-4.1" {
|
||||
t.Fatalf("usageModelIDAt(%d) = %q, want the run-level model", index, got)
|
||||
}
|
||||
}
|
||||
msg.usageModelIDs = []string{"gpt-4.1-mini"}
|
||||
if got := msg.usageModelIDAt(0); got != "gpt-4.1-mini" {
|
||||
t.Fatalf("usageModelIDAt(0) = %q, want the per-event model", got)
|
||||
}
|
||||
if got := msg.usageModelIDAt(1); got != "gpt-4.1" {
|
||||
t.Fatalf("usageModelIDAt(1) = %q, want the run-level fallback past the record", got)
|
||||
}
|
||||
}
|
||||
|
||||
// AND THE PRICE THAT COMES BACK OUT OF THE SESSION LOG. The live record and the
|
||||
// persisted event are two representations of one usage event, and only the
|
||||
// first carried the model. `zero usage report` rebuilds cost from the persisted
|
||||
// payload, falling back to the session's own model when the event does not name
|
||||
// one, so an escalated run was priced end to end at the model it started on.
|
||||
// This drives the real persistence path and reconstructs the report from it.
|
||||
func TestEscalatedRunPricesPersistedUsageAtEachModel(t *testing.T) {
|
||||
registry, err := modelregistry.DefaultRegistry()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
store := testSessionStore(t)
|
||||
session, err := store.Create(sessions.CreateInput{SessionID: "escalation_cost"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
const startModel = "claude-haiku-4.5"
|
||||
before := zeroruntime.Usage{InputTokens: 10, OutputTokens: 1}
|
||||
after := zeroruntime.Usage{InputTokens: 20, OutputTokens: 2}
|
||||
starting := &scriptedProvider{scripts: [][]zeroruntime.StreamEvent{{
|
||||
{Type: zeroruntime.StreamEventUsage, Usage: before},
|
||||
{Type: zeroruntime.StreamEventToolCallStart, ToolCallID: "escalate", ToolName: "escalate_model"},
|
||||
{Type: zeroruntime.StreamEventToolCallDelta, ToolCallID: "escalate", ArgumentsFragment: `{"reason":"harder than it looked"}`},
|
||||
{Type: zeroruntime.StreamEventToolCallEnd, ToolCallID: "escalate"},
|
||||
{Type: zeroruntime.StreamEventDone},
|
||||
}}}
|
||||
escalated := &scriptedProvider{scripts: [][]zeroruntime.StreamEvent{{
|
||||
{Type: zeroruntime.StreamEventText, Content: "Done on the stronger model."},
|
||||
{Type: zeroruntime.StreamEventUsage, Usage: after},
|
||||
{Type: zeroruntime.StreamEventDone},
|
||||
}}}
|
||||
toolRegistry := tools.NewRegistry()
|
||||
toolRegistry.Register(tools.NewEscalateModelTool())
|
||||
|
||||
switchedTo := ""
|
||||
m := newModel(context.Background(), Options{
|
||||
ProviderName: "anthropic",
|
||||
ModelName: startModel,
|
||||
Provider: starting,
|
||||
Registry: toolRegistry,
|
||||
SessionStore: store,
|
||||
AllowEscalation: true,
|
||||
ProviderProfile: config.ProviderProfile{Model: startModel},
|
||||
NewProvider: func(profile config.ProviderProfile) (zeroruntime.Provider, error) {
|
||||
switchedTo = profile.Model
|
||||
return escalated, nil
|
||||
},
|
||||
})
|
||||
m.activeSession = session
|
||||
m.agentOptions.Model = startModel
|
||||
|
||||
msg := execCmd(m.runAgentWithOptions(1, context.Background(), "do the hard thing", nil, tuiAgentRunOptions{}))
|
||||
response, ok := msg.(agentResponseMsg)
|
||||
if !ok || response.err != nil {
|
||||
t.Fatalf("run returned %T err=%v", msg, response.err)
|
||||
}
|
||||
if switchedTo == "" || switchedTo == startModel {
|
||||
t.Fatalf("SETUP INVALID: no escalation happened (switched to %q)", switchedTo)
|
||||
}
|
||||
|
||||
// Persist through the same path a real run uses, then read the events back.
|
||||
m, rows := m.appendSessionEvents(response.sessionEvents)
|
||||
for _, row := range rows {
|
||||
if row.kind == rowError {
|
||||
t.Fatalf("session record error: %s", row.text)
|
||||
}
|
||||
}
|
||||
events, err := store.ReadEvents(session.SessionID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
persisted := []string{}
|
||||
for _, event := range events {
|
||||
if event.Type != sessions.EventUsage {
|
||||
continue
|
||||
}
|
||||
var payload struct {
|
||||
Model string `json:"model"`
|
||||
}
|
||||
if err := json.Unmarshal(event.Payload, &payload); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
persisted = append(persisted, payload.Model)
|
||||
}
|
||||
if want := []string{startModel, switchedTo}; !reflect.DeepEqual(persisted, want) {
|
||||
t.Fatalf("persisted usage models = %v, want %v", persisted, want)
|
||||
}
|
||||
|
||||
// The report prices from those payloads. Its session metadata names only the
|
||||
// starting model, which is the fallback that used to price both events.
|
||||
metadata := []sessions.Metadata{{SessionID: session.SessionID, ModelID: startModel}}
|
||||
report, err := usage.BuildReport(events, metadata, ®istry, 0)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
// Expected cost comes from the catalog, one completion at each model.
|
||||
expected := 0.0
|
||||
for _, priced := range []struct {
|
||||
modelID string
|
||||
usage zeroruntime.Usage
|
||||
}{{startModel, before}, {switchedTo, after}} {
|
||||
model, err := registry.Require(priced.modelID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
cost, err := modelregistry.CalculateCost(model, priced.usage)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
expected += cost.TotalCost
|
||||
}
|
||||
// And the wrong answer, for contrast: both events at the starting model.
|
||||
startingOnly := 0.0
|
||||
startingEntry, err := registry.Require(startModel)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for _, event := range []zeroruntime.Usage{before, after} {
|
||||
cost, err := modelregistry.CalculateCost(startingEntry, event)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
startingOnly += cost.TotalCost
|
||||
}
|
||||
if expected == startingOnly {
|
||||
t.Fatal("SETUP INVALID: the two models price this usage identically, so the assertion below proves nothing")
|
||||
}
|
||||
if difference := report.Total.TotalCost - expected; difference > 1e-12 || difference < -1e-12 {
|
||||
t.Fatalf("report cost = %v, want %v (pricing both events at %s gives %v)", report.Total.TotalCost, expected, startModel, startingOnly)
|
||||
}
|
||||
}
|
||||
|
||||
// A run without escalation persists no model on its usage events: the model
|
||||
// cannot change mid-run, so the session-wide identity is the whole story and
|
||||
// the payload stays compact, exactly as exec does it.
|
||||
func TestUnescalatedRunLeavesTheUsagePayloadAlone(t *testing.T) {
|
||||
store := testSessionStore(t)
|
||||
session, err := store.Create(sessions.CreateInput{SessionID: "no_escalation_usage"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
provider := &scriptedProvider{scripts: [][]zeroruntime.StreamEvent{{
|
||||
{Type: zeroruntime.StreamEventText, Content: "done"},
|
||||
{Type: zeroruntime.StreamEventUsage, Usage: zeroruntime.Usage{InputTokens: 10, OutputTokens: 1}},
|
||||
{Type: zeroruntime.StreamEventDone},
|
||||
}}}
|
||||
m := newModel(context.Background(), Options{
|
||||
ProviderName: "anthropic",
|
||||
ModelName: "claude-haiku-4.5",
|
||||
Provider: provider,
|
||||
Registry: tools.NewRegistry(),
|
||||
SessionStore: store,
|
||||
})
|
||||
m.activeSession = session
|
||||
|
||||
msg := execCmd(m.runAgentWithOptions(1, context.Background(), "hello", nil, tuiAgentRunOptions{}))
|
||||
response, ok := msg.(agentResponseMsg)
|
||||
if !ok || response.err != nil {
|
||||
t.Fatalf("run returned %T err=%v", msg, response.err)
|
||||
}
|
||||
usageEvents := 0
|
||||
for _, event := range response.sessionEvents {
|
||||
if event.Type != sessions.EventUsage {
|
||||
continue
|
||||
}
|
||||
usageEvents++
|
||||
payload, ok := event.Payload.(map[string]any)
|
||||
if !ok {
|
||||
t.Fatalf("usage payload = %T, want map", event.Payload)
|
||||
}
|
||||
if _, present := payload["model"]; present {
|
||||
t.Errorf("a non-escalation run recorded a model on its usage payload: %v", payload)
|
||||
}
|
||||
}
|
||||
if usageEvents == 0 {
|
||||
t.Fatal("SETUP INVALID: the run recorded no usage event")
|
||||
}
|
||||
}
|
||||
+25
-80
@@ -70,24 +70,21 @@ 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 + /<skill-name>
|
||||
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
|
||||
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 + /<skill-name>
|
||||
userConfigPath string
|
||||
doctorUserConfigPath string
|
||||
projectConfigPath string
|
||||
gitBranch string
|
||||
providerName string
|
||||
modelName string
|
||||
modelCatalog modelregistry.Registry
|
||||
providerProfile config.ProviderProfile
|
||||
savedProviders []config.ProviderProfile
|
||||
provider zeroruntime.Provider
|
||||
newProvider func(config.ProviderProfile) (zeroruntime.Provider, error)
|
||||
newTurnSessionProvider func(config.ProviderProfile, zeroruntime.Provider) zeroruntime.TurnSessionProvider
|
||||
probeProviderHealth func(context.Context, providerhealth.Options) providerhealth.Result
|
||||
@@ -661,16 +658,10 @@ type agentUsageMsg struct {
|
||||
}
|
||||
|
||||
type agentResponseMsg struct {
|
||||
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.
|
||||
runID int
|
||||
rows []transcriptRow
|
||||
usageEvents []zeroruntime.Usage
|
||||
usageModelID string
|
||||
usageModelIDs []string
|
||||
sessionEvents []pendingSessionEvent
|
||||
specReview *pendingSpecReviewPrompt
|
||||
err error
|
||||
@@ -683,16 +674,6 @@ 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
|
||||
@@ -1018,7 +999,6 @@ 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,
|
||||
@@ -2566,7 +2546,7 @@ func (m model) updateModel(msg tea.Msg) (tea.Model, tea.Cmd) {
|
||||
continue
|
||||
}
|
||||
var usageRows []transcriptRow
|
||||
m, usageRows = m.recordUsageEvent(msg.usageModelIDAt(index), event)
|
||||
m, usageRows = m.recordUsageEvent(msg.usageModelID, event)
|
||||
for _, row := range usageRows {
|
||||
m.transcript = appendTranscriptRow(m.transcript, row)
|
||||
}
|
||||
@@ -2653,7 +2633,7 @@ func (m model) updateModel(msg tea.Msg) (tea.Model, tea.Cmd) {
|
||||
continue
|
||||
}
|
||||
var usageRows []transcriptRow
|
||||
m, usageRows = m.recordUsageEvent(msg.usageModelIDAt(index), event)
|
||||
m, usageRows = m.recordUsageEvent(msg.usageModelID, event)
|
||||
for _, row := range usageRows {
|
||||
m.transcript = appendTranscriptRow(m.transcript, row)
|
||||
}
|
||||
@@ -5447,9 +5427,6 @@ 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)
|
||||
@@ -5519,27 +5496,6 @@ func (m model) runAgentWithOptions(runID int, runCtx context.Context, prompt str
|
||||
options.ContextWindowFor = func(modelID string) int {
|
||||
return modelregistry.AgentContextWindow(m.modelContextWindow(modelID))
|
||||
}
|
||||
// And make that switch reachable, when the operator asked for it. The
|
||||
// consequences of an escalation were already handled here (the window
|
||||
// above, and the summarizer resolved against the active profile) while
|
||||
// nothing on this surface could cause one: escalate_model was registered
|
||||
// only by exec.
|
||||
//
|
||||
// BUILT FROM THE ACTIVE PROFILE, NOT THE STARTUP ONE. A TUI session can
|
||||
// change models with /model, so escalating from the profile captured at
|
||||
// launch would switch from whatever the session began with rather than
|
||||
// from what is in force now, and would carry that stale profile's base URL
|
||||
// and credential with it. m.providerProfile tracks the switches, which is
|
||||
// why this is built per turn rather than once in the caller.
|
||||
if m.allowEscalation {
|
||||
options.ModelSwitcher, options.ModelSessionSwitcher = providers.EscalationSwitchers(
|
||||
m.providerProfile, m.provider, m.newProvider,
|
||||
// Usage attribution follows the switch, as exec reassigns its
|
||||
// currentModel: every usage event after a real escalation is billed
|
||||
// to the escalated model, not the one the run started on.
|
||||
func(modelID string) { usageModelID = modelID },
|
||||
)
|
||||
}
|
||||
|
||||
// Post-edit self-correction is on by default in the TUI but kept FAST: it
|
||||
// runs LSP diagnostics over the changed files only — cheap, change-scoped,
|
||||
@@ -5905,20 +5861,9 @@ func (m model) runAgentWithOptions(runID int, runCtx context.Context, prompt str
|
||||
onUsage := options.OnUsage
|
||||
options.OnUsage = func(event zeroruntime.Usage) {
|
||||
usageEvents = append(usageEvents, event)
|
||||
usageModelIDs = append(usageModelIDs, usageModelID)
|
||||
payload := usage.EventUsagePayload(event)
|
||||
// AND ON THE PERSISTED EVENT TOO, not only the in-memory record: the
|
||||
// report reconstructs cost from the payload, falling back to the
|
||||
// session-wide model, so an escalated run would be priced entirely at
|
||||
// the model it started on. Written only under escalation, which is the
|
||||
// only way the model in force can change mid-run, matching what exec
|
||||
// records under the same flag.
|
||||
if m.allowEscalation {
|
||||
payload["model"] = usageModelID
|
||||
}
|
||||
sessionEvents = append(sessionEvents, pendingSessionEvent{
|
||||
Type: sessions.EventUsage,
|
||||
Payload: payload,
|
||||
Payload: usage.EventUsagePayload(event),
|
||||
})
|
||||
m.sendAgentUsage(runID, usageModelID, event)
|
||||
if onUsage != nil {
|
||||
@@ -5933,7 +5878,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, usageModelIDs: usageModelIDs, sessionEvents: sessionEvents, err: err, goalAware: goalAwareRun, turnTools: toolCalls, turnElapsed: m.activeTurnElapsed(started)}
|
||||
return agentResponseMsg{runID: runID, rows: rows, usageEvents: usageEvents, usageModelID: usageModelID, 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 == "" {
|
||||
@@ -5943,10 +5888,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, usageModelIDs: usageModelIDs, sessionEvents: sessionEvents, err: err, goalAware: goalAwareRun, turnTools: toolCalls, turnElapsed: m.activeTurnElapsed(started)}
|
||||
return agentResponseMsg{runID: runID, rows: rows, usageEvents: usageEvents, usageModelID: usageModelID, 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, usageModelIDs: usageModelIDs, sessionEvents: sessionEvents, specReview: specReview, goalAware: goalAwareRun, turnTools: toolCalls, turnElapsed: m.activeTurnElapsed(started)}
|
||||
return agentResponseMsg{runID: runID, rows: rows, usageEvents: usageEvents, usageModelID: usageModelID, sessionEvents: sessionEvents, specReview: specReview, goalAware: goalAwareRun, turnTools: toolCalls, turnElapsed: m.activeTurnElapsed(started)}
|
||||
}
|
||||
flushReasoning(m.now())
|
||||
elapsed := m.activeTurnElapsed(started)
|
||||
@@ -5967,7 +5912,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, usageModelIDs: usageModelIDs, sessionEvents: sessionEvents, goalAware: goalAwareRun, turnTools: toolCalls, turnElapsed: elapsed, ttft: firstTokenElapsed}
|
||||
return agentResponseMsg{runID: runID, rows: rows, usageEvents: usageEvents, usageModelID: usageModelID, sessionEvents: sessionEvents, goalAware: goalAwareRun, turnTools: toolCalls, turnElapsed: elapsed, ttft: firstTokenElapsed}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -70,13 +70,7 @@ 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
|
||||
// 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
|
||||
LoadSkills func() []skills.Skill
|
||||
PermissionMode agent.PermissionMode
|
||||
ReasoningEffort modelregistry.ReasoningEffort
|
||||
ResponseStyle string
|
||||
|
||||
Reference in new issue
Block a user