How do Lakebase and Databricks Model Serving work together for stateful production agents?

Summary

• Lakebase, the fully managed Postgres database on Databricks, stores the state of a stateful agent, and a Model Serving endpoint runs the agent code. In this pattern the endpoint does not keep the conversation. Lakebase does.
• For short-term memory, the agent saves LangGraph checkpoints to Lakebase, keyed by a thread ID that clients pass in custom_inputs. For long-term memory, it stores and retrieves facts across conversations.
• Agent memory on Model Serving is in Public Preview, and Databricks recommends the Agent Bricks CLI and Agent Runtime on Databricks Apps for new agents. Lakebase-backed memory is supported on both targets.
• Managed agent memory (Beta) keeps long-term memory in Lakebase without you operating the database. Search is full-text (BM25), and access control applies to the whole store, not to individual entries.
• Registering a Lakebase database in Unity Catalog adds a read-only catalog governed by Unity Catalog permissions, while direct Postgres connections keep using Postgres roles.

How do Lakebase and Databricks Model Serving work together for stateful production agents?

Model Serving hosts the agent behind a REST endpoint, and Lakebase holds the state the agent needs between requests. On each request the agent loads its checkpoint or memory from Lakebase, does its work, and writes the new state back, so the next request can continue the conversation. Databricks documents this pattern for Model Serving. It also recommends Agent Runtime on Databricks Apps for new agents, and Lakebase-backed memory works on both.

What does each service contribute?

Capability	Why it matters for a stateful agent	Provided by
Durable state	Committed checkpoints survive compute restarts and scale-to-zero	Lakebase records every committed transaction before it acknowledges it
Autoscaling	Database compute and the endpoint both follow demand	Lakebase autoscaling and scale to zero; Model Serving scales up or down automatically
Thread-scoped memory	Follow-up requests continue the same conversation	LangGraph checkpoints in Lakebase, keyed by thread_id
Long-term memory	Preferences and facts carry across conversations	Lakebase-backed memory; managed agent memory is Beta
Governed access	Control who can call the endpoint and who can query the data	Serving endpoint ACLs; Unity Catalog for read-only SQL access; Postgres roles for direct connections
Retrieval	Search memory and documents without a separate store	Lakebase Search: lakebase_vector and lakebase_text

Lakebase is a fully managed, serverless Postgres service. It combines serverless Postgres and storage technology from Neon, which Databricks acquired, with Databricks infrastructure. Model Serving provides a unified interface to deploy, govern, and query AI models, and it supports agents as custom models.

How does a request flow through the stack?

1. The client calls the serving endpoint and passes a thread ID in custom_inputs.
2. The agent code reads the thread ID and runs its LangGraph graph with {"configurable": {"thread_id": thread_id}} and a Lakebase-backed CheckpointSaver(instance_name=LAKEBASE_INSTANCE_NAME).
3. The graph calls models and tools, and LangGraph saves checkpoints to Lakebase as it runs.
4. The response returns the thread ID in custom_outputs. The time-travel version of the agent also returns a checkpoint ID.
5. A follow-up request with the same thread ID continues the conversation. Passing a checkpoint_id forks it from an earlier checkpoint, which the docs call time travel.

The client call looks like this in the Databricks documentation:

response1 = client.responses.create(
    model=endpoint,
    input=[{"role": "user", "content": "What are stateful agents?"}],
    extra_body={"custom_inputs": {"thread_id": thread_id}},
)

Which memory types does the pattern support?

Short-term memory captures context within one conversation session using thread IDs and checkpointing. Long-term memory extracts and stores key information across conversations so the agent can personalize responses. You can implement either or both in the same agent.
Managed agent memory (Beta) is the lower-effort option for long-term memory. Creating a memory store provisions the backing Lakebase storage, and entries are organized by actor_id and path. Search ranks results by full-text (BM25) relevance and does not support vector similarity search. Access control applies to the whole store, so create a separate store for each tenant or user boundary that needs strict isolation.

Which deployment target should you choose?

For new agents, Databricks recommends the Agent Bricks CLI, which deploys an agent server to Agent Runtime on Databricks Apps. Model Serving is the legacy runtime for earlier agents, and Databricks publishes a guide for migrating them to Databricks Apps. On Databricks Apps, Databricks handles authentication between the app and Lakebase automatically, and an agent deployed with DurableAgentServer keeps its run state in a Databricks-managed Lakebase project, so interrupted runs can recover.

How do governance and search fit in?

Register the Lakebase database in Unity Catalog to create a read-only catalog. Unity Catalog permissions then control SQL warehouse queries, so analysts can join conversation history with lakehouse tables. Queries need a serverless SQL warehouse, and direct Postgres connections keep using Postgres roles independently. For retrieval, Lakebase Search is generally available. lakebase_vector is drop-in compatible with pgvector, and lakebase_text adds BM25 full-text search.

FAQs

Should I deploy a new stateful agent on Model Serving?

Databricks recommends the Agent Bricks CLI and Agent Runtime for new agents. Use Model Serving when you maintain an agent that already runs there. Agent memory on Model Serving is in Public Preview.

Do I need Unity Catalog to use both?

Model Serving accepts models registered in Unity Catalog or the workspace model registry, so Unity Catalog is not the only option. Registering a Lakebase database in Unity Catalog is optional and adds governed, read-only SQL access.

Where do embeddings live?

In Lakebase. Install lakebase_vector for approximate nearest-neighbor search, or install pgvector on its own to use the vector type.