From f7a8ffd0926deef5f72be57b9d40d1901c0057cb Mon Sep 17 00:00:00 2001 From: Yashraj Shukla Date: Sun, 5 Jul 2026 13:28:54 +0000 Subject: [PATCH 1/2] fix(go-adk): add per-call session isolation for Agent tools Signed-off-by: Yashraj Shukla --- go/adk/pkg/agent/agent.go | 7 +- go/adk/pkg/tools/remote_a2a_tool.go | 62 +++- go/adk/pkg/tools/remote_a2a_tool_test.go | 31 ++ go/api/adk/types.go | 6 + .../config/crd/bases/kagent.dev_agents.yaml | 22 ++ .../crd/bases/kagent.dev_sandboxagents.yaml | 22 ++ go/api/v1alpha2/agent_types.go | 21 ++ go/api/v1alpha2/zz_generated.deepcopy.go | 5 + .../controller/translator/agent/compiler.go | 9 +- .../agent_with_isolated_session_tool.yaml | 49 +++ .../agent_with_isolated_session_tool.json | 301 ++++++++++++++++++ .../templates/kagent.dev_agents.yaml | 22 ++ .../templates/kagent.dev_sandboxagents.yaml | 22 ++ .../kagent-adk/src/kagent/adk/types.py | 5 + 14 files changed, 565 insertions(+), 19 deletions(-) create mode 100644 go/core/internal/controller/translator/agent/testdata/inputs/agent_with_isolated_session_tool.yaml create mode 100644 go/core/internal/controller/translator/agent/testdata/outputs/agent_with_isolated_session_tool.json diff --git a/go/adk/pkg/agent/agent.go b/go/adk/pkg/agent/agent.go index 6d6bab4aa..c18f32317 100644 --- a/go/adk/pkg/agent/agent.go +++ b/go/adk/pkg/agent/agent.go @@ -65,10 +65,15 @@ func CreateGoogleADKAgentWithSubagentSessionIDs(ctx context.Context, agentConfig log.Info("Skipping remote agent with empty URL", "name", remoteAgent.Name) continue } - remoteTool, sessionID, err := tools.NewKAgentRemoteA2ATool(remoteAgent.Name, remoteAgent.Description, remoteAgent.Url, nil, remoteAgent.Headers, propagateToken) + remoteTool, sessionID, err := tools.NewKAgentRemoteA2ATool(remoteAgent.Name, remoteAgent.Description, remoteAgent.Url, nil, remoteAgent.Headers, propagateToken, remoteAgent.IsolateSessions) if err != nil { return nil, nil, fmt.Errorf("failed to create remote A2A tool for %s: %w", remoteAgent.Name, err) } + // Isolated tools mint a fresh context_id per call, so there is no + // single pre-known session id to stamp function_call parts with; + // NewKAgentRemoteA2ATool returns "" in that case and we skip the + // map entry. The UI instead links each sub-agent card via the + // per-call subagent_session_id in the tool's function_response. if sessionID != "" { subagentSessionIDs[remoteAgent.Name] = sessionID } diff --git a/go/adk/pkg/tools/remote_a2a_tool.go b/go/adk/pkg/tools/remote_a2a_tool.go index 705c5c9a6..b04d29fb9 100644 --- a/go/adk/pkg/tools/remote_a2a_tool.go +++ b/go/adk/pkg/tools/remote_a2a_tool.go @@ -149,30 +149,44 @@ type remoteA2AState struct { initOnce sync.Once initErr error + // lastContextID is the stable A2A context_id used for every call to this + // sub-agent when isolateSessions is false (the default): all calls land + // in one shared sub-agent session, giving stateful sub-agents session + // continuity across calls. lastContextID string + + // isolateSessions mints a fresh context_id per call (see nextContextID) + // instead of reusing lastContextID, so each call runs in its own isolated + // sub-agent session. Required for parallel fan-out: without it, N + // parallel calls in one turn collapse into a single shared sub-agent + // session. See go/api/v1alpha2.Tool.IsolateSessions. + isolateSessions bool } // NewKAgentRemoteA2ATool creates a function tool that calls a remote A2A agent and // propagates HITL state. It returns: // - the tool.Tool to register with the agent config -// - the initial A2A context/session ID for subagent session stamping +// - the initial A2A context/session ID for subagent session stamping, or "" +// when isolateSessions is true (there is no single pre-known session id +// to stamp in that case; see remoteA2AState.isolateSessions) // // The agent card is fetched lazily from baseURL/.well-known/agent.json. // If httpClient is nil, a default client is created. The client's transport is // wrapped with otelhttp to propagate W3C trace context to subagents. -func NewKAgentRemoteA2ATool(name, description, baseURL string, httpClient *http.Client, extraHeaders map[string]string, propagateToken bool) (tool.Tool, string, error) { +func NewKAgentRemoteA2ATool(name, description, baseURL string, httpClient *http.Client, extraHeaders map[string]string, propagateToken, isolateSessions bool) (tool.Tool, string, error) { if httpClient == nil { httpClient = &http.Client{} } httpClient = withOTelTransport(httpClient) state := &remoteA2AState{ - name: name, - description: description, - baseURL: baseURL, - httpClient: httpClient, - extraHeaders: extraHeaders, - propagateToken: propagateToken, - lastContextID: a2atype.NewContextID(), + name: name, + description: description, + baseURL: baseURL, + httpClient: httpClient, + extraHeaders: extraHeaders, + propagateToken: propagateToken, + lastContextID: a2atype.NewContextID(), + isolateSessions: isolateSessions, } ft, err := functiontool.New(functiontool.Config{ Name: name, @@ -183,9 +197,23 @@ func NewKAgentRemoteA2ATool(name, description, baseURL string, httpClient *http. if err != nil { return nil, "", fmt.Errorf("failed to create remote A2A function tool for %s: %w", name, err) } + if state.isolateSessions { + return ft, "", nil + } return ft, state.lastContextID, nil } +// nextContextID returns the A2A context_id to stamp on the next outbound +// call: a fresh id when isolateSessions is enabled (isolated per-call +// session), or the tool's stable lastContextID otherwise (shared session for +// the lifetime of the tool). +func (s *remoteA2AState) nextContextID() string { + if s.isolateSessions { + return a2atype.NewContextID() + } + return s.lastContextID +} + // ensureClient lazily resolves the agent card and initialises the A2A client. // Initialization is protected by sync.Once to avoid races under concurrent use. func (s *remoteA2AState) ensureClient(ctx context.Context) (*a2aclient.Client, error) { @@ -257,11 +285,12 @@ func (s *remoteA2AState) handleFirstCall(ctx adkagent.Context, requestText strin return map[string]any{"error": err.Error()}, nil } + contextID := s.nextContextID() message := a2atype.NewMessage( a2atype.MessageRoleUser, a2atype.TextPart{Text: requestText}, ) - message.ContextID = s.lastContextID + message.ContextID = contextID sendCtx := context.WithValue(ctx, userIDContextKey{}, ctx.UserID()) sendCtx = context.WithValue(sendCtx, parentContextIDContextKey{}, ctx.SessionID()) @@ -271,7 +300,7 @@ func (s *remoteA2AState) handleFirstCall(ctx adkagent.Context, requestText strin return map[string]any{"error": fmt.Sprintf("Remote agent '%s' request failed: %v", s.name, err)}, nil } - return s.processResult(ctx, result) + return s.processResult(ctx, contextID, result) } // handleResume is Phase 2: forward the user's decision to the remote agent's pending task. @@ -322,7 +351,7 @@ func (s *remoteA2AState) handleResume(ctx adkagent.Context) (map[string]any, err return map[string]any{"error": fmt.Sprintf("Remote agent '%s' resume failed: %v", subagentName, err)}, nil } - ret, retErr := s.processResult(ctx, result) + ret, retErr := s.processResult(ctx, contextID, result) // Prefer the context_id from the confirmation payload (the original subagent // session) over the pre-generated one. Mirrors Python's: // "subagent_session_id": context_id or self._last_context_id @@ -337,7 +366,12 @@ func (s *remoteA2AState) handleResume(ctx adkagent.Context) (map[string]any, err } // processResult converts a SendMessageResult into a tool return value. -func (s *remoteA2AState) processResult(ctx adkagent.Context, result a2atype.SendMessageResult) (map[string]any, error) { +// contextID is the A2A context_id this call was sent under (from +// nextContextID, or the confirmation payload on resume) and is reported back +// as subagent_session_id so the UI's AgentCallDisplay links the card to the +// session that actually ran the call — critical when isolateSessions is true +// and every call has a different id. +func (s *remoteA2AState) processResult(ctx adkagent.Context, contextID string, result a2atype.SendMessageResult) (map[string]any, error) { switch r := result.(type) { case *a2atype.Message: return map[string]any{"result": extractTextFromMessage(r)}, nil @@ -358,7 +392,7 @@ func (s *remoteA2AState) processResult(ctx adkagent.Context, result a2atype.Send text := extractTextFromTask(r) ret := map[string]any{ "result": text, - "subagent_session_id": s.lastContextID, + "subagent_session_id": contextID, } if usage := extractUsageFromTask(r); usage != nil { ret["kagent_usage_metadata"] = usage diff --git a/go/adk/pkg/tools/remote_a2a_tool_test.go b/go/adk/pkg/tools/remote_a2a_tool_test.go index 14bcd8f3a..7b2c0687a 100644 --- a/go/adk/pkg/tools/remote_a2a_tool_test.go +++ b/go/adk/pkg/tools/remote_a2a_tool_test.go @@ -124,3 +124,34 @@ func assertSingleHeader(t *testing.T, req *a2aclient.Request, key, want string) t.Errorf("%s: got %q, want %q", key, got[0], want) } } + +// TestNextContextID_IsolateSessions covers the EP#2137 fix: isolated tools +// mint a fresh context_id per call so parallel/serial calls to the same +// sub-agent land in independent sessions, while non-isolated tools keep +// reusing one context_id for session continuity. +func TestNextContextID_IsolateSessions(t *testing.T) { + t.Run("isolated: each call gets a distinct, non-empty context_id", func(t *testing.T) { + s := &remoteA2AState{isolateSessions: true, lastContextID: "stable-id"} + + first := s.nextContextID() + second := s.nextContextID() + + if first == "" || second == "" { + t.Fatalf("expected non-empty context ids, got %q and %q", first, second) + } + if first == second { + t.Errorf("expected distinct context ids for isolated calls, got the same id %q twice", first) + } + }) + + t.Run("not isolated: every call reuses the stable lastContextID", func(t *testing.T) { + s := &remoteA2AState{isolateSessions: false, lastContextID: "stable-id"} + + first := s.nextContextID() + second := s.nextContextID() + + if first != "stable-id" || second != "stable-id" { + t.Errorf("expected both calls to reuse lastContextID %q, got %q and %q", "stable-id", first, second) + } + }) +} diff --git a/go/api/adk/types.go b/go/api/adk/types.go index bcc587cf4..dc262d4e7 100644 --- a/go/api/adk/types.go +++ b/go/api/adk/types.go @@ -379,6 +379,12 @@ type RemoteAgentConfig struct { Url string `json:"url"` Headers map[string]string `json:"headers,omitempty"` Description string `json:"description,omitempty"` + // IsolateSessions requests a fresh A2A context_id (and therefore a fresh + // sub-agent session) on every call to this remote agent, instead of the + // default single shared session per tool lifetime. Honored by the Go + // declarative runtime (go/adk/pkg/tools/remote_a2a_tool.go); accepted by + // the Python config model for schema parity only (python/packages/kagent-adk). + IsolateSessions bool `json:"isolate_sessions,omitempty"` } // EmbeddingConfig is the embedding model config for memory tools. diff --git a/go/api/config/crd/bases/kagent.dev_agents.yaml b/go/api/config/crd/bases/kagent.dev_agents.yaml index 15b08f0ca..e2423097c 100644 --- a/go/api/config/crd/bases/kagent.dev_agents.yaml +++ b/go/api/config/crd/bases/kagent.dev_agents.yaml @@ -13222,6 +13222,26 @@ spec: rule: (has(self.value) && !has(self.valueFrom)) || (!has(self.value) && has(self.valueFrom)) type: array + isolateSessions: + description: |- + IsolateSessions controls per-call session isolation for Agent-type tools. + Only valid when Type is Agent. + + When unset or false (default), every call this agent makes to the + referenced sub-agent reuses the same A2A context_id, so all calls land + in one shared sub-agent session (session continuity for stateful + sub-agents). + + When true, each call mints a fresh context_id, so every invocation runs + in its own isolated sub-agent session. This is required for parallel + fan-out to a sub-agent: without it, N parallel calls in one turn + collapse into a single shared sub-agent session instead of N + independent ones. + + Cross-turn/conversation continuity for stateful sub-agents does not + depend on this flag; it rides the x-kagent-root-context-id header, + which stays stable regardless of IsolateSessions. + type: boolean mcpServer: properties: allowedHeaders: @@ -13292,6 +13312,8 @@ spec: rule: '!(has(self.agent) && self.type != ''Agent'')' - message: type.agent must be specified for Agent filter.type rule: '!(!has(self.agent) && self.type == ''Agent'')' + - message: isolateSessions can only be set when type is Agent + rule: '!(has(self.isolateSessions) && self.type != ''Agent'')' maxItems: 20 type: array type: object diff --git a/go/api/config/crd/bases/kagent.dev_sandboxagents.yaml b/go/api/config/crd/bases/kagent.dev_sandboxagents.yaml index 8dd8560b8..36a1379d6 100644 --- a/go/api/config/crd/bases/kagent.dev_sandboxagents.yaml +++ b/go/api/config/crd/bases/kagent.dev_sandboxagents.yaml @@ -10879,6 +10879,26 @@ spec: rule: (has(self.value) && !has(self.valueFrom)) || (!has(self.value) && has(self.valueFrom)) type: array + isolateSessions: + description: |- + IsolateSessions controls per-call session isolation for Agent-type tools. + Only valid when Type is Agent. + + When unset or false (default), every call this agent makes to the + referenced sub-agent reuses the same A2A context_id, so all calls land + in one shared sub-agent session (session continuity for stateful + sub-agents). + + When true, each call mints a fresh context_id, so every invocation runs + in its own isolated sub-agent session. This is required for parallel + fan-out to a sub-agent: without it, N parallel calls in one turn + collapse into a single shared sub-agent session instead of N + independent ones. + + Cross-turn/conversation continuity for stateful sub-agents does not + depend on this flag; it rides the x-kagent-root-context-id header, + which stays stable regardless of IsolateSessions. + type: boolean mcpServer: properties: allowedHeaders: @@ -10949,6 +10969,8 @@ spec: rule: '!(has(self.agent) && self.type != ''Agent'')' - message: type.agent must be specified for Agent filter.type rule: '!(!has(self.agent) && self.type == ''Agent'')' + - message: isolateSessions can only be set when type is Agent + rule: '!(has(self.isolateSessions) && self.type != ''Agent'')' maxItems: 20 type: array type: object diff --git a/go/api/v1alpha2/agent_types.go b/go/api/v1alpha2/agent_types.go index 230368759..5f60d7f28 100644 --- a/go/api/v1alpha2/agent_types.go +++ b/go/api/v1alpha2/agent_types.go @@ -508,6 +508,7 @@ const ( // +kubebuilder:validation:XValidation:message="type.mcpServer must be specified for McpServer filter.type",rule="!(!has(self.mcpServer) && self.type == 'McpServer')" // +kubebuilder:validation:XValidation:message="type.agent must be nil if the type is not Agent",rule="!(has(self.agent) && self.type != 'Agent')" // +kubebuilder:validation:XValidation:message="type.agent must be specified for Agent filter.type",rule="!(!has(self.agent) && self.type == 'Agent')" +// +kubebuilder:validation:XValidation:message="isolateSessions can only be set when type is Agent",rule="!(has(self.isolateSessions) && self.type != 'Agent')" type Tool struct { // +optional Type ToolProviderType `json:"type,omitempty"` @@ -516,6 +517,26 @@ type Tool struct { // +optional Agent *TypedReference `json:"agent,omitempty"` + // IsolateSessions controls per-call session isolation for Agent-type tools. + // Only valid when Type is Agent. + // + // When unset or false (default), every call this agent makes to the + // referenced sub-agent reuses the same A2A context_id, so all calls land + // in one shared sub-agent session (session continuity for stateful + // sub-agents). + // + // When true, each call mints a fresh context_id, so every invocation runs + // in its own isolated sub-agent session. This is required for parallel + // fan-out to a sub-agent: without it, N parallel calls in one turn + // collapse into a single shared sub-agent session instead of N + // independent ones. + // + // Cross-turn/conversation continuity for stateful sub-agents does not + // depend on this flag; it rides the x-kagent-root-context-id header, + // which stays stable regardless of IsolateSessions. + // +optional + IsolateSessions *bool `json:"isolateSessions,omitempty"` + // HeadersFrom specifies a list of configuration values to be added as // headers to requests sent to the Tool from this agent. The value of // each header is resolved from either a Secret or ConfigMap in the same diff --git a/go/api/v1alpha2/zz_generated.deepcopy.go b/go/api/v1alpha2/zz_generated.deepcopy.go index 98d89ec39..88c9bc936 100644 --- a/go/api/v1alpha2/zz_generated.deepcopy.go +++ b/go/api/v1alpha2/zz_generated.deepcopy.go @@ -1949,6 +1949,11 @@ func (in *Tool) DeepCopyInto(out *Tool) { *out = new(TypedReference) **out = **in } + if in.IsolateSessions != nil { + in, out := &in.IsolateSessions, &out.IsolateSessions + *out = new(bool) + **out = **in + } if in.HeadersFrom != nil { in, out := &in.HeadersFrom, &out.HeadersFrom *out = make([]ValueRef, len(*in)) diff --git a/go/core/internal/controller/translator/agent/compiler.go b/go/core/internal/controller/translator/agent/compiler.go index eba707b33..6354f7fed 100644 --- a/go/core/internal/controller/translator/agent/compiler.go +++ b/go/core/internal/controller/translator/agent/compiler.go @@ -359,10 +359,11 @@ func (a *adkApiTranslator) translateInlineAgent(ctx context.Context, agent v1alp } cfg.RemoteAgents = append(cfg.RemoteAgents, adk.RemoteAgentConfig{ - Name: utils.ConvertToPythonIdentifier(utils.GetObjectRef(toolAgent)), - Url: targetURL, - Headers: headers, - Description: toolSpec.Description, + Name: utils.ConvertToPythonIdentifier(utils.GetObjectRef(toolAgent)), + Url: targetURL, + Headers: headers, + Description: toolSpec.Description, + IsolateSessions: tool.IsolateSessions != nil && *tool.IsolateSessions, }) default: return nil, nil, nil, fmt.Errorf("unknown agent type: %s", toolSpec.Type) diff --git a/go/core/internal/controller/translator/agent/testdata/inputs/agent_with_isolated_session_tool.yaml b/go/core/internal/controller/translator/agent/testdata/inputs/agent_with_isolated_session_tool.yaml new file mode 100644 index 000000000..74aa6c41d --- /dev/null +++ b/go/core/internal/controller/translator/agent/testdata/inputs/agent_with_isolated_session_tool.yaml @@ -0,0 +1,49 @@ +operation: translateAgent +targetObject: coordinator-agent +namespace: test +objects: + - apiVersion: v1 + kind: Secret + metadata: + name: openai-secret + namespace: test + data: + api-key: c2stdGVzdC1hcGkta2V5 # base64 encoded "sk-test-api-key" + - apiVersion: kagent.dev/v1alpha2 + kind: ModelConfig + metadata: + name: nested-model + namespace: test + spec: + provider: OpenAI + model: gpt-4o + apiKeySecret: openai-secret + apiKeySecretKey: api-key + - apiVersion: kagent.dev/v1alpha2 + kind: Agent + metadata: + name: worker-agent + namespace: test + spec: + type: Declarative + declarative: + description: A worker agent that can be called in parallel by the coordinator + systemMessage: You are a worker agent. Complete the assigned task and report back. + modelConfig: nested-model + tools: [] + - apiVersion: kagent.dev/v1alpha2 + kind: Agent + metadata: + name: coordinator-agent + namespace: test + spec: + type: Declarative + declarative: + description: A coordinator agent that fans out isolated parallel calls to a worker + systemMessage: You are a coordinating agent that delegates tasks to a worker agent, potentially in parallel. + modelConfig: nested-model + tools: + - type: Agent + agent: + name: worker-agent + isolateSessions: true diff --git a/go/core/internal/controller/translator/agent/testdata/outputs/agent_with_isolated_session_tool.json b/go/core/internal/controller/translator/agent/testdata/outputs/agent_with_isolated_session_tool.json new file mode 100644 index 000000000..d1db9984d --- /dev/null +++ b/go/core/internal/controller/translator/agent/testdata/outputs/agent_with_isolated_session_tool.json @@ -0,0 +1,301 @@ +{ + "agentCard": { + "capabilities": { + "streaming": true + }, + "defaultInputModes": [ + "text" + ], + "defaultOutputModes": [ + "text" + ], + "description": "", + "name": "coordinator_agent", + "skills": null, + "supportedInterfaces": [ + { + "protocolBinding": "JSONRPC", + "protocolVersion": "0.3", + "url": "http://coordinator-agent.test:8080" + }, + { + "protocolBinding": "JSONRPC", + "protocolVersion": "1.0", + "url": "http://coordinator-agent.test:8080" + } + ], + "version": "" + }, + "config": { + "description": "", + "instruction": "You are a coordinating agent that delegates tasks to a worker agent, potentially in parallel.", + "model": { + "base_url": "", + "model": "gpt-4o", + "type": "openai" + }, + "remote_agents": [ + { + "isolate_sessions": true, + "name": "test__NS__worker_agent", + "url": "http://worker-agent.test:8080" + } + ], + "stream": false + }, + "manifest": [ + { + "apiVersion": "v1", + "kind": "Secret", + "metadata": { + "labels": { + "app": "kagent", + "app.kubernetes.io/managed-by": "kagent", + "app.kubernetes.io/name": "coordinator-agent", + "app.kubernetes.io/part-of": "kagent", + "kagent": "coordinator-agent" + }, + "name": "coordinator-agent", + "namespace": "test", + "ownerReferences": [ + { + "apiVersion": "kagent.dev/v1alpha2", + "blockOwnerDeletion": true, + "controller": true, + "kind": "Agent", + "name": "coordinator-agent", + "uid": "" + } + ] + }, + "stringData": { + "agent-card.json": "{\n \"defaultInputModes\": [\n \"text\"\n ],\n \"defaultOutputModes\": [\n \"text\"\n ],\n \"description\": \"\",\n \"name\": \"coordinator_agent\",\n \"version\": \"\",\n \"skills\": [],\n \"capabilities\": {\n \"streaming\": true\n },\n \"supportedInterfaces\": [\n {\n \"url\": \"http://coordinator-agent.test:8080\",\n \"protocolBinding\": \"JSONRPC\",\n \"protocolVersion\": \"0.3\"\n },\n {\n \"url\": \"http://coordinator-agent.test:8080\",\n \"protocolBinding\": \"JSONRPC\",\n \"protocolVersion\": \"1.0\"\n }\n ],\n \"url\": \"http://coordinator-agent.test:8080\",\n \"protocolVersion\": \"0.3\",\n \"preferredTransport\": \"JSONRPC\"\n}", + "config.json": "{\"model\":{\"type\":\"openai\",\"model\":\"gpt-4o\",\"base_url\":\"\"},\"description\":\"\",\"instruction\":\"You are a coordinating agent that delegates tasks to a worker agent, potentially in parallel.\",\"remote_agents\":[{\"name\":\"test__NS__worker_agent\",\"url\":\"http://worker-agent.test:8080\",\"isolate_sessions\":true}],\"stream\":false}" + } + }, + { + "apiVersion": "v1", + "kind": "ServiceAccount", + "metadata": { + "labels": { + "app": "kagent", + "app.kubernetes.io/managed-by": "kagent", + "app.kubernetes.io/name": "coordinator-agent", + "app.kubernetes.io/part-of": "kagent", + "kagent": "coordinator-agent" + }, + "name": "coordinator-agent", + "namespace": "test", + "ownerReferences": [ + { + "apiVersion": "kagent.dev/v1alpha2", + "blockOwnerDeletion": true, + "controller": true, + "kind": "Agent", + "name": "coordinator-agent", + "uid": "" + } + ] + } + }, + { + "apiVersion": "apps/v1", + "kind": "Deployment", + "metadata": { + "labels": { + "app": "kagent", + "app.kubernetes.io/managed-by": "kagent", + "app.kubernetes.io/name": "coordinator-agent", + "app.kubernetes.io/part-of": "kagent", + "kagent": "coordinator-agent" + }, + "name": "coordinator-agent", + "namespace": "test", + "ownerReferences": [ + { + "apiVersion": "kagent.dev/v1alpha2", + "blockOwnerDeletion": true, + "controller": true, + "kind": "Agent", + "name": "coordinator-agent", + "uid": "" + } + ] + }, + "spec": { + "selector": { + "matchLabels": { + "app": "kagent", + "kagent": "coordinator-agent" + } + }, + "strategy": { + "rollingUpdate": { + "maxSurge": 1, + "maxUnavailable": 0 + }, + "type": "RollingUpdate" + }, + "template": { + "metadata": { + "annotations": { + "kagent.dev/config-hash": "12596745338766631722" + }, + "labels": { + "app": "kagent", + "app.kubernetes.io/managed-by": "kagent", + "app.kubernetes.io/name": "coordinator-agent", + "app.kubernetes.io/part-of": "kagent", + "kagent": "coordinator-agent" + } + }, + "spec": { + "containers": [ + { + "args": [ + "--host", + "0.0.0.0", + "--port", + "8080", + "--filepath", + "/config" + ], + "env": [ + { + "name": "OPENAI_API_KEY", + "valueFrom": { + "secretKeyRef": { + "key": "api-key", + "name": "openai-secret" + } + } + }, + { + "name": "KAGENT_NAMESPACE", + "valueFrom": { + "fieldRef": { + "fieldPath": "metadata.namespace" + } + } + }, + { + "name": "KAGENT_NAME", + "value": "coordinator-agent" + }, + { + "name": "KAGENT_URL", + "value": "http://kagent-controller.kagent:8083" + } + ], + "image": "cr.kagent.dev/kagent-dev/kagent/app@sha256:test-app", + "imagePullPolicy": "IfNotPresent", + "name": "kagent", + "ports": [ + { + "containerPort": 8080, + "name": "http" + } + ], + "readinessProbe": { + "httpGet": { + "path": "/.well-known/agent-card.json", + "port": "http" + }, + "initialDelaySeconds": 15, + "periodSeconds": 15, + "timeoutSeconds": 15 + }, + "resources": { + "limits": { + "cpu": "2", + "memory": "1Gi" + }, + "requests": { + "cpu": "100m", + "memory": "384Mi" + } + }, + "volumeMounts": [ + { + "mountPath": "/config", + "name": "config" + }, + { + "mountPath": "/var/run/secrets/tokens", + "name": "kagent-token" + } + ] + } + ], + "serviceAccountName": "coordinator-agent", + "volumes": [ + { + "name": "config", + "secret": { + "secretName": "coordinator-agent" + } + }, + { + "name": "kagent-token", + "projected": { + "sources": [ + { + "serviceAccountToken": { + "audience": "kagent", + "expirationSeconds": 3600, + "path": "kagent-token" + } + } + ] + } + } + ] + } + } + }, + "status": {} + }, + { + "apiVersion": "v1", + "kind": "Service", + "metadata": { + "labels": { + "app": "kagent", + "app.kubernetes.io/managed-by": "kagent", + "app.kubernetes.io/name": "coordinator-agent", + "app.kubernetes.io/part-of": "kagent", + "kagent": "coordinator-agent" + }, + "name": "coordinator-agent", + "namespace": "test", + "ownerReferences": [ + { + "apiVersion": "kagent.dev/v1alpha2", + "blockOwnerDeletion": true, + "controller": true, + "kind": "Agent", + "name": "coordinator-agent", + "uid": "" + } + ] + }, + "spec": { + "ports": [ + { + "name": "http", + "port": 8080, + "targetPort": 8080 + } + ], + "selector": { + "app": "kagent", + "kagent": "coordinator-agent" + }, + "type": "ClusterIP" + }, + "status": { + "loadBalancer": {} + } + } + ] +} \ No newline at end of file diff --git a/helm/kagent-crds/templates/kagent.dev_agents.yaml b/helm/kagent-crds/templates/kagent.dev_agents.yaml index 15b08f0ca..e2423097c 100644 --- a/helm/kagent-crds/templates/kagent.dev_agents.yaml +++ b/helm/kagent-crds/templates/kagent.dev_agents.yaml @@ -13222,6 +13222,26 @@ spec: rule: (has(self.value) && !has(self.valueFrom)) || (!has(self.value) && has(self.valueFrom)) type: array + isolateSessions: + description: |- + IsolateSessions controls per-call session isolation for Agent-type tools. + Only valid when Type is Agent. + + When unset or false (default), every call this agent makes to the + referenced sub-agent reuses the same A2A context_id, so all calls land + in one shared sub-agent session (session continuity for stateful + sub-agents). + + When true, each call mints a fresh context_id, so every invocation runs + in its own isolated sub-agent session. This is required for parallel + fan-out to a sub-agent: without it, N parallel calls in one turn + collapse into a single shared sub-agent session instead of N + independent ones. + + Cross-turn/conversation continuity for stateful sub-agents does not + depend on this flag; it rides the x-kagent-root-context-id header, + which stays stable regardless of IsolateSessions. + type: boolean mcpServer: properties: allowedHeaders: @@ -13292,6 +13312,8 @@ spec: rule: '!(has(self.agent) && self.type != ''Agent'')' - message: type.agent must be specified for Agent filter.type rule: '!(!has(self.agent) && self.type == ''Agent'')' + - message: isolateSessions can only be set when type is Agent + rule: '!(has(self.isolateSessions) && self.type != ''Agent'')' maxItems: 20 type: array type: object diff --git a/helm/kagent-crds/templates/kagent.dev_sandboxagents.yaml b/helm/kagent-crds/templates/kagent.dev_sandboxagents.yaml index 8dd8560b8..36a1379d6 100644 --- a/helm/kagent-crds/templates/kagent.dev_sandboxagents.yaml +++ b/helm/kagent-crds/templates/kagent.dev_sandboxagents.yaml @@ -10879,6 +10879,26 @@ spec: rule: (has(self.value) && !has(self.valueFrom)) || (!has(self.value) && has(self.valueFrom)) type: array + isolateSessions: + description: |- + IsolateSessions controls per-call session isolation for Agent-type tools. + Only valid when Type is Agent. + + When unset or false (default), every call this agent makes to the + referenced sub-agent reuses the same A2A context_id, so all calls land + in one shared sub-agent session (session continuity for stateful + sub-agents). + + When true, each call mints a fresh context_id, so every invocation runs + in its own isolated sub-agent session. This is required for parallel + fan-out to a sub-agent: without it, N parallel calls in one turn + collapse into a single shared sub-agent session instead of N + independent ones. + + Cross-turn/conversation continuity for stateful sub-agents does not + depend on this flag; it rides the x-kagent-root-context-id header, + which stays stable regardless of IsolateSessions. + type: boolean mcpServer: properties: allowedHeaders: @@ -10949,6 +10969,8 @@ spec: rule: '!(has(self.agent) && self.type != ''Agent'')' - message: type.agent must be specified for Agent filter.type rule: '!(!has(self.agent) && self.type == ''Agent'')' + - message: isolateSessions can only be set when type is Agent + rule: '!(has(self.isolateSessions) && self.type != ''Agent'')' maxItems: 20 type: array type: object diff --git a/python/packages/kagent-adk/src/kagent/adk/types.py b/python/packages/kagent-adk/src/kagent/adk/types.py index 867ac317d..f97e15cb4 100644 --- a/python/packages/kagent-adk/src/kagent/adk/types.py +++ b/python/packages/kagent-adk/src/kagent/adk/types.py @@ -238,6 +238,11 @@ class RemoteAgentConfig(BaseModel): headers: dict[str, Any] | None = None timeout: float = DEFAULT_TIMEOUT description: str = "" + # isolate_sessions: accepted for schema parity with the Go declarative + # runtime (see go/api/v1alpha2.Tool.IsolateSessions). The Python low-level + # tool (KAgentRemoteA2AToolset / _remote_a2a_tool.py) does not yet honor + # this flag — it only affects agents running on runtime: go. + isolate_sessions: bool = False class BaseLLM(BaseModel): From cee06704d293c99d5484d0373b63316d56782a34 Mon Sep 17 00:00:00 2001 From: Yashraj Shukla Date: Fri, 10 Jul 2026 01:49:07 +0000 Subject: [PATCH 2/2] fix(go-adk): use typed response + always report subagent_session_id Addresses review feedback on #2153 Signed-off-by: Yashraj Shukla --- go/adk/pkg/agent/agent.go | 24 ++- go/adk/pkg/tools/remote_a2a_tool.go | 164 +++++++++++------- go/adk/pkg/tools/remote_a2a_tool_test.go | 70 ++++++-- .../agent_with_isolated_session_tool.json | 2 +- 4 files changed, 173 insertions(+), 87 deletions(-) diff --git a/go/adk/pkg/agent/agent.go b/go/adk/pkg/agent/agent.go index c18f32317..c77907ac0 100644 --- a/go/adk/pkg/agent/agent.go +++ b/go/adk/pkg/agent/agent.go @@ -57,6 +57,20 @@ func CreateGoogleADKAgentWithSubagentSessionIDs(ctx context.Context, agentConfig } toolsets := mcp.CreateToolsets(ctx, agentConfig.HttpTools, agentConfig.SseTools, propagateToken, dynamicHeaderProvider) mcpAppToolNames := mcp.MCPAppToolNamesFromToolsets(toolsets) + + // subagentSessionIDs is threaded through to the executor for stamping a + // pre-known session id onto outbound function_call events (see + // go/adk/pkg/a2a/executor.go and converter.go's stampSubagentSessionID). + // Remote agent tools no longer report a constructor-time session id + // (NewKAgentRemoteA2ATool returns only an error now) because a single + // pre-known id doesn't fit isolateSessions, where the real session is + // per invocation. So this map stays empty today; stampSubagentSessionID + // already no-ops on an empty map. The per-call session id is instead + // reported in each tool's own response as subagent_session_id, which the + // UI reads as the single source of truth (see remoteA2AResponse). We are + // leaving the executor plumbing in place rather than removing it, since + // something else may want to populate this map in the future and ripping + // it out is a separate, bigger cleanup than this fix. subagentSessionIDs := make(map[string]string) var remoteAgentTools []tool.Tool @@ -65,18 +79,10 @@ func CreateGoogleADKAgentWithSubagentSessionIDs(ctx context.Context, agentConfig log.Info("Skipping remote agent with empty URL", "name", remoteAgent.Name) continue } - remoteTool, sessionID, err := tools.NewKAgentRemoteA2ATool(remoteAgent.Name, remoteAgent.Description, remoteAgent.Url, nil, remoteAgent.Headers, propagateToken, remoteAgent.IsolateSessions) + remoteTool, err := tools.NewKAgentRemoteA2ATool(remoteAgent.Name, remoteAgent.Description, remoteAgent.Url, nil, remoteAgent.Headers, propagateToken, remoteAgent.IsolateSessions) if err != nil { return nil, nil, fmt.Errorf("failed to create remote A2A tool for %s: %w", remoteAgent.Name, err) } - // Isolated tools mint a fresh context_id per call, so there is no - // single pre-known session id to stamp function_call parts with; - // NewKAgentRemoteA2ATool returns "" in that case and we skip the - // map entry. The UI instead links each sub-agent card via the - // per-call subagent_session_id in the tool's function_response. - if sessionID != "" { - subagentSessionIDs[remoteAgent.Name] = sessionID - } remoteAgentTools = append(remoteAgentTools, remoteTool) log.Info("Wired remote A2A agent tool", "name", remoteAgent.Name, "url", remoteAgent.Url) } diff --git a/go/adk/pkg/tools/remote_a2a_tool.go b/go/adk/pkg/tools/remote_a2a_tool.go index b04d29fb9..81e1c202a 100644 --- a/go/adk/pkg/tools/remote_a2a_tool.go +++ b/go/adk/pkg/tools/remote_a2a_tool.go @@ -149,31 +149,54 @@ type remoteA2AState struct { initOnce sync.Once initErr error - // lastContextID is the stable A2A context_id used for every call to this + // sharedContextID is the stable A2A context_id used for every call to this // sub-agent when isolateSessions is false (the default): all calls land // in one shared sub-agent session, giving stateful sub-agents session - // continuity across calls. - lastContextID string + // continuity across calls. Unused when isolateSessions is true — each + // call mints its own id instead (see contextIDForCall). + sharedContextID string - // isolateSessions mints a fresh context_id per call (see nextContextID) - // instead of reusing lastContextID, so each call runs in its own isolated + // isolateSessions mints a fresh context_id per call (see contextIDForCall) + // instead of reusing sharedContextID, so each call runs in its own isolated // sub-agent session. Required for parallel fan-out: without it, N // parallel calls in one turn collapse into a single shared sub-agent // session. See go/api/v1alpha2.Tool.IsolateSessions. isolateSessions bool } -// NewKAgentRemoteA2ATool creates a function tool that calls a remote A2A agent and -// propagates HITL state. It returns: -// - the tool.Tool to register with the agent config -// - the initial A2A context/session ID for subagent session stamping, or "" -// when isolateSessions is true (there is no single pre-known session id -// to stamp in that case; see remoteA2AState.isolateSessions) +// remoteA2AResponse is the typed return value for every remote A2A tool +// invocation. Using one shared struct (instead of ad-hoc map[string]any +// literals per branch) means every response path — success, input_required, +// and failure — carries the same fields, so a field like SubagentSessionID +// can't be silently forgotten in one branch while present in another. +// functiontool.New infers the tool's output schema from this type. +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"` +} + +// NewKAgentRemoteA2ATool creates a function tool that calls a remote A2A agent +// and propagates HITL state. +// +// It intentionally does not return a session id: the constructor-time +// context_id used to be pre-stamped onto outbound function_call events so +// the UI could link the Activity panel before a response arrived, but that +// model only works when a tool has exactly one session for its whole +// lifetime — it breaks down for isolateSessions, where the real session is +// per invocation, not per tool instance. Every call now reports its own +// actual context_id back as SubagentSessionID in the tool's response (see +// remoteA2AResponse), which is the single source of truth the UI reads from +// for both isolated and non-isolated tools alike. // // The agent card is fetched lazily from baseURL/.well-known/agent.json. // If httpClient is nil, a default client is created. The client's transport is // wrapped with otelhttp to propagate W3C trace context to subagents. -func NewKAgentRemoteA2ATool(name, description, baseURL string, httpClient *http.Client, extraHeaders map[string]string, propagateToken, isolateSessions bool) (tool.Tool, string, error) { +func NewKAgentRemoteA2ATool(name, description, baseURL string, httpClient *http.Client, extraHeaders map[string]string, propagateToken, isolateSessions bool) (tool.Tool, error) { if httpClient == nil { httpClient = &http.Client{} } @@ -185,33 +208,30 @@ func NewKAgentRemoteA2ATool(name, description, baseURL string, httpClient *http. httpClient: httpClient, extraHeaders: extraHeaders, propagateToken: propagateToken, - lastContextID: a2atype.NewContextID(), + sharedContextID: a2atype.NewContextID(), isolateSessions: isolateSessions, } ft, err := functiontool.New(functiontool.Config{ Name: name, Description: description, - }, func(ctx adkagent.Context, in remoteA2AInput) (map[string]any, error) { + }, func(ctx adkagent.Context, in remoteA2AInput) (remoteA2AResponse, error) { return state.run(ctx, in.Request) }) if err != nil { - return nil, "", fmt.Errorf("failed to create remote A2A function tool for %s: %w", name, err) + return nil, fmt.Errorf("failed to create remote A2A function tool for %s: %w", name, err) } - if state.isolateSessions { - return ft, "", nil - } - return ft, state.lastContextID, nil + return ft, nil } -// nextContextID returns the A2A context_id to stamp on the next outbound +// contextIDForCall returns the A2A context_id to stamp on the next outbound // call: a fresh id when isolateSessions is enabled (isolated per-call -// session), or the tool's stable lastContextID otherwise (shared session for -// the lifetime of the tool). -func (s *remoteA2AState) nextContextID() string { +// session), or the tool's stable sharedContextID otherwise (shared session +// for the lifetime of the tool). +func (s *remoteA2AState) contextIDForCall() string { if s.isolateSessions { return a2atype.NewContextID() } - return s.lastContextID + return s.sharedContextID } // ensureClient lazily resolves the agent card and initialises the A2A client. @@ -267,7 +287,7 @@ func (s *remoteA2AState) ensureClient(ctx context.Context) (*a2aclient.Client, e } // run dispatches to handleResume or handleFirstCall based on ToolConfirmation presence. -func (s *remoteA2AState) run(ctx adkagent.Context, requestText string) (map[string]any, error) { +func (s *remoteA2AState) run(ctx adkagent.Context, requestText string) (remoteA2AResponse, error) { if ctx.ToolConfirmation() != nil { return s.handleResume(ctx) } @@ -275,17 +295,17 @@ func (s *remoteA2AState) run(ctx adkagent.Context, requestText string) (map[stri } // handleFirstCall is Phase 1: send the request to the remote agent. -func (s *remoteA2AState) handleFirstCall(ctx adkagent.Context, requestText string) (map[string]any, error) { +func (s *remoteA2AState) handleFirstCall(ctx adkagent.Context, requestText string) (remoteA2AResponse, error) { if requestText == "" { - return map[string]any{"error": "missing or empty 'request' argument"}, nil + return remoteA2AResponse{Error: "missing or empty 'request' argument"}, nil } client, err := s.ensureClient(ctx) if err != nil { - return map[string]any{"error": err.Error()}, nil + return remoteA2AResponse{Error: err.Error()}, nil } - contextID := s.nextContextID() + contextID := s.contextIDForCall() message := a2atype.NewMessage( a2atype.MessageRoleUser, a2atype.TextPart{Text: requestText}, @@ -297,14 +317,14 @@ func (s *remoteA2AState) handleFirstCall(ctx adkagent.Context, requestText strin result, err := client.SendMessage(sendCtx, &a2atype.MessageSendParams{Message: message}) if err != nil { slog.Error("Remote agent request failed", "tool", s.name, "error", err) - return map[string]any{"error": fmt.Sprintf("Remote agent '%s' request failed: %v", s.name, err)}, nil + return remoteA2AResponse{Error: fmt.Sprintf("Remote agent '%s' request failed: %v", s.name, err)}, nil } return s.processResult(ctx, contextID, result) } // handleResume is Phase 2: forward the user's decision to the remote agent's pending task. -func (s *remoteA2AState) handleResume(ctx adkagent.Context) (map[string]any, error) { +func (s *remoteA2AState) handleResume(ctx adkagent.Context) (remoteA2AResponse, error) { confirmation := ctx.ToolConfirmation() payload, _ := confirmation.Payload.(map[string]any) hitlPayload := a2a.ParseHitlConfirmationPayload(payload) @@ -318,7 +338,7 @@ func (s *remoteA2AState) handleResume(ctx adkagent.Context) (map[string]any, err if taskID == "" { slog.Error("Resume for remote agent but no task_id in confirmation payload", "tool", s.name) - return map[string]any{"error": fmt.Sprintf("Cannot resume remote agent '%s': missing task context.", subagentName)}, nil + return remoteA2AResponse{Error: fmt.Sprintf("Cannot resume remote agent '%s': missing task context.", subagentName)}, nil } decisionData := buildDecisionData(confirmation.Confirmed, hitlPayload) @@ -340,7 +360,7 @@ func (s *remoteA2AState) handleResume(ctx adkagent.Context) (map[string]any, err client, err := s.ensureClient(ctx) if err != nil { - return map[string]any{"error": err.Error()}, nil + return remoteA2AResponse{Error: err.Error()}, nil } sendCtx := context.WithValue(ctx, userIDContextKey{}, ctx.UserID()) @@ -348,68 +368,77 @@ func (s *remoteA2AState) handleResume(ctx adkagent.Context) (map[string]any, err result, err := client.SendMessage(sendCtx, &a2atype.MessageSendParams{Message: message}) if err != nil { slog.Error("Remote agent resume failed", "tool", subagentName, "error", err) - return map[string]any{"error": fmt.Sprintf("Remote agent '%s' resume failed: %v", subagentName, err)}, nil + return remoteA2AResponse{Error: fmt.Sprintf("Remote agent '%s' resume failed: %v", subagentName, err)}, nil } - ret, retErr := s.processResult(ctx, contextID, result) - // Prefer the context_id from the confirmation payload (the original subagent - // session) over the pre-generated one. Mirrors Python's: - // "subagent_session_id": context_id or self._last_context_id - if retErr == nil && ret != nil { - sessionID := contextID - if sessionID == "" { - sessionID = s.lastContextID - } - ret["subagent_session_id"] = sessionID - } - return ret, retErr + // contextID here is whatever the confirmation payload carried (the + // original subagent session from the paused task). It is always the + // correct id to report — unlike before, there is no fallback to a + // construction-time id: for an isolated tool, sharedContextID was never + // actually used in any A2A call, so falling back to it would report a + // session id that doesn't correspond to any real subagent activity. + return s.processResult(ctx, contextID, result) } // processResult converts a SendMessageResult into a tool return value. // contextID is the A2A context_id this call was sent under (from -// nextContextID, or the confirmation payload on resume) and is reported back -// as subagent_session_id so the UI's AgentCallDisplay links the card to the -// session that actually ran the call — critical when isolateSessions is true -// and every call has a different id. -func (s *remoteA2AState) processResult(ctx adkagent.Context, contextID string, result a2atype.SendMessageResult) (map[string]any, error) { +// contextIDForCall, or the confirmation payload on resume) and is reported +// back as SubagentSessionID on every branch — success, input_required, and +// failure alike — so the UI's AgentCallDisplay can always link the card to +// the session that actually ran the call. This is the single source of +// truth the UI reads from; there is no separate constructor-time id to fall +// back on (see NewKAgentRemoteA2ATool). +func (s *remoteA2AState) processResult(ctx adkagent.Context, contextID string, result a2atype.SendMessageResult) (remoteA2AResponse, error) { switch r := result.(type) { case *a2atype.Message: - return map[string]any{"result": extractTextFromMessage(r)}, nil + return remoteA2AResponse{ + Result: extractTextFromMessage(r), + SubagentSessionID: contextID, + }, nil case *a2atype.Task: switch r.Status.State { case a2atype.TaskStateInputRequired: - return s.handleInputRequired(ctx, r), nil + return s.handleInputRequired(ctx, r, contextID), nil case a2atype.TaskStateFailed: text := extractTextFromTask(r) if text == "" { text = fmt.Sprintf("Remote agent '%s' failed.", s.name) } - return map[string]any{"error": text}, nil + return remoteA2AResponse{ + Error: text, + SubagentSessionID: contextID, + }, nil default: // completed — include sub-agent's final LLM usage from task.metadata // so the parent can display it on the AgentCall card in the UI. // Mirrors Python's _extract_usage_from_task(task). - text := extractTextFromTask(r) - ret := map[string]any{ - "result": text, - "subagent_session_id": contextID, + ret := remoteA2AResponse{ + Result: extractTextFromTask(r), + SubagentSessionID: contextID, } if usage := extractUsageFromTask(r); usage != nil { - ret["kagent_usage_metadata"] = usage + ret.KAgentUsageMetadata = usage } return ret, nil } default: - return map[string]any{"error": fmt.Sprintf("Remote agent '%s' returned no result.", s.name)}, nil + return remoteA2AResponse{ + Error: fmt.Sprintf("Remote agent '%s' returned no result.", s.name), + SubagentSessionID: contextID, + }, nil } } // handleInputRequired pauses parent agent execution via RequestConfirmation. -func (s *remoteA2AState) handleInputRequired(ctx adkagent.Context, task *a2atype.Task) map[string]any { +// contextID is reported back as SubagentSessionID so the UI can link the +// pending Activity panel to the paused subagent session before the human +// decision is forwarded and a final result comes back. +func (s *remoteA2AState) handleInputRequired(ctx adkagent.Context, task *a2atype.Task, contextID string) remoteA2AResponse { if task == nil { slog.Error("Subagent returned input_required without task", "tool", s.name) - return map[string]any{ - "error": fmt.Sprintf("Remote agent '%s' returned input_required without task context.", s.name), + return remoteA2AResponse{ + Error: fmt.Sprintf("Remote agent '%s' returned input_required without task context.", s.name), + SubagentSessionID: contextID, } } @@ -446,10 +475,11 @@ func (s *remoteA2AState) handleInputRequired(ctx adkagent.Context, task *a2atype if err := ctx.RequestConfirmation(hint, confirmPayload.ToMap()); err != nil { slog.Error("Failed to request confirmation", "tool", s.name, "error", err) } - return map[string]any{ - "status": "pending", - "waiting_for": "subagent_approval", - "subagent": s.name, + return remoteA2AResponse{ + Status: "pending", + WaitingFor: "subagent_approval", + Subagent: s.name, + SubagentSessionID: contextID, } } diff --git a/go/adk/pkg/tools/remote_a2a_tool_test.go b/go/adk/pkg/tools/remote_a2a_tool_test.go index 7b2c0687a..fa94137f2 100644 --- a/go/adk/pkg/tools/remote_a2a_tool_test.go +++ b/go/adk/pkg/tools/remote_a2a_tool_test.go @@ -4,6 +4,7 @@ import ( "context" "testing" + a2atype "github.com/a2aproject/a2a-go/a2a" "github.com/a2aproject/a2a-go/a2aclient" "github.com/a2aproject/a2a-go/a2asrv" ) @@ -125,16 +126,16 @@ func assertSingleHeader(t *testing.T, req *a2aclient.Request, key, want string) } } -// TestNextContextID_IsolateSessions covers the EP#2137 fix: isolated tools +// TestContextIDForCall_IsolateSessions covers the EP#2137 fix: isolated tools // mint a fresh context_id per call so parallel/serial calls to the same // sub-agent land in independent sessions, while non-isolated tools keep // reusing one context_id for session continuity. -func TestNextContextID_IsolateSessions(t *testing.T) { +func TestContextIDForCall_IsolateSessions(t *testing.T) { t.Run("isolated: each call gets a distinct, non-empty context_id", func(t *testing.T) { - s := &remoteA2AState{isolateSessions: true, lastContextID: "stable-id"} + s := &remoteA2AState{isolateSessions: true, sharedContextID: "stable-id"} - first := s.nextContextID() - second := s.nextContextID() + first := s.contextIDForCall() + second := s.contextIDForCall() if first == "" || second == "" { t.Fatalf("expected non-empty context ids, got %q and %q", first, second) @@ -144,14 +145,63 @@ func TestNextContextID_IsolateSessions(t *testing.T) { } }) - t.Run("not isolated: every call reuses the stable lastContextID", func(t *testing.T) { - s := &remoteA2AState{isolateSessions: false, lastContextID: "stable-id"} + t.Run("not isolated: every call reuses the stable sharedContextID", func(t *testing.T) { + s := &remoteA2AState{isolateSessions: false, sharedContextID: "stable-id"} - first := s.nextContextID() - second := s.nextContextID() + first := s.contextIDForCall() + second := s.contextIDForCall() if first != "stable-id" || second != "stable-id" { - t.Errorf("expected both calls to reuse lastContextID %q, got %q and %q", "stable-id", first, second) + t.Errorf("expected both calls to reuse sharedContextID %q, got %q and %q", "stable-id", first, second) + } + }) +} + +// TestProcessResult_SetsSubagentSessionIDOnEveryBranch covers the review +// feedback on #2153: subagent_session_id must be present in the response for +// every result shape (direct Message, completed Task, input_required Task, +// failed Task, and the unrecognised-result fallback) — not just the +// completed-Task branch — since it is the UI's only source of truth for +// linking the AgentCallDisplay Activity panel to the correct subagent +// session, especially when isolateSessions means every call has a distinct id. +func TestProcessResult_SetsSubagentSessionIDOnEveryBranch(t *testing.T) { + const contextID = "call-specific-context-id" + s := &remoteA2AState{name: "worker"} + ctx := context.Background() + + t.Run("direct Message result", func(t *testing.T) { + msg := &a2atype.Message{Parts: a2atype.ContentParts{a2atype.TextPart{Text: "hi"}}} + resp, err := s.processResult(nil, contextID, msg) + _ = ctx + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if resp.SubagentSessionID != contextID { + t.Errorf("SubagentSessionID = %q, want %q", resp.SubagentSessionID, contextID) + } + }) + + t.Run("failed Task result", func(t *testing.T) { + task := &a2atype.Task{Status: a2atype.TaskStatus{State: a2atype.TaskStateFailed}} + resp, err := s.processResult(nil, contextID, task) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if resp.SubagentSessionID != contextID { + t.Errorf("SubagentSessionID = %q, want %q", resp.SubagentSessionID, contextID) + } + if resp.Error == "" { + t.Errorf("expected a non-empty Error for a failed task") + } + }) + + t.Run("unrecognised result type", func(t *testing.T) { + resp, err := s.processResult(nil, contextID, nil) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if resp.SubagentSessionID != contextID { + t.Errorf("SubagentSessionID = %q, want %q", resp.SubagentSessionID, contextID) } }) } diff --git a/go/core/internal/controller/translator/agent/testdata/outputs/agent_with_isolated_session_tool.json b/go/core/internal/controller/translator/agent/testdata/outputs/agent_with_isolated_session_tool.json index d1db9984d..67dc549b5 100644 --- a/go/core/internal/controller/translator/agent/testdata/outputs/agent_with_isolated_session_tool.json +++ b/go/core/internal/controller/translator/agent/testdata/outputs/agent_with_isolated_session_tool.json @@ -187,7 +187,7 @@ "value": "http://kagent-controller.kagent:8083" } ], - "image": "cr.kagent.dev/kagent-dev/kagent/app@sha256:test-app", + "image": "ghcr.io/kagent-dev/kagent/app@sha256:test-app", "imagePullPolicy": "IfNotPresent", "name": "kagent", "ports": [