Agent-System/internal/agentx/loop.go

package agentx

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

	"chat-agent-system/internal/agentx/llm"
	"chat-agent-system/internal/agentx/session"
	"chat-agent-system/internal/agentx/tools"
)

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
}

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

// 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
	for turn := 0; turn < cfg.MaxTurns; turn++ {
		if ctx.Err() != nil {
			return Result{Text: last.Content, Messages: messages, StopReason: llm.StopAborted, Err: ctx.Err()}, nil
		}
		emit(cfg, Event{Type: "turn_start"})
		end, errEv := consume(ctx, cfg, messages)
		if errEv != nil {
			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 {
			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})
		return Result{Text: last.Content, Messages: messages, StopReason: last.StopReason, Usage: usage}, nil
	}
	return Result{
		Text:       last.Content,
		Messages:   messages,
		StopReason: llm.StopError,
		Usage:      usage,
		Err:        errString("max turns exceeded"),
	}, nil
}

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 emit(cfg Config, ev Event) {
	if cfg.OnEvent != nil {
		cfg.OnEvent(ev)
	}
}

type errString string

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