feat(browser_pools): typed acquire helper distinguishing got / timed-out

src/kernel/lib/browser_pools/__init__.py

@@ -0,0 +1,19 @@
+from __future__ import annotations
+
+from .acquire import (
+    Acquired,
+    TimedOut,
+    PoolNotFound,
+    AcquireResult,
+    acquire,
+    acquire_async,
+)
+
+__all__ = [
+    "Acquired",
+    "TimedOut",
+    "PoolNotFound",
+    "AcquireResult",
+    "acquire",
+    "acquire_async",
+]

src/kernel/lib/browser_pools/acquire.py

@@ -0,0 +1,98 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Union
+from dataclasses import dataclass
+from typing_extensions import TypeAlias
+
+import httpx
+
+from ..._types import Body, Omit, Query, Headers, NotGiven, omit, not_given
+from ..._exceptions import NotFoundError
+from ...types.browser_pool_acquire_response import BrowserPoolAcquireResponse
+
+if TYPE_CHECKING:
+    from ..._client import Kernel, AsyncKernel
+
+
+@dataclass
+class Acquired:
+    """A browser was leased from the pool."""
+
+    browser: BrowserPoolAcquireResponse
+
+
+@dataclass
+class TimedOut:
+    """The long poll expired before a browser became available. Retry to keep waiting."""
+
+
+@dataclass
+class PoolNotFound:
+    """No pool exists with the given id or name."""
+
+
+AcquireResult: TypeAlias = Union[Acquired, TimedOut, PoolNotFound]
+
+
+def acquire(
+    client: "Kernel",
+    id_or_name: str,
+    *,
+    acquire_timeout_seconds: Union[int, Omit] = omit,
+    extra_headers: Union[Headers, None] = None,
+    extra_query: Union[Query, None] = None,
+    extra_body: Union[Body, None] = None,
+    timeout: Union[float, httpx.Timeout, None, NotGiven] = not_given,
+) -> AcquireResult:
+    """Long-polling acquire that surfaces the HTTP outcome as a typed result.
+
+    Returns one of:
+
+    * :class:`Acquired` — a browser was leased from the pool.
+    * :class:`TimedOut` — the long poll expired without a browser becoming available.
+      Retry to keep waiting.
+    * :class:`PoolNotFound` — no pool exists with the given id or name.
+
+    Other API errors (auth, server errors, etc.) still raise.
+    """
+    try:
+        raw = client.browser_pools.with_raw_response.acquire(
+            id_or_name,
+            acquire_timeout_seconds=acquire_timeout_seconds,
+            extra_headers=extra_headers,
+            extra_query=extra_query,
+            extra_body=extra_body,
+            timeout=timeout,
+        )
+    except NotFoundError:
+        return PoolNotFound()
+    if raw.http_response.status_code == 204:
+        return TimedOut()
+    return Acquired(browser=raw.parse())
+
+
+async def acquire_async(
+    client: "AsyncKernel",
+    id_or_name: str,
+    *,
+    acquire_timeout_seconds: Union[int, Omit] = omit,
+    extra_headers: Union[Headers, None] = None,
+    extra_query: Union[Query, None] = None,
+    extra_body: Union[Body, None] = None,
+    timeout: Union[float, httpx.Timeout, None, NotGiven] = not_given,
+) -> AcquireResult:
+    """Async variant of :func:`acquire`."""
+    try:
+        raw = await client.browser_pools.with_raw_response.acquire(
+            id_or_name,
+            acquire_timeout_seconds=acquire_timeout_seconds,
+            extra_headers=extra_headers,
+            extra_query=extra_query,
+            extra_body=extra_body,
+            timeout=timeout,
+        )
+    except NotFoundError:
+        return PoolNotFound()
+    if raw.http_response.status_code == 204:
+        return TimedOut()
+    return Acquired(browser=await raw.parse())

tests/test_browser_pools_typed_acquire.py

@@ -0,0 +1,68 @@
+from __future__ import annotations
+
+import os
+
+import httpx
+import respx
+import pytest
+
+from kernel import Kernel, AsyncKernel
+from kernel.lib.browser_pools import Acquired, TimedOut, PoolNotFound, acquire, acquire_async
+
+base_url = os.environ.get("TEST_API_BASE_URL", "http://127.0.0.1:4010")
+api_key = "sk-123"
+
+
+def _browser_payload() -> dict[str, object]:
+    return {
+        "session_id": "sess-1",
+        "base_url": "http://browser-session.test/browser/kernel",
+        "cdp_ws_url": "wss://browser-session.test/browser/cdp?jwt=t",
+        "webdriver_ws_url": "wss://x",
+        "created_at": "2020-01-01T00:00:00Z",
+        "headless": True,
+        "stealth": False,
+        "timeout_seconds": 60,
+    }
+
+
+@respx.mock
+def test_acquire_returns_acquired_on_200() -> None:
+    respx.post(f"{base_url}/browser_pools/my-pool/acquire").mock(
+        return_value=httpx.Response(200, json=_browser_payload())
+    )
+    with Kernel(base_url=base_url, api_key=api_key, _strict_response_validation=True) as client:
+        result = acquire(client, "my-pool")
+
+    assert isinstance(result, Acquired)
+    assert result.browser.session_id == "sess-1"
+
+
+@respx.mock
+def test_acquire_returns_timed_out_on_204() -> None:
+    respx.post(f"{base_url}/browser_pools/my-pool/acquire").mock(return_value=httpx.Response(204))
+    with Kernel(base_url=base_url, api_key=api_key, _strict_response_validation=True) as client:
+        result = acquire(client, "my-pool")
+
+    assert isinstance(result, TimedOut)
+
+
+@respx.mock
+def test_acquire_returns_pool_not_found_on_404() -> None:
+    respx.post(f"{base_url}/browser_pools/missing/acquire").mock(
+        return_value=httpx.Response(404, json={"code": "not_found", "message": "pool not found"})
+    )
+    with Kernel(base_url=base_url, api_key=api_key, _strict_response_validation=True) as client:
+        result = acquire(client, "missing")
+
+    assert isinstance(result, PoolNotFound)
+
+
+@pytest.mark.asyncio
+@respx.mock
+async def test_acquire_async_returns_timed_out_on_204() -> None:
+    respx.post(f"{base_url}/browser_pools/my-pool/acquire").mock(return_value=httpx.Response(204))
+    async with AsyncKernel(base_url=base_url, api_key=api_key, _strict_response_validation=True) as client:
+        result = await acquire_async(client, "my-pool")
+
+    assert isinstance(result, TimedOut)