Skip to content

fix(go-adk): add per-call session isolation for Agent tools#2153

Open
yashrajshuklaaa wants to merge 2 commits into
kagent-dev:mainfrom
yashrajshuklaaa:fix/2137-isolate-sessions
Open

fix(go-adk): add per-call session isolation for Agent tools#2153
yashrajshuklaaa wants to merge 2 commits into
kagent-dev:mainfrom
yashrajshuklaaa:fix/2137-isolate-sessions

Conversation

@yashrajshuklaaa

Copy link
Copy Markdown

Summary

Fixes #2137

The bug : every call a parent agent makes to a sub-agent (Agent tool) reuses the same A2A context_id since it's minted once when the tool is built. worker just uses that context_id as its session id directly ( executor.go: sessionID := reqCtx.ContextID ) so if a coordinator fires off N parallel calls to the same sub-agent in one turn , they all pile into a single shared session instead of getting their own.

Fix is an opt-in IsolateSessions flag on Agent-type tools :

spec:
  declarative:
    tools:
      - type: Agent
        agent:
          name: worker
        isolateSessions: true   # each call to worker gets its own session

Default behavior (flag unset/false) doesn't change
one context_id for the tool's whole lifetime so stateful sub-agents keep their session continuity. with the flag on we mint a fresh context_id on every call so each invocation is isolated.

Doesn't touch the x-kagent-root-context-id header stuff. that's still what carries cross-turn continuity and it stays stable either way.

What changed

  • agent_types.go: added Tool.IsolateSessions *bool plus a CEL rule so it can only be set when type: Agent
  • adk/types.go: RemoteAgentConfig.IsolateSessions bool
  • compiler.go: passes the flag through into RemoteAgentConfig
  • agent.go: forwards the flag to NewKAgentRemoteA2ATool. also skips adding an entry to subagentSessionIDs when isolated since there's no single session id to pre-stamp function_call parts with anymore
  • remote_a2a_tool.go: added isolateSessions plus a small nextContextID() helper that either returns the stable id or mints a new one. This gets reported back as subagent_session_id
  • Python types.py: added isolate_sessions for schema parity only, it doesn't do anything on that side yet (matches what the issue scoped out)
  • Regenerated CRDs/deepcopy with make controller-manifests
  • Added a test for the new context-id logic and a golden fixture (agent_with_isolated_session_tool) to check it flows through the whole translation pipeline

About the UI

Isolated tools don't have one fixed session id, so the executor can't pre-populate the stamp map for them at startup. Instead the UI grabs the session id per call from subagent_session_id in the function_response. AgentCallDisplay already reads that field, so nothing new needed there.

Not doing in this PR

  • Not a concurrency limiter, that's separate (max-concurrency.md)
  • HITL resume is unaffected, it still uses the context_id from the confirmation payload
  • Only the Go runtime respects this flag right now. Python accepts it for config parity but the low-level tool doesn't use it yet

Tests

Ran go test ./adk/pkg/tools/... ./adk/pkg/agent/... ./api/... and the golden translator tests. Everything passes.
Confirmed isolate_sessions: true shows up correctly in the generated config.json for the new fixture.

One small thing I noticed

In handleResume, processResult now sets subagent_session_id from the contextID I pass in, but there's older code right after it that sets the same key again with a fallback to lastContextID. Not wrong, just a bit redundant now since it writes the same value twice in the normal case. Left it as is since it's harmless but flagging it in case someone wants it cleaned up.

Copilot AI review requested due to automatic review settings July 5, 2026 13:41
@github-actions github-actions Bot added the bug Something isn't working label Jul 5, 2026

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Adds an opt-in isolateSessions flag for Agent-type tools so each sub-agent invocation can use a fresh A2A context_id (and therefore a distinct sub-agent session), fixing the “parallel fan-out collapses into one shared worker session” bug described in #2137.

Changes:

  • Introduces Tool.IsolateSessions (CRD + CEL validation) and threads it through translation into RemoteAgentConfig.IsolateSessions.
  • Updates the Go remote A2A tool to mint per-call context_id when isolation is enabled and to avoid pre-stamping a single subagent session id for isolated tools.
  • Adds translation fixture coverage and a unit test for the new context-id selection logic; adds Python schema parity field.

Reviewed changes

Copilot reviewed 13 out of 14 changed files in this pull request and generated 1 comment.

Show a summary per file
File Description
python/packages/kagent-adk/src/kagent/adk/types.py Adds isolate_sessions field for config/schema parity (no runtime behavior change in Python).
helm/kagent-crds/templates/kagent.dev_sandboxagents.yaml Regenerates CRD template to include isolateSessions and CEL rule.
helm/kagent-crds/templates/kagent.dev_agents.yaml Regenerates CRD template to include isolateSessions and CEL rule.
go/core/internal/controller/translator/agent/testdata/outputs/agent_with_isolated_session_tool.json Adds golden output asserting isolate_sessions: true is emitted into config.json.
go/core/internal/controller/translator/agent/testdata/inputs/agent_with_isolated_session_tool.yaml Adds translator input fixture exercising isolateSessions: true.
go/core/internal/controller/translator/agent/compiler.go Threads tool.IsolateSessions into RemoteAgentConfig.IsolateSessions.
go/api/v1alpha2/zz_generated.deepcopy.go Regenerates deepcopy to include Tool.IsolateSessions.
go/api/v1alpha2/agent_types.go Adds Tool.IsolateSessions *bool and CEL validation restricting it to Agent tools.
go/api/config/crd/bases/kagent.dev_sandboxagents.yaml Regenerates CRD base to include isolateSessions schema + validation.
go/api/config/crd/bases/kagent.dev_agents.yaml Regenerates CRD base to include isolateSessions schema + validation.
go/api/adk/types.go Adds RemoteAgentConfig.IsolateSessions (isolate_sessions) to runtime config.
go/adk/pkg/tools/remote_a2a_tool.go Implements per-call context id minting for isolated tools and reports back per-call session id.
go/adk/pkg/tools/remote_a2a_tool_test.go Adds unit test for nextContextID() isolation semantics.
go/adk/pkg/agent/agent.go Passes isolation flag into tool creation and skips pre-stamp map entry when isolated.
Files not reviewed (1)
  • go/api/v1alpha2/zz_generated.deepcopy.go: Generated file

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Comment thread go/adk/pkg/tools/remote_a2a_tool.go Outdated
Comment on lines 372 to 376
// session that actually ran the call — critical when isolateSessions is true
// and every call has a different id.
func (s *remoteA2AState) processResult(ctx adkagent.ToolContext, contextID string, result a2atype.SendMessageResult) (map[string]any, error) {
switch r := result.(type) {
case *a2atype.Message:
@yashrajshuklaaa

Copy link
Copy Markdown
Author

PTAL @EItanya

Comment thread go/adk/pkg/tools/remote_a2a_tool.go
Comment thread go/adk/pkg/tools/remote_a2a_tool.go Outdated

@EItanya EItanya left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks so much for the contribution, however I did find a bug we need to fix and clean up

I think the current approach fixes the session-collision issue, but it also highlights an existing design smell: we have two different mechanisms for linking an AgentCall card to subagent activity.

Today we pre-generate a subagent context_id when constructing the remote tool, return that from NewKAgentRemoteA2ATool, store it in subagentSessionIDs, and stamp it onto the function_call metadata so the UI can show the nested Activity panel. With isolateSessions, that model no longer fits because the actual subagent session is per invocation, not per tool instance. Returning "" from the constructor for isolated tools is a sign that the constructor-time session ID is the wrong abstraction.

I’d suggest simplifying this so both shared and isolated modes use the same UI path:

  1. Stop treating the constructor-time context ID as UI metadata.
  2. Don’t return a subagent session ID from NewKAgentRemoteA2ATool.
  3. Don’t special-case isolated tools in CreateGoogleADKAgentWithSubagentSessionIDs.
  4. Always return the actual invocation’s subagent session ID from the tool result:
return remoteA2AResponse{
    Result:            text,
    SubagentSessionID: contextID,
}
  1. Also include it for input_required / HITL:
return remoteA2AResponse{
    Status:            "pending",
    WaitingFor:        "subagent_approval",
    Subagent:          s.name,
    SubagentSessionID: task.ContextID,
}

Then the UI can use one source of truth for both modes: function_response.response.subagent_session_id. That ID is the actual remote A2A context_id used by the subagent session.

This also lets the Go code be simpler:

type remoteA2AState struct {
    // ...
    sharedContextID string
    isolateSessions bool
}

func (s *remoteA2AState) contextIDForCall() string {
    if s.isolateSessions {
        return a2atype.NewContextID()
    }
    return s.sharedContextID
}

lastContextID should probably be renamed to sharedContextID, since it is not really “last”; in shared mode it is a stable session ID, and in isolated mode it should not be used.

While touching this path, I’d also prefer replacing the ad-hoc map[string]any tool responses with a typed response struct. Right now fields like result, error, status, waiting_for, subagent, subagent_session_id, and kagent_usage_metadata are implicit string keys spread across several branches. That makes it easy to forget subagent_session_id in one path, which is exactly the kind of issue this PR risks for input_required.

Something like:

type remoteA2AResponse struct {
    Result              string         `json:"result,omitempty"`
    Error               string         `json:"error,omitempty"`
    Status              string         `json:"status,omitempty"`
    WaitingFor          string         `json:"waiting_for,omitempty"`
    Subagent            string         `json:"subagent,omitempty"`
    SubagentSessionID   string         `json:"subagent_session_id,omitempty"`
    KAgentUsageMetadata map[string]any `json:"kagent_usage_metadata,omitempty"`
}

Each path can return the same typed shape, and subagent_session_id becomes an explicit field rather than a magic map key. That would make this change safer and easier to review.

The one behavior change is that the nested Activity panel would appear once the tool response/pending response exists, rather than being pre-linked from the initial function_call event. I think that is the cleaner and more honest model. If we need live subagent activity before a response exists, that should be modeled explicitly with a per-invocation “tool started”/running event that carries the generated context ID. Constructor-time pre-stamping only works for shared sessions and is exactly what breaks down here.

Signed-off-by: Yashraj Shukla <shuklayashraj68@gmail.com>
Addresses review feedback on kagent-dev#2153

Signed-off-by: Yashraj Shukla <shuklayashraj68@gmail.com>
@yashrajshuklaaa yashrajshuklaaa force-pushed the fix/2137-isolate-sessions branch from 6a90619 to babc9cb Compare July 10, 2026 02:07
@yashrajshuklaaa yashrajshuklaaa requested a review from a team as a code owner July 10, 2026 02:07
@yashrajshuklaaa

Copy link
Copy Markdown
Author

Thanks for the thorough review, makes sense. Pushed a follow up commit for this:

  • NewKAgentRemoteA2ATool doesn't return a session id anymore. Dropped the constructor time context_id as UI metadata entirely, per your suggestion.
  • Every response path now sets SubagentSessionID consistently (Message, completed Task, input_required, failed Task, and the unrecognised result fallback), using a typed remoteA2AResponse struct instead of ad hoc map[string]any literals. functiontool.New is generic so it infers the output schema straight from the struct, no manual marshaling needed.
  • Renamed lastContextID to sharedContextID and nextContextID() to contextIDForCall(), matching what you proposed.
  • agent.go no longer pulls a session id from the constructor. Left the subagentSessionIDs/stampSubagentSessionID plumbing in the executor as is (just always gets an empty map now) with a comment explaining why, instead of ripping it out. Felt like a bigger cleanup outside this issue's scope tbh. Happy to do that in a follow up PR if you'd rather see it gone now.
  • Checked the UI side before touching this, ToolCallDisplay.tsx already falls back to reading subagent_session_id from the function_response when it's missing from the function_call's DataPart metadata (~line 294). So no UI changes needed. Only behavior change is the Activity panel links once a response/pending state exists instead of pre-linking from the initial function_call, which matches what you described as the cleaner model.

Also added TestProcessResult_SetsSubagentSessionIDOnEveryBranch to cover the branches that were missing the field before.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

bug Something isn't working

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[BUG] Subagent session isolation

4 participants