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

[yashrajshuklaaa]

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]

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.

[yashrajshuklaaa]

PTAL @EItanya

[EItanya]

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

5. 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.

[yashrajshuklaaa]

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.

[dimetron]

Thank you @yashrajshuklaaa for taking care of #2137 some additional changes I've discovered during my testing

Additional Required changes

1. [required] Response struct is missing failure-linkage and structured fields

The PR struct carries only result/error/status/waiting_for/subagent/ subagent_session_id/kagent_usage_metadata. For failed/canceled/rejected
children the parent also needs:

• subagent_task_id, subagent_task_state — link to the exact failed task.
• error_type, error_step, error_message, remediation_hint — child-provided
  structured failure passed through instead of buried in text, so the coordinator
  can act without re-parsing.

Without these, a failed sub-agent is not debuggable from the parent session
(the original motivation for the fan-out investigation). Add them to the struct
and populate on the terminal-failure branch.

2. [required] processResult must switch explicitly on terminal state

The PR keeps the old shape where any *Task state other than input_required
and failed falls through to "completed". A working/submitted task after the
call ends is not success. Switch explicitly:
completed → result; failed/canceled/rejected → typed error;
input_required → HITL; anything else → subagent_non_terminal error.

[dimetron]

@yashrajshuklaaa just for the reference #2202

[yashrajshuklaaa]

These are good additions, thanks for testing it out @dimetron. I think it's a separate concern from what #2137 asked for (session isolation vs structured failure reporting) so my plan is to pick these up in a focused follow up PR once this one lands and link back to this thread when I open it.
@EItanya curious what you think too , happy to fold it into this PR instead if you'd rather see it land together

[yashrajshuklaaa]

Also keeps git log/blame cleaner, a commit called "add per-call session isolation" quietly also adding a failure reporting schema would be confusing to trace back later.

[yashrajshuklaaa]

@EItanya PTAL

[supreme-gg-gg]

@yashrajshuklaaa the changes look good overall, pls resolve the merge conflict and I'll take another look / approve. A few follow ups:

1. I agree we need to address @dimetron's comments regarding terminal state handling and failure fields
2. I think it's a fair tradeoff to drop constructor-time pre-stamping of session ID for a single response-based UI path, with that, we should get rid of the dead plumbings like stampSubagentSessionID, subsagentSessionIDs threading in agent.go, executor.go, converter.go

There's a regression regarding how HITL activity is supposed to work that is quite important to address. Without pre-stamped IDs, subagent activity panels will not open until a function_response supplies the session ID. HITL events surfaced from subagents returns this function_response containing status: "pending" + subagent_session_id, but it is currently skipped in the UI since it wasn't needed before due to pre-stamping IDs. One of the key motivations of having subagent activity viewing is for the user to know what the subagent is doing before approving its tool calls, which no longer works.

kagent/ui/src/lib/messageHandlers.ts

Lines 940 to 950 in eed7287

	} else if (partType === "function_response") {
	// Skip internal HITL markers: the before_tool_callback stub and
	// the ask_user first-invocation pending stub.
	const responseData = (data as { response?: Record<string, unknown> })?.response;
	const responseStatus = responseData?.status as string | undefined;
	if (responseStatus === "confirmation_requested" || responseStatus === "pending") {
	continue;
	}
	const toolData = data as unknown as ToolResponseData;
	const source = getSourceFromMetadata(adkMetadata, defaultAgentSource);
	processFunctionResponsePart(toolData, statusUpdate.contextId, statusUpdate.taskId, source);

A simple fix that I came up with is to check if it comes from a pending subagent and do not skip it:

             } else if (partType === "function_response") {
               const responseData = (data as { response?: Record<string, unknown> })?.response;
               const responseStatus = responseData?.status as string | undefined;
-              if (responseStatus === "confirmation_requested" || responseStatus === "pending") {
+              const toolData = data as unknown as ToolResponseData;
+              const isPendingAgentSession =
+                responseStatus === "pending" &&
+                isAgentToolName(toolData.name) &&
+                typeof responseData?.subagent_session_id === "string";
+              if (
+                (responseStatus === "confirmation_requested" || responseStatus === "pending") &&
+                !isPendingAgentSession
+              ) {
                 continue;
               }
-              const toolData = data as unknown as ToolResponseData;
               const source = getSourceFromMetadata(adkMetadata, defaultAgentSource);
               processFunctionResponsePart(toolData, statusUpdate.contextId, statusUpdate.taskId, source);
             }

I've tried this out and it preserves the old behaviour where when a subagent requires an approval, the user can view its activity.

[EItanya]

Hey there, unit tests are failing

[yashrajshuklaaa]

@supreme-gg-gg @EItanya thanks for the review!

[yashrajshuklaaa]

@supreme-gg-gg @EItanya removed the dead plumbing
stampSubagentSessionID / stampSubagentSessionIDOnDataPart gone from converter.go, SubagentSessionIDs dropped from executor.go, CreateGoogleADKAgentWithSubagentSessionIDs now just returns (agent.Agent, error). build/vet/test all passing. PTAL

[supreme-gg-gg]

This looks great! One nit and this is good to merge

[supreme-gg-gg]

lgtm, thanks for the PR!

[EItanya]

This will need to be backported to make it into 0.10