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 80774883..b49bcd7e 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 @@ -482,6 +482,12 @@ func rewriteCodexSpawnAgentTools(payload []byte, toolPaths []string, models []co return updated } +// HasCodexMultiAgentV2NamespaceConflict reports whether the request defines +// the reserved optimized namespace, which must remain untouched. +func HasCodexMultiAgentV2NamespaceConflict(payload []byte) bool { + return hasCodexOptimizedCollaborationConflict(payload) +} + func hasCodexOptimizedCollaborationConflict(payload []byte) bool { if codexToolsHaveOptimizedCollaborationConflict(gjson.GetBytes(payload, "tools")) { return true diff --git a/internal/runtime/executor/codex_websockets_execute.go b/internal/runtime/executor/codex_websockets_execute.go index 5243bc7a..f8c51ef0 100644 --- a/internal/runtime/executor/codex_websockets_execute.go +++ b/internal/runtime/executor/codex_websockets_execute.go @@ -64,6 +64,7 @@ func (e *CodexWebsocketsExecutor) Execute(ctx context.Context, auth *cliproxyaut } body = sanitizeOpenAIResponsesReasoningEncryptedContent(ctx, "codex websockets executor", body) body = normalizeCodexWebsocketParallelToolCalls(body, opts.Headers) + multiAgentV2Conflict := helps.HasCodexMultiAgentV2NamespaceConflict(body) body, optimizeMultiAgentV2 := helps.OptimizeCodexMultiAgentV2RequestForAuth(ctx, opts.Headers, body, e.cfg, auth, baseModel) body, replayScope, errReplay := applyCodexReasoningReplayCacheRequired(ctx, from, req, opts, body) if errReplay != nil { @@ -182,6 +183,7 @@ func (e *CodexWebsocketsExecutor) Execute(ctx context.Context, auth *cliproxyaut sess.clearActive(conn, readCh) }() } + restoreMultiAgentV2 := !multiAgentV2Conflict && (optimizeMultiAgentV2 || sess.isMultiAgentV2Optimized(conn)) if errSend := writeCodexWebsocketMessage(sess, conn, wsReqBody); errSend != nil { errSend = mapCodexWebsocketWriteError(sess, conn, errSend) @@ -215,6 +217,7 @@ func (e *CodexWebsocketsExecutor) Execute(ctx context.Context, auth *cliproxyaut return resp, errBind } readCh = sess.activate(conn) + restoreMultiAgentV2 = !multiAgentV2Conflict && (optimizeMultiAgentV2 || sess.isMultiAgentV2Optimized(conn)) wsReqBodyRetry := buildCodexWebsocketRequestBody(upstreamBody) helps.RecordAPIWebsocketRequest(ctx, e.cfg, helps.UpstreamRequestLog{ URL: wsURL, @@ -248,6 +251,10 @@ func (e *CodexWebsocketsExecutor) Execute(ctx context.Context, auth *cliproxyaut } } + if optimizeMultiAgentV2 { + sess.markMultiAgentV2Optimized(conn) + } + outputItemsByIndex := make(map[int64][]byte) var outputItemsFallback [][]byte for { @@ -279,7 +286,7 @@ func (e *CodexWebsocketsExecutor) Execute(ctx context.Context, auth *cliproxyaut reporter.MarkFirstResponseByte() payload = applyCodexIdentityConfuseResponsePayload(payload, identityState) helps.AppendAPIWebsocketResponse(ctx, e.cfg, payload) - payload = helps.RestoreCodexMultiAgentV2Response(payload, optimizeMultiAgentV2) + payload = helps.RestoreCodexMultiAgentV2Response(payload, restoreMultiAgentV2) if wsErr, ok := parseCodexWebsocketError(payload); ok { if sess != nil { diff --git a/internal/runtime/executor/codex_websockets_session.go b/internal/runtime/executor/codex_websockets_session.go index 76cb29ac..9c721ca7 100644 --- a/internal/runtime/executor/codex_websockets_session.go +++ b/internal/runtime/executor/codex_websockets_session.go @@ -51,14 +51,15 @@ type codexWebsocketSession struct { reqMu sync.Mutex - connMu sync.Mutex - conn *websocket.Conn - connCloser *websocketConnectionCloser - wsURL string - authID string - lifecycleBindMu sync.Mutex - lifecycle cliproxyexecutor.ExecutionLifecycle - lifecycleModel string + connMu sync.Mutex + conn *websocket.Conn + connCloser *websocketConnectionCloser + wsURL string + authID string + multiAgentV2OptimizedConn *websocket.Conn + lifecycleBindMu sync.Mutex + lifecycle cliproxyexecutor.ExecutionLifecycle + lifecycleModel string writeMu sync.Mutex @@ -163,6 +164,26 @@ func (s *codexWebsocketSession) writeMessage(conn *websocket.Conn, msgType int, return conn.WriteMessage(msgType, payload) } +func (s *codexWebsocketSession) markMultiAgentV2Optimized(conn *websocket.Conn) { + if s == nil || conn == nil { + return + } + s.connMu.Lock() + if s.conn == conn { + s.multiAgentV2OptimizedConn = conn + } + s.connMu.Unlock() +} + +func (s *codexWebsocketSession) isMultiAgentV2Optimized(conn *websocket.Conn) bool { + if s == nil || conn == nil { + return false + } + s.connMu.Lock() + defer s.connMu.Unlock() + return s.conn == conn && s.multiAgentV2OptimizedConn == conn +} + // sendTerminalWebsocketRead reports whether it invalidated a full channel's connection before waiting. func sendTerminalWebsocketRead(ch chan<- codexWebsocketRead, done <-chan struct{}, event codexWebsocketRead, invalidate func()) bool { select { @@ -272,6 +293,7 @@ func (s *codexWebsocketSession) detachConnection(conn *websocket.Conn, lifecycle closer = s.connCloser s.conn = nil s.connCloser = nil + s.multiAgentV2OptimizedConn = nil if s.readerConn == conn { s.readerConn = nil } @@ -346,6 +368,7 @@ func detachMismatchedWebsocketSessionConn(sess *codexWebsocketSession, authID st sess.lifecycleModel = "" sess.conn = nil sess.connCloser = nil + sess.multiAgentV2OptimizedConn = nil if sess.readerConn == conn { sess.readerConn = nil } @@ -505,6 +528,7 @@ func (e *CodexWebsocketsExecutor) ensureUpstreamConn(ctx context.Context, auth * } sess.conn = conn sess.connCloser = closer + sess.multiAgentV2OptimizedConn = nil sess.wsURL = wsURL sess.authID = authID sess.readerConn = conn @@ -602,6 +626,7 @@ func (e *CodexWebsocketsExecutor) invalidateUpstreamConnWithNotify(sess *codexWe sess.lifecycleModel = "" sess.conn = nil sess.connCloser = nil + sess.multiAgentV2OptimizedConn = nil if sess.readerConn == conn { sess.readerConn = nil } @@ -693,6 +718,7 @@ func closeCodexWebsocketSession(sess *codexWebsocketSession, reason string) { sess.lifecycleModel = "" sess.conn = nil sess.connCloser = nil + sess.multiAgentV2OptimizedConn = nil if sess.readerConn == conn { sess.readerConn = nil } diff --git a/internal/runtime/executor/codex_websockets_spawn_agent_test.go b/internal/runtime/executor/codex_websockets_spawn_agent_test.go index d0a2fc33..fc3771b6 100644 --- a/internal/runtime/executor/codex_websockets_spawn_agent_test.go +++ b/internal/runtime/executor/codex_websockets_spawn_agent_test.go @@ -4,6 +4,8 @@ import ( "fmt" "net/http" "net/http/httptest" + "strings" + "sync/atomic" "testing" "github.com/gorilla/websocket" @@ -15,6 +17,122 @@ import ( "github.com/tidwall/gjson" ) +func TestCodexWebsocketsExecutorRestoresMultiAgentV2NamespaceAcrossIncrementalTurns(t *testing.T) { + for _, tt := range []struct { + name string + stream bool + }{ + {name: "execute"}, + {name: "stream", stream: true}, + } { + t.Run(tt.name, func(t *testing.T) { + upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }} + capturedPayload := make(chan []byte, 3) + var connectionCount atomic.Int32 + var requestCount atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, request *http.Request) { + connectionCount.Add(1) + conn, errUpgrade := upgrader.Upgrade(w, request, nil) + if errUpgrade != nil { + t.Errorf("upgrade websocket: %v", errUpgrade) + return + } + defer func() { _ = conn.Close() }() + + for { + _, payload, errRead := conn.ReadMessage() + if errRead != nil { + return + } + capturedPayload <- append([]byte(nil), payload...) + turn := requestCount.Add(1) + completed := []byte(fmt.Sprintf(`{"type":"response.completed","response":{"id":"resp_%d","object":"response","status":"completed","output":[{"type":"function_call","name":"spawn_agent","namespace":"collaboration-optimize","arguments":"{}","call_id":"call_%d"}]}}`, turn, turn)) + if errWrite := conn.WriteMessage(websocket.TextMessage, completed); errWrite != nil { + t.Errorf("write websocket response: %v", errWrite) + return + } + if turn == 3 { + return + } + } + })) + t.Cleanup(server.Close) + + executor := NewCodexWebsocketsExecutor(&config.Config{Codex: config.CodexConfig{OptimizeMultiAgentV2: true}}) + const executionSessionID = "multi-agent-v2-incremental" + t.Cleanup(func() { executor.CloseExecutionSession(executionSessionID) }) + auth := &cliproxyauth.Auth{ + ID: "codex-test", + Attributes: map[string]string{ + "base_url": server.URL, + "api_key": "test", + }, + } + opts := cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FromString("openai-response"), + ResponseFormat: sdktranslator.FromString("openai-response"), + Headers: http.Header{"User-Agent": []string{"overridden-client/1.0"}}, + Metadata: map[string]any{ + cliproxyexecutor.ExecutionSessionMetadataKey: executionSessionID, + }, + } + execute := func(payload []byte) []byte { + t.Helper() + req := cliproxyexecutor.Request{Model: "gpt-5.4", Payload: payload} + if !tt.stream { + response, errExecute := executor.Execute(codexSpawnAgentTestContext(), auth, req, opts) + if errExecute != nil { + t.Fatalf("Execute() error = %v", errExecute) + } + return response.Payload + } + + result, errExecute := executor.ExecuteStream(codexSpawnAgentTestContext(), auth, req, opts) + if errExecute != nil { + t.Fatalf("ExecuteStream() error = %v", errExecute) + } + var responsePayload []byte + for chunk := range result.Chunks { + if chunk.Err != nil { + t.Fatalf("stream chunk error = %v", chunk.Err) + } + responsePayload = append(responsePayload, chunk.Payload...) + } + return responsePayload + } + + firstClientPayload := execute(codexSpawnAgentTestPayload()) + firstUpstreamPayload := <-capturedPayload + if namespace := gjson.GetBytes(firstUpstreamPayload, "input.0.tools.0.name").String(); namespace != "collaboration-optimize" { + t.Fatalf("first upstream namespace = %q, want collaboration-optimize", namespace) + } + assertCodexSpawnAgentClientNamespace(t, firstClientPayload) + + secondRequest := []byte(`{"model":"gpt-5.4","previous_response_id":"resp_1","input":[{"type":"function_call_output","call_id":"call_1","output":"done"}]}`) + secondClientPayload := execute(secondRequest) + secondUpstreamPayload := <-capturedPayload + if strings.Contains(string(secondUpstreamPayload), "collaboration") || strings.Contains(string(secondUpstreamPayload), "spawn_agent") { + t.Fatalf("incremental upstream request unexpectedly contains collaboration tools: %s", secondUpstreamPayload) + } + assertCodexSpawnAgentClientNamespace(t, secondClientPayload) + + conflictingRequest := []byte(`{"model":"gpt-5.4","tools":[{"type":"namespace","name":"collaboration-optimize","tools":[{"type":"function","name":"spawn_agent","description":"User-defined tool."}]}],"input":[{"type":"message","role":"user","content":"use the user-defined namespace"}]}`) + conflictingClientPayload := execute(conflictingRequest) + conflictingUpstreamPayload := <-capturedPayload + if namespace := gjson.GetBytes(conflictingUpstreamPayload, "tools.0.name").String(); namespace != "collaboration-optimize" { + t.Fatalf("conflicting upstream namespace = %q, want collaboration-optimize", namespace) + } + if !strings.Contains(string(conflictingClientPayload), `"namespace":"collaboration-optimize"`) { + t.Fatalf("user-defined collaboration-optimize namespace was rewritten: %s", conflictingClientPayload) + } + + if got := connectionCount.Load(); got != 1 { + t.Fatalf("upstream websocket connections = %d, want 1", got) + } + }) + } +} + func TestCodexWebsocketsExecutorOptimizeMultiAgentV2(t *testing.T) { modelID := "codex-websocket-spawn-agent-test-model" clientID := "codex-websocket-spawn-agent-test-client" diff --git a/internal/runtime/executor/codex_websockets_stream.go b/internal/runtime/executor/codex_websockets_stream.go index 505081f0..84ea8495 100644 --- a/internal/runtime/executor/codex_websockets_stream.go +++ b/internal/runtime/executor/codex_websockets_stream.go @@ -61,6 +61,7 @@ func (e *CodexWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *clipr } body = sanitizeOpenAIResponsesReasoningEncryptedContent(ctx, "codex websockets executor", body) body = normalizeCodexWebsocketParallelToolCalls(body, opts.Headers) + multiAgentV2Conflict := helps.HasCodexMultiAgentV2NamespaceConflict(body) body, optimizeMultiAgentV2 := helps.OptimizeCodexMultiAgentV2RequestForAuth(ctx, opts.Headers, body, e.cfg, auth, baseModel) body, replayScope, errReplay := applyCodexReasoningReplayCacheRequired(ctx, from, req, opts, body) if errReplay != nil { @@ -183,6 +184,7 @@ func (e *CodexWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *clipr if sess != nil { readCh = sess.activate(conn) } + restoreMultiAgentV2 := !multiAgentV2Conflict && (optimizeMultiAgentV2 || sess.isMultiAgentV2Optimized(conn)) if errSend := writeCodexWebsocketMessage(sess, conn, wsReqBody); errSend != nil { errSend = mapCodexWebsocketWriteError(sess, conn, errSend) @@ -223,6 +225,7 @@ func (e *CodexWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *clipr return nil, errBind } readCh = sess.activate(conn) + restoreMultiAgentV2 = !multiAgentV2Conflict && (optimizeMultiAgentV2 || sess.isMultiAgentV2Optimized(conn)) wsReqBodyRetry := buildCodexWebsocketRequestBody(upstreamBody) helps.RecordAPIWebsocketRequest(ctx, e.cfg, helps.UpstreamRequestLog{ URL: wsURL, @@ -255,6 +258,10 @@ func (e *CodexWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *clipr } } + if optimizeMultiAgentV2 { + sess.markMultiAgentV2Optimized(conn) + } + out := make(chan cliproxyexecutor.StreamChunk) go func() { terminateReason := "completed" @@ -336,7 +343,7 @@ func (e *CodexWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *clipr reporter.MarkFirstResponseByte() payload = applyCodexIdentityConfuseResponsePayload(payload, identityState) helps.AppendAPIWebsocketResponse(ctx, e.cfg, payload) - payload = helps.RestoreCodexMultiAgentV2Response(payload, optimizeMultiAgentV2) + payload = helps.RestoreCodexMultiAgentV2Response(payload, restoreMultiAgentV2) if wsErr, ok := parseCodexWebsocketError(payload); ok { terminateReason = "upstream_error" diff --git a/internal/runtime/executor/helps/codex_multi_agent_v2.go b/internal/runtime/executor/helps/codex_multi_agent_v2.go index 981177ba..4e2209f8 100644 --- a/internal/runtime/executor/helps/codex_multi_agent_v2.go +++ b/internal/runtime/executor/helps/codex_multi_agent_v2.go @@ -97,6 +97,12 @@ func TranslateRequestWithAPIKeyModelCompatibility(ctx context.Context, headers h return thinking.ApplySummaryConfigForModel(translated, to.String(), model, summaryConfig) } +// HasCodexMultiAgentV2NamespaceConflict reports whether the request defines +// the reserved optimized namespace, which must remain untouched. +func HasCodexMultiAgentV2NamespaceConflict(payload []byte) bool { + return multiagentv2.HasCodexMultiAgentV2NamespaceConflict(payload) +} + // OptimizeCodexMultiAgentV2Request rewrites an eligible spawn_agent request and // reports whether the collaboration namespace was renamed for upstream use. func OptimizeCodexMultiAgentV2Request(ctx context.Context, headers http.Header, payload []byte, cfg *config.Config) ([]byte, bool) {