diff --git a/internal/api/server_routes.go b/internal/api/server_routes.go index 5a48162cef..7bf7ff70b9 100644 --- a/internal/api/server_routes.go +++ b/internal/api/server_routes.go @@ -1077,12 +1077,19 @@ func formatHomeGeminiModel(entry homeModelEntry) map[string]any { if displayName == "" { displayName = entry.id } - return map[string]any{ + model := map[string]any{ "name": name, "displayName": displayName, "description": displayName, "supportedGenerationMethods": []string{"generateContent"}, } + if entry.contextLength > 0 { + model["inputTokenLimit"] = entry.contextLength + } + if entry.maxCompletionTokens > 0 { + model["outputTokenLimit"] = entry.maxCompletionTokens + } + return model } func homeGeminiModelMatches(entry homeModelEntry, action string) bool { @@ -1176,10 +1183,17 @@ func decodeHomeModels(raw []byte) ([]homeModelEntry, error) { return nil, fmt.Errorf("home models payload contains an error envelope") } + sections := make([]string, 0, len(bySection)) + for section := range bySection { + sections = append(sections, section) + } + sort.Strings(sections) + indexByID := make(map[string]int) out := make([]homeModelEntry, 0, 256) hasModelRecords := false - for section, models := range bySection { + for _, section := range sections { + models := bySection[section] if models == nil { return nil, fmt.Errorf("home models section %q is null", section) } @@ -1201,12 +1215,7 @@ func decodeHomeModels(raw []byte) ([]homeModelEntry, error) { Provider: provider, NativeCapabilities: nativeCapabilities, } - if index, ok := indexByID[id]; ok { - out[index].providers = appendUniqueHomeProvider(out[index].providers, provider) - out[index].nativeCapabilityRoutes = append(out[index].nativeCapabilityRoutes, route) - continue - } - + created := homeModelInt64Value(model, "created") ownedBy, _ := model["owned_by"].(string) ownedBy = strings.TrimSpace(ownedBy) displayName, _ := model["display_name"].(string) @@ -1215,17 +1224,47 @@ func decodeHomeModels(raw []byte) ([]homeModelEntry, error) { displayName, _ = model["displayName"].(string) displayName = strings.TrimSpace(displayName) } + contextLength := int(homeModelInt64Value(model, "context_length", "contextLength", "inputTokenLimit", "max_input_tokens")) + maxContextLength := int(homeModelInt64Value(model, "max_context_length", "maxContextLength")) + maxCompletionTokens := int(homeModelInt64Value(model, "max_completion_tokens", "maxCompletionTokens", "outputTokenLimit", "max_tokens")) thinking := homeModelThinkingSupport(model) + if index, ok := indexByID[id]; ok { + out[index].providers = appendUniqueHomeProvider(out[index].providers, provider) + out[index].nativeCapabilityRoutes = append(out[index].nativeCapabilityRoutes, route) + if out[index].created <= 0 && created > 0 { + out[index].created = created + } + if out[index].ownedBy == "" && ownedBy != "" { + out[index].ownedBy = ownedBy + } + if out[index].displayName == "" && displayName != "" { + out[index].displayName = displayName + } + if out[index].contextLength <= 0 && contextLength > 0 { + out[index].contextLength = contextLength + } + if out[index].maxContextLength <= 0 && maxContextLength > 0 { + out[index].maxContextLength = maxContextLength + } + if out[index].maxCompletionTokens <= 0 && maxCompletionTokens > 0 { + out[index].maxCompletionTokens = maxCompletionTokens + } + if out[index].thinking == nil && thinking != nil { + out[index].thinking = thinking + } + continue + } + indexByID[id] = len(out) out = append(out, homeModelEntry{ id: id, - created: homeModelInt64Value(model, "created"), + created: created, ownedBy: ownedBy, displayName: displayName, - contextLength: int(homeModelInt64Value(model, "context_length", "contextLength", "inputTokenLimit", "max_input_tokens")), - maxContextLength: int(homeModelInt64Value(model, "max_context_length", "maxContextLength")), - maxCompletionTokens: int(homeModelInt64Value(model, "max_completion_tokens", "maxCompletionTokens", "outputTokenLimit", "max_tokens")), + contextLength: contextLength, + maxContextLength: maxContextLength, + maxCompletionTokens: maxCompletionTokens, thinking: thinking, providers: appendUniqueHomeProvider(nil, provider), nativeCapabilityRoutes: []registry.NativeCapabilityRoute{route}, diff --git a/internal/api/server_test.go b/internal/api/server_test.go index 69e6919f44..0a03008484 100644 --- a/internal/api/server_test.go +++ b/internal/api/server_test.go @@ -2814,6 +2814,99 @@ func TestDecodeHomeModelsKeepsTokenMetadata(t *testing.T) { } } +func TestHomeGeminiModelsIncludeTokenLimits(t *testing.T) { + withLimits := formatHomeGeminiModel(homeModelEntry{ + id: "gemini-3.8-flash", + displayName: "Gemini 3.8 Flash", + contextLength: 1048576, + maxCompletionTokens: 65536, + }) + if got := withLimits["name"]; got != "models/gemini-3.8-flash" { + t.Fatalf("name = %v, want models/gemini-3.8-flash", got) + } + if got := withLimits["displayName"]; got != "Gemini 3.8 Flash" { + t.Fatalf("displayName = %v, want Gemini 3.8 Flash", got) + } + if got := withLimits["inputTokenLimit"]; got != 1048576 { + t.Fatalf("inputTokenLimit = %v, want 1048576", got) + } + if got := withLimits["outputTokenLimit"]; got != 65536 { + t.Fatalf("outputTokenLimit = %v, want 65536", got) + } + + withoutLimits := formatHomeGeminiModel(homeModelEntry{ + id: "gemini-no-limits", + }) + if _, ok := withoutLimits["inputTokenLimit"]; ok { + t.Fatalf("expected inputTokenLimit to be omitted when unknown, got %v", withoutLimits) + } + if _, ok := withoutLimits["outputTokenLimit"]; ok { + t.Fatalf("expected outputTokenLimit to be omitted when unknown, got %v", withoutLimits) + } + + server := newTestServer(t) + server.cfg.Home.Enabled = true + engine := gin.New() + engine.GET("/v1beta/models", server.geminiModelsHandler(nil)) + engine.GET("/v1beta/models/*action", server.geminiGetHandler(nil)) + + client := newHomeCatalogClient(t, `{ + "a-custom": [ + { + "id": "gemini-3.8-flash" + } + ], + "gemini": [ + { + "id": "gemini-3.8-flash", + "display_name": "Gemini 3.8 Flash", + "context_length": 1048576, + "max_completion_tokens": 65536 + } + ] + }`) + previousHome := home.Current() + home.SetCurrent(client) + t.Cleanup(func() { home.SetCurrent(previousHome) }) + + listRecorder := httptest.NewRecorder() + engine.ServeHTTP(listRecorder, httptest.NewRequest(http.MethodGet, "/v1beta/models", nil)) + if listRecorder.Code != http.StatusOK { + t.Fatalf("GET /v1beta/models status = %d; body=%s", listRecorder.Code, listRecorder.Body.String()) + } + var listPayload struct { + Models []map[string]any `json:"models"` + } + if errUnmarshal := json.Unmarshal(listRecorder.Body.Bytes(), &listPayload); errUnmarshal != nil { + t.Fatalf("unmarshal /v1beta/models response: %v", errUnmarshal) + } + if len(listPayload.Models) != 1 { + t.Fatalf("models len = %d, want 1; body=%s", len(listPayload.Models), listRecorder.Body.String()) + } + if got := homeModelInt64Value(listPayload.Models[0], "inputTokenLimit"); got != 1048576 { + t.Fatalf("list inputTokenLimit = %d, want 1048576; model=%v", got, listPayload.Models[0]) + } + if got := homeModelInt64Value(listPayload.Models[0], "outputTokenLimit"); got != 65536 { + t.Fatalf("list outputTokenLimit = %d, want 65536; model=%v", got, listPayload.Models[0]) + } + + getRecorder := httptest.NewRecorder() + engine.ServeHTTP(getRecorder, httptest.NewRequest(http.MethodGet, "/v1beta/models/gemini-3.8-flash", nil)) + if getRecorder.Code != http.StatusOK { + t.Fatalf("GET /v1beta/models/gemini-3.8-flash status = %d; body=%s", getRecorder.Code, getRecorder.Body.String()) + } + var getPayload map[string]any + if errUnmarshal := json.Unmarshal(getRecorder.Body.Bytes(), &getPayload); errUnmarshal != nil { + t.Fatalf("unmarshal /v1beta/models/gemini-3.8-flash response: %v", errUnmarshal) + } + if got := homeModelInt64Value(getPayload, "inputTokenLimit"); got != 1048576 { + t.Fatalf("get inputTokenLimit = %d, want 1048576; model=%v", got, getPayload) + } + if got := homeModelInt64Value(getPayload, "outputTokenLimit"); got != 65536 { + t.Fatalf("get outputTokenLimit = %d, want 65536; model=%v", got, getPayload) + } +} + func TestHomeCodexModels_MaxContextLength(t *testing.T) { entries, errDecode := decodeHomeModels([]byte(`{ "codex": [ diff --git a/internal/registry/model_registry.go b/internal/registry/model_registry.go index c64a5048b5..0c0a84319d 100644 --- a/internal/registry/model_registry.go +++ b/internal/registry/model_registry.go @@ -1803,11 +1803,19 @@ func (r *ModelRegistry) convertModelToMap(model *ModelInfo, handlerType string) if model.Description != "" { result["description"] = model.Description } - if model.InputTokenLimit > 0 { - result["inputTokenLimit"] = model.InputTokenLimit + inputTokenLimit := model.InputTokenLimit + if inputTokenLimit <= 0 { + inputTokenLimit = model.ContextLength } - if model.OutputTokenLimit > 0 { - result["outputTokenLimit"] = model.OutputTokenLimit + if inputTokenLimit > 0 { + result["inputTokenLimit"] = inputTokenLimit + } + outputTokenLimit := model.OutputTokenLimit + if outputTokenLimit <= 0 { + outputTokenLimit = model.MaxCompletionTokens + } + if outputTokenLimit > 0 { + result["outputTokenLimit"] = outputTokenLimit } if len(model.SupportedGenerationMethods) > 0 { result["supportedGenerationMethods"] = append([]string(nil), model.SupportedGenerationMethods...) diff --git a/internal/registry/model_registry_cache_test.go b/internal/registry/model_registry_cache_test.go index fb49e1f4ac..3a6011b4b5 100644 --- a/internal/registry/model_registry_cache_test.go +++ b/internal/registry/model_registry_cache_test.go @@ -68,6 +68,70 @@ func TestGetAvailableModelsClaudeIncludesTokenLimits(t *testing.T) { } } +func TestGetAvailableModelsGeminiIncludesTokenLimits(t *testing.T) { + r := newTestModelRegistry() + r.RegisterClient("client-1", "gemini", []*ModelInfo{ + { + ID: "gemini-3.8-flash", + Name: "models/gemini-3.8-flash", + DisplayName: "Gemini 3.8 Flash", + ContextLength: 1048576, + MaxCompletionTokens: 65536, + }, + { + ID: "gemini-2.5-pro", + Name: "models/gemini-2.5-pro", + DisplayName: "Gemini 2.5 Pro", + InputTokenLimit: 1048576, + OutputTokenLimit: 65536, + }, + { + ID: "gemini-no-limits", + Name: "models/gemini-no-limits", + }, + }) + + models := r.GetAvailableModels("gemini") + byName := make(map[string]map[string]any, len(models)) + for _, m := range models { + name, _ := m["name"].(string) + byName[name] = m + } + + fromContext, ok := byName["models/gemini-3.8-flash"] + if !ok { + t.Fatalf("expected models/gemini-3.8-flash in available models, got %v", byName) + } + if got := fromContext["inputTokenLimit"]; got != 1048576 { + t.Fatalf("gemini-3.8-flash inputTokenLimit = %v, want 1048576", got) + } + if got := fromContext["outputTokenLimit"]; got != 65536 { + t.Fatalf("gemini-3.8-flash outputTokenLimit = %v, want 65536", got) + } + + fromExplicit, ok := byName["models/gemini-2.5-pro"] + if !ok { + t.Fatalf("expected models/gemini-2.5-pro in available models, got %v", byName) + } + if got := fromExplicit["inputTokenLimit"]; got != 1048576 { + t.Fatalf("gemini-2.5-pro inputTokenLimit = %v, want 1048576", got) + } + if got := fromExplicit["outputTokenLimit"]; got != 65536 { + t.Fatalf("gemini-2.5-pro outputTokenLimit = %v, want 65536", got) + } + + noLimits, ok := byName["models/gemini-no-limits"] + if !ok { + t.Fatalf("expected models/gemini-no-limits in available models, got %v", byName) + } + if _, exists := noLimits["inputTokenLimit"]; exists { + t.Fatalf("expected inputTokenLimit to be omitted when unknown, got %v", noLimits) + } + if _, exists := noLimits["outputTokenLimit"]; exists { + t.Fatalf("expected outputTokenLimit to be omitted when unknown, got %v", noLimits) + } +} + func TestGetAvailableModelsInvalidatesCacheOnRegistryChanges(t *testing.T) { r := newTestModelRegistry() r.RegisterClient("client-1", "OpenAI", []*ModelInfo{{ID: "m1", OwnedBy: "team-a", DisplayName: "Model One"}})