refactor: replace boolean endInvocation flag with atomic.Bool to enable cross-context state propagation

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
 }
 

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")
+	}
+}

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()

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
 				}
 			}

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")
+	}
+}

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)
 	}()