fix: recover summary limits and sanitize agentic tool history

This commit is contained in:
Sycun committed 2026-10-12 16:49:22 +08:00
1 parent deda642bd6
commit 4f0a3049cb
7 files changed
+488 -14

No files matched your search

@@ -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
}
@@ -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")
}
}
@@ -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))
@@ -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))
}
}
@@ -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 <analysis> block and output only <summary>...</summary>. 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 {
@@ -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)
}
@@ -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")
}
}
}
})
}
}
}