From 0a1d3ea9b83190dd4f4a8d3ff45c9860981e665a Mon Sep 17 00:00:00 2001 From: "hanzhi.421" Date: Mon, 24 Aug 2026 16:36:54 +0800 Subject: [PATCH] feat(a2a): add session-aware AgentKit skill agent --- tests/a2a/test_agentkit_remote_skill_agent.py | 419 +++++++++++++++ tests/tools/builtin_tools/test_agentkit.py | 122 +++++ veadk/a2a/__init__.py | 5 + veadk/a2a/agentkit_remote_skill_agent.py | 486 ++++++++++++++++++ veadk/a2a/remote_ve_agent.py | 60 ++- veadk/tools/builtin_tools/_agentkit.py | 471 +++++++++++++---- 6 files changed, 1428 insertions(+), 135 deletions(-) create mode 100644 tests/a2a/test_agentkit_remote_skill_agent.py create mode 100644 veadk/a2a/agentkit_remote_skill_agent.py diff --git a/tests/a2a/test_agentkit_remote_skill_agent.py b/tests/a2a/test_agentkit_remote_skill_agent.py new file mode 100644 index 000000000..563ed643c --- /dev/null +++ b/tests/a2a/test_agentkit_remote_skill_agent.py @@ -0,0 +1,419 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# 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 __future__ import annotations + +import asyncio +import json +from types import SimpleNamespace +from unittest.mock import AsyncMock, patch + +import pytest +import requests +from a2a.types import Message, Task, TaskState, TaskStatus +from google.adk.events import Event + +from veadk.a2a.agentkit_remote_skill_agent import ( + _AGENTKIT_SESSION_ID_METADATA_KEY, + _SessionBoundRemoteSkillAgent, + AgentkitRemoteSkillAgent, +) +from veadk.tools.builtin_tools._agentkit import AgentKitSessionLease + + +def _lease(session_id: str, endpoint: str) -> AgentKitSessionLease: + return AgentKitSessionLease( + tool_id="tool-1", + logical_user_session_id="skills_user_session", + user_session_id="skills_user_session", + session_id=session_id, + status="Ready", + endpoint=endpoint, + internal_endpoint="", + created_at="2026-08-24T01:00:00+00:00", + expire_at="2026-08-24T02:00:00+00:00", + ) + + +def _agent_card() -> dict: + return { + "name": "skill-sandbox", + "description": "Skill sandbox", + "url": "http://127.0.0.1:8000/a2a", + "version": "1.0.0", + "capabilities": {}, + "defaultInputModes": ["text/plain"], + "defaultOutputModes": ["text/plain"], + "skills": [], + } + + +def _response( + status_code: int, + *, + json_body: object | None = None, + body: bytes = b"", + content_type: str = "application/json", +) -> requests.Response: + response = requests.Response() + response.status_code = status_code + response.headers["Content-Type"] = content_type + response._content = ( + json.dumps(json_body).encode("utf-8") if json_body is not None else body + ) + return response + + +def test_constructor_does_not_resolve_agentkit_session() -> None: + with patch( + "veadk.a2a.agentkit_remote_skill_agent.ensure_agentkit_session_lease" + ) as ensure: + AgentkitRemoteSkillAgent(name="skills", tool_id="tool-1") + + ensure.assert_not_called() + + +def test_resolve_lease_uses_adk_session_as_logical_key() -> None: + captured = {} + expected = _lease("session-1", "https://sandbox.test") + + def fake_ensure(**kwargs): + captured.update(kwargs) + return expected + + ctx = SimpleNamespace( + user_id="user-1", + session=SimpleNamespace(id="adk-session-1", state={"key": "value"}), + ) + agent = AgentkitRemoteSkillAgent( + name="skills", + tool_id="tool-1", + request_timeout=900, + expiry_buffer=90, + ) + + with patch( + "veadk.a2a.agentkit_remote_skill_agent.ensure_agentkit_session_lease", + side_effect=fake_ensure, + ): + result = asyncio.run(agent._resolve_lease(ctx)) # type: ignore[arg-type] + + assert result is expected + assert captured["tool_user_session_id"] == "skills_user-1_adk-session-1" + assert captured["min_remaining_seconds"] == 990 + assert captured["tool_state"] == {"key": "value"} + + +def test_delegate_is_rebuilt_when_physical_session_changes() -> None: + created = [] + + class FakeDelegate: + def __init__(self, **kwargs): + created.append(kwargs) + self._httpx_client = SimpleNamespace(headers={}) + + async def cleanup(self): + return None + + ctx = SimpleNamespace(credential_service=None) + agent = AgentkitRemoteSkillAgent(name="skills", tool_id="tool-1") + first = _lease( + "session-1", + "https://sandbox.test/?faasInstanceName=first", + ) + second = _lease( + "session-2", + "https://sandbox.test/?faasInstanceName=second", + ) + + async def run(): + with ( + patch( + "veadk.a2a.agentkit_remote_skill_agent._SessionBoundRemoteSkillAgent", + FakeDelegate, + ), + patch.object( + agent, + "_wait_for_agent_card", + new=AsyncMock(), + ) as wait_for_agent_card, + ): + first_delegate = await agent._delegate_for(first, ctx) # type: ignore[arg-type] + reused_delegate = await agent._delegate_for(first, ctx) # type: ignore[arg-type] + second_delegate = await agent._delegate_for(second, ctx) # type: ignore[arg-type] + return first_delegate, reused_delegate, second_delegate, wait_for_agent_card + + first_delegate, reused_delegate, second_delegate, wait_for_agent_card = asyncio.run( + run() + ) + + assert first_delegate is reused_delegate + assert second_delegate is not first_delegate + assert len(created) == 2 + assert created[0]["rpc_url"] == ("https://sandbox.test/a2a?faasInstanceName=first") + assert created[1]["rpc_url"] == ("https://sandbox.test/a2a?faasInstanceName=second") + assert wait_for_agent_card.await_count == 2 + + +def test_wait_for_agent_card_retries_502_and_preserves_endpoint_query() -> None: + agent = AgentkitRemoteSkillAgent( + name="skills", + tool_id="tool-1", + a2a_ready_timeout=10, + a2a_ready_poll_interval=0.01, + ) + responses = [ + _response(502, body=b"bad gateway", content_type="text/html"), + _response(200, json_body=_agent_card()), + ] + + async def run() -> None: + with ( + patch( + "veadk.a2a.agentkit_remote_skill_agent.requests.get", + side_effect=responses, + ) as get, + patch( + "veadk.a2a.agentkit_remote_skill_agent.asyncio.sleep", + new=AsyncMock(), + ) as sleep, + ): + await agent._wait_for_agent_card( + endpoint="https://sandbox.test/?faasInstanceName=instance-1", + headers={"inbound_auth": "token"}, + ) + + assert get.call_count == 2 + assert get.call_args_list[0].args[0] == ( + "https://sandbox.test/.well-known/agent-card.json" + "?faasInstanceName=instance-1" + ) + assert get.call_args_list[0].kwargs["headers"] == {"inbound_auth": "token"} + sleep.assert_awaited_once() + + asyncio.run(run()) + + +def test_wait_for_agent_card_retries_non_json_response() -> None: + agent = AgentkitRemoteSkillAgent( + name="skills", + tool_id="tool-1", + a2a_ready_timeout=10, + a2a_ready_poll_interval=0.01, + ) + + async def run() -> None: + with ( + patch( + "veadk.a2a.agentkit_remote_skill_agent.requests.get", + side_effect=[ + _response(200, body=b"starting", content_type="text/plain"), + _response(200, json_body=_agent_card()), + ], + ), + patch( + "veadk.a2a.agentkit_remote_skill_agent.asyncio.sleep", + new=AsyncMock(), + ) as sleep, + ): + await agent._wait_for_agent_card( + endpoint="https://sandbox.test", + headers={}, + ) + + sleep.assert_awaited_once() + + asyncio.run(run()) + + +def test_wait_for_agent_card_fails_fast_for_404() -> None: + agent = AgentkitRemoteSkillAgent( + name="skills", + tool_id="tool-1", + a2a_ready_timeout=10, + ) + + async def run() -> None: + with patch( + "veadk.a2a.agentkit_remote_skill_agent.requests.get", + return_value=_response(404, body=b"not found", content_type="text/plain"), + ) as get: + with pytest.raises(RuntimeError, match=r"HTTP 404"): + await agent._wait_for_agent_card( + endpoint="https://sandbox.test", + headers={}, + ) + get.assert_called_once() + + asyncio.run(run()) + + +def test_wait_for_agent_card_timeout_reports_redacted_response_summary() -> None: + agent = AgentkitRemoteSkillAgent( + name="skills", + tool_id="tool-1", + a2a_ready_timeout=5, + ) + + async def run() -> None: + with ( + patch( + "veadk.a2a.agentkit_remote_skill_agent.requests.get", + return_value=_response( + 503, + body=b"temporarily unavailable", + content_type="text/plain", + ), + ), + patch( + "veadk.a2a.agentkit_remote_skill_agent._monotonic", + side_effect=[0, 0, 6], + ), + ): + with pytest.raises( + TimeoutError, + match=( + r"HTTP 503, content-type=text/plain, " + r"body-bytes=23" + ), + ): + await agent._wait_for_agent_card( + endpoint="https://sandbox.test/?Authorization=secret", + headers={}, + ) + + asyncio.run(run()) + + +def test_session_bound_delegate_rejects_stale_a2a_context() -> None: + response = SimpleNamespace(json=lambda: _agent_card()) + with patch("veadk.a2a.remote_ve_agent.requests.get", return_value=response): + delegate = _SessionBoundRemoteSkillAgent( + agentkit_session_id="session-current", + poll_interval=2, + max_poll_interval=16, + name="skills", + url="https://sandbox.test", + rpc_url="https://sandbox.test/a2a", + ) + + current = Event( + author="skills", + custom_metadata={ + "a2a:response": True, + _AGENTKIT_SESSION_ID_METADATA_KEY: "session-current", + }, + ) + stale = Event( + author="skills", + custom_metadata={ + "a2a:response": True, + _AGENTKIT_SESSION_ID_METADATA_KEY: "session-old", + }, + ) + + assert delegate._is_remote_response(current) + assert not delegate._is_remote_response(stale) + asyncio.run(delegate.cleanup()) + + +def test_session_bound_delegate_polls_non_blocking_task_to_completion() -> None: + response = SimpleNamespace(json=lambda: _agent_card()) + working = Task( + id="task-1", + contextId="context-1", + status=TaskStatus(state=TaskState.working), + ) + completed = Task( + id="task-1", + contextId="context-1", + status=TaskStatus(state=TaskState.completed), + ) + + async def run(): + with patch("veadk.a2a.remote_ve_agent.requests.get", return_value=response): + delegate = _SessionBoundRemoteSkillAgent( + agentkit_session_id="session-current", + poll_interval=2, + max_poll_interval=16, + name="skills", + url="https://sandbox.test", + rpc_url="https://sandbox.test/a2a", + ) + await delegate._ensure_resolved() + client = delegate._a2a_client + client._transport.send_message = AsyncMock(return_value=working) + client.get_task = AsyncMock(return_value=completed) + delegate._configure_polling_client() + + with patch( + "veadk.a2a.agentkit_remote_skill_agent.asyncio.sleep", + new=AsyncMock(), + ): + results = [ + item + async for item in client.send_message( + request=Message( + messageId="message-1", + role="user", + parts=[{"kind": "text", "text": "run skill"}], + ), + request_metadata={"user_id": "user-1"}, + ) + ] + await delegate.cleanup() + return client, results + + client, results = asyncio.run(run()) + + assert client._config.polling is True + assert results == [(completed, None)] + send_params = client._transport.send_message.await_args.args[0] + assert send_params.configuration.blocking is False + assert send_params.metadata == {"user_id": "user-1"} + client.get_task.assert_awaited_once() + query = client.get_task.await_args.args[0] + assert query.id == "task-1" + assert query.history_length == 20 + + +def test_run_tags_events_with_physical_agentkit_session() -> None: + lease = _lease("session-1", "https://sandbox.test") + + class FakeDelegate: + async def _run_async_impl(self, _ctx): + yield Event(author="skills") + + async def fake_resolve(self, _ctx): + return lease + + async def fake_delegate(self, _lease, _ctx): + return FakeDelegate() + + agent = AgentkitRemoteSkillAgent(name="skills", tool_id="tool-1") + ctx = SimpleNamespace(invocation_id="invocation-1", branch=None) + + async def run(): + with ( + patch.object(AgentkitRemoteSkillAgent, "_resolve_lease", fake_resolve), + patch.object(AgentkitRemoteSkillAgent, "_delegate_for", fake_delegate), + ): + return [ + event + async for event in agent._run_async_impl(ctx) # type: ignore[arg-type] + ] + + events = asyncio.run(run()) + + assert events[0].custom_metadata == {_AGENTKIT_SESSION_ID_METADATA_KEY: "session-1"} diff --git a/tests/tools/builtin_tools/test_agentkit.py b/tests/tools/builtin_tools/test_agentkit.py index 76f26b837..f7d07cdcc 100644 --- a/tests/tools/builtin_tools/test_agentkit.py +++ b/tests/tools/builtin_tools/test_agentkit.py @@ -18,6 +18,7 @@ import sys import types import unittest +from datetime import datetime, timedelta, timezone from pathlib import Path from unittest.mock import patch @@ -553,6 +554,127 @@ def get_session(self, _request): wait_until_ready=True, ) + def test_reuses_ready_session_with_enough_remaining_lifetime(self): + expires_at = (datetime.now(timezone.utc) + timedelta(hours=1)).isoformat() + reusable = types.SimpleNamespace( + session_id="session-existing", + user_session_id="user-session-1", + status="Ready", + created_at="2026-08-24T01:00:00+00:00", + expire_at=expires_at, + ) + + class FakeClient: + def list_sessions(self, request): + self.request = request + return types.SimpleNamespace(session_infos=[reusable], next_token=None) + + def create_session(self, _request): + raise AssertionError("a reusable Session must not be replaced") + + result = self.agentkit_module._get_or_create_agentkit_session( + client=FakeClient(), + tool_id="tool-1", + tool_user_session_id="user-session-1", + ttl=1800, + min_remaining_seconds=600, + ) + + self.assertIs(result, reusable) + + def test_rotates_near_expiry_session_with_physical_user_session_id(self): + expires_at = (datetime.now(timezone.utc) + timedelta(seconds=30)).isoformat() + old_session = types.SimpleNamespace( + session_id="session-old", + user_session_id="user-session-1", + status="Ready", + created_at="2026-08-24T01:00:00+00:00", + expire_at=expires_at, + ) + captured = {} + + class FakeClient: + def list_sessions(self, _request): + return types.SimpleNamespace( + session_infos=[old_session], next_token=None + ) + + def create_session(self, request): + captured["user_session_id"] = request.user_session_id + return types.SimpleNamespace( + session_id="session-new", + user_session_id=request.user_session_id, + ) + + result = self.agentkit_module._get_or_create_agentkit_session( + client=FakeClient(), + tool_id="tool-1", + tool_user_session_id="user-session-1", + ttl=1800, + min_remaining_seconds=600, + ) + + self.assertEqual(result.session_id, "session-new") + self.assertRegex( + captured["user_session_id"], r"^user-session-1_r_[0-9a-f]{12}$" + ) + + def test_list_sessions_follows_next_token(self): + requests = [] + matching = types.SimpleNamespace( + session_id="session-2", + user_session_id="user-session-1_r_123456789abc", + ) + + class FakeClient: + def list_sessions(self, request): + requests.append(request) + if len(requests) == 1: + return types.SimpleNamespace(session_infos=[], next_token="page-2") + return types.SimpleNamespace(session_infos=[matching], next_token=None) + + result = self.agentkit_module._list_agentkit_sessions( + client=FakeClient(), + tool_id="tool-1", + physical_user_session_id_base="user-session-1", + ) + + self.assertEqual(result, [matching]) + self.assertEqual(len(requests), 2) + self.assertIsNone(requests[0].next_token) + self.assertEqual(requests[1].next_token, "page-2") + + def test_recovers_session_after_ambiguous_create_failure(self): + list_calls = 0 + recovered = types.SimpleNamespace( + session_id="session-created", + user_session_id="user-session-1", + status="Starting", + created_at="2026-08-24T01:00:00+00:00", + expire_at=(datetime.now(timezone.utc) + timedelta(minutes=30)).isoformat(), + ) + + class FakeClient: + def list_sessions(self, _request): + nonlocal list_calls + list_calls += 1 + sessions = [] if list_calls == 1 else [recovered] + return types.SimpleNamespace(session_infos=sessions, next_token=None) + + def create_session(self, _request): + raise TimeoutError("CreateSession response was lost") + + result = self.agentkit_module._get_or_create_agentkit_session( + client=FakeClient(), + tool_id="tool-1", + tool_user_session_id="user-session-1", + ttl=1800, + min_remaining_seconds=600, + ) + + self.assertIs(result, recovered) + self.assertEqual(list_calls, 2) + if __name__ == "__main__": unittest.main() diff --git a/veadk/a2a/__init__.py b/veadk/a2a/__init__.py index 7f463206f..d0ee22019 100644 --- a/veadk/a2a/__init__.py +++ b/veadk/a2a/__init__.py @@ -11,3 +11,8 @@ # 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 veadk.a2a.agentkit_remote_skill_agent import AgentkitRemoteSkillAgent +from veadk.a2a.remote_ve_agent import RemoteVeAgent + +__all__ = ["AgentkitRemoteSkillAgent", "RemoteVeAgent"] diff --git a/veadk/a2a/agentkit_remote_skill_agent.py b/veadk/a2a/agentkit_remote_skill_agent.py new file mode 100644 index 000000000..3fa3dcd3b --- /dev/null +++ b/veadk/a2a/agentkit_remote_skill_agent.py @@ -0,0 +1,486 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# 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. + +"""A Session-aware A2A agent for ephemeral AgentKit Skill sandboxes.""" + +from __future__ import annotations + +import asyncio +import time +from types import MethodType +from typing import AsyncGenerator, Literal, Optional + +import requests +from a2a.types import ( + AgentCard, + MessageSendConfiguration, + MessageSendParams, + Task, + TaskQueryParams, + TaskState, +) +from google.adk.agents.base_agent import BaseAgent +from google.adk.agents.callback_context import CallbackContext +from google.adk.agents.invocation_context import InvocationContext +from google.adk.events.event import Event +from google.adk.utils.context_utils import Aclosing +from pydantic import PrivateAttr + +from veadk.a2a.remote_ve_agent import RemoteVeAgent, _url_with_path +from veadk.tools.builtin_tools._agentkit import ( + AgentKitSessionLease, + ensure_agentkit_session_lease, + resolve_agentkit_tool_id, +) +from veadk.utils.auth import build_auth_config +from veadk.utils.logger import get_logger + + +logger = get_logger(__name__) + +_AGENTKIT_SESSION_ID_METADATA_KEY = "veadk:agentkit_session_id" +_INBOUND_AUTH_HEADER = "inbound_auth" +_A2A_HISTORY_LENGTH = 20 +_A2A_POLLING_STATES = frozenset({TaskState.submitted, TaskState.working}) +_AGENT_CARD_PATH = "/.well-known/agent-card.json" +_AGENT_CARD_RETRY_STATUS_CODES = frozenset({502, 503, 504}) +_AGENT_CARD_REQUEST_TIMEOUT = 10.0 +_monotonic = time.monotonic + + +def _credential_token_value(credential: object | None) -> str | None: + if credential is None: + return None + api_key = getattr(credential, "api_key", None) + if api_key: + return str(api_key) + http = getattr(credential, "http", None) + http_credentials = getattr(http, "credentials", None) if http else None + http_token = getattr(http_credentials, "token", None) if http_credentials else None + if http_token: + return str(http_token) + oauth2 = getattr(credential, "oauth2", None) + access_token = getattr(oauth2, "access_token", None) if oauth2 else None + return str(access_token) if access_token else None + + +async def _inbound_auth_token(ctx: InvocationContext) -> str | None: + if not ctx.credential_service: + return None + auth_config = build_auth_config( + credential_key="inbound_auth", + auth_method="header", + header_scheme="bearer", + ) + credential = await ctx.credential_service.load_credential( + auth_config=auth_config, + callback_context=CallbackContext(ctx), + ) + return _credential_token_value(credential) + + +def _agentkit_request_metadata(ctx: InvocationContext, _message) -> dict[str, str]: + return { + "user_id": ctx.user_id, + "session_id": ctx.session.id, + } + + +class _SessionBoundRemoteSkillAgent(RemoteVeAgent): + """RemoteVeAgent that only resumes A2A context from its physical Session.""" + + _bound_agentkit_session_id: str = PrivateAttr() + _poll_interval: float = PrivateAttr() + _max_poll_interval: float = PrivateAttr() + _polling_configured: bool = PrivateAttr(default=False) + + def __init__( + self, + *, + agentkit_session_id: str, + poll_interval: float, + max_poll_interval: float, + **kwargs, + ): + super().__init__(**kwargs) + self._a2a_request_meta_provider = _agentkit_request_metadata + self._bound_agentkit_session_id = agentkit_session_id + self._poll_interval = poll_interval + self._max_poll_interval = max_poll_interval + + def _is_remote_response(self, event: Event) -> bool: + if not super()._is_remote_response(event): + return False + metadata = event.custom_metadata or {} + return ( + metadata.get(_AGENTKIT_SESSION_ID_METADATA_KEY) + == self._bound_agentkit_session_id + ) + + async def _pre_run(self, ctx: InvocationContext) -> None: + await super()._pre_run(ctx) + self._configure_polling_client() + + def _configure_polling_client(self) -> None: + """Adapt ADK's A2A client to Skill Sandbox non-blocking task execution.""" + if self._polling_configured: + return + client = self._a2a_client + client._config.polling = True + request_timeout = float(self._timeout) + poll_interval = self._poll_interval + max_poll_interval = self._max_poll_interval + + async def send_message_with_polling( + _client, + request, + *, + context=None, + request_metadata=None, + ): + deadline = asyncio.get_running_loop().time() + request_timeout + configuration = MessageSendConfiguration( + accepted_output_modes=client._config.accepted_output_modes, + blocking=False, + push_notification_config=( + client._config.push_notification_configs[0] + if client._config.push_notification_configs + else None + ), + ) + params = MessageSendParams( + message=request, + configuration=configuration, + metadata=request_metadata, + ) + response = await client._transport.send_message( + params, + context=context, + ) + result = (response, None) if isinstance(response, Task) else response + await client.consume(result, client._card) + if not isinstance(response, Task) or ( + response.status.state not in _A2A_POLLING_STATES + ): + yield result + return + + current_task = response + current_interval = poll_interval + while current_task.status.state in _A2A_POLLING_STATES: + remaining = deadline - asyncio.get_running_loop().time() + if remaining <= 0: + raise TimeoutError( + f"Timed out while waiting for A2A task {response.id}" + ) + await asyncio.sleep(min(current_interval, remaining)) + current_task = await client.get_task( + TaskQueryParams( + id=response.id, + historyLength=_A2A_HISTORY_LENGTH, + ), + context=context, + ) + current_interval = min( + current_interval * 2, + max_poll_interval, + ) + await client.consume((current_task, None), client._card) + yield current_task, None + + client.send_message = MethodType(send_message_with_polling, client) + self._polling_configured = True + + +class AgentkitRemoteSkillAgent(BaseAgent): + """Connect to an A2A agent hosted by an expiring AgentKit Session. + + ``tool_user_session_id`` is a stable logical key. The corresponding physical + AgentKit Session, endpoint, Agent Card, and A2A client are resolved lazily and + replaced when the Session no longer has enough remaining lifetime. + """ + + tool_id: Optional[str] = None + tool_user_session_id: Optional[str] = None + ttl: int = 1800 + request_timeout: int = 1800 + expiry_buffer: int = 60 + ready_timeout: float = 120 + a2a_ready_timeout: float = 120 + a2a_ready_poll_interval: float = 2 + poll_interval: float = 2 + max_poll_interval: float = 16 + prefer_internal_endpoint: bool = False + rpc_path: str = "/a2a" + auth_method: Literal["header", "querystring"] | None = "header" + + _delegates: dict[str, _SessionBoundRemoteSkillAgent] = PrivateAttr( + default_factory=dict + ) + _leases: dict[str, AgentKitSessionLease] = PrivateAttr(default_factory=dict) + _delegate_use_counts: dict[str, int] = PrivateAttr(default_factory=dict) + _delegate_lock: asyncio.Lock = PrivateAttr(default_factory=asyncio.Lock) + + def __init__( + self, + name: str, + *, + tool_id: Optional[str] = None, + tool_user_session_id: Optional[str] = None, + description: str = "", + ttl: int = 1800, + request_timeout: int = 1800, + expiry_buffer: int = 60, + ready_timeout: float = 120, + a2a_ready_timeout: float = 120, + a2a_ready_poll_interval: float = 2, + poll_interval: float = 2, + max_poll_interval: float = 16, + prefer_internal_endpoint: bool = False, + rpc_path: str = "/a2a", + auth_method: Literal["header", "querystring"] | None = "header", + **kwargs, + ) -> None: + super().__init__(name=name, description=description, **kwargs) + if not 60 <= ttl <= 86400: + raise ValueError("ttl must be between 60 and 86400 seconds") + if request_timeout <= 0: + raise ValueError("request_timeout must be greater than 0") + if expiry_buffer < 0: + raise ValueError("expiry_buffer must be greater than or equal to 0") + if request_timeout + expiry_buffer >= 86400: + raise ValueError( + "request_timeout plus expiry_buffer must be less than 86400 seconds" + ) + if ready_timeout < 0: + raise ValueError("ready_timeout must be greater than or equal to 0") + if a2a_ready_timeout <= 0: + raise ValueError("a2a_ready_timeout must be greater than 0") + if a2a_ready_poll_interval <= 0: + raise ValueError("a2a_ready_poll_interval must be greater than 0") + if poll_interval <= 0: + raise ValueError("poll_interval must be greater than 0") + if max_poll_interval < poll_interval: + raise ValueError( + "max_poll_interval must be greater than or equal to poll_interval" + ) + if not rpc_path.strip(): + raise ValueError("rpc_path must not be empty") + + self.tool_id = tool_id + self.tool_user_session_id = tool_user_session_id + self.ttl = ttl + self.request_timeout = request_timeout + self.expiry_buffer = expiry_buffer + self.ready_timeout = ready_timeout + self.a2a_ready_timeout = a2a_ready_timeout + self.a2a_ready_poll_interval = a2a_ready_poll_interval + self.poll_interval = poll_interval + self.max_poll_interval = max_poll_interval + self.prefer_internal_endpoint = prefer_internal_endpoint + self.rpc_path = rpc_path + self.auth_method = auth_method + + def _logical_user_session_id(self, ctx: InvocationContext) -> str: + if self.tool_user_session_id: + return self.tool_user_session_id + return f"{self.name}_{ctx.user_id}_{ctx.session.id}" + + async def _resolve_lease(self, ctx: InvocationContext) -> AgentKitSessionLease: + tool_id = self.tool_id or resolve_agentkit_tool_id("AGENTKIT_TOOL_ID_SKILLS") + session_state = getattr(ctx.session, "state", None) + tool_state = dict(session_state) if isinstance(session_state, dict) else None + return await asyncio.to_thread( + ensure_agentkit_session_lease, + tool_id=tool_id, + tool_user_session_id=self._logical_user_session_id(ctx), + tool_state=tool_state, + ttl=self.ttl, + min_remaining_seconds=self.request_timeout + self.expiry_buffer, + wait_until_ready=True, + ready_timeout=self.ready_timeout, + ) + + @staticmethod + def _agent_card_response_summary(response: requests.Response) -> str: + content_type = response.headers.get("Content-Type", "unknown").split(";", 1)[0] + return ( + f"HTTP {response.status_code}, content-type={content_type}, " + f"body-bytes={len(response.content)}" + ) + + async def _wait_for_agent_card( + self, + *, + endpoint: str, + headers: dict[str, str], + ) -> None: + """Wait until the Session data plane serves a valid A2A Agent Card.""" + url = _url_with_path(endpoint, _AGENT_CARD_PATH) + deadline = _monotonic() + self.a2a_ready_timeout + last_result = "no response" + + while True: + remaining = deadline - _monotonic() + if remaining <= 0: + raise TimeoutError( + "Timed out waiting for AgentKit A2A Agent Card; " + f"last result: {last_result}" + ) + + try: + response = await asyncio.to_thread( + requests.get, + url, + headers=headers, + timeout=min(_AGENT_CARD_REQUEST_TIMEOUT, remaining), + ) + except requests.RequestException as exc: + last_result = f"request failed ({type(exc).__name__})" + else: + last_result = self._agent_card_response_summary(response) + if response.status_code == 200: + try: + AgentCard.model_validate(response.json()) + except ValueError: + last_result = f"{last_result}, invalid JSON" + else: + logger.debug("AgentKit A2A Agent Card is ready") + return + elif response.status_code not in _AGENT_CARD_RETRY_STATUS_CODES: + raise RuntimeError( + f"AgentKit A2A Agent Card request failed: {last_result}" + ) + + remaining = deadline - _monotonic() + if remaining <= 0: + raise TimeoutError( + "Timed out waiting for AgentKit A2A Agent Card; " + f"last result: {last_result}" + ) + await asyncio.sleep(min(self.a2a_ready_poll_interval, remaining)) + + async def _delegate_for( + self, + lease: AgentKitSessionLease, + ctx: InvocationContext, + ) -> _SessionBoundRemoteSkillAgent: + inbound_auth = await _inbound_auth_token(ctx) + endpoint = lease.select_endpoint( + prefer_internal_endpoint=self.prefer_internal_endpoint + ) + if not endpoint: + raise RuntimeError(f"AgentKit session {lease.session_id} has no endpoint") + extra_headers = {_INBOUND_AUTH_HEADER: inbound_auth} if inbound_auth else {} + + async with self._delegate_lock: + needs_agent_card = lease.session_id not in self._delegates + if needs_agent_card: + await self._wait_for_agent_card( + endpoint=endpoint, + headers=extra_headers, + ) + + stale_delegate: _SessionBoundRemoteSkillAgent | None = None + async with self._delegate_lock: + delegate = self._delegates.get(lease.session_id) + if delegate is None: + delegate = await asyncio.to_thread( + _SessionBoundRemoteSkillAgent, + agentkit_session_id=lease.session_id, + poll_interval=self.poll_interval, + max_poll_interval=self.max_poll_interval, + name=self.name, + url=endpoint, + rpc_url=_url_with_path(endpoint, self.rpc_path), + auth_method=self.auth_method, + extra_headers=extra_headers, + timeout=float(self.request_timeout), + ) + self._delegates[lease.session_id] = delegate + logger.info( + "Bound AgentKit A2A agent %s to Session %s", + self.name, + lease.session_id, + ) + elif inbound_auth: + delegate._httpx_client.headers[_INBOUND_AUTH_HEADER] = inbound_auth + else: + delegate._httpx_client.headers.pop(_INBOUND_AUTH_HEADER, None) + + previous_lease = self._leases.get(lease.logical_user_session_id) + self._leases[lease.logical_user_session_id] = lease + self._delegate_use_counts[lease.session_id] = ( + self._delegate_use_counts.get(lease.session_id, 0) + 1 + ) + if ( + previous_lease + and previous_lease.session_id != lease.session_id + and self._delegate_use_counts.get(previous_lease.session_id, 0) == 0 + ): + stale_delegate = self._delegates.pop(previous_lease.session_id, None) + + if stale_delegate: + await stale_delegate.cleanup() + return delegate + + async def _release_delegate(self, session_id: str) -> None: + stale_delegate: _SessionBoundRemoteSkillAgent | None = None + async with self._delegate_lock: + remaining = max(self._delegate_use_counts.get(session_id, 1) - 1, 0) + if remaining: + self._delegate_use_counts[session_id] = remaining + return + self._delegate_use_counts.pop(session_id, None) + if all(lease.session_id != session_id for lease in self._leases.values()): + stale_delegate = self._delegates.pop(session_id, None) + + if stale_delegate: + await stale_delegate.cleanup() + + async def _run_async_impl( + self, + ctx: InvocationContext, + ) -> AsyncGenerator[Event, None]: + try: + lease = await self._resolve_lease(ctx) + delegate = await self._delegate_for(lease, ctx) + except Exception as exc: + yield Event( + author=self.name, + error_message=f"Failed to initialize AgentKit A2A Session: {exc}", + invocation_id=ctx.invocation_id, + branch=ctx.branch, + ) + return + + try: + async with Aclosing(delegate._run_async_impl(ctx)) as agen: + async for event in agen: + event.custom_metadata = event.custom_metadata or {} + event.custom_metadata[_AGENTKIT_SESSION_ID_METADATA_KEY] = ( + lease.session_id + ) + yield event + finally: + await self._release_delegate(lease.session_id) + + async def cleanup(self) -> None: + """Close every A2A HTTP client created for physical Sessions.""" + delegates = list(self._delegates.values()) + self._delegates.clear() + self._leases.clear() + self._delegate_use_counts.clear() + for delegate in delegates: + await delegate.cleanup() diff --git a/veadk/a2a/remote_ve_agent.py b/veadk/a2a/remote_ve_agent.py index 1eddedaa9..d21f02111 100644 --- a/veadk/a2a/remote_ve_agent.py +++ b/veadk/a2a/remote_ve_agent.py @@ -15,6 +15,7 @@ import json import functools from typing import AsyncGenerator, Literal, Optional +from urllib.parse import urlsplit, urlunsplit from a2a.client.base_client import BaseClient import httpx @@ -35,6 +36,17 @@ AGENT_CARD_WELL_KNOWN_PATH = "/.well-known/agent-card.json" +def _url_with_path(base_url: str, path: str) -> str: + """Append a path while preserving Session endpoint query authentication.""" + parts = urlsplit(base_url) + base_path = parts.path.rstrip("/") + relative_path = path.lstrip("/") + joined_path = f"{base_path}/{relative_path}" if base_path else f"/{relative_path}" + return urlunsplit( + (parts.scheme, parts.netloc, joined_path, parts.query, parts.fragment) + ) + + def _convert_agent_card_dict_to_obj(agent_card_dict: dict) -> AgentCard: agent_card_json_str = json.dumps(agent_card_dict, ensure_ascii=False, indent=2) agent_card_object = AgentCard.model_validate_json(str(agent_card_json_str)) @@ -138,6 +150,9 @@ def __init__( auth_token: Optional[str] = None, auth_method: Literal["header", "querystring"] | None = None, httpx_client: Optional[httpx.AsyncClient] = None, + rpc_url: Optional[str] = None, + extra_headers: Optional[dict[str, str]] = None, + timeout: float = 600, ): # Determine the effective URL for the agent and handle conflicts. effective_url = url @@ -155,12 +170,12 @@ def __init__( "Could not determine agent URL. Please provide the `url` parameter or an `httpx_client` with a configured `base_url`." ) - req_headers = {} + req_headers = dict(extra_headers or {}) req_params = {} if auth_token: if auth_method == "header": - req_headers = {"Authorization": f"Bearer {auth_token}"} + req_headers["Authorization"] = f"Bearer {auth_token}" elif auth_method == "querystring": req_params = {"token": auth_token} elif auth_method: @@ -169,45 +184,42 @@ def __init__( ) agent_card_dict = requests.get( - effective_url + AGENT_CARD_WELL_KNOWN_PATH, + _url_with_path(effective_url, AGENT_CARD_WELL_KNOWN_PATH), headers=req_headers, params=req_params, ).json() # replace agent_card_url with actual host - agent_card_dict["url"] = effective_url + agent_card_dict["url"] = rpc_url or effective_url agent_card_object = _convert_agent_card_dict_to_obj(agent_card_dict) - logger.debug(f"Agent card of {name}: {agent_card_object}") + logger.debug("Loaded Agent card for %s", name) client_was_provided = httpx_client is not None client_to_use = httpx_client if client_was_provided: # If a client was provided, update it with auth info - if auth_token: - if auth_method == "header": - client_to_use.headers.update(req_headers) - elif auth_method == "querystring": - new_params = dict(client_to_use.params) - new_params.update(req_params) - client_to_use.params = new_params + if req_headers: + client_to_use.headers.update(req_headers) + if auth_token and auth_method == "querystring": + new_params = dict(client_to_use.params) + new_params.update(req_params) + client_to_use.params = new_params else: # If no client was provided, create a new one with auth info - if auth_token: - if auth_method == "header": - client_to_use = httpx.AsyncClient( - base_url=effective_url, headers=req_headers, timeout=600 - ) - elif auth_method == "querystring": - client_to_use = httpx.AsyncClient( - base_url=effective_url, params=req_params, timeout=600 - ) - else: # No auth, no client provided - client_to_use = httpx.AsyncClient(base_url=effective_url, timeout=600) + client_to_use = httpx.AsyncClient( + base_url=effective_url, + headers=req_headers, + params=req_params, + timeout=timeout, + ) super().__init__( - name=name, agent_card=agent_card_object, httpx_client=client_to_use + name=name, + agent_card=agent_card_object, + httpx_client=client_to_use, + timeout=timeout, ) # The parent class sets _httpx_client_needs_cleanup based on whether diff --git a/veadk/tools/builtin_tools/_agentkit.py b/veadk/tools/builtin_tools/_agentkit.py index b779f0dba..4d443cbb9 100644 --- a/veadk/tools/builtin_tools/_agentkit.py +++ b/veadk/tools/builtin_tools/_agentkit.py @@ -12,9 +12,14 @@ # See the License for the specific language governing permissions and # limitations under the License. +import hashlib import json import os +import re +import threading import time +from dataclasses import dataclass +from datetime import datetime, timezone from typing import Any, Optional from veadk.auth.veauth.utils import get_credential_from_vefaas_iam @@ -28,10 +33,46 @@ _SESSION_READY_TIMEOUT = 120.0 _SESSION_POLL_INTERVAL = 1.0 _SESSION_TERMINAL_STATUSES = frozenset({"failed", "terminating", "terminated"}) +_SESSION_REUSABLE_STATUSES = frozenset({"starting", "ready"}) +_SESSION_LIST_PAGE_SIZE = 100 +_SESSION_USER_ID_MAX_LENGTH = 200 +_SESSION_ROTATION_SUFFIX_LENGTH = 15 +_SESSION_LOCK_STRIPE_COUNT = 64 _AGENTKIT_REQUEST_CONNECT_TIMEOUT = 10.0 _AGENTKIT_REQUEST_MIN_READ_TIMEOUT = 60.0 _AGENTKIT_REQUEST_TIMEOUT_BUFFER = 30.0 +_session_locks = tuple(threading.Lock() for _ in range(_SESSION_LOCK_STRIPE_COUNT)) + + +@dataclass(frozen=True) +class AgentKitSessionLease: + """A resolved AgentKit Session and its session-scoped data-plane endpoints.""" + + tool_id: str + logical_user_session_id: str + user_session_id: str + session_id: str + status: str + endpoint: str + internal_endpoint: str + created_at: str + expire_at: str + + def select_endpoint(self, *, prefer_internal_endpoint: bool = False) -> str: + if prefer_internal_endpoint: + return self.internal_endpoint or self.endpoint + return self.endpoint or self.internal_endpoint + + def remaining_seconds(self, *, now: datetime | None = None) -> float | None: + expires_at = _parse_agentkit_timestamp(self.expire_at) + if expires_at is None: + return None + current = now or datetime.now(timezone.utc) + if current.tzinfo is None: + current = current.replace(tzinfo=timezone.utc) + return (expires_at - current.astimezone(timezone.utc)).total_seconds() + def _agentkit_request_timeout(operation_timeout: int) -> tuple[float, float]: """Keep the synchronous request alive longer than the tool operation.""" @@ -44,6 +85,91 @@ def _agentkit_request_timeout(operation_timeout: int) -> tuple[float, float]: ) +def _parse_agentkit_timestamp(value: object) -> datetime | None: + if not isinstance(value, str) or not value.strip(): + return None + normalized = value.strip() + if normalized.endswith("Z"): + normalized = normalized[:-1] + "+00:00" + try: + parsed = datetime.fromisoformat(normalized) + except ValueError: + logger.warning("Invalid AgentKit Session timestamp: %s", value) + return None + if parsed.tzinfo is None: + parsed = parsed.replace(tzinfo=timezone.utc) + return parsed.astimezone(timezone.utc) + + +def _safe_agentkit_user_session_id(logical_user_session_id: str) -> str: + """Return a contract-compliant, stable base for physical UserSessionIds.""" + if not logical_user_session_id: + raise ValueError("tool_user_session_id must not be empty") + normalized = re.sub(r"[^A-Za-z0-9_-]", "_", logical_user_session_id) + digest = hashlib.sha256(logical_user_session_id.encode("utf-8")).hexdigest()[:12] + if normalized != logical_user_session_id: + normalized = f"{normalized}_{digest}" + max_base_length = _SESSION_USER_ID_MAX_LENGTH - _SESSION_ROTATION_SUFFIX_LENGTH + if len(normalized) > max_base_length: + normalized = f"{normalized[: max_base_length - 13]}_{digest}" + return normalized + + +def _session_lock(tool_id: str, logical_user_session_id: str) -> threading.Lock: + digest = hashlib.sha256( + f"{tool_id}\0{logical_user_session_id}".encode("utf-8") + ).digest() + index = int.from_bytes(digest[:4], "big") % _SESSION_LOCK_STRIPE_COUNT + return _session_locks[index] + + +def _session_has_enough_time( + session: object, + *, + min_remaining_seconds: float, + now: datetime, +) -> bool: + expires_at = _parse_agentkit_timestamp(getattr(session, "expire_at", None)) + if expires_at is None: + return min_remaining_seconds <= 0 + return (expires_at - now).total_seconds() > min_remaining_seconds + + +def _session_is_reusable( + session: object, + *, + physical_user_session_id_base: str, + min_remaining_seconds: float, + now: datetime, +) -> bool: + user_session_id = getattr(session, "user_session_id", None) + if not isinstance(user_session_id, str) or not ( + user_session_id == physical_user_session_id_base + or user_session_id.startswith(f"{physical_user_session_id_base}_r_") + ): + return False + status = (getattr(session, "status", None) or "").strip().lower() + return status in _SESSION_REUSABLE_STATUSES and _session_has_enough_time( + session, + min_remaining_seconds=min_remaining_seconds, + now=now, + ) + + +def _rotated_user_session_id( + physical_user_session_id_base: str, + sessions: list[object], +) -> str: + existing_ids = sorted( + str(getattr(session, "session_id", "") or "") for session in sessions + ) + generation = hashlib.sha256("\n".join(existing_ids).encode("utf-8")).hexdigest()[ + :12 + ] + suffix = f"_r_{generation}" + return f"{physical_user_session_id_base[: _SESSION_USER_ID_MAX_LENGTH - len(suffix)]}{suffix}" + + def resolve_agentkit_tool_id(*preferred_env_names: str) -> str: """Resolve the first configured AgentKit tool id with AGENTKIT_TOOL_ID fallback.""" for env_name in [*preferred_env_names, "AGENTKIT_TOOL_ID"]: @@ -242,34 +368,27 @@ def _get_or_create_agentkit_session( tool_id: str, tool_user_session_id: str, ttl: int, + min_remaining_seconds: float = 0, ): - """Return the newest reusable session for ``tool_user_session_id`` or create one.""" + """Return a reusable physical Session for a stable logical session key.""" from agentkit.sdk.tools import types as tools_types - try: - listing = client.list_sessions( - tools_types.ListSessionsRequest( - ToolId=tool_id, - Filters=[ - tools_types.FiltersItemForListSessions( - Name="UserSessionId", - Values=[tool_user_session_id], - ) - ], - PageSize=20, - ) - ) - except Exception as exc: # noqa: BLE001 - logger.debug(f"AgentKit ListSessions failed, falling back to create: {exc}") - listing = None - - candidates = getattr(listing, "session_infos", None) or [] + physical_user_session_id_base = _safe_agentkit_user_session_id(tool_user_session_id) + candidates = _list_agentkit_sessions( + client=client, + tool_id=tool_id, + physical_user_session_id_base=physical_user_session_id_base, + ) + now = datetime.now(timezone.utc) reusable = [ info for info in candidates - if getattr(info, "user_session_id", None) == tool_user_session_id - and (getattr(info, "status", None) or "").strip().lower() - not in _SESSION_TERMINAL_STATUSES + if _session_is_reusable( + info, + physical_user_session_id_base=physical_user_session_id_base, + min_remaining_seconds=min_remaining_seconds, + now=now, + ) ] if reusable: reusable.sort( @@ -278,122 +397,252 @@ def _get_or_create_agentkit_session( chosen = reusable[0] logger.debug( f"Reusing AgentKit session {getattr(chosen, 'session_id', None)} " - f"for UserSessionId={tool_user_session_id}" + f"for logical UserSessionId={tool_user_session_id}" ) return chosen - return client.create_session( - tools_types.CreateSessionRequest( - ToolId=tool_id, - UserSessionId=tool_user_session_id, - Ttl=ttl, + physical_user_session_id = ( + physical_user_session_id_base + if not candidates + else _rotated_user_session_id( + physical_user_session_id_base, + candidates, ) ) + try: + return client.create_session( + tools_types.CreateSessionRequest( + ToolId=tool_id, + UserSessionId=physical_user_session_id, + Ttl=ttl, + ) + ) + except Exception: + # CreateSession may have succeeded even if its response was lost. Recover + # the physical Session before deciding whether the operation failed. + refreshed = _list_agentkit_sessions( + client=client, + tool_id=tool_id, + physical_user_session_id_base=physical_user_session_id_base, + ) + recovered = [ + info + for info in refreshed + if getattr(info, "user_session_id", None) == physical_user_session_id + and (getattr(info, "status", None) or "").strip().lower() + in _SESSION_REUSABLE_STATUSES + ] + if recovered: + recovered.sort( + key=lambda info: getattr(info, "created_at", "") or "", + reverse=True, + ) + return recovered[0] + raise -def ensure_agentkit_session_endpoint( +def _list_agentkit_sessions( + *, + client, + tool_id: str, + physical_user_session_id_base: str, +) -> list[object]: + """List every physical Session associated with one logical key.""" + if not hasattr(client, "list_sessions"): + # Compatibility for older clients and lightweight test doubles. Current + # AgentKit SDK versions expose ListSessions and use the paginated path. + return [] + + from agentkit.sdk.tools import types as tools_types + + sessions: list[object] = [] + next_token: str | None = None + seen_tokens: set[str] = set() + while True: + request_kwargs: dict[str, object] = { + "ToolId": tool_id, + "Filters": [ + tools_types.FiltersItemForListSessions( + NameContains="UserSessionId", + Values=[physical_user_session_id_base], + ) + ], + "MaxResults": _SESSION_LIST_PAGE_SIZE, + } + if next_token: + request_kwargs["NextToken"] = next_token + listing = client.list_sessions( + tools_types.ListSessionsRequest(**request_kwargs) + ) + for info in getattr(listing, "session_infos", None) or []: + user_session_id = getattr(info, "user_session_id", None) + if user_session_id == physical_user_session_id_base or ( + isinstance(user_session_id, str) + and user_session_id.startswith(f"{physical_user_session_id_base}_r_") + ): + sessions.append(info) + + next_token = getattr(listing, "next_token", None) or None + if not next_token: + return sessions + if next_token in seen_tokens: + raise RuntimeError("AgentKit ListSessions returned a repeated NextToken") + seen_tokens.add(next_token) + + +def _agentkit_session_lease( + *, + session: object, + fallback_session: object | None, + tool_id: str, + logical_user_session_id: str, +) -> AgentKitSessionLease: + def value(name: str) -> str: + current = getattr(session, name, None) + fallback = getattr(fallback_session, name, None) if fallback_session else None + return str(current or fallback or "") + + return AgentKitSessionLease( + tool_id=tool_id, + logical_user_session_id=logical_user_session_id, + user_session_id=value("user_session_id") + or _safe_agentkit_user_session_id(logical_user_session_id), + session_id=value("session_id"), + status=value("status"), + endpoint=value("endpoint"), + internal_endpoint=value("internal_endpoint"), + created_at=value("created_at"), + expire_at=value("expire_at"), + ) + + +def ensure_agentkit_session_lease( *, tool_id: str, tool_user_session_id: str, tool_state: Optional[dict[str, Any]] = None, ttl: int = 1800, - prefer_internal_endpoint: bool = False, - wait_until_ready: bool = False, + min_remaining_seconds: float = 0, + wait_until_ready: bool = True, ready_timeout: float = _SESSION_READY_TIMEOUT, poll_interval: float = _SESSION_POLL_INTERVAL, -) -> str: - """Create or reuse an AgentKit tool session and return its endpoint.""" +) -> AgentKitSessionLease: + """Resolve a live Session lease for a stable logical UserSessionId.""" from agentkit.sdk.tools import types as tools_types from agentkit.sdk.tools.client import AgentkitToolsClient - if wait_until_ready: - if ready_timeout < 0: - raise ValueError("ready_timeout must be greater than or equal to 0") - if poll_interval <= 0: - raise ValueError("poll_interval must be greater than 0") - + if not 60 <= ttl <= 86400: + raise ValueError("ttl must be between 60 and 86400 seconds") + if min_remaining_seconds < 0: + raise ValueError("min_remaining_seconds must be greater than or equal to 0") + if min_remaining_seconds >= 86400: + raise ValueError("min_remaining_seconds must be less than 86400 seconds") + if ready_timeout < 0: + raise ValueError("ready_timeout must be greater than or equal to 0") + if poll_interval <= 0: + raise ValueError("poll_interval must be greater than 0") + + required_ttl = max(ttl, int(min_remaining_seconds) + 1) _, region, _, _ = get_agentkit_endpoint_config() ak, sk, header = get_agentkit_credentials(tool_state) - session_token = header.get("X-Security-Token", "") client = AgentkitToolsClient( access_key=ak, secret_key=sk, region=region, - session_token=session_token, - ) - session = _get_or_create_agentkit_session( - client=client, - tool_id=tool_id, - tool_user_session_id=tool_user_session_id, - ttl=ttl, + session_token=header.get("X-Security-Token", ""), ) - if not wait_until_ready: - public_endpoint = getattr(session, "endpoint", None) - internal_endpoint = getattr(session, "internal_endpoint", None) - endpoint = ( - internal_endpoint or public_endpoint - if prefer_internal_endpoint - else public_endpoint or internal_endpoint - ) - if endpoint: - return endpoint - session_id = session.session_id + with _session_lock(tool_id, tool_user_session_id): + session = _get_or_create_agentkit_session( + client=client, + tool_id=tool_id, + tool_user_session_id=tool_user_session_id, + ttl=required_ttl, + min_remaining_seconds=min_remaining_seconds, + ) + session_id = getattr(session, "session_id", None) if not session_id: - return "" - current_session = client.get_session( - tools_types.GetSessionRequest( - ToolId=tool_id, - SessionId=session_id, + raise RuntimeError("AgentKit CreateSession response is missing SessionId") + + if not wait_until_ready: + lease = _agentkit_session_lease( + session=session, + fallback_session=None, + tool_id=tool_id, + logical_user_session_id=tool_user_session_id, ) - ) - if prefer_internal_endpoint: - return current_session.internal_endpoint or current_session.endpoint or "" - return current_session.endpoint or current_session.internal_endpoint or "" + if lease.endpoint or lease.internal_endpoint: + return lease + + deadline = time.monotonic() + ready_timeout + last_status = "Unknown" + while True: + current_session = client.get_session( + tools_types.GetSessionRequest( + ToolId=tool_id, + SessionId=session_id, + ) + ) + status = (getattr(current_session, "status", None) or "").strip() + last_status = status or "Unknown" + logger.debug("AgentKit session %s status: %s", session_id, last_status) + normalized_status = status.lower() + if normalized_status == "ready": + lease = _agentkit_session_lease( + session=current_session, + fallback_session=session, + tool_id=tool_id, + logical_user_session_id=tool_user_session_id, + ) + if not lease.select_endpoint(): + raise RuntimeError( + f"AgentKit session {session_id} is Ready but has no endpoint" + ) + if not _session_has_enough_time( + current_session, + min_remaining_seconds=min_remaining_seconds, + now=datetime.now(timezone.utc), + ): + raise RuntimeError( + f"AgentKit session {session_id} became Ready without enough " + "remaining lifetime" + ) + return lease + if normalized_status in _SESSION_TERMINAL_STATUSES: + raise RuntimeError( + f"AgentKit session {session_id} entered terminal status {last_status}" + ) - session_id = session.session_id - if not session_id: - raise RuntimeError("AgentKit CreateSession response is missing SessionId") + remaining = deadline - time.monotonic() + if remaining <= 0: + raise TimeoutError( + f"Timed out waiting for AgentKit session {session_id} to become " + f"Ready; last status: {last_status}" + ) + time.sleep(min(poll_interval, remaining)) - deadline = time.monotonic() + ready_timeout - last_status = "Unknown" - while True: - current_session = client.get_session( - tools_types.GetSessionRequest( - ToolId=tool_id, - SessionId=session_id, - ) - ) - status = (getattr(current_session, "status", None) or "").strip() - last_status = status or "Unknown" - logger.debug(f"AgentKit session {session_id} status: {last_status}") - normalized_status = status.lower() - if normalized_status == "ready": - public_endpoint = getattr(current_session, "endpoint", None) or getattr( - session, "endpoint", None - ) - internal_endpoint = getattr( - current_session, "internal_endpoint", None - ) or getattr(session, "internal_endpoint", None) - endpoint = ( - internal_endpoint or public_endpoint - if prefer_internal_endpoint - else public_endpoint or internal_endpoint - ) - if endpoint: - return endpoint - raise RuntimeError( - f"AgentKit session {session_id} is Ready but has no endpoint" - ) - if normalized_status in _SESSION_TERMINAL_STATUSES: - raise RuntimeError( - f"AgentKit session {session_id} entered terminal status {last_status}" - ) - remaining = deadline - time.monotonic() - if remaining <= 0: - raise TimeoutError( - f"Timed out waiting for AgentKit session {session_id} to become " - f"Ready; last status: {last_status}" - ) - time.sleep(min(poll_interval, remaining)) +def ensure_agentkit_session_endpoint( + *, + tool_id: str, + tool_user_session_id: str, + tool_state: Optional[dict[str, Any]] = None, + ttl: int = 1800, + prefer_internal_endpoint: bool = False, + wait_until_ready: bool = False, + ready_timeout: float = _SESSION_READY_TIMEOUT, + poll_interval: float = _SESSION_POLL_INTERVAL, + min_remaining_seconds: float = 0, +) -> str: + """Create or reuse an AgentKit tool session and return its endpoint.""" + lease = ensure_agentkit_session_lease( + tool_id=tool_id, + tool_user_session_id=tool_user_session_id, + tool_state=tool_state, + ttl=ttl, + min_remaining_seconds=min_remaining_seconds, + wait_until_ready=wait_until_ready, + ready_timeout=ready_timeout, + poll_interval=poll_interval, + ) + return lease.select_endpoint(prefer_internal_endpoint=prefer_internal_endpoint)