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 b70b2b9ff..3b20ecd6c 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 @@ -27,6 +27,15 @@ const ( codexOptimizedCollaborationNamePrefix = codexOptimizedCollaborationNamespace + "__" ) +// codexCollaborationMessageTools are the collaboration tool names whose +// parameters.properties.message.encrypted field must be stripped so that +// message content remains readable by the proxy. +var codexCollaborationMessageTools = map[string]struct{}{ + "spawn_agent": {}, + "send_message": {}, + "followup_task": {}, +} + type codexSpawnAgentModel struct { id string description string @@ -73,6 +82,7 @@ func OptimizeCodexMultiAgentV2Request(ctx context.Context, headers http.Header, return payload, false } updated := rewriteCodexAgentMessageContent(payload) + updated = removeCodexCollaborationMessageEncryption(updated, codexCollaborationMessageToolPaths(updated)) toolPaths := codexSpawnAgentToolPaths(updated) if len(toolPaths) == 0 || hasCodexOptimizedCollaborationConflict(updated) { return updated, false @@ -609,8 +619,19 @@ func rewriteCodexAgentMessageContent(payload []byte) []byte { } func codexSpawnAgentToolPaths(payload []byte) []string { - paths := make([]string, 0, 1) - collectCodexSpawnAgentToolPaths(gjson.GetBytes(payload, "tools"), "tools", &paths) + return codexToolPathsByNames(payload, map[string]struct{}{"spawn_agent": {}}) +} + +// codexCollaborationMessageToolPaths discovers function tools named +// spawn_agent, send_message, or followup_task inside top-level tools arrays and +// input[].additional_tools arrays, including nested namespace tools. +func codexCollaborationMessageToolPaths(payload []byte) []string { + return codexToolPathsByNames(payload, codexCollaborationMessageTools) +} + +func codexToolPathsByNames(payload []byte, names map[string]struct{}) []string { + paths := make([]string, 0, len(names)) + collectCodexToolPathsByNames(gjson.GetBytes(payload, "tools"), "tools", &paths, names) input := gjson.GetBytes(payload, "input") if input.IsArray() { @@ -618,28 +639,45 @@ func codexSpawnAgentToolPaths(payload []byte) []string { if strings.TrimSpace(item.Get("type").String()) != "additional_tools" { continue } - collectCodexSpawnAgentToolPaths(item.Get("tools"), fmt.Sprintf("input.%d.tools", index), &paths) + collectCodexToolPathsByNames(item.Get("tools"), fmt.Sprintf("input.%d.tools", index), &paths, names) } } return paths } -func collectCodexSpawnAgentToolPaths(tools gjson.Result, path string, paths *[]string) { +func collectCodexToolPathsByNames(tools gjson.Result, path string, paths *[]string, names map[string]struct{}) { if !tools.IsArray() { return } for index, tool := range tools.Array() { toolPath := fmt.Sprintf("%s.%d", path, index) toolType := strings.TrimSpace(tool.Get("type").String()) - if toolType == "function" && strings.TrimSpace(tool.Get("name").String()) == "spawn_agent" { - *paths = append(*paths, toolPath) + if toolType == "function" { + if _, ok := names[strings.TrimSpace(tool.Get("name").String())]; ok { + *paths = append(*paths, toolPath) + } } if toolType == "namespace" { - collectCodexSpawnAgentToolPaths(tool.Get("tools"), toolPath+".tools", paths) + collectCodexToolPathsByNames(tool.Get("tools"), toolPath+".tools", paths, names) } } } +// removeCodexCollaborationMessageEncryption deletes the +// parameters.properties.message.encrypted field from each discovered +// collaboration message tool so the proxy can read the plaintext message. +func removeCodexCollaborationMessageEncryption(payload []byte, toolPaths []string) []byte { + updated := payload + for _, toolPath := range toolPaths { + var errDelete error + updated, errDelete = sjson.DeleteBytes(updated, toolPath+".parameters.properties.message.encrypted") + if errDelete != nil { + return payload + } + } + return updated +} + func formatCodexSpawnAgentModels(models []codexSpawnAgentModel) string { var modelList strings.Builder for _, model := range models { 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 08622de90..112c32650 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 @@ -587,3 +587,187 @@ func TestCodexClientUserAgentPrefersGinRequest(t *testing.T) { t.Fatalf("codexClientUserAgent() = %q, want gin request User-Agent", got) } } + +func TestCodexCollaborationMessageToolPathsFindsAllThreeTools(t *testing.T) { + t.Parallel() + + payload := []byte(`{ + "tools":[ + {"type":"namespace","name":"collaboration","tools":[ + {"type":"function","name":"spawn_agent","parameters":{"properties":{"message":{"encrypted":true}}}}, + {"type":"function","name":"send_message","parameters":{"properties":{"message":{"encrypted":true}}}}, + {"type":"function","name":"followup_task","parameters":{"properties":{"message":{"encrypted":true}}}}, + {"type":"function","name":"unrelated_tool","parameters":{"properties":{"message":{"encrypted":true}}}} + ]} + ] + }`) + paths := codexCollaborationMessageToolPaths(payload) + wantCount := 3 + if len(paths) != wantCount { + t.Fatalf("path count = %d, want %d; paths=%v", len(paths), wantCount, paths) + } +} + +func TestCodexCollaborationMessageToolPathsAdditionalTools(t *testing.T) { + t.Parallel() + + payload := []byte(`{ + "input":[ + {"type":"additional_tools","role":"developer","tools":[ + {"type":"namespace","name":"collaboration","tools":[ + {"type":"function","name":"send_message","parameters":{"properties":{"message":{"encrypted":true}}}}, + {"type":"function","name":"followup_task","parameters":{"properties":{"message":{"encrypted":true}}}} + ]} + ]} + ] + }`) + paths := codexCollaborationMessageToolPaths(payload) + if len(paths) != 2 { + t.Fatalf("path count = %d, want 2; paths=%v", len(paths), paths) + } +} + +func TestRemoveCodexCollaborationMessageEncryptionAllTools(t *testing.T) { + t.Parallel() + + payload := []byte(`{ + "tools":[ + {"type":"namespace","name":"collaboration","tools":[ + {"type":"function","name":"spawn_agent","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}}}} + ]} + ] + }`) + paths := codexCollaborationMessageToolPaths(payload) + got := removeCodexCollaborationMessageEncryption(payload, paths) + + for _, toolPath := range []string{ + "tools.0.tools.0", + "tools.0.tools.1", + "tools.0.tools.2", + } { + if encrypted := gjson.GetBytes(got, toolPath+".parameters.properties.message.encrypted"); encrypted.Exists() { + t.Fatalf("%s.parameters.properties.message.encrypted was not removed: %s", toolPath, encrypted.Raw) + } + if msgType := gjson.GetBytes(got, toolPath+".parameters.properties.message.type").String(); msgType != "string" { + t.Fatalf("%s.parameters.properties.message.type changed: %q", toolPath, msgType) + } + } +} + +func TestRemoveCodexCollaborationMessageEncryptionPreservesUnrelatedEncryptedFields(t *testing.T) { + t.Parallel() + + payload := []byte(`{ + "tools":[ + {"type":"function","name":"send_message","parameters":{"properties":{"message":{"type":"string","encrypted":true},"data":{"encrypted":"keep-me"}}}}, + {"type":"function","name":"unrelated_tool","parameters":{"properties":{"message":{"encrypted":true}}}} + ] + }`) + paths := codexCollaborationMessageToolPaths(payload) + got := removeCodexCollaborationMessageEncryption(payload, paths) + + if encrypted := gjson.GetBytes(got, "tools.0.parameters.properties.message.encrypted"); encrypted.Exists() { + t.Fatalf("send_message message.encrypted was not removed: %s", encrypted.Raw) + } + if dataEncrypted := gjson.GetBytes(got, "tools.0.parameters.properties.data.encrypted").String(); dataEncrypted != "keep-me" { + t.Fatalf("unrelated data.encrypted was changed: %q", dataEncrypted) + } + if unrelatedEncrypted := gjson.GetBytes(got, "tools.1.parameters.properties.message.encrypted"); !unrelatedEncrypted.Exists() { + t.Fatalf("unrelated tool message.encrypted was removed: %s", got) + } +} + +func TestOptimizeCodexMultiAgentV2RequestRemovesEncryptionWithoutSpawnAgent(t *testing.T) { + t.Parallel() + + payload := []byte(`{ + "tools":[ + {"type":"namespace","name":"collaboration","tools":[ + {"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}} + got, optimized := OptimizeCodexMultiAgentV2Request(context.Background(), headers, payload, cfg) + + if optimized { + t.Fatal("namespace was unexpectedly optimized without spawn_agent") + } + for _, path := range []string{"tools.0.tools.0", "tools.0.tools.1"} { + if encrypted := gjson.GetBytes(got, path+".parameters.properties.message.encrypted"); encrypted.Exists() { + t.Fatalf("%s.parameters.properties.message.encrypted was not removed: %s", path, encrypted.Raw) + } + } +} + +func TestOptimizeCodexMultiAgentV2RequestRemovesEncryptionInAdditionalTools(t *testing.T) { + t.Parallel() + + payload := []byte(`{ + "input":[ + {"type":"additional_tools","role":"developer","tools":[ + {"type":"namespace","name":"collaboration","tools":[ + {"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-tui/0.145.0"}} + cfg := &config.Config{Codex: config.CodexConfig{OptimizeMultiAgentV2: true}} + got, _ := OptimizeCodexMultiAgentV2Request(context.Background(), headers, payload, cfg) + + for _, path := range []string{"input.0.tools.0.tools.0", "input.0.tools.0.tools.1"} { + if encrypted := gjson.GetBytes(got, path+".parameters.properties.message.encrypted"); encrypted.Exists() { + t.Fatalf("%s.parameters.properties.message.encrypted was not removed: %s", path, encrypted.Raw) + } + } +} + +func TestOptimizeCodexMultiAgentV2RequestRemovesEncryptionFromAllThreeToolsWithSpawnAgent(t *testing.T) { + t.Parallel() + + payload := []byte(`{ + "tools":[ + {"type":"namespace","name":"collaboration","tools":[ + {"type":"function","name":"spawn_agent","description":"Spawns an agent.","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}} + got, optimized := OptimizeCodexMultiAgentV2Request(context.Background(), headers, payload, cfg) + + if !optimized { + t.Fatal("collaboration namespace was not optimized with spawn_agent present") + } + for _, path := range []string{"tools.0.tools.0", "tools.0.tools.1", "tools.0.tools.2"} { + if encrypted := gjson.GetBytes(got, path+".parameters.properties.message.encrypted"); encrypted.Exists() { + t.Fatalf("%s.parameters.properties.message.encrypted was not removed: %s", path, encrypted.Raw) + } + } + if namespace := gjson.GetBytes(got, "tools.0.name").String(); namespace != codexOptimizedCollaborationNamespace { + t.Fatalf("namespace = %q, want %q", namespace, codexOptimizedCollaborationNamespace) + } +} + +func TestRemoveCodexCollaborationMessageEncryptionNoOpWithoutEncrypted(t *testing.T) { + t.Parallel() + + payload := []byte(`{ + "tools":[ + {"type":"function","name":"send_message","parameters":{"type":"object","properties":{"message":{"type":"string"}}}} + ] + }`) + paths := codexCollaborationMessageToolPaths(payload) + got := removeCodexCollaborationMessageEncryption(payload, paths) + if string(got) != string(payload) { + t.Fatalf("payload changed when no encrypted field existed: %s", got) + } +}