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"
)

// 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 只在 type=usage 时有意义。
	PromptTokens     int
	CompletionTokens 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)
	// Log 可选。nil 时丢弃日志，不改变循环结果。
	Log *logsys.Logger
	// ExecutionID 只用于日志字段，不参与控制流。
	ExecutionID string
	// Budget 可选。任一上限命中即停止并返回 budget_exceeded 错误（StopError）；
	// 零值字段不设限。截止时刻与外层 ctx 先到者生效。
	Budget *Budget
}

// 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
	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 || 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
	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"})
		end, errEv := consume(ctx, cfg, messages)
		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})

		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)
			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 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
		}
		emit(cfg, Event{Type: "agent_end", Message: last})
		emitUsage(cfg, usage)
		return Result{Text: last.Content, Messages: messages, StopReason: last.StopReason, Usage: 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})
		}
	}
	var end llm.Event
	sawEnd := false
	for ev := range cfg.Stream(ctx, cfg.Model, messages, decls, llm.StreamOptions{}) {
		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})
	res := tools.Result{Content: "no tools configured", IsError: true}
	if cfg.Tools != nil {
		res = cfg.Tools.Execute(ctx, call.Name, json.RawMessage(call.Arguments))
	}
	emit(cfg, Event{Type: "tool_execution_end", Tool: call.Name, Text: res.Content})
	return llm.Message{
		Role:       llm.RoleTool,
		Name:       call.Name,
		ToolCallID: call.ID,
		Content:    res.Content,
		IsError:    res.IsError,
	}
}

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})
}

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) }