wehub-resource-sync
2150 lines
63 KiB
Go
2150 lines
63 KiB
Go
/*
|
|
* Copyright 2025 CloudWeGo Authors
|
|
*
|
|
* Licensed under the Apache License, Version 2.0 (the "License");
|
|
* you may not use this file except in compliance with the License.
|
|
* You may obtain a copy of the License at
|
|
*
|
|
* http://www.apache.org/licenses/LICENSE-2.0
|
|
*
|
|
* Unless required by applicable law or agreed to in writing, software
|
|
* distributed under the License is distributed on an "AS IS" BASIS,
|
|
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
* See the License for the specific language governing permissions and
|
|
* limitations under the License.
|
|
*/
|
|
|
|
package adk
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"sync"
|
|
"sync/atomic"
|
|
"testing"
|
|
|
|
"github.com/stretchr/testify/assert"
|
|
|
|
"github.com/cloudwego/eino/components/model"
|
|
"github.com/cloudwego/eino/components/tool"
|
|
"github.com/cloudwego/eino/compose"
|
|
"github.com/cloudwego/eino/schema"
|
|
)
|
|
|
|
type interruptTestToolsHandler struct {
|
|
*BaseChatModelAgentMiddleware
|
|
tools []tool.BaseTool
|
|
}
|
|
|
|
func TestPreprocessADKCheckpoint(t *testing.T) {
|
|
t.Run("no-op when missing markers", func(t *testing.T) {
|
|
in := []byte("random")
|
|
out := preprocessADKCheckpoint(append([]byte(nil), in...))
|
|
assert.Equal(t, in, out)
|
|
})
|
|
|
|
t.Run("rewrite legacy name for v0.8.0-v0.8.3", func(t *testing.T) {
|
|
const (
|
|
lenPrefixedReactStateName = "\x15" + stateGobNameV07
|
|
lenPrefixedCompatName = "\x15" + stateGobNameV080
|
|
lenPrefixedStateSerializationName = "\x12stateSerialization"
|
|
)
|
|
|
|
in := []byte(lenPrefixedReactStateName + "xxx" + lenPrefixedStateSerializationName + "yyy")
|
|
out := preprocessADKCheckpoint(append([]byte(nil), in...))
|
|
assert.True(t, bytes.Contains(out, []byte(lenPrefixedCompatName)))
|
|
assert.False(t, bytes.Contains(out, []byte(lenPrefixedReactStateName)))
|
|
})
|
|
}
|
|
|
|
func (h *interruptTestToolsHandler) BeforeAgent(ctx context.Context, runCtx *ChatModelAgentContext) (context.Context, *ChatModelAgentContext, error) {
|
|
runCtx.Tools = append(runCtx.Tools, h.tools...)
|
|
return ctx, runCtx, nil
|
|
}
|
|
|
|
func TestSaveAgentEventWrapper(t *testing.T) {
|
|
sr, sw := schema.Pipe[Message](1)
|
|
sw.Send(schema.UserMessage("test"), nil)
|
|
sw.Close()
|
|
sr = sr.Copy(2)[1]
|
|
|
|
w := &agentEventWrapper{
|
|
AgentEvent: &AgentEvent{
|
|
Output: &AgentOutput{
|
|
MessageOutput: &MessageVariant{
|
|
IsStreaming: true,
|
|
MessageStream: sr,
|
|
},
|
|
},
|
|
RunPath: []RunStep{
|
|
{
|
|
"a1",
|
|
},
|
|
{
|
|
"a2",
|
|
},
|
|
},
|
|
},
|
|
mu: sync.Mutex{},
|
|
concatenatedMessage: nil,
|
|
}
|
|
|
|
_, err := getMessageFromWrappedEvent(w)
|
|
assert.NoError(t, err)
|
|
|
|
buf, err := w.GobEncode()
|
|
assert.NoError(t, err)
|
|
assert.NoError(t, err)
|
|
|
|
w1 := &agentEventWrapper{}
|
|
err = w1.GobDecode(buf)
|
|
assert.NoError(t, err)
|
|
}
|
|
|
|
func TestInterruptFunctionsPopulateInterruptContextsImmediately(t *testing.T) {
|
|
ctx := context.Background()
|
|
ctx, _ = initRunCtx(ctx, "TestAgent", &AgentInput{Messages: []Message{}})
|
|
ctx = AppendAddressSegment(ctx, AddressSegmentAgent, "TestAgent")
|
|
|
|
t.Run("Interrupt populates InterruptContexts", func(t *testing.T) {
|
|
event := Interrupt(ctx, "test info")
|
|
assert.NotNil(t, event.Action)
|
|
assert.NotNil(t, event.Action.Interrupted)
|
|
assert.NotNil(t, event.Action.Interrupted.InterruptContexts)
|
|
assert.Equal(t, 1, len(event.Action.Interrupted.InterruptContexts))
|
|
assert.Equal(t, "test info", event.Action.Interrupted.InterruptContexts[0].Info)
|
|
assert.True(t, event.Action.Interrupted.InterruptContexts[0].IsRootCause)
|
|
assert.Equal(t, Address{
|
|
{Type: AddressSegmentAgent, ID: "TestAgent"},
|
|
}, event.Action.Interrupted.InterruptContexts[0].Address)
|
|
})
|
|
|
|
t.Run("StatefulInterrupt populates InterruptContexts", func(t *testing.T) {
|
|
event := StatefulInterrupt(ctx, "stateful info", "my state")
|
|
assert.NotNil(t, event.Action)
|
|
assert.NotNil(t, event.Action.Interrupted)
|
|
assert.NotNil(t, event.Action.Interrupted.InterruptContexts)
|
|
assert.Equal(t, 1, len(event.Action.Interrupted.InterruptContexts))
|
|
assert.Equal(t, "stateful info", event.Action.Interrupted.InterruptContexts[0].Info)
|
|
assert.True(t, event.Action.Interrupted.InterruptContexts[0].IsRootCause)
|
|
})
|
|
|
|
t.Run("CompositeInterrupt populates InterruptContexts with filtered parent chain", func(t *testing.T) {
|
|
subCtx := AppendAddressSegment(ctx, AddressSegmentAgent, "SubAgent")
|
|
subEvent := Interrupt(subCtx, "sub info")
|
|
event := CompositeInterrupt(ctx, "composite info", "composite state", subEvent.Action.internalInterrupted)
|
|
assert.NotNil(t, event.Action)
|
|
assert.NotNil(t, event.Action.Interrupted)
|
|
assert.NotNil(t, event.Action.Interrupted.InterruptContexts)
|
|
assert.Equal(t, 1, len(event.Action.Interrupted.InterruptContexts))
|
|
|
|
rootCause := event.Action.Interrupted.InterruptContexts[0]
|
|
assert.Equal(t, "sub info", rootCause.Info)
|
|
assert.True(t, rootCause.IsRootCause)
|
|
assert.Equal(t, Address{
|
|
{Type: AddressSegmentAgent, ID: "TestAgent"},
|
|
{Type: AddressSegmentAgent, ID: "SubAgent"},
|
|
}, rootCause.Address)
|
|
|
|
assert.NotNil(t, rootCause.Parent, "Parent should not be nil for composite interrupt")
|
|
assert.Equal(t, "composite info", rootCause.Parent.Info)
|
|
assert.Equal(t, Address{
|
|
{Type: AddressSegmentAgent, ID: "TestAgent"},
|
|
}, rootCause.Parent.Address)
|
|
})
|
|
|
|
t.Run("Address only contains agent/tool segments", func(t *testing.T) {
|
|
event := Interrupt(ctx, "test info")
|
|
addr := event.Action.Interrupted.InterruptContexts[0].Address
|
|
for _, seg := range addr {
|
|
assert.True(t, seg.Type == AddressSegmentAgent || seg.Type == AddressSegmentTool,
|
|
"Address should only contain agent/tool segments, got: %s", seg.Type)
|
|
}
|
|
})
|
|
}
|
|
|
|
func TestSimpleInterrupt(t *testing.T) {
|
|
data := "hello world"
|
|
agent := &myAgent{
|
|
runFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {
|
|
iter, generator := NewAsyncIteratorPair[*AgentEvent]()
|
|
generator.Send(&AgentEvent{
|
|
Output: &AgentOutput{
|
|
MessageOutput: &MessageVariant{
|
|
IsStreaming: true,
|
|
Message: nil,
|
|
MessageStream: schema.StreamReaderFromArray([]Message{
|
|
schema.UserMessage("hello "),
|
|
schema.UserMessage("world"),
|
|
}),
|
|
},
|
|
},
|
|
})
|
|
intEvent := Interrupt(ctx, data)
|
|
intEvent.Action.Interrupted.Data = data
|
|
generator.Send(intEvent)
|
|
generator.Close()
|
|
return iter
|
|
},
|
|
resumeFn: func(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] {
|
|
assert.True(t, info.WasInterrupted)
|
|
assert.Nil(t, info.InterruptState)
|
|
assert.True(t, info.EnableStreaming)
|
|
assert.Equal(t, data, info.Data)
|
|
|
|
assert.True(t, info.IsResumeTarget)
|
|
iter, generator := NewAsyncIteratorPair[*AgentEvent]()
|
|
generator.Close()
|
|
return iter
|
|
},
|
|
}
|
|
store := newMyStore()
|
|
ctx := context.Background()
|
|
runner := NewRunner(ctx, RunnerConfig{
|
|
Agent: agent,
|
|
EnableStreaming: true,
|
|
CheckPointStore: store,
|
|
})
|
|
iter := runner.Query(ctx, "hello world", WithCheckPointID("1"))
|
|
_, ok := iter.Next()
|
|
assert.True(t, ok)
|
|
interruptEvent, ok := iter.Next()
|
|
assert.True(t, ok)
|
|
assert.Equal(t, data, interruptEvent.Action.Interrupted.Data)
|
|
assert.NotEmpty(t, interruptEvent.Action.Interrupted.InterruptContexts[0].ID)
|
|
assert.True(t, interruptEvent.Action.Interrupted.InterruptContexts[0].IsRootCause)
|
|
assert.Equal(t, data, interruptEvent.Action.Interrupted.InterruptContexts[0].Info)
|
|
assert.Equal(t, Address{{Type: AddressSegmentAgent, ID: "myAgent"}},
|
|
interruptEvent.Action.Interrupted.InterruptContexts[0].Address)
|
|
_, ok = iter.Next()
|
|
assert.False(t, ok)
|
|
|
|
iter, err := runner.ResumeWithParams(ctx, "1", &ResumeParams{
|
|
Targets: map[string]any{
|
|
interruptEvent.Action.Interrupted.InterruptContexts[0].ID: nil,
|
|
},
|
|
})
|
|
assert.NoError(t, err)
|
|
_, ok = iter.Next()
|
|
assert.False(t, ok)
|
|
}
|
|
|
|
func TestMultiAgentInterrupt(t *testing.T) {
|
|
ctx := context.Background()
|
|
sa1 := &myAgent{
|
|
name: "sa1",
|
|
runFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {
|
|
iter, generator := NewAsyncIteratorPair[*AgentEvent]()
|
|
generator.Send(&AgentEvent{
|
|
AgentName: "sa1",
|
|
Action: &AgentAction{
|
|
TransferToAgent: &TransferToAgentAction{
|
|
DestAgentName: "sa2",
|
|
},
|
|
},
|
|
})
|
|
generator.Close()
|
|
return iter
|
|
},
|
|
}
|
|
sa2 := &myAgent{
|
|
name: "sa2",
|
|
runFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {
|
|
iter, generator := NewAsyncIteratorPair[*AgentEvent]()
|
|
intEvent := StatefulInterrupt(ctx, "hello world", "temp state")
|
|
intEvent.Action.Interrupted.Data = "hello world"
|
|
generator.Send(intEvent)
|
|
generator.Close()
|
|
return iter
|
|
},
|
|
resumeFn: func(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] {
|
|
assert.NotNil(t, info)
|
|
assert.Equal(t, info.Data, "hello world")
|
|
|
|
assert.True(t, info.WasInterrupted)
|
|
assert.NotNil(t, info.InterruptState)
|
|
assert.Equal(t, "temp state", info.InterruptState)
|
|
|
|
assert.True(t, info.IsResumeTarget)
|
|
assert.NotNil(t, info.ResumeData)
|
|
assert.Equal(t, "resume data", info.ResumeData)
|
|
|
|
iter, generator := NewAsyncIteratorPair[*AgentEvent]()
|
|
generator.Send(&AgentEvent{
|
|
AgentName: "sa2",
|
|
Output: &AgentOutput{
|
|
MessageOutput: &MessageVariant{Message: schema.UserMessage(info.ResumeData.(string))},
|
|
},
|
|
})
|
|
generator.Close()
|
|
return iter
|
|
},
|
|
}
|
|
a, err := SetSubAgents(ctx, sa1, []Agent{sa2})
|
|
assert.NoError(t, err)
|
|
runner := NewRunner(ctx, RunnerConfig{
|
|
Agent: a,
|
|
EnableStreaming: false,
|
|
CheckPointStore: newMyStore(),
|
|
})
|
|
iter := runner.Query(ctx, "", WithCheckPointID("1"))
|
|
event, ok := iter.Next()
|
|
assert.True(t, ok)
|
|
assert.NotNil(t, event.Action.TransferToAgent)
|
|
event, ok = iter.Next()
|
|
assert.True(t, ok)
|
|
assert.NotNil(t, event.Action.Interrupted)
|
|
assert.Equal(t, 1, len(event.Action.Interrupted.InterruptContexts))
|
|
assert.Equal(t, "hello world", event.Action.Interrupted.InterruptContexts[0].Info)
|
|
assert.True(t, event.Action.Interrupted.InterruptContexts[0].IsRootCause)
|
|
assert.Equal(t, Address{
|
|
{Type: AddressSegmentAgent, ID: "sa1"},
|
|
{Type: AddressSegmentAgent, ID: "sa2"},
|
|
}, event.Action.Interrupted.InterruptContexts[0].Address)
|
|
assert.NotEmpty(t, event.Action.Interrupted.InterruptContexts[0].ID)
|
|
|
|
interruptID := event.Action.Interrupted.InterruptContexts[0].ID
|
|
_, ok = iter.Next()
|
|
assert.False(t, ok)
|
|
|
|
iter, err = runner.ResumeWithParams(ctx, "1", &ResumeParams{
|
|
Targets: map[string]any{
|
|
interruptID: "resume data",
|
|
},
|
|
})
|
|
assert.NoError(t, err)
|
|
event, ok = iter.Next()
|
|
assert.True(t, ok)
|
|
assert.Equal(t, event.Output.MessageOutput.Message.Content, "resume data")
|
|
_, ok = iter.Next()
|
|
assert.False(t, ok)
|
|
}
|
|
|
|
func TestWorkflowInterrupt(t *testing.T) {
|
|
ctx := context.Background()
|
|
sa1 := &myAgent{
|
|
name: "sa1",
|
|
runFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {
|
|
iter, generator := NewAsyncIteratorPair[*AgentEvent]()
|
|
|
|
intEvent := Interrupt(ctx, "sa1 interrupt data")
|
|
intEvent.Action.Interrupted.Data = "sa1 interrupt data"
|
|
generator.Send(intEvent)
|
|
generator.Close()
|
|
return iter
|
|
},
|
|
resumeFn: func(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] {
|
|
assert.Equal(t, info.InterruptInfo.Data, "sa1 interrupt data")
|
|
assert.True(t, info.WasInterrupted)
|
|
assert.Nil(t, info.InterruptState)
|
|
assert.True(t, info.IsResumeTarget)
|
|
assert.Equal(t, "resume sa1", info.ResumeData)
|
|
iter, generator := NewAsyncIteratorPair[*AgentEvent]()
|
|
generator.Close()
|
|
return iter
|
|
},
|
|
} // interrupt once
|
|
sa2 := &myAgent{
|
|
name: "sa2",
|
|
runFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {
|
|
iter, generator := NewAsyncIteratorPair[*AgentEvent]()
|
|
|
|
intEvent := StatefulInterrupt(ctx, "sa2 interrupt data", "sa2 interrupt")
|
|
intEvent.Action.Interrupted.Data = "sa2 interrupt data"
|
|
generator.Send(intEvent)
|
|
generator.Close()
|
|
return iter
|
|
},
|
|
resumeFn: func(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] {
|
|
assert.Equal(t, info.InterruptInfo.Data, "sa2 interrupt data")
|
|
assert.True(t, info.WasInterrupted)
|
|
assert.NotNil(t, info.InterruptState)
|
|
assert.Equal(t, "sa2 interrupt", info.InterruptState)
|
|
|
|
assert.True(t, info.IsResumeTarget)
|
|
assert.NotNil(t, info.ResumeData)
|
|
assert.Equal(t, "resume sa2", info.ResumeData)
|
|
iter, generator := NewAsyncIteratorPair[*AgentEvent]()
|
|
generator.Close()
|
|
return iter
|
|
},
|
|
} // interrupt once
|
|
sa3 := &myAgent{
|
|
name: "sa3",
|
|
runFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {
|
|
iter, generator := NewAsyncIteratorPair[*AgentEvent]()
|
|
generator.Send(&AgentEvent{
|
|
AgentName: "sa3",
|
|
Output: &AgentOutput{
|
|
MessageOutput: &MessageVariant{
|
|
Message: schema.UserMessage("sa3 completed"),
|
|
},
|
|
},
|
|
})
|
|
generator.Close()
|
|
return iter
|
|
},
|
|
} // won't interrupt
|
|
sa4 := &myAgent{
|
|
name: "sa4",
|
|
runFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {
|
|
iter, generator := NewAsyncIteratorPair[*AgentEvent]()
|
|
generator.Send(&AgentEvent{
|
|
AgentName: "sa4",
|
|
Output: &AgentOutput{
|
|
MessageOutput: &MessageVariant{
|
|
Message: schema.UserMessage("sa4 completed"),
|
|
},
|
|
},
|
|
})
|
|
generator.Close()
|
|
return iter
|
|
},
|
|
} // won't interrupt
|
|
|
|
firstInterruptEvent := &AgentEvent{
|
|
AgentName: "sa1",
|
|
RunPath: []RunStep{{"sequential"}, {"sa1"}},
|
|
Action: &AgentAction{
|
|
Interrupted: &InterruptInfo{
|
|
Data: &WorkflowInterruptInfo{
|
|
OrigInput: &AgentInput{
|
|
Messages: []Message{schema.UserMessage("hello world")},
|
|
},
|
|
SequentialInterruptIndex: 0,
|
|
SequentialInterruptInfo: &InterruptInfo{
|
|
Data: "sa1 interrupt data",
|
|
},
|
|
LoopIterations: 0,
|
|
},
|
|
InterruptContexts: []*InterruptCtx{
|
|
{
|
|
ID: "agent:sequential;agent:sa1",
|
|
Info: "sa1 interrupt data",
|
|
Address: Address{
|
|
{
|
|
ID: "sequential",
|
|
Type: AddressSegmentAgent,
|
|
},
|
|
{
|
|
ID: "sa1",
|
|
Type: AddressSegmentAgent,
|
|
},
|
|
},
|
|
IsRootCause: true,
|
|
Parent: &InterruptCtx{
|
|
ID: "agent:sequential",
|
|
Info: "Sequential workflow interrupted",
|
|
Address: Address{
|
|
{
|
|
ID: "sequential",
|
|
Type: AddressSegmentAgent,
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}
|
|
_ = firstInterruptEvent
|
|
secondInterruptEvent := &AgentEvent{
|
|
AgentName: "sa2",
|
|
RunPath: []RunStep{{"sequential"}, {"sa1"}, {"sa2"}},
|
|
Action: &AgentAction{
|
|
Interrupted: &InterruptInfo{
|
|
Data: &WorkflowInterruptInfo{
|
|
OrigInput: &AgentInput{
|
|
Messages: []Message{schema.UserMessage("hello world")},
|
|
},
|
|
SequentialInterruptIndex: 1,
|
|
SequentialInterruptInfo: &InterruptInfo{
|
|
Data: "sa2 interrupt data",
|
|
},
|
|
},
|
|
InterruptContexts: []*InterruptCtx{
|
|
{
|
|
ID: "agent:sequential;agent:sa1;agent:sa2",
|
|
Info: "sa2 interrupt data",
|
|
Address: Address{
|
|
{
|
|
ID: "sequential",
|
|
Type: AddressSegmentAgent,
|
|
},
|
|
{
|
|
ID: "sa2",
|
|
Type: AddressSegmentAgent,
|
|
},
|
|
},
|
|
IsRootCause: true,
|
|
Parent: &InterruptCtx{
|
|
ID: "agent:sequential",
|
|
Info: "Sequential workflow interrupted",
|
|
Address: Address{
|
|
{
|
|
ID: "sequential",
|
|
Type: AddressSegmentAgent,
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}
|
|
_ = secondInterruptEvent
|
|
messageEvents := []*AgentEvent{
|
|
{
|
|
AgentName: "sa3",
|
|
RunPath: []RunStep{{"sequential"}, {"sa1"}, {"sa2"}, {"sa3"}},
|
|
Output: &AgentOutput{
|
|
MessageOutput: &MessageVariant{
|
|
Message: schema.UserMessage("sa3 completed"),
|
|
},
|
|
},
|
|
},
|
|
{
|
|
AgentName: "sa4",
|
|
RunPath: []RunStep{{"sequential"}, {"sa1"}, {"sa2"}, {"sa3"}, {"sa4"}},
|
|
Output: &AgentOutput{
|
|
MessageOutput: &MessageVariant{
|
|
Message: schema.UserMessage("sa4 completed"),
|
|
},
|
|
},
|
|
},
|
|
}
|
|
_ = messageEvents
|
|
|
|
t.Run("test sequential workflow agent", func(t *testing.T) {
|
|
|
|
// sequential
|
|
a, err := NewSequentialAgent(ctx, &SequentialAgentConfig{
|
|
Name: "sequential",
|
|
Description: "sequential agent",
|
|
SubAgents: []Agent{sa1, sa2, sa3, sa4},
|
|
})
|
|
assert.NoError(t, err)
|
|
runner := NewRunner(ctx, RunnerConfig{
|
|
Agent: a,
|
|
CheckPointStore: newMyStore(),
|
|
})
|
|
var events []*AgentEvent
|
|
iter := runner.Query(ctx, "hello world", WithCheckPointID("sequential-1"))
|
|
for {
|
|
event, ok := iter.Next()
|
|
if !ok {
|
|
break
|
|
}
|
|
events = append(events, event)
|
|
}
|
|
|
|
assert.Equal(t, 1, len(events))
|
|
assert.Equal(t, firstInterruptEvent.AgentName, events[0].AgentName)
|
|
assert.Equal(t, firstInterruptEvent.RunPath, events[0].RunPath)
|
|
assert.True(t, events[0].Action.Interrupted.InterruptContexts[0].EqualsWithoutID(firstInterruptEvent.Action.Interrupted.InterruptContexts[0]))
|
|
interruptID1 := events[0].Action.Interrupted.InterruptContexts[0].ID
|
|
events = []*AgentEvent{}
|
|
|
|
// Resume after sa1 interrupt
|
|
iter, err = runner.ResumeWithParams(ctx, "sequential-1", &ResumeParams{
|
|
Targets: map[string]any{
|
|
interruptID1: "resume sa1",
|
|
},
|
|
})
|
|
assert.NoError(t, err)
|
|
for {
|
|
event, ok := iter.Next()
|
|
if !ok {
|
|
break
|
|
}
|
|
events = append(events, event)
|
|
}
|
|
|
|
assert.Equal(t, 1, len(events))
|
|
assert.Equal(t, secondInterruptEvent.AgentName, events[0].AgentName)
|
|
assert.Equal(t, secondInterruptEvent.RunPath, events[0].RunPath)
|
|
assert.True(t, events[0].Action.Interrupted.InterruptContexts[0].
|
|
EqualsWithoutID(secondInterruptEvent.Action.Interrupted.InterruptContexts[0]))
|
|
interruptID2 := events[0].Action.Interrupted.InterruptContexts[0].ID
|
|
events = []*AgentEvent{}
|
|
|
|
// Resume after sa2 interrupt
|
|
iter, err = runner.ResumeWithParams(ctx, "sequential-1", &ResumeParams{
|
|
Targets: map[string]any{
|
|
interruptID2: "resume sa2",
|
|
},
|
|
})
|
|
assert.NoError(t, err)
|
|
for {
|
|
event, ok := iter.Next()
|
|
if !ok {
|
|
break
|
|
}
|
|
events = append(events, event)
|
|
}
|
|
|
|
assert.Equal(t, 2, len(events))
|
|
assert.Equal(t, messageEvents, events)
|
|
})
|
|
|
|
t.Run("test loop workflow agent", func(t *testing.T) {
|
|
// loop
|
|
a, err := NewLoopAgent(ctx, &LoopAgentConfig{
|
|
Name: "loop",
|
|
SubAgents: []Agent{sa1, sa2, sa3, sa4},
|
|
MaxIterations: 2,
|
|
})
|
|
assert.NoError(t, err)
|
|
runner := NewRunner(ctx, RunnerConfig{
|
|
Agent: a,
|
|
CheckPointStore: newMyStore(),
|
|
})
|
|
var events []*AgentEvent
|
|
iter := runner.Query(ctx, "hello world", WithCheckPointID("loop-1"))
|
|
for {
|
|
event, ok := iter.Next()
|
|
if !ok {
|
|
break
|
|
}
|
|
events = append(events, event)
|
|
}
|
|
|
|
loopFirstInterruptEvent := &AgentEvent{
|
|
AgentName: "sa1",
|
|
RunPath: []RunStep{{"loop"}, {"sa1"}},
|
|
Action: &AgentAction{
|
|
Interrupted: &InterruptInfo{
|
|
Data: &WorkflowInterruptInfo{
|
|
OrigInput: &AgentInput{
|
|
Messages: []Message{schema.UserMessage("hello world")},
|
|
},
|
|
SequentialInterruptIndex: 0,
|
|
SequentialInterruptInfo: &InterruptInfo{
|
|
Data: "sa1 interrupt data",
|
|
},
|
|
LoopIterations: 0,
|
|
},
|
|
InterruptContexts: []*InterruptCtx{
|
|
{
|
|
ID: "agent:loop;agent:sa1",
|
|
Info: "sa1 interrupt data",
|
|
Address: Address{
|
|
{
|
|
ID: "loop",
|
|
Type: AddressSegmentAgent,
|
|
},
|
|
{
|
|
ID: "sa1",
|
|
Type: AddressSegmentAgent,
|
|
},
|
|
},
|
|
IsRootCause: true,
|
|
Parent: &InterruptCtx{
|
|
ID: "agent:loop",
|
|
Info: "Loop workflow interrupted",
|
|
Address: Address{
|
|
{
|
|
ID: "loop",
|
|
Type: AddressSegmentAgent,
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}
|
|
assert.Equal(t, 1, len(events))
|
|
assert.Equal(t, loopFirstInterruptEvent.AgentName, events[0].AgentName)
|
|
assert.Equal(t, loopFirstInterruptEvent.RunPath, events[0].RunPath)
|
|
assert.True(t, events[0].Action.Interrupted.InterruptContexts[0].EqualsWithoutID(loopFirstInterruptEvent.Action.Interrupted.InterruptContexts[0]))
|
|
loopInterruptID1 := events[0].Action.Interrupted.InterruptContexts[0].ID
|
|
events = []*AgentEvent{}
|
|
|
|
// Resume after sa1 interrupt
|
|
iter, err = runner.ResumeWithParams(ctx, "loop-1", &ResumeParams{
|
|
Targets: map[string]any{
|
|
loopInterruptID1: "resume sa1",
|
|
},
|
|
})
|
|
assert.NoError(t, err)
|
|
for {
|
|
event, ok := iter.Next()
|
|
if !ok {
|
|
break
|
|
}
|
|
events = append(events, event)
|
|
}
|
|
|
|
loopSecondInterruptEvent := &AgentEvent{
|
|
AgentName: "sa2",
|
|
RunPath: []RunStep{{"loop"}, {"sa1"}, {"sa2"}},
|
|
Action: &AgentAction{
|
|
Interrupted: &InterruptInfo{
|
|
Data: &WorkflowInterruptInfo{
|
|
OrigInput: &AgentInput{
|
|
Messages: []Message{schema.UserMessage("hello world")},
|
|
},
|
|
SequentialInterruptIndex: 1,
|
|
SequentialInterruptInfo: &InterruptInfo{
|
|
Data: "sa2 interrupt data",
|
|
},
|
|
LoopIterations: 0,
|
|
},
|
|
InterruptContexts: []*InterruptCtx{
|
|
{
|
|
ID: "agent:loop;agent:sa1;agent:sa2",
|
|
Info: "sa2 interrupt data",
|
|
Address: Address{
|
|
{
|
|
ID: "loop",
|
|
Type: AddressSegmentAgent,
|
|
},
|
|
{
|
|
ID: "sa2",
|
|
Type: AddressSegmentAgent,
|
|
},
|
|
},
|
|
IsRootCause: true,
|
|
Parent: &InterruptCtx{
|
|
ID: "agent:loop",
|
|
Info: "Loop workflow interrupted",
|
|
Address: Address{
|
|
{
|
|
ID: "loop",
|
|
Type: AddressSegmentAgent,
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}
|
|
assert.Equal(t, 1, len(events))
|
|
assert.Equal(t, loopSecondInterruptEvent.AgentName, events[0].AgentName)
|
|
assert.Equal(t, loopSecondInterruptEvent.RunPath, events[0].RunPath)
|
|
assert.True(t, events[0].Action.Interrupted.InterruptContexts[0].EqualsWithoutID(loopSecondInterruptEvent.Action.Interrupted.InterruptContexts[0]))
|
|
loopInterruptID2 := events[0].Action.Interrupted.InterruptContexts[0].ID
|
|
events = []*AgentEvent{}
|
|
|
|
// Resume after sa2 interrupt
|
|
iter, err = runner.ResumeWithParams(ctx, "loop-1", &ResumeParams{
|
|
Targets: map[string]any{
|
|
loopInterruptID2: "resume sa2",
|
|
},
|
|
})
|
|
assert.NoError(t, err)
|
|
for {
|
|
event, ok := iter.Next()
|
|
if !ok {
|
|
break
|
|
}
|
|
events = append(events, event)
|
|
}
|
|
|
|
loopThirdInterruptEvent := &AgentEvent{
|
|
AgentName: "sa1",
|
|
RunPath: []RunStep{{"loop"}, {"sa1"}, {"sa2"}, {"sa3"}, {"sa4"}, {"sa1"}},
|
|
Action: &AgentAction{
|
|
Interrupted: &InterruptInfo{
|
|
Data: &WorkflowInterruptInfo{
|
|
OrigInput: &AgentInput{
|
|
Messages: []Message{schema.UserMessage("hello world")},
|
|
},
|
|
SequentialInterruptIndex: 0,
|
|
SequentialInterruptInfo: &InterruptInfo{
|
|
Data: "sa1 interrupt data",
|
|
},
|
|
LoopIterations: 1,
|
|
},
|
|
InterruptContexts: []*InterruptCtx{
|
|
{
|
|
ID: "agent:loop;agent:sa1;agent:sa2;agent:sa3;agent:sa4;agent:sa1",
|
|
Info: "sa1 interrupt data",
|
|
Address: Address{
|
|
{
|
|
ID: "loop",
|
|
Type: AddressSegmentAgent,
|
|
},
|
|
{
|
|
ID: "sa1",
|
|
Type: AddressSegmentAgent,
|
|
},
|
|
},
|
|
IsRootCause: true,
|
|
Parent: &InterruptCtx{
|
|
ID: "agent:loop",
|
|
Info: "Loop workflow interrupted",
|
|
Address: Address{
|
|
{
|
|
ID: "loop",
|
|
Type: AddressSegmentAgent,
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}
|
|
|
|
loopFourthInterruptEvent := &AgentEvent{
|
|
AgentName: "sa2",
|
|
RunPath: []RunStep{{"loop"}, {"sa1"}, {"sa2"}, {"sa3"}, {"sa4"}, {"sa1"}, {"sa2"}},
|
|
Action: &AgentAction{
|
|
Interrupted: &InterruptInfo{
|
|
Data: &WorkflowInterruptInfo{
|
|
OrigInput: &AgentInput{
|
|
Messages: []Message{schema.UserMessage("hello world")},
|
|
},
|
|
SequentialInterruptIndex: 1,
|
|
SequentialInterruptInfo: &InterruptInfo{
|
|
Data: "sa2 interrupt data",
|
|
},
|
|
LoopIterations: 1,
|
|
},
|
|
InterruptContexts: []*InterruptCtx{
|
|
{
|
|
ID: "agent:loop;agent:sa1;agent:sa2;agent:sa3;agent:sa4;agent:sa1;agent:sa2",
|
|
Info: "sa2 interrupt data",
|
|
Address: Address{
|
|
{
|
|
ID: "loop",
|
|
Type: AddressSegmentAgent,
|
|
},
|
|
{
|
|
ID: "sa2",
|
|
Type: AddressSegmentAgent,
|
|
},
|
|
},
|
|
IsRootCause: true,
|
|
Parent: &InterruptCtx{
|
|
ID: "agent:loop",
|
|
Info: "Loop workflow interrupted",
|
|
Address: Address{
|
|
{
|
|
ID: "loop",
|
|
Type: AddressSegmentAgent,
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}
|
|
|
|
loopMessageEvents := []*AgentEvent{
|
|
{
|
|
AgentName: "sa3",
|
|
RunPath: []RunStep{{"loop"}, {"sa1"}, {"sa2"}, {"sa3"}},
|
|
Output: &AgentOutput{
|
|
MessageOutput: &MessageVariant{
|
|
Message: schema.UserMessage("sa3 completed"),
|
|
},
|
|
},
|
|
},
|
|
{
|
|
AgentName: "sa4",
|
|
RunPath: []RunStep{{"loop"}, {"sa1"}, {"sa2"}, {"sa3"}, {"sa4"}},
|
|
Output: &AgentOutput{
|
|
MessageOutput: &MessageVariant{
|
|
Message: schema.UserMessage("sa4 completed"),
|
|
},
|
|
},
|
|
},
|
|
loopThirdInterruptEvent,
|
|
}
|
|
assert.Equal(t, 3, len(events))
|
|
// Check the first two message events
|
|
assert.Equal(t, loopMessageEvents[0].AgentName, events[0].AgentName)
|
|
assert.Equal(t, loopMessageEvents[0].RunPath, events[0].RunPath)
|
|
assert.Equal(t, loopMessageEvents[0].Output.MessageOutput.Message.Content, events[0].Output.MessageOutput.Message.Content)
|
|
|
|
assert.Equal(t, loopMessageEvents[1].AgentName, events[1].AgentName)
|
|
assert.Equal(t, loopMessageEvents[1].RunPath, events[1].RunPath)
|
|
assert.Equal(t, loopMessageEvents[1].Output.MessageOutput.Message.Content, events[1].Output.MessageOutput.Message.Content)
|
|
|
|
// Check the third interrupt event using EqualsWithoutID
|
|
assert.Equal(t, loopMessageEvents[2].AgentName, events[2].AgentName)
|
|
assert.Equal(t, loopMessageEvents[2].RunPath, events[2].RunPath)
|
|
assert.True(t, events[2].Action.Interrupted.InterruptContexts[0].EqualsWithoutID(loopMessageEvents[2].Action.Interrupted.InterruptContexts[0]))
|
|
loopInterruptID3 := events[2].Action.Interrupted.InterruptContexts[0].ID
|
|
events = []*AgentEvent{}
|
|
|
|
// Resume after third interrupt
|
|
iter, err = runner.ResumeWithParams(ctx, "loop-1", &ResumeParams{
|
|
Targets: map[string]any{
|
|
loopInterruptID3: "resume sa1",
|
|
},
|
|
})
|
|
assert.NoError(t, err)
|
|
for {
|
|
event, ok := iter.Next()
|
|
if !ok {
|
|
break
|
|
}
|
|
events = append(events, event)
|
|
}
|
|
assert.Equal(t, 1, len(events))
|
|
assert.Equal(t, loopFourthInterruptEvent.AgentName, events[0].AgentName)
|
|
assert.Equal(t, loopFourthInterruptEvent.RunPath, events[0].RunPath)
|
|
assert.True(t, events[0].Action.Interrupted.InterruptContexts[0].EqualsWithoutID(loopFourthInterruptEvent.Action.Interrupted.InterruptContexts[0]))
|
|
loopInterruptID4 := events[0].Action.Interrupted.InterruptContexts[0].ID
|
|
events = []*AgentEvent{}
|
|
|
|
// Resume after fourth interrupt
|
|
iter, err = runner.ResumeWithParams(ctx, "loop-1", &ResumeParams{
|
|
Targets: map[string]any{
|
|
loopInterruptID4: "resume sa2",
|
|
},
|
|
})
|
|
assert.NoError(t, err)
|
|
for {
|
|
event, ok := iter.Next()
|
|
if !ok {
|
|
break
|
|
}
|
|
events = append(events, event)
|
|
}
|
|
loopFinalMessageEvents := []*AgentEvent{
|
|
{
|
|
AgentName: "sa3",
|
|
RunPath: []RunStep{{"loop"}, {"sa1"}, {"sa2"}, {"sa3"}, {"sa4"}, {"sa1"}, {"sa2"}, {"sa3"}},
|
|
Output: &AgentOutput{
|
|
MessageOutput: &MessageVariant{
|
|
Message: schema.UserMessage("sa3 completed"),
|
|
},
|
|
},
|
|
},
|
|
{
|
|
AgentName: "sa4",
|
|
RunPath: []RunStep{{"loop"}, {"sa1"}, {"sa2"}, {"sa3"}, {"sa4"}, {"sa1"}, {"sa2"}, {"sa3"}, {"sa4"}},
|
|
Output: &AgentOutput{
|
|
MessageOutput: &MessageVariant{
|
|
Message: schema.UserMessage("sa4 completed"),
|
|
},
|
|
},
|
|
},
|
|
}
|
|
assert.Equal(t, 2, len(events))
|
|
assert.Equal(t, loopFinalMessageEvents, events)
|
|
})
|
|
|
|
t.Run("test parallel workflow agent", func(t *testing.T) {
|
|
// parallel
|
|
a, err := NewParallelAgent(ctx, &ParallelAgentConfig{
|
|
Name: "parallel agent",
|
|
SubAgents: []Agent{sa1, sa2, sa3, sa4},
|
|
})
|
|
assert.NoError(t, err)
|
|
runner := NewRunner(ctx, RunnerConfig{
|
|
Agent: a,
|
|
CheckPointStore: newMyStore(),
|
|
})
|
|
iter := runner.Query(ctx, "hello world", WithCheckPointID("1"))
|
|
var (
|
|
events []*AgentEvent
|
|
interruptEvent *AgentEvent
|
|
)
|
|
|
|
for {
|
|
event, ok := iter.Next()
|
|
if !ok {
|
|
break
|
|
}
|
|
if event.Action != nil && event.Action.Interrupted != nil {
|
|
interruptEvent = event
|
|
continue
|
|
}
|
|
events = append(events, event)
|
|
}
|
|
assert.Equal(t, 2, len(events))
|
|
|
|
// Debug: Print actual events to see what we're getting
|
|
for i, event := range events {
|
|
t.Logf("Event %d: AgentName=%s, RunPath=%v, Output=%v", i, event.AgentName, event.RunPath, event.Output)
|
|
}
|
|
|
|
// Define parallel message events separately
|
|
parallelMessageEvents := []*AgentEvent{
|
|
{
|
|
AgentName: "sa4",
|
|
RunPath: []RunStep{{"parallel agent"}, {"sa4"}},
|
|
Output: &AgentOutput{
|
|
MessageOutput: &MessageVariant{
|
|
Message: schema.UserMessage("sa4 completed"),
|
|
},
|
|
},
|
|
},
|
|
{
|
|
AgentName: "sa3",
|
|
RunPath: []RunStep{{"parallel agent"}, {"sa3"}},
|
|
Output: &AgentOutput{
|
|
MessageOutput: &MessageVariant{
|
|
Message: schema.UserMessage("sa3 completed"),
|
|
},
|
|
},
|
|
},
|
|
}
|
|
|
|
assert.Contains(t, events, parallelMessageEvents[0])
|
|
assert.Contains(t, events, parallelMessageEvents[1])
|
|
|
|
assert.NotNil(t, interruptEvent)
|
|
assert.Equal(t, "parallel agent", interruptEvent.AgentName)
|
|
assert.Equal(t, []RunStep{{"parallel agent"}}, interruptEvent.RunPath)
|
|
assert.NotNil(t, interruptEvent.Action.Interrupted)
|
|
wii, ok := interruptEvent.Action.Interrupted.Data.(*WorkflowInterruptInfo)
|
|
assert.True(t, ok)
|
|
assert.Equal(t, 2, len(wii.ParallelInterruptInfo))
|
|
|
|
var sa1Found, sa2Found bool
|
|
for _, info := range wii.ParallelInterruptInfo {
|
|
switch info.Data {
|
|
case "sa1 interrupt data":
|
|
sa1Found = true
|
|
case "sa2 interrupt data":
|
|
sa2Found = true
|
|
}
|
|
}
|
|
assert.True(t, sa1Found)
|
|
assert.True(t, sa2Found)
|
|
|
|
var sa1InfoFound, sa2InfoFound bool
|
|
for _, ctx := range interruptEvent.Action.Interrupted.InterruptContexts {
|
|
switch ctx.Info {
|
|
case "sa1 interrupt data":
|
|
sa1InfoFound = true
|
|
case "sa2 interrupt data":
|
|
sa2InfoFound = true
|
|
}
|
|
}
|
|
|
|
assert.Equal(t, 2, len(interruptEvent.Action.Interrupted.InterruptContexts))
|
|
assert.True(t, sa1InfoFound)
|
|
assert.True(t, sa2InfoFound)
|
|
|
|
var parallelInterruptID1, parallelInterruptID2 string
|
|
for _, ctx := range interruptEvent.Action.Interrupted.InterruptContexts {
|
|
switch ctx.Info {
|
|
case "sa1 interrupt data":
|
|
parallelInterruptID1 = ctx.ID
|
|
case "sa2 interrupt data":
|
|
parallelInterruptID2 = ctx.ID
|
|
}
|
|
}
|
|
assert.NotEmpty(t, parallelInterruptID1)
|
|
assert.NotEmpty(t, parallelInterruptID2)
|
|
|
|
iter, err = runner.ResumeWithParams(ctx, "1", &ResumeParams{
|
|
Targets: map[string]any{
|
|
parallelInterruptID1: "resume sa1",
|
|
parallelInterruptID2: "resume sa2",
|
|
},
|
|
})
|
|
assert.NoError(t, err)
|
|
_, ok = iter.Next()
|
|
assert.False(t, ok)
|
|
})
|
|
}
|
|
|
|
func TestChatModelInterrupt(t *testing.T) {
|
|
ctx := context.Background()
|
|
a, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{
|
|
Name: "name",
|
|
Description: "description",
|
|
Instruction: "instruction",
|
|
Model: &myModel{
|
|
validator: func(i int, messages []*schema.Message) bool {
|
|
if i > 0 && (len(messages) != 4 || messages[2].Content != "new user message") {
|
|
return false
|
|
}
|
|
return true
|
|
},
|
|
messages: []*schema.Message{
|
|
schema.AssistantMessage("", []schema.ToolCall{
|
|
{
|
|
ID: "1",
|
|
Function: schema.FunctionCall{
|
|
Name: "tool1",
|
|
Arguments: "arguments",
|
|
},
|
|
},
|
|
}),
|
|
schema.AssistantMessage("completed", nil),
|
|
},
|
|
},
|
|
ToolsConfig: ToolsConfig{
|
|
ToolsNodeConfig: compose.ToolsNodeConfig{
|
|
Tools: []tool.BaseTool{&myTool1{}},
|
|
},
|
|
},
|
|
})
|
|
assert.NoError(t, err)
|
|
runner := NewRunner(ctx, RunnerConfig{
|
|
Agent: a,
|
|
CheckPointStore: newMyStore(),
|
|
})
|
|
iter := runner.Query(ctx, "hello world", WithCheckPointID("1"))
|
|
_, ok := iter.Next()
|
|
assert.True(t, ok)
|
|
event, ok := iter.Next()
|
|
assert.True(t, ok)
|
|
assert.NoError(t, event.Err)
|
|
assert.NotNil(t, event.Action.Interrupted)
|
|
assert.Equal(t, 1, len(event.Action.Interrupted.InterruptContexts))
|
|
assert.Equal(t, Address{
|
|
{Type: AddressSegmentAgent, ID: "name"},
|
|
{Type: AddressSegmentTool, ID: "tool1", SubID: "1"},
|
|
}, event.Action.Interrupted.InterruptContexts[0].Address)
|
|
|
|
var (
|
|
chatModelAgentID string
|
|
toolID string
|
|
)
|
|
|
|
intCtx := event.Action.Interrupted.InterruptContexts[0]
|
|
for intCtx != nil {
|
|
if intCtx.Address[len(intCtx.Address)-1].Type == AddressSegmentTool {
|
|
toolID = intCtx.ID
|
|
} else if intCtx.Address[len(intCtx.Address)-1].Type == AddressSegmentAgent {
|
|
chatModelAgentID = intCtx.ID
|
|
}
|
|
intCtx = intCtx.Parent
|
|
}
|
|
|
|
_, ok = iter.Next()
|
|
assert.False(t, ok)
|
|
|
|
iter, err = runner.ResumeWithParams(ctx, "1", &ResumeParams{
|
|
Targets: map[string]any{
|
|
chatModelAgentID: &ChatModelAgentResumeData{
|
|
HistoryModifier: func(ctx context.Context, history []Message) []Message {
|
|
history[2].Content = "new user message"
|
|
return history
|
|
},
|
|
},
|
|
toolID: "tool resume result",
|
|
},
|
|
})
|
|
assert.NoError(t, err)
|
|
event, ok = iter.Next()
|
|
assert.True(t, ok)
|
|
assert.NoError(t, event.Err)
|
|
assert.Equal(t, event.Output.MessageOutput.Message.Content, "tool resume result")
|
|
event, ok = iter.Next()
|
|
assert.True(t, ok)
|
|
assert.NoError(t, event.Err)
|
|
assert.Equal(t, event.Output.MessageOutput.Message.Content, "completed")
|
|
}
|
|
|
|
func TestChatModelAgentToolInterrupt(t *testing.T) {
|
|
sa := &myAgent{
|
|
runFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {
|
|
iter, generator := NewAsyncIteratorPair[*AgentEvent]()
|
|
intAct := Interrupt(ctx, "hello world")
|
|
intAct.Action.Interrupted.Data = "hello world"
|
|
generator.Send(intAct)
|
|
generator.Close()
|
|
return iter
|
|
},
|
|
resumeFn: func(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] {
|
|
assert.NotNil(t, info)
|
|
assert.False(t, info.EnableStreaming)
|
|
|
|
if !info.IsResumeTarget {
|
|
iter, generator := NewAsyncIteratorPair[*AgentEvent]()
|
|
intAct := Interrupt(ctx, "interrupt again")
|
|
intAct.Action.Interrupted.Data = "interrupt again"
|
|
generator.Send(intAct)
|
|
generator.Close()
|
|
return iter
|
|
}
|
|
|
|
assert.NotNil(t, info.ResumeData)
|
|
assert.Equal(t, "resume sa", info.ResumeData)
|
|
|
|
iter, generator := NewAsyncIteratorPair[*AgentEvent]()
|
|
generator.Send(&AgentEvent{Output: &AgentOutput{MessageOutput: &MessageVariant{Message: schema.UserMessage(fmt.Sprintf("my agent completed with data %s", info.ResumeData))}}})
|
|
generator.Close()
|
|
return iter
|
|
},
|
|
}
|
|
ctx := context.Background()
|
|
a, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{
|
|
Name: "name",
|
|
Description: "description",
|
|
Instruction: "instruction",
|
|
Model: &myModel{
|
|
messages: []*schema.Message{
|
|
schema.AssistantMessage("", []schema.ToolCall{
|
|
{
|
|
ID: "1",
|
|
Function: schema.FunctionCall{
|
|
Name: "myAgent",
|
|
Arguments: "{\"request\":\"123\"}",
|
|
},
|
|
},
|
|
}),
|
|
schema.AssistantMessage("completed", nil),
|
|
},
|
|
},
|
|
ToolsConfig: ToolsConfig{
|
|
ToolsNodeConfig: compose.ToolsNodeConfig{
|
|
Tools: []tool.BaseTool{NewAgentTool(ctx, sa)},
|
|
},
|
|
},
|
|
})
|
|
assert.NoError(t, err)
|
|
runner := NewRunner(ctx, RunnerConfig{
|
|
Agent: a,
|
|
CheckPointStore: newMyStore(),
|
|
})
|
|
|
|
iter := runner.Query(ctx, "hello world", WithCheckPointID("1"))
|
|
_, ok := iter.Next()
|
|
assert.True(t, ok)
|
|
event, ok := iter.Next()
|
|
assert.True(t, ok)
|
|
assert.NoError(t, event.Err)
|
|
assert.NotNil(t, event.Action.Interrupted)
|
|
_, ok = iter.Next()
|
|
assert.False(t, ok)
|
|
|
|
iter, err = runner.Resume(ctx, "1")
|
|
assert.NoError(t, err)
|
|
event, ok = iter.Next()
|
|
assert.True(t, ok)
|
|
assert.NoError(t, event.Err)
|
|
assert.NotNil(t, event.Action.Interrupted)
|
|
assert.Equal(t, 1, len(event.Action.Interrupted.InterruptContexts))
|
|
for _, ctx := range event.Action.Interrupted.InterruptContexts {
|
|
if ctx.IsRootCause {
|
|
assert.Equal(t, Address{
|
|
{Type: AddressSegmentAgent, ID: "name"},
|
|
{Type: AddressSegmentTool, ID: "myAgent", SubID: "1"},
|
|
{Type: AddressSegmentAgent, ID: "myAgent"},
|
|
}, ctx.Address)
|
|
assert.Equal(t, "interrupt again", ctx.Info)
|
|
}
|
|
}
|
|
|
|
var toolInterruptID string
|
|
for _, ctx := range event.Action.Interrupted.InterruptContexts {
|
|
if ctx.IsRootCause {
|
|
toolInterruptID = ctx.ID
|
|
break
|
|
}
|
|
}
|
|
assert.NotEmpty(t, toolInterruptID)
|
|
|
|
_, ok = iter.Next()
|
|
assert.False(t, ok)
|
|
|
|
iter, err = runner.ResumeWithParams(ctx, "1", &ResumeParams{
|
|
Targets: map[string]any{
|
|
toolInterruptID: "resume sa",
|
|
},
|
|
})
|
|
assert.NoError(t, err)
|
|
event, ok = iter.Next()
|
|
assert.True(t, ok)
|
|
assert.NoError(t, event.Err)
|
|
assert.Equal(t, event.Output.MessageOutput.Message.Content, "my agent completed with data resume sa")
|
|
event, ok = iter.Next()
|
|
assert.True(t, ok)
|
|
assert.NoError(t, event.Err)
|
|
assert.Equal(t, event.Output.MessageOutput.Message.Content, "completed")
|
|
_, ok = iter.Next()
|
|
assert.False(t, ok)
|
|
}
|
|
|
|
func newMyStore() *myStore {
|
|
return &myStore{
|
|
m: map[string][]byte{},
|
|
}
|
|
}
|
|
|
|
type myStore struct {
|
|
m map[string][]byte
|
|
}
|
|
|
|
func (m *myStore) Set(_ context.Context, key string, value []byte) error {
|
|
m.m[key] = value
|
|
return nil
|
|
}
|
|
|
|
func (m *myStore) Get(_ context.Context, key string) ([]byte, bool, error) {
|
|
v, ok := m.m[key]
|
|
return v, ok, nil
|
|
}
|
|
|
|
type myAgentOptions struct {
|
|
value string
|
|
}
|
|
|
|
func withValue(value string) AgentRunOption {
|
|
return WrapImplSpecificOptFn(func(t *myAgentOptions) {
|
|
t.value = value
|
|
})
|
|
}
|
|
|
|
type myAgent struct {
|
|
name string
|
|
runFn func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent]
|
|
resumeFn func(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*AgentEvent]
|
|
}
|
|
|
|
func (m *myAgent) Name(_ context.Context) string {
|
|
if len(m.name) > 0 {
|
|
return m.name
|
|
}
|
|
return "myAgent"
|
|
}
|
|
|
|
func (m *myAgent) Description(_ context.Context) string {
|
|
return "myAgent description"
|
|
}
|
|
|
|
func (m *myAgent) Run(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {
|
|
return m.runFn(ctx, input, options...)
|
|
}
|
|
|
|
func (m *myAgent) Resume(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] {
|
|
return m.resumeFn(ctx, info, opts...)
|
|
}
|
|
|
|
type myModel struct {
|
|
times int
|
|
messages []*schema.Message
|
|
validator func(int, []*schema.Message) bool
|
|
}
|
|
|
|
func (m *myModel) Generate(_ context.Context, input []*schema.Message, _ ...model.Option) (*schema.Message, error) {
|
|
if m.validator != nil && !m.validator(m.times, input) {
|
|
return nil, errors.New("invalid input")
|
|
}
|
|
if m.times >= len(m.messages) {
|
|
return nil, errors.New("exceeded max number of messages")
|
|
}
|
|
t := m.times
|
|
m.times++
|
|
return m.messages[t], nil
|
|
}
|
|
|
|
func (m *myModel) Stream(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) {
|
|
panic("implement me")
|
|
}
|
|
|
|
func (m *myModel) WithTools(_ []*schema.ToolInfo) (model.ToolCallingChatModel, error) {
|
|
return m, nil
|
|
}
|
|
|
|
type myTool1 struct{}
|
|
|
|
func (m *myTool1) Info(_ context.Context) (*schema.ToolInfo, error) {
|
|
return &schema.ToolInfo{
|
|
Name: "tool1",
|
|
Desc: "desc",
|
|
}, nil
|
|
}
|
|
|
|
func (m *myTool1) InvokableRun(ctx context.Context, _ string, _ ...tool.Option) (string, error) {
|
|
if wasInterrupted, _, _ := tool.GetInterruptState[any](ctx); !wasInterrupted {
|
|
return "", tool.Interrupt(ctx, nil)
|
|
}
|
|
|
|
if isResumeFlow, hasResumeData, data := tool.GetResumeContext[string](ctx); !isResumeFlow {
|
|
return "", tool.Interrupt(ctx, nil)
|
|
} else if hasResumeData {
|
|
return data, nil
|
|
}
|
|
|
|
return "result", nil
|
|
}
|
|
|
|
func TestCyclicalAgentInterrupt(t *testing.T) {
|
|
ctx := context.Background()
|
|
|
|
var agentA, agentB, agentC Agent
|
|
|
|
// agentC interrupts
|
|
agentC = &myAgent{
|
|
name: "C",
|
|
runFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {
|
|
iter, generator := NewAsyncIteratorPair[*AgentEvent]()
|
|
intAct := Interrupt(ctx, "interrupt from C")
|
|
generator.Send(intAct)
|
|
generator.Close()
|
|
return iter
|
|
},
|
|
resumeFn: func(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] {
|
|
assert.True(t, info.IsResumeTarget)
|
|
assert.NotNil(t, info.ResumeData)
|
|
assert.Equal(t, "resume C", info.ResumeData)
|
|
iter, generator := NewAsyncIteratorPair[*AgentEvent]()
|
|
generator.Send(&AgentEvent{
|
|
AgentName: "C",
|
|
Output: &AgentOutput{
|
|
MessageOutput: &MessageVariant{Message: schema.UserMessage("C completed")},
|
|
},
|
|
})
|
|
generator.Close()
|
|
return iter
|
|
},
|
|
}
|
|
|
|
// agentB transfers back to its parent A
|
|
agentB = &myAgent{
|
|
name: "B",
|
|
runFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {
|
|
iter, generator := NewAsyncIteratorPair[*AgentEvent]()
|
|
generator.Send(&AgentEvent{
|
|
AgentName: "B",
|
|
Action: &AgentAction{
|
|
TransferToAgent: &TransferToAgentAction{
|
|
DestAgentName: "A", // Transfer back to parent
|
|
},
|
|
},
|
|
})
|
|
generator.Close()
|
|
return iter
|
|
},
|
|
}
|
|
|
|
// agentA is the parent, orchestrating the A->B->A->C flow
|
|
agentA = &myAgent{
|
|
name: "A",
|
|
runFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {
|
|
runCtx := getRunCtx(ctx)
|
|
iter, generator := NewAsyncIteratorPair[*AgentEvent]()
|
|
|
|
// If the last agent was B, we are in the A->B->A path, so transfer to C.
|
|
// Otherwise, it's the first run, transfer to B.
|
|
dest := "B"
|
|
if len(runCtx.RunPath) > 1 && runCtx.RunPath[len(runCtx.RunPath)-2].agentName == "B" {
|
|
dest = "C"
|
|
}
|
|
|
|
generator.Send(&AgentEvent{
|
|
AgentName: "A",
|
|
Action: &AgentAction{
|
|
TransferToAgent: &TransferToAgentAction{
|
|
DestAgentName: dest,
|
|
},
|
|
},
|
|
})
|
|
generator.Close()
|
|
return iter
|
|
},
|
|
}
|
|
|
|
// Set up the hierarchy: A is parent of B and C.
|
|
agentA, err := SetSubAgents(ctx, agentA, []Agent{agentB, agentC})
|
|
assert.NoError(t, err)
|
|
|
|
// Run the test
|
|
runner := NewRunner(ctx, RunnerConfig{
|
|
Agent: agentA,
|
|
CheckPointStore: newMyStore(),
|
|
})
|
|
iter := runner.Query(ctx, "start", WithCheckPointID("cyclical-1"))
|
|
|
|
var events []*AgentEvent
|
|
for {
|
|
event, ok := iter.Next()
|
|
if !ok {
|
|
break
|
|
}
|
|
events = append(events, event)
|
|
}
|
|
|
|
// We expect 3 transfer events (A->B, B->A, A->C) and 1 interrupt event from C.
|
|
assert.Equal(t, 4, len(events))
|
|
|
|
interruptEvent := events[3]
|
|
assert.NotNil(t, interruptEvent.Action.Interrupted)
|
|
assert.Equal(t, "C", interruptEvent.AgentName)
|
|
|
|
// Check the interrupt context
|
|
assert.Equal(t, 1, len(interruptEvent.Action.Interrupted.InterruptContexts))
|
|
interruptCtx := interruptEvent.Action.Interrupted.InterruptContexts[0]
|
|
assert.True(t, interruptCtx.IsRootCause)
|
|
assert.Equal(t, "interrupt from C", interruptCtx.Info)
|
|
|
|
expectedAddr := Address{
|
|
{Type: AddressSegmentAgent, ID: "A"},
|
|
{Type: AddressSegmentAgent, ID: "B"},
|
|
{Type: AddressSegmentAgent, ID: "A"},
|
|
{Type: AddressSegmentAgent, ID: "C"},
|
|
}
|
|
assert.Equal(t, expectedAddr, interruptCtx.Address)
|
|
assert.NotEmpty(t, interruptCtx.ID)
|
|
|
|
// Resume the execution
|
|
iter, err = runner.ResumeWithParams(ctx, "cyclical-1", &ResumeParams{
|
|
Targets: map[string]any{
|
|
interruptCtx.ID: "resume C",
|
|
},
|
|
})
|
|
assert.NoError(t, err)
|
|
|
|
events = []*AgentEvent{}
|
|
for {
|
|
event, ok := iter.Next()
|
|
if !ok {
|
|
break
|
|
}
|
|
events = append(events, event)
|
|
}
|
|
|
|
// We expect one output event from C
|
|
assert.Equal(t, 1, len(events))
|
|
assert.Equal(t, "C completed", events[0].Output.MessageOutput.Message.Content)
|
|
}
|
|
|
|
// myStatefulTool is a tool that can interrupt and has internal state to track invocations.
|
|
|
|
type myStatefulTool struct {
|
|
name string
|
|
t *testing.T
|
|
}
|
|
|
|
func (m *myStatefulTool) Info(_ context.Context) (*schema.ToolInfo, error) {
|
|
return &schema.ToolInfo{
|
|
Name: m.name,
|
|
Desc: "desc",
|
|
}, nil
|
|
}
|
|
|
|
type myStatefulToolState struct {
|
|
InterruptCount int
|
|
}
|
|
|
|
func init() {
|
|
schema.Register[myStatefulToolState]()
|
|
}
|
|
|
|
func (m *myStatefulTool) InvokableRun(ctx context.Context, _ string, _ ...tool.Option) (string, error) {
|
|
wasInterrupted, hasState, state := tool.GetInterruptState[myStatefulToolState](ctx)
|
|
if !wasInterrupted {
|
|
return "", tool.StatefulInterrupt(ctx, fmt.Sprintf("interrupt from %s", m.name), myStatefulToolState{InterruptCount: 1})
|
|
}
|
|
|
|
isResumeFlow, hasResumeData, data := tool.GetResumeContext[string](ctx)
|
|
if !isResumeFlow || !hasResumeData {
|
|
assert.True(m.t, hasState, "tool %s should have interrupt state on resume", m.name)
|
|
return "", tool.StatefulInterrupt(ctx, fmt.Sprintf("interrupt from %s", m.name), myStatefulToolState{InterruptCount: state.InterruptCount + 1})
|
|
}
|
|
|
|
return data, nil
|
|
}
|
|
|
|
func TestChatModelParallelToolInterruptAndResume(t *testing.T) {
|
|
ctx := context.Background()
|
|
|
|
toolA := &myStatefulTool{name: "toolA", t: t}
|
|
toolB := &myStatefulTool{name: "toolB", t: t}
|
|
|
|
chatModel, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{
|
|
Name: "ParallelToolAgent",
|
|
Description: "An agent that uses parallel tools",
|
|
Model: &myModel{
|
|
messages: []*schema.Message{
|
|
// 1. First model response: call toolA and toolB in parallel
|
|
schema.AssistantMessage("", []schema.ToolCall{
|
|
{ID: "1", Function: schema.FunctionCall{Name: "toolA", Arguments: "{}"}},
|
|
{ID: "2", Function: schema.FunctionCall{Name: "toolB", Arguments: "{}"}},
|
|
}),
|
|
// 2. Second model response (after tools are resumed): call them again to check state
|
|
schema.AssistantMessage("", []schema.ToolCall{
|
|
{ID: "3", Function: schema.FunctionCall{Name: "toolA", Arguments: "{}"}},
|
|
{ID: "4", Function: schema.FunctionCall{Name: "toolB", Arguments: "{}"}},
|
|
}),
|
|
// 3. Final completion
|
|
schema.AssistantMessage("all done", nil),
|
|
},
|
|
},
|
|
ToolsConfig: ToolsConfig{
|
|
ToolsNodeConfig: compose.ToolsNodeConfig{
|
|
Tools: []tool.BaseTool{toolA, toolB},
|
|
},
|
|
},
|
|
})
|
|
assert.NoError(t, err)
|
|
|
|
runner := NewRunner(ctx, RunnerConfig{
|
|
Agent: chatModel,
|
|
CheckPointStore: newMyStore(),
|
|
})
|
|
|
|
// 1. Initial query -> parallel interrupt from toolA and toolB
|
|
iter := runner.Query(ctx, "start", WithCheckPointID("parallel-tool-test-1"))
|
|
normalEvents, interruptEvent := consumeUntilInterrupt(iter)
|
|
|
|
assert.Equal(t, 1, len(normalEvents))
|
|
assert.NotNil(t, interruptEvent)
|
|
assert.Equal(t, 2, len(interruptEvent.Action.Interrupted.InterruptContexts),
|
|
"should have 2 interrupts")
|
|
|
|
var toolAInterruptID, toolBInterruptID string
|
|
for _, info := range interruptEvent.Action.Interrupted.InterruptContexts {
|
|
switch info.Info {
|
|
case "interrupt from toolA":
|
|
toolAInterruptID = info.ID
|
|
assert.True(t, info.IsRootCause)
|
|
case "interrupt from toolB":
|
|
toolBInterruptID = info.ID
|
|
assert.True(t, info.IsRootCause)
|
|
}
|
|
}
|
|
assert.NotEmpty(t, toolAInterruptID)
|
|
assert.NotEmpty(t, toolBInterruptID)
|
|
|
|
// 2. Resume, targeting only toolA. toolB should re-interrupt.
|
|
iter, err = runner.ResumeWithParams(ctx, "parallel-tool-test-1", &ResumeParams{
|
|
Targets: map[string]any{
|
|
toolAInterruptID: "toolA resumed",
|
|
},
|
|
})
|
|
assert.NoError(t, err)
|
|
_, interruptEvent = consumeUntilInterrupt(iter)
|
|
|
|
assert.NotNil(t, interruptEvent, "expected a re-interrupt from toolB")
|
|
assert.Equal(t, 1, len(interruptEvent.Action.Interrupted.InterruptContexts),
|
|
"should have 1 remaining interrupts")
|
|
|
|
var rootCause *InterruptCtx
|
|
for _, info := range interruptEvent.Action.Interrupted.InterruptContexts {
|
|
if info.IsRootCause {
|
|
rootCause = info
|
|
break
|
|
}
|
|
}
|
|
|
|
if rootCause == nil {
|
|
t.Fatal("expected a root cause interrupt from toolB")
|
|
}
|
|
assert.Equal(t, "interrupt from toolB", rootCause.Info)
|
|
toolBReInterruptID := rootCause.ID
|
|
|
|
// 3. Resume the re-interrupted toolB. The agent should then call the tools again.
|
|
iter, err = runner.ResumeWithParams(ctx, "parallel-tool-test-1", &ResumeParams{
|
|
Targets: map[string]any{
|
|
toolBReInterruptID: "toolB resumed",
|
|
},
|
|
})
|
|
assert.NoError(t, err)
|
|
|
|
// 4. Consume all final events. The internal assertions in the tools will check the wasInterrupted flag.
|
|
// We expect to see the results of the second tool calls, and then the final agent completion.
|
|
finalEvents, interruptEvent := consumeUntilInterrupt(iter)
|
|
assert.Equal(t, 2, len(finalEvents))
|
|
assert.NotNil(t, interruptEvent)
|
|
}
|
|
|
|
// TestNestedChatModelAgentWithAgentTool verifies that the shouldFire method correctly prevents
|
|
// duplicate event firing in nested ChatModelAgent scenarios (ChatModelAgent -> AgentTool -> ChatModelAgent).
|
|
// This ensures that only the inner agent's cbHandler fires, not the outer agent's.
|
|
func TestNestedChatModelAgentWithAgentTool(t *testing.T) {
|
|
ctx := context.Background()
|
|
|
|
// Create an interruptible tool for the inner agent
|
|
innerTool := &myStatefulTool{name: "innerTool", t: t}
|
|
|
|
// Create the inner ChatModelAgent that will be wrapped by AgentTool
|
|
innerAgent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{
|
|
Name: "InnerAgent",
|
|
Description: "Inner agent with interruptible tool",
|
|
Model: &myModel{
|
|
messages: []*schema.Message{
|
|
schema.AssistantMessage("", []schema.ToolCall{
|
|
{ID: "1", Function: schema.FunctionCall{Name: "innerTool", Arguments: "{}"}},
|
|
}),
|
|
schema.AssistantMessage("inner agent completed", nil),
|
|
},
|
|
},
|
|
ToolsConfig: ToolsConfig{
|
|
ToolsNodeConfig: compose.ToolsNodeConfig{
|
|
Tools: []tool.BaseTool{innerTool},
|
|
},
|
|
},
|
|
})
|
|
assert.NoError(t, err)
|
|
|
|
// Wrap the inner agent in an AgentTool
|
|
agentTool := NewAgentTool(ctx, innerAgent)
|
|
|
|
// Create the outer ChatModelAgent that uses the AgentTool
|
|
outerAgent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{
|
|
Name: "OuterAgent",
|
|
Description: "Outer agent with AgentTool containing inner agent",
|
|
Model: &myModel{
|
|
messages: []*schema.Message{
|
|
schema.AssistantMessage("", []schema.ToolCall{
|
|
{ID: "1", Function: schema.FunctionCall{Name: "InnerAgent", Arguments: "{}"}},
|
|
}),
|
|
schema.AssistantMessage("outer agent completed", nil),
|
|
},
|
|
},
|
|
ToolsConfig: ToolsConfig{
|
|
ToolsNodeConfig: compose.ToolsNodeConfig{
|
|
Tools: []tool.BaseTool{agentTool},
|
|
},
|
|
},
|
|
})
|
|
assert.NoError(t, err)
|
|
|
|
runner := NewRunner(ctx, RunnerConfig{
|
|
Agent: outerAgent,
|
|
CheckPointStore: newMyStore(),
|
|
})
|
|
|
|
// Run the query - this should trigger the nested agent structure
|
|
iter := runner.Query(ctx, "start", WithCheckPointID("nested-agent-test-1"))
|
|
|
|
// Collect all events to verify no duplicates
|
|
var allEvents []*AgentEvent
|
|
var interruptEvent *AgentEvent
|
|
|
|
for {
|
|
event, ok := iter.Next()
|
|
if !ok {
|
|
break
|
|
}
|
|
|
|
if event.Action != nil && event.Action.Interrupted != nil {
|
|
assert.Nil(t, interruptEvent)
|
|
interruptEvent = event
|
|
}
|
|
|
|
allEvents = append(allEvents, event)
|
|
}
|
|
|
|
if interruptEvent == nil {
|
|
t.Fatal("expected an interrupt event")
|
|
}
|
|
|
|
// Verify we got exactly one interrupt event (not duplicated)
|
|
assert.NotNil(t, interruptEvent, "should have an interrupt event")
|
|
assert.Equal(t, 1, len(interruptEvent.Action.Interrupted.InterruptContexts),
|
|
"should have exactly one interrupt context")
|
|
|
|
// Verify the interrupt comes from the inner tool, not duplicated
|
|
interruptCtx := interruptEvent.Action.Interrupted.InterruptContexts[0]
|
|
assert.True(t, interruptCtx.IsRootCause, "interrupt should be root cause")
|
|
assert.Equal(t, "interrupt from innerTool", interruptCtx.Info)
|
|
|
|
// Verify the address path shows the correct nested structure
|
|
expectedAddress := Address{
|
|
{Type: AddressSegmentAgent, ID: "OuterAgent"},
|
|
{Type: AddressSegmentTool, ID: "InnerAgent", SubID: "1"},
|
|
{Type: AddressSegmentAgent, ID: "InnerAgent"},
|
|
{Type: AddressSegmentTool, ID: "innerTool", SubID: "1"},
|
|
}
|
|
assert.Equal(t, expectedAddress, interruptCtx.Address,
|
|
"interrupt address should show correct nested structure")
|
|
|
|
// Verify no duplicate events by checking agent names in events
|
|
var agentNames []string
|
|
for _, event := range allEvents {
|
|
if event.AgentName != "" {
|
|
agentNames = append(agentNames, event.AgentName)
|
|
}
|
|
}
|
|
|
|
// Should only have events from the outer agent (the inner agent's events should be handled
|
|
// by the AgentTool and not duplicated by the outer agent's cbHandler)
|
|
for _, name := range agentNames {
|
|
assert.Equal(t, "OuterAgent", name,
|
|
"all events should come from OuterAgent, not duplicated from InnerAgent")
|
|
}
|
|
|
|
// Now resume the interrupt
|
|
interruptID := interruptCtx.ID
|
|
iter, err = runner.ResumeWithParams(ctx, "nested-agent-test-1", &ResumeParams{
|
|
Targets: map[string]any{
|
|
interruptID: "resume inner tool",
|
|
},
|
|
})
|
|
assert.NoError(t, err)
|
|
|
|
// Collect final events after resume
|
|
var finalEvents []*AgentEvent
|
|
for {
|
|
event, ok := iter.Next()
|
|
if !ok {
|
|
break
|
|
}
|
|
finalEvents = append(finalEvents, event)
|
|
}
|
|
|
|
// Verify completion events
|
|
assert.Greater(t, len(finalEvents), 0, "should have completion events after resume")
|
|
|
|
// Check that we get the expected completion messages
|
|
var foundInnerCompletion, foundOuterCompletion bool
|
|
for _, event := range finalEvents {
|
|
if event.Output != nil && event.Output.MessageOutput != nil {
|
|
if event.Output.MessageOutput.Message != nil {
|
|
content := event.Output.MessageOutput.Message.Content
|
|
switch content {
|
|
case "inner agent completed":
|
|
foundInnerCompletion = true
|
|
case "outer agent completed":
|
|
foundOuterCompletion = true
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
assert.True(t, foundInnerCompletion, "should have inner agent completion")
|
|
assert.True(t, foundOuterCompletion, "should have outer agent completion")
|
|
}
|
|
|
|
// consumeUntilInterrupt consumes events from the iterator until an interrupt is found or it's exhausted.
|
|
func consumeUntilInterrupt(iter *AsyncIterator[*AgentEvent]) (normalEvents []*AgentEvent, interruptEvent *AgentEvent) {
|
|
for {
|
|
event, ok := iter.Next()
|
|
if !ok {
|
|
break
|
|
}
|
|
if event.Action != nil && event.Action.Interrupted != nil {
|
|
interruptEvent = event
|
|
continue
|
|
}
|
|
normalEvents = append(normalEvents, event)
|
|
}
|
|
return
|
|
}
|
|
|
|
type returnDirectlyTool struct {
|
|
name string
|
|
}
|
|
|
|
func (t *returnDirectlyTool) Info(_ context.Context) (*schema.ToolInfo, error) {
|
|
return &schema.ToolInfo{
|
|
Name: t.name,
|
|
Desc: "A tool that returns directly",
|
|
}, nil
|
|
}
|
|
|
|
func (t *returnDirectlyTool) InvokableRun(_ context.Context, _ string, _ ...tool.Option) (string, error) {
|
|
return "return directly result", nil
|
|
}
|
|
|
|
type interruptingTool struct {
|
|
name string
|
|
}
|
|
|
|
func (i *interruptingTool) Info(_ context.Context) (*schema.ToolInfo, error) {
|
|
return &schema.ToolInfo{
|
|
Name: i.name,
|
|
Desc: "A tool that interrupts",
|
|
}, nil
|
|
}
|
|
|
|
func (i *interruptingTool) InvokableRun(ctx context.Context, _ string, _ ...tool.Option) (string, error) {
|
|
if wasInterrupted, _, _ := compose.GetInterruptState[any](ctx); !wasInterrupted {
|
|
return "", compose.Interrupt(ctx, "interrupt data")
|
|
}
|
|
|
|
if isResumeFlow, hasResumeData, data := compose.GetResumeContext[string](ctx); isResumeFlow && hasResumeData {
|
|
return data, nil
|
|
}
|
|
|
|
return "resumed without data", nil
|
|
}
|
|
|
|
type twoToolCallModel struct {
|
|
returnDirectlyToolName string
|
|
interruptingToolName string
|
|
callCount int
|
|
receivedTools []*schema.ToolInfo
|
|
mu sync.Mutex
|
|
}
|
|
|
|
func (m *twoToolCallModel) Generate(_ context.Context, _ []*schema.Message, opts ...model.Option) (*schema.Message, error) {
|
|
m.mu.Lock()
|
|
m.callCount++
|
|
callNum := m.callCount
|
|
options := model.GetCommonOptions(&model.Options{}, opts...)
|
|
if options.Tools != nil {
|
|
m.receivedTools = options.Tools
|
|
}
|
|
m.mu.Unlock()
|
|
|
|
if callNum == 1 {
|
|
return &schema.Message{
|
|
Role: schema.Assistant,
|
|
Content: "",
|
|
ToolCalls: []schema.ToolCall{
|
|
{
|
|
ID: "call_return_directly",
|
|
Type: "function",
|
|
Function: schema.FunctionCall{
|
|
Name: m.returnDirectlyToolName,
|
|
Arguments: "{}",
|
|
},
|
|
},
|
|
{
|
|
ID: "call_interrupting",
|
|
Type: "function",
|
|
Function: schema.FunctionCall{
|
|
Name: m.interruptingToolName,
|
|
Arguments: "{}",
|
|
},
|
|
},
|
|
},
|
|
}, nil
|
|
}
|
|
return schema.AssistantMessage("final response", nil), nil
|
|
}
|
|
|
|
func (m *twoToolCallModel) Stream(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) {
|
|
panic("not implemented")
|
|
}
|
|
|
|
func (m *twoToolCallModel) WithTools(_ []*schema.ToolInfo) (model.ToolCallingChatModel, error) {
|
|
return m, nil
|
|
}
|
|
|
|
func (m *twoToolCallModel) GetReceivedTools() []*schema.ToolInfo {
|
|
m.mu.Lock()
|
|
defer m.mu.Unlock()
|
|
return m.receivedTools
|
|
}
|
|
|
|
type dynamicTool struct {
|
|
name string
|
|
}
|
|
|
|
func (t *dynamicTool) Info(_ context.Context) (*schema.ToolInfo, error) {
|
|
return &schema.ToolInfo{
|
|
Name: t.name,
|
|
Desc: "A dynamically added tool",
|
|
}, nil
|
|
}
|
|
|
|
func (t *dynamicTool) InvokableRun(_ context.Context, _ string, _ ...tool.Option) (string, error) {
|
|
return "dynamic tool result", nil
|
|
}
|
|
|
|
func TestReturnDirectlyEventSentAfterResume(t *testing.T) {
|
|
ctx := context.Background()
|
|
|
|
returnDirectlyToolName := "return_directly_tool"
|
|
interruptingToolName := "interrupting_tool"
|
|
dynamicToolName := "dynamic_tool"
|
|
|
|
mdl := &twoToolCallModel{
|
|
returnDirectlyToolName: returnDirectlyToolName,
|
|
interruptingToolName: interruptingToolName,
|
|
}
|
|
|
|
agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{
|
|
Name: "TestAgent",
|
|
Description: "Test agent for return directly + interrupt",
|
|
Model: mdl,
|
|
ToolsConfig: ToolsConfig{
|
|
ToolsNodeConfig: compose.ToolsNodeConfig{
|
|
Tools: []tool.BaseTool{
|
|
&returnDirectlyTool{name: returnDirectlyToolName},
|
|
&interruptingTool{name: interruptingToolName},
|
|
},
|
|
},
|
|
ReturnDirectly: map[string]bool{
|
|
returnDirectlyToolName: true,
|
|
},
|
|
},
|
|
Handlers: []ChatModelAgentMiddleware{
|
|
&interruptTestToolsHandler{tools: []tool.BaseTool{&dynamicTool{name: dynamicToolName}}},
|
|
},
|
|
})
|
|
assert.NoError(t, err)
|
|
|
|
store := newMyStore()
|
|
runner := NewRunner(ctx, RunnerConfig{
|
|
Agent: agent,
|
|
EnableStreaming: false,
|
|
CheckPointStore: store,
|
|
})
|
|
|
|
iter := runner.Query(ctx, "test input", WithCheckPointID("test_checkpoint"))
|
|
|
|
var interruptEvent *AgentEvent
|
|
for {
|
|
event, ok := iter.Next()
|
|
if !ok {
|
|
break
|
|
}
|
|
if event.Action != nil && event.Action.Interrupted != nil {
|
|
interruptEvent = event
|
|
}
|
|
}
|
|
|
|
assert.NotNil(t, interruptEvent, "Should have an interrupt event")
|
|
assert.NotEmpty(t, interruptEvent.Action.Interrupted.InterruptContexts)
|
|
|
|
receivedToolsBeforeResume := mdl.GetReceivedTools()
|
|
var hasDynamicToolBeforeResume bool
|
|
for _, ti := range receivedToolsBeforeResume {
|
|
if ti.Name == dynamicToolName {
|
|
hasDynamicToolBeforeResume = true
|
|
}
|
|
}
|
|
assert.True(t, hasDynamicToolBeforeResume, "Dynamic tool should be in tool list before interrupt")
|
|
|
|
interruptID := interruptEvent.Action.Interrupted.InterruptContexts[0].ID
|
|
resumeIter, err := runner.ResumeWithParams(ctx, "test_checkpoint", &ResumeParams{
|
|
Targets: map[string]any{
|
|
interruptID: "resume data",
|
|
},
|
|
})
|
|
assert.NoError(t, err)
|
|
|
|
var resumeEvents []*AgentEvent
|
|
for {
|
|
event, ok := resumeIter.Next()
|
|
if !ok {
|
|
break
|
|
}
|
|
resumeEvents = append(resumeEvents, event)
|
|
}
|
|
|
|
var hasReturnDirectlyEvent bool
|
|
for _, e := range resumeEvents {
|
|
if e.Output != nil && e.Output.MessageOutput != nil {
|
|
if e.Output.MessageOutput.Role == schema.Tool && e.Output.MessageOutput.ToolName == returnDirectlyToolName {
|
|
hasReturnDirectlyEvent = true
|
|
}
|
|
}
|
|
}
|
|
assert.True(t, hasReturnDirectlyEvent, "ReturnDirectlyEvent should be sent after resume")
|
|
|
|
receivedToolsAfterResume := mdl.GetReceivedTools()
|
|
var hasDynamicToolAfterResume bool
|
|
for _, ti := range receivedToolsAfterResume {
|
|
if ti.Name == dynamicToolName {
|
|
hasDynamicToolAfterResume = true
|
|
}
|
|
}
|
|
assert.True(t, hasDynamicToolAfterResume, "Dynamic tool should be in tool list after resume (bc.toolUpdated path)")
|
|
}
|
|
|
|
// streamErrorThenToolCallModel simulates a model that:
|
|
// - On the first Stream call: emits several good chunks then an error (triggering retry)
|
|
// - On the second Stream call (retry): returns a tool call message (success)
|
|
type streamErrorThenToolCallModel struct {
|
|
callCount int32
|
|
toolCallName string
|
|
}
|
|
|
|
func (m *streamErrorThenToolCallModel) Generate(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) {
|
|
return schema.AssistantMessage("final answer", nil), nil
|
|
}
|
|
|
|
func (m *streamErrorThenToolCallModel) Stream(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) {
|
|
count := atomic.AddInt32(&m.callCount, 1)
|
|
|
|
sr, sw := schema.Pipe[*schema.Message](10)
|
|
go func() {
|
|
defer sw.Close()
|
|
if count == 1 {
|
|
// First call: emit good chunks then error
|
|
sw.Send(schema.AssistantMessage("chunk1", nil), nil)
|
|
sw.Send(schema.AssistantMessage("chunk2", nil), nil)
|
|
sw.Send(schema.AssistantMessage("chunk3", nil), nil)
|
|
sw.Send(nil, errRetryAble)
|
|
return
|
|
}
|
|
// Second call (retry): return tool call
|
|
sw.Send(schema.AssistantMessage("", []schema.ToolCall{{
|
|
ID: "call-1",
|
|
Function: schema.FunctionCall{Name: m.toolCallName, Arguments: "{}"},
|
|
}}), nil)
|
|
}()
|
|
return sr, nil
|
|
}
|
|
|
|
func (m *streamErrorThenToolCallModel) WithTools(_ []*schema.ToolInfo) (model.ToolCallingChatModel, error) {
|
|
return m, nil
|
|
}
|
|
|
|
// TestStreamRetryThenToolInterruptCheckpoint reproduces a bug where:
|
|
// 1. ChatModelAgent with ModelRetryConfig.MaxRetries = 2
|
|
// 2. First model Stream call emits good chunks then a retryable error
|
|
// 3. Retry succeeds, model returns a tool call
|
|
// 4. The tool triggers an interrupt, causing the Runner to save a checkpoint
|
|
// 5. The checkpoint save fails because the first (failed) model call's stream event
|
|
// is in the session, and when MessageVariant.GobEncode consumes the stream,
|
|
// it hits the error chunk and returns an encoding error.
|
|
func TestStreamRetryThenToolInterruptCheckpoint(t *testing.T) {
|
|
ctx := context.Background()
|
|
|
|
interruptToolName := "interrupt_tool"
|
|
mdl := &streamErrorThenToolCallModel{toolCallName: interruptToolName}
|
|
|
|
interruptTool := &interruptingTool{name: interruptToolName}
|
|
|
|
agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{
|
|
Name: "RetryInterruptAgent",
|
|
Description: "Agent that retries model then tool interrupts",
|
|
Instruction: "You are a test agent.",
|
|
Model: mdl,
|
|
ModelRetryConfig: &ModelRetryConfig{
|
|
MaxRetries: 2,
|
|
IsRetryAble: func(ctx context.Context, err error) bool { return errors.Is(err, errRetryAble) },
|
|
},
|
|
ToolsConfig: ToolsConfig{
|
|
ToolsNodeConfig: compose.ToolsNodeConfig{
|
|
Tools: []tool.BaseTool{interruptTool},
|
|
},
|
|
},
|
|
})
|
|
assert.NoError(t, err)
|
|
|
|
store := newMyStore()
|
|
runner := NewRunner(ctx, RunnerConfig{
|
|
Agent: agent,
|
|
EnableStreaming: true,
|
|
CheckPointStore: store,
|
|
})
|
|
|
|
iter := runner.Run(ctx, []Message{schema.UserMessage("test query")}, WithCheckPointID("retry_interrupt_ckpt"))
|
|
|
|
var events []*AgentEvent
|
|
var checkpointErr error
|
|
for {
|
|
event, ok := iter.Next()
|
|
if !ok {
|
|
break
|
|
}
|
|
events = append(events, event)
|
|
if event.Err != nil {
|
|
checkpointErr = event.Err
|
|
t.Logf("event error: %v", event.Err)
|
|
}
|
|
}
|
|
|
|
// The bug: checkpoint save fails because the failed stream's error chunk
|
|
// is encountered during gob encoding of the session events.
|
|
// If the bug is fixed, checkpointErr should be nil and we should see an interrupt event.
|
|
assert.NoError(t, checkpointErr, "checkpoint save should not fail due to failed stream's error in session")
|
|
|
|
var hasInterrupt bool
|
|
for _, event := range events {
|
|
if event.Action != nil && event.Action.Interrupted != nil {
|
|
hasInterrupt = true
|
|
}
|
|
}
|
|
assert.True(t, hasInterrupt, "should receive an interrupt event from the tool")
|
|
|
|
// Verify the model was called twice (first call errored, second succeeded)
|
|
assert.Equal(t, int32(2), atomic.LoadInt32(&mdl.callCount), "model should be called exactly twice (1 failure + 1 retry)")
|
|
}
|