Skip to content

Simplify Mason agent demo with SDK durable runtime - #479

Draft
shivam5 wants to merge 18 commits into
databricks:mainfrom
shivam5:demo_with_runtime
Draft

Simplify Mason agent demo with SDK durable runtime#479
shivam5 wants to merge 18 commits into
databricks:mainfrom
shivam5:demo_with_runtime

Conversation

@shivam5

@shivam5 shivam5 commented Aug 27, 2026

Copy link
Copy Markdown
Contributor

What changed

This is the end-to-end Mason + Agent Platform demo built from #474, #478, and the SDK-owned entrypoint idea in #476.

  • Add DatabricksDurableApp, which owns /invocations, background execution, polling, recovery heartbeats, and durable SSE replay.
  • Keep the framework-native JSON payload unchanged. Runtime context moves to headers:
    • Idempotency-Key for the durable invocation ID.
    • X-Routing-Key for the agent session ID.
  • Use an in-memory durability store automatically for local development.
  • Switch the runtime and LangGraph checkpointer to the managed Session Store's Lakebase project when AGENT_SESSION_STORE is wired at deploy time.
  • Add mason dev.
  • Let bare mason deploy infer the current project directory, deployment name, and the profile seeded by mason init.
  • Replace the template-owned FastAPI server, background store, MCP/memory/tracing demo plumbing, and HITL sample with one SDK entrypoint, one checkpointer helper, and one sample tool.

Compared with the current #478 template, this removes 1,041 lines and adds 256 lines.

Demo

uv tool install --editable ~/databricks-ai-bridge-worktrees/demo_with_runtime/integrations/mason

mason init --framework langgraph \
  --profile e2-dogfood \
  --repo https://github.com/shivam5/databricks-ai-bridge.git \
  --ref demo_with_runtime \
  ~/my-agent

cd ~/my-agent
mason dev
uv run pytest
mason deploy

Durable deployed state remains opt-in for the demo:

mason deploy --with-session-store my-agent-sessions --create-stores

Validation

Validated on August 27, 2026:

  • 32 focused durable runtime/app/store tests pass.
  • 53 Mason unit tests pass.
  • Ruff check and format pass for the changed Python files.
  • ty check passes for the runtime and app code.
  • A fresh template scaffolded from this remote branch installs successfully and passes 8 hermetic tests.
  • The generated app was manually exercised with a stub LangGraph backend:
    • plain sync payload + header context;
    • persisted SSE ending in [DONE];
    • background submission and polling;
    • 409 on conflicting reuse of an idempotency key.

Remaining live check

The selected e2-dogfood profile currently has an invalid OAuth refresh token, so a fresh mason dev startup and mason deploy cannot complete until running:

databricks auth login --profile e2-dogfood

A model-backed sync invocation succeeded earlier from the already-running current template, before the cached credential expired.

Draft note

The generated template temporarily installs databricks-ai-bridge[agent-server] from this fork branch so the end-to-end demo can consume the runtime before it is released. Replace that direct reference with the merged/released package before landing.

elainewang-db and others added 18 commits August 25, 2026 20:53
@elainewang-db Isaac
Add mason init to scaffold a local agent project from app-templates
b7b8140 
Sparse-clones a basic template (--framework openai|langgraph) out of the
app-templates repo into a local directory, ready for `mason deploy --source`.
Repo/ref are overridable so a template can be fetched from a fork/branch before
it merges to the canonical repo.

Co-authored-by: Isaac <no-reply@databricks.com>
@elainewang-db Isaac
[templates] Add agent-langgraph-scratch template
4b567fd 
A from-scratch FastAPI LangGraph agent backend for Databricks Apps: native
LangGraph wire on both ends, in-memory checkpointer, streaming + background,
MCP loading, and human-in-the-loop tool approval (HumanInTheLoopMiddleware).

Co-authored-by: Isaac <no-reply@databricks.com>
@elainewang-db Isaac
[templates] Drop phantom AGENT_SESSION_STORE / SQLite config
0a443b4 
This scratch template uses an in-process LangGraph checkpointer; durability is a
code swap (PostgresSaver in session_store.py), not an env var. Remove the unused
AGENT_SESSION_STORE / LOCAL_SESSION_DB_PATH references from .env.example and app.yaml.

Co-authored-by: Isaac <no-reply@databricks.com>
@elainewang-db Isaac
mason init: --profile seeds a local .env so start-server runs immedia…
0c591ee 
…tely

Without a profile, the scaffolded template has no .env (it's gitignored) and
`uv run start-server` fails on missing model auth. --profile (or the -p / `mason
login` default) copies .env.example -> .env with DATABRICKS_CONFIG_PROFILE set,
making the project runnable out of the box. Never clobbers an existing .env.

Co-authored-by: Isaac <no-reply@databricks.com>
@elainewang-db Isaac
[templates] Fail fast with a clear message when Databricks auth is mi…
8d1c0d7 
…ssing

configure() now resolves a WorkspaceClient at startup, so a missing/invalid
profile raises an actionable RuntimeError (naming the profile tried + how to fix)
instead of a generic SDK error buried in the first request's traceback.

Co-authored-by: Isaac <no-reply@databricks.com>
@elainewang-db Isaac
mason init: surface the profile/.env step in Next steps when no profi…
a071916 
…le was set

Co-authored-by: Isaac <no-reply@databricks.com>
@elainewang-db Isaac
[templates] Address review feedback on agent-langgraph-scratch
7da998e 
Apply James's PR databricks#265 review:
- Move template under integrations/mason/templates/.
- Rename server/ -> runtime/ (start_server.py -> main.py, app.py -> runtime.py).
- Drop the /responses route + alias; keep only /invocations, and rename the
  background poll route to /invocations/{invocation_id} (resp_ -> inv_ ids) to
  decouple from the Responses API.
- Read the session id from the X-Routing-Key header instead of the request body.
- Remove the thin mason/wire/ module (a misnomer); inline session-id + event
  serialization into agent.py, keeping runtime.py SDK-agnostic.
- Rename BackgroundRuns.create() -> start(); drop the stray mlflow.message.format
  "openai" attribute (this template is LangChain-native).
- Trim the tutorial-style runtime docstring; point deploy docs at `mason deploy`.
- Note the local OSS MLflow tracking-server option in .env.example.

Co-authored-by: Isaac <no-reply@databricks.com>
@elainewang-db Isaac
[templates] Fix HITL: return the session id we checkpoint under
c562a80 
invoke_handler generated a fresh session id for the checkpointer thread (via
stream_handler) and a *different* one for the returned session_id, so a request
without a session id got back an id that was never persisted — the client
couldn't resume the paused run. Pin the id into the request once, up front.

Co-authored-by: Isaac <no-reply@databricks.com>
@elainewang-db Isaac
[templates] Add durable DatabricksMemorySaver via AGENT_SESSION_STORE
62ea603 
When AGENT_SESSION_STORE names a managed Session Store, checkpointer() now
returns a DatabricksMemorySaver — LangGraph's PostgresSaver over the store's
service-managed Lakebase — so full graph state, including paused HITL runs,
survives restarts and replicas. Unset keeps the in-process InMemorySaver.

checkpointer() is now async (opens/caches the Lakebase pool once). Resolving a
Session Store name to its backing Lakebase instance is a pending fast-follow in
the Session Store API; until then LAKEBASE_INSTANCE_NAME supplies it, and the
saver raises a clear error if the store is set without it.

Co-authored-by: Isaac <no-reply@databricks.com>
@elainewang-db Isaac
[templates] Use sync CheckpointSaver; keep checkpointer() sync
d990ad8 
Switch the durable path from AsyncCheckpointSaver to the sync CheckpointSaver
(both databricks-langchain[memory]; the sync one runs fine under astream via
LangGraph's threadpool). checkpointer() is sync again — no async ripple into
create_agent_graph or the tests — and the lru_cache reuses one saver/pool.

Co-authored-by: Isaac <no-reply@databricks.com>
@elainewang-db Isaac
[templates] Derive durable checkpointer from AGENT_SESSION_STORE alone
d61ae95 
Use the verified Session Store provisioning convention (shared per-workspace
project databricks-internal-lakebase-agent-session-store, default "production"
branch, autoscaling) so the durable CheckpointSaver is built from the store name
with no extra connection config — drops LAKEBASE_INSTANCE_NAME.

Known gap, documented in session_store.py: databricks-langchain's pool hardcodes
the default database, so checkpoints land in the project's default DB, not the
per-store database that holds the session's message items. Co-locate once the
pool supports selecting a database.

Co-authored-by: Isaac <no-reply@databricks.com>
@elainewang-db Isaac
[templates] Bundle databricks-langchain[memory]; clarify durable path…
afd288b 
… is deployed-app-only

Add the [memory] extra to databricks-langchain so the durable checkpointer's
Lakebase pool (langgraph-checkpoint-postgres + psycopg) is installed.

The managed Session Store's Lakebase grants the deployed app's service principal,
not human users, so the durable path fails under local user credentials. Wrap the
Postgres auth failure in a clear RuntimeError explaining this and pointing to the
in-process default for local dev; document it in README / AGENTS / .env.example.

Co-authored-by: Isaac <no-reply@databricks.com>
@elainewang-db Isaac
[templates] Fix async/sync checkpointer mismatch (durable path)
d3c2ab1 
The agent drives the graph with astream (async), but the durable path used the
sync CheckpointSaver, whose aget_tuple/aput/aput_writes are unimplemented on the
base class — so an async run raised NotImplementedError. Switch to
AsyncCheckpointSaver (AsyncPostgresSaver) and make checkpointer() async again;
open/cache one saver per process. InMemorySaver already implements both, so the
local default was unaffected.

Co-authored-by: Isaac <no-reply@databricks.com>
@shivam5
Merge branch 'pr474' into mason-demo
49600a0
@shivam5
Add SDK-provided durable agent runtime
d5913b4
@shivam5
Integrate Mason LangGraph template demo
f642822
@shivam5
Simplify Mason demo with durable app runtime
0cb5d22
@shivam5
Fix template build with runtime Git dependency
ceced67
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@shivam5 @elainewang-db