backend/internal/agent/run.go

package agent

import (
	"context"
	"fmt"
	"strings"

	"code.littlelan.cn/CarrotAssistant/backend/internal/llm"
	"code.littlelan.cn/CarrotAssistant/backend/internal/skill"
)

// Event is what Run yields to its caller (the SSE handler). Exactly one field
// is meaningful per event; the Kind tells which.
type Event struct {
	Kind EventKind

	// Text is set for KindToken.
	Text string
	// ToolCall is set for KindToolCall (a tool is about to be executed).
	ToolCall ToolCallEvent
	// ToolResult is set for KindToolResult (the outcome of a tool call).
	ToolResult ToolResultEvent
	// Pending is set for KindConfirmation (a mutating call awaits user OK).
	Pending ConfirmationEvent
	// Message is set for KindError.
	Message string
}

type EventKind string

const (
	KindToken        EventKind = "token"         // assistant text delta
	KindToolCall     EventKind = "tool_call"     // a tool was requested
	KindToolResult   EventKind = "tool_result"   // tool finished
	KindConfirmation EventKind = "confirmation"  // needs user confirm
	KindDone         EventKind = "done"          // turn finished
	KindError        EventKind = "error"         // fatal error
)

// ToolCallEvent describes a tool the model asked to run.
type ToolCallEvent struct {
	ID        string         `json:"id"`
	Name      string         `json:"name"`
	Arguments map[string]any `json:"arguments"`
}

// ToolResultEvent describes a tool outcome.
type ToolResultEvent struct {
	CallID string `json:"call_id"`
	// Result is the content sent back to the model. For confirmations it is
	// the pending-action description.
	Result string `json:"result"`
	// Denied / Pending flag the non-executed outcomes.
	Denied  bool `json:"denied,omitempty"`
	Pending bool `json:"pending,omitempty"`
}

// ConfirmationEvent carries the call that needs approval.
type ConfirmationEvent struct {
	CallID    string         `json:"call_id"`
	Name      string         `json:"name"`
	Arguments map[string]any `json:"arguments"`
	Detail    string         `json:"detail"` // human-readable summary
}

// Runner drives a single conversation turn. It is constructed per request
// and is not safe for reuse.
type Runner struct {
	Client  *llm.Client
	Skill   *skill.Skill
	Tools   map[string]skill.Tool // name -> tool, for lookup during execution
	History []llm.Message         // prior messages (system + conversation)
	OnEvent func(Event)
	pending *pendingCall // set when Run paused for confirmation
	startLen int         // length of History at construction; messages beyond this are new this turn
}

// NewRunner builds a runner for one skill. History should already include the
// system prompt as its first message.
func NewRunner(client *llm.Client, sk *skill.Skill, history []llm.Message, onEvent func(Event)) *Runner {
	tools := map[string]skill.Tool{}
	for _, t := range sk.Tools {
		tools[t.Name] = t
	}
	if onEvent == nil {
		onEvent = func(Event) {}
	}
	return &Runner{Client: client, Skill: sk, Tools: tools, History: history, OnEvent: onEvent, startLen: len(history)}
}

// HistoryTail returns only the messages this Runner added during Run/Resume —
// i.e. History[startLen:]. Callers use it to persist just the new turn,
// avoiding double-saving the input history.
func (r *Runner) HistoryTail() []llm.Message {
	if r.startLen >= len(r.History) {
		return nil
	}
	return r.History[r.startLen:]
}

// Run executes the assistant loop until the model produces a final text
// answer or a hard limit is hit. If a mutating tool needs confirmation, Run
// returns a *PendingConfirmation sentinel so the caller can resume later via
// ResumeAfterConfirmation.
//
// userMessage, when non-empty, is appended as the latest user turn before
// running. When empty, the caller is expected to have already included the
// user message in History (e.g. loaded from storage).
func (r *Runner) Run(ctx context.Context, userMessage string) error {
	history := append([]llm.Message{}, r.History...)
	if strings.TrimSpace(userMessage) != "" {
		history = append(history, llm.Message{Role: llm.RoleUser, Content: userMessage})
		r.startLen = len(r.History) // user msg is "new" this turn
	}
	// When userMessage is empty, History already holds the user message;
	// startLen stays as set at construction so HistoryTail excludes it.

	for round := 0; round < MaxToolRounds; round++ {
		assistant, err := r.streamOnce(ctx, history)
		if err != nil {
			r.OnEvent(Event{Kind: KindError, Message: err.Error()})
			return err
		}
		history = append(history, assistant)

		if len(assistant.ToolCalls) == 0 {
			// Final answer delivered. Persist the turn by handing the caller
			// the updated history through a done event.
			r.History = history
			r.OnEvent(Event{Kind: KindDone})
			return nil
		}

		// Execute each requested tool call.
		for _, call := range assistant.ToolCalls {
			tool, ok := r.Tools[call.Name]
			if !ok {
				// Model hallucinated a tool; feed the error back and continue.
				history = append(history, llm.Message{
					Role: llm.RoleTool, ToolCallID: call.ID, Name: call.Name,
					Content: "[工具不存在: " + call.Name + "]",
				})
				continue
			}
			out, err := Execute(ctx, tool, r.Skill.Permissions, call.Arguments)
			if err != nil {
				history = append(history, llm.Message{
					Role: llm.RoleTool, ToolCallID: call.ID, Name: call.Name,
					Content: "[工具执行错误: " + err.Error() + "]",
				})
				continue
			}

			if out.Denied {
				r.OnEvent(Event{Kind: KindToolResult, ToolResult: ToolResultEvent{
					CallID: call.ID, Result: out.Result, Denied: true,
				}})
				history = append(history, llm.Message{
					Role: llm.RoleTool, ToolCallID: call.ID, Name: call.Name,
					Content: out.Result,
				})
				continue
			}
			if out.NeedsConfirmation {
				// Pause the loop: persist intermediate state and surface the
				// request. The SSE handler will hold the connection / resume.
				r.OnEvent(Event{Kind: KindConfirmation, Pending: ConfirmationEvent{
					CallID: call.ID, Name: call.Name, Arguments: call.Arguments,
					Detail: out.Result,
				}})
				// Stash resume context for the handler.
				r.pending = &pendingCall{
					call:     call,
					tool:     tool,
					history:  history,
				}
				return ErrPendingConfirmation
			}

			r.OnEvent(Event{Kind: KindToolResult, ToolResult: ToolResultEvent{
				CallID: call.ID, Result: truncatePreview(out.Result),
			}})
			history = append(history, llm.Message{
				Role: llm.RoleTool, ToolCallID: call.ID, Name: call.Name,
				Content: out.Result,
			})
		}
		// Loop again: the model now sees tool results and may answer or call more.
	}

	r.History = history
	r.OnEvent(Event{Kind: KindError, Message: "已达到工具调用轮数上限"})
	return ErrTooManyRounds
}

// pendingCall holds state needed to resume after a user confirms a mutating
// tool call. Set when Run returns ErrPendingConfirmation.
type pendingCall struct {
	call    llm.ToolCall
	tool    skill.Tool
	history []llm.Message
}

// ErrPendingConfirmation signals that Run paused awaiting user confirmation.
// ResumeAfterConfirmation continues from where it stopped.
var ErrPendingConfirmation = fmt.Errorf("pending confirmation")

// ErrTooManyRounds signals the loop hit MaxToolRounds.
var ErrTooManyRounds = fmt.Errorf("too many tool rounds")

// PendingCall returns the call awaiting confirmation, if any.
func (r *Runner) PendingCall() (llm.ToolCall, bool) {
	if r.pending == nil {
		return llm.ToolCall{}, false
	}
	return r.pending.call, true
}

// ResumeAfterConfirmation continues the loop after the user approves or denies
// a paused mutating tool call. On approve, the tool executes and the loop
// continues; on deny, a "user declined" tool message is appended and the loop
// continues (the model can then answer without the data).
func (r *Runner) ResumeAfterConfirmation(ctx context.Context, approved bool) error {
	if r.pending == nil {
		return fmt.Errorf("no pending confirmation")
	}
	pc := r.pending
	r.pending = nil
	history := pc.history

	if !approved {
		history = append(history, llm.Message{
			Role: llm.RoleTool, ToolCallID: pc.call.ID, Name: pc.call.Name,
			Content: "[用户拒绝了该操作]",
		})
	} else {
		out, err := Execute(ctx, pc.tool, r.Skill.Permissions, pc.call.Arguments)
		if err != nil {
			history = append(history, llm.Message{
				Role: llm.RoleTool, ToolCallID: pc.call.ID, Name: pc.call.Name,
				Content: "[工具执行错误: " + err.Error() + "]",
			})
		} else {
			r.OnEvent(Event{Kind: KindToolResult, ToolResult: ToolResultEvent{
				CallID: pc.call.ID, Result: truncatePreview(out.Result),
			}})
			history = append(history, llm.Message{
				Role: llm.RoleTool, ToolCallID: pc.call.ID, Name: pc.call.Name,
				Content: out.Result,
			})
		}
	}

	// Continue the loop with the remaining round budget.
	for round := 0; round < MaxToolRounds; round++ {
		assistant, err := r.streamOnce(ctx, history)
		if err != nil {
			r.OnEvent(Event{Kind: KindError, Message: err.Error()})
			return err
		}
		history = append(history, assistant)
		if len(assistant.ToolCalls) == 0 {
			r.History = history
			r.OnEvent(Event{Kind: KindDone})
			return nil
		}
		for _, call := range assistant.ToolCalls {
			tool, ok := r.Tools[call.Name]
			if !ok {
				history = append(history, llm.Message{
					Role: llm.RoleTool, ToolCallID: call.ID, Name: call.Name,
					Content: "[工具不存在: " + call.Name + "]",
				})
				continue
			}
			out, err := Execute(ctx, tool, r.Skill.Permissions, call.Arguments)
			if err != nil {
				history = append(history, llm.Message{
					Role: llm.RoleTool, ToolCallID: call.ID, Name: call.Name,
					Content: "[工具执行错误: " + err.Error() + "]",
				})
				continue
			}
			if out.Denied {
				history = append(history, llm.Message{
					Role: llm.RoleTool, ToolCallID: call.ID, Name: call.Name,
					Content: out.Result,
				})
				continue
			}
			if out.NeedsConfirmation {
				r.OnEvent(Event{Kind: KindConfirmation, Pending: ConfirmationEvent{
					CallID: call.ID, Name: call.Name, Arguments: call.Arguments,
					Detail: out.Result,
				}})
				r.pending = &pendingCall{call: call, tool: tool, history: history}
				return ErrPendingConfirmation
			}
			r.OnEvent(Event{Kind: KindToolResult, ToolResult: ToolResultEvent{
				CallID: call.ID, Result: truncatePreview(out.Result),
			}})
			history = append(history, llm.Message{
				Role: llm.RoleTool, ToolCallID: call.ID, Name: call.Name,
				Content: out.Result,
			})
		}
	}
	r.History = history
	r.OnEvent(Event{Kind: KindError, Message: "已达到工具调用轮数上限"})
	return ErrTooManyRounds
}

// streamOnce calls the LLM with the current history, streaming tokens and
// surfacing the assembled tool calls. Returns the assistant message.
func (r *Runner) streamOnce(ctx context.Context, history []llm.Message) (llm.Message, error) {
	funcs := ToFunctions(r.Skill.Tools)
	msg, err := r.Client.RunStream(ctx, history, funcs, func(ev llm.StreamEvent) {
		switch {
		case ev.Delta != "":
			r.OnEvent(Event{Kind: KindToken, Text: ev.Delta})
		case len(ev.ToolCalls) > 0:
			for _, c := range ev.ToolCalls {
				r.OnEvent(Event{Kind: KindToolCall, ToolCall: ToolCallEvent{
					ID: c.ID, Name: c.Name, Arguments: c.Arguments,
				}})
			}
		}
	})
	if err != nil {
		return llm.Message{}, err
	}
	return msg, nil
}

// truncatePreview keeps the SSE event payload small; the full result still
// goes into history for the model.
func truncatePreview(s string) string {
	const preview = 500
	if len(s) <= preview {
		return s
	}
	return s[:preview] + "…"
}