Address code review feedback

- Create constants.py with shared EXTERNAL_TOKEN_SESSION_KEY constant
- Create test_fixtures.py with shared test helper functions
- Update agent.py and agent_to_a2a.py to use shared constant
- Make comment about header provider neutral (not referencing specific header)
- Remove useless test_external_token_stored_in_session test
- Simplify tests by removing unnecessary async decorators
- Move imports to top of test files instead of inside functions
- Update test_a2a_starlette.py to use shared test fixtures
- All 14 tests passing
- All linters passing (ruff, mypy)

Co-authored-by: g3force <779094+g3force@users.noreply.github.com>

Author: Copilot

adk/agenticlayer/agent.py

@@ -7,42 +7,39 @@
 import httpx
 from a2a.client import A2ACardResolver
 from a2a.utils.constants import AGENT_CARD_WELL_KNOWN_PATH
-from google.adk.agents.readonly_context import ReadonlyContext
 from google.adk.agents import BaseAgent, LlmAgent
 from google.adk.agents.llm_agent import ToolUnion
+from google.adk.agents.readonly_context import ReadonlyContext
 from google.adk.agents.remote_a2a_agent import RemoteA2aAgent
 from google.adk.tools.agent_tool import AgentTool
 from google.adk.tools.mcp_tool import StreamableHTTPConnectionParams
 from google.adk.tools.mcp_tool.mcp_toolset import McpToolset
 from httpx_retries import Retry, RetryTransport
 
 from agenticlayer.config import InteractionType, McpTool, SubAgent
+from agenticlayer.constants import EXTERNAL_TOKEN_SESSION_KEY
 
 logger = logging.getLogger(__name__)
 
-# Key used to retrieve the external token from the ADK session state
-# This must match the key used in agent_to_a2a.py
-_EXTERNAL_TOKEN_SESSION_KEY = "__external_token__"
-
 
 def _get_mcp_headers_from_session(readonly_context: ReadonlyContext) -> dict[str, str]:
     """Header provider function for MCP tools that retrieves token from ADK session.
 
     This function is called by the ADK when MCP tools are invoked. It reads the
-    X-External-Token from the session state where it was stored during request
+    external token from the session state where it was stored during request
     processing by TokenCapturingA2aAgentExecutor.
 
     Args:
         readonly_context: The ADK ReadonlyContext providing access to the session
 
     Returns:
         A dictionary of headers to include in MCP tool requests.
-        If a token is stored in the session, includes the X-External-Token header.
+        If a token is stored in the session, includes it in the headers.
     """
     # Access the session state through the readonly context
     # The session state is a dict that can contain the external token
     if readonly_context and readonly_context.session:
-        external_token = readonly_context.session.state.get(_EXTERNAL_TOKEN_SESSION_KEY)
+        external_token = readonly_context.session.state.get(EXTERNAL_TOKEN_SESSION_KEY)
         if external_token:
             return {"X-External-Token": external_token}
     return {}
@@ -132,7 +129,7 @@ def load_tools(self, mcp_tools: list[McpTool]) -> list[ToolUnion]:
                         url=str(tool.url),
                         timeout=tool.timeout,
                     ),
-                    # Pass header provider that retrieves X-External-Token from ADK session
+                    # Provide header provider to inject session-stored token into tool requests
                     header_provider=_get_mcp_headers_from_session,
                 )
             )

adk/agenticlayer/agent_to_a2a.py

@@ -29,12 +29,10 @@
 from .agent import AgentFactory
 from .callback_tracer_plugin import CallbackTracerPlugin
 from .config import McpTool, SubAgent
+from .constants import EXTERNAL_TOKEN_SESSION_KEY
 
 logger = logging.getLogger(__name__)
 
-# Key used to store the external token in the ADK session state
-_EXTERNAL_TOKEN_SESSION_KEY = "__external_token__"
-
 
 class TokenCapturingA2aAgentExecutor(A2aAgentExecutor):
     """Custom A2A agent executor that captures and stores the X-External-Token header.
@@ -76,7 +74,7 @@ async def _prepare_session(
             if external_token:
                 # Store the token in the session state with a private key
                 # The session state is mutable and changes are persisted automatically
-                session.state[_EXTERNAL_TOKEN_SESSION_KEY] = external_token
+                session.state[EXTERNAL_TOKEN_SESSION_KEY] = external_token
                 logger.debug("Stored external token in session %s", session.id)
 
         return session

adk/agenticlayer/constants.py

@@ -0,0 +1,4 @@
+"""Constants shared across the agenticlayer package."""
+
+# Key used to store the external token in the ADK session state
+EXTERNAL_TOKEN_SESSION_KEY = "__external_token__"

adk/tests/test_a2a_starlette.py

@@ -1,5 +1,4 @@
 import contextlib
-import uuid
 from collections.abc import AsyncIterator
 from typing import Any
 
@@ -13,12 +12,13 @@
 from agenticlayer.config import InteractionType, McpTool, SubAgent
 from asgi_lifespan import LifespanManager
 from google.adk.agents.llm_agent import LlmAgent
-from google.adk.models.lite_llm import LiteLlm
 from httpx import Response
 from httpx_retries import Retry
 from pydantic import AnyHttpUrl
 from starlette.testclient import TestClient
 
+from .test_fixtures import create_agent, create_send_message_request
+
 
 def create_mock_agent_card(
     agent_name: str,
@@ -39,39 +39,6 @@ def create_mock_agent_card(
     }
 
 
-def create_send_message_request(
-    message_text: str = "Hello, agent!",
-) -> dict[str, Any]:
-    """Helper function to create a valid A2A send message request."""
-    message_id = str(uuid.uuid4())
-    context_id = str(uuid.uuid4())
-    return {
-        "jsonrpc": "2.0",
-        "id": 1,
-        "method": "message/send",
-        "params": {
-            "message": {
-                "role": "user",
-                "parts": [{"kind": "text", "text": message_text}],
-                "messageId": message_id,
-                "contextId": context_id,
-            },
-            "metadata": {},
-        },
-    }
-
-
-def create_agent(
-    name: str = "test_agent",
-) -> LlmAgent:
-    return LlmAgent(
-        name=name,
-        model=LiteLlm(model="gemini/gemini-2.5-flash"),
-        description="Test agent",
-        instruction="You are a test agent.",
-    )
-
-
 @pytest_asyncio.fixture
 def app_factory() -> Any:
     @contextlib.asynccontextmanager

adk/tests/test_external_token.py

@@ -1,102 +1,19 @@
 """Tests for external token passing to MCP tools via ADK session."""
 
-import uuid
-from typing import Any
-
-import pytest
-from agenticlayer.agent import AgentFactory
-from agenticlayer.agent_to_a2a import _EXTERNAL_TOKEN_SESSION_KEY, to_a2a
-from agenticlayer.config import McpTool
-from asgi_lifespan import LifespanManager
-from google.adk.agents.llm_agent import LlmAgent
-from google.adk.models.lite_llm import LiteLlm
-from httpx_retries import Retry
-from pydantic import AnyHttpUrl
-from starlette.testclient import TestClient
-
-
-def create_send_message_request(
-    message_text: str = "Hello, agent!",
-) -> dict[str, Any]:
-    """Helper function to create a valid A2A send message request."""
-    message_id = str(uuid.uuid4())
-    context_id = str(uuid.uuid4())
-    return {
-        "jsonrpc": "2.0",
-        "id": 1,
-        "method": "message/send",
-        "params": {
-            "message": {
-                "role": "user",
-                "parts": [{"kind": "text", "text": message_text}],
-                "messageId": message_id,
-                "contextId": context_id,
-            },
-            "metadata": {},
-        },
-    }
-
-
-def create_agent(
-    name: str = "test_agent",
-) -> LlmAgent:
-    return LlmAgent(
-        name=name,
-        model=LiteLlm(model="gemini/gemini-2.5-flash"),
-        description="Test agent",
-        instruction="You are a test agent.",
-    )
-
-
-@pytest.mark.asyncio
-async def test_external_token_stored_in_session() -> None:
-    """Test that X-External-Token header is captured and stored in ADK session state."""
-    # Given: An agent with a tool
-    agent = create_agent()
-    tools = [McpTool(name="test_tool", url=AnyHttpUrl("http://tool-1.local/mcp"))]
-    test_token = "test-bearer-token-12345"
-
-    # When: Creating an app and sending a request with X-External-Token header
-    rpc_url = "http://localhost:80/"
-    app = to_a2a(
-        agent=agent,
-        rpc_url=rpc_url,
-        tools=tools,
-        agent_factory=AgentFactory(retry=Retry(total=2)),
-    )
-
-    async with LifespanManager(app) as manager:
-        client = TestClient(manager.app)
-
-        # Send a request with the X-External-Token header
-        response = client.post(
-            "",
-            json=create_send_message_request(),
-            headers={"X-External-Token": test_token},
-        )
+from agenticlayer.agent import _get_mcp_headers_from_session
+from agenticlayer.constants import EXTERNAL_TOKEN_SESSION_KEY
+from google.adk.sessions.session import Session
 
-        # Then: The request should succeed
-        assert response.status_code == 200
 
-        # Note: We cannot directly verify the session state from the test client
-        # because the session is internal to the ADK executor. However, we can
-        # verify that the app starts correctly and processes the request, which
-        # means our custom executor is working.
-
-
-@pytest.mark.asyncio
-async def test_header_provider_retrieves_token_from_session() -> None:
+def test_header_provider_retrieves_token_from_session() -> None:
     """Test that the header provider function can retrieve token from session state."""
-    from agenticlayer.agent import _get_mcp_headers_from_session
-    from google.adk.sessions.session import Session
-
     # Given: A session with an external token stored
     test_token = "test-api-token-xyz"
     session = Session(
         id="test-session",
         app_name="test-app",
         user_id="test-user",
-        state={_EXTERNAL_TOKEN_SESSION_KEY: test_token},
+        state={EXTERNAL_TOKEN_SESSION_KEY: test_token},
         events=[],
         last_update_time=0.0,
     )
@@ -115,12 +32,8 @@ def __init__(self, session):
     assert headers == {"X-External-Token": test_token}
 
 
-@pytest.mark.asyncio
-async def test_header_provider_returns_empty_when_no_token() -> None:
+def test_header_provider_returns_empty_when_no_token() -> None:
     """Test that the header provider returns empty dict when no token is present."""
-    from agenticlayer.agent import _get_mcp_headers_from_session
-    from google.adk.sessions.session import Session
-
     # Given: A session without an external token
     session = Session(
         id="test-session",
@@ -145,13 +58,11 @@ def __init__(self, session):
     assert headers == {}
 
 
-@pytest.mark.asyncio
-async def test_header_provider_handles_none_context() -> None:
+def test_header_provider_handles_none_context() -> None:
     """Test that the header provider safely handles None context."""
-    from agenticlayer.agent import _get_mcp_headers_from_session
-
     # When: Calling the header provider with None
     headers = _get_mcp_headers_from_session(None)
 
     # Then: The headers should be empty (no exception)
     assert headers == {}
+

adk/tests/test_fixtures.py

@@ -0,0 +1,41 @@
+"""Common test fixtures and helper functions for ADK tests."""
+
+import uuid
+from typing import Any
+
+from google.adk.agents.llm_agent import LlmAgent
+from google.adk.models.lite_llm import LiteLlm
+
+
+def create_send_message_request(
+    message_text: str = "Hello, agent!",
+) -> dict[str, Any]:
+    """Helper function to create a valid A2A send message request."""
+    message_id = str(uuid.uuid4())
+    context_id = str(uuid.uuid4())
+    return {
+        "jsonrpc": "2.0",
+        "id": 1,
+        "method": "message/send",
+        "params": {
+            "message": {
+                "role": "user",
+                "parts": [{"kind": "text", "text": message_text}],
+                "messageId": message_id,
+                "contextId": context_id,
+            },
+            "metadata": {},
+        },
+    }
+
+
+def create_agent(
+    name: str = "test_agent",
+) -> LlmAgent:
+    """Helper function to create a test agent."""
+    return LlmAgent(
+        name=name,
+        model=LiteLlm(model="gemini/gemini-2.5-flash"),
+        description="Test agent",
+        instruction="You are a test agent.",
+    )