fix(gemini): populate token limits in Gemini model catalog responses

- Populate inputTokenLimit and outputTokenLimit in Home Gemini model responses when positive token limits are available, and backfill missing metadata across duplicate Home catalog entries.
- Fall back to ContextLength and MaxCompletionTokens when converting registry models for the Gemini handler so catalog entries using generic limit fields expose Gemini token limits.

Closes: #6914
This commit is contained in:
Luis Pater
2026-10-23 17:09:18 +08:00
parent e5434f65c8
commit 83cb1b95eb
4 changed files with 220 additions and 16 deletions

View File

@@ -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},

View File

@@ -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": [

View File

@@ -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...)

View File

@@ -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"}})