databricks-mason 0.1.3.dev0__tar.gz → 0.1.5.dev0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (61) hide show
  1. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/PKG-INFO +136 -36
  2. databricks_mason-0.1.5.dev0/README.md +313 -0
  3. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/pyproject.toml +15 -5
  4. databricks_mason-0.1.5.dev0/src/databricks_mason/__init__.py +52 -0
  5. databricks_mason-0.1.3.dev0/src/databricks_mason/client.py → databricks_mason-0.1.5.dev0/src/databricks_mason/_api_client.py +63 -41
  6. databricks_mason-0.1.5.dev0/src/databricks_mason/_pagination.py +15 -0
  7. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/src/databricks_mason/agent_project.py +96 -9
  8. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/src/databricks_mason/auth.py +4 -4
  9. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/src/databricks_mason/cli.py +5 -5
  10. databricks_mason-0.1.5.dev0/src/databricks_mason/client.py +28 -0
  11. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/src/databricks_mason/deploy.py +250 -123
  12. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/src/databricks_mason/dev.py +58 -46
  13. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/src/databricks_mason/help.py +5 -2
  14. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/src/databricks_mason/init.py +201 -31
  15. databricks_mason-0.1.5.dev0/src/databricks_mason/lakebase_durability_store.py +78 -0
  16. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/src/databricks_mason/memory.py +35 -15
  17. databricks_mason-0.1.5.dev0/src/databricks_mason/memory_store.py +356 -0
  18. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/src/databricks_mason/render.py +29 -1
  19. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/src/databricks_mason/runtime/__init__.py +8 -2
  20. databricks_mason-0.1.5.dev0/src/databricks_mason/runtime/durability/__init__.py +1 -0
  21. databricks_mason-0.1.5.dev0/src/databricks_mason/runtime/durability/app.py +227 -0
  22. databricks_mason-0.1.5.dev0/src/databricks_mason/runtime/durability/attempt.py +140 -0
  23. databricks_mason-0.1.5.dev0/src/databricks_mason/runtime/durability/recovery.py +112 -0
  24. databricks_mason-0.1.5.dev0/src/databricks_mason/runtime/durability/runtime.py +167 -0
  25. databricks_mason-0.1.5.dev0/src/databricks_mason/runtime/durability/store.py +783 -0
  26. databricks_mason-0.1.5.dev0/src/databricks_mason/runtime/durability/types.py +189 -0
  27. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/src/databricks_mason/runtime/tool_manifest.py +25 -12
  28. databricks_mason-0.1.5.dev0/src/databricks_mason/session_store.py +358 -0
  29. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/src/databricks_mason/sessions.py +22 -7
  30. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/src/databricks_mason/store_access.py +6 -1
  31. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/src/databricks_mason/tracing.py +10 -4
  32. databricks_mason-0.1.3.dev0/README.md +0 -215
  33. databricks_mason-0.1.3.dev0/src/databricks_mason/__init__.py +0 -89
  34. databricks_mason-0.1.3.dev0/src/databricks_mason/runtime/background.py +0 -37
  35. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/.gitignore +0 -0
  36. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/NOTICE +0 -0
  37. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/src/databricks_mason/errors.py +0 -0
  38. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/src/databricks_mason/langgraph/__init__.py +0 -0
  39. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/src/databricks_mason/langgraph/mcp.py +0 -0
  40. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/src/databricks_mason/langgraph/memory.py +0 -0
  41. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/src/databricks_mason/langgraph/session_store.py +0 -0
  42. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/src/databricks_mason/mcp.py +0 -0
  43. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/src/databricks_mason/memory_store_access.py +0 -0
  44. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/src/databricks_mason/models.py +0 -0
  45. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/src/databricks_mason/openai/__init__.py +0 -0
  46. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/src/databricks_mason/openai/mcp.py +0 -0
  47. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/src/databricks_mason/openai/memory.py +0 -0
  48. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/src/databricks_mason/openai/sessions.py +0 -0
  49. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/src/databricks_mason/project_config.py +0 -0
  50. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/src/databricks_mason/py.typed +0 -0
  51. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/src/databricks_mason/runtime/session_store_client.py +0 -0
  52. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/src/databricks_mason/runtime/tracing.py +0 -0
  53. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/src/databricks_mason/runtime/workspace.py +0 -0
  54. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/src/databricks_mason/sandbox.py +0 -0
  55. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/src/databricks_mason/session_store_access.py +0 -0
  56. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/src/databricks_mason/templates/python_tool_langgraph.py +0 -0
  57. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/src/databricks_mason/templates/python_tool_test.py +0 -0
  58. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/src/databricks_mason/templates/sandbox_mcp.py +0 -0
  59. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/src/databricks_mason/templates/sandbox_mcp_langgraph.py +0 -0
  60. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/src/databricks_mason/timefmt.py +0 -0
  61. {databricks_mason-0.1.3.dev0 → databricks_mason-0.1.5.dev0}/src/databricks_mason/tools.py +0 -0
@@ -1,12 +1,12 @@
1
1
  Metadata-Version: 2.5
2
2
  Name: databricks-mason
3
- Version: 0.1.3.dev0
3
+ Version: 0.1.5.dev0
4
4
  Summary: Databricks integration for Mason
5
5
  Author-email: Databricks <agent-feedback@databricks.com>
6
6
  License-File: NOTICE
7
7
  Requires-Python: >=3.10
8
8
  Requires-Dist: click>=8.1
9
- Requires-Dist: databricks-sdk>=0.49
9
+ Requires-Dist: databricks-sdk>=0.94.0
10
10
  Requires-Dist: psycopg[binary]>=3.1
11
11
  Requires-Dist: pyyaml>=6.0
12
12
  Requires-Dist: rich>=13.7
@@ -14,6 +14,7 @@ Requires-Dist: tomli>=2.0
14
14
  Requires-Dist: tomlkit>=0.13
15
15
  Provides-Extra: runtime
16
16
  Requires-Dist: databricks-agents>=1.9.3; extra == 'runtime'
17
+ Requires-Dist: databricks-ai-bridge[memory]>=0.21.0; extra == 'runtime'
17
18
  Requires-Dist: databricks-langchain>=0.17.0; extra == 'runtime'
18
19
  Requires-Dist: fastapi>=0.129.0; extra == 'runtime'
19
20
  Requires-Dist: langchain-mcp-adapters>=0.2.1; extra == 'runtime'
@@ -21,17 +22,18 @@ Requires-Dist: langchain>=1.0.0; extra == 'runtime'
21
22
  Requires-Dist: langgraph>=1.1.0; extra == 'runtime'
22
23
  Requires-Dist: mlflow>=3.10.1; extra == 'runtime'
23
24
  Requires-Dist: opentelemetry-exporter-otlp-proto-grpc>=1.25.0; extra == 'runtime'
24
- Requires-Dist: uuid-utils>=0.10.0; extra == 'runtime'
25
+ Requires-Dist: uvicorn>=0.20.0; extra == 'runtime'
25
26
  Provides-Extra: runtime-openai
26
27
  Requires-Dist: databricks-agents>=1.9.3; extra == 'runtime-openai'
28
+ Requires-Dist: databricks-ai-bridge[memory]>=0.21.0; extra == 'runtime-openai'
27
29
  Requires-Dist: databricks-openai>=0.13.0; extra == 'runtime-openai'
28
30
  Requires-Dist: fastapi>=0.129.0; extra == 'runtime-openai'
29
31
  Requires-Dist: mlflow>=3.10.1; extra == 'runtime-openai'
30
32
  Requires-Dist: openai-agents>=0.7.0; extra == 'runtime-openai'
31
33
  Requires-Dist: opentelemetry-exporter-otlp-proto-grpc>=1.25.0; extra == 'runtime-openai'
32
- Requires-Dist: uuid-utils>=0.10.0; extra == 'runtime-openai'
34
+ Requires-Dist: uvicorn>=0.20.0; extra == 'runtime-openai'
33
35
  Provides-Extra: tracing
34
- Requires-Dist: mlflow[databricks]>=3.9.0; extra == 'tracing'
36
+ Requires-Dist: mlflow[databricks]>=3.10.1; extra == 'tracing'
35
37
  Description-Content-Type: text/markdown
36
38
 
37
39
  # `databricks-mason`
@@ -62,6 +64,12 @@ For tracing commands, install Mason with tracing extras:
62
64
  pip install 'databricks-mason[tracing]'
63
65
  ```
64
66
 
67
+ For the SDK-hosted durable agent application, install the runtime extra:
68
+
69
+ ```sh
70
+ pip install 'databricks-mason[runtime]'
71
+ ```
72
+
65
73
  ## Shell completion
66
74
  Add this to `~/.zshrc`:
67
75
  ```sh
@@ -90,28 +98,117 @@ You can also pass the global `--profile/-p` option before an individual command,
90
98
 
91
99
  ## Python SDK
92
100
 
93
- The same memory and session APIs are available programmatically through
94
- `MasonClient`, which authenticates exactly like the CLI (a `.databrickscfg` profile
95
- or the SDK's default resolution):
101
+ `MasonClient` adds a small resource-oriented layer over the Mason API. Pass it an
102
+ authenticated Databricks `WorkspaceClient`, or omit the argument to use the
103
+ Databricks SDK's default authentication resolution:
96
104
 
97
105
  ```python
106
+ from databricks.sdk import WorkspaceClient
98
107
  from databricks_mason import MasonClient
99
108
 
100
- client = MasonClient(profile="my-workspace") # or MasonClient() for default auth
109
+ mason = MasonClient(WorkspaceClient(profile="my-workspace"))
110
+
111
+ session_store = mason.session_stores.create("support-agent-sessions")
112
+ session = session_store.add(actor_id="customer-123", session_id="case-456")
113
+ session.append_items(
114
+ [
115
+ {"type": "message", "role": "user", "content": "I need help with my cluster."},
116
+ {"type": "message", "role": "assistant", "content": "Let's take a look."},
117
+ ]
118
+ )
119
+
120
+ memory_store = mason.memory_stores.create("coding-agent-memory")
121
+ memory = memory_store.add(
122
+ actor_id="alice",
123
+ path="/preferences/style.md",
124
+ content="The user prefers concise answers.",
125
+ )
126
+ results = memory_store.search(
127
+ actor_id="alice",
128
+ query="response preferences",
129
+ limit=10,
130
+ )
131
+ memory = memory.update(content="The user prefers very concise answers.")
132
+ memory.delete()
133
+ ```
134
+
135
+ The root collections manage stores: `mason.memory_stores.create/get/list` and
136
+ `mason.session_stores.create/get/list`. A returned store owns operations on its
137
+ contents, such as `memory_store.add()`, `memory_store.get("memory-id")`,
138
+ `memory_store.list()`, and `memory_store.search()`, or `session_store.add()`,
139
+ `session_store.get("session-id")`, and `session_store.list()`. Returned memories,
140
+ sessions, and stores own their `update()` and `delete()` operations.
141
+
142
+ All `list()` methods return iterators that automatically consume server pages. List
143
+ `page_size` and search `limit` values must be between 1 and 100. `session.list_items()`
144
+ also auto-pages. `session.fork(...)` creates an independent copy, optionally through
145
+ a specific item. Deleting a session with descendants requires
146
+ `session.delete(force=True)` to cascade the deletion.
147
+
148
+ The resource layer intentionally does not mirror every API method. Its private
149
+ transport will be replaced by the generated `WorkspaceClient.mason` service when that
150
+ is released, without changing this public surface. Deployment, sandbox, tracing, and
151
+ the existing CLI commands remain separate.
101
152
 
102
- store = client.create_memory_store("my-store")
103
- print(store.name, store.display_name) # typed attribute access
153
+ ## Agent application
104
154
 
105
- client.create_memory_entry("my-store", actor_id="alice", path="/notes/1.md", content="hi")
106
- for entry in client.list_memory_entries("my-store", actor_id="alice").entries:
107
- print(entry.path, entry.content)
155
+ `AgentApp` provides Mason's invocation HTTP contract, including foreground, streaming, background,
156
+ polling, and event endpoints. By default its state is process-local. Set `durable_runtime=True` to
157
+ use Lakebase persistence, heartbeats, and crash recovery after deployment:
158
+
159
+ ```python
160
+ from databricks_mason import AgentApp, DurableAgentContext
161
+
162
+ app = AgentApp(durable_runtime=True)
163
+
164
+
165
+ @app.invoke
166
+ async def invoke(input: object, context: DurableAgentContext) -> object:
167
+ return await run_agent(input, session_id=context.session_id)
168
+
169
+
170
+ @app.on_recovery
171
+ async def recover(input: object, context: DurableAgentContext) -> object:
172
+ return await recover_agent(input, session_id=context.session_id)
108
173
  ```
109
174
 
110
- Each method maps to one `/api/agents/v1` operation. Responses come back as typed
111
- models (`MemoryStore`, `Session`, `SessionItemList`, ...) that expose attribute
112
- accessors (`store.name`) while remaining plain dicts underneath — so `store["name"]`,
113
- `json.dumps(store)`, and any new server-side fields keep working. API errors raise
114
- `databricks_mason.AgentCliError`. Deployment, sandbox, and tracing remain CLI-only.
175
+ The Mason server exposes `POST /api/invocations`, `GET /api/invocations/{invocation_id}`, and
176
+ `GET /api/invocations/{invocation_id}/events?after={cursor}`. Databricks Apps bearer-token requests
177
+ must use `/api/` routes
178
+ ([Apps documentation](https://docs.databricks.com/aws/en/dev-tools/databricks-apps/connect-local)).
179
+ The client supplies a UUID `id`, which is also the idempotency key for every invocation mode:
180
+
181
+ - foreground sync returns `200` with the result under `output`;
182
+ - background sync returns `202` with a status URL;
183
+ - foreground streaming returns `200` server-sent events; and
184
+ - background streaming returns `202` with status and event URLs.
185
+
186
+ `input` and `output` may be any JSON value. Transport fields are not passed to the callback. A
187
+ top-level `session_id` is rejected, but a framework template may carry its own stable application
188
+ session inside `input`. Polling uses only the invocation ID and relies on Databricks Apps
189
+ authentication. Without the durable runtime, request state and events exist only in the serving
190
+ process and horizontally scaled clients need sticky routing. With the durable runtime, Mason
191
+ persists the input, attempt status, heartbeats, lifecycle events, application events, and output.
192
+
193
+ Durability is enabled by default for both framework templates. Mason writes the durability binding
194
+ to `agent.toml`, and `mason deploy` then attaches one Lakebase database for runtime durability,
195
+ chosen in this order:
196
+
197
+ 1. Reuse the configured Session Store's Lakebase database.
198
+ 2. Otherwise reuse or provision a dedicated `<app>-durability` Lakebase project.
199
+
200
+ Mason adds only its `databricks_mason_runtime_<app-hash>` schema and tables to the selected database,
201
+ giving each app one owned schema. A replacement worker claims a stale heartbeat and calls the
202
+ `@app.on_recovery` handler. If that handler is omitted, startup warns that automatic crash recovery
203
+ is disabled; register the same function for both decorators when replaying the initial invocation is
204
+ safe. Agent checkpoint restoration and idempotent external side effects remain the developer's
205
+ responsibility.
206
+
207
+ Bare `mason init`, `--framework langgraph`, and `--framework openai` scaffold `AgentApp` with its
208
+ durable runtime enabled. Pass `--no-durable-runtime` for the same Mason HTTP contract with
209
+ process-local state and no Lakebase provisioning. Pass `--server custom` for a minimal FastAPI
210
+ server with one foreground `/invocations` route and no Mason `AgentApp`. Use `--disable-chat-app`
211
+ independently for API-only Mason server output.
115
212
 
116
213
  ## Commands
117
214
 
@@ -119,15 +216,19 @@ accessors (`store.name`) while remaining plain dicts underneath — so `store["n
119
216
  mason [-p <profile>] [-o text|json]
120
217
  login [--profile P]
121
218
  logout
122
- init [--framework openai|langgraph] [--disable-chat-app]
219
+ init [--framework openai|langgraph] [--server mason|custom]
220
+ [--no-durable-runtime] [--disable-chat-app]
123
221
  [--profile P] [--repo URL] [--ref REF] [directory]
124
222
  dev [--source PATH] [--prepare-environment] [--app-port PORT]
125
- [--memory/-m N] [--session/-s N]
126
- [--with-traces C.S] [--no-create-stores]
223
+ [--with-traces C.S]
127
224
  memory
225
+ bind STORE [--source PATH] [--no-create-stores]
226
+ unbind [--source PATH]
128
227
  stores create | list | get | update | delete
129
228
  entries create | get | list | search | update | delete
130
229
  sessions create | list | get | update | delete | fork
230
+ bind STORE [--source PATH] [--no-create-stores]
231
+ unbind [--source PATH]
131
232
  stores create | list | get | update | delete
132
233
  items list | append | pop | clear
133
234
  tracing
@@ -141,9 +242,7 @@ mason [-p <profile>] [-o text|json]
141
242
  add uc-function FUNCTION [--name NAME] [--source PATH]
142
243
  add python NAME [--source PATH]
143
244
  list [--source PATH]
144
- deploy <name> --source PATH [--memory/-m N]
145
- [--session/-s N] [--actor-id ID]
146
- [--with-traces C.S] [--no-create-stores]
245
+ deploy <name> --source PATH [--with-traces C.S] [--instances N]
147
246
  deployments list | get | logs | start | stop | delete
148
247
  ```
149
248
 
@@ -222,28 +321,29 @@ mason init --framework langgraph \
222
321
  --profile <profile> \
223
322
  ./my-agent
224
323
  cd ./my-agent
225
- uv run start-server
324
+ mason dev
226
325
  ```
227
326
 
228
327
  The chat app includes synchronous, SSE streaming, background polling, Session Store, Memory Store,
229
328
  and HITL resume UI. The framework-specific overlay adds `ui/`, `runtime/ui.py`, the UI-enabled
230
329
  `runtime/main.py`, and UI tests.
231
330
 
232
- For the full deployed demo, connect both managed stores:
331
+ For the full deployed demo, bind both managed stores, then deploy:
233
332
 
234
333
  ```sh
235
- mason --profile <profile> deploy mason-agent-demo --source . \
236
- --session mason-demo-sessions \
237
- --memory mason-demo-memory \
238
- --actor-id alice
334
+ mason sessions bind mason-demo-sessions
335
+ mason memory bind mason-demo-memory
336
+ mason --profile <profile> deploy mason-agent-demo --source .
239
337
  ```
240
338
 
241
- (Missing stores are created automatically; pass `--no-create-stores` to require they already exist.)
339
+ (Binding creates a missing store automatically; pass `--no-create-stores` to require it already
340
+ exists. The agent reads the bound stores from `agent.toml` at runtime; `deploy` grants the app's
341
+ service principal access to them.)
242
342
 
243
- The Databricks Apps `__Host-databricks-app-router` cookie is both the sticky routing key and the
244
- application session id. The browser sends it automatically; API clients must reuse it as a cookie.
245
- Request bodies never carry `session_id`. A localhost-only `mason-local-session` cookie provides the
246
- same behavior outside Databricks Apps. TODO: move to `X-Routing-Key` when Apps supports it.
343
+ The chat UI generates a stable application session UUID in browser local storage, places it inside
344
+ the durable invocation's opaque `input`, and creates a fresh invocation UUID per turn. The
345
+ `__Host-databricks-app-router` cookie remains independent: API clients may reuse it for sticky
346
+ replica routing, but it is neither authentication nor the template's application session state.
247
347
 
248
348
  The generated `README.md` documents every request the client makes: config discovery, sync and SSE
249
349
  invocations, background submission and polling, session transcript loading, HITL resume, and memory
@@ -0,0 +1,313 @@
1
+ # `databricks-mason`
2
+
3
+ Mason is an experimental CLI for Databricks custom agent preview APIs and
4
+ deployments. It manages memory, sessions, tracing, and deployments from one
5
+ authenticated command.
6
+
7
+ > The underlying APIs are in preview and may need workspace enablement.
8
+
9
+ ## Installation
10
+
11
+ From PyPI:
12
+
13
+ ```sh
14
+ pip install databricks-mason
15
+ ```
16
+
17
+ From source:
18
+
19
+ ```sh
20
+ pip install 'git+https://github.com/databricks/databricks-ai-bridge.git#subdirectory=integrations/mason'
21
+ ```
22
+
23
+ For tracing commands, install Mason with tracing extras:
24
+
25
+ ```sh
26
+ pip install 'databricks-mason[tracing]'
27
+ ```
28
+
29
+ For the SDK-hosted durable agent application, install the runtime extra:
30
+
31
+ ```sh
32
+ pip install 'databricks-mason[runtime]'
33
+ ```
34
+
35
+ ## Shell completion
36
+ Add this to `~/.zshrc`:
37
+ ```sh
38
+ eval "$(_MASON_COMPLETE=zsh_source mason)"
39
+ ```
40
+
41
+ ## Authentication
42
+
43
+ Mason uses [Databricks authentication](https://docs.databricks.com/aws/en/dev-tools/cli/authentication).
44
+ Ask Mason to authenticate and remember a named profile:
45
+
46
+ ```sh
47
+ mason login --profile <profile>
48
+ mason sessions stores list
49
+ ```
50
+
51
+ `mason login` validates existing credentials first. If credentials are missing or rejected in
52
+ an interactive terminal, Mason runs `databricks auth login --profile <profile>`, revalidates the
53
+ profile, and stores the selection in `~/.mason/config.json`. This browser-based setup requires
54
+ the Databricks CLI. In non-interactive environments, authenticate the profile before running
55
+ Mason. `mason logout` forgets the saved selection without revoking the underlying credentials.
56
+
57
+ If Databricks SDK default authentication is already configured, you can skip `mason login`.
58
+ You can also pass the global `--profile/-p` option before an individual command, for example
59
+ `mason --profile <profile> mcp list`. Use `--output json` for scripting.
60
+
61
+ ## Python SDK
62
+
63
+ `MasonClient` adds a small resource-oriented layer over the Mason API. Pass it an
64
+ authenticated Databricks `WorkspaceClient`, or omit the argument to use the
65
+ Databricks SDK's default authentication resolution:
66
+
67
+ ```python
68
+ from databricks.sdk import WorkspaceClient
69
+ from databricks_mason import MasonClient
70
+
71
+ mason = MasonClient(WorkspaceClient(profile="my-workspace"))
72
+
73
+ session_store = mason.session_stores.create("support-agent-sessions")
74
+ session = session_store.add(actor_id="customer-123", session_id="case-456")
75
+ session.append_items(
76
+ [
77
+ {"type": "message", "role": "user", "content": "I need help with my cluster."},
78
+ {"type": "message", "role": "assistant", "content": "Let's take a look."},
79
+ ]
80
+ )
81
+
82
+ memory_store = mason.memory_stores.create("coding-agent-memory")
83
+ memory = memory_store.add(
84
+ actor_id="alice",
85
+ path="/preferences/style.md",
86
+ content="The user prefers concise answers.",
87
+ )
88
+ results = memory_store.search(
89
+ actor_id="alice",
90
+ query="response preferences",
91
+ limit=10,
92
+ )
93
+ memory = memory.update(content="The user prefers very concise answers.")
94
+ memory.delete()
95
+ ```
96
+
97
+ The root collections manage stores: `mason.memory_stores.create/get/list` and
98
+ `mason.session_stores.create/get/list`. A returned store owns operations on its
99
+ contents, such as `memory_store.add()`, `memory_store.get("memory-id")`,
100
+ `memory_store.list()`, and `memory_store.search()`, or `session_store.add()`,
101
+ `session_store.get("session-id")`, and `session_store.list()`. Returned memories,
102
+ sessions, and stores own their `update()` and `delete()` operations.
103
+
104
+ All `list()` methods return iterators that automatically consume server pages. List
105
+ `page_size` and search `limit` values must be between 1 and 100. `session.list_items()`
106
+ also auto-pages. `session.fork(...)` creates an independent copy, optionally through
107
+ a specific item. Deleting a session with descendants requires
108
+ `session.delete(force=True)` to cascade the deletion.
109
+
110
+ The resource layer intentionally does not mirror every API method. Its private
111
+ transport will be replaced by the generated `WorkspaceClient.mason` service when that
112
+ is released, without changing this public surface. Deployment, sandbox, tracing, and
113
+ the existing CLI commands remain separate.
114
+
115
+ ## Agent application
116
+
117
+ `AgentApp` provides Mason's invocation HTTP contract, including foreground, streaming, background,
118
+ polling, and event endpoints. By default its state is process-local. Set `durable_runtime=True` to
119
+ use Lakebase persistence, heartbeats, and crash recovery after deployment:
120
+
121
+ ```python
122
+ from databricks_mason import AgentApp, DurableAgentContext
123
+
124
+ app = AgentApp(durable_runtime=True)
125
+
126
+
127
+ @app.invoke
128
+ async def invoke(input: object, context: DurableAgentContext) -> object:
129
+ return await run_agent(input, session_id=context.session_id)
130
+
131
+
132
+ @app.on_recovery
133
+ async def recover(input: object, context: DurableAgentContext) -> object:
134
+ return await recover_agent(input, session_id=context.session_id)
135
+ ```
136
+
137
+ The Mason server exposes `POST /api/invocations`, `GET /api/invocations/{invocation_id}`, and
138
+ `GET /api/invocations/{invocation_id}/events?after={cursor}`. Databricks Apps bearer-token requests
139
+ must use `/api/` routes
140
+ ([Apps documentation](https://docs.databricks.com/aws/en/dev-tools/databricks-apps/connect-local)).
141
+ The client supplies a UUID `id`, which is also the idempotency key for every invocation mode:
142
+
143
+ - foreground sync returns `200` with the result under `output`;
144
+ - background sync returns `202` with a status URL;
145
+ - foreground streaming returns `200` server-sent events; and
146
+ - background streaming returns `202` with status and event URLs.
147
+
148
+ `input` and `output` may be any JSON value. Transport fields are not passed to the callback. A
149
+ top-level `session_id` is rejected, but a framework template may carry its own stable application
150
+ session inside `input`. Polling uses only the invocation ID and relies on Databricks Apps
151
+ authentication. Without the durable runtime, request state and events exist only in the serving
152
+ process and horizontally scaled clients need sticky routing. With the durable runtime, Mason
153
+ persists the input, attempt status, heartbeats, lifecycle events, application events, and output.
154
+
155
+ Durability is enabled by default for both framework templates. Mason writes the durability binding
156
+ to `agent.toml`, and `mason deploy` then attaches one Lakebase database for runtime durability,
157
+ chosen in this order:
158
+
159
+ 1. Reuse the configured Session Store's Lakebase database.
160
+ 2. Otherwise reuse or provision a dedicated `<app>-durability` Lakebase project.
161
+
162
+ Mason adds only its `databricks_mason_runtime_<app-hash>` schema and tables to the selected database,
163
+ giving each app one owned schema. A replacement worker claims a stale heartbeat and calls the
164
+ `@app.on_recovery` handler. If that handler is omitted, startup warns that automatic crash recovery
165
+ is disabled; register the same function for both decorators when replaying the initial invocation is
166
+ safe. Agent checkpoint restoration and idempotent external side effects remain the developer's
167
+ responsibility.
168
+
169
+ Bare `mason init`, `--framework langgraph`, and `--framework openai` scaffold `AgentApp` with its
170
+ durable runtime enabled. Pass `--no-durable-runtime` for the same Mason HTTP contract with
171
+ process-local state and no Lakebase provisioning. Pass `--server custom` for a minimal FastAPI
172
+ server with one foreground `/invocations` route and no Mason `AgentApp`. Use `--disable-chat-app`
173
+ independently for API-only Mason server output.
174
+
175
+ ## Commands
176
+
177
+ ```text
178
+ mason [-p <profile>] [-o text|json]
179
+ login [--profile P]
180
+ logout
181
+ init [--framework openai|langgraph] [--server mason|custom]
182
+ [--no-durable-runtime] [--disable-chat-app]
183
+ [--profile P] [--repo URL] [--ref REF] [directory]
184
+ dev [--source PATH] [--prepare-environment] [--app-port PORT]
185
+ [--with-traces C.S]
186
+ memory
187
+ bind STORE [--source PATH] [--no-create-stores]
188
+ unbind [--source PATH]
189
+ stores create | list | get | update | delete
190
+ entries create | get | list | search | update | delete
191
+ sessions create | list | get | update | delete | fork
192
+ bind STORE [--source PATH] [--no-create-stores]
193
+ unbind [--source PATH]
194
+ stores create | list | get | update | delete
195
+ items list | append | pop | clear
196
+ tracing
197
+ setup --catalog C --schema S [--experiment E]
198
+ list | get | instrument
199
+ mcp
200
+ list [--schema CATALOG.SCHEMA]
201
+ tools
202
+ add sandbox --scope SCOPE [--scope SCOPE ...] [--source PATH]
203
+ add mcp SERVICE [--name NAME] [--source PATH]
204
+ add uc-function FUNCTION [--name NAME] [--source PATH]
205
+ add python NAME [--source PATH]
206
+ list [--source PATH]
207
+ deploy <name> --source PATH [--with-traces C.S] [--instances N]
208
+ deployments list | get | logs | start | stop | delete
209
+ ```
210
+
211
+ ## Command help
212
+
213
+ Use the conventional help flag at any command level. Every command's help includes runnable
214
+ examples:
215
+
216
+ ```sh
217
+ mason --help
218
+ mason deploy --help
219
+ mason sessions items append --help
220
+ ```
221
+
222
+ For the shortest path from a blank directory to a running and deployed agent:
223
+
224
+ ```sh
225
+ mason login --profile <profile>
226
+ mason init my-agent
227
+ cd my-agent
228
+ mason dev
229
+ mason deploy my-agent
230
+ ```
231
+
232
+ ## Agent tools
233
+
234
+ `mason init` writes portable tool intent to `agent.toml` and template provenance to
235
+ `.mason/project.toml`. The manifest runtime is currently implemented only by the in-repository
236
+ `agent-langgraph` template; `mason tools add` fails explicitly for other frameworks until they
237
+ provide an adapter at the same runtime seam.
238
+
239
+ Remote tools update only `agent.toml`; they do not generate framework source. The LangGraph runtime
240
+ loads the manifest and materializes its native MCP tools when the agent runs, so a direct manifest
241
+ edit and a CLI edit have the same behavior:
242
+
243
+ ```sh
244
+ mason tools add sandbox --scope table:samples.nyctaxi.trips
245
+ mason tools add mcp system.ai.web_search
246
+ mason tools add uc-function catalog.schema.lookup_ticket
247
+ mason tools add python lookup-ticket
248
+ mason tools remove mcp system.ai.web_search
249
+ mason tools list
250
+ ```
251
+
252
+ For MCP services, the remove command accepts the same service name as the add command. You can also
253
+ remove any binding by the ID shown in `mason tools list`, for example `mason tools remove
254
+ web_search`. Removal updates only `agent.toml`; Python source and test files remain user-owned.
255
+
256
+ Discover the MCP Services available to your user before adding one. By default Mason lists the
257
+ Databricks-managed services in `system.ai`; pass `--schema catalog.schema` for another Unity Catalog
258
+ schema. Text output includes a copyable add command, while `--output json` returns normalized service
259
+ records for scripts:
260
+
261
+ ```sh
262
+ mason mcp list
263
+ mason mcp list --schema main.tools
264
+ ```
265
+
266
+ The Python command additionally creates user-owned `agent/tools/<name>.py` and
267
+ `tests/tools/test_<name>.py` files using the LangGraph-native `@tool` decorator. `mason dev` and
268
+ `mason deploy` preserve `agent.toml`; they do not generate or patch agent source.
269
+
270
+ Sandbox scopes default to read-only access. Repeat `--scope` to allow more than one resource, use
271
+ `volume:` or `workspace:` for those resource types, and use `--permission read_write` only when the
272
+ agent needs writes. Every sandbox call carries this fixed downscope in MCP `_meta`, outside the tool
273
+ arguments controlled by the model.
274
+
275
+ ## Initialize the chat app demo
276
+
277
+ The chat app is a LangGraph-specific init overlay, not a command that mutates an existing project.
278
+ It is included by default for `--framework langgraph`; pass `--disable-chat-app` to scaffold the
279
+ API-only backend instead.
280
+
281
+ ```sh
282
+ mason init --framework langgraph \
283
+ --profile <profile> \
284
+ ./my-agent
285
+ cd ./my-agent
286
+ mason dev
287
+ ```
288
+
289
+ The chat app includes synchronous, SSE streaming, background polling, Session Store, Memory Store,
290
+ and HITL resume UI. The framework-specific overlay adds `ui/`, `runtime/ui.py`, the UI-enabled
291
+ `runtime/main.py`, and UI tests.
292
+
293
+ For the full deployed demo, bind both managed stores, then deploy:
294
+
295
+ ```sh
296
+ mason sessions bind mason-demo-sessions
297
+ mason memory bind mason-demo-memory
298
+ mason --profile <profile> deploy mason-agent-demo --source .
299
+ ```
300
+
301
+ (Binding creates a missing store automatically; pass `--no-create-stores` to require it already
302
+ exists. The agent reads the bound stores from `agent.toml` at runtime; `deploy` grants the app's
303
+ service principal access to them.)
304
+
305
+ The chat UI generates a stable application session UUID in browser local storage, places it inside
306
+ the durable invocation's opaque `input`, and creates a fresh invocation UUID per turn. The
307
+ `__Host-databricks-app-router` cookie remains independent: API clients may reuse it for sticky
308
+ replica routing, but it is neither authentication nor the template's application session state.
309
+
310
+ The generated `README.md` documents every request the client makes: config discovery, sync and SSE
311
+ invocations, background submission and polling, session transcript loading, HITL resume, and memory
312
+ entry operations. Capability colors are automatic from `/api/demo/config`; only the
313
+ sync/streaming/background transport selector is manual.
@@ -1,6 +1,6 @@
1
1
  [project]
2
2
  name = "databricks-mason"
3
- version = "0.1.3.dev0"
3
+ version = "0.1.5.dev0"
4
4
  description = "Databricks integration for Mason"
5
5
  authors = [
6
6
  { name="Databricks", email="agent-feedback@databricks.com" },
@@ -9,7 +9,7 @@ readme = "README.md"
9
9
  requires-python = ">=3.10"
10
10
  dependencies = [
11
11
  "click>=8.1",
12
- "databricks-sdk>=0.49",
12
+ "databricks-sdk>=0.94.0",
13
13
  "psycopg[binary]>=3.1",
14
14
  "PyYAML>=6.0",
15
15
  "rich>=13.7",
@@ -19,7 +19,7 @@ dependencies = [
19
19
 
20
20
  [project.optional-dependencies]
21
21
  tracing = [
22
- "mlflow[databricks]>=3.9.0",
22
+ "mlflow[databricks]>=3.10.1",
23
23
  ]
24
24
  # The agent-side runtime helpers a deployed agent imports. Kept as extras so a plain
25
25
  # `pip install databricks-mason` (the CLI) stays light; each template depends on the extra for its
@@ -27,24 +27,26 @@ tracing = [
27
27
  # OpenAI Agents SDK adapter (databricks_mason.openai). Both carry the shared framework-neutral stack
28
28
  # (databricks_mason.runtime); the framework SDKs differ, so an agent installs only the one it uses.
29
29
  runtime = [
30
+ "databricks-ai-bridge[memory]>=0.21.0",
30
31
  "databricks-langchain>=0.17.0",
31
32
  "langgraph>=1.1.0",
32
33
  "langchain>=1.0.0",
33
34
  "langchain-mcp-adapters>=0.2.1",
34
35
  "fastapi>=0.129.0",
35
36
  "mlflow>=3.10.1",
36
- "uuid-utils>=0.10.0",
37
37
  "opentelemetry-exporter-otlp-proto-grpc>=1.25.0",
38
38
  "databricks-agents>=1.9.3",
39
+ "uvicorn>=0.20.0",
39
40
  ]
40
41
  runtime-openai = [
42
+ "databricks-ai-bridge[memory]>=0.21.0",
41
43
  "openai-agents>=0.7.0",
42
44
  "databricks-openai>=0.13.0",
43
45
  "fastapi>=0.129.0",
44
46
  "mlflow>=3.10.1",
45
- "uuid-utils>=0.10.0",
46
47
  "opentelemetry-exporter-otlp-proto-grpc>=1.25.0",
47
48
  "databricks-agents>=1.9.3",
49
+ "uvicorn>=0.20.0",
48
50
  ]
49
51
 
50
52
  [project.scripts]
@@ -57,7 +59,12 @@ dev = [
57
59
  { include-group = "tests" },
58
60
  ]
59
61
  tests = [
62
+ "databricks-ai-bridge[memory]>=0.21.0",
63
+ "fastapi>=0.129.0",
64
+ "httpx>=0.28.1",
60
65
  "pytest==9.0.2",
66
+ "pytest-asyncio==1.3.0",
67
+ "uvicorn>=0.20.0",
61
68
  ]
62
69
 
63
70
  [build-system]
@@ -75,6 +82,9 @@ include = [
75
82
  [tool.hatch.build.targets.wheel]
76
83
  packages = ["src/databricks_mason"]
77
84
 
85
+ [tool.uv.sources]
86
+ databricks-ai-bridge = { path = "../..", editable = true }
87
+
78
88
  [tool.ruff]
79
89
  include = ["pyproject.toml", "src/**/*.py", "tests/**/*.py"]
80
90
  extend = "../../pyproject.toml"