Add the ability to click through to sub-agent session

go/internal/database/client.go

@@ -491,6 +491,27 @@ func (c *clientImpl) DeleteCheckpoint(userID, threadID string) error {
 	})
 }
 
+func (c *clientImpl) FindSessionByParentFunctionCallID(functionCallID string) (*dbpkg.Session, error) {
+	var event dbpkg.Event
+	// Use a wildcard between the key and value to handle potential whitespace (e.g. ": " vs ":")
+	searchPattern := fmt.Sprintf("%%\"parent_function_call_id\"%%\"%s\"%%", functionCallID)
+	err := c.db.Where("data LIKE ?", searchPattern).First(&event).Error
+	if err != nil {
+		if errors.Is(err, gorm.ErrRecordNotFound) {
+			return nil, gorm.ErrRecordNotFound
+		}
+		return nil, fmt.Errorf("failed to find event by parent function call ID: %w", err)
+	}
+
+	var session dbpkg.Session
+	err = c.db.Where("id = ?", event.SessionID).First(&session).Error
+	if err != nil {
+		return nil, fmt.Errorf("failed to find session for event: %w", err)
+	}
+
+	return &session, nil
+}
+
 // CrewAI methods
 
 // StoreCrewAIMemory stores CrewAI agent memory

go/internal/database/fake/client.go

@@ -741,6 +741,26 @@ func (c *InMemoryFakeClient) ListWrites(userID, threadID, checkpointNS, checkpoi
 	return writes[start:end], nil
 }
 
+func (c *InMemoryFakeClient) FindSessionByParentFunctionCallID(functionCallID string) (*database.Session, error) {
+	c.mu.RLock()
+	defer c.mu.RUnlock()
+
+	searchPattern1 := fmt.Sprintf(`"parent_function_call_id":"%s"`, functionCallID)
+	searchPattern2 := fmt.Sprintf(`"parent_function_call_id": "%s"`, functionCallID)
+
+	for _, event := range c.events {
+		if strings.Contains(event.Data, searchPattern1) || strings.Contains(event.Data, searchPattern2) {
+			for _, session := range c.sessions {
+				if session.ID == event.SessionID {
+					return session, nil
+				}
+			}
+			return nil, gorm.ErrRecordNotFound
+		}
+	}
+	return nil, gorm.ErrRecordNotFound
+}
+
 // CrewAI methods
 
 // StoreCrewAIMemory stores CrewAI agent memory

go/internal/httpserver/handlers/sessions.go

@@ -1,6 +1,7 @@
 package handlers
 
 import (
+	stdErrors "errors"
 	"fmt"
 	"net/http"
 	"strconv"
@@ -10,6 +11,7 @@ import (
 	"github.com/kagent-dev/kagent/go/internal/utils"
 	"github.com/kagent-dev/kagent/go/pkg/client/api"
 	"github.com/kagent-dev/kagent/go/pkg/database"
+	"gorm.io/gorm"
 	ctrllog "sigs.k8s.io/controller-runtime/pkg/log"
 	"trpc.group/trpc-go/trpc-a2a-go/protocol"
 )
@@ -395,6 +397,32 @@ func (h *SessionsHandler) HandleAddEventToSession(w ErrorResponseWriter, r *http
 	RespondWithJSON(w, http.StatusCreated, data)
 }
 
+func (h *SessionsHandler) HandleFindSessionByFunctionCall(w ErrorResponseWriter, r *http.Request) {
+	log := ctrllog.FromContext(r.Context()).WithName("sessions-handler").WithValues("operation", "find-by-function-call")
+
+	functionCallID := r.URL.Query().Get("function_call_id")
+	if functionCallID == "" {
+		w.RespondWithError(errors.NewBadRequestError("function_call_id query parameter is required", nil))
+		return
+	}
+	log = log.WithValues("function_call_id", functionCallID)
+
+	log.V(1).Info("Finding session by parent function call ID")
+	session, err := h.DatabaseService.FindSessionByParentFunctionCallID(functionCallID)
+	if err != nil {
+		if stdErrors.Is(err, gorm.ErrRecordNotFound) {
+			w.RespondWithError(errors.NewNotFoundError("Session not found for function call ID", err))
+			return
+		}
+		w.RespondWithError(errors.NewInternalServerError("Failed to find session by function call ID", err))
+		return
+	}
+
+	log.Info("Successfully found session", "sessionID", session.ID)
+	data := api.NewResponse(session, "Successfully found session", false)
+	RespondWithJSON(w, http.StatusOK, data)
+}
+
 func getUserID(r *http.Request) (string, error) {
 	log := ctrllog.Log.WithName("http-helpers")
 

go/internal/httpserver/handlers/sessions_test.go

@@ -269,11 +269,9 @@ func TestSessionsHandler(t *testing.T) {
 			userID := "test-user"
 			sessionID := "test-session"
 
-			// Create test session
 			agentID := "1"
 			createTestSession(dbClient, sessionID, userID, agentID)
 
-			// Create events with different timestamps
 			event1 := &database.Event{
 				ID:        "event-1",
 				SessionID: sessionID,
@@ -522,4 +520,58 @@ func TestSessionsHandler(t *testing.T) {
 			assert.NotNil(t, responseRecorder.errorReceived)
 		})
 	})
+
+	t.Run("HandleFindSessionByFunctionCall", func(t *testing.T) {
+		t.Run("Success", func(t *testing.T) {
+			handler, dbClient, responseRecorder := setupHandler()
+			userID := "test-user"
+			sessionID := "test-session"
+			functionCallID := "call-123"
+
+			agentID := "1"
+			session := createTestSession(dbClient, sessionID, userID, agentID)
+
+			event := &database.Event{
+				ID:        "event-1",
+				SessionID: sessionID,
+				UserID:    userID,
+				Data:      `{"parent_function_call_id": "` + functionCallID + `"}`,
+			}
+			dbClient.StoreEvents(event)
+
+			req := httptest.NewRequest("GET", "/api/sessions/find-by-function-call?function_call_id="+functionCallID, nil)
+
+			handler.HandleFindSessionByFunctionCall(responseRecorder, req)
+
+			assert.Equal(t, http.StatusOK, responseRecorder.Code)
+
+			var response api.StandardResponse[*database.Session]
+			err := json.Unmarshal(responseRecorder.Body.Bytes(), &response)
+			require.NoError(t, err)
+			assert.Equal(t, session.ID, response.Data.ID)
+		})
+
+		t.Run("SessionNotFound", func(t *testing.T) {
+			handler, _, responseRecorder := setupHandler()
+			functionCallID := "non-existent-call"
+
+			req := httptest.NewRequest("GET", "/api/sessions/find-by-function-call?function_call_id="+functionCallID, nil)
+
+			handler.HandleFindSessionByFunctionCall(responseRecorder, req)
+
+			assert.Equal(t, http.StatusNotFound, responseRecorder.Code)
+			assert.NotNil(t, responseRecorder.errorReceived)
+		})
+
+		t.Run("MissingFunctionCallID", func(t *testing.T) {
+			handler, _, responseRecorder := setupHandler()
+
+			req := httptest.NewRequest("GET", "/api/sessions/find-by-function-call", nil)
+
+			handler.HandleFindSessionByFunctionCall(responseRecorder, req)
+
+			assert.Equal(t, http.StatusBadRequest, responseRecorder.Code)
+			assert.NotNil(t, responseRecorder.errorReceived)
+		})
+	})
 }

go/internal/httpserver/server.go

@@ -167,6 +167,7 @@ func (s *HTTPServer) setupRoutes() {
 	s.router.HandleFunc(APIPathSessions, adaptHandler(s.handlers.Sessions.HandleListSessions)).Methods(http.MethodGet)
 	s.router.HandleFunc(APIPathSessions, adaptHandler(s.handlers.Sessions.HandleCreateSession)).Methods(http.MethodPost)
 	s.router.HandleFunc(APIPathSessions+"/agent/{namespace}/{name}", adaptHandler(s.handlers.Sessions.HandleGetSessionsForAgent)).Methods(http.MethodGet)
+	s.router.HandleFunc(APIPathSessions+"/find-by-function-call", adaptHandler(s.handlers.Sessions.HandleFindSessionByFunctionCall)).Methods(http.MethodGet)
 	s.router.HandleFunc(APIPathSessions+"/{session_id}", adaptHandler(s.handlers.Sessions.HandleGetSession)).Methods(http.MethodGet)
 	s.router.HandleFunc(APIPathSessions+"/{session_id}/tasks", adaptHandler(s.handlers.Sessions.HandleListTasksForSession)).Methods(http.MethodGet)
 	s.router.HandleFunc(APIPathSessions+"/{session_id}", adaptHandler(s.handlers.Sessions.HandleDeleteSession)).Methods(http.MethodDelete)

go/pkg/database/client.go

@@ -55,6 +55,9 @@ type Client interface {
 	ListEventsForSession(sessionID, userID string, options QueryOptions) ([]*Event, error)
 	ListPushNotifications(taskID string) ([]*protocol.TaskPushNotificationConfig, error)
 
+	// Lookup methods
+	FindSessionByParentFunctionCallID(functionCallID string) (*Session, error)
+
 	// Helper methods
 	RefreshToolsForServer(serverName string, groupKind string, tools ...*v1alpha2.MCPTool) error
 

python/packages/kagent-adk/src/kagent/adk/_agent_executor.py

@@ -33,6 +33,7 @@
 from typing_extensions import override
 
 from kagent.core.a2a import TaskResultAggregator, get_kagent_metadata_key
+
 from kagent.core.tracing._span_processor import (
     clear_kagent_span_attributes,
     set_kagent_span_attributes,
@@ -262,11 +263,14 @@ async def _handle_request(
         # ensure the session exists
         session = await self._prepare_session(context, run_args, runner)
 
-        # set request headers to session state
+        # set request headers and A2A request metadata to session state
         headers = context.call_context.state.get("headers", {})
-        state_changes = {
+        state_changes: dict[str, Any] = {
             "headers": headers,
         }
+        request_metadata = context.metadata
+        if request_metadata:
+            state_changes["a2a_request_metadata"] = request_metadata
 
         actions_with_update = EventActions(state_delta=state_changes)
         system_event = Event(

python/packages/kagent-adk/src/kagent/adk/_sub_agent_session_plugin.py

@@ -0,0 +1,141 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Plugin that captures A2A sub-agent session metadata.
+
+When a parent agent delegates to a remote sub-agent via AgentTool, this plugin:
+1. Detects the sub-agent's context_id from A2A event metadata (on_event_callback)
+2. Embeds the context_id in the tool result for historical/stored access (after_tool_callback)
+
+This plugin is automatically propagated to child runners via AgentTool's
+include_plugins=True default, so on_event_callback fires on the child runner's
+events where A2A metadata is present.
+"""
+
+from __future__ import annotations
+
+import contextvars
+import logging
+from dataclasses import dataclass
+from typing import Any, Optional, TYPE_CHECKING
+
+from google.adk.plugins.base_plugin import BasePlugin
+from google.adk.tools.agent_tool import AgentTool
+
+if TYPE_CHECKING:
+    from google.adk.agents.invocation_context import InvocationContext
+    from google.adk.events.event import Event
+    from google.adk.tools.base_tool import BaseTool
+    from google.adk.tools.tool_context import ToolContext
+
+logger = logging.getLogger("kagent_adk." + __name__)
+
+
+@dataclass
+class _ToolCallState:
+    agent_tool_name: str | None = None
+    function_call_id: str | None = None
+    captured_context_id: str | None = None
+    captured_task_id: str | None = None
+
+
+_current_tool_call: contextvars.ContextVar[_ToolCallState | None] = contextvars.ContextVar(
+    "_current_tool_call", default=None
+)
+
+
+def get_current_function_call_id() -> str | None:
+    """Return the function_call_id of the currently executing AgentTool call.
+
+    This is used by the a2a_request_meta_provider to inject the parent's
+    function_call_id into outgoing A2A requests so the sub-agent can store it.
+    """
+    tc = _current_tool_call.get(None)
+    return tc.function_call_id if tc else None
+
+
+class SubAgentSessionPlugin(BasePlugin):
+    """Captures A2A sub-agent session context_id and embeds it in tool results."""
+
+    def __init__(self):
+        super().__init__(name="sub_agent_session")
+
+    async def before_tool_callback(
+        self,
+        *,
+        tool: BaseTool,
+        tool_args: dict[str, Any],
+        tool_context: ToolContext,
+    ) -> Optional[dict]:
+        if isinstance(tool, AgentTool):
+            _current_tool_call.set(
+                _ToolCallState(
+                    agent_tool_name=tool.agent.name if hasattr(tool, "agent") else tool.name,
+                    function_call_id=tool_context.function_call_id,
+                )
+            )
+        return None
+
+    async def on_event_callback(
+        self,
+        *,
+        invocation_context: InvocationContext,
+        event: Event,
+    ) -> Optional[Event]:
+        if not event.custom_metadata:
+            return None
+
+        tc = _current_tool_call.get(None)
+        if tc is None:
+            return None
+
+        context_id = event.custom_metadata.get("a2a:context_id")
+        task_id = event.custom_metadata.get("a2a:task_id")
+
+        if not context_id and not task_id:
+            return None
+
+        if context_id and not tc.captured_context_id:
+            tc.captured_context_id = context_id
+        if task_id and not tc.captured_task_id:
+            tc.captured_task_id = task_id
+
+        return None
+
+    async def after_tool_callback(
+        self,
+        *,
+        tool: BaseTool,
+        tool_args: dict[str, Any],
+        tool_context: ToolContext,
+        result: dict,
+    ) -> Optional[dict]:
+        if not isinstance(tool, AgentTool):
+            return None
+
+        tc = _current_tool_call.get(None)
+        if tc is None or (not tc.captured_context_id and not tc.captured_task_id):
+            return None
+
+        if isinstance(result, str):
+            result = {"result": result}
+
+        if isinstance(result, dict):
+            if tc.captured_context_id:
+                result["a2a:context_id"] = tc.captured_context_id
+            if tc.captured_task_id:
+                result["a2a:task_id"] = tc.captured_task_id
+            return result
+
+        return None