From 3a9e5dc7c3a9f3f719508a79bd52f06924cbfa04 Mon Sep 17 00:00:00 2001 From: Ed1s0nZ Date: Mon, 14 Sep 2026 16:16:56 +0800 Subject: [PATCH] fix: stream summaries and validate completion metadata --- .../multiagent/eino_agentic_summarize_test.go | 5 +- internal/multiagent/eino_summarize.go | 3 +- .../multiagent/eino_summarize_model_guard.go | 113 ++++++- .../eino_summarize_model_guard_test.go | 296 ++++++++++++++++++ internal/multiagent/eino_summarize_test.go | 2 +- 5 files changed, 410 insertions(+), 9 deletions(-) diff --git a/internal/multiagent/eino_agentic_summarize_test.go b/internal/multiagent/eino_agentic_summarize_test.go index 16b81e7a..5ffb7c8b 100644 --- a/internal/multiagent/eino_agentic_summarize_test.go +++ b/internal/multiagent/eino_agentic_summarize_test.go @@ -10,6 +10,7 @@ import ( "github.com/cloudwego/eino/adk" "github.com/cloudwego/eino/schema" + schemaopenai "github.com/cloudwego/eino/schema/openai" ) func TestNewEinoAgenticSummarizationMiddlewareCompactsWithNativeTypedMiddleware(t *testing.T) { @@ -18,7 +19,8 @@ func TestNewEinoAgenticSummarizationMiddlewareCompactsWithNativeTypedMiddleware( emit := false summaryModel := &capturingAgenticChatModel{ output: &schema.AgenticMessage{ - Role: schema.AgenticRoleTypeAssistant, + Role: schema.AgenticRoleTypeAssistant, + ResponseMeta: &schema.AgenticResponseMeta{OpenAIExtension: &schemaopenai.ResponseMetaExtension{Status: schemaopenai.ResponseStatusCompleted}}, ContentBlocks: []*schema.ContentBlock{schema.NewContentBlock(&schema.AssistantGenText{Text: `检查历史 ## 1. 授权范围与约束 @@ -246,6 +248,7 @@ func TestAppendEinoAgenticChatModelTailMiddlewaresIncludesTypedSummarization(t * func agenticAssistantTextMessage(text string) *schema.AgenticMessage { return &schema.AgenticMessage{ Role: schema.AgenticRoleTypeAssistant, + ResponseMeta: &schema.AgenticResponseMeta{OpenAIExtension: &schemaopenai.ResponseMetaExtension{Status: schemaopenai.ResponseStatusCompleted}}, ContentBlocks: []*schema.ContentBlock{schema.NewContentBlock(&schema.AssistantGenText{Text: text})}, } } diff --git a/internal/multiagent/eino_summarize.go b/internal/multiagent/eino_summarize.go index b7cd7e41..50678478 100644 --- a/internal/multiagent/eino_summarize.go +++ b/internal/multiagent/eino_summarize.go @@ -291,7 +291,8 @@ func newEinoSummarizationMiddleware( return mw, nil } -// newEinoSummarizationModelOptions applies only to summarization Generate calls +// newEinoSummarizationModelOptions applies only to summary requests (streamed +// internally by the summary model guard while exposing Generate to Eino) // on the shared main model. Summary generation should be plain-text and cheap: // strip provider reasoning/thinking controls so DeepSeek/OpenAI-compatible // endpoints do not spend the reserved output budget on invisible reasoning. diff --git a/internal/multiagent/eino_summarize_model_guard.go b/internal/multiagent/eino_summarize_model_guard.go index 3a770d95..c0b9757d 100644 --- a/internal/multiagent/eino_summarize_model_guard.go +++ b/internal/multiagent/eino_summarize_model_guard.go @@ -4,10 +4,13 @@ import ( "context" "errors" "fmt" + "io" "strings" + "github.com/cloudwego/eino-ext/components/model/agenticopenai" "github.com/cloudwego/eino/components/model" "github.com/cloudwego/eino/schema" + schemaopenai "github.com/cloudwego/eino/schema/openai" ) type einoSummarizationModelError struct { @@ -56,12 +59,22 @@ func newNonEmptySummaryChatModel(base model.BaseChatModel) model.BaseChatModel { } func (m *nonEmptySummaryChatModel) Generate(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) { - out, err := m.base.Generate(ctx, input, opts...) + if err := ctx.Err(); err != nil { + return nil, newEinoSummarizationModelError(err) + } + stream, err := m.base.Stream(ctx, input, opts...) if err != nil { - return out, newEinoSummarizationModelError(err) + return nil, newEinoSummarizationModelError(err) + } + out, err := collectSummaryStream(ctx, stream, schema.ConcatMessages) + if err != nil { + return nil, newEinoSummarizationModelError(err) } if strings.TrimSpace(classicAssistantTextContent(out)) == "" { - return out, newEinoSummarizationModelError(fmt.Errorf("summary content is empty: %s", classicSummaryEmptyDiagnostics(out))) + return nil, newEinoSummarizationModelError(fmt.Errorf("summary content is empty: %s", classicSummaryEmptyDiagnostics(out))) + } + if err := validateClassicSummaryCompletion(out); err != nil { + return nil, newEinoSummarizationModelError(err) } return out, nil } @@ -79,12 +92,22 @@ func newNonEmptyAgenticSummaryModel(base model.BaseModel[*schema.AgenticMessage] } func (m *nonEmptyAgenticSummaryModel) Generate(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.AgenticMessage, error) { - out, err := m.base.Generate(ctx, input, opts...) + if err := ctx.Err(); err != nil { + return nil, newEinoSummarizationModelError(err) + } + stream, err := m.base.Stream(ctx, input, opts...) if err != nil { - return out, newEinoSummarizationModelError(err) + return nil, newEinoSummarizationModelError(err) + } + out, err := collectSummaryStream(ctx, stream, schema.ConcatAgenticMessages) + if err != nil { + return nil, newEinoSummarizationModelError(err) } if strings.TrimSpace(agenticAssistantTextContent(out)) == "" { - return out, newEinoSummarizationModelError(fmt.Errorf("summary content is empty: %s", agenticSummaryEmptyDiagnostics(out))) + return nil, newEinoSummarizationModelError(fmt.Errorf("summary content is empty: %s", agenticSummaryEmptyDiagnostics(out))) + } + if err := validateAgenticSummaryCompletion(out); err != nil { + return nil, newEinoSummarizationModelError(err) } return out, nil } @@ -212,3 +235,81 @@ func agenticSummaryEmptyDiagnostics(msg *schema.AgenticMessage) string { } return strings.Join(fields, " ") } + +// Summarization expects a complete message, but providers such as Claude require +// streaming for large output budgets. Only finalize after a clean EOF: a partial +// summary must never replace the original conversation when the stream fails. +// Recv relies on the underlying model honoring ctx to unblock network reads; +// do not start detached receive goroutines, which can leak on stalled providers. +func collectSummaryStream[T any](ctx context.Context, stream *schema.StreamReader[*T], concat func([]*T) (*T, error)) (*T, error) { + if stream == nil { + return nil, fmt.Errorf("summary model returned nil stream") + } + defer stream.Close() + var chunks []*T + for { + if err := ctx.Err(); err != nil { + return nil, err + } + chunk, err := stream.Recv() + if ctxErr := ctx.Err(); ctxErr != nil { + return nil, ctxErr + } + if errors.Is(err, io.EOF) { + break + } + if err != nil { + return nil, err + } + if chunk != nil { + chunks = append(chunks, chunk) + } + } + if len(chunks) == 0 { + return nil, fmt.Errorf("summary content is empty: model returned no chunks") + } + return concat(chunks) +} + +// EOF only means the transport ended. Require positive completion metadata before +// allowing the middleware to replace conversation history with the summary. +func validateClassicSummaryCompletion(msg *schema.Message) error { + reason := "" + if msg.ResponseMeta != nil { + reason = msg.ResponseMeta.FinishReason + } + if reason != "stop" || len(msg.ToolCalls) != 0 { + return fmt.Errorf("summary did not complete: finish_reason=%q tool_calls=%d", reason, len(msg.ToolCalls)) + } + return nil +} + +func validateAgenticSummaryCompletion(msg *schema.AgenticMessage) error { + if meta := msg.ResponseMeta; meta != nil { + if ext, ok := meta.Extension.(*agenticopenai.ChatResponseMetaExtension); ok && ext != nil { + if ext.FinishReason == "stop" { + return nil + } + return fmt.Errorf("summary did not complete: openai chat finish_reason=%q", ext.FinishReason) + } + if ext := meta.ClaudeExtension; ext != nil { + if ext.StopReason == "end_turn" { + return nil + } + return fmt.Errorf("summary did not complete: claude stop_reason=%q", ext.StopReason) + } + if ext := meta.OpenAIExtension; ext != nil { + if ext.Status == schemaopenai.ResponseStatusCompleted && ext.Error == nil && ext.IncompleteDetails == nil { + return nil + } + return fmt.Errorf("summary did not complete: openai status=%q", ext.Status) + } + if ext := meta.GeminiExtension; ext != nil { + if ext.FinishReason == "STOP" { + return nil + } + return fmt.Errorf("summary did not complete: gemini finish_reason=%q", ext.FinishReason) + } + } + return fmt.Errorf("summary did not complete: missing completion metadata") +} diff --git a/internal/multiagent/eino_summarize_model_guard_test.go b/internal/multiagent/eino_summarize_model_guard_test.go index b52f7c71..70d701d4 100644 --- a/internal/multiagent/eino_summarize_model_guard_test.go +++ b/internal/multiagent/eino_summarize_model_guard_test.go @@ -2,11 +2,21 @@ package multiagent import ( "context" + "encoding/json" + "errors" + "fmt" + "net/http" + "net/http/httptest" "strings" "testing" + "time" + + "cyberstrike-ai/internal/config" "github.com/cloudwego/eino/components/model" "github.com/cloudwego/eino/schema" + "github.com/cloudwego/eino/schema/claude" + schemaopenai "github.com/cloudwego/eino/schema/openai" ) type guardClassicSummaryModel struct { @@ -99,3 +109,289 @@ func TestNonEmptyAgenticSummaryModelReportsEmptyContentDiagnostics(t *testing.T) } } } + +// Generate deliberately reproduces the SDK rejection; summaries must use Stream. +type streamingSummaryTestModel[T any] struct { + stream *schema.StreamReader[T] + err error + input []T + opts []model.Option +} + +func (m *streamingSummaryTestModel[T]) Generate(context.Context, []T, ...model.Option) (T, error) { + var zero T + return zero, errors.New("streaming is required for operations that may take longer than 10 minutes") +} +func (m *streamingSummaryTestModel[T]) Stream(_ context.Context, input []T, opts ...model.Option) (*schema.StreamReader[T], error) { + m.input, m.opts = input, opts + return m.stream, m.err +} + +func TestSummaryGenerateUsesStream(t *testing.T) { + t.Run("classic", func(t *testing.T) { + tail := schema.AssistantMessage("摘要", nil) + tail.ResponseMeta = &schema.ResponseMeta{FinishReason: "stop", Usage: &schema.TokenUsage{TotalTokens: 42}} + base := &streamingSummaryTestModel[*schema.Message]{stream: schema.StreamReaderFromArray([]*schema.Message{schema.AssistantMessage("完整", nil), nil, tail})} + input := []*schema.Message{schema.UserMessage("history")} + out, err := newNonEmptySummaryChatModel(base).Generate(context.Background(), input, model.WithMaxTokens(64000)) + if err != nil { + t.Fatal(err) + } + if out.Content != "完整摘要" || out.ResponseMeta.Usage.TotalTokens != 42 { + t.Fatalf("lost streamed content or usage: %+v", out) + } + if base.input[0] != input[0] || *model.GetCommonOptions(nil, base.opts...).MaxTokens != 64000 { + t.Fatal("input/options not forwarded") + } + }) + t.Run("agentic", func(t *testing.T) { + chunk := func(text string) *schema.AgenticMessage { + block := schema.NewContentBlock(&schema.AssistantGenText{Text: text}) + block.StreamingMeta = &schema.StreamingMeta{Index: 0} + return &schema.AgenticMessage{Role: schema.AgenticRoleTypeAssistant, ContentBlocks: []*schema.ContentBlock{block}} + } + tail := chunk("摘要") + tail.ResponseMeta = &schema.AgenticResponseMeta{ClaudeExtension: &claude.ResponseMetaExtension{StopReason: "end_turn"}, TokenUsage: &schema.TokenUsage{TotalTokens: 42}} + base := &streamingSummaryTestModel[*schema.AgenticMessage]{stream: schema.StreamReaderFromArray([]*schema.AgenticMessage{chunk("完整"), nil, tail})} + input := []*schema.AgenticMessage{chunk("history")} + out, err := newNonEmptyAgenticSummaryModel(base).Generate(context.Background(), input, model.WithMaxTokens(64000)) + if err != nil { + t.Fatal(err) + } + if agenticAssistantTextContent(out) != "完整摘要" || out.ResponseMeta.TokenUsage.TotalTokens != 42 { + t.Fatalf("lost streamed content or usage: %+v", out) + } + if base.input[0] != input[0] || *model.GetCommonOptions(nil, base.opts...).MaxTokens != 64000 { + t.Fatal("input/options not forwarded") + } + }) +} + +func TestSummaryStreamFailuresDoNotReturnPartialSummary(t *testing.T) { + failure := errors.New("connection reset") + for _, kind := range []string{"start", "receive", "empty", "nil", "cancel"} { + t.Run(kind, func(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + base := &streamingSummaryTestModel[*schema.Message]{} + switch kind { + case "start": + base.err = failure + case "receive": + reader, writer := schema.Pipe[*schema.Message](2) + writer.Send(schema.AssistantMessage("partial", nil), nil) + writer.Send(nil, failure) + writer.Close() + base.stream = reader + case "empty": + base.stream = schema.StreamReaderFromArray([]*schema.Message{}) + case "cancel": + base.stream = schema.StreamReaderFromArray([]*schema.Message{schema.AssistantMessage("partial", nil)}) + cancel() + } + out, err := newNonEmptySummaryChatModel(base).Generate(ctx, nil) + if err == nil || out != nil { + t.Fatalf("out=%+v err=%v", out, err) + } + var wrapped *einoSummarizationModelError + if !errors.As(err, &wrapped) { + t.Fatalf("missing summary error wrapper: %v", err) + } + if (kind == "start" || kind == "receive") && !errors.Is(err, failure) { + t.Fatal("lost original error") + } + if kind == "cancel" && !errors.Is(err, context.Canceled) { + t.Fatal("lost cancellation") + } + }) + } +} + +func TestClaudeSummaryLargeBudgetStreamsThroughNativeSDK(t *testing.T) { + for _, mode := range []string{"complete", "truncated", "early_eof", "cancel_read"} { + t.Run(mode, func(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + started := make(chan struct{}) + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var body struct { + Stream bool `json:"stream"` + MaxTokens int `json:"max_tokens"` + } + if err := json.NewDecoder(r.Body).Decode(&body); err != nil { + t.Error(err) + w.WriteHeader(400) + return + } + if !body.Stream || body.MaxTokens != 64000 { + t.Errorf("unexpected request: %+v", body) + w.WriteHeader(400) + return + } + w.Header().Set("Content-Type", "text/event-stream") + events := []string{ + `{"type":"message_start","message":{"id":"msg_test","type":"message","role":"assistant","content":[],"model":"claude-sonnet-4-20250514","stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":10,"output_tokens":0}}}`, + `{"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}`, + `{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"完整"}}`, + `{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"摘要"}}`, + `{"type":"content_block_stop","index":0}`, + `{"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"output_tokens":2}}`, + `{"type":"message_stop"}`, + } + switch mode { + case "truncated": + events[5] = strings.ReplaceAll(events[5], "end_turn", "max_tokens") + case "early_eof", "cancel_read": + events = events[:4] + } + for _, event := range events { + var header struct { + Type string `json:"type"` + } + _ = json.Unmarshal([]byte(event), &header) + fmt.Fprintf(w, "event: %s\ndata: %s\n\n", header.Type, event) + } + if mode == "cancel_read" { + w.(http.Flusher).Flush() + close(started) + <-r.Context().Done() + } + })) + defer server.Close() + factory := newEinoAgenticChatModelFactory(server.Client(), nil, nil) + native, err := factory(ctx, config.OpenAIConfig{Provider: "claude", APIKey: "test-key", BaseURL: server.URL, Model: "claude-sonnet-4-20250514"}, einoModelModeNormal) + if err != nil { + t.Fatal(err) + } + input := EinoMessagesToAgentic([]*schema.Message{schema.UserMessage("summarize history")}) + opts := newEinoSummarizationModelOptions(64000, "claude-sonnet-4-20250514", "agentic", nil, nil) + if _, err = native.Generate(ctx, input, opts...); err == nil || !strings.Contains(err.Error(), "streaming is required") { + t.Fatalf("expected original SDK rejection, got %v", err) + } + if mode == "cancel_read" { + go func() { + select { + case <-started: + cancel() + case <-ctx.Done(): + } + }() + // Bound the test even if cancellation stops propagating to the HTTP request. + timer := time.AfterFunc(5*time.Second, cancel) + defer timer.Stop() + } + out, err := newNonEmptyAgenticSummaryModel(native).Generate(ctx, input, opts...) + if mode != "complete" { + if out != nil || err == nil { + t.Fatalf("accepted partial summary: out=%+v err=%v", out, err) + } + if mode == "cancel_read" && !errors.Is(err, context.Canceled) { + t.Fatalf("lost cancellation: %v", err) + } + return + } + if err != nil { + t.Fatal(err) + } + if agenticAssistantTextContent(out) != "完整摘要" { + t.Fatalf("unexpected summary: %+v", out) + } + if out.ResponseMeta == nil || out.ResponseMeta.TokenUsage == nil || out.ResponseMeta.TokenUsage.CompletionTokens != 2 { + t.Fatalf("missing usage: %+v", out.ResponseMeta) + } + }) + } +} + +func TestSummaryCompletionValidation(t *testing.T) { + for _, reason := range []string{"stop", "", "length", "content_filter", "tool_calls", "unknown"} { + t.Run("classic/"+reason, func(t *testing.T) { + msg := schema.AssistantMessage("partial or complete summary", nil) + msg.ResponseMeta = &schema.ResponseMeta{FinishReason: reason} + out, err := newNonEmptySummaryChatModel(&guardClassicSummaryModel{out: msg}).Generate(context.Background(), nil) + if reason == "stop" { + if err != nil { + t.Fatal(err) + } + return + } + if out != nil || err == nil { + t.Fatalf("out=%+v err=%v", out, err) + } + }) + } + for _, reason := range []string{"end_turn", "", "max_tokens", "stop_sequence", "tool_use", "pause_turn", "refusal"} { + t.Run("claude/"+reason, func(t *testing.T) { + msg := &schema.AgenticMessage{Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: []*schema.ContentBlock{schema.NewContentBlock(&schema.AssistantGenText{Text: "summary"})}, + ResponseMeta: &schema.AgenticResponseMeta{ClaudeExtension: &claude.ResponseMetaExtension{StopReason: reason}}} + out, err := newNonEmptyAgenticSummaryModel(&guardAgenticSummaryModel{out: msg}).Generate(context.Background(), nil) + if reason == "end_turn" { + if err != nil { + t.Fatal(err) + } + return + } + if out != nil || err == nil { + t.Fatalf("out=%+v err=%v", out, err) + } + }) + } + for _, status := range []schemaopenai.ResponseStatus{"completed", "incomplete", "failed", "cancelled", "in_progress", ""} { + t.Run("openai/"+string(status), func(t *testing.T) { + msg := &schema.AgenticMessage{Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: []*schema.ContentBlock{schema.NewContentBlock(&schema.AssistantGenText{Text: "summary"})}, + ResponseMeta: &schema.AgenticResponseMeta{OpenAIExtension: &schemaopenai.ResponseMetaExtension{Status: status}}} + out, err := newNonEmptyAgenticSummaryModel(&guardAgenticSummaryModel{out: msg}).Generate(context.Background(), nil) + if status == "completed" { + if err != nil { + t.Fatal(err) + } + return + } + if out != nil || err == nil { + t.Fatalf("out=%+v err=%v", out, err) + } + }) + } +} + +// Exercise the actual Chat Completions adapter: its finish reason is stored in +// ResponseMeta.Extension, unlike the OpenAI Responses API's OpenAIExtension. +func TestOpenAIChatSummaryStreamCompletion(t *testing.T) { + for _, reason := range []string{"stop", "length", ""} { + t.Run(reason, func(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var body struct { + Stream bool `json:"stream"` + } + if err := json.NewDecoder(r.Body).Decode(&body); err != nil || !body.Stream { + t.Errorf("expected streamed request: %+v, %v", body, err) + w.WriteHeader(http.StatusBadRequest) + return + } + w.Header().Set("Content-Type", "text/event-stream") + fmt.Fprint(w, "data: {\"id\":\"test\",\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\",\"content\":\"summary\"},\"finish_reason\":null}]}\n\n") + if reason != "" { + fmt.Fprintf(w, "data: {\"id\":\"test\",\"choices\":[{\"index\":0,\"delta\":{},\"finish_reason\":%q}]}\n\n", reason) + fmt.Fprint(w, "data: [DONE]\n\n") + } + })) + defer server.Close() + factory := newEinoAgenticChatModelFactory(server.Client(), nil, nil) + native, err := factory(context.Background(), config.OpenAIConfig{Provider: "openai", APIKey: "test-key", BaseURL: server.URL, Model: "gpt-4o"}, einoModelModeNormal) + if err != nil { + t.Fatal(err) + } + out, err := newNonEmptyAgenticSummaryModel(native).Generate(context.Background(), []*schema.AgenticMessage{schema.UserAgenticMessage("summarize")}) + if reason == "stop" { + if err != nil || agenticAssistantTextContent(out) != "summary" { + t.Fatalf("out=%+v err=%v", out, err) + } + } else if err == nil || out != nil { + t.Fatalf("accepted incomplete summary: out=%+v err=%v", out, err) + } + }) + } +} diff --git a/internal/multiagent/eino_summarize_test.go b/internal/multiagent/eino_summarize_test.go index 0997c863..cbd32354 100644 --- a/internal/multiagent/eino_summarize_test.go +++ b/internal/multiagent/eino_summarize_test.go @@ -392,7 +392,7 @@ func TestEinoSummarizationMiddlewareRetriesWhenSummaryModelReturnsEmpty(t *testi emit := false summaryModel := &capturingClassicChatModel{outputs: []*schema.Message{ schema.AssistantMessage("", nil), - schema.AssistantMessage("有效摘要：继续验证 SQL 注入路径", nil), + {Role: schema.Assistant, Content: "有效摘要：继续验证 SQL 注入路径", ResponseMeta: &schema.ResponseMeta{FinishReason: "stop"}}, }} appCfg := &config.Config{} appCfg.OpenAI.Model = "gpt-4o"