fix(go-adk): add per-call session isolation for Agent tools#2153
fix(go-adk): add per-call session isolation for Agent tools#2153yashrajshuklaaa wants to merge 2 commits into
Conversation
There was a problem hiding this comment.
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 intoRemoteAgentConfig.IsolateSessions. - Updates the Go remote A2A tool to mint per-call
context_idwhen 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.
| // 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: |
|
PTAL @EItanya |
EItanya
left a comment
There was a problem hiding this comment.
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:
- Stop treating the constructor-time context ID as UI metadata.
- Don’t return a subagent session ID from
NewKAgentRemoteA2ATool. - Don’t special-case isolated tools in
CreateGoogleADKAgentWithSubagentSessionIDs. - Always return the actual invocation’s subagent session ID from the tool result:
return remoteA2AResponse{
Result: text,
SubagentSessionID: contextID,
}- 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>
6a90619 to
babc9cb
Compare
|
Thanks for the thorough review, makes sense. Pushed a follow up commit for this:
Also added TestProcessResult_SetsSubagentSessionIDOnEveryBranch to cover the branches that were missing the field before. |
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 :
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: addedTool.IsolateSessions *boolplus a CEL rule so it can only be set whentype: Agentadk/types.go:RemoteAgentConfig.IsolateSessions boolcompiler.go: passes the flag through intoRemoteAgentConfigagent.go: forwards the flag toNewKAgentRemoteA2ATool. also skips adding an entry tosubagentSessionIDswhen isolated since there's no single session id to pre-stamp function_call parts with anymoreremote_a2a_tool.go: addedisolateSessionsplus a smallnextContextID()helper that either returns the stable id or mints a new one. This gets reported back assubagent_session_idtypes.py: addedisolate_sessionsfor schema parity only, it doesn't do anything on that side yet (matches what the issue scoped out)make controller-manifestsagent_with_isolated_session_tool) to check it flows through the whole translation pipelineAbout 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_idin the function_response.AgentCallDisplayalready reads that field, so nothing new needed there.Not doing in this PR
max-concurrency.md)Tests
Ran
go test ./adk/pkg/tools/... ./adk/pkg/agent/... ./api/...and the golden translator tests. Everything passes.Confirmed
isolate_sessions: trueshows up correctly in the generatedconfig.jsonfor the new fixture.One small thing I noticed
In
handleResume,processResultnow setssubagent_session_idfrom the contextID I pass in, but there's older code right after it that sets the same key again with a fallback tolastContextID. 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.