Agent-Native/runtime/loop.go

package runtime

import (
	"context"
	"encoding/json"
	"fmt"
	"strings"
	"sync"
	"time"

	"agent-native/runtime/llm"
	"agent-native/runtime/session"
	"agent-native/runtime/tools"

	"logging-system/logsys"
)

const (
	ExecSequential = "sequential"
	ExecParallel   = "parallel"

	maxStopContinuations   = 3
	maxLengthContinuations = 3
	maxRepeatReminders     = 3
)

const repeatReminderText = "你已连续 3 次以相同参数调用同一工具。如果结果不符合预期，请更换参数或做法，不要再重复这一调用。"
const lengthContinuePrompt = "上一段回复因为达到长度上限被截断。请从断点接着写，不要重复已经写出的内容。"

// Event is a lifecycle notice for logs and UI.
type Event struct {
	Type    string
	Message llm.Message
	Tool    string
	Text    string
	// Detail carries tool arguments/results or reasoning deltas for
	// process-visualization consumers; log writers may ignore it.
	Detail string
	// PromptTokens / CompletionTokens / TotalTokens 在 type=usage 或
	// type=message_end 时有意义，供调用方累计用量。
	PromptTokens     int
	CompletionTokens int
	TotalTokens      int
}

// Config drives one agent session.
type Config struct {
	Model         string
	SystemPrompt  string
	Tools         *tools.Registry
	Stream        llm.StreamFunc
	Session       *session.Store
	MaxTurns      int
	ToolExecution string
	GetSteering   func() []llm.Message
	GetFollowUp   func() []llm.Message
	OnEvent       func(Event)
	// ThinkingLevel is forwarded on each provider request.
	ThinkingLevel string
	// Resolve, when set, overrides Model and ThinkingLevel at the start of each turn.
	Resolve func() (model, thinking string)
	// Log 可选。nil 时丢弃日志，不改变循环结果。
	Log *logsys.Logger
	// ExecutionID 只用于日志字段，不参与控制流。
	ExecutionID string
	// Budget 可选。任一上限命中即停止并返回 budget_exceeded 错误（StopError）；
	// 零值字段不设限。截止时刻与外层 ctx 先到者生效。
	Budget *Budget
	// History 是会话转录为空时的初始转录。以 system 开头时，不再前置
	// SystemPrompt 和工具清单。不回写进 Session。
	History []llm.Message
	// Stop 在模型不再调用工具、且本轮没有 follow-up 时询问是否续跑。
	// nil 表示不询问。内核最多接受 3 次 Continue。
	Stop func(last llm.Message, info StopInfo) StopDecision
	// BeforeTool 返回错误时不执行工具，但仍写一条 IsError 的 tool 结果。
	BeforeTool func(call llm.ToolCall) error
	// UnknownTool 处理注册表里没有的调用。nil 时仍返回 unknown tool。
	UnknownTool func(call llm.ToolCall) tools.Result
	// AfterTool 只观察工具结果，不能改写。
	AfterTool func(call llm.ToolCall, res llm.Message)
	// Reminders 在每次调用模型前提供提醒。正文包在 system-reminder 里，
	// 只进入本次请求，不写进 Session。
	Reminders func(turn int) []string
	// RepeatGuard 为真时，同名且参数按键排序后相同的调用连续 3 次，注入一条提醒。
	// 每个 run 最多 3 条，不中断循环。
	RepeatGuard bool
	// ContinueOnLength 为真时，FinishReason=length 且没有工具调用就续写，最多 3 次。
	// 续写提示不写进 Session。第 4 次截断返回错误。
	ContinueOnLength bool
}

// StopInfo 是 Stop 钩子看到的本 run 进度。
type StopInfo struct {
	Continuations int
	ToolCalls     int
}

// StopDecision 决定模型停下之后是否再跑一轮。
type StopDecision struct {
	Continue    bool
	Replacement string
	Reason      string
	Notice      string
}

// Budget 是循环层的执行预算；与 exec.Budget 字段同构，由适配层换算。
type Budget struct {
	MaxPromptTokens     int
	MaxCompletionTokens int
	MaxToolCalls        int
	Deadline            time.Time
}

// Result is the final assistant text and the transcript produced by one run.
type Result struct {
	Text       string
	Messages   []llm.Message
	StopReason llm.StopReason
	Usage      llm.Usage
	Notice     string
	Err        error
	Retryable  bool
}

// Run executes the turn loop until the model stops, follow-ups drain, or the turn cap is hit.
func Run(ctx context.Context, cfg Config, user llm.Message) (Result, error) {
	if cfg.Stream == nil {
		return Result{StopReason: llm.StopError, Err: errString("stream func is required")}, nil
	}
	if cfg.MaxTurns <= 0 {
		cfg.MaxTurns = 12
	}
	if cfg.ToolExecution == "" {
		cfg.ToolExecution = ExecSequential
	}
	var messages []llm.Message
	if cfg.Session != nil {
		loaded, err := cfg.Session.Load()
		if err != nil {
			return Result{}, err
		}
		messages = loaded
	}
	if len(messages) == 0 && len(cfg.History) > 0 {
		messages = append([]llm.Message{}, cfg.History...)
	}
	if len(messages) == 0 || messages[0].Role != llm.RoleSystem {
		sys := llm.Message{Role: llm.RoleSystem, Content: systemText(cfg)}
		messages = append([]llm.Message{sys}, messages...)
		if err := persist(cfg, sys); err != nil {
			return Result{}, err
		}
	}
	if err := persist(cfg, user); err != nil {
		return Result{}, err
	}
	messages = append(messages, user)
	emit(cfg, Event{Type: "agent_start"})

	var last llm.Message
	var usage llm.Usage
	toolCalls := 0
	stopContinues := 0
	lengthContinues := 0
	lengthNudge := false
	var repeats repeatState
	for turn := 0; turn < cfg.MaxTurns; turn++ {
		if ctx.Err() != nil {
			cfg.logger().Error("runtime turn aborted", ctx.Err())
			return Result{Text: last.Content, Messages: messages, StopReason: llm.StopAborted, Err: ctx.Err()}, nil
		}
		if reason := budgetExceeded(cfg.Budget, usage, toolCalls, time.Now()); reason != "" {
			return budgetStop(cfg, last, messages, usage, reason), nil
		}
		emit(cfg, Event{Type: "turn_start"})
		emit(cfg, Event{Type: "message_start"})
		request := withEphemeral(messages, turnReminders(cfg, turn, &repeats, &lengthNudge))
		end, errEv := consume(ctx, cfg, request)
		if errEv != nil {
			cfg.logger().Error("runtime turn failed", errEv, "retryable", end.Retryable)
			return Result{Text: end.Message.Content, Messages: messages, StopReason: end.Message.StopReason, Err: errEv, Retryable: end.Retryable}, nil
		}
		last = end.Message
		usage.PromptTokens += end.Usage.PromptTokens
		usage.CompletionTokens += end.Usage.CompletionTokens
		usage.TotalTokens += end.Usage.TotalTokens
		if err := persist(cfg, last); err != nil {
			return Result{}, err
		}
		messages = append(messages, last)
		emit(cfg, Event{Type: "message_end", Message: last, PromptTokens: end.Usage.PromptTokens, CompletionTokens: end.Usage.CompletionTokens, TotalTokens: end.Usage.TotalTokens})

		if last.StopReason == llm.StopToolCalls && len(last.ToolCalls) > 0 {
			if cfg.Budget != nil && cfg.Budget.MaxToolCalls > 0 && toolCalls+len(last.ToolCalls) > cfg.Budget.MaxToolCalls {
				reason := fmt.Sprintf("budget_exceeded: tool calls %d+%d > max %d", toolCalls, len(last.ToolCalls), cfg.Budget.MaxToolCalls)
				return budgetStop(cfg, last, messages, usage, reason), nil
			}
			toolCalls += len(last.ToolCalls)
			if cfg.RepeatGuard {
				repeats.observe(last.ToolCalls)
			}
			results := executeTools(ctx, cfg, last)
			for _, res := range results {
				if err := persist(cfg, res); err != nil {
					return Result{}, err
				}
				messages = append(messages, res)
			}
			if steered := drain(cfg.GetSteering); len(steered) > 0 {
				for _, m := range steered {
					if err := persist(cfg, m); err != nil {
						return Result{}, err
					}
					messages = append(messages, m)
				}
			}
			emit(cfg, Event{Type: "turn_end", Message: last})
			continue
		}
		emit(cfg, Event{Type: "turn_end", Message: last})
		if cfg.ContinueOnLength && end.FinishReason == "length" && len(last.ToolCalls) == 0 {
			if lengthContinues >= maxLengthContinuations {
				errCap := errString("length continuation exceeded")
				cfg.logger().Error("runtime length continuation exceeded", errCap)
				return Result{Text: last.Content, Messages: messages, StopReason: llm.StopError, Usage: usage, Err: errCap}, nil
			}
			lengthContinues++
			lengthNudge = true
			continue
		}
		if follow := drain(cfg.GetFollowUp); len(follow) > 0 {
			for _, m := range follow {
				if err := persist(cfg, m); err != nil {
					return Result{}, err
				}
				messages = append(messages, m)
			}
			continue
		}
		if cfg.Stop != nil {
			decision := cfg.Stop(last, StopInfo{Continuations: stopContinues, ToolCalls: toolCalls})
			if decision.Continue && stopContinues < maxStopContinuations {
				stopContinues++
				if decision.Replacement != "" && len(messages) > 0 {
					messages[len(messages)-1].Content = decision.Replacement
					last.Content = decision.Replacement
				}
				if decision.Reason != "" {
					reasonMsg := llm.Message{Role: llm.RoleUser, Content: decision.Reason}
					if err := persist(cfg, reasonMsg); err != nil {
						return Result{}, err
					}
					messages = append(messages, reasonMsg)
				}
				continue
			}
			return settle(cfg, last, messages, usage, decision.Notice), nil
		}
		return settle(cfg, last, messages, usage, ""), nil
	}
	errCap := errString("max turns exceeded")
	cfg.logger().Error("runtime turn cap hit", errCap)
	return Result{
		Text:       last.Content,
		Messages:   messages,
		StopReason: llm.StopError,
		Usage:      usage,
		Err:        errCap,
	}, nil
}

func (cfg Config) logger() *logsys.Logger {
	if cfg.Log == nil {
		return logsys.Nop()
	}
	log := cfg.Log.Module("runtime")
	if cfg.ExecutionID == "" {
		return log
	}
	return log.With("execution_id", cfg.ExecutionID)
}

func systemText(cfg Config) string {
	var b strings.Builder
	b.WriteString(cfg.SystemPrompt)
	if cfg.Tools == nil {
		return b.String()
	}
	decls := cfg.Tools.Decls()
	if len(decls) == 0 {
		return b.String()
	}
	if b.Len() > 0 {
		b.WriteString("\n\n")
	}
	b.WriteString("Tools available in this workspace:")
	for _, d := range decls {
		b.WriteString("\n- ")
		b.WriteString(d.Name)
		b.WriteString(": ")
		b.WriteString(d.Description)
	}
	return b.String()
}

func consume(ctx context.Context, cfg Config, messages []llm.Message) (llm.Event, error) {
	var decls []llm.ToolDecl
	if cfg.Tools != nil {
		for _, d := range cfg.Tools.Decls() {
			decls = append(decls, llm.ToolDecl{Name: d.Name, Description: d.Description, Parameters: d.Parameters})
		}
	}
	model := cfg.Model
	opts := llm.StreamOptions{ThinkingLevel: cfg.ThinkingLevel}
	if cfg.Resolve != nil {
		if m, th := cfg.Resolve(); m != "" || th != "" {
			if m != "" {
				model = m
			}
			if th != "" {
				opts.ThinkingLevel = th
			}
		}
	}
	var end llm.Event
	sawEnd := false
	for ev := range cfg.Stream(ctx, model, messages, decls, opts) {
		switch ev.Kind {
		case llm.KindTextDelta:
			emit(cfg, Event{Type: "text_delta", Text: ev.Text})
		case llm.KindReasoningDelta:
			emit(cfg, Event{Type: "reasoning_delta", Text: ev.ReasoningDelta})
		case llm.KindToolCallDelta:
			emit(cfg, Event{Type: "tool_call_delta", Tool: ev.ToolName, Detail: ev.ArgumentsDelta})
		case llm.KindError:
			return ev, ev.Err
		case llm.KindMessageEnd:
			end = ev
			sawEnd = true
		}
	}
	if !sawEnd {
		return llm.Event{Message: llm.Message{StopReason: llm.StopError}}, errString("stream ended without a message")
	}
	return end, nil
}

func executeTools(ctx context.Context, cfg Config, assistant llm.Message) []llm.Message {
	if cfg.ToolExecution == ExecParallel && len(assistant.ToolCalls) > 1 {
		return executeParallel(ctx, cfg, assistant)
	}
	out := make([]llm.Message, 0, len(assistant.ToolCalls))
	for _, call := range assistant.ToolCalls {
		out = append(out, runOne(ctx, cfg, call))
	}
	return out
}

func executeParallel(ctx context.Context, cfg Config, assistant llm.Message) []llm.Message {
	out := make([]llm.Message, len(assistant.ToolCalls))
	var wg sync.WaitGroup
	for i, call := range assistant.ToolCalls {
		wg.Add(1)
		go func(i int, call llm.ToolCall) {
			defer wg.Done()
			out[i] = runOne(ctx, cfg, call)
		}(i, call)
	}
	wg.Wait()
	return out
}

func runOne(ctx context.Context, cfg Config, call llm.ToolCall) llm.Message {
	emit(cfg, Event{Type: "tool_execution_start", Tool: call.Name, Detail: call.Arguments})
	if cfg.BeforeTool != nil {
		if err := cfg.BeforeTool(call); err != nil {
			msg := llm.Message{Role: llm.RoleTool, Name: call.Name, ToolCallID: call.ID, Content: err.Error(), IsError: true}
			emit(cfg, Event{Type: "tool_execution_end", Tool: call.Name, Text: msg.Content})
			if cfg.AfterTool != nil {
				cfg.AfterTool(call, msg)
			}
			return msg
		}
	}
	emit(cfg, Event{Type: "tool_execution_update", Tool: call.Name, Detail: call.Arguments})
	res := tools.Result{Content: "no tools configured", IsError: true}
	known := cfg.Tools != nil && cfg.Tools.Has(call.Name)
	switch {
	case known:
		res = cfg.Tools.Execute(ctx, call.Name, json.RawMessage(call.Arguments))
	case cfg.UnknownTool != nil:
		res = cfg.UnknownTool(call)
	case cfg.Tools != nil:
		res = cfg.Tools.Execute(ctx, call.Name, json.RawMessage(call.Arguments))
	}
	msg := llm.Message{
		Role:       llm.RoleTool,
		Name:       call.Name,
		ToolCallID: call.ID,
		Content:    res.Content,
		IsError:    res.IsError,
	}
	emit(cfg, Event{Type: "tool_execution_end", Tool: call.Name, Text: res.Content})
	if cfg.AfterTool != nil {
		cfg.AfterTool(call, msg)
	}
	return msg
}

func settle(cfg Config, last llm.Message, messages []llm.Message, usage llm.Usage, notice string) Result {
	emit(cfg, Event{Type: "agent_end", Message: last})
	emit(cfg, Event{Type: "agent_settled", Message: last})
	emitUsage(cfg, usage)
	return Result{Text: last.Content, Messages: messages, StopReason: last.StopReason, Usage: usage, Notice: notice}
}

func reminderMessage(text string) llm.Message {
	return llm.Message{Role: llm.RoleUser, Content: "<system-reminder>\n" + text + "\n</system-reminder>"}
}

func turnReminders(cfg Config, turn int, repeats *repeatState, lengthNudge *bool) []llm.Message {
	var extra []llm.Message
	if cfg.Reminders != nil {
		for _, text := range cfg.Reminders(turn) {
			if strings.TrimSpace(text) == "" {
				continue
			}
			extra = append(extra, reminderMessage(text))
		}
	}
	if lengthNudge != nil && *lengthNudge {
		extra = append(extra, llm.Message{Role: llm.RoleUser, Content: lengthContinuePrompt})
		*lengthNudge = false
	}
	if cfg.RepeatGuard && repeats != nil {
		due := repeats.streak/3 - repeats.reminders
		room := maxRepeatReminders - repeats.reminders
		if due > room {
			due = room
		}
		for i := 0; i < due; i++ {
			extra = append(extra, reminderMessage(repeatReminderText))
			repeats.reminders++
		}
	}
	return extra
}

func withEphemeral(messages, extra []llm.Message) []llm.Message {
	if len(extra) == 0 {
		return messages
	}
	out := make([]llm.Message, 0, len(messages)+len(extra))
	out = append(out, messages...)
	out = append(out, extra...)
	return out
}

type repeatState struct {
	sig       string
	streak    int
	reminders int
}

func (s *repeatState) observe(calls []llm.ToolCall) {
	for _, call := range calls {
		sig := call.Name + "\n" + canonicalArgs(call.Arguments)
		if sig == s.sig {
			s.streak++
			continue
		}
		s.sig = sig
		s.streak = 1
	}
}

func canonicalArgs(raw string) string {
	var v any
	if err := json.Unmarshal([]byte(raw), &v); err != nil {
		return raw
	}
	b, err := json.Marshal(v)
	if err != nil {
		return raw
	}
	return string(b)
}

func persist(cfg Config, msg llm.Message) error {
	if cfg.Session == nil {
		return nil
	}
	return cfg.Session.Append(msg)
}

func drain(fn func() []llm.Message) []llm.Message {
	if fn == nil {
		return nil
	}
	return fn()
}

func emitUsage(cfg Config, usage llm.Usage) {
	emit(cfg, Event{Type: "usage", PromptTokens: usage.PromptTokens, CompletionTokens: usage.CompletionTokens, TotalTokens: usage.TotalTokens})
}

func emit(cfg Config, ev Event) {
	if cfg.OnEvent != nil {
		cfg.OnEvent(ev)
	}
}

// budgetExceeded 返回命中的预算说明；空串表示预算仍充足（或未设预算）。
// 只在「即将消耗更多预算」的节点调用（轮次顶端 / 工具批执行前）。
func budgetExceeded(b *Budget, usage llm.Usage, toolCalls int, now time.Time) string {
	if b == nil {
		return ""
	}
	if !b.Deadline.IsZero() && now.After(b.Deadline) {
		return fmt.Sprintf("budget_exceeded: deadline %s passed", b.Deadline.Format(time.RFC3339))
	}
	if b.MaxPromptTokens > 0 && usage.PromptTokens > b.MaxPromptTokens {
		return fmt.Sprintf("budget_exceeded: prompt tokens %d > max %d", usage.PromptTokens, b.MaxPromptTokens)
	}
	if b.MaxCompletionTokens > 0 && usage.CompletionTokens > b.MaxCompletionTokens {
		return fmt.Sprintf("budget_exceeded: completion tokens %d > max %d", usage.CompletionTokens, b.MaxCompletionTokens)
	}
	if b.MaxToolCalls > 0 && toolCalls > b.MaxToolCalls {
		return fmt.Sprintf("budget_exceeded: tool calls %d > max %d", toolCalls, b.MaxToolCalls)
	}
	return ""
}

// budgetStop 以预算超限收尾：发事件、留日志、按 StopError 返回（不可重试）。
func budgetStop(cfg Config, last llm.Message, messages []llm.Message, usage llm.Usage, reason string) Result {
	emit(cfg, Event{Type: "budget_exceeded", Text: reason})
	err := errString(reason)
	cfg.logger().Error("runtime budget exceeded", err)
	return Result{Text: last.Content, Messages: messages, StopReason: llm.StopError, Usage: usage, Err: err}
}

type errString string

func (e errString) Error() string { return string(e) }