databricks-agentbricks 0.2.0__py3-none-any.whl

databricks_agentkit/runtime/auth.py

@@ -0,0 +1,149 @@
+"""Transient authentication for request-bound, user-authorized tool calls."""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import os
+from collections.abc import Mapping
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, NoReturn, SupportsIndex
+
+from databricks_agentkit.runtime.store import RUNTIME_STORE_LOCAL_ENV
+
+if TYPE_CHECKING:
+    from databricks.sdk import WorkspaceClient
+
+
+class AuthError(RuntimeError):
+    """A credential-free error suitable for returning to the caller."""
+
+    def __init__(
+        self, code: str, message: str, status_code: int = 401, integration_id: str | None = None
+    ) -> None:
+        super().__init__(message)
+        self.code = code
+        self.status_code = status_code
+        self.integration_id = integration_id
+
+    def payload(self) -> dict[str, str]:
+        result = {"code": self.code, "message": str(self)}
+        if self.integration_id is not None:
+            result["integration_id"] = self.integration_id
+        return result
+
+
+@dataclass(frozen=True)
+class InvocationAuthPolicy:
+    """Whether an invocation needs credentials scoped to its active execution attempt."""
+
+    user_tools: tuple[str, ...] = ()
+
+    @property
+    def requires_user(self) -> bool:
+        return bool(self.user_tools)
+
+    @classmethod
+    def from_manifest(cls, *, allow_missing: bool = False) -> InvocationAuthPolicy:
+        """Infer request-user tools from ``agent.toml`` in the active project."""
+        from databricks_agentkit.runtime.tool_manifest import load_tools, project_root, tomllib
+
+        try:
+            root = project_root()
+        except RuntimeError:
+            if allow_missing and not os.getenv("AGENTBRICKS_PROJECT_ROOT"):
+                return cls()
+            raise
+        with (root / "agent.toml").open("rb") as source:
+            document = tomllib.load(source)
+        tools = load_tools(expected_framework=document.get("agent", {}).get("framework", ""))
+        return cls(tuple(tool.id for tool in tools if tool.auth == "user"))
+
+
+class RequestAuthContext:
+    """Private request credentials, never part of a persisted invocation payload.
+
+    Deployed Apps must only be reachable through the trusted Apps ingress, which supplies the
+    forwarded identity headers. Local development deliberately ignores those headers.
+    """
+
+    __slots__ = (
+        "_token",
+        "_principal",
+        "_host",
+        "_app",
+        "_workspace",
+        "_local",
+        "_clients",
+        "_closed",
+    )
+
+    def __init__(self, *, token: str | None, principal: str, local: bool) -> None:
+        self._token = token
+        self._principal = principal
+        self._host = os.getenv("DATABRICKS_HOST", "").rstrip("/")
+        self._app = os.getenv("DATABRICKS_APP_NAME", "local")
+        self._workspace = os.getenv("DATABRICKS_WORKSPACE_ID", "")
+        self._local = local
+        self._clients: dict[str, WorkspaceClient] = {}
+        self._closed = False
+
+    @classmethod
+    def from_headers(cls, headers: Mapping[str, str]) -> RequestAuthContext:
+        local_runtime = os.getenv(RUNTIME_STORE_LOCAL_ENV, "").lower() == "true"
+        if local_runtime or not os.getenv("DATABRICKS_APP_NAME"):
+            return cls(token=None, principal="local-developer", local=True)
+        normalized = {name.lower(): value for name, value in headers.items()}
+        token = normalized.get("x-forwarded-access-token", "").strip()
+        principal = normalized.get("x-forwarded-user", "").strip()
+        if not token:
+            raise AuthError("MCP_USER_AUTH_REQUIRED", "Missing request-user access token")
+        if not principal:
+            raise AuthError("MCP_USER_IDENTITY_MISSING", "Missing trusted request-user identity")
+        return cls(token=token, principal=principal, local=False)
+
+    def __repr__(self) -> str:
+        return f"RequestAuthContext(closed={self._closed})"
+
+    def __reduce_ex__(self, protocol: SupportsIndex, /) -> NoReturn:
+        raise TypeError("RequestAuthContext cannot be serialized")
+
+    def namespace(self, kind: str, value: str) -> str:
+        material = json.dumps(
+            [self._host, self._workspace, self._app, self._principal, kind, value],
+            separators=(",", ":"),
+        )
+        return hashlib.sha256(material.encode()).hexdigest()
+
+    def client_for(self, mode: str) -> WorkspaceClient:
+        if self._closed:
+            raise AuthError(
+                "MCP_USER_AUTH_EXPIRED", "Request-user authentication is no longer active"
+            )
+        if mode not in ("user", "app"):
+            raise ValueError("auth must be 'user' or 'app'")
+        if mode not in self._clients:
+            from databricks_agentkit.runtime.workspace import workspace_client, workspace_headers
+
+            if mode == "app" or self._local:
+                self._clients[mode] = workspace_client()
+            else:
+                from databricks.sdk import WorkspaceClient
+
+                if not self._host:
+                    raise AuthError(
+                        "MCP_USER_AUTH_CONFIGURATION", "Workspace host is not configured", 500
+                    )
+                self._clients[mode] = WorkspaceClient(
+                    host=self._host,
+                    token=self._token,
+                    auth_type="pat",
+                    custom_headers=workspace_headers(),
+                )
+        return self._clients[mode]
+
+    def close(self) -> None:
+        """Drop credentials and prohibit further client resolution."""
+        self._closed = True
+        self._token = None
+        self._clients.clear()

databricks_agentkit/runtime/durability/__init__.py

@@ -0,0 +1 @@
+"""Durability helpers for the canonical AgentKit runtime."""

databricks_agentkit/runtime/durability/execution.py

@@ -0,0 +1,133 @@
+"""Durable scheduling and execution of fenced invocation attempts."""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+from collections.abc import Callable
+
+from databricks_agentkit.runtime.durability.heartbeat import Heartbeat
+from databricks_agentkit.runtime.durability.recovery import RecoveryScheduler
+from databricks_agentkit.runtime.durability.store import DurableRuntimeStore
+from databricks_agentkit.runtime.execution import AttemptExecution, InvocationExecutor
+from databricks_agentkit.runtime.types import Invocation, InvocationStatus
+
+logger = logging.getLogger(__name__)
+
+
+class DurableInvocationExecutor(InvocationExecutor):
+    """Claim durable attempts and wrap shared execution with heartbeats."""
+
+    def __init__(
+        self,
+        execution: AttemptExecution,
+        *,
+        runtime_store: DurableRuntimeStore,
+        recovery_enabled: Callable[[], bool],
+        heartbeat_seconds: float,
+        stale_seconds: float,
+        scan_seconds: float,
+    ) -> None:
+        if stale_seconds <= heartbeat_seconds:
+            raise ValueError("stale_seconds must be greater than heartbeat_seconds")
+        self._execution = execution
+        self._runtime_store = runtime_store
+        self._recovery_enabled = recovery_enabled
+        self._heartbeat_seconds = heartbeat_seconds
+        self._stale_seconds = stale_seconds
+        self._scan_seconds = scan_seconds
+        self._tasks: dict[str, asyncio.Task[None]] = {}
+        self._queued_scanner: asyncio.Task[None] | None = None
+        self._recovery_scheduler: RecoveryScheduler | None = None
+
+    async def start(self) -> None:
+        self._queued_scanner = asyncio.create_task(
+            self._scan_queued_loop(),
+            name="databricks-durable-runtime-queued-scanner",
+        )
+        self._recovery_scheduler = RecoveryScheduler(
+            scheduler=self,
+            runtime_store=self._runtime_store,
+            stale_seconds=self._stale_seconds,
+            scan_seconds=self._scan_seconds,
+        )
+        self._recovery_scheduler.start(recover=self._recovery_enabled())
+
+    async def stop(self) -> None:
+        if self._queued_scanner is not None:
+            self._queued_scanner.cancel()
+        if self._recovery_scheduler is not None:
+            await self._recovery_scheduler.stop()
+            self._recovery_scheduler = None
+        await asyncio.gather(
+            *([self._queued_scanner] if self._queued_scanner is not None else []),
+            return_exceptions=True,
+        )
+        self._queued_scanner = None
+        tasks = list(self._tasks.values())
+        for task in tasks:
+            task.cancel()
+        await asyncio.gather(*tasks, return_exceptions=True)
+        self._tasks.clear()
+
+    def ensure_scheduled(self, state: Invocation) -> None:
+        """Schedule a newly queued invocation on this worker."""
+        if state.status != InvocationStatus.QUEUED:
+            return
+        self._schedule(state.invocation_id, recovery=False)
+
+    def ensure_recovery_scheduled(self, invocation_id: str) -> None:
+        """Schedule stale active work discovered by the recovery scanner."""
+        self._schedule(invocation_id, recovery=True)
+
+    async def _scan_queued_loop(self) -> None:
+        """Continuously schedule persisted first attempts, including after process restart."""
+        while True:
+            try:
+                for invocation_id in await self._runtime_store.queued_invocation_ids():
+                    self._schedule(invocation_id, recovery=False)
+            except asyncio.CancelledError:
+                raise
+            except Exception:
+                logger.exception("Databricks durable runtime queued invocation scan failed")
+            await asyncio.sleep(self._scan_seconds)
+
+    def _schedule(self, invocation_id: str, *, recovery: bool) -> None:
+        current = self._tasks.get(invocation_id)
+        if current is not None and not current.done():
+            return
+        task = asyncio.create_task(
+            self._run(invocation_id, recovery=recovery),
+            name=f"durable-invocation-{invocation_id}",
+        )
+        self._tasks[invocation_id] = task
+        task.add_done_callback(lambda completed: self._discard_task(invocation_id, completed))
+
+    async def _run(self, invocation_id: str, *, recovery: bool) -> None:
+        try:
+            if recovery:
+                claimed = await self._runtime_store.claim_recoverable(
+                    invocation_id,
+                    self._stale_seconds,
+                )
+            else:
+                claimed = await self._runtime_store.claim(invocation_id)
+        except Exception:
+            logger.exception("Failed to claim durable invocation: %s", invocation_id)
+            return
+        if claimed is None:
+            return
+
+        async with Heartbeat(
+            runtime_store=self._runtime_store,
+            invocation_id=invocation_id,
+            attempt=claimed.attempt,
+            heartbeat_seconds=self._heartbeat_seconds,
+        ):
+            await self._execution.run(claimed)
+
+    def _discard_task(self, invocation_id: str, completed: asyncio.Task[None]) -> None:
+        if self._tasks.get(invocation_id) is completed:
+            self._tasks.pop(invocation_id, None)
+        if not completed.cancelled():
+            completed.exception()

databricks_agentkit/runtime/durability/heartbeat.py

@@ -0,0 +1,64 @@
+"""Heartbeat lifecycle for a durable invocation attempt."""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+
+from databricks_agentkit.runtime.durability.store import DurableRuntimeStore
+
+logger = logging.getLogger(__name__)
+
+
+class Heartbeat:
+    """Refresh attempt ownership until execution finishes or ownership changes."""
+
+    def __init__(
+        self,
+        *,
+        runtime_store: DurableRuntimeStore,
+        invocation_id: str,
+        attempt: int,
+        heartbeat_seconds: float,
+    ) -> None:
+        if heartbeat_seconds <= 0:
+            raise ValueError("heartbeat_seconds must be positive")
+        self._runtime_store = runtime_store
+        self._invocation_id = invocation_id
+        self._attempt = attempt
+        self._heartbeat_seconds = heartbeat_seconds
+        self._task: asyncio.Task[None] | None = None
+
+    async def __aenter__(self) -> Heartbeat:
+        self._task = asyncio.create_task(
+            self._run(),
+            name=f"durable-heartbeat-{self._invocation_id}-{self._attempt}",
+        )
+        return self
+
+    async def __aexit__(self, exc_type, exc, traceback) -> None:
+        if self._task is None:
+            return
+        self._task.cancel()
+        await asyncio.gather(self._task, return_exceptions=True)
+        self._task = None
+
+    async def _run(self) -> None:
+        while True:
+            try:
+                owns_attempt = await self._runtime_store.heartbeat(
+                    self._invocation_id,
+                    self._attempt,
+                )
+            except Exception:
+                logger.warning(
+                    "Durable heartbeat failed: %s attempt=%d",
+                    self._invocation_id,
+                    self._attempt,
+                    exc_info=True,
+                )
+                await asyncio.sleep(self._heartbeat_seconds)
+                continue
+            if not owns_attempt:
+                return
+            await asyncio.sleep(self._heartbeat_seconds)