diff --git a/agent/agent.go b/agent/agent.go index 037e1c3fd..7cb605539 100644 --- a/agent/agent.go +++ b/agent/agent.go @@ -18,6 +18,7 @@ import ( "context" "fmt" "iter" + "sync/atomic" "go.opentelemetry.io/otel/trace" "google.golang.org/genai" @@ -368,7 +369,7 @@ type invocationContext struct { isolationScope string userContent *genai.Content runConfig *RunConfig - endInvocation bool + endInvocation *atomic.Bool } // Apply implements [InvocationContext]. @@ -432,10 +433,16 @@ func (c *invocationContext) RunConfig() *RunConfig { } func (c *invocationContext) EndInvocation() { - c.endInvocation = true + if c.endInvocation != nil { + c.endInvocation.Store(true) + } } func (c *invocationContext) Ended() bool { + return c.endInvocation != nil && c.endInvocation.Load() +} + +func (c *invocationContext) EndInvocationPtr() *atomic.Bool { return c.endInvocation } diff --git a/agent/agent_test.go b/agent/agent_test.go index 7a498d30c..e98eec27c 100644 --- a/agent/agent_test.go +++ b/agent/agent_test.go @@ -17,6 +17,7 @@ package agent import ( "context" "iter" + "sync/atomic" "testing" "github.com/google/go-cmp/cmp" @@ -122,9 +123,10 @@ func TestAgentCallbacks(t *testing.T) { } ctx := &invocationContext{ - Context: t.Context(), - agent: testAgent, - session: &mockSession{sessionID: "test-session"}, + Context: t.Context(), + agent: testAgent, + session: &mockSession{sessionID: "test-session"}, + endInvocation: new(atomic.Bool), } var gotEvents []*session.Event for event, err := range testAgent.Run(ctx) { @@ -168,10 +170,13 @@ func TestEndInvocation_EndsBeforeMainCall(t *testing.T) { t.Fatalf("failed to create agent: %v", err) } + endedBool := new(atomic.Bool) + endedBool.Store(true) + ctx := &invocationContext{ Context: t.Context(), agent: testAgent, - endInvocation: true, + endInvocation: endedBool, session: &mockSession{sessionID: "test-session"}, } for _, err := range testAgent.Run(ctx) { @@ -204,9 +209,10 @@ func TestEndInvocation_EndsAfterMainCall(t *testing.T) { } ctx := &invocationContext{ - Context: t.Context(), - agent: testAgent, - session: &mockSession{sessionID: "test-session"}, + Context: t.Context(), + agent: testAgent, + session: &mockSession{sessionID: "test-session"}, + endInvocation: new(atomic.Bool), } var gotEvents []*session.Event for event, err := range testAgent.Run(ctx) { @@ -262,9 +268,10 @@ type testKey struct{} func TestWithContext(t *testing.T) { baseCtx := t.Context() inv := &invocationContext{ - Context: baseCtx, - invocationID: "test", - branch: "branch", + Context: baseCtx, + invocationID: "test", + branch: "branch", + endInvocation: new(atomic.Bool), } key := testKey{} @@ -274,7 +281,13 @@ func TestWithContext(t *testing.T) { if got.Value(key) != val { t.Errorf("WithContext() did not update context") } - if diff := cmp.Diff(inv, got, cmp.AllowUnexported(invocationContext{}), cmpopts.IgnoreFields(invocationContext{}, "Context")); diff != "" { + atomicBoolComparer := cmp.Comparer(func(a, b *atomic.Bool) bool { + if a == nil || b == nil { + return a == b + } + return a.Load() == b.Load() + }) + if diff := cmp.Diff(inv, got, cmp.AllowUnexported(invocationContext{}), cmpopts.IgnoreFields(invocationContext{}, "Context"), atomicBoolComparer); diff != "" { t.Errorf("WithContext() params mismatch (-want +got):\n%s", diff) } } @@ -351,3 +364,75 @@ func TestFindAgent(t *testing.T) { }) } } + +func TestSubAgentCallbackEndInvocationPropagated(t *testing.T) { + t.Parallel() + + subAgentCalled := false + + subAgent, err := New(Config{ + Name: "sub_agent", + BeforeAgentCallbacks: []BeforeAgentCallback{ + func(ctx Context) (*genai.Content, error) { + subAgentCalled = true + return genai.NewContentFromText("sub agent ended", genai.RoleModel), nil + }, + }, + Run: func(InvocationContext) iter.Seq2[*session.Event, error] { + return func(func(*session.Event, error) bool) {} + }, + }) + if err != nil { + t.Fatalf("failed to create sub_agent: %v", err) + } + + parentRunCalled := false + parentEndedAfterSubAgent := false + + parentAgent, err := New(Config{ + Name: "parent_agent", + SubAgents: []Agent{subAgent}, + Run: func(ctx InvocationContext) iter.Seq2[*session.Event, error] { + parentRunCalled = true + return func(yield func(*session.Event, error) bool) { + for _, sa := range ctx.Agent().SubAgents() { + for event, err := range sa.Run(ctx) { + if !yield(event, err) { + return + } + } + } + // Demonstrates that when sub-agent BeforeAgentCallback returns non-nil content + // (which calls ctx.EndInvocation() on the sub-agent context), the endInvocation + // flag propagates to the parent context via *atomic.Bool. + parentEndedAfterSubAgent = ctx.Ended() + } + }, + }) + if err != nil { + t.Fatalf("failed to create parent_agent: %v", err) + } + + parentCtx := &invocationContext{ + Context: t.Context(), + agent: parentAgent, + session: &mockSession{sessionID: "test-session"}, + endInvocation: new(atomic.Bool), + } + + for _, err := range parentAgent.Run(parentCtx) { + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + } + + if !parentRunCalled { + t.Error("expected parent run to be called") + } + if !subAgentCalled { + t.Error("expected sub_agent BeforeAgentCallback to be called") + } + if !parentEndedAfterSubAgent { + t.Error("expected parentEndedAfterSubAgent to be true after sub-agent callback ended invocation") + } +} diff --git a/agent/common_context.go b/agent/common_context.go index a78268178..687ba39f4 100644 --- a/agent/common_context.go +++ b/agent/common_context.go @@ -18,6 +18,7 @@ import ( "context" "fmt" "iter" + "sync/atomic" "time" "google.golang.org/genai" @@ -220,6 +221,13 @@ func (c *commonContext) Ended() bool { return c.invocationContext.Ended() } +func (c *commonContext) EndInvocationPtr() *atomic.Bool { + if carrier, ok := c.invocationContext.(interface{ EndInvocationPtr() *atomic.Bool }); ok { + return carrier.EndInvocationPtr() + } + return nil +} + // IsolationScope implements [InvocationContext]. func (c *commonContext) IsolationScope() string { return c.invocationContext.IsolationScope() diff --git a/agent/workflowagents/loopagent/agent.go b/agent/workflowagents/loopagent/agent.go index 9ecb358a0..5354b9a3a 100644 --- a/agent/workflowagents/loopagent/agent.go +++ b/agent/workflowagents/loopagent/agent.go @@ -79,6 +79,9 @@ func (a *loopAgent) Run(ctx agent.InvocationContext) iter.Seq2[*session.Event, e for { shouldExit := false for _, subAgent := range ctx.Agent().SubAgents() { + if ctx.Ended() { + return + } for event, err := range subAgent.Run(ctx) { // TODO: ensure consistency -- if there's an error, return and close iterator, verify everywhere in ADK. if !yield(event, err) { @@ -89,7 +92,7 @@ func (a *loopAgent) Run(ctx agent.InvocationContext) iter.Seq2[*session.Event, e shouldExit = true } } - if shouldExit { + if shouldExit || ctx.Ended() { return } } diff --git a/agent/workflowagents/loopagent/agent_test.go b/agent/workflowagents/loopagent/agent_test.go index eac99bd06..6be43d128 100644 --- a/agent/workflowagents/loopagent/agent_test.go +++ b/agent/workflowagents/loopagent/agent_test.go @@ -376,3 +376,81 @@ func (f *FakeLLM) GenerateContent(ctx context.Context, req *model.LLMRequest, st } } } + +func TestLoopAgent_SubAgentEndInvocationPropagated(t *testing.T) { + subAgent1Called := false + subAgent2Called := false + + subAgent1, err := agent.New(agent.Config{ + Name: "sub_agent_1", + Run: func(ctx agent.InvocationContext) iter.Seq2[*session.Event, error] { + return func(yield func(*session.Event, error) bool) { + subAgent1Called = true + ctx.EndInvocation() + yield(&session.Event{Author: "sub_agent_1"}, nil) + } + }, + }) + if err != nil { + t.Fatal(err) + } + + subAgent2, err := agent.New(agent.Config{ + Name: "sub_agent_2", + Run: func(ctx agent.InvocationContext) iter.Seq2[*session.Event, error] { + return func(yield func(*session.Event, error) bool) { + subAgent2Called = true + yield(&session.Event{Author: "sub_agent_2"}, nil) + } + }, + }) + if err != nil { + t.Fatal(err) + } + + loopAgent, err := loopagent.New(loopagent.Config{ + MaxIterations: 2, + AgentConfig: agent.Config{ + Name: "loop_agent", + SubAgents: []agent.Agent{subAgent1, subAgent2}, + }, + }) + if err != nil { + t.Fatal(err) + } + + sessionService := session.InMemoryService() + ctx := t.Context() + _, err = sessionService.Create(ctx, &session.CreateRequest{ + AppName: "test_app", + UserID: "user_id", + SessionID: "session_id", + }) + if err != nil { + t.Fatal(err) + } + + r, err := runner.New(runner.Config{ + AppName: "test_app", + Agent: loopAgent, + SessionService: sessionService, + }) + if err != nil { + t.Fatal(err) + } + + for _, err := range r.Run(ctx, "user_id", "session_id", genai.NewContentFromText("hello", genai.RoleUser), agent.RunConfig{}) { + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + } + + if !subAgent1Called { + t.Error("expected sub_agent_1 to be called") + } + // Demonstrates that calling EndInvocation() in sub_agent_1 + // ends the invocation across the shared *atomic.Bool context state, preventing sub_agent_2 from running. + if subAgent2Called { + t.Error("expected sub_agent_2 NOT to be called because sub_agent_1 ended the invocation") + } +} diff --git a/agent/workflowagents/parallelagent/agent.go b/agent/workflowagents/parallelagent/agent.go index 5d99e6a42..aa28b4bf0 100644 --- a/agent/workflowagents/parallelagent/agent.go +++ b/agent/workflowagents/parallelagent/agent.go @@ -18,6 +18,7 @@ package parallelagent import ( "fmt" "iter" + "sync/atomic" "golang.org/x/sync/errgroup" @@ -73,28 +74,40 @@ func run(ctx agent.InvocationContext) iter.Seq2[*session.Event, error] { resultsChan = make(chan result) ) + var parentEndInv *atomic.Bool + if carrier, ok := ctx.(interface{ EndInvocationPtr() *atomic.Bool }); ok { + parentEndInv = carrier.EndInvocationPtr() + } + var anySiblingEnded atomic.Bool + for _, sa := range ctx.Agent().SubAgents() { branch := fmt.Sprintf("%s.%s", curAgent.Name(), sa.Name()) if ctx.Branch() != "" { branch = fmt.Sprintf("%s.%s", ctx.Branch(), branch) } + siblingEndInv := new(atomic.Bool) subAgent := sa errGroup.Go(func() error { subCtx := icontext.NewInvocationContext(errGroupCtx, icontext.InvocationContextParams{ - Artifacts: ctx.Artifacts(), - Memory: ctx.Memory(), - Session: ctx.Session(), - Branch: branch, - Agent: subAgent, - UserContent: ctx.UserContent(), - RunConfig: ctx.RunConfig(), - InvocationID: ctx.InvocationID(), + Artifacts: ctx.Artifacts(), + Memory: ctx.Memory(), + Session: ctx.Session(), + Branch: branch, + Agent: subAgent, + UserContent: ctx.UserContent(), + RunConfig: ctx.RunConfig(), + EndInvocation: siblingEndInv, + InvocationID: ctx.InvocationID(), }) if err := runSubAgent(subCtx, subAgent, resultsChan, doneChan); err != nil { return fmt.Errorf("failed to run sub-agent %q: %w", subAgent.Name(), err) } + if siblingEndInv.Load() { + anySiblingEnded.Store(true) + } + return nil }) } @@ -106,6 +119,9 @@ func run(ctx agent.InvocationContext) iter.Seq2[*session.Event, error] { case <-doneChan: } } + if anySiblingEnded.Load() && parentEndInv != nil { + parentEndInv.Store(true) + } close(resultsChan) }() diff --git a/agent/workflowagents/parallelagent/agent_test.go b/agent/workflowagents/parallelagent/agent_test.go index 4f24d87b6..97451eb21 100644 --- a/agent/workflowagents/parallelagent/agent_test.go +++ b/agent/workflowagents/parallelagent/agent_test.go @@ -33,6 +33,7 @@ import ( "google.golang.org/adk/v2/agent/llmagent" "google.golang.org/adk/v2/agent/workflowagents/loopagent" "google.golang.org/adk/v2/agent/workflowagents/parallelagent" + icontext "google.golang.org/adk/v2/internal/context" "google.golang.org/adk/v2/internal/httprr" "google.golang.org/adk/v2/internal/testutil" "google.golang.org/adk/v2/model" @@ -535,3 +536,67 @@ func TestParallelAgent_StateSync(t *testing.T) { t.Fatalf("expected state value 'test_value', got %v", gotValue) } } + +func TestParallelAgent_IsolatedEndInvocation(t *testing.T) { + enderAgent, err := agent.New(agent.Config{ + Name: "ender", + Run: func(ctx agent.InvocationContext) iter.Seq2[*session.Event, error] { + return func(yield func(*session.Event, error) bool) { + ctx.EndInvocation() + yield(&session.Event{Author: "ender"}, nil) + } + }, + }) + if err != nil { + t.Fatal(err) + } + + workerAgent, err := agent.New(agent.Config{ + Name: "worker", + Run: func(ctx agent.InvocationContext) iter.Seq2[*session.Event, error] { + return func(yield func(*session.Event, error) bool) { + time.Sleep(1000 * time.Millisecond) + if ctx.Ended() { + t.Errorf("worker's ctx.Ended() should be false even if ender called EndInvocation") + return + } + yield(&session.Event{Author: "worker"}, nil) + } + }, + }) + if err != nil { + t.Fatal(err) + } + + parallelAgent, err := parallelagent.New(parallelagent.Config{ + AgentConfig: agent.Config{ + Name: "parallel", + SubAgents: []agent.Agent{enderAgent, workerAgent}, + }, + }) + if err != nil { + t.Fatal(err) + } + + parentCtx := icontext.NewInvocationContext(t.Context(), icontext.InvocationContextParams{ + Agent: parallelAgent, + }) + + var authors []string + for event, err := range parallelAgent.Run(parentCtx) { + if err != nil { + t.Fatal(err) + } + authors = append(authors, event.Author) + } + + slices.Sort(authors) + wantAuthors := []string{"ender", "worker"} + if !slices.Equal(authors, wantAuthors) { + t.Fatalf("got authors %v, want %v", authors, wantAuthors) + } + + if !parentCtx.Ended() { + t.Fatalf("expected parentCtx.Ended() to be true after parallelAgent completed") + } +} diff --git a/agent/workflowagents/sequentialagent/agent.go b/agent/workflowagents/sequentialagent/agent.go index 517a3b16b..f42b3608e 100644 --- a/agent/workflowagents/sequentialagent/agent.go +++ b/agent/workflowagents/sequentialagent/agent.go @@ -78,6 +78,9 @@ type sequentialAgent struct{} func (a *sequentialAgent) Run(ctx agent.InvocationContext) iter.Seq2[*session.Event, error] { return func(yield func(*session.Event, error) bool) { for _, subAgent := range ctx.Agent().SubAgents() { + if ctx.Ended() { + return + } for event, err := range subAgent.Run(ctx) { // TODO: ensure consistency -- if there's an error, return and close iterator, verify everywhere in ADK. if !yield(event, err) { @@ -166,6 +169,9 @@ func (a *sequentialAgent) RunLive(ctx agent.InvocationContext) (agent.LiveSessio wrappedIter := func(yield func(*session.Event, error) bool) { for _, subAgent := range subAgents { + if ctx.Ended() { + return + } liveAgent, ok := subAgent.(interface { RunLive(ctx agent.InvocationContext) (agent.LiveSession, iter.Seq2[*session.Event, error], error) }) diff --git a/agent/workflowagents/sequentialagent/agent_test.go b/agent/workflowagents/sequentialagent/agent_test.go index 7e937717c..fe1ef1495 100644 --- a/agent/workflowagents/sequentialagent/agent_test.go +++ b/agent/workflowagents/sequentialagent/agent_test.go @@ -374,6 +374,7 @@ type mockInvocationContext struct { agent agent.Agent invocationID string ctx context.Context + ended bool } func (m *mockInvocationContext) Agent() agent.Agent { @@ -384,6 +385,10 @@ func (m *mockInvocationContext) InvocationID() string { return m.invocationID } +func (m *mockInvocationContext) Ended() bool { + return m.ended +} + func (m *mockInvocationContext) Context() context.Context { return m.ctx } @@ -571,3 +576,162 @@ func TestSequentialAgent_RunLive_SequentialOrchestration(t *testing.T) { t.Errorf("expected sub_agent_2 session to be closed at the end") } } + +func TestSequentialAgent_SubAgentEndInvocationPropagated(t *testing.T) { + subAgent1Called := false + subAgent2Called := false + + subAgent1, err := llmagent.New(llmagent.Config{ + Name: "sub_agent_1", + Model: &FakeLLM{id: 1}, + BeforeAgentCallbacks: []agent.BeforeAgentCallback{ + func(ctx agent.Context) (*genai.Content, error) { + subAgent1Called = true + return genai.NewContentFromText("sub agent 1 ended", genai.RoleModel), nil + }, + }, + }) + if err != nil { + t.Fatal(err) + } + + subAgent2, err := llmagent.New(llmagent.Config{ + Name: "sub_agent_2", + Model: &FakeLLM{id: 2}, + BeforeAgentCallbacks: []agent.BeforeAgentCallback{ + func(ctx agent.Context) (*genai.Content, error) { + subAgent2Called = true + return nil, nil + }, + }, + }) + if err != nil { + t.Fatal(err) + } + + seqAgent := newSequentialAgent(t, []agent.Agent{subAgent1, subAgent2}, "seq_agent") + + sessionService := session.InMemoryService() + ctx := context.Background() + _, err = sessionService.Create(ctx, &session.CreateRequest{ + AppName: "test_app", + UserID: "user_id", + SessionID: "session_id", + }) + if err != nil { + t.Fatal(err) + } + + r, err := runner.New(runner.Config{ + AppName: "test_app", + Agent: seqAgent, + SessionService: sessionService, + }) + if err != nil { + t.Fatal(err) + } + + for _, err := range r.Run(ctx, "user_id", "session_id", genai.NewContentFromText("hello", genai.RoleUser), agent.RunConfig{}) { + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + } + + if !subAgent1Called { + t.Error("expected sub_agent_1 callback to be called") + } + + // Demonstrates that calling EndInvocation() in sub_agent_1 callback (via non-nil BeforeAgentCallback return) + // ends the invocation across the shared *atomic.Bool context state, preventing sub_agent_2 from running. + if subAgent2Called { + t.Error("expected sub_agent_2 callback NOT to be called because sub_agent_1 ended the invocation") + } +} + +func TestSequentialAgent_RunLive_Ended(t *testing.T) { + ctx := t.Context() + + subAgent1Called := false + subAgent2Called := false + + agent1 := mustAgent(agent.New(agent.Config{Name: "sub_agent_1"})) + liveAgent1 := &mockLiveAgent{ + Agent: agent1, + runLiveFn: func(ctx agent.InvocationContext) (agent.LiveSession, iter.Seq2[*session.Event, error], error) { + subAgent1Called = true + iterFn := func(yield func(*session.Event, error) bool) { + ev := session.NewEvent(ctx, ctx.InvocationID()) + ev.Author = "sub_agent_1" + yield(ev, nil) + } + return &dummyLiveSession{sendChan: make(chan agent.LiveRequest, 1)}, iterFn, nil + }, + } + + agent2 := mustAgent(agent.New(agent.Config{Name: "sub_agent_2"})) + liveAgent2 := &mockLiveAgent{ + Agent: agent2, + runLiveFn: func(ctx agent.InvocationContext) (agent.LiveSession, iter.Seq2[*session.Event, error], error) { + subAgent2Called = true + iterFn := func(yield func(*session.Event, error) bool) { + ev := session.NewEvent(ctx, ctx.InvocationID()) + ev.Author = "sub_agent_2" + yield(ev, nil) + } + return &dummyLiveSession{sendChan: make(chan agent.LiveRequest, 1)}, iterFn, nil + }, + } + + seqAgent, err := sequentialagent.New(sequentialagent.Config{ + AgentConfig: agent.Config{ + Name: "seq_agent", + SubAgents: []agent.Agent{liveAgent1, liveAgent2}, + }, + }) + if err != nil { + t.Fatalf("failed to create sequential agent: %v", err) + } + + invCtx := &mockInvocationContext{ + agent: seqAgent, + invocationID: "test_inv_id", + ctx: ctx, + } + + liveAgent, ok := seqAgent.(interface { + RunLive(ctx agent.InvocationContext) (agent.LiveSession, iter.Seq2[*session.Event, error], error) + }) + if !ok { + t.Fatalf("sequential agent does not implement RunLive") + } + + _, seqIter, err := liveAgent.RunLive(invCtx) + if err != nil { + t.Fatalf("RunLive failed: %v", err) + } + + next, stop := iter.Pull2(seqIter) + defer stop() + + // Consume first sub-agent event + ev1, err1, ok := next() + if !ok || err1 != nil { + t.Fatalf("expected first event, got ok=%v, err=%v", ok, err1) + } + if ev1.Author != "sub_agent_1" { + t.Errorf("expected event from sub_agent_1, got %s", ev1.Author) + } + if !subAgent1Called { + t.Errorf("expected sub_agent_1 to be called") + } + + invCtx.ended = true + + _, _, ok = next() + if ok { + t.Errorf("expected iterator to end when ctx.Ended() is true") + } + if subAgent2Called { + t.Errorf("expected sub_agent_2 not to be called after ctx.Ended() became true") + } +} diff --git a/internal/context/context_test.go b/internal/context/context_test.go index 1d8bf6b37..f84efc571 100644 --- a/internal/context/context_test.go +++ b/internal/context/context_test.go @@ -17,6 +17,7 @@ package context import ( "context" "strings" + "sync/atomic" "testing" "github.com/google/go-cmp/cmp" @@ -62,7 +63,13 @@ func TestWithContext(t *testing.T) { if got.Value(key) != val { t.Errorf("WithContext() did not update context") } - if diff := cmp.Diff(inv, got, cmp.AllowUnexported(InvocationContext{}), cmpopts.IgnoreFields(InvocationContext{}, "Context")); diff != "" { + atomicBoolComparer := cmp.Comparer(func(a, b *atomic.Bool) bool { + if a == nil || b == nil { + return a == b + } + return a.Load() == b.Load() + }) + if diff := cmp.Diff(inv, got, cmp.AllowUnexported(InvocationContext{}), cmpopts.IgnoreFields(InvocationContext{}, "Context"), atomicBoolComparer); diff != "" { t.Errorf("WithContext() mismatch (-want +got):\n%s", diff) } } diff --git a/internal/context/invocation_context.go b/internal/context/invocation_context.go index c31d1e30d..8f44cc25c 100644 --- a/internal/context/invocation_context.go +++ b/internal/context/invocation_context.go @@ -16,6 +16,7 @@ package context import ( "context" + "sync/atomic" "google.golang.org/genai" @@ -35,7 +36,7 @@ type InvocationContextParams struct { UserContent *genai.Content RunConfig *agent.RunConfig - EndInvocation bool + EndInvocation *atomic.Bool InvocationID string LiveSessionResumptionHandle string Path string @@ -47,6 +48,9 @@ func NewInvocationContext(ctx context.Context, params InvocationContextParams) a if params.InvocationID == "" { params.InvocationID = "e-" + platform.NewUUID(ctx) } + if params.EndInvocation == nil { + params.EndInvocation = new(atomic.Bool) + } return &InvocationContext{ Context: ctx, params: params, @@ -96,10 +100,16 @@ func (c *InvocationContext) RunConfig() *agent.RunConfig { } func (c *InvocationContext) EndInvocation() { - c.params.EndInvocation = true + if c.params.EndInvocation != nil { + c.params.EndInvocation.Store(true) + } } func (c *InvocationContext) Ended() bool { + return c.params.EndInvocation != nil && c.params.EndInvocation.Load() +} + +func (c *InvocationContext) EndInvocationPtr() *atomic.Bool { return c.params.EndInvocation } diff --git a/workflow/agent_node.go b/workflow/agent_node.go index 80d70eb02..125819f1c 100644 --- a/workflow/agent_node.go +++ b/workflow/agent_node.go @@ -18,6 +18,7 @@ import ( "encoding/json" "fmt" "iter" + "sync/atomic" "github.com/google/jsonschema-go/jsonschema" "google.golang.org/genai" @@ -101,6 +102,10 @@ func (n *AgentNode) Run(ctx agent.Context, input any) iter.Seq2[*session.Event, // activation's branch; the scheduler assigns sub-branches at // fan-out, and the LLM flow's history filter scopes events // by branch prefix. + var endInvPtr *atomic.Bool + if carrier, ok := ctx.(interface{ EndInvocationPtr() *atomic.Bool }); ok { + endInvPtr = carrier.EndInvocationPtr() + } params := internalcontext.InvocationContextParams{ Artifacts: ctx.Artifacts(), Memory: ctx.Memory(), @@ -110,7 +115,7 @@ func (n *AgentNode) Run(ctx agent.Context, input any) iter.Seq2[*session.Event, Agent: n.agent, UserContent: userContent, RunConfig: ctx.RunConfig(), - EndInvocation: ctx.Ended(), + EndInvocation: endInvPtr, InvocationID: ctx.InvocationID(), } agentCtx := internalcontext.NewInvocationContext(ctx, params) diff --git a/workflow/agent_node_test.go b/workflow/agent_node_test.go index 0698eabc8..56d25db94 100644 --- a/workflow/agent_node_test.go +++ b/workflow/agent_node_test.go @@ -19,6 +19,7 @@ import ( "errors" "iter" "strings" + "sync/atomic" "testing" "time" @@ -831,3 +832,62 @@ func TestAgentNode_AutomaticOutputExtraction(t *testing.T) { t.Errorf("expected automatically extracted output %q, got %q", want, got) } } + +type endInvContext struct { + agent.Context + ptr *atomic.Bool +} + +func (c *endInvContext) EndInvocationPtr() *atomic.Bool { + return c.ptr +} + +func TestAgentNode_EndInvocationPropagated(t *testing.T) { + called := false + myAgent, err := agent.New(agent.Config{ + Name: "test_agent", + Run: func(ctx agent.InvocationContext) iter.Seq2[*session.Event, error] { + return func(yield func(*session.Event, error) bool) { + called = true + if ctx.Ended() { + t.Errorf("expected ctx.Ended() to be false initially") + } + ctx.EndInvocation() + yield(&session.Event{Author: "test_agent"}, nil) + } + }, + }) + if err != nil { + t.Fatal(err) + } + + node, err := NewAgentNode(myAgent, defaultNodeConfig) + if err != nil { + t.Fatal(err) + } + + ptr := new(atomic.Bool) + ptr.Store(false) + + mockCtx := newMockCtx(t) + mockCtx.sess = &mockSession{id: "test-session"} + + wrappedCtx := &endInvContext{ + Context: agent.NewContext(mockCtx), + ptr: ptr, + } + + for _, err := range node.Run(wrappedCtx, nil) { + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + } + + if !called { + t.Error("expected agent to be called") + } + + if !ptr.Load() { + t.Error("expected parent EndInvocation ptr to be updated to true") + } +}