databricks-mason 0.1.4.dev0__tar.gz → 0.1.5.dev0__tar.gz

{databricks_mason-0.1.4.dev0 → databricks_mason-0.1.5.dev0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: databricks-mason
-Version: 0.1.4.dev0
+Version: 0.1.5.dev0
 Summary: Databricks integration for Mason
 Author-email: Databricks <agent-feedback@databricks.com>
 License-File: NOTICE
@@ -14,6 +14,7 @@ Requires-Dist: tomli>=2.0
 Requires-Dist: tomlkit>=0.13
 Provides-Extra: runtime
 Requires-Dist: databricks-agents>=1.9.3; extra == 'runtime'
+Requires-Dist: databricks-ai-bridge[memory]>=0.21.0; extra == 'runtime'
 Requires-Dist: databricks-langchain>=0.17.0; extra == 'runtime'
 Requires-Dist: fastapi>=0.129.0; extra == 'runtime'
 Requires-Dist: langchain-mcp-adapters>=0.2.1; extra == 'runtime'
@@ -21,15 +22,16 @@ Requires-Dist: langchain>=1.0.0; extra == 'runtime'
 Requires-Dist: langgraph>=1.1.0; extra == 'runtime'
 Requires-Dist: mlflow>=3.10.1; extra == 'runtime'
 Requires-Dist: opentelemetry-exporter-otlp-proto-grpc>=1.25.0; extra == 'runtime'
-Requires-Dist: uuid-utils>=0.10.0; extra == 'runtime'
+Requires-Dist: uvicorn>=0.20.0; extra == 'runtime'
 Provides-Extra: runtime-openai
 Requires-Dist: databricks-agents>=1.9.3; extra == 'runtime-openai'
+Requires-Dist: databricks-ai-bridge[memory]>=0.21.0; extra == 'runtime-openai'
 Requires-Dist: databricks-openai>=0.13.0; extra == 'runtime-openai'
 Requires-Dist: fastapi>=0.129.0; extra == 'runtime-openai'
 Requires-Dist: mlflow>=3.10.1; extra == 'runtime-openai'
 Requires-Dist: openai-agents>=0.7.0; extra == 'runtime-openai'
 Requires-Dist: opentelemetry-exporter-otlp-proto-grpc>=1.25.0; extra == 'runtime-openai'
-Requires-Dist: uuid-utils>=0.10.0; extra == 'runtime-openai'
+Requires-Dist: uvicorn>=0.20.0; extra == 'runtime-openai'
 Provides-Extra: tracing
 Requires-Dist: mlflow[databricks]>=3.10.1; extra == 'tracing'
 Description-Content-Type: text/markdown
@@ -62,6 +64,12 @@ For tracing commands, install Mason with tracing extras:
 pip install 'databricks-mason[tracing]'
 ```
 
+For the SDK-hosted durable agent application, install the runtime extra:
+
+```sh
+pip install 'databricks-mason[runtime]'
+```
+
 ## Shell completion
 Add this to `~/.zshrc`:
 ```sh
@@ -142,13 +150,74 @@ transport will be replaced by the generated `WorkspaceClient.mason` service when
 is released, without changing this public surface. Deployment, sandbox, tracing, and
 the existing CLI commands remain separate.
 
+## Agent application
+
+`AgentApp` provides Mason's invocation HTTP contract, including foreground, streaming, background,
+polling, and event endpoints. By default its state is process-local. Set `durable_runtime=True` to
+use Lakebase persistence, heartbeats, and crash recovery after deployment:
+
+```python
+from databricks_mason import AgentApp, DurableAgentContext
+
+app = AgentApp(durable_runtime=True)
+
+
+@app.invoke
+async def invoke(input: object, context: DurableAgentContext) -> object:
+    return await run_agent(input, session_id=context.session_id)
+
+
+@app.on_recovery
+async def recover(input: object, context: DurableAgentContext) -> object:
+    return await recover_agent(input, session_id=context.session_id)
+```
+
+The Mason server exposes `POST /api/invocations`, `GET /api/invocations/{invocation_id}`, and
+`GET /api/invocations/{invocation_id}/events?after={cursor}`. Databricks Apps bearer-token requests
+must use `/api/` routes
+([Apps documentation](https://docs.databricks.com/aws/en/dev-tools/databricks-apps/connect-local)).
+The client supplies a UUID `id`, which is also the idempotency key for every invocation mode:
+
+- foreground sync returns `200` with the result under `output`;
+- background sync returns `202` with a status URL;
+- foreground streaming returns `200` server-sent events; and
+- background streaming returns `202` with status and event URLs.
+
+`input` and `output` may be any JSON value. Transport fields are not passed to the callback. A
+top-level `session_id` is rejected, but a framework template may carry its own stable application
+session inside `input`. Polling uses only the invocation ID and relies on Databricks Apps
+authentication. Without the durable runtime, request state and events exist only in the serving
+process and horizontally scaled clients need sticky routing. With the durable runtime, Mason
+persists the input, attempt status, heartbeats, lifecycle events, application events, and output.
+
+Durability is enabled by default for both framework templates. Mason writes the durability binding
+to `agent.toml`, and `mason deploy` then attaches one Lakebase database for runtime durability,
+chosen in this order:
+
+1. Reuse the configured Session Store's Lakebase database.
+2. Otherwise reuse or provision a dedicated `<app>-durability` Lakebase project.
+
+Mason adds only its `databricks_mason_runtime_<app-hash>` schema and tables to the selected database,
+giving each app one owned schema. A replacement worker claims a stale heartbeat and calls the
+`@app.on_recovery` handler. If that handler is omitted, startup warns that automatic crash recovery
+is disabled; register the same function for both decorators when replaying the initial invocation is
+safe. Agent checkpoint restoration and idempotent external side effects remain the developer's
+responsibility.
+
+Bare `mason init`, `--framework langgraph`, and `--framework openai` scaffold `AgentApp` with its
+durable runtime enabled. Pass `--no-durable-runtime` for the same Mason HTTP contract with
+process-local state and no Lakebase provisioning. Pass `--server custom` for a minimal FastAPI
+server with one foreground `/invocations` route and no Mason `AgentApp`. Use `--disable-chat-app`
+independently for API-only Mason server output.
+
 ## Commands
 
 ```text
 mason [-p <profile>] [-o text|json]
   login        [--profile P]
   logout
-  init         [--framework openai|langgraph] [--disable-chat-app]
+  init         [--framework openai|langgraph] [--server mason|custom]
+               [--no-durable-runtime] [--disable-chat-app]
                [--profile P] [--repo URL] [--ref REF] [directory]
   dev          [--source PATH] [--prepare-environment] [--app-port PORT]
                [--with-traces C.S]
@@ -173,7 +242,7 @@ mason [-p <profile>] [-o text|json]
     add uc-function  FUNCTION [--name NAME] [--source PATH]
     add python       NAME [--source PATH]
     list             [--source PATH]
-  deploy       <name> --source PATH [--with-traces C.S]
+  deploy       <name> --source PATH [--with-traces C.S] [--instances N]
   deployments  list | get | logs | start | stop | delete
 ```
 
@@ -252,7 +321,7 @@ mason init --framework langgraph \
   --profile <profile> \
   ./my-agent
 cd ./my-agent
-uv run start-server
+mason dev
 ```
 
 The chat app includes synchronous, SSE streaming, background polling, Session Store, Memory Store,
@@ -271,10 +340,10 @@ mason --profile <profile> deploy mason-agent-demo --source .
 exists. The agent reads the bound stores from `agent.toml` at runtime; `deploy` grants the app's
 service principal access to them.)
 
-The Databricks Apps `__Host-databricks-app-router` cookie is both the sticky routing key and the
-application session id. The browser sends it automatically; API clients must reuse it as a cookie.
-Request bodies never carry `session_id`. A localhost-only `mason-local-session` cookie provides the
-same behavior outside Databricks Apps. TODO: move to `X-Routing-Key` when Apps supports it.
+The chat UI generates a stable application session UUID in browser local storage, places it inside
+the durable invocation's opaque `input`, and creates a fresh invocation UUID per turn. The
+`__Host-databricks-app-router` cookie remains independent: API clients may reuse it for sticky
+replica routing, but it is neither authentication nor the template's application session state.
 
 The generated `README.md` documents every request the client makes: config discovery, sync and SSE
 invocations, background submission and polling, session transcript loading, HITL resume, and memory

{databricks_mason-0.1.4.dev0 → databricks_mason-0.1.5.dev0}/README.md

@@ -26,6 +26,12 @@ For tracing commands, install Mason with tracing extras:
 pip install 'databricks-mason[tracing]'
 ```
 
+For the SDK-hosted durable agent application, install the runtime extra:
+
+```sh
+pip install 'databricks-mason[runtime]'
+```
+
 ## Shell completion
 Add this to `~/.zshrc`:
 ```sh
@@ -106,13 +112,74 @@ transport will be replaced by the generated `WorkspaceClient.mason` service when
 is released, without changing this public surface. Deployment, sandbox, tracing, and
 the existing CLI commands remain separate.
 
+## Agent application
+
+`AgentApp` provides Mason's invocation HTTP contract, including foreground, streaming, background,
+polling, and event endpoints. By default its state is process-local. Set `durable_runtime=True` to
+use Lakebase persistence, heartbeats, and crash recovery after deployment:
+
+```python
+from databricks_mason import AgentApp, DurableAgentContext
+
+app = AgentApp(durable_runtime=True)
+
+
+@app.invoke
+async def invoke(input: object, context: DurableAgentContext) -> object:
+    return await run_agent(input, session_id=context.session_id)
+
+
+@app.on_recovery
+async def recover(input: object, context: DurableAgentContext) -> object:
+    return await recover_agent(input, session_id=context.session_id)
+```
+
+The Mason server exposes `POST /api/invocations`, `GET /api/invocations/{invocation_id}`, and
+`GET /api/invocations/{invocation_id}/events?after={cursor}`. Databricks Apps bearer-token requests
+must use `/api/` routes
+([Apps documentation](https://docs.databricks.com/aws/en/dev-tools/databricks-apps/connect-local)).
+The client supplies a UUID `id`, which is also the idempotency key for every invocation mode:
+
+- foreground sync returns `200` with the result under `output`;
+- background sync returns `202` with a status URL;
+- foreground streaming returns `200` server-sent events; and
+- background streaming returns `202` with status and event URLs.
+
+`input` and `output` may be any JSON value. Transport fields are not passed to the callback. A
+top-level `session_id` is rejected, but a framework template may carry its own stable application
+session inside `input`. Polling uses only the invocation ID and relies on Databricks Apps
+authentication. Without the durable runtime, request state and events exist only in the serving
+process and horizontally scaled clients need sticky routing. With the durable runtime, Mason
+persists the input, attempt status, heartbeats, lifecycle events, application events, and output.
+
+Durability is enabled by default for both framework templates. Mason writes the durability binding
+to `agent.toml`, and `mason deploy` then attaches one Lakebase database for runtime durability,
+chosen in this order:
+
+1. Reuse the configured Session Store's Lakebase database.
+2. Otherwise reuse or provision a dedicated `<app>-durability` Lakebase project.
+
+Mason adds only its `databricks_mason_runtime_<app-hash>` schema and tables to the selected database,
+giving each app one owned schema. A replacement worker claims a stale heartbeat and calls the
+`@app.on_recovery` handler. If that handler is omitted, startup warns that automatic crash recovery
+is disabled; register the same function for both decorators when replaying the initial invocation is
+safe. Agent checkpoint restoration and idempotent external side effects remain the developer's
+responsibility.
+
+Bare `mason init`, `--framework langgraph`, and `--framework openai` scaffold `AgentApp` with its
+durable runtime enabled. Pass `--no-durable-runtime` for the same Mason HTTP contract with
+process-local state and no Lakebase provisioning. Pass `--server custom` for a minimal FastAPI
+server with one foreground `/invocations` route and no Mason `AgentApp`. Use `--disable-chat-app`
+independently for API-only Mason server output.
+
 ## Commands
 
 ```text
 mason [-p <profile>] [-o text|json]
   login        [--profile P]
   logout
-  init         [--framework openai|langgraph] [--disable-chat-app]
+  init         [--framework openai|langgraph] [--server mason|custom]
+               [--no-durable-runtime] [--disable-chat-app]
                [--profile P] [--repo URL] [--ref REF] [directory]
   dev          [--source PATH] [--prepare-environment] [--app-port PORT]
                [--with-traces C.S]
@@ -137,7 +204,7 @@ mason [-p <profile>] [-o text|json]
     add uc-function  FUNCTION [--name NAME] [--source PATH]
     add python       NAME [--source PATH]
     list             [--source PATH]
-  deploy       <name> --source PATH [--with-traces C.S]
+  deploy       <name> --source PATH [--with-traces C.S] [--instances N]
   deployments  list | get | logs | start | stop | delete
 ```
 
@@ -216,7 +283,7 @@ mason init --framework langgraph \
   --profile <profile> \
   ./my-agent
 cd ./my-agent
-uv run start-server
+mason dev
 ```
 
 The chat app includes synchronous, SSE streaming, background polling, Session Store, Memory Store,
@@ -235,10 +302,10 @@ mason --profile <profile> deploy mason-agent-demo --source .
 exists. The agent reads the bound stores from `agent.toml` at runtime; `deploy` grants the app's
 service principal access to them.)
 
-The Databricks Apps `__Host-databricks-app-router` cookie is both the sticky routing key and the
-application session id. The browser sends it automatically; API clients must reuse it as a cookie.
-Request bodies never carry `session_id`. A localhost-only `mason-local-session` cookie provides the
-same behavior outside Databricks Apps. TODO: move to `X-Routing-Key` when Apps supports it.
+The chat UI generates a stable application session UUID in browser local storage, places it inside
+the durable invocation's opaque `input`, and creates a fresh invocation UUID per turn. The
+`__Host-databricks-app-router` cookie remains independent: API clients may reuse it for sticky
+replica routing, but it is neither authentication nor the template's application session state.
 
 The generated `README.md` documents every request the client makes: config discovery, sync and SSE
 invocations, background submission and polling, session transcript loading, HITL resume, and memory

{databricks_mason-0.1.4.dev0 → databricks_mason-0.1.5.dev0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "databricks-mason"
-version = "0.1.4.dev0"
+version = "0.1.5.dev0"
 description = "Databricks integration for Mason"
 authors = [
     { name="Databricks", email="agent-feedback@databricks.com" },
@@ -27,24 +27,26 @@ tracing = [
 # OpenAI Agents SDK adapter (databricks_mason.openai). Both carry the shared framework-neutral stack
 # (databricks_mason.runtime); the framework SDKs differ, so an agent installs only the one it uses.
 runtime = [
+    "databricks-ai-bridge[memory]>=0.21.0",
     "databricks-langchain>=0.17.0",
     "langgraph>=1.1.0",
     "langchain>=1.0.0",
     "langchain-mcp-adapters>=0.2.1",
     "fastapi>=0.129.0",
     "mlflow>=3.10.1",
-    "uuid-utils>=0.10.0",
     "opentelemetry-exporter-otlp-proto-grpc>=1.25.0",
     "databricks-agents>=1.9.3",
+    "uvicorn>=0.20.0",
 ]
 runtime-openai = [
+    "databricks-ai-bridge[memory]>=0.21.0",
     "openai-agents>=0.7.0",
     "databricks-openai>=0.13.0",
     "fastapi>=0.129.0",
     "mlflow>=3.10.1",
-    "uuid-utils>=0.10.0",
     "opentelemetry-exporter-otlp-proto-grpc>=1.25.0",
     "databricks-agents>=1.9.3",
+    "uvicorn>=0.20.0",
 ]
 
 [project.scripts]
@@ -57,7 +59,12 @@ dev = [
     { include-group = "tests" },
 ]
 tests = [
+    "databricks-ai-bridge[memory]>=0.21.0",
+    "fastapi>=0.129.0",
+    "httpx>=0.28.1",
     "pytest==9.0.2",
+    "pytest-asyncio==1.3.0",
+    "uvicorn>=0.20.0",
 ]
 
 [build-system]
@@ -75,6 +82,9 @@ include = [
 [tool.hatch.build.targets.wheel]
 packages = ["src/databricks_mason"]
 
+[tool.uv.sources]
+databricks-ai-bridge = { path = "../..", editable = true }
+
 [tool.ruff]
 include = ["pyproject.toml", "src/**/*.py", "tests/**/*.py"]
 extend = "../../pyproject.toml"

{databricks_mason-0.1.4.dev0 → databricks_mason-0.1.5.dev0}/src/databricks_mason/__init__.py

@@ -8,6 +8,8 @@ from databricks_mason.session_store import Session, SessionItem, SessionStore
 
 if TYPE_CHECKING:
     from databricks_mason.runtime import (
+        AgentApp,
+        DurableAgentContext,
         configure_tracing,
         tag_session,
         workspace_client,
@@ -22,6 +24,8 @@ __all__ = [
     "Session",
     "SessionItem",
     "SessionStore",
+    "AgentApp",
+    "DurableAgentContext",
     "configure_tracing",
     "tag_session",
     "workspace_client",
@@ -29,7 +33,14 @@ __all__ = [
 ]
 
 _RUNTIME_REEXPORTS = frozenset(
-    {"configure_tracing", "tag_session", "workspace_client", "workspace_headers"}
+    {
+        "AgentApp",
+        "DurableAgentContext",
+        "configure_tracing",
+        "tag_session",
+        "workspace_client",
+        "workspace_headers",
+    }
 )
 
 

{databricks_mason-0.1.4.dev0 → databricks_mason-0.1.5.dev0}/src/databricks_mason/agent_project.py

@@ -18,6 +18,7 @@ from databricks_mason.errors import AgentCliError
 from databricks_mason.runtime.tool_manifest import MEMORY_STORE_TABLE, SESSION_STORE_TABLE
 
 _SCHEMA_VERSION = 1
+_DURABILITY_TABLE = "durability"
 _SUPPORTED_FRAMEWORKS = {"langgraph", "openai"}
 _SUPPORTED_SCOPE_KINDS = {"table", "volume", "workspace"}
 _SUPPORTED_PERMISSIONS = {"read_only", "read_write"}
@@ -201,6 +202,16 @@ def _store_name_from_manifest(value: object, table: str) -> str | None:
     return _required_string(cast(Mapping[str, Any], value).get("name"), f"[{table}] name")
 
 
+def _durability_from_manifest(value: object) -> bool:
+    if value is None:
+        return False
+    if not isinstance(value, Mapping):
+        raise AgentCliError("agent.toml [durability] must be a table.")
+    if cast(Mapping[str, Any], value).get("enabled") is not True:
+        raise AgentCliError("agent.toml [durability] must set enabled = true.")
+    return True
+
+
 def _store_id_from_manifest(value: object) -> str | None:
     """Read the optional bare store ``id`` from a ``[memory_store]`` table, or None if absent."""
     if not isinstance(value, Mapping):
@@ -296,6 +307,8 @@ class AgentProject:
         memory_store: str | None = None,
         session_store: str | None = None,
         memory_store_id: str | None = None,
+        deployment_name: str | None = None,
+        durability_enabled: bool = False,
     ) -> None:
         self.root = root
         self.path = root / "agent.toml"
@@ -307,6 +320,9 @@ class AgentProject:
         self.memory_store = memory_store
         self.session_store = session_store
         self.memory_store_id = memory_store_id
+        # The deployment's base name (`mason deploy` prefixes it with `mason-`); None until named.
+        self.deployment_name = deployment_name
+        self.durability_enabled = durability_enabled
 
     @classmethod
     def load(cls, root: pathlib.Path | str) -> "AgentProject":
@@ -333,6 +349,11 @@ class AgentProject:
         framework = _required_string(agent.get("framework"), "agent.framework")
         if framework not in _SUPPORTED_FRAMEWORKS:
             raise AgentCliError(f"Unsupported Mason framework {framework!r}.")
+        deployment_name = agent.get("deployment_name")
+        if deployment_name is not None and not (
+            isinstance(deployment_name, str) and deployment_name
+        ):
+            raise AgentCliError("agent.toml [agent] deployment_name must be a non-empty string.")
         raw_tools = document.get("tools", [])
         if not isinstance(raw_tools, list):
             raise AgentCliError("agent.toml tools must be an array of tables.")
@@ -347,6 +368,7 @@ class AgentProject:
         session_store = _store_name_from_manifest(
             document.get(SESSION_STORE_TABLE), SESSION_STORE_TABLE
         )
+        durability_enabled = _durability_from_manifest(document.get(_DURABILITY_TABLE))
         return cls(
             project_root,
             document,
@@ -355,10 +377,18 @@ class AgentProject:
             memory_store,
             session_store,
             memory_store_id,
+            str(deployment_name) if deployment_name is not None else None,
+            durability_enabled,
         )
 
     @classmethod
-    def create(cls, root: pathlib.Path | str, *, framework: str) -> "AgentProject":
+    def create(
+        cls,
+        root: pathlib.Path | str,
+        *,
+        framework: str,
+        durability_enabled: bool = False,
+    ) -> "AgentProject":
         if framework not in _SUPPORTED_FRAMEWORKS:
             raise AgentCliError(f"Unsupported Mason framework {framework!r}.")
         project_root = pathlib.Path(root).expanduser().resolve()
@@ -368,7 +398,29 @@ class AgentProject:
         agent = tomlkit.table()
         agent.add("framework", framework)
         document.add("agent", agent)
-        return cls(project_root, document, framework, [])
+        if durability_enabled:
+            durability = tomlkit.table()
+            durability.add("enabled", True)
+            document.add(_DURABILITY_TABLE, durability)
+        return cls(
+            project_root,
+            document,
+            framework,
+            [],
+            durability_enabled=durability_enabled,
+        )
+
+    def set_deployment_name(self, name: str) -> bool:
+        """Record the deployment's base name under [agent].deployment_name. True if it changed."""
+        name = _required_string(name, "[agent] deployment_name")
+        if self.deployment_name == name:
+            return False
+        agent = self._document.get("agent")
+        if not isinstance(agent, Mapping):
+            raise AgentCliError("agent.toml must declare an [agent] table.")
+        agent["deployment_name"] = name
+        self.deployment_name = name
+        return True
 
     def add_tool(self, spec: ToolSpec) -> bool:
         for existing in self.tools: