From 280626a3060d8da1a9b088380110328540d6b51c Mon Sep 17 00:00:00 2001 From: Shivam Mittal Date: Fri, 2 Oct 2026 22:44:52 +0000 Subject: [PATCH 1/3] Fix local managed state and authoritative chat history --- integrations/agentbricks/README.md | 42 ++- integrations/agentbricks/cli.md | 3 +- .../src/databricks_agentbricks/cli/dev.py | 133 ++++++-- .../src/databricks_agentbricks/cli/memory.py | 29 +- .../databricks_agentbricks/cli/sessions.py | 29 +- .../agent-langgraph/AGENTKIT_CONTRACT.md | 5 +- .../agent-openai/AGENTKIT_CONTRACT.md | 5 +- .../templates/agent-openai/AGENTS.md | 6 +- .../templates/ui/agent-langgraph/CHAT_APP.md | 13 + .../ui/agent-langgraph/runtime/ui.py | 172 +++++++--- .../ui/agent-langgraph/tests/test_demo_ui.py | 12 +- .../templates/ui/agent-langgraph/ui/app.js | 42 +-- .../templates/ui/agent-openai/CHAT_APP.md | 13 + .../templates/ui/agent-openai/runtime/ui.py | 140 +++++--- .../agentbricks/tests/unit_tests/dev_test.py | 125 ++++++- .../unit_tests/template_state_history_test.py | 308 ++++++++++++++++++ 16 files changed, 888 insertions(+), 189 deletions(-) create mode 100644 integrations/agentbricks/tests/unit_tests/template_state_history_test.py diff --git a/integrations/agentbricks/README.md b/integrations/agentbricks/README.md index ae8ab835d..624bc3e1c 100644 --- a/integrations/agentbricks/README.md +++ b/integrations/agentbricks/README.md @@ -423,6 +423,44 @@ agentbricks memory bind support-agent-memory agentbricks deploy my-agent ``` +### Testing state before deployment + +A **store** is a workspace resource. A **binding** is the store name saved in `agent.toml`. +A **session** is one conversation inside a session store. Creating a store does not bind it, +and saving a binding does not verify that the store exists or that your profile can use it. + +Use separate development stores when testing locally. For a new pair of stores: + +```sh +agentbricks --profile sessions stores create --name my-agent-dev-sessions +agentbricks --profile memory stores create --name my-agent-dev-memory +agentbricks sessions bind my-agent-dev-sessions +agentbricks memory bind my-agent-dev-memory +agentbricks --profile dev --workspace-stores +``` + +For existing stores, skip the create commands. Binding the same names again is safe. `dev +--workspace-stores` resolves both bindings in the selected workspace and verifies read access before +starting the app; it creates no remote resources or grants. The agent uses your profile's credentials +for store operations. Write permissions are checked when a request writes data. Missing stores, +permission errors, and unavailable services stop startup instead of silently disabling state. + +| State | Default `dev` | `dev --workspace-stores` | Deployed managed server | +| --- | --- | --- | --- | +| Conversation history | In-process; lost on restart | Durable when a session store is bound | Durable when a session store is bound | +| Cross-conversation memory | Off | Enabled when a memory store is bound | Enabled when a memory store is bound | +| Invocation status, event replay, background work | In-process; lost on restart | In-process; lost on restart | Durable Runtime Store | +| OpenAI pending approval `RunState` | In-process | In-process | In-process | + +The chat UI reads the same ordered transcript/checkpoint used by the agent, so API-created turns +also appear when the UI opens that conversation. With app-auth agents, session/actor identifiers +are application-defined and stores remain shared at the store permission boundary. When the +project requires request-user auth, the runtime namespaces session and actor identifiers by trusted +request identity; the UI uses the same mapping and keeps the public browser session id stable. +Local development uses a single local-developer identity and does not simulate multiple signed-in +users. Use the same `session_id` and `input.actor` when comparing UI and API history. Changing the +auth mode changes the namespace; it does not migrate existing history. + Memory and session stores are independent resources: deleting one never affects the other. ## Commands @@ -971,8 +1009,8 @@ agentbricks --profile deploy agent-bricks-agent-demo --source . (`bind` declares the store name in `agent.toml`; `agentbricks deploy` creates any declared-but-missing store and grants the app's service principal access to it. The memory store id flows to the runtime -via the `AGENT_MEMORY_STORE` env var that `deploy` injects; `agentbricks dev` runs locally with memory off -and does not inject it. The id is not persisted in `agent.toml`.) +via the `AGENT_MEMORY_STORE` env var that `deploy` injects; `agentbricks dev --workspace-stores` +resolves and injects it after validating access. Plain `agentbricks dev` keeps memory off. The id is not persisted in `agent.toml`.) The chat UI generates a stable application session UUID in browser local storage, sends it as the invocation's top-level `session_id`, and creates a fresh invocation UUID per turn. The chat app also diff --git a/integrations/agentbricks/cli.md b/integrations/agentbricks/cli.md index e6f46af75..1cde69db2 100644 --- a/integrations/agentbricks/cli.md +++ b/integrations/agentbricks/cli.md @@ -168,7 +168,7 @@ Auth uses your Databricks profile (`-p` / `agentbricks login`), and the agent re Under the hood this wraps `databricks apps run-local`: it reads the command + env from `app.yaml` and runs the app the way the Apps runtime would, so local behavior matches a deployment. The environment is built on the first run and reused after; pass `--prepare-environment` to force a rebuild (e.g. after changing dependencies). -Everything runs locally: `agentbricks dev` is a local deployment that does not depend on a Databricks workspace for its resources. Tracing goes to a local MLflow tracking server (sqlite-backed, under the existing `.agentbricks/` state directory) so traces are recorded on your machine with no workspace experiment or setup - open the printed Traces URL to view them (`agentbricks tracing unbind` doesn't affect dev; it only stops the deployed agent's tracing). Long-term memory is off and conversation history is in-process (not durable): the memory/session stores bound with `agentbricks memory/sessions bind` are created and used only when you `agentbricks deploy`, not here. So there's nothing to provision and no service-principal grant to make; that all happens at `agentbricks deploy` time. +By default, state is local: tracing goes to a local MLflow tracking server (sqlite-backed, under `.agentbricks/`) so traces are recorded on your machine with no workspace experiment or setup - open the printed Traces URL to view them (`agentbricks tracing unbind` doesn't affect dev; it only stops the deployed agent's tracing). Long-term memory is off and conversation history is in-process, lost on restart. Pass `--workspace-stores` to exercise existing bound stores before deploying, using your selected profile's credentials. This reads and writes real workspace data; use development stores. Missing or inaccessible bindings stop startup with an actionable error. No stores or grants are created by dev. Execution state (background runs and event replay) remains in-process in both modes, even when conversation history and memory are durable. ``` agentbricks dev [options] @@ -182,6 +182,7 @@ _Options_ | `--source ` | path | `.` | no | Local source directory to run (containing app.yaml). Defaults to the current directory. | | `--prepare-environment`, `--no-prepare-environment` | flag | - | no | Build the app's environment with uv before running. Default: build only if no .venv exists yet, and reuse it otherwise. Requires uv. | | `--app-port ` | integer | - | no | Port to run the app on (default 8000). | +| `--workspace-stores` | flag | false | no | Use existing memory/session stores bound in agent.toml with your selected profile. Validates access without creating stores; local runs can read and write their data. | ### `agentbricks memory` diff --git a/integrations/agentbricks/src/databricks_agentbricks/cli/dev.py b/integrations/agentbricks/src/databricks_agentbricks/cli/dev.py index ad2a26d20..bb41058f8 100644 --- a/integrations/agentbricks/src/databricks_agentbricks/cli/dev.py +++ b/integrations/agentbricks/src/databricks_agentbricks/cli/dev.py @@ -18,6 +18,7 @@ from databricks_agentbricks import render from databricks_agentbricks.cli.deploy import ( _load_project, + _resolve_memory_store, resource_bindings, ) from databricks_agentbricks.cli.endpoint_examples import print_agent_invoke_command @@ -26,6 +27,7 @@ from databricks_agentbricks.errors import AgentCliError from databricks_agentbricks.project_config import require_managed_tool_support from databricks_agentbricks.project_types import AgentServer +from databricks_agentbricks.render import field from databricks_agentkit.runtime.store import RUNTIME_STORE_LOCAL_ENV from databricks_agentkit.runtime.tool_manifest import MEMORY_STORE_ENV, SESSION_STORE_ENV @@ -45,9 +47,8 @@ # (which MLflow resolves ahead of MLFLOW_EXPERIMENT_NAME) would point the local agent at an # experiment that doesn't exist on the local sqlite server. dev re-adds the local # MLFLOW_TRACKING_URI / MLFLOW_EXPERIMENT_NAME itself. -# - stores: dev never uses the workspace memory/session stores (memory off, sessions in-process), -# so a prior deploy's AGENT_MEMORY_STORE / AGENT_SESSION_STORE must not quietly pull dev onto -# them. dev re-adds nothing here - the runtime falls back to its local defaults. +# - stores: only --workspace-stores opts into the current project's workspace bindings. Strip a +# prior deploy's ids first, then re-add the ids verified for the selected profile, if requested. # (The Runtime Store's lakebase env needs no strip: dev already forces RUNTIME_STORE_LOCAL=true.) _DEPLOY_TRACING_ENVS = frozenset( {"MLFLOW_TRACKING_URI", "MLFLOW_EXPERIMENT_ID", "MLFLOW_TRACING_DESTINATION"} @@ -69,12 +70,19 @@ "exists yet, and reuse it otherwise. Requires uv.", ) @click.option("--app-port", type=int, default=None, help="Port to run the app on (default 8000).") +@click.option( + "--workspace-stores", + is_flag=True, + help="Use existing memory/session stores bound in agent.toml with your selected profile. " + "Validates access without creating stores; local runs can read and write their data.", +) @click.pass_obj def dev( obj, source: str, prepare_environment: Optional[bool], app_port: Optional[int], + workspace_stores: bool, ) -> None: """Run your agent locally so you can try it before deploying. @@ -90,14 +98,15 @@ def dev( deployment. The environment is built on the first run and reused after; pass `--prepare-environment` to force a rebuild (e.g. after changing dependencies). - Everything runs locally: `agentbricks dev` is a local deployment that does not depend on a Databricks - workspace for its resources. Tracing goes to a local MLflow tracking server (sqlite-backed, under `.agentbricks/`) + By default, state is local: tracing goes to a local MLflow tracking server (sqlite-backed, under `.agentbricks/`) so traces are recorded on your machine with no workspace experiment or setup - open the printed Traces URL to view them (`agentbricks tracing unbind` doesn't affect dev; it only stops the deployed - agent's tracing). Long-term memory is off and conversation history is in-process (not durable): - the memory/session stores bound with `agentbricks memory/sessions bind` are created and used only when you - `agentbricks deploy`, not here. So there's nothing to provision and no service-principal grant to make; - that all happens at `agentbricks deploy` time. + agent's tracing). Long-term memory is off and conversation history is in-process, lost on restart. + Pass --workspace-stores to exercise existing bound stores before deploying, using your selected + profile's credentials. This reads and writes real workspace data; use development stores. + Missing or inaccessible bindings stop startup with an actionable error. No stores or grants + are created by dev. Execution state (background runs and event replay) remains in-process in + both modes, even when conversation history and memory are durable. """ source_dir = pathlib.Path(source) app_yaml = source_dir / "app.yaml" @@ -111,29 +120,32 @@ def dev( if project is not None and project.tools: require_managed_tool_support(source_dir) - # `agentbricks dev` is a fully local sandbox: the Runtime Store and tracing already run locally, and - # memory/sessions follow suit here. Dev never reaches the workspace for stores (so it stays fast - # and works offline): long-term memory is off and conversation history is in-process (not - # durable), regardless of any binding. Stores are created and used only by `agentbricks deploy`; the - # deploy-written store env is stripped from the dev manifest (see `_dev_entry_point`) so a prior - # deploy can't quietly pull dev onto the workspace stores. Read the bindings only to name them. + # Keep the default local sandbox. Workspace state is an explicit opt-in, validated before + # starting either process and without modifying app.yaml or provisioning remote resources. memory_store, session_store, trace_experiment = resource_bindings(source_dir) - if memory_store: + local_env: dict[str, str] = {} + if workspace_stores: + local_env.update(_workspace_store_env(obj, memory_store, session_store)) + elif memory_store: render.console().print( f"[dim]Memory store '{memory_store}' is bound but `agentbricks dev` runs with " - "long-term memory off. Run `agentbricks deploy` to use bound store.[/]" + "long-term memory off. Use `agentbricks dev --workspace-stores` for an existing store, " + "or `agentbricks deploy` to provision it.[/]" ) - if session_store: + if session_store and not workspace_stores: render.console().print( f"[dim]Session store '{session_store}' is bound but `agentbricks dev` keeps " - "conversation history in-process (not durable). Run `agentbricks deploy` to use bound store.[/]" + "conversation history in-process (not durable). Use `agentbricks dev --workspace-stores` " + "for an existing store, or `agentbricks deploy` to provision it.[/]" ) if trace_experiment: render.console().print( f"[dim]Tracing experiment '{trace_experiment}' is bound but `agentbricks dev` " "traces to a local MLflow server. Run `agentbricks deploy` to trace to the bound experiment.[/]" ) - local_env: dict[str, str] = {} + render.console().print( + "[dim]Execution state is in-process: pending runs and replay events are lost on restart.[/]" + ) # Local tracing: start a local MLflow tracking server backed by sqlite under .agentbricks/ and point the # agent at it via the dev-only manifest — for any project, regardless of framework/server. An agent # that uses MLflow (autolog or `start_trace`) then traces to it; it's harmless for one that doesn't. @@ -225,7 +237,8 @@ def _announce_local_url( next_steps=[ f"Open {base} to chat with your agent", tool_step, - ("agentbricks memory bind ", "Attach a memory / session store"), + ("agentbricks memory bind ", "Declare a long-term memory store"), + ("agentbricks sessions bind ", "Declare a conversation history store"), (f"agentbricks deploy {deploy_name}", "Deploy it to Databricks"), ], ) @@ -260,6 +273,84 @@ def _announce_local_url( ) +def _workspace_store_env( + obj, memory_store: str | None, session_store: str | None +) -> dict[str, str]: + """Resolve existing bindings with the same profile used by run-local; never provision.""" + if not memory_store and not session_store: + raise AgentCliError( + "No memory or session store is bound in agent.toml.", + hint="Run `agentbricks sessions bind ` and/or `agentbricks memory bind `, " + "then create the named stores or select existing ones before using --workspace-stores.", + ) + try: + client = obj.client() + except Exception as exc: + raise AgentCliError( + "Could not authenticate for workspace stores.", + hint="Select your workspace with `agentbricks --profile dev --workspace-stores`. " + "Check that profile with `databricks auth describe --profile `.", + ) from exc + + env: dict[str, str] = {} + for kind, name in (("memory", memory_store), ("session", session_store)): + if not name: + continue + try: + if kind == "memory": + store = _resolve_memory_store(client, name) + if store is None: + raise AgentCliError("Store not found.", error_code="NOT_FOUND") + resource_name = field(store, "name") or "" + store_id = resource_name.removeprefix("memory-stores/") + if not store_id: + raise AgentCliError("Store response has no id.", error_code="INVALID_RESPONSE") + # Listing resolves a display name; GET verifies access to the resolved resource. + client.get_memory_store(store_id) + env[MEMORY_STORE_ENV] = store_id + else: + client.get_session_store(name) + env[SESSION_STORE_ENV] = name + except AgentCliError as exc: + code = exc.error_code + if code in {"NOT_FOUND", "RESOURCE_DOES_NOT_EXIST"}: + hint = ( + f"Check the selected profile and binding. If the store is new, create it with " + f"`agentbricks --profile {kind} stores create --name {name}`, " + "or run `agentbricks deploy` to provision the declared stores. " + "A store absent from your listing may instead require an access grant from its owner." + ) + elif code == "PERMISSION_DENIED": + hint = "Ask the store owner for access, or bind a development store you can access." + elif code in {"UNAUTHENTICATED", "UNAUTHORIZED", "INVALID_ACCESS_TOKEN"}: + hint = "Refresh credentials with `databricks auth login --profile `." + else: + hint = exc.hint or ( + "Check the selected profile, workspace connectivity, and store availability, " + "then retry. No local server or workspace store was created." + ) + raise AgentCliError( + f"Could not verify bound {kind} store '{name}' for local development.", + error_code=code, + hint=hint, + ) from exc + + render.success( + "Using workspace stores for local development", + fields={ + "Workspace": client.host, + "Profile": obj.profile or "configured credentials", + "Conversation history": session_store or "In-process (lost on restart)", + "Long-term memory": memory_store or "Off", + }, + ) + render.console().print( + "[dim]Store read access verified. Agent requests use your credentials and can write to " + "these stores; write permissions are checked when used.[/]" + ) + return env + + def _dev_entry_point( app_yaml: pathlib.Path, extra_env: dict[str, str] | None = None ) -> pathlib.Path: diff --git a/integrations/agentbricks/src/databricks_agentbricks/cli/memory.py b/integrations/agentbricks/src/databricks_agentbricks/cli/memory.py index cefc0e3eb..18c12d267 100644 --- a/integrations/agentbricks/src/databricks_agentbricks/cli/memory.py +++ b/integrations/agentbricks/src/databricks_agentbricks/cli/memory.py @@ -3,6 +3,7 @@ from __future__ import annotations import pathlib +import shlex import sys from typing import Any @@ -111,19 +112,26 @@ def memory_bind(obj, store: str, source: pathlib.Path) -> None: project = AgentProject.load(source) project.bind_memory_store(store) project.write() + source_arg = "" if source == pathlib.Path(".") else f" --source {shlex.quote(str(source))}" if obj.output == "json": render.emit_json({"memory_store": store, "manifest": str(project.path)}) return render.success( f"Bound memory store '{store}'", - fields={"agent.toml": str(project.path)}, + fields={"agent.toml": str(project.path), "Store readiness": "Not checked (binding only)"}, next_steps=[ ( - f"agentbricks memory stores create --name {store}", - "Create the store now without deploying", + f"agentbricks memory stores create --name {shlex.quote(store)}", + "Create it if missing; skip for an existing store", + ), + ( + f"agentbricks dev --workspace-stores{source_arg}", + "Validate and use existing bound stores locally", + ), + ( + f"agentbricks deploy {source_arg}", + "Create it if missing and grant the app access", ), - ("agentbricks dev", "Re-run to pick up the store locally"), - ("agentbricks deploy ", "Create it if missing and grant the app access"), ], ) @@ -227,14 +235,11 @@ def stores_create(obj, display_name, description) -> None: fields={"Store ID": store_id, "Name": field(data, "name")}, next_steps=[ ( - f"agentbricks memory entries create --store {store_id} --actor-id --path

", - "Add a memory entry for an actor", - ), - (f"agentbricks memory stores get {store_id}", "View this store's details"), - ( - f"agentbricks memory bind {display_name}", - "Bind this store to the agent (wired in on dev/deploy)", + f"agentbricks memory bind {shlex.quote(display_name)}", + "Bind this store to the project", ), + ("agentbricks dev --workspace-stores", "Use the bound store locally"), + (f"agentbricks memory stores get {shlex.quote(store_id)}", "View this store's details"), ], ) diff --git a/integrations/agentbricks/src/databricks_agentbricks/cli/sessions.py b/integrations/agentbricks/src/databricks_agentbricks/cli/sessions.py index 4a12e84af..7000a2779 100644 --- a/integrations/agentbricks/src/databricks_agentbricks/cli/sessions.py +++ b/integrations/agentbricks/src/databricks_agentbricks/cli/sessions.py @@ -4,6 +4,7 @@ import json import pathlib +import shlex import sys from typing import Any, Optional @@ -77,19 +78,26 @@ def sessions_bind(obj, store: str, source: pathlib.Path) -> None: project = AgentProject.load(source) project.bind_session_store(store) project.write() + source_arg = "" if source == pathlib.Path(".") else f" --source {shlex.quote(str(source))}" if obj.output == "json": render.emit_json({"session_store": store, "manifest": str(project.path)}) return render.success( f"Bound session store '{store}'", - fields={"agent.toml": str(project.path)}, + fields={"agent.toml": str(project.path), "Store readiness": "Not checked (binding only)"}, next_steps=[ ( - f"agentbricks sessions stores create --name {store}", - "Create the store now without deploying", + f"agentbricks sessions stores create --name {shlex.quote(store)}", + "Create it if missing; skip for an existing store", + ), + ( + f"agentbricks dev --workspace-stores{source_arg}", + "Validate and use existing bound stores locally", + ), + ( + f"agentbricks deploy {source_arg}", + "Create it if missing and grant the app access", ), - ("agentbricks dev", "Re-run to pick up the store locally"), - ("agentbricks deploy ", "Create it if missing and grant the app access"), ], ) @@ -157,15 +165,12 @@ def stores_create(obj, name, description, metadata) -> None: f"Created session store '{name}'", fields={"Store ID": field(data, "session_store_id")}, next_steps=[ + (f"agentbricks sessions bind {shlex.quote(name)}", "Bind this store to the project"), ( - f"agentbricks sessions create --store {name} --actor-id ", - "Start a session for an actor", - ), - (f"agentbricks sessions stores get {name}", "View this store's details"), - ( - f"agentbricks sessions bind {name}", - "Bind this store to the agent (wired in on dev/deploy)", + "agentbricks dev --workspace-stores", + "Use the bound store locally; the agent creates conversations", ), + (f"agentbricks sessions stores get {shlex.quote(name)}", "View this store's details"), ], ) diff --git a/integrations/agentbricks/src/databricks_agentbricks/templates/agent-langgraph/AGENTKIT_CONTRACT.md b/integrations/agentbricks/src/databricks_agentbricks/templates/agent-langgraph/AGENTKIT_CONTRACT.md index 927587490..60ddcf30d 100644 --- a/integrations/agentbricks/src/databricks_agentbricks/templates/agent-langgraph/AGENTKIT_CONTRACT.md +++ b/integrations/agentbricks/src/databricks_agentbricks/templates/agent-langgraph/AGENTKIT_CONTRACT.md @@ -28,8 +28,9 @@ distribution, supply the real command in `app.yaml`, load configuration before a on the app port. The Agent Bricks CLI and runtime find `agent.toml` from the working directory or `AGENTBRICKS_PROJECT_ROOT`. Keep credentials out of `app.yaml`. -Store binding commands declare intent. `dev` and `deploy` resolve or provision declared stores and -supply runtime config; deploy grants app access. The resolved Memory Store ID reaches the runtime +Store binding commands declare intent. `dev --workspace-stores` validates existing declared stores +and supplies runtime config using the selected profile; plain `dev` keeps state in-process and +memory off. `deploy` provisions missing declared stores and grants app access. The resolved Memory Store ID reaches the runtime as `AGENT_MEMORY_STORE`, rather than being stored as an ID in the manifest. Do not hardcode values that make later bindings ineffective. diff --git a/integrations/agentbricks/src/databricks_agentbricks/templates/agent-openai/AGENTKIT_CONTRACT.md b/integrations/agentbricks/src/databricks_agentbricks/templates/agent-openai/AGENTKIT_CONTRACT.md index 1e3f04637..95ed5f719 100644 --- a/integrations/agentbricks/src/databricks_agentbricks/templates/agent-openai/AGENTKIT_CONTRACT.md +++ b/integrations/agentbricks/src/databricks_agentbricks/templates/agent-openai/AGENTKIT_CONTRACT.md @@ -28,8 +28,9 @@ distribution, supply the real command in `app.yaml`, load configuration before a on the app port. The Agent Bricks CLI and runtime find `agent.toml` from the working directory or `AGENTBRICKS_PROJECT_ROOT`. Keep credentials out of `app.yaml`. -Store binding commands declare intent. `dev` and `deploy` resolve or provision declared stores and -supply runtime config; deploy grants app access. The resolved Memory Store ID reaches the runtime +Store binding commands declare intent. `dev --workspace-stores` validates existing declared stores +and supplies runtime config using the selected profile; plain `dev` keeps state in-process and +memory off. `deploy` provisions missing declared stores and grants app access. The resolved Memory Store ID reaches the runtime as `AGENT_MEMORY_STORE`, rather than being stored as an ID in the manifest. Do not hardcode values that make later bindings ineffective. diff --git a/integrations/agentbricks/src/databricks_agentbricks/templates/agent-openai/AGENTS.md b/integrations/agentbricks/src/databricks_agentbricks/templates/agent-openai/AGENTS.md index 11f122edb..35b22975b 100644 --- a/integrations/agentbricks/src/databricks_agentbricks/templates/agent-openai/AGENTS.md +++ b/integrations/agentbricks/src/databricks_agentbricks/templates/agent-openai/AGENTS.md @@ -62,8 +62,10 @@ Agents SDK does not expose checkpoint continuation. - Invocation state/events: in-memory in `agentbricks dev`; Lakebase when `agentbricks deploy` attaches a Runtime Store. -- Conversation transcript: in-process in `agentbricks dev`; managed Session Store when bound, on `agentbricks deploy`. -- Long-term memory: off in `agentbricks dev`; managed Memory Store when bound, on `agentbricks deploy`. +- Conversation transcript: in-process in plain `agentbricks dev`; managed Session Store when bound, + with `dev --workspace-stores` or `deploy`. +- Long-term memory: off in plain `agentbricks dev`; managed Memory Store when bound, + with `dev --workspace-stores` or `deploy`. - OpenAI HITL `RunState`: process-local even with Session Store; it does not survive worker loss. - Recovery: replay the persisted application input against the same session. diff --git a/integrations/agentbricks/src/databricks_agentbricks/templates/ui/agent-langgraph/CHAT_APP.md b/integrations/agentbricks/src/databricks_agentbricks/templates/ui/agent-langgraph/CHAT_APP.md index a6e6bc87f..4d22533bf 100644 --- a/integrations/agentbricks/src/databricks_agentbricks/templates/ui/agent-langgraph/CHAT_APP.md +++ b/integrations/agentbricks/src/databricks_agentbricks/templates/ui/agent-langgraph/CHAT_APP.md @@ -47,3 +47,16 @@ can be listed because there is no shared session index. Transcript responses include only user, assistant, tool, system, and human-decision message items; checkpoint fragments remain in Session Store but are never returned to the chat UI. + +## Local state and history + +Plain `agentbricks dev` keeps conversation state in-process and memory off. To use existing bound +workspace stores before deployment, run `agentbricks --profile dev --workspace-stores`. +It verifies read access before startup without provisioning resources. Requests can write to those +stores using your credentials; use development stores. Session history and memory then survive +process restart, but local invocation status, background runs, and replay events do not. + +The UI reads the agent's authoritative transcript/checkpoint, including turns submitted through the +API. For request-user-auth projects, state routes use the same identity namespace as invocation +routes, while the public browser session id remains unchanged. App-auth agents retain application +actor partitioning and store-level access. Local development uses one local-developer identity. diff --git a/integrations/agentbricks/src/databricks_agentbricks/templates/ui/agent-langgraph/runtime/ui.py b/integrations/agentbricks/src/databricks_agentbricks/templates/ui/agent-langgraph/runtime/ui.py index 24ab29fb0..fa3b75a3d 100644 --- a/integrations/agentbricks/src/databricks_agentbricks/templates/ui/agent-langgraph/runtime/ui.py +++ b/integrations/agentbricks/src/databricks_agentbricks/templates/ui/agent-langgraph/runtime/ui.py @@ -16,8 +16,12 @@ from pydantic import BaseModel, Field from databricks_agentkit import workspace_client +from databricks_agentkit.runtime.auth import AuthError, RequestAuthContext from databricks_agentkit.runtime.model_services import list_ai_gateway_model_services -from databricks_agentkit.runtime.store import runtime_store_is_persistent_environment +from databricks_agentkit.runtime.store import ( + RUNTIME_STORE_LOCAL_ENV, + runtime_store_is_persistent_environment, +) _UI_ROOT = Path(__file__).resolve().parent.parent / "ui" _INSTANCE_ID = uuid.uuid4().hex[:12] # identifies this process in the UI @@ -75,6 +79,8 @@ def _request_actor(request: Request) -> str: session views here list exactly what the agent reads/writes for the current user. Falls back to ``"agent"`` locally / when unauthenticated. """ + if os.getenv(RUNTIME_STORE_LOCAL_ENV, "").lower() == "true": + return "agent" for header in _USER_HEADERS: if value := request.headers.get(header): return value @@ -89,6 +95,25 @@ def _request_session_id(request: Request) -> str: return str(session_id) +def _state_identity( + request: Request, *, session_id: str | None = None, actor: str | None = None +) -> tuple[str, str, bool]: + """Match the invocation runtime/adapter's identity without changing browser session ids.""" + session_id = session_id or _request_session_id(request) + actor = actor or _request_actor(request) + policy = getattr(request.app, "auth_policy", None) + if policy is None or not policy.requires_user: + return session_id, actor, False + try: + auth = RequestAuthContext.from_headers(request.headers) + except AuthError as exc: + raise HTTPException(status_code=exc.status_code, detail=exc.payload()) from exc + try: + return auth.namespace("session", session_id), auth.namespace("actor", actor), True + finally: + auth.close() + + def _is_deployed() -> bool: app_url = os.getenv("DATABRICKS_APP_URL", "") is_local = app_url.startswith(("http://localhost", "http://127.0.0.1")) @@ -224,7 +249,12 @@ def search_memory_entries(self, actor: str, request: MemorySearchRequest) -> dic }, ) - def ensure_session(self, actor: str, session_id: str) -> dict: + def ensure_session( + self, actor: str, session_id: str, public_session_id: str | None = None + ) -> dict: + metadata = {"client": "agentbricks-demo-ui"} + if public_session_id: + metadata["public_session_id"] = public_session_id try: return self._do( "POST", @@ -232,7 +262,7 @@ def ensure_session(self, actor: str, session_id: str) -> dict: query={"session_id": session_id}, body={ "actor_id": actor, - "metadata": {"client": "agentbricks-demo-ui"}, + "metadata": metadata, }, ) except Exception as exc: @@ -270,11 +300,21 @@ def append_session_items(self, session_id: str, items: list[dict[str, Any]]) -> ) def list_session_items(self, session_id: str) -> dict: - return self._do( - "GET", - f"{_AGENTS_API}/session-stores/{_session_store()}/sessions/{session_id}/items", - query={"order_by": "create_time asc", "page_size": 100}, - ) + items = [] + page_token = None + while True: + query = {"order_by": "create_time asc", "page_size": 100} + if page_token: + query["page_token"] = page_token + page = self._do( + "GET", + f"{_AGENTS_API}/session-stores/{_session_store()}/sessions/{session_id}/items", + query=query, + ) + items.extend(page.get("session_items", [])) + page_token = page.get("next_page_token") + if not page_token: + return {"session_items": items} @lru_cache(maxsize=1) @@ -297,7 +337,8 @@ def _require_memory() -> None: if not _memory_store(): raise HTTPException( status_code=503, - detail="No memory store configured. Run `agentbricks memory bind `.", + detail="No memory store configured. Bind an existing store with `agentbricks memory bind `, " + "then restart with `agentbricks dev --workspace-stores` or deploy.", ) @@ -305,16 +346,23 @@ def _require_session() -> None: if not _session_store(): raise HTTPException( status_code=503, - detail="No session store configured. Run `agentbricks sessions bind `.", + detail="No session store configured. Bind an existing store with `agentbricks sessions bind `, " + "then restart with `agentbricks dev --workspace-stores` or deploy.", ) -async def _checkpoint_history(session_id: str, actor: str) -> dict[str, Any]: +async def _checkpoint_history( + session_id: str, actor: str, *, workspace_client_for=None +) -> dict[str, Any]: from agent.agent import create_agent_graph from databricks_agentkit.langgraph.session_store import thread_config - graph = await create_agent_graph(actor) + # Use the graph's public state view: it applies pending writes and the graph's reducers, + # including messages completed in parallel with a paused approval. This constructs graph/tool + # definitions but never invokes the graph, a model, or a tool. + kwargs = {"workspace_client_for": workspace_client_for} if workspace_client_for else {} + graph = await create_agent_graph(actor, **kwargs) snapshot = await graph.aget_state(thread_config(session_id, actor)) values = snapshot.values if isinstance(snapshot.values, dict) else {} items = [] @@ -334,16 +382,29 @@ async def _checkpoint_history(session_id: str, actor: str) -> dict[str, Any]: return {"session_id": session_id, "session_items": items, "interrupts": interrupts} -def _chat_sessions(result: dict[str, Any]) -> list[dict[str, Any]]: +def _chat_sessions( + result: dict[str, Any], + *, + user_scoped: bool = False, + current_effective_id: str | None = None, + current_public_id: str | None = None, +) -> list[dict[str, Any]]: sessions = [] for session in result.get("sessions", []): if not isinstance(session, dict): continue metadata = session.get("metadata") metadata = metadata if isinstance(metadata, dict) else {} - if metadata.get("public_session_id"): - continue - sessions.append(session) + public_id = metadata.get("public_session_id") + if user_scoped and session.get("session_id") == current_effective_id: + # API-created sessions may predate UI metadata. We can map the current known id. + public_id = current_public_id + if user_scoped: + # An opaque runtime id cannot be sent back as a public id (it would be hashed twice). + if isinstance(public_id, str) and public_id: + sessions.append({**session, "session_id": public_id}) + elif not public_id: + sessions.append(session) return sessions @@ -429,12 +490,8 @@ async def demo_models() -> dict: @app.post("/api/demo/memory/entries", include_in_schema=False) async def create_memory_entry(request: Request, payload: MemoryEntryRequest) -> dict: _require_memory() - return await _managed_call( - _state_client().create_memory_entry, - _request_actor(request), - payload, - _request_session_id(request), - ) + session_id, actor, _ = _state_identity(request) + return await _managed_call(_state_client().create_memory_entry, actor, payload, session_id) @app.get("/api/demo/memory/entries", include_in_schema=False) async def list_memory_entries( @@ -444,31 +501,34 @@ async def list_memory_entries( ) -> dict: # The UI can browse another actor's memories by passing ?actor=; default to the viewer. _require_memory() + _, effective_actor, _ = _state_identity(request, actor=actor) return await _managed_call( - _state_client().list_memory_entries, actor or _request_actor(request), path_prefix + _state_client().list_memory_entries, effective_actor, path_prefix ) @app.post("/api/demo/memory/search", include_in_schema=False) async def search_memory_entries(request: Request, payload: MemorySearchRequest) -> dict: # payload.actor lets the UI search another actor's memories; default to the viewer. _require_memory() - return await _managed_call( - _state_client().search_memory_entries, payload.actor or _request_actor(request), payload - ) + _, actor, _ = _state_identity(request, actor=payload.actor) + return await _managed_call(_state_client().search_memory_entries, actor, payload) @app.post("/api/demo/sessions", include_in_schema=False) async def ensure_session(request: Request) -> dict: _require_session() - return await _managed_call( - _state_client().ensure_session, - _request_actor(request), - _request_session_id(request), + session_id, actor, user_scoped = _state_identity(request) + args = ( + (actor, session_id, _request_session_id(request)) + if user_scoped + else (actor, session_id) ) + result = await _managed_call(_state_client().ensure_session, *args) + return {**result, "session_id": _request_session_id(request)} @app.get("/api/demo/sessions", include_in_schema=False) async def list_sessions(request: Request) -> dict: session_id = _request_session_id(request) - actor = _request_actor(request) + effective_id, actor, user_scoped = _state_identity(request) if not _session_store(): return { "sessions": [ @@ -484,7 +544,12 @@ async def list_sessions(request: Request) -> dict: result = await _managed_call(_state_client().list_sessions, actor) return { **result, - "sessions": _chat_sessions(result), + "sessions": _chat_sessions( + result, + user_scoped=user_scoped, + current_effective_id=effective_id, + current_public_id=session_id, + ), "current_session_id": session_id, "managed": True, } @@ -492,8 +557,9 @@ async def list_sessions(request: Request) -> dict: @app.post("/api/demo/sessions/{session_id}/open", include_in_schema=False) async def open_session(request: Request, session_id: str) -> JSONResponse: _require_session() - session = await _managed_call(_state_client().get_session, session_id) - if session.get("actor_id") != _request_actor(request): + effective_id, actor, _ = _state_identity(request, session_id=session_id) + session = await _managed_call(_state_client().get_session, effective_id) + if session.get("actor_id") != actor: raise HTTPException(status_code=403, detail="Session belongs to another actor.") return JSONResponse( { @@ -506,21 +572,37 @@ async def open_session(request: Request, session_id: str) -> JSONResponse: @app.get("/api/demo/session", include_in_schema=False) async def get_session(request: Request) -> dict: _require_session() - return await _managed_call(_state_client().get_session, _request_session_id(request)) + session_id, _, _ = _state_identity(request) + result = await _managed_call(_state_client().get_session, session_id) + return {**result, "session_id": _request_session_id(request)} @app.post("/api/demo/session/items", include_in_schema=False) async def append_session_items(request: Request, payload: SessionItemsRequest) -> dict: _require_session() - return await _managed_call( - _state_client().append_session_items, - _request_session_id(request), - payload.items, - ) + session_id, _, _ = _state_identity(request) + return await _managed_call(_state_client().append_session_items, session_id, payload.items) @app.get("/api/demo/session/items", include_in_schema=False) async def list_session_items(request: Request) -> dict: - session_id = _request_session_id(request) - if _session_store(): - result = await _managed_call(_state_client().list_session_items, session_id) - return _chat_session_items(result) - return await _checkpoint_history(session_id, _request_actor(request)) + session_id, actor, user_scoped = _state_identity(request) + # Managed items are serialized checkpoints, not a second browser-written transcript. + auth = RequestAuthContext.from_headers(request.headers) if user_scoped else None + try: + kwargs = {"workspace_client_for": auth.client_for} if auth else {} + result = await _checkpoint_history(session_id, actor, **kwargs) + return {**result, "session_id": _request_session_id(request)} + except HTTPException: + raise + except Exception as exc: + code = getattr(exc, "error_code", None) + detail = "Could not load conversation history" + if code: + detail += f" ({code})" + raise HTTPException( + status_code=502, + detail=detail + ". Check access to the bound session store and configured tools " + "using the selected workspace credentials, then retry.", + ) from exc + finally: + if auth: + auth.close() diff --git a/integrations/agentbricks/src/databricks_agentbricks/templates/ui/agent-langgraph/tests/test_demo_ui.py b/integrations/agentbricks/src/databricks_agentbricks/templates/ui/agent-langgraph/tests/test_demo_ui.py index 7689d62ac..ef2921429 100644 --- a/integrations/agentbricks/src/databricks_agentbricks/templates/ui/agent-langgraph/tests/test_demo_ui.py +++ b/integrations/agentbricks/src/databricks_agentbricks/templates/ui/agent-langgraph/tests/test_demo_ui.py @@ -122,7 +122,7 @@ def _client(monkeypatch, *, configured=False, history=False, session_id="routing else: monkeypatch.delenv("AGENT_MEMORY_STORE", raising=False) monkeypatch.delenv("AGENT_SESSION_STORE", raising=False) - if history: + if history or configured: monkeypatch.setattr(ui, "_checkpoint_history", _session_history) # Keep model discovery deterministic and offline (no AI Gateway listing call). monkeypatch.setattr(ui, "_default_model", lambda: "system.ai.claude-sonnet-4-5") @@ -448,13 +448,7 @@ def model_dump(self): class Snapshot: values = {"messages": [Message()]} - tasks = [ - type( - "Task", - (), - {"interrupts": [_FakeInterrupt({"approval": True}, "int-1")]}, - )() - ] + tasks = [type("Task", (), {"interrupts": [_FakeInterrupt({"approval": True}, "int-1")]})()] class FakeAgent: async def aget_state(self, config): @@ -549,7 +543,7 @@ def test_managed_memory_and_session_routes(monkeypatch): assert [ item["data"]["content"] for item in client.get("/api/demo/session/items").json()["session_items"] - ] == ["s1", "saved reply"] + ] == ["s1", "checkpoint reply"] opened = client.post("/api/demo/sessions/s2/open") assert opened.json() == { diff --git a/integrations/agentbricks/src/databricks_agentbricks/templates/ui/agent-langgraph/ui/app.js b/integrations/agentbricks/src/databricks_agentbricks/templates/ui/agent-langgraph/ui/app.js index 6af362c9f..baad3927e 100644 --- a/integrations/agentbricks/src/databricks_agentbricks/templates/ui/agent-langgraph/ui/app.js +++ b/integrations/agentbricks/src/databricks_agentbricks/templates/ui/agent-langgraph/ui/app.js @@ -973,24 +973,6 @@ async function refreshSessionView({ hydrateChat = false } = {}) { await refreshSessions(); } -async function recordSessionItems(items) { - if (!state.config?.session.managed || !items.length) return; - try { - const sessionId = await ensureManagedSession(); - const response = await fetch(demoUrl("/api/demo/session/items"), { - method: "POST", - headers: { "Content-Type": "application/json", ...routingHeaders() }, - body: JSON.stringify({ items }), - }); - const result = await jsonResponse(response); - addEvent("session.items.append", result); - await refreshSessionView(); - } catch (error) { - stateMessage(elements.sessionItems, error instanceof Error ? error.message : String(error), "error"); - addEvent("session.error", { message: String(error) }); - } -} - async function invokeSync(payload) { const response = await fetch("/api/invocations", { method: "POST", @@ -1096,16 +1078,7 @@ async function sendText(text, mode = state.mode) { setBusy(true, mode === "background" ? "Starting background run" : mode === "streaming" ? "Streaming" : "Running"); try { await dispatch({ messages: [{ role: "user", content }] }, mode); - const items = [{ role: "user", content, transport: mode, instance_id: state.instanceId }]; - if (state.lastAssistantText) { - items.push({ - role: "assistant", - content: state.lastAssistantText, - transport: mode, - instance_id: state.instanceId, - }); - } - await recordSessionItems(items); + await refreshSessionView(); return state.lastAssistantText; } catch (error) { finishDraft(); @@ -1130,18 +1103,7 @@ async function resume(decision) { setBusy(true, "Resuming"); try { await dispatch(payload, "streaming"); - const items = [ - { role: "human_decision", content: decision, instance_id: state.instanceId }, - ]; - if (state.lastAssistantText) { - items.push({ - role: "assistant", - content: state.lastAssistantText, - transport: "streaming", - instance_id: state.instanceId, - }); - } - await recordSessionItems(items); + await refreshSessionView(); } catch (error) { appendError(error); } finally { diff --git a/integrations/agentbricks/src/databricks_agentbricks/templates/ui/agent-openai/CHAT_APP.md b/integrations/agentbricks/src/databricks_agentbricks/templates/ui/agent-openai/CHAT_APP.md index ea9b3c6eb..8e0301661 100644 --- a/integrations/agentbricks/src/databricks_agentbricks/templates/ui/agent-openai/CHAT_APP.md +++ b/integrations/agentbricks/src/databricks_agentbricks/templates/ui/agent-openai/CHAT_APP.md @@ -52,3 +52,16 @@ Human-in-the-loop pauses are **in-process only**: a paused run (the Agents SDK ` memory by `agent/agent.py`, not in the session transcript, so — unlike the LangGraph template — it is not durable even with a managed Session Store, and the unmanaged history path never reports pending interrupts. Resume a pause on the same process that created it. + +## Local state and history + +Plain `agentbricks dev` keeps conversation state in-process and memory off. To use existing bound +workspace stores before deployment, run `agentbricks --profile dev --workspace-stores`. +It verifies read access before startup without provisioning resources. Requests can write to those +stores using your credentials; use development stores. Session history and memory then survive +process restart, but local invocation status, background runs, and replay events do not. + +The UI reads the agent's authoritative transcript/checkpoint, including turns submitted through the +API. For request-user-auth projects, state routes use the same identity namespace as invocation +routes, while the public browser session id remains unchanged. App-auth agents retain application +actor partitioning and store-level access. Local development uses one local-developer identity. diff --git a/integrations/agentbricks/src/databricks_agentbricks/templates/ui/agent-openai/runtime/ui.py b/integrations/agentbricks/src/databricks_agentbricks/templates/ui/agent-openai/runtime/ui.py index 5054dc5c6..3ddd7c398 100644 --- a/integrations/agentbricks/src/databricks_agentbricks/templates/ui/agent-openai/runtime/ui.py +++ b/integrations/agentbricks/src/databricks_agentbricks/templates/ui/agent-openai/runtime/ui.py @@ -16,8 +16,12 @@ from pydantic import BaseModel, Field from databricks_agentkit import workspace_client +from databricks_agentkit.runtime.auth import AuthError, RequestAuthContext from databricks_agentkit.runtime.model_services import list_ai_gateway_model_services -from databricks_agentkit.runtime.store import runtime_store_is_persistent_environment +from databricks_agentkit.runtime.store import ( + RUNTIME_STORE_LOCAL_ENV, + runtime_store_is_persistent_environment, +) _UI_ROOT = Path(__file__).resolve().parent.parent / "ui" _INSTANCE_ID = uuid.uuid4().hex[:12] # identifies this process in the UI @@ -75,6 +79,8 @@ def _request_actor(request: Request) -> str: session views here list exactly what the agent reads/writes for the current user. Falls back to ``"agent"`` locally / when unauthenticated. """ + if os.getenv(RUNTIME_STORE_LOCAL_ENV, "").lower() == "true": + return "agent" for header in _USER_HEADERS: if value := request.headers.get(header): return value @@ -89,6 +95,25 @@ def _request_session_id(request: Request) -> str: return str(session_id) +def _state_identity( + request: Request, *, session_id: str | None = None, actor: str | None = None +) -> tuple[str, str, bool]: + """Match the invocation runtime/adapter's identity without changing browser session ids.""" + session_id = session_id or _request_session_id(request) + actor = actor or _request_actor(request) + policy = getattr(request.app, "auth_policy", None) + if policy is None or not policy.requires_user: + return session_id, actor, False + try: + auth = RequestAuthContext.from_headers(request.headers) + except AuthError as exc: + raise HTTPException(status_code=exc.status_code, detail=exc.payload()) from exc + try: + return auth.namespace("session", session_id), auth.namespace("actor", actor), True + finally: + auth.close() + + def _is_deployed() -> bool: app_url = os.getenv("DATABRICKS_APP_URL", "") is_local = app_url.startswith(("http://localhost", "http://127.0.0.1")) @@ -224,7 +249,12 @@ def search_memory_entries(self, actor: str, request: MemorySearchRequest) -> dic }, ) - def ensure_session(self, actor: str, session_id: str) -> dict: + def ensure_session( + self, actor: str, session_id: str, public_session_id: str | None = None + ) -> dict: + metadata = {"client": "agentbricks-demo-ui"} + if public_session_id: + metadata["public_session_id"] = public_session_id try: return self._do( "POST", @@ -232,7 +262,7 @@ def ensure_session(self, actor: str, session_id: str) -> dict: query={"session_id": session_id}, body={ "actor_id": actor, - "metadata": {"client": "agentbricks-demo-ui"}, + "metadata": metadata, }, ) except Exception as exc: @@ -270,11 +300,21 @@ def append_session_items(self, session_id: str, items: list[dict[str, Any]]) -> ) def list_session_items(self, session_id: str) -> dict: - return self._do( - "GET", - f"{_AGENTS_API}/session-stores/{_session_store()}/sessions/{session_id}/items", - query={"order_by": "create_time asc", "page_size": 100}, - ) + items = [] + page_token = None + while True: + query = {"order_by": "create_time asc", "page_size": 100} + if page_token: + query["page_token"] = page_token + page = self._do( + "GET", + f"{_AGENTS_API}/session-stores/{_session_store()}/sessions/{session_id}/items", + query=query, + ) + items.extend(page.get("session_items", [])) + page_token = page.get("next_page_token") + if not page_token: + return {"session_items": items} @lru_cache(maxsize=1) @@ -297,7 +337,8 @@ def _require_memory() -> None: if not _memory_store(): raise HTTPException( status_code=503, - detail="No memory store configured. Run `agentbricks memory bind `.", + detail="No memory store configured. Bind an existing store with `agentbricks memory bind `, " + "then restart with `agentbricks dev --workspace-stores` or deploy.", ) @@ -305,7 +346,8 @@ def _require_session() -> None: if not _session_store(): raise HTTPException( status_code=503, - detail="No session store configured. Run `agentbricks sessions bind `.", + detail="No session store configured. Bind an existing store with `agentbricks sessions bind `, " + "then restart with `agentbricks dev --workspace-stores` or deploy.", ) @@ -327,16 +369,29 @@ async def _local_history(session_id: str) -> dict[str, Any]: return {"session_id": session_id, "session_items": items, "interrupts": []} -def _chat_sessions(result: dict[str, Any]) -> list[dict[str, Any]]: +def _chat_sessions( + result: dict[str, Any], + *, + user_scoped: bool = False, + current_effective_id: str | None = None, + current_public_id: str | None = None, +) -> list[dict[str, Any]]: sessions = [] for session in result.get("sessions", []): if not isinstance(session, dict): continue metadata = session.get("metadata") metadata = metadata if isinstance(metadata, dict) else {} - if metadata.get("public_session_id"): - continue - sessions.append(session) + public_id = metadata.get("public_session_id") + if user_scoped and session.get("session_id") == current_effective_id: + # API-created sessions may predate UI metadata. We can map the current known id. + public_id = current_public_id + if user_scoped: + # An opaque runtime id cannot be sent back as a public id (it would be hashed twice). + if isinstance(public_id, str) and public_id: + sessions.append({**session, "session_id": public_id}) + elif not public_id: + sessions.append(session) return sessions @@ -422,12 +477,8 @@ async def demo_models() -> dict: @app.post("/api/demo/memory/entries", include_in_schema=False) async def create_memory_entry(request: Request, payload: MemoryEntryRequest) -> dict: _require_memory() - return await _managed_call( - _state_client().create_memory_entry, - _request_actor(request), - payload, - _request_session_id(request), - ) + session_id, actor, _ = _state_identity(request) + return await _managed_call(_state_client().create_memory_entry, actor, payload, session_id) @app.get("/api/demo/memory/entries", include_in_schema=False) async def list_memory_entries( @@ -437,31 +488,34 @@ async def list_memory_entries( ) -> dict: # The UI can browse another actor's memories by passing ?actor=; default to the viewer. _require_memory() + _, effective_actor, _ = _state_identity(request, actor=actor) return await _managed_call( - _state_client().list_memory_entries, actor or _request_actor(request), path_prefix + _state_client().list_memory_entries, effective_actor, path_prefix ) @app.post("/api/demo/memory/search", include_in_schema=False) async def search_memory_entries(request: Request, payload: MemorySearchRequest) -> dict: # payload.actor lets the UI search another actor's memories; default to the viewer. _require_memory() - return await _managed_call( - _state_client().search_memory_entries, payload.actor or _request_actor(request), payload - ) + _, actor, _ = _state_identity(request, actor=payload.actor) + return await _managed_call(_state_client().search_memory_entries, actor, payload) @app.post("/api/demo/sessions", include_in_schema=False) async def ensure_session(request: Request) -> dict: _require_session() - return await _managed_call( - _state_client().ensure_session, - _request_actor(request), - _request_session_id(request), + session_id, actor, user_scoped = _state_identity(request) + args = ( + (actor, session_id, _request_session_id(request)) + if user_scoped + else (actor, session_id) ) + result = await _managed_call(_state_client().ensure_session, *args) + return {**result, "session_id": _request_session_id(request)} @app.get("/api/demo/sessions", include_in_schema=False) async def list_sessions(request: Request) -> dict: session_id = _request_session_id(request) - actor = _request_actor(request) + effective_id, actor, user_scoped = _state_identity(request) if not _session_store(): return { "sessions": [ @@ -477,7 +531,12 @@ async def list_sessions(request: Request) -> dict: result = await _managed_call(_state_client().list_sessions, actor) return { **result, - "sessions": _chat_sessions(result), + "sessions": _chat_sessions( + result, + user_scoped=user_scoped, + current_effective_id=effective_id, + current_public_id=session_id, + ), "current_session_id": session_id, "managed": True, } @@ -485,8 +544,9 @@ async def list_sessions(request: Request) -> dict: @app.post("/api/demo/sessions/{session_id}/open", include_in_schema=False) async def open_session(request: Request, session_id: str) -> JSONResponse: _require_session() - session = await _managed_call(_state_client().get_session, session_id) - if session.get("actor_id") != _request_actor(request): + effective_id, actor, _ = _state_identity(request, session_id=session_id) + session = await _managed_call(_state_client().get_session, effective_id) + if session.get("actor_id") != actor: raise HTTPException(status_code=403, detail="Session belongs to another actor.") return JSONResponse( { @@ -499,21 +559,21 @@ async def open_session(request: Request, session_id: str) -> JSONResponse: @app.get("/api/demo/session", include_in_schema=False) async def get_session(request: Request) -> dict: _require_session() - return await _managed_call(_state_client().get_session, _request_session_id(request)) + session_id, _, _ = _state_identity(request) + result = await _managed_call(_state_client().get_session, session_id) + return {**result, "session_id": _request_session_id(request)} @app.post("/api/demo/session/items", include_in_schema=False) async def append_session_items(request: Request, payload: SessionItemsRequest) -> dict: _require_session() - return await _managed_call( - _state_client().append_session_items, - _request_session_id(request), - payload.items, - ) + session_id, _, _ = _state_identity(request) + return await _managed_call(_state_client().append_session_items, session_id, payload.items) @app.get("/api/demo/session/items", include_in_schema=False) async def list_session_items(request: Request) -> dict: - session_id = _request_session_id(request) + session_id, _, _ = _state_identity(request) if _session_store(): result = await _managed_call(_state_client().list_session_items, session_id) return _chat_session_items(result) - return await _local_history(session_id) + result = await _local_history(session_id) + return {**result, "session_id": _request_session_id(request)} diff --git a/integrations/agentbricks/tests/unit_tests/dev_test.py b/integrations/agentbricks/tests/unit_tests/dev_test.py index 5a0697c89..d4ed333f3 100644 --- a/integrations/agentbricks/tests/unit_tests/dev_test.py +++ b/integrations/agentbricks/tests/unit_tests/dev_test.py @@ -591,7 +591,7 @@ def test_dev_notes_local_stores_when_bound(tmp_path: pathlib.Path): out = " ".join(result.output.split()) # collapse rich line-wrapping assert "Memory store 'mem' is bound" in out assert "Session store 'sess' is bound" in out - assert "Run `agentbricks deploy` to use bound store" in out + assert "`agentbricks dev --workspace-stores`" in out def test_dev_notes_bound_tracing_experiment(tmp_path: pathlib.Path): @@ -606,3 +606,126 @@ def test_dev_notes_bound_tracing_experiment(tmp_path: pathlib.Path): out = " ".join(result.output.split()) # collapse rich line-wrapping assert "Tracing experiment '/Shared/agentbricks_traces/mine' is bound" in out assert "Run `agentbricks deploy` to trace to the bound experiment" in out + + +class _WorkspaceStoresCtx(_Ctx): + def __init__(self, client): + super().__init__(profile="chosen-workspace") + self._client = client + + def client(self): + return self._client + + +def _workspace_store_client(): + client = mock.Mock(host="https://selected-workspace.example") + client.list_memory_stores.return_value = { + "managed_memory_stores": [{"display_name": "memory", "name": "memory-stores/memory-id"}] + } + client.get_memory_store.return_value = {"name": "memory-stores/memory-id"} + client.get_session_store.return_value = {"session_store_name": "sessions"} + return client + + +def test_workspace_stores_validates_and_wires_only_current_bindings(tmp_path, monkeypatch): + original = { + "command": ["start-server"], + "env": [{"name": "AGENT_MEMORY_STORE", "value": "old-workspace-id"}], + } + (tmp_path / "app.yaml").write_text(yaml.safe_dump(original)) + _write_agent_manifest(tmp_path, memory="memory", session="sessions") + client = _workspace_store_client() + client.list_memory_stores.side_effect = [ + {"managed_memory_stores": [], "next_page_token": "page2"}, + client.list_memory_stores.return_value, + ] + captured = {} + + def run_local(args, profile, **kwargs): + assert profile == "chosen-workspace" + client.get_memory_store.assert_called_once_with("memory-id") + client.get_session_store.assert_called_once_with("sessions") + captured.update(yaml.safe_load((tmp_path / "app.agentbricksdev.yaml").read_text())) + + monkeypatch.setattr(dev_mod, "_databricks", run_local) + result = CliRunner().invoke( + dev_mod.dev, + ["--source", str(tmp_path), "--workspace-stores"], + obj=_WorkspaceStoresCtx(client), + ) + assert result.exit_code == 0, result.output + assert client.list_memory_stores.call_args_list[-1].kwargs["page_token"] == "page2" + env = {item["name"]: item["value"] for item in captured["env"]} + assert env["AGENT_MEMORY_STORE"] == "memory-id" + assert env["AGENT_SESSION_STORE"] == "sessions" + assert env["DATABRICKS_AGENTBRICKS_RUNTIME_STORE_LOCAL"] == "true" + assert yaml.safe_load((tmp_path / "app.yaml").read_text()) == original + assert not (tmp_path / "app.agentbricksdev.yaml").exists() + client.create_memory_store.assert_not_called() + client.create_session_store.assert_not_called() + output = " ".join(result.output.split()) + assert "chosen-workspace" in output + assert "write permissions are checked when used" in output + assert "pending runs and replay events are lost on restart" in output + + +@pytest.mark.parametrize( + "code,hint", + [ + ("NOT_FOUND", "stores create"), + ("PERMISSION_DENIED", "Ask the store owner"), + ("UNAUTHENTICATED", "auth login"), + ("UNAVAILABLE", "workspace connectivity"), + ], +) +def test_workspace_store_failures_stop_before_starting_servers(tmp_path, code, hint): + (tmp_path / "app.yaml").write_text("command: [start-server]\n") + _write_agent_manifest(tmp_path, session="sessions") + client = _workspace_store_client() + client.get_session_store.side_effect = AgentCliError("secret-sentinel", error_code=code) + with ( + mock.patch.object(dev_mod, "_databricks") as runner, + mock.patch.object(dev_mod, "start_local_tracing_server") as tracing, + ): + result = CliRunner().invoke( + dev_mod.dev, + ["--source", str(tmp_path), "--workspace-stores"], + obj=_WorkspaceStoresCtx(client), + ) + assert result.exit_code != 0 + output = " ".join(result.output.split()) + assert hint in output + assert "secret-sentinel" not in output + runner.assert_not_called() + tracing.assert_not_called() + assert not (tmp_path / "app.agentbricksdev.yaml").exists() + + +def test_workspace_stores_requires_at_least_one_binding(tmp_path): + (tmp_path / "app.yaml").write_text("command: [start-server]\n") + client = _workspace_store_client() + result = CliRunner().invoke( + dev_mod.dev, + ["--source", str(tmp_path), "--workspace-stores"], + obj=_WorkspaceStoresCtx(client), + ) + assert result.exit_code != 0 + assert "No memory or session store is bound" in result.output + assert client.mock_calls == [] + + +def test_workspace_stores_missing_memory_does_not_create_or_fall_back(tmp_path): + (tmp_path / "app.yaml").write_text("command: [start-server]\n") + _write_agent_manifest(tmp_path, memory="memory") + client = _workspace_store_client() + client.list_memory_stores.return_value = {"managed_memory_stores": []} + with mock.patch.object(dev_mod, "_databricks") as runner: + result = CliRunner().invoke( + dev_mod.dev, + ["--source", str(tmp_path), "--workspace-stores"], + obj=_WorkspaceStoresCtx(client), + ) + assert result.exit_code != 0 + assert "Could not verify bound memory store" in result.output + runner.assert_not_called() + client.create_memory_store.assert_not_called() diff --git a/integrations/agentbricks/tests/unit_tests/template_state_history_test.py b/integrations/agentbricks/tests/unit_tests/template_state_history_test.py new file mode 100644 index 000000000..2bc0cc9dc --- /dev/null +++ b/integrations/agentbricks/tests/unit_tests/template_state_history_test.py @@ -0,0 +1,308 @@ +"""State views use the agent's backend and request identity, including API-created turns. + +These tests exercise state HTTP routing and real framework state. Only the managed Session +Store transport is replaced; no model or workspace is required for deterministic regressions. +""" + +import importlib.util +import json +import sys +from pathlib import Path +from types import ModuleType, SimpleNamespace +from typing import Any, cast +from unittest.mock import Mock +from uuid import uuid4 + +import httpx +import pytest + +from databricks_agentkit import DurableAgentServer +from databricks_agentkit.runtime.auth import InvocationAuthPolicy, RequestAuthContext +from databricks_agentkit.runtime.session_store_client import ( + Session, + SessionItem, + SessionStoreClient, +) +from databricks_agentkit.runtime.store import RUNTIME_STORE_LOCAL_ENV, InMemoryRuntimeStore + + +@pytest.fixture(params=["openai", "langgraph"]) +def ui(request, monkeypatch): + framework = request.param + pytest.importorskip("agents" if framework == "openai" else "langgraph") + path = ( + Path(__file__).parents[2] + / f"src/databricks_agentbricks/templates/ui/agent-{framework}/runtime/ui.py" + ) + spec = importlib.util.spec_from_file_location(f"state_ui_{framework}", path) + assert spec and spec.loader + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + monkeypatch.delenv("AGENT_MEMORY_STORE", raising=False) + monkeypatch.delenv("AGENT_SESSION_STORE", raising=False) + monkeypatch.delenv(RUNTIME_STORE_LOCAL_ENV, raising=False) + monkeypatch.setenv("DATABRICKS_APP_NAME", "state-test") + monkeypatch.setenv("DATABRICKS_HOST", "https://state-test.example") + return framework, module + + +class _SessionTransport(SessionStoreClient): + """In-memory remote transport, retaining serialized items across saver instances.""" + + def __init__(self): + self.items = {} + + def get_session(self, *, session_id): + return Session("test", session_id, "actor") + + def append_items(self, session, *, items): + self.items.setdefault(session.session_id, []).extend(json.loads(json.dumps(items))) + + def list_items(self, session, *, order_by=None): + assert order_by == "create_time asc" + return iter( + SessionItem(str(i), item) + for i, item in enumerate(self.items.get(session.session_id, [])) + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("user_auth", [False, True]) +async def test_api_turns_and_ui_read_the_same_ordered_history(ui, monkeypatch, user_auth): + framework, module = ui + session_id = str(uuid4()) + headers = { + "x-forwarded-user": "user-id", + "x-forwarded-email": "alice", + "x-forwarded-access-token": "test-token", + } + if framework == "langgraph": + from langchain_core.runnables import RunnableConfig + from langgraph.graph import END, START, MessagesState, StateGraph + + import databricks_agentkit.langgraph.session_store as stores + + # Exercise serialized managed checkpoints, not browser-written duplicate message items. + transport = _SessionTransport() + saver = stores.DatabricksSessionStoreSaver("test", client=transport) + monkeypatch.setattr(stores, "checkpointer", lambda: saver) + monkeypatch.setenv("AGENT_SESSION_STORE", "test") + monkeypatch.setattr( + module, + "_state_client", + Mock(side_effect=AssertionError("history must read checkpoints")), + ) + builder = StateGraph(cast(Any, MessagesState)) + builder.add_node( + "reply", + lambda state: {"messages": [("assistant", "reply " + state["messages"][-1].content)]}, + ) + builder.add_edge(START, "reply") + builder.add_edge("reply", END) + graph = builder.compile(checkpointer=saver) + _install_graph(monkeypatch, graph) + + async def save(context, text, actor): + await graph.ainvoke( + {"messages": [("user", text)]}, + cast(RunnableConfig, stores.thread_config(context.session_id, actor)), + ) + else: + from databricks_agentkit.openai.sessions import session_store + + async def save(context, text, actor): + await session_store(context.session_id, actor).add_items( + [ + {"role": "user", "content": text}, + {"role": "assistant", "content": "reply " + text}, + ] + ) + + app = DurableAgentServer( + runtime_store=InMemoryRuntimeStore(), + auth_policy=InvocationAuthPolicy(user_required=user_auth), + ) + + # Persist turns through the real framework backend with the runtime's documented mapping. + actor = "alice" + effective_session = session_id + if user_auth: + auth = RequestAuthContext.from_headers(headers) + actor = auth.namespace("actor", actor) + effective_session = auth.namespace("session", session_id) + auth.close() + context = SimpleNamespace(session_id=effective_session) + for text in ("first API turn", "second API turn"): + await save(context, text, actor) + module.install_ui(app) + async with httpx.AsyncClient( + transport=httpx.ASGITransport(app=app), base_url="http://test", headers=headers + ) as client: + result = await client.get("/api/demo/session/items", params={"session_id": session_id}) + assert result.status_code == 200, result.text + assert [item["data"]["content"] for item in result.json()["session_items"]] == [ + "first API turn", + "reply first API turn", + "second API turn", + "reply second API turn", + ] + assert result.json()["session_id"] == session_id + + +def test_managed_history_paginates_in_write_order(ui): + _, module = ui + client = object.__new__(module._ManagedStateClient) + client._do = Mock( + side_effect=[ + {"session_items": [{"item_id": "1"}], "next_page_token": "page-2"}, + {"session_items": [{"item_id": "2"}]}, + ] + ) + assert client.list_session_items("conversation")["session_items"] == [ + {"item_id": "1"}, + {"item_id": "2"}, + ] + assert client._do.call_args_list[-1].kwargs["query"] == { + "order_by": "create_time asc", + "page_size": 100, + "page_token": "page-2", + } + + +@pytest.mark.asyncio +async def test_user_state_routes_require_identity_and_preserve_public_session_ids(ui, monkeypatch): + _, module = ui + app = DurableAgentServer( + runtime_store=InMemoryRuntimeStore(), auth_policy=InvocationAuthPolicy(user_required=True) + ) + module.install_ui(app) + monkeypatch.setenv("AGENT_SESSION_STORE", "test") + monkeypatch.setenv("AGENT_MEMORY_STORE", "test") + headers = { + "x-forwarded-user": "user-id", + "x-forwarded-email": "alice", + "x-forwarded-access-token": "test-token", + } + auth = RequestAuthContext.from_headers(headers) + effective_session = auth.namespace("session", "public-session") + effective_actor = auth.namespace("actor", "alice") + auth.close() + state = Mock() + state.ensure_session.return_value = { + "session_id": effective_session, + "actor_id": effective_actor, + } + state.get_session.return_value = state.ensure_session.return_value + state.list_sessions.return_value = { + "sessions": [ + { + **state.ensure_session.return_value, + "metadata": {"public_session_id": "public-session"}, + } + ] + } + state.list_memory_entries.return_value = {"managed_memory_entries": []} + monkeypatch.setattr(module, "_state_client", lambda: state) + async with httpx.AsyncClient( + transport=httpx.ASGITransport(app=app), + base_url="http://test", + params={"session_id": "public-session"}, + ) as client: + denied = await client.get("/api/demo/session") + assert denied.status_code == 401 + state.get_session.assert_not_called() + client.headers.update(headers) + created = await client.post("/api/demo/sessions") + assert created.json()["session_id"] == "public-session" + state.ensure_session.assert_called_once_with( + effective_actor, effective_session, "public-session" + ) + listed = await client.get("/api/demo/sessions") + assert listed.json()["sessions"][0]["session_id"] == "public-session" + opened = await client.post("/api/demo/sessions/public-session/open") + assert opened.json()["session_id"] == "public-session" + state.get_session.assert_called_once_with(effective_session) + await client.get("/api/demo/memory/entries") + state.list_memory_entries.assert_called_once_with(effective_actor, None) + state.get_session.return_value = {"actor_id": "another-actor"} + assert (await client.post("/api/demo/sessions/public-session/open")).status_code == 403 + + +def _install_graph(monkeypatch, graph): + async def create_graph(actor, **kwargs): + return graph + + monkeypatch.setitem(sys.modules, "agent", ModuleType("agent")) + monkeypatch.setitem( + sys.modules, "agent.agent", SimpleNamespace(create_agent_graph=create_graph) + ) + + +@pytest.mark.asyncio +async def test_langgraph_history_includes_messages_pending_parallel_approval(ui, monkeypatch): + framework, module = ui + if framework != "langgraph": + pytest.skip("LangGraph checkpoint representation") + from langgraph.checkpoint.memory import InMemorySaver + from langgraph.graph import END, START, MessagesState, StateGraph + from langgraph.types import interrupt + + builder = StateGraph(cast(Any, MessagesState)) + builder.add_node("reply", lambda state: {"messages": [("assistant", "parallel reply")]}) + builder.add_node("approval", lambda state: interrupt({"approval": True})) + builder.add_edge(START, "reply") + builder.add_edge(START, "approval") + builder.add_edge("reply", END) + builder.add_edge("approval", END) + graph = builder.compile(checkpointer=InMemorySaver()) + _install_graph(monkeypatch, graph) + assert await module._checkpoint_history("new", "alice") == { + "session_id": "new", + "session_items": [], + "interrupts": [], + } + await graph.ainvoke( + {"messages": [("user", "hello")]}, + {"configurable": {"thread_id": "new", "actor_id": "alice"}}, + ) + result = await module._checkpoint_history("new", "alice") + assert [item["data"]["content"] for item in result["session_items"]] == [ + "hello", + "parallel reply", + ] + assert [item["value"] for item in result["interrupts"]] == [{"approval": True}] + + +@pytest.mark.asyncio +async def test_langgraph_history_errors_are_actionable_without_raw_exception(ui, monkeypatch): + framework, module = ui + if framework != "langgraph": + pytest.skip("LangGraph checkpoint representation") + from databricks.sdk.errors import PermissionDenied + + async def unavailable(*args, **kwargs): + raise PermissionDenied("secret-sentinel", error_code="PERMISSION_DENIED") + + monkeypatch.setattr(module, "_checkpoint_history", unavailable) + app = DurableAgentServer( + runtime_store=InMemoryRuntimeStore(), auth_policy=InvocationAuthPolicy() + ) + module.install_ui(app) + async with httpx.AsyncClient( + transport=httpx.ASGITransport(app=app), base_url="http://test" + ) as client: + response = await client.get("/api/demo/session/items", params={"session_id": "test"}) + assert response.status_code == 502 + assert "PERMISSION_DENIED" in response.text + assert "Check access to the bound session store" in response.text + assert "secret-sentinel" not in response.text + + +def test_user_session_listing_maps_current_api_session_without_ui_metadata(ui): + _, module = ui + assert module._chat_sessions( + {"sessions": [{"session_id": "private-current"}, {"session_id": "private-unmapped"}]}, + user_scoped=True, + current_effective_id="private-current", + current_public_id="public-current", + ) == [{"session_id": "public-current"}] From 748df2dec754317f6ebb47da4de5d0dd2768f740 Mon Sep 17 00:00:00 2001 From: Shivam Mittal Date: Fri, 2 Oct 2026 22:46:34 +0000 Subject: [PATCH 2/3] Correct missing session store recovery command --- integrations/agentbricks/src/databricks_agentbricks/cli/dev.py | 3 ++- integrations/agentbricks/tests/unit_tests/dev_test.py | 2 +- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/integrations/agentbricks/src/databricks_agentbricks/cli/dev.py b/integrations/agentbricks/src/databricks_agentbricks/cli/dev.py index bb41058f8..4dc5c1521 100644 --- a/integrations/agentbricks/src/databricks_agentbricks/cli/dev.py +++ b/integrations/agentbricks/src/databricks_agentbricks/cli/dev.py @@ -314,9 +314,10 @@ def _workspace_store_env( except AgentCliError as exc: code = exc.error_code if code in {"NOT_FOUND", "RESOURCE_DOES_NOT_EXIST"}: + command_group = "sessions" if kind == "session" else "memory" hint = ( f"Check the selected profile and binding. If the store is new, create it with " - f"`agentbricks --profile {kind} stores create --name {name}`, " + f"`agentbricks --profile {command_group} stores create --name {name}`, " "or run `agentbricks deploy` to provision the declared stores. " "A store absent from your listing may instead require an access grant from its owner." ) diff --git a/integrations/agentbricks/tests/unit_tests/dev_test.py b/integrations/agentbricks/tests/unit_tests/dev_test.py index d4ed333f3..347909c9d 100644 --- a/integrations/agentbricks/tests/unit_tests/dev_test.py +++ b/integrations/agentbricks/tests/unit_tests/dev_test.py @@ -672,7 +672,7 @@ def run_local(args, profile, **kwargs): @pytest.mark.parametrize( "code,hint", [ - ("NOT_FOUND", "stores create"), + ("NOT_FOUND", "sessions stores create"), ("PERMISSION_DENIED", "Ask the store owner"), ("UNAUTHENTICATED", "auth login"), ("UNAVAILABLE", "workspace connectivity"), From ea8aa068337546d7005db65a3a4a5c7372e1200a Mon Sep 17 00:00:00 2001 From: Shivam Mittal Date: Fri, 2 Oct 2026 22:47:53 +0000 Subject: [PATCH 3/3] Support scaffold helper imports in state regression tests --- .../tests/unit_tests/template_state_history_test.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/integrations/agentbricks/tests/unit_tests/template_state_history_test.py b/integrations/agentbricks/tests/unit_tests/template_state_history_test.py index 2bc0cc9dc..6d1bd9aab 100644 --- a/integrations/agentbricks/tests/unit_tests/template_state_history_test.py +++ b/integrations/agentbricks/tests/unit_tests/template_state_history_test.py @@ -29,6 +29,13 @@ @pytest.fixture(params=["openai", "langgraph"]) def ui(request, monkeypatch): framework = request.param + # A scaffold may bundle the discovery helper under runtime/ for compatibility with the + # released SDK. Load raw template files with the same helper without creating a scaffold. + from databricks_agentkit.runtime import model_services + + if "runtime" not in sys.modules: + monkeypatch.setitem(sys.modules, "runtime", ModuleType("runtime")) + monkeypatch.setitem(sys.modules, "runtime.model_services", model_services) pytest.importorskip("agents" if framework == "openai" else "langgraph") path = ( Path(__file__).parents[2]