Prototype SDK-provided durable entrypoint

[shivam5]

What changed

This adds a small SDK-provided durable runtime and FastAPI application to databricks-mason, plus one new minimal LangGraph example named durability-app.

from databricks_mason import DurableAgentApp

server = DurableAgentApp(run_agent, on_resume=run_agent)
app = server.app

The runtime library lives under Mason's existing package:

integrations/mason/src/databricks_mason/runtime/
  app.py
  runtime.py
  store.py
  types.py

• runtime.py owns durable execution, heartbeats, attempt fencing, recovery, and event replay.
• store.py provides the in-memory and Lakebase-backed stores.
• app.py is the thin HTTP adapter over the runtime.
• types.py contains the public callback/context/store contracts.

Public imports are available from either databricks_mason or databricks_mason.runtime.

Template scope

This PR does not modify the existing LangGraph, OpenAI, or chat UI templates. Their effective diff against origin/main is empty.

It adds only:

integrations/mason/templates/durability-app/
  agent/agent.py
  runtime/main.py
  tests/test_app.py
  app.yaml
  pyproject.toml
  README.md

The app is intentionally deterministic and has no model dependency: one LangGraph node optionally waits and returns Processed: <message>. runtime/main.py only constructs DurableAgentApp and exports its FastAPI app.

Bare mason init selects durability-app. Explicit --framework langgraph and --framework openai retain the existing templates and behavior. Migration/cleanup of those existing templates is deferred to a separate follow-up PR.

HTTP contract

• The client supplies the invocation id.
• __Host-databricks-app-router is the only routing/session cookie.
• There are no run/session response headers, forwarded-user identity headers, or local fallback session cookie.
• POST /invocations and POST /api/invocations support foreground, background, and streaming modes.
• GET /invocations/{id} and GET /invocations/{id}/events have /api aliases.
• SSE replays persisted events and ends when execution reaches a terminal state.
• There is no /api/healthz route.
• Recovery scanning runs only when on_resume is configured.

Lakebase durability

Only projects scaffolded from durability-app automatically configure runtime durability.

1. If mason deploy --session <store> is supplied, Mason reuses that Session Store's Lakebase database.
2. Otherwise Mason reuses or provisions one dedicated <app>-durability Lakebase project.

Mason adds only the databricks_mason_runtime schema and its executions / execution_events tables to the selected database. The selected endpoint is injected as DATABRICKS_MASON_RUNTIME_ENDPOINT. Local execution falls back to an in-memory store. --no-create-stores prevents dedicated fallback provisioning.

Live end-to-end test (e2-dogfood)

Validated commit bcb2017b using a disposable copy of the new template. The temporary source included the locally built Mason wheel and uncommitted startup assertions; none of that test-only wiring is part of this PR.

Build and deploy

cd integrations/mason
uv build

uv run mason --profile e2-dogfood --output json deploy \
  p476d-bcb2017 \
  --source /tmp/pr476-live-bcb2017.fH0pcU

databricks apps get mason-p476d-bcb2017 \
  --profile e2-dogfood \
  --output json

The deployment succeeded and attached exactly one Postgres resource:

projects/mason-p476d-bcb2017-durability/branches/production/databases/databricks-postgres

Mason injected:

DATABRICKS_MASON_RUNTIME_ENDPOINT=projects/mason-p476d-bcb2017-durability/branches/production/endpoints/primary

Exercise foreground, background, replay, and persistence

The disposable app exercised the real FastAPI routes from inside the deployed process using ASGI transport and a fixed __Host-databricks-app-router cookie. It verified:

• foreground execution;
• identical request-ID replay returning the original result;
• request-ID reuse with a different input returning 409;
• background submission and polling;
• persisted started and completed events;
• SSE replay using the first event cursor;
• absence of run/session response headers.

Lakebase was inspected using:

ENDPOINT=projects/mason-p476d-bcb2017-durability/branches/production/endpoints/primary

databricks postgres get-endpoint "$ENDPOINT" \
  --profile e2-dogfood \
  --output json

databricks postgres generate-database-credential "$ENDPOINT" \
  --profile e2-dogfood \
  --output json

Before restart, databricks_mason_runtime contained:

background  COMPLETED  attempt=1  result="Processed: background"
foreground  COMPLETED  attempt=1  result="Processed: foreground"
recovery    ACTIVE     attempt=1

The schema contained exactly:

databricks_mason_runtime.execution_events
databricks_mason_runtime.executions

Stop/start recovery

databricks apps stop mason-p476d-bcb2017 \
  --profile e2-dogfood \
  --timeout 10m \
  --output json

sleep 15

databricks apps start mason-p476d-bcb2017 \
  --profile e2-dogfood \
  --timeout 10m \
  --output json

After restart, the same persisted row was:

recovery  COMPLETED  attempt=2  result="Processed: recovery"  recovered=true

execution_events contained stage=recovered, attempt=2, followed by stage=completed, attempt=2. A separate recovery-marker foreground invocation also completed after restart.

Cleanup

databricks apps delete mason-p476d-bcb2017 \
  --profile e2-dogfood \
  --output json

databricks workspace delete \
  /Workspace/Users/shivam.mittal@databricks.com/mason_deployments/mason-p476d-bcb2017 \
  --recursive \
  --profile e2-dogfood \
  --output json

databricks postgres delete-project \
  projects/mason-p476d-bcb2017-durability \
  --purge \
  --profile e2-dogfood \
  --timeout 20m \
  --output json

Follow-up get calls confirmed the temporary App, workspace source, and Lakebase project no longer exist.

Validation

Validated on the latest origin/main (5e75997ca8f51a526a742d95fb0ecfdefad5fc44):

• Mason unit tests: 345 passed
• Standalone durability-app tests: 3 passed
• ruff check
• ruff format --check
• ty check
• uv lock --check
• git diff --check
• Mason wheel and standalone template wheel build successfully
• Wheel contains agent_durability_store.py and runtime/{app,runtime,store,types}.py
• Existing template directories have no effective diff against origin/main
• Live e2-dogfood deployment, Lakebase persistence, and stop/start recovery passed

[shivam5]

Interface comparison set:

• Prototype transport-neutral durable runtime library #475 — transport-neutral library; developer owns the server
• Prototype SDK-provided durable entrypoint #476 — AgentCore-style @app.entrypoint; SDK supplies the app
• Prototype generic JSON durable agent server #477 — supplied generic JSON server
• Generalize long-running agent recovery and replay #467 — Responses-specific LongRunningAgentServer

The first three use the same tiny progress agent and the same run/heartbeat/event durability semantics so the developer-facing differences are directly comparable.

[elainewang-db]

Can we make a new template with using our new durable runtime for now? I think we should keep the old templates with the purpose of illustrating how people can think about writing their existing agents with our databricks_mason sdk. we can point mason init to pull in the new template by default, wdyt?

In a follow up, we really should spend some time organizing the templates and pulling out the common and sdk agnostic pieces, so that the template folders themselves only have the minimum of what they need and the common pieces are all shared and copied in by mason init.

[shivam5]

Addressed Elaine’s template-scope feedback: this PR now leaves all existing LangGraph/OpenAI/chat templates unchanged and adds only the minimal durability-app default. I moved the existing-template migration and cleanup into draft follow-up #550.

---

This comment was generated with GitHub MCP.

[junchoi-db]

qq: when execution fails -- do we also cleanup the heartbeat task?