CyberStrikeAI/internal/multiagent/eino_turn_loop_runtime.go

package multiagent

import (
	"context"
	"strings"
	"time"

	"github.com/cloudwego/eino/adk"
	"github.com/cloudwego/eino/schema"
)

const (
	einoTurnLoopInterruptPreemptTimeout = 3 * time.Second
	einoTurnLoopIdleStop                = 250 * time.Millisecond
)

// EinoTurnLoopItem is the conversation-level input unit consumed by an Eino
// TurnLoop. The item is gob-friendly so it can be checkpointed by TurnLoop when
// a CheckPointStore is configured.
type EinoTurnLoopItem struct {
	Messages []*schema.Message
	Kind     string
	Note     string
}

// EinoTurnLoopRuntime wraps Eino's native TurnLoop with the semantics this
// project needs: persistent per-conversation runtime, user-supplement preempt,
// and graceful idle shutdown.
type EinoTurnLoopRuntime struct {
	loop             *adk.TurnLoop[EinoTurnLoopItem, *schema.Message]
	interruptTimeout time.Duration
}

type EinoTurnLoopRuntimeConfig struct {
	Agent            adk.Agent
	InitialMessages  []*schema.Message
	Store            adk.CheckPointStore
	CheckpointID     string
	EnableStreaming  bool
	PrepareAgent     func(context.Context, *adk.TurnLoop[EinoTurnLoopItem, *schema.Message], []EinoTurnLoopItem) (adk.Agent, error)
	OnAgentEvents    func(context.Context, *adk.TurnContext[EinoTurnLoopItem, *schema.Message], *adk.AsyncIterator[*adk.AgentEvent]) error
	InterruptTimeout time.Duration
}

func NewEinoTurnLoopRuntime(cfg EinoTurnLoopRuntimeConfig) *EinoTurnLoopRuntime {
	timeout := cfg.InterruptTimeout
	if timeout <= 0 {
		timeout = einoTurnLoopInterruptPreemptTimeout
	}
	enableStreaming := cfg.EnableStreaming
	prepareAgent := cfg.PrepareAgent
	if prepareAgent == nil {
		prepareAgent = func(context.Context, *adk.TurnLoop[EinoTurnLoopItem, *schema.Message], []EinoTurnLoopItem) (adk.Agent, error) {
			return cfg.Agent, nil
		}
	}
	loop := adk.NewTurnLoop[EinoTurnLoopItem, *schema.Message](adk.TurnLoopConfig[EinoTurnLoopItem, *schema.Message]{
		Store:        cfg.Store,
		CheckpointID: cfg.CheckpointID,
		GenInput: func(ctx context.Context, _ *adk.TurnLoop[EinoTurnLoopItem, *schema.Message], items []EinoTurnLoopItem) (*adk.GenInputResult[EinoTurnLoopItem, *schema.Message], error) {
			msgs := mergeEinoTurnLoopMessages(items)
			return &adk.GenInputResult[EinoTurnLoopItem, *schema.Message]{
				RunCtx: ctx,
				Input: &adk.AgentInput{
					Messages:        msgs,
					EnableStreaming: enableStreaming,
				},
				Consumed: items,
			}, nil
		},
		GenResume: func(ctx context.Context, _ *adk.TurnLoop[EinoTurnLoopItem, *schema.Message], interruptedItems, unhandledItems, newItems []EinoTurnLoopItem) (*adk.GenResumeResult[EinoTurnLoopItem, *schema.Message], error) {
			consumed := make([]EinoTurnLoopItem, 0, len(interruptedItems)+len(newItems))
			consumed = append(consumed, interruptedItems...)
			consumed = append(consumed, newItems...)
			remaining := append([]EinoTurnLoopItem(nil), unhandledItems...)
			return &adk.GenResumeResult[EinoTurnLoopItem, *schema.Message]{
				RunCtx:    ctx,
				Consumed:  consumed,
				Remaining: remaining,
			}, nil
		},
		PrepareAgent:  prepareAgent,
		OnAgentEvents: cfg.OnAgentEvents,
	})
	if len(cfg.InitialMessages) > 0 {
		loop.Push(EinoTurnLoopItem{Kind: "initial", Messages: cloneSchemaMessages(cfg.InitialMessages)})
	}
	return &EinoTurnLoopRuntime{loop: loop, interruptTimeout: timeout}
}

func (r *EinoTurnLoopRuntime) Run(ctx context.Context) {
	if r == nil || r.loop == nil {
		return
	}
	r.loop.Run(ctx)
}

func (r *EinoTurnLoopRuntime) PushInterruptContinue(note string) bool {
	if r == nil || r.loop == nil {
		return false
	}
	item := EinoTurnLoopItem{
		Kind:     "interrupt_continue",
		Note:     strings.TrimSpace(note),
		Messages: []*schema.Message{schema.UserMessage(formatInterruptContinuePrompt(note))},
	}
	ok, ack := r.loop.Push(item, adk.WithPreemptTimeout[EinoTurnLoopItem, *schema.Message](adk.AnySafePoint, r.interruptTimeout))
	if ack != nil {
		go func() { <-ack }()
	}
	return ok
}

func (r *EinoTurnLoopRuntime) StopImmediate(cause string) {
	if r == nil || r.loop == nil {
		return
	}
	r.loop.Stop(adk.WithImmediate(), adk.WithStopCause(cause))
}

func (r *EinoTurnLoopRuntime) StopWhenIdle() {
	if r == nil || r.loop == nil {
		return
	}
	r.loop.Stop(adk.UntilIdleFor(einoTurnLoopIdleStop))
}

func (r *EinoTurnLoopRuntime) Wait() *adk.TurnLoopExitState[EinoTurnLoopItem, *schema.Message] {
	if r == nil || r.loop == nil {
		return nil
	}
	return r.loop.Wait()
}

func mergeEinoTurnLoopMessages(items []EinoTurnLoopItem) []*schema.Message {
	var msgs []*schema.Message
	for _, item := range items {
		msgs = append(msgs, cloneSchemaMessages(item.Messages)...)
	}
	return msgs
}

func formatInterruptContinuePrompt(note string) string {
	note = strings.TrimSpace(note)
	if note == "" {
		return "用户请求中断当前推理并继续。请基于已经完成的步骤继续，不要重复已完成工具调用。"
	}
	return "用户请求中断当前推理并补充上下文后继续：\n" + note +
		"\n\n请基于已经完成的步骤继续，不要重复已完成工具调用。"
}

func cloneSchemaMessages(in []*schema.Message) []*schema.Message {
	if len(in) == 0 {
		return nil
	}
	out := make([]*schema.Message, 0, len(in))
	for _, msg := range in {
		if msg == nil {
			continue
		}
		cp := *msg
		if len(msg.ToolCalls) > 0 {
			cp.ToolCalls = append([]schema.ToolCall(nil), msg.ToolCalls...)
		}
		if len(msg.MultiContent) > 0 {
			cp.MultiContent = append([]schema.ChatMessagePart(nil), msg.MultiContent...)
		}
		if len(msg.UserInputMultiContent) > 0 {
			cp.UserInputMultiContent = append([]schema.MessageInputPart(nil), msg.UserInputMultiContent...)
		}
		if len(msg.AssistantGenMultiContent) > 0 {
			cp.AssistantGenMultiContent = append([]schema.MessageOutputPart(nil), msg.AssistantGenMultiContent...)
		}
		cp.Extra = cloneAnyMap(msg.Extra)
		out = append(out, &cp)
	}
	return out
}