Handle encrypted collaboration message schemas

- Generalize spawn_agent tool path discovery into codexToolPathsByNames driven by a name set, and add codexCollaborationMessageTools covering spawn_agent, send_message, and followup_task. - Strip parameters.properties.message.encrypted from all collaboration message tools (top-level tools and input[].additional_tools, including nested namespace tools) via removeCodexCollaborationMessageEncryption, invoked unconditionally before the spawn-agent optimization decision. - Extend optimize_multi_agent_v2_test.go with coverage for the new tool path discovery, encrypted-field removal semantics, preservation of unrelated encrypted fields, and end-to-end OptimizeCodexMultiAgentV2Request behavior across namespace and additional_tools payloads. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-09-03 06:35:00 +08:00 · 2026-08-03 04:40:13 +08:00
parent 41fc5e1346
commit 1cf9a45ca2
2 changed files with 229 additions and 7 deletions
--- 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 {
--- 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)
+	}
+}