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.
- llm_waterfall/LICENSE +21 -0
- llm_waterfall/__init__.py +53 -0
- llm_waterfall/adapters/__init__.py +36 -0
- llm_waterfall/adapters/anthropic.py +105 -0
- llm_waterfall/adapters/aws_mantle.py +47 -0
- llm_waterfall/adapters/azure_openai.py +71 -0
- llm_waterfall/adapters/base.py +51 -0
- llm_waterfall/adapters/bedrock.py +309 -0
- llm_waterfall/adapters/openai.py +130 -0
- llm_waterfall/classify.py +184 -0
- llm_waterfall/pricing.py +110 -0
- llm_waterfall/py.typed +0 -0
- llm_waterfall/types.py +295 -0
- llm_waterfall/waterfall.py +255 -0
- wmo/__init__.py +38 -0
- wmo/agents/__init__.py +7 -0
- wmo/agents/default.py +29 -0
- wmo/agents/meta.py +55 -0
- wmo/agents/optimizer.py +55 -0
- wmo/agents/project.py +928 -0
- wmo/cli/__init__.py +5 -0
- wmo/cli/agent_session.py +1123 -0
- wmo/cli/app.py +2489 -0
- wmo/cli/e2b_cmds.py +212 -0
- wmo/cli/eval_closed_loop.py +207 -0
- wmo/cli/harness_app.py +1147 -0
- wmo/cli/harness_distill.py +659 -0
- wmo/cli/hosted_session.py +880 -0
- wmo/cli/ingest_cmd.py +165 -0
- wmo/cli/model_roles.py +82 -0
- wmo/cli/platform_cmds.py +372 -0
- wmo/cli/route_app.py +274 -0
- wmo/cli/session_state.py +243 -0
- wmo/cli/ui.py +1107 -0
- wmo/cli/workspace_sync.py +504 -0
- wmo/config/__init__.py +60 -0
- wmo/config/card.py +129 -0
- wmo/config/config.py +367 -0
- wmo/config/dotenv.py +67 -0
- wmo/config/settings.py +128 -0
- wmo/config/store.py +177 -0
- wmo/conftest.py +19 -0
- wmo/connect/__init__.py +88 -0
- wmo/connect/apps.py +78 -0
- wmo/connect/brave.py +284 -0
- wmo/connect/connector.py +79 -0
- wmo/connect/credentials.py +164 -0
- wmo/connect/github.py +321 -0
- wmo/connect/google.py +627 -0
- wmo/connect/notion.py +790 -0
- wmo/connect/oauth.py +461 -0
- wmo/connect/slack.py +555 -0
- wmo/connect/store.py +199 -0
- wmo/connect/types.py +156 -0
- wmo/core/__init__.py +21 -0
- wmo/core/parsing.py +281 -0
- wmo/core/render.py +271 -0
- wmo/core/text.py +40 -0
- wmo/core/types.py +116 -0
- wmo/distill/__init__.py +14 -0
- wmo/distill/agents.py +140 -0
- wmo/distill/config.py +1006 -0
- wmo/distill/cost.py +437 -0
- wmo/distill/data.py +921 -0
- wmo/distill/deadlines.py +254 -0
- wmo/distill/fake_tinker.py +734 -0
- wmo/distill/gate.py +122 -0
- wmo/distill/loop.py +3499 -0
- wmo/distill/renderers.py +399 -0
- wmo/distill/rendering.py +620 -0
- wmo/distill/rollouts.py +726 -0
- wmo/distill/samples.py +195 -0
- wmo/distill/store.py +829 -0
- wmo/distill/teacher.py +714 -0
- wmo/distill/tokens.py +535 -0
- wmo/distill/tracking.py +552 -0
- wmo/distill/tripwire.py +411 -0
- wmo/distill/xtoken/byte_offsets.py +152 -0
- wmo/distill/xtoken/chunks.py +457 -0
- wmo/distill/xtoken/prompt_logprobs.py +475 -0
- wmo/distill/xtoken/teacher_render.py +346 -0
- wmo/engine/__init__.py +28 -0
- wmo/engine/autoconfig.py +367 -0
- wmo/engine/build.py +346 -0
- wmo/engine/demo.py +77 -0
- wmo/engine/eval_suites.py +245 -0
- wmo/engine/grounding.py +491 -0
- wmo/engine/knowledge.py +291 -0
- wmo/engine/loader.py +36 -0
- wmo/engine/play.py +92 -0
- wmo/engine/prompts.py +99 -0
- wmo/engine/replay.py +443 -0
- wmo/engine/reporting.py +58 -0
- wmo/engine/workspace.py +468 -0
- wmo/engine/world_model.py +568 -0
- wmo/env/__init__.py +22 -0
- wmo/env/base.py +121 -0
- wmo/env/closed_loop.py +229 -0
- wmo/env/episode.py +107 -0
- wmo/env/llm_agent.py +93 -0
- wmo/env/scenarios.py +73 -0
- wmo/evals/__init__.py +52 -0
- wmo/evals/agreement.py +110 -0
- wmo/evals/base.py +45 -0
- wmo/evals/closed_loop.py +480 -0
- wmo/evals/failover.py +96 -0
- wmo/evals/gold.py +127 -0
- wmo/evals/grid.py +394 -0
- wmo/evals/grid_plot.py +205 -0
- wmo/evals/harbor/__init__.py +27 -0
- wmo/evals/harbor/agent.py +573 -0
- wmo/evals/harbor/ctrf.py +171 -0
- wmo/evals/harbor/e2b_environment.py +587 -0
- wmo/evals/harbor/e2b_template_policy.py +144 -0
- wmo/evals/harbor/scorer.py +875 -0
- wmo/evals/harbor/tasks.py +140 -0
- wmo/evals/open_loop.py +194 -0
- wmo/evals/tasks.py +53 -0
- wmo/harness/__init__.py +51 -0
- wmo/harness/code_runtime.py +288 -0
- wmo/harness/create.py +1191 -0
- wmo/harness/delta.py +220 -0
- wmo/harness/doc.py +556 -0
- wmo/harness/e2b_ledger.py +342 -0
- wmo/harness/e2b_reap.py +476 -0
- wmo/harness/e2b_sandbox.py +350 -0
- wmo/harness/environment.py +35 -0
- wmo/harness/live_session.py +543 -0
- wmo/harness/mutate.py +343 -0
- wmo/harness/pi_e2b.py +1710 -0
- wmo/harness/pi_entry/entry.ts +268 -0
- wmo/harness/pi_entry/runner_frames.ts +92 -0
- wmo/harness/pi_entry/runner_live.ts +587 -0
- wmo/harness/pi_entry/runner_service.ts +270 -0
- wmo/harness/pi_entry/runner_stdio.ts +374 -0
- wmo/harness/pi_entry/runner_termination.ts +142 -0
- wmo/harness/pi_local.py +262 -0
- wmo/harness/pi_runtime.py +495 -0
- wmo/harness/pi_vendor.py +65 -0
- wmo/harness/population.py +509 -0
- wmo/harness/project_proposer.py +569 -0
- wmo/harness/proposer.py +977 -0
- wmo/harness/runner_link.py +619 -0
- wmo/harness/runtime.py +389 -0
- wmo/harness/scoring.py +247 -0
- wmo/harness/skills.py +116 -0
- wmo/harness/source_tree.py +319 -0
- wmo/harness/store.py +176 -0
- wmo/harness/tools.py +105 -0
- wmo/harness/vendor/manifest.sha256 +58 -0
- wmo/harness/vendor/pi-agent/CHANGELOG.md +556 -0
- wmo/harness/vendor/pi-agent/LICENSE +21 -0
- wmo/harness/vendor/pi-agent/README.md +488 -0
- wmo/harness/vendor/pi-agent/VENDOR.md +39 -0
- wmo/harness/vendor/pi-agent/docs/agent-harness.md +486 -0
- wmo/harness/vendor/pi-agent/docs/durable-harness.md +212 -0
- wmo/harness/vendor/pi-agent/docs/hooks.md +445 -0
- wmo/harness/vendor/pi-agent/docs/models.md +966 -0
- wmo/harness/vendor/pi-agent/docs/observability.md +376 -0
- wmo/harness/vendor/pi-agent/package.json +60 -0
- wmo/harness/vendor/pi-agent/src/agent-loop.ts +748 -0
- wmo/harness/vendor/pi-agent/src/agent.ts +575 -0
- wmo/harness/vendor/pi-agent/src/harness/agent-harness.ts +1029 -0
- wmo/harness/vendor/pi-agent/src/harness/compaction/branch-summarization.ts +261 -0
- wmo/harness/vendor/pi-agent/src/harness/compaction/compaction.ts +747 -0
- wmo/harness/vendor/pi-agent/src/harness/compaction/utils.ts +144 -0
- wmo/harness/vendor/pi-agent/src/harness/env/nodejs.ts +550 -0
- wmo/harness/vendor/pi-agent/src/harness/messages.ts +164 -0
- wmo/harness/vendor/pi-agent/src/harness/prompt-templates.ts +267 -0
- wmo/harness/vendor/pi-agent/src/harness/session/jsonl-repo.ts +177 -0
- wmo/harness/vendor/pi-agent/src/harness/session/jsonl-storage.ts +293 -0
- wmo/harness/vendor/pi-agent/src/harness/session/memory-repo.ts +50 -0
- wmo/harness/vendor/pi-agent/src/harness/session/memory-storage.ts +131 -0
- wmo/harness/vendor/pi-agent/src/harness/session/repo-utils.ts +51 -0
- wmo/harness/vendor/pi-agent/src/harness/session/session.ts +267 -0
- wmo/harness/vendor/pi-agent/src/harness/session/uuid.ts +54 -0
- wmo/harness/vendor/pi-agent/src/harness/skills.ts +375 -0
- wmo/harness/vendor/pi-agent/src/harness/system-prompt.ts +34 -0
- wmo/harness/vendor/pi-agent/src/harness/types.ts +836 -0
- wmo/harness/vendor/pi-agent/src/harness/utils/shell-output.ts +135 -0
- wmo/harness/vendor/pi-agent/src/harness/utils/truncate.ts +344 -0
- wmo/harness/vendor/pi-agent/src/index.ts +44 -0
- wmo/harness/vendor/pi-agent/src/node.ts +2 -0
- wmo/harness/vendor/pi-agent/src/proxy.ts +367 -0
- wmo/harness/vendor/pi-agent/src/types.ts +428 -0
- wmo/harness/vendor/pi-agent/test/agent-loop.test.ts +1351 -0
- wmo/harness/vendor/pi-agent/test/agent.test.ts +699 -0
- wmo/harness/vendor/pi-agent/test/e2e.test.ts +404 -0
- wmo/harness/vendor/pi-agent/test/harness/agent-harness-stream.test.ts +213 -0
- wmo/harness/vendor/pi-agent/test/harness/agent-harness.test.ts +608 -0
- wmo/harness/vendor/pi-agent/test/harness/compaction.test.ts +655 -0
- wmo/harness/vendor/pi-agent/test/harness/nodejs-env.test.ts +321 -0
- wmo/harness/vendor/pi-agent/test/harness/prompt-templates.test.ts +90 -0
- wmo/harness/vendor/pi-agent/test/harness/repo.test.ts +68 -0
- wmo/harness/vendor/pi-agent/test/harness/resource-formatting.test.ts +24 -0
- wmo/harness/vendor/pi-agent/test/harness/session-test-utils.ts +55 -0
- wmo/harness/vendor/pi-agent/test/harness/session-uuid.test.ts +50 -0
- wmo/harness/vendor/pi-agent/test/harness/session.test.ts +156 -0
- wmo/harness/vendor/pi-agent/test/harness/skills.test.ts +116 -0
- wmo/harness/vendor/pi-agent/test/harness/storage.test.ts +299 -0
- wmo/harness/vendor/pi-agent/test/harness/system-prompt.test.ts +66 -0
- wmo/harness/vendor/pi-agent/test/harness/truncate.test.ts +169 -0
- wmo/harness/vendor/pi-agent/test/scratch/simple.ts +72 -0
- wmo/harness/vendor/pi-agent/test/utils/calculate.ts +32 -0
- wmo/harness/vendor/pi-agent/test/utils/get-current-time.ts +46 -0
- wmo/harness/vendor/pi-agent/tsconfig.build.json +13 -0
- wmo/harness/vendor/pi-agent/vitest.config.ts +19 -0
- wmo/harness/vendor/pi-agent/vitest.harness.config.ts +28 -0
- wmo/harness/vendor/vendor_pi.sh +59 -0
- wmo/harness/workspace_patch.py +270 -0
- wmo/ingest/__init__.py +47 -0
- wmo/ingest/adapter.py +72 -0
- wmo/ingest/base.py +114 -0
- wmo/ingest/braintrust.py +339 -0
- wmo/ingest/detect.py +126 -0
- wmo/ingest/langfuse.py +291 -0
- wmo/ingest/langsmith.py +444 -0
- wmo/ingest/mastra.py +330 -0
- wmo/ingest/messages.py +170 -0
- wmo/ingest/normalize.py +679 -0
- wmo/ingest/otel_genai.py +69 -0
- wmo/ingest/otel_writer.py +100 -0
- wmo/ingest/phoenix.py +150 -0
- wmo/ingest/postgres.py +246 -0
- wmo/ingest/posthog.py +320 -0
- wmo/ingest/quality.py +28 -0
- wmo/ingest/stream.py +209 -0
- wmo/ingest/testdata/sample_otlp.json +60 -0
- wmo/ingest/testdata/sample_spans.jsonl +3 -0
- wmo/optimize/__init__.py +25 -0
- wmo/optimize/base.py +143 -0
- wmo/optimize/gepa.py +806 -0
- wmo/optimize/judge.py +262 -0
- wmo/optimize/judge_quality.py +359 -0
- wmo/optimize/knn.py +468 -0
- wmo/optimize/numeric.py +152 -0
- wmo/optimize/outcomes.py +103 -0
- wmo/optimize/policy.py +669 -0
- wmo/optimize/report.py +231 -0
- wmo/optimize/reward.py +129 -0
- wmo/optimize/routing.py +373 -0
- wmo/platform/__init__.py +6 -0
- wmo/platform/auth.py +115 -0
- wmo/platform/client.py +551 -0
- wmo/platform/credentials.py +126 -0
- wmo/platform/transfer.py +158 -0
- wmo/providers/__init__.py +40 -0
- wmo/providers/_bedrock_chat.py +155 -0
- wmo/providers/_openai_common.py +182 -0
- wmo/providers/_responses_common.py +472 -0
- wmo/providers/anthropic.py +134 -0
- wmo/providers/azure_openai.py +296 -0
- wmo/providers/base.py +300 -0
- wmo/providers/bedrock.py +312 -0
- wmo/providers/models.py +205 -0
- wmo/providers/openai.py +143 -0
- wmo/providers/openai_responses.py +240 -0
- wmo/providers/pool.py +170 -0
- wmo/providers/registry.py +73 -0
- wmo/providers/retry.py +151 -0
- wmo/providers/tinker.py +936 -0
- wmo/providers/waterfall.py +336 -0
- wmo/research/__init__.py +81 -0
- wmo/research/ablation.py +133 -0
- wmo/research/concurrency_plot.py +523 -0
- wmo/research/concurrency_run.py +240 -0
- wmo/research/concurrency_scaling.py +270 -0
- wmo/research/gepa_scaling.py +274 -0
- wmo/research/pipeline.py +198 -0
- wmo/research/scaling_split.py +82 -0
- wmo/research/scenario_fidelity.py +198 -0
- wmo/research/scenario_recovery.py +92 -0
- wmo/research/seed_stability.py +90 -0
- wmo/research/trace_scaling.py +348 -0
- wmo/retrieval/__init__.py +6 -0
- wmo/retrieval/embedders.py +105 -0
- wmo/retrieval/leakfree.py +52 -0
- wmo/retrieval/retriever.py +173 -0
- wmo/scenarios/__init__.py +58 -0
- wmo/scenarios/builder.py +152 -0
- wmo/scenarios/mining/__init__.py +27 -0
- wmo/scenarios/mining/clustering.py +171 -0
- wmo/scenarios/mining/facets.py +226 -0
- wmo/scenarios/mining/selection.py +220 -0
- wmo/scenarios/synthesis/__init__.py +6 -0
- wmo/scenarios/synthesis/scenario_set.py +63 -0
- wmo/scenarios/synthesis/synthesizer.py +85 -0
- wmo/scenarios/verification/__init__.py +17 -0
- wmo/scenarios/verification/judge.py +97 -0
- wmo/scenarios/verification/verify.py +135 -0
- wmo/serving/__init__.py +5 -0
- wmo/serving/builds.py +451 -0
- wmo/serving/chat.py +878 -0
- wmo/serving/endpoint_config.py +64 -0
- wmo/serving/savings.py +250 -0
- wmo/serving/server.py +553 -0
- wmo/serving/traces_source.py +206 -0
- wmo/telemetry.py +213 -0
- wmo/tracking/__init__.py +36 -0
- wmo/tracking/clock.py +24 -0
- wmo/tracking/metered.py +125 -0
- wmo/tracking/pricing.py +99 -0
- wmo/tracking/store.py +31 -0
- wmo/tracking/tracker.py +149 -0
- world_model_optimizer-0.2.0.dist-info/METADATA +203 -0
- world_model_optimizer-0.2.0.dist-info/RECORD +308 -0
- world_model_optimizer-0.2.0.dist-info/WHEEL +4 -0
- 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
|