|
1 | 1 | """Chat endpoints with SSE streaming support.""" |
2 | 2 |
|
3 | 3 | import asyncio |
| 4 | +import hmac |
4 | 5 | import json |
5 | 6 | import logging |
6 | 7 | import math |
| 8 | +import os |
7 | 9 | import time |
8 | 10 | import uuid |
9 | 11 |
|
10 | 12 | from fastapi import APIRouter, Request |
| 13 | +from fastapi.responses import JSONResponse |
11 | 14 | from starlette.responses import StreamingResponse |
12 | 15 |
|
13 | 16 | from pydantic_ai.messages import ( |
|
31 | 34 | ) |
32 | 35 | from models import ChatRequest, ChatResponse |
33 | 36 | from prompts.system import SYSTEM_PROMPT |
| 37 | +from services.org_pulse_client import PerRequestClient |
34 | 38 | from structured_logging import log_event |
35 | 39 | from tools import ALL_TOOLS, retry_events |
36 | 40 |
|
|
50 | 54 | "Please contact an administrator." |
51 | 55 | ) |
52 | 56 |
|
| 57 | +_AUTH_MISSING_MESSAGE = ( |
| 58 | + "Authentication headers missing. The chatbot must be accessed through the main application." |
| 59 | +) |
| 60 | + |
| 61 | + |
| 62 | +_EXPECTED_PROXY_SECRET = os.environ.get("PROXY_AUTH_SECRET", "") |
| 63 | + |
| 64 | + |
| 65 | +def _extract_auth(http_request: Request) -> tuple[str, str] | None: |
| 66 | + """Extract and validate proxy auth headers. Returns (proxy_secret, user_email) or None.""" |
| 67 | + proxy_secret = http_request.headers.get("x-proxy-secret", "").strip() |
| 68 | + user_email = http_request.headers.get("x-forwarded-email", "").strip() |
| 69 | + if not proxy_secret or not user_email: |
| 70 | + return None |
| 71 | + if _EXPECTED_PROXY_SECRET and not hmac.compare_digest(proxy_secret, _EXPECTED_PROXY_SECRET): |
| 72 | + return None |
| 73 | + return proxy_secret, user_email |
| 74 | + |
53 | 75 |
|
54 | 76 | # --------------------------------------------------------------------------- |
55 | 77 | # Helpers |
@@ -315,15 +337,20 @@ def _llm_error_message(exc: Exception) -> str: |
315 | 337 | async def chat(request: ChatRequest, http_request: Request): |
316 | 338 | request_id = uuid.uuid4().hex[:12] |
317 | 339 |
|
| 340 | + auth = _extract_auth(http_request) |
| 341 | + if not auth: |
| 342 | + return JSONResponse(status_code=401, content={"error": _AUTH_MISSING_MESSAGE}) |
| 343 | + proxy_secret, user_email = auth |
| 344 | + |
318 | 345 | agent = http_request.app.state.agent |
319 | 346 | if not agent: |
320 | 347 | return ChatResponse(message=_NOT_CONFIGURED_MESSAGE, trace={"request_id": request_id}) |
321 | 348 |
|
322 | 349 | retries: list[dict] = [] |
323 | 350 | retry_events.set(retries) |
324 | 351 |
|
325 | | - org_pulse_client = http_request.app.state.org_pulse_client |
326 | | - org_pulse_client.clear_cache() |
| 352 | + shared_client = http_request.app.state.org_pulse_client |
| 353 | + org_pulse_client = PerRequestClient(shared_client, proxy_secret, user_email) |
327 | 354 | message_history = _convert_history(request.history) |
328 | 355 | turn = _count_turns(request.history) + 1 |
329 | 356 |
|
@@ -455,6 +482,11 @@ async def chat(request: ChatRequest, http_request: Request): |
455 | 482 | async def chat_stream(request: ChatRequest, http_request: Request): |
456 | 483 | request_id = uuid.uuid4().hex[:12] |
457 | 484 |
|
| 485 | + auth = _extract_auth(http_request) |
| 486 | + if not auth: |
| 487 | + return JSONResponse(status_code=401, content={"error": _AUTH_MISSING_MESSAGE}) |
| 488 | + proxy_secret, user_email = auth |
| 489 | + |
458 | 490 | agent = http_request.app.state.agent |
459 | 491 | if not agent: |
460 | 492 | async def _not_configured(): |
@@ -515,8 +547,8 @@ async def _refusal_stream(): |
515 | 547 | else: |
516 | 548 | run_toolsets, tool_selection, selector_ms = await selector_coro |
517 | 549 |
|
518 | | - org_pulse_client = http_request.app.state.org_pulse_client |
519 | | - org_pulse_client.clear_cache() |
| 550 | + shared_client = http_request.app.state.org_pulse_client |
| 551 | + org_pulse_client = PerRequestClient(shared_client, proxy_secret, user_email) |
520 | 552 | message_history = _convert_history(request.history) |
521 | 553 | turn = _count_turns(request.history) + 1 |
522 | 554 |
|
|
0 commit comments