feat(client): add proxy_domain + timeout kwargs, fix default proxy host
Some checks failed
ci/woodpecker/push/unit Pipeline was successful
ci/woodpecker/pr/unit Pipeline was successful
ci/woodpecker/pr/code-runner Pipeline failed
ci/woodpecker/pr/integration Pipeline was successful

- WrennClient/AsyncWrennClient accept proxy_domain= and timeout= kwargs
- WRENN_PROXY_DOMAIN env var supported
- Default proxy host: app.wrenn.dev -> wrenn.dev (was port-id.app.wrenn.dev)
- Custom base_url preserves host verbatim (with port)
- Default timeout: httpx.Timeout(30.0, connect=10.0)
- _build_proxy_url/_build_http_proxy_url take optional proxy_domain
- code_runner proxy + WS URL builders thread proxy_domain through
This commit is contained in:
2026-05-21 01:46:04 +06:00
parent db48e3cfbf
commit 8a62b6207c
12 changed files with 242 additions and 37 deletions

View File

@ -192,6 +192,29 @@ Jupyter kernel.
`httpx.AsyncClient` must be closed via `await close()` or
`async with`.
## Client Config
`WrennClient` / `AsyncWrennClient` accept:
- `api_key` — falls back to `WRENN_API_KEY`.
- `base_url` — falls back to `WRENN_BASE_URL`, then `DEFAULT_BASE_URL`
(`https://app.wrenn.dev/api`).
- `proxy_domain` — host suffix for capsule proxy URLs
(`{port}-{capsule_id}.<domain>`). Resolution:
1. explicit `proxy_domain=` kwarg
2. `WRENN_PROXY_DOMAIN` env
3. `wrenn.dev` when `base_url` host == `app.wrenn.dev` exactly
4. else `base_url` host (with port) verbatim
Exact match in step 3 is intentional: staging/other Wrenn envs keep
their host so they don't accidentally collapse to prod `wrenn.dev`.
- `timeout``httpx.Timeout | float | None`. Default
`httpx.Timeout(30.0, connect=10.0)`. Helper `_resolve_timeout`
centralizes the float-or-Timeout coercion.
`_build_proxy_url` / `_build_http_proxy_url` in `wrenn.capsule` now take
an optional `proxy_domain` arg. When omitted they fall back to the
`base_url` host (legacy behavior, preserved for direct callers/tests).
Production call sites pass `self._client._proxy_domain`.
### Tests
- `tests/test_code_runner_unit.py` — pure unit tests (respx + mocked

View File

@ -26,10 +26,31 @@ Optionally override the API base URL:
export WRENN_BASE_URL="https://app.wrenn.dev/api" # default
```
For self-hosted deployments you can also override the capsule proxy domain
(used to build `{port}-{capsule_id}.<domain>` URLs returned by
`Capsule.get_url`):
```bash
export WRENN_PROXY_DOMAIN="wrenn.example.com"
```
Resolution order: explicit `proxy_domain=` kwarg → `WRENN_PROXY_DOMAIN` env →
`wrenn.dev` when `base_url` is the default `app.wrenn.dev` host, else the
`base_url` host (with port) verbatim.
You can also pass credentials directly:
```python
from wrenn import Capsule
from wrenn import WrennClient, Capsule
# WrennClient also accepts a timeout (httpx.Timeout or float seconds).
# Default: 30s read/write/pool, 10s connect.
client = WrennClient(
api_key="wrn_...",
base_url="https://...",
proxy_domain="wrenn.example.com", # optional override
timeout=30.0, # optional override
)
capsule = Capsule(api_key="wrn_...", base_url="https://...")
```

View File

@ -1,5 +1,7 @@
from __future__ import annotations
DEFAULT_BASE_URL = "https://app.wrenn.dev/api"
DEFAULT_PROXY_DOMAIN = "wrenn.dev"
ENV_API_KEY = "WRENN_API_KEY"
ENV_BASE_URL = "WRENN_BASE_URL"
ENV_PROXY_DOMAIN = "WRENN_PROXY_DOMAIN"

View File

@ -434,7 +434,12 @@ class AsyncCapsule:
WebSocket access, see the lower-level ``_build_proxy_url``
helper or the ``pty()`` API.
"""
return _build_http_proxy_url(self._client._base_url, self._id, port)
return _build_http_proxy_url(
self._client._base_url,
self._id,
port,
self._client._proxy_domain,
)
# ── Snapshots ───────────────────────────────────────────────

View File

@ -20,27 +20,48 @@ from wrenn.models import Status, Template
from wrenn.pty import PtySession
def _build_proxy_url(base_url: str, capsule_id: str | None, port: int) -> str:
"""Build the WebSocket proxy URL (``ws://`` / ``wss://``)."""
def _build_proxy_url(
base_url: str,
capsule_id: str | None,
port: int,
proxy_domain: str | None = None,
) -> str:
"""Build the WebSocket proxy URL (``ws://`` / ``wss://``).
Scheme is derived from ``base_url``. The host portion comes from
``proxy_domain`` if provided; otherwise falls back to the ``base_url``
host (with port).
"""
parsed = httpx.URL(base_url)
host = parsed.host
if parsed.port:
host = f"{host}:{parsed.port}"
if proxy_domain:
host = proxy_domain
else:
host = parsed.host
if parsed.port:
host = f"{host}:{parsed.port}"
scheme = "ws" if parsed.scheme == "http" else "wss"
return f"{scheme}://{port}-{capsule_id}.{host}"
def _build_http_proxy_url(base_url: str, capsule_id: str | None, port: int) -> str:
def _build_http_proxy_url(
base_url: str,
capsule_id: str | None,
port: int,
proxy_domain: str | None = None,
) -> str:
"""Build the HTTP proxy URL (``http://`` / ``https://``).
The capsule's API base URL typically carries an ``/api`` path suffix
(e.g. ``https://app.wrenn.dev/api``). The proxy host is derived from
the URL's host only — any path is discarded.
Scheme is derived from ``base_url``. The host portion comes from
``proxy_domain`` if provided; otherwise falls back to the ``base_url``
host (with port). Any path on ``base_url`` is discarded.
"""
parsed = httpx.URL(base_url)
host = parsed.host
if parsed.port:
host = f"{host}:{parsed.port}"
if proxy_domain:
host = proxy_domain
else:
host = parsed.host
if parsed.port:
host = f"{host}:{parsed.port}"
scheme = "http" if parsed.scheme in ("http", "ws") else "https"
return f"{scheme}://{port}-{capsule_id}.{host}"
@ -526,7 +547,12 @@ class Capsule:
WebSocket access, see the lower-level ``_build_proxy_url``
helper or the ``pty()`` API.
"""
return _build_http_proxy_url(self._client._base_url, self._id, port)
return _build_http_proxy_url(
self._client._base_url,
self._id,
port,
self._client._proxy_domain,
)
# ── Snapshots ───────────────────────────────────────────────

View File

@ -4,7 +4,13 @@ import os
import httpx
from wrenn._config import DEFAULT_BASE_URL, ENV_API_KEY, ENV_BASE_URL
from wrenn._config import (
DEFAULT_BASE_URL,
DEFAULT_PROXY_DOMAIN,
ENV_API_KEY,
ENV_BASE_URL,
ENV_PROXY_DOMAIN,
)
from wrenn.exceptions import handle_response
from wrenn.models import (
@ -15,6 +21,7 @@ from wrenn.models import (
)
_LONG_TIMEOUT = httpx.Timeout(60.0)
_DEFAULT_TIMEOUT = httpx.Timeout(30.0, connect=10.0)
def _resolve_api_key(api_key: str | None) -> str:
@ -26,6 +33,36 @@ def _resolve_api_key(api_key: str | None) -> str:
return resolved
def _resolve_timeout(
timeout: httpx.Timeout | float | None,
) -> httpx.Timeout:
if timeout is None:
return _DEFAULT_TIMEOUT
if isinstance(timeout, httpx.Timeout):
return timeout
return httpx.Timeout(timeout)
def _resolve_proxy_domain(base_url: str, override: str | None) -> str:
"""Resolve proxy host suffix for ``{port}-{capsule_id}.<domain>`` URLs.
Precedence: explicit ``override`` arg, ``WRENN_PROXY_DOMAIN`` env, then
``wrenn.dev`` only when ``base_url`` is the default Wrenn host
(``app.wrenn.dev``). Otherwise the ``base_url`` host (with port) is used
verbatim — appropriate for local dev or custom deployments.
"""
resolved = override or os.environ.get(ENV_PROXY_DOMAIN)
if resolved:
return resolved
parsed = httpx.URL(base_url)
host = parsed.host
if host == "app.wrenn.dev":
return DEFAULT_PROXY_DOMAIN
if parsed.port:
return f"{host}:{parsed.port}"
return host
class CapsulesResource:
"""Sync capsule control-plane operations."""
@ -394,18 +431,28 @@ class WrennClient:
Args:
api_key: API key (``wrn_...``). Falls back to ``WRENN_API_KEY`` env var.
base_url: Wrenn API base URL.
proxy_domain: Host suffix for capsule proxy URLs
(``{port}-{capsule_id}.<domain>``). Falls back to
``WRENN_PROXY_DOMAIN`` env, then ``wrenn.dev`` when ``base_url``
is the default ``app.wrenn.dev`` host, else the ``base_url`` host.
timeout: HTTP timeout. Accepts ``httpx.Timeout``, a float (seconds),
or ``None`` for the default (30s read/write/pool, 10s connect).
"""
def __init__(
self,
api_key: str | None = None,
base_url: str | None = None,
proxy_domain: str | None = None,
timeout: httpx.Timeout | float | None = None,
) -> None:
self._api_key = _resolve_api_key(api_key)
self._base_url = base_url or os.environ.get(ENV_BASE_URL, DEFAULT_BASE_URL)
self._proxy_domain = _resolve_proxy_domain(self._base_url, proxy_domain)
self._http = httpx.Client(
base_url=self._base_url,
headers={"X-API-Key": self._api_key},
timeout=_resolve_timeout(timeout),
)
self.capsules = CapsulesResource(self._http)
@ -440,18 +487,28 @@ class AsyncWrennClient:
Args:
api_key: API key (``wrn_...``). Falls back to ``WRENN_API_KEY`` env var.
base_url: Wrenn API base URL. Falls back to ``WRENN_BASE_URL`` env var.
proxy_domain: Host suffix for capsule proxy URLs
(``{port}-{capsule_id}.<domain>``). Falls back to
``WRENN_PROXY_DOMAIN`` env, then ``wrenn.dev`` when ``base_url``
is the default ``app.wrenn.dev`` host, else the ``base_url`` host.
timeout: HTTP timeout. Accepts ``httpx.Timeout``, a float (seconds),
or ``None`` for the default (30s read/write/pool, 10s connect).
"""
def __init__(
self,
api_key: str | None = None,
base_url: str | None = None,
proxy_domain: str | None = None,
timeout: httpx.Timeout | float | None = None,
) -> None:
self._api_key = _resolve_api_key(api_key)
self._base_url = base_url or os.environ.get(ENV_BASE_URL, DEFAULT_BASE_URL)
self._proxy_domain = _resolve_proxy_domain(self._base_url, proxy_domain)
self._http = httpx.AsyncClient(
base_url=self._base_url,
headers={"X-API-Key": self._api_key},
timeout=_resolve_timeout(timeout),
)
self.capsules = AsyncCapsulesResource(self._http)

View File

@ -45,7 +45,12 @@ def build_execute_request(code: str) -> dict:
}
def build_ws_url(base_url: str, capsule_id: str, kernel_id: str) -> str:
def build_ws_url(
base_url: str,
capsule_id: str,
kernel_id: str,
proxy_domain: str | None = None,
) -> str:
"""Build the Jupyter kernel WebSocket URL for the given capsule."""
proxy = _build_proxy_url(base_url, capsule_id, 8888)
proxy = _build_proxy_url(base_url, capsule_id, 8888, proxy_domain)
return f"{proxy}/api/kernels/{kernel_id}/channels"

View File

@ -110,7 +110,12 @@ class AsyncCapsule(BaseAsyncCapsule):
def _get_proxy_client(self) -> httpx.AsyncClient:
if self._proxy_client is None:
url = _build_http_proxy_url(self._client._base_url, self._id, 8888)
url = _build_http_proxy_url(
self._client._base_url,
self._id,
8888,
self._client._proxy_domain,
)
self._proxy_client = httpx.AsyncClient(
base_url=url,
headers={"X-API-Key": self._client._api_key},
@ -196,7 +201,12 @@ class AsyncCapsule(BaseAsyncCapsule):
"non-Python kernelspec."
)
kernel_id = await self._ensure_kernel(jupyter_timeout=jupyter_timeout)
ws_url = build_ws_url(self._client._base_url, self._id, kernel_id)
ws_url = build_ws_url(
self._client._base_url,
self._id,
kernel_id,
self._client._proxy_domain,
)
msg = build_execute_request(code)
msg_id = msg["header"]["msg_id"]

View File

@ -138,7 +138,12 @@ class Capsule(BaseCapsule):
def _get_proxy_client(self) -> httpx.Client:
if self._proxy_client is None:
url = _build_http_proxy_url(self._client._base_url, self._id, 8888)
url = _build_http_proxy_url(
self._client._base_url,
self._id,
8888,
self._client._proxy_domain,
)
self._proxy_client = httpx.Client(
base_url=url,
headers={"X-API-Key": self._client._api_key},
@ -231,7 +236,12 @@ class Capsule(BaseCapsule):
"non-Python kernelspec."
)
kernel_id = self._ensure_kernel(jupyter_timeout=jupyter_timeout)
ws_url = build_ws_url(self._client._base_url, self._id, kernel_id)
ws_url = build_ws_url(
self._client._base_url,
self._id,
kernel_id,
self._client._proxy_domain,
)
msg = build_execute_request(code)
msg_id = msg["header"]["msg_id"]

View File

@ -45,6 +45,16 @@ class TestBuildHttpProxyUrl:
url = _build_http_proxy_url("https://api.example.com:9443", "sb-1", 80)
assert url == "https://80-sb-1.api.example.com:9443"
def test_proxy_domain_override_http(self):
url = _build_http_proxy_url(
"https://app.wrenn.dev/api", "cl-abc", 8080, "wrenn.dev"
)
assert url == "https://8080-cl-abc.wrenn.dev"
def test_proxy_domain_override_ws(self):
url = _build_proxy_url("https://app.wrenn.dev/api", "cl-abc", 8888, "wrenn.dev")
assert url == "wss://8888-cl-abc.wrenn.dev"
class TestCapsuleCreate:
@respx.mock
@ -222,7 +232,7 @@ class TestGetUrlPublic:
202, json={"id": "cl-99", "status": "starting"}
)
cap = Capsule(api_key=API_KEY, base_url=BASE)
assert cap.get_url(8080) == "https://8080-cl-99.app.wrenn.dev"
assert cap.get_url(8080) == "https://8080-cl-99.wrenn.dev"
@respx.mock
def test_sync_get_url_localhost(self):
@ -242,7 +252,7 @@ class TestGetUrlPublic:
202, json={"id": "cl-async", "status": "starting"}
)
cap = await AsyncCapsule.create(api_key=API_KEY, base_url=BASE)
assert cap.get_url(5000) == "https://5000-cl-async.app.wrenn.dev"
assert cap.get_url(5000) == "https://5000-cl-async.wrenn.dev"
await cap._client.aclose()

View File

@ -261,3 +261,39 @@ class TestAsyncClient:
)
with pytest.raises(WrennNotFoundError):
await async_client.capsules.get("nope")
class TestClientResolution:
def test_default_base_url_strips_app_subdomain(self):
with WrennClient(api_key="wrn_test1234567890abcdef12345678") as c:
assert c._proxy_domain == "wrenn.dev"
def test_custom_base_url_preserves_host(self):
with WrennClient(
api_key="wrn_test1234567890abcdef12345678",
base_url="http://localhost:8080/api",
) as c:
assert c._proxy_domain == "localhost:8080"
def test_explicit_proxy_domain_wins(self):
with WrennClient(
api_key="wrn_test1234567890abcdef12345678",
base_url="https://app.wrenn.dev/api",
proxy_domain="custom.example.com",
) as c:
assert c._proxy_domain == "custom.example.com"
def test_env_proxy_domain(self, monkeypatch):
monkeypatch.setenv("WRENN_PROXY_DOMAIN", "env.example.com")
with WrennClient(api_key="wrn_test1234567890abcdef12345678") as c:
assert c._proxy_domain == "env.example.com"
def test_default_timeout(self):
with WrennClient(api_key="wrn_test1234567890abcdef12345678") as c:
t = c._http.timeout
assert t.connect == 10.0
assert t.read == 30.0
def test_timeout_float_override(self):
with WrennClient(api_key="wrn_test1234567890abcdef12345678", timeout=5.0) as c:
assert c._http.timeout.connect == 5.0

View File

@ -263,7 +263,7 @@ class TestEnsureKernel:
@respx.mock
def test_creates_kernel_with_wrenn_name_when_none_exist(self):
c = _make_capsule()
proxy_base = "https://8888-sb-1.app.wrenn.dev"
proxy_base = "https://8888-sb-1.wrenn.dev"
list_route = respx.get(f"{proxy_base}/api/kernels").respond(200, json=[])
create_route = respx.post(f"{proxy_base}/api/kernels").respond(
201, json={"id": "k-new", "name": "wrenn"}
@ -279,7 +279,7 @@ class TestEnsureKernel:
@respx.mock
def test_reuses_existing_wrenn_kernel(self):
c = _make_capsule()
proxy_base = "https://8888-sb-1.app.wrenn.dev"
proxy_base = "https://8888-sb-1.wrenn.dev"
respx.get(f"{proxy_base}/api/kernels").respond(
200,
json=[
@ -295,7 +295,7 @@ class TestEnsureKernel:
@respx.mock
def test_creates_when_only_other_kernels_exist(self):
c = _make_capsule()
proxy_base = "https://8888-sb-1.app.wrenn.dev"
proxy_base = "https://8888-sb-1.wrenn.dev"
respx.get(f"{proxy_base}/api/kernels").respond(
200, json=[{"id": "k-other", "name": "python3"}]
)
@ -308,7 +308,7 @@ class TestEnsureKernel:
@respx.mock
def test_caches_kernel_id(self):
c = _make_capsule()
proxy_base = "https://8888-sb-1.app.wrenn.dev"
proxy_base = "https://8888-sb-1.wrenn.dev"
route = respx.get(f"{proxy_base}/api/kernels").respond(
200, json=[{"id": "k-1", "name": "wrenn"}]
)
@ -322,7 +322,7 @@ class TestEnsureKernel:
202, json={"id": "sb-1", "status": "starting"}
)
c = Capsule(kernel="python3", api_key=API_KEY, base_url=BASE)
proxy_base = "https://8888-sb-1.app.wrenn.dev"
proxy_base = "https://8888-sb-1.wrenn.dev"
respx.get(f"{proxy_base}/api/kernels").respond(200, json=[])
create = respx.post(f"{proxy_base}/api/kernels").respond(
201, json={"id": "k-py", "name": "python3"}
@ -334,7 +334,7 @@ class TestEnsureKernel:
@respx.mock
def test_retries_on_5xx_then_succeeds(self):
c = _make_capsule()
proxy_base = "https://8888-sb-1.app.wrenn.dev"
proxy_base = "https://8888-sb-1.wrenn.dev"
responses = [
httpx.Response(503),
httpx.Response(200, json=[{"id": "k-1", "name": "wrenn"}]),
@ -347,7 +347,7 @@ class TestEnsureKernel:
@respx.mock
def test_raises_on_4xx(self):
c = _make_capsule()
proxy_base = "https://8888-sb-1.app.wrenn.dev"
proxy_base = "https://8888-sb-1.wrenn.dev"
respx.get(f"{proxy_base}/api/kernels").respond(401)
with pytest.raises(httpx.HTTPStatusError):
c._ensure_kernel(jupyter_timeout=2)
@ -355,7 +355,7 @@ class TestEnsureKernel:
@respx.mock
def test_timeout_raises(self):
c = _make_capsule()
proxy_base = "https://8888-sb-1.app.wrenn.dev"
proxy_base = "https://8888-sb-1.wrenn.dev"
respx.get(f"{proxy_base}/api/kernels").respond(503)
with patch("time.sleep"):
with pytest.raises(TimeoutError):
@ -813,7 +813,7 @@ class TestAsyncEnsureKernel:
@respx.mock
async def test_async_creates_kernel_when_none_exist(self):
c = _make_async_capsule()
proxy_base = "https://8888-sb-1.app.wrenn.dev"
proxy_base = "https://8888-sb-1.wrenn.dev"
list_route = respx.get(f"{proxy_base}/api/kernels").respond(200, json=[])
create_route = respx.post(f"{proxy_base}/api/kernels").respond(
201, json={"id": "k-new", "name": "wrenn"}
@ -829,7 +829,7 @@ class TestAsyncEnsureKernel:
@respx.mock
async def test_async_reuses_existing_wrenn_kernel(self):
c = _make_async_capsule()
proxy_base = "https://8888-sb-1.app.wrenn.dev"
proxy_base = "https://8888-sb-1.wrenn.dev"
respx.get(f"{proxy_base}/api/kernels").respond(
200,
json=[
@ -847,7 +847,7 @@ class TestAsyncEnsureKernel:
@respx.mock
async def test_async_retries_on_5xx_then_succeeds(self):
c = _make_async_capsule()
proxy_base = "https://8888-sb-1.app.wrenn.dev"
proxy_base = "https://8888-sb-1.wrenn.dev"
responses = [
httpx.Response(503),
httpx.Response(200, json=[{"id": "k-1", "name": "wrenn"}]),
@ -867,7 +867,7 @@ class TestAsyncEnsureKernel:
@respx.mock
async def test_async_raises_on_4xx(self):
c = _make_async_capsule()
proxy_base = "https://8888-sb-1.app.wrenn.dev"
proxy_base = "https://8888-sb-1.wrenn.dev"
respx.get(f"{proxy_base}/api/kernels").respond(401)
with pytest.raises(httpx.HTTPStatusError):
await c._ensure_kernel(jupyter_timeout=2)
@ -877,7 +877,7 @@ class TestAsyncEnsureKernel:
@respx.mock
async def test_async_caches_kernel_id(self):
c = _make_async_capsule()
proxy_base = "https://8888-sb-1.app.wrenn.dev"
proxy_base = "https://8888-sb-1.wrenn.dev"
route = respx.get(f"{proxy_base}/api/kernels").respond(
200, json=[{"id": "k-1", "name": "wrenn"}]
)