databricks-mason 0.1.0.dev0__tar.gz → 0.1.1.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 (52) hide show
  1. databricks_mason-0.1.1.dev0/NOTICE +39 -0
  2. databricks_mason-0.1.1.dev0/PKG-INFO +216 -0
  3. databricks_mason-0.1.1.dev0/README.md +189 -0
  4. {databricks_mason-0.1.0.dev0 → databricks_mason-0.1.1.dev0}/pyproject.toml +24 -2
  5. databricks_mason-0.1.1.dev0/src/databricks_mason/__init__.py +89 -0
  6. databricks_mason-0.1.1.dev0/src/databricks_mason/agent_project.py +380 -0
  7. {databricks_mason-0.1.0.dev0 → databricks_mason-0.1.1.dev0}/src/databricks_mason/auth.py +2 -2
  8. {databricks_mason-0.1.0.dev0 → databricks_mason-0.1.1.dev0}/src/databricks_mason/cli.py +13 -5
  9. databricks_mason-0.1.1.dev0/src/databricks_mason/client.py +434 -0
  10. databricks_mason-0.1.1.dev0/src/databricks_mason/deploy.py +534 -0
  11. databricks_mason-0.1.1.dev0/src/databricks_mason/dev.py +160 -0
  12. databricks_mason-0.1.1.dev0/src/databricks_mason/init.py +283 -0
  13. databricks_mason-0.1.1.dev0/src/databricks_mason/langgraph/__init__.py +95 -0
  14. databricks_mason-0.1.1.dev0/src/databricks_mason/langgraph/long_running.py +59 -0
  15. databricks_mason-0.1.1.dev0/src/databricks_mason/langgraph/mcp.py +112 -0
  16. databricks_mason-0.1.1.dev0/src/databricks_mason/langgraph/memory.py +63 -0
  17. databricks_mason-0.1.1.dev0/src/databricks_mason/langgraph/recovery.py +247 -0
  18. databricks_mason-0.1.1.dev0/src/databricks_mason/langgraph/session_store.py +469 -0
  19. databricks_mason-0.1.1.dev0/src/databricks_mason/mcp.py +89 -0
  20. databricks_mason-0.1.1.dev0/src/databricks_mason/memory_store_access.py +30 -0
  21. databricks_mason-0.1.1.dev0/src/databricks_mason/models.py +270 -0
  22. databricks_mason-0.1.1.dev0/src/databricks_mason/project_config.py +141 -0
  23. databricks_mason-0.1.1.dev0/src/databricks_mason/py.typed +0 -0
  24. databricks_mason-0.1.1.dev0/src/databricks_mason/runtime/__init__.py +45 -0
  25. databricks_mason-0.1.1.dev0/src/databricks_mason/runtime/background.py +37 -0
  26. databricks_mason-0.1.1.dev0/src/databricks_mason/runtime/durability.py +377 -0
  27. databricks_mason-0.1.1.dev0/src/databricks_mason/runtime/session_store_client.py +120 -0
  28. databricks_mason-0.1.1.dev0/src/databricks_mason/runtime/tool_manifest.py +149 -0
  29. databricks_mason-0.1.1.dev0/src/databricks_mason/runtime/tracing.py +52 -0
  30. databricks_mason-0.1.1.dev0/src/databricks_mason/runtime/workspace.py +27 -0
  31. databricks_mason-0.1.1.dev0/src/databricks_mason/sandbox.py +822 -0
  32. databricks_mason-0.1.1.dev0/src/databricks_mason/session_store_access.py +18 -0
  33. databricks_mason-0.1.1.dev0/src/databricks_mason/store_access.py +167 -0
  34. databricks_mason-0.1.1.dev0/src/databricks_mason/templates/mcp_runtime_langgraph.py +96 -0
  35. databricks_mason-0.1.1.dev0/src/databricks_mason/templates/python_tool_langgraph.py +9 -0
  36. databricks_mason-0.1.1.dev0/src/databricks_mason/templates/python_tool_test.py +12 -0
  37. databricks_mason-0.1.1.dev0/src/databricks_mason/templates/sandbox_mcp.py +39 -0
  38. databricks_mason-0.1.1.dev0/src/databricks_mason/templates/sandbox_mcp_langgraph.py +39 -0
  39. databricks_mason-0.1.1.dev0/src/databricks_mason/templates/tool_manifest_runtime.py +143 -0
  40. databricks_mason-0.1.1.dev0/src/databricks_mason/tools.py +307 -0
  41. {databricks_mason-0.1.0.dev0 → databricks_mason-0.1.1.dev0}/src/databricks_mason/tracing.py +100 -15
  42. databricks_mason-0.1.0.dev0/PKG-INFO +0 -80
  43. databricks_mason-0.1.0.dev0/README.md +0 -65
  44. databricks_mason-0.1.0.dev0/src/databricks_mason/__init__.py +0 -1
  45. databricks_mason-0.1.0.dev0/src/databricks_mason/client.py +0 -298
  46. databricks_mason-0.1.0.dev0/src/databricks_mason/deploy.py +0 -325
  47. {databricks_mason-0.1.0.dev0 → databricks_mason-0.1.1.dev0}/.gitignore +0 -0
  48. {databricks_mason-0.1.0.dev0 → databricks_mason-0.1.1.dev0}/src/databricks_mason/errors.py +0 -0
  49. {databricks_mason-0.1.0.dev0 → databricks_mason-0.1.1.dev0}/src/databricks_mason/memory.py +0 -0
  50. {databricks_mason-0.1.0.dev0 → databricks_mason-0.1.1.dev0}/src/databricks_mason/render.py +0 -0
  51. {databricks_mason-0.1.0.dev0 → databricks_mason-0.1.1.dev0}/src/databricks_mason/sessions.py +0 -0
  52. {databricks_mason-0.1.0.dev0 → databricks_mason-0.1.1.dev0}/src/databricks_mason/timefmt.py +0 -0
@@ -0,0 +1,39 @@
1
+ Copyright (c) 2026 Databricks, Inc. All rights reserved.
2
+
3
+ This Software is the property of Databricks, Inc. ("Databricks") and no license, right, or authorization to install, use, copy, modify, or distribute this Software is granted except pursuant to Databricks' prior written authorization. Any use of this Software without such written authorization from Databricks is strictly prohibited.
4
+
5
+ # MIT License
6
+
7
+ This Software contains code from the following open source projects, licensed under the MIT license (https://opensource.org/licenses/MIT).
8
+
9
+ Package: PyYAML
10
+ Project URL: https://github.com/yaml/pyyaml
11
+ Copyright: 2017-2021 Ingy döt Net, 2006-2016 Kirill Simonov
12
+
13
+ Package: rich
14
+ Project URL: https://github.com/Textualize/rich
15
+ Copyright: 2020 Will McGugan
16
+
17
+ Package: tomli
18
+ Project URL: https://github.com/hukkin/tomli
19
+ Copyright: 2021 Taneli Hukkinen
20
+
21
+ Package: tomlkit
22
+ Project URL: https://github.com/python-poetry/tomlkit
23
+ Copyright: 2018 Sébastien Eustace
24
+
25
+ # Apache License 2.0
26
+
27
+ This Software contains code from the following open source projects, licensed under the Apache License 2.0 license (https://www.apache.org/licenses/LICENSE-2.0).
28
+
29
+ Package: databricks-sdk
30
+ Project URL: https://github.com/databricks/databricks-sdk-py
31
+ Copyright: 2023 Databricks, Inc.
32
+
33
+ # BSD 3-Clause License
34
+
35
+ This Software contains code from the following open source projects, licensed under the BSD 3-Clause License (https://opensource.org/license/bsd-3-clause).
36
+
37
+ Package: click
38
+ Project URL: https://github.com/pallets/click
39
+ Copyright: 2014 Pallets
@@ -0,0 +1,216 @@
1
+ Metadata-Version: 2.5
2
+ Name: databricks-mason
3
+ Version: 0.1.1.dev0
4
+ Summary: Databricks integration for Mason
5
+ Author-email: Databricks <agent-feedback@databricks.com>
6
+ License-File: NOTICE
7
+ Requires-Python: >=3.10
8
+ Requires-Dist: click>=8.1
9
+ Requires-Dist: databricks-sdk>=0.49
10
+ Requires-Dist: pyyaml>=6.0
11
+ Requires-Dist: rich>=13.7
12
+ Requires-Dist: tomli>=2.0
13
+ Requires-Dist: tomlkit>=0.13
14
+ Provides-Extra: runtime
15
+ Requires-Dist: databricks-agents>=1.9.3; extra == 'runtime'
16
+ Requires-Dist: databricks-langchain>=0.17.0; extra == 'runtime'
17
+ Requires-Dist: fastapi>=0.129.0; extra == 'runtime'
18
+ Requires-Dist: langchain-mcp-adapters>=0.2.1; extra == 'runtime'
19
+ Requires-Dist: langchain>=1.0.0; extra == 'runtime'
20
+ Requires-Dist: langgraph>=1.1.0; extra == 'runtime'
21
+ Requires-Dist: mlflow>=3.10.1; extra == 'runtime'
22
+ Requires-Dist: opentelemetry-exporter-otlp-proto-grpc>=1.25.0; extra == 'runtime'
23
+ Requires-Dist: uuid-utils>=0.10.0; extra == 'runtime'
24
+ Provides-Extra: tracing
25
+ Requires-Dist: mlflow[databricks]>=3.9.0; extra == 'tracing'
26
+ Description-Content-Type: text/markdown
27
+
28
+ # `databricks-mason`
29
+
30
+ Mason is an experimental CLI for Databricks custom agent preview APIs and
31
+ deployments. It manages memory, sessions, tracing, and deployments from one
32
+ authenticated command.
33
+
34
+ > The underlying APIs are in preview and may need workspace enablement.
35
+
36
+ ## Installation
37
+
38
+ From PyPI:
39
+
40
+ ```sh
41
+ pip install databricks-mason
42
+ ```
43
+
44
+ From source:
45
+
46
+ ```sh
47
+ pip install 'git+https://github.com/databricks/databricks-ai-bridge.git#subdirectory=integrations/mason'
48
+ ```
49
+
50
+ For tracing commands, install Mason with tracing extras:
51
+
52
+ ```sh
53
+ pip install 'databricks-mason[tracing]'
54
+ ```
55
+
56
+ ## Authentication
57
+
58
+ 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:
61
+
62
+ ```sh
63
+ databricks auth login --profile <profile>
64
+ mason login --profile <profile>
65
+ mason sessions stores list
66
+ ```
67
+
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.
73
+
74
+ ## Python SDK
75
+
76
+ The same memory and session APIs are available programmatically through
77
+ `MasonClient`, which authenticates exactly like the CLI (a `.databrickscfg` profile
78
+ or the SDK's default resolution):
79
+
80
+ ```python
81
+ from databricks_mason import MasonClient
82
+
83
+ client = MasonClient(profile="my-workspace") # or MasonClient() for default auth
84
+
85
+ store = client.create_memory_store("my-store")
86
+ print(store.name, store.display_name) # typed attribute access
87
+
88
+ client.create_memory_entry("my-store", actor_id="alice", path="/notes/1.md", content="hi")
89
+ for entry in client.list_memory_entries("my-store", actor_id="alice").entries:
90
+ print(entry.path, entry.content)
91
+ ```
92
+
93
+ Each method maps to one `/api/agents/v1` operation. Responses come back as typed
94
+ models (`MemoryStore`, `Session`, `SessionItemList`, ...) that expose attribute
95
+ accessors (`store.name`) while remaining plain dicts underneath — so `store["name"]`,
96
+ `json.dumps(store)`, and any new server-side fields keep working. API errors raise
97
+ `databricks_mason.AgentCliError`. Deployment, sandbox, and tracing remain CLI-only.
98
+
99
+ ## Commands
100
+
101
+ ```text
102
+ mason [-p <profile>] [-o text|json]
103
+ login [--profile P]
104
+ logout
105
+ init [--framework openai|langgraph] [--enable-chat-app]
106
+ [--profile P] [--repo URL] [--ref REF] [directory]
107
+ memory
108
+ stores create | list | get | update | delete
109
+ entries create | get | list | search | update | delete
110
+ sessions create | list | get | update | delete | fork
111
+ stores create | list | get | update | delete
112
+ items list | append | pop | clear
113
+ tracing
114
+ setup --catalog C --schema S [--experiment E]
115
+ list | get | instrument
116
+ mcp
117
+ list [--schema CATALOG.SCHEMA]
118
+ init [--framework openai|langgraph] [--profile P] [DIRECTORY]
119
+ tools
120
+ add sandbox --scope SCOPE [--scope SCOPE ...] [--source PATH]
121
+ add mcp SERVICE [--name NAME] [--source PATH]
122
+ add uc-function FUNCTION [--name NAME] [--source PATH]
123
+ add python NAME [--source PATH]
124
+ 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]
128
+ deployments list | get | logs | start | stop | delete
129
+ ```
130
+
131
+ ## Agent tools
132
+
133
+ `mason init` writes portable tool intent to `agent.toml` and template provenance to
134
+ `.mason/project.toml`. The manifest runtime is currently implemented only by the in-repository
135
+ `agent-langgraph` template; `mason tools add` fails explicitly for other frameworks until they
136
+ provide an adapter at the same runtime seam.
137
+
138
+ Remote tools update only `agent.toml`; they do not generate framework source. The LangGraph runtime
139
+ loads the manifest and materializes its native MCP tools when the agent runs, so a direct manifest
140
+ edit and a CLI edit have the same behavior:
141
+
142
+ ```sh
143
+ mason tools add sandbox --scope table:samples.nyctaxi.trips
144
+ mason tools add mcp system.ai.web_search
145
+ mason tools add uc-function catalog.schema.lookup_ticket
146
+ mason tools add python lookup-ticket
147
+ mason tools list
148
+ ```
149
+
150
+ Discover the MCP Services available to your user before adding one. By default Mason lists the
151
+ Databricks-managed services in `system.ai`; pass `--schema catalog.schema` for another Unity Catalog
152
+ schema. Text output includes a copyable add command, while `--output json` returns normalized service
153
+ records for scripts:
154
+
155
+ ```sh
156
+ mason mcp list
157
+ mason mcp list --schema main.tools
158
+ ```
159
+
160
+ The Python command additionally creates user-owned `agent/tools/<name>.py` and
161
+ `tests/tools/test_<name>.py` files using the LangGraph-native `@tool` decorator. `mason dev` and
162
+ `mason deploy` preserve `agent.toml`; they do not generate or patch agent source.
163
+
164
+ Sandbox scopes default to read-only access. Repeat `--scope` to allow more than one resource, use
165
+ `volume:` or `workspace:` for those resource types, and use `--permission read_write` only when the
166
+ agent needs writes. Every sandbox call carries this fixed downscope in MCP `_meta`, outside the tool
167
+ arguments controlled by the model.
168
+
169
+ ## Initialize the chat app demo
170
+
171
+ The chat app is a LangGraph-specific init overlay, not a command that mutates an existing project:
172
+
173
+ ```sh
174
+ mason init --framework langgraph --enable-chat-app \
175
+ --profile <profile> \
176
+ ./my-agent
177
+ cd ./my-agent
178
+ uv run start-server
179
+ ```
180
+
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.
186
+
187
+ For the full deployed demo, connect both managed stores:
188
+
189
+ ```sh
190
+ 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
195
+ ```
196
+
197
+ The Databricks Apps `__Host-databricks-app-router` cookie is both the sticky routing key and the
198
+ application session id. The browser sends it automatically; API clients must reuse it as a cookie.
199
+ Request bodies never carry `session_id`. A localhost-only `mason-local-session` cookie provides the
200
+ same behavior outside Databricks Apps. TODO: move to `X-Routing-Key` when Apps supports it.
201
+
202
+ 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.
@@ -0,0 +1,189 @@
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
+ ## Authentication
30
+
31
+ 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:
34
+
35
+ ```sh
36
+ databricks auth login --profile <profile>
37
+ mason login --profile <profile>
38
+ mason sessions stores list
39
+ ```
40
+
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.
46
+
47
+ ## Python SDK
48
+
49
+ The same memory and session APIs are available programmatically through
50
+ `MasonClient`, which authenticates exactly like the CLI (a `.databrickscfg` profile
51
+ or the SDK's default resolution):
52
+
53
+ ```python
54
+ from databricks_mason import MasonClient
55
+
56
+ client = MasonClient(profile="my-workspace") # or MasonClient() for default auth
57
+
58
+ store = client.create_memory_store("my-store")
59
+ print(store.name, store.display_name) # typed attribute access
60
+
61
+ client.create_memory_entry("my-store", actor_id="alice", path="/notes/1.md", content="hi")
62
+ for entry in client.list_memory_entries("my-store", actor_id="alice").entries:
63
+ print(entry.path, entry.content)
64
+ ```
65
+
66
+ Each method maps to one `/api/agents/v1` operation. Responses come back as typed
67
+ models (`MemoryStore`, `Session`, `SessionItemList`, ...) that expose attribute
68
+ accessors (`store.name`) while remaining plain dicts underneath — so `store["name"]`,
69
+ `json.dumps(store)`, and any new server-side fields keep working. API errors raise
70
+ `databricks_mason.AgentCliError`. Deployment, sandbox, and tracing remain CLI-only.
71
+
72
+ ## Commands
73
+
74
+ ```text
75
+ mason [-p <profile>] [-o text|json]
76
+ login [--profile P]
77
+ logout
78
+ init [--framework openai|langgraph] [--enable-chat-app]
79
+ [--profile P] [--repo URL] [--ref REF] [directory]
80
+ memory
81
+ stores create | list | get | update | delete
82
+ entries create | get | list | search | update | delete
83
+ sessions create | list | get | update | delete | fork
84
+ stores create | list | get | update | delete
85
+ items list | append | pop | clear
86
+ tracing
87
+ setup --catalog C --schema S [--experiment E]
88
+ list | get | instrument
89
+ mcp
90
+ list [--schema CATALOG.SCHEMA]
91
+ init [--framework openai|langgraph] [--profile P] [DIRECTORY]
92
+ tools
93
+ add sandbox --scope SCOPE [--scope SCOPE ...] [--source PATH]
94
+ add mcp SERVICE [--name NAME] [--source PATH]
95
+ add uc-function FUNCTION [--name NAME] [--source PATH]
96
+ add python NAME [--source PATH]
97
+ 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]
101
+ deployments list | get | logs | start | stop | delete
102
+ ```
103
+
104
+ ## Agent tools
105
+
106
+ `mason init` writes portable tool intent to `agent.toml` and template provenance to
107
+ `.mason/project.toml`. The manifest runtime is currently implemented only by the in-repository
108
+ `agent-langgraph` template; `mason tools add` fails explicitly for other frameworks until they
109
+ provide an adapter at the same runtime seam.
110
+
111
+ Remote tools update only `agent.toml`; they do not generate framework source. The LangGraph runtime
112
+ loads the manifest and materializes its native MCP tools when the agent runs, so a direct manifest
113
+ edit and a CLI edit have the same behavior:
114
+
115
+ ```sh
116
+ mason tools add sandbox --scope table:samples.nyctaxi.trips
117
+ mason tools add mcp system.ai.web_search
118
+ mason tools add uc-function catalog.schema.lookup_ticket
119
+ mason tools add python lookup-ticket
120
+ mason tools list
121
+ ```
122
+
123
+ Discover the MCP Services available to your user before adding one. By default Mason lists the
124
+ Databricks-managed services in `system.ai`; pass `--schema catalog.schema` for another Unity Catalog
125
+ schema. Text output includes a copyable add command, while `--output json` returns normalized service
126
+ records for scripts:
127
+
128
+ ```sh
129
+ mason mcp list
130
+ mason mcp list --schema main.tools
131
+ ```
132
+
133
+ The Python command additionally creates user-owned `agent/tools/<name>.py` and
134
+ `tests/tools/test_<name>.py` files using the LangGraph-native `@tool` decorator. `mason dev` and
135
+ `mason deploy` preserve `agent.toml`; they do not generate or patch agent source.
136
+
137
+ Sandbox scopes default to read-only access. Repeat `--scope` to allow more than one resource, use
138
+ `volume:` or `workspace:` for those resource types, and use `--permission read_write` only when the
139
+ agent needs writes. Every sandbox call carries this fixed downscope in MCP `_meta`, outside the tool
140
+ arguments controlled by the model.
141
+
142
+ ## Initialize the chat app demo
143
+
144
+ The chat app is a LangGraph-specific init overlay, not a command that mutates an existing project:
145
+
146
+ ```sh
147
+ mason init --framework langgraph --enable-chat-app \
148
+ --profile <profile> \
149
+ ./my-agent
150
+ cd ./my-agent
151
+ uv run start-server
152
+ ```
153
+
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.
159
+
160
+ For the full deployed demo, connect both managed stores:
161
+
162
+ ```sh
163
+ 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
168
+ ```
169
+
170
+ The Databricks Apps `__Host-databricks-app-router` cookie is both the sticky routing key and the
171
+ application session id. The browser sends it automatically; API clients must reuse it as a cookie.
172
+ Request bodies never carry `session_id`. A localhost-only `mason-local-session` cookie provides the
173
+ same behavior outside Databricks Apps. TODO: move to `X-Routing-Key` when Apps supports it.
174
+
175
+ 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.
@@ -1,24 +1,39 @@
1
1
  [project]
2
2
  name = "databricks-mason"
3
- version = "0.1.0.dev0"
3
+ version = "0.1.1.dev0"
4
4
  description = "Databricks integration for Mason"
5
5
  authors = [
6
6
  { name="Databricks", email="agent-feedback@databricks.com" },
7
7
  ]
8
8
  readme = "README.md"
9
- license = { text="Apache-2.0" }
10
9
  requires-python = ">=3.10"
11
10
  dependencies = [
12
11
  "click>=8.1",
13
12
  "databricks-sdk>=0.49",
14
13
  "PyYAML>=6.0",
15
14
  "rich>=13.7",
15
+ "tomli>=2.0",
16
+ "tomlkit>=0.13",
16
17
  ]
17
18
 
18
19
  [project.optional-dependencies]
19
20
  tracing = [
20
21
  "mlflow[databricks]>=3.9.0",
21
22
  ]
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]`.
26
+ runtime = [
27
+ "databricks-langchain>=0.17.0",
28
+ "langgraph>=1.1.0",
29
+ "langchain>=1.0.0",
30
+ "langchain-mcp-adapters>=0.2.1",
31
+ "fastapi>=0.129.0",
32
+ "mlflow>=3.10.1",
33
+ "uuid-utils>=0.10.0",
34
+ "opentelemetry-exporter-otlp-proto-grpc>=1.25.0",
35
+ "databricks-agents>=1.9.3",
36
+ ]
22
37
 
23
38
  [project.scripts]
24
39
  mason = "databricks_mason.cli:main"
@@ -39,7 +54,10 @@ build-backend = "hatchling.build"
39
54
 
40
55
  [tool.hatch.build]
41
56
  include = [
57
+ "NOTICE",
42
58
  "src/databricks_mason/*",
59
+ "src/databricks_mason/py.typed",
60
+ "src/databricks_mason/templates/**",
43
61
  ]
44
62
 
45
63
  [tool.hatch.build.targets.wheel]
@@ -57,3 +75,7 @@ root = ["./src", "./tests"]
57
75
 
58
76
  [tool.ty.src]
59
77
  include = ["./src", "./tests"]
78
+ # The files under templates/ are code-as-data — scaffolding snippets written into a generated agent
79
+ # project, not importable package modules. They import scaffold-relative paths (agent.mason.*,
80
+ # agent.mcps) that never resolve in the package, so type-checking them here is meaningless.
81
+ exclude = ["./src/databricks_mason/templates"]
@@ -0,0 +1,89 @@
1
+ """Databricks integration for Mason.
2
+
3
+ Mason ships a `mason` CLI and this Python SDK over the same agents/v1 preview APIs.
4
+ Construct a `MasonClient` and call one method per API operation:
5
+
6
+ from databricks_mason import MasonClient
7
+
8
+ client = MasonClient(profile="my-workspace")
9
+ client.create_memory_store("my-store")
10
+ for entry in client.list_memory_entries("my-store", actor_id="alice").get("entries", []):
11
+ ...
12
+
13
+ Auth resolves through the Databricks SDK: pass a `.databrickscfg` profile or rely on
14
+ its default resolution. API errors surface as `AgentCliError`.
15
+
16
+ The framework-neutral runtime helpers (``configure_tracing``, ``tag_session``, ``workspace_client``,
17
+ ``workspace_headers``) are also re-exported here for convenience — they live in
18
+ :mod:`databricks_mason.runtime` and are resolved lazily (PEP 562) so a plain CLI ``import
19
+ databricks_mason`` does not pull in the tracing module's ``mlflow`` dependency.
20
+ """
21
+
22
+ from typing import TYPE_CHECKING
23
+
24
+ from databricks_mason.client import MasonClient, memory_entry_path, memory_store_path
25
+ from databricks_mason.errors import AgentCliError
26
+ from databricks_mason.models import (
27
+ MemoryEntry,
28
+ MemoryEntryList,
29
+ MemorySearchHit,
30
+ MemorySearchResult,
31
+ MemoryStore,
32
+ MemoryStoreList,
33
+ PoppedSessionItem,
34
+ Session,
35
+ SessionItem,
36
+ SessionItemList,
37
+ SessionList,
38
+ SessionStore,
39
+ SessionStoreList,
40
+ StorageBackend,
41
+ )
42
+
43
+ if TYPE_CHECKING:
44
+ from databricks_mason.runtime import (
45
+ configure_tracing,
46
+ tag_session,
47
+ workspace_client,
48
+ workspace_headers,
49
+ )
50
+
51
+ __all__ = [
52
+ "MasonClient",
53
+ "AgentCliError",
54
+ "memory_store_path",
55
+ "memory_entry_path",
56
+ "MemoryStore",
57
+ "MemoryStoreList",
58
+ "MemoryEntry",
59
+ "MemoryEntryList",
60
+ "MemorySearchResult",
61
+ "MemorySearchHit",
62
+ "StorageBackend",
63
+ "SessionStore",
64
+ "SessionStoreList",
65
+ "Session",
66
+ "SessionList",
67
+ "SessionItem",
68
+ "SessionItemList",
69
+ "PoppedSessionItem",
70
+ # Framework-neutral runtime helpers (lazily re-exported from databricks_mason.runtime).
71
+ "configure_tracing",
72
+ "tag_session",
73
+ "workspace_client",
74
+ "workspace_headers",
75
+ ]
76
+
77
+ # Neutral runtime helpers, re-exported lazily so the light CLI import path stays free of the agent
78
+ # stack (mlflow, etc.). Everything else above is light and imported eagerly.
79
+ _RUNTIME_REEXPORTS = frozenset(
80
+ {"configure_tracing", "tag_session", "workspace_client", "workspace_headers"}
81
+ )
82
+
83
+
84
+ def __getattr__(name: str) -> object:
85
+ if name in _RUNTIME_REEXPORTS:
86
+ import importlib
87
+
88
+ return getattr(importlib.import_module("databricks_mason.runtime"), name)
89
+ raise AttributeError(f"module {__name__!r} has no attribute {name!r}")