world-model-optimizer 0.2.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (308) hide show
  1. llm_waterfall/LICENSE +21 -0
  2. llm_waterfall/__init__.py +53 -0
  3. llm_waterfall/adapters/__init__.py +36 -0
  4. llm_waterfall/adapters/anthropic.py +105 -0
  5. llm_waterfall/adapters/aws_mantle.py +47 -0
  6. llm_waterfall/adapters/azure_openai.py +71 -0
  7. llm_waterfall/adapters/base.py +51 -0
  8. llm_waterfall/adapters/bedrock.py +309 -0
  9. llm_waterfall/adapters/openai.py +130 -0
  10. llm_waterfall/classify.py +184 -0
  11. llm_waterfall/pricing.py +110 -0
  12. llm_waterfall/py.typed +0 -0
  13. llm_waterfall/types.py +295 -0
  14. llm_waterfall/waterfall.py +255 -0
  15. wmo/__init__.py +38 -0
  16. wmo/agents/__init__.py +7 -0
  17. wmo/agents/default.py +29 -0
  18. wmo/agents/meta.py +55 -0
  19. wmo/agents/optimizer.py +55 -0
  20. wmo/agents/project.py +928 -0
  21. wmo/cli/__init__.py +5 -0
  22. wmo/cli/agent_session.py +1123 -0
  23. wmo/cli/app.py +2489 -0
  24. wmo/cli/e2b_cmds.py +212 -0
  25. wmo/cli/eval_closed_loop.py +207 -0
  26. wmo/cli/harness_app.py +1147 -0
  27. wmo/cli/harness_distill.py +659 -0
  28. wmo/cli/hosted_session.py +880 -0
  29. wmo/cli/ingest_cmd.py +165 -0
  30. wmo/cli/model_roles.py +82 -0
  31. wmo/cli/platform_cmds.py +372 -0
  32. wmo/cli/route_app.py +274 -0
  33. wmo/cli/session_state.py +243 -0
  34. wmo/cli/ui.py +1107 -0
  35. wmo/cli/workspace_sync.py +504 -0
  36. wmo/config/__init__.py +60 -0
  37. wmo/config/card.py +129 -0
  38. wmo/config/config.py +367 -0
  39. wmo/config/dotenv.py +67 -0
  40. wmo/config/settings.py +128 -0
  41. wmo/config/store.py +177 -0
  42. wmo/conftest.py +19 -0
  43. wmo/connect/__init__.py +88 -0
  44. wmo/connect/apps.py +78 -0
  45. wmo/connect/brave.py +284 -0
  46. wmo/connect/connector.py +79 -0
  47. wmo/connect/credentials.py +164 -0
  48. wmo/connect/github.py +321 -0
  49. wmo/connect/google.py +627 -0
  50. wmo/connect/notion.py +790 -0
  51. wmo/connect/oauth.py +461 -0
  52. wmo/connect/slack.py +555 -0
  53. wmo/connect/store.py +199 -0
  54. wmo/connect/types.py +156 -0
  55. wmo/core/__init__.py +21 -0
  56. wmo/core/parsing.py +281 -0
  57. wmo/core/render.py +271 -0
  58. wmo/core/text.py +40 -0
  59. wmo/core/types.py +116 -0
  60. wmo/distill/__init__.py +14 -0
  61. wmo/distill/agents.py +140 -0
  62. wmo/distill/config.py +1006 -0
  63. wmo/distill/cost.py +437 -0
  64. wmo/distill/data.py +921 -0
  65. wmo/distill/deadlines.py +254 -0
  66. wmo/distill/fake_tinker.py +734 -0
  67. wmo/distill/gate.py +122 -0
  68. wmo/distill/loop.py +3499 -0
  69. wmo/distill/renderers.py +399 -0
  70. wmo/distill/rendering.py +620 -0
  71. wmo/distill/rollouts.py +726 -0
  72. wmo/distill/samples.py +195 -0
  73. wmo/distill/store.py +829 -0
  74. wmo/distill/teacher.py +714 -0
  75. wmo/distill/tokens.py +535 -0
  76. wmo/distill/tracking.py +552 -0
  77. wmo/distill/tripwire.py +411 -0
  78. wmo/distill/xtoken/byte_offsets.py +152 -0
  79. wmo/distill/xtoken/chunks.py +457 -0
  80. wmo/distill/xtoken/prompt_logprobs.py +475 -0
  81. wmo/distill/xtoken/teacher_render.py +346 -0
  82. wmo/engine/__init__.py +28 -0
  83. wmo/engine/autoconfig.py +367 -0
  84. wmo/engine/build.py +346 -0
  85. wmo/engine/demo.py +77 -0
  86. wmo/engine/eval_suites.py +245 -0
  87. wmo/engine/grounding.py +491 -0
  88. wmo/engine/knowledge.py +291 -0
  89. wmo/engine/loader.py +36 -0
  90. wmo/engine/play.py +92 -0
  91. wmo/engine/prompts.py +99 -0
  92. wmo/engine/replay.py +443 -0
  93. wmo/engine/reporting.py +58 -0
  94. wmo/engine/workspace.py +468 -0
  95. wmo/engine/world_model.py +568 -0
  96. wmo/env/__init__.py +22 -0
  97. wmo/env/base.py +121 -0
  98. wmo/env/closed_loop.py +229 -0
  99. wmo/env/episode.py +107 -0
  100. wmo/env/llm_agent.py +93 -0
  101. wmo/env/scenarios.py +73 -0
  102. wmo/evals/__init__.py +52 -0
  103. wmo/evals/agreement.py +110 -0
  104. wmo/evals/base.py +45 -0
  105. wmo/evals/closed_loop.py +480 -0
  106. wmo/evals/failover.py +96 -0
  107. wmo/evals/gold.py +127 -0
  108. wmo/evals/grid.py +394 -0
  109. wmo/evals/grid_plot.py +205 -0
  110. wmo/evals/harbor/__init__.py +27 -0
  111. wmo/evals/harbor/agent.py +573 -0
  112. wmo/evals/harbor/ctrf.py +171 -0
  113. wmo/evals/harbor/e2b_environment.py +587 -0
  114. wmo/evals/harbor/e2b_template_policy.py +144 -0
  115. wmo/evals/harbor/scorer.py +875 -0
  116. wmo/evals/harbor/tasks.py +140 -0
  117. wmo/evals/open_loop.py +194 -0
  118. wmo/evals/tasks.py +53 -0
  119. wmo/harness/__init__.py +51 -0
  120. wmo/harness/code_runtime.py +288 -0
  121. wmo/harness/create.py +1191 -0
  122. wmo/harness/delta.py +220 -0
  123. wmo/harness/doc.py +556 -0
  124. wmo/harness/e2b_ledger.py +342 -0
  125. wmo/harness/e2b_reap.py +476 -0
  126. wmo/harness/e2b_sandbox.py +350 -0
  127. wmo/harness/environment.py +35 -0
  128. wmo/harness/live_session.py +543 -0
  129. wmo/harness/mutate.py +343 -0
  130. wmo/harness/pi_e2b.py +1710 -0
  131. wmo/harness/pi_entry/entry.ts +268 -0
  132. wmo/harness/pi_entry/runner_frames.ts +92 -0
  133. wmo/harness/pi_entry/runner_live.ts +587 -0
  134. wmo/harness/pi_entry/runner_service.ts +270 -0
  135. wmo/harness/pi_entry/runner_stdio.ts +374 -0
  136. wmo/harness/pi_entry/runner_termination.ts +142 -0
  137. wmo/harness/pi_local.py +262 -0
  138. wmo/harness/pi_runtime.py +495 -0
  139. wmo/harness/pi_vendor.py +65 -0
  140. wmo/harness/population.py +509 -0
  141. wmo/harness/project_proposer.py +569 -0
  142. wmo/harness/proposer.py +977 -0
  143. wmo/harness/runner_link.py +619 -0
  144. wmo/harness/runtime.py +389 -0
  145. wmo/harness/scoring.py +247 -0
  146. wmo/harness/skills.py +116 -0
  147. wmo/harness/source_tree.py +319 -0
  148. wmo/harness/store.py +176 -0
  149. wmo/harness/tools.py +105 -0
  150. wmo/harness/vendor/manifest.sha256 +58 -0
  151. wmo/harness/vendor/pi-agent/CHANGELOG.md +556 -0
  152. wmo/harness/vendor/pi-agent/LICENSE +21 -0
  153. wmo/harness/vendor/pi-agent/README.md +488 -0
  154. wmo/harness/vendor/pi-agent/VENDOR.md +39 -0
  155. wmo/harness/vendor/pi-agent/docs/agent-harness.md +486 -0
  156. wmo/harness/vendor/pi-agent/docs/durable-harness.md +212 -0
  157. wmo/harness/vendor/pi-agent/docs/hooks.md +445 -0
  158. wmo/harness/vendor/pi-agent/docs/models.md +966 -0
  159. wmo/harness/vendor/pi-agent/docs/observability.md +376 -0
  160. wmo/harness/vendor/pi-agent/package.json +60 -0
  161. wmo/harness/vendor/pi-agent/src/agent-loop.ts +748 -0
  162. wmo/harness/vendor/pi-agent/src/agent.ts +575 -0
  163. wmo/harness/vendor/pi-agent/src/harness/agent-harness.ts +1029 -0
  164. wmo/harness/vendor/pi-agent/src/harness/compaction/branch-summarization.ts +261 -0
  165. wmo/harness/vendor/pi-agent/src/harness/compaction/compaction.ts +747 -0
  166. wmo/harness/vendor/pi-agent/src/harness/compaction/utils.ts +144 -0
  167. wmo/harness/vendor/pi-agent/src/harness/env/nodejs.ts +550 -0
  168. wmo/harness/vendor/pi-agent/src/harness/messages.ts +164 -0
  169. wmo/harness/vendor/pi-agent/src/harness/prompt-templates.ts +267 -0
  170. wmo/harness/vendor/pi-agent/src/harness/session/jsonl-repo.ts +177 -0
  171. wmo/harness/vendor/pi-agent/src/harness/session/jsonl-storage.ts +293 -0
  172. wmo/harness/vendor/pi-agent/src/harness/session/memory-repo.ts +50 -0
  173. wmo/harness/vendor/pi-agent/src/harness/session/memory-storage.ts +131 -0
  174. wmo/harness/vendor/pi-agent/src/harness/session/repo-utils.ts +51 -0
  175. wmo/harness/vendor/pi-agent/src/harness/session/session.ts +267 -0
  176. wmo/harness/vendor/pi-agent/src/harness/session/uuid.ts +54 -0
  177. wmo/harness/vendor/pi-agent/src/harness/skills.ts +375 -0
  178. wmo/harness/vendor/pi-agent/src/harness/system-prompt.ts +34 -0
  179. wmo/harness/vendor/pi-agent/src/harness/types.ts +836 -0
  180. wmo/harness/vendor/pi-agent/src/harness/utils/shell-output.ts +135 -0
  181. wmo/harness/vendor/pi-agent/src/harness/utils/truncate.ts +344 -0
  182. wmo/harness/vendor/pi-agent/src/index.ts +44 -0
  183. wmo/harness/vendor/pi-agent/src/node.ts +2 -0
  184. wmo/harness/vendor/pi-agent/src/proxy.ts +367 -0
  185. wmo/harness/vendor/pi-agent/src/types.ts +428 -0
  186. wmo/harness/vendor/pi-agent/test/agent-loop.test.ts +1351 -0
  187. wmo/harness/vendor/pi-agent/test/agent.test.ts +699 -0
  188. wmo/harness/vendor/pi-agent/test/e2e.test.ts +404 -0
  189. wmo/harness/vendor/pi-agent/test/harness/agent-harness-stream.test.ts +213 -0
  190. wmo/harness/vendor/pi-agent/test/harness/agent-harness.test.ts +608 -0
  191. wmo/harness/vendor/pi-agent/test/harness/compaction.test.ts +655 -0
  192. wmo/harness/vendor/pi-agent/test/harness/nodejs-env.test.ts +321 -0
  193. wmo/harness/vendor/pi-agent/test/harness/prompt-templates.test.ts +90 -0
  194. wmo/harness/vendor/pi-agent/test/harness/repo.test.ts +68 -0
  195. wmo/harness/vendor/pi-agent/test/harness/resource-formatting.test.ts +24 -0
  196. wmo/harness/vendor/pi-agent/test/harness/session-test-utils.ts +55 -0
  197. wmo/harness/vendor/pi-agent/test/harness/session-uuid.test.ts +50 -0
  198. wmo/harness/vendor/pi-agent/test/harness/session.test.ts +156 -0
  199. wmo/harness/vendor/pi-agent/test/harness/skills.test.ts +116 -0
  200. wmo/harness/vendor/pi-agent/test/harness/storage.test.ts +299 -0
  201. wmo/harness/vendor/pi-agent/test/harness/system-prompt.test.ts +66 -0
  202. wmo/harness/vendor/pi-agent/test/harness/truncate.test.ts +169 -0
  203. wmo/harness/vendor/pi-agent/test/scratch/simple.ts +72 -0
  204. wmo/harness/vendor/pi-agent/test/utils/calculate.ts +32 -0
  205. wmo/harness/vendor/pi-agent/test/utils/get-current-time.ts +46 -0
  206. wmo/harness/vendor/pi-agent/tsconfig.build.json +13 -0
  207. wmo/harness/vendor/pi-agent/vitest.config.ts +19 -0
  208. wmo/harness/vendor/pi-agent/vitest.harness.config.ts +28 -0
  209. wmo/harness/vendor/vendor_pi.sh +59 -0
  210. wmo/harness/workspace_patch.py +270 -0
  211. wmo/ingest/__init__.py +47 -0
  212. wmo/ingest/adapter.py +72 -0
  213. wmo/ingest/base.py +114 -0
  214. wmo/ingest/braintrust.py +339 -0
  215. wmo/ingest/detect.py +126 -0
  216. wmo/ingest/langfuse.py +291 -0
  217. wmo/ingest/langsmith.py +444 -0
  218. wmo/ingest/mastra.py +330 -0
  219. wmo/ingest/messages.py +170 -0
  220. wmo/ingest/normalize.py +679 -0
  221. wmo/ingest/otel_genai.py +69 -0
  222. wmo/ingest/otel_writer.py +100 -0
  223. wmo/ingest/phoenix.py +150 -0
  224. wmo/ingest/postgres.py +246 -0
  225. wmo/ingest/posthog.py +320 -0
  226. wmo/ingest/quality.py +28 -0
  227. wmo/ingest/stream.py +209 -0
  228. wmo/ingest/testdata/sample_otlp.json +60 -0
  229. wmo/ingest/testdata/sample_spans.jsonl +3 -0
  230. wmo/optimize/__init__.py +25 -0
  231. wmo/optimize/base.py +143 -0
  232. wmo/optimize/gepa.py +806 -0
  233. wmo/optimize/judge.py +262 -0
  234. wmo/optimize/judge_quality.py +359 -0
  235. wmo/optimize/knn.py +468 -0
  236. wmo/optimize/numeric.py +152 -0
  237. wmo/optimize/outcomes.py +103 -0
  238. wmo/optimize/policy.py +669 -0
  239. wmo/optimize/report.py +231 -0
  240. wmo/optimize/reward.py +129 -0
  241. wmo/optimize/routing.py +373 -0
  242. wmo/platform/__init__.py +6 -0
  243. wmo/platform/auth.py +115 -0
  244. wmo/platform/client.py +551 -0
  245. wmo/platform/credentials.py +126 -0
  246. wmo/platform/transfer.py +158 -0
  247. wmo/providers/__init__.py +40 -0
  248. wmo/providers/_bedrock_chat.py +155 -0
  249. wmo/providers/_openai_common.py +182 -0
  250. wmo/providers/_responses_common.py +472 -0
  251. wmo/providers/anthropic.py +134 -0
  252. wmo/providers/azure_openai.py +296 -0
  253. wmo/providers/base.py +300 -0
  254. wmo/providers/bedrock.py +312 -0
  255. wmo/providers/models.py +205 -0
  256. wmo/providers/openai.py +143 -0
  257. wmo/providers/openai_responses.py +240 -0
  258. wmo/providers/pool.py +170 -0
  259. wmo/providers/registry.py +73 -0
  260. wmo/providers/retry.py +151 -0
  261. wmo/providers/tinker.py +936 -0
  262. wmo/providers/waterfall.py +336 -0
  263. wmo/research/__init__.py +81 -0
  264. wmo/research/ablation.py +133 -0
  265. wmo/research/concurrency_plot.py +523 -0
  266. wmo/research/concurrency_run.py +240 -0
  267. wmo/research/concurrency_scaling.py +270 -0
  268. wmo/research/gepa_scaling.py +274 -0
  269. wmo/research/pipeline.py +198 -0
  270. wmo/research/scaling_split.py +82 -0
  271. wmo/research/scenario_fidelity.py +198 -0
  272. wmo/research/scenario_recovery.py +92 -0
  273. wmo/research/seed_stability.py +90 -0
  274. wmo/research/trace_scaling.py +348 -0
  275. wmo/retrieval/__init__.py +6 -0
  276. wmo/retrieval/embedders.py +105 -0
  277. wmo/retrieval/leakfree.py +52 -0
  278. wmo/retrieval/retriever.py +173 -0
  279. wmo/scenarios/__init__.py +58 -0
  280. wmo/scenarios/builder.py +152 -0
  281. wmo/scenarios/mining/__init__.py +27 -0
  282. wmo/scenarios/mining/clustering.py +171 -0
  283. wmo/scenarios/mining/facets.py +226 -0
  284. wmo/scenarios/mining/selection.py +220 -0
  285. wmo/scenarios/synthesis/__init__.py +6 -0
  286. wmo/scenarios/synthesis/scenario_set.py +63 -0
  287. wmo/scenarios/synthesis/synthesizer.py +85 -0
  288. wmo/scenarios/verification/__init__.py +17 -0
  289. wmo/scenarios/verification/judge.py +97 -0
  290. wmo/scenarios/verification/verify.py +135 -0
  291. wmo/serving/__init__.py +5 -0
  292. wmo/serving/builds.py +451 -0
  293. wmo/serving/chat.py +878 -0
  294. wmo/serving/endpoint_config.py +64 -0
  295. wmo/serving/savings.py +250 -0
  296. wmo/serving/server.py +553 -0
  297. wmo/serving/traces_source.py +206 -0
  298. wmo/telemetry.py +213 -0
  299. wmo/tracking/__init__.py +36 -0
  300. wmo/tracking/clock.py +24 -0
  301. wmo/tracking/metered.py +125 -0
  302. wmo/tracking/pricing.py +99 -0
  303. wmo/tracking/store.py +31 -0
  304. wmo/tracking/tracker.py +149 -0
  305. world_model_optimizer-0.2.0.dist-info/METADATA +203 -0
  306. world_model_optimizer-0.2.0.dist-info/RECORD +308 -0
  307. world_model_optimizer-0.2.0.dist-info/WHEEL +4 -0
  308. world_model_optimizer-0.2.0.dist-info/entry_points.txt +2 -0
wmo/connect/oauth.py ADDED
@@ -0,0 +1,461 @@
1
+ """Provider-agnostic OAuth building blocks: PKCE, loopback and device flows, refresh.
2
+
3
+ Connectors compose these instead of reimplementing OAuth per service. Everything user-visible is
4
+ injected (`open_url`, `present`, `sleep`, an httpx `transport`), so this module never prints and
5
+ tests never touch the network or wait on real clocks.
6
+ """
7
+
8
+ from __future__ import annotations
9
+
10
+ import base64
11
+ import hashlib
12
+ import logging
13
+ import re
14
+ import secrets
15
+ import threading
16
+ import time
17
+ from collections.abc import Callable
18
+ from datetime import UTC, datetime, timedelta
19
+ from http.server import BaseHTTPRequestHandler, HTTPServer
20
+ from typing import cast
21
+ from urllib.parse import parse_qsl, urlencode, urlsplit
22
+
23
+ import httpx
24
+ from pydantic import BaseModel, Field
25
+
26
+ from wmo.connect.credentials import save_connector_auth
27
+ from wmo.connect.types import ConnectError, ConnectorAuth, transport_errors
28
+ from wmo.core.types import JsonObject, JsonValue
29
+
30
+ logger = logging.getLogger(__name__)
31
+
32
+ # Refresh when the access token expires within this window (seconds).
33
+ REFRESH_LEEWAY_SECONDS = 60.0
34
+
35
+ _TOKEN_TIMEOUT_SECONDS = 30.0
36
+
37
+ _DEVICE_GRANT = "urn:ietf:params:oauth:grant-type:device_code"
38
+
39
+ _CALLBACK_HTML = (
40
+ '<!doctype html><html><body style="font-family: sans-serif; padding: 2rem;">'
41
+ "<p>Authorization received. You can close this tab and return to the terminal.</p>"
42
+ "</body></html>"
43
+ )
44
+
45
+
46
+ class OAuthApp(BaseModel):
47
+ """One provider's OAuth application: endpoints, client credential, default scopes.
48
+
49
+ Attributes:
50
+ name: Provider name (matches the connector name, e.g. "github").
51
+ client_id: The registered OAuth client id.
52
+ client_secret: Client secret, for providers whose token endpoint requires one.
53
+ auth_url: Browser authorization endpoint.
54
+ token_url: Token exchange/refresh endpoint.
55
+ device_url: Device authorization endpoint (RFC 8628), when the provider has one.
56
+ scopes: Default scopes requested when a flow is run without an explicit list.
57
+ extra_auth_params: Additional authorize-URL query params (e.g. Google's
58
+ access_type=offline).
59
+ """
60
+
61
+ name: str
62
+ client_id: str
63
+ client_secret: str | None = None
64
+ auth_url: str
65
+ token_url: str
66
+ device_url: str | None = None
67
+ scopes: list[str] = Field(default_factory=list)
68
+ extra_auth_params: dict[str, str] = Field(default_factory=dict)
69
+
70
+
71
+ def pkce_challenge() -> tuple[str, str]:
72
+ """A fresh PKCE S256 pair (RFC 7636): (verifier, challenge).
73
+
74
+ The verifier is 43-128 chars of URL-safe randomness; the challenge is the unpadded
75
+ base64url-encoded SHA-256 of the verifier.
76
+ """
77
+ verifier = secrets.token_urlsafe(64)
78
+ digest = hashlib.sha256(verifier.encode("ascii")).digest()
79
+ challenge = base64.urlsafe_b64encode(digest).rstrip(b"=").decode("ascii")
80
+ return verifier, challenge
81
+
82
+
83
+ class LoopbackServer(HTTPServer):
84
+ """A single-use localhost server that captures the OAuth redirect."""
85
+
86
+ def __init__(self) -> None:
87
+ super().__init__(("127.0.0.1", 0), _LoopbackHandler)
88
+ self.timeout = 0.1 # handle_request wakes this often so the serve loop can exit
89
+ self.callback_params: dict[str, str] | None = None
90
+ self.received = threading.Event()
91
+
92
+
93
+ class _LoopbackHandler(BaseHTTPRequestHandler):
94
+ """Answers the provider redirect with a close-this-tab page; ignores stray requests."""
95
+
96
+ def do_GET(self) -> None:
97
+ server = cast(LoopbackServer, self.server)
98
+ params = dict(parse_qsl(urlsplit(self.path).query))
99
+ if not params:
100
+ # Stray request (favicon and friends), not the callback: keep waiting.
101
+ self.send_response(404)
102
+ self.end_headers()
103
+ return
104
+ body = _CALLBACK_HTML.encode("utf-8")
105
+ self.send_response(200)
106
+ self.send_header("Content-Type", "text/html; charset=utf-8")
107
+ self.send_header("Content-Length", str(len(body)))
108
+ self.end_headers()
109
+ self.wfile.write(body)
110
+ server.callback_params = params
111
+ server.received.set()
112
+
113
+ def log_message(self, format: str, *args: object) -> None:
114
+ logger.debug("oauth loopback server: %s", format % args)
115
+
116
+
117
+ def serve_until(server: LoopbackServer, deadline: float) -> None:
118
+ """Handle requests until the callback arrives or the deadline passes."""
119
+ while not server.received.is_set() and time.monotonic() < deadline:
120
+ server.handle_request()
121
+
122
+
123
+ def run_loopback_flow(
124
+ app: OAuthApp,
125
+ *,
126
+ scopes: list[str] | None = None,
127
+ open_url: Callable[[str], None],
128
+ timeout: float = 300.0,
129
+ transport: httpx.BaseTransport | None = None,
130
+ ) -> ConnectorAuth:
131
+ """Run the browser authorization-code flow against a localhost redirect (PKCE S256).
132
+
133
+ Binds an ephemeral `http://127.0.0.1:<port>/callback` server in a thread, hands the
134
+ authorize URL to `open_url` (the CLI layer prints/opens it; this module never does), waits
135
+ for exactly one callback, verifies the `state`, then exchanges the code at `app.token_url`.
136
+
137
+ Args:
138
+ app: The provider's OAuth application config.
139
+ scopes: Scopes to request; defaults to `app.scopes`.
140
+ open_url: Called once with the authorize URL (open a browser, print it, ...).
141
+ timeout: Seconds to wait for the browser callback.
142
+ transport: Injected httpx transport for the token exchange (tests); None = real network.
143
+
144
+ Returns:
145
+ The normalized oauth-kind credential.
146
+
147
+ Raises:
148
+ ConnectError: On denied consent, a state mismatch, a callback timeout, or a failed
149
+ token exchange.
150
+ """
151
+ requested = list(scopes) if scopes is not None else list(app.scopes)
152
+ verifier, challenge = pkce_challenge()
153
+ state = secrets.token_urlsafe(16)
154
+ server = LoopbackServer()
155
+ deadline = time.monotonic() + timeout
156
+ thread = threading.Thread(
157
+ target=serve_until, args=(server, deadline), name="wmo-oauth-loopback", daemon=True
158
+ )
159
+ port = int(server.server_address[1])
160
+ redirect_uri = f"http://127.0.0.1:{port}/callback"
161
+ try:
162
+ thread.start()
163
+ params: dict[str, str] = {
164
+ "client_id": app.client_id,
165
+ "redirect_uri": redirect_uri,
166
+ "response_type": "code",
167
+ "state": state,
168
+ "code_challenge": challenge,
169
+ "code_challenge_method": "S256",
170
+ **app.extra_auth_params,
171
+ }
172
+ if requested:
173
+ params["scope"] = " ".join(requested)
174
+ open_url(f"{app.auth_url}?{urlencode(params)}")
175
+ received = server.received.wait(timeout)
176
+ finally:
177
+ server.received.set()
178
+ thread.join(2.0)
179
+ server.server_close()
180
+ if not received:
181
+ raise ConnectError(
182
+ f"timed out after {timeout:g}s waiting for the {app.name} OAuth callback; "
183
+ "re-run the command and approve access in the browser"
184
+ )
185
+ callback = server.callback_params or {}
186
+ if "error" in callback:
187
+ raise ConnectError(
188
+ f"{app.name} authorization failed: {_describe_callback_error(callback)}; "
189
+ "re-run the command and approve access"
190
+ )
191
+ if callback.get("state") != state:
192
+ raise ConnectError(
193
+ f"OAuth state mismatch in the {app.name} callback (a stale browser tab or a "
194
+ "forged request); re-run the command and use the freshly opened tab"
195
+ )
196
+ code = callback.get("code")
197
+ if not code:
198
+ raise ConnectError(
199
+ f"the {app.name} callback carried no authorization code; re-run the command"
200
+ )
201
+ data = {
202
+ "grant_type": "authorization_code",
203
+ "code": code,
204
+ "redirect_uri": redirect_uri,
205
+ "client_id": app.client_id,
206
+ "code_verifier": verifier,
207
+ }
208
+ if app.client_secret:
209
+ data["client_secret"] = app.client_secret
210
+ response, payload = _token_request(app.token_url, data, transport=transport)
211
+ _raise_on_token_error(app, response, payload, doing="token exchange")
212
+ return _auth_from_token_response(payload, requested)
213
+
214
+
215
+ def run_device_flow(
216
+ app: OAuthApp,
217
+ *,
218
+ scopes: list[str] | None = None,
219
+ present: Callable[[str, str], None],
220
+ timeout: float = 900.0,
221
+ transport: httpx.BaseTransport | None = None,
222
+ sleep: Callable[[float], None] = time.sleep,
223
+ ) -> ConnectorAuth:
224
+ """Run the RFC 8628 device authorization flow (for headless/SSH sessions).
225
+
226
+ Requests a device code, hands `(verification_uri, user_code)` to `present`, then polls
227
+ `app.token_url` honoring the server's `interval` and `slow_down` (+5s) responses.
228
+
229
+ Args:
230
+ app: The provider's OAuth application config (must define `device_url`).
231
+ scopes: Scopes to request; defaults to `app.scopes`.
232
+ present: Called once with the verification URI and user code to show the user.
233
+ timeout: Polling budget in seconds, measured against the injected `sleep` amounts.
234
+ transport: Injected httpx transport (tests); None = real network.
235
+ sleep: Injected sleeper so tests never wait on real clocks.
236
+
237
+ Returns:
238
+ The normalized oauth-kind credential.
239
+
240
+ Raises:
241
+ ConnectError: When the provider has no device endpoint, the user denies access, the
242
+ device code expires, or polling exceeds `timeout`.
243
+ """
244
+ if not app.device_url:
245
+ raise ConnectError(
246
+ f"{app.name} defines no device authorization endpoint; "
247
+ "use the browser flow (run_loopback_flow) instead"
248
+ )
249
+ requested = list(scopes) if scopes is not None else list(app.scopes)
250
+ data = {"client_id": app.client_id}
251
+ if requested:
252
+ data["scope"] = " ".join(requested)
253
+ response, payload = _token_request(app.device_url, data, transport=transport)
254
+ device_code = payload.get("device_code")
255
+ if response.status_code != 200 or not isinstance(device_code, str) or not device_code:
256
+ raise ConnectError(
257
+ f"device authorization at {app.device_url} failed "
258
+ f"(HTTP {response.status_code}): {response.text[:200]}; "
259
+ "check the OAuth app's client id and retry"
260
+ )
261
+ user_code = str(payload.get("user_code") or "")
262
+ verification_uri = str(payload.get("verification_uri") or payload.get("verification_url") or "")
263
+ interval_value = payload.get("interval")
264
+ interval = float(interval_value) if isinstance(interval_value, int | float) else 5.0
265
+
266
+ present(verification_uri, user_code)
267
+
268
+ poll = {"grant_type": _DEVICE_GRANT, "device_code": device_code, "client_id": app.client_id}
269
+ if app.client_secret:
270
+ poll["client_secret"] = app.client_secret
271
+ elapsed = 0.0
272
+ while True:
273
+ if elapsed + interval > timeout:
274
+ raise ConnectError(
275
+ f"timed out after {timeout:g}s waiting for {app.name} device authorization; "
276
+ "re-run the command and enter the code sooner"
277
+ )
278
+ sleep(interval)
279
+ elapsed += interval
280
+ response, payload = _token_request(app.token_url, poll, transport=transport)
281
+ error = payload.get("error")
282
+ if error == "authorization_pending":
283
+ continue
284
+ if error == "slow_down":
285
+ interval += 5.0 # RFC 8628 section 3.5
286
+ continue
287
+ if error == "access_denied":
288
+ raise ConnectError(
289
+ f"{app.name} authorization was denied; re-run the command and approve access"
290
+ )
291
+ if error == "expired_token":
292
+ raise ConnectError(
293
+ f"the {app.name} device code expired before it was approved; "
294
+ "re-run the command for a fresh code"
295
+ )
296
+ _raise_on_token_error(app, response, payload, doing="device token poll")
297
+ return _auth_from_token_response(payload, requested)
298
+
299
+
300
+ def refresh_auth(
301
+ app: OAuthApp,
302
+ auth: ConnectorAuth,
303
+ *,
304
+ transport: httpx.BaseTransport | None = None,
305
+ ) -> ConnectorAuth:
306
+ """Exchange the refresh token for a fresh access token.
307
+
308
+ Keeps the old refresh token when the response omits one (providers commonly rotate only the
309
+ access token) and carries the stored `account`/`extra` identity forward.
310
+
311
+ Raises:
312
+ ConnectError: When no refresh token is stored or the provider rejects the refresh.
313
+ """
314
+ if not auth.refresh_token:
315
+ raise ConnectError(
316
+ f"no refresh token stored for {app.name}; the connection must be reauthorized"
317
+ )
318
+ data = {
319
+ "grant_type": "refresh_token",
320
+ "refresh_token": auth.refresh_token,
321
+ "client_id": app.client_id,
322
+ }
323
+ if app.client_secret:
324
+ data["client_secret"] = app.client_secret
325
+ response, payload = _token_request(app.token_url, data, transport=transport)
326
+ _raise_on_token_error(app, response, payload, doing="token refresh")
327
+ refreshed = _auth_from_token_response(payload, auth.scopes)
328
+ return refreshed.model_copy(
329
+ update={
330
+ "refresh_token": refreshed.refresh_token or auth.refresh_token,
331
+ "account": auth.account,
332
+ "extra": auth.extra,
333
+ }
334
+ )
335
+
336
+
337
+ def ensure_fresh(
338
+ app: OAuthApp,
339
+ name: str,
340
+ auth: ConnectorAuth,
341
+ *,
342
+ transport: httpx.BaseTransport | None = None,
343
+ ) -> ConnectorAuth:
344
+ """Refresh `auth` iff it expires within `REFRESH_LEEWAY_SECONDS` and can be refreshed.
345
+
346
+ A refreshed credential is persisted via `save_connector_auth(name, ...)` so the next
347
+ invocation starts fresh. Credentials without a refresh token or expiry pass through
348
+ unchanged (an unparseable `expires_at` is treated as already expired).
349
+ """
350
+ if not auth.refresh_token or not auth.expires_at:
351
+ return auth
352
+ expires = _parse_expiry(auth.expires_at)
353
+ if expires is not None:
354
+ remaining = (expires - datetime.now(UTC)).total_seconds()
355
+ if remaining > REFRESH_LEEWAY_SECONDS:
356
+ return auth
357
+ refreshed = refresh_auth(app, auth, transport=transport)
358
+ save_connector_auth(name, refreshed)
359
+ return refreshed
360
+
361
+
362
+ def _parse_expiry(value: str) -> datetime | None:
363
+ """Parse an ISO-8601 expiry; naive timestamps are taken as UTC, garbage becomes None."""
364
+ try:
365
+ parsed = datetime.fromisoformat(value.replace("Z", "+00:00"))
366
+ except ValueError:
367
+ return None
368
+ return parsed if parsed.tzinfo else parsed.replace(tzinfo=UTC)
369
+
370
+
371
+ def _token_request(
372
+ url: str, data: dict[str, str], *, transport: httpx.BaseTransport | None
373
+ ) -> tuple[httpx.Response, JsonObject]:
374
+ """POST a form-encoded OAuth request and parse the JSON object response (or {}).
375
+
376
+ Raises:
377
+ ConnectError: When the endpoint is unreachable (DNS, refused connection, timeout).
378
+ """
379
+ host = httpx.URL(url).host or url
380
+ with httpx.Client(timeout=_TOKEN_TIMEOUT_SECONDS, transport=transport) as client:
381
+ with transport_errors(host):
382
+ response = client.post(url, data=data, headers={"Accept": "application/json"})
383
+ try:
384
+ raw = response.json()
385
+ except ValueError:
386
+ raw = None
387
+ payload = cast(JsonObject, raw) if isinstance(raw, dict) else {}
388
+ return response, payload
389
+
390
+
391
+ def _raise_on_token_error(
392
+ app: OAuthApp, response: httpx.Response, payload: JsonObject, *, doing: str
393
+ ) -> None:
394
+ """Turn OAuth error payloads and non-200 responses into actionable ConnectErrors."""
395
+ error = payload.get("error")
396
+ if error:
397
+ raise ConnectError(
398
+ f"{app.name} {doing} failed: {_describe_oauth_error(payload)}; "
399
+ "the connection must be reauthorized (check the OAuth app's client id/secret if it "
400
+ "keeps failing)"
401
+ )
402
+ if response.status_code != 200:
403
+ raise ConnectError(
404
+ f"{app.name} {doing} returned HTTP {response.status_code}: {response.text[:200]}; "
405
+ "check the OAuth app configuration and retry"
406
+ )
407
+
408
+
409
+ def _describe_oauth_error(payload: JsonObject) -> str:
410
+ """'error: error_description' when a description is present, else just the error code."""
411
+ error = payload.get("error")
412
+ description = payload.get("error_description")
413
+ return f"{error}: {description}" if description else str(error)
414
+
415
+
416
+ def _describe_callback_error(callback: dict[str, str]) -> str:
417
+ """Same shape as `_describe_oauth_error` for redirect-callback query params."""
418
+ error = callback.get("error", "unknown_error")
419
+ description = callback.get("error_description")
420
+ return f"{error}: {description}" if description else error
421
+
422
+
423
+ def _auth_from_token_response(payload: JsonObject, fallback_scopes: list[str]) -> ConnectorAuth:
424
+ """Normalize an OAuth token response into a `ConnectorAuth`.
425
+
426
+ Handles the provider quirks the harness has to care about: `expires_in` as int or numeric
427
+ string (turned into an absolute ISO-8601 `expires_at`), `scope` as a space or comma
428
+ separated string or a list, and responses that omit scopes entirely (falls back to what was
429
+ requested).
430
+ """
431
+ access_token = payload.get("access_token")
432
+ if not isinstance(access_token, str) or not access_token:
433
+ raise ConnectError(
434
+ "the token response carried no access_token; "
435
+ "check the OAuth app's client id/secret and retry"
436
+ )
437
+ refresh_token = payload.get("refresh_token")
438
+ expires_in = payload.get("expires_in")
439
+ if isinstance(expires_in, str) and expires_in.isdigit():
440
+ expires_in = int(expires_in)
441
+ expires_at: str | None = None
442
+ if isinstance(expires_in, int | float) and not isinstance(expires_in, bool):
443
+ moment = datetime.now(UTC) + timedelta(seconds=float(expires_in))
444
+ expires_at = moment.isoformat(timespec="seconds")
445
+ scopes = _parse_scopes(payload.get("scope"))
446
+ return ConnectorAuth(
447
+ kind="oauth",
448
+ access_token=access_token,
449
+ refresh_token=refresh_token if isinstance(refresh_token, str) else None,
450
+ expires_at=expires_at,
451
+ scopes=scopes if scopes is not None else list(fallback_scopes),
452
+ )
453
+
454
+
455
+ def _parse_scopes(value: JsonValue | None) -> list[str] | None:
456
+ """Scope fields arrive as a space/comma separated string or a list; None means absent."""
457
+ if isinstance(value, str):
458
+ return [scope for scope in re.split(r"[\s,]+", value) if scope]
459
+ if isinstance(value, list):
460
+ return [scope for scope in value if isinstance(scope, str)]
461
+ return None