Skip to content

Commit 2753d9f

Browse files
committed
feat: add validation for Claude streaming responses
- Implemented `validateClaudeStreamingResponse` to ensure upstream streaming data integrity. - Added new tests to verify response validation, including empty streams, error events, incomplete streams, and valid streams. - Integrated validation logic into the Claude executor's streaming handler, returning detailed errors for malformed upstream data. Fixed: #2193
1 parent 672fdd1 commit 2753d9f

2 files changed

Lines changed: 169 additions & 0 deletions

File tree

internal/runtime/executor/claude_executor.go

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -285,6 +285,10 @@ func (e *ClaudeExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, r
285285
}
286286
helps.AppendAPIResponseChunk(ctx, e.cfg, data)
287287
if stream {
288+
if errValidate := validateClaudeStreamingResponse(data); errValidate != nil {
289+
helps.RecordAPIResponseError(ctx, e.cfg, errValidate)
290+
return resp, errValidate
291+
}
288292
lines := bytes.Split(data, []byte("\n"))
289293
for _, line := range lines {
290294
if detail, ok := helps.ParseClaudeStreamUsage(line); ok {
@@ -533,6 +537,64 @@ func (e *ClaudeExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.A
533537
return &cliproxyexecutor.StreamResult{Headers: httpResp.Header.Clone(), Chunks: out}, nil
534538
}
535539

540+
func validateClaudeStreamingResponse(data []byte) error {
541+
scanner := bufio.NewScanner(bytes.NewReader(data))
542+
scanner.Buffer(nil, 52_428_800)
543+
544+
hasData := false
545+
hasMessageStart := false
546+
hasMessageDelta := false
547+
548+
for scanner.Scan() {
549+
line := bytes.TrimSpace(scanner.Bytes())
550+
if len(line) == 0 || !bytes.HasPrefix(line, []byte("data:")) {
551+
continue
552+
}
553+
payload := bytes.TrimSpace(line[len("data:"):])
554+
if len(payload) == 0 || bytes.Equal(payload, []byte("[DONE]")) {
555+
continue
556+
}
557+
hasData = true
558+
if !gjson.ValidBytes(payload) {
559+
return statusErr{code: http.StatusBadGateway, msg: "claude executor: upstream returned malformed stream data"}
560+
}
561+
562+
root := gjson.ParseBytes(payload)
563+
switch root.Get("type").String() {
564+
case "error":
565+
message := strings.TrimSpace(root.Get("error.message").String())
566+
if message == "" {
567+
message = strings.TrimSpace(root.Get("error.type").String())
568+
}
569+
if message == "" {
570+
message = "unknown upstream error"
571+
}
572+
return statusErr{code: http.StatusBadGateway, msg: "claude executor: upstream returned error event: " + message}
573+
case "message_start":
574+
message := root.Get("message")
575+
if strings.TrimSpace(message.Get("id").String()) == "" || strings.TrimSpace(message.Get("model").String()) == "" {
576+
return statusErr{code: http.StatusBadGateway, msg: "claude executor: upstream stream message_start is missing id or model"}
577+
}
578+
hasMessageStart = true
579+
case "message_delta":
580+
hasMessageDelta = true
581+
}
582+
}
583+
if errScan := scanner.Err(); errScan != nil {
584+
return errScan
585+
}
586+
if !hasData {
587+
return statusErr{code: http.StatusBadGateway, msg: "claude executor: upstream returned empty stream response"}
588+
}
589+
if !hasMessageStart {
590+
return statusErr{code: http.StatusBadGateway, msg: "claude executor: upstream stream response is missing message_start"}
591+
}
592+
if !hasMessageDelta {
593+
return statusErr{code: http.StatusBadGateway, msg: "claude executor: upstream stream response ended before message completion"}
594+
}
595+
return nil
596+
}
597+
536598
func (e *ClaudeExecutor) CountTokens(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
537599
baseModel := thinking.ParseSuffix(req.Model).ModelName
538600

internal/runtime/executor/claude_executor_test.go

Lines changed: 107 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -936,6 +936,113 @@ func TestClaudeExecutor_GeneratesNewUserIDByDefault(t *testing.T) {
936936
}
937937
}
938938

939+
func TestClaudeExecutor_ExecuteOpenAINonStreamRejectsEmptyClaudeStream(t *testing.T) {
940+
_, err := executeOpenAIChatCompletionThroughClaude(t, "")
941+
if err == nil {
942+
t.Fatal("Execute error = nil, want empty stream error")
943+
}
944+
assertStatusErr(t, err, http.StatusBadGateway)
945+
if !strings.Contains(err.Error(), "empty stream response") {
946+
t.Fatalf("Execute error = %q, want empty stream response", err.Error())
947+
}
948+
}
949+
950+
func TestClaudeExecutor_ExecuteOpenAINonStreamRejectsClaudeErrorEvent(t *testing.T) {
951+
body := `data: {"type":"error","error":{"type":"overloaded_error","message":"upstream overloaded"}}` + "\n"
952+
_, err := executeOpenAIChatCompletionThroughClaude(t, body)
953+
if err == nil {
954+
t.Fatal("Execute error = nil, want upstream error event")
955+
}
956+
assertStatusErr(t, err, http.StatusBadGateway)
957+
if !strings.Contains(err.Error(), "upstream overloaded") {
958+
t.Fatalf("Execute error = %q, want upstream overloaded", err.Error())
959+
}
960+
}
961+
962+
func TestClaudeExecutor_ExecuteOpenAINonStreamRejectsIncompleteClaudeStream(t *testing.T) {
963+
body := strings.Join([]string{
964+
`data: {"type":"message_start","message":{"id":"msg_123","model":"claude-3-5-sonnet-20241022"}}`,
965+
`data: {"type":"message_stop"}`,
966+
``,
967+
}, "\n")
968+
969+
_, err := executeOpenAIChatCompletionThroughClaude(t, body)
970+
if err == nil {
971+
t.Fatal("Execute error = nil, want incomplete stream error")
972+
}
973+
assertStatusErr(t, err, http.StatusBadGateway)
974+
if !strings.Contains(err.Error(), "ended before message completion") {
975+
t.Fatalf("Execute error = %q, want incomplete stream error", err.Error())
976+
}
977+
}
978+
979+
func TestClaudeExecutor_ExecuteOpenAINonStreamConvertsValidClaudeStream(t *testing.T) {
980+
body := strings.Join([]string{
981+
`event: message_start`,
982+
`data: {"type":"message_start","message":{"id":"msg_123","model":"claude-3-5-sonnet-20241022"}}`,
983+
`event: content_block_delta`,
984+
`data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"ok"}}`,
985+
`event: message_delta`,
986+
`data: {"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"input_tokens":2,"output_tokens":1}}`,
987+
`event: message_stop`,
988+
`data: {"type":"message_stop"}`,
989+
``,
990+
}, "\n")
991+
992+
resp, err := executeOpenAIChatCompletionThroughClaude(t, body)
993+
if err != nil {
994+
t.Fatalf("Execute error: %v", err)
995+
}
996+
if got := gjson.GetBytes(resp.Payload, "id").String(); got != "msg_123" {
997+
t.Fatalf("response id = %q, want msg_123; payload=%s", got, string(resp.Payload))
998+
}
999+
if got := gjson.GetBytes(resp.Payload, "model").String(); got != "claude-3-5-sonnet-20241022" {
1000+
t.Fatalf("response model = %q, want claude-3-5-sonnet-20241022", got)
1001+
}
1002+
if got := gjson.GetBytes(resp.Payload, "choices.0.message.content").String(); got != "ok" {
1003+
t.Fatalf("response content = %q, want ok", got)
1004+
}
1005+
if got := gjson.GetBytes(resp.Payload, "usage.total_tokens").Int(); got != 3 {
1006+
t.Fatalf("usage.total_tokens = %d, want 3", got)
1007+
}
1008+
}
1009+
1010+
func executeOpenAIChatCompletionThroughClaude(t *testing.T, upstreamBody string) (cliproxyexecutor.Response, error) {
1011+
t.Helper()
1012+
1013+
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
1014+
w.Header().Set("Content-Type", "text/event-stream")
1015+
_, _ = w.Write([]byte(upstreamBody))
1016+
}))
1017+
defer server.Close()
1018+
1019+
executor := NewClaudeExecutor(&config.Config{})
1020+
auth := &cliproxyauth.Auth{Attributes: map[string]string{
1021+
"api_key": "key-123",
1022+
"base_url": server.URL,
1023+
}}
1024+
payload := []byte(`{"model":"claude-3-5-sonnet-20241022","messages":[{"role":"user","content":"hi"}]}`)
1025+
1026+
return executor.Execute(context.Background(), auth, cliproxyexecutor.Request{
1027+
Model: "claude-3-5-sonnet-20241022",
1028+
Payload: payload,
1029+
}, cliproxyexecutor.Options{
1030+
SourceFormat: sdktranslator.FromString("openai"),
1031+
})
1032+
}
1033+
1034+
func assertStatusErr(t *testing.T, err error, want int) {
1035+
t.Helper()
1036+
1037+
status, ok := err.(interface{ StatusCode() int })
1038+
if !ok {
1039+
t.Fatalf("error %T does not expose StatusCode", err)
1040+
}
1041+
if got := status.StatusCode(); got != want {
1042+
t.Fatalf("StatusCode() = %d, want %d", got, want)
1043+
}
1044+
}
1045+
9391046
func TestStripClaudeToolPrefixFromResponse_NestedToolReference(t *testing.T) {
9401047
input := []byte(`{"content":[{"type":"tool_result","tool_use_id":"toolu_123","content":[{"type":"tool_reference","tool_name":"proxy_mcp__nia__manage_resource"}]}]}`)
9411048
out := stripClaudeToolPrefixFromResponse(input, "proxy_")

0 commit comments

Comments
 (0)