gatewaysdk 0.3.2__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 (376) hide show
  1. gatewaysdk/__init__.py +540 -0
  2. gatewaysdk/adapter/__init__.py +15 -0
  3. gatewaysdk/adapter/base.py +94 -0
  4. gatewaysdk/adapter/messages.py +270 -0
  5. gatewaysdk/adapter/triplet.py +1026 -0
  6. gatewaysdk/algorithm/__init__.py +39 -0
  7. gatewaysdk/algorithm/apo/__init__.py +5 -0
  8. gatewaysdk/algorithm/apo/apo.py +898 -0
  9. gatewaysdk/algorithm/apo/prompts/apply_edit_variant01.poml +22 -0
  10. gatewaysdk/algorithm/apo/prompts/apply_edit_variant02.poml +18 -0
  11. gatewaysdk/algorithm/apo/prompts/text_gradient_variant01.poml +18 -0
  12. gatewaysdk/algorithm/apo/prompts/text_gradient_variant02.poml +16 -0
  13. gatewaysdk/algorithm/apo/prompts/text_gradient_variant03.poml +107 -0
  14. gatewaysdk/algorithm/base.py +262 -0
  15. gatewaysdk/algorithm/decorator.py +264 -0
  16. gatewaysdk/algorithm/evals/__init__.py +7 -0
  17. gatewaysdk/algorithm/evals/evals.py +217 -0
  18. gatewaysdk/algorithm/fast.py +250 -0
  19. gatewaysdk/algorithm/gepa/__init__.py +61 -0
  20. gatewaysdk/algorithm/gepa/adapter.py +495 -0
  21. gatewaysdk/algorithm/gepa/gepa.py +570 -0
  22. gatewaysdk/algorithm/gepa/lib/__init__.py +18 -0
  23. gatewaysdk/algorithm/gepa/lib/adapters/README.md +12 -0
  24. gatewaysdk/algorithm/gepa/lib/adapters/__init__.py +0 -0
  25. gatewaysdk/algorithm/gepa/lib/adapters/anymaths_adapter/README.md +341 -0
  26. gatewaysdk/algorithm/gepa/lib/adapters/anymaths_adapter/__init__.py +1 -0
  27. gatewaysdk/algorithm/gepa/lib/adapters/anymaths_adapter/anymaths_adapter.py +174 -0
  28. gatewaysdk/algorithm/gepa/lib/adapters/anymaths_adapter/requirements.txt +1 -0
  29. gatewaysdk/algorithm/gepa/lib/adapters/default_adapter/README.md +0 -0
  30. gatewaysdk/algorithm/gepa/lib/adapters/default_adapter/__init__.py +0 -0
  31. gatewaysdk/algorithm/gepa/lib/adapters/default_adapter/default_adapter.py +209 -0
  32. gatewaysdk/algorithm/gepa/lib/adapters/dspy_adapter/README.md +7 -0
  33. gatewaysdk/algorithm/gepa/lib/adapters/dspy_adapter/__init__.py +0 -0
  34. gatewaysdk/algorithm/gepa/lib/adapters/dspy_adapter/dspy_adapter.py +307 -0
  35. gatewaysdk/algorithm/gepa/lib/adapters/dspy_full_program_adapter/README.md +99 -0
  36. gatewaysdk/algorithm/gepa/lib/adapters/dspy_full_program_adapter/dspy_program_proposal_signature.py +137 -0
  37. gatewaysdk/algorithm/gepa/lib/adapters/dspy_full_program_adapter/full_program_adapter.py +268 -0
  38. gatewaysdk/algorithm/gepa/lib/adapters/generic_rag_adapter/GEPA_RAG.md +621 -0
  39. gatewaysdk/algorithm/gepa/lib/adapters/generic_rag_adapter/__init__.py +56 -0
  40. gatewaysdk/algorithm/gepa/lib/adapters/generic_rag_adapter/evaluation_metrics.py +226 -0
  41. gatewaysdk/algorithm/gepa/lib/adapters/generic_rag_adapter/generic_rag_adapter.py +496 -0
  42. gatewaysdk/algorithm/gepa/lib/adapters/generic_rag_adapter/rag_pipeline.py +238 -0
  43. gatewaysdk/algorithm/gepa/lib/adapters/generic_rag_adapter/vector_store_interface.py +212 -0
  44. gatewaysdk/algorithm/gepa/lib/adapters/generic_rag_adapter/vector_stores/__init__.py +2 -0
  45. gatewaysdk/algorithm/gepa/lib/adapters/generic_rag_adapter/vector_stores/chroma_store.py +196 -0
  46. gatewaysdk/algorithm/gepa/lib/adapters/generic_rag_adapter/vector_stores/lancedb_store.py +422 -0
  47. gatewaysdk/algorithm/gepa/lib/adapters/generic_rag_adapter/vector_stores/milvus_store.py +409 -0
  48. gatewaysdk/algorithm/gepa/lib/adapters/generic_rag_adapter/vector_stores/qdrant_store.py +368 -0
  49. gatewaysdk/algorithm/gepa/lib/adapters/generic_rag_adapter/vector_stores/weaviate_store.py +418 -0
  50. gatewaysdk/algorithm/gepa/lib/adapters/mcp_adapter/README.md +552 -0
  51. gatewaysdk/algorithm/gepa/lib/adapters/mcp_adapter/__init__.py +37 -0
  52. gatewaysdk/algorithm/gepa/lib/adapters/mcp_adapter/mcp_adapter.py +699 -0
  53. gatewaysdk/algorithm/gepa/lib/adapters/mcp_adapter/mcp_client.py +364 -0
  54. gatewaysdk/algorithm/gepa/lib/adapters/terminal_bench_adapter/README.md +9 -0
  55. gatewaysdk/algorithm/gepa/lib/adapters/terminal_bench_adapter/__init__.py +0 -0
  56. gatewaysdk/algorithm/gepa/lib/adapters/terminal_bench_adapter/terminal_bench_adapter.py +217 -0
  57. gatewaysdk/algorithm/gepa/lib/api.py +382 -0
  58. gatewaysdk/algorithm/gepa/lib/core/__init__.py +0 -0
  59. gatewaysdk/algorithm/gepa/lib/core/adapter.py +180 -0
  60. gatewaysdk/algorithm/gepa/lib/core/data_loader.py +74 -0
  61. gatewaysdk/algorithm/gepa/lib/core/engine.py +379 -0
  62. gatewaysdk/algorithm/gepa/lib/core/result.py +233 -0
  63. gatewaysdk/algorithm/gepa/lib/core/state.py +636 -0
  64. gatewaysdk/algorithm/gepa/lib/examples/__init__.py +0 -0
  65. gatewaysdk/algorithm/gepa/lib/examples/aime.py +24 -0
  66. gatewaysdk/algorithm/gepa/lib/examples/anymaths-bench/eval_default.py +111 -0
  67. gatewaysdk/algorithm/gepa/lib/examples/anymaths-bench/prompt-templates/instruction_prompt.txt +9 -0
  68. gatewaysdk/algorithm/gepa/lib/examples/anymaths-bench/prompt-templates/optimal_prompt.txt +24 -0
  69. gatewaysdk/algorithm/gepa/lib/examples/anymaths-bench/train_anymaths.py +177 -0
  70. gatewaysdk/algorithm/gepa/lib/examples/dspy_full_program_evolution/arc_agi.ipynb +25705 -0
  71. gatewaysdk/algorithm/gepa/lib/examples/dspy_full_program_evolution/example.ipynb +348 -0
  72. gatewaysdk/algorithm/gepa/lib/examples/mcp_adapter/__init__.py +4 -0
  73. gatewaysdk/algorithm/gepa/lib/examples/mcp_adapter/mcp_optimization_example.py +456 -0
  74. gatewaysdk/algorithm/gepa/lib/examples/rag_adapter/RAG_GUIDE.md +613 -0
  75. gatewaysdk/algorithm/gepa/lib/examples/rag_adapter/__init__.py +9 -0
  76. gatewaysdk/algorithm/gepa/lib/examples/rag_adapter/rag_optimization.py +820 -0
  77. gatewaysdk/algorithm/gepa/lib/examples/rag_adapter/requirements-rag.txt +29 -0
  78. gatewaysdk/algorithm/gepa/lib/examples/terminal-bench/prompt-templates/instruction_prompt.txt +16 -0
  79. gatewaysdk/algorithm/gepa/lib/examples/terminal-bench/prompt-templates/terminus.txt +9 -0
  80. gatewaysdk/algorithm/gepa/lib/examples/terminal-bench/train_terminus.py +161 -0
  81. gatewaysdk/algorithm/gepa/lib/gepa_utils.py +117 -0
  82. gatewaysdk/algorithm/gepa/lib/logging/__init__.py +0 -0
  83. gatewaysdk/algorithm/gepa/lib/logging/experiment_tracker.py +187 -0
  84. gatewaysdk/algorithm/gepa/lib/logging/logger.py +75 -0
  85. gatewaysdk/algorithm/gepa/lib/logging/utils.py +103 -0
  86. gatewaysdk/algorithm/gepa/lib/proposer/__init__.py +0 -0
  87. gatewaysdk/algorithm/gepa/lib/proposer/base.py +31 -0
  88. gatewaysdk/algorithm/gepa/lib/proposer/merge.py +357 -0
  89. gatewaysdk/algorithm/gepa/lib/proposer/reflective_mutation/__init__.py +0 -0
  90. gatewaysdk/algorithm/gepa/lib/proposer/reflective_mutation/base.py +49 -0
  91. gatewaysdk/algorithm/gepa/lib/proposer/reflective_mutation/reflective_mutation.py +176 -0
  92. gatewaysdk/algorithm/gepa/lib/py.typed +0 -0
  93. gatewaysdk/algorithm/gepa/lib/strategies/__init__.py +0 -0
  94. gatewaysdk/algorithm/gepa/lib/strategies/batch_sampler.py +77 -0
  95. gatewaysdk/algorithm/gepa/lib/strategies/candidate_selector.py +50 -0
  96. gatewaysdk/algorithm/gepa/lib/strategies/component_selector.py +36 -0
  97. gatewaysdk/algorithm/gepa/lib/strategies/eval_policy.py +64 -0
  98. gatewaysdk/algorithm/gepa/lib/strategies/instruction_proposal.py +126 -0
  99. gatewaysdk/algorithm/gepa/lib/utils/__init__.py +10 -0
  100. gatewaysdk/algorithm/gepa/lib/utils/stop_condition.py +196 -0
  101. gatewaysdk/algorithm/gepa/tracing.py +105 -0
  102. gatewaysdk/algorithm/utils.py +177 -0
  103. gatewaysdk/algorithm/verl/__init__.py +5 -0
  104. gatewaysdk/algorithm/verl/interface.py +202 -0
  105. gatewaysdk/automations.py +111 -0
  106. gatewaysdk/benchmark_hub/README.md +88 -0
  107. gatewaysdk/benchmark_hub/__init__.py +119 -0
  108. gatewaysdk/benchmark_hub/_archive.py +189 -0
  109. gatewaysdk/benchmark_hub/_tar.py +6 -0
  110. gatewaysdk/benchmark_hub/client.py +992 -0
  111. gatewaysdk/benchmark_hub/dispatch_shard.py +34 -0
  112. gatewaysdk/benchmark_hub/eval_config.py +20 -0
  113. gatewaysdk/benchmark_hub/evals.py +351 -0
  114. gatewaysdk/benchmark_hub/harbor_adapter.py +673 -0
  115. gatewaysdk/benchmark_hub/path_utils.py +81 -0
  116. gatewaysdk/benchmark_hub/save_utils.py +101 -0
  117. gatewaysdk/benchmark_hub/verifiers_adapter.py +327 -0
  118. gatewaysdk/benchmark_hub/versioning.py +72 -0
  119. gatewaysdk/build.py +515 -0
  120. gatewaysdk/cli/__init__.py +58 -0
  121. gatewaysdk/cli/agent_runner.py +132 -0
  122. gatewaysdk/cli/http_client.py +115 -0
  123. gatewaysdk/cli/platform.py +4086 -0
  124. gatewaysdk/cli/prometheus.py +115 -0
  125. gatewaysdk/cli/release_gate.py +215 -0
  126. gatewaysdk/cli/store.py +131 -0
  127. gatewaysdk/cli/vllm.py +29 -0
  128. gatewaysdk/client.py +406 -0
  129. gatewaysdk/config.py +348 -0
  130. gatewaysdk/connectors/__init__.py +25 -0
  131. gatewaysdk/connectors/client.py +203 -0
  132. gatewaysdk/connectors/skill.py +25 -0
  133. gatewaysdk/connectors/template.py +106 -0
  134. gatewaysdk/context.py +606 -0
  135. gatewaysdk/emitter/__init__.py +43 -0
  136. gatewaysdk/emitter/annotation.py +370 -0
  137. gatewaysdk/emitter/exception.py +54 -0
  138. gatewaysdk/emitter/message.py +61 -0
  139. gatewaysdk/emitter/object.py +117 -0
  140. gatewaysdk/emitter/reward.py +320 -0
  141. gatewaysdk/env_var.py +156 -0
  142. gatewaysdk/environment/__init__.py +108 -0
  143. gatewaysdk/environment/_bundle.py +281 -0
  144. gatewaysdk/environment/_harbor.py +276 -0
  145. gatewaysdk/environment/_materialize.py +97 -0
  146. gatewaysdk/environment/_tar.py +33 -0
  147. gatewaysdk/environment/_world.py +346 -0
  148. gatewaysdk/environment/core.py +527 -0
  149. gatewaysdk/environment/errors.py +58 -0
  150. gatewaysdk/environment/runtime.py +150 -0
  151. gatewaysdk/environment/schema/__init__.py +35 -0
  152. gatewaysdk/environment/schema/__main__.py +225 -0
  153. gatewaysdk/environment/schema/_slack_fidelity.py +116 -0
  154. gatewaysdk/environment/schema/_slack_scenario.py +183 -0
  155. gatewaysdk/environment/schema/api.py +1628 -0
  156. gatewaysdk/environment/schema/batch.py +755 -0
  157. gatewaysdk/environment/schema/compiler.py +615 -0
  158. gatewaysdk/environment/schema/conform.py +428 -0
  159. gatewaysdk/environment/schema/connector.py +777 -0
  160. gatewaysdk/environment/schema/connectors/apple-business-manager/handlers.py +17 -0
  161. gatewaysdk/environment/schema/connectors/apple-business-manager/parity.json +28 -0
  162. gatewaysdk/environment/schema/connectors/apple-business-manager/provenance.json +233 -0
  163. gatewaysdk/environment/schema/connectors/apple-business-manager/scope.toml +112 -0
  164. gatewaysdk/environment/schema/connectors/apple-business-manager/world.json +1380 -0
  165. gatewaysdk/environment/schema/connectors/base.json +28 -0
  166. gatewaysdk/environment/schema/connectors/custom/handlers.py +18 -0
  167. gatewaysdk/environment/schema/connectors/custom/world.json +123 -0
  168. gatewaysdk/environment/schema/connectors/github/handlers.py +81 -0
  169. gatewaysdk/environment/schema/connectors/github/parity.json +37 -0
  170. gatewaysdk/environment/schema/connectors/github/provenance.json +134 -0
  171. gatewaysdk/environment/schema/connectors/github/scope.toml +158 -0
  172. gatewaysdk/environment/schema/connectors/github/world.json +5869 -0
  173. gatewaysdk/environment/schema/connectors/google-calendar/handlers.py +23 -0
  174. gatewaysdk/environment/schema/connectors/google-calendar/world.json +209 -0
  175. gatewaysdk/environment/schema/connectors/google-drive/handlers.py +18 -0
  176. gatewaysdk/environment/schema/connectors/google-drive/world.json +286 -0
  177. gatewaysdk/environment/schema/connectors/jamf/handlers.py +34 -0
  178. gatewaysdk/environment/schema/connectors/jamf/parity.json +49 -0
  179. gatewaysdk/environment/schema/connectors/jamf/provenance.json +181 -0
  180. gatewaysdk/environment/schema/connectors/jamf/scope.toml +128 -0
  181. gatewaysdk/environment/schema/connectors/jamf/world.json +2218 -0
  182. gatewaysdk/environment/schema/connectors/jira/handlers.py +28 -0
  183. gatewaysdk/environment/schema/connectors/jira/world.json +389 -0
  184. gatewaysdk/environment/schema/connectors/kandji/handlers.py +38 -0
  185. gatewaysdk/environment/schema/connectors/kandji/parity.json +29 -0
  186. gatewaysdk/environment/schema/connectors/kandji/provenance.json +152 -0
  187. gatewaysdk/environment/schema/connectors/kandji/scope.toml +120 -0
  188. gatewaysdk/environment/schema/connectors/kandji/world.json +765 -0
  189. gatewaysdk/environment/schema/connectors/linear/handlers.py +20 -0
  190. gatewaysdk/environment/schema/connectors/linear/world.json +363 -0
  191. gatewaysdk/environment/schema/connectors/microsoft-teams/fidelity.json +45 -0
  192. gatewaysdk/environment/schema/connectors/microsoft-teams/handlers.py +1 -0
  193. gatewaysdk/environment/schema/connectors/microsoft-teams/provenance.json +851 -0
  194. gatewaysdk/environment/schema/connectors/microsoft-teams/scope.md +231 -0
  195. gatewaysdk/environment/schema/connectors/microsoft-teams/teams-types.json +1710 -0
  196. gatewaysdk/environment/schema/connectors/microsoft-teams/world.json +858 -0
  197. gatewaysdk/environment/schema/connectors/netsuite/handlers.py +32 -0
  198. gatewaysdk/environment/schema/connectors/netsuite/parity.json +50 -0
  199. gatewaysdk/environment/schema/connectors/netsuite/provenance.json +137 -0
  200. gatewaysdk/environment/schema/connectors/netsuite/scope.toml +78 -0
  201. gatewaysdk/environment/schema/connectors/netsuite/world.json +16243 -0
  202. gatewaysdk/environment/schema/connectors/salesforce/handlers.py +806 -0
  203. gatewaysdk/environment/schema/connectors/salesforce/parity.json +46 -0
  204. gatewaysdk/environment/schema/connectors/salesforce/provenance.json +143 -0
  205. gatewaysdk/environment/schema/connectors/salesforce/scope.toml +174 -0
  206. gatewaysdk/environment/schema/connectors/salesforce/world.json +4695 -0
  207. gatewaysdk/environment/schema/connectors/slack/MODEL.md +236 -0
  208. gatewaysdk/environment/schema/connectors/slack/capabilities.json +1039 -0
  209. gatewaysdk/environment/schema/connectors/slack/handlers.py +762 -0
  210. gatewaysdk/environment/schema/connectors/slack/parity.json +65 -0
  211. gatewaysdk/environment/schema/connectors/slack/provenance.json +294 -0
  212. gatewaysdk/environment/schema/connectors/slack/scope.toml +44 -0
  213. gatewaysdk/environment/schema/connectors/slack/world.json +4993 -0
  214. gatewaysdk/environment/schema/connectors/workday/handlers.py +75 -0
  215. gatewaysdk/environment/schema/connectors/workday/parity.json +36 -0
  216. gatewaysdk/environment/schema/connectors/workday/provenance.json +179 -0
  217. gatewaysdk/environment/schema/connectors/workday/scope.toml +180 -0
  218. gatewaysdk/environment/schema/connectors/workday/world.json +806 -0
  219. gatewaysdk/environment/schema/data.py +593 -0
  220. gatewaysdk/environment/schema/extract.py +263 -0
  221. gatewaysdk/environment/schema/host.py +2437 -0
  222. gatewaysdk/environment/schema/host_surfaces.py +275 -0
  223. gatewaysdk/environment/schema/platform.py +1272 -0
  224. gatewaysdk/environment/schema/scaffold.py +250 -0
  225. gatewaysdk/environment/schema/skill.py +118 -0
  226. gatewaysdk/environment/schema/skills/agents/connector-template-author.md +202 -0
  227. gatewaysdk/environment/schema/skills/connector-schema-authoring/CONNECTORS.md +331 -0
  228. gatewaysdk/environment/schema/skills/connector-schema-authoring/CONVENTIONS.md +100 -0
  229. gatewaysdk/environment/schema/skills/connector-schema-authoring/SKILL.md +313 -0
  230. gatewaysdk/environment/schema/skills/world-data-ingestion/ROWS.md +223 -0
  231. gatewaysdk/environment/schema/skills/world-data-ingestion/SKILL.md +293 -0
  232. gatewaysdk/environment/schema/skills/worlds-getting-started/SKILL.md +390 -0
  233. gatewaysdk/environment/schema/slack.py +757 -0
  234. gatewaysdk/environment/schema/snapshot.py +479 -0
  235. gatewaysdk/environment/schema/store.py +904 -0
  236. gatewaysdk/environment/schema/tasks.py +525 -0
  237. gatewaysdk/environment/schema/tools_world.py +322 -0
  238. gatewaysdk/environment/schema/validation.py +577 -0
  239. gatewaysdk/environment/schema/worker.py +78 -0
  240. gatewaysdk/environment/schema/workers.py +500 -0
  241. gatewaysdk/environment/schema/world_tests.py +336 -0
  242. gatewaysdk/environments/__init__.py +24 -0
  243. gatewaysdk/environments/client.py +486 -0
  244. gatewaysdk/environments/types.py +318 -0
  245. gatewaysdk/execution/__init__.py +15 -0
  246. gatewaysdk/execution/base.py +64 -0
  247. gatewaysdk/execution/client_server.py +443 -0
  248. gatewaysdk/execution/events.py +69 -0
  249. gatewaysdk/execution/inter_process.py +16 -0
  250. gatewaysdk/execution/shared_memory.py +282 -0
  251. gatewaysdk/experiments/__init__.py +90 -0
  252. gatewaysdk/experiments/assignment.py +177 -0
  253. gatewaysdk/experiments/client.py +877 -0
  254. gatewaysdk/experiments/exposure.py +222 -0
  255. gatewaysdk/experiments/types.py +81 -0
  256. gatewaysdk/importers/__init__.py +49 -0
  257. gatewaysdk/importers/_normalize.py +91 -0
  258. gatewaysdk/importers/builder.py +173 -0
  259. gatewaysdk/importers/client.py +95 -0
  260. gatewaysdk/importers/langsmith.py +196 -0
  261. gatewaysdk/importers/recipes.py +249 -0
  262. gatewaysdk/instrumentation/__init__.py +300 -0
  263. gatewaysdk/instrumentation/agentops.py +314 -0
  264. gatewaysdk/instrumentation/agentops_langchain.py +45 -0
  265. gatewaysdk/instrumentation/base.py +119 -0
  266. gatewaysdk/instrumentation/litellm.py +83 -0
  267. gatewaysdk/instrumentation/registry.py +273 -0
  268. gatewaysdk/instrumentation/vllm.py +81 -0
  269. gatewaysdk/instrumentation/weave.py +500 -0
  270. gatewaysdk/integrations/__init__.py +15 -0
  271. gatewaysdk/integrations/gateway/__init__.py +11 -0
  272. gatewaysdk/integrations/gateway/client.py +171 -0
  273. gatewaysdk/integrations/tool_access.py +549 -0
  274. gatewaysdk/litagent/__init__.py +11 -0
  275. gatewaysdk/litagent/decorator.py +536 -0
  276. gatewaysdk/litagent/litagent.py +252 -0
  277. gatewaysdk/llm_proxy.py +1742 -0
  278. gatewaysdk/logging.py +370 -0
  279. gatewaysdk/memory.py +278 -0
  280. gatewaysdk/personas.py +186 -0
  281. gatewaysdk/platform/__init__.py +17 -0
  282. gatewaysdk/platform/builder.py +221 -0
  283. gatewaysdk/platform/compatibility.py +67 -0
  284. gatewaysdk/platform/manifest.py +185 -0
  285. gatewaysdk/platform/orchestrator.py +901 -0
  286. gatewaysdk/platform/registry.py +122 -0
  287. gatewaysdk/platform/worker.py +864 -0
  288. gatewaysdk/replay/__init__.py +1032 -0
  289. gatewaysdk/replay/pytest.py +56 -0
  290. gatewaysdk/reward.py +7 -0
  291. gatewaysdk/run.py +1781 -0
  292. gatewaysdk/runner/__init__.py +11 -0
  293. gatewaysdk/runner/agent.py +878 -0
  294. gatewaysdk/runner/base.py +182 -0
  295. gatewaysdk/runner/legacy.py +309 -0
  296. gatewaysdk/security.py +700 -0
  297. gatewaysdk/semconv.py +170 -0
  298. gatewaysdk/server.py +399 -0
  299. gatewaysdk/sessions.py +282 -0
  300. gatewaysdk/store/__init__.py +45 -0
  301. gatewaysdk/store/base.py +908 -0
  302. gatewaysdk/store/client_server.py +2093 -0
  303. gatewaysdk/store/collection/__init__.py +30 -0
  304. gatewaysdk/store/collection/base.py +587 -0
  305. gatewaysdk/store/collection/memory.py +970 -0
  306. gatewaysdk/store/collection/mongo.py +1412 -0
  307. gatewaysdk/store/collection_based.py +1823 -0
  308. gatewaysdk/store/gateway.py +983 -0
  309. gatewaysdk/store/gateway_listener.py +465 -0
  310. gatewaysdk/store/listener.py +58 -0
  311. gatewaysdk/store/memory.py +396 -0
  312. gatewaysdk/store/mongo.py +165 -0
  313. gatewaysdk/store/redis_stream.py +517 -0
  314. gatewaysdk/store/sqlite.py +3 -0
  315. gatewaysdk/store/threading.py +370 -0
  316. gatewaysdk/store/utils.py +142 -0
  317. gatewaysdk/tracer/__init__.py +14 -0
  318. gatewaysdk/tracer/base.py +286 -0
  319. gatewaysdk/tracer/dummy.py +106 -0
  320. gatewaysdk/tracer/otel.py +559 -0
  321. gatewaysdk/tracing/__init__.py +110 -0
  322. gatewaysdk/tracing/api.py +808 -0
  323. gatewaysdk/tracing/attributes.py +9 -0
  324. gatewaysdk/tracing/context.py +272 -0
  325. gatewaysdk/tracing/exporters/__init__.py +10 -0
  326. gatewaysdk/tracing/exporters/gateway.py +228 -0
  327. gatewaysdk/tracing/identity.py +288 -0
  328. gatewaysdk/tracing/init.py +620 -0
  329. gatewaysdk/tracing/instrumentors/__init__.py +15 -0
  330. gatewaysdk/tracing/instrumentors/claude_agent_sdk.py +766 -0
  331. gatewaysdk/tracing/instrumentors/instrumentation_principles.md +294 -0
  332. gatewaysdk/tracing/instrumentors/registry.py +352 -0
  333. gatewaysdk/tracing/mapping.py +729 -0
  334. gatewaysdk/tracing/processors.py +393 -0
  335. gatewaysdk/tracing/push.py +617 -0
  336. gatewaysdk/tracing/push_models.py +247 -0
  337. gatewaysdk/tracing/semconv.py +294 -0
  338. gatewaysdk/tracing/span_builder.py +356 -0
  339. gatewaysdk/trainer/__init__.py +6 -0
  340. gatewaysdk/trainer/init_utils.py +263 -0
  341. gatewaysdk/trainer/legacy.py +359 -0
  342. gatewaysdk/trainer/registry.py +12 -0
  343. gatewaysdk/trainer/trainer.py +638 -0
  344. gatewaysdk/types/__init__.py +63 -0
  345. gatewaysdk/types/core.py +556 -0
  346. gatewaysdk/types/resources.py +204 -0
  347. gatewaysdk/types/tracer.py +515 -0
  348. gatewaysdk/types/tracing.py +162 -0
  349. gatewaysdk/users.py +251 -0
  350. gatewaysdk/utils/__init__.py +1 -0
  351. gatewaysdk/utils/id.py +18 -0
  352. gatewaysdk/utils/metrics.py +1025 -0
  353. gatewaysdk/utils/otel.py +550 -0
  354. gatewaysdk/utils/otlp.py +556 -0
  355. gatewaysdk/utils/redact.py +22 -0
  356. gatewaysdk/utils/server_launcher.py +1045 -0
  357. gatewaysdk/utils/system_snapshot.py +90 -0
  358. gatewaysdk/verl/__init__.py +8 -0
  359. gatewaysdk/verl/__main__.py +6 -0
  360. gatewaysdk/verl/async_server.py +46 -0
  361. gatewaysdk/verl/config.yaml +27 -0
  362. gatewaysdk/verl/daemon.py +1154 -0
  363. gatewaysdk/verl/dataset.py +44 -0
  364. gatewaysdk/verl/entrypoint.py +248 -0
  365. gatewaysdk/verl/trainer.py +549 -0
  366. gatewaysdk/world_browser.py +748 -0
  367. gatewaysdk/world_data.py +453 -0
  368. gatewaysdk/world_sessions.py +978 -0
  369. gatewaysdk/world_tasks.py +285 -0
  370. gatewaysdk/world_tools.py +115 -0
  371. gatewaysdk/worlds.py +639 -0
  372. gatewaysdk-0.3.2.dist-info/METADATA +236 -0
  373. gatewaysdk-0.3.2.dist-info/RECORD +376 -0
  374. gatewaysdk-0.3.2.dist-info/WHEEL +4 -0
  375. gatewaysdk-0.3.2.dist-info/entry_points.txt +5 -0
  376. gatewaysdk-0.3.2.dist-info/licenses/LICENSE +19 -0
@@ -0,0 +1,2437 @@
1
+ """The World Host: one long-lived process that serves many world sessions.
2
+
3
+ Where the pod engine boots a container per session, the host keeps each
4
+ world VERSION cached once and opens a SESSION as a copy-on-write clone of that
5
+ version's snapshot — a file operation, not a scheduling one. One process
6
+ answers for every session it owns, routed by the ``w-<sid>[-api]`` Host header
7
+ the platform already mints (``worldSurfaces.sessionOrigins``), so there is
8
+ nothing per session in Kubernetes or the load balancer.
9
+
10
+ Two planes on one HTTP server:
11
+
12
+ * the **control plane** (``/__host/*``, behind ``X-Host-Secret``) is the
13
+ platform's half: open, close, reset, seed, state, grade, call. It is
14
+ synchronous — a ``seed`` or ``grade`` completes in the response, not over a
15
+ poll — because the host holds the world in memory.
16
+ * the **data plane** (everything else) is the end user's half: the world's own
17
+ vendor routes, selected by session and checked against that session's token,
18
+ answered in-process by :class:`WorldApp`.
19
+
20
+ Every session carries a lease. Control calls and data-plane requests renew it;
21
+ a session the platform stops touching expires on the host and is unlinked, so
22
+ a missed reconcile tick on the platform is recovered, not leaked.
23
+
24
+ A session's state is one sqlite file, so it is also its own durable form: the
25
+ platform asks the host to **export** a session (a consistent copy, hashed) and
26
+ to **upload** that copy to a presigned URL the platform holds the key for, and
27
+ later **opens** the same session id on any host with ``restore`` naming the
28
+ export, which is how a session survives the host that held it.
29
+
30
+ The host serves schema worlds (a ``gateway-world/1`` tree :class:`WorldState`
31
+ can open) with the ``tools`` and ``api`` surfaces. A world with no code of its
32
+ own is answered in this process. A world that brings an ``api/handlers.py`` is
33
+ placed on a **version worker** (:mod:`workers`): a subprocess of this host, one
34
+ per version, reached over a unix socket, with rlimits and a scrubbed
35
+ environment, running the same :class:`WorldHost` with ``run_handlers`` on. The
36
+ supervisor keeps the registry, the leases and the token gate; the worker keeps
37
+ the world. The **ui** surface is the version's prebuilt static app, served from
38
+ the cache by the supervisor with its ``api_path`` rewritten onto the session's
39
+ api; the **browser** surface is a websocket tunnel to the pool Chromium the
40
+ platform leased for the session (:mod:`host_surfaces`). Custom-server worlds
41
+ stay on pods; this module refuses a version it cannot open rather than
42
+ guessing.
43
+ """
44
+
45
+ from __future__ import annotations
46
+
47
+ import base64
48
+ import fcntl
49
+ import hashlib
50
+ import hmac
51
+ import http.client
52
+ import json
53
+ import os
54
+ import re
55
+ import shutil
56
+ import signal
57
+ import subprocess
58
+ import sys
59
+ import tarfile
60
+ import threading
61
+ import time
62
+ import urllib.error
63
+ import urllib.parse
64
+ import urllib.request
65
+ import uuid
66
+ from collections import Counter
67
+ from collections.abc import Callable, Iterable, Iterator
68
+ from contextlib import contextmanager
69
+ from dataclasses import dataclass, field
70
+ from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
71
+ from pathlib import Path
72
+ from typing import Any, Protocol
73
+
74
+ from .._tar import safe_tar_extract
75
+ from . import snapshot as snapshots
76
+ from .api import (
77
+ SCENARIOS,
78
+ ApiSettings,
79
+ Handler,
80
+ Router,
81
+ WorldApp,
82
+ WorldState,
83
+ load_api_settings,
84
+ load_handlers,
85
+ send_response,
86
+ )
87
+ from .compiler import check_bundle
88
+ from .connector import Spec, load_spec
89
+ from .host_surfaces import (
90
+ HOP_HEADER,
91
+ UiSurface,
92
+ api_target,
93
+ cookie_token,
94
+ load_ui_surface,
95
+ matches,
96
+ mint_live_grant,
97
+ query_token,
98
+ session_cookie,
99
+ static_file,
100
+ tunnel,
101
+ verify_live_grant,
102
+ without_query_token,
103
+ )
104
+ from .tasks import (
105
+ TaskSpec,
106
+ evaluate_assertions,
107
+ parse_task_spec,
108
+ unknown_task_entities,
109
+ unknown_task_tools,
110
+ write_task_dir,
111
+ )
112
+ from .tools_world import ToolsVersion, ToolsWorld, is_tools_bundle, load_tools_version
113
+ from .validation import ContractError, fail
114
+ from .workers import (
115
+ DEFAULT_WORKER_BUDGET_MB,
116
+ DEFAULT_WORKER_MEMORY_MB,
117
+ Worker,
118
+ WorkerCapacity,
119
+ WorkerGone,
120
+ WorkerPool,
121
+ confined_command,
122
+ verifier_environment,
123
+ )
124
+
125
+ #: The world tree inside an api bundle, or the bundle root for a bare tree.
126
+ BUNDLED_WORLD = "server/world"
127
+ API_BUNDLE = "gateway-api-world.toml"
128
+ VERIFY_TIMEOUT_S = 30.0
129
+ #: How long a session lives with no control call or request renewing it.
130
+ DEFAULT_LEASE_S = 3600.0
131
+ #: How often the host sweeps expired leases.
132
+ LEASE_SWEEP_S = 30.0
133
+ #: How long an owner lookup answer is trusted before asking the platform again.
134
+ #: An owner that does not answer a proxied request is asked about again at once.
135
+ OWNER_CACHE_S = 60.0
136
+ #: How long "no owner" is trusted: long enough that a client hammering a closed
137
+ #: session costs the platform one call at a time, short enough that a session
138
+ #: recovered onto another host is reached at once.
139
+ OWNER_MISS_S = 2.0
140
+ #: How many sessions one host takes before the platform places elsewhere.
141
+ DEFAULT_MAX_SESSIONS = 1500
142
+ #: The error code (and a 404) a control call gets for a session this host does
143
+ #: not hold. The platform treats exactly this answer as "the host lost it".
144
+ UNKNOWN_SESSION = "unknown_session"
145
+ #: How often a host tells the platform it is alive and what it holds.
146
+ HEARTBEAT_S = 10.0
147
+ #: How long a draining host waits for its sessions to close before it exits.
148
+ DEFAULT_DRAIN_S = 45.0
149
+ #: Response headers a relayed answer does NOT keep: the framing is this
150
+ #: connection's, everything else the world set (its ``page_headers``, a
151
+ #: handler's own headers) passes through, from a worker and from an owner alike.
152
+ FRAMING_HEADERS = frozenset({"content-length", "connection", "transfer-encoding", "server", "date"})
153
+ #: Hop-by-hop headers (RFC 9110 §7.6.1) an owner's answer does not carry past
154
+ #: the proxying replica.
155
+ HOP_BY_HOP_HEADERS = frozenset(
156
+ {"keep-alive", "proxy-authenticate", "proxy-authorization", "te", "trailer", "upgrade", "proxy-connection"}
157
+ )
158
+ PROXIED_DROP_HEADERS = FRAMING_HEADERS | HOP_BY_HOP_HEADERS
159
+
160
+
161
+ class _NoRedirect(urllib.request.HTTPRedirectHandler):
162
+ """An owner's redirect is the client's to follow, not the proxy's: the ui
163
+ token exchange answers 303 with the session cookie, and following it here
164
+ would drop the cookie and land on the owner's 401."""
165
+
166
+ def redirect_request(self, req, fp, code, msg, headers, newurl):
167
+ return None
168
+
169
+
170
+ #: The opener the one-hop proxy uses: relays the owner's answer as is, a
171
+ #: redirect included, rather than following it.
172
+ _PROXY_OPENER = urllib.request.build_opener(_NoRedirect)
173
+ #: How long the grant the host mints towards the relay lives: the viewer's own
174
+ #: grant was already checked, this one only has to survive the handshake.
175
+ LIVE_RELAY_GRANT_MS = 60_000
176
+ #: What a session export is: the session's sqlite file, as the store wrote it.
177
+ EXPORT_CONTENT_TYPE = "application/vnd.sqlite3"
178
+ #: How long one export upload or restore download may take.
179
+ TRANSFER_TIMEOUT_S = 300.0
180
+ #: The most a version bundle may weigh, compressed and extracted, before the
181
+ #: host refuses it: the cache is one shared disk and a build's memory is the
182
+ #: supervisor's.
183
+ DEFAULT_MAX_BUNDLE_MB = 2048
184
+ MAX_BUNDLE_MEMBERS = 200_000
185
+ #: The largest request body the host reads, on either plane. A body is held
186
+ #: in memory whole (the seed rows, a tool call's args, a vendor route's
187
+ #: payload), and the memory is the supervisor's, shared by every session on
188
+ #: the replica; a client cannot be let size that allocation.
189
+ DEFAULT_MAX_BODY_MB = 32
190
+
191
+
192
+ # --------------------------------------------------------------------------- #
193
+ # Reflink: a session's state is a copy-on-write clone of the version snapshot.
194
+ # --------------------------------------------------------------------------- #
195
+
196
+ _FICLONE = 0x40049409 # ioctl: clone one file's extents into another (btrfs/XFS)
197
+
198
+
199
+ def clone_file(src: Path, dst: Path) -> None:
200
+ """Copy ``src`` to ``dst`` cheaply. A reflink (``FICLONE``) shares the
201
+ version snapshot's blocks copy-on-write, so opening a session costs the
202
+ same whatever the snapshot's size; a filesystem without reflink falls back
203
+ to a plain copy, which is what the pod engine does today."""
204
+ dst.parent.mkdir(parents=True, exist_ok=True)
205
+ try:
206
+ with open(src, "rb") as s, open(dst, "wb") as d:
207
+ fcntl.ioctl(d.fileno(), _FICLONE, s.fileno())
208
+ return
209
+ except OSError:
210
+ pass # not supported here (tmpfs, ext4, cross-device) — copy instead
211
+ shutil.copyfile(src, dst)
212
+
213
+
214
+ OBJECT_STORE_ANY = "any"
215
+ _ORIGIN_HOST = re.compile(r"^[A-Za-z0-9.-]+(:[0-9]+)?$")
216
+
217
+
218
+ def parse_origins(value: str | None) -> frozenset[str] | None:
219
+ """``WORLDS_HOST_OBJECT_STORE``: the origins (``scheme://host[:port]``,
220
+ comma-separated) the platform's presigned URLs may point at, or the word
221
+ ``any`` for a dev host reading ``file://`` bundles (``None``: every URL
222
+ admitted). Unset is refused rather than read as ``any``: a deployed host
223
+ with no allowlist is a host whose secret can aim it anywhere, and a
224
+ template's placeholder (``https://<bucket>.s3...``) is refused by name."""
225
+ if value is None or not value.strip():
226
+ fail(
227
+ "host",
228
+ "WORLDS_HOST_OBJECT_STORE is unset; name the object store's origins "
229
+ f"(https://bucket.s3.amazonaws.com) or {OBJECT_STORE_ANY!r} for a dev host",
230
+ )
231
+ raise AssertionError # fail() never returns
232
+ if value.strip().lower() == OBJECT_STORE_ANY:
233
+ return None
234
+ origins = set()
235
+ for item in value.split(","):
236
+ item = item.strip()
237
+ if not item:
238
+ continue
239
+ parts = urllib.parse.urlsplit(item)
240
+ if parts.scheme not in ("http", "https") or not parts.netloc or parts.path not in ("", "/"):
241
+ fail(
242
+ "host",
243
+ f"WORLDS_HOST_OBJECT_STORE entries are origins like https://bucket.s3.amazonaws.com, not {item!r}",
244
+ )
245
+ if not _ORIGIN_HOST.match(parts.netloc):
246
+ fail(
247
+ "host",
248
+ f"WORLDS_HOST_OBJECT_STORE names a placeholder, not a host: {item!r}; set the real bucket origin",
249
+ )
250
+ origins.add(f"{parts.scheme}://{parts.netloc}".lower())
251
+ return frozenset(origins)
252
+
253
+
254
+ def check_object_store_url(url: object, origins: frozenset[str] | None, what: str) -> str:
255
+ """The presigned URL the control plane handed us, once it is a string on one
256
+ of the object store's origins. The control plane is trusted with the host
257
+ secret, not with the host's network: a URL pointing anywhere else (the
258
+ metadata service, a cluster Service, the internet) is refused before a
259
+ byte moves, so a leaked secret cannot turn a session's state or the host's
260
+ reach into someone else's download."""
261
+ if not isinstance(url, str) or not url:
262
+ fail("host", f"{what} must be a presigned url")
263
+ raise AssertionError # fail() never returns
264
+ if origins is None:
265
+ return url
266
+ parts = urllib.parse.urlsplit(url)
267
+ origin = f"{parts.scheme}://{parts.netloc}".lower()
268
+ if origin not in origins:
269
+ fail("host", f"{what} must be on the object store ({', '.join(sorted(origins))}), not {origin}")
270
+ return url
271
+
272
+
273
+ def file_digest(path: Path) -> tuple[str, int]:
274
+ """The hex sha256 and size of a file, streamed."""
275
+ digest = hashlib.sha256()
276
+ size = 0
277
+ with open(path, "rb") as fh:
278
+ for chunk in iter(lambda: fh.read(1 << 20), b""):
279
+ digest.update(chunk)
280
+ size += len(chunk)
281
+ return digest.hexdigest(), size
282
+
283
+
284
+ def download_verified(url: str, sha256: str, dest: Path) -> Path:
285
+ """Fetch ``url`` into ``dest`` and refuse it unless its bytes hash to
286
+ ``sha256``; a mismatch leaves nothing behind."""
287
+ dest.parent.mkdir(parents=True, exist_ok=True)
288
+ staging = dest.with_name(f".{dest.name}.{uuid.uuid4().hex}")
289
+ try:
290
+ try:
291
+ with (
292
+ urllib.request.urlopen(urllib.request.Request(url), timeout=TRANSFER_TIMEOUT_S) as r,
293
+ open(staging, "wb") as out,
294
+ ):
295
+ shutil.copyfileobj(r, out)
296
+ except Exception as error: # noqa: BLE001 - name the stage in the message
297
+ fail("host", f"could not download the session export: {error}")
298
+ actual, _ = file_digest(staging)
299
+ if actual != sha256.lower():
300
+ fail("host", f"session export checksum mismatch: expected {sha256}, got {actual}")
301
+ os.replace(staging, dest)
302
+ finally:
303
+ staging.unlink(missing_ok=True)
304
+ return dest
305
+
306
+
307
+ def upload_file(url: str, path: Path, sha256: str, size: int, content_type: str = EXPORT_CONTENT_TYPE) -> None:
308
+ """PUT a file to a presigned URL, sending exactly the headers the platform
309
+ signed for it: content type, length and the base64 sha256 the object store
310
+ verifies on receipt."""
311
+ headers = {
312
+ "Content-Type": content_type,
313
+ "Content-Length": str(size),
314
+ "x-amz-checksum-sha256": base64.b64encode(bytes.fromhex(sha256)).decode("ascii"),
315
+ }
316
+ with open(path, "rb") as fh:
317
+ request = urllib.request.Request(url, data=fh, method="PUT", headers=headers)
318
+ try:
319
+ with urllib.request.urlopen(request, timeout=TRANSFER_TIMEOUT_S) as response:
320
+ response.read()
321
+ except urllib.error.HTTPError as error:
322
+ fail("host", f"the export upload was refused: HTTP {error.code} {error.read()[:200]!r}")
323
+ except urllib.error.URLError as error:
324
+ fail("host", f"could not upload the export: {error.reason}")
325
+
326
+
327
+ # --------------------------------------------------------------------------- #
328
+ # Locating and reading a world tree.
329
+ # --------------------------------------------------------------------------- #
330
+
331
+
332
+ def world_tree(bundle: Path) -> Path:
333
+ """The gateway-world tree inside an extracted bundle. ``bench push`` packages
334
+ a schema world as its api bundle with the tree under ``server/world``; a bare
335
+ tree is its own root."""
336
+ nested = bundle / BUNDLED_WORLD
337
+ if (bundle / API_BUNDLE).is_file() and (nested / "schema/world.json").is_file():
338
+ return nested
339
+ if (bundle / "schema/world.json").is_file():
340
+ return bundle
341
+ fail(
342
+ "host",
343
+ f"{bundle} holds no schema world the host can open (no {BUNDLED_WORLD}/schema/world.json "
344
+ "nor schema/world.json); custom-server worlds run on the pod engine",
345
+ )
346
+ raise AssertionError # fail() never returns
347
+
348
+
349
+ def _scenario_dir(bundle: Path, tree: Path, task: str) -> Path | None:
350
+ """A task's scenario directory, checked at the bundle root first (where
351
+ ``build_api_bundle`` writes them) then in the tree."""
352
+ for root in (bundle / SCENARIOS / task, tree / SCENARIOS / task):
353
+ if (root / "scenario.toml").is_file():
354
+ return root
355
+ return None
356
+
357
+
358
+ # --------------------------------------------------------------------------- #
359
+ # Per-version cache: bundle + snapshot + parsed spec, fetched once per host.
360
+ # --------------------------------------------------------------------------- #
361
+
362
+
363
+ @dataclass
364
+ class CachedVersion:
365
+ """One cached version, of either kind the host serves: an api world (a
366
+ connector ``spec`` and ``settings``, sessions cloned from ``snapshot``) or a
367
+ mock-tools world (``tools``: handlers, tasks and verifiers, sessions copied
368
+ from its template). Exactly one of ``spec`` and ``tools`` is set."""
369
+
370
+ version_id: str
371
+ bundle: Path
372
+ tree: Path
373
+ spec: Spec | None
374
+ settings: ApiSettings | None
375
+ snapshot: Path | None
376
+ #: The world's own code, imported once per version — only in a process
377
+ #: that runs handlers (a worker). None in the supervisor.
378
+ handlers: dict[str, Handler] | None = None
379
+ #: The version's prebuilt static ui, when its manifest declares one.
380
+ ui: UiSurface | None = None
381
+ #: A mock-tools version; None for an api world.
382
+ tools: ToolsVersion | None = None
383
+
384
+ @property
385
+ def needs_worker(self) -> bool:
386
+ """Whether the version runs world-authored code: a mock-tools world
387
+ always (its tools ARE handlers), an api world when it declares some."""
388
+ if self.tools is not None:
389
+ return True
390
+ return self.settings is not None and self.settings.handlers is not None
391
+
392
+ def router(self) -> Router:
393
+ """A fresh router per session: routes are compiled from the shared spec,
394
+ but the oauth2 tokens a world mints must stay that session's own."""
395
+ assert self.spec is not None and self.settings is not None
396
+ return Router(self.spec, self.settings)
397
+
398
+
399
+ class VersionCache:
400
+ """Every world version this host has seen, on local disk. The first session
401
+ of a version pays for the download and extract; the rest clone its
402
+ snapshot. Nothing here is per session, so a hot version is cheap to serve
403
+ many times over."""
404
+
405
+ def __init__(self, root: Path, *, run_handlers: bool = False, max_bundle_mb: int = DEFAULT_MAX_BUNDLE_MB):
406
+ self.root = Path(root)
407
+ self.root.mkdir(parents=True, exist_ok=True)
408
+ #: Whether this process imports a version's ``handlers.py``. Off in the
409
+ #: supervisor, on in a version worker.
410
+ self.run_handlers = run_handlers
411
+ self.max_bundle_bytes = max_bundle_mb * (1 << 20)
412
+ self._lock = threading.Lock()
413
+ self._versions: dict[str, CachedVersion] = {}
414
+ #: One build at a time per version: the first open of a version does the
415
+ #: download, concurrent opens of the same version wait for it.
416
+ self._building: dict[str, threading.Lock] = {}
417
+
418
+ def ensure(
419
+ self,
420
+ version_id: str,
421
+ bundle_url: str,
422
+ snapshot: dict[str, str] | None,
423
+ ) -> CachedVersion:
424
+ with self._lock:
425
+ hit = self._versions.get(version_id)
426
+ if hit is not None:
427
+ return hit
428
+ building = self._building.setdefault(version_id, threading.Lock())
429
+ with building:
430
+ with self._lock:
431
+ hit = self._versions.get(version_id)
432
+ if hit is not None:
433
+ return hit
434
+ try:
435
+ cached = self._build(version_id, bundle_url, snapshot)
436
+ with self._lock:
437
+ self._versions[version_id] = cached
438
+ return cached
439
+ finally:
440
+ # Success or failure, the build slot is over: a failed build
441
+ # leaves nothing behind, and the next open of the version starts
442
+ # a fresh one.
443
+ with self._lock:
444
+ self._building.pop(version_id, None)
445
+
446
+ def has(self, version_id: str) -> bool:
447
+ with self._lock:
448
+ return version_id in self._versions
449
+
450
+ def version_ids(self) -> list[str]:
451
+ with self._lock:
452
+ return sorted(self._versions)
453
+
454
+ def _build(self, version_id: str, bundle_url: str, snapshot: dict[str, str] | None) -> CachedVersion:
455
+ bundle = self.root / version_id / "bundle"
456
+ if not bundle.is_dir():
457
+ self._download(bundle_url, bundle)
458
+ if is_tools_bundle(bundle):
459
+ tools = load_tools_version(
460
+ bundle, self.root / version_id / "template.db", snapshot, run_handlers=self.run_handlers
461
+ )
462
+ return CachedVersion(
463
+ version_id=version_id,
464
+ bundle=bundle,
465
+ tree=bundle,
466
+ spec=None,
467
+ settings=None,
468
+ snapshot=None,
469
+ tools=tools,
470
+ )
471
+ tree = world_tree(bundle)
472
+ contract = check_bundle(tree)
473
+ if snapshot:
474
+ snapshots.fetch(tree, contract, snapshot["url"], snapshot["sha256"], snapshot["digest"])
475
+ spec = load_spec(tree)
476
+ settings = load_api_settings(tree, spec)
477
+ handlers = load_handlers(settings) if self.run_handlers and settings.handlers is not None else None
478
+ snap = snapshots.default_snapshot(tree)
479
+ if not snap.is_file():
480
+ fail(
481
+ "host",
482
+ f"version {version_id} has no compiled snapshot (db/state.sqlite); "
483
+ "publish a data snapshot or pin one on the version",
484
+ )
485
+ return CachedVersion(
486
+ version_id=version_id,
487
+ bundle=bundle,
488
+ tree=tree,
489
+ spec=spec,
490
+ settings=settings,
491
+ snapshot=snap,
492
+ handlers=handlers,
493
+ ui=load_ui_surface(bundle),
494
+ )
495
+
496
+ def _download(self, url: str, dest: Path) -> None:
497
+ """Fetch and extract into a sibling staging directory, then rename it into
498
+ place: ``dest`` exists only once it is complete, so a download that dies
499
+ halfway is never mistaken for a cached bundle on the next open."""
500
+ dest.parent.mkdir(parents=True, exist_ok=True)
501
+ staging = dest.parent / f".{dest.name}.{uuid.uuid4().hex}"
502
+ tgz = staging.with_suffix(".tar.gz")
503
+ limit = self.max_bundle_bytes
504
+ try:
505
+ try:
506
+ with (
507
+ urllib.request.urlopen(urllib.request.Request(url), timeout=TRANSFER_TIMEOUT_S) as r,
508
+ open(tgz, "wb") as out,
509
+ ):
510
+ copied = 0
511
+ for chunk in iter(lambda: r.read(1 << 20), b""):
512
+ copied += len(chunk)
513
+ if copied > limit:
514
+ raise ValueError(f"the bundle is larger than {limit} bytes")
515
+ out.write(chunk)
516
+ except Exception as error: # noqa: BLE001 - name the stage in the message
517
+ fail("host", f"could not download the bundle from {url}: {error}")
518
+ staging.mkdir()
519
+ with tarfile.open(tgz, "r:gz") as t:
520
+ members = t.getmembers()
521
+ if len(members) > MAX_BUNDLE_MEMBERS:
522
+ fail("host", f"the bundle holds {len(members)} files, more than the {MAX_BUNDLE_MEMBERS} allowed")
523
+ extracted = sum(m.size for m in members)
524
+ if extracted > limit:
525
+ fail("host", f"the bundle extracts to {extracted} bytes, more than the {limit} allowed")
526
+ safe_tar_extract(t, staging)
527
+ os.replace(staging, dest)
528
+ finally:
529
+ tgz.unlink(missing_ok=True)
530
+ shutil.rmtree(staging, ignore_errors=True)
531
+
532
+
533
+ # --------------------------------------------------------------------------- #
534
+ # Session state: the cached snapshot, cloned into this session's own file.
535
+ # --------------------------------------------------------------------------- #
536
+
537
+
538
+ class HostWorldState(WorldState):
539
+ """A session's world, opened by cloning the version's cached snapshot into a
540
+ file on the host's own volume — the reflink is the whole per-session cost.
541
+ Otherwise a plain :class:`WorldState`."""
542
+
543
+ def __init__(self, cached: CachedVersion, session_dir: Path):
544
+ self._session_dir = Path(session_dir)
545
+ self._session_dir.mkdir(parents=True, exist_ok=True)
546
+ self._cached = cached
547
+ self._database: Path | None = None
548
+ super().__init__(cached.tree)
549
+
550
+ @property
551
+ def session_dir(self) -> Path:
552
+ return self._session_dir
553
+
554
+ def _open(self, snapshot: Path | None, *, initial: bool = True) -> None:
555
+ # A snapshot start is a clone of the version's cached snapshot into this
556
+ # session's file; the legacy no-snapshot path stays as the base class runs it.
557
+ if snapshot is None:
558
+ return super()._open(None, initial=initial)
559
+ if self.store is not None:
560
+ self.store.close()
561
+ if self._database is not None:
562
+ self._database.unlink(missing_ok=True)
563
+ database = self._session_dir / f"{uuid.uuid4().hex}.db"
564
+ clone_file(snapshot, database)
565
+ self.store = snapshots.SQLiteStore.open(database, self.contract, audit=False)
566
+ self._database = database
567
+
568
+ def export(self, destination: Path) -> Path:
569
+ """The session's state as one sqlite file: the store's own audited
570
+ online backup, consistent while requests keep writing."""
571
+ assert self.store is not None
572
+ with self._lock:
573
+ return self.store.snapshot(destination)
574
+
575
+ def restore(self, source: Path) -> None:
576
+ """Continue from a file :meth:`export` wrote. The file is audited row by
577
+ row on open: it came over the network, not from this store."""
578
+ with self._lock:
579
+ if self.store is not None:
580
+ self.store.close()
581
+ if self._database is not None:
582
+ self._database.unlink(missing_ok=True)
583
+ database = self._session_dir / f"{uuid.uuid4().hex}.db"
584
+ clone_file(source, database)
585
+ self.store = snapshots.SQLiteStore.open(database, self.contract, audit=True)
586
+ self._database = database
587
+
588
+ def close(self) -> None:
589
+ super().close()
590
+ shutil.rmtree(self._session_dir, ignore_errors=True)
591
+
592
+
593
+ @dataclass
594
+ class Export:
595
+ """A session's state copied to a file on this host, hashed, waiting for the
596
+ platform to hand over the presigned URL it goes to."""
597
+
598
+ export_id: str
599
+ path: Path
600
+ sha256: str
601
+ bytes: int
602
+
603
+ def describe(self) -> dict:
604
+ return {
605
+ "exportId": self.export_id,
606
+ "sha256": self.sha256,
607
+ "bytes": self.bytes,
608
+ "contentType": EXPORT_CONTENT_TYPE,
609
+ }
610
+
611
+
612
+ @dataclass
613
+ class Session:
614
+ session_id: str
615
+ version_id: str
616
+ #: The bundled task (a ``scenarios/<task>`` of the version), or None for a
617
+ #: session opened on a task of its own.
618
+ task: str | None
619
+ #: The api world's connector spec; None for a mock-tools world, which has
620
+ #: no http surface (its data plane is 404, its tools are ``call``).
621
+ spec: Spec | None
622
+ token: str | None
623
+ bundle: Path
624
+ tree: Path
625
+ #: Unix time the lease runs out unless renewed.
626
+ lease_until: float
627
+ #: The world, when this process serves it. None for a session placed on a
628
+ #: worker, whose state lives in that process.
629
+ state: HostWorldState | None = None
630
+ app: WorldApp | None = None
631
+ #: The mock-tools world, when this process (a worker) serves one.
632
+ tools: ToolsWorld | None = None
633
+ #: The version worker serving this session, or None for in-process.
634
+ worker: Worker | None = None
635
+ #: The version's static ui, served by the supervisor for this session.
636
+ ui: UiSurface | None = None
637
+ #: The pool browser leased for this session: its relay's ``ws://ip:port``
638
+ #: and the token that relay answers to. None when the session has no
639
+ #: browser surface.
640
+ browser_url: str | None = None
641
+ browser_token: str | None = None
642
+ #: For an oauth2 world: the bearer this world minted for the host itself,
643
+ #: presented on behalf of a caller who authenticated with the session token.
644
+ grant: str | None = None
645
+ #: Requests inside the world right now. A session is only torn down once
646
+ #: this reaches zero, so a close or an expiry never pulls the database out
647
+ #: from under a handler mid-request.
648
+ inflight: int = 0
649
+ #: Set once the session has left the registry; the last request out closes
650
+ #: the state.
651
+ closing: bool = False
652
+ #: Unix time the session was opened; its session-seconds are counted on close.
653
+ opened_at: float = field(default_factory=time.time)
654
+ #: Unix time of the last request that used the world: a data-plane call or
655
+ #: a control operation (seed, reset, grade, call, state). A renewal or an
656
+ #: export is the platform looking, not the session working, and does not
657
+ #: move it. Reported on renew so the platform's idle clock follows traffic
658
+ #: that never passes through it.
659
+ touched_at: float = field(default_factory=time.time)
660
+ #: The last export taken and not yet uploaded, if any.
661
+ export: Export | None = None
662
+ #: Set when the session was opened from an export rather than the snapshot.
663
+ restored: bool = False
664
+ #: The task supplied at open, when the session has one instead of a
665
+ #: bundled task: its seed is applied after the clone, its grader answers
666
+ #: ``grade``, and it lives under ``<session>/task/`` (:mod:`tasks`).
667
+ task_spec: TaskSpec | None = None
668
+
669
+ @property
670
+ def task_key(self) -> tuple[str | None, dict | None]:
671
+ """What an open retry must match: the bundled name or the spec itself."""
672
+ return self.task, self.task_spec.raw if self.task_spec is not None else None
673
+
674
+ def renew(self, lease_s: float) -> None:
675
+ self.lease_until = time.time() + lease_s
676
+
677
+ @property
678
+ def seconds(self) -> float:
679
+ return max(0.0, time.time() - self.opened_at)
680
+
681
+
682
+ def browser_lease(browser: dict[str, str] | None) -> tuple[str | None, str | None]:
683
+ """The ``browser`` field of an open request: the leased relay's ``url`` and
684
+ ``token``, both required when either is given, or ``(None, None)``."""
685
+ if not browser:
686
+ return None, None
687
+ url, token = browser.get("url"), browser.get("token")
688
+ if not isinstance(url, str) or not url.startswith("ws://") or not isinstance(token, str) or not token:
689
+ fail("host", "browser must be {url: 'ws://host:port', token: <relay token>}")
690
+ raise AssertionError # fail() never returns
691
+ return url.rstrip("/"), token
692
+
693
+
694
+ # --------------------------------------------------------------------------- #
695
+ # The host: the session registry and the control operations.
696
+ # --------------------------------------------------------------------------- #
697
+
698
+
699
+ class WorldHost:
700
+ """The session registry and the control operations, in one of two roles:
701
+
702
+ * the **supervisor** (the deployed host): serves code-free worlds in
703
+ process and places worlds with handlers on its :class:`WorkerPool`;
704
+ * a **worker** (``run_handlers=True``, ``pinned_version=...``): serves one
705
+ version, handlers included, for its supervisor.
706
+ """
707
+
708
+ def __init__(
709
+ self,
710
+ cache_root: Path,
711
+ *,
712
+ lease_s: float = DEFAULT_LEASE_S,
713
+ advertise_url: str | None = None,
714
+ workers: WorkerPool | None = None,
715
+ versions_root: Path | None = None,
716
+ run_handlers: bool = False,
717
+ pinned_version: str | None = None,
718
+ limits: dict[str, int] | None = None,
719
+ max_sessions: int = DEFAULT_MAX_SESSIONS,
720
+ max_bundle_mb: int = DEFAULT_MAX_BUNDLE_MB,
721
+ object_store: frozenset[str] | None = None,
722
+ ):
723
+ #: Where every presigned URL the control plane hands this host must
724
+ #: point (bundle, snapshot, export upload, restore); None admits any.
725
+ self.object_store = object_store
726
+ self.cache = VersionCache(
727
+ versions_root or Path(cache_root) / "versions", run_handlers=run_handlers, max_bundle_mb=max_bundle_mb
728
+ )
729
+ self.sessions_root = Path(cache_root) / "sessions"
730
+ # Sessions live in this process; whatever a previous one left on disk is
731
+ # unreachable, so a start begins with an empty sessions tree.
732
+ shutil.rmtree(self.sessions_root, ignore_errors=True)
733
+ self.sessions_root.mkdir(parents=True, exist_ok=True)
734
+ self.lease_s = lease_s
735
+ #: The address the platform pins a session to, so control calls reach
736
+ #: this replica and not the Service's round-robin. None = one replica.
737
+ self.advertise_url = advertise_url
738
+ self.workers = workers
739
+ self.run_handlers = run_handlers
740
+ self.pinned_version = pinned_version
741
+ #: The rlimits this process runs under (a worker's), for ``/healthz``.
742
+ self.limits = limits
743
+ #: Sessions this host admits before ``open`` answers 503 and the
744
+ #: platform places elsewhere; the heartbeat advertises it as capacity.
745
+ self.max_sessions = max_sessions
746
+ #: Set on SIGTERM: no new sessions, existing ones served until they
747
+ #: close or the drain grace runs out.
748
+ self.draining = False
749
+ self.started_at = time.time()
750
+ self._lock = threading.Lock()
751
+ self._sessions: dict[str, Session] = {}
752
+ #: Ids of sessions being opened: each holds one of ``max_sessions`` from
753
+ #: the moment its open is admitted, so concurrent opens cannot all pass
754
+ #: the capacity check and land past the limit.
755
+ self._opening: set[str] = set()
756
+ #: Opens the platform closed before they finished: it gave up waiting
757
+ #: and told this host to forget the session, so the session is closed
758
+ #: the moment it exists rather than held until its lease runs out.
759
+ self._abandoned: set[str] = set()
760
+ #: Lifetime counts for ``/metrics``: opens, closes by reason, requests
761
+ #: by surface, session-seconds of closed sessions.
762
+ self._counters: Counter[str] = Counter()
763
+ self._session_seconds = 0.0
764
+
765
+ # -- lifecycle ---------------------------------------------------------- #
766
+
767
+ def open(
768
+ self,
769
+ session_id: str,
770
+ version_id: str,
771
+ bundle_url: str,
772
+ task: str | None,
773
+ token: str | None,
774
+ snapshot: dict[str, str] | None,
775
+ browser: dict[str, str] | None = None,
776
+ restore: dict[str, str] | None = None,
777
+ task_spec: Any = None,
778
+ ) -> dict:
779
+ """Open ``session_id`` on ``version_id``, either on the bundled ``task``
780
+ or on ``task_spec``, a session task of its own (:mod:`tasks`); exactly
781
+ one of the two."""
782
+ if self.pinned_version is not None and version_id != self.pinned_version:
783
+ fail("host", f"this worker serves version {self.pinned_version} only, not {version_id}")
784
+ if (task is None) == (task_spec is None):
785
+ fail("host", "an open names a bundled task or carries a task spec, one of the two")
786
+ if task is not None and not isinstance(task, str):
787
+ fail("host", "task is the bundled task's name")
788
+ spec = parse_task_spec(task_spec) if task_spec is not None else None
789
+ with self._lock:
790
+ existing = self._sessions.get(session_id)
791
+ if existing is None:
792
+ # A retry of an open this host already served is answered even
793
+ # while draining: refusing it would send the platform to open
794
+ # the same id elsewhere while this copy holds the live writes.
795
+ if self.draining:
796
+ fail("draining", "this host is draining and takes no new sessions")
797
+ self._admit(session_id)
798
+ if existing is not None:
799
+ # The platform retried an open that already succeeded: the same
800
+ # session is the answer, never a rebuild that drops its writes.
801
+ wanted = (version_id, (task, spec.raw if spec is not None else None), token)
802
+ if (existing.version_id, existing.task_key, existing.token) != wanted:
803
+ fail("host", f"session {session_id} is already open as another version, task or token")
804
+ return self._describe(existing)
805
+ try:
806
+ described = self._open(session_id, version_id, bundle_url, task, spec, token, snapshot, browser, restore)
807
+ except Exception:
808
+ self.tally("open_failures")
809
+ with self._lock:
810
+ self._opening.discard(session_id)
811
+ self._abandoned.discard(session_id)
812
+ raise
813
+ with self._lock:
814
+ self._opening.discard(session_id)
815
+ abandoned = session_id in self._abandoned
816
+ self._abandoned.discard(session_id)
817
+ if abandoned:
818
+ self.close(session_id)
819
+ fail(UNKNOWN_SESSION, f"session {session_id} was closed while it opened")
820
+ return described
821
+
822
+ def _admit(self, session_id: str) -> None:
823
+ """Reserve one of this host's ``max_sessions`` for ``session_id`` while it
824
+ opens. Under ``_lock``."""
825
+ if session_id in self._opening:
826
+ fail("host", f"session {session_id} is being opened already")
827
+ if len(self._sessions) + len(self._opening) >= self.max_sessions:
828
+ fail("at_capacity", f"this host holds its {self.max_sessions} sessions; place elsewhere")
829
+ self._opening.add(session_id)
830
+
831
+ def _open(
832
+ self,
833
+ session_id: str,
834
+ version_id: str,
835
+ bundle_url: str,
836
+ task: str | None,
837
+ spec: TaskSpec | None,
838
+ token: str | None,
839
+ snapshot: dict[str, str] | None,
840
+ browser: dict[str, str] | None,
841
+ restore: dict[str, str] | None,
842
+ ) -> dict:
843
+ cached = self.cache.ensure(version_id, *self._trusted_version(bundle_url, snapshot))
844
+ if spec is not None and (cached.tools is not None or cached.spec is None):
845
+ fail("host", "a session task needs an api world; a mock-tools world takes its bundled tasks only")
846
+ if cached.spec is not None:
847
+ unsupported = unsupported_auth(cached.spec)
848
+ if unsupported:
849
+ fail("host", unsupported)
850
+ if spec is not None:
851
+ self._check_task_tools(spec, cached.spec)
852
+ browser_url, browser_token = browser_lease(browser)
853
+ if cached.needs_worker and not self.run_handlers:
854
+ if self.workers is None:
855
+ fail(
856
+ "host",
857
+ f"version {version_id} runs world-authored code ({_code_of(cached)}); "
858
+ "this host runs none in its process and has no version workers",
859
+ )
860
+ return self._open_on_worker(
861
+ cached, session_id, bundle_url, task, spec, token, browser_url, browser_token, restore
862
+ )
863
+ if cached.tools is not None:
864
+ assert task is not None
865
+ return self._open_tools(cached, cached.tools, session_id, task, token, restore)
866
+ assert cached.snapshot is not None and cached.spec is not None
867
+ state = HostWorldState(cached, self.sessions_root / session_id)
868
+ try:
869
+ if spec is not None:
870
+ self._check_task_entities(spec, state)
871
+ write_task_dir(spec, state.session_dir)
872
+ self._apply_task_seed(state, spec)
873
+ else:
874
+ assert task is not None
875
+ self._apply_scenario(cached, state, task)
876
+ restored = self._restore_into(restore, self.sessions_root / session_id, state.restore)
877
+ except Exception:
878
+ state.close()
879
+ raise
880
+ # ``handlers={}`` for a code-free world: WorldApp would otherwise import
881
+ # the module itself, and only a worker may do that.
882
+ app = WorldApp(
883
+ state,
884
+ cached.router(),
885
+ token,
886
+ handlers=cached.handlers or {},
887
+ allowed=frozenset(spec.tools) if spec is not None and spec.tools is not None else None,
888
+ )
889
+ session = Session(
890
+ session_id=session_id,
891
+ version_id=version_id,
892
+ task=task,
893
+ spec=cached.spec,
894
+ token=token,
895
+ bundle=cached.bundle,
896
+ tree=cached.tree,
897
+ lease_until=time.time() + self.lease_s,
898
+ state=state,
899
+ app=app,
900
+ ui=cached.ui,
901
+ browser_url=browser_url,
902
+ browser_token=browser_token,
903
+ restored=restored,
904
+ task_spec=spec,
905
+ )
906
+ if cached.spec.scheme == "oauth2":
907
+ # The host's own client-credentials grant, so a caller admitted by
908
+ # the session token is presented to the world as a granted client.
909
+ issued = app.router.issue_token(
910
+ {"grant_type": "client_credentials", "client_secret": token or ""},
911
+ token,
912
+ )
913
+ session.grant = issued["access_token"] if issued else None
914
+ self._register(session)
915
+ return self._describe(session)
916
+
917
+ def _open_tools(
918
+ self,
919
+ cached: CachedVersion,
920
+ tools: ToolsVersion,
921
+ session_id: str,
922
+ task: str,
923
+ token: str | None,
924
+ restore: dict[str, str] | None,
925
+ ) -> dict:
926
+ """A mock-tools session in this (worker) process: a copy of the version's
927
+ template, executed through its handlers. No http surface, so no ui and
928
+ no browser: the agent reaches it through ``call``."""
929
+ world = ToolsWorld(tools, self.sessions_root / session_id, task)
930
+ try:
931
+ restored = self._restore_into(restore, self.sessions_root / session_id, world.restore)
932
+ except Exception:
933
+ world.close()
934
+ raise
935
+ session = Session(
936
+ session_id=session_id,
937
+ version_id=cached.version_id,
938
+ task=task,
939
+ spec=None,
940
+ token=token,
941
+ bundle=cached.bundle,
942
+ tree=cached.tree,
943
+ lease_until=time.time() + self.lease_s,
944
+ tools=world,
945
+ restored=restored,
946
+ )
947
+ self._register(session)
948
+ return self._describe(session)
949
+
950
+ def _open_on_worker(
951
+ self,
952
+ cached: CachedVersion,
953
+ session_id: str,
954
+ bundle_url: str,
955
+ task: str | None,
956
+ spec: TaskSpec | None,
957
+ token: str | None,
958
+ browser_url: str | None,
959
+ browser_token: str | None,
960
+ restore: dict[str, str] | None,
961
+ ) -> dict:
962
+ """Place the session on the version's worker. The worker opens it from
963
+ the bundle this process already cached (no snapshot to fetch: the file
964
+ is on disk), and the supervisor records only where it went. The ui and
965
+ browser surfaces are the supervisor's own (static files, a tunnel), so
966
+ the worker hears nothing of them. An export to restore from is fetched
967
+ and verified here too, then handed to the worker as a local file.
968
+ A worker every one of whose peers holds sessions cannot be started
969
+ past the pool's memory budget: that is ``at_capacity``, and the
970
+ platform places the session on another replica."""
971
+ assert self.workers is not None
972
+ try:
973
+ worker = self.workers.get(cached.version_id)
974
+ except WorkerCapacity as error:
975
+ fail("at_capacity", str(error))
976
+ # The session is the worker's from here, so another version's open
977
+ # cannot stop this worker as session-less while it is being placed.
978
+ with self._lock:
979
+ worker.sessions.add(session_id)
980
+ source = None
981
+ try:
982
+ source = self._restore_source(restore, self._worker_session_dir(worker, session_id))
983
+ self._worker_control(
984
+ worker,
985
+ "POST",
986
+ "/__host/sessions",
987
+ {
988
+ "sessionId": session_id,
989
+ "versionId": cached.version_id,
990
+ "bundleUrl": bundle_url,
991
+ "task": task,
992
+ "taskSpec": spec.raw if spec is not None else None,
993
+ "token": token,
994
+ "restore": {"path": str(source)} if source is not None else None,
995
+ },
996
+ )
997
+ except BaseException:
998
+ with self._lock:
999
+ worker.sessions.discard(session_id)
1000
+ if not worker.sessions:
1001
+ worker.idle_since = time.time()
1002
+ raise
1003
+ finally:
1004
+ if source is not None:
1005
+ source.unlink(missing_ok=True)
1006
+ session = Session(
1007
+ session_id=session_id,
1008
+ version_id=cached.version_id,
1009
+ task=task,
1010
+ spec=cached.spec,
1011
+ token=token,
1012
+ bundle=cached.bundle,
1013
+ tree=cached.tree,
1014
+ lease_until=time.time() + self.lease_s,
1015
+ worker=worker,
1016
+ ui=cached.ui,
1017
+ browser_url=browser_url,
1018
+ browser_token=browser_token,
1019
+ restored=source is not None,
1020
+ task_spec=spec,
1021
+ )
1022
+ self._register(session)
1023
+ return self._describe(session)
1024
+
1025
+ def _worker_session_dir(self, worker: Worker, session_id: str) -> Path:
1026
+ """Where a worker keeps a session's files: its own ``WorldHost`` runs
1027
+ under ``workers/<version>`` of this host's cache root (:mod:`worker`)."""
1028
+ assert self.workers is not None
1029
+ return self.workers.root / "workers" / worker.version_id / "sessions" / session_id
1030
+
1031
+ def _restore_into(self, restore: dict[str, str] | None, session_dir: Path, load: Callable[[Path], None]) -> bool:
1032
+ """Load the export ``restore`` names into a freshly opened world, then
1033
+ drop the fetched file. False when there was nothing to restore."""
1034
+ source = self._restore_source(restore, session_dir)
1035
+ if source is None:
1036
+ return False
1037
+ try:
1038
+ load(source)
1039
+ finally:
1040
+ source.unlink(missing_ok=True)
1041
+ return True
1042
+
1043
+ def _restore_source(self, restore: dict[str, str] | None, session_dir: Path) -> Path | None:
1044
+ """The export a session continues from, as a local file, or None for a
1045
+ fresh session. The platform names it by ``{url, sha256}`` and the file is
1046
+ fetched into the session's directory and verified; a supervisor names it
1047
+ to its worker by ``{path}``, already verified, and only the supervisor
1048
+ reaches a worker's control socket."""
1049
+ if not restore:
1050
+ return None
1051
+ path = restore.get("path")
1052
+ if path is not None:
1053
+ if not self.run_handlers:
1054
+ fail("host", "restore by path is a worker's; name the export by {url, sha256}")
1055
+ source = Path(path)
1056
+ if not source.is_file():
1057
+ fail("host", f"restore file {path} is missing")
1058
+ return source
1059
+ url, sha256 = restore.get("url"), restore.get("sha256")
1060
+ if not isinstance(url, str) or not isinstance(sha256, str) or len(sha256) != 64:
1061
+ fail("host", "restore must be {url: <presigned export url>, sha256: <hex digest>}")
1062
+ raise AssertionError # fail() never returns
1063
+ check_object_store_url(url, self.object_store, "restore.url")
1064
+ return download_verified(url, sha256, session_dir / "restore.sqlite")
1065
+
1066
+ def _trusted_version(self, bundle_url: str, snapshot: dict[str, str] | None) -> tuple[str, dict[str, str] | None]:
1067
+ check_object_store_url(bundle_url, self.object_store, "bundleUrl")
1068
+ if snapshot:
1069
+ check_object_store_url(snapshot.get("url"), self.object_store, "snapshot.url")
1070
+ return bundle_url, snapshot
1071
+
1072
+ def _register(self, session: Session) -> None:
1073
+ with self._lock:
1074
+ replaced = self._retire(self._sessions.pop(session.session_id, None))
1075
+ self._sessions[session.session_id] = session
1076
+ self._counters["opened"] += 1
1077
+ if replaced is not None:
1078
+ self._finish(replaced, reason="replaced")
1079
+
1080
+ def _describe(self, session: Session) -> dict:
1081
+ return {
1082
+ "sessionId": session.session_id,
1083
+ "versionId": session.version_id,
1084
+ "task": session.task,
1085
+ "taskName": session.task if session.task_spec is None else session.task_spec.name,
1086
+ "taskSource": "bundled" if session.task_spec is None else "inline",
1087
+ "tools": (
1088
+ list(session.task_spec.tools)
1089
+ if session.task_spec is not None and session.task_spec.tools is not None
1090
+ else None
1091
+ ),
1092
+ "ready": True,
1093
+ "hostUrl": self.advertise_url,
1094
+ "leaseSeconds": self.lease_s,
1095
+ "leaseUntil": session.lease_until,
1096
+ "workerPid": session.worker.pid if session.worker is not None else None,
1097
+ "ui": {"entry": session.ui.entry, "apiPath": session.ui.api_path} if session.ui is not None else None,
1098
+ "browser": session.browser_url is not None,
1099
+ "restored": session.restored,
1100
+ "openedAt": session.opened_at,
1101
+ "sessionSeconds": session.seconds,
1102
+ "touchedAt": session.touched_at,
1103
+ }
1104
+
1105
+ def close(self, session_id: str) -> dict:
1106
+ with self._lock:
1107
+ session = self._sessions.pop(session_id, None)
1108
+ idle = self._retire(session)
1109
+ if session is None and session_id in self._opening:
1110
+ self._abandoned.add(session_id)
1111
+ if idle is not None:
1112
+ self._finish(idle, reason="closed")
1113
+ if session is None:
1114
+ return {"closed": True, "openedAt": None, "sessionSeconds": 0.0}
1115
+ return {"closed": True, "openedAt": session.opened_at, "sessionSeconds": session.seconds}
1116
+
1117
+ # -- durability: a session's state as a file the platform keeps ---------- #
1118
+
1119
+ def export(self, session_id: str, known_sha256: str | None = None) -> dict:
1120
+ """Copy the session's state to a file on this host and hash it. The copy
1121
+ is consistent (sqlite's online backup under the store's lock) and waits
1122
+ for :meth:`upload_export`; taking a new one discards the last. A copy
1123
+ that hashes to ``known_sha256`` (the platform's last export of the
1124
+ session) is thrown away and reported ``unchanged``: nothing to upload."""
1125
+ with self.using(session_id, touch=False) as session:
1126
+ if session.worker is not None:
1127
+ answer = self._worker_control(session.worker, "POST", f"/__host/sessions/{session_id}/export", {})
1128
+ path = self._worker_export_path(session, answer)
1129
+ else:
1130
+ exports = self.sessions_root / session_id / "exports"
1131
+ shutil.rmtree(exports, ignore_errors=True)
1132
+ path = exports / f"{uuid.uuid4().hex}.sqlite"
1133
+ if session.tools is not None:
1134
+ session.tools.export(path)
1135
+ else:
1136
+ assert session.state is not None
1137
+ session.state.export(path)
1138
+ sha256, size = file_digest(path)
1139
+ if known_sha256 is not None and sha256 == known_sha256:
1140
+ path.unlink(missing_ok=True)
1141
+ with self._lock:
1142
+ session.export = None
1143
+ self._counters["exports_unchanged"] += 1
1144
+ return {"exportId": None, "sha256": sha256, "bytes": size, "unchanged": True}
1145
+ export = Export(export_id=path.stem, path=path, sha256=sha256, bytes=size)
1146
+ with self._lock:
1147
+ session.export = export
1148
+ self._counters["exports"] += 1
1149
+ described = export.describe()
1150
+ if self.run_handlers:
1151
+ described["path"] = str(path) # for the supervisor, which uploads on the worker's behalf
1152
+ return described
1153
+
1154
+ def _worker_export_path(self, session: Session, answer: dict) -> Path:
1155
+ """The file a worker's export landed in, admitted only from inside that
1156
+ worker's directory for this session: the worker runs tenant code, so
1157
+ what it names is checked, and the bytes are hashed here, not trusted."""
1158
+ assert session.worker is not None
1159
+ exports = self._worker_session_dir(session.worker, session.session_id) / "exports"
1160
+ path = Path(str(answer.get("path") or ""))
1161
+ try:
1162
+ resolved = path.resolve(strict=True)
1163
+ resolved.relative_to(exports.resolve())
1164
+ except (OSError, ValueError):
1165
+ fail("host", f"the worker's export is not where a worker's exports go: {path}")
1166
+ raise AssertionError # fail() never returns
1167
+ return resolved
1168
+
1169
+ def upload_export(self, session_id: str, export_id: str, url: str) -> dict:
1170
+ """PUT the export :meth:`export` took to the presigned URL the platform
1171
+ signed for its hash and size; the local copy goes once it is up."""
1172
+ check_object_store_url(url, self.object_store, "url")
1173
+ with self.using(session_id, touch=False) as session:
1174
+ export = session.export
1175
+ if export is None or export.export_id != export_id:
1176
+ fail("host", f"no export {export_id} for session {session_id}; take one first")
1177
+ raise AssertionError # fail() never returns
1178
+ upload_file(url, export.path, export.sha256, export.bytes)
1179
+ export.path.unlink(missing_ok=True)
1180
+ with self._lock:
1181
+ session.export = None
1182
+ self._counters["exports_uploaded"] += 1
1183
+ return {"uploaded": True, **export.describe()}
1184
+
1185
+ def renew(self, session_id: str) -> dict:
1186
+ with self.using(session_id, touch=False) as session:
1187
+ return self._describe(session)
1188
+
1189
+ def expire_leases(self, now: float | None = None) -> list[str]:
1190
+ """Close every session whose lease has run out; returns their ids. The
1191
+ check and the removal happen under one lock, so a renewal that lands
1192
+ in between keeps its session, and a session with a request in flight
1193
+ is left for the next sweep (the request renewed it anyway)."""
1194
+ now = time.time() if now is None else now
1195
+ with self._lock:
1196
+ expired = [s for s in self._sessions.values() if s.lease_until <= now and s.inflight == 0]
1197
+ for session in expired:
1198
+ del self._sessions[session.session_id]
1199
+ self._retire(session)
1200
+ for session in expired:
1201
+ self._finish(session, reason="expired")
1202
+ return [s.session_id for s in expired]
1203
+
1204
+ def _retire(self, session: Session | None) -> Session | None:
1205
+ """Mark a session that has left the registry; returns it when nobody is
1206
+ inside it, so the caller closes the state now. Otherwise the last
1207
+ request out does (see :meth:`_release`). Call under the lock."""
1208
+ if session is None:
1209
+ return None
1210
+ session.closing = True
1211
+ return session if session.inflight == 0 else None
1212
+
1213
+ def _finish(self, session: Session, *, reason: str = "closed") -> None:
1214
+ """Free a session that has left the registry and has nobody inside it:
1215
+ close its state here, or tell its worker to."""
1216
+ with self._lock:
1217
+ self._counters[f"closed.{reason}"] += 1
1218
+ self._session_seconds += session.seconds
1219
+ if session.worker is None:
1220
+ if session.tools is not None:
1221
+ session.tools.close()
1222
+ return
1223
+ assert session.state is not None
1224
+ session.state.close()
1225
+ return
1226
+ worker = session.worker
1227
+ with self._lock:
1228
+ worker.sessions.discard(session.session_id)
1229
+ if not worker.sessions:
1230
+ worker.idle_since = time.time()
1231
+ try:
1232
+ self._worker_control(worker, "DELETE", f"/__host/sessions/{session.session_id}", None)
1233
+ except (WorkerGone, ContractError):
1234
+ pass # the worker is gone, or already forgot it: either way the session is over
1235
+
1236
+ def _worker_control(self, worker: Worker, method: str, path: str, body: dict | None) -> dict:
1237
+ """One control call to a worker, its answer as this host would give it:
1238
+ the worker's refusal is re-raised as ours, and a worker that has died
1239
+ takes every session it held with it."""
1240
+ raw = json.dumps(body).encode("utf-8") if body is not None else b""
1241
+ try:
1242
+ status, _, payload = worker.request(method, path, {"content-type": "application/json"}, raw, control=True)
1243
+ except WorkerGone:
1244
+ self.lose_worker(worker)
1245
+ raise
1246
+ answer = json.loads(payload or b"null")
1247
+ if status != 200:
1248
+ detail = answer if isinstance(answer, dict) else {}
1249
+ fail(str(detail.get("error") or "host"), str(detail.get("detail") or answer))
1250
+ return answer if isinstance(answer, dict) else {"result": answer}
1251
+
1252
+ def lose_worker(self, worker: Worker) -> list[str]:
1253
+ """A worker exited: forget it and every session it held. Returns the
1254
+ session ids, which are now gone from this host (the platform learns of
1255
+ it on the next renew)."""
1256
+ if self.workers is not None:
1257
+ self.workers.drop(worker.version_id)
1258
+ with self._lock:
1259
+ lost = [sid for sid, s in self._sessions.items() if s.worker is worker]
1260
+ for sid in lost:
1261
+ session = self._sessions.pop(sid)
1262
+ session.closing = True
1263
+ self._counters["closed.lost"] += 1
1264
+ self._session_seconds += session.seconds
1265
+ worker.sessions.clear()
1266
+ return lost
1267
+
1268
+ def reap_workers(self) -> list[str]:
1269
+ """Stop version workers that have had no session for a while."""
1270
+ return self.workers.reap_idle() if self.workers is not None else []
1271
+
1272
+ def session_count(self) -> int:
1273
+ with self._lock:
1274
+ return len(self._sessions)
1275
+
1276
+ def worker_count(self) -> int:
1277
+ return len(self.workers._workers) if self.workers is not None else 0
1278
+
1279
+ # -- fleet: warming, draining, what the platform and prometheus see ------- #
1280
+
1281
+ def warm(self, version_id: str, bundle_url: str, snapshot: dict[str, str] | None) -> dict:
1282
+ """Cache a version before any session asks for it, so the first open
1283
+ pays nothing. Idempotent: a version already here is reported, not
1284
+ rebuilt."""
1285
+ already = self.cache.has(version_id)
1286
+ if not already:
1287
+ self.cache.ensure(version_id, *self._trusted_version(bundle_url, snapshot))
1288
+ return {"versionId": version_id, "cached": True, "alreadyCached": already}
1289
+
1290
+ def drain(self) -> None:
1291
+ """Stop taking sessions; the ones here are served until they close."""
1292
+ self.draining = True
1293
+
1294
+ def tally(self, key: str) -> None:
1295
+ with self._lock:
1296
+ self._counters[key] += 1
1297
+
1298
+ def stats(self) -> dict:
1299
+ """What this host is and holds, for the heartbeat and ``/__host/stats``:
1300
+ the placement inputs (address, load, capacity, draining, which versions
1301
+ are already cached) plus the lifetime counters."""
1302
+ with self._lock:
1303
+ sessions = len(self._sessions)
1304
+ versions_in_use = sorted({s.version_id for s in self._sessions.values()})
1305
+ counters = dict(self._counters)
1306
+ session_seconds = self._session_seconds
1307
+ return {
1308
+ "hostUrl": self.advertise_url,
1309
+ "sessions": sessions,
1310
+ "capacity": self.max_sessions,
1311
+ "draining": self.draining,
1312
+ "workers": self.worker_count(),
1313
+ "versions": self.cache.version_ids(),
1314
+ "versionsInUse": versions_in_use,
1315
+ "uptimeSeconds": time.time() - self.started_at,
1316
+ "counters": counters,
1317
+ "sessionSeconds": session_seconds,
1318
+ "pid": os.getpid(),
1319
+ "limits": self.limits,
1320
+ }
1321
+
1322
+ def metrics_text(self) -> str:
1323
+ """The prometheus exposition of :meth:`stats`. ``worlds_host_sessions``
1324
+ is the number the fleet scales on (sessions per host, not cpu)."""
1325
+ stats = self.stats()
1326
+ counters: dict[str, int] = stats["counters"]
1327
+ lines = [
1328
+ "# TYPE worlds_host_sessions gauge",
1329
+ f"worlds_host_sessions {stats['sessions']}",
1330
+ "# TYPE worlds_host_capacity gauge",
1331
+ f"worlds_host_capacity {stats['capacity']}",
1332
+ "# TYPE worlds_host_draining gauge",
1333
+ f"worlds_host_draining {int(stats['draining'])}",
1334
+ "# TYPE worlds_host_workers gauge",
1335
+ f"worlds_host_workers {stats['workers']}",
1336
+ "# TYPE worlds_host_versions_cached gauge",
1337
+ f"worlds_host_versions_cached {len(stats['versions'])}",
1338
+ "# TYPE worlds_host_uptime_seconds gauge",
1339
+ f"worlds_host_uptime_seconds {stats['uptimeSeconds']:.0f}",
1340
+ "# TYPE worlds_host_sessions_opened_total counter",
1341
+ f"worlds_host_sessions_opened_total {counters.get('opened', 0)}",
1342
+ "# TYPE worlds_host_open_failures_total counter",
1343
+ f"worlds_host_open_failures_total {counters.get('open_failures', 0)}",
1344
+ "# TYPE worlds_host_session_seconds_total counter",
1345
+ f"worlds_host_session_seconds_total {stats['sessionSeconds']:.0f}",
1346
+ "# TYPE worlds_host_exports_total counter",
1347
+ f"worlds_host_exports_total {counters.get('exports', 0)}",
1348
+ "# TYPE worlds_host_exports_uploaded_total counter",
1349
+ f"worlds_host_exports_uploaded_total {counters.get('exports_uploaded', 0)}",
1350
+ "# TYPE worlds_host_sessions_closed_total counter",
1351
+ ]
1352
+ for reason in ("closed", "expired", "lost", "replaced"):
1353
+ lines.append(
1354
+ f'worlds_host_sessions_closed_total{{reason="{reason}"}} {counters.get(f"closed.{reason}", 0)}'
1355
+ )
1356
+ lines.append("# TYPE worlds_host_requests_total counter")
1357
+ for surface in ("api", "ui", "browser", "proxied"):
1358
+ lines.append(f'worlds_host_requests_total{{surface="{surface}"}} {counters.get(f"requests.{surface}", 0)}')
1359
+ return "\n".join(lines) + "\n"
1360
+
1361
+ def _apply_scenario(self, cached: CachedVersion, state: HostWorldState, task: str) -> None:
1362
+ """Start the session in its task's scenario. The compiled snapshot the
1363
+ session is cloned from IS the base state, so a scenario with no seed
1364
+ overlay (the synthesized ``default``) needs nothing; one with an overlay
1365
+ applies it on top of the pinned snapshot the way the runner does."""
1366
+ scenario = _scenario_dir(cached.bundle, cached.tree, task)
1367
+ if scenario is None:
1368
+ if task in state.scenarios():
1369
+ state.apply_scenario(task) # a tree-authored scenario
1370
+ return
1371
+ fail(
1372
+ "host",
1373
+ f"no scenario '{task}' in this world; it has {', '.join(state.scenarios()) or 'default'}",
1374
+ )
1375
+ overlay = snapshots.scenario_overlay(scenario)
1376
+ if not overlay:
1377
+ return # base state; the cloned snapshot already holds it
1378
+ digest = state.digest
1379
+ if digest is None:
1380
+ fail("host", "a scenario overlay needs a compiled snapshot to apply against")
1381
+ state.seed({"base": digest, "rows": overlay, "mode": "replace"})
1382
+
1383
+ @staticmethod
1384
+ def _check_task_entities(spec: TaskSpec, state: HostWorldState) -> None:
1385
+ """A task is refused for a world that does not define every entity it
1386
+ names. The seed would fail on its own; an assertion over an unknown
1387
+ entity would not, selecting no rows, so `absent` or `count: eq 0`
1388
+ checks would pass on a typo and grade a task that can never be graded."""
1389
+ unknown = unknown_task_entities(spec, state.contract.entities)
1390
+ if unknown:
1391
+ defined = ", ".join(sorted(state.contract.entities)) or "none"
1392
+ fail(
1393
+ "task",
1394
+ f"the task names {', '.join(unknown)}, which this world does not define; its entities are {defined}",
1395
+ "/task",
1396
+ )
1397
+
1398
+ @staticmethod
1399
+ def _check_task_tools(spec: TaskSpec, world: Spec) -> None:
1400
+ """A task's allow-list names tools the world serves, every one of them:
1401
+ a misspelt name would narrow the list without anyone noticing."""
1402
+ unknown = unknown_task_tools(spec, world.operations)
1403
+ if unknown:
1404
+ served = ", ".join(sorted(world.operations)) or "none"
1405
+ fail(
1406
+ "task",
1407
+ f"the task allows {', '.join(unknown)}, which this world does not serve; its tools are {served}",
1408
+ "/task/tools",
1409
+ )
1410
+
1411
+ def _apply_task_seed(self, state: HostWorldState, spec: TaskSpec) -> None:
1412
+ """Start the session in its own task's state: the version's snapshot
1413
+ (already cloned) with the task's rows on top, as a session seed lands
1414
+ them. A task without a seed starts on the base data."""
1415
+ if spec.seed is None:
1416
+ return
1417
+ state.seed_live(spec.seed["rows"], spec.seed["mode"])
1418
+
1419
+ @contextmanager
1420
+ def using(self, session_id: str, *, touch: bool = True) -> Iterator[Session]:
1421
+ """The session for the duration of one request: its lease renewed (any
1422
+ traffic keeps it alive), its ``touched_at`` moved unless ``touch`` is
1423
+ off (a renewal or export observing the session, not using it), and its
1424
+ teardown held off until the block ends."""
1425
+ with self._lock:
1426
+ session = self._sessions.get(session_id)
1427
+ if session is not None:
1428
+ session.renew(self.lease_s)
1429
+ if touch:
1430
+ session.touched_at = time.time()
1431
+ session.inflight += 1
1432
+ if session is None:
1433
+ fail(UNKNOWN_SESSION, f"no session {session_id} on this host")
1434
+ raise AssertionError # fail() never returns
1435
+ try:
1436
+ yield session
1437
+ finally:
1438
+ self._release(session)
1439
+
1440
+ def _release(self, session: Session) -> None:
1441
+ with self._lock:
1442
+ session.inflight -= 1
1443
+ last_out = session.closing and session.inflight == 0
1444
+ if last_out:
1445
+ self._finish(session)
1446
+
1447
+ def owns(self, session_id: str) -> bool:
1448
+ with self._lock:
1449
+ return session_id in self._sessions
1450
+
1451
+ # -- control operations (the platform's rendezvous, run in-process) ----- #
1452
+
1453
+ def reset(self, session_id: str) -> dict:
1454
+ with self.using(session_id) as session:
1455
+ if session.worker is not None:
1456
+ return self._worker_control(session.worker, "POST", f"/__host/sessions/{session_id}/reset", {})
1457
+ if session.tools is not None:
1458
+ return session.tools.reset()
1459
+ assert session.state is not None
1460
+ # Back to the task's starting state: re-clone the snapshot, then
1461
+ # re-apply the task's seed or the scenario overlay if it had one.
1462
+ session.state.reset()
1463
+ if session.task_spec is not None:
1464
+ self._apply_task_seed(session.state, session.task_spec)
1465
+ else:
1466
+ assert session.task is not None
1467
+ self._apply_scenario(self.cache._versions[session.version_id], session.state, session.task)
1468
+ return {"reset": True}
1469
+
1470
+ def seed(self, session_id: str, args: dict) -> dict:
1471
+ rows = args.get("rows")
1472
+ if not isinstance(rows, dict):
1473
+ fail("host", "seed needs args.rows as {entity: [row, ...]}")
1474
+ mode = args.get("mode", "append")
1475
+ with self.using(session_id) as session:
1476
+ if session.worker is not None:
1477
+ return self._worker_control(
1478
+ session.worker, "POST", f"/__host/sessions/{session_id}/seed", {"args": args}
1479
+ )
1480
+ if session.tools is not None:
1481
+ return session.tools.seed(args)
1482
+ assert session.state is not None
1483
+ # Onto the live state, the way the runner's session seed lands: an
1484
+ # append upserts by key, a replace swaps the named entities.
1485
+ session.state.seed_live(rows, mode)
1486
+ return {
1487
+ "seeded": sum(len(v) for v in rows.values()),
1488
+ "mode": mode,
1489
+ "entity_counts": {
1490
+ entity: len(session.state.rows(entity))
1491
+ for entity in rows
1492
+ if entity in session.state.contract.entities
1493
+ },
1494
+ }
1495
+
1496
+ def state(self, session_id: str) -> dict:
1497
+ with self.using(session_id) as session:
1498
+ if session.worker is not None:
1499
+ return self._worker_control(session.worker, "GET", f"/__host/sessions/{session_id}/state", None)
1500
+ if session.tools is not None:
1501
+ return session.tools.state()
1502
+ assert session.state is not None
1503
+ return {**session.state.dump(), "snapshot": session.state.describe_snapshot()}
1504
+
1505
+ def grade(self, session_id: str) -> dict:
1506
+ with self.using(session_id) as session:
1507
+ if session.worker is not None:
1508
+ return self._worker_control(session.worker, "POST", f"/__host/sessions/{session_id}/grade", {})
1509
+ if session.tools is not None:
1510
+ return session.tools.grade()
1511
+ return self._grade(session)
1512
+
1513
+ def _grade(self, session: Session) -> dict:
1514
+ assert session.state is not None
1515
+ state = {**session.state.dump(), "snapshot": session.state.describe_snapshot()}
1516
+ spec = session.task_spec
1517
+ if spec is not None:
1518
+ # The session's own task: its assertions are evaluated here, pure
1519
+ # Python over the dump; its verify.py runs as a bundled one would.
1520
+ if spec.grader is None:
1521
+ return {"reward": None, "rewards": {}}
1522
+ if spec.grader["kind"] == "assertions":
1523
+ return evaluate_assertions(spec.grader["checks"], state)
1524
+ if spec.graded_on_platform:
1525
+ # A rubric is an LLM judge's brief. The host has no model, no
1526
+ # key and no business calling one from a tenant's session: the
1527
+ # platform grades it from the state and the calls log.
1528
+ return {"reward": None, "rewards": {}, "pending": True, "grader": "rubric"}
1529
+ verify = session.state.session_dir / "task" / "checks" / "verify.py"
1530
+ return self._run_verifier(session, verify, state, spec.grader.get("timeoutSeconds") or VERIFY_TIMEOUT_S)
1531
+ assert session.task is not None
1532
+ scenario = _scenario_dir(session.bundle, session.tree, session.task)
1533
+ verify = scenario / "checks" / "verify.py" if scenario else None
1534
+ if verify is None or not verify.is_file():
1535
+ return {"reward": None, "rewards": {}}
1536
+ return self._run_verifier(session, verify, state, VERIFY_TIMEOUT_S)
1537
+
1538
+ def _run_verifier(self, session: Session, verify: Path, state: dict, timeout_s: float) -> dict:
1539
+ assert session.state is not None
1540
+ if not verify.is_file():
1541
+ fail("host", f"the task's verify.py is missing at {verify}")
1542
+ dump = json.dumps(state)
1543
+ # The verifier is the world's code: it runs confined like a worker, with
1544
+ # a scrubbed environment, and is told the state on stdin. It sees no
1545
+ # secret of this process, be it the host's or a worker's.
1546
+ home = session.state.session_dir / "verify"
1547
+ home.mkdir(parents=True, exist_ok=True)
1548
+ try:
1549
+ proc = subprocess.run(
1550
+ confined_command(verify),
1551
+ input=dump,
1552
+ check=False,
1553
+ capture_output=True,
1554
+ text=True,
1555
+ timeout=timeout_s,
1556
+ cwd=str(verify.parent),
1557
+ env=verifier_environment(home=home),
1558
+ )
1559
+ except subprocess.TimeoutExpired:
1560
+ fail("host", f"verify.py timed out after {timeout_s:g}s")
1561
+ if proc.returncode != 0:
1562
+ fail(
1563
+ "host",
1564
+ f"verify.py exited {proc.returncode}: {proc.stderr.strip()[-800:]}",
1565
+ )
1566
+ try:
1567
+ rewards = json.loads(proc.stdout.strip().splitlines()[-1])
1568
+ except (json.JSONDecodeError, IndexError):
1569
+ fail(
1570
+ "host",
1571
+ f"verify.py printed no rewards JSON: {proc.stdout.strip()[-400:]}",
1572
+ )
1573
+ if not isinstance(rewards, dict):
1574
+ fail("host", "verify.py rewards output is not an object")
1575
+ # The runner's shape (``api_runner.grade_scenario``): the scalar under
1576
+ # ``reward``, every other key the verifier printed under ``rewards``.
1577
+ reward = rewards.get("reward")
1578
+ return {
1579
+ "reward": float(reward) if isinstance(reward, (int, float)) else None,
1580
+ "rewards": {
1581
+ k: (float(v) if isinstance(v, (int, float)) else None) for k, v in rewards.items() if k != "reward"
1582
+ },
1583
+ }
1584
+
1585
+ def call(self, session_id: str, tool: str, args: dict) -> dict:
1586
+ """Invoke one of the world's vendor routes by operation name, in process.
1587
+ The arg shape mirrors the runner's ``execute_route``: ``path``/``query``/
1588
+ ``body``, or leftover top-level keys as the body."""
1589
+ with self.using(session_id) as session:
1590
+ if session.worker is not None:
1591
+ return self._worker_control(
1592
+ session.worker,
1593
+ "POST",
1594
+ f"/__host/sessions/{session_id}/call",
1595
+ {"tool": tool, "args": args},
1596
+ )
1597
+ if session.tools is not None:
1598
+ return session.tools.call(tool, args)
1599
+ return self._call(session, tool, args)
1600
+
1601
+ def _call(self, session: Session, tool: str, args: dict) -> dict:
1602
+ assert session.app is not None and session.spec is not None
1603
+ operation = session.spec.operations.get(tool)
1604
+ if operation is None:
1605
+ known = ", ".join(sorted(session.spec.operations)) or "(none)"
1606
+ fail("host", f"unknown tool '{tool}'; this world serves: {known}")
1607
+ if session.task_spec is not None and not session.task_spec.allows_tool(tool):
1608
+ allowed = ", ".join(session.task_spec.tools or ())
1609
+ fail("tool_not_allowed", f"tool '{tool}' is not allowed by this session's task; it allows: {allowed}")
1610
+ path = operation.path # type: ignore[union-attr]
1611
+ path_params = args.get("path") if isinstance(args.get("path"), dict) else {}
1612
+ for key, value in (path_params or {}).items():
1613
+ path = path.replace("{" + str(key) + "}", urllib.parse.quote(str(value), safe=""))
1614
+ query = args.get("query") if isinstance(args.get("query"), dict) else {}
1615
+ if query:
1616
+ path = f"{path}?" + urllib.parse.urlencode(query)
1617
+ body = args.get("body") if isinstance(args.get("body"), dict) else None
1618
+ if body is None:
1619
+ leftovers = {k: v for k, v in args.items() if k not in ("path", "query", "body")}
1620
+ body = leftovers or None
1621
+ method = operation.method # type: ignore[union-attr]
1622
+ raw = json.dumps(body).encode("utf-8") if body is not None and method != "GET" else b""
1623
+ base = urllib.parse.urlsplit(session.spec.base_url).path.rstrip("/")
1624
+ headers, credential_query = world_credential(session)
1625
+ target = with_query(f"{base}{path}", credential_query)
1626
+ status, answer, _ = session.app.handle(method, target, headers, raw, "127.0.0.1")
1627
+ return {"status": status, "body": answer}
1628
+
1629
+
1630
+ # --------------------------------------------------------------------------- #
1631
+ # The HTTP server: control plane by secret, data plane by Host header.
1632
+ # --------------------------------------------------------------------------- #
1633
+
1634
+
1635
+ def public_host(host_header: str | None, worlds_domain: str) -> bool:
1636
+ """True when ``Host`` names something under the worlds domain: a session,
1637
+ or any other label the wildcard ALB forwarded. That is the public side.
1638
+ Prometheus, the kubelet and the ALB health check address the pod itself
1639
+ (``ip:port``), which is not."""
1640
+ name = (host_header or "").split(":")[0].rstrip(".").lower()
1641
+ domain = worlds_domain.rstrip(".").lower()
1642
+ return bool(domain) and (name == domain or name.endswith("." + domain))
1643
+
1644
+
1645
+ def session_from_host(host_header: str | None) -> tuple[str, str] | None:
1646
+ """``w-<sid>-api.<domain>`` / ``w-<sid>.<domain>`` → (session id, surface).
1647
+ Returns None when the host does not name a session."""
1648
+ if not host_header:
1649
+ return None
1650
+ label = host_header.split(":")[0].split(".")[0]
1651
+ if not label.startswith("w-"):
1652
+ return None
1653
+ rest = label[2:]
1654
+ if rest.endswith("-api"):
1655
+ return rest[: -len("-api")], "api"
1656
+ if rest.endswith("-cdp"):
1657
+ return rest[: -len("-cdp")], "browser"
1658
+ return rest, "ui"
1659
+
1660
+
1661
+ #: Auth schemes the host can present a world's own credential for. Everything
1662
+ #: in ``connector.SCHEMES`` is here; the list is explicit so a new scheme is
1663
+ #: refused at open rather than admitted with no credential.
1664
+ SUPPORTED_AUTH_SCHEMES = ("none", "header", "bearer", "basic", "query", "oauth2")
1665
+
1666
+ #: A caller that satisfies the host's session-token gate has already been
1667
+ #: authorized; the host then presents the credential the WORLD itself expects,
1668
+ #: so the emulated vendor's own auth check passes. The session token is that
1669
+ #: credential; with no session token a sentinel satisfies the world's
1670
+ #: non-empty check (the host's gate is the real access control either way).
1671
+ _WORLD_AUTH_SENTINEL = "gateway-session"
1672
+
1673
+
1674
+ def _code_of(cached: CachedVersion) -> str:
1675
+ """What world-authored code a version runs, for a refusal message."""
1676
+ if cached.tools is not None:
1677
+ return "mock-tools handlers"
1678
+ assert cached.settings is not None and cached.settings.handlers is not None
1679
+ return f"api handlers module {cached.settings.handlers.name}"
1680
+
1681
+
1682
+ def unsupported_auth(spec: Spec) -> str | None:
1683
+ if spec.scheme in SUPPORTED_AUTH_SCHEMES:
1684
+ return None
1685
+ return (
1686
+ f"the host cannot present auth.scheme '{spec.scheme}' to a world; it serves {', '.join(SUPPORTED_AUTH_SCHEMES)}"
1687
+ )
1688
+
1689
+
1690
+ def world_credential(session: Session) -> tuple[dict[str, str], dict[str, str]]:
1691
+ """(headers, query) carrying the world's own credential for this session.
1692
+ A ``query`` world reads its key from the query string, so that half is a
1693
+ query parameter; an ``oauth2`` world accepts only a bearer it minted, so
1694
+ the host presents its own grant."""
1695
+ spec = session.spec
1696
+ assert spec is not None, "a mock-tools session has no http surface to present a credential on"
1697
+ presented = session.token or _WORLD_AUTH_SENTINEL
1698
+ if spec.scheme == "none":
1699
+ return {}, {}
1700
+ if spec.scheme == "header":
1701
+ return {(spec.auth_name or "x-api-key").lower(): presented}, {}
1702
+ if spec.scheme == "bearer":
1703
+ return {"authorization": f"Bearer {presented}"}, {}
1704
+ if spec.scheme == "oauth2":
1705
+ return {"authorization": f"Bearer {session.grant or ''}"}, {}
1706
+ if spec.scheme == "basic":
1707
+ import base64 as _b64
1708
+
1709
+ return {"authorization": "Basic " + _b64.b64encode(f"{presented}:".encode()).decode()}, {}
1710
+ if spec.scheme == "query":
1711
+ return {}, {spec.auth_name or "api_key": presented}
1712
+ fail("host", unsupported_auth(spec) or spec.scheme)
1713
+ raise AssertionError # fail() never returns
1714
+
1715
+
1716
+ def with_query(target: str, extra: dict[str, str]) -> str:
1717
+ """``target`` with ``extra`` merged into its query string."""
1718
+ if not extra:
1719
+ return target
1720
+ url = urllib.parse.urlsplit(target)
1721
+ query = {k: v[0] for k, v in urllib.parse.parse_qs(url.query, keep_blank_values=True).items()}
1722
+ query.update(extra)
1723
+ return urllib.parse.urlunsplit(url._replace(query=urllib.parse.urlencode(query)))
1724
+
1725
+
1726
+ def bearer_of(headers: dict[str, str]) -> str | None:
1727
+ auth = headers.get("authorization", "")
1728
+ return auth[7:] if auth.lower().startswith("bearer ") else None
1729
+
1730
+
1731
+ def vendor_shaped_token(spec: Spec, headers: dict[str, str], token: str) -> bool:
1732
+ """Whether the request carries ``token`` the way the world's own auth scheme
1733
+ carries a credential: the declared header for ``header``, the user of a
1734
+ Basic authorization for ``basic``. A ``query`` scheme is not honoured here (a
1735
+ token in the URL outlives the request in logs); ``bearer`` is the gate's
1736
+ own shape and ``oauth2`` has its grant."""
1737
+ if spec.scheme == "header":
1738
+ presented = headers.get((spec.auth_name or "x-api-key").lower())
1739
+ return presented is not None and hmac.compare_digest(presented, token)
1740
+ if spec.scheme == "basic":
1741
+ auth = headers.get("authorization", "")
1742
+ if not auth.lower().startswith("basic "):
1743
+ return False
1744
+ import base64 as _b64
1745
+
1746
+ try:
1747
+ decoded = _b64.b64decode(auth[6:].strip().encode("ascii"), validate=True).decode("utf-8")
1748
+ except (ValueError, UnicodeDecodeError):
1749
+ return False
1750
+ user = decoded.split(":", 1)[0]
1751
+ return hmac.compare_digest(user, token)
1752
+ return False
1753
+
1754
+
1755
+ def admits(session: Session, path: str, headers: dict[str, str]) -> str | None:
1756
+ """How the host's gate admits a data-plane request, or None to refuse.
1757
+
1758
+ ``"session"``: the caller presented the session token as a bearer and the
1759
+ host will present the world's credential for it. Only the header counts: a
1760
+ token in the URL would outlive the request in access logs and referrers.
1761
+ ``"world"``: an oauth2 caller doing the vendor's own dance — the grant
1762
+ request (the world checks ``client_secret`` against the session token) or
1763
+ a bearer this session's world minted — whose headers pass through as sent.
1764
+ """
1765
+ if session.token is None:
1766
+ return "session"
1767
+ bearer = bearer_of(headers)
1768
+ if bearer is not None and hmac.compare_digest(bearer, session.token):
1769
+ return "session"
1770
+ if session.spec is not None and vendor_shaped_token(session.spec, headers, session.token):
1771
+ # The vendor's own wire shape carrying the session token: a client
1772
+ # written against the real vendor (`x-api-key: …`, a Basic user) works
1773
+ # unchanged with the session token as its key. The same secret as the
1774
+ # bearer, in the header the world declares — never in the URL.
1775
+ return "session"
1776
+ if session.spec is not None and session.spec.scheme == "oauth2":
1777
+ if session.app is None:
1778
+ return "world" # the world's tokens live in its worker, which runs this same gate # the world's tokens live in its worker, which runs this same gate
1779
+ token_path = session.app.settings.token_path
1780
+ if token_path and session.app._strip(path) == token_path:
1781
+ return "world"
1782
+ if bearer is not None and session.app.router.authorized(headers, {}, session.token):
1783
+ return "world"
1784
+ return None
1785
+
1786
+
1787
+ class OwnerLookup(Protocol):
1788
+ """Which host owns a session this replica does not: its control-plane
1789
+ URL, or None for a closed or unknown one. ``fresh`` bypasses whatever the
1790
+ lookup remembers, for a caller that has just found the remembered owner
1791
+ dead."""
1792
+
1793
+ def __call__(self, sid: str, *, fresh: bool = False) -> str | None: ...
1794
+
1795
+
1796
+ def platform_owner_lookup(platform_url: str, secret: str | None, self_url: str | None) -> OwnerLookup:
1797
+ """Ask the platform which host owns a session the caller reached us for.
1798
+ ``GET /api/public/world-sessions/{sid}/host`` is host-secret authenticated
1799
+ and answers ``{hostUrl}``; a 404 is a closed or unknown session. Answers are
1800
+ held briefly so a busy session does not cost one platform call per request.
1801
+ Our own address answers None: a session we do not hold but the platform
1802
+ says is ours is gone, not one hop away."""
1803
+ cache: dict[str, tuple[float, str | None]] = {}
1804
+ lock = threading.Lock()
1805
+ base = platform_url.rstrip("/")
1806
+
1807
+ def lookup(sid: str, *, fresh: bool = False) -> str | None:
1808
+ now = time.time()
1809
+ if not fresh:
1810
+ with lock:
1811
+ hit = cache.get(sid)
1812
+ if hit and hit[0] > now:
1813
+ return hit[1]
1814
+ owner: str | None = None
1815
+ req = urllib.request.Request(
1816
+ f"{base}/api/public/world-sessions/{urllib.parse.quote(sid, safe='')}/host",
1817
+ headers={"X-Host-Secret": secret or ""},
1818
+ )
1819
+ try:
1820
+ with urllib.request.urlopen(req, timeout=5) as r:
1821
+ answer = json.loads(r.read() or b"{}")
1822
+ url = answer.get("hostUrl") if isinstance(answer, dict) else None
1823
+ owner = url if isinstance(url, str) and url else None
1824
+ except urllib.error.HTTPError as e:
1825
+ if e.code != 404:
1826
+ raise
1827
+ if owner and self_url and owner.rstrip("/") == self_url.rstrip("/"):
1828
+ owner = None
1829
+ with lock:
1830
+ cache[sid] = (now + (OWNER_CACHE_S if owner else OWNER_MISS_S), owner)
1831
+ return owner
1832
+
1833
+ return lookup
1834
+
1835
+
1836
+ class OwnerUnreachable(Exception):
1837
+ """The owner a request was proxied to did not answer: no connection, or
1838
+ none in time. Not an answer of any status."""
1839
+
1840
+
1841
+ class BodyRefused(Exception):
1842
+ """A request body the host will not read: the answer's status and error."""
1843
+
1844
+ def __init__(self, status: int, error: str):
1845
+ super().__init__(error)
1846
+ self.status = status
1847
+ self.error = error
1848
+
1849
+
1850
+ def make_host_handler(
1851
+ host: WorldHost,
1852
+ secret: str,
1853
+ worlds_domain: str,
1854
+ owner_lookup: OwnerLookup | None = None,
1855
+ *,
1856
+ max_body_mb: int = DEFAULT_MAX_BODY_MB,
1857
+ secure_cookies: bool = True,
1858
+ ):
1859
+ """The request handler class for one host. ``secret`` is required: the
1860
+ control plane is the only trust boundary, so a host with no secret has no
1861
+ control plane at all rather than an open one. ``secure_cookies`` is off
1862
+ only for a host whose public surface is reached over plain http (a dev
1863
+ host); it is the deployment's word, not the request's, since a client
1864
+ writes the forwarding headers a request arrives with."""
1865
+ if not secret:
1866
+ raise ValueError("the world host needs a control-plane secret (WORLDS_HOST_SECRET)")
1867
+ if max_body_mb <= 0:
1868
+ raise ValueError(f"the request body cap must be positive (WORLDS_HOST_MAX_BODY_MB={max_body_mb})")
1869
+ max_body_bytes = max_body_mb * (1 << 20)
1870
+
1871
+ class Handler(BaseHTTPRequestHandler):
1872
+ server_version = "gateway-world-host/1"
1873
+
1874
+ def log_message(self, format: str, *args: object) -> None:
1875
+ pass
1876
+
1877
+ # -- plumbing ------------------------------------------------------- #
1878
+
1879
+ def _read(self) -> bytes:
1880
+ declared = self.headers.get("Content-Length")
1881
+ if declared is None:
1882
+ return b""
1883
+ try:
1884
+ length = int(declared)
1885
+ except ValueError:
1886
+ raise BodyRefused(400, "Content-Length is not a number") from None
1887
+ if length < 0:
1888
+ raise BodyRefused(400, "Content-Length is negative")
1889
+ if length > max_body_bytes:
1890
+ raise BodyRefused(413, f"request body over {max_body_mb} MiB")
1891
+ return self.rfile.read(length) if length else b""
1892
+
1893
+ def _reply(self, status: int, body: Any, extra: dict[str, str] | None = None) -> None:
1894
+ send_response(self, status, body, extra, "application/json")
1895
+
1896
+ def _headers(self) -> dict[str, str]:
1897
+ return {k.lower(): v for k, v in self.headers.items()}
1898
+
1899
+ # -- dispatch ------------------------------------------------------- #
1900
+
1901
+ def do_GET(self):
1902
+ self._dispatch("GET")
1903
+
1904
+ def do_POST(self):
1905
+ self._dispatch("POST")
1906
+
1907
+ def do_PUT(self):
1908
+ self._dispatch("PUT")
1909
+
1910
+ def do_PATCH(self):
1911
+ self._dispatch("PATCH")
1912
+
1913
+ def do_DELETE(self):
1914
+ self._dispatch("DELETE")
1915
+
1916
+ def _dispatch(self, method: str) -> None:
1917
+ url = urllib.parse.urlsplit(self.path)
1918
+ # The public side is whatever arrives under the worlds domain. It
1919
+ # gets sessions and nothing about the host: the probes' answer is
1920
+ # a bare ok, the counters and the rest live on the pod address
1921
+ # (prometheus) or behind the control-plane secret (/__host/stats).
1922
+ if url.path == "/healthz":
1923
+ return self._reply(200, {"ok": True})
1924
+ if url.path == "/metrics" and not public_host(self.headers.get("Host"), worlds_domain):
1925
+ payload = host.metrics_text().encode("utf-8")
1926
+ return self._relay(
1927
+ 200, [("Content-Type", "text/plain; version=0.0.4; charset=utf-8")], payload, drop=FRAMING_HEADERS
1928
+ )
1929
+ try:
1930
+ if url.path.startswith("/__host/"):
1931
+ return self._control(method, url)
1932
+ return self._data(method, url)
1933
+ except BodyRefused as refused:
1934
+ # The body was not read, so the connection cannot carry another
1935
+ # request: answer and hang up.
1936
+ self.close_connection = True
1937
+ return self._reply(refused.status, {"error": refused.error}, {"Connection": "close"})
1938
+
1939
+ # -- control plane -------------------------------------------------- #
1940
+
1941
+ def _control(self, method: str, url) -> None:
1942
+ if not hmac.compare_digest(self.headers.get("X-Host-Secret") or "", secret):
1943
+ return self._reply(403, {"error": "forbidden"})
1944
+ raw = self._read()
1945
+ try:
1946
+ body = json.loads(raw) if raw else {}
1947
+ except ValueError as error:
1948
+ return self._reply(400, {"error": "malformed json", "detail": str(error)})
1949
+ if not isinstance(body, dict):
1950
+ return self._reply(400, {"error": "malformed json", "detail": "the body must be a json object"})
1951
+ parts = url.path.strip("/").split("/") # __host, sessions, [sid], [verb]
1952
+ try:
1953
+ if method == "GET" and parts == ["__host", "stats"]:
1954
+ return self._reply(200, host.stats())
1955
+ if method == "POST" and parts == ["__host", "drain"]:
1956
+ host.drain()
1957
+ return self._reply(200, {"draining": True, "sessions": host.session_count()})
1958
+ if method == "GET" and parts == ["__host", "versions"]:
1959
+ return self._reply(200, {"versions": host.cache.version_ids()})
1960
+ if method == "POST" and parts == ["__host", "versions"]:
1961
+ return self._reply(
1962
+ 200,
1963
+ host.warm(body["versionId"], body["bundleUrl"], body.get("snapshot")),
1964
+ )
1965
+ if method == "POST" and parts == ["__host", "sessions"]:
1966
+ return self._reply(
1967
+ 200,
1968
+ host.open(
1969
+ session_id=body["sessionId"],
1970
+ version_id=body["versionId"],
1971
+ bundle_url=body["bundleUrl"],
1972
+ task=body.get("task"),
1973
+ token=body.get("token"),
1974
+ snapshot=body.get("snapshot"),
1975
+ browser=body.get("browser"),
1976
+ restore=body.get("restore"),
1977
+ task_spec=body.get("taskSpec"),
1978
+ ),
1979
+ )
1980
+ if len(parts) >= 3 and parts[:2] == ["__host", "sessions"]:
1981
+ sid = parts[2]
1982
+ verb = parts[3] if len(parts) > 3 else None
1983
+ if method == "DELETE" and verb is None:
1984
+ return self._reply(200, host.close(sid))
1985
+ if method == "POST" and verb == "renew":
1986
+ return self._reply(200, host.renew(sid))
1987
+ if method == "POST" and verb == "reset":
1988
+ return self._reply(200, host.reset(sid))
1989
+ if method == "POST" and verb == "seed":
1990
+ return self._reply(200, host.seed(sid, body.get("args") or {}))
1991
+ if method == "POST" and verb == "export" and len(parts) == 4:
1992
+ known = body.get("knownSha256")
1993
+ return self._reply(200, host.export(sid, known if isinstance(known, str) else None))
1994
+ if method == "POST" and verb == "export" and len(parts) == 5:
1995
+ url = body.get("url")
1996
+ if not isinstance(url, str) or not url.startswith(("https://", "http://")):
1997
+ return self._reply(400, {"error": "url must be the presigned upload url"})
1998
+ return self._reply(200, host.upload_export(sid, parts[4], url))
1999
+ if method == "GET" and verb == "state":
2000
+ return self._reply(200, host.state(sid))
2001
+ if method == "POST" and verb == "grade":
2002
+ return self._reply(200, host.grade(sid))
2003
+ if method == "POST" and verb == "call":
2004
+ tool = body.get("tool")
2005
+ if not isinstance(tool, str):
2006
+ return self._reply(400, {"error": "tool must be a string"})
2007
+ return self._reply(200, host.call(sid, tool, body.get("args") or {}))
2008
+ return self._reply(404, {"error": "unknown control route"})
2009
+ except ContractError as error:
2010
+ code = error.detail.get("code", "host")
2011
+ # A host that will not take the session is not a bad request:
2012
+ # the platform places it on another replica. A session this
2013
+ # host does not hold is the one answer the platform may read
2014
+ # as "lost"; every other refusal is about the request.
2015
+ if code in ("draining", "at_capacity"):
2016
+ status = 503
2017
+ elif code == UNKNOWN_SESSION:
2018
+ status = 404
2019
+ else:
2020
+ status = 400
2021
+ return self._reply(status, {"error": code, "detail": str(error)})
2022
+ except KeyError as error:
2023
+ return self._reply(400, {"error": "missing field", "field": str(error)})
2024
+ except WorkerGone as error:
2025
+ return self._reply(502, {"error": "worker exited", "detail": str(error)})
2026
+
2027
+ # -- data plane ----------------------------------------------------- #
2028
+
2029
+ def _data(self, method: str, url) -> None:
2030
+ headers = self._headers()
2031
+ named = session_from_host(headers.get("host"))
2032
+ if named is None:
2033
+ return self._reply(404, {"error": "no session named by Host"})
2034
+ sid, surface = named
2035
+ if not host.owns(sid):
2036
+ # Another host owns it (or it is closed). Proxy one hop when the
2037
+ # platform can point us at the owner; a request that already
2038
+ # hopped once is not sent on again.
2039
+ hopped = HOP_HEADER in headers
2040
+ owner = owner_lookup(sid) if owner_lookup and not hopped else None
2041
+ if owner and owner_lookup:
2042
+ host.tally("requests.proxied")
2043
+ return self._proxy(sid, owner, owner_lookup, method, headers)
2044
+ return self._reply(404, {"error": "unknown or closed session"})
2045
+ host.tally(f"requests.{surface}")
2046
+ try:
2047
+ with host.using(sid) as session:
2048
+ if session.spec is None:
2049
+ return self._reply(
2050
+ 404, {"error": "this world has no http surface; reach its tools through call"}
2051
+ )
2052
+ if surface == "browser":
2053
+ self._serve_browser(session, url, headers)
2054
+ elif surface == "ui":
2055
+ self._serve_ui(session, method, url, headers)
2056
+ else:
2057
+ self._serve(session, method, url, headers)
2058
+ except ContractError:
2059
+ # Closed between the ownership check and the lookup.
2060
+ self._reply(404, {"error": "unknown or closed session"})
2061
+
2062
+ def _serve(self, session: Session, method: str, url, headers: dict[str, str]) -> None:
2063
+ admitted = admits(session, url.path, headers)
2064
+ if admitted is None:
2065
+ return self._reply(401, {"error": "invalid session token"})
2066
+ self._serve_api(session, method, self.path, headers, admitted)
2067
+
2068
+ def _serve_api(
2069
+ self, session: Session, method: str, target: str, headers: dict[str, str], admitted: str
2070
+ ) -> None:
2071
+ """Answer ``target`` (a path plus query) from the session's world,
2072
+ in this process or its worker, once the gate has admitted it."""
2073
+ sid = session.session_id
2074
+ raw = self._read()
2075
+ if session.worker is not None:
2076
+ # The gate passed here; the worker runs it again and presents the
2077
+ # world its credential, so the request goes over as it came.
2078
+ try:
2079
+ status, answer_headers, payload = session.worker.request(method, target, headers, raw)
2080
+ except WorkerGone as error:
2081
+ host.lose_worker(session.worker)
2082
+ return self._reply(502, {"error": "worker exited", "detail": str(error)})
2083
+ return self._relay(status, answer_headers, payload, drop=FRAMING_HEADERS)
2084
+ assert session.app is not None
2085
+ world_headers = headers
2086
+ if admitted == "session":
2087
+ # The host's token gate is the access control; present the world
2088
+ # its OWN expected credential so the emulated vendor auth passes.
2089
+ credential_headers, credential_query = world_credential(session)
2090
+ world_headers = {**headers, **credential_headers}
2091
+ target = with_query(target, credential_query)
2092
+ status, body, extra = session.app.handle(method, target, world_headers, raw, f"w-{sid}-api.{worlds_domain}")
2093
+ self._reply(status, body, extra)
2094
+
2095
+ # -- the ui surface -------------------------------------------------- #
2096
+
2097
+ def _serve_ui(self, session: Session, method: str, url, headers: dict[str, str]) -> None:
2098
+ ui = session.ui
2099
+ if ui is None:
2100
+ return self._reply(404, {"error": "this world has no static ui"})
2101
+ token = session.token
2102
+ if token is not None:
2103
+ if matches(query_token(url.query), token):
2104
+ # The way in: trade the one-time ``?token=`` for the cookie
2105
+ # and land on the same page without it.
2106
+ landing = without_query_token(url)
2107
+ return self._reply(
2108
+ 303,
2109
+ {"location": landing},
2110
+ {
2111
+ "Location": landing,
2112
+ "Set-Cookie": session_cookie(token, secure=secure_cookies, max_age_s=host.lease_s),
2113
+ },
2114
+ )
2115
+ if not matches(cookie_token(headers), token) and not matches(bearer_of(headers), token):
2116
+ return self._reply(401, {"error": "invalid session token"})
2117
+ route = api_target(ui, url.path)
2118
+ if route is not None:
2119
+ # The SPA's same-origin api call: strip the prefix and answer it
2120
+ # as the api surface would for a caller holding the session
2121
+ # token (the cookie stood in for the bearer here).
2122
+ target = urllib.parse.urlunsplit(("", "", route, url.query, ""))
2123
+ api_headers = {k: v for k, v in headers.items() if k not in ("cookie", "authorization")}
2124
+ if token is not None:
2125
+ api_headers["authorization"] = f"Bearer {token}"
2126
+ return self._serve_api(session, method, target, api_headers, "session")
2127
+ if method != "GET":
2128
+ return self._reply(405, {"error": "the static ui answers GET only"})
2129
+ status, static_headers, payload = static_file(ui, url.path)
2130
+ self._relay(status, static_headers.items(), payload, drop=FRAMING_HEADERS)
2131
+
2132
+ # -- the browser surface --------------------------------------------- #
2133
+
2134
+ def _serve_browser(self, session: Session, url, headers: dict[str, str]) -> None:
2135
+ """``/cdp/<session token>`` and ``/live/<grant>`` on the session's
2136
+ cdp host: check the credential against the SESSION, then tunnel
2137
+ the websocket to the leased relay with the POOL credential the
2138
+ relay knows. The relay's token never leaves the host."""
2139
+ if session.browser_url is None or session.browser_token is None:
2140
+ return self._reply(404, {"error": "this session has no browser surface"})
2141
+ parts = url.path.strip("/").split("/")
2142
+ if len(parts) != 2 or parts[0] not in ("cdp", "live"):
2143
+ return self._reply(404, {"error": "browser routes are /cdp/<token> and /live/<grant>"})
2144
+ kind, credential = parts
2145
+ token = session.token or ""
2146
+ if kind == "cdp":
2147
+ if not matches(credential, token):
2148
+ return self._reply(401, {"error": "invalid session token"})
2149
+ # The relay decodes the path segment, so a token with '/' or
2150
+ # '=' (openssl rand -base64) must travel percent-encoded.
2151
+ upstream_path = f"/cdp/{urllib.parse.quote(session.browser_token, safe='')}"
2152
+ else:
2153
+ scope = verify_live_grant(credential, token)
2154
+ if scope is None:
2155
+ return self._reply(401, {"error": "invalid or expired live grant"})
2156
+ upstream_path = f"/live/{mint_live_grant(session.browser_token, LIVE_RELAY_GRANT_MS, scope)}"
2157
+ if headers.get("upgrade", "").lower() != "websocket":
2158
+ return self._reply(400, {"error": "the browser surface is a websocket"})
2159
+ self.close_connection = True
2160
+ try:
2161
+ tunnel(self.connection, session.browser_url, upstream_path, headers)
2162
+ except OSError as error:
2163
+ self._reply(502, {"error": "browser unreachable", "detail": str(error)})
2164
+
2165
+ def _proxy(self, sid: str, owner: str, lookup: OwnerLookup, method: str, headers: dict[str, str]) -> None:
2166
+ """Send the request on to ``owner``. An owner that does not answer
2167
+ at all may be a host that died with the session since recovered
2168
+ elsewhere, so the platform is asked once more, uncached, and the
2169
+ request goes to the owner it names now; only then is it a 502.
2170
+ The retry stays one hop (the hop header goes with it)."""
2171
+ websocket = headers.get("upgrade", "").lower() == "websocket"
2172
+ raw = b"" if websocket else self._read()
2173
+ try:
2174
+ return self._forward(owner, method, headers, raw, websocket)
2175
+ except OwnerUnreachable as stale:
2176
+ current = lookup(sid, fresh=True)
2177
+ if not current or current.rstrip("/") == owner.rstrip("/"):
2178
+ return self._reply(502, {"error": "owner unreachable", "detail": str(stale)})
2179
+ host.tally("requests.rerouted")
2180
+ try:
2181
+ self._forward(current, method, headers, raw, websocket)
2182
+ except OwnerUnreachable as error:
2183
+ self._reply(502, {"error": "owner unreachable", "detail": str(error)})
2184
+
2185
+ def _forward(self, owner: str, method: str, headers: dict[str, str], raw: bytes, websocket: bool) -> None:
2186
+ """One attempt at ``owner``; raises :class:`OwnerUnreachable` when it
2187
+ did not answer, before anything has been written to the client."""
2188
+ if websocket:
2189
+ # A browser session on another host: the owner runs the
2190
+ # credential check; this hop only carries bytes. ``tunnel``
2191
+ # fails before the 101 or not at all.
2192
+ owner_ws = "ws://" + owner.split("://", 1)[-1].rstrip("/")
2193
+ self.close_connection = True
2194
+ try:
2195
+ tunnel(self.connection, owner_ws, self.path, {**headers, HOP_HEADER: "1"}, host=headers.get("host"))
2196
+ except OSError as error:
2197
+ raise OwnerUnreachable(str(error)) from error
2198
+ return
2199
+ target = f"{owner.rstrip('/')}{self.path}"
2200
+ forward = {k: v for k, v in headers.items() if k not in ("host", "content-length")}
2201
+ forward["host"] = headers.get("host", "")
2202
+ forward[HOP_HEADER] = "1"
2203
+ req = urllib.request.Request(target, data=raw or None, method=method, headers=forward)
2204
+ try:
2205
+ with _PROXY_OPENER.open(req, timeout=60) as r:
2206
+ self._relay(r.status, r.headers.items(), r.read(), drop=PROXIED_DROP_HEADERS)
2207
+ except urllib.error.HTTPError as e:
2208
+ self._relay(e.code, e.headers.items(), e.read(), drop=PROXIED_DROP_HEADERS)
2209
+ except (OSError, http.client.HTTPException) as error:
2210
+ # URLError is an OSError: no connection, or none in time.
2211
+ raise OwnerUnreachable(str(error)) from error
2212
+
2213
+ def _relay(
2214
+ self,
2215
+ status: int,
2216
+ headers: Iterable[tuple[str, str]],
2217
+ payload: bytes,
2218
+ *,
2219
+ drop: frozenset[str],
2220
+ ) -> None:
2221
+ """Write another process's answer through as it came: same status,
2222
+ same body bytes, every header the world set except ``drop`` (the
2223
+ framing, which is this connection's, and from an owner the
2224
+ hop-by-hop headers too). The world's contract, its ``page_headers``
2225
+ included, reads the same from every replica."""
2226
+ self.send_response(status)
2227
+ content_type = None
2228
+ for name, value in headers:
2229
+ lower = name.lower()
2230
+ if lower == "content-type":
2231
+ content_type = value
2232
+ elif lower not in drop:
2233
+ self.send_header(name, value)
2234
+ self.send_header("Content-Type", content_type or "application/json")
2235
+ self.send_header("Content-Length", str(len(payload)))
2236
+ self.end_headers()
2237
+ self.wfile.write(payload)
2238
+
2239
+ return Handler
2240
+
2241
+
2242
+ def start_lease_sweeper(world_host: WorldHost, interval_s: float = LEASE_SWEEP_S) -> threading.Event:
2243
+ """Expire leases every ``interval_s`` until the returned event is set."""
2244
+ stop = threading.Event()
2245
+
2246
+ def sweep() -> None:
2247
+ while not stop.wait(interval_s):
2248
+ for sid in world_host.expire_leases():
2249
+ print(
2250
+ f"world host: lease expired, closed session {sid}",
2251
+ file=sys.stderr,
2252
+ flush=True,
2253
+ )
2254
+ for version_id in world_host.reap_workers():
2255
+ print(
2256
+ f"world host: stopped idle worker for version {version_id}",
2257
+ file=sys.stderr,
2258
+ flush=True,
2259
+ )
2260
+
2261
+ threading.Thread(target=sweep, name="world-host-leases", daemon=True).start()
2262
+ return stop
2263
+
2264
+
2265
+ def heartbeat_once(world_host: WorldHost, platform_url: str, secret: str, timeout_s: float = 5.0) -> bool:
2266
+ """Tell the platform what this host is and holds (:meth:`WorldHost.stats`).
2267
+ The platform's registry is what places new sessions: a host it has not
2268
+ heard from lately gets none. Returns whether the platform took it."""
2269
+ body = json.dumps(world_host.stats()).encode("utf-8")
2270
+ req = urllib.request.Request(
2271
+ f"{platform_url.rstrip('/')}/api/public/world-hosts",
2272
+ data=body,
2273
+ method="POST",
2274
+ headers={"content-type": "application/json", "x-host-secret": secret},
2275
+ )
2276
+ try:
2277
+ with urllib.request.urlopen(req, timeout=timeout_s) as r:
2278
+ return 200 <= r.status < 300
2279
+ except (urllib.error.URLError, OSError, ValueError) as error:
2280
+ print(f"world host: heartbeat failed: {error}", file=sys.stderr, flush=True)
2281
+ return False
2282
+
2283
+
2284
+ def start_heartbeat(
2285
+ world_host: WorldHost, platform_url: str, secret: str, interval_s: float = HEARTBEAT_S
2286
+ ) -> threading.Event:
2287
+ """Heartbeat every ``interval_s`` until the returned event is set; the
2288
+ first beat goes out at once so a fresh replica is placeable immediately,
2289
+ and a last one goes out on stop so the registry sees it draining."""
2290
+ stop = threading.Event()
2291
+
2292
+ def beat() -> None:
2293
+ heartbeat_once(world_host, platform_url, secret)
2294
+ while not stop.wait(interval_s):
2295
+ heartbeat_once(world_host, platform_url, secret)
2296
+ heartbeat_once(world_host, platform_url, secret)
2297
+
2298
+ threading.Thread(target=beat, name="world-host-heartbeat", daemon=True).start()
2299
+ return stop
2300
+
2301
+
2302
+ class HostServer(ThreadingHTTPServer):
2303
+ """The listening socket plus what shutdown has to stop: the lease sweeper,
2304
+ the heartbeat and the version workers."""
2305
+
2306
+ def __init__(
2307
+ self,
2308
+ address: tuple[str, int],
2309
+ handler: type[BaseHTTPRequestHandler],
2310
+ world_host: WorldHost,
2311
+ *,
2312
+ heartbeat: tuple[str, str] | None = None,
2313
+ heartbeat_s: float = HEARTBEAT_S,
2314
+ drain_s: float = DEFAULT_DRAIN_S,
2315
+ ):
2316
+ super().__init__(address, handler)
2317
+ self.world_host = world_host
2318
+ self.lease_sweeper = start_lease_sweeper(world_host)
2319
+ #: ``(platform url, secret)`` to report to; None runs unregistered (one
2320
+ #: replica behind WORLDS_HOST_URL needs no registry).
2321
+ self.heartbeat = start_heartbeat(world_host, *heartbeat, interval_s=heartbeat_s) if heartbeat else None
2322
+ self.drain_s = drain_s
2323
+
2324
+ def drain(self) -> None:
2325
+ """SIGTERM: refuse new sessions, announce it, serve what is here until
2326
+ it closes or the grace runs out, then stop. Runs off the signal thread
2327
+ so ``serve_forever`` keeps answering meanwhile."""
2328
+ self.world_host.drain()
2329
+ if self.heartbeat is not None:
2330
+ self.heartbeat.set() # its last beat carries ``draining: true``
2331
+
2332
+ def wait_then_stop() -> None:
2333
+ deadline = time.time() + self.drain_s
2334
+ while self.world_host.session_count() > 0 and time.time() < deadline:
2335
+ time.sleep(0.5)
2336
+ self.shutdown()
2337
+
2338
+ threading.Thread(target=wait_then_stop, name="world-host-drain", daemon=True).start()
2339
+
2340
+
2341
+ def serve(cache_root: str | Path, *, port: int = 8080, host_addr: str = "0.0.0.0") -> HostServer:
2342
+ """The host as deployed, configured from the environment:
2343
+
2344
+ ``WORLDS_HOST_SECRET`` gates the control plane; ``WORLDS_INGRESS_DOMAIN`` is
2345
+ the domain the session hostnames live under; ``WORLDS_HOST_ADVERTISE_URL`` is
2346
+ this replica's own address (``http://<pod ip>:<port>``), returned on open so
2347
+ the platform pins the session to it; ``WORLDS_PLATFORM_URL`` lets a
2348
+ non-owner replica ask the platform who owns a session and proxy one hop;
2349
+ ``WORLDS_HOST_LEASE_S`` is how long an untouched session lives;
2350
+ ``WORLDS_HOST_MAX_SESSIONS`` is how many it takes before answering 503;
2351
+ ``WORLDS_HOST_HEARTBEAT_S`` is how often it reports to the platform (with
2352
+ ``WORLDS_PLATFORM_URL`` and an advertised address, else no heartbeat);
2353
+ ``WORLDS_HOST_DRAIN_S`` is how long a SIGTERM waits for sessions to close;
2354
+ ``WORLDS_WORKER_MEMORY_MB`` caps each version worker's address space and
2355
+ ``WORLDS_WORKER_BUDGET_MB`` all of them together (their count is the
2356
+ quotient; one more version answers 503 ``at_capacity``),
2357
+ ``WORLDS_HOST_OBJECT_STORE`` is where every presigned URL it is handed
2358
+ may point (required; ``any`` for a dev host),
2359
+ ``WORLDS_HOST_MAX_BUNDLE_MB`` is the largest version bundle the cache
2360
+ accepts, ``WORLDS_HOST_MAX_BODY_MB`` the largest request body it reads,
2361
+ ``WORLDS_HOST_PLAIN_HTTP=1`` says the public surface is reached over plain
2362
+ http (a dev host: the ui cookie is not marked ``Secure``) and
2363
+ ``WORLDS_HOST_WORKERS=0`` runs without workers (handler worlds refused)."""
2364
+ secret = os.environ.get("WORLDS_HOST_SECRET") or ""
2365
+ if not secret:
2366
+ raise SystemExit("world host: WORLDS_HOST_SECRET is unset; refusing to start without a control-plane secret")
2367
+ try:
2368
+ object_store = parse_origins(os.environ.get("WORLDS_HOST_OBJECT_STORE"))
2369
+ except ContractError as error:
2370
+ raise SystemExit(f"world host: {error}; refusing to start") from error
2371
+ advertise = os.environ.get("WORLDS_HOST_ADVERTISE_URL") or None
2372
+ platform_url = os.environ.get("WORLDS_PLATFORM_URL") or None
2373
+ lease_s = float(os.environ.get("WORLDS_HOST_LEASE_S") or DEFAULT_LEASE_S)
2374
+ worlds_domain = os.environ.get("WORLDS_INGRESS_DOMAIN", "worlds.local")
2375
+ workers = None
2376
+ if os.environ.get("WORLDS_HOST_WORKERS", "1") != "0":
2377
+ workers = WorkerPool(
2378
+ Path(cache_root),
2379
+ worlds_domain=worlds_domain,
2380
+ memory_mb=int(os.environ.get("WORLDS_WORKER_MEMORY_MB") or DEFAULT_WORKER_MEMORY_MB),
2381
+ budget_mb=int(os.environ.get("WORLDS_WORKER_BUDGET_MB") or DEFAULT_WORKER_BUDGET_MB),
2382
+ )
2383
+ world_host = WorldHost(
2384
+ Path(cache_root),
2385
+ lease_s=lease_s,
2386
+ advertise_url=advertise,
2387
+ workers=workers,
2388
+ max_sessions=int(os.environ.get("WORLDS_HOST_MAX_SESSIONS") or DEFAULT_MAX_SESSIONS),
2389
+ max_bundle_mb=int(os.environ.get("WORLDS_HOST_MAX_BUNDLE_MB") or DEFAULT_MAX_BUNDLE_MB),
2390
+ object_store=object_store,
2391
+ )
2392
+ handler = make_host_handler(
2393
+ world_host,
2394
+ secret=secret,
2395
+ worlds_domain=worlds_domain,
2396
+ owner_lookup=platform_owner_lookup(platform_url, secret, advertise) if platform_url else None,
2397
+ max_body_mb=int(os.environ.get("WORLDS_HOST_MAX_BODY_MB") or DEFAULT_MAX_BODY_MB),
2398
+ secure_cookies=os.environ.get("WORLDS_HOST_PLAIN_HTTP", "0") != "1",
2399
+ )
2400
+ return HostServer(
2401
+ (host_addr, port),
2402
+ handler,
2403
+ world_host,
2404
+ heartbeat=(platform_url, secret) if platform_url and advertise else None,
2405
+ heartbeat_s=float(os.environ.get("WORLDS_HOST_HEARTBEAT_S") or HEARTBEAT_S),
2406
+ drain_s=float(os.environ.get("WORLDS_HOST_DRAIN_S") or DEFAULT_DRAIN_S),
2407
+ )
2408
+
2409
+
2410
+ def main(argv: list[str] | None = None) -> int:
2411
+ import argparse
2412
+
2413
+ parser = argparse.ArgumentParser(prog="gateway-world-host")
2414
+ parser.add_argument("--cache-root", default=os.environ.get("WORLDS_HOST_CACHE", "/cache"))
2415
+ parser.add_argument("--port", type=int, default=int(os.environ.get("PORT") or 8080))
2416
+ args = parser.parse_args(argv)
2417
+ server = serve(args.cache_root, port=args.port)
2418
+ print(
2419
+ f"world host listening on :{args.port}, cache {args.cache_root}",
2420
+ file=sys.stderr,
2421
+ flush=True,
2422
+ )
2423
+ signal.signal(signal.SIGTERM, lambda *_: server.drain())
2424
+ try:
2425
+ server.serve_forever()
2426
+ finally:
2427
+ server.lease_sweeper.set()
2428
+ if server.heartbeat is not None:
2429
+ server.heartbeat.set()
2430
+ if server.world_host.workers is not None:
2431
+ server.world_host.workers.stop_all()
2432
+ server.server_close()
2433
+ return 0
2434
+
2435
+
2436
+ if __name__ == "__main__":
2437
+ raise SystemExit(main())