Merge pull request #192 from cnjack/fix/langfuse-usage-and-grok-image-support

fix: record Langfuse usage details and inherit Grok image support

Author: cnjack

internal/model/chatmodel.go

@@ -42,6 +42,7 @@ type TokenUsage struct {
 	lastCached     int64
 	lastReasoning  int64
 	lastCacheWrite int64
+	lastModel      string
 	// cacheSeen is set (sticky) once the provider returns a prompt_tokens_details
 	// object, so CacheObserved can report "caching supported" even on a 0-hit
 	// turn. Cleared by Reset (a session boundary), never by ResetContext.
@@ -212,6 +213,7 @@ func (t *TokenUsage) Reset() {
 	atomic.StoreInt64(&t.turnBaseCached, 0)
 	t.mu.Lock()
 	t.byModel = nil
+	t.lastModel = ""
 	t.mu.Unlock()
 }
 
@@ -269,9 +271,17 @@ func (t *TokenUsage) AddByModel(model string, prompt, completion, total int) {
 		t.byModel = make(map[string]int64)
 	}
 	t.byModel[model] += int64(total)
+	t.lastModel = model
 	t.mu.Unlock()
 }
 
+// GetLastModel returns the model name of the last recorded API call.
+func (t *TokenUsage) GetLastModel() string {
+	t.mu.RLock()
+	defer t.mu.RUnlock()
+	return t.lastModel
+}
+
 // GetByModel returns a snapshot of per-model token totals.
 func (t *TokenUsage) GetByModel() map[string]int64 {
 	t.mu.RLock()

internal/model/registry.go

@@ -272,11 +272,165 @@ func (r *ModelRegistry) MergeConfigProviders(providers map[string]*config.Provid
 			if cm.Context > 0 {
 				rm.Limit = &ModelLimit{Context: cm.Context}
 			}
+			// Managed rows persist only what the live catalog last wrote. A newer
+			// Grok/Codex ID often lands before the static registry lists it, so
+			// omitted capabilities inherit from the closest baked-in sibling
+			// (grok-4.6 ← grok-4.5) instead of rendering as text-only.
+			if cm.Managed {
+				applyRelatedManagedModelDefaults(rm, RelatedRegistryModel(prov, cm.ID))
+			}
 			prov.Models[cm.ID] = rm
 		}
 	}
 }
 
+// RelatedRegistryModel returns a baked-in sibling that shares a versioned
+// family key with modelID. grok-4.6 matches grok-4.5; grok-imagine-image
+// does not. Used when a live managed catalog lists an ID the static
+// registry has not been updated to include yet.
+//
+// Candidates come from immutable generatedProviders, never from the live
+// provider.Models map, so an earlier custom row cannot pollute inheritance.
+func RelatedRegistryModel(provider *RegistryProvider, modelID string) *RegistryModel {
+	if provider == nil {
+		return nil
+	}
+	baked := generatedProviders[provider.ID]
+	if baked == nil {
+		return nil
+	}
+	return relatedBakedInModel(baked.Models, modelID)
+}
+
+func relatedBakedInModel(models map[string]*RegistryModel, modelID string) *RegistryModel {
+	key := managedModelFamilyKey(modelID)
+	if key == "" || models == nil {
+		return nil
+	}
+	targetParts := managedModelVersionParts(modelID)
+	var related *RegistryModel
+	var relatedDist int
+	for _, candidate := range models {
+		if candidate == nil || candidate.ID == modelID || managedModelFamilyKey(candidate.ID) != key {
+			continue
+		}
+		dist := managedModelVersionDistance(managedModelVersionParts(candidate.ID), targetParts)
+		if !preferRelatedModel(related, relatedDist, candidate, dist) {
+			continue
+		}
+		related = candidate
+		relatedDist = dist
+	}
+	return related
+}
+
+func preferRelatedModel(current *RegistryModel, currentDist int, candidate *RegistryModel, candidateDist int) bool {
+	if current == nil {
+		return true
+	}
+	if candidateDist != currentDist {
+		return candidateDist < currentDist
+	}
+	if candidate.DefaultEnabled != current.DefaultEnabled {
+		return candidate.DefaultEnabled
+	}
+	return candidate.ID < current.ID
+}
+
+func managedModelFamilyKey(id string) string {
+	id = strings.ToLower(strings.TrimSpace(id))
+	for i := 0; i < len(id); i++ {
+		if id[i] != '-' || i+1 >= len(id) || id[i+1] < '0' || id[i+1] > '9' {
+			continue
+		}
+		end := i + 1
+		for end < len(id) && id[end] >= '0' && id[end] <= '9' {
+			end++
+		}
+		return id[:end]
+	}
+	return ""
+}
+
+func managedModelVersionParts(id string) []int {
+	id = strings.ToLower(strings.TrimSpace(id))
+	i := 0
+	for i < len(id) {
+		if id[i] == '-' && i+1 < len(id) && id[i+1] >= '0' && id[i+1] <= '9' {
+			i++
+			break
+		}
+		i++
+	}
+	var parts []int
+	for i < len(id) {
+		if id[i] < '0' || id[i] > '9' {
+			if id[i] == '.' || id[i] == '-' {
+				i++
+				continue
+			}
+			break
+		}
+		n := 0
+		for i < len(id) && id[i] >= '0' && id[i] <= '9' {
+			n = n*10 + int(id[i]-'0')
+			i++
+		}
+		parts = append(parts, n)
+	}
+	return parts
+}
+
+func managedModelVersionDistance(a, b []int) int {
+	n := len(a)
+	if len(b) > n {
+		n = len(b)
+	}
+	dist := 0
+	weights := [...]int{1_000_000, 1_000, 1}
+	for i := 0; i < n; i++ {
+		av, bv := 0, 0
+		if i < len(a) {
+			av = a[i]
+		}
+		if i < len(b) {
+			bv = b[i]
+		}
+		d := av - bv
+		if d < 0 {
+			d = -d
+		}
+		w := 1
+		if i < len(weights) {
+			w = weights[i]
+		}
+		dist += d * w
+	}
+	return dist
+}
+
+func applyRelatedManagedModelDefaults(dst *RegistryModel, related *RegistryModel) {
+	if dst == nil || related == nil {
+		return
+	}
+	if related.Attachment {
+		dst.Attachment = true
+		if dst.Modalities == nil && related.Modalities != nil {
+			dst.Modalities = deepCopyModel(related).Modalities
+		}
+	}
+	if !dst.Reasoning && related.Reasoning {
+		dst.Reasoning = true
+		if len(dst.ReasoningOptions) == 0 && len(related.ReasoningOptions) > 0 {
+			dst.ReasoningOptions = append([]ReasoningOption(nil), related.ReasoningOptions...)
+		}
+	}
+	if dst.Limit == nil && related.Limit != nil && related.Limit.Context > 0 {
+		limitCopy := *related.Limit
+		dst.Limit = &limitCopy
+	}
+}
+
 // Load returns the provider/model data.
 func (r *ModelRegistry) Load() (map[string]*RegistryProvider, error) {
 	return r.providers, nil

internal/model/registry_test.go

@@ -2,6 +2,8 @@ package model
 
 import (
 	"testing"
+
+	"github.com/cnjack/jcode/internal/config"
 )
 
 // TestLookupModel tests model lookup from the generated registry.
@@ -83,3 +85,88 @@ func TestHasProvider(t *testing.T) {
 		t.Error("Expected HasProvider to return false for non-existent provider")
 	}
 }
+
+func TestMergeConfigProvidersInheritsGrokFamilyImageSupport(t *testing.T) {
+	registry := NewModelRegistryWithConfig(&config.Config{
+		Providers: map[string]*config.ProviderConfig{
+			"xai": {
+				CustomModels: []config.CustomModelConfig{{
+					ID: "grok-4.6", Name: "Grok 4.6", ToolCall: true, Managed: true,
+				}},
+			},
+		},
+	})
+	_, got, ok := registry.LookupModel("xai", "grok-4.6")
+	if !ok || got == nil {
+		t.Fatal("managed grok-4.6 was not merged")
+	}
+	if !got.Attachment || !got.SupportsImageInput() {
+		t.Fatalf("grok-4.6 should inherit grok-4.5 image support: %#v", got)
+	}
+	if !got.Reasoning || got.Limit == nil || got.Limit.Context <= 0 {
+		t.Fatalf("grok-4.6 should inherit grok-4.5 reasoning/context: %#v", got)
+	}
+}
+
+func TestRelatedRegistryModelDoesNotMatchImagineIDs(t *testing.T) {
+	provider := NewModelRegistry().GetProvider("xai")
+	if RelatedRegistryModel(provider, "grok-imagine-image") != nil {
+		t.Fatal("image-generation ids must not inherit chat-family metadata")
+	}
+	related := RelatedRegistryModel(provider, "grok-4.6")
+	if related == nil || related.ID != "grok-4.5" {
+		t.Fatalf("related grok-4.6 = %#v", related)
+	}
+}
+
+func TestRelatedRegistryModelIgnoresLiveCustomSiblings(t *testing.T) {
+	registry := NewModelRegistry()
+	provider := registry.GetProvider("xai")
+	if provider == nil {
+		t.Fatal("xai missing")
+	}
+	provider.Models["grok-4.7"] = &RegistryModel{
+		ID: "grok-4.7", Name: "Custom Grok 4.7", DefaultEnabled: true,
+	}
+	related := RelatedRegistryModel(provider, "grok-4.6")
+	if related == nil || related.ID != "grok-4.5" {
+		t.Fatalf("related grok-4.6 should stay on baked-in grok-4.5, got %#v", related)
+	}
+}
+
+func TestRelatedBakedInModelPicksClosestEnabledSibling(t *testing.T) {
+	models := map[string]*RegistryModel{
+		"grok-4.1": {ID: "grok-4.1", DefaultEnabled: true},
+		"grok-4.5": {ID: "grok-4.5", DefaultEnabled: true},
+		"grok-4.8": {ID: "grok-4.8", DefaultEnabled: false},
+	}
+	got := relatedBakedInModel(models, "grok-4.6")
+	if got == nil || got.ID != "grok-4.5" {
+		t.Fatalf("closest sibling for grok-4.6 = %#v", got)
+	}
+
+	models["grok-4.5-preview"] = &RegistryModel{ID: "grok-4.5-preview", DefaultEnabled: true}
+	got = relatedBakedInModel(models, "grok-4.5-new")
+	if got == nil || got.ID != "grok-4.5" {
+		t.Fatalf("tied versions should pick lexicographically first enabled id, got %#v", got)
+	}
+}
+
+func TestMergeConfigProvidersCopiesModalitiesWhenLiveAttachmentIsSet(t *testing.T) {
+	registry := NewModelRegistryWithConfig(&config.Config{
+		Providers: map[string]*config.ProviderConfig{
+			"xai": {
+				CustomModels: []config.CustomModelConfig{{
+					ID: "grok-4.6", Name: "Grok 4.6", ToolCall: true, Managed: true, Attachment: true,
+				}},
+			},
+		},
+	})
+	_, got, ok := registry.LookupModel("xai", "grok-4.6")
+	if !ok || got == nil {
+		t.Fatal("managed grok-4.6 was not merged")
+	}
+	if !got.Attachment || !got.SupportsImageInput() {
+		t.Fatalf("live attachment should still copy related modalities: %#v", got)
+	}
+}

internal/model/token_usage_test.go

@@ -1,6 +1,10 @@
 package model
 
-import "testing"
+import (
+	"testing"
+
+	openai "github.com/sashabaranov/go-openai"
+)
 
 func TestTokenUsage_AddAndGetFull(t *testing.T) {
 	u := &TokenUsage{}
@@ -136,6 +140,9 @@ func TestTokenUsage_Reset(t *testing.T) {
 	u := &TokenUsage{}
 	u.Add(AddParams{Prompt: 100, Completion: 20, Total: 120, Cached: 80, Reasoning: 5})
 	u.AddByModel("m", 100, 20, 120)
+	if u.GetLastModel() != "m" {
+		t.Errorf("GetLastModel() = %q, want m", u.GetLastModel())
+	}
 	u.Reset()
 	if got := u.GetFull(); got.PromptTokens != 0 || got.CachedTokens != 0 || got.CallCount != 0 {
 		t.Errorf("after Reset GetFull() = %+v, want zero", got)
@@ -146,4 +153,34 @@ func TestTokenUsage_Reset(t *testing.T) {
 	if u.CacheObserved() {
 		t.Errorf("after Reset CacheObserved() should be false")
 	}
+	if u.GetLastModel() != "" {
+		t.Errorf("after Reset GetLastModel() = %q, want empty", u.GetLastModel())
+	}
+}
+
+func TestExtractUsage_ReadsAPITotalAndDetails(t *testing.T) {
+	got := extractUsage(openai.Usage{
+		PromptTokens:     19,
+		CompletionTokens: 10,
+		TotalTokens:      29,
+		PromptTokensDetails: &openai.PromptTokensDetails{
+			CachedTokens: 4,
+		},
+		CompletionTokensDetails: &openai.CompletionTokensDetails{
+			ReasoningTokens: 3,
+		},
+	})
+	if got.Prompt != 19 || got.Completion != 10 || got.Total != 29 || got.Cached != 4 || got.Reasoning != 3 {
+		t.Errorf("extractUsage() = %+v, want prompt/completion/total from API and cached/reasoning from details", got)
+	}
+	if !got.CacheDetailsPresent {
+		t.Error("CacheDetailsPresent = false, want true when prompt_tokens_details is present")
+	}
+}
+
+func TestExtractUsage_DerivesTotalWhenAPIOmitsIt(t *testing.T) {
+	got := extractUsage(openai.Usage{PromptTokens: 19, CompletionTokens: 10})
+	if got.Total != 29 {
+		t.Errorf("extractUsage() Total = %d, want 19+10 when API total_tokens is 0", got.Total)
+	}
 }