From 4d68ca8a63494e9df2572a0e35bdf900a147d5f0 Mon Sep 17 00:00:00 2001 From: Luis Pater Date: Sat, 22 Aug 2026 20:58:35 +0800 Subject: [PATCH] fix(codex): convert Grok client keepalive SSE frames to comments - Add Grok client detection via User-Agent (including Gin context fallback) in a new `grokbuild` helper. - Transform keepalive SSE `event`/`data` frames into `: keepalive` comments when streaming to Grok clients. - Keep keepalive frames untouched for non-Grok clients while preserving existing stream translation behavior. Closes: #5171 --- internal/client/grokbuild/keepalive.go | 84 ++++++ internal/client/grokbuild/keepalive_test.go | 156 ++++++++++ ...codex_executor_grokbuild_keepalive_test.go | 280 ++++++++++++++++++ .../runtime/executor/codex_executor_stream.go | 11 +- 4 files changed, 529 insertions(+), 2 deletions(-) create mode 100644 internal/client/grokbuild/keepalive.go create mode 100644 internal/client/grokbuild/keepalive_test.go create mode 100644 internal/runtime/executor/codex_executor_grokbuild_keepalive_test.go diff --git a/internal/client/grokbuild/keepalive.go b/internal/client/grokbuild/keepalive.go new file mode 100644 index 000000000..d7318c822 --- /dev/null +++ b/internal/client/grokbuild/keepalive.go @@ -0,0 +1,84 @@ +package grokbuild + +import ( + "bytes" + "context" + "net/http" + "slices" + "strings" + + "github.com/gin-gonic/gin" + "github.com/tidwall/gjson" +) + +var keepaliveSSEComment = []byte(": keepalive\n\n") + +// KeepaliveSSEComment returns the standard SSE comment used for keepalive. +func KeepaliveSSEComment() []byte { + return bytes.Clone(keepaliveSSEComment) +} + +// IsGrokClientUserAgent checks if the user agent contains "grok-pager" or "grok-shell". +func IsGrokClientUserAgent(userAgent string) bool { + ua := strings.ToLower(userAgent) + return strings.Contains(ua, "grok-pager") || strings.Contains(ua, "grok-shell") +} + +// IsGrokClientHeaders checks if the provided HTTP headers indicate a Grok client. +func IsGrokClientHeaders(headers http.Header) bool { + if headers == nil { + return false + } + for key, values := range headers { + if strings.EqualFold(key, "User-Agent") { + if slices.ContainsFunc(values, IsGrokClientUserAgent) { + return true + } + } + } + return false +} + +// IsGrokClientContext checks if either the context (e.g. Gin context) or headers indicate a Grok client. +func IsGrokClientContext(ctx context.Context, headers http.Header) bool { + if ctx != nil { + if ginCtx, ok := ctx.Value("gin").(*gin.Context); ok && ginCtx != nil && ginCtx.Request != nil { + if IsGrokClientHeaders(ginCtx.Request.Header) { + return true + } + } + } + return IsGrokClientHeaders(headers) +} + +// IsKeepalivePayload reports whether a JSON payload has type "keepalive". +func IsKeepalivePayload(payload []byte) bool { + return gjson.GetBytes(payload, "type").String() == "keepalive" +} + +// IsKeepaliveSSELine reports whether an SSE line represents a keepalive event or data frame. +func IsKeepaliveSSELine(line []byte) bool { + trimmed := bytes.TrimSpace(line) + if bytes.HasPrefix(trimmed, []byte("event:")) { + eventName := bytes.TrimSpace(trimmed[6:]) + return bytes.Equal(eventName, []byte("keepalive")) + } + if bytes.HasPrefix(trimmed, []byte("data:")) { + data := bytes.TrimSpace(trimmed[5:]) + return IsKeepalivePayload(data) + } + return false +} + +// TransformKeepaliveSSELine transforms a keepalive SSE line into an SSE comment line +// when isGrokClient is true. If the line is not a keepalive line or isGrokClient is false, +// it returns the original line and false. +func TransformKeepaliveSSELine(line []byte, isGrokClient bool) ([]byte, bool) { + if !isGrokClient { + return line, false + } + if IsKeepaliveSSELine(line) { + return bytes.Clone(keepaliveSSEComment), true + } + return line, false +} diff --git a/internal/client/grokbuild/keepalive_test.go b/internal/client/grokbuild/keepalive_test.go new file mode 100644 index 000000000..b4e5d0468 --- /dev/null +++ b/internal/client/grokbuild/keepalive_test.go @@ -0,0 +1,156 @@ +package grokbuild + +import ( + "bytes" + "context" + "net/http" + "net/http/httptest" + "testing" + + "github.com/gin-gonic/gin" +) + +func TestIsGrokClientUserAgent(t *testing.T) { + tests := []struct { + ua string + want bool + }{ + {"grok-shell/0.2.119 (macos; aarch64)", true}, + {"grok-pager/1.0.5 grok-shell/1.0.5 (linux; x86_64)", true}, + {"grok-pager/1.0.5", true}, + {"GROK-PAGER/1.0", true}, + {"GROK-SHELL/1.0", true}, + {"curl/8.7.1", false}, + {"openai-python/1.0.0", false}, + {"", false}, + } + for _, tc := range tests { + if got := IsGrokClientUserAgent(tc.ua); got != tc.want { + t.Errorf("IsGrokClientUserAgent(%q) = %v, want %v", tc.ua, got, tc.want) + } + } +} + +func TestIsGrokClientHeaders(t *testing.T) { + tests := []struct { + name string + headers http.Header + want bool + }{ + { + name: "User-Agent with grok-pager", + headers: http.Header{"User-Agent": []string{"grok-pager/1.0.5"}}, + want: true, + }, + { + name: "case insensitive header name", + headers: http.Header{"user-agent": []string{"grok-shell/0.2"}}, + want: true, + }, + { + name: "unrelated user agent", + headers: http.Header{"User-Agent": []string{"curl/8.7.1"}}, + want: false, + }, + { + name: "nil headers", + headers: nil, + want: false, + }, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + if got := IsGrokClientHeaders(tc.headers); got != tc.want { + t.Errorf("IsGrokClientHeaders() = %v, want %v", got, tc.want) + } + }) + } +} + +func TestIsGrokClientContext(t *testing.T) { + gin.SetMode(gin.TestMode) + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil) + c.Request.Header.Set("User-Agent", "grok-pager/1.0.5 grok-shell/1.0.5") + + ctx := context.WithValue(context.Background(), "gin", c) + if !IsGrokClientContext(ctx, nil) { + t.Error("expected IsGrokClientContext to detect gin context user agent") + } + + plainCtx := context.Background() + headers := http.Header{"User-Agent": []string{"grok-shell/1.0"}} + if !IsGrokClientContext(plainCtx, headers) { + t.Error("expected IsGrokClientContext to detect headers when gin context is absent") + } +} + +func TestIsKeepalivePayload(t *testing.T) { + tests := []struct { + payload []byte + want bool + }{ + {[]byte(`{"type":"keepalive","sequence_number":3}`), true}, + {[]byte(`{"type":"keepalive"}`), true}, + {[]byte(`{"type":"response.created"}`), false}, + {[]byte(`{"type":"response.reasoning.delta"}`), false}, + {[]byte(``), false}, + } + for _, tc := range tests { + if got := IsKeepalivePayload(tc.payload); got != tc.want { + t.Errorf("IsKeepalivePayload(%s) = %v, want %v", string(tc.payload), got, tc.want) + } + } +} + +func TestIsKeepaliveSSELine(t *testing.T) { + tests := []struct { + line []byte + want bool + }{ + {[]byte("event: keepalive"), true}, + {[]byte("event: keepalive\n"), true}, + {[]byte(" event: keepalive "), true}, + {[]byte(`data: {"type":"keepalive","sequence_number":3}`), true}, + {[]byte(`data: {"type":"keepalive"}`), true}, + {[]byte("event: response.created"), false}, + {[]byte("event: keepalive-other"), false}, + {[]byte(`data: {"type":"response.created"}`), false}, + {[]byte(""), false}, + } + for _, tc := range tests { + if got := IsKeepaliveSSELine(tc.line); got != tc.want { + t.Errorf("IsKeepaliveSSELine(%s) = %v, want %v", string(tc.line), got, tc.want) + } + } +} + +func TestTransformKeepaliveSSELine(t *testing.T) { + comment := KeepaliveSSEComment() + + // Grok client: keepalive line is transformed + got, ok := TransformKeepaliveSSELine([]byte("event: keepalive"), true) + if !ok || !bytes.Equal(got, comment) { + t.Errorf("TransformKeepaliveSSELine(event: keepalive, true) = %q, %v, want %q, true", string(got), ok, string(comment)) + } + + got, ok = TransformKeepaliveSSELine([]byte(`data: {"type":"keepalive","sequence_number":3}`), true) + if !ok || !bytes.Equal(got, comment) { + t.Errorf("TransformKeepaliveSSELine(data: keepalive, true) = %q, %v, want %q, true", string(got), ok, string(comment)) + } + + // Grok client: normal line is untouched + normalLine := []byte(`data: {"type":"response.created"}`) + got, ok = TransformKeepaliveSSELine(normalLine, true) + if ok || !bytes.Equal(got, normalLine) { + t.Errorf("TransformKeepaliveSSELine(normalLine, true) = %q, %v, want unchanged, false", string(got), ok) + } + + // Non-Grok client: keepalive line is untouched + keepaliveLine := []byte("event: keepalive") + got, ok = TransformKeepaliveSSELine(keepaliveLine, false) + if ok || !bytes.Equal(got, keepaliveLine) { + t.Errorf("TransformKeepaliveSSELine(event: keepalive, false) = %q, %v, want unchanged, false", string(got), ok) + } +} diff --git a/internal/runtime/executor/codex_executor_grokbuild_keepalive_test.go b/internal/runtime/executor/codex_executor_grokbuild_keepalive_test.go new file mode 100644 index 000000000..027f65bad --- /dev/null +++ b/internal/runtime/executor/codex_executor_grokbuild_keepalive_test.go @@ -0,0 +1,280 @@ +package executor + +import ( + "bytes" + "context" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/gin-gonic/gin" + "github.com/router-for-me/CLIProxyAPI/v7/internal/config" + _ "github.com/router-for-me/CLIProxyAPI/v7/internal/translator" + cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" + sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator" +) + +func TestCodexExecutorExecuteStream_GrokBuildConvertsKeepaliveToSSEComment(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + _, _ = w.Write([]byte("event: response.created\n")) + _, _ = w.Write([]byte(`data: {"type":"response.created","response":{"id":"resp_1","model":"gpt-5.6-luna"}}` + "\n\n")) + _, _ = w.Write([]byte("event: keepalive\n")) + _, _ = w.Write([]byte(`data: {"type":"keepalive","sequence_number":3}` + "\n\n")) + _, _ = w.Write([]byte("event: response.completed\n")) + _, _ = w.Write([]byte(`data: {"type":"response.completed","response":{"id":"resp_1","status":"completed","output":[]}}` + "\n\n")) + })) + defer server.Close() + + executor := NewCodexExecutor(&config.Config{}) + auth := &cliproxyauth.Auth{Attributes: map[string]string{ + "base_url": server.URL, + "api_key": "test", + }} + + tests := []struct { + name string + userAgent string + }{ + { + name: "Grok Build with grok-pager and grok-shell", + userAgent: "grok-pager/1.0.5 grok-shell/1.0.5 (linux; x86_64)", + }, + { + name: "Grok Shell only", + userAgent: "grok-shell/0.2.119 (macos; aarch64)", + }, + { + name: "Grok Pager only", + userAgent: "grok-pager/1.0.5 (linux; x86_64)", + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + res, err := executor.ExecuteStream(context.Background(), auth, cliproxyexecutor.Request{ + Model: "gpt-5.6-luna", + Payload: []byte(`{"model":"gpt-5.6-luna","input":"test"}`), + }, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FromString("openai-response"), + Stream: true, + Headers: http.Header{"User-Agent": []string{tc.userAgent}}, + }) + if err != nil { + t.Fatalf("ExecuteStream error: %v", err) + } + + var fullOutput bytes.Buffer + timeout := time.After(3 * time.Second) + done := false + for !done { + select { + case chunk, ok := <-res.Chunks: + if !ok { + done = true + break + } + if chunk.Err != nil { + t.Fatalf("unexpected chunk error: %v", chunk.Err) + } + fullOutput.Write(chunk.Payload) + case <-timeout: + t.Fatal("timed out reading stream chunks") + } + } + + outputStr := fullOutput.String() + if strings.Contains(outputStr, `{"type":"keepalive"`) || strings.Contains(outputStr, "event: keepalive") { + t.Fatalf("output must not contain keepalive event/data frame, got:\n%s", outputStr) + } + if !strings.Contains(outputStr, ": keepalive") { + t.Fatalf("output must contain ': keepalive' SSE comment, got:\n%s", outputStr) + } + if !strings.Contains(outputStr, "response.created") || !strings.Contains(outputStr, "response.completed") { + t.Fatalf("output missing normal lifecycle events, got:\n%s", outputStr) + } + }) + } +} + +func TestCodexExecutorExecuteStream_GrokBuildWithBuffering(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + _, _ = w.Write([]byte("event: response.created\n")) + _, _ = w.Write([]byte(`data: {"type":"response.created","response":{"id":"resp_1","model":"gpt-5.6-luna"}}` + "\n\n")) + _, _ = w.Write([]byte("event: keepalive\n")) + _, _ = w.Write([]byte(`data: {"type":"keepalive","sequence_number":3}` + "\n\n")) + _, _ = w.Write([]byte("event: response.completed\n")) + _, _ = w.Write([]byte(`data: {"type":"response.completed","response":{"id":"resp_1","status":"completed","output":[]}}` + "\n\n")) + })) + defer server.Close() + + cfg := &config.Config{} + cfg.Codex.StreamBootstrapBuffering = true + executor := NewCodexExecutor(cfg) + auth := &cliproxyauth.Auth{Attributes: map[string]string{ + "base_url": server.URL, + "api_key": "test", + }} + + res, err := executor.ExecuteStream(context.Background(), auth, cliproxyexecutor.Request{ + Model: "gpt-5.6-luna", + Payload: []byte(`{"model":"gpt-5.6-luna","input":"test"}`), + }, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FromString("openai-response"), + Stream: true, + Headers: http.Header{"User-Agent": []string{"grok-shell/1.0.5"}}, + }) + if err != nil { + t.Fatalf("ExecuteStream error: %v", err) + } + + var fullOutput bytes.Buffer + timeout := time.After(3 * time.Second) + done := false + for !done { + select { + case chunk, ok := <-res.Chunks: + if !ok { + done = true + break + } + if chunk.Err != nil { + t.Fatalf("unexpected chunk error: %v", chunk.Err) + } + fullOutput.Write(chunk.Payload) + case <-timeout: + t.Fatal("timed out reading stream chunks") + } + } + + outputStr := fullOutput.String() + if strings.Contains(outputStr, `{"type":"keepalive"`) || strings.Contains(outputStr, "event: keepalive") { + t.Fatalf("output must not contain keepalive event/data frame, got:\n%s", outputStr) + } + if !strings.Contains(outputStr, ": keepalive") { + t.Fatalf("output must contain ': keepalive' SSE comment, got:\n%s", outputStr) + } +} + +func TestCodexExecutorExecuteStream_GrokBuildDetectedFromGinContext(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + _, _ = w.Write([]byte("event: response.created\n")) + _, _ = w.Write([]byte(`data: {"type":"response.created","response":{"id":"resp_1","model":"gpt-5.6-luna"}}` + "\n\n")) + _, _ = w.Write([]byte("event: keepalive\n")) + _, _ = w.Write([]byte(`data: {"type":"keepalive","sequence_number":3}` + "\n\n")) + _, _ = w.Write([]byte("event: response.completed\n")) + _, _ = w.Write([]byte(`data: {"type":"response.completed","response":{"id":"resp_1","status":"completed","output":[]}}` + "\n\n")) + })) + defer server.Close() + + executor := NewCodexExecutor(&config.Config{}) + auth := &cliproxyauth.Auth{Attributes: map[string]string{ + "base_url": server.URL, + "api_key": "test", + }} + + gin.SetMode(gin.TestMode) + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil) + c.Request.Header.Set("User-Agent", "grok-pager/1.0.5") + ctx := context.WithValue(context.Background(), "gin", c) + + res, err := executor.ExecuteStream(ctx, auth, cliproxyexecutor.Request{ + Model: "gpt-5.6-luna", + Payload: []byte(`{"model":"gpt-5.6-luna","input":"test"}`), + }, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FromString("openai-response"), + Stream: true, + }) + if err != nil { + t.Fatalf("ExecuteStream error: %v", err) + } + + var fullOutput bytes.Buffer + timeout := time.After(3 * time.Second) + done := false + for !done { + select { + case chunk, ok := <-res.Chunks: + if !ok { + done = true + break + } + if chunk.Err != nil { + t.Fatalf("unexpected chunk error: %v", chunk.Err) + } + fullOutput.Write(chunk.Payload) + case <-timeout: + t.Fatal("timed out reading stream chunks") + } + } + + outputStr := fullOutput.String() + if strings.Contains(outputStr, `{"type":"keepalive"`) || strings.Contains(outputStr, "event: keepalive") { + t.Fatalf("output must not contain keepalive event/data frame, got:\n%s", outputStr) + } + if !strings.Contains(outputStr, ": keepalive") { + t.Fatalf("output must contain ': keepalive' SSE comment, got:\n%s", outputStr) + } +} + +func TestCodexExecutorExecuteStream_NonGrokClientKeepsVerbatim(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + _, _ = w.Write([]byte("event: response.created\n")) + _, _ = w.Write([]byte(`data: {"type":"response.created","response":{"id":"resp_1","model":"gpt-5.6-luna"}}` + "\n\n")) + _, _ = w.Write([]byte("event: keepalive\n")) + _, _ = w.Write([]byte(`data: {"type":"keepalive","sequence_number":3}` + "\n\n")) + _, _ = w.Write([]byte("event: response.completed\n")) + _, _ = w.Write([]byte(`data: {"type":"response.completed","response":{"id":"resp_1","status":"completed","output":[]}}` + "\n\n")) + })) + defer server.Close() + + executor := NewCodexExecutor(&config.Config{}) + auth := &cliproxyauth.Auth{Attributes: map[string]string{ + "base_url": server.URL, + "api_key": "test", + }} + + res, err := executor.ExecuteStream(context.Background(), auth, cliproxyexecutor.Request{ + Model: "gpt-5.6-luna", + Payload: []byte(`{"model":"gpt-5.6-luna","input":"test"}`), + }, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FromString("openai-response"), + Stream: true, + Headers: http.Header{"User-Agent": []string{"curl/8.7.1"}}, + }) + if err != nil { + t.Fatalf("ExecuteStream error: %v", err) + } + + var fullOutput bytes.Buffer + timeout := time.After(3 * time.Second) + done := false + for !done { + select { + case chunk, ok := <-res.Chunks: + if !ok { + done = true + break + } + if chunk.Err != nil { + t.Fatalf("unexpected chunk error: %v", chunk.Err) + } + fullOutput.Write(chunk.Payload) + case <-timeout: + t.Fatal("timed out reading stream chunks") + } + } + + outputStr := fullOutput.String() + if !strings.Contains(outputStr, `{"type":"keepalive"`) && !strings.Contains(outputStr, "event: keepalive") { + t.Fatalf("expected verbatim keepalive for non-Grok client, got:\n%s", outputStr) + } +} diff --git a/internal/runtime/executor/codex_executor_stream.go b/internal/runtime/executor/codex_executor_stream.go index d6e4918f3..aedf5ae9d 100644 --- a/internal/runtime/executor/codex_executor_stream.go +++ b/internal/runtime/executor/codex_executor_stream.go @@ -8,6 +8,7 @@ import ( "net/http" "strings" + "github.com/router-for-me/CLIProxyAPI/v7/internal/client/grokbuild" "github.com/router-for-me/CLIProxyAPI/v7/internal/config" "github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor/helps" "github.com/router-for-me/CLIProxyAPI/v7/internal/thinking" @@ -38,6 +39,7 @@ func (e *CodexExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Au from := opts.SourceFormat responseFormat := cliproxyexecutor.ResponseFormatOrSource(opts) + isGrokClient := grokbuild.IsGrokClientContext(ctx, opts.Headers) to := sdktranslator.FromString("codex") originalPayloadSource := req.Payload if len(opts.OriginalRequest) > 0 { @@ -163,7 +165,10 @@ func (e *CodexExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Au isHandshake := false terminalSuccess := false - if bytes.HasPrefix(line, dataTag) { + if transformed, ok := grokbuild.TransformKeepaliveSSELine(translatedLine, isGrokClient); ok { + translatedLine = transformed + isHandshake = true + } else if bytes.HasPrefix(line, dataTag) { data := bytes.TrimSpace(line[5:]) data = helps.RestoreCodexMultiAgentV2Response(data, optimizeMultiAgentV2) translatedLine = append([]byte("data: "), data...) @@ -287,7 +292,9 @@ func (e *CodexExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Au translatedLine := bytes.Clone(line) terminalSuccess := false - if bytes.HasPrefix(line, dataTag) { + if transformed, ok := grokbuild.TransformKeepaliveSSELine(translatedLine, isGrokClient); ok { + translatedLine = transformed + } else if bytes.HasPrefix(line, dataTag) { data := bytes.TrimSpace(line[5:]) data = helps.RestoreCodexMultiAgentV2Response(data, optimizeMultiAgentV2) translatedLine = append([]byte("data: "), data...)