diff --git a/internal/acp/agent.go b/internal/acp/agent.go index c355f471a..16c7110cd 100644 --- a/internal/acp/agent.go +++ b/internal/acp/agent.go @@ -11,6 +11,9 @@ import ( "github.com/Gitlawb/zero/internal/agent" "github.com/Gitlawb/zero/internal/config" + "github.com/Gitlawb/zero/internal/providercatalog" + "github.com/Gitlawb/zero/internal/providermodelcatalog" + "github.com/Gitlawb/zero/internal/providermodeldiscovery" "github.com/Gitlawb/zero/internal/sandbox" "github.com/Gitlawb/zero/internal/sessions" "github.com/Gitlawb/zero/internal/tools" @@ -23,9 +26,10 @@ import ( // means the editor only hosts the thread — ZERO owns BYOK and telemetry-free // operation. type Deps struct { - ResolveConfig func(workspaceRoot string, overrides config.Overrides) (config.ResolvedConfig, error) - NewProvider func(profile config.ProviderProfile) (zeroruntime.Provider, error) - RunAgent func(ctx context.Context, prompt string, provider zeroruntime.Provider, opts agent.Options) (agent.Result, error) + ResolveConfig func(workspaceRoot string, overrides config.Overrides) (config.ResolvedConfig, error) + DiscoverModels func(context.Context, config.ProviderProfile) ([]providermodeldiscovery.Model, error) + NewProvider func(profile config.ProviderProfile) (zeroruntime.Provider, error) + RunAgent func(ctx context.Context, prompt string, provider zeroruntime.Provider, opts agent.Options) (agent.Result, error) // BuildWorkspace builds the SCOPED tool registry and the sandbox engine for a // validated workspace root, so ACP shell tools (bash/exec_command) are confined // exactly like the exec surface — never run unconfined on the host. @@ -60,13 +64,16 @@ type acpSession struct { // turnMu serializes prompt turns for one session: concurrent session/prompt // calls run one at a time so they can't interleave history or clobber the // single cancel slot. - turnMu sync.Mutex + turnMu sync.Mutex + modelMu sync.Mutex - mu sync.Mutex - mode agent.PermissionMode - model string // override; "" => config default - cancel context.CancelFunc - history []turnRecord + mu sync.Mutex + mode agent.PermissionMode + model string + models []SessionConfigOptionValue + restrictModels bool + cancel context.CancelFunc + history []turnRecord } // NewAgent builds the ACP server and registers its method handlers on conn. @@ -120,7 +127,7 @@ func (a *Agent) handleInitialize(_ context.Context, params json.RawMessage) (any // ---- session lifecycle ---- -func (a *Agent) handleSessionNew(_ context.Context, params json.RawMessage) (any, error) { +func (a *Agent) handleSessionNew(ctx context.Context, params json.RawMessage) (any, error) { var p NewSessionParams if err := json.Unmarshal(params, &p); err != nil { return nil, RPCError(codeInvalidParams, "invalid session/new params") @@ -129,18 +136,23 @@ func (a *Agent) handleSessionNew(_ context.Context, params json.RawMessage) (any if err != nil { return nil, RPCError(codeInvalidParams, err.Error()) } - meta, err := a.deps.Store.Create(sessions.CreateInput{Title: "ACP session", Cwd: root}) + model, models, restrictModels, err := a.resolveModelChoices(ctx, root) + if err != nil { + return nil, RPCError(codeInternalError, "config: "+err.Error()) + } + meta, err := a.deps.Store.Create(sessions.CreateInput{Title: "ACP session", Cwd: root, ModelID: model}) if err != nil { return nil, RPCError(codeInternalError, "create session: "+err.Error()) } - sess := a.registerSession(meta.SessionID, root, nil) + sess := a.registerSession(meta.SessionID, root, nil, model, models, restrictModels) return NewSessionResult{ - SessionID: sess.id, - Modes: a.modeState(sess), + SessionID: sess.id, + ConfigOptions: a.configOptions(sess), + Modes: a.modeState(sess), }, nil } -func (a *Agent) handleSessionLoad(_ context.Context, params json.RawMessage) (any, error) { +func (a *Agent) handleSessionLoad(ctx context.Context, params json.RawMessage) (any, error) { var p LoadSessionParams if err := json.Unmarshal(params, &p); err != nil { return nil, RPCError(codeInvalidParams, "invalid session/load params") @@ -161,7 +173,17 @@ func (a *Agent) handleSessionLoad(_ context.Context, params json.RawMessage) (an // a half-initialized session (registerSession sets history under the lock and // reuses an already-live session rather than orphaning its in-flight turn). history, historyErr := a.loadHistory(meta.SessionID) - sess := a.registerSession(meta.SessionID, root, history) + model, models, restrictModels, err := a.resolveModelChoices(ctx, root) + if err != nil { + return nil, RPCError(codeInternalError, "config: "+err.Error()) + } + if persistedModel := strings.TrimSpace(meta.ModelID); persistedModel != "" && (!restrictModels || modelChoiceExists(models, persistedModel)) { + model = persistedModel + if !modelChoiceExists(models, persistedModel) { + models = append(models, SessionConfigOptionValue{Value: persistedModel, Name: persistedModel}) + } + } + sess := a.registerSession(meta.SessionID, root, history, model, models, restrictModels) a.warnPersistence( ¬ifier{conn: a.conn, sessionID: sess.id}, "load session history", @@ -169,7 +191,8 @@ func (a *Agent) handleSessionLoad(_ context.Context, params json.RawMessage) (an historyErr, ) return LoadSessionResult{ - Modes: a.modeState(sess), + ConfigOptions: a.configOptions(sess), + Modes: a.modeState(sess), }, nil } @@ -354,11 +377,27 @@ func (a *Agent) handleSetConfigOption(_ context.Context, params json.RawMessage) if sess == nil { return nil, RPCError(codeInvalidParams, "unknown session: "+p.SessionID) } - if p.ConfigID != configIDModel { + switch p.ConfigID { + case configIDModel: + model := strings.TrimSpace(p.Value) + if err := a.updateModel(sess, model, sess.restrictModels); err != nil { + return nil, err + } + case configIDMode: + mode := agent.PermissionMode(p.Value) + switch mode { + case agent.PermissionModeAuto, agent.PermissionModeAsk: + sess.setMode(mode) + (¬ifier{conn: a.conn, sessionID: sess.id}).currentMode(string(mode)) + case agent.PermissionModeUnsafe: + return nil, RPCError(codeInvalidParams, "mode not permitted over ACP: "+p.Value) + default: + return nil, RPCError(codeInvalidParams, "unknown mode: "+p.Value) + } + default: return nil, RPCError(codeInvalidParams, "unknown config option: "+p.ConfigID) } - sess.setModel(p.Value) - return SetSessionConfigOptionResult{}, nil + return SetSessionConfigOptionResult{ConfigOptions: a.configOptions(sess)}, nil } func (a *Agent) handleZeroSetModel(_ context.Context, params json.RawMessage) (any, error) { @@ -370,8 +409,24 @@ func (a *Agent) handleZeroSetModel(_ context.Context, params json.RawMessage) (a if sess == nil { return nil, RPCError(codeInvalidParams, "unknown session: "+p.SessionID) } - sess.setModel(p.Model) - return ZeroSetModelResult{Model: p.Model}, nil + model := strings.TrimSpace(p.Model) + if err := a.updateModel(sess, model, false); err != nil { + return nil, err + } + return ZeroSetModelResult{Model: model}, nil +} + +func (a *Agent) updateModel(sess *acpSession, model string, restrictModels bool) error { + sess.modelMu.Lock() + defer sess.modelMu.Unlock() + if restrictModels && !sess.hasModel(model) { + return RPCError(codeInvalidParams, "unknown model: "+model) + } + if _, err := a.deps.Store.UpdateModel(sess.id, model); err != nil { + return RPCError(codeInternalError, "save model selection: "+err.Error()) + } + sess.setModel(model) + return nil } func (a *Agent) handleCancel(_ context.Context, params json.RawMessage) { @@ -398,6 +453,76 @@ func (a *Agent) modeState(s *acpSession) *SessionModeState { } } +// resolveModelChoices advertises authenticated live models when discovery is +// available. The configured model is always retained as the sole fallback. +func (a *Agent) resolveModelChoices(ctx context.Context, cwd string) (string, []SessionConfigOptionValue, bool, error) { + resolved, err := a.deps.ResolveConfig(cwd, config.Overrides{}) + if err != nil { + return "", nil, false, err + } + selected := strings.TrimSpace(resolved.Provider.Model) + options := make([]SessionConfigOptionValue, 0, 8) + seen := make(map[string]bool) + add := func(id, description string) { + id = strings.TrimSpace(id) + if id == "" || seen[id] { + return + } + seen[id] = true + options = append(options, SessionConfigOptionValue{Value: id, Name: id, Description: description}) + } + add(selected, "") + descriptor, knownProvider := providercatalog.Get(resolved.Provider.CatalogID) + restrictModels := knownProvider && !descriptor.Custom + if a.deps.DiscoverModels != nil { + discovered, discoverErr := a.deps.DiscoverModels(ctx, resolved.Provider) + if ctx.Err() != nil { + return "", nil, false, ctx.Err() + } + if discoverErr == nil && len(discovered) > 0 { + for _, model := range discovered { + if providermodelcatalog.ModelIDAllowedForProvider(resolved.Provider.CatalogID, model.ID) { + add(model.ID, model.Description) + } + } + } + } + return selected, options, restrictModels, nil +} + +func modelChoiceExists(models []SessionConfigOptionValue, model string) bool { + for _, option := range models { + if option.Value == model { + return true + } + } + return false +} + +func (a *Agent) configOptions(s *acpSession) []SessionConfigOption { + model, models, mode := s.configState() + return []SessionConfigOption{{ + ID: configIDModel, + Name: "Model", + Description: "Model used for this session.", + Category: configCategoryModel, + Type: configOptionTypeSelect, + CurrentValue: model, + Options: models, + }, { + ID: configIDMode, + Name: "Mode", + Description: "Permission mode used for this session.", + Category: configCategoryMode, + Type: configOptionTypeSelect, + CurrentValue: string(mode), + Options: []SessionConfigOptionValue{ + {Value: string(agent.PermissionModeAuto), Name: "Auto", Description: "Run safe tools automatically; ask before risky ones."}, + {Value: string(agent.PermissionModeAsk), Name: "Ask", Description: "Ask before every tool that changes state."}, + }, + }} +} + // ---- persistence + continuity ---- func (a *Agent) persistTurn(sess *acpSession, user, assistant string) error { @@ -525,13 +650,13 @@ func promptImages(blocks []ContentBlock) []zeroruntime.ImageBlock { // session is returned unchanged rather than orphaning its turn or resetting its // mode/model. history is set BEFORE publishing so no concurrent prompt can read a // half-initialized session. -func (a *Agent) registerSession(id, cwd string, history []turnRecord) *acpSession { +func (a *Agent) registerSession(id, cwd string, history []turnRecord, model string, models []SessionConfigOptionValue, restrictModels bool) *acpSession { a.mu.Lock() defer a.mu.Unlock() if existing := a.sessions[id]; existing != nil { return existing } - sess := &acpSession{id: id, cwd: cwd, mode: agent.PermissionModeAuto, history: history} + sess := &acpSession{id: id, cwd: cwd, mode: agent.PermissionModeAuto, model: model, models: models, restrictModels: restrictModels, history: history} a.sessions[id] = sess return sess } @@ -570,7 +695,18 @@ func (s *acpSession) currentMode() agent.PermissionMode { } func (s *acpSession) setModel(model string) { + model = strings.TrimSpace(model) s.mu.Lock() + found := false + for _, option := range s.models { + if option.Value == model { + found = true + break + } + } + if !found && model != "" { + s.models = append(s.models, SessionConfigOptionValue{Value: model, Name: model}) + } s.model = model s.mu.Unlock() } @@ -581,6 +717,24 @@ func (s *acpSession) currentModel() string { return s.model } +func (s *acpSession) hasModel(model string) bool { + model = strings.TrimSpace(model) + s.mu.Lock() + defer s.mu.Unlock() + for _, option := range s.models { + if option.Value == model { + return true + } + } + return false +} + +func (s *acpSession) configState() (string, []SessionConfigOptionValue, agent.PermissionMode) { + s.mu.Lock() + defer s.mu.Unlock() + return s.model, append([]SessionConfigOptionValue(nil), s.models...), s.mode +} + func (s *acpSession) appendHistory(rec turnRecord) { s.mu.Lock() s.history = append(s.history, rec) diff --git a/internal/acp/agent_test.go b/internal/acp/agent_test.go index d3d7c3f5f..a9e07eccc 100644 --- a/internal/acp/agent_test.go +++ b/internal/acp/agent_test.go @@ -3,6 +3,7 @@ package acp import ( "context" "encoding/json" + "errors" "fmt" "io" "os" @@ -13,6 +14,7 @@ import ( "github.com/Gitlawb/zero/internal/agent" "github.com/Gitlawb/zero/internal/config" + "github.com/Gitlawb/zero/internal/providermodeldiscovery" "github.com/Gitlawb/zero/internal/sandbox" "github.com/Gitlawb/zero/internal/sessions" "github.com/Gitlawb/zero/internal/tools" @@ -137,6 +139,12 @@ func TestACPEndToEndPrompt(t *testing.T) { if newRes.Modes == nil || newRes.Modes.CurrentModeID != string(agent.PermissionModeAuto) { t.Fatalf("expected auto mode, got %+v", newRes.Modes) } + if len(newRes.ConfigOptions) != 2 || newRes.ConfigOptions[0].ID != configIDModel || newRes.ConfigOptions[0].CurrentValue != "fake-model" { + t.Fatalf("model config option = %+v, want fake-model fallback", newRes.ConfigOptions) + } + if newRes.ConfigOptions[1].ID != configIDMode || newRes.ConfigOptions[1].CurrentValue != string(agent.PermissionModeAuto) { + t.Fatalf("mode config option = %+v", newRes.ConfigOptions[1]) + } // session/prompt var promptRes PromptResult @@ -156,6 +164,201 @@ func TestACPEndToEndPrompt(t *testing.T) { } } +func TestACPModelConfigOptionsCatalogSelectionAndLoad(t *testing.T) { + deps := testDeps(t) + deps.ResolveConfig = func(_ string, o config.Overrides) (config.ResolvedConfig, error) { + model := "gpt-5.5" + if o.Provider.Model != "" { + model = o.Provider.Model + } + return config.ResolvedConfig{Provider: config.ProviderProfile{ + Name: "ChatGPT", CatalogID: "chatgpt", Model: model, + }}, nil + } + deps.DiscoverModels = func(_ context.Context, _ config.ProviderProfile) ([]providermodeldiscovery.Model, error) { + return []providermodeldiscovery.Model{ + {ID: " gpt-5.4-mini ", Description: "Fast"}, + {ID: "gpt-5.5", Description: "duplicate configured model"}, + }, nil + } + h := newHarness(t, deps) + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + var created NewSessionResult + if err := h.client.Call(ctx, MethodSessionNew, NewSessionParams{Cwd: t.TempDir()}, &created); err != nil { + t.Fatalf("session/new: %v", err) + } + option := created.ConfigOptions[0] + if option.CurrentValue != "gpt-5.5" || len(option.Options) < 2 { + t.Fatalf("new model option = %+v", option) + } + for _, choice := range option.Options { + if choice.Name != choice.Value { + t.Fatalf("model choice name = %q, want model id %q", choice.Name, choice.Value) + } + } + var selected SetSessionConfigOptionResult + if err := h.client.Call(ctx, MethodSessionSetConfigOption, SetSessionConfigOptionParams{ + SessionID: created.SessionID, ConfigID: configIDModel, Value: " gpt-5.4-mini ", + }, &selected); err != nil { + t.Fatalf("set_config_option: %v", err) + } + if got := selected.ConfigOptions[0].CurrentValue; got != "gpt-5.4-mini" { + t.Fatalf("selected model = %q", got) + } + if err := h.client.Call(ctx, MethodSessionSetConfigOption, SetSessionConfigOptionParams{ + SessionID: created.SessionID, ConfigID: configIDModel, Value: "not-advertised", + }, &SetSessionConfigOptionResult{}); err == nil { + t.Fatal("unknown standard model selection was accepted") + } + h.stop() + h = newHarness(t, deps) + defer h.stop() + var loaded LoadSessionResult + if err := h.client.Call(ctx, MethodSessionLoad, LoadSessionParams{SessionID: created.SessionID}, &loaded); err != nil { + t.Fatalf("session/load: %v", err) + } + if len(loaded.ConfigOptions) != 2 || loaded.ConfigOptions[0].CurrentValue != "gpt-5.4-mini" { + t.Fatalf("load model option = %+v", loaded.ConfigOptions) + } + if loaded.ConfigOptions[1].CurrentValue != string(agent.PermissionModeAuto) { + t.Fatalf("load mode option = %+v", loaded.ConfigOptions[1]) + } +} + +func TestACPModelDiscoveryFailureUsesConfiguredFallbackOnly(t *testing.T) { + deps := testDeps(t) + deps.ResolveConfig = func(_ string, _ config.Overrides) (config.ResolvedConfig, error) { + return config.ResolvedConfig{Provider: config.ProviderProfile{ + Name: "ChatGPT", CatalogID: "chatgpt", Model: " configured-model ", + }}, nil + } + deps.DiscoverModels = func(context.Context, config.ProviderProfile) ([]providermodeldiscovery.Model, error) { + return nil, fmt.Errorf("discovery unavailable") + } + h := newHarness(t, deps) + defer h.stop() + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + var created NewSessionResult + if err := h.client.Call(ctx, MethodSessionNew, NewSessionParams{Cwd: t.TempDir()}, &created); err != nil { + t.Fatalf("session/new: %v", err) + } + option := created.ConfigOptions[0] + if option.CurrentValue != "configured-model" || len(option.Options) != 1 || option.Options[0].Value != "configured-model" { + t.Fatalf("fallback model option = %+v", option) + } +} + +func TestACPCustomProviderAllowsUnadvertisedModel(t *testing.T) { + deps := testDeps(t) + deps.ResolveConfig = func(_ string, o config.Overrides) (config.ResolvedConfig, error) { + model := "configured-model" + if o.Provider.Model != "" { + model = o.Provider.Model + } + return config.ResolvedConfig{Provider: config.ProviderProfile{ + Name: "Custom", CatalogID: "custom-openai-compatible", Model: model, + }}, nil + } + h := newHarness(t, deps) + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + var created NewSessionResult + if err := h.client.Call(ctx, MethodSessionNew, NewSessionParams{Cwd: t.TempDir()}, &created); err != nil { + t.Fatalf("session/new: %v", err) + } + var selected SetSessionConfigOptionResult + if err := h.client.Call(ctx, MethodSessionSetConfigOption, SetSessionConfigOptionParams{ + SessionID: created.SessionID, ConfigID: configIDModel, Value: " vendor/model ", + }, &selected); err != nil { + t.Fatalf("set custom model: %v", err) + } + if got := selected.ConfigOptions[0].CurrentValue; got != "vendor/model" { + t.Fatalf("custom model = %q", got) + } + h.stop() + h = newHarness(t, deps) + defer h.stop() + var loaded LoadSessionResult + if err := h.client.Call(ctx, MethodSessionLoad, LoadSessionParams{SessionID: created.SessionID}, &loaded); err != nil { + t.Fatalf("session/load: %v", err) + } + option := loaded.ConfigOptions[0] + if option.CurrentValue != "vendor/model" || !modelChoiceExists(option.Options, "vendor/model") { + t.Fatalf("loaded custom model option = %+v", option) + } +} + +func TestACPModelDiscoveryFiltersProviderIncompatibleModels(t *testing.T) { + a := &Agent{deps: Deps{ + ResolveConfig: func(string, config.Overrides) (config.ResolvedConfig, error) { + return config.ResolvedConfig{Provider: config.ProviderProfile{ + CatalogID: "opencode-go-anthropic-compatible", Model: "minimax-m3", + }}, nil + }, + DiscoverModels: func(context.Context, config.ProviderProfile) ([]providermodeldiscovery.Model, error) { + return []providermodeldiscovery.Model{{ID: "qwen3.7-plus"}, {ID: "claude-sonnet-4.5"}}, nil + }, + }} + _, options, restricted, err := a.resolveModelChoices(context.Background(), t.TempDir()) + if err != nil { + t.Fatal(err) + } + if !restricted || !modelChoiceExists(options, "qwen3.7-plus") || modelChoiceExists(options, "claude-sonnet-4.5") { + t.Fatalf("filtered model options = %+v, restricted=%v", options, restricted) + } +} + +func TestACPModelDiscoveryHonorsCancellation(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + cancel() + a := &Agent{deps: Deps{ + ResolveConfig: func(string, config.Overrides) (config.ResolvedConfig, error) { + return config.ResolvedConfig{Provider: config.ProviderProfile{Model: "configured-model"}}, nil + }, + DiscoverModels: func(context.Context, config.ProviderProfile) ([]providermodeldiscovery.Model, error) { + return nil, context.Canceled + }, + }} + if _, _, _, err := a.resolveModelChoices(ctx, t.TempDir()); !errors.Is(err, context.Canceled) { + t.Fatalf("resolve error = %v, want context canceled", err) + } +} + +func TestACPConfigOptionWireSchema(t *testing.T) { + b, err := json.Marshal(SessionConfigOption{ + ID: "model", Name: "Model", Description: "desc", Category: "model", + Type: configOptionTypeSelect, CurrentValue: "m1", + Options: []SessionConfigOptionValue{{Value: "m1", Name: "m1", Description: "choice"}}, + }) + if err != nil { + t.Fatal(err) + } + var wire map[string]any + if err := json.Unmarshal(b, &wire); err != nil { + t.Fatal(err) + } + for _, key := range []string{"type", "category", "currentValue", "options"} { + if _, ok := wire[key]; !ok { + t.Errorf("wire field %q absent: %s", key, b) + } + } + if _, ok := wire["value"]; ok { + t.Errorf("obsolete option value field present: %s", b) + } + if _, ok := wire["values"]; ok { + t.Errorf("obsolete values field present: %s", b) + } + choice := wire["options"].([]any)[0].(map[string]any) + if choice["value"] != "m1" { + t.Errorf("options[].value = %#v", choice["value"]) + } +} + func TestACPUnknownSessionPromptErrors(t *testing.T) { h := newHarness(t, testDeps(t)) defer h.stop() @@ -181,6 +384,15 @@ func TestACPSetModeUpdatesSession(t *testing.T) { if err := h.client.Call(ctx, MethodSessionSetMode, SetSessionModeParams{SessionID: newRes.SessionID, ModeID: string(agent.PermissionModeAsk)}, &SetSessionModeResult{}); err != nil { t.Fatalf("set_mode ask: %v", err) } + var configured SetSessionConfigOptionResult + if err := h.client.Call(ctx, MethodSessionSetConfigOption, SetSessionConfigOptionParams{ + SessionID: newRes.SessionID, ConfigID: configIDMode, Value: string(agent.PermissionModeAuto), + }, &configured); err != nil { + t.Fatalf("set_config_option mode: %v", err) + } + if got := configured.ConfigOptions[1].CurrentValue; got != string(agent.PermissionModeAuto) { + t.Fatalf("configured mode = %q", got) + } // Unsafe must be rejected over ACP — a client can't self-grant no-prompt host access. if err := h.client.Call(ctx, MethodSessionSetMode, SetSessionModeParams{SessionID: newRes.SessionID, ModeID: string(agent.PermissionModeUnsafe)}, &SetSessionModeResult{}); err == nil { t.Fatal("expected Unsafe mode to be rejected over ACP") diff --git a/internal/acp/types.go b/internal/acp/types.go index d507ac303..340c9be88 100644 --- a/internal/acp/types.go +++ b/internal/acp/types.go @@ -337,19 +337,22 @@ type SetSessionModeParams struct { type SetSessionModeResult struct{} -// ---- session config options (model selection) ---- +// ---- session config options ---- type SessionConfigOptionValue struct { - ID string `json:"id"` - Name string `json:"name"` + Value string `json:"value"` + Name string `json:"name"` + Description string `json:"description,omitempty"` } type SessionConfigOption struct { - ID string `json:"id"` - Name string `json:"name"` - Description string `json:"description,omitempty"` - Value string `json:"value"` - Values []SessionConfigOptionValue `json:"values,omitempty"` + ID string `json:"id"` + Name string `json:"name"` + Description string `json:"description,omitempty"` + Category string `json:"category,omitempty"` + Type string `json:"type"` + CurrentValue string `json:"currentValue"` + Options []SessionConfigOptionValue `json:"options"` } type SetSessionConfigOptionParams struct { @@ -373,6 +376,11 @@ type ZeroSetModelResult struct { Model string `json:"model"` } -// configIDModel is the SessionConfigOption id ZERO uses to expose model choice -// through the standard session/set_config_option method. -const configIDModel = "model" +const ( + configIDModel = "model" + configIDMode = "mode" + + configOptionTypeSelect = "select" + configCategoryModel = "model" + configCategoryMode = "mode" +) diff --git a/internal/cli/acp.go b/internal/cli/acp.go index ec146699a..dfade3b32 100644 --- a/internal/cli/acp.go +++ b/internal/cli/acp.go @@ -1,6 +1,7 @@ package cli import ( + "context" "fmt" "io" "os" @@ -9,6 +10,7 @@ import ( "github.com/Gitlawb/zero/internal/acp" "github.com/Gitlawb/zero/internal/agent" "github.com/Gitlawb/zero/internal/config" + "github.com/Gitlawb/zero/internal/providermodeldiscovery" "github.com/Gitlawb/zero/internal/sandbox" "github.com/Gitlawb/zero/internal/tools" ) @@ -43,6 +45,9 @@ func runACP(args []string, stdout io.Writer, stderr io.Writer, deps appDeps) int conn := acp.NewConn(deps.stdin, stdout) acp.NewAgent(conn, acp.Deps{ ResolveConfig: deps.resolveConfig, + DiscoverModels: func(ctx context.Context, profile config.ProviderProfile) ([]providermodeldiscovery.Model, error) { + return defaultDiscoverProviderModels(ctx, discoveryCredentialProfile(profile)) + }, // deps.newProvider is wrapped in fillAppDeps to apply the stored API key, // so ACP is authenticated for apiKeyStored profiles like every other // surface — no ACP-specific credential handling needed. diff --git a/internal/cli/provider_models.go b/internal/cli/provider_models.go index b09eca50d..542ad1207 100644 --- a/internal/cli/provider_models.go +++ b/internal/cli/provider_models.go @@ -10,6 +10,7 @@ import ( "github.com/Gitlawb/zero/internal/config" "github.com/Gitlawb/zero/internal/providermodelcatalog" "github.com/Gitlawb/zero/internal/providermodeldiscovery" + "github.com/Gitlawb/zero/internal/providers" ) type providerModelsOptions struct { @@ -137,7 +138,12 @@ func discoveryCredentialProfile(profile config.ProviderProfile) config.ProviderP // the provider's model-listing endpoint with no curated-catalog merge or // coding-model filtering, so a custom provider's full model list is returned. func defaultDiscoverProviderModels(ctx context.Context, profile config.ProviderProfile) ([]providermodeldiscovery.Model, error) { - return providermodeldiscovery.Discover(ctx, profile, providermodeldiscovery.Options{}) + resolver, loginKey := oauthLoginForProfile(profile) + return providermodeldiscovery.Discover(ctx, profile, providermodeldiscovery.Options{ + OAuthResolver: resolver, + CodexAccountResolver: providers.CodexAccountResolverForLogin(loginKey), + UserAgent: userAgent(), + }) } func parseProviderModelsArgs(args []string) (providerModelsOptions, bool, error) { diff --git a/internal/providermodeldiscovery/discovery.go b/internal/providermodeldiscovery/discovery.go index 278bf550f..76a7174a6 100644 --- a/internal/providermodeldiscovery/discovery.go +++ b/internal/providermodeldiscovery/discovery.go @@ -14,6 +14,7 @@ import ( "github.com/Gitlawb/zero/internal/config" "github.com/Gitlawb/zero/internal/providercatalog" "github.com/Gitlawb/zero/internal/providermodelcatalog" + "github.com/Gitlawb/zero/internal/providers/openai" "github.com/Gitlawb/zero/internal/providers/providerio" "github.com/Gitlawb/zero/internal/redaction" ) @@ -35,9 +36,12 @@ type Model struct { } type Options struct { - HTTPClient *http.Client - ModelsDevURL string - OpenGatewayURL string + HTTPClient *http.Client + ModelsDevURL string + OpenGatewayURL string + OAuthResolver providerio.TokenResolver + CodexAccountResolver openai.CodexAccountResolver + UserAgent string } func DiscoverCatalog(ctx context.Context, provider providercatalog.Descriptor, profile config.ProviderProfile, options Options) ([]Model, error) { @@ -188,6 +192,18 @@ func discoverOpenAIModels(ctx context.Context, profile config.ProviderProfile, o if err != nil { return nil, err } + var configure func(*http.Request) + if providercatalog.NormalizeID(profile.CatalogID) == "chatgpt" { + configure = func(request *http.Request) { + account := "" + if options.CodexAccountResolver != nil { + if resolved, ok, resolveErr := options.CodexAccountResolver(request.Context()); resolveErr == nil && ok { + account = resolved + } + } + openai.ApplyCodexHeaders(request, account, options.UserAgent) + } + } return fetchProviderModels(ctx, endpoint, profile, options, providerio.AuthHeaders{ APIKey: profile.APIKey, DefaultAuthHeader: "Authorization", @@ -196,7 +212,7 @@ func discoverOpenAIModels(ctx context.Context, profile config.ProviderProfile, o AuthScheme: profile.AuthScheme, AuthHeaderValue: profile.AuthHeaderValue, CustomHeaders: providerio.CopyHeaders(profile.CustomHeaders), - }, nil) + }, configure) } func discoverAnthropicModels(ctx context.Context, profile config.ProviderProfile, options Options) ([]Model, error) { @@ -217,25 +233,16 @@ func discoverAnthropicModels(ctx context.Context, profile config.ProviderProfile } func fetchProviderModels(ctx context.Context, endpoint string, profile config.ProviderProfile, options Options, auth providerio.AuthHeaders, configure func(*http.Request)) ([]Model, error) { - request, err := http.NewRequestWithContext(ctx, http.MethodGet, endpoint, nil) - if err != nil { - return nil, err - } - // Authenticate via either an APIKey or a raw auth-header value / custom - // headers, matching how the live providers build their requests - // (internal/providers/providerio). Honoring AuthHeaderValue keeps discovery - // consistent with the credential-present logic elsewhere. - providerio.ApplyAuthHeaders(request, auth) - request.Header.Set("Accept", "application/json") - if configure != nil { - configure(request) - } - client := options.HTTPClient if client == nil { client = &http.Client{Timeout: 10 * time.Second} } - response, err := client.Do(request) + response, err := providerio.SendWithAuthRetry(ctx, client, http.MethodGet, endpoint, nil, auth, options.OAuthResolver, func(request *http.Request) { + request.Header.Set("Accept", "application/json") + if configure != nil { + configure(request) + } + }, 1) if err != nil { return nil, redactDiscoveryError(err, profile) } diff --git a/internal/providermodeldiscovery/discovery_test.go b/internal/providermodeldiscovery/discovery_test.go index ece03bffa..c69eafdf3 100644 --- a/internal/providermodeldiscovery/discovery_test.go +++ b/internal/providermodeldiscovery/discovery_test.go @@ -52,6 +52,63 @@ func TestDiscoverOpenAICompatibleModelsFetchesModelsEndpoint(t *testing.T) { } } +func TestDiscoverChatGPTModelsUsesOAuthAndCodexHeaders(t *testing.T) { + var requests int + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + requests++ + if got := r.URL.Path; got != "/backend-api/codex/models" { + t.Errorf("path = %q, want Codex models endpoint", got) + } + wantToken := "Bearer old-token" + wantAccount := "old-account" + if requests == 2 { + wantToken = "Bearer refreshed-token" + wantAccount = "refreshed-account" + } + if got := r.Header.Get("Authorization"); got != wantToken { + t.Errorf("request %d Authorization = %q, want %q", requests, got, wantToken) + } + if got := r.Header.Get("chatgpt-account-id"); got != wantAccount { + t.Errorf("request %d chatgpt-account-id = %q, want %q", requests, got, wantAccount) + } + if got := r.Header.Get("originator"); got != "codex_cli_rs" { + t.Errorf("originator = %q", got) + } + if requests == 1 { + w.WriteHeader(http.StatusUnauthorized) + return + } + _, _ = w.Write([]byte(`{"data":[{"id":"gpt-5.4"}]}`)) + })) + defer server.Close() + + models, err := Discover(context.Background(), config.ProviderProfile{ + CatalogID: "chatgpt", + ProviderKind: config.ProviderKindOpenAICompatible, + BaseURL: server.URL + "/backend-api/codex", + }, Options{ + HTTPClient: server.Client(), + OAuthResolver: func(_ context.Context, force bool) (string, string, bool, error) { + if force { + return "Authorization", "Bearer refreshed-token", true, nil + } + return "Authorization", "Bearer old-token", true, nil + }, + CodexAccountResolver: func(context.Context) (string, bool, error) { + if requests == 0 { + return "old-account", true, nil + } + return "refreshed-account", true, nil + }, + }) + if err != nil { + t.Fatalf("Discover returned error: %v", err) + } + if requests != 2 || len(models) != 1 || models[0].ID != "gpt-5.4" { + t.Fatalf("requests = %d, models = %#v", requests, models) + } +} + func TestDiscoverAIMLAPIModelsSendsAuthAndCustomHeadersWithoutAttribution(t *testing.T) { server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { for header, want := range map[string]string{ diff --git a/internal/providers/factory.go b/internal/providers/factory.go index b0ed75c5b..7faf36305 100644 --- a/internal/providers/factory.go +++ b/internal/providers/factory.go @@ -310,13 +310,6 @@ func isCodexCatalog(profile config.ProviderProfile, _ resolvedProfile) bool { // login than the bearer — a mismatch the backend rejects. func newCodexProvider(profile config.ProviderProfile, resolved resolvedProfile, options Options) (zeroruntime.Provider, error) { accountKey := options.OAuthLoginKey - resolver := openai.CodexAccountResolver(func(ctx context.Context) (string, bool, error) { - account := codexAccountForKey(accountKey) - if account == "" { - return "", false, nil - } - return account, true, nil - }) return openai.NewCodexProvider(openai.CodexOptions{ Options: openai.Options{ BaseURL: resolved.baseURL, @@ -337,7 +330,7 @@ func newCodexProvider(profile config.ProviderProfile, resolved resolvedProfile, // override. The codex provider's constructor derives the // `/responses` endpoint from BaseURL, so the factory stays out of // the path. - AccountResolver: resolver, + AccountResolver: CodexAccountResolverForLogin(accountKey), }) } @@ -360,3 +353,14 @@ func codexAccountForKey(key string) string { } return strings.TrimSpace(token.Account) } + +// CodexAccountResolverForLogin returns the same per-request account resolver +// used by the runtime provider. Auxiliary Codex requests use this rather than +// independently selecting or parsing an OAuth login, which could mismatch the +// bearer selected by oauthLoginForProfile. +func CodexAccountResolverForLogin(key string) openai.CodexAccountResolver { + return func(context.Context) (string, bool, error) { + account := codexAccountForKey(key) + return account, account != "", nil + } +} diff --git a/internal/providers/openai/codex.go b/internal/providers/openai/codex.go index bd37f70dd..a9688f3f8 100644 --- a/internal/providers/openai/codex.go +++ b/internal/providers/openai/codex.go @@ -23,6 +23,19 @@ const ( codexOriginatorHeader = "originator" ) +// ApplyCodexHeaders applies the non-bearer headers required by ChatGPT's Codex +// backend. It is shared by runtime completions and auxiliary Codex endpoints +// such as live model discovery so their authentication cannot drift apart. +func ApplyCodexHeaders(req *http.Request, accountID, userAgent string) { + req.Header.Set(codexOriginatorHeader, codexDefaultOriginator) + if accountID = strings.TrimSpace(accountID); accountID != "" { + req.Header.Set(codexAccountHeader, accountID) + } + if userAgent = strings.TrimSpace(userAgent); userAgent != "" { + req.Header.Set("User-Agent", userAgent) + } +} + // CodexAccountResolver returns the `chatgpt_account_id` claim for the bearer // that is about to be sent on a request. It is invoked once per request // (including the 401-refresh retry) so the value can be re-derived from the @@ -168,15 +181,13 @@ func (p *CodexProvider) StreamCompletion(ctx context.Context, request zeroruntim // openai provider. It sets the three Codex-required headers; the bearer is // applied separately by the openai provider's auth path. func (p *CodexProvider) injectCodexHeaders(req *http.Request) { - req.Header.Set(codexOriginatorHeader, p.originator) - if account, ok, err := p.resolveAccount(req.Context()); err == nil && ok && account != "" { - req.Header.Set(codexAccountHeader, account) + account := "" + if resolved, ok, err := p.resolveAccount(req.Context()); err == nil && ok && resolved != "" { + account = strings.TrimSpace(resolved) } - // Branded User-Agent overrides the openai provider's default. Set last - // so a caller that supplies a different UserAgent in custom-headers is - // still respected (the openai provider's setExtra already ran before us). - if p.userAgent != "" { - req.Header.Set("User-Agent", p.userAgent) + ApplyCodexHeaders(req, account, p.userAgent) + if p.originator != codexDefaultOriginator { + req.Header.Set(codexOriginatorHeader, p.originator) } } diff --git a/internal/sessions/session_title_test.go b/internal/sessions/session_title_test.go index e2123d283..61e54a9a3 100644 --- a/internal/sessions/session_title_test.go +++ b/internal/sessions/session_title_test.go @@ -75,3 +75,33 @@ func TestUpdateTitle(t *testing.T) { t.Fatal("expected an invalid session id to be rejected") } } + +func TestUpdateModel(t *testing.T) { + store := newTitleTestStore(t) + session, err := store.Create(CreateInput{ModelID: "model-a"}) + if err != nil { + t.Fatalf("create: %v", err) + } + if _, err := store.AppendEvent(session.SessionID, AppendEventInput{Type: EventMessage, Payload: map[string]any{"role": "user"}}); err != nil { + t.Fatalf("append: %v", err) + } + before, err := store.Get(session.SessionID) + if err != nil || before == nil { + t.Fatalf("get before: %v", err) + } + + updated, err := store.UpdateModel(session.SessionID, " model-b ") + if err != nil { + t.Fatalf("update model: %v", err) + } + if updated.ModelID != "model-b" { + t.Fatalf("model = %q, want model-b", updated.ModelID) + } + if updated.UpdatedAt != before.UpdatedAt || updated.EventCount != before.EventCount { + t.Fatalf("model update changed activity metadata: before=%+v after=%+v", before, updated) + } + persisted, err := store.Get(session.SessionID) + if err != nil || persisted == nil || persisted.ModelID != "model-b" { + t.Fatalf("persisted model: metadata=%+v err=%v", persisted, err) + } +} diff --git a/internal/sessions/store.go b/internal/sessions/store.go index 3eb98f14e..2f8854de9 100644 --- a/internal/sessions/store.go +++ b/internal/sessions/store.go @@ -712,6 +712,33 @@ func (store *Store) UpdateTitle(sessionID string, title string) (Metadata, error return session, nil } +// UpdateModel replaces a session's selected model without changing its activity +// timestamp or event counters. An empty model clears the session override. +func (store *Store) UpdateModel(sessionID string, modelID string) (Metadata, error) { + if !ValidSessionID(sessionID) { + return Metadata{}, fmt.Errorf("invalid zero session id %q", sessionID) + } + modelID = strings.TrimSpace(modelID) + unlock, err := store.lockSession(sessionID) + if err != nil { + return Metadata{}, err + } + defer unlock() + + session, err := store.readMetadata(sessionID) + if err != nil { + return Metadata{}, err + } + if session.ModelID == modelID { + return session, nil + } + session.ModelID = modelID + if err := store.writeMetadata(session); err != nil { + return Metadata{}, err + } + return session, nil +} + func (store *Store) ReadEvents(sessionID string) ([]Event, error) { if !ValidSessionID(sessionID) { return nil, fmt.Errorf("invalid zero session id %q", sessionID)