javadandCursor
b93ecf1691
Keep source, config, and deploy files. Leave runtime data, uploads, and logs out of git. Co-authored-by: Cursor <cursoragent@cursor.com>
109 lines
3.2 KiB
Python
109 lines
3.2 KiB
Python
"""Auth middleware must not break concurrent SSE when another client authenticates."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
from fastapi import FastAPI, Request
|
|
from fastapi.responses import StreamingResponse
|
|
from httpx import ASGITransport, AsyncClient
|
|
|
|
from kamidso.api.deps import KamidsoAuthMiddleware, require_user
|
|
from kamidso.api.tokens import encode_jwt
|
|
from kamidso.contracts.user import CurrentUser
|
|
from kamidso.settings.flags import load_kamidso_settings
|
|
|
|
|
|
def _token_for(user: CurrentUser, settings) -> str:
|
|
return encode_jwt(
|
|
{
|
|
"sub": user.id,
|
|
"email": user.email,
|
|
"role": user.role,
|
|
"workspace_id": user.workspace_id,
|
|
},
|
|
settings.jwt_secret,
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sse_keeps_streaming_while_another_client_hits_auth():
|
|
settings = load_kamidso_settings()
|
|
user_a = CurrentUser(
|
|
id="user_a",
|
|
email="a@example.com",
|
|
name="A",
|
|
role="user",
|
|
workspace_id="ws_a",
|
|
)
|
|
user_b = CurrentUser(
|
|
id="user_b",
|
|
email="b@example.com",
|
|
name="B",
|
|
role="user",
|
|
workspace_id="ws_b",
|
|
)
|
|
token_a = _token_for(user_a, settings)
|
|
token_b = _token_for(user_b, settings)
|
|
|
|
class _Workspace:
|
|
async def get_user(self, user_id: str):
|
|
return user_a if user_id == user_a.id else user_b
|
|
|
|
release = asyncio.Event()
|
|
chunks_sent = 0
|
|
|
|
app = FastAPI()
|
|
app.state.kamidso = SimpleNamespace(settings=settings, workspace=_Workspace())
|
|
app.add_middleware(KamidsoAuthMiddleware)
|
|
|
|
@app.get("/api/kamidso/v1/auth/session")
|
|
async def session(request: Request):
|
|
user = require_user(request)
|
|
return {"user": user.to_public_dict()}
|
|
|
|
@app.get("/api/kamidso/v1/chat-stream")
|
|
async def chat_stream(request: Request):
|
|
require_user(request)
|
|
|
|
async def gen():
|
|
nonlocal chunks_sent
|
|
for i in range(6):
|
|
chunks_sent += 1
|
|
yield f"data: {{\"n\": {i}}}\n\n"
|
|
if i == 1:
|
|
await release.wait()
|
|
else:
|
|
await asyncio.sleep(0.01)
|
|
|
|
return StreamingResponse(gen(), media_type="text/event-stream")
|
|
|
|
transport = ASGITransport(app=app)
|
|
async with AsyncClient(transport=transport, base_url="http://test") as client:
|
|
stream_task = asyncio.create_task(
|
|
client.get(
|
|
"/api/kamidso/v1/chat-stream",
|
|
headers={"Authorization": f"Bearer {token_a}"},
|
|
timeout=5.0,
|
|
)
|
|
)
|
|
# Let the stream emit the first chunk, then hit auth as another user.
|
|
await asyncio.sleep(0.05)
|
|
session = await client.get(
|
|
"/api/kamidso/v1/auth/session",
|
|
headers={"Authorization": f"Bearer {token_b}"},
|
|
timeout=5.0,
|
|
)
|
|
assert session.status_code == 200
|
|
assert session.json()["user"]["id"] == "user_b"
|
|
release.set()
|
|
response = await stream_task
|
|
|
|
assert response.status_code == 200
|
|
body = response.text
|
|
assert '"n": 0' in body
|
|
assert '"n": 5' in body
|
|
assert chunks_sent == 6
|