diff --git a/internal/multiagent/agentic_tool_call_arguments_sanitizer_middleware.go b/internal/multiagent/agentic_tool_call_arguments_sanitizer_middleware.go new file mode 100644 index 00000000..ad204bf3 --- /dev/null +++ b/internal/multiagent/agentic_tool_call_arguments_sanitizer_middleware.go @@ -0,0 +1,75 @@ +package multiagent + +import ( + "context" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/schema" + "go.uber.org/zap" +) + +// Apply the classic history repair to native function-call blocks too. Providers +// may parse arguments embedded in otherwise valid request JSON and reject the +// next turn if a previous model response contained truncated arguments. Repair +// only historical input; never repair or execute an incomplete new tool call. +type agenticToolCallArgumentsSanitizerMiddleware struct { + *adk.TypedBaseChatModelAgentMiddleware[*schema.AgenticMessage] + logger *zap.Logger + phase string +} + +func newAgenticToolCallArgumentsSanitizerMiddleware(logger *zap.Logger, phase string) adk.TypedChatModelAgentMiddleware[*schema.AgenticMessage] { + return &agenticToolCallArgumentsSanitizerMiddleware{ + TypedBaseChatModelAgentMiddleware: &adk.TypedBaseChatModelAgentMiddleware[*schema.AgenticMessage]{}, + logger: logger, phase: phase, + } +} + +func (m *agenticToolCallArgumentsSanitizerMiddleware) BeforeModelRewriteState( + ctx context.Context, state *adk.TypedChatModelAgentState[*schema.AgenticMessage], + _ *adk.TypedModelContext[*schema.AgenticMessage], +) (context.Context, *adk.TypedChatModelAgentState[*schema.AgenticMessage], error) { + if m == nil || state == nil { + return ctx, state, nil + } + var out []*schema.AgenticMessage + repaired := 0 + for i, msg := range state.Messages { + if msg == nil || msg.Role != schema.AgenticRoleTypeAssistant { + continue + } + var blocks []*schema.ContentBlock + for j, block := range msg.ContentBlocks { + if block == nil || block.FunctionToolCall == nil || validToolArgumentsJSONObject(block.FunctionToolCall.Arguments) { + continue + } + if blocks == nil { + blocks = append([]*schema.ContentBlock(nil), msg.ContentBlocks...) + } + clonedBlock := *block + clonedCall := *block.FunctionToolCall + clonedCall.Arguments = repairedMalformedToolArguments + clonedBlock.FunctionToolCall = &clonedCall + blocks[j] = &clonedBlock + repaired++ + } + if blocks == nil { + continue + } + if out == nil { + out = append([]*schema.AgenticMessage(nil), state.Messages...) + } + clonedMessage := *msg + clonedMessage.ContentBlocks = blocks + out[i] = &clonedMessage + } + if repaired == 0 { + return ctx, state, nil + } + if m.logger != nil { + m.logger.Warn("eino agentic malformed tool-call arguments repaired before model call", zap.String("phase", m.phase), zap.Int("repaired_calls", repaired)) + } + clonedState := *state + clonedState.Messages = out + return ctx, &clonedState, nil +} diff --git a/internal/multiagent/agentic_tool_call_arguments_sanitizer_middleware_test.go b/internal/multiagent/agentic_tool_call_arguments_sanitizer_middleware_test.go new file mode 100644 index 00000000..e1878e02 --- /dev/null +++ b/internal/multiagent/agentic_tool_call_arguments_sanitizer_middleware_test.go @@ -0,0 +1,214 @@ +package multiagent + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "cyberstrike-ai/internal/config" + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/schema" +) + +func TestAgenticToolArgumentsSanitizerPreservesHistory(t *testing.T) { + for _, args := range []string{`{"command":"unterminated`, `[]`, `null`, ``, `{"command":`} { + t.Run(args, func(t *testing.T) { + good := schema.NewContentBlock(&schema.FunctionToolCall{CallID: "good", Name: "exec", Arguments: `{"command":"pwd"}`}) + bad := schema.NewContentBlock(&schema.FunctionToolCall{CallID: "bad", Name: "exec", Arguments: args}) + text := schema.NewContentBlock(&schema.AssistantGenText{Text: "keep findings"}) + user := schema.UserAgenticMessage("latest instruction") + result := schema.UserAgenticMessage("execution failed") + msg := &schema.AgenticMessage{Role: schema.AgenticRoleTypeAssistant, ContentBlocks: []*schema.ContentBlock{nil, text, good, bad}} + state := &adk.TypedChatModelAgentState[*schema.AgenticMessage]{Messages: []*schema.AgenticMessage{nil, msg, result, user}} + mw := newAgenticToolCallArgumentsSanitizerMiddleware(nil, "test") + _, out, err := mw.BeforeModelRewriteState(context.Background(), state, nil) + if err != nil { + t.Fatal(err) + } + if out == state || out.Messages[1] == msg || out.Messages[1].ContentBlocks[3] == bad { + t.Fatal("repair must copy history") + } + repaired := out.Messages[1].ContentBlocks[3].FunctionToolCall + if repaired.Arguments != `{}` || repaired.Name != "exec" || repaired.CallID != "bad" { + t.Fatalf("repair changed call identity: %+v", repaired) + } + if bad.FunctionToolCall.Arguments != args { + t.Fatal("mutated original arguments") + } + if out.Messages[1].ContentBlocks[1] != text || out.Messages[1].ContentBlocks[2] != good || out.Messages[2] != result || out.Messages[3] != user { + t.Fatal("changed unrelated history") + } + _, unchanged, err := mw.BeforeModelRewriteState(context.Background(), out, nil) + if err != nil || unchanged != out { + t.Fatal("valid history should remain unchanged") + } + }) + } +} + +// A real SDK round trip reproduces a gateway that parses tool argument strings +// on subsequent turns: valid outer JSON alone does not prevent this HTTP 400. +func TestAgenticToolArgumentsSanitizerPreventsMultiTurnHTTP400(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var body struct { + Messages []struct { + Role string `json:"role"` + ToolCalls []struct { + ID string `json:"id"` + Function struct { + Arguments string `json:"arguments"` + } `json:"function"` + } `json:"tool_calls"` + } `json:"messages"` + } + if err := json.NewDecoder(r.Body).Decode(&body); err != nil { + t.Error(err) + w.WriteHeader(400) + return + } + calls := 0 + for _, msg := range body.Messages { + for _, call := range msg.ToolCalls { + calls++ + if !validToolArgumentsJSONObject(call.Function.Arguments) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(400) + fmt.Fprint(w, `{"error":{"message":"Unterminated string starting at: line 1 column 13 (char 12)","type":"invalid_request_error"}}`) + return + } + if call.ID != "broken" || call.Function.Arguments != `{}` { + t.Errorf("unexpected history call: %+v", call) + } + } + } + if calls != 1 { + t.Errorf("history calls=%d, want 1", calls) + } + w.Header().Set("Content-Type", "text/event-stream") + fmt.Fprint(w, "data: {\"id\":\"test\",\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\",\"content\":\"continued\"},\"finish_reason\":null}]}\n\ndata: {\"id\":\"test\",\"choices\":[{\"index\":0,\"delta\":{},\"finish_reason\":\"stop\"}]}\n\ndata: [DONE]\n\n") + })) + defer server.Close() + ctx := context.Background() + base, err := newEinoAgenticChatModelFactory(server.Client(), nil, nil)(ctx, config.OpenAIConfig{Provider: "openai", APIKey: "test-key", BaseURL: server.URL, Model: "space-bunny-free"}, einoModelModeNormal) + if err != nil { + t.Fatal(err) + } + broken := schema.AssistantMessage("", []schema.ToolCall{{ID: "broken", Type: "function", Function: schema.FunctionCall{Name: "exec", Arguments: `{"command":"unterminated`}}}) + input := []*schema.Message{schema.UserMessage("start"), broken, schema.ToolMessage("Invalid JSON arguments; tool was not executed", "broken"), schema.UserMessage("continue")} + for _, repair := range []bool{false, true} { + t.Run(fmt.Sprint(repair), func(t *testing.T) { + var handlers []adk.TypedChatModelAgentMiddleware[*schema.AgenticMessage] + if repair { + handlers = appendEinoAgenticChatModelTailMiddlewares(nil, einoChatModelTailConfig{phase: "test"}) + } + agent, err := newEinoAgenticChatModelAgentAdapter(ctx, einoAgenticChatModelAgentConfig{Name: "test", Model: base, Handlers: handlers}) + if err != nil { + t.Fatal(err) + } + it := agent.Run(ctx, &adk.AgentInput{Messages: input, EnableStreaming: true}) + var runErr error + var answer strings.Builder + for { + ev, ok := it.Next() + if !ok { + break + } + if ev.Err != nil { + runErr = ev.Err + break + } + if ev.Output != nil && ev.Output.MessageOutput != nil && ev.Output.MessageOutput.MessageStream != nil { + stream := ev.Output.MessageOutput.MessageStream + for { + msg, err := stream.Recv() + if err != nil { + if !errors.Is(err, io.EOF) { + runErr = err + } + break + } + if msg != nil { + answer.WriteString(msg.Content) + } + } + stream.Close() + } + } + if repair && runErr != nil { + t.Fatalf("repaired history rejected: %v", runErr) + } + if repair && !strings.Contains(answer.String(), "continued") { + t.Fatalf("agent did not continue: %q", answer.String()) + } + if !repair && runErr == nil { + t.Fatal("expected HTTP 400 without sanitizer") + } + if broken.ToolCalls[0].Function.Arguments != `{"command":"unterminated` { + t.Fatal("original history mutated") + } + }) + } +} + +// Summarization makes its own model call and can rewrite retained history, so +// guard both its input and the final model input rather than just one boundary. +type agenticSummaryArgumentsProbe struct { + *adk.TypedBaseChatModelAgentMiddleware[*schema.AgenticMessage] + checked bool +} + +func (p *agenticSummaryArgumentsProbe) BeforeModelRewriteState(ctx context.Context, state *adk.TypedChatModelAgentState[*schema.AgenticMessage], _ *adk.TypedModelContext[*schema.AgenticMessage]) (context.Context, *adk.TypedChatModelAgentState[*schema.AgenticMessage], error) { + for _, msg := range state.Messages { + for _, call := range agenticFunctionToolCalls(msg) { + if !validToolArgumentsJSONObject(call.Arguments) { + return ctx, state, fmt.Errorf("invalid arguments reached summarization") + } + p.checked = true + } + } + out := *state + out.Messages = append([]*schema.AgenticMessage(nil), state.Messages...) + for i, msg := range out.Messages { + if len(agenticFunctionToolCalls(msg)) == 0 { + continue + } + calls := []*schema.FunctionToolCall{{CallID: "broken", Name: "exec", Arguments: `{"command":"truncated`}} + out.Messages[i] = cloneAgenticMessageWithCalls(msg, calls) + } + return ctx, &out, nil +} + +func TestAgenticToolArgumentsSanitizerAroundSummarization(t *testing.T) { + probe := &agenticSummaryArgumentsProbe{TypedBaseChatModelAgentMiddleware: &adk.TypedBaseChatModelAgentMiddleware[*schema.AgenticMessage]{}} + history := EinoMessagesToAgentic([]*schema.Message{ + schema.AssistantMessage("", []schema.ToolCall{{ID: "broken", Type: "function", Function: schema.FunctionCall{Name: "exec", Arguments: `{"command":"truncated`}}}), + schema.ToolMessage("invalid JSON; not executed", "broken"), + schema.UserMessage("continue"), + }) + state := &adk.TypedChatModelAgentState[*schema.AgenticMessage]{Messages: history} + handlers := appendEinoAgenticChatModelTailMiddlewares(nil, einoChatModelTailConfig{agenticSummarization: probe}) + ctx := context.Background() + for _, handler := range handlers { + var err error + ctx, state, err = handler.BeforeModelRewriteState(ctx, state, nil) + if err != nil { + t.Fatal(err) + } + } + if !probe.checked { + t.Fatal("summary probe did not inspect retained call") + } + calls := agenticFunctionToolCalls(state.Messages[0]) + if len(calls) != 1 || calls[0].Arguments != `{}` { + t.Fatalf("final input contains malformed call: %+v", calls) + } + if agenticFunctionToolCalls(history[0])[0].Arguments != `{"command":"truncated` { + t.Fatal("mutated original history") + } +} diff --git a/internal/multiagent/eino_agentic_chat_model_tail_middleware.go b/internal/multiagent/eino_agentic_chat_model_tail_middleware.go index 738a9b27..7ee7f3eb 100644 --- a/internal/multiagent/eino_agentic_chat_model_tail_middleware.go +++ b/internal/multiagent/eino_agentic_chat_model_tail_middleware.go @@ -10,19 +10,20 @@ import ( ) // appendEinoAgenticChatModelTailMiddlewares appends protocol-neutral handlers for -// TypedChatModelAgent[*schema.AgenticMessage]. Classic ReAct history repair -// handlers stay on the schema.Message path because AgenticMessage has native -// content blocks for function calls/results. +// TypedChatModelAgent[*schema.AgenticMessage]. History repair operates +// directly on native function-call/result blocks without converting protocols. func appendEinoAgenticChatModelTailMiddlewares( handlers []adk.TypedChatModelAgentMiddleware[*schema.AgenticMessage], cfg einoChatModelTailConfig, ) []adk.TypedChatModelAgentMiddleware[*schema.AgenticMessage] { handlers = append(handlers, newAgenticSystemMessageNormalizerMiddleware(cfg.logger, cfg.phase)) handlers = append(handlers, newAgenticContinuationUserDedupMiddleware(cfg.logger, cfg.phase)) + handlers = append(handlers, newAgenticToolCallArgumentsSanitizerMiddleware(cfg.logger, cfg.phase+"_pre_summarization")) if cfg.agenticSummarization != nil { handlers = append(handlers, newAgenticToolPairReconcilerMiddleware(cfg.logger, cfg.phase+"_pre_summarization")) handlers = append(handlers, cfg.agenticSummarization) } + handlers = append(handlers, newAgenticToolCallArgumentsSanitizerMiddleware(cfg.logger, cfg.phase)) handlers = append(handlers, newAgenticToolPairReconcilerMiddleware(cfg.logger, cfg.phase)) if !cfg.skipOrphanPruner { handlers = append(handlers, newAgenticOrphanToolPrunerMiddleware(cfg.logger, cfg.phase)) diff --git a/internal/multiagent/eino_agentic_chat_model_tail_middleware_test.go b/internal/multiagent/eino_agentic_chat_model_tail_middleware_test.go index 9b6d9d5c..22e135aa 100644 --- a/internal/multiagent/eino_agentic_chat_model_tail_middleware_test.go +++ b/internal/multiagent/eino_agentic_chat_model_tail_middleware_test.go @@ -106,8 +106,8 @@ func TestAppendEinoAgenticChatModelTailMiddlewares(t *testing.T) { phase: "agentic", trace: holder, }) - // system + continuation + reconciler + orphan_pruner + trace - if len(handlers) != 5 { - t.Fatalf("handlers = %d, want system + continuation + reconciler + orphan_pruner + trace", len(handlers)) + // system + continuation + two argument sanitizers + reconciler + orphan_pruner + trace + if len(handlers) != 7 { + t.Fatalf("handlers = %d, want system + continuation + two argument sanitizers + reconciler + orphan_pruner + trace", len(handlers)) } } diff --git a/internal/multiagent/eino_summarize_model_guard.go b/internal/multiagent/eino_summarize_model_guard.go index c0b9757d..d735acf3 100644 --- a/internal/multiagent/eino_summarize_model_guard.go +++ b/internal/multiagent/eino_summarize_model_guard.go @@ -50,6 +50,20 @@ func einoSummarizationModelRawErrorText(err error) (string, bool) { return strings.TrimSpace(summaryErr.err.Error()), true } +// Output-limit failures need a different prompt, rather than repeating the same +// overlong summary or accepting truncated text as replacement conversation history. +type einoSummaryOutputLimitError struct{ err error } + +func (e *einoSummaryOutputLimitError) Error() string { return e.err.Error() } +func (e *einoSummaryOutputLimitError) Unwrap() error { return e.err } + +const conciseSummaryRetryInstruction = `The previous summarization request exhausted its output token budget. Generate a new, substantially shorter, complete summary from the original conversation above. Omit the block and output only .... Preserve authorization, user constraints, confirmed findings, essential technical details, current progress and next steps. Merge repetitions and summarize long logs and asset lists. Do not repeat the full transcript or call tools. Finish all sections and the closing tag within the available output budget.` + +func shouldRetryConciseSummary(ctx context.Context, err error) bool { + var limitErr *einoSummaryOutputLimitError + return ctx.Err() == nil && errors.As(err, &limitErr) +} + type nonEmptySummaryChatModel struct { base model.BaseChatModel } @@ -59,6 +73,16 @@ 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.generateOnce(ctx, input, opts...) + if !shouldRetryConciseSummary(ctx, err) { + return out, err + } + // Copy the slice: retry instructions must not mutate caller-owned history. + retryInput := append(append([]*schema.Message(nil), input...), schema.UserMessage(conciseSummaryRetryInstruction)) + return m.generateOnce(ctx, retryInput, opts...) +} + +func (m *nonEmptySummaryChatModel) generateOnce(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) { if err := ctx.Err(); err != nil { return nil, newEinoSummarizationModelError(err) } @@ -70,11 +94,15 @@ func (m *nonEmptySummaryChatModel) Generate(ctx context.Context, input []*schema if err != nil { return nil, newEinoSummarizationModelError(err) } + completionErr := validateClassicSummaryCompletion(out) + if shouldRetryConciseSummary(ctx, completionErr) { + return nil, newEinoSummarizationModelError(completionErr) + } if strings.TrimSpace(classicAssistantTextContent(out)) == "" { return nil, newEinoSummarizationModelError(fmt.Errorf("summary content is empty: %s", classicSummaryEmptyDiagnostics(out))) } - if err := validateClassicSummaryCompletion(out); err != nil { - return nil, newEinoSummarizationModelError(err) + if completionErr != nil { + return nil, newEinoSummarizationModelError(completionErr) } return out, nil } @@ -92,6 +120,16 @@ 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.generateOnce(ctx, input, opts...) + if !shouldRetryConciseSummary(ctx, err) { + return out, err + } + // Copy the slice: retry instructions must not mutate caller-owned history. + retryInput := append(append([]*schema.AgenticMessage(nil), input...), schema.UserAgenticMessage(conciseSummaryRetryInstruction)) + return m.generateOnce(ctx, retryInput, opts...) +} + +func (m *nonEmptyAgenticSummaryModel) generateOnce(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.AgenticMessage, error) { if err := ctx.Err(); err != nil { return nil, newEinoSummarizationModelError(err) } @@ -103,11 +141,15 @@ func (m *nonEmptyAgenticSummaryModel) Generate(ctx context.Context, input []*sch if err != nil { return nil, newEinoSummarizationModelError(err) } + completionErr := validateAgenticSummaryCompletion(out) + if shouldRetryConciseSummary(ctx, completionErr) { + return nil, newEinoSummarizationModelError(completionErr) + } if strings.TrimSpace(agenticAssistantTextContent(out)) == "" { return nil, newEinoSummarizationModelError(fmt.Errorf("summary content is empty: %s", agenticSummaryEmptyDiagnostics(out))) } - if err := validateAgenticSummaryCompletion(out); err != nil { - return nil, newEinoSummarizationModelError(err) + if completionErr != nil { + return nil, newEinoSummarizationModelError(completionErr) } return out, nil } @@ -278,6 +320,9 @@ func validateClassicSummaryCompletion(msg *schema.Message) error { if msg.ResponseMeta != nil { reason = msg.ResponseMeta.FinishReason } + if reason == "length" && len(msg.ToolCalls) == 0 { + return &einoSummaryOutputLimitError{err: fmt.Errorf("summary did not complete: finish_reason=%q tool_calls=%d", reason, len(msg.ToolCalls))} + } if reason != "stop" || len(msg.ToolCalls) != 0 { return fmt.Errorf("summary did not complete: finish_reason=%q tool_calls=%d", reason, len(msg.ToolCalls)) } @@ -290,13 +335,21 @@ func validateAgenticSummaryCompletion(msg *schema.AgenticMessage) error { if ext.FinishReason == "stop" { return nil } - return fmt.Errorf("summary did not complete: openai chat finish_reason=%q", ext.FinishReason) + err := fmt.Errorf("summary did not complete: openai chat finish_reason=%q", ext.FinishReason) + if ext.FinishReason == "length" { + return &einoSummaryOutputLimitError{err: err} + } + return err } 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) + err := fmt.Errorf("summary did not complete: claude stop_reason=%q", ext.StopReason) + if ext.StopReason == "max_tokens" { + return &einoSummaryOutputLimitError{err: err} + } + return err } if ext := meta.OpenAIExtension; ext != nil { if ext.Status == schemaopenai.ResponseStatusCompleted && ext.Error == nil && ext.IncompleteDetails == nil { diff --git a/internal/multiagent/eino_summarize_model_guard_test.go b/internal/multiagent/eino_summarize_model_guard_test.go index 8fe2da2f..7cb562b3 100644 --- a/internal/multiagent/eino_summarize_model_guard_test.go +++ b/internal/multiagent/eino_summarize_model_guard_test.go @@ -8,6 +8,7 @@ import ( "net/http" "net/http/httptest" "strings" + "sync/atomic" "testing" "time" @@ -208,13 +209,15 @@ func TestSummaryStreamFailuresDoNotReturnPartialSummary(t *testing.T) { } func TestClaudeSummaryLargeBudgetStreamsThroughNativeSDK(t *testing.T) { - for _, mode := range []string{"complete", "truncated", "early_eof", "cancel_read"} { + for _, mode := range []string{"complete", "recovers", "truncated", "early_eof", "cancel_read"} { t.Run(mode, func(t *testing.T) { ctx, cancel := context.WithCancel(context.Background()) defer cancel() started := make(chan struct{}) + var requests atomic.Int32 server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + request := requests.Add(1) var body struct { Stream bool `json:"stream"` MaxTokens int `json:"max_tokens"` @@ -240,6 +243,10 @@ func TestClaudeSummaryLargeBudgetStreamsThroughNativeSDK(t *testing.T) { `{"type":"message_stop"}`, } switch mode { + case "recovers": + if request == 1 { + events[5] = strings.ReplaceAll(events[5], "end_turn", "max_tokens") + } case "truncated": events[5] = strings.ReplaceAll(events[5], "end_turn", "max_tokens") case "early_eof", "cancel_read": @@ -283,7 +290,14 @@ func TestClaudeSummaryLargeBudgetStreamsThroughNativeSDK(t *testing.T) { defer timer.Stop() } out, err := newNonEmptyAgenticSummaryModel(native).Generate(ctx, input, opts...) - if mode != "complete" { + wantRequests := int32(1) + if mode == "recovers" || mode == "truncated" { + wantRequests = 2 + } + if requests.Load() != wantRequests { + t.Fatalf("requests=%d, want %d", requests.Load(), wantRequests) + } + if mode != "complete" && mode != "recovers" { if out != nil || err == nil { t.Fatalf("accepted partial summary: out=%+v err=%v", out, err) } diff --git a/internal/multiagent/eino_summarize_payload_test.go b/internal/multiagent/eino_summarize_payload_test.go index 186b7793..21e7e17c 100644 --- a/internal/multiagent/eino_summarize_payload_test.go +++ b/internal/multiagent/eino_summarize_payload_test.go @@ -3,9 +3,11 @@ package multiagent import ( "context" "encoding/json" + "fmt" "net/http" "net/http/httptest" "strings" + "sync/atomic" "testing" "cyberstrike-ai/internal/config" @@ -148,3 +150,118 @@ func TestEinoSummarizationModelOptionsSetOnlyMaxCompletionTokens(t *testing.T) { t.Fatalf("max_completion_tokens = %#v, want %d", body["max_completion_tokens"], outputReserve) } } + +// Exercise both production factories and their real SDK options, not just a +// fake model: a truncated summary must be replaced only by a complete retry. +func TestSummaryOutputLimitConciseRetry(t *testing.T) { + for _, kind := range []string{"classic", "agentic"} { + for _, tc := range []struct { + name string + firstReason string + retryReason string + firstText string + wantCalls int + wantErr bool + }{ + {"complete", "stop", "", "complete summary", 1, false}, + {"recovers", "length", "stop", "truncated summary", 2, false}, + {"reasoning_only_recovers", "length", "stop", "", 2, false}, + {"still_truncated", "length", "length", "truncated summary", 2, true}, + {"retry_filtered", "length", "content_filter", "truncated summary", 2, true}, + {"filtered", "content_filter", "", "filtered summary", 1, true}, + {"missing_completion", "", "", "partial summary", 1, true}, + } { + t.Run(kind+"/"+tc.name, func(t *testing.T) { + var calls atomic.Int32 + bodies := make(chan map[string]any, 3) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + n := calls.Add(1) + var body map[string]any + if err := json.NewDecoder(r.Body).Decode(&body); err != nil { + t.Errorf("decode request: %v", err) + w.WriteHeader(400) + return + } + bodies <- body + if r.Header.Get("X-CyberStrike-Summarization") != "1" { + t.Error("missing summary request header") + } + text, reason := tc.firstText, tc.firstReason + if n > 1 { + text, reason = "complete retry summary", tc.retryReason + } + w.Header().Set("Content-Type", "text/event-stream") + fmt.Fprintf(w, "data: {\"id\":\"test\",\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\",\"content\":%q},\"finish_reason\":null}]}\n\n", text) + 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() + oa := config.OpenAIConfig{Provider: "openai", APIKey: "test-key", BaseURL: server.URL, Model: "deepseek-chat", MaxCompletionTokens: 128} + const reserve = 512 + opts := newEinoSummarizationModelOptions(reserve, oa.Model, kind, &oa, nil) + ctx := context.Background() + var err error + var text string + if kind == "classic" { + base, modelErr := newEinoToolCallingChatModelFactory(server.Client(), nil, nil)(ctx, oa, einoModelModeNormal) + if modelErr != nil { + t.Fatal(modelErr) + } + input := []*schema.Message{schema.UserMessage("original conversation")} + out, genErr := newNonEmptySummaryChatModel(base).Generate(ctx, input, opts...) + err = genErr + text = classicAssistantTextContent(out) + if len(input) != 1 || input[0].Content != "original conversation" { + t.Fatal("mutated original history") + } + } else { + base, modelErr := newEinoAgenticChatModelFactory(server.Client(), nil, nil)(ctx, oa, einoModelModeNormal) + if modelErr != nil { + t.Fatal(modelErr) + } + input := []*schema.AgenticMessage{schema.UserAgenticMessage("original conversation")} + out, genErr := newNonEmptyAgenticSummaryModel(base).Generate(ctx, input, opts...) + err = genErr + text = agenticAssistantTextContent(out) + if len(input) != 1 { + t.Fatal("mutated original history") + } + } + if (err != nil) != tc.wantErr { + t.Fatalf("error = %v, wantErr = %v", err, tc.wantErr) + } + if tc.wantErr && text != "" { + t.Fatalf("accepted partial summary %q", text) + } + if !tc.wantErr && tc.wantCalls == 2 && text != "complete retry summary" { + t.Fatalf("summary = %q", text) + } + if int(calls.Load()) != tc.wantCalls { + t.Fatalf("calls = %d, want %d", calls.Load(), tc.wantCalls) + } + for i := 0; i < tc.wantCalls; i++ { + body := <-bodies + if body["max_completion_tokens"] != float64(reserve) { + t.Fatalf("output budget changed: %#v", body) + } + if _, ok := body["max_tokens"]; ok { + t.Fatal("duplicate output limits") + } + if _, ok := body["reasoning_effort"]; ok { + t.Fatal("reasoning enabled for summary") + } + msgs := body["messages"].([]any) + if i == 1 { + last := msgs[len(msgs)-1].(map[string]any) + if last["role"] != "user" || !strings.Contains(fmt.Sprint(last["content"]), conciseSummaryRetryInstruction) { + t.Fatalf("missing concise retry instruction: %#v", last) + } + if strings.Contains(fmt.Sprint(msgs), "truncated summary") { + t.Fatal("partial response leaked into retry history") + } + } + } + }) + } + } +}