diff --git a/internal/fantasy/agent.go b/internal/fantasy/agent.go index 426beff16..87902653a 100644 --- a/internal/fantasy/agent.go +++ b/internal/fantasy/agent.go @@ -140,9 +140,13 @@ type agentSettings struct { headers map[string]string providerOptions ProviderOptions - // TODO: add support for provider tools - tools []AgentTool - maxRetries *int + tools []AgentTool + toolRuntime ToolRuntime + maxRetries *int + + transitionObservers []ReActTransitionObserver + stepObservers []ReActStepObserver + toolResultObservers []ReActToolResultObserver model LanguageModel @@ -469,7 +473,12 @@ func (a *agent) Generate(ctx context.Context, opts AgentCall) (*AgentResult, err } } - toolResults, err := a.executeTools(ctx, stepTools, stepToolCalls, nil) + var toolResults []ToolResultContent + if a.settings.toolRuntime != nil { + toolResults, err = a.settings.toolRuntime.Execute(ctx, stepTools, stepToolCalls, nil) + } else { + toolResults, err = a.executeTools(ctx, stepTools, stepToolCalls, nil) + } // Build step content with validated tool calls and tool results stepContent := []Content{} @@ -486,9 +495,11 @@ func (a *agent) Generate(ctx context.Context, opts AgentCall) (*AgentResult, err stepContent = append(stepContent, content) } } - // Add tool results - for _, result := range toolResults { - stepContent = append(stepContent, result) + for _, tr := range toolResults { + stepContent = append(stepContent, tr) + for _, obs := range a.settings.toolResultObservers { + obs.OnReActToolResult(ctx, len(steps), tr) + } } currentStepMessages := toResponseMessages(stepContent) responseMessages = append(responseMessages, currentStepMessages...) @@ -504,6 +515,11 @@ func (a *agent) Generate(ctx context.Context, opts AgentCall) (*AgentResult, err Messages: currentStepMessages, } steps = append(steps, stepResult) + + for _, obs := range a.settings.stepObservers { + obs.OnReActStep(ctx, len(steps)-1, stepResult) + } + shouldStop := isStopConditionMet(opts.StopWhen, steps) if shouldStop || err != nil || len(stepToolCalls) == 0 || result.FinishReason != FinishReasonToolCalls { @@ -865,7 +881,10 @@ func (a *agent) Stream(ctx context.Context, opts AgentStreamCall) (*AgentResult, steps = append(steps, result.StepResult) totalUsage = addUsage(totalUsage, result.StepResult.Usage) - // Call step finished callback + for _, obs := range a.settings.stepObservers { + obs.OnReActStep(ctx, len(steps)-1, result.StepResult) + } + if opts.OnStepFinish != nil { _ = opts.OnStepFinish(result.StepResult) } @@ -1095,6 +1114,39 @@ func WithOnRetry(callback OnRetryCallback) AgentOption { } } +// WithToolRuntime sets a custom ToolRuntime (DAG, parallel, offloading, etc.) +// that replaces the default sequential tool execution. When set, the agent +// delegates all tool execution to this runtime. +func WithToolRuntime(rt ToolRuntime) AgentOption { + return func(s *agentSettings) { + s.toolRuntime = rt + } +} + +// WithTransitionObserver appends a ReActTransitionObserver that fires on +// every FSM state transition during Generate/Stream. +func WithTransitionObserver(o ReActTransitionObserver) AgentOption { + return func(s *agentSettings) { + s.transitionObservers = append(s.transitionObservers, o) + } +} + +// WithStepObserver appends a ReActStepObserver that fires after each +// completed step. +func WithStepObserver(o ReActStepObserver) AgentOption { + return func(s *agentSettings) { + s.stepObservers = append(s.stepObservers, o) + } +} + +// WithToolResultObserver appends a ReActToolResultObserver that fires for +// every tool result produced during execution. +func WithToolResultObserver(o ReActToolResultObserver) AgentOption { + return func(s *agentSettings) { + s.toolResultObservers = append(s.toolResultObservers, o) + } +} + // processStepStream processes a single step's stream and returns the step result. func (a *agent) processStepStream(ctx context.Context, stream StreamResponse, opts AgentStreamCall, _ []StepResult, stepTools []AgentTool) (stepExecutionResult, error) { var stepContent []Content @@ -1112,12 +1164,14 @@ func (a *agent) processStepStream(ctx context.Context, stream StreamResponse, op } activeReasoningContent := make(map[string]reasoningContent) - // Set up concurrent tool execution + useRuntimeBatch := a.settings.toolRuntime != nil + + // Set up concurrent tool execution (used only when no ToolRuntime is set) type toolExecutionRequest struct { toolCall ToolCallContent parallel bool } - toolChan := make(chan toolExecutionRequest, 10) + var toolChan chan toolExecutionRequest var toolExecutionWg sync.WaitGroup var toolStateMu sync.Mutex toolResults := make([]ToolResultContent, 0) @@ -1129,17 +1183,32 @@ func (a *agent) processStepStream(ctx context.Context, stream StreamResponse, op toolMap[tool.Info().Name] = tool } - // Semaphores for controlling parallelism - parallelSem := make(chan struct{}, 5) - var sequentialMu sync.Mutex + if !useRuntimeBatch { + toolChan = make(chan toolExecutionRequest, 10) - // Single coordinator goroutine that dispatches tools - toolExecutionWg.Go(func() { - for req := range toolChan { - if req.parallel { - parallelSem <- struct{}{} - toolExecutionWg.Go(func() { - defer func() { <-parallelSem }() + // Semaphores for controlling parallelism + parallelSem := make(chan struct{}, 5) + var sequentialMu sync.Mutex + + // Single coordinator goroutine that dispatches tools + toolExecutionWg.Go(func() { + for req := range toolChan { + if req.parallel { + parallelSem <- struct{}{} + toolExecutionWg.Go(func() { + defer func() { <-parallelSem }() + result, isCriticalError := a.executeSingleTool(ctx, toolMap, req.toolCall, opts.OnToolResult) + toolStateMu.Lock() + toolResults = append(toolResults, result) + if isCriticalError && toolExecutionErr == nil { + if errorResult, ok := result.Result.(ToolResultOutputContentError); ok && errorResult.Error != nil { + toolExecutionErr = errorResult.Error + } + } + toolStateMu.Unlock() + }) + } else { + sequentialMu.Lock() result, isCriticalError := a.executeSingleTool(ctx, toolMap, req.toolCall, opts.OnToolResult) toolStateMu.Lock() toolResults = append(toolResults, result) @@ -1149,22 +1218,11 @@ func (a *agent) processStepStream(ctx context.Context, stream StreamResponse, op } } toolStateMu.Unlock() - }) - } else { - sequentialMu.Lock() - result, isCriticalError := a.executeSingleTool(ctx, toolMap, req.toolCall, opts.OnToolResult) - toolStateMu.Lock() - toolResults = append(toolResults, result) - if isCriticalError && toolExecutionErr == nil { - if errorResult, ok := result.Result.(ToolResultOutputContentError); ok && errorResult.Error != nil { - toolExecutionErr = errorResult.Error - } + sequentialMu.Unlock() } - toolStateMu.Unlock() - sequentialMu.Unlock() } - } - }) + }) + } // Process stream parts for part := range stream { @@ -1320,16 +1378,14 @@ func (a *agent) processStepStream(ctx context.Context, stream StreamResponse, op } } - // Determine if tool can run in parallel - isParallel := false - if tool, exists := toolMap[validatedToolCall.ToolName]; exists { - isParallel = tool.Info().Parallel + if !useRuntimeBatch { + isParallel := false + if tool, exists := toolMap[validatedToolCall.ToolName]; exists { + isParallel = tool.Info().Parallel + } + toolChan <- toolExecutionRequest{toolCall: validatedToolCall, parallel: isParallel} } - // Send tool call to execution channel - toolChan <- toolExecutionRequest{toolCall: validatedToolCall, parallel: isParallel} - - // Clean up active tool call delete(activeToolCalls, part.ID) case StreamPartTypeSource: @@ -1364,19 +1420,26 @@ func (a *agent) processStepStream(ctx context.Context, stream StreamResponse, op } } - // Close the tool execution channel and wait for all executions to complete - close(toolChan) - toolExecutionWg.Wait() - - // Check for tool execution errors - if toolExecutionErr != nil { - return stepExecutionResult{}, toolExecutionErr + if useRuntimeBatch { + if len(stepToolCalls) > 0 { + var err error + toolResults, err = a.settings.toolRuntime.Execute(ctx, stepTools, stepToolCalls, opts.OnToolResult) + if err != nil { + return stepExecutionResult{}, err + } + } + } else { + close(toolChan) + toolExecutionWg.Wait() + if toolExecutionErr != nil { + return stepExecutionResult{}, toolExecutionErr + } } - // Add tool results to content if any - if len(toolResults) > 0 { - for _, result := range toolResults { - stepContent = append(stepContent, result) + for _, tr := range toolResults { + stepContent = append(stepContent, tr) + for _, obs := range a.settings.toolResultObservers { + obs.OnReActToolResult(ctx, 0, tr) } } diff --git a/internal/fantasy/react_fsm.go b/internal/fantasy/react_fsm.go new file mode 100644 index 000000000..35db33f5c --- /dev/null +++ b/internal/fantasy/react_fsm.go @@ -0,0 +1,84 @@ +package fantasy + +import ( + "sync" + "time" +) + +// ReActState represents the execution state of the agent loop. +// +// NOTE: This is intentionally generic and lives in the vendored fantasy module. +// Budgetsmith-specific persistence and instrumentation should be implemented via +// interfaces/hooks in the Budgetsmith repo. +type ReActState string + +const ( + ReActStateInit ReActState = "init" + ReActStatePrepareStep ReActState = "prepare_step" + ReActStateLLMCall ReActState = "llm_call" + ReActStateToolValidation ReActState = "tool_validation" + ReActStateToolExecution ReActState = "tool_execution" + ReActStateAppendMessages ReActState = "append_messages" + ReActStateStopCheck ReActState = "stop_check" + ReActStateDone ReActState = "done" + ReActStateError ReActState = "error" +) + +// ReActTrigger is the discrete event that causes a state transition. +type ReActTrigger string + +const ( + ReActTriggerStart ReActTrigger = "start" + ReActTriggerPrepared ReActTrigger = "prepared" + ReActTriggerLLMResponded ReActTrigger = "llm_responded" + ReActTriggerToolsValidated ReActTrigger = "tools_validated" + ReActTriggerToolsExecuted ReActTrigger = "tools_executed" + ReActTriggerMessagesAppended ReActTrigger = "messages_appended" + ReActTriggerStopConditionMet ReActTrigger = "stop_condition_met" + ReActTriggerContinue ReActTrigger = "continue" + ReActTriggerFinished ReActTrigger = "finished" + ReActTriggerErrored ReActTrigger = "errored" + ReActTriggerRecoveredContinue ReActTrigger = "recovered_continue" +) + +// ReActTransition is a single transition taken by the agent loop. +type ReActTransition struct { + From ReActState `json:"from"` + To ReActState `json:"to"` + Trigger ReActTrigger `json:"trigger"` + At time.Time `json:"at"` + StepIndex int `json:"step_index"` + Meta map[string]any `json:"meta,omitempty"` + Error string `json:"error,omitempty"` +} + +// ReActTransitionLog is an append-only log of state transitions. +// It is safe for concurrent use. +type ReActTransitionLog struct { + mu sync.Mutex + list []ReActTransition +} + +func NewReActTransitionLog() *ReActTransitionLog { + return &ReActTransitionLog{} +} + +func (l *ReActTransitionLog) Append(t ReActTransition) { + if l == nil { + return + } + l.mu.Lock() + defer l.mu.Unlock() + l.list = append(l.list, t) +} + +func (l *ReActTransitionLog) Snapshot() []ReActTransition { + if l == nil { + return nil + } + l.mu.Lock() + defer l.mu.Unlock() + out := make([]ReActTransition, len(l.list)) + copy(out, l.list) + return out +} diff --git a/internal/fantasy/react_fsm_machine.go b/internal/fantasy/react_fsm_machine.go new file mode 100644 index 000000000..980a2aeac --- /dev/null +++ b/internal/fantasy/react_fsm_machine.go @@ -0,0 +1,121 @@ +package fantasy + +import ( + "context" + "time" + + "github.com/qmuntal/stateless" +) + +// reactFSM is a thin wrapper around a stateless.StateMachine that emits +// transitions to a log and an optional observer. +// +// TODO: For now, this is used by Agent.Generate (Generate-first). Streaming parity is +// implemented later. +type reactFSM struct { + sm *stateless.StateMachine + log *ReActTransitionLog + observer ReActTransitionObserver + + // stepIndex is a pointer to the current step index for the in-flight call. + // It is intentionally owned by the caller; transitions capture its value at + // emission time. + stepIndex *int +} + +func newReActFSM(observer ReActTransitionObserver, stepIndex *int) *reactFSM { + f := &reactFSM{ + sm: stateless.NewStateMachine(ReActStateInit), + log: NewReActTransitionLog(), + observer: observer, + stepIndex: stepIndex, + } + + // Be permissive: instrumentation should not break core behavior. + f.sm.OnUnhandledTrigger(func(context.Context, stateless.State, stateless.Trigger, []string) error { + return nil + }) + + // Emit transitions. + f.sm.OnTransitioned(func(ctx context.Context, tr stateless.Transition) { + t := ReActTransition{ + At: time.Now().UTC(), + } + if from, ok := tr.Source.(ReActState); ok { + t.From = from + } + if to, ok := tr.Destination.(ReActState); ok { + t.To = to + } + if trig, ok := tr.Trigger.(ReActTrigger); ok { + t.Trigger = trig + } + if f.stepIndex != nil { + t.StepIndex = *f.stepIndex + } + + f.log.Append(t) + if f.observer != nil { + f.observer.OnReActTransition(ctx, t) + } + }) + + // State graph for Generate(). + f.configure() + + return f +} + +func (f *reactFSM) configure() { + if f == nil || f.sm == nil { + return + } + + f.sm.Configure(ReActStateInit). + Permit(ReActTriggerStart, ReActStatePrepareStep). + Permit(ReActTriggerErrored, ReActStateError) + + f.sm.Configure(ReActStatePrepareStep). + Permit(ReActTriggerPrepared, ReActStateLLMCall). + Permit(ReActTriggerErrored, ReActStateError) + + f.sm.Configure(ReActStateLLMCall). + Permit(ReActTriggerLLMResponded, ReActStateToolValidation). + Permit(ReActTriggerErrored, ReActStateError) + + f.sm.Configure(ReActStateToolValidation). + Permit(ReActTriggerToolsValidated, ReActStateToolExecution). + Permit(ReActTriggerErrored, ReActStateError) + + f.sm.Configure(ReActStateToolExecution). + Permit(ReActTriggerToolsExecuted, ReActStateAppendMessages). + Permit(ReActTriggerErrored, ReActStateError) + + f.sm.Configure(ReActStateAppendMessages). + Permit(ReActTriggerMessagesAppended, ReActStateStopCheck). + Permit(ReActTriggerErrored, ReActStateError) + + f.sm.Configure(ReActStateStopCheck). + Permit(ReActTriggerStopConditionMet, ReActStateDone). + Permit(ReActTriggerFinished, ReActStateDone). + Permit(ReActTriggerContinue, ReActStatePrepareStep). + Permit(ReActTriggerErrored, ReActStateError) + + f.sm.Configure(ReActStateError). + Permit(ReActTriggerFinished, ReActStateDone). + Permit(ReActTriggerRecoveredContinue, ReActStatePrepareStep) +} + +func (f *reactFSM) Fire(ctx context.Context, trigger ReActTrigger) { + if f == nil || f.sm == nil { + return + } + _ = f.sm.FireCtx(ctx, trigger) +} + +func (f *reactFSM) SnapshotTransitions() []ReActTransition { + if f == nil { + return nil + } + return f.log.Snapshot() +} diff --git a/internal/fantasy/react_fsm_observer.go b/internal/fantasy/react_fsm_observer.go new file mode 100644 index 000000000..e5dad4515 --- /dev/null +++ b/internal/fantasy/react_fsm_observer.go @@ -0,0 +1,19 @@ +package fantasy + +import "context" + +// ReActTransitionObserver can observe transitions taken by the ReAct state machine. +// +// This is a generic hook intended for callers (like Budgetsmith) to persist or +// debug agent execution without coupling the fantasy module to any particular +// storage layer. +type ReActTransitionObserver interface { + OnReActTransition(ctx context.Context, t ReActTransition) +} + +// ReActTransitionObserverFunc is a functional adapter for ReActTransitionObserver. +type ReActTransitionObserverFunc func(ctx context.Context, t ReActTransition) + +func (f ReActTransitionObserverFunc) OnReActTransition(ctx context.Context, t ReActTransition) { + f(ctx, t) +} diff --git a/internal/fantasy/react_step_observer.go b/internal/fantasy/react_step_observer.go new file mode 100644 index 000000000..4190dbda4 --- /dev/null +++ b/internal/fantasy/react_step_observer.go @@ -0,0 +1,18 @@ +package fantasy + +import "context" + +// ReActStepObserver can observe the completion of each agent step in Generate(). +// +// This is primarily intended for callers to persist step snapshots or to build +// debugging/telemetry around multi-step execution. +type ReActStepObserver interface { + OnReActStep(ctx context.Context, stepIndex int, step StepResult) +} + +// ReActStepObserverFunc is a functional adapter for ReActStepObserver. +type ReActStepObserverFunc func(ctx context.Context, stepIndex int, step StepResult) + +func (f ReActStepObserverFunc) OnReActStep(ctx context.Context, stepIndex int, step StepResult) { + f(ctx, stepIndex, step) +} diff --git a/internal/fantasy/react_tool_result_observer.go b/internal/fantasy/react_tool_result_observer.go new file mode 100644 index 000000000..a70bd7898 --- /dev/null +++ b/internal/fantasy/react_tool_result_observer.go @@ -0,0 +1,19 @@ +package fantasy + +import "context" + +// ReActToolResultObserver can observe tool results produced during Generate(). +// +// This hook is intended for callers to persist tool outputs (e.g., filesystem +// offloading + DB indexing) without coupling the fantasy module to a storage +// layer. +type ReActToolResultObserver interface { + OnReActToolResult(ctx context.Context, stepIndex int, result ToolResultContent) +} + +// ReActToolResultObserverFunc is a functional adapter for ReActToolResultObserver. +type ReActToolResultObserverFunc func(ctx context.Context, stepIndex int, result ToolResultContent) + +func (f ReActToolResultObserverFunc) OnReActToolResult(ctx context.Context, stepIndex int, result ToolResultContent) { + f(ctx, stepIndex, result) +} diff --git a/internal/fantasy/tool_dag.go b/internal/fantasy/tool_dag.go new file mode 100644 index 000000000..e548fd9ef --- /dev/null +++ b/internal/fantasy/tool_dag.go @@ -0,0 +1,148 @@ +package fantasy + +import ( + "encoding/json" + "errors" + "fmt" + "strings" +) + +// ToolDAGNode is a node in a tool execution dependency DAG. +type ToolDAGNode struct { + ID string + ToolCall ToolCallContent + Dependencies []string +} + +// ToolDAG is a dependency graph where edges point from dependency -> dependent. +type ToolDAG struct { + Nodes map[string]*ToolDAGNode + Edges map[string][]string +} + +// BuildToolDAG builds a dependency graph for tool calls by scanning JSON inputs +// for $tool. references. +// +// Supported reference forms: +// - "$tool." +// - "$tool.." +func BuildToolDAG(toolCalls []ToolCallContent) (*ToolDAG, error) { + if len(toolCalls) == 0 { + return &ToolDAG{Nodes: map[string]*ToolDAGNode{}, Edges: map[string][]string{}}, nil + } + + nodes := make(map[string]*ToolDAGNode, len(toolCalls)) + for _, tc := range toolCalls { + if strings.TrimSpace(tc.ToolCallID) == "" { + return nil, errors.New("tool call id is empty") + } + if _, exists := nodes[tc.ToolCallID]; exists { + return nil, fmt.Errorf("duplicate tool call id: %s", tc.ToolCallID) + } + nodes[tc.ToolCallID] = &ToolDAGNode{ + ID: tc.ToolCallID, + ToolCall: tc, + Dependencies: nil, + } + } + + edges := make(map[string][]string, len(toolCalls)) + + for _, tc := range toolCalls { + deps, err := extractToolDependencies(tc.Input) + if err != nil { + return nil, fmt.Errorf("parse tool dependencies for %s: %w", tc.ToolCallID, err) + } + + // Filter + validate deps. + filtered := make([]string, 0, len(deps)) + seen := make(map[string]struct{}, len(deps)) + for _, dep := range deps { + dep = strings.TrimSpace(dep) + if dep == "" || dep == tc.ToolCallID { + continue + } + if _, ok := nodes[dep]; !ok { + return nil, fmt.Errorf("tool call %s depends on unknown tool call id %s", tc.ToolCallID, dep) + } + if _, ok := seen[dep]; ok { + continue + } + seen[dep] = struct{}{} + filtered = append(filtered, dep) + } + + nodes[tc.ToolCallID].Dependencies = filtered + for _, dep := range filtered { + edges[dep] = append(edges[dep], tc.ToolCallID) + } + } + + return &ToolDAG{ + Nodes: nodes, + Edges: edges, + }, nil +} + +func extractToolDependencies(input string) ([]string, error) { + input = strings.TrimSpace(input) + if input == "" { + return nil, nil + } + + var v any + if err := json.Unmarshal([]byte(input), &v); err != nil { + // If the tool input isn't valid JSON, treat it as having no dependencies. + return nil, nil + } + + var out []string + walkJSON(v, func(s string) { + for _, id := range extractToolRefIDsFromString(s) { + out = append(out, id) + } + }) + return out, nil +} + +func walkJSON(v any, visitString func(string)) { + switch t := v.(type) { + case map[string]any: + for _, vv := range t { + walkJSON(vv, visitString) + } + case []any: + for _, vv := range t { + walkJSON(vv, visitString) + } + case string: + visitString(t) + default: + // ignore numbers/bools/null + } +} + +func extractToolRefIDsFromString(s string) []string { + // Minimal and strict: only consider strings that start with "$tool.". + // More complex embedding (e.g. "... $tool.x ...") can be added later if needed. + s = strings.TrimSpace(s) + if !strings.HasPrefix(s, "$tool.") { + return nil + } + + rest := strings.TrimPrefix(s, "$tool.") + if rest == "" { + return nil + } + + // ID is the first path segment. + id := rest + if idx := strings.IndexByte(rest, '.'); idx >= 0 { + id = rest[:idx] + } + id = strings.TrimSpace(id) + if id == "" { + return nil + } + return []string{id} +} diff --git a/internal/fantasy/tool_runtime.go b/internal/fantasy/tool_runtime.go new file mode 100644 index 000000000..7b33a4d8c --- /dev/null +++ b/internal/fantasy/tool_runtime.go @@ -0,0 +1,32 @@ +package fantasy + +import "context" + +// ToolRuntime controls how tool calls are executed (sequential, parallel, DAG, etc). +// +// The default behavior is sequential execution identical to the legacy agent +// implementation. +type ToolRuntime interface { + Execute(ctx context.Context, tools []AgentTool, toolCalls []ToolCallContent, toolResultCallback func(result ToolResultContent) error) ([]ToolResultContent, error) +} + +// ToolRuntimeMetrics captures lightweight execution counters. +type ToolRuntimeMetrics struct { + Queued int + InFlightParallel int + BarrierWaits int +} + +// ToolRuntimeLogEvent is an optional structured log emitted by runtimes. +type ToolRuntimeLogEvent struct { + Event string + ToolCallID string + ToolName string + Detail string +} + +// ToolRuntimeMetricsFunc emits metrics. +type ToolRuntimeMetricsFunc = func(ToolRuntimeMetrics) + +// ToolRuntimeLogFunc emits structured log events. +type ToolRuntimeLogFunc = func(ToolRuntimeLogEvent) diff --git a/internal/fantasy/tool_runtime_dag.go b/internal/fantasy/tool_runtime_dag.go new file mode 100644 index 000000000..2eb04a0b5 --- /dev/null +++ b/internal/fantasy/tool_runtime_dag.go @@ -0,0 +1,363 @@ +package fantasy + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "sort" + "strconv" + "strings" + "sync" +) + +// DAGToolRuntime executes tool calls according to an explicit dependency DAG. +// +// Dependencies are discovered via BuildToolDAG (using $tool. references in +// JSON tool inputs). Independent nodes may be executed concurrently, subject to +// MaxConcurrency and tool parallel-safety. +type DAGToolRuntime struct { + MaxConcurrency int + + // Metrics emits optional runtime metrics. + Metrics ToolRuntimeMetricsFunc + + // Log emits structured runtime events. + Log ToolRuntimeLogFunc +} + +func (r DAGToolRuntime) Execute(ctx context.Context, tools []AgentTool, toolCalls []ToolCallContent, toolResultCallback func(result ToolResultContent) error) ([]ToolResultContent, error) { + if len(toolCalls) == 0 { + return nil, nil + } + + metrics := func(m ToolRuntimeMetrics) { + if r.Metrics != nil { + r.Metrics(m) + } + } + logEvent := func(e ToolRuntimeLogEvent) { + if r.Log != nil { + r.Log(e) + } + } + + maxConc := r.MaxConcurrency + if maxConc <= 0 { + maxConc = 4 + } + + dag, err := BuildToolDAG(toolCalls) + if err != nil { + return nil, err + } + + toolMap := make(map[string]AgentTool, len(tools)) + for _, t := range tools { + toolMap[t.Info().Name] = t + } + + idToIndex := make(map[string]int, len(toolCalls)) + for i, tc := range toolCalls { + idToIndex[tc.ToolCallID] = i + } + + indegree := make(map[string]int, len(dag.Nodes)) + for id, n := range dag.Nodes { + indegree[id] = len(n.Dependencies) + } + + executed := make(map[string]bool, len(dag.Nodes)) + doneResults := make(map[string]ToolResultContent, len(dag.Nodes)) + results := make([]ToolResultContent, len(toolCalls)) + + isParallelSafeTool := func(tc ToolCallContent) bool { + if tc.Invalid { + return false + } + t, ok := toolMap[tc.ToolName] + if !ok { + return false + } + return t.Info().Parallel + } + + remaining := len(dag.Nodes) + barrierWaits := 0 + for remaining > 0 { + var readyIDs []string + for id := range dag.Nodes { + if executed[id] { + continue + } + if indegree[id] == 0 { + readyIDs = append(readyIDs, id) + } + } + + if len(readyIDs) == 0 { + return nil, errors.New("tool dependency cycle detected (no ready nodes)") + } + + sort.Slice(readyIDs, func(i, j int) bool { + return idToIndex[readyIDs[i]] < idToIndex[readyIDs[j]] + }) + + metrics(ToolRuntimeMetrics{ + Queued: remaining, + InFlightParallel: 0, + BarrierWaits: barrierWaits, + }) + + // If any ready node is non-parallel-safe, execute the earliest one as a barrier. + barrierID := "" + for _, id := range readyIDs { + n := dag.Nodes[id] + if n == nil { + continue + } + if !isParallelSafeTool(n.ToolCall) { + barrierID = id + break + } + } + + if barrierID != "" { + n := dag.Nodes[barrierID] + barrierWaits++ + logEvent(ToolRuntimeLogEvent{Event: "barrier_start", ToolCallID: n.ToolCall.ToolCallID, ToolName: n.ToolCall.ToolName}) + res, critical, execErr := executeDAGNode(ctx, toolMap, n.ToolCall, doneResults) + if execErr != nil { + return nil, execErr + } + if toolResultCallback != nil { + _ = toolResultCallback(res) + } + logEvent(ToolRuntimeLogEvent{Event: "barrier_finish", ToolCallID: n.ToolCall.ToolCallID, ToolName: n.ToolCall.ToolName}) + if critical { + if errorResult, ok := res.Result.(ToolResultOutputContentError); ok && errorResult.Error != nil { + return nil, errorResult.Error + } + return nil, errors.New("critical tool error") + } + + executed[barrierID] = true + doneResults[barrierID] = res + results[idToIndex[barrierID]] = res + remaining-- + + for _, dep := range dag.Edges[barrierID] { + indegree[dep]-- + } + continue + } + + // Parallel wave: execute up to maxConc ready nodes concurrently. + if len(readyIDs) > maxConc { + readyIDs = readyIDs[:maxConc] + } + + type outcome struct { + id string + res ToolResultContent + critical bool + execErr error + } + + outcomes := make([]outcome, len(readyIDs)) + var wg sync.WaitGroup + wg.Add(len(readyIDs)) + for i, id := range readyIDs { + i, id := i, id + tc := dag.Nodes[id].ToolCall + go func() { + defer wg.Done() + logEvent(ToolRuntimeLogEvent{Event: "dispatch", ToolCallID: tc.ToolCallID, ToolName: tc.ToolName}) + res, critical, execErr := executeDAGNode(ctx, toolMap, tc, doneResults) + outcomes[i] = outcome{ + id: id, + res: res, + critical: critical, + execErr: execErr, + } + logEvent(ToolRuntimeLogEvent{Event: "finish", ToolCallID: tc.ToolCallID, ToolName: tc.ToolName}) + }() + } + wg.Wait() + + // Commit results in deterministic order, and abort on first critical error in that order. + for i := range outcomes { + o := outcomes[i] + if o.execErr != nil { + return nil, o.execErr + } + if toolResultCallback != nil { + _ = toolResultCallback(o.res) + } + if o.critical { + if errorResult, ok := o.res.Result.(ToolResultOutputContentError); ok && errorResult.Error != nil { + return nil, errorResult.Error + } + return nil, errors.New("critical tool error") + } + + executed[o.id] = true + doneResults[o.id] = o.res + results[idToIndex[o.id]] = o.res + remaining-- + + for _, dep := range dag.Edges[o.id] { + indegree[dep]-- + } + } + } + + return results, nil +} + +func executeDAGNode(ctx context.Context, toolMap map[string]AgentTool, toolCall ToolCallContent, prior map[string]ToolResultContent) (ToolResultContent, bool, error) { + resolvedInput, err := resolveToolRefsInInput(toolCall.Input, prior) + if err != nil { + return ToolResultContent{}, false, err + } + tc := toolCall + tc.Input = resolvedInput + res, critical := executeSingleToolCompat(ctx, toolMap, tc, nil) + return res, critical, nil +} + +func resolveToolRefsInInput(input string, results map[string]ToolResultContent) (string, error) { + input = strings.TrimSpace(input) + if input == "" { + return input, nil + } + + var v any + if err := json.Unmarshal([]byte(input), &v); err != nil { + // Not JSON; nothing to resolve. + return input, nil + } + + updated, err := rewriteJSONToolRefs(v, results) + if err != nil { + return "", err + } + b, err := json.Marshal(updated) + if err != nil { + return "", err + } + return string(b), nil +} + +func rewriteJSONToolRefs(v any, results map[string]ToolResultContent) (any, error) { + switch t := v.(type) { + case map[string]any: + out := make(map[string]any, len(t)) + for k, vv := range t { + rv, err := rewriteJSONToolRefs(vv, results) + if err != nil { + return nil, err + } + out[k] = rv + } + return out, nil + case []any: + out := make([]any, len(t)) + for i, vv := range t { + rv, err := rewriteJSONToolRefs(vv, results) + if err != nil { + return nil, err + } + out[i] = rv + } + return out, nil + case string: + s := strings.TrimSpace(t) + if !strings.HasPrefix(s, "$tool.") { + return t, nil + } + val, err := resolveToolRefValue(s, results) + if err != nil { + return nil, err + } + return val, nil + default: + return v, nil + } +} + +func resolveToolRefValue(ref string, results map[string]ToolResultContent) (any, error) { + rest := strings.TrimPrefix(strings.TrimSpace(ref), "$tool.") + if rest == "" { + return nil, fmt.Errorf("invalid tool ref: %q", ref) + } + + parts := strings.Split(rest, ".") + depID := strings.TrimSpace(parts[0]) + if depID == "" { + return nil, fmt.Errorf("invalid tool ref: %q", ref) + } + + r, ok := results[depID] + if !ok { + return nil, fmt.Errorf("tool ref %q not available yet", ref) + } + + base := "" + switch v := r.Result.(type) { + case ToolResultOutputContentText: + base = v.Text + case ToolResultOutputContentMedia: + if v.Text != "" { + base = v.Text + } else { + base = v.Data + } + case ToolResultOutputContentError: + if v.Error != nil { + base = v.Error.Error() + } + default: + base = fmt.Sprint(r.Result) + } + + // No path: substitute as string. + if len(parts) == 1 { + return base, nil + } + + // Path resolution: interpret base as JSON and walk. + var cur any + if err := json.Unmarshal([]byte(base), &cur); err != nil { + return nil, fmt.Errorf("tool ref %q path requires JSON output, got non-JSON", ref) + } + + for _, seg := range parts[1:] { + seg = strings.TrimSpace(seg) + if seg == "" { + return nil, fmt.Errorf("invalid tool ref path in %q", ref) + } + + switch typed := cur.(type) { + case map[string]any: + nv, ok := typed[seg] + if !ok { + return nil, fmt.Errorf("tool ref %q missing key %q", ref, seg) + } + cur = nv + case []any: + idx, err := strconv.Atoi(seg) + if err != nil { + return nil, fmt.Errorf("tool ref %q array segment %q is not an int", ref, seg) + } + if idx < 0 || idx >= len(typed) { + return nil, fmt.Errorf("tool ref %q array index %d out of range", ref, idx) + } + cur = typed[idx] + default: + return nil, fmt.Errorf("tool ref %q path segment %q on non-container", ref, seg) + } + } + + return cur, nil +} diff --git a/internal/fantasy/tool_runtime_dag_test.go b/internal/fantasy/tool_runtime_dag_test.go new file mode 100644 index 000000000..ffbc17168 --- /dev/null +++ b/internal/fantasy/tool_runtime_dag_test.go @@ -0,0 +1,214 @@ +package fantasy + +import ( + "context" + "fmt" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +func TestDAGToolRuntime_IndependentToolsRunConcurrently(t *testing.T) { + t.Parallel() + + started := make(chan string, 2) + release := make(chan struct{}) + + tool := NewParallelAgentTool("p", "parallel tool", func(ctx context.Context, _ struct{}, call ToolCall) (ToolResponse, error) { + select { + case started <- call.ID: + default: + } + select { + case <-ctx.Done(): + return NewTextErrorResponse(ctx.Err().Error()), nil + case <-release: + } + return NewTextResponse(call.ID), nil + }) + + rt := DAGToolRuntime{MaxConcurrency: 2} + + toolCalls := []ToolCallContent{ + {ToolCallID: "a", ToolName: "p", Input: `{}`}, + {ToolCallID: "b", ToolName: "p", Input: `{}`}, + } + + var ( + res []ToolResultContent + err error + wg sync.WaitGroup + ) + wg.Add(1) + go func() { + defer wg.Done() + res, err = rt.Execute(context.Background(), []AgentTool{tool}, toolCalls, nil) + }() + + // Both tools should start before we release. + got1 := <-started + got2 := <-started + require.NotEqual(t, got1, got2) + + close(release) + wg.Wait() + + require.NoError(t, err) + require.Len(t, res, 2) +} + +func TestDAGToolRuntime_DependenciesWaitAndInputIsResolved(t *testing.T) { + t.Parallel() + + startA := make(chan struct{}, 1) + startB := make(chan struct{}, 1) + releaseA := make(chan struct{}) + + toolA := NewParallelAgentTool("a", "tool a", func(ctx context.Context, _ struct{}, _ ToolCall) (ToolResponse, error) { + startA <- struct{}{} + select { + case <-ctx.Done(): + return NewTextErrorResponse(ctx.Err().Error()), nil + case <-releaseA: + } + // JSON output that downstream can path into. + return NewTextResponse(`{"x":1}`), nil + }) + + type bInput struct { + Val int `json:"val"` + } + toolB := NewParallelAgentTool("b", "tool b", func(_ context.Context, in bInput, _ ToolCall) (ToolResponse, error) { + startB <- struct{}{} + return NewTextResponse(string(rune('0' + in.Val))), nil + }) + + rt := DAGToolRuntime{MaxConcurrency: 4} + toolCalls := []ToolCallContent{ + {ToolCallID: "callA", ToolName: "a", Input: `{}`}, + {ToolCallID: "callB", ToolName: "b", Input: `{"val":"$tool.callA.x"}`}, + } + + var ( + res []ToolResultContent + err error + wg sync.WaitGroup + ) + wg.Add(1) + go func() { + defer wg.Done() + res, err = rt.Execute(context.Background(), []AgentTool{toolA, toolB}, toolCalls, nil) + }() + + <-startA + + // B must not start until A is released. + select { + case <-startB: + t.Fatalf("dependent tool started before dependency completed") + case <-time.After(30 * time.Millisecond): + } + + close(releaseA) + + <-startB + wg.Wait() + + require.NoError(t, err) + require.Len(t, res, 2) + + // B should have received val=1 and returned "1". + require.Equal(t, "callB", res[1].ToolCallID) + require.Equal(t, "1", res[1].Result.(ToolResultOutputContentText).Text) +} + +func TestDAGToolRuntime_CycleDetected(t *testing.T) { + t.Parallel() + + tool := NewParallelAgentTool("p", "tool", func(_ context.Context, _ struct{}, _ ToolCall) (ToolResponse, error) { + return NewTextResponse("ok"), nil + }) + + rt := DAGToolRuntime{MaxConcurrency: 4} + toolCalls := []ToolCallContent{ + {ToolCallID: "a", ToolName: "p", Input: `{"x":"$tool.b"}`}, + {ToolCallID: "b", ToolName: "p", Input: `{"x":"$tool.a"}`}, + } + + res, err := rt.Execute(context.Background(), []AgentTool{tool}, toolCalls, nil) + require.Error(t, err) + require.Nil(t, res) +} + +func TestDAGToolRuntime_OnToolResultSerialized(t *testing.T) { + t.Parallel() + + var inFlight atomic.Int32 + var orderMu sync.Mutex + var order []string + + tool := NewParallelAgentTool("p", "tool", func(_ context.Context, _ struct{}, call ToolCall) (ToolResponse, error) { + return NewTextResponse(call.ID), nil + }) + + toolCalls := []ToolCallContent{ + {ToolCallID: "a", ToolName: "p", Input: `{}`}, + {ToolCallID: "b", ToolName: "p", Input: `{}`}, + } + + rt := DAGToolRuntime{MaxConcurrency: 2} + cb := func(res ToolResultContent) error { + if inFlight.Add(1) != 1 { + return fmt.Errorf("callback executed concurrently") + } + orderMu.Lock() + order = append(order, res.ToolCallID) + orderMu.Unlock() + time.Sleep(5 * time.Millisecond) + inFlight.Add(-1) + return nil + } + + res, err := rt.Execute(context.Background(), []AgentTool{tool}, toolCalls, cb) + require.NoError(t, err) + require.Len(t, res, 2) + + orderMu.Lock() + defer orderMu.Unlock() + require.Equal(t, []string{"a", "b"}, order) +} + +func TestDAGToolRuntime_MetricsAndLogHooks(t *testing.T) { + t.Parallel() + + var metricsCalled bool + var logCalled bool + + tool := NewParallelAgentTool("p", "tool", func(_ context.Context, _ struct{}, call ToolCall) (ToolResponse, error) { + return NewTextResponse(call.ID), nil + }) + + rt := DAGToolRuntime{ + MaxConcurrency: 1, + Metrics: func(m ToolRuntimeMetrics) { + metricsCalled = true + require.GreaterOrEqual(t, m.Queued, 0) + }, + Log: func(e ToolRuntimeLogEvent) { + logCalled = true + require.NotEmpty(t, e.Event) + }, + } + + toolCalls := []ToolCallContent{ + {ToolCallID: "a", ToolName: "p", Input: `{}`}, + } + res, err := rt.Execute(context.Background(), []AgentTool{tool}, toolCalls, nil) + require.NoError(t, err) + require.Len(t, res, 1) + require.True(t, metricsCalled) + require.True(t, logCalled) +} diff --git a/internal/fantasy/tool_runtime_parallel.go b/internal/fantasy/tool_runtime_parallel.go new file mode 100644 index 000000000..711a3eafe --- /dev/null +++ b/internal/fantasy/tool_runtime_parallel.go @@ -0,0 +1,214 @@ +package fantasy + +import ( + "context" + "errors" + "sync" +) + +// ParallelToolRuntime executes tool calls concurrently when tools opt-in via +// ToolInfo.Parallel. +// +// Ordering is preserved: results are returned in the same order as toolCalls. +// Non-parallel-safe tools act as barriers and run sequentially. +type ParallelToolRuntime struct { + // MaxConcurrency limits concurrent tool execution within a parallel batch. + // If <= 0, a safe default is used. + MaxConcurrency int + + // Metrics emits optional runtime metrics. + Metrics ToolRuntimeMetricsFunc + + // Log emits structured runtime events. + Log ToolRuntimeLogFunc +} + +func (r ParallelToolRuntime) Execute(ctx context.Context, tools []AgentTool, toolCalls []ToolCallContent, toolResultCallback func(result ToolResultContent) error) ([]ToolResultContent, error) { + if len(toolCalls) == 0 { + return nil, nil + } + + metrics := func(m ToolRuntimeMetrics) { + if r.Metrics != nil { + r.Metrics(m) + } + } + logEvent := func(e ToolRuntimeLogEvent) { + if r.Log != nil { + r.Log(e) + } + } + + maxConc := r.MaxConcurrency + if maxConc <= 0 { + maxConc = 4 + } + + // Quick tool lookup. + toolMap := make(map[string]AgentTool, len(tools)) + for _, t := range tools { + toolMap[t.Info().Name] = t + } + + results := make([]ToolResultContent, len(toolCalls)) + + isParallelSafe := func(tc ToolCallContent) bool { + if tc.Invalid { + return false + } + t, ok := toolMap[tc.ToolName] + if !ok { + return false + } + return t.Info().Parallel + } + + sem := make(chan struct{}, maxConc) + inFlight := 0 + barrierWaits := 0 + i := 0 + emit := func() { + metrics(ToolRuntimeMetrics{Queued: len(toolCalls) - i, InFlightParallel: inFlight, BarrierWaits: barrierWaits}) + } + for i < len(toolCalls) { + if !isParallelSafe(toolCalls[i]) { + barrierWaits++ + emit() + logEvent(ToolRuntimeLogEvent{Event: "barrier_start", ToolCallID: toolCalls[i].ToolCallID, ToolName: toolCalls[i].ToolName}) + res, critical := executeSingleToolCompat(ctx, toolMap, toolCalls[i], toolResultCallback) + logEvent(ToolRuntimeLogEvent{Event: "barrier_finish", ToolCallID: toolCalls[i].ToolCallID, ToolName: toolCalls[i].ToolName}) + results[i] = res + if critical { + if errorResult, ok := res.Result.(ToolResultOutputContentError); ok && errorResult.Error != nil { + return nil, errorResult.Error + } + return nil, errors.New("critical tool error") + } + i++ + continue + } + + // Collect a contiguous batch of parallel-safe tool calls. + start := i + for i < len(toolCalls) && isParallelSafe(toolCalls[i]) { + i++ + } + end := i + + type outcome struct { + res ToolResultContent + critical bool + } + outcomes := make([]outcome, end-start) + + var wg sync.WaitGroup + for bi := start; bi < end; bi++ { + localIndex := bi - start + tc := toolCalls[bi] + wg.Add(1) + go func() { + defer wg.Done() + logEvent(ToolRuntimeLogEvent{Event: "dispatch", ToolCallID: tc.ToolCallID, ToolName: tc.ToolName}) + sem <- struct{}{} + inFlight++ + emit() + defer func() { + <-sem + inFlight-- + emit() + }() + + res, critical := executeSingleToolCompat(ctx, toolMap, tc, nil) + outcomes[localIndex] = outcome{res: res, critical: critical} + logEvent(ToolRuntimeLogEvent{Event: "finish", ToolCallID: tc.ToolCallID, ToolName: tc.ToolName}) + }() + } + wg.Wait() + + // Emit callback and copy results in deterministic order. + for bi := start; bi < end; bi++ { + o := outcomes[bi-start] + results[bi] = o.res + if toolResultCallback != nil { + _ = toolResultCallback(o.res) + } + if o.critical { + if errorResult, ok := o.res.Result.(ToolResultOutputContentError); ok && errorResult.Error != nil { + return nil, errorResult.Error + } + return nil, errors.New("critical tool error") + } + } + } + + return results, nil +} + +// executeSingleToolCompat mirrors the legacy sequential agent tool execution +// semantics, but is packaged as a helper so tool runtimes can share it. +func executeSingleToolCompat(ctx context.Context, toolMap map[string]AgentTool, toolCall ToolCallContent, toolResultCallback func(result ToolResultContent) error) (ToolResultContent, bool) { + result := ToolResultContent{ + ToolCallID: toolCall.ToolCallID, + ToolName: toolCall.ToolName, + ProviderExecuted: false, + } + + // Skip invalid tool calls - create error result (not critical). + if toolCall.Invalid { + result.Result = ToolResultOutputContentError{ + Error: toolCall.ValidationError, + } + if toolResultCallback != nil { + _ = toolResultCallback(result) + } + return result, false + } + + tool, exists := toolMap[toolCall.ToolName] + if !exists { + result.Result = ToolResultOutputContentError{ + Error: errors.New("Error: Tool not found: " + toolCall.ToolName), + } + if toolResultCallback != nil { + _ = toolResultCallback(result) + } + return result, false + } + + toolResult, err := tool.Run(ctx, ToolCall{ + ID: toolCall.ToolCallID, + Name: toolCall.ToolName, + Input: toolCall.Input, + }) + if err != nil { + result.Result = ToolResultOutputContentError{ + Error: err, + } + result.ClientMetadata = toolResult.Metadata + if toolResultCallback != nil { + _ = toolResultCallback(result) + } + return result, true + } + + result.ClientMetadata = toolResult.Metadata + if toolResult.IsError { + result.Result = ToolResultOutputContentError{ + Error: errors.New(toolResult.Content), + } + } else if toolResult.Type == "image" || toolResult.Type == "media" { + result.Result = ToolResultOutputContentMedia{ + Data: string(toolResult.Data), + MediaType: toolResult.MediaType, + Text: toolResult.Content, + } + } else { + result.Result = ToolResultOutputContentText{ + Text: toolResult.Content, + } + } + if toolResultCallback != nil { + _ = toolResultCallback(result) + } + return result, false +} diff --git a/internal/fantasy/tool_runtime_parallel_test.go b/internal/fantasy/tool_runtime_parallel_test.go new file mode 100644 index 000000000..977a0ba85 --- /dev/null +++ b/internal/fantasy/tool_runtime_parallel_test.go @@ -0,0 +1,156 @@ +package fantasy + +import ( + "context" + "errors" + "sync" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +func TestParallelToolRuntime_OrderAndCallbackDeterminism(t *testing.T) { + t.Parallel() + + type input struct { + DelayMs int `json:"delay_ms"` + Value string `json:"value"` + } + + tool := NewParallelAgentTool("p", "parallel tool", func(ctx context.Context, in input, _ ToolCall) (ToolResponse, error) { + if in.DelayMs > 0 { + select { + case <-ctx.Done(): + return NewTextErrorResponse(ctx.Err().Error()), nil + case <-time.After(time.Duration(in.DelayMs) * time.Millisecond): + } + } + return NewTextResponse(in.Value), nil + }) + + runtime := ParallelToolRuntime{MaxConcurrency: 3} + + var cbMu sync.Mutex + var cbOrder []string + cb := func(res ToolResultContent) error { + cbMu.Lock() + defer cbMu.Unlock() + cbOrder = append(cbOrder, res.ToolCallID) + return nil + } + + toolCalls := []ToolCallContent{ + {ToolCallID: "c1", ToolName: "p", Input: `{"delay_ms":50,"value":"a"}`}, + {ToolCallID: "c2", ToolName: "p", Input: `{"delay_ms":10,"value":"b"}`}, + {ToolCallID: "c3", ToolName: "p", Input: `{"delay_ms":30,"value":"c"}`}, + } + + results, err := runtime.Execute(context.Background(), []AgentTool{tool}, toolCalls, cb) + require.NoError(t, err) + require.Len(t, results, 3) + + require.Equal(t, "c1", results[0].ToolCallID) + require.Equal(t, "a", results[0].Result.(ToolResultOutputContentText).Text) + require.Equal(t, "c2", results[1].ToolCallID) + require.Equal(t, "b", results[1].Result.(ToolResultOutputContentText).Text) + require.Equal(t, "c3", results[2].ToolCallID) + require.Equal(t, "c", results[2].Result.(ToolResultOutputContentText).Text) + + cbMu.Lock() + require.Equal(t, []string{"c1", "c2", "c3"}, cbOrder) + cbMu.Unlock() +} + +func TestParallelToolRuntime_BarrierForNonParallelTools(t *testing.T) { + t.Parallel() + + parallel := NewParallelAgentTool("p", "parallel tool", func(_ context.Context, _ struct{}, call ToolCall) (ToolResponse, error) { + _ = call + return NewTextResponse("p"), nil + }) + seq := NewAgentTool("s", "sequential tool", func(_ context.Context, _ struct{}, call ToolCall) (ToolResponse, error) { + _ = call + return NewTextResponse("s"), nil + }) + + runtime := ParallelToolRuntime{MaxConcurrency: 8} + + var mu sync.Mutex + var order []string + cb := func(res ToolResultContent) error { + mu.Lock() + defer mu.Unlock() + order = append(order, res.ToolCallID) + return nil + } + + toolCalls := []ToolCallContent{ + {ToolCallID: "p1", ToolName: "p", Input: `{}`}, + {ToolCallID: "p2", ToolName: "p", Input: `{}`}, + {ToolCallID: "s1", ToolName: "s", Input: `{}`}, + {ToolCallID: "p3", ToolName: "p", Input: `{}`}, + } + + results, err := runtime.Execute(context.Background(), []AgentTool{parallel, seq}, toolCalls, cb) + require.NoError(t, err) + require.Len(t, results, 4) + + mu.Lock() + require.Equal(t, []string{"p1", "p2", "s1", "p3"}, order) + mu.Unlock() +} + +func TestParallelToolRuntime_CriticalErrorPropagation(t *testing.T) { + t.Parallel() + + tool := NewParallelAgentTool("p", "parallel tool", func(_ context.Context, _ struct{}, call ToolCall) (ToolResponse, error) { + if call.ID == "bad" { + return ToolResponse{}, errors.New("boom") + } + return NewTextResponse("ok"), nil + }) + + runtime := ParallelToolRuntime{MaxConcurrency: 4} + + toolCalls := []ToolCallContent{ + {ToolCallID: "good", ToolName: "p", Input: `{}`}, + {ToolCallID: "bad", ToolName: "p", Input: `{}`}, + } + + results, err := runtime.Execute(context.Background(), []AgentTool{tool}, toolCalls, nil) + require.Error(t, err) + require.Nil(t, results) +} + +func TestParallelToolRuntime_MetricsAndLogHooks(t *testing.T) { + t.Parallel() + + var metricsCalled bool + var logCalled bool + + tool := NewParallelAgentTool("p", "tool", func(_ context.Context, _ struct{}, call ToolCall) (ToolResponse, error) { + return NewTextResponse(call.ID), nil + }) + + rt := ParallelToolRuntime{ + MaxConcurrency: 2, + Metrics: func(m ToolRuntimeMetrics) { + metricsCalled = true + require.GreaterOrEqual(t, m.Queued, 0) + }, + Log: func(e ToolRuntimeLogEvent) { + logCalled = true + require.NotEmpty(t, e.Event) + }, + } + + toolCalls := []ToolCallContent{ + {ToolCallID: "a", ToolName: "p", Input: `{}`}, + } + res, err := rt.Execute(context.Background(), []AgentTool{tool}, toolCalls, nil) + require.NoError(t, err) + require.Len(t, res, 1) + require.True(t, metricsCalled) + require.True(t, logCalled) +}