databricks-agentbricks 0.2.0__py3-none-any.whl

databricks_agentkit/runtime/durability/recovery.py

@@ -0,0 +1,70 @@
+"""Detection and scheduling of recoverable durable attempts."""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+from typing import TYPE_CHECKING, Protocol
+
+if TYPE_CHECKING:
+    from databricks_agentkit.runtime.durability.store import DurableRuntimeStore
+
+logger = logging.getLogger(__name__)
+
+
+class RecoverySchedulerTarget(Protocol):
+    """Schedule one invocation after recovery finds it eligible."""
+
+    def ensure_recovery_scheduled(self, invocation_id: str) -> None:
+        """Schedule a queued or stale invocation for a recovery attempt."""
+        ...
+
+
+class RecoveryScheduler:
+    """Find stale active attempts when recovery is enabled."""
+
+    def __init__(
+        self,
+        scheduler: RecoverySchedulerTarget,
+        *,
+        runtime_store: "DurableRuntimeStore",
+        stale_seconds: float,
+        scan_seconds: float,
+    ) -> None:
+        if scan_seconds <= 0:
+            raise ValueError("scan_seconds must be positive")
+        self._scheduler = scheduler
+        self._runtime_store = runtime_store
+        self._stale_seconds = stale_seconds
+        self._scan_seconds = scan_seconds
+        self._scanner: asyncio.Task[None] | None = None
+
+    def start(self, *, recover: bool) -> None:
+        """Enable scheduling and optionally start proactive stale-work scanning."""
+        if recover:
+            self._scanner = asyncio.create_task(
+                self._scan_loop(),
+                name="databricks-durable-runtime-scanner",
+            )
+
+    async def stop(self) -> None:
+        """Stop scanning, leaving active rows recoverable elsewhere."""
+        if self._scanner is not None:
+            self._scanner.cancel()
+        await asyncio.gather(
+            *([self._scanner] if self._scanner is not None else []),
+            return_exceptions=True,
+        )
+        self._scanner = None
+
+    async def _scan_loop(self) -> None:
+        while True:
+            try:
+                invocation_ids = await self._runtime_store.stale_invocation_ids(self._stale_seconds)
+                for invocation_id in invocation_ids:
+                    self._scheduler.ensure_recovery_scheduled(invocation_id)
+            except asyncio.CancelledError:
+                raise
+            except Exception:
+                logger.exception("Databricks durable runtime recovery scan failed")
+            await asyncio.sleep(self._scan_seconds)

databricks_agentkit/runtime/durability/store.py

@@ -0,0 +1,33 @@
+"""Runtime Store capabilities required for heartbeat-based recovery."""
+
+from __future__ import annotations
+
+from typing import Protocol, runtime_checkable
+
+from databricks_agentkit.runtime.store import RuntimeStore
+from databricks_agentkit.runtime.types import Invocation
+
+
+@runtime_checkable
+class DurableRuntimeStore(RuntimeStore, Protocol):
+    """Extend Runtime Store operations with durable recovery coordination."""
+
+    async def queued_invocation_ids(self) -> list[str]:
+        """List invocations waiting to begin their first attempt."""
+        ...
+
+    async def stale_invocation_ids(self, stale_seconds: float) -> list[str]:
+        """List active invocations whose heartbeat is stale."""
+        ...
+
+    async def claim_recoverable(
+        self,
+        invocation_id: str,
+        stale_seconds: float,
+    ) -> Invocation | None:
+        """Replace a stale active attempt and return the new active attempt."""
+        ...
+
+    async def heartbeat(self, invocation_id: str, attempt: int) -> bool:
+        """Refresh an active attempt's lease while the caller owns it."""
+        ...

databricks_agentkit/runtime/execution.py

@@ -0,0 +1,160 @@
+"""Process-local scheduling and shared agent attempt execution."""
+
+from __future__ import annotations
+
+import asyncio
+import copy
+import json
+import logging
+from typing import Protocol, cast
+
+from databricks_agentkit.runtime.store import RuntimeStore
+from databricks_agentkit.runtime.types import (
+    Invocation,
+    InvocationAttemptContext,
+    InvocationExecutorFn,
+    InvocationStatus,
+    JsonObject,
+    JsonValue,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def copy_json_value(value: JsonValue, name: str) -> JsonValue:
+    """Return a detached JSON value or reject a value that cannot be persisted."""
+    try:
+        return cast(JsonValue, json.loads(json.dumps(value, allow_nan=False)))
+    except (TypeError, ValueError) as exc:
+        raise TypeError(f"{name} must be JSON serializable") from exc
+
+
+def _copy_json_object(value: JsonObject, name: str) -> JsonObject:
+    copied = copy_json_value(value, name)
+    if not isinstance(copied, dict):
+        raise TypeError(f"{name} must be a JSON object")
+    return copied
+
+
+class InvocationExecutor(Protocol):
+    """Schedule invocation attempts for one runtime process."""
+
+    async def start(self) -> None:
+        """Start any background scheduling owned by this executor."""
+        ...
+
+    async def stop(self) -> None:
+        """Stop scheduled work and release executor resources."""
+        ...
+
+    def ensure_scheduled(self, state: Invocation) -> None:
+        """Arrange for a newly accepted invocation to run on this process."""
+        ...
+
+
+class AttemptExecution:
+    """Call the agent hook and commit one already-claimed attempt."""
+
+    def __init__(self, execute_fn: InvocationExecutorFn, *, runtime_store: RuntimeStore) -> None:
+        self._execute_fn = execute_fn
+        self._runtime_store = runtime_store
+
+    async def run(self, claimed: Invocation) -> None:
+        invocation_id = claimed.invocation_id
+        try:
+
+            async def emit(event: JsonObject) -> int:
+                sequence_number = await self._runtime_store.append_event(
+                    invocation_id,
+                    claimed.attempt,
+                    _copy_json_object(event, "event"),
+                )
+                if sequence_number is None:
+                    raise RuntimeError(
+                        f"invocation {invocation_id!r} no longer owns attempt {claimed.attempt}"
+                    )
+                return sequence_number
+
+            response = await self._execute_fn(
+                copy.deepcopy(claimed.request),
+                InvocationAttemptContext(
+                    invocation_id=invocation_id,
+                    attempt=claimed.attempt,
+                    _emit=emit,
+                ),
+            )
+            response = copy_json_value(response, "executor response")
+            completed = await self._runtime_store.complete(
+                invocation_id,
+                claimed.attempt,
+                response,
+            )
+            if not completed:
+                logger.info(
+                    "Skipped completion after Runtime Store ownership changed: %s attempt=%d",
+                    invocation_id,
+                    claimed.attempt,
+                )
+        except asyncio.CancelledError:
+            raise
+        except Exception:
+            logger.exception(
+                "Invocation execution failed: %s attempt=%d",
+                invocation_id,
+                claimed.attempt,
+            )
+            try:
+                await self._runtime_store.fail(invocation_id, claimed.attempt)
+            except Exception:
+                logger.exception(
+                    "Failed to persist invocation failure: %s attempt=%d",
+                    invocation_id,
+                    claimed.attempt,
+                )
+
+
+class LocalInvocationExecutor(InvocationExecutor):
+    """Run background invocations inside the current process without heartbeats."""
+
+    def __init__(self, execution: AttemptExecution, *, runtime_store: RuntimeStore) -> None:
+        self._execution = execution
+        self._runtime_store = runtime_store
+        self._tasks: dict[str, asyncio.Task[None]] = {}
+
+    async def start(self) -> None:
+        pass
+
+    async def stop(self) -> 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:
+        if state.status != InvocationStatus.QUEUED:
+            return
+        current = self._tasks.get(state.invocation_id)
+        if current is not None and not current.done():
+            return
+        task = asyncio.create_task(
+            self._run(state.invocation_id),
+            name=f"invocation-{state.invocation_id}",
+        )
+        self._tasks[state.invocation_id] = task
+        task.add_done_callback(lambda completed: self._discard_task(state.invocation_id, completed))
+
+    async def _run(self, invocation_id: str) -> None:
+        try:
+            claimed = await self._runtime_store.claim(invocation_id)
+        except Exception:
+            logger.exception("Failed to claim invocation: %s", invocation_id)
+            return
+        if claimed is not None:
+            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/genie.py

@@ -0,0 +1,224 @@
+"""Framework-neutral tools for a configured Genie Agent's Chat-mode conversation API."""
+
+from __future__ import annotations
+
+import asyncio
+import re
+import time
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Any
+
+from databricks_agentkit.runtime.workspace import workspace_client
+
+if TYPE_CHECKING:
+    from databricks.sdk import WorkspaceClient
+
+_IDENTIFIER = re.compile(r"^[a-f0-9]{32}$")
+_TERMINAL = {"COMPLETED", "FAILED", "CANCELLED", "QUERY_RESULT_EXPIRED"}
+
+
+async def _run_in_thread(function: Callable[..., Any], *args: Any) -> Any:
+    task = asyncio.create_task(asyncio.to_thread(function, *args))
+    try:
+        return await asyncio.shield(task)
+    except asyncio.CancelledError as cancelled:
+        while not task.done():
+            try:
+                await asyncio.shield(task)
+            except asyncio.CancelledError:
+                continue
+            except Exception:
+                break
+        if not task.cancelled():
+            task.exception()
+        raise cancelled
+
+
+def _identifier(value: str, name: str) -> str:
+    if not isinstance(value, str) or not _IDENTIFIER.fullmatch(value):
+        raise ValueError(f"{name} must be a 32-character lowercase hexadecimal ID.")
+    return value
+
+
+class GenieAgent:
+    """Call one configured space without allowing the model to choose credentials or scope."""
+
+    def __init__(
+        self,
+        space_id: str,
+        *,
+        wait_seconds: float = 120.0,
+        auth: str = "app",
+        workspace_client_for: Callable[[str], WorkspaceClient] | None = None,
+    ) -> None:
+        self.space_id = _identifier(space_id, "space_id")
+        if not 0 < wait_seconds <= 600:
+            raise ValueError("wait_seconds must be greater than zero and at most 600.")
+        if auth not in ("user", "app"):
+            raise ValueError("auth must be 'user' or 'app'.")
+        self.wait_seconds = wait_seconds
+        self.auth = auth
+        self.workspace_client_for = workspace_client_for
+
+    def _client(self) -> WorkspaceClient:
+        if self.workspace_client_for is not None:
+            return self.workspace_client_for(self.auth)
+        return workspace_client()
+
+    async def ask(self, question: str, conversation_id: str | None = None) -> dict[str, Any]:
+        """Ask the configured Genie Agent a data question, or continue a conversation.
+
+        Waits up to two minutes for an answer. Preserve returned attachments, SQL and deep_link.
+        If timed_out is true and IDs are returned, call the matching poll tool; do not resubmit.
+        If indeterminate_submission is true, submission may still complete but no message ID was
+        received: do not resubmit automatically. NOT_SUBMITTED means setup timed out before sending
+        the question. FAILED, CANCELLED and QUERY_RESULT_EXPIRED are
+        terminal, not successful answers. Use query_result with a query attachment's attachment_id
+        to read its rows. No workspace or space selection is needed.
+        """
+        if not isinstance(question, str) or not question.strip():
+            raise ValueError("question must not be empty.")
+        if conversation_id is not None:
+            _identifier(conversation_id, "conversation_id")
+        deadline = time.monotonic() + self.wait_seconds
+        try:
+            client = await asyncio.wait_for(
+                _run_in_thread(self._client), timeout=deadline - time.monotonic()
+            )
+        except asyncio.TimeoutError:
+            return self._submission_timeout(conversation_id, indeterminate=False)
+        try:
+            if conversation_id is None:
+                operation = await asyncio.wait_for(
+                    _run_in_thread(client.genie.start_conversation, self.space_id, question),
+                    timeout=deadline - time.monotonic(),
+                )
+                conversation_id = operation.response.conversation_id
+            else:
+                operation = await asyncio.wait_for(
+                    _run_in_thread(
+                        client.genie.create_message, self.space_id, conversation_id, question
+                    ),
+                    timeout=deadline - time.monotonic(),
+                )
+        except asyncio.TimeoutError:
+            return self._submission_timeout(conversation_id, indeterminate=True)
+        return await self._wait(
+            client, conversation_id, operation.response.message_id, deadline=deadline
+        )
+
+    def _submission_timeout(
+        self, conversation_id: str | None, *, indeterminate: bool
+    ) -> dict[str, Any]:
+        return {
+            "space_id": self.space_id,
+            **({"conversation_id": conversation_id} if conversation_id is not None else {}),
+            "status": "INDETERMINATE_SUBMISSION" if indeterminate else "NOT_SUBMITTED",
+            "timed_out": True,
+            "indeterminate_submission": indeterminate,
+            "warning": (
+                "Submission may still complete, but no message ID was received. "
+                "Do not resubmit automatically."
+                if indeterminate
+                else "Client setup timed out before submission; the question was not submitted."
+            ),
+        }
+
+    async def poll(self, conversation_id: str, message_id: str) -> dict[str, Any]:
+        """Resume waiting for a previously submitted Genie Agent message without resubmitting it.
+
+        Waits up to two minutes, polling with backoff. If timed_out is true, retain the IDs to poll
+        again later. Terminal failures and expired results are returned explicitly, never rerun.
+        """
+        _identifier(conversation_id, "conversation_id")
+        _identifier(message_id, "message_id")
+        deadline = time.monotonic() + self.wait_seconds
+        try:
+            client = await asyncio.wait_for(
+                _run_in_thread(self._client), timeout=deadline - time.monotonic()
+            )
+        except asyncio.TimeoutError:
+            return {
+                "space_id": self.space_id,
+                "conversation_id": conversation_id,
+                "message_id": message_id,
+                "status": "IN_PROGRESS",
+                "timed_out": True,
+            }
+        return await self._wait(client, conversation_id, message_id, deadline=deadline)
+
+    def _reference(self, client: Any, conversation_id: str, message_id: str) -> dict[str, Any]:
+        _identifier(conversation_id, "conversation_id")
+        _identifier(message_id, "message_id")
+        return {
+            "space_id": self.space_id,
+            "conversation_id": conversation_id,
+            "message_id": message_id,
+            "deep_link": f"{client.config.host.rstrip('/')}/genie/spaces/{self.space_id}"
+            f"/conversations/{conversation_id}",
+        }
+
+    async def _wait(
+        self, client: Any, conversation_id: str, message_id: str, *, deadline: float
+    ) -> dict[str, Any]:
+        reference = self._reference(client, conversation_id, message_id)
+        delay = 2.0
+        message: dict[str, Any] = {"status": "IN_PROGRESS"}
+        while (remaining := deadline - time.monotonic()) > 0:
+            try:
+                response = await asyncio.wait_for(
+                    _run_in_thread(
+                        client.genie.get_message, self.space_id, conversation_id, message_id
+                    ),
+                    timeout=remaining,
+                )
+            except asyncio.TimeoutError:
+                break
+            message = response.as_dict()
+            if message.get("status") in _TERMINAL:
+                return {**message, **reference}
+            remaining = deadline - time.monotonic()
+            if remaining > 0:
+                await asyncio.sleep(min(delay, remaining))
+            delay = min(delay * 2, 30.0)
+        return {**message, **reference, "timed_out": True}
+
+    async def query_result(
+        self, conversation_id: str, message_id: str, attachment_id: str
+    ) -> dict[str, Any]:
+        """Read up to 100 rows and column schema for a Genie Agent query attachment.
+
+        Use the conversation_id, message_id and query attachment's attachment_id from ask/poll.
+        Check status before interpreting rows. When truncated is true, deep_link opens the full
+        result in Databricks. This never re-executes an expired query or downloads external links.
+        """
+        _identifier(conversation_id, "conversation_id")
+        _identifier(message_id, "message_id")
+        _identifier(attachment_id, "attachment_id")
+        client = await _run_in_thread(self._client)
+        response = await _run_in_thread(
+            client.genie.get_message_attachment_query_result,
+            self.space_id,
+            conversation_id,
+            message_id,
+            attachment_id,
+        )
+        statement = response.as_dict().get("statement_response") or {}
+        manifest = statement.get("manifest") or {}
+        result = statement.get("result") or {}
+        rows = result.get("data_array") or []
+        total = manifest.get("total_row_count")
+        return {
+            **self._reference(client, conversation_id, message_id),
+            "attachment_id": attachment_id,
+            "status": statement.get("status") or {"state": "UNKNOWN"},
+            "columns": (manifest.get("schema") or {}).get("columns", []),
+            "rows": rows[:100],
+            "total_row_count": total,
+            "truncated": bool(
+                len(rows) > 100
+                or manifest.get("truncated")
+                or result.get("next_chunk_index") is not None
+                or (total is not None and total > len(rows[:100]))
+            ),
+        }

databricks_agentkit/runtime/mcp_auth.py

@@ -0,0 +1,114 @@
+"""Shared authentication helpers for framework-specific MCP adapters."""
+
+from __future__ import annotations
+
+import os
+from collections.abc import Callable
+from typing import Any
+
+from databricks.sdk import WorkspaceClient
+
+from databricks_agentkit.runtime.auth import AuthError
+
+WorkspaceClientResolver = Callable[[str], WorkspaceClient]
+
+
+def resolve_mcp_workspace_client(
+    mode: str,
+    integration_id: str,
+    workspace_client_for: WorkspaceClientResolver | None,
+    default_client: Callable[[], WorkspaceClient],
+) -> WorkspaceClient:
+    """Resolve one manifest binding to its request-user or App client.
+
+    Request-aware templates pass ``RequestAuthContext.client_for`` as ``workspace_client_for``.
+    Databricks Apps sets ``DATABRICKS_APP_NAME`` in deployed compute, so a user-auth binding there
+    fails closed when invocation plumbing did not supply that request-owned resolver. Local runs
+    retain the environment-authenticated client fallback.
+    """
+    if workspace_client_for is None and mode == "user" and os.getenv("DATABRICKS_APP_NAME"):
+        raise AuthError(
+            "MCP_USER_AUTHORIZATION_MISSING",
+            "This deployed MCP integration requires request-user authorization.",
+            integration_id=integration_id,
+        )
+    try:
+        return workspace_client_for(mode) if workspace_client_for else default_client()
+    except Exception as error:
+        raise mcp_auth_error(error, integration_id) or AuthError(
+            "MCP_CLIENT_CONFIGURATION_FAILED",
+            "Could not configure the MCP client.",
+            500,
+            integration_id,
+        ) from None
+
+
+def mcp_auth_error(
+    error: BaseException, integration_id: str, seen: set[int] | None = None
+) -> AuthError | None:
+    """Classify nested credential failures without exposing upstream error details."""
+    from databricks.sdk.errors import PermissionDenied, Unauthenticated
+
+    seen = set() if seen is None else seen
+    if id(error) in seen:
+        return None
+    seen.add(id(error))
+    if isinstance(error, AuthError):
+        return (
+            error
+            if error.integration_id
+            else AuthError(error.code, str(error), error.status_code, integration_id)
+        )
+    nested_errors = [
+        *getattr(error, "exceptions", ()),
+        error.__cause__,
+        error.__context__,
+    ]
+    for nested in nested_errors:
+        if nested is not None and (classified := mcp_auth_error(nested, integration_id, seen)):
+            return classified
+    status = getattr(error, "status_code", None) or getattr(
+        getattr(error, "response", None), "status_code", None
+    )
+    code = getattr(getattr(error, "error", None), "code", None)
+    if isinstance(error, (PermissionError, PermissionDenied)) or status == 403:
+        return AuthError("MCP_PERMISSION_DENIED", "MCP permission denied.", 403, integration_id)
+    if isinstance(error, Unauthenticated) or status == 401:
+        return AuthError(
+            "MCP_USER_AUTHORIZATION_INVALID",
+            "MCP authorization was rejected.",
+            401,
+            integration_id,
+        )
+    if code == -32042:
+        return AuthError(
+            "MCP_AUTHORIZATION_REQUIRED",
+            "Authorize the configured service in Databricks before retrying.",
+            401,
+            integration_id,
+        )
+    return None
+
+
+def mcp_tool_error(result: Any, integration_id: str) -> AuthError | None:
+    """Classify structured MCP authorization failures returned as tool results."""
+    structured = getattr(result, "structuredContent", None)
+    detail = structured.get("error", structured) if isinstance(structured, dict) else {}
+    code = detail.get("code") or detail.get("error_code") if isinstance(detail, dict) else None
+    if code in (403, "PERMISSION_DENIED", "MCP_PERMISSION_DENIED"):
+        return AuthError("MCP_PERMISSION_DENIED", "MCP permission denied.", 403, integration_id)
+    if code in (401, "UNAUTHENTICATED", "MCP_USER_AUTHORIZATION_INVALID"):
+        return AuthError(
+            "MCP_USER_AUTHORIZATION_INVALID",
+            "MCP authorization was rejected.",
+            401,
+            integration_id,
+        )
+    if code == -32042:
+        return AuthError(
+            "MCP_AUTHORIZATION_REQUIRED",
+            "Authorize the configured service in Databricks.",
+            401,
+            integration_id,
+        )
+    return None