CyberStrikeAI/internal/multiagent/eino_summarize_model_guard.go

package multiagent

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 {
	err error
}

func newEinoSummarizationModelError(err error) error {
	if err == nil {
		return nil
	}
	var existing *einoSummarizationModelError
	if errors.As(err, &existing) {
		return err
	}
	return &einoSummarizationModelError{err: err}
}

func (e *einoSummarizationModelError) Error() string {
	if e == nil || e.err == nil {
		return "summarization model error"
	}
	return "summarization model error: " + e.err.Error()
}

func (e *einoSummarizationModelError) Unwrap() error {
	if e == nil {
		return nil
	}
	return e.err
}

func einoSummarizationModelRawErrorText(err error) (string, bool) {
	var summaryErr *einoSummarizationModelError
	if !errors.As(err, &summaryErr) || summaryErr == nil || summaryErr.err == nil {
		return "", false
	}
	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
}

func newNonEmptySummaryChatModel(base model.BaseChatModel) model.BaseChatModel {
	return &nonEmptySummaryChatModel{base: base}
}

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)
	}
	stream, err := m.base.Stream(ctx, input, opts...)
	if err != nil {
		return nil, newEinoSummarizationModelError(err)
	}
	out, err := collectSummaryStream(ctx, stream, schema.ConcatMessages)
	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 completionErr != nil {
		return nil, newEinoSummarizationModelError(completionErr)
	}
	return out, nil
}

func (m *nonEmptySummaryChatModel) Stream(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) {
	return m.base.Stream(ctx, input, opts...)
}

type nonEmptyAgenticSummaryModel struct {
	base model.BaseModel[*schema.AgenticMessage]
}

func newNonEmptyAgenticSummaryModel(base model.BaseModel[*schema.AgenticMessage]) model.BaseModel[*schema.AgenticMessage] {
	return &nonEmptyAgenticSummaryModel{base: base}
}

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)
	}
	stream, err := m.base.Stream(ctx, input, opts...)
	if err != nil {
		return nil, newEinoSummarizationModelError(err)
	}
	out, err := collectSummaryStream(ctx, stream, schema.ConcatAgenticMessages)
	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 completionErr != nil {
		return nil, newEinoSummarizationModelError(completionErr)
	}
	return out, nil
}

func (m *nonEmptyAgenticSummaryModel) Stream(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.StreamReader[*schema.AgenticMessage], error) {
	return m.base.Stream(ctx, input, opts...)
}

func classicAssistantTextContent(msg *schema.Message) string {
	if msg == nil || msg.Role != schema.Assistant {
		return ""
	}
	parts := make([]string, 0, len(msg.AssistantGenMultiContent))
	for _, part := range msg.AssistantGenMultiContent {
		if part.Type == schema.ChatMessagePartTypeText && part.Text != "" {
			parts = append(parts, part.Text)
		}
	}
	if len(parts) > 0 {
		return strings.Join(parts, "\n")
	}
	return msg.Content
}

func agenticAssistantTextContent(msg *schema.AgenticMessage) string {
	if msg == nil || msg.Role != schema.AgenticRoleTypeAssistant {
		return ""
	}
	parts := make([]string, 0, len(msg.ContentBlocks))
	for _, block := range msg.ContentBlocks {
		if block != nil && block.AssistantGenText != nil {
			parts = append(parts, block.AssistantGenText.Text)
		}
	}
	return strings.Join(parts, "\n")
}

func classicSummaryEmptyDiagnostics(msg *schema.Message) string {
	if msg == nil {
		return "model returned nil message"
	}
	var textParts, reasoningParts, otherParts int
	var multiTextRunes, multiReasoningRunes int
	for _, part := range msg.AssistantGenMultiContent {
		switch part.Type {
		case schema.ChatMessagePartTypeText:
			textParts++
			multiTextRunes += len([]rune(strings.TrimSpace(part.Text)))
		case schema.ChatMessagePartTypeReasoning:
			reasoningParts++
			if part.Reasoning != nil {
				multiReasoningRunes += len([]rune(strings.TrimSpace(part.Reasoning.Text)))
			}
		default:
			otherParts++
		}
	}
	reasoningRunes := len([]rune(strings.TrimSpace(msg.ReasoningContent))) + multiReasoningRunes
	fields := []string{
		fmt.Sprintf("role=%s", msg.Role),
		fmt.Sprintf("content_runes=%d", len([]rune(strings.TrimSpace(msg.Content)))+multiTextRunes),
		fmt.Sprintf("reasoning_runes=%d", reasoningRunes),
		fmt.Sprintf("text_parts=%d", textParts),
		fmt.Sprintf("reasoning_parts=%d", reasoningParts),
		fmt.Sprintf("other_parts=%d", otherParts),
		fmt.Sprintf("tool_calls=%d", len(msg.ToolCalls)),
	}
	if msg.ResponseMeta != nil {
		fields = append(fields, fmt.Sprintf("finish_reason=%q", msg.ResponseMeta.FinishReason))
		if usage := msg.ResponseMeta.Usage; usage != nil {
			fields = append(fields,
				fmt.Sprintf("prompt_tokens=%d", usage.PromptTokens),
				fmt.Sprintf("completion_tokens=%d", usage.CompletionTokens),
				fmt.Sprintf("total_tokens=%d", usage.TotalTokens),
				fmt.Sprintf("reasoning_tokens=%d", usage.CompletionTokensDetails.ReasoningTokens),
			)
		}
	}
	if reasoningRunes > 0 {
		fields = append(fields, "hint=模型返回了 reasoning_content 但没有返回可作为摘要正文的 content；请检查 DeepSeek thinking 是否已在摘要请求中关闭")
	}
	return strings.Join(fields, " ")
}

func agenticSummaryEmptyDiagnostics(msg *schema.AgenticMessage) string {
	if msg == nil {
		return "model returned nil agentic message"
	}
	var textBlocks, reasoningBlocks, otherBlocks int
	var textRunes, reasoningRunes int
	for _, block := range msg.ContentBlocks {
		if block == nil {
			continue
		}
		switch {
		case block.AssistantGenText != nil:
			textBlocks++
			textRunes += len([]rune(strings.TrimSpace(block.AssistantGenText.Text)))
		case block.Reasoning != nil:
			reasoningBlocks++
			reasoningRunes += len([]rune(strings.TrimSpace(block.Reasoning.Text)))
		default:
			otherBlocks++
		}
	}
	fields := []string{
		fmt.Sprintf("role=%s", msg.Role),
		fmt.Sprintf("content_runes=%d", textRunes),
		fmt.Sprintf("reasoning_runes=%d", reasoningRunes),
		fmt.Sprintf("text_blocks=%d", textBlocks),
		fmt.Sprintf("reasoning_blocks=%d", reasoningBlocks),
		fmt.Sprintf("other_blocks=%d", otherBlocks),
	}
	if msg.ResponseMeta != nil && msg.ResponseMeta.TokenUsage != nil {
		usage := msg.ResponseMeta.TokenUsage
		fields = append(fields,
			fmt.Sprintf("prompt_tokens=%d", usage.PromptTokens),
			fmt.Sprintf("completion_tokens=%d", usage.CompletionTokens),
			fmt.Sprintf("total_tokens=%d", usage.TotalTokens),
			fmt.Sprintf("reasoning_tokens=%d", usage.CompletionTokensDetails.ReasoningTokens),
		)
	}
	if reasoningRunes > 0 {
		fields = append(fields, "hint=模型返回了 reasoning block 但没有返回可作为摘要正文的 text block；请检查 DeepSeek thinking 是否已在摘要请求中关闭")
	}
	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 == "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))
	}
	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
			}
			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
			}
			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 {
				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")
}