diff --git a/internal/extbridge/context.go b/internal/extbridge/context.go index 25cbfaf9..c3cdd621 100644 --- a/internal/extbridge/context.go +++ b/internal/extbridge/context.go @@ -249,10 +249,10 @@ func BaseContext(ctx context.Context, kitInstance *kit.Kit) extensions.Context { // inspecting the active model's limits or pricing would force // every extension to pass a model string back in by hand. // - // Resolve from configuration rather than from the extension - // Context: Context.Model is a bare display ID ("claude-opus-5"), - // which the registry cannot address unambiguously because many - // providers serve the same ID. + // Resolve from the Kit rather than from the extension Context. + // Until the first model switch, Context.Model is a bare display + // ID ("claude-opus-5"). The registry cannot address it + // unambiguously, because many providers serve the same ID. if model == "" { model = kitInstance.GetCurrentModel() } @@ -261,11 +261,14 @@ func BaseContext(ctx context.Context, kitInstance *kit.Kit) extensions.Context { CheckModelAvailable: func(model string) bool { return kit.CheckModelAvailable(model) }, + // The provider and the model ID come from the Kit too. + // Context.Model is only the model ID until the first switch, and + // a model ID can contain "/", as in openrouter/anthropic/claude-opus-5. GetCurrentProvider: func() string { - return kit.GetCurrentProvider(kitInstance.Extensions().GetContext().Model) + return kit.GetCurrentProvider(kitInstance.GetCurrentModel()) }, GetCurrentModelID: func() string { - return kit.GetCurrentModelID(kitInstance.Extensions().GetContext().Model) + return kit.GetCurrentModelID(kitInstance.GetCurrentModel()) }, GetTerminalSize: func() (int, int) { // Read through to the runner rather than snapshotting on Context: diff --git a/internal/extbridge/model_capabilities_test.go b/internal/extbridge/model_capabilities_test.go index 1033e6d8..9200fd75 100644 --- a/internal/extbridge/model_capabilities_test.go +++ b/internal/extbridge/model_capabilities_test.go @@ -85,6 +85,46 @@ func TestBaseContext_GetModelCapabilitiesResolvesActiveModel(t *testing.T) { } } +// TestBaseContext_CurrentModelFollowsSwitch verifies that the bridge reports +// the model in use through GetModelCapabilities(""), GetCurrentProvider and +// GetCurrentModelID. That is the configured model at the start, and the new +// model after SetModel. It does not depend on Context.Model, which is only +// the model ID until the first switch in the TUI. The switch goes to another +// provider, so a bridge that keeps the provider from startup fails the test. +func TestBaseContext_CurrentModelFollowsSwitch(t *testing.T) { + ctx := context.Background() + host := newTestKit(t) + extCtx := BaseContext(ctx, host) + + check := func(when, provider, modelID string) { + t.Helper() + caps, errStr := extCtx.GetModelCapabilities("") + if errStr != "" || caps.Provider != provider || caps.ModelID != modelID { + t.Errorf("GetModelCapabilities(\"\") %s = %s/%s %q; want %s/%s", + when, caps.Provider, caps.ModelID, errStr, provider, modelID) + } + if got := extCtx.GetCurrentProvider(); got != provider { + t.Errorf("GetCurrentProvider() %s = %q; want %q", when, got, provider) + } + if got := extCtx.GetCurrentModelID(); got != modelID { + t.Errorf("GetCurrentModelID() %s = %q; want %q", when, got, modelID) + } + if got, want := host.GetCurrentModel(), provider+"/"+modelID; got != want { + t.Errorf("GetCurrentModel() %s = %q; want %q", when, got, want) + } + } + + check("at the start", "anthropic", "claude-sonnet-4-5-20250929") + // The key from newTestKit is bound to the provider of the startup model. + // The new provider takes its key from the environment. The dummy value + // lets the provider construct; the test sends no request. + t.Setenv("OPENAI_API_KEY", "sk-test-not-a-real-key") + if err := host.SetModel(ctx, "openai/gpt-4o-mini"); err != nil { + t.Fatalf("SetModel: %v", err) + } + check("after SetModel", "openai", "gpt-4o-mini") +} + // TestBaseContext_GetModelCapabilitiesUnknownModel verifies the empty-model // fallback does not mask genuine lookup failures. func TestBaseContext_GetModelCapabilitiesUnknownModel(t *testing.T) { diff --git a/pkg/kit/kit.go b/pkg/kit/kit.go index cda14893..cf47ee9e 100644 --- a/pkg/kit/kit.go +++ b/pkg/kit/kit.go @@ -53,8 +53,15 @@ type Kit struct { closeDone chan struct{} closeErr error - agent *agent.Agent - session SessionManager + agent *agent.Agent + session SessionManager + // modelString is the model in use, as "provider/model". setModelLocked + // is its only writer after New, and it writes with modelMu held. Reads + // that can run on another goroutine go through GetModelString, which + // takes modelMu. They must not take promptOptsMu: applyPromptOptions + // holds it for a whole PromptWithOptions turn, and extension code that + // runs inside the turn reads the model. + modelMu sync.RWMutex modelString string // endpointProvider is the provider the provider-url, provider-api-key // and provider-wire overrides belong to: the provider of the model that @@ -898,7 +905,9 @@ func (m *Kit) setModelLocked(ctx context.Context, modelString string) error { // block runtime changes while the agent waits for MCP loading. m.applyComposedSystemPrompt() + m.modelMu.Lock() m.modelString = modelString + m.modelMu.Unlock() m.activeModelFromFactory = fromFactory // Keep the selected mode during a same-model provider rebuild. A model // switch uses the new model's default mode. @@ -1047,12 +1056,18 @@ func environmentSection(cwd string, bare bool) string { } // GetCurrentModel returns the fully qualified "provider/model" string for the -// model currently in use, as resolved from configuration. Unlike the model -// name carried on the extension Context — which is a bare display ID — this is -// always registry-addressable, so it can be passed straight to -// GetModelCapabilities. +// model currently in use: the configured model until SetModel switches it. +// The model name carried on the extension Context is a bare display ID until +// the first switch. This string is always registry-addressable, so it can be +// passed straight to GetModelCapabilities. func (m *Kit) GetCurrentModel() string { - if m == nil || m.v == nil { + if m == nil { + return "" + } + if model := m.GetModelString(); model != "" { + return model + } + if m.v == nil { return "" } return m.v.GetString("model") @@ -1147,7 +1162,7 @@ func (m *Kit) ExecuteCompletion(ctx context.Context, req CompleteRequest) (Compl if req.Model == "" { // Reuse the active agent's model. llmModel = m.agent.GetModel() - usedModel = m.modelString + usedModel = m.GetModelString() skipMaxOutputTokens = m.agent.SkipMaxOutputTokens() closer = func() {} // nothing to clean up // A factory-backed model gets the options its factory returned: Kit @@ -2521,12 +2536,13 @@ func (m *Kit) Subagent(ctx context.Context, cfg SubagentConfig) (*SubagentResult // Resolve model: fall back to parent's model, and inherit the parent's // provider when only a bare model name is given (e.g. "claude-haiku" // instead of "anthropic/claude-haiku"). This avoids provider guessing. + parentModel := m.GetModelString() model := cfg.Model if model == "" { - model = m.modelString + model = parentModel } else if !strings.Contains(model, "/") { // Bare model name — prepend parent's provider. - if parts := strings.SplitN(m.modelString, "/", 2); len(parts) == 2 { + if parts := strings.SplitN(parentModel, "/", 2); len(parts) == 2 { model = parts[0] + "/" + model } } @@ -2536,7 +2552,7 @@ func (m *Kit) Subagent(ctx context.Context, cfg SubagentConfig) (*SubagentResult // gives the calling agent immediate feedback it can act on — e.g. // correcting a typo — instead of waiting for a full Kit.New() cycle // that silently falls back to the parent model. - if model != m.modelString && !m.hasInstanceProvider(model) { + if model != parentModel && !m.hasInstanceProvider(model) { if err := models.GetGlobalRegistry().ValidateModelString(model); err != nil { return nil, fmt.Errorf("invalid subagent model %q: %w", model, err) } @@ -3745,7 +3761,10 @@ func (m *Kit) ClearSession() { // GetModelString returns the current model string identifier (e.g., // "anthropic/claude-sonnet-4-5-20250929" or "openai/gpt-4") being used by the agent. +// It is safe to call while SetModel runs on another goroutine. func (m *Kit) GetModelString() string { + m.modelMu.RLock() + defer m.modelMu.RUnlock() return m.modelString } @@ -3753,7 +3772,7 @@ func (m *Kit) GetModelString() string { // (capabilities, pricing, limits). Returns nil if the model is not in the // registry — this is expected for new models or custom fine-tunes. func (m *Kit) GetModelInfo() *ModelInfo { - provider, modelID, err := ParseModelString(m.modelString) + provider, modelID, err := ParseModelString(m.GetModelString()) if err != nil { return nil } @@ -3805,7 +3824,7 @@ func (m *Kit) SetThinkingLevel(ctx context.Context, level string) error { // // Returns the full vocabulary when no model is selected yet. func (m *Kit) SupportedThinkingLevels() []string { - provider, modelName, err := ParseModelString(m.modelString) + provider, modelName, err := ParseModelString(m.GetModelString()) if err != nil { return ThinkingLevels() } diff --git a/pkg/kit/model_modes.go b/pkg/kit/model_modes.go index 77efb646..bc1b4ef3 100644 --- a/pkg/kit/model_modes.go +++ b/pkg/kit/model_modes.go @@ -24,13 +24,14 @@ type ModelMode struct { // ModelModes returns catalog modes, including unavailable modes and their reasons. func (m *Kit) ModelModes() []ModelMode { - info := models.LookupModelForSettings(m.modelString) + modelString := m.GetModelString() + info := models.LookupModelForSettings(modelString) if info == nil { return nil } result := make([]ModelMode, 0, len(info.Modes)) for _, mode := range info.Modes { - reason := models.ModeUnsupportedReason(m.modelString, mode.Name, info) + reason := models.ModeUnsupportedReason(modelString, mode.Name, info) cost := info.EffectiveCost(mode.Name) result = append(result, ModelMode{Name: mode.Name, Experimental: true, Supported: reason == "", UnsupportedReason: reason, InputCost: cost.Input, OutputCost: cost.Output}) } diff --git a/pkg/kit/model_switch_prompt_test.go b/pkg/kit/model_switch_prompt_test.go index fb014bd8..faa98d2e 100644 --- a/pkg/kit/model_switch_prompt_test.go +++ b/pkg/kit/model_switch_prompt_test.go @@ -6,6 +6,7 @@ import ( "path/filepath" "strings" "testing" + "time" "github.com/spf13/viper" ) @@ -190,3 +191,57 @@ func TestSetModelConcurrentWithPromptRefresh(t *testing.T) { } <-done } + +// TestSetModelConcurrentWithModelReads switches the model while another +// goroutine reads it through the accessors that extensions and the TUI call. +// Run with -race: the accessors must read the model string under its lock. +func TestSetModelConcurrentWithModelReads(t *testing.T) { + requireCatalogModel(t, testCatalogModel) + k := newBareTestKit(t, false) + + done := make(chan struct{}) + go func() { + defer close(done) + for range 20 { + _ = k.GetModelString() + _ = k.GetCurrentModel() + _ = k.GetModelInfo() + _ = k.SupportedThinkingLevels() + _ = k.ModelModes() + } + }() + for i := range 4 { + model := testCatalogModel + if i%2 == 1 { + model = "openai/gpt-4o-mini" + } + if err := k.SetModel(context.Background(), model); err != nil { + t.Errorf("SetModel(%s): %v", model, err) + break + } + } + <-done +} + +// TestModelAccessorsDoNotWaitForPromptOptions holds promptOptsMu, as +// applyPromptOptions does for a whole PromptWithOptions turn, and reads the +// model from another goroutine. Extension code reads the model inside the +// turn, in the LLM usage event for example, so a read that waits for +// promptOptsMu would block the turn. +func TestModelAccessorsDoNotWaitForPromptOptions(t *testing.T) { + k := newBareTestKit(t, false) + + k.promptOptsMu.Lock() + defer k.promptOptsMu.Unlock() + + got := make(chan string, 1) + go func() { got <- k.GetCurrentModel() }() + select { + case model := <-got: + if model != "openai/gpt-4o-mini" { + t.Errorf("GetCurrentModel() = %q; want %q", model, "openai/gpt-4o-mini") + } + case <-time.After(10 * time.Second): + t.Fatal("GetCurrentModel() waits for promptOptsMu") + } +} diff --git a/www/pages/sdk/api-reference.md b/www/pages/sdk/api-reference.md index 4bc26447..e2a341e0 100644 --- a/www/pages/sdk/api-reference.md +++ b/www/pages/sdk/api-reference.md @@ -465,7 +465,7 @@ Custom decision providers (an application-supplied `DecisionModel`): - [`Kit.GetAvailableModels`](https://pkg.go.dev/github.com/mark3labs/kit/pkg/kit#Kit.GetAvailableModels) — GetAvailableModels returns a list of known models from the registry. - [`Kit.GetBufferedDebugMessages`](https://pkg.go.dev/github.com/mark3labs/kit/pkg/kit#Kit.GetBufferedDebugMessages) — GetBufferedDebugMessages returns any debug messages that were buffered during initialization, then clears the buffer. - [`Kit.GetContextFiles`](https://pkg.go.dev/github.com/mark3labs/kit/pkg/kit#Kit.GetContextFiles) — GetContextFiles returns the context files. -- [`Kit.GetCurrentModel`](https://pkg.go.dev/github.com/mark3labs/kit/pkg/kit#Kit.GetCurrentModel) — GetCurrentModel returns the fully qualified "provider/model" string for the model currently in use, as resolved from configuration. +- [`Kit.GetCurrentModel`](https://pkg.go.dev/github.com/mark3labs/kit/pkg/kit#Kit.GetCurrentModel) — GetCurrentModel returns the fully qualified "provider/model" string for the model currently in use: the configured model until SetModel switches it. - [`Kit.GetExtensionToolCount`](https://pkg.go.dev/github.com/mark3labs/kit/pkg/kit#Kit.GetExtensionToolCount) — GetExtensionToolCount returns the number of tools registered by extensions. - [`Kit.GetExtraTools`](https://pkg.go.dev/github.com/mark3labs/kit/pkg/kit#Kit.GetExtraTools) — GetExtraTools returns a snapshot of the native extra tools that were added via AddTools / SetExtraTools / RemoveTools or Options.ExtraTools. - [`Kit.GetLoadedServerNames`](https://pkg.go.dev/github.com/mark3labs/kit/pkg/kit#Kit.GetLoadedServerNames) — GetLoadedServerNames returns the names of successfully loaded MCP servers.