From 98c98d66be1a17ae3b30ccd060848f1944f3765b Mon Sep 17 00:00:00 2001 From: Luis Pater Date: Sat, 15 Aug 2026 03:49:56 +0800 Subject: [PATCH] fix(codex): cache multi-agent spawn-agent model data and invalidate on updates - Add registry generation tracking and a Codex catalog revision accessor to provide stable cache keys. - Cache parsed model templates and rendered spawn-agent markdown, and reuse them across requests when revision/generation are unchanged. - Invalidate/recompute caches when model registrations or catalog content change, and separate rewrite handling for spawn tool descriptions vs. message encryption stripping. Closes: #4967 --- .../optimize_multi_agent_v2.go | 187 ++++++++++++++++-- .../optimize_multi_agent_v2_test.go | 168 ++++++++++++++++ internal/registry/codex_client_models.go | 7 + internal/registry/model_registry.go | 10 + 4 files changed, 353 insertions(+), 19 deletions(-) diff --git a/internal/client/codex/optimize-multi-agent-v2/optimize_multi_agent_v2.go b/internal/client/codex/optimize-multi-agent-v2/optimize_multi_agent_v2.go index b49bcd7e..69046d4e 100644 --- a/internal/client/codex/optimize-multi-agent-v2/optimize_multi_agent_v2.go +++ b/internal/client/codex/optimize-multi-agent-v2/optimize_multi_agent_v2.go @@ -9,6 +9,7 @@ import ( "net/url" "sort" "strings" + "sync" "github.com/gin-gonic/gin" "github.com/router-for-me/CLIProxyAPI/v7/internal/config" @@ -86,14 +87,22 @@ func PrepareCodexMultiAgentV2Tools(ctx context.Context, headers http.Header, pay return payload, false } - updated := removeCodexCollaborationMessageEncryption(payload, codexCollaborationMessageToolPaths(payload)) - toolPaths := codexSpawnAgentToolPaths(updated) - if len(toolPaths) == 0 || hasCodexOptimizedCollaborationConflict(updated) { - return updated, true + toolPaths := codexSpawnAgentToolPaths(payload) + messageToolPaths := codexCollaborationMessageToolPaths(payload) + if len(toolPaths) == 0 && len(messageToolPaths) == 0 { + return payload, true + } + if hasCodexOptimizedCollaborationConflict(payload) { + return removeCodexCollaborationMessageEncryption(payload, messageToolPaths), true + } + + var models []codexSpawnAgentModel + var formattedMarkdown string + if len(toolPaths) > 0 { + models, formattedMarkdown = codexSpawnAgentModelsAndMarkdownForRequest(ctx, headers, homeEnabled) } - models := codexSpawnAgentModelsForRequest(ctx, headers, homeEnabled) - updated = rewriteCodexSpawnAgentTools(updated, toolPaths, models) + updated := rewriteCodexCollaborationTools(payload, messageToolPaths, toolPaths, models, formattedMarkdown) return updated, true } @@ -179,14 +188,128 @@ func isCodexMultiAgentClient(userAgent string) bool { return IsCodexClientUserAgent(userAgent) } -func codexSpawnAgentModelsForRequest(ctx context.Context, headers http.Header, homeEnabled bool) []codexSpawnAgentModel { - availableModels := registry.GetGlobalRegistry().GetAvailableModels("openai") +var ( + codexCatalogTemplatesMu sync.RWMutex + codexCatalogTemplatesLoaded bool + codexCatalogTemplatesRevision uint64 + codexCatalogTemplates map[string]map[string]any + codexCatalogDefaultTemplate map[string]any + + codexSpawnAgentCacheMu sync.RWMutex + codexSpawnAgentCacheRevision uint64 + codexSpawnAgentCacheGeneration uint64 + codexSpawnAgentCachedModels []codexSpawnAgentModel + codexSpawnAgentCachedMarkdown string +) + +func loadCodexCatalogTemplates() (map[string]map[string]any, map[string]any, uint64, error) { + currentRevision := registry.GetCodexClientModelsRevision() + + codexCatalogTemplatesMu.RLock() + if codexCatalogTemplatesLoaded && codexCatalogTemplatesRevision == currentRevision { + templates := codexCatalogTemplates + defaultTemplate := codexCatalogDefaultTemplate + codexCatalogTemplatesMu.RUnlock() + return templates, defaultTemplate, currentRevision, nil + } + codexCatalogTemplatesMu.RUnlock() + + codexCatalogTemplatesMu.Lock() + defer codexCatalogTemplatesMu.Unlock() + if codexCatalogTemplatesLoaded && codexCatalogTemplatesRevision == currentRevision { + return codexCatalogTemplates, codexCatalogDefaultTemplate, currentRevision, nil + } + + raw, revision := registry.GetCodexClientModelsSnapshot() + + var catalog codexClientModelsCatalog + errUnmarshal := json.Unmarshal(raw, &catalog) + if errUnmarshal != nil || len(catalog.Models) == 0 { + codexCatalogTemplatesLoaded = true + codexCatalogTemplatesRevision = revision + codexCatalogTemplates = nil + codexCatalogDefaultTemplate = nil + return nil, nil, revision, errUnmarshal + } + + templates := make(map[string]map[string]any, len(catalog.Models)) + var defaultTemplate map[string]any + for _, model := range catalog.Models { + modelID := mapString(model, "slug") + if modelID == "" { + continue + } + templates[modelID] = model + if modelID == "gpt-5.5" { + defaultTemplate = model + } + } + + codexCatalogTemplatesLoaded = true + codexCatalogTemplatesRevision = revision + codexCatalogTemplates = templates + codexCatalogDefaultTemplate = defaultTemplate + return templates, defaultTemplate, revision, nil +} + +func codexSpawnAgentModelsAndMarkdownForRequest(ctx context.Context, headers http.Header, homeEnabled bool) ([]codexSpawnAgentModel, string) { if homeEnabled { - availableModels = codexHomeAvailableModels(ctx, headers) + availableModels := codexHomeAvailableModels(ctx, headers) + templates, defaultTemplate, _, errLoad := loadCodexCatalogTemplates() + if errLoad != nil || defaultTemplate == nil { + return nil, "" + } + models := codexSpawnAgentModelsFromTemplates(availableModels, templates, defaultTemplate, func(modelID string) *registry.ModelInfo { + return registry.LookupModelInfo(modelID) + }) + formatted := formatCodexSpawnAgentModels(models) + return models, formatted } - return codexSpawnAgentModelsFromSources(availableModels, registry.GetCodexClientModelsJSON(), func(modelID string) *registry.ModelInfo { + + currentRevision := registry.GetCodexClientModelsRevision() + currentGeneration := registry.GetGlobalRegistry().GetGeneration() + + codexSpawnAgentCacheMu.RLock() + if codexSpawnAgentCachedModels != nil && codexSpawnAgentCacheRevision == currentRevision && codexSpawnAgentCacheGeneration == currentGeneration { + models := codexSpawnAgentCachedModels + markdown := codexSpawnAgentCachedMarkdown + codexSpawnAgentCacheMu.RUnlock() + return models, markdown + } + codexSpawnAgentCacheMu.RUnlock() + + templates, defaultTemplate, _, errLoad := loadCodexCatalogTemplates() + if errLoad != nil || defaultTemplate == nil { + return nil, "" + } + + availableModels := registry.GetGlobalRegistry().GetAvailableModels("openai") + lookup := func(modelID string) *registry.ModelInfo { return registry.LookupModelInfo(modelID) - }) + } + models := codexSpawnAgentModelsFromTemplates(availableModels, templates, defaultTemplate, lookup) + formatted := formatCodexSpawnAgentModels(models) + + codexSpawnAgentCacheMu.Lock() + if currentRevision == registry.GetCodexClientModelsRevision() && currentGeneration == registry.GetGlobalRegistry().GetGeneration() { + codexSpawnAgentCacheRevision = currentRevision + codexSpawnAgentCacheGeneration = currentGeneration + codexSpawnAgentCachedModels = models + codexSpawnAgentCachedMarkdown = formatted + } + codexSpawnAgentCacheMu.Unlock() + + return models, formatted +} + +func codexSpawnAgentModelsForRequest(ctx context.Context, headers http.Header, homeEnabled bool) []codexSpawnAgentModel { + models, _ := codexSpawnAgentModelsAndMarkdownForRequest(ctx, headers, homeEnabled) + return models +} + +func formatCodexSpawnAgentModelsForRequest(ctx context.Context, headers http.Header, homeEnabled bool) string { + _, formatted := codexSpawnAgentModelsAndMarkdownForRequest(ctx, headers, homeEnabled) + return formatted } func codexHomeAvailableModels(ctx context.Context, headers http.Header) []map[string]any { @@ -272,6 +395,14 @@ func codexSpawnAgentModelsFromSources(availableModels []map[string]any, catalogJ return nil } + return codexSpawnAgentModelsFromTemplates(availableModels, templates, defaultTemplate, lookupModel) +} + +func codexSpawnAgentModelsFromTemplates(availableModels []map[string]any, templates map[string]map[string]any, defaultTemplate map[string]any, lookupModel func(string) *registry.ModelInfo) []codexSpawnAgentModel { + if defaultTemplate == nil { + return nil + } + seen := make(map[string]struct{}, len(availableModels)) templateModels := make([]codexSpawnAgentModel, 0, len(availableModels)) synthesizedModels := make([]codexSpawnAgentModel, 0, len(availableModels)) @@ -454,12 +585,18 @@ func rewriteCodexSpawnAgentDescription(payload []byte, models []codexSpawnAgentM } func rewriteCodexSpawnAgentTools(payload []byte, toolPaths []string, models []codexSpawnAgentModel) []byte { - if len(toolPaths) == 0 { + return rewriteCodexCollaborationTools(payload, toolPaths, toolPaths, models, "") +} + +func rewriteCodexCollaborationTools(payload []byte, messageToolPaths, spawnAgentToolPaths []string, models []codexSpawnAgentModel, modelList string) []byte { + if len(messageToolPaths) == 0 && len(spawnAgentToolPaths) == 0 { return payload } - modelList := formatCodexSpawnAgentModels(models) + if modelList == "" && len(models) > 0 { + modelList = formatCodexSpawnAgentModels(models) + } updated := payload - for _, toolPath := range toolPaths { + for _, toolPath := range spawnAgentToolPaths { descriptionPath := toolPath + ".description" description := gjson.GetBytes(updated, descriptionPath) if description.Type == gjson.String && modelList != "" { @@ -472,11 +609,16 @@ func rewriteCodexSpawnAgentTools(payload []byte, toolPaths []string, models []co } } } + } - var errDelete error - updated, errDelete = sjson.DeleteBytes(updated, toolPath+".parameters.properties.message.encrypted") - if errDelete != nil { - return payload + for _, toolPath := range messageToolPaths { + encryptedPath := toolPath + ".parameters.properties.message.encrypted" + if gjson.GetBytes(updated, encryptedPath).Exists() { + var errDelete error + updated, errDelete = sjson.DeleteBytes(updated, encryptedPath) + if errDelete != nil { + return payload + } } } return updated @@ -724,8 +866,12 @@ func collectCodexToolPathsByNames(tools gjson.Result, path string, paths *[]stri func removeCodexCollaborationMessageEncryption(payload []byte, toolPaths []string) []byte { updated := payload for _, toolPath := range toolPaths { + encryptedPath := toolPath + ".parameters.properties.message.encrypted" + if !gjson.GetBytes(updated, encryptedPath).Exists() { + continue + } var errDelete error - updated, errDelete = sjson.DeleteBytes(updated, toolPath+".parameters.properties.message.encrypted") + updated, errDelete = sjson.DeleteBytes(updated, encryptedPath) if errDelete != nil { return payload } @@ -812,6 +958,9 @@ func replaceCodexSpawnAgentModels(description, modelList string) string { } func removeCodexSpawnAgentModelSections(description string) (string, string) { + if !strings.Contains(description, codexSpawnAgentModelsHeading) { + return description, "" + } lines := strings.SplitAfter(description, "\n") var cleaned strings.Builder headingIndent := "" diff --git a/internal/client/codex/optimize-multi-agent-v2/optimize_multi_agent_v2_test.go b/internal/client/codex/optimize-multi-agent-v2/optimize_multi_agent_v2_test.go index 01a64484..08f7779c 100644 --- a/internal/client/codex/optimize-multi-agent-v2/optimize_multi_agent_v2_test.go +++ b/internal/client/codex/optimize-multi-agent-v2/optimize_multi_agent_v2_test.go @@ -839,3 +839,171 @@ func TestRemoveCodexCollaborationMessageEncryptionNoOpWithoutEncrypted(t *testin t.Fatalf("payload changed when no encrypted field existed: %s", got) } } + +func TestCodexSpawnAgentModelsCacheInvalidation(t *testing.T) { + modelRegistry := registry.GetGlobalRegistry() + clientID1 := "cache-invalidation-client-1" + clientID2 := "cache-invalidation-client-2" + + // 1. Initial registration + modelRegistry.RegisterClient(clientID1, "openai", []*registry.ModelInfo{ + { + ID: "test-spawn-model-alpha", + DisplayName: "Test Spawn Model Alpha", + Description: "Initial description.", + Thinking: ®istry.ThinkingSupport{ + Levels: []string{"low", "medium"}, + }, + }, + }) + t.Cleanup(func() { + modelRegistry.UnregisterClient(clientID1) + modelRegistry.UnregisterClient(clientID2) + }) + + formatted1 := formatCodexSpawnAgentModelsForRequest(context.Background(), nil, false) + if !strings.Contains(formatted1, "test-spawn-model-alpha") { + t.Fatalf("expected initial markdown to contain test-spawn-model-alpha, got: %s", formatted1) + } + if !strings.Contains(formatted1, "Reasoning efforts: low, medium") { + t.Fatalf("expected initial reasoning efforts low, medium, got: %s", formatted1) + } + + // 2. Cache hit returns identical content + formattedHit := formatCodexSpawnAgentModelsForRequest(context.Background(), nil, false) + if formattedHit != formatted1 { + t.Fatalf("cache hit expected identical output, got %s vs %s", formattedHit, formatted1) + } + + // 3. Registering second model invalidates cache + modelRegistry.RegisterClient(clientID2, "openai", []*registry.ModelInfo{ + { + ID: "test-spawn-model-beta", + DisplayName: "Test Spawn Model Beta", + Description: "Second model.", + }, + }) + + formatted2 := formatCodexSpawnAgentModelsForRequest(context.Background(), nil, false) + if !strings.Contains(formatted2, "test-spawn-model-beta") { + t.Fatalf("expected cache invalidation to include test-spawn-model-beta, got: %s", formatted2) + } + + // 4. Modifying model thinking levels invalidates cache + modelRegistry.RegisterClient(clientID1, "openai", []*registry.ModelInfo{ + { + ID: "test-spawn-model-alpha", + DisplayName: "Test Spawn Model Alpha", + Description: "Initial description.", + Thinking: ®istry.ThinkingSupport{ + Levels: []string{"low", "medium", "high", "max"}, + }, + }, + }) + + formatted3 := formatCodexSpawnAgentModelsForRequest(context.Background(), nil, false) + if !strings.Contains(formatted3, "low, medium (default), high, max") { + t.Fatalf("expected updated thinking levels to reflect in markdown, got: %s", formatted3) + } + + // 5. Unregistering client invalidates cache + modelRegistry.UnregisterClient(clientID2) + formatted4 := formatCodexSpawnAgentModelsForRequest(context.Background(), nil, false) + if strings.Contains(formatted4, "test-spawn-model-beta") { + t.Fatalf("expected test-spawn-model-beta to be removed after unregistering, got: %s", formatted4) + } +} + +func BenchmarkCodexSpawnAgentModelsForRequest(b *testing.B) { + modelRegistry := registry.GetGlobalRegistry() + clientID := "bench-client-models" + modelRegistry.RegisterClient(clientID, "openai", []*registry.ModelInfo{ + { + ID: "gpt-5.5", + DisplayName: "Default model", + Description: "Default model description.", + }, + { + ID: "claude-3-7-sonnet", + DisplayName: "Claude 3.7 Sonnet", + Description: "Claude model description.", + }, + }) + b.Cleanup(func() { + modelRegistry.UnregisterClient(clientID) + }) + + ctx := context.Background() + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + codexSpawnAgentModelsForRequest(ctx, nil, false) + } +} + +func BenchmarkPrepareCodexMultiAgentV2Tools(b *testing.B) { + modelRegistry := registry.GetGlobalRegistry() + clientID := "bench-client-prepare" + modelRegistry.RegisterClient(clientID, "openai", []*registry.ModelInfo{ + { + ID: "gpt-5.5", + DisplayName: "Default model", + Description: "Default model description.", + }, + }) + b.Cleanup(func() { + modelRegistry.UnregisterClient(clientID) + }) + + payload := []byte(`{ + "tools":[ + {"type":"namespace","name":"collaboration","tools":[ + {"type":"function","name":"spawn_agent","description":"Spawns an agent.\n","parameters":{"type":"object","properties":{"message":{"type":"string","encrypted":true}}}}, + {"type":"function","name":"send_message","parameters":{"type":"object","properties":{"message":{"type":"string","encrypted":true}}}}, + {"type":"function","name":"followup_task","parameters":{"type":"object","properties":{"message":{"type":"string","encrypted":true}}}} + ]} + ] + }`) + headers := http.Header{"User-Agent": []string{"Codex Desktop/0.146.0-alpha.3"}} + ctx := context.Background() + + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + PrepareCodexMultiAgentV2Tools(ctx, headers, payload, true, false) + } +} + +func BenchmarkOptimizeCodexMultiAgentV2Request(b *testing.B) { + modelRegistry := registry.GetGlobalRegistry() + clientID := "bench-client-opt" + modelRegistry.RegisterClient(clientID, "openai", []*registry.ModelInfo{ + { + ID: "gpt-5.5", + DisplayName: "Default model", + Description: "Default model description.", + }, + }) + b.Cleanup(func() { + modelRegistry.UnregisterClient(clientID) + }) + + payload := []byte(`{ + "tools":[ + {"type":"namespace","name":"collaboration","tools":[ + {"type":"function","name":"spawn_agent","description":"Spawns an agent.\n","parameters":{"type":"object","properties":{"message":{"type":"string","encrypted":true}}}}, + {"type":"function","name":"send_message","parameters":{"type":"object","properties":{"message":{"type":"string","encrypted":true}}}}, + {"type":"function","name":"followup_task","parameters":{"type":"object","properties":{"message":{"type":"string","encrypted":true}}}} + ]} + ] + }`) + headers := http.Header{"User-Agent": []string{"Codex Desktop/0.146.0-alpha.3"}} + cfg := &config.Config{Codex: config.CodexConfig{OptimizeMultiAgentV2: true}} + ctx := context.Background() + + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + OptimizeCodexMultiAgentV2Request(ctx, headers, payload, cfg) + } +} diff --git a/internal/registry/codex_client_models.go b/internal/registry/codex_client_models.go index 8e601f11..370abf29 100644 --- a/internal/registry/codex_client_models.go +++ b/internal/registry/codex_client_models.go @@ -39,6 +39,13 @@ func GetCodexClientModelsJSON() []byte { return data } +// GetCodexClientModelsRevision returns the current revision of the Codex client model catalog. +func GetCodexClientModelsRevision() uint64 { + codexClientCatalogStore.mu.RLock() + defer codexClientCatalogStore.mu.RUnlock() + return codexClientCatalogStore.revision +} + // GetCodexClientModelsSnapshot returns a consistent catalog copy and revision. // The revision changes only when validated catalog content changes. func GetCodexClientModelsSnapshot() ([]byte, uint64) { diff --git a/internal/registry/model_registry.go b/internal/registry/model_registry.go index 2b5035bd..ed904bc2 100644 --- a/internal/registry/model_registry.go +++ b/internal/registry/model_registry.go @@ -151,6 +151,8 @@ type ModelRegistry struct { mutex *sync.RWMutex // availableModelsCache stores per-handler snapshots for GetAvailableModels. availableModelsCache map[string]availableModelsCacheEntry + // generation tracks changes to model registrations and availability. + generation uint64 // hook is an optional callback sink for model registration changes hook ModelRegistryHook } @@ -180,12 +182,20 @@ func (r *ModelRegistry) ensureAvailableModelsCacheLocked() { } func (r *ModelRegistry) invalidateAvailableModelsCacheLocked() { + r.generation++ if len(r.availableModelsCache) == 0 { return } clear(r.availableModelsCache) } +// GetGeneration returns the current generation counter of model registrations. +func (r *ModelRegistry) GetGeneration() uint64 { + r.mutex.RLock() + defer r.mutex.RUnlock() + return r.generation +} + // LookupModelInfo searches dynamic registry (provider-specific > global) then static definitions. func LookupModelInfo(modelID string, provider ...string) *ModelInfo { modelID = strings.TrimSpace(modelID) -- 2.51.2