From 968205df97fca83ab999f53cd18d3f098548a631 Mon Sep 17 00:00:00 2001 From: Dobes Vandermeer Date: Fri, 20 Feb 2026 10:02:49 -0800 Subject: [PATCH] Inject agent session info into context It can be useful for agents to know their own session ID to provide to tools that might want to call back to the agent, generate a link to the agent session, or just have a stable identifier to use for something. Signed-off-by: Dobes Vandermeer --- .../packages/kagent-adk/src/kagent/adk/cli.py | 6 ++- .../kagent-adk/src/kagent/adk/tools/README.md | 1 + .../src/kagent/adk/tools/__init__.py | 3 ++ .../src/kagent/adk/tools/session_tool.py | 51 ++++++++++++++++++ .../src/kagent/adk/tools/skills_toolset.py | 7 ++- .../tests/unittests/test_session_tool.py | 52 +++++++++++++++++++ 6 files changed, 116 insertions(+), 4 deletions(-) create mode 100644 python/packages/kagent-adk/src/kagent/adk/tools/session_tool.py create mode 100644 python/packages/kagent-adk/tests/unittests/test_session_tool.py diff --git a/python/packages/kagent-adk/src/kagent/adk/cli.py b/python/packages/kagent-adk/src/kagent/adk/cli.py index 838d70134..bb75c5e09 100644 --- a/python/packages/kagent-adk/src/kagent/adk/cli.py +++ b/python/packages/kagent-adk/src/kagent/adk/cli.py @@ -15,7 +15,7 @@ from kagent.core import KAgentConfig, configure_logging, configure_tracing from . import AgentConfig, KAgentApp -from .tools import add_skills_tool_to_agent +from .tools import add_session_tool, add_skills_tool_to_agent logger = logging.getLogger(__name__) logging.getLogger("google_adk.google.adk.tools.base_authenticated_tool").setLevel(logging.ERROR) @@ -74,7 +74,7 @@ def static( def root_agent_factory() -> BaseAgent: root_agent = agent_config.to_agent(app_cfg.name, sts_integration) - + add_session_tool(root_agent) maybe_add_skills(root_agent) return root_agent @@ -149,6 +149,7 @@ def root_agent_factory() -> BaseAgent: if sts_integration: add_to_agent(sts_integration, root_agent) + add_session_tool(root_agent) maybe_add_skills(root_agent) return root_agent @@ -213,6 +214,7 @@ async def test_agent(agent_config: AgentConfig, agent_card: AgentCard, task: str def root_agent_factory() -> BaseAgent: root_agent = agent_config.to_agent(app_cfg.name, sts_integration) + add_session_tool(root_agent) maybe_add_skills(root_agent) return root_agent diff --git a/python/packages/kagent-adk/src/kagent/adk/tools/README.md b/python/packages/kagent-adk/src/kagent/adk/tools/README.md index d997d9efb..54131dc9e 100644 --- a/python/packages/kagent-adk/src/kagent/adk/tools/README.md +++ b/python/packages/kagent-adk/src/kagent/adk/tools/README.md @@ -123,6 +123,7 @@ description: Analyze CSV/Excel files | **ReadFile** | Read files with line numbers | `read_file("skills/data-analysis/config.json")` | | **WriteFile** | Create/overwrite files | `write_file("outputs/report.pdf", data)` | | **EditFile** | Precise string replacements | `edit_file("script.py", old="x", new="y")` | +| **SessionInfo**| Auto-injects session details into system instructions | *(automatic — no tool call needed)* | ### Working Directory Structure diff --git a/python/packages/kagent-adk/src/kagent/adk/tools/__init__.py b/python/packages/kagent-adk/src/kagent/adk/tools/__init__.py index 062f17e8e..ba88ac024 100644 --- a/python/packages/kagent-adk/src/kagent/adk/tools/__init__.py +++ b/python/packages/kagent-adk/src/kagent/adk/tools/__init__.py @@ -1,5 +1,6 @@ from .bash_tool import BashTool from .file_tools import EditFileTool, ReadFileTool, WriteFileTool +from .session_tool import SessionInfoTool, add_session_tool from .skill_tool import SkillsTool from .skills_plugin import add_skills_tool_to_agent from .skills_toolset import SkillsToolset @@ -8,8 +9,10 @@ "SkillsTool", "SkillsToolset", "BashTool", + "SessionInfoTool", "EditFileTool", "ReadFileTool", "WriteFileTool", + "add_session_tool", "add_skills_tool_to_agent", ] diff --git a/python/packages/kagent-adk/src/kagent/adk/tools/session_tool.py b/python/packages/kagent-adk/src/kagent/adk/tools/session_tool.py new file mode 100644 index 000000000..bf5b7e116 --- /dev/null +++ b/python/packages/kagent-adk/src/kagent/adk/tools/session_tool.py @@ -0,0 +1,51 @@ +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING + +from google.adk.agents import BaseAgent, LlmAgent +from google.adk.tools.base_tool import BaseTool +from google.adk.tools.tool_context import ToolContext +from typing_extensions import override + +if TYPE_CHECKING: + from google.adk.models.llm_request import LlmRequest + +logger = logging.getLogger("kagent_adk." + __name__) + + +def add_session_tool(agent: BaseAgent) -> None: + if not isinstance(agent, LlmAgent): + return + existing_tool_names = {getattr(t, "name", None) for t in agent.tools} + if "get_session_info" not in existing_tool_names: + agent.tools.append(SessionInfoTool()) + logger.debug(f"Added session info tool to agent: {agent.name}") + + +class SessionInfoTool(BaseTool): + """Tool for retrieving information about the current agent session.""" + + def __init__(self): + super().__init__( + name="get_session_info", + description="Get information about the current agent session, including the session ID.", + ) + + @override + async def process_llm_request( + self, + *, + tool_context: ToolContext, + llm_request: LlmRequest, + ) -> None: + session = tool_context.session + if not session: + return + info = ( + "kagent session:\n" + f"- session_id: {session.id or 'N/A'}\n" + f"- user_id: {session.user_id or 'N/A'}\n" + f"- app_name: {session.app_name or 'N/A'}" + ) + llm_request.append_instructions([info]) diff --git a/python/packages/kagent-adk/src/kagent/adk/tools/skills_toolset.py b/python/packages/kagent-adk/src/kagent/adk/tools/skills_toolset.py index 5653b20dd..9d337065d 100644 --- a/python/packages/kagent-adk/src/kagent/adk/tools/skills_toolset.py +++ b/python/packages/kagent-adk/src/kagent/adk/tools/skills_toolset.py @@ -13,7 +13,7 @@ from google.adk.tools import BaseTool from google.adk.tools.base_toolset import BaseToolset -from ..tools import BashTool, EditFileTool, ReadFileTool, WriteFileTool +from ..tools import BashTool, EditFileTool, ReadFileTool, SessionInfoTool, WriteFileTool from .skill_tool import SkillsTool logger = logging.getLogger("kagent_adk." + __name__) @@ -28,6 +28,7 @@ class SkillsToolset(BaseToolset): 3. WriteFileTool - Write/create files 4. EditFileTool - Edit files with precise replacements 5. BashTool - Execute shell commands + 6. SessionInfoTool - Inject session information into system instructions Skills provide specialized domain knowledge and scripts that the agent can use to solve complex tasks. The toolset enables discovery of available skills, @@ -51,13 +52,14 @@ def __init__(self, skills_directory: str | Path): self.write_file_tool = WriteFileTool() self.edit_file_tool = EditFileTool() self.bash_tool = BashTool(skills_directory) + self.session_info_tool = SessionInfoTool() @override async def get_tools(self, readonly_context: Optional[ReadonlyContext] = None) -> List[BaseTool]: """Get all skills tools. Returns: - List containing all skills tools: skills, read, write, edit, and bash. + List containing all skills tools: skills, read, write, edit, bash, and session info. """ return [ self.skills_tool, @@ -65,4 +67,5 @@ async def get_tools(self, readonly_context: Optional[ReadonlyContext] = None) -> self.write_file_tool, self.edit_file_tool, self.bash_tool, + self.session_info_tool, ] diff --git a/python/packages/kagent-adk/tests/unittests/test_session_tool.py b/python/packages/kagent-adk/tests/unittests/test_session_tool.py new file mode 100644 index 000000000..1e82041be --- /dev/null +++ b/python/packages/kagent-adk/tests/unittests/test_session_tool.py @@ -0,0 +1,52 @@ +# 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. + +from unittest.mock import Mock + +import pytest +from google.adk.models.llm_request import LlmRequest + +from kagent.adk.tools.session_tool import SessionInfoTool + + +class TestSessionInfoTool: + @pytest.mark.asyncio + async def test_session_info_tool(self): + tool = SessionInfoTool() + + context = Mock() + context.session = Mock() + context.session.id = "session-123" + context.session.user_id = "user-456" + context.session.app_name = "test-app" + + llm_request = LlmRequest() + await tool.process_llm_request(tool_context=context, llm_request=llm_request) + + assert "session-123" in llm_request.config.system_instruction + assert "user-456" in llm_request.config.system_instruction + assert "test-app" in llm_request.config.system_instruction + assert "get_session_info" not in llm_request.tools_dict + + @pytest.mark.asyncio + async def test_session_info_tool_none_session(self): + tool = SessionInfoTool() + + context = Mock() + context.session = None + + llm_request = LlmRequest() + await tool.process_llm_request(tool_context=context, llm_request=llm_request) + + assert llm_request.config.system_instruction is None