agent.go

package core

import (
	"context"
	"fmt"
	"sync"
)

// StreamConsumer is called for each delta when Run uses streaming. When nil, Run uses
// Complete only. Thinking deltas are for display only; the consumer must not persist them.
type StreamConsumer func(d Delta)

// Agent is the central unit of execution. It wraps a system prompt, a model configuration,
// a set of tools, optional long-term memory, and a callback chain.
//
// An Agent is safe for concurrent use: multiple goroutines may call Run simultaneously,
// each with its own Conversation. Internal state is read-only after construction.
type Agent struct {
	// ID is a unique identifier for this agent instance, used in memory entries and logs.
	ID string

	// Instructions is the system prompt prepended to every conversation on Run.
	// If the conversation already starts with a system message, Instructions is not prepended.
	Instructions string

	// Model holds the provider and generation parameters.
	Model ModelConfig

	// Tools is the set of tools available to this agent.
	// The agent builds an internal name→tool map at the start of each Run call.
	Tools []Tool

	// Memory is the optional long-term memory store. Nil disables long-term memory.
	Memory MemoryStore

	// Callbacks is the lifecycle hook chain. A nil slice is a valid no-op chain.
	Callbacks CallbackChain

	// MaxTurns caps LLM sampling rounds per Run; 0 means unlimited (default).
	MaxTurns int

	// QueryEventSink receives query-loop observability events when non-nil.
	QueryEventSink func(QueryEvent)

	// QueryLoopHooks wires optional query-loop extension points (permission view, stop hooks).
	QueryLoopHooks QueryLoopHooks

	// QueryTracking optionally supplies chain identity for sub-agent runs; nil uses a root chain per Run.
	QueryTracking *QueryTracking
}

// Run executes the agent loop for the given conversation via the Claude-parity query loop
// (prepare → model → recoveries → tools → prefetch stubs; see query_loop_*.go).
//
// It prepends the system prompt (if Instructions is set and not already present), then iterates:
//  1. BeforeLLMCall callbacks
//  2. Provider completion (Stream when streamConsumer is non-nil, else Complete)
//  3. AfterLLMCall callbacks
//  4. If tool calls present: dispatch concurrently, fire per-tool callbacks, append results, continue
//  5. If no tool calls: stop
//
// When streamConsumer is non-nil, Run uses Provider.Stream and invokes it for each delta;
// only the final assistant message (from Delta.Final) is appended to the conversation.
// Run modifies conv in place and also returns it.
// ctx cancellation aborts the loop at the next blocking call.
func (a *Agent) Run(ctx context.Context, conv *Conversation, streamConsumer StreamConsumer) (*Conversation, error) {
	if a.Instructions != "" {
		if len(conv.Messages) == 0 || conv.Messages[0].Role != RoleSystem {
			conv.Messages = append([]Message{{Role: RoleSystem, Content: a.Instructions}}, conv.Messages...)
		}
	}

	// Build name→tool map once per Run to avoid per-iteration allocations.
	toolMap := make(map[string]Tool, len(a.Tools))
	for _, t := range a.Tools {
		toolMap[t.Name()] = t
	}

	if err := a.Callbacks.BeforeAgentLoop(ctx, conv); err != nil {
		return nil, err
	}

	dispatch := func(ctx context.Context, calls []ToolCall) []ToolResult {
		return a.dispatchTools(ctx, calls, toolMap)
	}
	err := runQueryLoop(ctx, &queryLoopParams{
		Agent:          a,
		Conv:           conv,
		StreamConsumer: streamConsumer,
		MaxTurns:       a.MaxTurns,
		EventSink:      a.QueryEventSink,
		Tracking:       a.QueryTracking,
		Hooks:          a.QueryLoopHooks,
		DispatchTools:  dispatch,
	})
	if err != nil {
		return nil, err
	}

	if err := a.Callbacks.AfterAgentLoop(ctx, conv); err != nil {
		return nil, err
	}

	return conv, nil
}

// dispatchTools executes all tool calls concurrently and collects results in call order.
// Each call fires BeforeToolExecution (which may mutate or reject the call) then AfterToolExecution.
func (a *Agent) dispatchTools(ctx context.Context, calls []ToolCall, toolMap map[string]Tool) []ToolResult {
	results := make([]ToolResult, len(calls))
	var wg sync.WaitGroup

	for i, call := range calls {
		wg.Add(1)
		// Capture loop variables.
		go func(i int, call ToolCall) {
			defer wg.Done()

			call, err := a.Callbacks.BeforeToolExecution(ctx, call)
			if err != nil {
				results[i] = ToolResult{CallID: call.ID, Name: call.Name, Err: err}
				return
			}

			t, ok := toolMap[call.Name]
			if !ok {
				results[i] = ToolResult{
					CallID: call.ID,
					Name:   call.Name,
					Err:    fmt.Errorf("unknown tool: %s", call.Name),
				}
				return
			}

			content, execErr := t.Execute(ctx, call.Params)
			result := ToolResult{CallID: call.ID, Name: call.Name, Content: content, Err: execErr}

			// AfterToolExecution errors are not fatal to the loop.
			_ = a.Callbacks.AfterToolExecution(ctx, call, result)

			results[i] = result
		}(i, call)
	}

	wg.Wait()
	return results
}