Update agent examples, quickstart and skills to support Lakebase Autoscaling by jennsun · Pull Request #151 · databricks/app-templates

jennsun · 2026-03-04T19:54:55Z

updated quickstart experience to select instance:

examples work

ex app: https://dbc-fb904e0e-bca7.staging.cloud.databricks.com/apps/agent-longtermj?o=4146361371977516

uses askuserquestion to get lakebase information upfront to run quickstart script

.scripts/source/test_quickstart.py

agent-langgraph-long-term-memory/pyproject.toml

dhruv0811

Looks good overall!! Left some comments for local testing cleanup and some clarifications on lakebase options.

.gitignore

agent-openai-agents-sdk-short-term-memory/databricks.yml

dhruv0811 · 2026-03-05T23:47:47Z

agent-openai-agents-sdk-short-term-memory/.env.example

+# TODO: Update with your Lakebase instance for session storage
+# Option 1: Provisioned instance (set instance name)
+# LAKEBASE_INSTANCE_NAME=
+# Option 2: Autoscaling instance (set project and branch)


Did we not have a third option with just autoscaling_endpoint?

the thinking here is that when users develop locally, they do so via passing in project/branch name. the autoscaling_endpoint is only used when we are in the app environment and we can directly read from PGENDPOINT (since it's not as user friendly to pass in) - hence adding this check: https://github.com/databricks/app-templates/pull/151/changes#diff-1e8ffff9ea9c91fc9dce7d8af129175d1fcfe357eabbcec14705feff13eeacacR53

agent-langgraph-short-term-memory/databricks.yml

agent-langgraph-short-term-memory/.gitignore

agent-langgraph-short-term-memory/agent_server/start_server.py

dhruv0811 · 2026-03-05T23:53:36Z

agent-langgraph-short-term-memory/agent_server/agent.py

        # Check for Lakebase access/connection errors
        if any(keyword in error_msg for keyword in ["permission"]):
            logger.error(f"Lakebase access error: {e}")
+            lakebase_desc = LAKEBASE_INSTANCE_NAME or LAKEBASE_AUTOSCALING_ENDPOINT or f"{LAKEBASE_AUTOSCALING_PROJECT}/{LAKEBASE_AUTOSCALING_BRANCH}"


Is it worth doing some validation like we did in the SDK to ensure that we only have one valid combination of variables set here? Say a user has project and branch set, but accidentally also set instance name, they should be notified rather than silently picking up instance name and short circuiting here.

since the sdk covers all these cases/will throw errors, I think adding validation here would be a bit redundant

agent-langgraph-long-term-memory/databricks.yml

jennsun · 2026-03-06T02:41:55Z

agent-openai-agents-sdk-short-term-memory/agent_server/agent.py

+LAKEBASE_AUTOSCALING_ENDPOINT = os.getenv("PGENDPOINT") if _is_app_env else None
+LAKEBASE_AUTOSCALING_PROJECT = os.getenv("LAKEBASE_AUTOSCALING_PROJECT") or None
+LAKEBASE_AUTOSCALING_BRANCH = os.getenv("LAKEBASE_AUTOSCALING_BRANCH") or None
+


bit of a chicken & egg prob here when DABs are not ready -

in order for us to be able to pass in PGENDPOINT as a var, we need the autoscaling instance as a resource on the app

DAB deployment will overwrite app resources, so it would remove manually-added postgres resource when we deploy

we will therefore use LAKEBASE_AUTOSCALING_PROJECT and LAKEBASE_AUTOSCALING_BRANCH as
static env vars supported by our agent SDK

if we want endpoint extension in future, code will look something like:

# Autoscaling params: in the app environment, PGENDPOINT is provided automatically; # for local dev, use project/branch names directly. _is_app_env = bool(os.getenv("DATABRICKS_APP_NAME")) LAKEBASE_AUTOSCALING_ENDPOINT = os.getenv("PGENDPOINT") if _is_app_env else None LAKEBASE_AUTOSCALING_PROJECT = os.getenv("LAKEBASE_AUTOSCALING_PROJECT") or None LAKEBASE_AUTOSCALING_BRANCH = os.getenv("LAKEBASE_AUTOSCALING_BRANCH") or None

dhruv0811

Added some comments based on a dry run with Claude.

agent-langgraph-long-term-memory/scripts/quickstart.py

.scripts/source/quickstart.py

agent-langgraph-long-term-memory/scripts/grant_lakebase_permissions.py

agent-langgraph-short-term-memory/databricks.yml

jennsun · 2026-03-11T00:34:30Z

agent-openai-agents-sdk-short-term-memory/app.yaml

+    value: "<your-project-name>"
+  - name: LAKEBASE_AUTOSCALING_BRANCH
+    value: "<your-branch-name>"
+  # Use for provisioned lakebase resource


adding the app.yaml back for the stateful agent examples because we need an app.yaml to deploy the app for the following flow to properly add postgres as resource:
databricks bundle deploy
│
│ Uploads source code (including app.yaml) to workspace
│ Sets app config from databricks.yml (command, env vars, resources)
│ ⚠️ Overwrites app resources → wipes manually-added postgres resource
│
▼
databricks bundle run <app_resource_name>
│
│ Starts/restarts the app using config from databricks.yml
│ App boots with LAKEBASE_AUTOSCALING_PROJECT/BRANCH env vars
│ But postgres resource is NOT attached (wiped by bundle deploy)
│
▼
databricks api patch /api/2.0/apps/ --json '{"resources": [...]}'
│
│ Re-adds postgres resource to the app
│ Grants SP access to Lakebase + injects PG* env vars for frontend
│ But running app process doesn't have the new env vars yet
│
▼
databricks api post /api/2.0/apps//deployments --json '{"source_code_path": "..."}'
│
│ Triggers a redeploy (like pressing "Deploy" in UI)
│ Reads app.yaml from source code for command + env vars
│ App restarts with postgres resource env vars now available
│
▼
App is running with all env vars + postgres resource ✓

bbqiu

couple comments:

inconsistencies between the oai stm, lg stm, and lg ltm databricks.yml files
quickstart should be replacing / deleting the set of env vars that are NOT relevant to the lakebase instance type selected
adding tests so we fully understand the behavior of quickstart.py
no need to exclude lakebase client perms grant script from all the templates

in a followup, let's also make sure to update + run the integration tests

jennsun · 2026-03-11T19:06:53Z

updates:

make sure app.yaml/databricks.yml is synced across stateful examples
quickstart will sync env variables/replace unneeded ones depending on if lakebase is provisioned or autoscaling instance
updated test quickstart script
made lakebase permission granting script sync across all templates
added test for granting lakebase permissions script

- Add autoscaling (project/branch) as an alternative to provisioned Lakebase instances across all 3 memory templates - Detect app environment via DATABRICKS_APP_NAME and use PGENDPOINT (auto-injected by platform) for autoscaling in deployed apps - Update quickstart to support creating new Lakebase instances or connecting to existing provisioned/autoscaling instances - Use autoscaling_endpoint parameter (renamed in databricks-ai-bridge) - Add databricks-ai-bridge git dependency for autoscaling support - Update tests to remove update_databricks_yml_lakebase references Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

.scripts/source/quickstart.py

bbqiu · 2026-03-12T01:53:50Z

.scripts/templates.py

    "agent-langgraph-short-term-memory": {
        "sdk": "langgraph",
        "bundle_name": "agent_langgraph_short_term_memory",
+        "has_memory": True,


can we just use the databricks.yml as a source of truth as to whether or not we have a database? this way we don't have to remember to change as many things

the openai long running agent example has a lakebase persistence layer but isn't using it as a "memory store" the way the other stateful agents are. so are we sure this is the best proxy?

https://github.com/databricks/app-templates/tree/main/agent-openai-agents-sdk-long-running-agent

bbqiu

before merging, could we please

manually test + add code tests for the happy paths for provisioned and autoscaling in the quickstart
a. should modify the .env.example and databricks.yml file.
verify that all templates are ready to deploy after running quickstart for both paths

bbqiu · 2026-03-12T02:06:56Z

.scripts/source/quickstart.py

+    # If --lakebase-provisioned-name was provided, use it directly
+    if provisioned_name:
+        print(f"Using provided provisioned Lakebase instance: {provisioned_name}")
+        if not validate_lakebase_instance(profile_name, provisioned_name):


can we have a similar check for autoscaling?

added validation for autoscaling instance:

…caling

jennsun · 2026-03-12T17:25:34Z

added provisioned + autoscaling quickstart tests and tested stateful agents (manually and with claude):

jennsun · 2026-03-12T19:04:20Z

.scripts/agent-integration-tests/helpers.py


-def run_quickstart(template_dir: Path, profile: str, lakebase: str | None = None):
-    """Run `uv run quickstart --profile <profile>`, optionally with --lakebase."""
+def run_quickstart(


TODO: will update agent-integration-tests in follow-up PR to cover autoscaling cases

agent-langgraph-long-term-memory/app.yaml

bbqiu · 2026-03-12T21:25:06Z

.scripts/source/quickstart.py

+        update_env_file("LAKEBASE_AUTOSCALING_BRANCH", autoscaling_branch)
+        update_env_file("LAKEBASE_INSTANCE_NAME", "")
+
+        update_env_file("PGUSER", username)


do we not need pghost when using autoscaling? afaict, it's not looking for PGENDPOINT

added back all the PG vars we need for frontend UI to work properly:

bbqiu

stamping to unblock, but please address comments before merging! i also ran all integration tests for the lakebase provisioned version of things and it worked!

thank you for getting this out so quicly!

non-blocking -- we can probably simplify the setup_lakebase script a bit, but it's fine as is for now

dhruv0811

Looks good overall. I ran a couple dry runs with claude, and noticed a few issues:

DELETE privilege doesnt exist
.env updates are creating new lines
Should be pinning latest sdk version since that breaks grant privileges script if not on latest
deploy sequence ordering?
some other local vars snuck into pyproject.toml I think

agent-langgraph-long-term-memory/app.yaml

.scripts/source/grant_lakebase_permissions.py

dhruv0811 · 2026-03-12T21:32:04Z

.scripts/source/quickstart.py

+        instance_info = validate_lakebase_instance(profile_name, provisioned_name)
+        if not instance_info:
+            sys.exit(1)
+        update_env_file("LAKEBASE_INSTANCE_NAME", provisioned_name)


It seems update_env_file appends to the env file and creates duplicate entries since it doesn't match the comment line in .env.example. This is more of a nit, but it'd be cleaner to properly replace the commented env var rather than add a new line with the var?

updated so it should replace instead!

.claude/skills/add-tools-langgraph/examples/lakebase-autoscaling.md

…update quickstart to replace env vars

…leases/tag/v0.79.0

jennsun mentioned this pull request Mar 4, 2026

Update agent examples to support Lakebase Autoscaling #147

Closed

jennsun changed the title ~~Update agent examples to support Lakebase Autoscaling~~ Update agent examples and quickstart to support Lakebase Autoscaling Mar 4, 2026

jennsun marked this pull request as ready for review March 4, 2026 21:44

jennsun force-pushed the lakebase-autoscaling-updates-v2 branch from 248b026 to e39267c Compare March 5, 2026 23:21

jennsun commented Mar 5, 2026

View reviewed changes

.scripts/source/test_quickstart.py Outdated Show resolved Hide resolved

jennsun commented Mar 5, 2026

View reviewed changes

agent-langgraph-long-term-memory/pyproject.toml Show resolved Hide resolved

dhruv0811 reviewed Mar 5, 2026

View reviewed changes

jennsun requested review from bbqiu and dhruv0811 March 6, 2026 02:36

jennsun commented Mar 6, 2026

View reviewed changes

dhruv0811 reviewed Mar 9, 2026

View reviewed changes

jennsun changed the title ~~Update agent examples and quickstart to support Lakebase Autoscaling~~ Update agent examples, quickstart and skills to support Lakebase Autoscaling Mar 10, 2026

jennsun commented Mar 11, 2026

View reviewed changes

bbqiu reviewed Mar 11, 2026

View reviewed changes

jennsun requested review from bbqiu and dhruv0811 March 11, 2026 21:48

jennsun and others added 13 commits March 11, 2026 16:44

rm dab postgres resource since not ready in apps yet

f2a8398

restore to template scaffolding code

e74d9b8

revert gitignore fixes

12cca49

.

faa24fc

revert dab blocks

0f1186e

.

b5631c7

revert quickstart test

d2e8416

autoscaling lakebase claude skills

b23405e

specify sp client id

cebed77

rm PGENDPOINT

c17d8dc

.

f0460b7

rm passing in endpoint

e9ba0c1

jennsun added 2 commits March 11, 2026 16:44

update test quickstart

0bfad4e

add tests for grant lakebase permissions

2cdc939

jennsun force-pushed the lakebase-autoscaling-updates-v2 branch from bf85d09 to 2cdc939 Compare March 11, 2026 23:45

jennsun added 2 commits March 11, 2026 17:02

add lakebase configs in quickstart skill

e12a92d

.

6814088

bbqiu reviewed Mar 12, 2026

View reviewed changes

.scripts/source/quickstart.py Show resolved Hide resolved

bbqiu reviewed Mar 12, 2026

View reviewed changes

.scripts/source/quickstart.py Show resolved Hide resolved

bbqiu reviewed Mar 12, 2026

View reviewed changes

.scripts/source/quickstart.py Show resolved Hide resolved

bbqiu reviewed Mar 12, 2026

View reviewed changes

jennsun added 5 commits March 11, 2026 19:13

revert pguser/pgdatabase/pghost vars needed for frontend

6d622e0

add lakebase in .yml file

ea6c45d

validate lakebase autoscaling instance exists

0302b49

rename lakebase creation in quickstart to make clearer it's for autos…

543e341

…caling

test happy path for provisioned + autoscaling

483c55d

jennsun requested a review from bbqiu March 12, 2026 17:26

jennsun commented Mar 12, 2026

View reviewed changes

bbqiu reviewed Mar 12, 2026

View reviewed changes

agent-langgraph-long-term-memory/app.yaml Show resolved Hide resolved

bbqiu reviewed Mar 12, 2026

View reviewed changes

bbqiu approved these changes Mar 12, 2026

View reviewed changes

dhruv0811 reviewed Mar 12, 2026

View reviewed changes

bbqiu mentioned this pull request Mar 12, 2026

agent templates: small integ test fixes #156

Open

jennsun added 5 commits March 12, 2026 14:45

restore template names for autoscaling app.yaml

1f13c01

rm delete/make sure env vars commented out from template are cleaned up

0abb2d9

update deploy order so we grant permissions before re-deploying app, …

97adc27

…update quickstart to replace env vars

pin min to 0.79.0: https://github.com/databricks/databricks-sdk-py/re…

1be8598

…leases/tag/v0.79.0

remove hardcoded values

9685dfd

jennsun merged commit 10c59f9 into databricks:main Mar 13, 2026

Conversation

jennsun commented Mar 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Uh oh!

dhruv0811 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

dhruv0811 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

bbqiu left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jennsun commented Mar 11, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

bbqiu left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jennsun Mar 12, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jennsun commented Mar 12, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

bbqiu left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

dhruv0811 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

jennsun commented Mar 4, 2026 •

edited

Loading

bbqiu left a comment •

edited

Loading

jennsun commented Mar 11, 2026 •

edited

Loading

jennsun Mar 12, 2026 •

edited

Loading

jennsun commented Mar 12, 2026 •

edited

Loading

bbqiu left a comment •

edited

Loading