Compare commits
7 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| ff6a5a0bba | |||
| 0b6b3a306f | |||
| 904261d9a1 | |||
| 3e0ee950d9 | |||
| 064eb83665 | |||
| b8f70637ba | |||
| a41ef572dd |
+185
-1
@@ -40,6 +40,10 @@ _RUNNER_VERSION = VERSION
|
||||
_RUNNER_CONFIG_HOME_ENV_VAR = "OMNIGENT_CONFIG_HOME"
|
||||
_DEFAULT_RUNNER_IDLE_TIMEOUT_S = 60 * 60
|
||||
_RUNNER_IDLE_MONITOR_MAX_POLL_INTERVAL_S = 60.0
|
||||
# Re-mint a managed runner's owner JWT this many seconds before it
|
||||
# expires, so a live session's HTTP callbacks never present an expired
|
||||
# token. Well under the server-side token TTL.
|
||||
_MANAGED_MINT_REFRESH_SKEW_S = 300.0
|
||||
_logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@@ -229,6 +233,14 @@ class _RunnerDatabricksAuth(httpx.Auth):
|
||||
if self._factory is not None:
|
||||
token = self._factory()
|
||||
if not token:
|
||||
if getattr(self._factory, "declined", False):
|
||||
# The server definitively refuses to mint for this runner
|
||||
# (managed mint factory hit HTTP 400/404 after install —
|
||||
# e.g. its construction probe lost a boot race to a
|
||||
# no-auth server). Bare requests are correct there; do
|
||||
# NOT fail closed or the runner bricks every callback.
|
||||
yield request
|
||||
return
|
||||
raise httpx.RequestError("Databricks token refresh returned no token")
|
||||
request.headers["Authorization"] = f"Bearer {token}"
|
||||
response = yield request
|
||||
@@ -382,15 +394,187 @@ def _make_auth_token_factory(
|
||||
return oidc_token
|
||||
return _sdk_token()
|
||||
|
||||
# Probe once to check if credentials are available.
|
||||
# Probe once to check if a user credential is available.
|
||||
try:
|
||||
if _factory() is not None:
|
||||
return _factory
|
||||
except (ValueError, OSError, ImportError):
|
||||
pass
|
||||
|
||||
# Managed-sandbox fallback: no user credential resolved (no stored
|
||||
# OIDC token, no Databricks config), but a managed runner still holds
|
||||
# its tunnel binding token. Authenticate its HTTP callbacks (and the
|
||||
# tunnel bearer) with a short-lived owner JWT the server mints against
|
||||
# that binding token — refreshed on demand, so there is no static
|
||||
# credential at rest and no fixed session-length cap.
|
||||
if resolved_server_url:
|
||||
try:
|
||||
binding_token = _runner_tunnel_binding_token_from_env()
|
||||
except RuntimeError:
|
||||
binding_token = None
|
||||
if binding_token is not None:
|
||||
return _make_managed_mint_factory(resolved_server_url, binding_token)
|
||||
return None
|
||||
|
||||
|
||||
def _make_managed_mint_factory(
|
||||
server_url: str,
|
||||
binding_token: str,
|
||||
) -> Callable[[], str | None] | None:
|
||||
"""Build a token factory that mints a managed runner's owner JWT.
|
||||
|
||||
For a server-managed sandbox runner with no user credential of its
|
||||
own: mint a short-lived owner JWT from ``POST /v1/runners/{id}/token``,
|
||||
authenticated by the runner's tunnel binding token, and cache it in
|
||||
memory. The cached token is reused until it nears expiry, then
|
||||
re-minted — so a managed session runs arbitrarily long without its
|
||||
auth expiring (no fixed session-length cap), and no long-lived
|
||||
credential is ever written to the sandbox environment.
|
||||
|
||||
The same factory feeds both the WS tunnel bearer and the httpx
|
||||
callback client (see :func:`_make_auth_token_factory` callers), so one
|
||||
credential authenticates every runner->server surface.
|
||||
|
||||
:param server_url: Omnigent server base URL, e.g.
|
||||
``"https://omnigent.example.com"``.
|
||||
:param binding_token: The runner's tunnel binding token (the sandbox's
|
||||
only credential), presented to the mint endpoint.
|
||||
:returns: A sync callable returning a fresh owner JWT, or ``None`` only
|
||||
when the server *definitively* will not mint for this runner (HTTP
|
||||
400 no-auth/header mode, or 404 older server without the endpoint) —
|
||||
the runner then sends unauthenticated requests, as it did before this
|
||||
fallback existed. A *transient* probe failure still installs the
|
||||
factory, which re-mints on the next callback (so a blip at boot does
|
||||
not leave the runner unauthenticated until process restart). If such
|
||||
a post-install mint then gets the definitive 400/404, the factory
|
||||
latches ``declined`` and returns ``None`` thereafter, and
|
||||
:class:`_RunnerDatabricksAuth` falls back to bare requests.
|
||||
"""
|
||||
from omnigent.runner.identity import token_bound_runner_id
|
||||
|
||||
runner_id = token_bound_runner_id(binding_token)
|
||||
mint_url = f"{server_url.rstrip('/')}/v1/runners/{runner_id}/token"
|
||||
|
||||
# Construction probe. Decline to install the factory ONLY when the
|
||||
# server definitively will not mint for this runner — HTTP 400 (no auth
|
||||
# provider / header mode) or 404 (an older server without the endpoint).
|
||||
# There the runner falls back to bare requests, which are correct on a
|
||||
# no-auth server. Every other outcome installs the factory: a success
|
||||
# seeds the cache; a transient failure (network blip, 5xx, timeout)
|
||||
# installs it anyway so the next callback re-mints, rather than leaving
|
||||
# the runner unauthenticated until process restart.
|
||||
factory = _ManagedMintTokenFactory(mint_url, server_url, binding_token)
|
||||
factory()
|
||||
if factory.declined:
|
||||
return None
|
||||
return factory
|
||||
|
||||
|
||||
class _ManagedMintTokenFactory:
|
||||
"""Callable that mints (and caches) a managed runner's owner JWT.
|
||||
|
||||
Each call returns the cached JWT until it nears expiry, then re-mints
|
||||
via :func:`_mint_managed_owner_token`. When a mint gets a *definitive*
|
||||
refusal (HTTP 400 no-auth/header mode, 404 older server), the
|
||||
:attr:`declined` latch is set and every subsequent call returns
|
||||
``None`` without touching the network —
|
||||
:meth:`_RunnerDatabricksAuth.auth_flow` reads the latch to send bare
|
||||
requests instead of failing closed. The latch matters when the
|
||||
construction probe loses a boot race (a connection error installs the
|
||||
factory, then the first real mint learns the server never mints).
|
||||
"""
|
||||
|
||||
def __init__(self, mint_url: str, server_url: str, binding_token: str) -> None:
|
||||
"""
|
||||
:param mint_url: Fully-qualified ``/v1/runners/{id}/token`` URL.
|
||||
:param server_url: Omnigent server base URL.
|
||||
:param binding_token: The runner's tunnel binding token.
|
||||
"""
|
||||
self._mint_url = mint_url
|
||||
self._server_url = server_url
|
||||
self._binding_token = binding_token
|
||||
self._cached_token: str | None = None
|
||||
self._cached_expires_at = 0.0
|
||||
self.declined = False
|
||||
|
||||
def __call__(self) -> str | None:
|
||||
"""Return a fresh owner JWT, or ``None``.
|
||||
|
||||
:returns: The cached or freshly-minted JWT; ``None`` after a
|
||||
definitive server decline (sets :attr:`declined`) or on a
|
||||
transient mint failure with no still-valid cached token.
|
||||
"""
|
||||
if self.declined:
|
||||
return None
|
||||
now = time.time()
|
||||
if (
|
||||
self._cached_token is not None
|
||||
and now < self._cached_expires_at - _MANAGED_MINT_REFRESH_SKEW_S
|
||||
):
|
||||
return self._cached_token
|
||||
try:
|
||||
token, expires_at = _mint_managed_owner_token(
|
||||
self._mint_url, self._server_url, self._binding_token
|
||||
)
|
||||
except httpx.HTTPStatusError as exc:
|
||||
if exc.response.status_code in (400, 404):
|
||||
self.declined = True
|
||||
return None
|
||||
return self._still_valid_cached_token(now)
|
||||
except (httpx.HTTPError, ValueError, KeyError, OSError):
|
||||
# Transient mint failure: keep serving the cached token while
|
||||
# it is still valid; otherwise report "no token" and let the
|
||||
# tunnel's / HTTP client's on-401 retry drive the next mint.
|
||||
return self._still_valid_cached_token(now)
|
||||
self._cached_token = token
|
||||
self._cached_expires_at = expires_at
|
||||
return token
|
||||
|
||||
def _still_valid_cached_token(self, now: float) -> str | None:
|
||||
"""Return the cached token if it hasn't expired outright.
|
||||
|
||||
:param now: Current epoch seconds.
|
||||
:returns: The cached token while still valid, else ``None``.
|
||||
"""
|
||||
if self._cached_token is not None and now < self._cached_expires_at:
|
||||
return self._cached_token
|
||||
return None
|
||||
|
||||
|
||||
def _mint_managed_owner_token(
|
||||
mint_url: str,
|
||||
server_url: str,
|
||||
binding_token: str,
|
||||
) -> tuple[str, float]:
|
||||
"""Mint one managed-runner owner JWT from the server.
|
||||
|
||||
:param mint_url: Fully-qualified ``/v1/runners/{id}/token`` URL.
|
||||
:param server_url: Server base URL, used for the Databricks workspace
|
||||
routing header (``X-Databricks-Org-Id``) when applicable.
|
||||
:param binding_token: The runner's tunnel binding token, sent as the
|
||||
``X-Omnigent-Runner-Tunnel-Token`` header to authenticate the mint.
|
||||
:returns: ``(jwt, expires_at_epoch_seconds)``.
|
||||
:raises httpx.HTTPError: On network failure or a non-2xx response.
|
||||
:raises KeyError: If the response is missing the expected fields.
|
||||
"""
|
||||
from omnigent.cli_auth import databricks_request_headers
|
||||
from omnigent.runner.identity import (
|
||||
OMNIGENT_INTERNAL_WS_ORIGIN,
|
||||
RUNNER_TUNNEL_TOKEN_HEADER,
|
||||
)
|
||||
|
||||
headers = {
|
||||
"Origin": OMNIGENT_INTERNAL_WS_ORIGIN,
|
||||
RUNNER_TUNNEL_TOKEN_HEADER: binding_token,
|
||||
**databricks_request_headers(server_url),
|
||||
}
|
||||
with httpx.Client(timeout=10.0) as client:
|
||||
response = client.post(mint_url, headers=headers)
|
||||
response.raise_for_status()
|
||||
payload = response.json()
|
||||
return payload["token"], float(payload["expires_at"])
|
||||
|
||||
|
||||
def _runner_tunnel_binding_token_from_env() -> str | None:
|
||||
"""Return the optional tunnel binding token from the environment.
|
||||
|
||||
|
||||
@@ -246,6 +246,27 @@ class AuthProvider(ABC):
|
||||
"""Return the authenticated user ID, or ``None``."""
|
||||
...
|
||||
|
||||
def mint_runner_token(self, user_id: str, ttl_seconds: int) -> str | None: # noqa: ARG002
|
||||
"""
|
||||
Mint a short-lived bearer a managed-sandbox runner presents as *user_id*.
|
||||
|
||||
A managed runner runs in a sandbox with no logged-in user
|
||||
credential of its own, so the server mints one for its HTTP
|
||||
callbacks when auth is enabled (see the
|
||||
``POST /v1/runners/{id}/token`` endpoint). Default: ``None`` — no
|
||||
minting (single-user / no-auth, or a provider whose identity is
|
||||
asserted externally and can't be minted server-side, e.g.
|
||||
header/proxy auth). The runner then authenticates with its tunnel
|
||||
binding token alone.
|
||||
|
||||
:param user_id: The session owner the runner acts as, e.g.
|
||||
``"alice@example.com"``.
|
||||
:param ttl_seconds: Token lifetime in seconds.
|
||||
:returns: A bearer token string, or ``None`` when this provider
|
||||
cannot mint one.
|
||||
"""
|
||||
return None
|
||||
|
||||
|
||||
class UnifiedAuthProvider(AuthProvider):
|
||||
"""Unified authentication provider that supports header-based,
|
||||
@@ -348,6 +369,38 @@ class UnifiedAuthProvider(AuthProvider):
|
||||
return self._check_cookie(request)
|
||||
return self._check_header(request)
|
||||
|
||||
def mint_runner_token(self, user_id: str, ttl_seconds: int) -> str | None:
|
||||
"""
|
||||
Mint a short-lived owner JWT for a managed-sandbox runner.
|
||||
|
||||
Accounts / OIDC modes sign a session JWT in the same HS256 format
|
||||
:meth:`_check_cookie` validates, so the runner can present it as
|
||||
``Authorization: Bearer <jwt>`` on its HTTP callbacks and resolve
|
||||
to *user_id*. Header/proxy mode returns ``None`` — identity there
|
||||
is asserted by the upstream proxy and can't be minted server-side.
|
||||
|
||||
:param user_id: The session owner the runner acts as, e.g.
|
||||
``"alice@example.com"``.
|
||||
:param ttl_seconds: Token lifetime in seconds.
|
||||
:returns: An HS256-signed JWT, or ``None`` for header mode, an
|
||||
empty/reserved user, or a missing cookie config.
|
||||
"""
|
||||
if not user_id or user_id in _RESERVED_USERS:
|
||||
return None
|
||||
if self._source not in ("oidc", "accounts"):
|
||||
return None
|
||||
cookie_config = self._oidc_config if self._source == "oidc" else self._accounts_config
|
||||
if cookie_config is None:
|
||||
return None
|
||||
from omnigent.server.oidc import mint_session_token
|
||||
|
||||
return mint_session_token(
|
||||
user_id,
|
||||
cookie_config.cookie_secret,
|
||||
ttl_seconds,
|
||||
self._source,
|
||||
)
|
||||
|
||||
def _check_cookie(self, request: HTTPConnection) -> str | None:
|
||||
"""Validate the session cookie or Bearer token and return the
|
||||
user ID.
|
||||
|
||||
+34
-8
@@ -47,6 +47,39 @@ def derive_code_challenge(code_verifier: str) -> str:
|
||||
# ── Session cookie helpers ───────────────────────────────────────
|
||||
|
||||
|
||||
def mint_session_token(
|
||||
user_id: str,
|
||||
cookie_secret: bytes,
|
||||
ttl_seconds: int,
|
||||
provider: str,
|
||||
) -> str:
|
||||
"""
|
||||
Mint a signed session JWT with a second-granularity lifetime.
|
||||
|
||||
The seconds-based core behind :func:`mint_session_cookie`. A managed
|
||||
runner needs a short-lived (sub-hour) owner token, which the hours-only
|
||||
cookie helper cannot express; both share this HS256 claim shape so the
|
||||
same validator (:meth:`UnifiedAuthProvider._check_cookie`) accepts
|
||||
either.
|
||||
|
||||
:param user_id: The authenticated user's email, e.g.
|
||||
``"alice@example.com"``.
|
||||
:param cookie_secret: HMAC key for HS256 signing.
|
||||
:param ttl_seconds: Token lifetime in seconds.
|
||||
:param provider: Identity provider name, e.g. ``"google"`` or
|
||||
``"accounts"``. Stored as an informational claim.
|
||||
:returns: An HS256-signed JWT string.
|
||||
"""
|
||||
now = int(time.time())
|
||||
payload = {
|
||||
"sub": user_id,
|
||||
"iat": now,
|
||||
"exp": now + ttl_seconds,
|
||||
"provider": provider,
|
||||
}
|
||||
return jwt.encode(payload, cookie_secret, algorithm="HS256")
|
||||
|
||||
|
||||
def mint_session_cookie(
|
||||
user_id: str,
|
||||
cookie_secret: bytes,
|
||||
@@ -63,14 +96,7 @@ def mint_session_cookie(
|
||||
or ``"github"``. Stored as an informational claim.
|
||||
:returns: An HS256-signed JWT string.
|
||||
"""
|
||||
now = int(time.time())
|
||||
payload = {
|
||||
"sub": user_id,
|
||||
"iat": now,
|
||||
"exp": now + (ttl_hours * 3600),
|
||||
"provider": provider,
|
||||
}
|
||||
return jwt.encode(payload, cookie_secret, algorithm="HS256")
|
||||
return mint_session_token(user_id, cookie_secret, ttl_hours * 3600, provider)
|
||||
|
||||
|
||||
def hmac_digest(token: str, secret: bytes) -> str:
|
||||
|
||||
@@ -23,6 +23,7 @@ from ipaddress import ip_address
|
||||
|
||||
from fastapi import APIRouter, Request, WebSocket, WebSocketDisconnect
|
||||
|
||||
from omnigent.errors import ErrorCode, OmnigentError
|
||||
from omnigent.runner.identity import RUNNER_TUNNEL_TOKEN_HEADER, token_bound_runner_id
|
||||
from omnigent.runner.transports.ws_tunnel.frames import (
|
||||
PingFrame,
|
||||
@@ -45,6 +46,12 @@ PING_MISS_THRESHOLD = 3
|
||||
RUNNER_ID_MISMATCH_CLOSE_CODE = 4004
|
||||
_ON_RUNNER_CONNECT_TIMEOUT_SEC = 30.0
|
||||
|
||||
# Lifetime of a managed runner's minted owner bearer (POST
|
||||
# /v1/runners/{id}/token). Short by design: the runner re-mints on demand
|
||||
# via its token factory, so a compromised sandbox's credential is usable
|
||||
# only briefly, while a live session refreshes indefinitely with no cap.
|
||||
_MANAGED_RUNNER_TOKEN_TTL_S = 1800
|
||||
|
||||
|
||||
def _is_loopback_websocket_client(ws: WebSocket) -> bool:
|
||||
"""Return whether the WebSocket peer is a loopback client.
|
||||
@@ -274,6 +281,66 @@ def create_runner_tunnel_router(
|
||||
result["error"] = error
|
||||
return result
|
||||
|
||||
@router.post("/runners/{runner_id}/token")
|
||||
async def mint_runner_owner_token(request: Request, runner_id: str) -> dict[str, str | int]:
|
||||
"""Mint a short-lived owner bearer for a managed-sandbox runner.
|
||||
|
||||
A managed sandbox runner has no user credential of its own; it
|
||||
presents its server-minted tunnel binding token
|
||||
(``X-Omnigent-Runner-Tunnel-Token``) and the server returns a
|
||||
short-lived owner JWT the runner then uses on its HTTP callbacks
|
||||
(which gate on ``require_user``). This is the HTTP analog of the
|
||||
runner tunnel's binding-token handshake: the same SHA-256 gate
|
||||
(``token_bound_runner_id(token) == runner_id``) and the same
|
||||
owner resolution (``resolve_managed_runner_owner``), minting a
|
||||
bearer instead of registering a tunnel.
|
||||
|
||||
The binding-token match is required unconditionally — the
|
||||
allow-list shortcut honored on some other runner-token checks is
|
||||
deliberately NOT accepted here, because this endpoint issues a
|
||||
full owner credential and managed sandboxes always run
|
||||
token-bound (no allow-list).
|
||||
|
||||
:param request: The incoming FastAPI request (carries the binding
|
||||
token header).
|
||||
:param runner_id: Token-bound runner id from the path.
|
||||
:returns: ``{"token": <jwt>, "expires_at": <epoch seconds>}``.
|
||||
:raises OmnigentError: 401 when the binding token is absent,
|
||||
doesn't match ``runner_id``, or resolves to no managed-launch
|
||||
owner; 400 when the active auth mode can't mint server-side
|
||||
(header/proxy, or no auth provider).
|
||||
"""
|
||||
if auth_provider is None:
|
||||
# No auth configured: the runner authenticates by binding
|
||||
# token alone and needs no bearer — minting is meaningless.
|
||||
raise OmnigentError(
|
||||
"managed-runner token minting requires an auth provider",
|
||||
code=ErrorCode.INVALID_INPUT,
|
||||
)
|
||||
token = (request.headers.get(RUNNER_TUNNEL_TOKEN_HEADER) or "").strip()
|
||||
if not token or token_bound_runner_id(token) != runner_id:
|
||||
raise OmnigentError("unauthenticated", code=ErrorCode.UNAUTHORIZED)
|
||||
owner: str | None = None
|
||||
if resolve_managed_runner_owner is not None:
|
||||
owner = await asyncio.to_thread(resolve_managed_runner_owner, runner_id)
|
||||
if owner is None:
|
||||
# No managed-launch record bound to this runner id: a peer
|
||||
# with a syntactically valid but unrecognized token. Refuse,
|
||||
# the same fail-closed posture as the tunnel handshake.
|
||||
raise OmnigentError("unauthenticated", code=ErrorCode.UNAUTHORIZED)
|
||||
bearer = auth_provider.mint_runner_token(owner, _MANAGED_RUNNER_TOKEN_TTL_S)
|
||||
if bearer is None:
|
||||
# oidc/accounts mint; header/proxy mode can't (identity is
|
||||
# asserted upstream). Signal clearly rather than 401.
|
||||
raise OmnigentError(
|
||||
"managed-runner token minting is unsupported in this auth mode",
|
||||
code=ErrorCode.INVALID_INPUT,
|
||||
)
|
||||
return {
|
||||
"token": bearer,
|
||||
"expires_at": int(time.time()) + _MANAGED_RUNNER_TOKEN_TTL_S,
|
||||
}
|
||||
|
||||
@router.websocket("/runners/{runner_id}/tunnel")
|
||||
async def tunnel(ws: WebSocket, runner_id: str) -> None:
|
||||
"""Accept a runner's outbound WebSocket tunnel.
|
||||
|
||||
@@ -6919,6 +6919,61 @@
|
||||
]
|
||||
}
|
||||
},
|
||||
"/v1/runners/{runner_id}/token": {
|
||||
"post": {
|
||||
"description": "Mint a short-lived owner bearer for a managed-sandbox runner.\n\nA managed sandbox runner has no user credential of its own; it\npresents its server-minted tunnel binding token\n(`X-Omnigent-Runner-Tunnel-Token`) and the server returns a\nshort-lived owner JWT the runner then uses on its HTTP callbacks\n(which gate on `require_user`). This is the HTTP analog of the\nrunner tunnel's binding-token handshake: the same SHA-256 gate\n(`token_bound_runner_id(token) == runner_id`) and the same\nowner resolution (`resolve_managed_runner_owner`), minting a\nbearer instead of registering a tunnel.\n\nThe binding-token match is required unconditionally \u2014 the\nallow-list shortcut honored on some other runner-token checks is\ndeliberately NOT accepted here, because this endpoint issues a\nfull owner credential and managed sandboxes always run\ntoken-bound (no allow-list).\n\n**Returns:** `{\"token\": <jwt>, \"expires_at\": <epoch seconds>}`.\n\n**Raises**\n\n- `OmnigentError` \u2014 401 when the binding token is absent, doesn't match `runner_id`, or resolves to no managed-launch owner; 400 when the active auth mode can't mint server-side (header/proxy, or no auth provider).",
|
||||
"operationId": "mint_runner_owner_token_v1_runners__runner_id__token_post",
|
||||
"parameters": [
|
||||
{
|
||||
"description": "Token-bound runner id from the path.",
|
||||
"in": "path",
|
||||
"name": "runner_id",
|
||||
"required": true,
|
||||
"schema": {
|
||||
"title": "Runner Id",
|
||||
"type": "string"
|
||||
}
|
||||
}
|
||||
],
|
||||
"responses": {
|
||||
"200": {
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"additionalProperties": {
|
||||
"anyOf": [
|
||||
{
|
||||
"type": "string"
|
||||
},
|
||||
{
|
||||
"type": "integer"
|
||||
}
|
||||
]
|
||||
},
|
||||
"title": "Response Mint Runner Owner Token V1 Runners Runner Id Token Post",
|
||||
"type": "object"
|
||||
}
|
||||
}
|
||||
},
|
||||
"description": "Successful Response"
|
||||
},
|
||||
"422": {
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"$ref": "#/components/schemas/HTTPValidationError"
|
||||
}
|
||||
}
|
||||
},
|
||||
"description": "Validation Error"
|
||||
}
|
||||
},
|
||||
"summary": "Mint Runner Owner Token",
|
||||
"tags": [
|
||||
"runners"
|
||||
]
|
||||
}
|
||||
},
|
||||
"/v1/sessions": {
|
||||
"get": {
|
||||
"description": "List sessions with cursor-based pagination.\n\nSessions are conversations with a non-`None` `agent_id`\n\u2014 i.e. those created via `POST /v1/sessions`.\nConversations without an agent binding are excluded.\n\n**Returns:** A `PaginatedList` of `SessionListItem`.",
|
||||
|
||||
@@ -0,0 +1,274 @@
|
||||
"""End-to-end proof of the managed-sandbox runner HTTP-auth fix (#357 HTTP half).
|
||||
|
||||
A server-managed sandbox runner has no user credential of its own — only its
|
||||
tunnel binding token. Under accounts/OIDC auth every runner->server HTTP
|
||||
callback gates on ``require_user``, so before this fix those callbacks went out
|
||||
bare and 401'd (the runner connected its tunnel but could never fetch its own
|
||||
agent spec). The Option-C fix has the runner mint a short-lived owner JWT from
|
||||
``POST /v1/runners/{id}/token`` (authenticated by its binding token) and present
|
||||
it as ``Authorization: Bearer`` on every callback.
|
||||
|
||||
This test proves that fix works against a **real** ``omnigent server``
|
||||
subprocess with accounts auth enabled, driving the runner's **real** outbound
|
||||
code over a **real** TCP socket — no transports are stubbed:
|
||||
|
||||
* ``_make_auth_token_factory`` -> ``_make_managed_mint_factory`` ->
|
||||
``_mint_managed_owner_token`` (a real ``httpx`` POST to the mint endpoint), and
|
||||
* ``_RunnerDatabricksAuth`` on a real ``httpx.AsyncClient`` GET.
|
||||
|
||||
The only thing simulated is the managed-sandbox *condition* — no ``omnigent
|
||||
login`` token and no Databricks config on disk — which is exactly what makes the
|
||||
fix necessary (and what a fresh sandbox actually looks like).
|
||||
|
||||
The differential is asserted in one test, so the mint is provably the cause of
|
||||
the flip:
|
||||
|
||||
* WITHOUT a minted token (``_RunnerDatabricksAuth(None)`` — precisely what
|
||||
``_make_auth_token_factory`` returns for a managed sandbox on ``main``):
|
||||
``GET /v1/sessions/{id}/agent/contents`` -> **401**.
|
||||
* WITH the fix: the same GET -> **200**, returning the agent bundle.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import signal
|
||||
import subprocess
|
||||
import time
|
||||
from collections.abc import Iterator
|
||||
from pathlib import Path
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
from omnigent.runner._entry import _make_auth_token_factory, _RunnerDatabricksAuth
|
||||
from omnigent.runner.identity import (
|
||||
OMNIGENT_INTERNAL_WS_ORIGIN,
|
||||
RUNNER_TUNNEL_BINDING_TOKEN_ENV_VAR,
|
||||
token_bound_runner_id,
|
||||
)
|
||||
from omnigent.server.oidc import mint_session_cookie
|
||||
from omnigent.stores.conversation_store.sqlalchemy_store import (
|
||||
SqlAlchemyConversationStore,
|
||||
)
|
||||
from tests._helpers.compat import apply_server_env, compat_server_cwd, server_executable
|
||||
from tests._helpers.live_server import find_free_port
|
||||
from tests.server.helpers import build_agent_bundle
|
||||
|
||||
# Repo root — this file lives at tests/e2e/<name>.py.
|
||||
_REPO_ROOT = Path(__file__).resolve().parents[2]
|
||||
|
||||
# 32-byte cookie secret (64 hex chars), shared between this test process and the
|
||||
# server subprocess so (a) the accounts cookie we mint for the owner validates
|
||||
# server-side and (b) the JWT the mint endpoint signs validates the same way.
|
||||
_COOKIE_SECRET_HEX = "a1b2c3d4e5f6a7b8c9d0e1f2a3b4c5d6e7f8a9b0c1d2e3f4a5b6c7d8e9f0a1b2"
|
||||
_OWNER = "alice@example.com"
|
||||
_BINDING_TOKEN = "e2e-managed-sandbox-binding-token"
|
||||
_SERVER_HEALTH_TIMEOUT_S = 40.0
|
||||
|
||||
|
||||
def _await_health(base_url: str, log_path: Path) -> None:
|
||||
"""Poll ``/health`` until the server answers 200, or fail with the log tail.
|
||||
|
||||
:param base_url: Server base URL, e.g. ``"http://localhost:58123"``.
|
||||
:param log_path: Server stdout/stderr log, tailed into the failure message.
|
||||
:returns: None.
|
||||
:raises RuntimeError: If the server doesn't answer within the deadline.
|
||||
"""
|
||||
deadline = time.monotonic() + _SERVER_HEALTH_TIMEOUT_S
|
||||
while time.monotonic() < deadline:
|
||||
try:
|
||||
if httpx.get(f"{base_url}/health", timeout=2).status_code == 200:
|
||||
return
|
||||
except httpx.HTTPError:
|
||||
# Expected while the server is still booting (connection refused /
|
||||
# reset); keep polling until the deadline.
|
||||
pass
|
||||
time.sleep(0.5)
|
||||
tail = log_path.read_text()[-3000:] if log_path.exists() else "(no log)"
|
||||
raise RuntimeError(f"accounts server did not become healthy. Log:\n{tail}")
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def accounts_server(tmp_path: Path) -> Iterator[tuple[str, str]]:
|
||||
"""Run a real ``omnigent server`` subprocess with accounts auth enabled.
|
||||
|
||||
Accounts mode is selected by ``OMNIGENT_AUTH_PROVIDER=accounts`` plus a
|
||||
shared cookie secret; the subprocess handles the full runtime lifecycle
|
||||
(migrations, DBOS, auth provider, permission store) exactly as a deployed
|
||||
server does. ``_maybe_prompt_first_admin`` no-ops without a TTY, so the
|
||||
boot is non-interactive.
|
||||
|
||||
Deliberately independent of the session-scoped ``live_server`` fixture
|
||||
(which spawns a server + runner pair for the harness matrix): this test
|
||||
needs an accounts-auth server and no LLM, so it owns its own subprocess.
|
||||
|
||||
:param tmp_path: Per-test temp dir for the DB, artifacts, and server log.
|
||||
:returns: ``(base_url, db_uri)`` — the running server's URL and the SQLite
|
||||
URI the test opens directly to bind the managed runner id.
|
||||
"""
|
||||
port = find_free_port()
|
||||
db_path = tmp_path / "e2e.db"
|
||||
db_uri = f"sqlite:///{db_path}"
|
||||
artifact_dir = tmp_path / "artifacts"
|
||||
artifact_dir.mkdir()
|
||||
log_path = tmp_path / "server.log"
|
||||
base_url = f"http://localhost:{port}"
|
||||
|
||||
env = {**os.environ}
|
||||
env["OMNIGENT_AUTH_PROVIDER"] = "accounts"
|
||||
env["OMNIGENT_ACCOUNTS_COOKIE_SECRET"] = _COOKIE_SECRET_HEX
|
||||
env["OMNIGENT_ACCOUNTS_BASE_URL"] = base_url
|
||||
# Force the accounts branch of the auth-source switch (an ambient OIDC
|
||||
# issuer in the environment would otherwise select oidc mode).
|
||||
env.pop("OMNIGENT_OIDC_ISSUER", None)
|
||||
# Import the server package from this worktree, not an installed copy.
|
||||
apply_server_env(env, _REPO_ROOT)
|
||||
|
||||
log_handle = open(log_path, "w") # noqa: SIM115 — handle lives for the subprocess
|
||||
proc = subprocess.Popen(
|
||||
[
|
||||
server_executable(),
|
||||
"-m",
|
||||
"omnigent.cli",
|
||||
"server",
|
||||
"--port",
|
||||
str(port),
|
||||
"--database-uri",
|
||||
db_uri,
|
||||
"--artifact-location",
|
||||
str(artifact_dir),
|
||||
],
|
||||
env=env,
|
||||
cwd=compat_server_cwd(),
|
||||
stdout=log_handle,
|
||||
stderr=subprocess.STDOUT,
|
||||
)
|
||||
try:
|
||||
_await_health(base_url, log_path)
|
||||
yield base_url, db_uri
|
||||
finally:
|
||||
proc.send_signal(signal.SIGTERM)
|
||||
try:
|
||||
proc.wait(timeout=10)
|
||||
except subprocess.TimeoutExpired:
|
||||
proc.kill()
|
||||
proc.wait(timeout=5)
|
||||
log_handle.close()
|
||||
|
||||
|
||||
async def _get_agent_contents(
|
||||
base_url: str,
|
||||
path: str,
|
||||
auth: _RunnerDatabricksAuth,
|
||||
) -> httpx.Response:
|
||||
"""Drive the runner's real callback client for one ``GET``.
|
||||
|
||||
Builds the same ``httpx.AsyncClient`` the runner uses for its server
|
||||
callbacks (``auth=_RunnerDatabricksAuth(...)``, sentinel ``Origin``,
|
||||
redirects off) and issues a single request over a real socket.
|
||||
|
||||
:param base_url: Live server base URL.
|
||||
:param path: Request path, e.g. ``"/v1/sessions/<id>/agent/contents"``.
|
||||
:param auth: The runner's httpx auth — ``_RunnerDatabricksAuth(None)`` for
|
||||
the bare (pre-fix) leg, or one wired to the managed-mint factory.
|
||||
:returns: The HTTP response.
|
||||
"""
|
||||
async with httpx.AsyncClient(
|
||||
base_url=base_url,
|
||||
auth=auth,
|
||||
headers={"Origin": OMNIGENT_INTERNAL_WS_ORIGIN},
|
||||
follow_redirects=False,
|
||||
timeout=30.0,
|
||||
) as client:
|
||||
return await client.get(path)
|
||||
|
||||
|
||||
def test_managed_runner_callback_authenticates_end_to_end(
|
||||
accounts_server: tuple[str, str],
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""A managed runner's HTTP callback 401s bare and 200s with a minted token.
|
||||
|
||||
End-to-end against a live accounts-auth server, driving the runner's real
|
||||
outbound auth code over a real socket. Reverting the runner-side managed
|
||||
mint tier (or the server-side mint endpoint) turns the 200 assertion back
|
||||
into a 401 — that is the exact gap this change closes.
|
||||
|
||||
:param accounts_server: ``(base_url, db_uri)`` from the live-server fixture.
|
||||
:param monkeypatch: Puts this process into the managed-sandbox posture
|
||||
(binding token + server URL present; no user credential resolvable).
|
||||
:returns: None.
|
||||
"""
|
||||
base_url, db_uri = accounts_server
|
||||
|
||||
# 1. Alice owns a real session. Her identity comes from a directly-minted
|
||||
# accounts cookie signed with the server's shared secret — the same JWT
|
||||
# the password login flow issues; only the password dance is skipped.
|
||||
# The session-create, agent registration, and owner grant are all real.
|
||||
owner_cookie = mint_session_cookie(_OWNER, bytes.fromhex(_COOKIE_SECRET_HEX), 8, "accounts")
|
||||
bundle = build_agent_bundle(name="e2e-managed-runner-agent")
|
||||
with httpx.Client(base_url=base_url, timeout=30.0) as http:
|
||||
create = http.post(
|
||||
"/v1/sessions",
|
||||
headers={
|
||||
"Authorization": f"Bearer {owner_cookie}",
|
||||
"Origin": OMNIGENT_INTERNAL_WS_ORIGIN,
|
||||
},
|
||||
data={"metadata": "{}"},
|
||||
files={"bundle": ("agent.tar.gz", bundle, "application/gzip")},
|
||||
)
|
||||
assert create.status_code in (200, 201), (create.status_code, create.text)
|
||||
session_id = create.json()["session_id"]
|
||||
|
||||
# 2. Bind a managed runner id to Alice's session — what the managed-launch
|
||||
# path does at spawn time via replace_runner_id. WAL journaling + a 20s
|
||||
# busy_timeout make this cross-process write safe against the running
|
||||
# server, which then resolves runner_id -> owner from this row.
|
||||
runner_id = token_bound_runner_id(_BINDING_TOKEN)
|
||||
SqlAlchemyConversationStore(db_uri).replace_runner_id(session_id, runner_id)
|
||||
|
||||
# 3. Put this process in a managed-sandbox posture: the runner holds ONLY
|
||||
# its binding token and the server URL — no omnigent-login token, no
|
||||
# Databricks config. Forcing both credential sources to miss is what a
|
||||
# fresh sandbox actually is, and it routes _make_auth_token_factory to
|
||||
# the managed-mint tier under test.
|
||||
from omnigent.inner.databricks_executor import DatabricksAuthError
|
||||
|
||||
def _no_databricks_creds(*args: object, **kwargs: object) -> tuple[object, str]:
|
||||
"""Stand in for _resolve_databricks_auth in a credential-less sandbox."""
|
||||
raise DatabricksAuthError("managed sandbox has no Databricks config")
|
||||
|
||||
monkeypatch.setenv("RUNNER_SERVER_URL", base_url)
|
||||
monkeypatch.setenv(RUNNER_TUNNEL_BINDING_TOKEN_ENV_VAR, _BINDING_TOKEN)
|
||||
monkeypatch.setattr("omnigent.cli_auth.load_token", lambda _url: None)
|
||||
monkeypatch.setattr(
|
||||
"omnigent.inner.databricks_executor._resolve_databricks_auth",
|
||||
_no_databricks_creds,
|
||||
)
|
||||
|
||||
contents_path = f"/v1/sessions/{session_id}/agent/contents"
|
||||
|
||||
# 4a. Pre-fix behavior: a managed sandbox on main resolves no credential,
|
||||
# so _make_auth_token_factory returns None and the callback goes out
|
||||
# with no bearer. require_user rejects it.
|
||||
bare = asyncio.run(_get_agent_contents(base_url, contents_path, _RunnerDatabricksAuth(None)))
|
||||
assert bare.status_code == 401, (bare.status_code, bare.text)
|
||||
|
||||
# 4b. With the fix: the factory installs the managed-mint tier, mints a
|
||||
# short-lived owner JWT from the binding token (a real POST to the mint
|
||||
# endpoint at construction), and presents it on the callback.
|
||||
factory = _make_auth_token_factory()
|
||||
assert factory is not None, (
|
||||
"managed-mint factory should install for a managed sandbox "
|
||||
"(binding token + server URL present, no user credential)"
|
||||
)
|
||||
authed = asyncio.run(
|
||||
_get_agent_contents(base_url, contents_path, _RunnerDatabricksAuth(factory))
|
||||
)
|
||||
assert authed.status_code == 200, (authed.status_code, authed.text)
|
||||
# The minted owner token resolved to Alice, who owns the session, so the
|
||||
# real agent bundle comes back.
|
||||
assert authed.headers.get("X-Agent-Name")
|
||||
assert authed.content[:2] == b"\x1f\x8b", "expected a gzip agent bundle body"
|
||||
@@ -8,6 +8,7 @@ import io
|
||||
import logging
|
||||
import os
|
||||
import tarfile
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
@@ -19,6 +20,8 @@ from omnigent.runner._entry import (
|
||||
_agent_cache_dest,
|
||||
_load_runner_idle_timeout_s_from_config,
|
||||
_make_auth_token_factory,
|
||||
_make_managed_mint_factory,
|
||||
_mint_managed_owner_token,
|
||||
_parent_is_orphaned,
|
||||
_parent_process_is_alive,
|
||||
_resolve_agent_spec_from_server,
|
||||
@@ -31,6 +34,7 @@ from omnigent.runner._entry import (
|
||||
_server_url_from_env,
|
||||
main,
|
||||
)
|
||||
from omnigent.runner.identity import RUNNER_TUNNEL_TOKEN_HEADER
|
||||
from omnigent.runner.transports.ws_tunnel.serve import RUNNER_TUNNEL_REJECTION_PREFIX
|
||||
|
||||
# Force-load the MCP streamable-http client before any test monkeypatches
|
||||
@@ -184,6 +188,302 @@ def test_make_auth_token_factory_returns_none_without_databricks_creds(
|
||||
assert _make_auth_token_factory() is None
|
||||
|
||||
|
||||
def test_make_auth_token_factory_uses_managed_mint_when_only_binding_token(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""A managed sandbox runner (binding token, no user creds) gets a factory.
|
||||
|
||||
With no stored OIDC token and no Databricks credentials, the factory
|
||||
would be ``None`` for a laptop runner — but a managed sandbox still
|
||||
holds its tunnel binding token, so the factory falls back to minting a
|
||||
short-lived owner JWT against it. This is what lets a managed runner's
|
||||
HTTP callbacks authenticate under OIDC/accounts auth.
|
||||
|
||||
:param monkeypatch: Pytest environment patch fixture.
|
||||
:returns: None.
|
||||
"""
|
||||
from omnigent.inner.databricks_executor import DatabricksAuthError
|
||||
|
||||
def _no_sdk(profile: str | None = None) -> tuple[Any, str]:
|
||||
"""Stand in for _resolve_databricks_auth with no credentials."""
|
||||
raise DatabricksAuthError("no Databricks credentials configured")
|
||||
|
||||
monkeypatch.setenv("RUNNER_SERVER_URL", "https://omnigent.example.com")
|
||||
monkeypatch.setenv("OMNIGENT_RUNNER_TUNNEL_BINDING_TOKEN", "managed-binding-token")
|
||||
monkeypatch.setattr("omnigent.cli_auth.load_token", lambda _url: None)
|
||||
monkeypatch.setattr("omnigent.inner.databricks_executor._resolve_databricks_auth", _no_sdk)
|
||||
monkeypatch.setattr(
|
||||
"omnigent.runner._entry._mint_managed_owner_token",
|
||||
lambda mint_url, server_url, binding_token: ("managed-jwt", time.time() + 1800),
|
||||
)
|
||||
|
||||
factory = _make_auth_token_factory()
|
||||
|
||||
assert factory is not None
|
||||
assert factory() == "managed-jwt"
|
||||
|
||||
|
||||
def test_make_auth_token_factory_none_without_creds_or_binding_token(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""No user creds AND no binding token → still ``None`` (unchanged posture).
|
||||
|
||||
The managed-mint fallback must not fire for a non-managed runner: with
|
||||
no binding token there is nothing to mint against, so the factory is
|
||||
``None`` exactly as before.
|
||||
|
||||
:param monkeypatch: Pytest environment patch fixture.
|
||||
:returns: None.
|
||||
"""
|
||||
from omnigent.inner.databricks_executor import DatabricksAuthError
|
||||
|
||||
def _no_sdk(profile: str | None = None) -> tuple[Any, str]:
|
||||
"""Stand in for _resolve_databricks_auth with no credentials."""
|
||||
raise DatabricksAuthError("no Databricks credentials configured")
|
||||
|
||||
monkeypatch.setenv("RUNNER_SERVER_URL", "https://omnigent.example.com")
|
||||
monkeypatch.delenv("OMNIGENT_RUNNER_TUNNEL_BINDING_TOKEN", raising=False)
|
||||
monkeypatch.setattr("omnigent.cli_auth.load_token", lambda _url: None)
|
||||
monkeypatch.setattr("omnigent.inner.databricks_executor._resolve_databricks_auth", _no_sdk)
|
||||
|
||||
assert _make_auth_token_factory() is None
|
||||
|
||||
|
||||
def test_managed_mint_factory_caches_token_until_refresh_skew(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""The factory caches a minted token and reuses it until near expiry.
|
||||
|
||||
A managed session makes many HTTP callbacks; re-minting on every one
|
||||
would hammer the server. The token is minted once and reused until it
|
||||
nears expiry.
|
||||
|
||||
:param monkeypatch: Pytest environment patch fixture.
|
||||
:returns: None.
|
||||
"""
|
||||
calls: list[int] = []
|
||||
|
||||
def _fake_mint(mint_url: str, server_url: str, binding_token: str) -> tuple[str, float]:
|
||||
"""Return a distinct token per call, expiring well beyond the skew."""
|
||||
calls.append(1)
|
||||
return (f"jwt-{len(calls)}", time.time() + 1800)
|
||||
|
||||
monkeypatch.setattr("omnigent.runner._entry._mint_managed_owner_token", _fake_mint)
|
||||
|
||||
# The construction probe mints jwt-1 once; the factory installs.
|
||||
factory = _make_managed_mint_factory("https://s.example.com", "btok")
|
||||
assert factory is not None
|
||||
|
||||
# Subsequent calls reuse the cached token — no new mint.
|
||||
assert factory() == "jwt-1"
|
||||
assert factory() == "jwt-1"
|
||||
assert len(calls) == 1
|
||||
|
||||
|
||||
def test_managed_mint_factory_serves_cached_token_when_refresh_fails(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""A transient mint failure serves the still-valid cached token.
|
||||
|
||||
A blip talking to the mint endpoint must not break in-flight
|
||||
callbacks: while the cached token is still valid, keep serving it and
|
||||
let the on-401 retry re-mint later.
|
||||
|
||||
:param monkeypatch: Pytest environment patch fixture.
|
||||
:returns: None.
|
||||
"""
|
||||
calls: list[int] = []
|
||||
|
||||
def _fake_mint(mint_url: str, server_url: str, binding_token: str) -> tuple[str, float]:
|
||||
"""First call mints a near-expiry token; the refresh attempt fails."""
|
||||
calls.append(1)
|
||||
if len(calls) == 1:
|
||||
# Expiry within the refresh skew → the next call attempts a re-mint.
|
||||
return ("jwt-1", time.time() + 250)
|
||||
raise httpx.ConnectError("mint endpoint unreachable")
|
||||
|
||||
monkeypatch.setattr("omnigent.runner._entry._mint_managed_owner_token", _fake_mint)
|
||||
|
||||
# Construction probe mints jwt-1 (near expiry); the factory installs.
|
||||
factory = _make_managed_mint_factory("https://s.example.com", "btok")
|
||||
assert factory is not None
|
||||
assert len(calls) == 1
|
||||
|
||||
# The token is within the refresh skew, so this call attempts a re-mint,
|
||||
# which fails — the still-valid cached token is served instead of erroring.
|
||||
assert factory() == "jwt-1"
|
||||
assert len(calls) == 2
|
||||
|
||||
|
||||
def test_managed_mint_factory_no_factory_when_server_definitively_refuses(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""A definitive no-mint (HTTP 400/404) installs no factory → bare requests.
|
||||
|
||||
HTTP 400 (no auth provider / header mode) and 404 (an older server
|
||||
without the endpoint) mean the server will never mint for this runner,
|
||||
so the runner must fall back to unauthenticated requests — correct on a
|
||||
no-auth server. No factory is installed.
|
||||
|
||||
:param monkeypatch: Pytest environment patch fixture.
|
||||
:returns: None.
|
||||
"""
|
||||
|
||||
def _refuses(mint_url: str, server_url: str, binding_token: str) -> tuple[str, float]:
|
||||
"""Reject the mint the way a no-auth / header-mode server does (400)."""
|
||||
request = httpx.Request("POST", mint_url)
|
||||
raise httpx.HTTPStatusError(
|
||||
"unsupported", request=request, response=httpx.Response(400, request=request)
|
||||
)
|
||||
|
||||
monkeypatch.setattr("omnigent.runner._entry._mint_managed_owner_token", _refuses)
|
||||
|
||||
assert _make_managed_mint_factory("https://s.example.com", "btok") is None
|
||||
|
||||
|
||||
def test_managed_mint_factory_installs_for_retry_on_transient_boot_failure(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""A transient probe failure still installs the factory (armed to retry).
|
||||
|
||||
If the mint endpoint has a blip at the instant the runner boots (network
|
||||
error, 5xx), the factory must still install so a later callback re-mints
|
||||
— otherwise the runner is left permanently unauthenticated until process
|
||||
restart.
|
||||
|
||||
:param monkeypatch: Pytest environment patch fixture.
|
||||
:returns: None.
|
||||
"""
|
||||
|
||||
def _blip(mint_url: str, server_url: str, binding_token: str) -> tuple[str, float]:
|
||||
"""A transient failure — the endpoint is momentarily unreachable."""
|
||||
raise httpx.ConnectError("mint endpoint unreachable at boot")
|
||||
|
||||
monkeypatch.setattr("omnigent.runner._entry._mint_managed_owner_token", _blip)
|
||||
|
||||
factory = _make_managed_mint_factory("https://s.example.com", "btok")
|
||||
assert factory is not None # installed despite the boot blip
|
||||
assert factory() is None # still can't mint, but it's armed to retry
|
||||
|
||||
|
||||
def test_managed_mint_factory_recovers_after_transient_boot_failure(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""After a transient boot blip, the factory re-mints on the next call.
|
||||
|
||||
Locks in the recovery guarantee: a one-time failure at construction does
|
||||
not disable auth — the very next callback mints successfully.
|
||||
|
||||
:param monkeypatch: Pytest environment patch fixture.
|
||||
:returns: None.
|
||||
"""
|
||||
calls: list[int] = []
|
||||
|
||||
def _fake_mint(mint_url: str, server_url: str, binding_token: str) -> tuple[str, float]:
|
||||
"""Fail the boot probe once, then mint successfully."""
|
||||
calls.append(1)
|
||||
if len(calls) == 1:
|
||||
raise httpx.ConnectError("boot blip")
|
||||
return ("jwt-recovered", time.time() + 1800)
|
||||
|
||||
monkeypatch.setattr("omnigent.runner._entry._mint_managed_owner_token", _fake_mint)
|
||||
|
||||
factory = _make_managed_mint_factory("https://s.example.com", "btok")
|
||||
assert factory is not None # installed despite the boot-probe failure
|
||||
assert factory() == "jwt-recovered" # first real callback re-mints
|
||||
assert len(calls) == 2 # probe (failed) + successful re-mint
|
||||
|
||||
|
||||
def test_managed_mint_factory_declines_at_request_time_and_auth_sends_bare(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""A post-install definitive 400 latches ``declined`` → bare requests.
|
||||
|
||||
The boot-race regression from CI: the runner starts before the server
|
||||
listens, so the construction probe hits a connection error (transient →
|
||||
factory installs), then every request-time mint gets the definitive
|
||||
HTTP 400 of a no-auth server. Without the latch, the factory returns
|
||||
``None`` forever and ``_RunnerDatabricksAuth`` fails closed — bricking
|
||||
every runner→server callback (``spec_resolver_failed``). With it, the
|
||||
first 400 flips the factory to declined and callbacks go out bare,
|
||||
exactly as if no factory had been installed.
|
||||
|
||||
:param monkeypatch: Pytest environment patch fixture.
|
||||
:returns: None.
|
||||
"""
|
||||
calls: list[int] = []
|
||||
|
||||
def _boot_blip_then_refuse(
|
||||
mint_url: str, server_url: str, binding_token: str
|
||||
) -> tuple[str, float]:
|
||||
"""Fail the boot probe with a connection error, then 400 every mint."""
|
||||
calls.append(1)
|
||||
if len(calls) == 1:
|
||||
raise httpx.ConnectError("server not listening yet")
|
||||
request = httpx.Request("POST", mint_url)
|
||||
raise httpx.HTTPStatusError(
|
||||
"no auth provider", request=request, response=httpx.Response(400, request=request)
|
||||
)
|
||||
|
||||
monkeypatch.setattr("omnigent.runner._entry._mint_managed_owner_token", _boot_blip_then_refuse)
|
||||
|
||||
factory = _make_managed_mint_factory("https://s.example.com", "btok")
|
||||
assert factory is not None # boot blip is transient → installed
|
||||
|
||||
auth = _RunnerDatabricksAuth(factory)
|
||||
request = httpx.Request("GET", "http://server/v1/agents/ag_1/download")
|
||||
sent = next(auth.auth_flow(request)) # must NOT raise (fail closed)
|
||||
assert "Authorization" not in sent.headers # bare request, like no factory
|
||||
|
||||
# The latch short-circuits: later callbacks never re-hit the endpoint.
|
||||
request2 = httpx.Request("GET", "http://server/v1/responses/turn_1")
|
||||
sent2 = next(auth.auth_flow(request2))
|
||||
assert "Authorization" not in sent2.headers
|
||||
assert len(calls) == 2 # probe blip + the single definitive 400
|
||||
|
||||
|
||||
def test_mint_managed_owner_token_posts_binding_token_and_parses_response(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""The mint call targets the right URL with the binding-token header.
|
||||
|
||||
Locks the runner->server contract: POST /v1/runners/{id}/token with
|
||||
the tunnel binding token in ``X-Omnigent-Runner-Tunnel-Token``,
|
||||
returning ``{"token", "expires_at"}``.
|
||||
|
||||
:param monkeypatch: Pytest environment patch fixture.
|
||||
:returns: None.
|
||||
"""
|
||||
captured: dict[str, str] = {}
|
||||
|
||||
def _handler(request: httpx.Request) -> httpx.Response:
|
||||
"""Capture the outgoing mint request and return a canned token."""
|
||||
captured["url"] = str(request.url)
|
||||
captured["method"] = request.method
|
||||
captured["binding_token"] = request.headers.get(RUNNER_TUNNEL_TOKEN_HEADER, "")
|
||||
return httpx.Response(200, json={"token": "owner-jwt", "expires_at": 1234567890})
|
||||
|
||||
real_client = httpx.Client
|
||||
|
||||
def _fake_client(**kwargs: Any) -> httpx.Client:
|
||||
"""Build a real sync client backed by the capturing MockTransport."""
|
||||
return real_client(transport=httpx.MockTransport(_handler), **kwargs)
|
||||
|
||||
monkeypatch.setattr("omnigent.runner._entry.httpx.Client", _fake_client)
|
||||
|
||||
token, expires_at = _mint_managed_owner_token(
|
||||
"https://s.example.com/v1/runners/runner_token_abc/token",
|
||||
"https://s.example.com",
|
||||
"the-binding-token",
|
||||
)
|
||||
|
||||
assert token == "owner-jwt"
|
||||
assert expires_at == 1234567890.0
|
||||
assert captured["method"] == "POST"
|
||||
assert captured["binding_token"] == "the-binding-token"
|
||||
assert captured["url"].endswith("/v1/runners/runner_token_abc/token")
|
||||
|
||||
|
||||
def test_runner_databricks_auth_injects_fresh_token_per_request() -> None:
|
||||
"""``_RunnerDatabricksAuth`` calls the factory on every request.
|
||||
|
||||
|
||||
@@ -11,9 +11,11 @@ from functools import partial
|
||||
import httpx
|
||||
import pytest
|
||||
from asgiref.testing import ApplicationCommunicator
|
||||
from fastapi import FastAPI
|
||||
from fastapi import FastAPI, Request
|
||||
from fastapi.responses import JSONResponse
|
||||
from starlette.requests import HTTPConnection
|
||||
|
||||
from omnigent.errors import OmnigentError
|
||||
from omnigent.runner import create_runner_app
|
||||
from omnigent.runner.identity import RUNNER_TUNNEL_TOKEN_HEADER, token_bound_runner_id
|
||||
from omnigent.runner.transports.ws_tunnel.frames import (
|
||||
@@ -980,3 +982,186 @@ async def test_ws_tunnel_loopback_unauthenticated_registers_as_local() -> None:
|
||||
await communicator.send_input({"type": "websocket.disconnect", "code": 1000})
|
||||
with contextlib.suppress(asyncio.TimeoutError):
|
||||
await communicator.wait(timeout=1.0)
|
||||
|
||||
|
||||
# ── Managed-runner token mint endpoint (POST /v1/runners/{id}/token) ──
|
||||
|
||||
|
||||
class _MintingAuthProvider(_CredentialHeaderAuthProvider):
|
||||
"""OIDC/accounts-style provider that also mints runner owner tokens.
|
||||
|
||||
Models the deployed contract where ``mint_runner_token`` returns a
|
||||
bearer. The real ``UnifiedAuthProvider`` signs a JWT; here a
|
||||
deterministic sentinel exercises the route without JWT machinery —
|
||||
the token round-trip itself is covered in
|
||||
``tests/server/test_accounts.py``.
|
||||
"""
|
||||
|
||||
def mint_runner_token(self, user_id: str, ttl_seconds: int) -> str | None:
|
||||
"""Return a deterministic sentinel bearer for *user_id*."""
|
||||
return f"minted-owner-token:{user_id}:{ttl_seconds}"
|
||||
|
||||
|
||||
def _mint_route_app(
|
||||
*,
|
||||
auth_provider: AuthProvider | None,
|
||||
resolve_managed_runner_owner: Callable[[str], str | None] | None,
|
||||
) -> FastAPI:
|
||||
"""Tunnel-route app with the ``OmnigentError`` -> HTTP handler installed.
|
||||
|
||||
The bare :func:`_tunnel_route_app` omits ``create_app``'s exception
|
||||
handler, so the mint endpoint's ``OmnigentError`` would surface as a
|
||||
raw 500. Install the same mapping here so the tests assert the real
|
||||
401 / 400 statuses the endpoint intends.
|
||||
|
||||
:param auth_provider: Auth provider wired into the route.
|
||||
:param resolve_managed_runner_owner: ``runner_id -> owner`` resolver.
|
||||
:returns: The FastAPI app (error handler installed).
|
||||
"""
|
||||
app = _tunnel_route_app(
|
||||
auth_provider=auth_provider,
|
||||
resolve_managed_runner_owner=resolve_managed_runner_owner,
|
||||
).app
|
||||
|
||||
@app.exception_handler(OmnigentError)
|
||||
async def _handle(request: Request, exc: OmnigentError) -> JSONResponse:
|
||||
"""Map the application error to its HTTP status (mirrors create_app)."""
|
||||
return JSONResponse(
|
||||
status_code=exc.http_status,
|
||||
content={"error": {"code": exc.code, "message": exc.message}},
|
||||
)
|
||||
|
||||
return app
|
||||
|
||||
|
||||
async def _post_mint_token(
|
||||
app: FastAPI,
|
||||
runner_id: str,
|
||||
*,
|
||||
token: str | None,
|
||||
) -> httpx.Response:
|
||||
"""POST the mint endpoint with an optional binding-token header.
|
||||
|
||||
:param app: The tunnel-route app under test.
|
||||
:param runner_id: Path runner id.
|
||||
:param token: Binding token for the ``X-Omnigent-Runner-Tunnel-Token``
|
||||
header, or ``None`` to omit it.
|
||||
:returns: The HTTP response.
|
||||
"""
|
||||
headers = {} if token is None else {RUNNER_TUNNEL_TOKEN_HEADER: token}
|
||||
async with httpx.AsyncClient(
|
||||
transport=httpx.ASGITransport(app=app),
|
||||
base_url="http://server",
|
||||
) as client:
|
||||
return await client.post(f"/v1/runners/{runner_id}/token", headers=headers)
|
||||
|
||||
|
||||
async def test_mint_token_endpoint_returns_owner_bearer_for_valid_binding_token() -> None:
|
||||
"""A valid binding token mints an owner bearer scoped to the launch owner.
|
||||
|
||||
The HTTP analog of the tunnel handshake: the managed runner presents
|
||||
its binding token and receives a short-lived owner JWT (here the
|
||||
provider's sentinel) plus an expiry, which it then uses on its HTTP
|
||||
callbacks.
|
||||
|
||||
:returns: None.
|
||||
"""
|
||||
token = "managed-runner-binding-token"
|
||||
runner_id = token_bound_runner_id(token)
|
||||
app = _mint_route_app(
|
||||
auth_provider=_MintingAuthProvider(),
|
||||
resolve_managed_runner_owner=(
|
||||
lambda rid: "owner@example.com" if rid == runner_id else None
|
||||
),
|
||||
)
|
||||
|
||||
response = await _post_mint_token(app, runner_id, token=token)
|
||||
|
||||
assert response.status_code == 200
|
||||
body = response.json()
|
||||
assert body["token"] == "minted-owner-token:owner@example.com:1800"
|
||||
assert isinstance(body["expires_at"], int)
|
||||
assert body["expires_at"] > 0
|
||||
|
||||
|
||||
async def test_mint_token_endpoint_rejects_unrecognized_token() -> None:
|
||||
"""A token with no managed-launch record is refused (fail closed).
|
||||
|
||||
An attacker-chosen token clears the SHA-256 gate for its *own*
|
||||
runner_id, but that id has no bound conversation, so the resolver
|
||||
returns ``None`` and minting is refused — the same posture as the
|
||||
tunnel handshake.
|
||||
|
||||
:returns: None.
|
||||
"""
|
||||
attacker_token = "attacker-chosen-token"
|
||||
runner_id = token_bound_runner_id(attacker_token)
|
||||
app = _mint_route_app(
|
||||
auth_provider=_MintingAuthProvider(),
|
||||
resolve_managed_runner_owner=lambda _rid: None,
|
||||
)
|
||||
|
||||
response = await _post_mint_token(app, runner_id, token=attacker_token)
|
||||
|
||||
assert response.status_code == 401
|
||||
|
||||
|
||||
async def test_mint_token_endpoint_rejects_runner_id_mismatch() -> None:
|
||||
"""A token that doesn't hash to the path runner_id is refused.
|
||||
|
||||
The SHA-256 binding gate runs before any owner lookup, so a token
|
||||
that maps to a different runner_id cannot mint for the path id even
|
||||
if that id has an owner.
|
||||
|
||||
:returns: None.
|
||||
"""
|
||||
app = _mint_route_app(
|
||||
auth_provider=_MintingAuthProvider(),
|
||||
resolve_managed_runner_owner=lambda _rid: "owner@example.com",
|
||||
)
|
||||
|
||||
response = await _post_mint_token(
|
||||
app, "runner_token_does_not_match", token="some-binding-token"
|
||||
)
|
||||
|
||||
assert response.status_code == 401
|
||||
|
||||
|
||||
async def test_mint_token_endpoint_missing_binding_token_rejected() -> None:
|
||||
"""A request without the binding-token header is refused.
|
||||
|
||||
:returns: None.
|
||||
"""
|
||||
runner_id = token_bound_runner_id("whatever")
|
||||
app = _mint_route_app(
|
||||
auth_provider=_MintingAuthProvider(),
|
||||
resolve_managed_runner_owner=lambda _rid: "owner@example.com",
|
||||
)
|
||||
|
||||
response = await _post_mint_token(app, runner_id, token=None)
|
||||
|
||||
assert response.status_code == 401
|
||||
|
||||
|
||||
async def test_mint_token_endpoint_header_mode_unsupported_returns_400() -> None:
|
||||
"""When the provider can't mint (header/proxy mode), the endpoint 400s.
|
||||
|
||||
The binding token is valid and the owner resolves, but header/proxy
|
||||
identity can't be minted server-side (``mint_runner_token`` returns
|
||||
``None``) — a clear 400, not a 401.
|
||||
|
||||
:returns: None.
|
||||
"""
|
||||
token = "managed-runner-binding-token"
|
||||
runner_id = token_bound_runner_id(token)
|
||||
app = _mint_route_app(
|
||||
# Base provider: mint_runner_token uses the ABC default (None).
|
||||
auth_provider=_CredentialHeaderAuthProvider(),
|
||||
resolve_managed_runner_owner=(
|
||||
lambda rid: "owner@example.com" if rid == runner_id else None
|
||||
),
|
||||
)
|
||||
|
||||
response = await _post_mint_token(app, runner_id, token=token)
|
||||
|
||||
assert response.status_code == 400
|
||||
|
||||
@@ -37,6 +37,7 @@ from omnigent.server.accounts_bootstrap import (
|
||||
from omnigent.server.accounts_config import AccountsConfig
|
||||
from omnigent.server.accounts_store import SqlAlchemyAccountStore
|
||||
from omnigent.server.auth import (
|
||||
AuthProvider,
|
||||
UnifiedAuthProvider,
|
||||
create_auth_provider,
|
||||
resolve_auth_source,
|
||||
@@ -363,6 +364,65 @@ def test_accounts_source_login_url_points_at_spa() -> None:
|
||||
assert provider.login_url == "/login"
|
||||
|
||||
|
||||
def test_mint_runner_token_round_trips_to_owner() -> None:
|
||||
"""A managed runner's minted owner token resolves back to the owner.
|
||||
|
||||
The sandbox runner has no login of its own, so the server mints an
|
||||
owner JWT it presents as ``Authorization: Bearer`` on its HTTP
|
||||
callbacks — ``get_user_id`` (the same check ``require_user`` applies)
|
||||
must resolve it to the owner, else every callback 401s.
|
||||
"""
|
||||
cfg = _make_accounts_config()
|
||||
provider = UnifiedAuthProvider(source="accounts", accounts_config=cfg)
|
||||
|
||||
token = provider.mint_runner_token("alice@example.com", 1800)
|
||||
assert token is not None
|
||||
|
||||
request = _FakeReq(headers={"Authorization": f"Bearer {token}"})
|
||||
assert provider.get_user_id(request) == "alice@example.com"
|
||||
|
||||
|
||||
def test_mint_runner_token_rejects_empty_and_reserved_owner() -> None:
|
||||
"""No token for an empty or reserved owner — never mint reserved-identity creds."""
|
||||
cfg = _make_accounts_config()
|
||||
provider = UnifiedAuthProvider(source="accounts", accounts_config=cfg)
|
||||
assert provider.mint_runner_token("", 1800) is None
|
||||
assert provider.mint_runner_token("local", 1800) is None
|
||||
|
||||
|
||||
def test_mint_runner_token_returns_none_for_header_source() -> None:
|
||||
"""Header/proxy auth can't be minted server-side, so it returns None.
|
||||
|
||||
Identity there is asserted by the upstream proxy; a managed runner
|
||||
can't synthesize it. The base ``AuthProvider`` default is also None.
|
||||
"""
|
||||
header_provider = UnifiedAuthProvider(source="header")
|
||||
assert header_provider.mint_runner_token("alice@example.com", 1800) is None
|
||||
|
||||
class _Base(AuthProvider):
|
||||
def get_user_id(self, request: object) -> str | None: # type: ignore[override]
|
||||
return None
|
||||
|
||||
assert _Base().mint_runner_token("alice@example.com", 1800) is None
|
||||
|
||||
|
||||
def test_mint_runner_token_expired_resolves_to_none() -> None:
|
||||
"""A short TTL genuinely expires: past its exp, get_user_id returns None.
|
||||
|
||||
This is what makes the managed-runner auth refreshable rather than a
|
||||
fixed cap — the token expires and the runner re-mints, so there is no
|
||||
static long-lived credential.
|
||||
"""
|
||||
cfg = _make_accounts_config()
|
||||
provider = UnifiedAuthProvider(source="accounts", accounts_config=cfg)
|
||||
|
||||
token = provider.mint_runner_token("alice@example.com", -1)
|
||||
assert token is not None
|
||||
|
||||
request = _FakeReq(headers={"Authorization": f"Bearer {token}"})
|
||||
assert provider.get_user_id(request) is None
|
||||
|
||||
|
||||
# ── resolve_auth_source (shared resolver used by every spawn path) ──
|
||||
|
||||
|
||||
|
||||
@@ -26,6 +26,7 @@ from omnigent.server.oidc import (
|
||||
generate_code_verifier,
|
||||
hmac_digest,
|
||||
mint_session_cookie,
|
||||
mint_session_token,
|
||||
)
|
||||
|
||||
# ── PKCE helpers ─────────────────────────────────────────────────
|
||||
@@ -114,6 +115,44 @@ def test_mint_session_cookie_rejected_with_wrong_secret() -> None:
|
||||
jwt.decode(token, wrong_secret, algorithms=["HS256"])
|
||||
|
||||
|
||||
def test_mint_session_token_produces_valid_jwt_with_seconds_ttl() -> None:
|
||||
"""The seconds-granularity primitive mints the same HS256 claim shape.
|
||||
|
||||
A managed runner needs a sub-hour owner token; ``mint_session_token``
|
||||
is the seconds-based core the hours-only ``mint_session_cookie``
|
||||
cannot express. The claims must match so the same validator accepts
|
||||
either.
|
||||
"""
|
||||
token = mint_session_token(
|
||||
user_id="alice@example.com",
|
||||
cookie_secret=_TEST_SECRET,
|
||||
ttl_seconds=120,
|
||||
provider="accounts",
|
||||
)
|
||||
payload = jwt.decode(token, _TEST_SECRET, algorithms=["HS256"])
|
||||
assert payload["sub"] == "alice@example.com"
|
||||
assert payload["provider"] == "accounts"
|
||||
# exp is ~120 seconds out, not hours.
|
||||
assert time.time() < payload["exp"] <= time.time() + 120 + 5
|
||||
|
||||
|
||||
def test_mint_session_token_short_ttl_expires() -> None:
|
||||
"""A past TTL yields a token the HS256 validator rejects as expired.
|
||||
|
||||
This is the property that removes the fixed session-length cap: the
|
||||
minted owner token genuinely expires, so the runner must (and does)
|
||||
re-mint — rather than holding one static long-lived credential.
|
||||
"""
|
||||
token = mint_session_token(
|
||||
user_id="alice@example.com",
|
||||
cookie_secret=_TEST_SECRET,
|
||||
ttl_seconds=-1,
|
||||
provider="accounts",
|
||||
)
|
||||
with pytest.raises(jwt.ExpiredSignatureError):
|
||||
jwt.decode(token, _TEST_SECRET, algorithms=["HS256"])
|
||||
|
||||
|
||||
def test_hmac_digest_is_deterministic() -> None:
|
||||
"""Same token + secret always produces the same digest.
|
||||
|
||||
|
||||
Reference in New Issue
Block a user