From 7f1b2b3f6e8ebc72f1449f89f4e893a6fe921f39 Mon Sep 17 00:00:00 2001 From: hkfires <10558748+hkfires@users.noreply.github.com> Date: Wed, 14 Jan 2026 16:30:28 +0800 Subject: [PATCH] fix(thinking): improve model lookup and validation --- internal/registry/model_registry.go | 15 ++++++++ .../runtime/executor/antigravity_executor.go | 6 ++-- internal/thinking/apply.go | 12 +++---- internal/thinking/provider/claude/apply.go | 11 ++---- internal/thinking/provider/codex/apply.go | 11 ++---- internal/thinking/provider/gemini/apply.go | 11 ++---- .../thinking/provider/gemini/apply_test.go | 34 +++++++------------ internal/thinking/provider/geminicli/apply.go | 11 ++---- .../thinking/provider/geminicli/apply_test.go | 34 +++++++------------ internal/thinking/provider/iflow/apply.go | 8 ++--- .../thinking/provider/iflow/apply_test.go | 30 ++++++---------- internal/thinking/provider/openai/apply.go | 11 ++---- .../thinking/provider/openai/apply_test.go | 21 ++++-------- 13 files changed, 78 insertions(+), 137 deletions(-) diff --git a/internal/registry/model_registry.go b/internal/registry/model_registry.go index c90f6f61..970c2dc9 100644 --- a/internal/registry/model_registry.go +++ b/internal/registry/model_registry.go @@ -132,6 +132,21 @@ func GetGlobalRegistry() *ModelRegistry { return globalRegistry } +// LookupModelInfo searches the dynamic registry first, then falls back to static model definitions. +// +// This helper exists because some code paths only have a model ID and still need Thinking and +// max completion token metadata even when the dynamic registry hasn't been populated. +func LookupModelInfo(modelID string) *ModelInfo { + modelID = strings.TrimSpace(modelID) + if modelID == "" { + return nil + } + if info := GetGlobalRegistry().GetModelInfo(modelID); info != nil { + return info + } + return LookupStaticModelInfo(modelID) +} + // SetHook sets an optional hook for observing model registration changes. func (r *ModelRegistry) SetHook(hook ModelRegistryHook) { if r == nil { diff --git a/internal/runtime/executor/antigravity_executor.go b/internal/runtime/executor/antigravity_executor.go index 1eb7b30f..90ebb53f 100644 --- a/internal/runtime/executor/antigravity_executor.go +++ b/internal/runtime/executor/antigravity_executor.go @@ -1458,7 +1458,7 @@ func generateProjectID() string { // - For Claude models: removing thinkingConfig if budget < minimum allowed func normalizeAntigravityThinking(model string, payload []byte, isClaude bool) []byte { payload = util.StripThinkingConfigIfUnsupported(model, payload) - modelInfo := registry.GetGlobalRegistry().GetModelInfo(model) + modelInfo := registry.LookupModelInfo(model) if modelInfo == nil || modelInfo.Thinking == nil { return payload } @@ -1501,7 +1501,7 @@ func antigravityEffectiveMaxTokens(model string, payload []byte) (max int, fromM if maxTok := gjson.GetBytes(payload, "request.generationConfig.maxOutputTokens"); maxTok.Exists() && maxTok.Int() > 0 { return int(maxTok.Int()), false } - if modelInfo := registry.GetGlobalRegistry().GetModelInfo(model); modelInfo != nil && modelInfo.MaxCompletionTokens > 0 { + if modelInfo := registry.LookupModelInfo(model); modelInfo != nil && modelInfo.MaxCompletionTokens > 0 { return modelInfo.MaxCompletionTokens, true } return 0, false @@ -1510,7 +1510,7 @@ func antigravityEffectiveMaxTokens(model string, payload []byte) (max int, fromM // antigravityMinThinkingBudget returns the minimum thinking budget for a model. // Falls back to -1 if no model info is found. func antigravityMinThinkingBudget(model string) int { - if modelInfo := registry.GetGlobalRegistry().GetModelInfo(model); modelInfo != nil && modelInfo.Thinking != nil { + if modelInfo := registry.LookupModelInfo(model); modelInfo != nil && modelInfo.Thinking != nil { return modelInfo.Thinking.Min } return -1 diff --git a/internal/thinking/apply.go b/internal/thinking/apply.go index 44566cab..415b721c 100644 --- a/internal/thinking/apply.go +++ b/internal/thinking/apply.go @@ -68,9 +68,11 @@ func IsUserDefinedModel(modelInfo *registry.ModelInfo) bool { // // Passthrough behavior (returns original body without error): // - Unknown provider (not in providerAppliers map) -// - modelInfo is nil (model not found in registry) // - modelInfo.Thinking is nil (model doesn't support thinking) // +// Note: Unknown models (modelInfo is nil) are treated as user-defined models: we skip +// validation and still apply the thinking config so the upstream can validate it. +// // Example: // // // With suffix - suffix config takes priority @@ -87,15 +89,13 @@ func ApplyThinking(body []byte, model string, provider string) ([]byte, error) { } // 2. Parse suffix and get modelInfo - // First try dynamic registry, then fall back to static lookup suffixResult := ParseSuffix(model) baseModel := suffixResult.ModelName - modelInfo := registry.GetGlobalRegistry().GetModelInfo(baseModel) - if modelInfo == nil { - modelInfo = registry.LookupStaticModelInfo(baseModel) - } + modelInfo := registry.LookupModelInfo(baseModel) // 3. Model capability check + // Unknown models are treated as user-defined so thinking config can still be applied. + // The upstream service is responsible for validating the configuration. if IsUserDefinedModel(modelInfo) { return applyUserDefinedModel(body, modelInfo, provider, suffixResult) } diff --git a/internal/thinking/provider/claude/apply.go b/internal/thinking/provider/claude/apply.go index 979ecd75..b7833072 100644 --- a/internal/thinking/provider/claude/apply.go +++ b/internal/thinking/provider/claude/apply.go @@ -53,18 +53,11 @@ func init() { // } // } func (a *Applier) Apply(body []byte, config thinking.ThinkingConfig, modelInfo *registry.ModelInfo) ([]byte, error) { - if modelInfo == nil { + if thinking.IsUserDefinedModel(modelInfo) { return applyCompatibleClaude(body, config) } if modelInfo.Thinking == nil { - if modelInfo.Type == "" { - modelID := modelInfo.ID - if modelID == "" { - modelID = "unknown" - } - return nil, thinking.NewThinkingErrorWithModel(thinking.ErrThinkingNotSupported, "thinking not supported for this model", modelID) - } - return applyCompatibleClaude(body, config) + return body, nil } // Only process ModeBudget and ModeNone; other modes pass through diff --git a/internal/thinking/provider/codex/apply.go b/internal/thinking/provider/codex/apply.go index 228bb6fe..3bed318b 100644 --- a/internal/thinking/provider/codex/apply.go +++ b/internal/thinking/provider/codex/apply.go @@ -44,18 +44,11 @@ func init() { // } // } func (a *Applier) Apply(body []byte, config thinking.ThinkingConfig, modelInfo *registry.ModelInfo) ([]byte, error) { - if modelInfo == nil { + if thinking.IsUserDefinedModel(modelInfo) { return applyCompatibleCodex(body, config) } if modelInfo.Thinking == nil { - if modelInfo.Type == "" { - modelID := modelInfo.ID - if modelID == "" { - modelID = "unknown" - } - return nil, thinking.NewThinkingErrorWithModel(thinking.ErrThinkingNotSupported, "thinking not supported for this model", modelID) - } - return applyCompatibleCodex(body, config) + return body, nil } // Only handle ModeLevel and ModeNone; other modes pass through unchanged. diff --git a/internal/thinking/provider/gemini/apply.go b/internal/thinking/provider/gemini/apply.go index bb574c31..c8560f19 100644 --- a/internal/thinking/provider/gemini/apply.go +++ b/internal/thinking/provider/gemini/apply.go @@ -59,18 +59,11 @@ func init() { // } // } func (a *Applier) Apply(body []byte, config thinking.ThinkingConfig, modelInfo *registry.ModelInfo) ([]byte, error) { - if modelInfo == nil { + if thinking.IsUserDefinedModel(modelInfo) { return a.applyCompatible(body, config) } if modelInfo.Thinking == nil { - if modelInfo.Type == "" { - modelID := modelInfo.ID - if modelID == "" { - modelID = "unknown" - } - return nil, thinking.NewThinkingErrorWithModel(thinking.ErrThinkingNotSupported, "thinking not supported for this model", modelID) - } - return a.applyCompatible(body, config) + return body, nil } if config.Mode != thinking.ModeBudget && config.Mode != thinking.ModeLevel && config.Mode != thinking.ModeNone && config.Mode != thinking.ModeAuto { diff --git a/internal/thinking/provider/gemini/apply_test.go b/internal/thinking/provider/gemini/apply_test.go index 1af2fa83..47c7e7ce 100644 --- a/internal/thinking/provider/gemini/apply_test.go +++ b/internal/thinking/provider/gemini/apply_test.go @@ -381,26 +381,21 @@ func TestGeminiApplyConflictingFields(t *testing.T) { } } -// TestGeminiApplyThinkingNotSupported tests error handling when modelInfo.Thinking is nil. +// TestGeminiApplyThinkingNotSupported tests passthrough handling when modelInfo.Thinking is nil. func TestGeminiApplyThinkingNotSupported(t *testing.T) { applier := NewApplier() config := thinking.ThinkingConfig{Mode: thinking.ModeBudget, Budget: 8192} + body := []byte(`{"generationConfig":{"thinkingConfig":{"thinkingBudget":8192}}}`) // Model with nil Thinking support modelInfo := ®istry.ModelInfo{ID: "gemini-unknown", Thinking: nil} - _, err := applier.Apply([]byte(`{}`), config, modelInfo) - if err == nil { - t.Fatal("Apply() expected error for nil Thinking, got nil") + got, err := applier.Apply(body, config, modelInfo) + if err != nil { + t.Fatalf("Apply() expected nil error for nil Thinking, got %v", err) } - - // Verify it's the correct error type - thinkErr, ok := err.(*thinking.ThinkingError) - if !ok { - t.Fatalf("Apply() error type = %T, want *thinking.ThinkingError", err) - } - if thinkErr.Code != thinking.ErrThinkingNotSupported { - t.Fatalf("Apply() error code = %v, want %v", thinkErr.Code, thinking.ErrThinkingNotSupported) + if string(got) != string(body) { + t.Fatalf("expected body unchanged, got %s", string(got)) } } @@ -462,17 +457,14 @@ func TestGeminiApplyEmptyModelID(t *testing.T) { applier := NewApplier() config := thinking.ThinkingConfig{Mode: thinking.ModeBudget, Budget: 8192} modelInfo := ®istry.ModelInfo{ID: "", Thinking: nil} + body := []byte(`{"generationConfig":{"thinkingConfig":{"thinkingBudget":8192}}}`) - _, err := applier.Apply([]byte(`{}`), config, modelInfo) - if err == nil { - t.Fatal("Apply() with empty modelID and nil Thinking should error") + got, err := applier.Apply(body, config, modelInfo) + if err != nil { + t.Fatalf("Apply() expected nil error, got %v", err) } - thinkErr, ok := err.(*thinking.ThinkingError) - if !ok { - t.Fatalf("Apply() error type = %T, want *thinking.ThinkingError", err) - } - if thinkErr.Model != "unknown" { - t.Fatalf("Apply() error model = %q, want %q", thinkErr.Model, "unknown") + if string(got) != string(body) { + t.Fatalf("expected body unchanged, got %s", string(got)) } } diff --git a/internal/thinking/provider/geminicli/apply.go b/internal/thinking/provider/geminicli/apply.go index eb6d82a4..c8887723 100644 --- a/internal/thinking/provider/geminicli/apply.go +++ b/internal/thinking/provider/geminicli/apply.go @@ -29,18 +29,11 @@ func init() { // Apply applies thinking configuration to Gemini CLI request body. func (a *Applier) Apply(body []byte, config thinking.ThinkingConfig, modelInfo *registry.ModelInfo) ([]byte, error) { - if modelInfo == nil { + if thinking.IsUserDefinedModel(modelInfo) { return a.applyCompatible(body, config) } if modelInfo.Thinking == nil { - if modelInfo.Type == "" { - modelID := modelInfo.ID - if modelID == "" { - modelID = "unknown" - } - return nil, thinking.NewThinkingErrorWithModel(thinking.ErrThinkingNotSupported, "thinking not supported for this model", modelID) - } - return a.applyCompatible(body, config) + return body, nil } if config.Mode != thinking.ModeBudget && config.Mode != thinking.ModeLevel && config.Mode != thinking.ModeNone && config.Mode != thinking.ModeAuto { diff --git a/internal/thinking/provider/geminicli/apply_test.go b/internal/thinking/provider/geminicli/apply_test.go index 6bf77dd2..e6900496 100644 --- a/internal/thinking/provider/geminicli/apply_test.go +++ b/internal/thinking/provider/geminicli/apply_test.go @@ -208,26 +208,21 @@ func TestGeminiCLIApplyConflictingFields(t *testing.T) { } } -// TestGeminiCLIApplyThinkingNotSupported tests error handling when modelInfo.Thinking is nil. +// TestGeminiCLIApplyThinkingNotSupported tests passthrough handling when modelInfo.Thinking is nil. func TestGeminiCLIApplyThinkingNotSupported(t *testing.T) { applier := NewApplier() config := thinking.ThinkingConfig{Mode: thinking.ModeBudget, Budget: 8192} + body := []byte(`{"request":{"generationConfig":{"thinkingConfig":{"thinkingBudget":8192}}}}`) // Model with nil Thinking support modelInfo := ®istry.ModelInfo{ID: "gemini-cli-unknown", Thinking: nil} - _, err := applier.Apply([]byte(`{}`), config, modelInfo) - if err == nil { - t.Fatal("Apply() expected error for nil Thinking, got nil") + got, err := applier.Apply(body, config, modelInfo) + if err != nil { + t.Fatalf("Apply() expected nil error for nil Thinking, got %v", err) } - - // Verify it's the correct error type - thinkErr, ok := err.(*thinking.ThinkingError) - if !ok { - t.Fatalf("Apply() error type = %T, want *thinking.ThinkingError", err) - } - if thinkErr.Code != thinking.ErrThinkingNotSupported { - t.Fatalf("Apply() error code = %v, want %v", thinkErr.Code, thinking.ErrThinkingNotSupported) + if string(got) != string(body) { + t.Fatalf("expected body unchanged, got %s", string(got)) } } @@ -252,17 +247,14 @@ func TestGeminiCLIApplyEmptyModelID(t *testing.T) { applier := NewApplier() config := thinking.ThinkingConfig{Mode: thinking.ModeBudget, Budget: 8192} modelInfo := ®istry.ModelInfo{ID: "", Thinking: nil} + body := []byte(`{"request":{"generationConfig":{"thinkingConfig":{"thinkingBudget":8192}}}}`) - _, err := applier.Apply([]byte(`{}`), config, modelInfo) - if err == nil { - t.Fatal("Apply() with empty modelID and nil Thinking should error") + got, err := applier.Apply(body, config, modelInfo) + if err != nil { + t.Fatalf("Apply() expected nil error, got %v", err) } - thinkErr, ok := err.(*thinking.ThinkingError) - if !ok { - t.Fatalf("Apply() error type = %T, want *thinking.ThinkingError", err) - } - if thinkErr.Model != "unknown" { - t.Fatalf("Apply() error model = %q, want %q", thinkErr.Model, "unknown") + if string(got) != string(body) { + t.Fatalf("expected body unchanged, got %s", string(got)) } } diff --git a/internal/thinking/provider/iflow/apply.go b/internal/thinking/provider/iflow/apply.go index 5bca94f2..da986d22 100644 --- a/internal/thinking/provider/iflow/apply.go +++ b/internal/thinking/provider/iflow/apply.go @@ -54,15 +54,11 @@ func init() { // "reasoning_split": true // } func (a *Applier) Apply(body []byte, config thinking.ThinkingConfig, modelInfo *registry.ModelInfo) ([]byte, error) { - if modelInfo == nil { + if thinking.IsUserDefinedModel(modelInfo) { return body, nil } if modelInfo.Thinking == nil { - modelID := modelInfo.ID - if modelID == "" { - modelID = "unknown" - } - return nil, thinking.NewThinkingErrorWithModel(thinking.ErrThinkingNotSupported, "thinking not supported for this model", modelID) + return body, nil } if isGLMModel(modelInfo.ID) { diff --git a/internal/thinking/provider/iflow/apply_test.go b/internal/thinking/provider/iflow/apply_test.go index f0c2a35b..9718c413 100644 --- a/internal/thinking/provider/iflow/apply_test.go +++ b/internal/thinking/provider/iflow/apply_test.go @@ -73,33 +73,23 @@ func TestApplyMissingThinkingSupport(t *testing.T) { applier := NewApplier() tests := []struct { - name string - modelID string - wantModel string + name string + modelID string }{ - {"model id", "glm-4.6", "glm-4.6"}, - {"empty model id", "", "unknown"}, + {"model id", "glm-4.6"}, + {"empty model id", ""}, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { modelInfo := ®istry.ModelInfo{ID: tt.modelID} - got, err := applier.Apply([]byte(`{"model":"`+tt.modelID+`"}`), thinking.ThinkingConfig{}, modelInfo) - if err == nil { - t.Fatalf("expected error, got nil") + body := []byte(`{"model":"` + tt.modelID + `"}`) + got, err := applier.Apply(body, thinking.ThinkingConfig{}, modelInfo) + if err != nil { + t.Fatalf("expected nil error, got %v", err) } - if got != nil { - t.Fatalf("expected nil body on error, got %s", string(got)) - } - thinkingErr, ok := err.(*thinking.ThinkingError) - if !ok { - t.Fatalf("expected ThinkingError, got %T", err) - } - if thinkingErr.Code != thinking.ErrThinkingNotSupported { - t.Fatalf("expected code %s, got %s", thinking.ErrThinkingNotSupported, thinkingErr.Code) - } - if thinkingErr.Model != tt.wantModel { - t.Fatalf("expected model %s, got %s", tt.wantModel, thinkingErr.Model) + if string(got) != string(body) { + t.Fatalf("expected body unchanged, got %s", string(got)) } }) } diff --git a/internal/thinking/provider/openai/apply.go b/internal/thinking/provider/openai/apply.go index aea1c055..eaad30ee 100644 --- a/internal/thinking/provider/openai/apply.go +++ b/internal/thinking/provider/openai/apply.go @@ -41,18 +41,11 @@ func init() { // "reasoning_effort": "high" // } func (a *Applier) Apply(body []byte, config thinking.ThinkingConfig, modelInfo *registry.ModelInfo) ([]byte, error) { - if modelInfo == nil { + if thinking.IsUserDefinedModel(modelInfo) { return applyCompatibleOpenAI(body, config) } if modelInfo.Thinking == nil { - if modelInfo.Type == "" { - modelID := modelInfo.ID - if modelID == "" { - modelID = "unknown" - } - return nil, thinking.NewThinkingErrorWithModel(thinking.ErrThinkingNotSupported, "thinking not supported for this model", modelID) - } - return applyCompatibleOpenAI(body, config) + return body, nil } // Only handle ModeLevel and ModeNone; other modes pass through unchanged. diff --git a/internal/thinking/provider/openai/apply_test.go b/internal/thinking/provider/openai/apply_test.go index 5be01e4e..1e348d9e 100644 --- a/internal/thinking/provider/openai/apply_test.go +++ b/internal/thinking/provider/openai/apply_test.go @@ -57,22 +57,13 @@ func TestApplyNilModelInfo(t *testing.T) { func TestApplyMissingThinkingSupport(t *testing.T) { applier := NewApplier() modelInfo := ®istry.ModelInfo{ID: "gpt-5.2"} - got, err := applier.Apply([]byte(`{"model":"gpt-5.2"}`), thinking.ThinkingConfig{}, modelInfo) - if err == nil { - t.Fatalf("expected error, got nil") + body := []byte(`{"model":"gpt-5.2"}`) + got, err := applier.Apply(body, thinking.ThinkingConfig{}, modelInfo) + if err != nil { + t.Fatalf("expected nil error, got %v", err) } - if got != nil { - t.Fatalf("expected nil body on error, got %s", string(got)) - } - thinkingErr, ok := err.(*thinking.ThinkingError) - if !ok { - t.Fatalf("expected ThinkingError, got %T", err) - } - if thinkingErr.Code != thinking.ErrThinkingNotSupported { - t.Fatalf("expected code %s, got %s", thinking.ErrThinkingNotSupported, thinkingErr.Code) - } - if thinkingErr.Model != "gpt-5.2" { - t.Fatalf("expected model gpt-5.2, got %s", thinkingErr.Model) + if string(got) != string(body) { + t.Fatalf("expected body unchanged, got %s", string(got)) } }