Overhaul Branch Logic from string based to token set based.

[dannovikov]

Link to Issue or Description of Change

1. Link to an existing issue (if applicable):

• Closes: Inconsistent Event Visibility in Parallel Agent Branching and LLM Context #3470

Problem:

The current approach for branching used in parallel agent for creating isolated parallel contexts for the subagents is broken and causes multiple intuitive architectures to fail. Both architectures are mentioned in the above issue.

A sequence of parallel agents doesn't intuitively pass context; each next parallel agent's subagents don't see any of the outputs of the previous parallel groups. The linked issue above also describes another failing architecture, the fork and join, or map/reduce style architecture, which work on their own, but not when nested inside parallel agents.

From my exploration, it was clear the string and prefix-match approach used for branching was insufficient.

Solution:

The proposed solution introduces a Branch class as the new branch representation, one which represents branches by sets of integer tokens.

When a ParallelAgent runs its subagents, it adds a unique token to each subagent's branch set. Then, the distinct branch sets are merged into a single set containing the union of those tokens, and this becomes the branch context going forward.

When determining if an event is visible to a branch, we check if the event's branch token set is a subset of the current branch. A global token factory ensures each new fork gets a unique token. This token-subset approach, instead of a string prefix approach, solves the two problem architectures noted in the above issue.

Example:
To explore this solution in action, consider the architecture:

Sequential[Parallel1[A, B], Parallel2[C,D], Parallel3[E, F]]

This architecture currently doesn't "work" -- each next parallel group can't see the events from the previous parallel group before them. So C, D can't see the events from A, B.

In the newly proposed branching paradigm:

We start with an empty branch set {} and start running the sequential agent.
When Parallel1 runs, it adds unique tokens to A's and B's branch set, so A runs on branch {1} and B runs on branch {2}.

After A and B finish, the branch sets are merged, and the main context's branch is updated to the union of the parallel branches - {1,2}.

When Parallel2 runs on branch {1, 2}, C and D are each given new branch tokens, so C runs on branch {1, 2, 3} and D runs on branch {1, 2, 4}. Both C and D can see the events that came on branch {1} and {2}, because those are subsets of each of their branches. After C and D finish, the current context's branch is updated to be their union - {1, 2, 3, 4}.

Now {1, 2, 3, 4} is the current context, and we run Parallel 3. E gets a new token in its token set, so it runs on {1, 2, 3, 4, 5}, and F runs on {1, 2, 3, 4, 6}. They can see all the events from before; the ones that came from A on branch {1}, B on branch {2}, C on branch {1, 2, 3}, and D on branch {1, 2, 4}, each of which are subsets of the current branches. Once E and F finish, their token sets are merged, such that the context's branch going forward is {1, 2, 3, 4, 5, 6}.

Token counter is reset on each call to runner.run_async, and the token sets are only useful for the current invocation. Events should always be visible across invocations, so if the invocation id of an event doesn't match the current invocation id, the event is always visible.

Testing Plan

All unit tests have been updated to use the new branching logic. I have added tests to test the two architectures mentioned in the issue, and confirmed that they both work as expected under the new system.

Unit Tests:

• I have added or updated unit tests for my change.
• All unit tests pass locally.

===== 3338 passed, 2437 warnings in 30.21s =====

Manual End-to-End (E2E) Tests:

Please run uv run pytest tests/unittests/agents/test_nested_agent_branch_visibility.py -vs to see the two problem architectures running correctly and a display of how the tokens are being assigned (shown below).

======================================================================
INTEGRATION TEST: Nested Parallel + Reduce (GitHub Issue #3470)
======================================================================

Architecture:
  Sequential[
    Parallel[
      Sequential[Parallel[Alice,Bob,Charlie], Reducer1],  ← Group 1
      Sequential[Parallel[David,Eve,Frank], Reducer2]     ← Group 2
    ],
    Final_Reducer  ← Sees all outputs
  ]


=== Token Distribution (Nested Parallel) ===
  Agent           Tokens                        
  --------------- ------------------------------
  Alice           | tokens=[1, 3]
  David           | tokens=[2, 6]
  Bob             | tokens=[1, 4]
  Eve             | tokens=[2, 7]
  Charlie         | tokens=[1, 5]
  Frank           | tokens=[2, 8]
  Reducer1        | tokens=[1, 3, 4, 5]
  Reducer2        | tokens=[2, 6, 7, 8]
  Final_Reducer   | tokens=[1, 2, 3, 4, 5, 6, 7, 8]
======================================================================

======================================================================
INTEGRATION TEST: Sequence of Parallels (GitHub Issue #3470)
======================================================================

Architecture:
  Sequential[
    Parallel1[Alice, Bob, Charlie],    ← Group 1
    Parallel2[David, Eve, Frank],      ← Group 2 (sees Group 1)
    Parallel3[Grace, Henry, Iris]      ← Group 3 (sees Groups 1 & 2)
  ]


=== Token Distribution ===
  Agent           Tokens                         Can See
  --------------- ------------------------------ ----------------------------------------
  --- Group 1 ---
  Alice           | tokens=[1]                          Root
  Bob             | tokens=[2]                          Root
  Charlie         | tokens=[3]                          Root
  --- Group 2 ---
  David           | tokens=[1, 2, 3, 4]                 Root, Group 1 (A,B,C)
  Eve             | tokens=[1, 2, 3, 5]                 Root, Group 1 (A,B,C)
  Frank           | tokens=[1, 2, 3, 6]                 Root, Group 1 (A,B,C)
  --- Group 3 ---
  Grace           | tokens=[1, 2, 3, 4, 5, 6, 7]        Root, Groups 1 & 2 (A-F)
  Henry           | tokens=[1, 2, 3, 4, 5, 6, 8]        Root, Groups 1 & 2 (A-F)
  Iris            | tokens=[1, 2, 3, 4, 5, 6, 9]        Root, Groups 1 & 2 (A-F)

Key Observations:
  ✓ Group 2 agents have tokens {1,2,3,...} - inherit from Group 1
  ✓ Group 3 agents have tokens {1,2,3,4,5,6,...} - inherit from Groups 1 & 2
  ✓ Each agent can see all events with token subsets
======================================================================

Checklist

• I have read the CONTRIBUTING.md document.
• I have performed a self-review of my own code.
• I have commented my code, particularly in hard-to-understand areas.
• I have added tests that prove my fix is effective or that my feature works.
• New and existing unit tests pass locally with my changes.
• I have manually tested my changes end-to-end.
• Any dependent changes have been merged and published in downstream modules.

[gemini-code-assist]

Summary of Changes

Hello @dannovikov, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request significantly improves the agent orchestration capabilities by replacing the existing string-based branching mechanism with a token set-based approach. This change resolves critical issues where parallel agents failed to correctly pass context in complex sequential and nested configurations. The new system ensures that agents can accurately determine event visibility based on their provenance, leading to more reliable and intuitive multi-agent workflows without unintended information leakage or missing context.

Highlights

• Branching Logic Overhaul: The pull request completely reworks the internal branching logic for parallel agents, moving from a string-based prefix matching approach to a more robust token set-based system.
• New Branch Representation: Introduces a new Branch class that represents branches as frozensets of integer tokens, along with a BranchTokenFactory to generate unique tokens.
• Enhanced Event Visibility: Event visibility is now determined by a subset relationship of these token sets, ensuring correct context passing in complex multi-agent architectures, such as sequential parallel agents and nested reducer patterns.
• Problem Architectures Fixed: Specifically addresses and fixes two previously failing architectures: a sequence of parallel agents and nested fork/join (map/reduce) patterns, as detailed in GitHub Issue Inconsistent Event Visibility in Parallel Agent Branching and LLM Context #3470.
• Comprehensive Testing: All existing unit tests have been updated, and new dedicated integration tests have been added to validate the correct behavior of the new branching logic for the previously problematic architectures.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request significantly overhauls the branching logic from a string-based prefix matching system to a more robust token-set based approach. This change effectively addresses known limitations in handling complex parallel and nested agent architectures, as detailed in the issue description. The introduction of the BranchTokenFactory and Branch classes provides a clear and efficient mechanism for tracking event provenance and determining visibility. The comprehensive unit and integration tests, including specific scenarios from the GitHub issue, demonstrate that the new system correctly handles the intended visibility rules. This is a well-designed and thoroughly tested solution that improves the flexibility and correctness of agent interactions.

[dannovikov]

@ryanaiagent @ankursharmas If you can please review this solution.

[ryanaiagent]

Hi @dannovikov , Thank you for your contribution! We appreciate you taking the time to submit this pull request.
Can you please fix the failing tests

[dannovikov]

Hi @dannovikov , Thank you for your contribution! We appreciate you taking the time to submit this pull request. Can you please fix the failing tests

@ryanaiagent Done!