diff --git a/src/opik_mcp/auth_context.py b/src/opik_mcp/auth_context.py index b82db10..bde9551 100644 --- a/src/opik_mcp/auth_context.py +++ b/src/opik_mcp/auth_context.py @@ -38,6 +38,18 @@ # any downstream call. ``None`` means "fall back to settings.comet_workspace". inbound_workspace: ContextVar[str | None] = ContextVar("inbound_workspace", default=None) +# OAuth-authorized workspace *name*, resolved from the opaque bearer via +# ``oauth_identity.resolve_workspace_name`` on the ``initialize`` handshake. +# Consumed ONLY by the instructions blob (``instructions.render_instructions``) +# so an agent can truthfully name the workspace it is operating against. Kept +# deliberately separate from ``inbound_workspace`` so this read-only display +# value never leaks into the outbound ``Comet-Workspace`` header / data routing +# (which stays token-derived server-side). ``None`` means "not resolved; fall +# back to the static settings workspace". +resolved_workspace_name: ContextVar[str | None] = ContextVar( + "resolved_workspace_name", default=None +) + def classify_bearer(auth_header: str) -> tuple[str, str]: """Classify a non-empty inbound ``Authorization`` header for BI analytics. diff --git a/src/opik_mcp/config.py b/src/opik_mcp/config.py index e8bfc47..6ec9ddc 100644 --- a/src/opik_mcp/config.py +++ b/src/opik_mcp/config.py @@ -95,6 +95,13 @@ class Settings(BaseSettings): # ingress URL where MCP hosts reach this opik-mcp instance. opik_mcp_resource_uri: str | None = None + # Timeout for the one-shot OAuth token introspection (POST /opik/auth-oauth) + # used to resolve the authorized workspace name for the per-session + # ``initialize`` instructions blob. Bounded so a slow/unreachable backend + # cannot stall the MCP handshake — the call fails soft (the blob falls back + # to the static workspace), so a short ceiling is safe. + opik_mcp_oauth_introspect_timeout_s: float = 5.0 + opik_mcp_log_level: str = "INFO" opik_mcp_transport: str = "stdio" opik_mcp_host: str = "127.0.0.1" diff --git a/src/opik_mcp/instructions.py b/src/opik_mcp/instructions.py index 946f378..13e8419 100644 --- a/src/opik_mcp/instructions.py +++ b/src/opik_mcp/instructions.py @@ -12,16 +12,21 @@ +60pp on smaller ones, so this is the highest-leverage single dial we have on tool selection quality. -Phase 1 ships static-template behavior — user_email and full per-session -context will land when the OAuth/identity path is wired up (Phase 2). For -now, ``workspace`` and ``opik_url`` come from ``Settings``. +The blob is rendered per session (see ``server.install_session_instructions``): +``workspace`` prefers the OAuth-authorized workspace for THIS session — the +inbound ``Comet-Workspace`` header, else the name introspected from the bearer +(``resolved_workspace_name``) — and only falls back to the static ``Settings`` +workspace for stdio / API-key installs. ``opik_url`` is the Opik **UI** base, +derived from ``Settings`` (the REST ``OPIK_URL`` minus its ``/api`` suffix). """ from __future__ import annotations from datetime import UTC, datetime +from opik_mcp.auth_context import inbound_workspace, resolved_workspace_name from opik_mcp.config import DEFAULT_WORKSPACE, Settings, get_settings +from opik_mcp.opik_client import opik_rest_base from opik_mcp.writes.registry import WRITE_OPERATIONS _TEMPLATE = """\ @@ -53,6 +58,23 @@ """ +def _opik_ui_url(s: Settings) -> str: + """Opik **UI** base URL for the blob, or a generic placeholder if unconfigured. + + Derived from :func:`opik_rest_base` — the single source of truth for where + Opik lives (``OPIK_URL`` override, else ``COMET_URL_OVERRIDE + "/opik/api"``) + — so the UI link can never drift from where REST calls actually go. That base + is the REST **API** base (``…/opik/api``); the UI lives at the same origin + without the trailing ``/api`` segment, so we strip it. + """ + base = opik_rest_base(s) + if base is None: + return "(Opik URL not configured)" + if base.endswith("/api"): + base = base[: -len("/api")] + return base + + def _render_default_project_clause(s: Settings) -> str: pname = s.opik_default_project_name if not pname: @@ -76,15 +98,19 @@ def render_instructions( ``user_email`` is omitted from the rendered text when unknown — better no claim than a stale claim. Same for ``opik_url`` — falls back to a generic placeholder if the config is partial. + + Workspace precedence (most to least authoritative for THIS session): an + explicit inbound ``Comet-Workspace`` header → the OAuth-introspected + ``resolved_workspace_name`` → the static ``Settings`` workspace → ``"default"``. """ s = settings if settings is not None else get_settings() - workspace = s.comet_workspace or DEFAULT_WORKSPACE - if s.opik_url: - opik_url = s.opik_url.rstrip("/") - elif s.comet_url_override: - opik_url = f"{s.comet_url_override.rstrip('/')}/opik" - else: - opik_url = "(Opik URL not configured)" + workspace = ( + inbound_workspace.get() + or resolved_workspace_name.get() + or s.comet_workspace + or DEFAULT_WORKSPACE + ) + opik_url = _opik_ui_url(s) user_clause = f" as {user_email}" if user_email else "" today = today if today is not None else datetime.now(UTC) diff --git a/src/opik_mcp/oauth_identity.py b/src/opik_mcp/oauth_identity.py new file mode 100644 index 0000000..58d89b1 --- /dev/null +++ b/src/opik_mcp/oauth_identity.py @@ -0,0 +1,83 @@ +"""Resolve the OAuth-authorized workspace name for the session-context blob. + +In OAuth-passthrough mode opik-mcp forwards an opaque ``opik_mcp_at_``-prefixed +bearer onward and lets opik-backend derive the workspace from the token row — it +never learns the workspace *name* itself (the host doesn't send ``Comet-Workspace`` +in OAuth mode; the name lives only in the token binding). That's fine for data +routing, but the per-session ``initialize`` instructions blob (``instructions.py``) +needs the human-readable name so an agent can truthfully say which workspace it is +operating against. + +opik-backend exposes a purpose-built introspection endpoint for exactly this — +``POST /opik/auth-oauth`` (``OAuthValidateTokenResource``) returns the identity a +bearer resolves to, including ``workspace_name``. We call it once per session, on +the ``initialize`` handshake, forwarding the inbound bearer verbatim. + +This is best-effort: any failure (unconfigured base, non-200, network error, +malformed body) returns ``None`` so the handshake never breaks — the blob simply +falls back to the static settings workspace. +""" + +from __future__ import annotations + +import logging + +import httpx + +from opik_mcp.config import Settings +from opik_mcp.opik_client import opik_rest_base + +logger = logging.getLogger("opik_mcp") + +# JAX-RS path of opik-backend's token-introspection endpoint +# (``OAuthConstants.OAUTH_VALIDATE_TOKEN_RESOURCE_BASE_PATH``). It is a sibling of +# the ``/v1/private/...`` REST routes at the backend root, so it hangs off the same +# REST base opik-mcp already uses for data calls. +_VALIDATE_TOKEN_PATH = "/opik/auth-oauth" + + +async def resolve_workspace_name(authorization: str, settings: Settings) -> str | None: + """Introspect an inbound OAuth bearer → its ``workspace_name``, or ``None``. + + ``authorization`` is the full inbound ``Authorization`` header value + (``"Bearer opik_mcp_at_…"``), forwarded verbatim — opik-backend re-validates + the token shape and resolves it server-side. Never raises; returns ``None`` on + any failure so the ``initialize`` handshake degrades gracefully. + """ + base = opik_rest_base(settings) + if base is None: + return None + url = f"{base}{_VALIDATE_TOKEN_PATH}" + headers = { + "Authorization": authorization, + "Content-Type": "application/json", + "Accept": "application/json", + } + try: + async with httpx.AsyncClient( + timeout=settings.opik_mcp_oauth_introspect_timeout_s + ) as client: + resp = await client.post(url, headers=headers) + if resp.status_code != 200: + logger.debug("workspace introspection: non-200 status %s", resp.status_code) + return None + body = resp.json() + except Exception: + # Best-effort by contract: this runs on the ``initialize`` handshake and + # must NEVER raise (a failure just falls back to the static workspace). + # Catch broadly on purpose — beyond httpx.HTTPError + the ValueError from + # resp.json() on a non-JSON body, httpx raises httpx.InvalidURL (a direct + # Exception subclass, NOT an HTTPError) for a malformed REST base, which + # would otherwise escape into the middleware and 500 the handshake. + # Telemetry-free debug only — a failed lookup isn't worth surfacing. + logger.debug("workspace introspection failed", exc_info=True) + return None + if not isinstance(body, dict): + return None + workspace_name = body.get("workspace_name") + if isinstance(workspace_name, str) and workspace_name: + return workspace_name + return None + + +__all__ = ["resolve_workspace_name"] diff --git a/src/opik_mcp/opik_client.py b/src/opik_mcp/opik_client.py index 26726ff..07e2eb2 100644 --- a/src/opik_mcp/opik_client.py +++ b/src/opik_mcp/opik_client.py @@ -662,19 +662,32 @@ def resolve_opik_config(settings: Settings) -> tuple[str, str, str | None]: # configured workspace, else "default" (Opik SDK convention). No hard # failure — lets local/OSS users run without setting a workspace. workspace = inbound_ws or settings.comet_workspace or DEFAULT_WORKSPACE - if settings.opik_url: - base = settings.opik_url.rstrip("/") - else: + base = opik_rest_base(settings) + if base is None: # ``comet_url_override`` has a non-empty default in ``Settings`` but # ``COMET_URL_OVERRIDE=""`` would override it to empty — defend against # that so we never POST to ``/opik/api`` (relative URL → wherever the # process happens to be). - if not settings.comet_url_override: - raise MissingConfigError("OPIK_URL or COMET_URL_OVERRIDE is required to call Opik REST") - base = f"{settings.comet_url_override.rstrip('/')}/opik/api" + raise MissingConfigError("OPIK_URL or COMET_URL_OVERRIDE is required to call Opik REST") return base, api_key, workspace +def opik_rest_base(settings: Settings) -> str | None: + """Resolve Opik's REST API base URL from settings, or ``None`` if unconfigured. + + Single source of truth for the rule: an explicit ``OPIK_URL`` override wins; + otherwise derive from ``COMET_URL_OVERRIDE + "/opik/api"``. Shared by + ``resolve_opik_config`` (which treats ``None`` as a fatal misconfig) and + ``oauth_identity.resolve_workspace_name`` (which treats ``None`` as "skip, + fall back to the static workspace"), so both agree on where Opik lives. + """ + if settings.opik_url: + return settings.opik_url.rstrip("/") + if settings.comet_url_override: + return f"{settings.comet_url_override.rstrip('/')}/opik/api" + return None + + def make_opik_client(settings: Settings) -> OpikClient: """Construct an ``OpikClient`` bound to the configured workspace.""" base_url, api_key, workspace = resolve_opik_config(settings) diff --git a/src/opik_mcp/server.py b/src/opik_mcp/server.py index fc401b1..ffa6b7c 100644 --- a/src/opik_mcp/server.py +++ b/src/opik_mcp/server.py @@ -33,10 +33,12 @@ classify_bearer, inbound_authorization, inbound_workspace, + resolved_workspace_name, settings_auth_mode, ) from opik_mcp.config import MissingConfigError, Settings, get_settings from opik_mcp.instructions import render_instructions +from opik_mcp.oauth_identity import resolve_workspace_name from opik_mcp.opik_client import make_opik_client, resolve_opik_config from opik_mcp.read_list import run_list, run_read from opik_mcp.read_list.registry import LISTABLE_TYPES, READABLE_TYPES @@ -640,11 +642,25 @@ async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) - auth_token = inbound_authorization.set(auth) workspace = request.headers.get("comet-workspace") workspace_token = inbound_workspace.set(workspace) + # On the session-creating request — the MCP ``initialize`` handshake is + # the only one without an ``Mcp-Session-Id`` — resolve the OAuth-authorized + # workspace NAME so the per-session instructions blob can name it. In OAuth + # mode the host sends no ``Comet-Workspace`` header and the token is opaque + # to us, so we introspect it here (once per session); tool calls carry a + # session id and skip this. Best-effort: a failure leaves the blob on its + # static fallback and never blocks the handshake. + resolved_token = None + if request.headers.get("mcp-session-id") is None and classify_bearer(auth)[0] == "oauth": + workspace_name = await resolve_workspace_name(auth, get_settings()) + if workspace_name: + resolved_token = resolved_workspace_name.set(workspace_name) try: return await call_next(request) finally: inbound_authorization.reset(auth_token) inbound_workspace.reset(workspace_token) + if resolved_token is not None: + resolved_workspace_name.reset(resolved_token) def _unauthorized(self) -> Response: # RFC 6750 §3 + RFC 9728: pointing MCP hosts at protected-resource @@ -1024,8 +1040,40 @@ async def _composed(app: Any) -> AsyncIterator[Any]: return _composed +def install_session_instructions(server: FastMCP) -> None: + """Render ``InitializeResult.instructions`` per session rather than once at boot. + + FastMCP captures the ``instructions`` string at construction time, so the blob + would be identical for every session and could never name the per-session + OAuth workspace. We wrap the lowlevel server's ``create_initialization_options`` + (invoked once per session, inside the session task that inherits the + ``initialize`` request's ContextVars) to re-render the blob with the workspace + resolved for that session. Mirrors ``install_tools_listed_emitter``'s in-place + handler swap. The boot-time static render stays on the options as a fallback. + """ + try: + lowlevel = server._mcp_server + except AttributeError: + logger.debug("install_session_instructions: mcp has no _mcp_server attribute") + return + original = lowlevel.create_initialization_options + + def create_initialization_options(*args: Any, **kwargs: Any) -> Any: + options = original(*args, **kwargs) + try: + options.instructions = render_instructions() + except Exception: + # A render hiccup must never break the initialize handshake — leave + # the boot-time static instructions already on the options in place. + logger.debug("per-session instructions render failed", exc_info=True) + return options + + lowlevel.create_initialization_options = create_initialization_options # type: ignore[method-assign] + + def build_app() -> ASGIApp: install_tools_listed_emitter(mcp) + install_session_instructions(mcp) s = get_settings() # Serve the transport at the configured path so it matches the advertised # resource URI behind a non-rewriting path-prefix proxy. Read at app-build diff --git a/tests/conftest.py b/tests/conftest.py index 94da4c9..2a073a1 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -60,6 +60,37 @@ def _reset_analytics_wrappers_state() -> Generator[None]: os.environ.pop(LIFECYCLE_SENTINEL, None) +@pytest.fixture(autouse=True) +def _disable_workspace_introspection() -> Generator[None]: + """Stub OAuth workspace introspection to a no-op by default. + + The ``initialize`` handshake resolves the authorized workspace name by + POSTing to opik-backend's ``/opik/auth-oauth`` (``server.resolve_workspace_name``). + Left live, every test that sends a session-less OAuth bearer to ``/mcp`` would + fire a real network call — slow, flaky, and able to land in another test's + ``@respx.mock`` window. Disabled by default (mirrors the analytics default + above); tests that exercise resolution ``monkeypatch.setattr`` this, and the + ``resolve_workspace_name`` unit tests call the real function directly. + + Uses a standalone ``pytest.MonkeyPatch()`` rather than the ``monkeypatch`` + *fixture* on purpose: depending on that fixture from an autouse fixture pulls + it into the autouse setup phase ahead of ``_reset_analytics_wrappers_state``, + inverting teardown order so its ``reset_analytics_for_tests()`` runs while a + test still has ``get_analytics`` patched to a non-lru_cache stub (AttributeError + on ``cache_info``). A self-owned instance keeps fixture ordering untouched. + """ + + async def _none(*_args: object, **_kwargs: object) -> None: + return None + + mp = pytest.MonkeyPatch() + mp.setattr("opik_mcp.server.resolve_workspace_name", _none) + try: + yield + finally: + mp.undo() + + # Session-scoped HTTP client over the real ASGI app. Shared across test # modules because the underlying FastMCP `StreamableHTTPSessionManager` is a # process-level singleton that may only be `.run()`'d once — letting each diff --git a/tests/test_http_auth.py b/tests/test_http_auth.py index ef3b495..12bd5f6 100644 --- a/tests/test_http_auth.py +++ b/tests/test_http_auth.py @@ -56,3 +56,34 @@ async def test_any_bearer_initializes(http_client: httpx.AsyncClient) -> None: ) assert r.status_code == 200 assert "opik-mcp" in r.text + + +@pytest.mark.anyio +async def test_initialize_names_oauth_workspace( + http_client: httpx.AsyncClient, monkeypatch: pytest.MonkeyPatch +) -> None: + """End-to-end (OPIK-7033): the per-session instructions blob in the + ``initialize`` result names the OAuth-authorized workspace, not "default". + + Drives the whole path — middleware introspection → ContextVar → per-session + ``create_initialization_options`` re-render — over the real ASGI app, with + only the backend introspection call stubbed. + """ + + async def fake_resolve(_auth: str, _settings: object) -> str: + return "andreicautisanu" + + monkeypatch.setattr("opik_mcp.server.resolve_workspace_name", fake_resolve) + r = await http_client.post( + "/mcp", + json=INITIALIZE, + headers={ + "Authorization": f"Bearer {OAUTH_ACCESS_TOKEN_PREFIX}tok", + "Accept": "application/json, text/event-stream", + }, + ) + assert r.status_code == 200 + # The blob is a JSON string inside the JSON-RPC result, so its inner quotes + # are backslash-escaped in the raw SSE body. + assert 'workspace \\"andreicautisanu\\"' in r.text + assert 'workspace \\"default\\"' not in r.text diff --git a/tests/test_instructions.py b/tests/test_instructions.py index 303d3d2..38e1b29 100644 --- a/tests/test_instructions.py +++ b/tests/test_instructions.py @@ -43,6 +43,44 @@ def test_render_uses_comet_url_override_when_opik_url_missing() -> None: assert "https://demo.comet.com/opik" in out +def test_render_strips_api_suffix_from_opik_url() -> None: + """OPIK_URL is the REST API base (…/opik/api); the blob must name the UI + base (…/opik) — the verbatim ``/api`` leak is OPIK-7033's defect #2.""" + s = _settings(opik_url="https://dev.comet.com/opik/api") + out = render_instructions(s) + assert "The Opik UI is at https://dev.comet.com/opik." in out + assert "https://dev.comet.com/opik/api" not in out + + +def test_render_prefers_resolved_workspace_over_settings() -> None: + """The OAuth-introspected workspace (set per session) outranks the static + env workspace — defect #1: the blob must name the authorized workspace.""" + from opik_mcp.auth_context import resolved_workspace_name + + token = resolved_workspace_name.set("andreicautisanu") + try: + out = render_instructions(_settings(comet_workspace="env-ws")) + finally: + resolved_workspace_name.reset(token) + assert 'workspace "andreicautisanu"' in out + assert 'workspace "env-ws"' not in out + + +def test_render_inbound_workspace_header_outranks_resolved() -> None: + """An explicit Comet-Workspace header (self-hosted / API-key) is the most + authoritative signal for the session.""" + from opik_mcp.auth_context import inbound_workspace, resolved_workspace_name + + t_header = inbound_workspace.set("header-ws") + t_resolved = resolved_workspace_name.set("resolved-ws") + try: + out = render_instructions(_settings(comet_workspace=None)) + finally: + inbound_workspace.reset(t_header) + resolved_workspace_name.reset(t_resolved) + assert 'workspace "header-ws"' in out + + def test_render_uses_default_workspace_when_unset() -> None: """With no workspace configured the tools operate against "default" (Opik SDK convention), so the LLM-facing context must say so rather than diff --git a/tests/test_oauth_identity.py b/tests/test_oauth_identity.py new file mode 100644 index 0000000..eaeef2f --- /dev/null +++ b/tests/test_oauth_identity.py @@ -0,0 +1,116 @@ +"""Unit tests for OAuth token → workspace-name introspection. + +``resolve_workspace_name`` POSTs the inbound bearer to opik-backend's +``/opik/auth-oauth`` introspection endpoint and pulls ``workspace_name`` out of +the ``ValidatedToken`` response. It is best-effort: any failure resolves to +``None`` so the ``initialize`` handshake never breaks. +""" + +import httpx +import pytest +import respx + +from opik_mcp.config import Settings +from opik_mcp.oauth_identity import resolve_workspace_name + +AUTH = "Bearer opik_mcp_at_abc123" + + +def _settings(**overrides: object) -> Settings: + base: dict[str, object] = {"opik_url": "https://opik.test/api"} + base.update(overrides) + return Settings(**base) # type: ignore[arg-type] + + +@pytest.fixture +def anyio_backend() -> str: + return "asyncio" + + +@pytest.mark.anyio +@respx.mock +async def test_resolves_workspace_name_on_200() -> None: + route = respx.post("https://opik.test/api/opik/auth-oauth").mock( + return_value=httpx.Response( + 200, + json={ + "user_name": "u", + "workspace_id": "ws-id", + "workspace_name": "andreicautisanu", + "resource": "https://opik.test/api/v1/mcp", + }, + ) + ) + ws = await resolve_workspace_name(AUTH, _settings()) + assert ws == "andreicautisanu" + assert route.called + # Inbound bearer is forwarded verbatim — opik-backend re-validates it. + assert route.calls.last.request.headers["authorization"] == AUTH + + +@pytest.mark.anyio +@respx.mock +async def test_returns_none_on_401() -> None: + respx.post("https://opik.test/api/opik/auth-oauth").mock(return_value=httpx.Response(401)) + assert await resolve_workspace_name(AUTH, _settings()) is None + + +@pytest.mark.anyio +async def test_returns_none_on_invalid_url() -> None: + """A malformed REST base must fail soft, not crash the handshake. + + ``httpx.InvalidURL`` is a direct ``Exception`` subclass (NOT an + ``httpx.HTTPError``), so it would escape a narrow ``except`` and 500 the + ``initialize`` request. The whitespace in the host forces ``InvalidURL`` at + request-build time (before any transport / respx mock is consulted). + """ + s = _settings(opik_url="http://exa mple.com/api") + assert await resolve_workspace_name(AUTH, s) is None + + +@pytest.mark.anyio +@respx.mock +async def test_returns_none_on_network_error() -> None: + respx.post("https://opik.test/api/opik/auth-oauth").mock( + side_effect=httpx.ConnectError("backend unreachable") + ) + assert await resolve_workspace_name(AUTH, _settings()) is None + + +@pytest.mark.anyio +@respx.mock +async def test_returns_none_on_non_json_body() -> None: + respx.post("https://opik.test/api/opik/auth-oauth").mock( + return_value=httpx.Response(200, text="not json") + ) + assert await resolve_workspace_name(AUTH, _settings()) is None + + +@pytest.mark.anyio +@respx.mock +async def test_returns_none_when_workspace_name_absent() -> None: + respx.post("https://opik.test/api/opik/auth-oauth").mock( + return_value=httpx.Response(200, json={"user_name": "u"}) + ) + assert await resolve_workspace_name(AUTH, _settings()) is None + + +@pytest.mark.anyio +@respx.mock +async def test_derives_url_from_comet_url_override() -> None: + route = respx.post("https://demo.comet.com/opik/api/opik/auth-oauth").mock( + return_value=httpx.Response(200, json={"workspace_name": "demo-ws"}) + ) + ws = await resolve_workspace_name( + AUTH, Settings(opik_url=None, comet_url_override="https://demo.comet.com/") + ) + assert ws == "demo-ws" + assert route.called + + +@pytest.mark.anyio +async def test_returns_none_when_base_unconfigured() -> None: + # No OPIK_URL and an explicitly empty COMET_URL_OVERRIDE → no base → skip + # without any network call. + s = Settings(opik_url=None, comet_url_override="") + assert await resolve_workspace_name(AUTH, s) is None diff --git a/tests/test_oauth_passthrough_mode.py b/tests/test_oauth_passthrough_mode.py index 44aea12..9ac5cd4 100644 --- a/tests/test_oauth_passthrough_mode.py +++ b/tests/test_oauth_passthrough_mode.py @@ -97,6 +97,89 @@ async def call_next(_r: Request) -> Response: assert captured["workspace"] == "my-team" +@pytest.mark.anyio +async def test_resolves_workspace_on_session_creating_oauth_request( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The ``initialize`` handshake (no Mcp-Session-Id) on an OAuth bearer + introspects the workspace name and exposes it via the ContextVar the + instructions blob reads — then resets it after the request.""" + from opik_mcp.auth_context import resolved_workspace_name + + async def fake_resolve(_auth: str, _settings: object) -> str: + return "andreicautisanu" + + monkeypatch.setattr("opik_mcp.server.resolve_workspace_name", fake_resolve) + mw = _build_middleware() + request = _make_request({"authorization": f"Bearer {OAUTH_ACCESS_TOKEN_PREFIX}abc"}) + + captured: dict[str, str | None] = {} + + async def call_next(_r: Request) -> Response: + captured["resolved"] = resolved_workspace_name.get() + return JSONResponse({"ok": True}) + + await mw.dispatch(request, call_next) + assert captured["resolved"] == "andreicautisanu" + # Reset after the request — no leakage to the next session. + assert resolved_workspace_name.get() is None + + +@pytest.mark.anyio +async def test_skips_resolution_when_session_already_exists( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Requests carrying an Mcp-Session-Id are tool calls, not the handshake — + they must not pay the introspection round-trip.""" + from opik_mcp.auth_context import resolved_workspace_name + + calls: list[str] = [] + + async def spy_resolve(auth: str, _settings: object) -> str: + calls.append(auth) + return "x" + + monkeypatch.setattr("opik_mcp.server.resolve_workspace_name", spy_resolve) + mw = _build_middleware() + request = _make_request( + { + "authorization": f"Bearer {OAUTH_ACCESS_TOKEN_PREFIX}abc", + "mcp-session-id": "session-123", + } + ) + + captured: dict[str, str | None] = {} + + async def call_next(_r: Request) -> Response: + captured["resolved"] = resolved_workspace_name.get() + return JSONResponse({"ok": True}) + + await mw.dispatch(request, call_next) + assert calls == [] + assert captured["resolved"] is None + + +@pytest.mark.anyio +async def test_skips_resolution_for_api_key_bearer(monkeypatch: pytest.MonkeyPatch) -> None: + """Only OAuth-prefixed bearers carry a token-bound workspace to resolve; + an API-key-shaped bearer keeps the legacy header/settings workspace path.""" + calls: list[str] = [] + + async def spy_resolve(auth: str, _settings: object) -> str: + calls.append(auth) + return "x" + + monkeypatch.setattr("opik_mcp.server.resolve_workspace_name", spy_resolve) + mw = _build_middleware() + request = _make_request({"authorization": "Bearer some-static-api-key"}) + + async def call_next(_r: Request) -> Response: + return JSONResponse({"ok": True}) + + await mw.dispatch(request, call_next) + assert calls == [] + + @pytest.mark.anyio async def test_missing_authorization_returns_401_with_www_authenticate() -> None: """The 401 must point hosts at the protected-resource metadata so they