mostlyright-data 0.9.0__py3-none-any.whl

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 (314) hide show
  1. mostlyright/data_harness/__init__.py +158 -0
  2. mostlyright/data_harness/acquisition/__init__.py +55 -0
  3. mostlyright/data_harness/acquisition/http.py +2773 -0
  4. mostlyright/data_harness/acquisition/parsing.py +809 -0
  5. mostlyright/data_harness/acquisition/ranges.py +495 -0
  6. mostlyright/data_harness/acquisition/result_download.py +360 -0
  7. mostlyright/data_harness/acquisition/retention_admission.py +248 -0
  8. mostlyright/data_harness/acquisition/sandbox.py +4888 -0
  9. mostlyright/data_harness/acquisition/url_policy.py +530 -0
  10. mostlyright/data_harness/agent_runtime.py +2743 -0
  11. mostlyright/data_harness/assets/logo-ink.svg +31 -0
  12. mostlyright/data_harness/backends/__init__.py +28 -0
  13. mostlyright/data_harness/backends/pandas_backend.py +350 -0
  14. mostlyright/data_harness/backends/polars_backend.py +366 -0
  15. mostlyright/data_harness/backends/protocol.py +124 -0
  16. mostlyright/data_harness/backends/reference.py +83 -0
  17. mostlyright/data_harness/backends/registry.py +55 -0
  18. mostlyright/data_harness/backends/restrictions.py +126 -0
  19. mostlyright/data_harness/canonical.py +333 -0
  20. mostlyright/data_harness/catalog_job.py +625 -0
  21. mostlyright/data_harness/cli.py +5398 -0
  22. mostlyright/data_harness/contracts.py +53 -0
  23. mostlyright/data_harness/coordinator.py +1307 -0
  24. mostlyright/data_harness/deploy.py +924 -0
  25. mostlyright/data_harness/deploy_target.py +312 -0
  26. mostlyright/data_harness/deployment_evidence.py +1067 -0
  27. mostlyright/data_harness/event_presentation.py +576 -0
  28. mostlyright/data_harness/events.py +2152 -0
  29. mostlyright/data_harness/fast_delimited.py +239 -0
  30. mostlyright/data_harness/fleet.py +237 -0
  31. mostlyright/data_harness/formats.py +236 -0
  32. mostlyright/data_harness/governors.py +1163 -0
  33. mostlyright/data_harness/hosted_bootstrap.py +972 -0
  34. mostlyright/data_harness/hosted_crawler.py +1115 -0
  35. mostlyright/data_harness/hosted_crawler_container_smoke.py +351 -0
  36. mostlyright/data_harness/hosted_crawler_fetch.py +423 -0
  37. mostlyright/data_harness/hosted_crawler_job.py +1277 -0
  38. mostlyright/data_harness/hosted_crawler_protocol.py +676 -0
  39. mostlyright/data_harness/hosted_dataset.py +1500 -0
  40. mostlyright/data_harness/hosted_deploy.py +3037 -0
  41. mostlyright/data_harness/hosted_handoff.py +62 -0
  42. mostlyright/data_harness/hosted_ingestion_contract.py +504 -0
  43. mostlyright/data_harness/hosted_ingestion_job.py +356 -0
  44. mostlyright/data_harness/hosted_ingestion_job_smoke.py +40 -0
  45. mostlyright/data_harness/hosted_session_container_smoke.py +194 -0
  46. mostlyright/data_harness/hosted_session_worker.py +3554 -0
  47. mostlyright/data_harness/hosted_session_worker_job_smoke.py +46 -0
  48. mostlyright/data_harness/hosted_worker.py +6784 -0
  49. mostlyright/data_harness/ingestion/__init__.py +56 -0
  50. mostlyright/data_harness/ingestion/contracts.py +461 -0
  51. mostlyright/data_harness/ingestion/faults.py +42 -0
  52. mostlyright/data_harness/ingestion/gcs_store.py +1162 -0
  53. mostlyright/data_harness/ingestion/spool.py +130 -0
  54. mostlyright/data_harness/ingestion/store.py +885 -0
  55. mostlyright/data_harness/key_seam.py +434 -0
  56. mostlyright/data_harness/linux_process_boundary.py +262 -0
  57. mostlyright/data_harness/local_contracts.py +2880 -0
  58. mostlyright/data_harness/local_search/__init__.py +5 -0
  59. mostlyright/data_harness/local_search/build_index.py +1087 -0
  60. mostlyright/data_harness/local_search/contracts.py +920 -0
  61. mostlyright/data_harness/local_search/query_trace.py +266 -0
  62. mostlyright/data_harness/local_search/retrieval.py +700 -0
  63. mostlyright/data_harness/local_search/sealed.py +474 -0
  64. mostlyright/data_harness/local_search/service.py +784 -0
  65. mostlyright/data_harness/nbrender/CONTRACT.md +212 -0
  66. mostlyright/data_harness/nbrender/__init__.py +12 -0
  67. mostlyright/data_harness/nbrender/chrome.py +359 -0
  68. mostlyright/data_harness/nbrender/code_body.py +266 -0
  69. mostlyright/data_harness/nbrender/document.py +407 -0
  70. mostlyright/data_harness/nbrender/frame.py +275 -0
  71. mostlyright/data_harness/nbrender/interactive.py +337 -0
  72. mostlyright/data_harness/nbrender/markdown_body.py +477 -0
  73. mostlyright/data_harness/nbrender/mr_components.py +134 -0
  74. mostlyright/data_harness/nbrender/outputs_data.py +595 -0
  75. mostlyright/data_harness/nbrender/outputs_rich.py +906 -0
  76. mostlyright/data_harness/nbrender/outputs_source.py +260 -0
  77. mostlyright/data_harness/nbrender/outputs_stage.py +176 -0
  78. mostlyright/data_harness/nbrender/outputs_text.py +400 -0
  79. mostlyright/data_harness/nbrender/parse.py +394 -0
  80. mostlyright/data_harness/nbrender/status.py +40 -0
  81. mostlyright/data_harness/nbrender/tokens.py +1295 -0
  82. mostlyright/data_harness/notebook.py +1710 -0
  83. mostlyright/data_harness/offline.py +2049 -0
  84. mostlyright/data_harness/operation_registry.py +1007 -0
  85. mostlyright/data_harness/operator_setup.py +239 -0
  86. mostlyright/data_harness/pipeline.py +6428 -0
  87. mostlyright/data_harness/plan_graph.py +2026 -0
  88. mostlyright/data_harness/preparation/__init__.py +104 -0
  89. mostlyright/data_harness/preparation/contracts.py +1017 -0
  90. mostlyright/data_harness/preparation/engine.py +221 -0
  91. mostlyright/data_harness/preparation/errors.py +14 -0
  92. mostlyright/data_harness/preparation/gates.py +751 -0
  93. mostlyright/data_harness/preparation/joins.py +574 -0
  94. mostlyright/data_harness/preparation/profile.py +384 -0
  95. mostlyright/data_harness/preparation/table.py +217 -0
  96. mostlyright/data_harness/preparation/transforms.py +568 -0
  97. mostlyright/data_harness/progress_events.py +534 -0
  98. mostlyright/data_harness/readers/__init__.py +46 -0
  99. mostlyright/data_harness/readers/containers.py +963 -0
  100. mostlyright/data_harness/readers/contracts.py +542 -0
  101. mostlyright/data_harness/readers/delimited.py +257 -0
  102. mostlyright/data_harness/readers/grib2/__init__.py +33 -0
  103. mostlyright/data_harness/readers/grib2/admission.py +722 -0
  104. mostlyright/data_harness/readers/grib2/decode.py +1009 -0
  105. mostlyright/data_harness/readers/grib2/geometry.py +1133 -0
  106. mostlyright/data_harness/readers/grib2/portable_math.py +501 -0
  107. mostlyright/data_harness/readers/json_tabular.py +485 -0
  108. mostlyright/data_harness/readers/registry.py +514 -0
  109. mostlyright/data_harness/readers/samples/README.md +110 -0
  110. mostlyright/data_harness/readers/samples/archive.gzip/1.0.0/cities_one_stream/cities.csv.gz +0 -0
  111. mostlyright/data_harness/readers/samples/archive.gzip/1.0.0/cities_one_stream/expected.json +24 -0
  112. mostlyright/data_harness/readers/samples/archive.gzip/1.1.0/cities_one_stream/cities.csv.gz +0 -0
  113. mostlyright/data_harness/readers/samples/archive.gzip/1.1.0/cities_one_stream/expected.json +24 -0
  114. mostlyright/data_harness/readers/samples/archive.tar/1.0.0/cities_beside_a_directory_entry/cities.tar +0 -0
  115. mostlyright/data_harness/readers/samples/archive.tar/1.0.0/cities_beside_a_directory_entry/expected.json +24 -0
  116. mostlyright/data_harness/readers/samples/archive.tar/1.1.0/cities_beside_a_directory_entry/cities.tar +0 -0
  117. mostlyright/data_harness/readers/samples/archive.tar/1.1.0/cities_beside_a_directory_entry/expected.json +24 -0
  118. mostlyright/data_harness/readers/samples/archive.zip/1.0.0/cities_beside_a_second_member/cities.zip +0 -0
  119. mostlyright/data_harness/readers/samples/archive.zip/1.0.0/cities_beside_a_second_member/expected.json +25 -0
  120. mostlyright/data_harness/readers/samples/archive.zip/1.1.0/dwd_semicolon_station_member/dwd-station.zip +0 -0
  121. mostlyright/data_harness/readers/samples/archive.zip/1.1.0/dwd_semicolon_station_member/expected.json +25 -0
  122. mostlyright/data_harness/readers/samples/archive.zip/1.2.0/dwd_semicolon_station_member/dwd-station.zip +0 -0
  123. mostlyright/data_harness/readers/samples/archive.zip/1.2.0/dwd_semicolon_station_member/expected.json +25 -0
  124. mostlyright/data_harness/readers/samples/delimited_text/1.0.0/an_ordinary_comma_separated_table/cities.csv +3 -0
  125. mostlyright/data_harness/readers/samples/delimited_text/1.0.0/an_ordinary_comma_separated_table/expected.json +23 -0
  126. mostlyright/data_harness/readers/samples/delimited_text/1.0.0/quoted_fields_holding_the_delimiter/cities.tsv +5 -0
  127. mostlyright/data_harness/readers/samples/delimited_text/1.0.0/quoted_fields_holding_the_delimiter/expected.json +25 -0
  128. mostlyright/data_harness/readers/samples/delimited_text/1.1.0/an_hourly_observation_table_served_as_plain_text/expected.json +30 -0
  129. mostlyright/data_harness/readers/samples/delimited_text/1.1.0/an_hourly_observation_table_served_as_plain_text/observations.csv +5 -0
  130. mostlyright/data_harness/readers/samples/json.tabular/1.0.0/nested_hourly_observations/expected.json +44 -0
  131. mostlyright/data_harness/readers/samples/json.tabular/1.0.0/nested_hourly_observations/stations.json +1 -0
  132. mostlyright/data_harness/readers/samples/json.tabular/1.1.0/an_observation_stream_served_as_plain_text/expected.json +48 -0
  133. mostlyright/data_harness/readers/samples/json.tabular/1.1.0/an_observation_stream_served_as_plain_text/observations.ndjson +4 -0
  134. mostlyright/data_harness/readers/samples/spreadsheet.xlsx/1.0.0/an_ordinary_table_beside_a_second_sheet/cities.xlsx +0 -0
  135. mostlyright/data_harness/readers/samples/spreadsheet.xlsx/1.0.0/an_ordinary_table_beside_a_second_sheet/expected.json +24 -0
  136. mostlyright/data_harness/readers/samples/spreadsheet.xlsx/1.0.0/shares_the_workbook_had_already_computed/expected.json +27 -0
  137. mostlyright/data_harness/readers/samples/spreadsheet.xlsx/1.0.0/shares_the_workbook_had_already_computed/shares.xlsx +0 -0
  138. mostlyright/data_harness/readers/samples/spreadsheet.xlsx/1.1.0/shares_the_workbook_had_already_computed/expected.json +27 -0
  139. mostlyright/data_harness/readers/samples/spreadsheet.xlsx/1.1.0/shares_the_workbook_had_already_computed/shares.xlsx +0 -0
  140. mostlyright/data_harness/readers/samples/weather.grib2/1.0.0/README.md +20 -0
  141. mostlyright/data_harness/readers/samples/weather.grib2/1.0.0/gfs_2m_temperature/expected.json +55 -0
  142. mostlyright/data_harness/readers/samples/weather.grib2/1.0.0/gfs_2m_temperature/gfs-2m-temperature.grib2 +0 -0
  143. mostlyright/data_harness/readers/samples/weather.grib2/1.0.0/hrrr_2m_temperature/expected.json +54 -0
  144. mostlyright/data_harness/readers/samples/weather.grib2/1.0.0/hrrr_2m_temperature/hrrr-2m-temperature.grib2 +0 -0
  145. mostlyright/data_harness/readers/samples/weather.grib2/1.0.0/hrrr_categorical_rain/expected.json +54 -0
  146. mostlyright/data_harness/readers/samples/weather.grib2/1.0.0/hrrr_categorical_rain/hrrr-categorical-rain.grib2 +0 -0
  147. mostlyright/data_harness/readers/samples/weather.grib2/2.0.0/hrrr_2m_temperature/expected.json +54 -0
  148. mostlyright/data_harness/readers/samples/weather.grib2/2.0.0/hrrr_2m_temperature/hrrr-2m-temperature.grib2 +0 -0
  149. mostlyright/data_harness/readers/samples.py +582 -0
  150. mostlyright/data_harness/readers/spreadsheet.py +803 -0
  151. mostlyright/data_harness/readers/tabular.py +510 -0
  152. mostlyright/data_harness/recipe.py +5321 -0
  153. mostlyright/data_harness/repair/__init__.py +78 -0
  154. mostlyright/data_harness/repair/adapters.py +274 -0
  155. mostlyright/data_harness/repair/contracts.py +872 -0
  156. mostlyright/data_harness/repair/coordinator.py +1099 -0
  157. mostlyright/data_harness/repair/errors.py +16 -0
  158. mostlyright/data_harness/review.py +2533 -0
  159. mostlyright/data_harness/rowset.py +283 -0
  160. mostlyright/data_harness/serving.py +1975 -0
  161. mostlyright/data_harness/serving_edge.py +590 -0
  162. mostlyright/data_harness/serving_http.py +1031 -0
  163. mostlyright/data_harness/session_probes.py +759 -0
  164. mostlyright/data_harness/signing.py +101 -0
  165. mostlyright/data_harness/source_discovery.py +898 -0
  166. mostlyright/data_harness/sources/__init__.py +209 -0
  167. mostlyright/data_harness/sources/_adapter_steps.py +213 -0
  168. mostlyright/data_harness/sources/adapters.py +1214 -0
  169. mostlyright/data_harness/sources/cadence.py +1428 -0
  170. mostlyright/data_harness/sources/cadence_emission.py +453 -0
  171. mostlyright/data_harness/sources/cadence_history.py +546 -0
  172. mostlyright/data_harness/sources/catalog/__init__.py +17 -0
  173. mostlyright/data_harness/sources/catalog/admission.py +477 -0
  174. mostlyright/data_harness/sources/catalog/authoring.py +1701 -0
  175. mostlyright/data_harness/sources/catalog/authoring_policy.py +701 -0
  176. mostlyright/data_harness/sources/catalog/authoring_shards.py +1217 -0
  177. mostlyright/data_harness/sources/catalog/bounded_io.py +231 -0
  178. mostlyright/data_harness/sources/catalog/channel.py +523 -0
  179. mostlyright/data_harness/sources/catalog/channel_client.py +296 -0
  180. mostlyright/data_harness/sources/catalog/contracts.py +825 -0
  181. mostlyright/data_harness/sources/catalog/coverage.py +137 -0
  182. mostlyright/data_harness/sources/catalog/delta.py +1340 -0
  183. mostlyright/data_harness/sources/catalog/embedding.py +532 -0
  184. mostlyright/data_harness/sources/catalog/entry_v2.py +1182 -0
  185. mostlyright/data_harness/sources/catalog/fill.py +3889 -0
  186. mostlyright/data_harness/sources/catalog/fill_partitions.py +459 -0
  187. mostlyright/data_harness/sources/catalog/fill_staging.py +1105 -0
  188. mostlyright/data_harness/sources/catalog/gating.py +374 -0
  189. mostlyright/data_harness/sources/catalog/generation_receipt.py +1607 -0
  190. mostlyright/data_harness/sources/catalog/harvest/__init__.py +7 -0
  191. mostlyright/data_harness/sources/catalog/harvest/ckan.py +384 -0
  192. mostlyright/data_harness/sources/catalog/harvest/datagov_v4.py +798 -0
  193. mostlyright/data_harness/sources/catalog/harvest/protocol.py +964 -0
  194. mostlyright/data_harness/sources/catalog/harvest/sdmx.py +445 -0
  195. mostlyright/data_harness/sources/catalog/harvest/stac.py +384 -0
  196. mostlyright/data_harness/sources/catalog/health.py +447 -0
  197. mostlyright/data_harness/sources/catalog/hosted_catalog.py +105 -0
  198. mostlyright/data_harness/sources/catalog/identity_history.py +1549 -0
  199. mostlyright/data_harness/sources/catalog/neural.py +1618 -0
  200. mostlyright/data_harness/sources/catalog/packed_catalog.py +2345 -0
  201. mostlyright/data_harness/sources/catalog/packed_retrieval.py +1517 -0
  202. mostlyright/data_harness/sources/catalog/packed_writer.py +2802 -0
  203. mostlyright/data_harness/sources/catalog/query_trace.py +1037 -0
  204. mostlyright/data_harness/sources/catalog/recommend.py +171 -0
  205. mostlyright/data_harness/sources/catalog/retrieval.py +230 -0
  206. mostlyright/data_harness/sources/catalog/retrieval_manifest.py +995 -0
  207. mostlyright/data_harness/sources/catalog/rights_decisions.py +254 -0
  208. mostlyright/data_harness/sources/catalog/sealed.py +560 -0
  209. mostlyright/data_harness/sources/catalog/search.py +230 -0
  210. mostlyright/data_harness/sources/catalog/streaming_delta.py +1097 -0
  211. mostlyright/data_harness/sources/catalog/update.py +891 -0
  212. mostlyright/data_harness/sources/collections.py +815 -0
  213. mostlyright/data_harness/sources/contracts.py +2223 -0
  214. mostlyright/data_harness/sources/deletion.py +761 -0
  215. mostlyright/data_harness/sources/fitness.py +162 -0
  216. mostlyright/data_harness/sources/governance.py +163 -0
  217. mostlyright/data_harness/sources/hosted.py +173 -0
  218. mostlyright/data_harness/sources/integration.py +218 -0
  219. mostlyright/data_harness/sources/range_reader.py +418 -0
  220. mostlyright/data_harness/sources/registry.py +514 -0
  221. mostlyright/data_harness/sources/rights_rule.py +59 -0
  222. mostlyright/data_harness/sources/source_cadence_vectors.v1.json +1 -0
  223. mostlyright/data_harness/sources/sports.py +521 -0
  224. mostlyright/data_harness/sources/stream.py +524 -0
  225. mostlyright/data_harness/sources/stream_connector.py +418 -0
  226. mostlyright/data_harness/sources/stream_recorder.py +1404 -0
  227. mostlyright/data_harness/studio_boundary.py +2019 -0
  228. mostlyright/data_harness/thin/__init__.py +37 -0
  229. mostlyright/data_harness/thin/acquire.py +1137 -0
  230. mostlyright/data_harness/thin/acquire_cancel.py +579 -0
  231. mostlyright/data_harness/thin/approvals.py +617 -0
  232. mostlyright/data_harness/thin/commands.py +406 -0
  233. mostlyright/data_harness/thin/download.py +194 -0
  234. mostlyright/data_harness/thin/narrative.py +589 -0
  235. mostlyright/data_harness/thin/parity.py +1070 -0
  236. mostlyright/data_harness/thin/propose.py +2759 -0
  237. mostlyright/data_harness/thin/research.py +1663 -0
  238. mostlyright/data_harness/thin/router.py +924 -0
  239. mostlyright/data_harness/thin/runs.py +519 -0
  240. mostlyright/data_harness/thin/session.py +281 -0
  241. mostlyright/data_harness/thin/stream.py +501 -0
  242. mostlyright/data_harness/thin/transport.py +187 -0
  243. mostlyright/data_harness/thin/vocabulary.py +368 -0
  244. mostlyright/data_harness/thin/workers.py +164 -0
  245. mostlyright/data_harness/ucum/TABLE-PIN.json +40 -0
  246. mostlyright/data_harness/ucum/ucum-subset.v1.json +632 -0
  247. mostlyright/data_harness/unit_flow.py +927 -0
  248. mostlyright/data_harness/units.py +572 -0
  249. mostlyright/data_harness/ux/__init__.py +9 -0
  250. mostlyright/data_harness/ux/approve.py +485 -0
  251. mostlyright/data_harness/ux/author_yaml.py +597 -0
  252. mostlyright/data_harness/ux/cloud_auth.py +447 -0
  253. mostlyright/data_harness/ux/commands/__init__.py +260 -0
  254. mostlyright/data_harness/ux/commands/approve.py +136 -0
  255. mostlyright/data_harness/ux/commands/auth.py +744 -0
  256. mostlyright/data_harness/ux/commands/author.py +79 -0
  257. mostlyright/data_harness/ux/commands/catalog_author.py +403 -0
  258. mostlyright/data_harness/ux/commands/catalog_fill.py +523 -0
  259. mostlyright/data_harness/ux/commands/catalog_harvest.py +545 -0
  260. mostlyright/data_harness/ux/commands/catalog_publish.py +1838 -0
  261. mostlyright/data_harness/ux/commands/catalog_search.py +71 -0
  262. mostlyright/data_harness/ux/commands/catalog_update.py +437 -0
  263. mostlyright/data_harness/ux/commands/deploy.py +134 -0
  264. mostlyright/data_harness/ux/commands/deploy_dataset.py +98 -0
  265. mostlyright/data_harness/ux/commands/deploy_plan.py +105 -0
  266. mostlyright/data_harness/ux/commands/deploy_status.py +104 -0
  267. mostlyright/data_harness/ux/commands/diff.py +74 -0
  268. mostlyright/data_harness/ux/commands/index.py +84 -0
  269. mostlyright/data_harness/ux/commands/inventory.py +47 -0
  270. mostlyright/data_harness/ux/commands/list_builds.py +143 -0
  271. mostlyright/data_harness/ux/commands/login.py +63 -0
  272. mostlyright/data_harness/ux/commands/peek.py +236 -0
  273. mostlyright/data_harness/ux/commands/plan_check.py +90 -0
  274. mostlyright/data_harness/ux/commands/preflight.py +97 -0
  275. mostlyright/data_harness/ux/commands/record.py +107 -0
  276. mostlyright/data_harness/ux/commands/review_setup.py +47 -0
  277. mostlyright/data_harness/ux/commands/search.py +440 -0
  278. mostlyright/data_harness/ux/commands/show.py +61 -0
  279. mostlyright/data_harness/ux/commands/whoami.py +37 -0
  280. mostlyright/data_harness/ux/credential_native.py +551 -0
  281. mostlyright/data_harness/ux/credential_store.py +1055 -0
  282. mostlyright/data_harness/ux/credentials.py +631 -0
  283. mostlyright/data_harness/ux/diffing.py +444 -0
  284. mostlyright/data_harness/ux/headline.py +671 -0
  285. mostlyright/data_harness/ux/hosted_acquisition.py +974 -0
  286. mostlyright/data_harness/ux/hosted_run_status.py +619 -0
  287. mostlyright/data_harness/ux/inventory.py +427 -0
  288. mostlyright/data_harness/ux/local_review.py +375 -0
  289. mostlyright/data_harness/ux/login.py +691 -0
  290. mostlyright/data_harness/ux/path_kind.py +147 -0
  291. mostlyright/data_harness/ux/peek.py +1000 -0
  292. mostlyright/data_harness/ux/plain_file.py +178 -0
  293. mostlyright/data_harness/ux/plan_check.py +311 -0
  294. mostlyright/data_harness/ux/preflight.py +918 -0
  295. mostlyright/data_harness/ux/readers.py +1124 -0
  296. mostlyright/data_harness/ux/remediation.py +2195 -0
  297. mostlyright/data_harness/ux/render.py +657 -0
  298. mostlyright/data_harness/ux/workload.py +1077 -0
  299. mostlyright/data_harness/viewer.py +3713 -0
  300. mostlyright/data_harness/visual_run/__init__.py +83 -0
  301. mostlyright/data_harness/visual_run/authoring.py +235 -0
  302. mostlyright/data_harness/visual_run/contracts.py +673 -0
  303. mostlyright/data_harness/visual_run/materialize.py +486 -0
  304. mostlyright/data_harness/visual_run/observations.py +874 -0
  305. mostlyright/data_harness/visual_run/query.py +259 -0
  306. mostlyright/data_harness/visual_run/reducer.py +280 -0
  307. mostlyright/data_harness/visual_run/sdk.py +892 -0
  308. mostlyright/data_harness/visual_run/store.py +584 -0
  309. mostlyright/data_harness/visual_run/transport.py +239 -0
  310. mostlyright/data_harness/watch.py +2999 -0
  311. mostlyright_data-0.9.0.dist-info/METADATA +607 -0
  312. mostlyright_data-0.9.0.dist-info/RECORD +314 -0
  313. mostlyright_data-0.9.0.dist-info/WHEEL +4 -0
  314. mostlyright_data-0.9.0.dist-info/entry_points.txt +12 -0
@@ -0,0 +1,3554 @@
1
+ """The warm worker behind one research session (ADR 0021).
2
+
3
+ A session worker is not a new kind of worker. Studio's ``RESEARCH-SESSIONS.md`` is explicit about
4
+ it: *"A session worker IS a producer-role attempt on the session run"*, with ``worker_kind =
5
+ "session"`` and two extra capability scopes, so everything it does after the lease --
6
+ ``attempt:start``, batch ``event:append``, ``attempt:heartbeat``, ``attempt:fail`` -- is an existing
7
+ fenced producer callback with the existing workspace/run/attempt/generation/fence binding. Only
8
+ four verbs are new -- ``probes:claim``, ``probes/{id}:result``, and the ``session-acquisitions``
9
+ pair that gives a probe its bytes -- and this module is the whole of the harness half of that
10
+ protocol.
11
+
12
+ Read this module against :mod:`mostlyright.data_harness.hosted_bootstrap`, which establishes the
13
+ same workload-identity and bounded-secret handling discipline.
14
+
15
+ **What is kept.** Ambient workload identity is read once, from the metadata server, into a
16
+ ``bytearray`` that is zeroed in a ``finally``. The capability that comes back is never a ``str``
17
+ this module holds: it goes straight into a :class:`~mostlyright.data_harness.hosted_worker.
18
+ CredentialLease`, which is non-serializable, redacted in every repr, and zeroed on close. There is
19
+ exactly one capability per attempt and exactly one attempt per session, so "one capability per
20
+ attempt" survives verbatim. Nothing is written to disk: the lease response is parsed out of a
21
+ bounded in-memory buffer that is zeroed as soon as the capability is inside its lease.
22
+
23
+ **What is different, and why it has to be.** The generic bootstrap reads a dispatch descriptor from
24
+ a file descriptor because producer and verifier jobs carry sealed job descriptions. A research
25
+ session Job instead receives three non-secret arguments and uses the atomically claimed dispatch
26
+ record in Studio as its bootstrap authority. A capability-free monitor starts an isolated child;
27
+ only that child exchanges and retains the capability in ``CredentialLease`` for the session
28
+ lifetime, never in argv, environment or a descriptor visible to another task.
29
+
30
+ **Request lifetime, and how scale-to-zero actually happens.** This is a Cloud Run *Job*, not a
31
+ service. Studio supplies only an opaque UUID4 dispatch id, an exact Studio exchange URL and the OIDC
32
+ audience as task arguments. The task reads its ambient workload identity, atomically exchanges the
33
+ dispatch id for the one fenced session lease, and runs :class:`SessionSupervisor` synchronously to
34
+ its terminal state. The Job task is therefore the unit Cloud Run retains: an overlapping session
35
+ gets a distinct task and cannot share a warm service instance or an in-memory capability with the
36
+ first.
37
+
38
+ The dispatch id is deliberately not a lease nonce. It is non-secret routing data, accepted exactly
39
+ once by Studio's transaction that creates and fences the attempt. The capability appears only in the
40
+ exchange response, immediately enters ``CredentialLease``, and is zeroed when the supervisor exits.
41
+ Job retry policy and task timeout are deployment controls, not a second leasing protocol: a retry
42
+ of an already-claimed dispatch receives Studio's typed duplicate/terminal refusal and must not run a
43
+ second supervisor.
44
+
45
+ **Exit is by stopping, never by completing.** There is no ``:complete`` here. Studio's contract
46
+ says why: ``:complete`` requires an admitted candidate, which a session never has, and after a
47
+ close the session run is already terminal so ``:fail`` cannot land either. ``:fail`` is reserved
48
+ for the one case where it is still meaningful -- the worker dying while its session is still live.
49
+
50
+ **Heartbeat is liveness, never activity.** It runs on its own timer thread and it deliberately does
51
+ not touch ``last_activity_at`` server-side. That is Studio's rule and the reason is the bill: a
52
+ worker heartbeats every thirty seconds, so counting a heartbeat as activity would mean no session
53
+ ever goes idle and scale-to-session never happens. Nothing in this module tries to keep a session
54
+ alive; keeping a session alive is the user's job, by using it.
55
+ """
56
+
57
+ from __future__ import annotations
58
+
59
+ import contextlib
60
+ import hashlib
61
+ import json
62
+ import logging
63
+ import math
64
+ import os
65
+ import re
66
+ import select
67
+ import signal
68
+ import subprocess
69
+ import sys
70
+ import tempfile
71
+ import threading
72
+ import time
73
+ from collections.abc import Callable, Mapping, Sequence
74
+ from dataclasses import dataclass, field, replace
75
+ from datetime import UTC, datetime
76
+ from http import HTTPStatus
77
+ from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
78
+ from pathlib import Path
79
+ from typing import Any, ClassVar, Protocol
80
+ from uuid import UUID
81
+
82
+ from mostlyright.data_harness import progress_events, session_probes
83
+ from mostlyright.data_harness.acquisition.parsing import ParseLimits
84
+ from mostlyright.data_harness.acquisition.result_download import (
85
+ ResultDownloadError,
86
+ validate_result_download,
87
+ verify_result_bytes,
88
+ )
89
+ from mostlyright.data_harness.canonical import canonical_json_bytes, sha256_bytes
90
+ from mostlyright.data_harness.hosted_bootstrap import (
91
+ BootstrapTransport,
92
+ UrlLibBootstrapTransport,
93
+ _https_url,
94
+ _metadata_identity,
95
+ )
96
+ from mostlyright.data_harness.hosted_crawler_protocol import (
97
+ HostedCrawlerProtocolError,
98
+ parse_crawler_result,
99
+ )
100
+ from mostlyright.data_harness.hosted_worker import (
101
+ MAX_CAPABILITY_BYTES,
102
+ CredentialLease,
103
+ HostedWorkerError,
104
+ _authenticated_client,
105
+ _idempotency,
106
+ _load_generated,
107
+ _new_transfer_client,
108
+ _ProducerAuthority,
109
+ _response,
110
+ _strict_json,
111
+ parse_capability,
112
+ )
113
+ from mostlyright.data_harness.linux_process_boundary import (
114
+ _child_subreaper_enabled,
115
+ _set_child_subreaper,
116
+ )
117
+
118
+ LOGGER = logging.getLogger(__name__)
119
+
120
+ #: The role this worker presents. Studio mints the attempt with ``worker_kind = "session"``; the
121
+ #: value is restated here only for the progress record that names the role.
122
+ SESSION_WORKER_ROLE = "session"
123
+
124
+ #: The one path Studio invokes, spelled by ``GoogleCloudRunSessionWorkerStarter``.
125
+ START_PATH = "/sessions:start"
126
+ #: An internal, descriptor-only command used by the retained Service rollback path. The service
127
+ #: process starts a fresh interpreter rather than forking Python from a request thread; the only
128
+ #: value passed in argv is an inherited descriptor number, never a nonce or capability.
129
+ LEGACY_BOOTSTRAP_FD_ARGUMENT = "--legacy-bootstrap-fd"
130
+ #: Liveness only. Cloud Run's startup probe needs an endpoint that answers before a session exists,
131
+ #: and answering it must never imply a session is warm -- ``warming`` is Studio's state to clear,
132
+ #: on a fenced attempt callback, and a 200 from here is not one.
133
+ HEALTH_PATHS = ("/livez", "/healthz")
134
+
135
+ #: Studio's start invocation carries its own contract version.
136
+ START_SCHEMA_VERSION = "3.0.0"
137
+ #: Cloud Run Job dispatches use the same schema as the lease response they mint. The dispatch
138
+ #: itself is deliberately only an opaque routing token; it is never a capability or lease nonce.
139
+ SESSION_DISPATCH_SCHEMA_VERSION = "3.0.0"
140
+ #: The only Studio URL a Job task may exchange its dispatch id against. The id is part of the
141
+ #: path so a task cannot turn an opaque value into a request to a neighbouring session.
142
+ SESSION_DISPATCH_EXCHANGE_PATH_TEMPLATE = (
143
+ "/internal/v3/session-worker/dispatches/{dispatch_id}:exchange"
144
+ )
145
+ #: Studio refuses an exchange after its one dispatch has already been consumed. These are normal
146
+ #: no-op terminal outcomes for a Cloud Run Job task, not failed executions: the task must not
147
+ #: invoke ``attempt:start`` after another task has claimed it or after Studio has closed it.
148
+ CONSUMED_SESSION_DISPATCH_CODES = frozenset(
149
+ {
150
+ "RESEARCH_SESSION_DISPATCH_ALREADY_CLAIMED",
151
+ "RESEARCH_SESSION_DISPATCH_CLOSED",
152
+ }
153
+ )
154
+ #: The lease path Studio publishes. Derived here rather than taken from the invocation; see
155
+ #: :func:`parse_start_invocation`.
156
+ LEASE_PATH_TEMPLATE = "/internal/v3/session-worker/sessions/{session_id}:lease"
157
+
158
+ #: A start invocation is five short fields. Anything larger is not one.
159
+ MAX_START_REQUEST_BYTES = 8 * 1024
160
+ #: A lease has no source bytes and must remain a small control-plane document. This separate
161
+ #: bound means a compromised exchange cannot make a long-lived Job hold an arbitrary response.
162
+ MAX_SESSION_DISPATCH_RESPONSE_BYTES = 64 * 1024
163
+ #: The service-to-fresh-interpreter start document is an already validated start request and must
164
+ #: stay no larger than the public request surface. The child accepts exactly one such document.
165
+ MAX_LEGACY_BOOTSTRAP_BYTES = MAX_START_REQUEST_BYTES
166
+ #: Service start must not acknowledge until its fresh monitor has installed the cancellation
167
+ #: handler that owns an inner capability child. This remains far below Studio's start request
168
+ #: timeout while rejecting a stuck/interrupted bootstrap rather than accepting a zombie session.
169
+ LEGACY_BOOTSTRAP_READY_TIMEOUT_SECONDS = 5.0
170
+ #: Studio bounds the nonce at 32..128 URL-safe characters.
171
+ MIN_NONCE_CHARS = 32
172
+ MAX_NONCE_CHARS = 128
173
+
174
+ #: How long the process waits for another session before exiting so Cloud Run can drop the
175
+ #: instance. Long enough that reopening a session is warm, short enough that a finished session is
176
+ #: not billed for a coffee break.
177
+ DEFAULT_SHUTDOWN_GRACE_SECONDS = 30.0
178
+ #: Pre-spawned no-network sandbox children. This is the first production caller of the warm pool
179
+ #: (#306), and the pool exists precisely for this: the follow-up-probe budget is 100-300 ms and a
180
+ #: cold interpreter spawn does not fit in it. On Linux under a bounded cgroup root
181
+ #: ``start_warm_pool`` returns 0 by design -- the count is read, not assumed, and a zero is not an
182
+ #: error, only a slower first probe.
183
+ DEFAULT_WARM_POOL_SIZE = 2
184
+ #: A probe is interactive. A probe that is not finished in this long is failed, not waited on.
185
+ DEFAULT_PROBE_TIMEOUT_SECONDS = 45.0
186
+ #: How long one session acquisition may take from the create call to a settled status. This is the
187
+ #: network half of a probe and it is bounded SEPARATELY from the Clean room ceiling above, because
188
+ #: the two fail for unrelated reasons and a client that is told "the crawl did not come back" can
189
+ #: act on that differently from "parsing did not finish".
190
+ DEFAULT_ACQUISITION_TIMEOUT_SECONDS = 90.0
191
+ #: How often the worker asks Studio whether its acquisition has settled. Fixed rather than backed
192
+ #: off: the whole wait is bounded above by one interactive ceiling, so a backoff would only make
193
+ #: the common fast case slower without making the slow case cheaper.
194
+ ACQUISITION_POLL_SECONDS = 1.0
195
+ #: The largest quarantine bundle a probe will pull into memory. The contract's own ceiling is 6 GiB
196
+ #: -- correct for a Build, absurd for a preview -- and the bundle is base64 inside canonical JSON,
197
+ #: so this is comfortably above `ProbeExecutor._max_source_bytes` (16 MiB) times the 4/3 encoding
198
+ #: expansion plus the receipt. A bundle over it is refused on its DECLARED size, before a byte is
199
+ #: transferred, which is the only place refusing is free.
200
+ MAX_PROBE_ACQUISITION_BUNDLE_BYTES = 32 * 1024 * 1024
201
+ #: Poll and heartbeat floors, used only when a lease omits them (it cannot -- both are required --
202
+ #: but a default beats a KeyError if Studio ever relaxes that).
203
+ FALLBACK_CLAIM_POLL_SECONDS = 1.0
204
+ FALLBACK_HEARTBEAT_SECONDS = 30.0
205
+ #: A refused claim poll is transient far more often than it is fatal. After this many in a row the
206
+ #: worker gives up on the session rather than polling a Studio that is saying no forever.
207
+ MAX_CONSECUTIVE_CLAIM_FAILURES = 10
208
+ #: Studio only lets a session shorten the ordinary thirty-minute producer lease. A later value
209
+ #: would make a compromised or mismatched control plane turn one session capability into an
210
+ #: unbounded worker lifetime, so it is a lease refusal rather than a value to clamp.
211
+ MAX_SESSION_CAPABILITY_REMAINING_SECONDS = 30 * 60
212
+ #: The isolated session child communicates only an expiry, identity and terminal summary to its
213
+ #: capability-free parent. It has no reason to emit a large record, and this ceiling makes the
214
+ #: parent-side parser a fixed-memory boundary if a child is compromised.
215
+ MAX_SESSION_SUPERVISOR_MESSAGE_BYTES = 128 * 1024
216
+
217
+ _UTC_TIMESTAMP = re.compile(
218
+ r"^[0-9]{4}-[0-9]{2}-[0-9]{2}T[0-9]{2}:[0-9]{2}:[0-9]{2}(?:\.[0-9]{1,6})?Z$"
219
+ )
220
+
221
+ #: The failure codes that mean the probe was REFUSED rather than attempted-and-broken. The
222
+ #: distinction is the only part of a failure a client can act on: a refused probe refuses the same
223
+ #: way every time, so the fix is a different probe -- or, for the source channel, somebody wiring
224
+ #: one -- while a broken probe may well not break twice. The set is written out rather than
225
+ #: inferred from the exception class, because the same class carries both halves.
226
+ PROBE_REFUSAL_CODES = frozenset(
227
+ {
228
+ # The open gap: no session-scoped source channel exists, so nothing was read.
229
+ "PROBE_SOURCE_CHANNEL_UNAVAILABLE",
230
+ # The session's own authority check, made before any read.
231
+ "PROBE_SOURCE_OUT_OF_SCOPE",
232
+ # A reader that answered for a different source. Refused before the bytes are used.
233
+ "PROBE_SOURCE_MISMATCH",
234
+ # The request body is not in the probe vocabulary.
235
+ "PROBE_REQUEST_INVALID",
236
+ # The request names a column, or an aggregate over a column, this source cannot answer.
237
+ # Discovered after the read rather than before it, and still a refusal of the REQUEST:
238
+ # re-asking it changes nothing.
239
+ "PROBE_COLUMN_UNKNOWN",
240
+ "PROBE_AGGREGATE_NOT_NUMERIC",
241
+ # Studio's own refusals on the session-acquisition door. The first is the session's
242
+ # source allowlist, checked before any lookup so a refusal is not an existence oracle;
243
+ # the second is "this source has no single active public.https registration", which is a
244
+ # fact about the source's registration and not about this probe.
245
+ "RESEARCH_SOURCE_FORBIDDEN",
246
+ "RESEARCH_SOURCE_UNACQUIRABLE",
247
+ # The acquired source is larger than a probe may hold. A property of the source.
248
+ "PROBE_SOURCE_TOO_LARGE",
249
+ }
250
+ )
251
+ #: The Clean room's own code for "did not finish in time" (`acquisition.sandbox`). It is distinct
252
+ #: from a source acquisition timing out: reducing a scan can help after bytes arrived, but cannot
253
+ #: make a delayed remote acquisition arrive sooner.
254
+ PROBE_TIMEOUT_CODES = frozenset({"SANDBOX_TIMEOUT", "PROBE_SOURCE_ACQUISITION_TIMEOUT"})
255
+ #: Studio's code for an acquisition whose own expiry passed -- the earliest of the crawler-session
256
+ #: TTL, the deployment's admission window and the worker's attempt lease. It is the one failure the
257
+ #: reader retries, because Studio releases the admission slot as it expires the acquisition, so
258
+ #: opening a fresh one is not a second spend. See :meth:`StudioSessionAcquisitionReader.read`.
259
+ ACQUISITION_EXPIRED_CODE = "CRAWLER_SESSION_EXPIRED"
260
+ #: A ceiling somebody else is holding, not a refusal of this request. Studio charges a session
261
+ #: acquisition to the principal who opened the session, against the same active-and-rate admission
262
+ #: limits a person acquiring by hand would spend (ADR 0015: two at once, four in sixty seconds), so
263
+ #: this arrives when the session is busy rather than when the probe is wrong. It is its own outcome
264
+ #: because the advice it carries is the opposite of a refusal's: ask again in a moment.
265
+ SESSION_ACQUISITION_THROTTLED = "SESSION_ACQUISITION_THROTTLED"
266
+ #: Studio's wire code is retained in the probe result. The progress/UI code is separate so an
267
+ #: acquisition admission window is not presented as model-token, billing, sign-in, or source-key
268
+ #: configuration trouble.
269
+ PROBE_THROTTLE_WIRE_CODES = frozenset({"AGENT_BUDGET_EXCEEDED"})
270
+ PROBE_THROTTLE_CODES = frozenset({SESSION_ACQUISITION_THROTTLED})
271
+ MAX_RETRY_AFTER_SECONDS = 3600
272
+
273
+ _ENV_PORT = "PORT"
274
+ _ENV_STUDIO_ORIGIN = "STUDIO_API_ORIGIN"
275
+ _ENV_LEASE_AUDIENCE = "MR_SESSION_WORKER_LEASE_AUDIENCE"
276
+ _ENV_WORKER_BOOTSTRAP_AUDIENCE = "STUDIO_WORKER_BOOTSTRAP_AUDIENCE"
277
+ _ENV_SHUTDOWN_GRACE = "MR_SESSION_WORKER_SHUTDOWN_GRACE_SECONDS"
278
+ _ENV_WARM_POOL = "MR_SESSION_WORKER_WARM_POOL_SIZE"
279
+ _ENV_PROBE_TIMEOUT = "MR_SESSION_WORKER_PROBE_TIMEOUT_SECONDS"
280
+ _ENV_ACQUISITION_TIMEOUT = "MR_SESSION_WORKER_ACQUISITION_TIMEOUT_SECONDS"
281
+ _ENV_EXTERNAL_NETWORK_POLICY = "MOSTLYRIGHT_EXTERNAL_NETWORK_POLICY_ATTESTATION"
282
+ _ENV_CRAWLER_RESULTS_BUCKET = "MOSTLYRIGHT_CRAWLER_RESULTS_BUCKET"
283
+
284
+
285
+ class SessionWorkerError(RuntimeError):
286
+ """A coded session-worker refusal. The code is what a probe or a session fails with."""
287
+
288
+ def __init__(self, code: str, detail: str, *, retry_after_seconds: int | None = None) -> None:
289
+ super().__init__(f"{code}: {detail}")
290
+ self.code = code
291
+ self.detail = detail
292
+ self.retry_after_seconds = retry_after_seconds
293
+
294
+
295
+ @dataclass(frozen=True)
296
+ class SessionDispatch:
297
+ """The complete, non-secret argument contract of one Cloud Run Job task.
298
+
299
+ Nothing from the resulting lease belongs here. In particular, keeping this object safe to
300
+ log is a guard against a future Terraform change accidentally moving a capability or nonce
301
+ into command arguments, which Cloud Run exposes as execution configuration.
302
+ """
303
+
304
+ dispatch_id: str
305
+ exchange_url: str
306
+ audience: str
307
+
308
+
309
+ # ==========================================================================================
310
+ # Configuration
311
+ # ==========================================================================================
312
+
313
+
314
+ @dataclass(frozen=True)
315
+ class SessionWorkerConfig:
316
+ """Everything the service reads from its environment, validated once at start."""
317
+
318
+ studio_origin: str
319
+ lease_audience: str
320
+ port: int = 8080
321
+ shutdown_grace_seconds: float = DEFAULT_SHUTDOWN_GRACE_SECONDS
322
+ warm_pool_size: int = DEFAULT_WARM_POOL_SIZE
323
+ probe_timeout_seconds: float = DEFAULT_PROBE_TIMEOUT_SECONDS
324
+ acquisition_timeout_seconds: float = DEFAULT_ACQUISITION_TIMEOUT_SECONDS
325
+ external_network_policy_attestation: str | None = None
326
+ crawler_results_bucket: str | None = None
327
+
328
+
329
+ def _positive_number(raw: str | None, default: float, label: str, *, maximum: float) -> float:
330
+ if raw is None or raw == "":
331
+ return default
332
+ try:
333
+ value = float(raw)
334
+ except ValueError as error:
335
+ raise SessionWorkerError("SESSION_WORKER_CONFIG_INVALID", f"{label} is not a number") from (
336
+ error
337
+ )
338
+ if not 0 < value <= maximum:
339
+ raise SessionWorkerError(
340
+ "SESSION_WORKER_CONFIG_INVALID", f"{label} must be in (0, {maximum}]"
341
+ )
342
+ return value
343
+
344
+
345
+ def config_from_environment(environment: Mapping[str, str] | None = None) -> SessionWorkerConfig:
346
+ """Build the config, refusing rather than defaulting anything that is a coordinate."""
347
+
348
+ env = os.environ if environment is None else environment
349
+ origin = env.get(_ENV_STUDIO_ORIGIN, "")
350
+ if not origin:
351
+ raise SessionWorkerError(
352
+ "SESSION_WORKER_CONFIG_INVALID", f"{_ENV_STUDIO_ORIGIN} is required"
353
+ )
354
+ origin = origin.rstrip("/")
355
+ _https_url(origin, _ENV_STUDIO_ORIGIN)
356
+ # Studio verifies the session worker against the same closed workload-identity audience as the
357
+ # producer bootstrap exchange. The shared Cloud Run template already supplies that coordinate;
358
+ # consuming it here keeps the identity token's ``aud`` exact while preserving the explicit
359
+ # Harness override for a deliberately split deployment. An older deployment without the shared
360
+ # coordinate remains compatible with its historical origin audience.
361
+ audience = (
362
+ env.get(_ENV_LEASE_AUDIENCE, "") or env.get(_ENV_WORKER_BOOTSTRAP_AUDIENCE, "") or origin
363
+ )
364
+ _https_url(
365
+ audience,
366
+ (
367
+ _ENV_LEASE_AUDIENCE
368
+ if env.get(_ENV_LEASE_AUDIENCE, "")
369
+ else _ENV_WORKER_BOOTSTRAP_AUDIENCE
370
+ if env.get(_ENV_WORKER_BOOTSTRAP_AUDIENCE, "")
371
+ else _ENV_STUDIO_ORIGIN
372
+ ),
373
+ )
374
+ raw_port = env.get(_ENV_PORT, "8080")
375
+ try:
376
+ port = int(raw_port)
377
+ except ValueError as error:
378
+ raise SessionWorkerError("SESSION_WORKER_CONFIG_INVALID", "PORT is not an integer") from (
379
+ error
380
+ )
381
+ if not 1 <= port <= 65535:
382
+ raise SessionWorkerError("SESSION_WORKER_CONFIG_INVALID", "PORT is out of range")
383
+ raw_pool = env.get(_ENV_WARM_POOL, "")
384
+ try:
385
+ warm_pool = DEFAULT_WARM_POOL_SIZE if raw_pool == "" else int(raw_pool)
386
+ except ValueError as error:
387
+ raise SessionWorkerError(
388
+ "SESSION_WORKER_CONFIG_INVALID", f"{_ENV_WARM_POOL} is not an integer"
389
+ ) from error
390
+ if not 0 <= warm_pool <= 8:
391
+ raise SessionWorkerError("SESSION_WORKER_CONFIG_INVALID", f"{_ENV_WARM_POOL} must be 0..8")
392
+ return SessionWorkerConfig(
393
+ studio_origin=origin,
394
+ lease_audience=audience,
395
+ port=port,
396
+ shutdown_grace_seconds=_positive_number(
397
+ env.get(_ENV_SHUTDOWN_GRACE),
398
+ DEFAULT_SHUTDOWN_GRACE_SECONDS,
399
+ _ENV_SHUTDOWN_GRACE,
400
+ maximum=600.0,
401
+ ),
402
+ warm_pool_size=warm_pool,
403
+ probe_timeout_seconds=_positive_number(
404
+ env.get(_ENV_PROBE_TIMEOUT),
405
+ DEFAULT_PROBE_TIMEOUT_SECONDS,
406
+ _ENV_PROBE_TIMEOUT,
407
+ maximum=300.0,
408
+ ),
409
+ acquisition_timeout_seconds=_positive_number(
410
+ env.get(_ENV_ACQUISITION_TIMEOUT),
411
+ DEFAULT_ACQUISITION_TIMEOUT_SECONDS,
412
+ _ENV_ACQUISITION_TIMEOUT,
413
+ maximum=600.0,
414
+ ),
415
+ external_network_policy_attestation=env.get(_ENV_EXTERNAL_NETWORK_POLICY) or None,
416
+ # The same name the hosted producer reads, because it is the same quarantine bucket and a
417
+ # second spelling would be a second thing to get wrong on one deployment. Absent means a
418
+ # Studio-local result authority, which `validate_result_download` then requires.
419
+ crawler_results_bucket=env.get(_ENV_CRAWLER_RESULTS_BUCKET) or None,
420
+ )
421
+
422
+
423
+ # ==========================================================================================
424
+ # Cloud Run Job dispatch bootstrap
425
+ # ==========================================================================================
426
+
427
+
428
+ def _dispatch_uuid(value: Any) -> str:
429
+ if not isinstance(value, str):
430
+ raise SessionWorkerError("SESSION_DISPATCH_INVALID", "dispatch id must be a string")
431
+ try:
432
+ parsed = UUID(value)
433
+ except ValueError as error:
434
+ raise SessionWorkerError("SESSION_DISPATCH_INVALID", "dispatch id is not a UUID") from error
435
+ if parsed.version != 4 or str(parsed) != value:
436
+ raise SessionWorkerError(
437
+ "SESSION_DISPATCH_INVALID", "dispatch id must be a canonical UUID4"
438
+ )
439
+ return value
440
+
441
+
442
+ def _dispatch_https_url(value: Any, label: str) -> str:
443
+ try:
444
+ return _https_url(value, label)
445
+ except HostedWorkerError as error:
446
+ raise SessionWorkerError("SESSION_DISPATCH_INVALID", error.detail) from error
447
+
448
+
449
+ def parse_session_dispatch(argv: Sequence[str]) -> SessionDispatch:
450
+ """Parse exactly the three non-secret Job arguments, in their deployment order.
451
+
452
+ A normal argument parser intentionally permits flags in any order and can silently admit
453
+ future options. Job overrides are a security boundary, so this parser instead spells the
454
+ only six argv cells Studio may create. It also keeps the invocation object incapable of
455
+ carrying a lease nonce or capability.
456
+ """
457
+
458
+ values = tuple(argv)
459
+ flags = (
460
+ "--session-dispatch-id",
461
+ "--session-dispatch-exchange-url",
462
+ "--session-dispatch-audience",
463
+ )
464
+ if len(values) != 6 or values[::2] != flags or any(not item for item in values[1::2]):
465
+ raise SessionWorkerError(
466
+ "SESSION_DISPATCH_ARGUMENTS_INVALID",
467
+ "expected --session-dispatch-id, --session-dispatch-exchange-url and "
468
+ "--session-dispatch-audience in that order",
469
+ )
470
+ dispatch_id, exchange_url, audience = values[1::2]
471
+ return SessionDispatch(
472
+ dispatch_id=_dispatch_uuid(dispatch_id),
473
+ exchange_url=_dispatch_https_url(exchange_url, "session dispatch exchange URL"),
474
+ audience=_dispatch_https_url(audience, "session dispatch audience"),
475
+ )
476
+
477
+
478
+ def _validate_session_dispatch(config: SessionWorkerConfig, dispatch: SessionDispatch) -> None:
479
+ """Bind a Job's opaque id to this deployment's one exchange endpoint.
480
+
481
+ The OIDC audience is intentionally supplied by Studio's Job override, but the URL itself must
482
+ remain derived from the locally pinned Studio origin. This closes request-target substitution
483
+ even if an operator can create an execution override.
484
+ """
485
+
486
+ expected = config.studio_origin + SESSION_DISPATCH_EXCHANGE_PATH_TEMPLATE.format(
487
+ dispatch_id=dispatch.dispatch_id
488
+ )
489
+ if dispatch.exchange_url != expected:
490
+ raise SessionWorkerError(
491
+ "SESSION_DISPATCH_URL_UNRECOGNIZED",
492
+ "session dispatch exchange URL is not this deployment's endpoint for this dispatch",
493
+ )
494
+ if dispatch.audience != config.studio_origin:
495
+ raise SessionWorkerError(
496
+ "SESSION_DISPATCH_AUDIENCE_UNRECOGNIZED",
497
+ "session dispatch audience is not this deployment's Studio audience",
498
+ )
499
+
500
+
501
+ def _dispatch_exchange_error(status: int, raw: bytearray) -> SessionWorkerError:
502
+ """Keep Studio's typed refusal without echoing a response that could contain a secret."""
503
+
504
+ try:
505
+ value = _strict_json(raw, "session dispatch exchange refusal")
506
+ except HostedWorkerError:
507
+ value = None
508
+ if isinstance(value, dict):
509
+ code = value.get("code")
510
+ if isinstance(code, str) and re.fullmatch(r"[A-Z][A-Z0-9_]{2,127}", code):
511
+ return SessionWorkerError(code, f"session dispatch exchange returned HTTP {status}")
512
+ return SessionWorkerError(
513
+ "SESSION_DISPATCH_EXCHANGE_FAILED", f"session dispatch exchange returned HTTP {status}"
514
+ )
515
+
516
+
517
+ def exchange_session_dispatch(
518
+ *,
519
+ config: SessionWorkerConfig,
520
+ dispatch: SessionDispatch,
521
+ transport: BootstrapTransport | None = None,
522
+ ) -> tuple[SessionLease, CredentialLease]:
523
+ """Redeem one dispatch through Studio and return the sole fenced session lease.
524
+
525
+ The metadata identity proves the Job's workload principal. The body is deliberately exact and
526
+ contains no secret: Studio's atomically claimed dispatch record is what binds that principal to
527
+ the resulting attempt. An already claimed or terminal dispatch is a typed non-200 response;
528
+ callers must treat it as terminal rather than retrying into a second supervisor.
529
+ """
530
+
531
+ _validate_session_dispatch(config, dispatch)
532
+ active_transport = transport or UrlLibBootstrapTransport()
533
+ identity = _metadata_identity(active_transport, dispatch.audience)
534
+ payload = canonical_json_bytes(
535
+ {
536
+ "schema_version": SESSION_DISPATCH_SCHEMA_VERSION,
537
+ "dispatch_id": dispatch.dispatch_id,
538
+ }
539
+ )
540
+ try:
541
+ token = identity.decode("ascii")
542
+ status, response = active_transport.request(
543
+ "POST",
544
+ dispatch.exchange_url,
545
+ {
546
+ "Authorization": f"Bearer {token}",
547
+ "Content-Type": "application/json",
548
+ },
549
+ payload,
550
+ MAX_SESSION_DISPATCH_RESPONSE_BYTES,
551
+ )
552
+ finally:
553
+ identity[:] = b"\x00" * len(identity)
554
+ raw = bytearray(response)
555
+ try:
556
+ if status != HTTPStatus.OK:
557
+ raise _dispatch_exchange_error(status, raw)
558
+ try:
559
+ lease_wire = _strict_json(raw, "session dispatch exchange response")
560
+ except HostedWorkerError as error:
561
+ raise SessionWorkerError(
562
+ "SESSION_DISPATCH_RESPONSE_INVALID", "session dispatch exchange response is invalid"
563
+ ) from error
564
+ try:
565
+ return parse_session_lease(lease_wire)
566
+ finally:
567
+ if isinstance(lease_wire, dict):
568
+ lease_wire.clear()
569
+ finally:
570
+ raw[:] = b"\x00" * len(raw)
571
+
572
+
573
+ # ==========================================================================================
574
+ # The start invocation
575
+ # ==========================================================================================
576
+
577
+
578
+ @dataclass(frozen=True)
579
+ class StartInvocation:
580
+ """The four non-secret coordinates of one start signal. The nonce is never in here."""
581
+
582
+ session_id: str
583
+ workspace_id: str
584
+ lease_url: str
585
+
586
+
587
+ def _uuid_text(value: Any, label: str) -> str:
588
+ if not isinstance(value, str):
589
+ raise SessionWorkerError("SESSION_START_INVALID", f"{label} must be a string")
590
+ try:
591
+ parsed = UUID(value)
592
+ except ValueError as error:
593
+ raise SessionWorkerError("SESSION_START_INVALID", f"{label} is not a UUID") from error
594
+ if str(parsed) != value:
595
+ raise SessionWorkerError("SESSION_START_INVALID", f"{label} is not canonical")
596
+ return value
597
+
598
+
599
+ def parse_start_invocation(raw: bytes, *, studio_origin: str) -> tuple[StartInvocation, bytearray]:
600
+ """Admit one start invocation and separate the nonce from everything else.
601
+
602
+ The nonce comes back as a ``bytearray`` rather than a ``str`` so the caller can zero it, and it
603
+ is never a field on :class:`StartInvocation` so that logging or repr'ing the invocation cannot
604
+ print it.
605
+
606
+ ``lease_url`` is **checked, never followed**. The URL arrives inside a request body, and a
607
+ request body is data: a worker that fetched whatever URL it was handed -- carrying a Google
608
+ identity token minted for the audience it was configured with -- would be an SSRF primitive
609
+ with a credential attached, reachable by anything that could reach this port. The value is
610
+ therefore compared against the URL derived from ``STUDIO_API_ORIGIN`` and the session id, and a
611
+ mismatch is a refusal. The call itself is then made through the generated client bound to
612
+ ``STUDIO_API_ORIGIN``, so even a check that somehow passed could not redirect the request.
613
+ """
614
+
615
+ if len(raw) > MAX_START_REQUEST_BYTES:
616
+ raise SessionWorkerError("SESSION_START_TOO_LARGE", "start invocation exceeds its budget")
617
+ try:
618
+ value = json.loads(raw.decode("utf-8"))
619
+ except (UnicodeDecodeError, ValueError) as error:
620
+ raise SessionWorkerError("SESSION_START_INVALID", "start invocation is not JSON") from error
621
+ if not isinstance(value, dict):
622
+ raise SessionWorkerError("SESSION_START_INVALID", "start invocation must be an object")
623
+ expected = {"schema_version", "session_id", "workspace_id", "lease_nonce", "lease_url"}
624
+ if set(value) != expected:
625
+ raise SessionWorkerError("SESSION_START_INVALID", "start invocation fields are not exact")
626
+ if value["schema_version"] != START_SCHEMA_VERSION:
627
+ raise SessionWorkerError("SESSION_START_INVALID", "start invocation schema is unsupported")
628
+ session_id = _uuid_text(value["session_id"], "session_id")
629
+ workspace_id = _uuid_text(value["workspace_id"], "workspace_id")
630
+ nonce = value["lease_nonce"]
631
+ if (
632
+ not isinstance(nonce, str)
633
+ or not MIN_NONCE_CHARS <= len(nonce) <= MAX_NONCE_CHARS
634
+ or not all(character.isalnum() or character in "_-" for character in nonce)
635
+ ):
636
+ raise SessionWorkerError("SESSION_START_INVALID", "lease nonce shape is invalid")
637
+ lease_url = value["lease_url"]
638
+ if not isinstance(lease_url, str):
639
+ raise SessionWorkerError("SESSION_START_INVALID", "lease_url must be a string")
640
+ derived = studio_origin.rstrip("/") + LEASE_PATH_TEMPLATE.format(session_id=session_id)
641
+ if lease_url != derived:
642
+ raise SessionWorkerError(
643
+ "SESSION_LEASE_URL_UNRECOGNIZED",
644
+ "lease_url is not this deployment's lease endpoint for this session",
645
+ )
646
+ return (
647
+ StartInvocation(session_id=session_id, workspace_id=workspace_id, lease_url=derived),
648
+ bytearray(nonce.encode("ascii")),
649
+ )
650
+
651
+
652
+ # ==========================================================================================
653
+ # The lease
654
+ # ==========================================================================================
655
+
656
+
657
+ @dataclass(frozen=True)
658
+ class SessionContext:
659
+ """The workspace coordinate the session probes against. Also the probe authority allowlist."""
660
+
661
+ dataset_id: str
662
+ question_id: str
663
+ table_id: str | None
664
+ source_ids: frozenset[str]
665
+
666
+
667
+ @dataclass(frozen=True)
668
+ class SessionLease:
669
+ """Everything the lease said except the capability, which lives in a ``CredentialLease``."""
670
+
671
+ session_id: str
672
+ workspace_id: UUID
673
+ run_id: UUID
674
+ attempt_id: UUID
675
+ generation: int
676
+ fence: int
677
+ capability_expires_at: datetime
678
+ callback_base_url: str
679
+ allowed_callback_paths: tuple[str, ...]
680
+ heartbeat_interval_seconds: float
681
+ claim_poll_interval_seconds: float
682
+ idle_timeout_seconds: float
683
+ context: SessionContext
684
+
685
+
686
+ def _positive_int(value: Any, label: str) -> int:
687
+ if type(value) is not int or value <= 0 or value > 2**53 - 1:
688
+ raise SessionWorkerError("SESSION_LEASE_INVALID", f"{label} is not a positive integer")
689
+ return value
690
+
691
+
692
+ def _parse_utc_timestamp(value: Any, label: str) -> datetime:
693
+ """Parse Studio's canonical UTC instant without accepting a local-time interpretation."""
694
+
695
+ if not isinstance(value, str) or _UTC_TIMESTAMP.fullmatch(value) is None:
696
+ raise SessionWorkerError(
697
+ "SESSION_LEASE_INVALID", f"{label} must be a canonical UTC timestamp"
698
+ )
699
+ try:
700
+ pattern = "%Y-%m-%dT%H:%M:%S.%fZ" if "." in value else "%Y-%m-%dT%H:%M:%SZ"
701
+ return datetime.strptime(value, pattern).replace(tzinfo=UTC)
702
+ except ValueError as error:
703
+ raise SessionWorkerError(
704
+ "SESSION_LEASE_INVALID", f"{label} is not a real UTC timestamp"
705
+ ) from error
706
+
707
+
708
+ def _wall_clock_now() -> datetime:
709
+ return datetime.now(UTC)
710
+
711
+
712
+ @dataclass(frozen=True)
713
+ class SessionDeadline:
714
+ """One capability boundary measured against both wall and monotonic clocks.
715
+
716
+ The wire instant is the authority, but a wall clock can jump backwards after the lease is
717
+ accepted. Capturing its remaining duration on a monotonic clock prevents that jump from
718
+ extending a session. Conversely, a forward wall-clock jump stops early, which is the safe
719
+ direction for a fenced credential.
720
+ """
721
+
722
+ expires_at: datetime
723
+ monotonic_expires_at: float
724
+ _wall_clock: Callable[[], datetime] = field(repr=False, compare=False)
725
+ _monotonic_clock: Callable[[], float] = field(repr=False, compare=False)
726
+
727
+ @classmethod
728
+ def from_lease(
729
+ cls,
730
+ lease: SessionLease,
731
+ *,
732
+ wall_clock: Callable[[], datetime] = _wall_clock_now,
733
+ monotonic_clock: Callable[[], float] = time.monotonic,
734
+ ) -> SessionDeadline:
735
+ return cls.from_expiry(
736
+ lease.capability_expires_at,
737
+ wall_clock=wall_clock,
738
+ monotonic_clock=monotonic_clock,
739
+ )
740
+
741
+ @classmethod
742
+ def from_expiry(
743
+ cls,
744
+ expires_at: datetime,
745
+ *,
746
+ wall_clock: Callable[[], datetime] = _wall_clock_now,
747
+ monotonic_clock: Callable[[], float] = time.monotonic,
748
+ ) -> SessionDeadline:
749
+ """Create a paired deadline from an already authenticated expiry instant."""
750
+
751
+ if not isinstance(expires_at, datetime) or expires_at.tzinfo is None:
752
+ raise SessionWorkerError(
753
+ "SESSION_LEASE_INVALID", "capability_expires_at is not an absolute instant"
754
+ )
755
+ expires_at = expires_at.astimezone(UTC)
756
+ now = cls._wall_now(wall_clock)
757
+ remaining = (expires_at - now).total_seconds()
758
+ if remaining <= 0:
759
+ raise SessionWorkerError(
760
+ "SESSION_CAPABILITY_EXPIRED",
761
+ "the session capability expired before the worker could use it",
762
+ )
763
+ if remaining > MAX_SESSION_CAPABILITY_REMAINING_SECONDS:
764
+ raise SessionWorkerError(
765
+ "SESSION_LEASE_EXPIRY_UNREASONABLE",
766
+ "the session capability exceeds the maximum producer lease",
767
+ )
768
+ monotonic_now = cls._monotonic_now(monotonic_clock)
769
+ return cls(
770
+ expires_at=expires_at,
771
+ monotonic_expires_at=monotonic_now + remaining,
772
+ _wall_clock=wall_clock,
773
+ _monotonic_clock=monotonic_clock,
774
+ )
775
+
776
+ @staticmethod
777
+ def _wall_now(clock: Callable[[], datetime]) -> datetime:
778
+ now = clock()
779
+ if not isinstance(now, datetime) or now.tzinfo is None:
780
+ raise SessionWorkerError(
781
+ "SESSION_CLOCK_INVALID", "the session worker wall clock is not an absolute instant"
782
+ )
783
+ return now.astimezone(UTC)
784
+
785
+ @staticmethod
786
+ def _monotonic_now(clock: Callable[[], float]) -> float:
787
+ value = float(clock())
788
+ if not math.isfinite(value):
789
+ raise SessionWorkerError(
790
+ "SESSION_CLOCK_INVALID", "the session worker monotonic clock is not finite"
791
+ )
792
+ return value
793
+
794
+ def remaining_seconds(self) -> float:
795
+ """Return the earliest safe deadline; a broken clock fails closed at zero."""
796
+
797
+ try:
798
+ wall_remaining = (self.expires_at - self._wall_now(self._wall_clock)).total_seconds()
799
+ monotonic_remaining = self.monotonic_expires_at - self._monotonic_now(
800
+ self._monotonic_clock
801
+ )
802
+ except (SessionWorkerError, TypeError, ValueError, OverflowError):
803
+ return 0.0
804
+ return min(wall_remaining, monotonic_remaining)
805
+
806
+ def expired(self) -> bool:
807
+ return self.remaining_seconds() <= 0
808
+
809
+ def bounded_wait(self, requested_seconds: float) -> float:
810
+ """Limit one interruptible wait to the remaining credential lifetime."""
811
+
812
+ return max(0.0, min(float(requested_seconds), self.remaining_seconds()))
813
+
814
+
815
+ def _linux_direct_child_pids() -> frozenset[int]:
816
+ """Return direct children across threads, or fail closed if Linux cannot enumerate them."""
817
+
818
+ if sys.platform != "linux":
819
+ return frozenset()
820
+ children: set[int] = set()
821
+ try:
822
+ for task in Path("/proc/self/task").iterdir():
823
+ raw = (task / "children").read_text(encoding="ascii").strip()
824
+ if raw:
825
+ children.update(int(value) for value in raw.split())
826
+ except (OSError, ValueError) as error:
827
+ raise SessionWorkerError(
828
+ "SESSION_PROCESS_ISOLATION_UNAVAILABLE",
829
+ "could not enumerate isolated-session descendants",
830
+ ) from error
831
+ return frozenset(children)
832
+
833
+
834
+ def parse_session_lease(value: Any) -> tuple[SessionLease, CredentialLease]:
835
+ """Admit one ``ResearchSessionLease`` and put the capability straight into a lease object.
836
+
837
+ The capability never becomes a module-level or dataclass-held ``str``. It is validated by
838
+ ``hosted_worker.parse_capability`` -- the same shape check the producer path applies -- and then
839
+ handed to ``CredentialLease``, which refuses to be pickled, refuses ``to_dict``, and renders as
840
+ ``<CredentialLease redacted state=open>`` wherever anything tries to print it.
841
+ """
842
+
843
+ if not isinstance(value, dict):
844
+ raise SessionWorkerError("SESSION_LEASE_INVALID", "lease must be an object")
845
+ required = {
846
+ "schema_version",
847
+ "session_id",
848
+ "workspace_id",
849
+ "run_id",
850
+ "attempt_id",
851
+ "generation",
852
+ "fence",
853
+ "capability",
854
+ "capability_expires_at",
855
+ "callback_base_url",
856
+ "allowed_callback_paths",
857
+ "heartbeat_interval_seconds",
858
+ "claim_poll_interval_seconds",
859
+ "idle_timeout_seconds",
860
+ "context",
861
+ }
862
+ if set(value) != required:
863
+ missing = sorted(required - set(value))
864
+ unexpected = sorted(set(value) - required)
865
+ detail = []
866
+ if missing:
867
+ detail.append(f"omits {', '.join(missing)}")
868
+ if unexpected:
869
+ detail.append(f"has unexpected {', '.join(unexpected)}")
870
+ raise SessionWorkerError("SESSION_LEASE_INVALID", f"lease {'; '.join(detail)}")
871
+ expires_at = _parse_utc_timestamp(value["capability_expires_at"], "capability_expires_at")
872
+ raw_capability = value["capability"]
873
+ if not isinstance(raw_capability, str) or len(raw_capability) > MAX_CAPABILITY_BYTES:
874
+ raise SessionWorkerError("SESSION_LEASE_INVALID", "lease capability is not admissible")
875
+ # `parse_capability` refuses anything that is not the exact `mr_cap_...` shape, so a lease that
876
+ # somehow carried a URL, a path, or a JWT never reaches an Authorization header.
877
+ token = parse_capability(raw_capability.encode("ascii", errors="replace"))
878
+ credential = CredentialLease(token.encode("ascii"), {"kind": "session-attempt-capability"})
879
+ del token, raw_capability
880
+
881
+ paths = value["allowed_callback_paths"]
882
+ if not isinstance(paths, list) or not paths:
883
+ credential.close()
884
+ raise SessionWorkerError("SESSION_LEASE_INVALID", "lease declares no callback paths")
885
+ for entry in paths:
886
+ if not isinstance(entry, str) or not entry.startswith("/internal/v3/producer/"):
887
+ credential.close()
888
+ raise SessionWorkerError(
889
+ "SESSION_LEASE_INVALID", "lease callback path is outside the Builder surface"
890
+ )
891
+ context = value["context"]
892
+ if not isinstance(context, dict):
893
+ credential.close()
894
+ raise SessionWorkerError("SESSION_LEASE_INVALID", "lease context must be an object")
895
+ source_ids = context.get("source_ids", [])
896
+ if not isinstance(source_ids, list) or len(source_ids) > 32:
897
+ credential.close()
898
+ raise SessionWorkerError("SESSION_LEASE_INVALID", "lease context source_ids is invalid")
899
+ try:
900
+ base_url = value["callback_base_url"]
901
+ _https_url(str(base_url).rstrip("/"), "callback_base_url")
902
+ lease = SessionLease(
903
+ session_id=_uuid_text(value["session_id"], "session_id"),
904
+ workspace_id=UUID(_uuid_text(value["workspace_id"], "workspace_id")),
905
+ run_id=UUID(_uuid_text(value["run_id"], "run_id")),
906
+ attempt_id=UUID(_uuid_text(value["attempt_id"], "attempt_id")),
907
+ generation=_positive_int(value["generation"], "generation"),
908
+ fence=_positive_int(value["fence"], "fence"),
909
+ capability_expires_at=expires_at,
910
+ callback_base_url=str(base_url).rstrip("/"),
911
+ allowed_callback_paths=tuple(paths),
912
+ heartbeat_interval_seconds=float(
913
+ _positive_int(value["heartbeat_interval_seconds"], "heartbeat_interval_seconds")
914
+ ),
915
+ claim_poll_interval_seconds=float(
916
+ _positive_int(value["claim_poll_interval_seconds"], "claim_poll_interval_seconds")
917
+ ),
918
+ idle_timeout_seconds=float(
919
+ _positive_int(value["idle_timeout_seconds"], "idle_timeout_seconds")
920
+ ),
921
+ context=SessionContext(
922
+ dataset_id=_uuid_text(context["dataset_id"], "context.dataset_id"),
923
+ question_id=_uuid_text(context["question_id"], "context.question_id"),
924
+ table_id=(
925
+ _uuid_text(context["table_id"], "context.table_id")
926
+ if context.get("table_id") is not None
927
+ else None
928
+ ),
929
+ source_ids=frozenset(
930
+ _uuid_text(item, "context.source_ids entry") for item in source_ids
931
+ ),
932
+ ),
933
+ )
934
+ except (KeyError, SessionWorkerError, HostedWorkerError, ValueError):
935
+ credential.close()
936
+ raise
937
+ return lease, credential
938
+
939
+
940
+ # ==========================================================================================
941
+ # Probe source access
942
+ # ==========================================================================================
943
+
944
+
945
+ @dataclass(frozen=True)
946
+ class ProbeSource:
947
+ """The bytes a probe looks at, and where they lawfully came from."""
948
+
949
+ source_id: str
950
+ content: bytes
951
+ media_type: str
952
+ data_format: str
953
+ filename: str
954
+ truncated: bool
955
+ origin: str
956
+
957
+
958
+ class ProbeSourceReader(Protocol):
959
+ """Where a probe's bytes come from. Injected, because which door is open is a deployment fact.
960
+
961
+ A session worker holds a *producer capability on a session run*, which is deliberately less
962
+ authority than a build attempt has, and that closes both of the older acquisition paths:
963
+
964
+ * ``producer/attempts/{id}/public-crawls`` (ADR 0015, worker side) requires
965
+ ``table_recipe_id``, ``recipe_version``, ``recipe_digest``, ``activation_epoch`` and
966
+ ``execution_scope_digest`` on the command, plus a frozen ``source_authority_digest`` from the
967
+ job's ``source_bindings``. A session run has none of the recipe-bound six -- Studio's own
968
+ design note says so, because probing is what a recipe gets written *from* -- so the command
969
+ cannot be constructed, let alone accepted.
970
+ * ``/v3/workspaces/{id}/public-acquisitions`` (ADR 0015, requester side) is
971
+ ``userBearer``-authenticated. The worker holds no user bearer and must not be given one.
972
+
973
+ Studio built a third door for exactly this authority (Studio #99), and
974
+ :class:`StudioSessionAcquisitionReader` is the harness half of it. It is now the production
975
+ default. :class:`UnavailableProbeSourceReader` stays as the honest reader for a deployment
976
+ whose Studio predates that door: a probe that cannot read its source is a failed probe, not a
977
+ failed session, and the session stays live either way.
978
+ """
979
+
980
+ def read(self, *, source_id: str, max_bytes: int) -> ProbeSource: ...
981
+
982
+
983
+ @dataclass(frozen=True)
984
+ class UnavailableProbeSourceReader:
985
+ """Refuses every read, with the reason, and never guesses.
986
+
987
+ No longer the production default -- :class:`StudioSessionAcquisitionReader` is -- and kept
988
+ rather than deleted for the two cases where refusing is the correct answer: a Studio that
989
+ predates the session-acquisition door, and a test that wants the refusal path without a fake
990
+ for the whole acquisition protocol.
991
+ """
992
+
993
+ reason: str = (
994
+ "no session-scoped source channel is wired: the producer public-crawl command requires "
995
+ "the recipe-bound coordinates a session run does not have, and the workspace "
996
+ "public-acquisition endpoint requires a user bearer this worker must not hold"
997
+ )
998
+
999
+ def read(self, *, source_id: str, max_bytes: int) -> ProbeSource:
1000
+ raise SessionWorkerError("PROBE_SOURCE_CHANNEL_UNAVAILABLE", self.reason)
1001
+
1002
+
1003
+ class StudioSessionAcquisitionReader:
1004
+ """The production reader, against Studio's third acquisition door (Studio #99).
1005
+
1006
+ **Why this endpoint and not one of the two that already existed.** A session worker holds a
1007
+ producer capability on a *session* run. `producer/attempts/{id}/public-crawls` validates its
1008
+ command against a frozen recipe activation -- ``table_recipe_id``, ``recipe_version``,
1009
+ ``recipe_digest``, ``activation_epoch``, ``execution_scope_digest`` and a
1010
+ ``source_authority_digest`` from the run's source bindings -- and a session run carries none of
1011
+ the six, because probing is what a recipe gets written *from*. `/v3/workspaces/{id}/public-
1012
+ acquisitions` is ``userBearer``-authenticated and this worker must never hold a user bearer.
1013
+ `session-acquisitions` is the door built for exactly this authority, with two scopes of its own.
1014
+
1015
+ **The worker names a source and nothing else, and that is load-bearing rather than tidy.** The
1016
+ command admits no additional properties, so there is no field through which a URL, an adapter,
1017
+ a reader pin, a limit or an egress attestation could arrive. Studio resolves the locator from
1018
+ the source's one active ``public.https`` connector configuration -- a registration a *person*
1019
+ made -- applies its own reviewed egress policy, and clamps the registered limits to the session
1020
+ ceilings. A worker therefore cannot point the Clean room anywhere, and this class could not
1021
+ make it do so if it tried, because it has nothing to send.
1022
+
1023
+ **What is verified before anything is parsed.** The bundle is the same generation-pinned
1024
+ quarantine object the Builder collects, and the same obligation comes with it: the storage
1025
+ observation is not the evidence. `validate_result_download` re-derives the transfer authority
1026
+ and the pinned length and integrity, `verify_result_bytes` checks the bytes against them, and
1027
+ `parse_crawler_result` then re-derives the nested ``normalized_content_digest``,
1028
+ ``normalized_size_bytes`` and ``result_digest`` from the document itself. Only after all three
1029
+ does a byte reach the sandbox.
1030
+
1031
+ **The wait is bounded and the loop is idempotent.** One create per source per attempt, under an
1032
+ Idempotency-Key derived from the source id, so a retried probe replays the first acquisition
1033
+ rather than spending a second admission slot -- the session's slots are charged to the person
1034
+ who opened it. Polling stops at `acquisition_timeout_seconds`; a probe is interactive and a
1035
+ crawl that has not settled by then is a timeout the client can be told about, not a thread the
1036
+ session should keep parked.
1037
+ """
1038
+
1039
+ def __init__(
1040
+ self,
1041
+ *,
1042
+ authority: Any,
1043
+ models: Any,
1044
+ lease: SessionLease,
1045
+ transfer: Any,
1046
+ results_bucket: str | None = None,
1047
+ timeout_seconds: float = DEFAULT_ACQUISITION_TIMEOUT_SECONDS,
1048
+ poll_seconds: float = ACQUISITION_POLL_SECONDS,
1049
+ sleep: Callable[[float], None] = time.sleep,
1050
+ monotonic: Callable[[], float] = time.monotonic,
1051
+ now: Callable[[], datetime] = lambda: datetime.now(UTC),
1052
+ ) -> None:
1053
+ self._authority = authority
1054
+ self._models = models
1055
+ self._lease = lease
1056
+ self._transfer = transfer
1057
+ self._results_bucket = results_bucket
1058
+ self._timeout = timeout_seconds
1059
+ self._poll = poll_seconds
1060
+ self._sleep = sleep
1061
+ self._monotonic = monotonic
1062
+ self._now = now
1063
+ # One entry per source, bumped only when that source's acquisition expired. It is what
1064
+ # makes the idempotency key stable for the ordinary case and re-openable for the one case
1065
+ # where replaying the old key would replay a corpse. See :meth:`read`.
1066
+ self._epochs: dict[str, int] = {}
1067
+
1068
+ # -- the read ----------------------------------------------------------------------------
1069
+
1070
+ def read(self, *, source_id: str, max_bytes: int) -> ProbeSource:
1071
+ try:
1072
+ session = self._settled(self._create(source_id), source_id)
1073
+ except SessionWorkerError as error:
1074
+ if error.code != ACQUISITION_EXPIRED_CODE:
1075
+ raise
1076
+ # An acquisition expires no later than the worker's attempt lease, and a warm session
1077
+ # outlives that comfortably -- the idle timeout alone is ten minutes. Without this, the
1078
+ # first source a long session reads becomes permanently unreadable the moment its
1079
+ # acquisition lapses, because the source-keyed idempotency key would keep replaying the
1080
+ # lapsed one. Studio releases the admission slot as it expires the acquisition, so
1081
+ # opening a fresh one under the next epoch is not a second spend. Exactly one retry:
1082
+ # a second expiry inside one probe is a clock problem, not a race.
1083
+ self._epochs[source_id] = self._epochs.get(source_id, 0) + 1
1084
+ session = self._settled(self._create(source_id), source_id)
1085
+ download = session.get("result_download")
1086
+ if not isinstance(download, Mapping):
1087
+ raise SessionWorkerError(
1088
+ "PROBE_SOURCE_RESULT_INVALID",
1089
+ "a ready session acquisition carried no result download",
1090
+ )
1091
+ raw = self._collect(download, session)
1092
+ return self._probe_source(raw, session, source_id, max_bytes)
1093
+
1094
+ # -- Studio ------------------------------------------------------------------------------
1095
+
1096
+ def _create(self, source_id: str) -> dict[str, Any]:
1097
+ command = self._models.CreateSessionScopedHostedAcquisitionCommand.from_dict(
1098
+ {
1099
+ "schema_version": START_SCHEMA_VERSION,
1100
+ "workspace_id": str(self._lease.workspace_id),
1101
+ "run_id": str(self._lease.run_id),
1102
+ "attempt_id": str(self._lease.attempt_id),
1103
+ "fence": self._lease.fence,
1104
+ "session_id": self._lease.session_id,
1105
+ "source_id": source_id,
1106
+ }
1107
+ )
1108
+ response = self._authority.create_session_acquisition.sync_detailed(
1109
+ attempt_id=self._lease.attempt_id,
1110
+ body=command,
1111
+ # KEYED ON THE SOURCE, deliberately, and not on a counter the way the claim poll
1112
+ # is. Two probes against one source inside one session must share the acquisition
1113
+ # they would otherwise both pay an admission slot for; a counter would turn a
1114
+ # follow-up question into a second crawl of the same bytes.
1115
+ idempotency_key=_idempotency(
1116
+ "session",
1117
+ self._lease.attempt_id,
1118
+ "acquire",
1119
+ f"{source_id}:{self._epochs.get(source_id, 0)}",
1120
+ ),
1121
+ )
1122
+ try:
1123
+ return _response("createSessionAcquisition", response, 201).to_dict()
1124
+ except HostedWorkerError as error:
1125
+ if (
1126
+ error.code in PROBE_THROTTLE_WIRE_CODES
1127
+ and int(response.status_code) == HTTPStatus.TOO_MANY_REQUESTS
1128
+ ):
1129
+ retry_after_seconds = _retry_after_seconds(getattr(response, "headers", {}))
1130
+ raise SessionWorkerError(
1131
+ error.code,
1132
+ error.detail,
1133
+ retry_after_seconds=retry_after_seconds,
1134
+ ) from error
1135
+ raise
1136
+
1137
+ def _settled(self, session: Mapping[str, Any], source_id: str) -> dict[str, Any]:
1138
+ deadline = self._monotonic() + self._timeout
1139
+ try:
1140
+ crawler_session_id = UUID(str(session.get("crawler_session_id")))
1141
+ except ValueError as error:
1142
+ raise SessionWorkerError(
1143
+ "PROBE_SOURCE_ACQUISITION_MISBOUND",
1144
+ "session acquisition has no canonical crawler session id",
1145
+ ) from error
1146
+ current = dict(session)
1147
+ while True:
1148
+ self._check(current, source_id)
1149
+ status = str(current.get("status", ""))
1150
+ if status == "result_ready":
1151
+ return current
1152
+ if status == "failed":
1153
+ raise SessionWorkerError(
1154
+ _acquisition_failure_code(current.get("failure_code")),
1155
+ "Studio session acquisition failed",
1156
+ )
1157
+ if self._monotonic() >= deadline:
1158
+ raise SessionWorkerError(
1159
+ "PROBE_SOURCE_ACQUISITION_TIMEOUT",
1160
+ f"session acquisition did not settle within {self._timeout:g}s",
1161
+ )
1162
+ self._sleep(self._poll)
1163
+ current = _response(
1164
+ "getSessionAcquisition",
1165
+ self._authority.get_session_acquisition.sync_detailed(
1166
+ attempt_id=self._lease.attempt_id,
1167
+ crawler_session_id=crawler_session_id,
1168
+ ),
1169
+ 200,
1170
+ ).to_dict()
1171
+
1172
+ def _check(self, session: Mapping[str, Any], source_id: str) -> None:
1173
+ """Refuse an acquisition that is not this attempt's, for this session, for this source.
1174
+
1175
+ Studio binds every one of these server-side, so this can only ever agree -- until the day
1176
+ it does not. Re-deriving them is what makes a mixed-up answer a refusal instead of bytes
1177
+ attributed to the wrong source, and it costs seven comparisons on a path that has just
1178
+ spent a network round trip.
1179
+ """
1180
+
1181
+ if (
1182
+ str(session.get("workspace_id")) != str(self._lease.workspace_id)
1183
+ or str(session.get("run_id")) != str(self._lease.run_id)
1184
+ or str(session.get("attempt_id")) != str(self._lease.attempt_id)
1185
+ or str(session.get("research_session_id")) != str(self._lease.session_id)
1186
+ or session.get("source_id") != source_id
1187
+ or session.get("producer_fence") != self._lease.fence
1188
+ or session.get("adapter_id") != "public.https"
1189
+ ):
1190
+ raise SessionWorkerError(
1191
+ "PROBE_SOURCE_ACQUISITION_MISBOUND",
1192
+ "session acquisition is not bound to this attempt, session and source",
1193
+ )
1194
+
1195
+ # -- the bundle --------------------------------------------------------------------------
1196
+
1197
+ def _collect(self, download: Mapping[str, Any], session: Mapping[str, Any]) -> bytes:
1198
+ declared = download.get("bundle_size_bytes")
1199
+ if type(declared) is not int or declared > MAX_PROBE_ACQUISITION_BUNDLE_BYTES:
1200
+ raise SessionWorkerError(
1201
+ "PROBE_SOURCE_TOO_LARGE",
1202
+ "the acquired bundle is larger than a probe may hold",
1203
+ )
1204
+ try:
1205
+ validated = validate_result_download(
1206
+ download,
1207
+ base_url=self._lease.callback_base_url,
1208
+ crawler_session_id=str(session["crawler_session_id"]),
1209
+ request_digest=str(session["request_digest"]),
1210
+ source_id=str(session["source_id"]),
1211
+ source_authority_digest=str(session["source_authority_digest"]),
1212
+ now=self._now(),
1213
+ expected_gcs_bucket=self._results_bucket,
1214
+ )
1215
+ except ResultDownloadError as error:
1216
+ raise SessionWorkerError("PROBE_SOURCE_DOWNLOAD_INVALID", str(error)) from error
1217
+ content = bytearray()
1218
+ with self._transfer.stream("GET", validated.url, headers={}) as response:
1219
+ if response.status_code != 200:
1220
+ raise SessionWorkerError(
1221
+ "PROBE_SOURCE_DOWNLOAD_FAILED",
1222
+ f"quarantine download returned HTTP {response.status_code}",
1223
+ )
1224
+ for chunk in response.iter_bytes():
1225
+ content.extend(chunk)
1226
+ if len(content) > validated.size:
1227
+ raise SessionWorkerError(
1228
+ "PROBE_SOURCE_DOWNLOAD_FAILED",
1229
+ "quarantine bytes exceed the generation-pinned size",
1230
+ )
1231
+ raw = bytes(content)
1232
+ try:
1233
+ verify_result_bytes(validated, raw)
1234
+ except ResultDownloadError as error:
1235
+ raise SessionWorkerError("PROBE_SOURCE_DOWNLOAD_FAILED", str(error)) from error
1236
+ return raw
1237
+
1238
+ def _probe_source(
1239
+ self, raw: bytes, session: Mapping[str, Any], source_id: str, max_bytes: int
1240
+ ) -> ProbeSource:
1241
+ document = _strict_json(raw, "session acquisition result")
1242
+ if not isinstance(document, dict) or canonical_json_bytes(document) != raw:
1243
+ raise SessionWorkerError(
1244
+ "PROBE_SOURCE_RESULT_INVALID",
1245
+ "the acquired result is not exact canonical JSON",
1246
+ )
1247
+ try:
1248
+ result = parse_crawler_result(document)
1249
+ except HostedCrawlerProtocolError as error:
1250
+ raise SessionWorkerError("PROBE_SOURCE_RESULT_INVALID", error.detail) from error
1251
+ receipt = result.acquisition_receipt
1252
+ if (
1253
+ result.source_id != source_id
1254
+ or result.request_digest != session.get("request_digest")
1255
+ or result.adapter_id != "public.https"
1256
+ or receipt.get("source_id") != source_id
1257
+ or receipt.get("normalized_content_sha256") != sha256_bytes(result.normalized_content)
1258
+ or receipt.get("normalized_size_bytes") != len(result.normalized_content)
1259
+ or receipt.get("egress_policy_attestation") != session.get("egress_policy_attestation")
1260
+ ):
1261
+ raise SessionWorkerError(
1262
+ "PROBE_SOURCE_RESULT_INVALID",
1263
+ "the acquired result is not bound to the acquisition that produced it",
1264
+ )
1265
+ if len(result.normalized_content) > max_bytes:
1266
+ raise SessionWorkerError(
1267
+ "PROBE_SOURCE_TOO_LARGE",
1268
+ "the acquired source is larger than a probe may hold",
1269
+ )
1270
+ media_type = receipt.get("normalized_media_type")
1271
+ data_format = receipt.get("normalized_data_format")
1272
+ filename = receipt.get("normalized_filename")
1273
+ if (
1274
+ not isinstance(media_type, str)
1275
+ or not isinstance(data_format, str)
1276
+ or not isinstance(filename, str)
1277
+ ):
1278
+ raise SessionWorkerError(
1279
+ "PROBE_SOURCE_RESULT_INVALID",
1280
+ "the acquisition receipt does not name its normalized shape",
1281
+ )
1282
+ return ProbeSource(
1283
+ source_id=source_id,
1284
+ content=result.normalized_content,
1285
+ media_type=media_type,
1286
+ data_format=data_format,
1287
+ filename=filename,
1288
+ # The bytes are whole or they are refused. Studio clamps the source to the session
1289
+ # ceiling before it crawls, and a truncated table would make every count a probe
1290
+ # reports quietly wrong, so there is no partial answer to flag here.
1291
+ truncated=False,
1292
+ origin="studio.session-acquisition",
1293
+ )
1294
+
1295
+
1296
+ def _acquisition_failure_code(value: Any) -> str:
1297
+ """Render Studio's ``failure_code`` as a code this worker will settle a probe under.
1298
+
1299
+ Studio's is free-form on the wire; anything that is not a code-shaped token becomes one stable
1300
+ value rather than reaching the probe record, and the timeline, as somebody's prose.
1301
+ """
1302
+
1303
+ if isinstance(value, str) and re.fullmatch(r"[A-Z][A-Z0-9_]{2,127}", value):
1304
+ return value
1305
+ return "PROBE_SOURCE_ACQUISITION_FAILED"
1306
+
1307
+
1308
+ def install_hosted_catalogue() -> Mapping[str, Any] | None:
1309
+ """Put this deployment's pinned sealed catalogue on the worker, when one is pinned.
1310
+
1311
+ The same hook the Builder calls, at start rather than at first query, so a worker that needs
1312
+ the catalogue discovers it cannot have one while it can still fail cleanly. ``None`` when no
1313
+ generation is pinned, which is an ordinary deployment and not a fault: probes then run against
1314
+ whatever the source reader provides.
1315
+
1316
+ Imported inside the function rather than at module scope. The import reaches the channel
1317
+ client and its digest verification, and a session that never touches the catalogue should not
1318
+ pay for loading them.
1319
+ """
1320
+
1321
+ from mostlyright.data_harness.sources.catalog.hosted_catalog import (
1322
+ ensure_hosted_catalog,
1323
+ )
1324
+
1325
+ receipt = ensure_hosted_catalog()
1326
+ return None if receipt is None else dict(receipt)
1327
+
1328
+
1329
+ # ==========================================================================================
1330
+ # Probe execution
1331
+ # ==========================================================================================
1332
+
1333
+
1334
+ class ProbeExecutor:
1335
+ """Run one probe inside the sandbox and answer with a vocabulary-valid result body.
1336
+
1337
+ Every probe is a *no-network* sandbox operation (``parse``), which is the only confinement the
1338
+ warm pool serves and therefore the only one that can meet the follow-up-probe budget. Bytes
1339
+ reach the executor already acquired, through the injected :class:`ProbeSourceReader`; nothing
1340
+ here retrieves anything, so a probe cannot cause an outbound request at all.
1341
+
1342
+ **The read is before the dispatch because all four kinds need it, not by accident.** When the
1343
+ read refuses -- against a Studio with no session-acquisition door, or for a source with no
1344
+ registration -- every kind refuses with it, and the obvious question is whether some kind could
1345
+ have been answered from metadata instead and let through. The answer is no, for each of them:
1346
+
1347
+ * ``sample_rows`` and ``profile_columns`` answer with cell values and per-column statistics
1348
+ computed over them.
1349
+ * ``evaluate_expression`` -- the one that looks like it might, since ``count`` with no
1350
+ predicate is just a number -- filters *rows* and aggregates *values*. Even the degenerate
1351
+ case would have to answer from ``declared_row_count``, which is what a publisher SAYS, and
1352
+ answering an observation with a claim is the substitution this whole lane exists to avoid.
1353
+ * ``source_inspect`` is the closest call. The sealed catalogue does carry ``declared_columns``,
1354
+ but they are provider-declared names with no types; the answer this kind gives is a
1355
+ ``schema_digest`` from the parse plus an ``inferred_type`` per column, both derived from the
1356
+ scanned cells, and neither is reachable without them.
1357
+
1358
+ So a read refusal is per-read and not per-kind, and there is no subset to exempt. What a caller
1359
+ gets instead is a *settled* probe carrying the refusal's own code -- reported to Studio so the
1360
+ probe record terminates, and narrated on the run stream as an ``outcome="refused"`` completion
1361
+ -- rather than a probe that starts and never ends.
1362
+ """
1363
+
1364
+ def __init__(
1365
+ self,
1366
+ *,
1367
+ sandbox: Any,
1368
+ reader: ProbeSourceReader,
1369
+ context: SessionContext,
1370
+ parse_limits: ParseLimits | None = None,
1371
+ max_source_bytes: int = 16 * 1024 * 1024,
1372
+ ) -> None:
1373
+ self._sandbox = sandbox
1374
+ self._reader = reader
1375
+ self._context = context
1376
+ self._limits = parse_limits or ParseLimits()
1377
+ self._max_source_bytes = max_source_bytes
1378
+
1379
+ def execute(self, *, probe_id: str, probe_kind: str, request: Any) -> dict[str, Any]:
1380
+ """Answer one probe.
1381
+
1382
+ Raises :class:`SessionWorkerError`, or
1383
+ :class:`~mostlyright.data_harness.session_probes.ProbeVocabularyError`, or whatever the
1384
+ Clean room raises -- every one of which carries a ``code`` that
1385
+ :func:`_probe_failure_code` settles the probe under.
1386
+ """
1387
+
1388
+ normalized = session_probes.validate_probe_request(probe_kind, request)
1389
+ source_id = normalized["source_id"]
1390
+ # The session's own context is the probe authority. A probe naming a source the session was
1391
+ # not opened against is refused here, before any read: the capability says which session
1392
+ # this is, and the session says which sources it may look at. Studio does not re-check this
1393
+ # -- the probe body is opaque to it -- so if this check is not here, it is nowhere.
1394
+ if self._context.source_ids and source_id not in self._context.source_ids:
1395
+ raise SessionWorkerError(
1396
+ "PROBE_SOURCE_OUT_OF_SCOPE",
1397
+ "probe names a source this session was not opened against",
1398
+ )
1399
+ source = self._reader.read(source_id=source_id, max_bytes=self._max_source_bytes)
1400
+ if source.source_id != source_id:
1401
+ raise SessionWorkerError(
1402
+ "PROBE_SOURCE_MISMATCH", "source reader answered for a different source"
1403
+ )
1404
+ result = self._sandbox.parse(
1405
+ request_id=_sandbox_request_id(probe_id),
1406
+ content=source.content,
1407
+ data_format=source.data_format,
1408
+ media_type=source.media_type,
1409
+ filename=source.filename,
1410
+ limits=self._limits,
1411
+ )
1412
+ parsed = getattr(result, "parsed", None)
1413
+ if parsed is None:
1414
+ raise SessionWorkerError("PROBE_PARSE_EMPTY", "the Clean room returned no table")
1415
+ envelope = {
1416
+ "schema_version": session_probes.PROBE_SCHEMA_VERSION,
1417
+ "probe_kind": probe_kind,
1418
+ "source_id": source_id,
1419
+ "content_sha256": hashlib.sha256(source.content).hexdigest(),
1420
+ "content_bytes": len(source.content),
1421
+ "data_format": parsed.data_format,
1422
+ "sandbox_policy_digest": result.policy_digest,
1423
+ "truncated": bool(source.truncated),
1424
+ }
1425
+ body = self._answer(probe_kind, normalized, parsed, envelope)
1426
+ return session_probes.validate_probe_result(probe_kind, body)
1427
+
1428
+ # -- per-kind answers -------------------------------------------------------------------
1429
+
1430
+ def _answer(
1431
+ self,
1432
+ probe_kind: str,
1433
+ request: Mapping[str, Any],
1434
+ parsed: Any,
1435
+ envelope: dict[str, Any],
1436
+ ) -> dict[str, Any]:
1437
+ columns: tuple[str, ...] = tuple(parsed.columns)
1438
+ rows: tuple[tuple[Any, ...], ...] = tuple(parsed.rows)
1439
+ if probe_kind == "sample_rows":
1440
+ return self._sample_rows(request, columns, rows, envelope)
1441
+ scanned = rows[: request["scan_rows"]]
1442
+ envelope["rows_scanned"] = len(scanned)
1443
+ if probe_kind == "source_inspect":
1444
+ return self._source_inspect(columns, scanned, parsed, envelope)
1445
+ if probe_kind == "profile_columns":
1446
+ return self._profile_columns(request, columns, scanned, envelope)
1447
+ return self._evaluate_expression(request, columns, scanned, envelope)
1448
+
1449
+ def _source_inspect(
1450
+ self,
1451
+ columns: Sequence[str],
1452
+ scanned: Sequence[Sequence[Any]],
1453
+ parsed: Any,
1454
+ envelope: dict[str, Any],
1455
+ ) -> dict[str, Any]:
1456
+ # Deliberately no cell values. `source_inspect` answers "what shape is this", and a client
1457
+ # that wants values asks `sample_rows` -- which is a separate probe, separately visible on
1458
+ # the session's event stream, so "who looked at the data" stays answerable.
1459
+ return {
1460
+ **envelope,
1461
+ "schema_digest": parsed.schema_digest,
1462
+ "columns": [
1463
+ {
1464
+ "name": name,
1465
+ "inferred_type": session_probes.inferred_type([row[index] for row in scanned]),
1466
+ }
1467
+ for index, name in enumerate(columns[: session_probes.MAX_PROFILE_COLUMNS])
1468
+ ],
1469
+ }
1470
+
1471
+ def _sample_rows(
1472
+ self,
1473
+ request: Mapping[str, Any],
1474
+ columns: Sequence[str],
1475
+ rows: Sequence[Sequence[Any]],
1476
+ envelope: dict[str, Any],
1477
+ ) -> dict[str, Any]:
1478
+ selected = self._selected_columns(request, columns)
1479
+ indices = [columns.index(name) for name in selected]
1480
+ window = rows[request["offset"] : request["offset"] + request["limit"]]
1481
+ envelope["rows_scanned"] = request["offset"] + len(window)
1482
+ body = {
1483
+ **envelope,
1484
+ "columns": list(selected),
1485
+ "rows": [
1486
+ [session_probes.preview_cell(row[index]) for index in indices] for row in window
1487
+ ],
1488
+ }
1489
+ # A preview that would not fit is trimmed from the end rather than refused: an interactive
1490
+ # caller asking for 200 wide rows should get the first N it can have, and `truncated` says
1491
+ # so. Refusing would make the useful answer unreachable without a second round trip.
1492
+ while (
1493
+ len(canonical_json_bytes(body)) > session_probes.PROBE_RESULT_MAX_BYTES and body["rows"]
1494
+ ):
1495
+ body["rows"].pop()
1496
+ body["truncated"] = True
1497
+ return body
1498
+
1499
+ def _profile_columns(
1500
+ self,
1501
+ request: Mapping[str, Any],
1502
+ columns: Sequence[str],
1503
+ scanned: Sequence[Sequence[Any]],
1504
+ envelope: dict[str, Any],
1505
+ ) -> dict[str, Any]:
1506
+ selected = self._selected_columns(request, columns)
1507
+ profiles: list[dict[str, Any]] = []
1508
+ for name in selected:
1509
+ index = columns.index(name)
1510
+ values = [row[index] for row in scanned]
1511
+ present = [value for value in values if value is not None]
1512
+ profile: dict[str, Any] = {
1513
+ "name": name,
1514
+ "inferred_type": session_probes.inferred_type(values),
1515
+ "non_null": len(present),
1516
+ "nulls": len(values) - len(present),
1517
+ # Exact over the scanned rows, and `rows_scanned` is in the envelope beside it, so
1518
+ # the number always arrives with its denominator.
1519
+ "distinct": len({session_probes.preview_cell(value) for value in present}),
1520
+ }
1521
+ ordered = self._orderable(present)
1522
+ if ordered:
1523
+ profile["minimum"] = session_probes.preview_cell(ordered[0])
1524
+ profile["maximum"] = session_probes.preview_cell(ordered[-1])
1525
+ profiles.append(profile)
1526
+ return {**envelope, "columns": profiles}
1527
+
1528
+ def _evaluate_expression(
1529
+ self,
1530
+ request: Mapping[str, Any],
1531
+ columns: Sequence[str],
1532
+ scanned: Sequence[Sequence[Any]],
1533
+ envelope: dict[str, Any],
1534
+ ) -> dict[str, Any]:
1535
+ expression = request["expression"]
1536
+ predicate = expression.get("where")
1537
+ aggregate = expression["aggregate"]
1538
+ column = expression.get("column")
1539
+ index = columns.index(column) if column in columns else None
1540
+ matched = 0
1541
+ values: list[Any] = []
1542
+ for row in scanned:
1543
+ record = dict(zip(columns, row, strict=False))
1544
+ if predicate is not None and not session_probes.evaluate_predicate(predicate, record):
1545
+ continue
1546
+ matched += 1
1547
+ if index is not None and row[index] is not None:
1548
+ values.append(row[index])
1549
+ body: dict[str, Any] = {
1550
+ **envelope,
1551
+ "aggregate": aggregate,
1552
+ "value": self._aggregate(aggregate, matched, values),
1553
+ "rows_matched": matched,
1554
+ }
1555
+ if column is not None:
1556
+ body["column"] = column
1557
+ return body
1558
+
1559
+ @staticmethod
1560
+ def _aggregate(aggregate: str, matched: int, values: Sequence[Any]) -> Any:
1561
+ if aggregate == "count":
1562
+ return matched
1563
+ if not values:
1564
+ # No value is `null`, never zero. A sum over nothing that reported 0 would be
1565
+ # indistinguishable from a sum over zeros, and an exploratory probe is exactly where
1566
+ # that distinction matters.
1567
+ return None
1568
+ numeric = [value for value in values if type(value) in {int, float}]
1569
+ if aggregate in {"min", "max"}:
1570
+ ordered = ProbeExecutor._orderable(values)
1571
+ if not ordered:
1572
+ return None
1573
+ return session_probes.preview_cell(ordered[0 if aggregate == "min" else -1])
1574
+ if len(numeric) != len(values):
1575
+ raise SessionWorkerError(
1576
+ "PROBE_AGGREGATE_NOT_NUMERIC",
1577
+ f"{aggregate} needs a numeric column; the scanned values are not all numbers",
1578
+ )
1579
+ total = sum(numeric)
1580
+ return session_probes.preview_cell(total if aggregate == "sum" else total / len(numeric))
1581
+
1582
+ @staticmethod
1583
+ def _orderable(values: Sequence[Any]) -> list[Any]:
1584
+ """Sort a column's present values, or answer empty when the column is not orderable.
1585
+
1586
+ Python 3 raises on ``1 < "a"``; a mixed column therefore has no extremes and says so,
1587
+ rather than the probe failing on data the caller cannot be expected to know the shape of.
1588
+ """
1589
+
1590
+ try:
1591
+ return sorted(values)
1592
+ except TypeError:
1593
+ return []
1594
+
1595
+ def _selected_columns(
1596
+ self, request: Mapping[str, Any], columns: Sequence[str]
1597
+ ) -> tuple[str, ...]:
1598
+ requested = request.get("columns")
1599
+ if requested is None:
1600
+ return tuple(columns[: session_probes.MAX_PROFILE_COLUMNS])
1601
+ unknown = [name for name in requested if name not in columns]
1602
+ if unknown:
1603
+ raise SessionWorkerError(
1604
+ "PROBE_COLUMN_UNKNOWN", f"source has no column named {unknown[0]!r}"
1605
+ )
1606
+ return tuple(requested)
1607
+
1608
+
1609
+ def _probe_failure_code(error: BaseException) -> str:
1610
+ """Render one probe failure as the stable code it should settle under.
1611
+
1612
+ A ``code`` attribute is the convention every coded failure in this package already follows --
1613
+ :class:`SessionWorkerError`, :class:`~mostlyright.data_harness.session_probes.
1614
+ ProbeVocabularyError`, :class:`~mostlyright.data_harness.hosted_worker.HostedWorkerError` and
1615
+ the Clean room's ``AcquisitionSecurityError`` all carry one -- so this reads the convention
1616
+ rather than a list of classes. Reading the class list instead is what lost ``SANDBOX_TIMEOUT``:
1617
+ the sandbox's error is declared in ``acquisition.url_policy`` and was in nobody's tuple.
1618
+
1619
+ The value is checked against the progress value domain before it is believed. A code reaches a
1620
+ timeline as a fact, and a fact that is prose is a record the vocabulary drops whole -- which
1621
+ would put the pair invariant back exactly where this change found it.
1622
+ """
1623
+
1624
+ code = getattr(error, "code", None)
1625
+ if isinstance(code, str) and progress_events.is_progress_fact_value(code):
1626
+ return code
1627
+ return "PROBE_FAILED"
1628
+
1629
+
1630
+ def _retry_after_seconds(headers: Any) -> int | None:
1631
+ """Accept Studio's bounded delta-seconds retry window, never prose or a guessed delay."""
1632
+
1633
+ if not isinstance(headers, Mapping):
1634
+ return None
1635
+ value = next(
1636
+ (
1637
+ candidate
1638
+ for key, candidate in headers.items()
1639
+ if isinstance(key, str) and key.lower() == "retry-after"
1640
+ ),
1641
+ None,
1642
+ )
1643
+ if not isinstance(value, str) or not re.fullmatch(r"[1-9][0-9]{0,3}", value):
1644
+ return None
1645
+ seconds = int(value)
1646
+ return seconds if seconds <= MAX_RETRY_AFTER_SECONDS else None
1647
+
1648
+
1649
+ def _probe_display_failure(error: BaseException, wire_code: str) -> tuple[str, int]:
1650
+ """Return a person-facing code and its server-stated wait while retaining wire identity."""
1651
+
1652
+ if wire_code not in PROBE_THROTTLE_WIRE_CODES:
1653
+ return wire_code, 0
1654
+ retry_after_seconds = getattr(error, "retry_after_seconds", None)
1655
+ if type(retry_after_seconds) is int and 1 <= retry_after_seconds <= MAX_RETRY_AFTER_SECONDS:
1656
+ return SESSION_ACQUISITION_THROTTLED, retry_after_seconds
1657
+ return SESSION_ACQUISITION_THROTTLED, 0
1658
+
1659
+
1660
+ def _probe_progress_outcome(failure_code: str | None) -> str:
1661
+ """Classify a settled probe for the unsealed timeline. Total over ``PROBE_OUTCOMES``.
1662
+
1663
+ Keyed on the code rather than on the exception, so the classification is a property of what the
1664
+ probe settled AS -- the same thing the client reads off the probe resource -- and not of how
1665
+ the worker happened to learn it.
1666
+ """
1667
+
1668
+ if failure_code is None:
1669
+ return "succeeded"
1670
+ if failure_code in PROBE_TIMEOUT_CODES:
1671
+ return "timed_out"
1672
+ if failure_code in PROBE_THROTTLE_CODES:
1673
+ return "throttled"
1674
+ if failure_code in PROBE_REFUSAL_CODES:
1675
+ return "refused"
1676
+ return "failed"
1677
+
1678
+
1679
+ # ==========================================================================================
1680
+ # The Studio authority
1681
+ # ==========================================================================================
1682
+
1683
+
1684
+ def _sandbox_request_id(probe_id: str) -> str:
1685
+ """Render a probe id as a sandbox request id.
1686
+
1687
+ ``CrawlerSandbox`` requires a lowercase bounded identifier that *starts with a letter*, and a
1688
+ UUID beginning with a digit does not -- which is roughly half of them. Passing the probe id
1689
+ straight through therefore refused every other probe with ``SANDBOX_IDENTIFIER``, and the
1690
+ prefix is what fixes it rather than a relaxation on the sandbox side. Anything outside the
1691
+ admitted alphabet is replaced rather than dropped, so two different probe ids cannot collapse
1692
+ into one request id.
1693
+ """
1694
+
1695
+ safe = "".join(
1696
+ character if character.islower() or character.isdigit() or character in "._-" else "-"
1697
+ for character in probe_id.lower()
1698
+ )
1699
+ return f"probe-{safe}"[:128]
1700
+
1701
+
1702
+ class _SessionAuthority(_ProducerAuthority):
1703
+ """The producer surface plus the two verbs a session adds, and the lease that precedes both.
1704
+
1705
+ Subclassed here rather than widened in ``hosted_worker`` so that the producer and verifier
1706
+ façades keep exactly the operation sets they had: a session worker gaining a verb must not
1707
+ silently widen what a build worker's façade will dispatch.
1708
+ """
1709
+
1710
+ _operations: ClassVar[dict[str, str]] = {
1711
+ **_ProducerAuthority._operations,
1712
+ "lease_session": "sessions.lease_research_session",
1713
+ "claim_probe": "producer.claim_research_probe",
1714
+ "report_probe_result": "producer.report_research_probe_result",
1715
+ "create_session_acquisition": "producer.create_session_acquisition",
1716
+ "get_session_acquisition": "producer.get_session_acquisition",
1717
+ }
1718
+
1719
+
1720
+ def _lease_session(
1721
+ *,
1722
+ config: SessionWorkerConfig,
1723
+ invocation: StartInvocation,
1724
+ nonce: bytearray,
1725
+ transport: BootstrapTransport,
1726
+ ) -> tuple[SessionLease, CredentialLease]:
1727
+ """Present ambient workload identity plus the nonce, and take the one lease on this session.
1728
+
1729
+ Two credentials are in play for exactly one call each and both are zeroed. The identity token
1730
+ proves *which fleet* is calling and is destroyed the moment the exchange returns. The nonce
1731
+ proves *which invocation*, and is zeroed by the caller as soon as this returns -- it is a
1732
+ one-time secret whose whole purpose is to stop a reviewed worker leasing a session it was not
1733
+ sent to, and a copy that outlives its one use is a copy that can be stolen.
1734
+ """
1735
+
1736
+ models, _, _, _ = _load_generated()
1737
+ identity = _metadata_identity(transport, config.lease_audience)
1738
+ try:
1739
+ client = _authenticated_client(config.studio_origin, identity.decode("ascii"))
1740
+ finally:
1741
+ identity[:] = b"\x00" * len(identity)
1742
+ try:
1743
+ authority = _SessionAuthority(client)
1744
+ command = models.ContractSessionLeaseCommand.from_dict(
1745
+ {
1746
+ "schema_version": START_SCHEMA_VERSION,
1747
+ "session_id": invocation.session_id,
1748
+ "lease_nonce": nonce.decode("ascii"),
1749
+ }
1750
+ )
1751
+ parsed = _response(
1752
+ "leaseResearchSession",
1753
+ authority.lease_session.sync_detailed(
1754
+ session_id=UUID(invocation.session_id), body=command
1755
+ ),
1756
+ 200,
1757
+ )
1758
+ wire = parsed.to_dict()
1759
+ finally:
1760
+ # The identity-bearing client is discarded whatever happens: everything after the lease
1761
+ # authenticates with the fenced capability, and an idle client holding a workload-identity
1762
+ # bearer for the rest of a session is authority nothing needs.
1763
+ client.get_httpx_client().close()
1764
+ lease, credential = parse_session_lease(wire)
1765
+ if lease.session_id != invocation.session_id or str(lease.workspace_id) != (
1766
+ invocation.workspace_id
1767
+ ):
1768
+ credential.close()
1769
+ raise SessionWorkerError(
1770
+ "SESSION_LEASE_MISMATCH", "the lease names a different session or workspace"
1771
+ )
1772
+ return lease, credential
1773
+
1774
+
1775
+ # ==========================================================================================
1776
+ # The supervisor
1777
+ # ==========================================================================================
1778
+
1779
+
1780
+ @dataclass
1781
+ class _SessionOutcome:
1782
+ session_id: str
1783
+ attempt_id: str | None = None
1784
+ probes_answered: int = 0
1785
+ probes_failed: int = 0
1786
+ terminal_state: str | None = None
1787
+ failure_code: str | None = None
1788
+ warm_pool_ready: int = 0
1789
+ catalogue: Mapping[str, Any] | None = None
1790
+ events: list[str] = field(default_factory=list)
1791
+
1792
+ def to_dict(self) -> dict[str, Any]:
1793
+ return {
1794
+ "status": "failed" if self.failure_code else "stopped",
1795
+ "role": SESSION_WORKER_ROLE,
1796
+ "session_id": self.session_id,
1797
+ "attempt_id": self.attempt_id,
1798
+ "probes_answered": self.probes_answered,
1799
+ "probes_failed": self.probes_failed,
1800
+ "terminal_state": self.terminal_state,
1801
+ "failure_code": self.failure_code,
1802
+ "warm_pool_ready": self.warm_pool_ready,
1803
+ }
1804
+
1805
+
1806
+ class _Heartbeat(threading.Thread):
1807
+ """Liveness on a timer, and nothing else.
1808
+
1809
+ A separate thread rather than a tick inside the claim loop, because the claim loop blocks for
1810
+ the length of a probe and a probe may run for the better part of a minute. A heartbeat that
1811
+ only fires between probes would go quiet during exactly the interval where a client most wants
1812
+ to know the worker is alive.
1813
+
1814
+ Every failure is swallowed. A missed heartbeat costs liveness resolution; the authority on
1815
+ whether this worker still holds its lease is the attempt lease expiry Studio enforces on every
1816
+ fenced call, and the claim loop discovers a lost lease there. Failing the session from this
1817
+ thread would be this thread deciding something it cannot see.
1818
+ """
1819
+
1820
+ def __init__(
1821
+ self,
1822
+ *,
1823
+ send: Callable[[], None],
1824
+ interval_seconds: float,
1825
+ stop: threading.Event,
1826
+ deadline: SessionDeadline | None = None,
1827
+ on_expired: Callable[[], None] | None = None,
1828
+ ) -> None:
1829
+ super().__init__(name="mr-session-heartbeat", daemon=True)
1830
+ self._send = send
1831
+ self._interval = max(1.0, float(interval_seconds))
1832
+ self._stop = stop
1833
+ self._deadline = deadline
1834
+ self._on_expired = on_expired
1835
+ self.sent = 0
1836
+
1837
+ def run(self) -> None:
1838
+ while not self._stop.is_set():
1839
+ wait_seconds = self._interval
1840
+ if self._deadline is not None:
1841
+ wait_seconds = self._deadline.bounded_wait(wait_seconds)
1842
+ if wait_seconds <= 0:
1843
+ self._expire()
1844
+ return
1845
+ if self._stop.wait(wait_seconds):
1846
+ return
1847
+ if self._deadline is not None and self._deadline.expired():
1848
+ self._expire()
1849
+ return
1850
+ try:
1851
+ self._send()
1852
+ except Exception:
1853
+ LOGGER.debug("session heartbeat failed", exc_info=True)
1854
+ continue
1855
+ self.sent += 1
1856
+
1857
+ def _expire(self) -> None:
1858
+ self._stop.set()
1859
+ if self._on_expired is not None:
1860
+ self._on_expired()
1861
+
1862
+
1863
+ class SessionSupervisor:
1864
+ """Lease, start, claim, answer, heartbeat, and stop on a terminal state.
1865
+
1866
+ In production Job mode, a capability-free parent supplies an opaque dispatch and this
1867
+ supervisor exchanges it only in its isolated child. The capability remains in that child
1868
+ inside ``CredentialLease`` only for the life of the session; it is not serialised, logged,
1869
+ written to disk or passed to the probe sandbox, and it is zeroed when the supervisor stops.
1870
+ ``invocation``/``nonce`` remain only as a compatibility seam for the retained rollback
1871
+ Service; Job mode never constructs or accepts a lease nonce.
1872
+ """
1873
+
1874
+ def __init__(
1875
+ self,
1876
+ *,
1877
+ config: SessionWorkerConfig,
1878
+ invocation: StartInvocation | None = None,
1879
+ nonce: bytearray | None = None,
1880
+ bootstrap: tuple[SessionLease, CredentialLease] | None = None,
1881
+ dispatch: SessionDispatch | None = None,
1882
+ transport: BootstrapTransport | None = None,
1883
+ sandbox_factory: Callable[[SessionWorkerConfig], Any] | None = None,
1884
+ reader_factory: Callable[[SessionLease], ProbeSourceReader] | None = None,
1885
+ lease_session: Callable[..., tuple[SessionLease, CredentialLease]] = _lease_session,
1886
+ sleep: Callable[[float], None] = time.sleep,
1887
+ wall_clock: Callable[[], datetime] = _wall_clock_now,
1888
+ monotonic_clock: Callable[[], float] = time.monotonic,
1889
+ process_isolation: bool = False,
1890
+ on_isolation_ready: Callable[[], None] | None = None,
1891
+ ) -> None:
1892
+ legacy = invocation is not None or nonce is not None
1893
+ modes = int(legacy) + int(bootstrap is not None) + int(dispatch is not None)
1894
+ if modes != 1 or (legacy and (invocation is None or nonce is None)):
1895
+ raise SessionWorkerError(
1896
+ "SESSION_SUPERVISOR_CONFIGURATION_INVALID",
1897
+ "provide exactly one legacy invocation, dispatch, or dispatch bootstrap",
1898
+ )
1899
+ self._config = config
1900
+ self._invocation = invocation
1901
+ self._nonce = nonce
1902
+ self._bootstrap = bootstrap
1903
+ self._dispatch = dispatch
1904
+ self._transport = transport or UrlLibBootstrapTransport()
1905
+ self._sandbox_factory = sandbox_factory or _default_sandbox
1906
+ # `None` rather than a default lambda, because the production reader cannot be built from
1907
+ # a lease alone: it needs the fenced authority and a transfer client, both of which exist
1908
+ # only once `_serve` has one. A test injecting a factory still overrides everything.
1909
+ self._reader_factory = reader_factory
1910
+ self._lease_session = lease_session
1911
+ self._sleep = sleep
1912
+ self._wall_clock = wall_clock
1913
+ self._monotonic_clock = monotonic_clock
1914
+ self._process_isolation = process_isolation
1915
+ self._on_isolation_ready = on_isolation_ready
1916
+ self._stop = threading.Event()
1917
+ self.outcome = _SessionOutcome(
1918
+ session_id=(
1919
+ bootstrap[0].session_id
1920
+ if bootstrap is not None
1921
+ else invocation.session_id
1922
+ if invocation is not None
1923
+ else ""
1924
+ )
1925
+ )
1926
+
1927
+ def stop(self) -> None:
1928
+ """Ask the claim loop to finish the probe it is on and then stop."""
1929
+
1930
+ self._stop.set()
1931
+
1932
+ @property
1933
+ def is_running(self) -> bool:
1934
+ """Whether this supervisor still owns a live session."""
1935
+
1936
+ return not self._stop.is_set() and self.outcome.terminal_state is None
1937
+
1938
+ def _stop_at_deadline(self, deadline: SessionDeadline) -> bool:
1939
+ """Settle the local process at the credential boundary, without a producer callback."""
1940
+
1941
+ if not deadline.expired():
1942
+ return False
1943
+ if self.outcome.terminal_state is None:
1944
+ self.outcome.terminal_state = "capability_expired"
1945
+ self._stop.set()
1946
+ return True
1947
+
1948
+ # -- the run ----------------------------------------------------------------------------
1949
+
1950
+ def run(self) -> dict[str, Any]:
1951
+ """Own one session end to end. Never raises; the outcome carries the failure code."""
1952
+
1953
+ if self._process_isolation:
1954
+ return self._run_isolated()
1955
+ return self._run_in_process()
1956
+
1957
+ def _run_isolated(self) -> dict[str, Any]:
1958
+ """Run the capability-bearing half in a child and enforce its expiry from this parent.
1959
+
1960
+ The parent receives only a session ID and absolute expiry, never the capability. It is
1961
+ therefore able to wait, kill and return exit-0 semantics even if a reader or native parser
1962
+ in the child never returns to Python or holds the child GIL. The child arms its own
1963
+ default-action ``SIGALRM`` before it acknowledges the lease, closing the small scheduling
1964
+ gap before this parent has installed its paired-clock monitor.
1965
+ """
1966
+
1967
+ if self._bootstrap is not None:
1968
+ raise SessionWorkerError(
1969
+ "SESSION_SUPERVISOR_CONFIGURATION_INVALID",
1970
+ "a capability-bearing bootstrap cannot enter process isolation",
1971
+ )
1972
+ if not hasattr(os, "fork") or not hasattr(signal, "setitimer"):
1973
+ raise SessionWorkerError(
1974
+ "SESSION_PROCESS_ISOLATION_UNAVAILABLE",
1975
+ "the session deadline requires POSIX process and interval-timer support",
1976
+ )
1977
+ previous_sigterm: Any | None = None
1978
+ if threading.current_thread() is threading.main_thread():
1979
+ previous_sigterm = signal.getsignal(signal.SIGTERM)
1980
+
1981
+ def request_stop(_signum: int, _frame: Any) -> None:
1982
+ # The monitor owns no capability. Do only the async-safe state change here; the
1983
+ # ordinary monitor path kills/reaps every tracked process group.
1984
+ self._stop.set()
1985
+
1986
+ # Install before the fork. A rejected/timeout legacy bootstrap must be able to signal
1987
+ # its outer monitor even if its inner child has already made a separate process group.
1988
+ signal.signal(signal.SIGTERM, request_stop)
1989
+ prior_subreaper = False
1990
+ subreaper_enabled = False
1991
+ baseline_children = frozenset()
1992
+ if sys.platform == "linux":
1993
+ try:
1994
+ prior_subreaper = _child_subreaper_enabled()
1995
+ _set_child_subreaper(True)
1996
+ subreaper_enabled = True
1997
+ # The monitor is dedicated to one session. A baseline still makes its teardown
1998
+ # exact if an embedding process happened to own a child before supervision began.
1999
+ baseline_children = _linux_direct_child_pids()
2000
+ except OSError as error:
2001
+ if subreaper_enabled:
2002
+ with contextlib.suppress(OSError):
2003
+ _set_child_subreaper(prior_subreaper)
2004
+ if previous_sigterm is not None:
2005
+ signal.signal(signal.SIGTERM, previous_sigterm)
2006
+ raise SessionWorkerError(
2007
+ "SESSION_PROCESS_ISOLATION_UNAVAILABLE",
2008
+ "could not establish isolated-session descendant supervision",
2009
+ ) from error
2010
+ try:
2011
+ child_to_parent_read, child_to_parent_write = os.pipe()
2012
+ parent_to_child_read, parent_to_child_write = os.pipe()
2013
+ except OSError as error:
2014
+ if subreaper_enabled:
2015
+ with contextlib.suppress(OSError):
2016
+ _set_child_subreaper(prior_subreaper)
2017
+ if previous_sigterm is not None:
2018
+ signal.signal(signal.SIGTERM, previous_sigterm)
2019
+ raise SessionWorkerError(
2020
+ "SESSION_PROCESS_ISOLATION_UNAVAILABLE",
2021
+ "could not allocate isolated-session control pipes",
2022
+ ) from error
2023
+ if self._stop.is_set():
2024
+ for descriptor in (
2025
+ child_to_parent_read,
2026
+ child_to_parent_write,
2027
+ parent_to_child_read,
2028
+ parent_to_child_write,
2029
+ ):
2030
+ with contextlib.suppress(OSError):
2031
+ os.close(descriptor)
2032
+ if subreaper_enabled:
2033
+ with contextlib.suppress(OSError):
2034
+ _set_child_subreaper(prior_subreaper)
2035
+ if previous_sigterm is not None:
2036
+ signal.signal(signal.SIGTERM, previous_sigterm)
2037
+ self.outcome.terminal_state = "service_stopped"
2038
+ return self.outcome.to_dict()
2039
+ try:
2040
+ child_pid = os.fork()
2041
+ except OSError as error:
2042
+ for descriptor in (
2043
+ child_to_parent_read,
2044
+ child_to_parent_write,
2045
+ parent_to_child_read,
2046
+ parent_to_child_write,
2047
+ ):
2048
+ os.close(descriptor)
2049
+ if subreaper_enabled:
2050
+ with contextlib.suppress(OSError):
2051
+ _set_child_subreaper(prior_subreaper)
2052
+ if previous_sigterm is not None:
2053
+ signal.signal(signal.SIGTERM, previous_sigterm)
2054
+ raise SessionWorkerError(
2055
+ "SESSION_PROCESS_ISOLATION_UNAVAILABLE", "could not start an isolated session child"
2056
+ ) from error
2057
+ if child_pid == 0: # pragma: no cover - exercised through the parent process contract.
2058
+ if previous_sigterm is not None:
2059
+ signal.signal(signal.SIGTERM, previous_sigterm)
2060
+ os.close(child_to_parent_read)
2061
+ os.close(parent_to_child_write)
2062
+ self._process_isolation = False
2063
+ try:
2064
+ # Give this one capability-bearing process its own group before it can create a
2065
+ # sandbox. The parent kills this group at expiry; separate sandbox groups are
2066
+ # reparented to the Linux subreaper and killed/reaped there.
2067
+ os.setsid()
2068
+ summary = self._run_in_process(
2069
+ on_lease=lambda lease, deadline: self._announce_isolated_lease(
2070
+ child_to_parent_write, parent_to_child_read, lease, deadline
2071
+ )
2072
+ )
2073
+ self._write_isolated_message(
2074
+ child_to_parent_write, {"kind": "summary", "summary": summary}
2075
+ )
2076
+ os._exit(0)
2077
+ except BaseException as error:
2078
+ # The parent gets a bounded, non-secret diagnostic instead of a traceback that
2079
+ # could retain the capability-bearing stack. ``_run_in_process`` normally
2080
+ # turns operational failures into a summary; this is only its own control plane.
2081
+ with contextlib.suppress(Exception):
2082
+ self._write_isolated_message(
2083
+ child_to_parent_write,
2084
+ {
2085
+ "kind": "error",
2086
+ "code": getattr(error, "code", "SESSION_CHILD_FAILED"),
2087
+ },
2088
+ )
2089
+ os._exit(1)
2090
+ finally:
2091
+ with contextlib.suppress(OSError):
2092
+ os.close(child_to_parent_write)
2093
+ with contextlib.suppress(OSError):
2094
+ os.close(parent_to_child_read)
2095
+
2096
+ os.close(child_to_parent_write)
2097
+ os.close(parent_to_child_read)
2098
+ # A fork duplicates a legacy nonce but not a capability (the child performs exchange).
2099
+ # Retire the parent's copy before it waits so the service process cannot retain it.
2100
+ if self._nonce is not None:
2101
+ self._nonce[:] = b"\x00" * len(self._nonce)
2102
+ try:
2103
+ if self._stop.is_set():
2104
+ self.outcome.terminal_state = "service_stopped"
2105
+ return self._terminate_isolated_child(child_pid)
2106
+ if self._on_isolation_ready is not None:
2107
+ try:
2108
+ self._on_isolation_ready()
2109
+ except Exception as error:
2110
+ self.outcome.failure_code = "SESSION_ISOLATION_READY_FAILED"
2111
+ LOGGER.warning(
2112
+ "isolated-session readiness acknowledgement failed: %s",
2113
+ type(error).__name__,
2114
+ )
2115
+ return self._terminate_isolated_child(child_pid)
2116
+ return self._monitor_isolated_child(
2117
+ child_pid, child_to_parent_read, parent_to_child_write
2118
+ )
2119
+ finally:
2120
+ if previous_sigterm is not None:
2121
+ signal.signal(signal.SIGTERM, previous_sigterm)
2122
+ with contextlib.suppress(OSError):
2123
+ os.close(child_to_parent_read)
2124
+ with contextlib.suppress(OSError):
2125
+ os.close(parent_to_child_write)
2126
+ if subreaper_enabled:
2127
+ try:
2128
+ self._kill_adopted_session_descendants(baseline_children)
2129
+ finally:
2130
+ _set_child_subreaper(prior_subreaper)
2131
+
2132
+ @staticmethod
2133
+ def _write_isolated_message(descriptor: int, payload: dict[str, Any]) -> None:
2134
+ raw = canonical_json_bytes(payload) + b"\n"
2135
+ if len(raw) > MAX_SESSION_SUPERVISOR_MESSAGE_BYTES:
2136
+ raise SessionWorkerError(
2137
+ "SESSION_CHILD_PROTOCOL_INVALID", "isolated-session message exceeds its byte limit"
2138
+ )
2139
+ view = memoryview(raw)
2140
+ try:
2141
+ while view:
2142
+ written = os.write(descriptor, view)
2143
+ if written <= 0:
2144
+ raise OSError("short isolated-session pipe write")
2145
+ view = view[written:]
2146
+ finally:
2147
+ view.release()
2148
+
2149
+ @staticmethod
2150
+ def _read_isolated_message(
2151
+ descriptor: int, buffer: bytearray, *, timeout: float | None
2152
+ ) -> dict[str, Any] | None:
2153
+ readable, _unused, _errors = select.select([descriptor], [], [], timeout)
2154
+ if not readable:
2155
+ return None
2156
+ chunk = os.read(descriptor, 64 * 1024)
2157
+ if not chunk:
2158
+ raise EOFError("isolated session child closed its control pipe")
2159
+ buffer.extend(chunk)
2160
+ if len(buffer) > MAX_SESSION_SUPERVISOR_MESSAGE_BYTES:
2161
+ raise SessionWorkerError(
2162
+ "SESSION_CHILD_PROTOCOL_INVALID", "isolated-session message exceeds its byte limit"
2163
+ )
2164
+ newline = buffer.find(b"\n")
2165
+ if newline < 0:
2166
+ return None
2167
+ raw = bytes(buffer[:newline])
2168
+ del buffer[: newline + 1]
2169
+ value = _strict_json(raw, "isolated-session message")
2170
+ if not isinstance(value, dict):
2171
+ raise SessionWorkerError(
2172
+ "SESSION_CHILD_PROTOCOL_INVALID", "isolated-session message is not an object"
2173
+ )
2174
+ return value
2175
+
2176
+ def _announce_isolated_lease(
2177
+ self,
2178
+ write_descriptor: int,
2179
+ read_descriptor: int,
2180
+ lease: SessionLease,
2181
+ deadline: SessionDeadline,
2182
+ ) -> None:
2183
+ remaining = deadline.remaining_seconds()
2184
+ if remaining <= 0:
2185
+ raise SessionWorkerError(
2186
+ "SESSION_CAPABILITY_EXPIRED", "the session capability expired before supervision"
2187
+ )
2188
+ # SIG_DFL is a kernel-enforced deadline: a child blocked in native code cannot delay it.
2189
+ signal.signal(signal.SIGALRM, signal.SIG_DFL)
2190
+ signal.setitimer(signal.ITIMER_REAL, remaining)
2191
+ expiry = lease.capability_expires_at.isoformat(timespec="microseconds").replace(
2192
+ "+00:00", "Z"
2193
+ )
2194
+ self._write_isolated_message(
2195
+ write_descriptor,
2196
+ {"kind": "lease", "session_id": lease.session_id, "capability_expires_at": expiry},
2197
+ )
2198
+ response = bytearray()
2199
+ while True:
2200
+ command = self._read_isolated_message(read_descriptor, response, timeout=None)
2201
+ if command is not None:
2202
+ break
2203
+ if command != {"kind": "start"}:
2204
+ raise SessionWorkerError(
2205
+ "SESSION_CHILD_PROTOCOL_INVALID", "isolated-session parent did not authorize start"
2206
+ )
2207
+
2208
+ def _monitor_isolated_child(
2209
+ self, child_pid: int, read_descriptor: int, write_descriptor: int
2210
+ ) -> dict[str, Any]:
2211
+ buffer = bytearray()
2212
+ deadline: SessionDeadline | None = None
2213
+ while True:
2214
+ if self._stop.is_set():
2215
+ if self.outcome.terminal_state is None:
2216
+ self.outcome.terminal_state = "service_stopped"
2217
+ return self._terminate_isolated_child(child_pid)
2218
+ timeout = 1.0 if deadline is None else deadline.bounded_wait(1.0)
2219
+ if deadline is not None and timeout <= 0:
2220
+ return self._expire_isolated_child(child_pid)
2221
+ try:
2222
+ message = self._read_isolated_message(read_descriptor, buffer, timeout=timeout)
2223
+ except EOFError:
2224
+ status = self._wait_for_isolated_child(child_pid)
2225
+ if deadline is not None and (
2226
+ deadline.expired()
2227
+ or (os.WIFSIGNALED(status) and os.WTERMSIG(status) == signal.SIGALRM)
2228
+ ):
2229
+ return self._deadline_summary()
2230
+ self.outcome.failure_code = "SESSION_CHILD_EXITED"
2231
+ return self.outcome.to_dict()
2232
+ if message is None:
2233
+ continue
2234
+ kind = message.get("kind")
2235
+ if kind == "lease":
2236
+ if deadline is not None or set(message) != {
2237
+ "kind",
2238
+ "session_id",
2239
+ "capability_expires_at",
2240
+ }:
2241
+ self.outcome.failure_code = "SESSION_CHILD_PROTOCOL_INVALID"
2242
+ return self._terminate_isolated_child(child_pid)
2243
+ try:
2244
+ session_id = _uuid_text(message["session_id"], "isolated session id")
2245
+ expires_at = _parse_utc_timestamp(
2246
+ message["capability_expires_at"], "isolated capability_expires_at"
2247
+ )
2248
+ deadline = SessionDeadline.from_expiry(
2249
+ expires_at,
2250
+ wall_clock=self._wall_clock,
2251
+ monotonic_clock=self._monotonic_clock,
2252
+ )
2253
+ except SessionWorkerError as error:
2254
+ if error.code == "SESSION_CAPABILITY_EXPIRED":
2255
+ return self._expire_isolated_child(child_pid)
2256
+ self.outcome.failure_code = error.code
2257
+ return self._terminate_isolated_child(child_pid)
2258
+ self.outcome.session_id = session_id
2259
+ self._write_isolated_message(write_descriptor, {"kind": "start"})
2260
+ continue
2261
+ if kind == "summary" and set(message) == {"kind", "summary"}:
2262
+ status = self._wait_for_isolated_child(child_pid)
2263
+ summary = message["summary"]
2264
+ if not isinstance(summary, dict):
2265
+ self.outcome.failure_code = "SESSION_CHILD_PROTOCOL_INVALID"
2266
+ return self.outcome.to_dict()
2267
+ if deadline is not None and (
2268
+ deadline.expired()
2269
+ or (os.WIFSIGNALED(status) and os.WTERMSIG(status) == signal.SIGALRM)
2270
+ ):
2271
+ return self._deadline_summary()
2272
+ return summary
2273
+ if (
2274
+ kind == "error"
2275
+ and set(message) == {"kind", "code"}
2276
+ and isinstance(message["code"], str)
2277
+ ):
2278
+ self.outcome.failure_code = message["code"]
2279
+ else:
2280
+ self.outcome.failure_code = "SESSION_CHILD_PROTOCOL_INVALID"
2281
+ return self._terminate_isolated_child(child_pid)
2282
+
2283
+ def _expire_isolated_child(self, child_pid: int) -> dict[str, Any]:
2284
+ self._kill_isolated_child_group(child_pid)
2285
+ self._wait_for_isolated_child(child_pid)
2286
+ return self._deadline_summary()
2287
+
2288
+ def _terminate_isolated_child(self, child_pid: int) -> dict[str, Any]:
2289
+ self._kill_isolated_child_group(child_pid)
2290
+ self._wait_for_isolated_child(child_pid)
2291
+ return self.outcome.to_dict()
2292
+
2293
+ @staticmethod
2294
+ def _kill_isolated_child_group(child_pid: int) -> None:
2295
+ """Kill the child process group, with a PID fallback before ``setsid`` can complete."""
2296
+
2297
+ try:
2298
+ os.killpg(child_pid, signal.SIGKILL)
2299
+ return
2300
+ except (PermissionError, ProcessLookupError):
2301
+ # Before the child's first control message it may not yet have made itself a group
2302
+ # leader. Some Darwin process-group races also report EPERM after the group leader has
2303
+ # exited. Killing the exact direct PID still prevents a bootstrap/protocol error from
2304
+ # leaving the capability exchange running.
2305
+ pass
2306
+ with contextlib.suppress(PermissionError, ProcessLookupError):
2307
+ os.kill(child_pid, signal.SIGKILL)
2308
+
2309
+ @staticmethod
2310
+ def _wait_for_isolated_child(child_pid: int) -> int:
2311
+ while True:
2312
+ try:
2313
+ _pid, status = os.waitpid(child_pid, 0)
2314
+ return status
2315
+ except InterruptedError:
2316
+ continue
2317
+ except ChildProcessError:
2318
+ return 0
2319
+
2320
+ @staticmethod
2321
+ def _kill_adopted_session_descendants(baseline_children: frozenset[int]) -> None:
2322
+ """Kill/reap Linux descendants adopted after the capability child dies.
2323
+
2324
+ Sandboxes intentionally create separate process groups so parser helpers cannot affect the
2325
+ coordinator. A force-killed capability child cannot run ``close_warm_pool``; as a Linux
2326
+ subreaper, this parent adopts those groups and kills their live leader/group before each
2327
+ reap. The monitor creates no other children, and the pre-fork baseline prevents touching
2328
+ a host process's pre-existing children if this code is embedded in a test or adapter.
2329
+ """
2330
+
2331
+ if sys.platform != "linux":
2332
+ return
2333
+ deadline = time.monotonic() + 5.0
2334
+ own_group = os.getpgrp()
2335
+ while True:
2336
+ children = _linux_direct_child_pids() - baseline_children
2337
+ if not children:
2338
+ return
2339
+ for child_pid in children:
2340
+ try:
2341
+ group = os.getpgid(child_pid)
2342
+ except ProcessLookupError:
2343
+ continue
2344
+ try:
2345
+ if group == own_group:
2346
+ # The dedicated capability child makes its own group before any sandbox
2347
+ # launch. A shared group is therefore never safe to signal wholesale.
2348
+ os.kill(child_pid, signal.SIGKILL)
2349
+ else:
2350
+ os.killpg(group, signal.SIGKILL)
2351
+ except ProcessLookupError:
2352
+ continue
2353
+ for child_pid in children:
2354
+ with contextlib.suppress(ChildProcessError):
2355
+ os.waitpid(child_pid, os.WNOHANG)
2356
+ if time.monotonic() >= deadline:
2357
+ raise SessionWorkerError(
2358
+ "SESSION_DESCENDANT_CLEANUP_FAILED",
2359
+ "isolated-session descendants did not exit after the capability child",
2360
+ )
2361
+ time.sleep(0.005)
2362
+
2363
+ def _deadline_summary(self) -> dict[str, Any]:
2364
+ self.outcome.terminal_state = "capability_expired"
2365
+ self._stop.set()
2366
+ return self.outcome.to_dict()
2367
+
2368
+ def _run_in_process(
2369
+ self, *, on_lease: Callable[[SessionLease, SessionDeadline], None] | None = None
2370
+ ) -> dict[str, Any]:
2371
+ """Run the capability-bearing half, optionally after an isolated-parent handshake."""
2372
+
2373
+ if self._bootstrap is not None:
2374
+ # Transfer ownership exactly once. A duplicate Job execution must obtain a fresh
2375
+ # bootstrap through Studio; it cannot re-run this supervisor with an old capability.
2376
+ lease, credential = self._bootstrap
2377
+ self._bootstrap = None
2378
+ elif self._dispatch is not None:
2379
+ try:
2380
+ lease, credential = exchange_session_dispatch(
2381
+ config=self._config, dispatch=self._dispatch, transport=self._transport
2382
+ )
2383
+ except (SessionWorkerError, HostedWorkerError) as error:
2384
+ code = getattr(error, "code", "SESSION_LEASE_FAILED")
2385
+ if code in CONSUMED_SESSION_DISPATCH_CODES:
2386
+ self.outcome.terminal_state = "dispatch_consumed"
2387
+ else:
2388
+ self.outcome.failure_code = code
2389
+ return self.outcome.to_dict()
2390
+ else:
2391
+ assert self._invocation is not None
2392
+ assert self._nonce is not None
2393
+ try:
2394
+ lease, credential = self._lease_session(
2395
+ config=self._config,
2396
+ invocation=self._invocation,
2397
+ nonce=self._nonce,
2398
+ transport=self._transport,
2399
+ )
2400
+ except (SessionWorkerError, HostedWorkerError) as error:
2401
+ self.outcome.failure_code = getattr(error, "code", "SESSION_LEASE_FAILED")
2402
+ return self.outcome.to_dict()
2403
+ finally:
2404
+ # Zeroed the instant the legacy lease exchange is over, success or failure. One
2405
+ # use, one lifetime. Job mode has no nonce at all.
2406
+ self._nonce[:] = b"\x00" * len(self._nonce)
2407
+ self.outcome.attempt_id = str(lease.attempt_id)
2408
+ try:
2409
+ with credential:
2410
+ deadline = SessionDeadline.from_lease(
2411
+ lease,
2412
+ wall_clock=self._wall_clock,
2413
+ monotonic_clock=self._monotonic_clock,
2414
+ )
2415
+ if on_lease is not None:
2416
+ on_lease(lease, deadline)
2417
+ try:
2418
+ self.outcome.catalogue = install_hosted_catalogue()
2419
+ except Exception as error:
2420
+ LOGGER.warning("hosted catalogue install failed: %s", type(error).__name__)
2421
+ self._serve(lease, credential, deadline)
2422
+ except (SessionWorkerError, HostedWorkerError) as error:
2423
+ if getattr(error, "code", None) == "SESSION_CAPABILITY_EXPIRED":
2424
+ # A dispatch that reaches its absolute capability boundary during bootstrap is a
2425
+ # normal no-op. There is no candidate to complete or session failure to report;
2426
+ # Studio owns the durable close/reap and the Job must not retry it as an error.
2427
+ self.outcome.terminal_state = "capability_expired"
2428
+ self._stop.set()
2429
+ else:
2430
+ self.outcome.failure_code = getattr(error, "code", "SESSION_WORKER_FAILED")
2431
+ except Exception as error:
2432
+ LOGGER.exception("session worker failed")
2433
+ self.outcome.failure_code = type(error).__name__
2434
+ return self.outcome.to_dict()
2435
+
2436
+ def _serve(
2437
+ self, lease: SessionLease, credential: CredentialLease, deadline: SessionDeadline
2438
+ ) -> None:
2439
+ models, _, _, _ = _load_generated()
2440
+ with credential.borrow() as view:
2441
+ client = _authenticated_client(lease.callback_base_url, bytes(view).decode("ascii"))
2442
+ authority = _SessionAuthority(client)
2443
+ sandbox = self._sandbox_factory(self._config)
2444
+ self.outcome.warm_pool_ready = _start_warm_pool(sandbox, self._config.warm_pool_size)
2445
+ transfer: Any | None = None
2446
+ if self._reader_factory is not None:
2447
+ reader = self._reader_factory(lease)
2448
+ else:
2449
+ # A SEPARATE client from the one carrying the capability. The quarantine download is
2450
+ # signed and must reach storage with no Authorization header on it; sharing the
2451
+ # authenticated client would put the attempt capability on a request to a bucket.
2452
+ transfer = _new_transfer_client()
2453
+ reader = StudioSessionAcquisitionReader(
2454
+ authority=authority,
2455
+ models=models,
2456
+ lease=lease,
2457
+ transfer=transfer,
2458
+ results_bucket=self._config.crawler_results_bucket,
2459
+ timeout_seconds=self._config.acquisition_timeout_seconds,
2460
+ sleep=self._sleep,
2461
+ )
2462
+ executor = ProbeExecutor(
2463
+ sandbox=sandbox,
2464
+ reader=reader,
2465
+ context=lease.context,
2466
+ )
2467
+ progress = _SessionProgress(
2468
+ authority=authority,
2469
+ models=models,
2470
+ lease=lease,
2471
+ deadline=deadline,
2472
+ )
2473
+ heartbeat = _Heartbeat(
2474
+ send=lambda: self._heartbeat(authority, models, lease, deadline),
2475
+ interval_seconds=lease.heartbeat_interval_seconds,
2476
+ stop=self._stop,
2477
+ deadline=deadline,
2478
+ on_expired=lambda: self._stop_at_deadline(deadline),
2479
+ )
2480
+ try:
2481
+ if self._stop_at_deadline(deadline):
2482
+ return
2483
+ # `attempt:start` is what clears `warming`. It is the first fenced callback and it is
2484
+ # sent before the warm pool is used for anything, so the session reports ready as soon
2485
+ # as a worker demonstrably exists rather than after the first probe.
2486
+ _response(
2487
+ "startProducerAttempt",
2488
+ authority.start_attempt.sync_detailed(
2489
+ attempt_id=lease.attempt_id,
2490
+ body=_session_callback(models, lease),
2491
+ idempotency_key=_idempotency("session", lease.attempt_id, "start"),
2492
+ ),
2493
+ 200,
2494
+ )
2495
+ if self._stop_at_deadline(deadline):
2496
+ return
2497
+ progress.emit(
2498
+ "attempt_started",
2499
+ {"role": SESSION_WORKER_ROLE, "stage_total": 0},
2500
+ )
2501
+ # Flushed here rather than left pending for the first probe. A session that opens and
2502
+ # then sits quiet is the ordinary case -- a person is reading before they ask anything
2503
+ # -- and its stream should say "the worker is here" immediately rather than at the
2504
+ # moment the first probe lands, which is the one moment nobody needs telling.
2505
+ progress.flush()
2506
+ if self._stop_at_deadline(deadline):
2507
+ return
2508
+ heartbeat.start()
2509
+ self._claim_loop(authority, models, lease, executor, progress, deadline)
2510
+ finally:
2511
+ self._stop.set()
2512
+ progress.emit(
2513
+ "attempt_finished",
2514
+ {"outcome": self.outcome.terminal_state or "stopped"},
2515
+ )
2516
+ progress.close()
2517
+ try:
2518
+ sandbox.close_warm_pool()
2519
+ except Exception:
2520
+ LOGGER.debug("warm pool close failed", exc_info=True)
2521
+ if transfer is not None:
2522
+ with contextlib.suppress(Exception):
2523
+ transfer.close()
2524
+ client.get_httpx_client().close()
2525
+
2526
+ def _claim_loop(
2527
+ self,
2528
+ authority: Any,
2529
+ models: Any,
2530
+ lease: SessionLease,
2531
+ executor: ProbeExecutor,
2532
+ progress: _SessionProgress,
2533
+ deadline: SessionDeadline,
2534
+ ) -> None:
2535
+ interval = lease.claim_poll_interval_seconds or FALLBACK_CLAIM_POLL_SECONDS
2536
+ failures = 0
2537
+ poll = 0
2538
+ while not self._stop.is_set():
2539
+ if self._stop_at_deadline(deadline):
2540
+ return
2541
+ poll += 1
2542
+ try:
2543
+ answer = _response(
2544
+ "claimResearchProbe",
2545
+ authority.claim_probe.sync_detailed(
2546
+ attempt_id=lease.attempt_id,
2547
+ body=_claim_command(models, lease),
2548
+ # Claim is NOT idempotent-keyed on a constant. Every poll must be able to
2549
+ # return a different probe; a fixed key would replay the first answer for
2550
+ # the life of the session and the worker would answer one probe forever.
2551
+ # A monotonic COUNTER, not a clock reading. Two polls inside one clock
2552
+ # tick would share a key, and a shared claim key means Studio replaying
2553
+ # the first answer -- the worker would re-answer one probe while the
2554
+ # queue grew behind it. A counter cannot collide.
2555
+ idempotency_key=_idempotency(
2556
+ "session", lease.attempt_id, "claim", str(poll)
2557
+ ),
2558
+ ),
2559
+ 200,
2560
+ ).to_dict()
2561
+ except HostedWorkerError as error:
2562
+ if self._stop_at_deadline(deadline):
2563
+ return
2564
+ # A stale fence is Studio saying this worker is no longer the session's worker.
2565
+ # There is no recovery from that and retrying would be a second worker arguing
2566
+ # with the first.
2567
+ if error.code in {"ATTEMPT_FENCE_STALE", "ATTEMPT_STALE", "AUTHORIZATION_DENIED"}:
2568
+ self.outcome.failure_code = error.code
2569
+ return
2570
+ failures += 1
2571
+ if failures >= MAX_CONSECUTIVE_CLAIM_FAILURES:
2572
+ self.outcome.failure_code = error.code
2573
+ return
2574
+ wait_seconds = deadline.bounded_wait(interval)
2575
+ if wait_seconds <= 0:
2576
+ self._stop_at_deadline(deadline)
2577
+ return
2578
+ self._sleep(wait_seconds)
2579
+ continue
2580
+ if self._stop_at_deadline(deadline):
2581
+ return
2582
+ failures = 0
2583
+ state = str(answer.get("session_state", ""))
2584
+ interval = float(answer.get("claim_poll_interval_seconds") or interval)
2585
+ if state in {"closed", "failed"}:
2586
+ # The whole of the exit protocol. No `:complete`, no `:fail`: the session is
2587
+ # already terminal server-side, so both would be refused, and Studio retires the
2588
+ # leftover attempt row on its next sweep.
2589
+ self.outcome.terminal_state = state
2590
+ return
2591
+ if not answer.get("claimed"):
2592
+ wait_seconds = deadline.bounded_wait(interval)
2593
+ if wait_seconds <= 0:
2594
+ self._stop_at_deadline(deadline)
2595
+ return
2596
+ self._sleep(wait_seconds)
2597
+ continue
2598
+ self._answer_probe(
2599
+ authority,
2600
+ models,
2601
+ lease,
2602
+ executor,
2603
+ progress,
2604
+ deadline,
2605
+ answer.get("probe"),
2606
+ )
2607
+
2608
+ def _answer_probe(
2609
+ self,
2610
+ authority: Any,
2611
+ models: Any,
2612
+ lease: SessionLease,
2613
+ executor: ProbeExecutor,
2614
+ progress: _SessionProgress,
2615
+ deadline: SessionDeadline,
2616
+ probe: Any,
2617
+ ) -> None:
2618
+ if self._stop_at_deadline(deadline):
2619
+ return
2620
+ if not isinstance(probe, dict):
2621
+ self.outcome.failure_code = "PROBE_CLAIM_INVALID"
2622
+ self._stop.set()
2623
+ return
2624
+ probe_id = str(probe.get("probe_id", ""))
2625
+ probe_kind = str(probe.get("probe_kind", ""))
2626
+ ordinal = probe.get("ordinal")
2627
+ ordinal = ordinal if type(ordinal) is int and ordinal > 0 else 1
2628
+ request = probe.get("request") or {}
2629
+ source_id = str(request.get("source_id", "")) if isinstance(request, dict) else ""
2630
+ # ONE condition for both halves of the pair. A watcher's whole use for this stream is that
2631
+ # a start it sees will be followed by an end, so the two records must be emittable or not
2632
+ # emittable together -- never a start whose end could not be spelled. `source_id` and
2633
+ # `probe_kind` are the two facts a settled record needs, and Studio's claim contract
2634
+ # requires both, so this is false only for a claim this worker could not answer anyway: a
2635
+ # body with no source, or a kind outside the four -- which `execute` refuses too, and which
2636
+ # a drift test keeps equal to the generated enum.
2637
+ narrated = bool(source_id) and probe_kind in session_probes.PROBE_KINDS
2638
+ failure_code: str | None = "PROBE_FAILED"
2639
+ wire_code = "PROBE_FAILED"
2640
+ display_code = "PROBE_FAILED"
2641
+ retry_after_seconds = 0
2642
+ result: dict[str, Any] | None = None
2643
+ try:
2644
+ if narrated:
2645
+ progress.emit(
2646
+ "source_probe_started",
2647
+ {"source_id": source_id, "probe_kind": probe_kind, "ordinal": ordinal},
2648
+ )
2649
+ # Flushed BEFORE the probe runs, not left to ride out with its own completion. A
2650
+ # probe may hold this thread for the whole `probe_timeout_seconds` ceiling, and a
2651
+ # start that arrives at the same instant as its end has told a watcher nothing it
2652
+ # did not already know.
2653
+ progress.flush()
2654
+ result = executor.execute(probe_id=probe_id, probe_kind=probe_kind, request=request)
2655
+ failure_code = None
2656
+ wire_code = progress_events.PROBE_OUTCOME_OK_CODE
2657
+ display_code = progress_events.PROBE_OUTCOME_OK_CODE
2658
+ except Exception as error:
2659
+ # Every exception, coded or not, and the code is KEPT when there is one. The two-clause
2660
+ # form this replaces named three classes and let everything else collapse to
2661
+ # `PROBE_FAILED` -- which silently swallowed the Clean room's own `SANDBOX_TIMEOUT`
2662
+ # (`AcquisitionSecurityError`, named in neither clause), turning the one failure a
2663
+ # client can act on into the one that says nothing.
2664
+ failure_code = _probe_failure_code(error)
2665
+ wire_code = failure_code
2666
+ display_code, retry_after_seconds = _probe_display_failure(error, wire_code)
2667
+ if failure_code == "PROBE_FAILED":
2668
+ LOGGER.warning("probe %s failed: %s", probe_id, type(error).__name__)
2669
+ if self._stop_at_deadline(deadline):
2670
+ return
2671
+ # Studio's `ContractProbeResultCommand.outcome` enum is two values. The four-value reading
2672
+ # below is a REFINEMENT for the unsealed timeline and never a second answer: the probe
2673
+ # record settles as `succeeded` or `failed`, exactly as before.
2674
+ settled = "succeeded" if failure_code is None else "failed"
2675
+ self._report(authority, models, lease, deadline, probe_id, settled, result, failure_code)
2676
+ if failure_code is None:
2677
+ self.outcome.probes_answered += 1
2678
+ else:
2679
+ self.outcome.probes_failed += 1
2680
+ if narrated:
2681
+ # EVERY terminal outcome, not only success. `source_probe_settled` travels as Studio's
2682
+ # `source_probe_completed` event type -- the exact type the research-session design
2683
+ # names for a settled probe -- so a client watching the existing per-run SSE stream
2684
+ # sees every probe end live, with no new event type and no second stream. Emitting
2685
+ # only on success (the shape this replaces) meant a watcher saw a start and never an
2686
+ # end, which with no source channel wired was every probe there is.
2687
+ #
2688
+ # `ordinal` and no denominator: how many probes a session will raise is not knowable
2689
+ # while one of them runs. The `source_acquire_*` pair this replaces demanded an
2690
+ # `index` and a `total`, so the worker sent the ordinal for both and every probe read
2691
+ # "N of N" -- a fraction a progress rail will happily draw and that never means
2692
+ # anything.
2693
+ progress.emit(
2694
+ "source_probe_settled",
2695
+ {
2696
+ "source_id": source_id,
2697
+ "probe_kind": probe_kind,
2698
+ "ordinal": ordinal,
2699
+ "outcome": _probe_progress_outcome(
2700
+ None if failure_code is None else display_code
2701
+ ),
2702
+ "code": wire_code,
2703
+ "display_code": display_code,
2704
+ "retry_after_seconds": retry_after_seconds,
2705
+ "bytes": len(canonical_json_bytes(result)) if result is not None else 0,
2706
+ },
2707
+ )
2708
+ progress.flush()
2709
+
2710
+ def _report(
2711
+ self,
2712
+ authority: Any,
2713
+ models: Any,
2714
+ lease: SessionLease,
2715
+ deadline: SessionDeadline,
2716
+ probe_id: str,
2717
+ outcome: str,
2718
+ result: Mapping[str, Any] | None,
2719
+ failure_code: str | None,
2720
+ ) -> None:
2721
+ if self._stop_at_deadline(deadline):
2722
+ return
2723
+ wire: dict[str, Any] = {
2724
+ "schema_version": START_SCHEMA_VERSION,
2725
+ "workspace_id": str(lease.workspace_id),
2726
+ "run_id": str(lease.run_id),
2727
+ "attempt_id": str(lease.attempt_id),
2728
+ "fence": lease.fence,
2729
+ "session_id": lease.session_id,
2730
+ "probe_id": probe_id,
2731
+ "outcome": outcome,
2732
+ }
2733
+ if result is not None:
2734
+ wire["result"] = dict(result)
2735
+ if failure_code is not None:
2736
+ wire["failure_code"] = failure_code
2737
+ try:
2738
+ _response(
2739
+ "reportResearchProbeResult",
2740
+ authority.report_probe_result.sync_detailed(
2741
+ attempt_id=lease.attempt_id,
2742
+ probe_id=UUID(probe_id),
2743
+ body=models.ContractProbeResultCommand.from_dict(wire),
2744
+ # Keyed on the probe, so a retried report is the same report. One answer per
2745
+ # probe is enforced server-side as well -- a second is
2746
+ # RESEARCH_PROBE_ALREADY_SETTLED -- and the key means an ordinary network
2747
+ # retry never becomes that.
2748
+ idempotency_key=_idempotency("session", lease.attempt_id, "result", probe_id),
2749
+ ),
2750
+ 200,
2751
+ )
2752
+ except HostedWorkerError as error:
2753
+ if error.code == "RESEARCH_PROBE_ALREADY_SETTLED":
2754
+ # Already answered, by us, before a retry. Nothing to do and nothing wrong: the
2755
+ # client has an answer, which is the only thing that matters.
2756
+ LOGGER.info("probe %s was already settled", probe_id)
2757
+ return
2758
+ raise
2759
+
2760
+ def _heartbeat(
2761
+ self,
2762
+ authority: Any,
2763
+ models: Any,
2764
+ lease: SessionLease,
2765
+ deadline: SessionDeadline | None = None,
2766
+ ) -> None:
2767
+ if deadline is not None and self._stop_at_deadline(deadline):
2768
+ return
2769
+ _response(
2770
+ "heartbeatProducerAttempt",
2771
+ authority.heartbeat_attempt.sync_detailed(
2772
+ attempt_id=lease.attempt_id,
2773
+ body=_session_callback(models, lease),
2774
+ idempotency_key=_idempotency("session", lease.attempt_id, "heartbeat"),
2775
+ ),
2776
+ 200,
2777
+ )
2778
+
2779
+
2780
+ def _session_callback(models: Any, lease: SessionLease) -> Any:
2781
+ return models.ContractAttemptCallback(
2782
+ schema_version=START_SCHEMA_VERSION,
2783
+ workspace_id=lease.workspace_id,
2784
+ run_id=lease.run_id,
2785
+ attempt_id=lease.attempt_id,
2786
+ fence=lease.fence,
2787
+ )
2788
+
2789
+
2790
+ def _claim_command(models: Any, lease: SessionLease) -> Any:
2791
+ return models.ContractProbeClaimCommand.from_dict(
2792
+ {
2793
+ "schema_version": START_SCHEMA_VERSION,
2794
+ "workspace_id": str(lease.workspace_id),
2795
+ "run_id": str(lease.run_id),
2796
+ "attempt_id": str(lease.attempt_id),
2797
+ "fence": lease.fence,
2798
+ "session_id": lease.session_id,
2799
+ }
2800
+ )
2801
+
2802
+
2803
+ class _SessionProgress:
2804
+ """Unsealed progress onto the session run's existing event log. Never fails a session.
2805
+
2806
+ The same promise ``hosted_worker._ProgressEmitter`` makes, kept for the same reason: a client
2807
+ that loses timeline resolution has lost a nicety, and a session that died because a progress
2808
+ append was refused has lost the thing it exists for. Records go out one at a time rather than
2809
+ batched, because a session's whole product claim is that a probe is visible immediately, and a
2810
+ 250 ms coalescing window is a quarter of the follow-up budget.
2811
+ """
2812
+
2813
+ def __init__(
2814
+ self,
2815
+ *,
2816
+ authority: Any,
2817
+ models: Any,
2818
+ lease: SessionLease,
2819
+ deadline: SessionDeadline | None = None,
2820
+ ) -> None:
2821
+ self._authority = authority
2822
+ self._models = models
2823
+ self._lease = lease
2824
+ self._deadline = deadline
2825
+ self._sequence = 0
2826
+ self._batch = 0
2827
+ self._pending: list[Any] = []
2828
+ self._armed = True
2829
+
2830
+ @property
2831
+ def armed(self) -> bool:
2832
+ return self._armed
2833
+
2834
+ def emit(self, name: str, facts: Mapping[str, Any]) -> None:
2835
+ if self._deadline is not None and self._deadline.expired():
2836
+ self._armed = False
2837
+ self._pending.clear()
2838
+ return
2839
+ if not self._armed:
2840
+ return
2841
+ if self._sequence >= progress_events.MAX_PROGRESS_EVENTS_PER_ATTEMPT:
2842
+ self._armed = False
2843
+ self._pending.clear()
2844
+ return
2845
+ try:
2846
+ self._sequence += 1
2847
+ payload, digest = progress_events.progress_payload_and_digest(name, facts)
2848
+ identifier = bytearray(
2849
+ hashlib.sha256(
2850
+ f"{self._lease.attempt_id}:{self._sequence}:{name}:{digest}".encode()
2851
+ ).digest()[:16]
2852
+ )
2853
+ identifier[6] = (identifier[6] & 0x0F) | 0x40
2854
+ identifier[8] = (identifier[8] & 0x3F) | 0x80
2855
+ self._pending.append(
2856
+ self._models.ContractProducerEventCommand(
2857
+ schema_version=START_SCHEMA_VERSION,
2858
+ workspace_id=self._lease.workspace_id,
2859
+ run_id=self._lease.run_id,
2860
+ attempt_id=self._lease.attempt_id,
2861
+ fence=self._lease.fence,
2862
+ client_event_id=UUID(bytes=bytes(identifier)),
2863
+ event_type=self._models.ContractProducerEventCommandEventType(
2864
+ progress_events.STUDIO_EVENT_TYPE[name]
2865
+ ),
2866
+ payload_digest=digest,
2867
+ payload=self._models.ContractProducerEventCommandPayload.from_dict(payload),
2868
+ )
2869
+ )
2870
+ except Exception:
2871
+ LOGGER.debug("session progress record dropped", exc_info=True)
2872
+
2873
+ def flush(self) -> None:
2874
+ if self._deadline is not None and self._deadline.expired():
2875
+ self._armed = False
2876
+ self._pending.clear()
2877
+ return
2878
+ if not self._pending:
2879
+ return
2880
+ batch, self._pending = self._pending, []
2881
+ self._batch += 1
2882
+ if not self._armed:
2883
+ return
2884
+ try:
2885
+ _response(
2886
+ "appendProducerEvents",
2887
+ self._authority.append_events.sync_detailed(
2888
+ attempt_id=self._lease.attempt_id,
2889
+ body=batch,
2890
+ idempotency_key=_idempotency(
2891
+ "session", self._lease.attempt_id, "event-progress", str(self._batch)
2892
+ ),
2893
+ ),
2894
+ 200,
2895
+ )
2896
+ except Exception:
2897
+ LOGGER.debug("session progress append failed", exc_info=True)
2898
+
2899
+ def close(self) -> None:
2900
+ try:
2901
+ self.flush()
2902
+ finally:
2903
+ self._armed = False
2904
+
2905
+
2906
+ def _start_warm_pool(sandbox: Any, size: int) -> int:
2907
+ """Ask for a warm pool and believe the answer.
2908
+
2909
+ ``start_warm_pool`` returns the number now ready, which on Linux under a bounded cgroup root is
2910
+ zero by design and is not an error -- the first probe simply spawns cold. The lazy-import gain
2911
+ that came with the pool (#306) applies either way, so a zero costs latency and nothing else.
2912
+ """
2913
+
2914
+ if size <= 0:
2915
+ return 0
2916
+ try:
2917
+ return int(sandbox.start_warm_pool(size=size))
2918
+ except Exception:
2919
+ LOGGER.info("warm sandbox pool unavailable; probes will spawn cold")
2920
+ return 0
2921
+
2922
+
2923
+ def _default_sandbox(config: SessionWorkerConfig) -> Any:
2924
+ from mostlyright.data_harness.acquisition.sandbox import (
2925
+ HOSTED_SESSION_CONTAINER_MEMORY_BYTES,
2926
+ CrawlerSandbox,
2927
+ )
2928
+
2929
+ staging_root = Path(tempfile.mkdtemp(prefix="mr-session-probe-")).resolve()
2930
+ return CrawlerSandbox(
2931
+ staging_root=staging_root,
2932
+ timeout_seconds=config.probe_timeout_seconds,
2933
+ external_network_policy_attestation=config.external_network_policy_attestation,
2934
+ hosted_session_container_memory_bytes=HOSTED_SESSION_CONTAINER_MEMORY_BYTES,
2935
+ )
2936
+
2937
+
2938
+ # ==========================================================================================
2939
+ # The HTTP surface
2940
+ # ==========================================================================================
2941
+
2942
+
2943
+ class _LegacySessionProcess:
2944
+ """A capability-free Service-side handle for one fresh legacy-session interpreter.
2945
+
2946
+ ``ThreadingHTTPServer`` calls :meth:`SessionWorkerService.start` from a request thread. It is
2947
+ unsafe to ``fork()`` there and then execute Python, HTTP clients or parsers: another thread
2948
+ may have held one of their runtime locks at the instant of the fork. This handle uses the
2949
+ platform's subprocess exec path only, then the fresh interpreter creates the monitored
2950
+ capability child from its main thread. The Service process therefore retains neither the
2951
+ lease capability nor a post-fork Python child.
2952
+ """
2953
+
2954
+ def __init__(self, *, process: Any, session_id: str) -> None:
2955
+ self._process = process
2956
+ self._session_id = session_id
2957
+ self._stopped = False
2958
+
2959
+ @classmethod
2960
+ def start(
2961
+ cls,
2962
+ *,
2963
+ config: SessionWorkerConfig,
2964
+ invocation: StartInvocation,
2965
+ nonce: bytearray,
2966
+ ) -> _LegacySessionProcess:
2967
+ """Pass one validated nonce over a pipe to a fresh interpreter, then zero local copies."""
2968
+
2969
+ try:
2970
+ read_descriptor, write_descriptor = os.pipe()
2971
+ except OSError as error:
2972
+ nonce[:] = b"\x00" * len(nonce)
2973
+ raise SessionWorkerError(
2974
+ "SESSION_LEGACY_PROCESS_START_FAILED",
2975
+ "could not allocate the isolated legacy session pipe",
2976
+ ) from error
2977
+ process: Any | None = None
2978
+ payload = bytearray()
2979
+ try:
2980
+ # No ``preexec_fn``: on POSIX that would run Python after a multithreaded fork. The
2981
+ # executable receives only a descriptor number; the nonce is written after exec over
2982
+ # an inherited private pipe and never appears in argv or the environment.
2983
+ process = subprocess.Popen(
2984
+ (
2985
+ sys.executable,
2986
+ "-m",
2987
+ "mostlyright.data_harness.hosted_session_worker",
2988
+ LEGACY_BOOTSTRAP_FD_ARGUMENT,
2989
+ str(read_descriptor),
2990
+ ),
2991
+ close_fds=True,
2992
+ pass_fds=(read_descriptor,),
2993
+ stdin=subprocess.DEVNULL,
2994
+ stdout=subprocess.PIPE,
2995
+ start_new_session=True,
2996
+ )
2997
+ os.close(read_descriptor)
2998
+ read_descriptor = -1
2999
+ payload = bytearray(
3000
+ canonical_json_bytes(
3001
+ {
3002
+ "schema_version": START_SCHEMA_VERSION,
3003
+ "session_id": invocation.session_id,
3004
+ "workspace_id": invocation.workspace_id,
3005
+ "lease_nonce": nonce.decode("ascii"),
3006
+ "lease_url": invocation.lease_url,
3007
+ }
3008
+ )
3009
+ )
3010
+ if len(payload) > MAX_LEGACY_BOOTSTRAP_BYTES:
3011
+ raise SessionWorkerError(
3012
+ "SESSION_LEGACY_BOOTSTRAP_INVALID",
3013
+ "legacy bootstrap exceeds its byte limit",
3014
+ )
3015
+ cls._write_all(write_descriptor, payload)
3016
+ cls._await_ready(process)
3017
+ except (OSError, UnicodeError, SessionWorkerError, subprocess.SubprocessError) as error:
3018
+ if process is not None:
3019
+ cls._terminate_process_group(process)
3020
+ raise SessionWorkerError(
3021
+ "SESSION_LEGACY_PROCESS_START_FAILED",
3022
+ "could not start the isolated legacy session process",
3023
+ ) from error
3024
+ finally:
3025
+ if read_descriptor >= 0:
3026
+ with contextlib.suppress(OSError):
3027
+ os.close(read_descriptor)
3028
+ with contextlib.suppress(OSError):
3029
+ os.close(write_descriptor)
3030
+ payload[:] = b"\x00" * len(payload)
3031
+ nonce[:] = b"\x00" * len(nonce)
3032
+ assert process is not None
3033
+ return cls(process=process, session_id=invocation.session_id)
3034
+
3035
+ @staticmethod
3036
+ def _write_all(descriptor: int, payload: bytearray) -> None:
3037
+ view = memoryview(payload)
3038
+ try:
3039
+ while view:
3040
+ written = os.write(descriptor, view)
3041
+ if written <= 0:
3042
+ raise OSError("short legacy bootstrap pipe write")
3043
+ view = view[written:]
3044
+ finally:
3045
+ view.release()
3046
+
3047
+ @staticmethod
3048
+ def _await_ready(process: Any) -> None:
3049
+ """Require the fresh bootstrap monitor to arm its cleanup handler before HTTP accepts."""
3050
+
3051
+ stdout = process.stdout
3052
+ if stdout is None:
3053
+ raise OSError("legacy bootstrap process has no control output")
3054
+ readable, _unused, _errors = select.select(
3055
+ [stdout.fileno()], [], [], LEGACY_BOOTSTRAP_READY_TIMEOUT_SECONDS
3056
+ )
3057
+ if not readable:
3058
+ raise OSError("legacy bootstrap monitor did not become ready in time")
3059
+ raw = stdout.readline(MAX_SESSION_SUPERVISOR_MESSAGE_BYTES + 1)
3060
+ if len(raw) > MAX_SESSION_SUPERVISOR_MESSAGE_BYTES or not raw.endswith(b"\n"):
3061
+ raise OSError("legacy bootstrap monitor readiness is invalid")
3062
+ try:
3063
+ message = _strict_json(raw[:-1], "legacy bootstrap readiness")
3064
+ except HostedWorkerError as error:
3065
+ raise OSError("legacy bootstrap monitor readiness is invalid") from error
3066
+ if message != {"kind": "ready"}:
3067
+ raise OSError("legacy bootstrap monitor did not acknowledge readiness")
3068
+
3069
+ @staticmethod
3070
+ def _terminate_process_group(process: Any) -> None:
3071
+ if process.poll() is not None:
3072
+ return
3073
+ # The fresh bootstrap installs its monitor handler *before* it forks the separate
3074
+ # capability group. SIGTERM therefore lets it kill/reap that group and its adopted
3075
+ # sandboxes during a readiness failure; SIGKILL is only the bounded fallback for a wedged
3076
+ # outer monitor.
3077
+ with contextlib.suppress(PermissionError, ProcessLookupError):
3078
+ os.killpg(process.pid, signal.SIGTERM)
3079
+ try:
3080
+ process.wait(timeout=1.0)
3081
+ return
3082
+ except subprocess.TimeoutExpired:
3083
+ pass
3084
+ with contextlib.suppress(PermissionError, ProcessLookupError):
3085
+ os.killpg(process.pid, signal.SIGKILL)
3086
+ with contextlib.suppress(Exception):
3087
+ process.wait(timeout=5)
3088
+
3089
+ @property
3090
+ def is_running(self) -> bool:
3091
+ return self._process.poll() is None
3092
+
3093
+ def stop(self) -> None:
3094
+ self._stopped = True
3095
+ if self._process.poll() is not None:
3096
+ return
3097
+ # The fresh bootstrap interpreter catches this signal in its capability-free monitor, then
3098
+ # kills the separate capability group and reaps its adopted sandbox groups before exiting.
3099
+ # The group here intentionally contains only that bootstrap interpreter, not its child.
3100
+ with contextlib.suppress(ProcessLookupError):
3101
+ os.killpg(self._process.pid, signal.SIGTERM)
3102
+
3103
+ def run(self) -> dict[str, Any]:
3104
+ raw, _unused_stderr = self._process.communicate()
3105
+ if (
3106
+ self._stopped
3107
+ and not raw
3108
+ and self._process.returncode is not None
3109
+ and self._process.returncode < 0
3110
+ ):
3111
+ return {
3112
+ "status": "stopped",
3113
+ "role": SESSION_WORKER_ROLE,
3114
+ "session_id": self._session_id,
3115
+ "attempt_id": None,
3116
+ "probes_answered": 0,
3117
+ "probes_failed": 0,
3118
+ "terminal_state": "service_stopped",
3119
+ "failure_code": None,
3120
+ "warm_pool_ready": 0,
3121
+ }
3122
+ if (
3123
+ self._process.returncode == 0
3124
+ and len(raw) <= MAX_SESSION_SUPERVISOR_MESSAGE_BYTES
3125
+ and raw.endswith(b"\n")
3126
+ and raw.count(b"\n") == 1
3127
+ ):
3128
+ try:
3129
+ summary = _strict_json(raw[:-1], "isolated legacy session summary")
3130
+ except HostedWorkerError:
3131
+ summary = None
3132
+ if isinstance(summary, dict) and summary.get("session_id") == self._session_id:
3133
+ return summary
3134
+ return {
3135
+ "status": "failed",
3136
+ "role": SESSION_WORKER_ROLE,
3137
+ "session_id": self._session_id,
3138
+ "attempt_id": None,
3139
+ "probes_answered": 0,
3140
+ "probes_failed": 0,
3141
+ "terminal_state": None,
3142
+ "failure_code": "SESSION_LEGACY_PROCESS_FAILED",
3143
+ "warm_pool_ready": 0,
3144
+ }
3145
+
3146
+
3147
+ class SessionWorkerService:
3148
+ """One process, at most one live session, and an exit that lets Cloud Run scale to zero.
3149
+
3150
+ ``max_instance_request_concurrency = 1`` means Cloud Run will not send this instance a second
3151
+ request while the start invocation is being handled, but the invocation returns in
3152
+ milliseconds, so a second start *can* arrive while the first session is still running. It is
3153
+ refused with 409 rather than served: one session per instance is the tofu's own comment ("A
3154
+ probe must not queue behind another session's CPU") and serving two here would silently break
3155
+ the latency claim the whole feature is built on. Cloud Run answers a 409 by placing the session
3156
+ on another instance, which is what should happen.
3157
+
3158
+ Sessions are served serially, though, and that matters: a client who closes one session and
3159
+ opens another within the grace window gets this warm instance instead of a second cold start.
3160
+ """
3161
+
3162
+ def __init__(
3163
+ self,
3164
+ config: SessionWorkerConfig,
3165
+ *,
3166
+ supervisor_factory: Callable[..., SessionSupervisor] | None = None,
3167
+ ) -> None:
3168
+ self._config = config
3169
+ self._supervisor_factory = supervisor_factory
3170
+ self._lock = threading.Lock()
3171
+ self._active: Any | None = None
3172
+ self._idle_since: float | None = time.monotonic()
3173
+ self._shutdown = threading.Event()
3174
+ self.completed: list[dict[str, Any]] = []
3175
+
3176
+ # -- the two verbs ----------------------------------------------------------------------
3177
+
3178
+ def start(self, raw: bytes) -> tuple[int, dict[str, Any]]:
3179
+ """Accept one start signal and return in milliseconds. See the module docstring."""
3180
+
3181
+ try:
3182
+ invocation, nonce = parse_start_invocation(
3183
+ raw, studio_origin=self._config.studio_origin
3184
+ )
3185
+ except SessionWorkerError as error:
3186
+ # The reason is returned because Studio's starter treats any non-2xx as a failed
3187
+ # session and an operator reading the log needs to know which check refused. The nonce
3188
+ # is not in the invocation object, so nothing here can echo it.
3189
+ return int(HTTPStatus.BAD_REQUEST), {"code": error.code, "detail": error.detail}
3190
+ with self._lock:
3191
+ if self._active is not None and self._active.is_running:
3192
+ nonce[:] = b"\x00" * len(nonce)
3193
+ return int(HTTPStatus.CONFLICT), {
3194
+ "code": "SESSION_WORKER_BUSY",
3195
+ "detail": "this instance already holds a live session",
3196
+ }
3197
+ if self._supervisor_factory is None:
3198
+ try:
3199
+ supervisor: Any = _LegacySessionProcess.start(
3200
+ config=self._config, invocation=invocation, nonce=nonce
3201
+ )
3202
+ except SessionWorkerError as error:
3203
+ return int(HTTPStatus.SERVICE_UNAVAILABLE), {
3204
+ "code": error.code,
3205
+ "detail": error.detail,
3206
+ }
3207
+ else:
3208
+ supervisor = self._supervisor_factory(
3209
+ config=self._config, invocation=invocation, nonce=nonce
3210
+ )
3211
+ self._active = supervisor
3212
+ self._idle_since = None
3213
+ threading.Thread(
3214
+ target=self._run_session,
3215
+ args=(supervisor,),
3216
+ name=f"mr-session-{invocation.session_id[:8]}",
3217
+ daemon=False,
3218
+ ).start()
3219
+ return int(HTTPStatus.ACCEPTED), {
3220
+ "accepted": True,
3221
+ "session_id": invocation.session_id,
3222
+ }
3223
+
3224
+ def live(self) -> tuple[int, dict[str, Any]]:
3225
+ """Liveness. Deliberately says nothing about whether a session is warm."""
3226
+
3227
+ return int(HTTPStatus.OK), {"status": "live", "role": SESSION_WORKER_ROLE}
3228
+
3229
+ # -- lifecycle --------------------------------------------------------------------------
3230
+
3231
+ def _run_session(self, supervisor: SessionSupervisor) -> None:
3232
+ try:
3233
+ summary = supervisor.run()
3234
+ finally:
3235
+ with self._lock:
3236
+ self._active = None
3237
+ self._idle_since = time.monotonic()
3238
+ self.completed.append(summary)
3239
+ sys.stdout.write(json.dumps(summary, sort_keys=True, separators=(",", ":")) + "\n")
3240
+ sys.stdout.flush()
3241
+
3242
+ def idle_seconds(self, *, now: float | None = None) -> float | None:
3243
+ with self._lock:
3244
+ if self._idle_since is None:
3245
+ return None
3246
+ return (time.monotonic() if now is None else now) - self._idle_since
3247
+
3248
+ def should_exit(self, *, now: float | None = None) -> bool:
3249
+ """True once the instance has been idle long enough to be worth giving back.
3250
+
3251
+ This is the whole of scale-to-zero on the worker side. Cloud Run would eventually reap an
3252
+ idle instance on its own, but 'eventually' is billed at ``cpu_idle = false`` rates, and the
3253
+ difference between roughly $21 and roughly $55 per user-month is the difference between the
3254
+ product's price working and not.
3255
+ """
3256
+
3257
+ idle = self.idle_seconds(now=now)
3258
+ return idle is not None and idle >= self._config.shutdown_grace_seconds
3259
+
3260
+ def stop(self) -> None:
3261
+ with self._lock:
3262
+ active = self._active
3263
+ if active is not None:
3264
+ active.stop()
3265
+ self._shutdown.set()
3266
+
3267
+
3268
+ class _Handler(BaseHTTPRequestHandler):
3269
+ """The smallest surface that satisfies the contract. Everything else is 404."""
3270
+
3271
+ protocol_version = "HTTP/1.1"
3272
+ service: SessionWorkerService
3273
+
3274
+ def version_string(self) -> str:
3275
+ return "mostlyright-session-worker"
3276
+
3277
+ def log_message(self, format: str, *args: Any) -> None:
3278
+ # The default writes the full request line to stderr. A request line is never secret here
3279
+ # (the nonce is in the body), but routing it through the logger keeps one place responsible
3280
+ # for what this service says out loud.
3281
+ LOGGER.info("%s %s", self.command, self.path)
3282
+
3283
+ def do_GET(self) -> None:
3284
+ if self.path in HEALTH_PATHS:
3285
+ self._respond(*self.service.live())
3286
+ return
3287
+ self._respond(int(HTTPStatus.NOT_FOUND), {"code": "NOT_FOUND"})
3288
+
3289
+ def do_POST(self) -> None:
3290
+ if self.path != START_PATH:
3291
+ self._respond(int(HTTPStatus.NOT_FOUND), {"code": "NOT_FOUND"})
3292
+ return
3293
+ try:
3294
+ length = int(self.headers.get("Content-Length", "0"))
3295
+ except ValueError:
3296
+ self._respond(int(HTTPStatus.BAD_REQUEST), {"code": "SESSION_START_INVALID"})
3297
+ return
3298
+ if length < 0 or length > MAX_START_REQUEST_BYTES:
3299
+ self._respond(
3300
+ int(HTTPStatus.REQUEST_ENTITY_TOO_LARGE), {"code": "SESSION_START_TOO_LARGE"}
3301
+ )
3302
+ return
3303
+ raw = self.rfile.read(length)
3304
+ self._respond(*self.service.start(raw))
3305
+
3306
+ def _respond(self, status: int, body: Mapping[str, Any]) -> None:
3307
+ encoded = json.dumps(body, sort_keys=True, separators=(",", ":")).encode("utf-8")
3308
+ self.send_response(status)
3309
+ self.send_header("Content-Type", "application/json")
3310
+ self.send_header("Content-Length", str(len(encoded)))
3311
+ self.send_header("Cache-Control", "no-store")
3312
+ self.end_headers()
3313
+ self.wfile.write(encoded)
3314
+
3315
+
3316
+ class _Server(ThreadingHTTPServer):
3317
+ daemon_threads = True
3318
+ allow_reuse_address = True
3319
+
3320
+
3321
+ def serve(
3322
+ config: SessionWorkerConfig,
3323
+ *,
3324
+ service: SessionWorkerService | None = None,
3325
+ poll_seconds: float = 0.5,
3326
+ ) -> int:
3327
+ """Run the service until the instance has been idle long enough to give back.
3328
+
3329
+ The exit is the point. A container that stays up holds a Cloud Run instance, and an instance
3330
+ that is held is billed; returning from here lets the platform drop it immediately rather than
3331
+ waiting on its own idle reaper.
3332
+ """
3333
+
3334
+ worker = service or SessionWorkerService(config)
3335
+ handler = type("_BoundHandler", (_Handler,), {"service": worker})
3336
+ shutdown_requested = threading.Event()
3337
+ previous_sigterm: Any | None = None
3338
+ if threading.current_thread() is threading.main_thread():
3339
+ previous_sigterm = signal.getsignal(signal.SIGTERM)
3340
+
3341
+ def request_shutdown(_signum: int, _frame: Any) -> None:
3342
+ # A signal handler must not take the Service lock. The main lifecycle loop calls
3343
+ # ``worker.stop`` and lets its bootstrap monitor kill/reap the isolated child.
3344
+ shutdown_requested.set()
3345
+
3346
+ signal.signal(signal.SIGTERM, request_shutdown)
3347
+ with _Server(("0.0.0.0", config.port), handler) as server:
3348
+ thread = threading.Thread(target=server.serve_forever, name="mr-session-http", daemon=True)
3349
+ thread.start()
3350
+ try:
3351
+ while not shutdown_requested.is_set() and not worker.should_exit():
3352
+ time.sleep(poll_seconds)
3353
+ finally:
3354
+ worker.stop()
3355
+ server.shutdown()
3356
+ if previous_sigterm is not None:
3357
+ signal.signal(signal.SIGTERM, previous_sigterm)
3358
+ return 0
3359
+
3360
+
3361
+ def run_session_job(
3362
+ *,
3363
+ config: SessionWorkerConfig,
3364
+ dispatch: SessionDispatch,
3365
+ transport: BootstrapTransport | None = None,
3366
+ sandbox_factory: Callable[[SessionWorkerConfig], Any] | None = None,
3367
+ reader_factory: Callable[[SessionLease], ProbeSourceReader] | None = None,
3368
+ sleep: Callable[[float], None] = time.sleep,
3369
+ wall_clock: Callable[[], datetime] = _wall_clock_now,
3370
+ monotonic_clock: Callable[[], float] = time.monotonic,
3371
+ process_isolation: bool = True,
3372
+ ) -> dict[str, Any]:
3373
+ """Exchange and synchronously supervise exactly one session in one Cloud Run Job task."""
3374
+
3375
+ # The command-line audience is the task's lease audience. Do not read a legacy audience
3376
+ # environment variable after the dispatch has been accepted: that would make a task's OIDC
3377
+ # proof differ from the deployment coordinate Studio minted for it.
3378
+ job_config = replace(config, lease_audience=dispatch.audience)
3379
+ if process_isolation:
3380
+ # The opaque dispatch is safe in the parent; the capability exchange is intentionally in
3381
+ # the child, so neither this monitor nor the Cloud Run task's launcher retain a bearer.
3382
+ supervisor = SessionSupervisor(
3383
+ config=job_config,
3384
+ dispatch=dispatch,
3385
+ transport=transport,
3386
+ sandbox_factory=sandbox_factory,
3387
+ reader_factory=reader_factory,
3388
+ sleep=sleep,
3389
+ wall_clock=wall_clock,
3390
+ monotonic_clock=monotonic_clock,
3391
+ process_isolation=True,
3392
+ )
3393
+ else:
3394
+ # Deterministic unit adapters are intentionally in-process; production callers use the
3395
+ # default above. This seam prevents forked copies of in-memory fakes from being mistaken
3396
+ # for an end-to-end Studio result.
3397
+ lease, credential = exchange_session_dispatch(
3398
+ config=job_config, dispatch=dispatch, transport=transport
3399
+ )
3400
+ supervisor = SessionSupervisor(
3401
+ config=job_config,
3402
+ bootstrap=(lease, credential),
3403
+ transport=transport,
3404
+ sandbox_factory=sandbox_factory,
3405
+ reader_factory=reader_factory,
3406
+ sleep=sleep,
3407
+ wall_clock=wall_clock,
3408
+ monotonic_clock=monotonic_clock,
3409
+ )
3410
+ return supervisor.run()
3411
+
3412
+
3413
+ def _read_legacy_bootstrap_descriptor(descriptor: int) -> bytearray:
3414
+ """Read exactly one bounded start request from the Service's private inherited pipe."""
3415
+
3416
+ if descriptor < 3:
3417
+ raise SessionWorkerError(
3418
+ "SESSION_LEGACY_BOOTSTRAP_INVALID", "legacy bootstrap descriptor is not inherited"
3419
+ )
3420
+ raw = bytearray()
3421
+ try:
3422
+ while True:
3423
+ chunk = os.read(descriptor, min(4096, MAX_LEGACY_BOOTSTRAP_BYTES + 1 - len(raw)))
3424
+ if not chunk:
3425
+ break
3426
+ raw.extend(chunk)
3427
+ if len(raw) > MAX_LEGACY_BOOTSTRAP_BYTES:
3428
+ raise SessionWorkerError(
3429
+ "SESSION_LEGACY_BOOTSTRAP_INVALID", "legacy bootstrap exceeds its byte limit"
3430
+ )
3431
+ finally:
3432
+ with contextlib.suppress(OSError):
3433
+ os.close(descriptor)
3434
+ if not raw:
3435
+ raise SessionWorkerError("SESSION_LEGACY_BOOTSTRAP_INVALID", "legacy bootstrap is empty")
3436
+ return raw
3437
+
3438
+
3439
+ def _legacy_bootstrap_child_main(arguments: Sequence[str]) -> int:
3440
+ """Start the rollback-session supervisor in a fresh main-thread interpreter.
3441
+
3442
+ This is intentionally an internal descriptor protocol rather than an alternate public Service
3443
+ command. It performs the old nonce lease exchange only inside a newly exec'd process, whose
3444
+ own :class:`SessionSupervisor` then forks the capability-bearing child from a safe main
3445
+ thread. Neither parent can be held past capability expiry by a reader or native parser.
3446
+ """
3447
+
3448
+ if (
3449
+ len(arguments) != 2
3450
+ or arguments[0] != LEGACY_BOOTSTRAP_FD_ARGUMENT
3451
+ or not arguments[1].isdecimal()
3452
+ ):
3453
+ sys.stderr.write("SESSION_LEGACY_BOOTSTRAP_INVALID: expected one inherited descriptor\n")
3454
+ return 2
3455
+ try:
3456
+ descriptor = int(arguments[1])
3457
+ raw = _read_legacy_bootstrap_descriptor(descriptor)
3458
+ config = config_from_environment()
3459
+ invocation, nonce = parse_start_invocation(raw, studio_origin=config.studio_origin)
3460
+ except (SessionWorkerError, HostedWorkerError, OSError) as error:
3461
+ code = getattr(error, "code", "SESSION_LEGACY_BOOTSTRAP_INVALID")
3462
+ sys.stderr.write(f"{code}: bootstrap refused\n")
3463
+ return 1
3464
+ finally:
3465
+ if "raw" in locals():
3466
+ raw[:] = b"\x00" * len(raw)
3467
+
3468
+ try:
3469
+ summary = SessionSupervisor(
3470
+ config=config,
3471
+ invocation=invocation,
3472
+ nonce=nonce,
3473
+ process_isolation=True,
3474
+ on_isolation_ready=lambda: SessionSupervisor._write_isolated_message(
3475
+ sys.stdout.fileno(), {"kind": "ready"}
3476
+ ),
3477
+ ).run()
3478
+ finally:
3479
+ nonce[:] = b"\x00" * len(nonce)
3480
+ try:
3481
+ SessionSupervisor._write_isolated_message(sys.stdout.fileno(), summary)
3482
+ except (OSError, SessionWorkerError):
3483
+ return 1
3484
+ return 0 if summary["status"] == "stopped" else 1
3485
+
3486
+
3487
+ def session_worker_job_main(argv: Sequence[str] | None = None) -> int:
3488
+ """The Cloud Run Job ``mr-data-session-worker-job`` entrypoint.
3489
+
3490
+ Its sole task-specific inputs are the opaque dispatch id, exact exchange URL and OIDC audience;
3491
+ the task never receives a nonce or a capability in argv or its environment.
3492
+ """
3493
+
3494
+ logging.basicConfig(
3495
+ level=logging.INFO, stream=sys.stderr, format="%(levelname)s %(name)s %(message)s"
3496
+ )
3497
+ arguments = tuple(sys.argv[1:] if argv is None else argv)
3498
+ try:
3499
+ dispatch = parse_session_dispatch(arguments)
3500
+ except SessionWorkerError as error:
3501
+ sys.stderr.write(f"{error.code}: {error.detail}\n")
3502
+ return 2
3503
+ try:
3504
+ config = config_from_environment()
3505
+ summary = run_session_job(config=config, dispatch=dispatch)
3506
+ except SessionWorkerError as error:
3507
+ if error.code in CONSUMED_SESSION_DISPATCH_CODES:
3508
+ LOGGER.info(
3509
+ "session dispatch %s was already consumed; exiting without work", error.code
3510
+ )
3511
+ return 0
3512
+ sys.stderr.write(f"{error.code}: {error.detail}\n")
3513
+ return 1
3514
+ except HostedWorkerError as error:
3515
+ sys.stderr.write(f"{error.code}: {error.detail}\n")
3516
+ return 1
3517
+ # A task that reaches Studio's terminal session state completed its work even when the state
3518
+ # is ``failed``: the dispatch was consumed and a retry must not create a second attempt.
3519
+ if summary["status"] == "stopped":
3520
+ return 0
3521
+ failure = summary.get("failure_code") or "SESSION_WORKER_FAILED"
3522
+ sys.stderr.write(f"{failure}: session job did not reach a terminal session state\n")
3523
+ return 1
3524
+
3525
+
3526
+ def session_worker_main(argv: Sequence[str] | None = None) -> int:
3527
+ """The legacy Cloud Run service entrypoint, retained unchanged for rollback."""
3528
+
3529
+ logging.basicConfig(
3530
+ level=logging.INFO, stream=sys.stderr, format="%(levelname)s %(name)s %(message)s"
3531
+ )
3532
+ # This executable is also the fresh interpreter launched by
3533
+ # ``_LegacySessionProcess``. In that path Python invokes this module with
3534
+ # ``--legacy-bootstrap-fd <fd>`` in ``sys.argv``; treating a default ``None``
3535
+ # as an empty sequence would incorrectly enter ``serve`` a second time and
3536
+ # contend with the parent service for ``PORT``. Keep an explicitly supplied
3537
+ # empty sequence meaningful for unit callers, but mirror the Job entrypoint
3538
+ # when invoked as a console script.
3539
+ arguments = tuple(sys.argv[1:] if argv is None else argv)
3540
+ if arguments and arguments[0] == LEGACY_BOOTSTRAP_FD_ARGUMENT:
3541
+ return _legacy_bootstrap_child_main(arguments)
3542
+ if arguments:
3543
+ sys.stderr.write("mr-data-session-worker takes no arguments\n")
3544
+ return 2
3545
+ try:
3546
+ config = config_from_environment()
3547
+ except (SessionWorkerError, HostedWorkerError) as error:
3548
+ sys.stderr.write(f"{error.code}: {error.detail}\n")
3549
+ return 1
3550
+ return serve(config)
3551
+
3552
+
3553
+ if __name__ == "__main__": # pragma: no cover - exercised as a console script
3554
+ raise SystemExit(session_worker_main())