feat(fantasy): add ReAct FSM, tool DAG, and parallel tool runtime
ReAct FSM: - react_fsm.go — ReActState enum + FSM state struct (step index, error, start/end times, last tool calls) - react_fsm_machine.go — FSM transition table with valid state edges, Transition() helper that validates and records transitions - react_fsm_observer.go — ObserverFSM interface for FSM lifecycle hooks (OnStateEnter, OnStateExit, OnTransition, OnError) - react_step_observer.go — StepObserver: per-step hook triggered before and after each ReAct iteration - react_tool_result_observer.go — ToolResultObserver: hook for capturing structured tool results at each execution step Tool DAG: - tool_dag.go — ToolDAGNode + ToolDAG; dependency resolution via topological sort; cycle detection; parallel batch extraction - tool_runtime.go — ToolRuntime interface + ExecToolResult value type Parallel runtime: - tool_runtime_parallel.go — ParallelToolRuntime: concurrent dispatch of tool calls respecting ToolInfo.Parallel flag; MaxConcurrency cap; barrier semantics for sequential tools; result ordering preserved - tool_runtime_dag.go — DAGToolRuntime: executes a ToolDAG by resolving dependency batches and dispatching each batch via ParallelToolRuntime Agent: - agent.go — wire FSM + runtime observers into agent loop init path; expose RunWithFSM entry point
This commit is contained in:
ZanzyTHEbar
parent
75c9e84309
commit
b240dfdaa2
12 changed files with 1504 additions and 53 deletions
|
|
@ -140,10 +140,14 @@ type agentSettings struct {
|
|||
headers map[string]string
|
||||
providerOptions ProviderOptions
|
||||
|
||||
// TODO: add support for provider tools
|
||||
tools []AgentTool
|
||||
toolRuntime ToolRuntime
|
||||
maxRetries *int
|
||||
|
||||
transitionObservers []ReActTransitionObserver
|
||||
stepObservers []ReActStepObserver
|
||||
toolResultObservers []ReActToolResultObserver
|
||||
|
||||
model LanguageModel
|
||||
|
||||
stopWhen []StopCondition
|
||||
|
|
@ -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,6 +1183,9 @@ func (a *agent) processStepStream(ctx context.Context, stream StreamResponse, op
|
|||
toolMap[tool.Info().Name] = tool
|
||||
}
|
||||
|
||||
if !useRuntimeBatch {
|
||||
toolChan = make(chan toolExecutionRequest, 10)
|
||||
|
||||
// Semaphores for controlling parallelism
|
||||
parallelSem := make(chan struct{}, 5)
|
||||
var sequentialMu sync.Mutex
|
||||
|
|
@ -1165,6 +1222,7 @@ func (a *agent) processStepStream(ctx context.Context, stream StreamResponse, op
|
|||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// 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
|
||||
if !useRuntimeBatch {
|
||||
isParallel := false
|
||||
if tool, exists := toolMap[validatedToolCall.ToolName]; exists {
|
||||
isParallel = tool.Info().Parallel
|
||||
}
|
||||
|
||||
// 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
|
||||
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()
|
||||
|
||||
// Check for tool execution errors
|
||||
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)
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
84
internal/fantasy/react_fsm.go
Normal file
84
internal/fantasy/react_fsm.go
Normal file
|
|
@ -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
|
||||
}
|
||||
121
internal/fantasy/react_fsm_machine.go
Normal file
121
internal/fantasy/react_fsm_machine.go
Normal file
|
|
@ -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()
|
||||
}
|
||||
19
internal/fantasy/react_fsm_observer.go
Normal file
19
internal/fantasy/react_fsm_observer.go
Normal file
|
|
@ -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)
|
||||
}
|
||||
18
internal/fantasy/react_step_observer.go
Normal file
18
internal/fantasy/react_step_observer.go
Normal file
|
|
@ -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)
|
||||
}
|
||||
19
internal/fantasy/react_tool_result_observer.go
Normal file
19
internal/fantasy/react_tool_result_observer.go
Normal file
|
|
@ -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)
|
||||
}
|
||||
148
internal/fantasy/tool_dag.go
Normal file
148
internal/fantasy/tool_dag.go
Normal file
|
|
@ -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.<toolCallID> references.
|
||||
//
|
||||
// Supported reference forms:
|
||||
// - "$tool.<id>"
|
||||
// - "$tool.<id>.<jsonPath...>"
|
||||
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}
|
||||
}
|
||||
32
internal/fantasy/tool_runtime.go
Normal file
32
internal/fantasy/tool_runtime.go
Normal file
|
|
@ -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)
|
||||
363
internal/fantasy/tool_runtime_dag.go
Normal file
363
internal/fantasy/tool_runtime_dag.go
Normal file
|
|
@ -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.<id> 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
|
||||
}
|
||||
214
internal/fantasy/tool_runtime_dag_test.go
Normal file
214
internal/fantasy/tool_runtime_dag_test.go
Normal file
|
|
@ -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)
|
||||
}
|
||||
214
internal/fantasy/tool_runtime_parallel.go
Normal file
214
internal/fantasy/tool_runtime_parallel.go
Normal file
|
|
@ -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
|
||||
}
|
||||
156
internal/fantasy/tool_runtime_parallel_test.go
Normal file
156
internal/fantasy/tool_runtime_parallel_test.go
Normal file
|
|
@ -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)
|
||||
}
|
||||
Loading…
Add table
Reference in a new issue