ZY-Agent/backend/tests/test_auth_middleware_workspace.py

"""AuthMiddleware injects the workspace ContextVar (Stage 0 PR4 T4.7).

After PR4 every authenticated request has a workspace bound on
``deerflow.runtime.workspace_context._current_workspace``. The middleware
populates it from the JWT's ``wid`` / ``role`` claims, mirrors what it
already does for ``user_context``, and tears both down in a single
``try/finally`` so leaks don't cross requests.

Legacy 4-field tokens (no ``wid``) are rejected upstream by
``decode_token`` (T4.6) — those should never reach the workspace
injection branch; this file pins that the 401 they trigger carries
``AuthErrorCode.WORKSPACE_REQUIRED``.
"""

from __future__ import annotations

from datetime import UTC, datetime, timedelta
from unittest.mock import AsyncMock, patch
from uuid import uuid4

import jwt
import pytest
from fastapi import FastAPI
from starlette.testclient import TestClient

from app.gateway.auth import create_access_token
from app.gateway.auth.config import get_auth_config
from app.gateway.auth.models import User
from app.gateway.auth_middleware import AuthMiddleware
from deerflow.runtime.workspace_context import get_current_workspace


@pytest.fixture(autouse=True)
def _stable_jwt_secret(monkeypatch):
    monkeypatch.setenv("AUTH_JWT_SECRET", "test-secret-key-for-jwt-testing-minimum-32-chars")
    yield


def _make_app() -> FastAPI:
    """App with AuthMiddleware + an inspect route that surfaces the contextvar."""
    app = FastAPI()
    app.add_middleware(AuthMiddleware)

    @app.get("/api/v1/auth/setup-status")  # public — never gates on wid
    async def setup_status():
        return {"needs_setup": False}

    @app.get("/api/models")  # protected — exercises wid injection
    async def inspect_workspace():
        ws = get_current_workspace()
        if ws is None:
            return {"workspace": None}
        return {"workspace": {"id": ws.id, "role": ws.role}}

    return app


def _make_user(uid: str) -> User:
    return User(id=uid, email="t@example.com", password_hash="hash", token_version=0)


def _make_legacy_token() -> str:
    """Encode a pre-PR4 JWT (no wid/role) directly."""
    now = datetime.now(UTC)
    payload = {
        "sub": str(uuid4()),
        "exp": now + timedelta(hours=1),
        "iat": now,
        "ver": 0,
    }
    return jwt.encode(payload, get_auth_config().jwt_secret, algorithm="HS256")


def test_new_jwt_injects_workspace_into_contextvar() -> None:
    """Cookie with wid+role → route observes the workspace via the contextvar."""
    uid = str(uuid4())
    token = create_access_token(uid, workspace_id="ws-abc", role="owner")

    with patch("app.gateway.deps.get_local_provider") as fn:
        fn.return_value.get_user = AsyncMock(return_value=_make_user(uid))
        client = TestClient(_make_app())
        res = client.get("/api/models", cookies={"access_token": token})

    assert res.status_code == 200, res.text
    assert res.json() == {"workspace": {"id": "ws-abc", "role": "owner"}}


def test_legacy_jwt_rejected_with_workspace_required() -> None:
    """No-wid tokens get 401 with AuthErrorCode.WORKSPACE_REQUIRED, not generic token_invalid."""
    client = TestClient(_make_app())
    res = client.get("/api/models", cookies={"access_token": _make_legacy_token()})

    assert res.status_code == 401
    assert res.json()["detail"]["code"] == "workspace_required"


def test_public_path_skips_workspace_check() -> None:
    """Public whitelist (e.g. /api/v1/auth/setup-status) does not require wid."""
    client = TestClient(_make_app())
    res = client.get("/api/v1/auth/setup-status")  # no cookie at all
    assert res.status_code == 200


@pytest.mark.no_auto_workspace
def test_workspace_contextvar_resets_between_requests() -> None:
    """After dispatch returns the contextvar must be clear (no leak across requests).

    Why we test this: if the try/finally is wired only for user_context but
    not workspace_context, two back-to-back requests can see each other's
    workspace under asyncio task switching.
    """
    uid = str(uuid4())
    token = create_access_token(uid, workspace_id="ws-first", role="owner")

    # First request resolves to ws-first
    with patch("app.gateway.deps.get_local_provider") as fn:
        fn.return_value.get_user = AsyncMock(return_value=_make_user(uid))
        client = TestClient(_make_app())
        res1 = client.get("/api/models", cookies={"access_token": token})
    assert res1.json() == {"workspace": {"id": "ws-first", "role": "owner"}}

    # Outside the request scope the contextvar must be empty again.
    assert get_current_workspace() is None

    # Second request with a different workspace must not see ws-first.
    token2 = create_access_token(uid, workspace_id="ws-second", role="owner")
    with patch("app.gateway.deps.get_local_provider") as fn:
        fn.return_value.get_user = AsyncMock(return_value=_make_user(uid))
        client = TestClient(_make_app())
        res2 = client.get("/api/models", cookies={"access_token": token2})
    assert res2.json() == {"workspace": {"id": "ws-second", "role": "owner"}}