kagent-adk 0.6.2__tar.gz → 0.6.4__tar.gz

{kagent_adk-0.6.2 → kagent_adk-0.6.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: kagent-adk
-Version: 0.6.2
+Version: 0.6.4
 Summary: kagent-adk is an sdk for integrating adk agents with kagent
 Requires-Python: >=3.12.11
 Requires-Dist: a2a-sdk>=0.2.16

{kagent_adk-0.6.2 → kagent_adk-0.6.4}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "kagent-adk"
-version = "0.6.2"
+version = "0.6.4"
 description = "kagent-adk is an sdk for integrating adk agents with kagent"
 readme = "README.md"
 requires-python = ">=3.12.11"

kagent_adk-0.6.4/src/kagent_adk/_token.py

@@ -0,0 +1,78 @@
+import logging  # noqa: I001
+import asyncio
+from contextlib import asynccontextmanager
+from typing import Any, Optional
+
+KAGENT_TOKEN_PATH = "/var/run/secrets/tokens/kagent-token"
+logger = logging.getLogger(__name__)
+
+
+class KAgentTokenService:
+    """Reads a k8s token from a file, and reloads it
+    periodically.
+    """
+
+    def __init__(self, app_name: str):
+        self.token = None
+        self.update_lock = asyncio.Lock()
+        self.update_task = None
+        self.app_name = app_name
+
+    def lifespan(self):
+        """Returns an async context manager to start the token update loop"""
+
+        @asynccontextmanager
+        async def _lifespan(app: Any):
+            await self._update_token_loop()
+            yield
+            self._drain()
+
+        return _lifespan
+
+    def event_hooks(self):
+        """Returns a dictionary of event hooks for the application
+        to use when creating the httpx.AsyncClient.
+        """
+        return {"request": [self._add_bearer_token]}
+
+    async def _update_token_loop(self) -> None:
+        self.token = await self._read_kagent_token()
+        # keep it updated - launch a background task to refresh it periodically
+        self.update_task = asyncio.create_task(self._refresh_token())
+
+    def _drain(self):
+        if self.update_task:
+            self.update_task.cancel()
+
+    async def _get_token(self) -> str | None:
+        async with self.update_lock:
+            return self.token
+
+    async def _read_kagent_token(self) -> str | None:
+        return await asyncio.to_thread(read_token)
+
+    async def _refresh_token(self):
+        while True:
+            await asyncio.sleep(60)  # Wait for 60 seconds before refreshing
+            token = await self._read_kagent_token()
+            if token is not None and token != self.token:
+                async with self.update_lock:
+                    self.token = token
+
+    async def _add_bearer_token(self, request):
+        # Your function to generate headers dynamically
+        token = await self._get_token()
+        headers = {"X-Agent-Name": self.app_name}
+        if token:
+            headers["Authorization"] = f"Bearer {token}"
+        request.headers.update(headers)
+
+
+def read_token() -> str | None:
+    try:
+        with open(KAGENT_TOKEN_PATH, "r", encoding="utf-8") as f:
+            token = f.read()
+            return token.strip()
+    except OSError as e:
+        logger.error(f"Error reading token from {KAGENT_TOKEN_PATH}: {e}")
+        return None

{kagent_adk-0.6.2 → kagent_adk-0.6.4}/src/kagent_adk/a2a.py

@@ -25,9 +25,7 @@ from google.genai import types
 from ._agent_executor import A2aAgentExecutor
 from ._session_service import KAgentSessionService
 from ._task_store import KAgentTaskStore
-
-# --- Constants ---
-USER_ID = "admin@kagent.dev"
+from ._token import KAgentTokenService
 
 # --- Configure Logging ---
 logger = logging.getLogger(__name__)
@@ -39,7 +37,7 @@ class KAgentUser(User):
 
     @property
     def is_authenticated(self) -> bool:
-        return False
+        return True
 
     @property
     def user_name(self) -> str:
@@ -51,9 +49,8 @@ class KAgentRequestContextBuilder(SimpleRequestContextBuilder):
     A request context builder that will be used to hack in the user_id for now.
     """
 
-    def __init__(self, user_id: str, task_store: TaskStore):
+    def __init__(self, task_store: TaskStore):
         super().__init__(task_store=task_store)
-        self.user_id = user_id
 
     async def build(
         self,
@@ -63,10 +60,12 @@ class KAgentRequestContextBuilder(SimpleRequestContextBuilder):
         task: Task | None = None,
         context: ServerCallContext | None = None,
     ) -> RequestContext:
-        if not context:
-            context = ServerCallContext(user=KAgentUser(user_id=self.user_id))
-        else:
-            context.user = KAgentUser(user_id=self.user_id)
+        if context:
+            # grab the user id from the header
+            headers = context.state.get("headers", {})
+            user_id = headers.get("x-user-id", None)
+            if user_id:
+                context.user = KAgentUser(user_id=user_id)
         request_context = await super().build(params, task_id, context_id, task, context)
         return request_context
 
@@ -101,7 +100,10 @@ class KAgentApp:
         self.agent_card = agent_card
 
     def build(self) -> FastAPI:
-        http_client = httpx.AsyncClient(base_url=kagent_url_override or self.kagent_url)
+        token_service = KAgentTokenService(self.app_name)
+        http_client = httpx.AsyncClient(  # TODO: add user  and agent headers
+            base_url=kagent_url_override or self.kagent_url, event_hooks=token_service.event_hooks()
+        )
         session_service = KAgentSessionService(http_client)
 
         def create_runner() -> Runner:
@@ -117,7 +119,7 @@ class KAgentApp:
 
         kagent_task_store = KAgentTaskStore(http_client)
 
-        request_context_builder = KAgentRequestContextBuilder(user_id=USER_ID, task_store=kagent_task_store)
+        request_context_builder = KAgentRequestContextBuilder(task_store=kagent_task_store)
         request_handler = DefaultRequestHandler(
             agent_executor=agent_executor,
             task_store=kagent_task_store,
@@ -130,7 +132,7 @@ class KAgentApp:
         )
 
         faulthandler.enable()
-        app = FastAPI()
+        app = FastAPI(lifespan=token_service.lifespan())
 
         # Health check/readiness probe
         app.add_route("/health", methods=["GET"], route=health_check)