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
This commit is contained in:
@@ -285,6 +285,10 @@ func (e *ClaudeExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, r
|
||||
}
|
||||
helps.AppendAPIResponseChunk(ctx, e.cfg, data)
|
||||
if stream {
|
||||
if errValidate := validateClaudeStreamingResponse(data); errValidate != nil {
|
||||
helps.RecordAPIResponseError(ctx, e.cfg, errValidate)
|
||||
return resp, errValidate
|
||||
}
|
||||
lines := bytes.Split(data, []byte("\n"))
|
||||
for _, line := range lines {
|
||||
if detail, ok := helps.ParseClaudeStreamUsage(line); ok {
|
||||
@@ -533,6 +537,64 @@ func (e *ClaudeExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.A
|
||||
return &cliproxyexecutor.StreamResult{Headers: httpResp.Header.Clone(), Chunks: out}, nil
|
||||
}
|
||||
|
||||
func validateClaudeStreamingResponse(data []byte) error {
|
||||
scanner := bufio.NewScanner(bytes.NewReader(data))
|
||||
scanner.Buffer(nil, 52_428_800)
|
||||
|
||||
hasData := false
|
||||
hasMessageStart := false
|
||||
hasMessageDelta := false
|
||||
|
||||
for scanner.Scan() {
|
||||
line := bytes.TrimSpace(scanner.Bytes())
|
||||
if len(line) == 0 || !bytes.HasPrefix(line, []byte("data:")) {
|
||||
continue
|
||||
}
|
||||
payload := bytes.TrimSpace(line[len("data:"):])
|
||||
if len(payload) == 0 || bytes.Equal(payload, []byte("[DONE]")) {
|
||||
continue
|
||||
}
|
||||
hasData = true
|
||||
if !gjson.ValidBytes(payload) {
|
||||
return statusErr{code: http.StatusBadGateway, msg: "claude executor: upstream returned malformed stream data"}
|
||||
}
|
||||
|
||||
root := gjson.ParseBytes(payload)
|
||||
switch root.Get("type").String() {
|
||||
case "error":
|
||||
message := strings.TrimSpace(root.Get("error.message").String())
|
||||
if message == "" {
|
||||
message = strings.TrimSpace(root.Get("error.type").String())
|
||||
}
|
||||
if message == "" {
|
||||
message = "unknown upstream error"
|
||||
}
|
||||
return statusErr{code: http.StatusBadGateway, msg: "claude executor: upstream returned error event: " + message}
|
||||
case "message_start":
|
||||
message := root.Get("message")
|
||||
if strings.TrimSpace(message.Get("id").String()) == "" || strings.TrimSpace(message.Get("model").String()) == "" {
|
||||
return statusErr{code: http.StatusBadGateway, msg: "claude executor: upstream stream message_start is missing id or model"}
|
||||
}
|
||||
hasMessageStart = true
|
||||
case "message_delta":
|
||||
hasMessageDelta = true
|
||||
}
|
||||
}
|
||||
if errScan := scanner.Err(); errScan != nil {
|
||||
return errScan
|
||||
}
|
||||
if !hasData {
|
||||
return statusErr{code: http.StatusBadGateway, msg: "claude executor: upstream returned empty stream response"}
|
||||
}
|
||||
if !hasMessageStart {
|
||||
return statusErr{code: http.StatusBadGateway, msg: "claude executor: upstream stream response is missing message_start"}
|
||||
}
|
||||
if !hasMessageDelta {
|
||||
return statusErr{code: http.StatusBadGateway, msg: "claude executor: upstream stream response ended before message completion"}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (e *ClaudeExecutor) CountTokens(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
|
||||
baseModel := thinking.ParseSuffix(req.Model).ModelName
|
||||
|
||||
|
||||
Reference in New Issue
Block a user