databricks-mason 0.1.1.dev0__tar.gz → 0.1.3.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 (55) hide show
  1. {databricks_mason-0.1.1.dev0 → databricks_mason-0.1.3.dev0}/PKG-INFO +74 -39
  2. {databricks_mason-0.1.1.dev0 → databricks_mason-0.1.3.dev0}/README.md +64 -38
  3. {databricks_mason-0.1.1.dev0 → databricks_mason-0.1.3.dev0}/pyproject.toml +16 -4
  4. {databricks_mason-0.1.1.dev0 → databricks_mason-0.1.3.dev0}/src/databricks_mason/agent_project.py +73 -3
  5. databricks_mason-0.1.3.dev0/src/databricks_mason/auth.py +139 -0
  6. {databricks_mason-0.1.1.dev0 → databricks_mason-0.1.3.dev0}/src/databricks_mason/cli.py +5 -0
  7. {databricks_mason-0.1.1.dev0 → databricks_mason-0.1.3.dev0}/src/databricks_mason/client.py +101 -20
  8. {databricks_mason-0.1.1.dev0 → databricks_mason-0.1.3.dev0}/src/databricks_mason/deploy.py +121 -41
  9. {databricks_mason-0.1.1.dev0 → databricks_mason-0.1.3.dev0}/src/databricks_mason/dev.py +58 -10
  10. databricks_mason-0.1.3.dev0/src/databricks_mason/errors.py +102 -0
  11. databricks_mason-0.1.3.dev0/src/databricks_mason/help.py +282 -0
  12. {databricks_mason-0.1.1.dev0 → databricks_mason-0.1.3.dev0}/src/databricks_mason/init.py +27 -47
  13. {databricks_mason-0.1.1.dev0 → databricks_mason-0.1.3.dev0}/src/databricks_mason/langgraph/__init__.py +3 -8
  14. databricks_mason-0.1.3.dev0/src/databricks_mason/langgraph/memory.py +67 -0
  15. {databricks_mason-0.1.1.dev0 → databricks_mason-0.1.3.dev0}/src/databricks_mason/langgraph/session_store.py +12 -12
  16. {databricks_mason-0.1.1.dev0 → databricks_mason-0.1.3.dev0}/src/databricks_mason/memory.py +178 -26
  17. databricks_mason-0.1.3.dev0/src/databricks_mason/openai/__init__.py +85 -0
  18. databricks_mason-0.1.3.dev0/src/databricks_mason/openai/mcp.py +94 -0
  19. databricks_mason-0.1.3.dev0/src/databricks_mason/openai/memory.py +67 -0
  20. databricks_mason-0.1.3.dev0/src/databricks_mason/openai/sessions.py +157 -0
  21. {databricks_mason-0.1.1.dev0 → databricks_mason-0.1.3.dev0}/src/databricks_mason/render.py +64 -5
  22. {databricks_mason-0.1.1.dev0 → databricks_mason-0.1.3.dev0}/src/databricks_mason/runtime/__init__.py +2 -2
  23. {databricks_mason-0.1.1.dev0 → databricks_mason-0.1.3.dev0}/src/databricks_mason/runtime/tool_manifest.py +47 -0
  24. {databricks_mason-0.1.1.dev0 → databricks_mason-0.1.3.dev0}/src/databricks_mason/sessions.py +143 -12
  25. {databricks_mason-0.1.1.dev0 → databricks_mason-0.1.3.dev0}/src/databricks_mason/store_access.py +50 -21
  26. {databricks_mason-0.1.1.dev0 → databricks_mason-0.1.3.dev0}/src/databricks_mason/timefmt.py +4 -4
  27. {databricks_mason-0.1.1.dev0 → databricks_mason-0.1.3.dev0}/src/databricks_mason/tools.py +90 -8
  28. {databricks_mason-0.1.1.dev0 → databricks_mason-0.1.3.dev0}/src/databricks_mason/tracing.py +20 -7
  29. databricks_mason-0.1.1.dev0/src/databricks_mason/auth.py +0 -83
  30. databricks_mason-0.1.1.dev0/src/databricks_mason/errors.py +0 -53
  31. databricks_mason-0.1.1.dev0/src/databricks_mason/langgraph/long_running.py +0 -59
  32. databricks_mason-0.1.1.dev0/src/databricks_mason/langgraph/memory.py +0 -63
  33. databricks_mason-0.1.1.dev0/src/databricks_mason/langgraph/recovery.py +0 -247
  34. databricks_mason-0.1.1.dev0/src/databricks_mason/runtime/durability.py +0 -377
  35. databricks_mason-0.1.1.dev0/src/databricks_mason/templates/mcp_runtime_langgraph.py +0 -96
  36. databricks_mason-0.1.1.dev0/src/databricks_mason/templates/tool_manifest_runtime.py +0 -143
  37. {databricks_mason-0.1.1.dev0 → databricks_mason-0.1.3.dev0}/.gitignore +0 -0
  38. {databricks_mason-0.1.1.dev0 → databricks_mason-0.1.3.dev0}/NOTICE +0 -0
  39. {databricks_mason-0.1.1.dev0 → databricks_mason-0.1.3.dev0}/src/databricks_mason/__init__.py +0 -0
  40. {databricks_mason-0.1.1.dev0 → databricks_mason-0.1.3.dev0}/src/databricks_mason/langgraph/mcp.py +0 -0
  41. {databricks_mason-0.1.1.dev0 → databricks_mason-0.1.3.dev0}/src/databricks_mason/mcp.py +0 -0
  42. {databricks_mason-0.1.1.dev0 → databricks_mason-0.1.3.dev0}/src/databricks_mason/memory_store_access.py +0 -0
  43. {databricks_mason-0.1.1.dev0 → databricks_mason-0.1.3.dev0}/src/databricks_mason/models.py +0 -0
  44. {databricks_mason-0.1.1.dev0 → databricks_mason-0.1.3.dev0}/src/databricks_mason/project_config.py +0 -0
  45. {databricks_mason-0.1.1.dev0 → databricks_mason-0.1.3.dev0}/src/databricks_mason/py.typed +0 -0
  46. {databricks_mason-0.1.1.dev0 → databricks_mason-0.1.3.dev0}/src/databricks_mason/runtime/background.py +0 -0
  47. {databricks_mason-0.1.1.dev0 → databricks_mason-0.1.3.dev0}/src/databricks_mason/runtime/session_store_client.py +0 -0
  48. {databricks_mason-0.1.1.dev0 → databricks_mason-0.1.3.dev0}/src/databricks_mason/runtime/tracing.py +0 -0
  49. {databricks_mason-0.1.1.dev0 → databricks_mason-0.1.3.dev0}/src/databricks_mason/runtime/workspace.py +0 -0
  50. {databricks_mason-0.1.1.dev0 → databricks_mason-0.1.3.dev0}/src/databricks_mason/sandbox.py +0 -0
  51. {databricks_mason-0.1.1.dev0 → databricks_mason-0.1.3.dev0}/src/databricks_mason/session_store_access.py +0 -0
  52. {databricks_mason-0.1.1.dev0 → databricks_mason-0.1.3.dev0}/src/databricks_mason/templates/python_tool_langgraph.py +0 -0
  53. {databricks_mason-0.1.1.dev0 → databricks_mason-0.1.3.dev0}/src/databricks_mason/templates/python_tool_test.py +0 -0
  54. {databricks_mason-0.1.1.dev0 → databricks_mason-0.1.3.dev0}/src/databricks_mason/templates/sandbox_mcp.py +0 -0
  55. {databricks_mason-0.1.1.dev0 → databricks_mason-0.1.3.dev0}/src/databricks_mason/templates/sandbox_mcp_langgraph.py +0 -0
@@ -1,12 +1,13 @@
1
1
  Metadata-Version: 2.5
2
2
  Name: databricks-mason
3
- Version: 0.1.1.dev0
3
+ Version: 0.1.3.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
9
  Requires-Dist: databricks-sdk>=0.49
10
+ Requires-Dist: psycopg[binary]>=3.1
10
11
  Requires-Dist: pyyaml>=6.0
11
12
  Requires-Dist: rich>=13.7
12
13
  Requires-Dist: tomli>=2.0
@@ -21,6 +22,14 @@ Requires-Dist: langgraph>=1.1.0; extra == 'runtime'
21
22
  Requires-Dist: mlflow>=3.10.1; extra == 'runtime'
22
23
  Requires-Dist: opentelemetry-exporter-otlp-proto-grpc>=1.25.0; extra == 'runtime'
23
24
  Requires-Dist: uuid-utils>=0.10.0; extra == 'runtime'
25
+ Provides-Extra: runtime-openai
26
+ Requires-Dist: databricks-agents>=1.9.3; extra == 'runtime-openai'
27
+ Requires-Dist: databricks-openai>=0.13.0; extra == 'runtime-openai'
28
+ Requires-Dist: fastapi>=0.129.0; extra == 'runtime-openai'
29
+ Requires-Dist: mlflow>=3.10.1; extra == 'runtime-openai'
30
+ Requires-Dist: openai-agents>=0.7.0; extra == 'runtime-openai'
31
+ Requires-Dist: opentelemetry-exporter-otlp-proto-grpc>=1.25.0; extra == 'runtime-openai'
32
+ Requires-Dist: uuid-utils>=0.10.0; extra == 'runtime-openai'
24
33
  Provides-Extra: tracing
25
34
  Requires-Dist: mlflow[databricks]>=3.9.0; extra == 'tracing'
26
35
  Description-Content-Type: text/markdown
@@ -53,23 +62,31 @@ For tracing commands, install Mason with tracing extras:
53
62
  pip install 'databricks-mason[tracing]'
54
63
  ```
55
64
 
65
+ ## Shell completion
66
+ Add this to `~/.zshrc`:
67
+ ```sh
68
+ eval "$(_MASON_COMPLETE=zsh_source mason)"
69
+ ```
70
+
56
71
  ## Authentication
57
72
 
58
73
  Mason uses [Databricks authentication](https://docs.databricks.com/aws/en/dev-tools/cli/authentication).
59
- If you do not already have credentials, authenticate a named profile first. You can
60
- then ask Mason to validate and remember that profile:
74
+ Ask Mason to authenticate and remember a named profile:
61
75
 
62
76
  ```sh
63
- databricks auth login --profile <profile>
64
77
  mason login --profile <profile>
65
78
  mason sessions stores list
66
79
  ```
67
80
 
68
- `mason login` does not create credentials; it stores the selected profile in
69
- `~/.mason/config.json`. `mason logout` forgets that selection without revoking the
70
- underlying credentials. If Databricks SDK default authentication is already configured,
71
- you can skip `mason login`. You can also pass `--profile/-p` for an individual command.
72
- Use `--output json` for scripting.
81
+ `mason login` validates existing credentials first. If credentials are missing or rejected in
82
+ an interactive terminal, Mason runs `databricks auth login --profile <profile>`, revalidates the
83
+ profile, and stores the selection in `~/.mason/config.json`. This browser-based setup requires
84
+ the Databricks CLI. In non-interactive environments, authenticate the profile before running
85
+ Mason. `mason logout` forgets the saved selection without revoking the underlying credentials.
86
+
87
+ If Databricks SDK default authentication is already configured, you can skip `mason login`.
88
+ You can also pass the global `--profile/-p` option before an individual command, for example
89
+ `mason --profile <profile> mcp list`. Use `--output json` for scripting.
73
90
 
74
91
  ## Python SDK
75
92
 
@@ -102,8 +119,11 @@ accessors (`store.name`) while remaining plain dicts underneath — so `store["n
102
119
  mason [-p <profile>] [-o text|json]
103
120
  login [--profile P]
104
121
  logout
105
- init [--framework openai|langgraph] [--enable-chat-app]
122
+ init [--framework openai|langgraph] [--disable-chat-app]
106
123
  [--profile P] [--repo URL] [--ref REF] [directory]
124
+ dev [--source PATH] [--prepare-environment] [--app-port PORT]
125
+ [--memory/-m N] [--session/-s N]
126
+ [--with-traces C.S] [--no-create-stores]
107
127
  memory
108
128
  stores create | list | get | update | delete
109
129
  entries create | get | list | search | update | delete
@@ -115,19 +135,39 @@ mason [-p <profile>] [-o text|json]
115
135
  list | get | instrument
116
136
  mcp
117
137
  list [--schema CATALOG.SCHEMA]
118
- init [--framework openai|langgraph] [--profile P] [DIRECTORY]
119
138
  tools
120
139
  add sandbox --scope SCOPE [--scope SCOPE ...] [--source PATH]
121
140
  add mcp SERVICE [--name NAME] [--source PATH]
122
141
  add uc-function FUNCTION [--name NAME] [--source PATH]
123
142
  add python NAME [--source PATH]
124
143
  list [--source PATH]
125
- deploy <name> --source PATH [--with-memory-store N]
126
- [--with-session-store N] [--actor-id ID]
127
- [--with-traces C.S] [--create-stores]
144
+ deploy <name> --source PATH [--memory/-m N]
145
+ [--session/-s N] [--actor-id ID]
146
+ [--with-traces C.S] [--no-create-stores]
128
147
  deployments list | get | logs | start | stop | delete
129
148
  ```
130
149
 
150
+ ## Command help
151
+
152
+ Use the conventional help flag at any command level. Every command's help includes runnable
153
+ examples:
154
+
155
+ ```sh
156
+ mason --help
157
+ mason deploy --help
158
+ mason sessions items append --help
159
+ ```
160
+
161
+ For the shortest path from a blank directory to a running and deployed agent:
162
+
163
+ ```sh
164
+ mason login --profile <profile>
165
+ mason init my-agent
166
+ cd my-agent
167
+ mason dev
168
+ mason deploy my-agent
169
+ ```
170
+
131
171
  ## Agent tools
132
172
 
133
173
  `mason init` writes portable tool intent to `agent.toml` and template provenance to
@@ -144,9 +184,14 @@ mason tools add sandbox --scope table:samples.nyctaxi.trips
144
184
  mason tools add mcp system.ai.web_search
145
185
  mason tools add uc-function catalog.schema.lookup_ticket
146
186
  mason tools add python lookup-ticket
187
+ mason tools remove mcp system.ai.web_search
147
188
  mason tools list
148
189
  ```
149
190
 
191
+ For MCP services, the remove command accepts the same service name as the add command. You can also
192
+ remove any binding by the ID shown in `mason tools list`, for example `mason tools remove
193
+ web_search`. Removal updates only `agent.toml`; Python source and test files remain user-owned.
194
+
150
195
  Discover the MCP Services available to your user before adding one. By default Mason lists the
151
196
  Databricks-managed services in `system.ai`; pass `--schema catalog.schema` for another Unity Catalog
152
197
  schema. Text output includes a copyable add command, while `--output json` returns normalized service
@@ -168,49 +213,39 @@ arguments controlled by the model.
168
213
 
169
214
  ## Initialize the chat app demo
170
215
 
171
- The chat app is a LangGraph-specific init overlay, not a command that mutates an existing project:
216
+ The chat app is a LangGraph-specific init overlay, not a command that mutates an existing project.
217
+ It is included by default for `--framework langgraph`; pass `--disable-chat-app` to scaffold the
218
+ API-only backend instead.
172
219
 
173
220
  ```sh
174
- mason init --framework langgraph --enable-chat-app \
221
+ mason init --framework langgraph \
175
222
  --profile <profile> \
176
223
  ./my-agent
177
224
  cd ./my-agent
178
225
  uv run start-server
179
226
  ```
180
227
 
181
- `--enable-chat-app` always includes synchronous, SSE streaming, background polling, Session Store,
182
- Memory Store, HITL resume, Start App, Stop App, heartbeat, and recovery UI. There are no separate
183
- stop/crash flags. The base agent owns `agent/mason/durability.py`, `agent/mason/recovery.py`, and
184
- `agent/mason/long_running.py`; the framework-specific overlay only adds `ui/`, `runtime/ui.py`, the
185
- UI-enabled `runtime/main.py`, and UI tests.
228
+ The chat app includes synchronous, SSE streaming, background polling, Session Store, Memory Store,
229
+ and HITL resume UI. The framework-specific overlay adds `ui/`, `runtime/ui.py`, the UI-enabled
230
+ `runtime/main.py`, and UI tests.
186
231
 
187
232
  For the full deployed demo, connect both managed stores:
188
233
 
189
234
  ```sh
190
235
  mason --profile <profile> deploy mason-agent-demo --source . \
191
- --with-session-store mason-demo-sessions \
192
- --with-memory-store mason-demo-memory \
193
- --actor-id alice \
194
- --create-stores
236
+ --session mason-demo-sessions \
237
+ --memory mason-demo-memory \
238
+ --actor-id alice
195
239
  ```
196
240
 
241
+ (Missing stores are created automatically; pass `--no-create-stores` to require they already exist.)
242
+
197
243
  The Databricks Apps `__Host-databricks-app-router` cookie is both the sticky routing key and the
198
244
  application session id. The browser sends it automatically; API clients must reuse it as a cookie.
199
245
  Request bodies never carry `session_id`. A localhost-only `mason-local-session` cookie provides the
200
246
  same behavior outside Databricks Apps. TODO: move to `X-Routing-Key` when Apps supports it.
201
247
 
202
248
  The generated `README.md` documents every request the client makes: config discovery, sync and SSE
203
- invocations, background submission and polling, session transcript loading, HITL resume, memory
204
- entry operations, and stop/start recovery. Capability colors are automatic from `/api/demo/config`;
205
- only the sync/streaming/background transport selector is manual.
206
-
207
- Start App runs `tool_step_1` through `tool_step_4` in a checkpointed sequence. Each completed output
208
- is committed before the next node. Stop App schedules `os._exit(86)`; Databricks Apps restarts the
209
- process, the browser waits for a new instance and a stale heartbeat, and then starts a new attempt
210
- with the same routing cookie. Completed tools are restored and skipped; an interrupted tool whose
211
- output was not committed can run again.
212
-
213
- The ownership log is intentionally demo-grade: Session Store records append-only attempts and
214
- heartbeats, but the claim is last-writer-wins rather than atomic (`atomic_claim: false`). Production
215
- durability also needs transactional ownership, server-side stale scanning, idempotent side effects,
216
- and durable event replay.
249
+ invocations, background submission and polling, session transcript loading, HITL resume, and memory
250
+ entry operations. Capability colors are automatic from `/api/demo/config`; only the
251
+ sync/streaming/background transport selector is manual.
@@ -26,23 +26,31 @@ For tracing commands, install Mason with tracing extras:
26
26
  pip install 'databricks-mason[tracing]'
27
27
  ```
28
28
 
29
+ ## Shell completion
30
+ Add this to `~/.zshrc`:
31
+ ```sh
32
+ eval "$(_MASON_COMPLETE=zsh_source mason)"
33
+ ```
34
+
29
35
  ## Authentication
30
36
 
31
37
  Mason uses [Databricks authentication](https://docs.databricks.com/aws/en/dev-tools/cli/authentication).
32
- If you do not already have credentials, authenticate a named profile first. You can
33
- then ask Mason to validate and remember that profile:
38
+ Ask Mason to authenticate and remember a named profile:
34
39
 
35
40
  ```sh
36
- databricks auth login --profile <profile>
37
41
  mason login --profile <profile>
38
42
  mason sessions stores list
39
43
  ```
40
44
 
41
- `mason login` does not create credentials; it stores the selected profile in
42
- `~/.mason/config.json`. `mason logout` forgets that selection without revoking the
43
- underlying credentials. If Databricks SDK default authentication is already configured,
44
- you can skip `mason login`. You can also pass `--profile/-p` for an individual command.
45
- Use `--output json` for scripting.
45
+ `mason login` validates existing credentials first. If credentials are missing or rejected in
46
+ an interactive terminal, Mason runs `databricks auth login --profile <profile>`, revalidates the
47
+ profile, and stores the selection in `~/.mason/config.json`. This browser-based setup requires
48
+ the Databricks CLI. In non-interactive environments, authenticate the profile before running
49
+ Mason. `mason logout` forgets the saved selection without revoking the underlying credentials.
50
+
51
+ If Databricks SDK default authentication is already configured, you can skip `mason login`.
52
+ You can also pass the global `--profile/-p` option before an individual command, for example
53
+ `mason --profile <profile> mcp list`. Use `--output json` for scripting.
46
54
 
47
55
  ## Python SDK
48
56
 
@@ -75,8 +83,11 @@ accessors (`store.name`) while remaining plain dicts underneath — so `store["n
75
83
  mason [-p <profile>] [-o text|json]
76
84
  login [--profile P]
77
85
  logout
78
- init [--framework openai|langgraph] [--enable-chat-app]
86
+ init [--framework openai|langgraph] [--disable-chat-app]
79
87
  [--profile P] [--repo URL] [--ref REF] [directory]
88
+ dev [--source PATH] [--prepare-environment] [--app-port PORT]
89
+ [--memory/-m N] [--session/-s N]
90
+ [--with-traces C.S] [--no-create-stores]
80
91
  memory
81
92
  stores create | list | get | update | delete
82
93
  entries create | get | list | search | update | delete
@@ -88,19 +99,39 @@ mason [-p <profile>] [-o text|json]
88
99
  list | get | instrument
89
100
  mcp
90
101
  list [--schema CATALOG.SCHEMA]
91
- init [--framework openai|langgraph] [--profile P] [DIRECTORY]
92
102
  tools
93
103
  add sandbox --scope SCOPE [--scope SCOPE ...] [--source PATH]
94
104
  add mcp SERVICE [--name NAME] [--source PATH]
95
105
  add uc-function FUNCTION [--name NAME] [--source PATH]
96
106
  add python NAME [--source PATH]
97
107
  list [--source PATH]
98
- deploy <name> --source PATH [--with-memory-store N]
99
- [--with-session-store N] [--actor-id ID]
100
- [--with-traces C.S] [--create-stores]
108
+ deploy <name> --source PATH [--memory/-m N]
109
+ [--session/-s N] [--actor-id ID]
110
+ [--with-traces C.S] [--no-create-stores]
101
111
  deployments list | get | logs | start | stop | delete
102
112
  ```
103
113
 
114
+ ## Command help
115
+
116
+ Use the conventional help flag at any command level. Every command's help includes runnable
117
+ examples:
118
+
119
+ ```sh
120
+ mason --help
121
+ mason deploy --help
122
+ mason sessions items append --help
123
+ ```
124
+
125
+ For the shortest path from a blank directory to a running and deployed agent:
126
+
127
+ ```sh
128
+ mason login --profile <profile>
129
+ mason init my-agent
130
+ cd my-agent
131
+ mason dev
132
+ mason deploy my-agent
133
+ ```
134
+
104
135
  ## Agent tools
105
136
 
106
137
  `mason init` writes portable tool intent to `agent.toml` and template provenance to
@@ -117,9 +148,14 @@ mason tools add sandbox --scope table:samples.nyctaxi.trips
117
148
  mason tools add mcp system.ai.web_search
118
149
  mason tools add uc-function catalog.schema.lookup_ticket
119
150
  mason tools add python lookup-ticket
151
+ mason tools remove mcp system.ai.web_search
120
152
  mason tools list
121
153
  ```
122
154
 
155
+ For MCP services, the remove command accepts the same service name as the add command. You can also
156
+ remove any binding by the ID shown in `mason tools list`, for example `mason tools remove
157
+ web_search`. Removal updates only `agent.toml`; Python source and test files remain user-owned.
158
+
123
159
  Discover the MCP Services available to your user before adding one. By default Mason lists the
124
160
  Databricks-managed services in `system.ai`; pass `--schema catalog.schema` for another Unity Catalog
125
161
  schema. Text output includes a copyable add command, while `--output json` returns normalized service
@@ -141,49 +177,39 @@ arguments controlled by the model.
141
177
 
142
178
  ## Initialize the chat app demo
143
179
 
144
- The chat app is a LangGraph-specific init overlay, not a command that mutates an existing project:
180
+ The chat app is a LangGraph-specific init overlay, not a command that mutates an existing project.
181
+ It is included by default for `--framework langgraph`; pass `--disable-chat-app` to scaffold the
182
+ API-only backend instead.
145
183
 
146
184
  ```sh
147
- mason init --framework langgraph --enable-chat-app \
185
+ mason init --framework langgraph \
148
186
  --profile <profile> \
149
187
  ./my-agent
150
188
  cd ./my-agent
151
189
  uv run start-server
152
190
  ```
153
191
 
154
- `--enable-chat-app` always includes synchronous, SSE streaming, background polling, Session Store,
155
- Memory Store, HITL resume, Start App, Stop App, heartbeat, and recovery UI. There are no separate
156
- stop/crash flags. The base agent owns `agent/mason/durability.py`, `agent/mason/recovery.py`, and
157
- `agent/mason/long_running.py`; the framework-specific overlay only adds `ui/`, `runtime/ui.py`, the
158
- UI-enabled `runtime/main.py`, and UI tests.
192
+ The chat app includes synchronous, SSE streaming, background polling, Session Store, Memory Store,
193
+ and HITL resume UI. The framework-specific overlay adds `ui/`, `runtime/ui.py`, the UI-enabled
194
+ `runtime/main.py`, and UI tests.
159
195
 
160
196
  For the full deployed demo, connect both managed stores:
161
197
 
162
198
  ```sh
163
199
  mason --profile <profile> deploy mason-agent-demo --source . \
164
- --with-session-store mason-demo-sessions \
165
- --with-memory-store mason-demo-memory \
166
- --actor-id alice \
167
- --create-stores
200
+ --session mason-demo-sessions \
201
+ --memory mason-demo-memory \
202
+ --actor-id alice
168
203
  ```
169
204
 
205
+ (Missing stores are created automatically; pass `--no-create-stores` to require they already exist.)
206
+
170
207
  The Databricks Apps `__Host-databricks-app-router` cookie is both the sticky routing key and the
171
208
  application session id. The browser sends it automatically; API clients must reuse it as a cookie.
172
209
  Request bodies never carry `session_id`. A localhost-only `mason-local-session` cookie provides the
173
210
  same behavior outside Databricks Apps. TODO: move to `X-Routing-Key` when Apps supports it.
174
211
 
175
212
  The generated `README.md` documents every request the client makes: config discovery, sync and SSE
176
- invocations, background submission and polling, session transcript loading, HITL resume, memory
177
- entry operations, and stop/start recovery. Capability colors are automatic from `/api/demo/config`;
178
- only the sync/streaming/background transport selector is manual.
179
-
180
- Start App runs `tool_step_1` through `tool_step_4` in a checkpointed sequence. Each completed output
181
- is committed before the next node. Stop App schedules `os._exit(86)`; Databricks Apps restarts the
182
- process, the browser waits for a new instance and a stale heartbeat, and then starts a new attempt
183
- with the same routing cookie. Completed tools are restored and skipped; an interrupted tool whose
184
- output was not committed can run again.
185
-
186
- The ownership log is intentionally demo-grade: Session Store records append-only attempts and
187
- heartbeats, but the claim is last-writer-wins rather than atomic (`atomic_claim: false`). Production
188
- durability also needs transactional ownership, server-side stale scanning, idempotent side effects,
189
- and durable event replay.
213
+ invocations, background submission and polling, session transcript loading, HITL resume, and memory
214
+ entry operations. Capability colors are automatic from `/api/demo/config`; only the
215
+ sync/streaming/background transport selector is manual.
@@ -1,6 +1,6 @@
1
1
  [project]
2
2
  name = "databricks-mason"
3
- version = "0.1.1.dev0"
3
+ version = "0.1.3.dev0"
4
4
  description = "Databricks integration for Mason"
5
5
  authors = [
6
6
  { name="Databricks", email="agent-feedback@databricks.com" },
@@ -10,6 +10,7 @@ requires-python = ">=3.10"
10
10
  dependencies = [
11
11
  "click>=8.1",
12
12
  "databricks-sdk>=0.49",
13
+ "psycopg[binary]>=3.1",
13
14
  "PyYAML>=6.0",
14
15
  "rich>=13.7",
15
16
  "tomli>=2.0",
@@ -20,9 +21,11 @@ dependencies = [
20
21
  tracing = [
21
22
  "mlflow[databricks]>=3.9.0",
22
23
  ]
23
- # The agent-side runtime helpers (databricks_mason.runtime) that a deployed agent imports. Kept as
24
- # an extra so a plain `pip install databricks-mason` (the CLI) stays light; the template depends on
25
- # `databricks-mason[runtime]`.
24
+ # The agent-side runtime helpers a deployed agent imports. Kept as extras so a plain
25
+ # `pip install databricks-mason` (the CLI) stays light; each template depends on the extra for its
26
+ # framework. `runtime` = the LangGraph adapter (databricks_mason.langgraph); `runtime-openai` = the
27
+ # OpenAI Agents SDK adapter (databricks_mason.openai). Both carry the shared framework-neutral stack
28
+ # (databricks_mason.runtime); the framework SDKs differ, so an agent installs only the one it uses.
26
29
  runtime = [
27
30
  "databricks-langchain>=0.17.0",
28
31
  "langgraph>=1.1.0",
@@ -34,6 +37,15 @@ runtime = [
34
37
  "opentelemetry-exporter-otlp-proto-grpc>=1.25.0",
35
38
  "databricks-agents>=1.9.3",
36
39
  ]
40
+ runtime-openai = [
41
+ "openai-agents>=0.7.0",
42
+ "databricks-openai>=0.13.0",
43
+ "fastapi>=0.129.0",
44
+ "mlflow>=3.10.1",
45
+ "uuid-utils>=0.10.0",
46
+ "opentelemetry-exporter-otlp-proto-grpc>=1.25.0",
47
+ "databricks-agents>=1.9.3",
48
+ ]
37
49
 
38
50
  [project.scripts]
39
51
  mason = "databricks_mason.cli:main"
@@ -15,6 +15,7 @@ from tomlkit import TOMLDocument
15
15
  from tomlkit.exceptions import ParseError
16
16
 
17
17
  from databricks_mason.errors import AgentCliError
18
+ from databricks_mason.runtime.tool_manifest import MEMORY_STORE_TABLE, SESSION_STORE_TABLE
18
19
 
19
20
  _SCHEMA_VERSION = 1
20
21
  _SUPPORTED_FRAMEWORKS = {"langgraph", "openai"}
@@ -191,6 +192,15 @@ def _required_string(value: object, description: str) -> str:
191
192
  return value
192
193
 
193
194
 
195
+ def _store_name_from_manifest(value: object, table: str) -> str | None:
196
+ """Read the ``name`` from a ``[memory_store]`` / ``[session_store]`` table, or None if absent."""
197
+ if value is None:
198
+ return None
199
+ if not isinstance(value, Mapping):
200
+ raise AgentCliError(f"agent.toml [{table}] must be a table.")
201
+ return _required_string(cast(Mapping[str, Any], value).get("name"), f"[{table}] name")
202
+
203
+
194
204
  def _scope_from_manifest(value: object) -> Scope:
195
205
  if not isinstance(value, Mapping):
196
206
  raise AgentCliError("Sandbox downscope entries must be TOML tables.")
@@ -275,12 +285,17 @@ class AgentProject:
275
285
  document: TOMLDocument,
276
286
  framework: str,
277
287
  tools: list[ToolSpec],
288
+ memory_store: str | None = None,
289
+ session_store: str | None = None,
278
290
  ) -> None:
279
291
  self.root = root
280
292
  self.path = root / "agent.toml"
281
293
  self._document = document
282
294
  self.framework = framework
283
295
  self.tools = tools
296
+ # Managed store bindings declared in agent.toml; None = unbound.
297
+ self.memory_store = memory_store
298
+ self.session_store = session_store
284
299
 
285
300
  @classmethod
286
301
  def load(cls, root: pathlib.Path | str) -> "AgentProject":
@@ -291,7 +306,8 @@ class AgentProject:
291
306
  except FileNotFoundError as exc:
292
307
  raise AgentCliError(
293
308
  f"Could not find agent.toml in {project_root}.",
294
- hint="Run `mason init` or use a legacy compatibility command.",
309
+ hint="This command needs a Mason project. Run `mason init` to create one, "
310
+ "or point at an existing project with --source <dir>.",
295
311
  ) from exc
296
312
  except (OSError, ParseError) as exc:
297
313
  raise AgentCliError(f"Could not read agent manifest at {path}: {exc}.") from exc
@@ -313,7 +329,13 @@ class AgentProject:
313
329
  ids = [tool.id for tool in tools]
314
330
  if len(ids) != len(set(ids)):
315
331
  raise AgentCliError("agent.toml tool ids must be unique.")
316
- return cls(project_root, document, framework, tools)
332
+ memory_store = _store_name_from_manifest(
333
+ document.get(MEMORY_STORE_TABLE), MEMORY_STORE_TABLE
334
+ )
335
+ session_store = _store_name_from_manifest(
336
+ document.get(SESSION_STORE_TABLE), SESSION_STORE_TABLE
337
+ )
338
+ return cls(project_root, document, framework, tools, memory_store, session_store)
317
339
 
318
340
  @classmethod
319
341
  def create(cls, root: pathlib.Path | str, *, framework: str) -> "AgentProject":
@@ -334,7 +356,17 @@ class AgentProject:
334
356
  continue
335
357
  if existing == spec:
336
358
  return False
337
- raise AgentCliError(f"Tool id {spec.id!r} already exists with different configuration.")
359
+
360
+ def _summary(s: ToolSpec) -> str:
361
+ src = s.source
362
+ return src.service or src.function or src.entrypoint or src.kind
363
+
364
+ raise AgentCliError(
365
+ f"Tool id {spec.id!r} already exists with a different configuration "
366
+ f"(existing: {_summary(existing)}; requested: {_summary(spec)}).",
367
+ hint="Use --name to add it under a different id, or remove the existing "
368
+ "tool from agent.toml first.",
369
+ )
338
370
  raw_tools = self._document.get("tools")
339
371
  if raw_tools is None:
340
372
  raw_tools = tomlkit.aot()
@@ -356,6 +388,44 @@ class AgentProject:
356
388
  del self.tools[index]
357
389
  return True
358
390
 
391
+ def bind_memory_store(self, name: str) -> bool:
392
+ """Declare the memory store binding in agent.toml. Returns True if it changed."""
393
+ return self._set_store(MEMORY_STORE_TABLE, name)
394
+
395
+ def bind_session_store(self, name: str) -> bool:
396
+ """Declare the session store binding in agent.toml. Returns True if it changed."""
397
+ return self._set_store(SESSION_STORE_TABLE, name)
398
+
399
+ def unbind_memory_store(self) -> bool:
400
+ """Remove the memory store binding from agent.toml. Returns True if it was present."""
401
+ return self._clear_store(MEMORY_STORE_TABLE)
402
+
403
+ def unbind_session_store(self) -> bool:
404
+ """Remove the session store binding from agent.toml. Returns True if it was present."""
405
+ return self._clear_store(SESSION_STORE_TABLE)
406
+
407
+ def _set_store(self, table: str, name: str) -> bool:
408
+ name = _required_string(name, f"[{table}] name")
409
+ if getattr(self, table) == name:
410
+ return False
411
+ existing = self._document.get(table)
412
+ if isinstance(existing, Mapping):
413
+ existing["name"] = name
414
+ else:
415
+ store_table = tomlkit.table()
416
+ store_table.add("name", name)
417
+ self._document.append(table, store_table)
418
+ setattr(self, table, name)
419
+ return True
420
+
421
+ def _clear_store(self, table: str) -> bool:
422
+ if getattr(self, table) is None:
423
+ return False
424
+ if table in self._document:
425
+ del self._document[table]
426
+ setattr(self, table, None)
427
+ return True
428
+
359
429
  def write(self) -> pathlib.Path:
360
430
  self.path.parent.mkdir(parents=True, exist_ok=True)
361
431
  temporary: pathlib.Path | None = None