Files
ai_agent/backend/app/core/middleware.py
T
furyhawk 8351e73d39 feat: add Zustand stores for conversation, file preview, sidebar, theme, and knowledge base selection
- Implemented `conversation-store` for managing conversations and messages.
- Created `file-preview-store` to handle file preview state.
- Added `sidebar-store` for sidebar visibility management.
- Developed `theme-store` for theme persistence and management.
- Introduced `kb-selection-store` for managing active knowledge base selections with persistence.

chore: define API and chat types

- Added types for API responses, authentication, chat messages, conversations, and projects.
- Defined interfaces for various entities including users, sessions, and message ratings.

build: configure TypeScript and testing setup

- Set up `tsconfig.json` for TypeScript configuration.
- Created `vitest.config.ts` for testing configuration with Vitest.
- Added `vitest.setup.ts` for global test setup including mocks for Next.js router and media queries.
- Configured Vercel deployment settings in `vercel.json`.
2026-06-11 16:54:43 +08:00

103 lines
3.6 KiB
Python

"""Application middleware."""
from typing import ClassVar
from uuid import uuid4
from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
from starlette.requests import Request
from starlette.responses import Response
from starlette.types import ASGIApp
class RequestIDMiddleware(BaseHTTPMiddleware):
"""Middleware that adds a unique request ID to each request.
The request ID is taken from the X-Request-ID header if present,
otherwise a new UUID is generated. The ID is added to the response
headers and is available in request.state.request_id.
"""
async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> Response:
"""Add request ID to request state and response headers."""
request_id = request.headers.get("X-Request-ID", str(uuid4()))
request.state.request_id = request_id
response = await call_next(request)
response.headers["X-Request-ID"] = request_id
return response
class SecurityHeadersMiddleware(BaseHTTPMiddleware):
"""Middleware that adds security headers to all responses.
This includes:
- Content-Security-Policy (CSP)
- X-Content-Type-Options
- X-Frame-Options
- X-XSS-Protection
- Referrer-Policy
- Permissions-Policy
Usage:
app.add_middleware(SecurityHeadersMiddleware)
# Or with custom CSP:
app.add_middleware(
SecurityHeadersMiddleware,
csp_directives={
"default-src": "'self'",
"script-src": "'self' 'unsafe-inline'",
}
)
"""
DEFAULT_CSP_DIRECTIVES: ClassVar[dict[str, str]] = {
"default-src": "'self'",
"script-src": "'self'",
"style-src": "'self' 'unsafe-inline'", # Allow inline styles for some UI libs
"img-src": "'self' data: https:",
"font-src": "'self' data:",
"connect-src": "'self'",
"frame-ancestors": "'none'",
"base-uri": "'self'",
"form-action": "'self'",
}
def __init__(
self,
app: ASGIApp,
csp_directives: dict[str, str] | None = None,
exclude_paths: set[str] | None = None,
) -> None:
super().__init__(app)
self.csp_directives = csp_directives or self.DEFAULT_CSP_DIRECTIVES
self.exclude_paths = exclude_paths or {"/docs", "/redoc", "/openapi.json"}
async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> Response:
"""Add security headers to the response."""
response = await call_next(request)
# Skip for docs/openapi endpoints which need different CSP
if request.url.path in self.exclude_paths:
return response
# Build CSP header
csp_value = "; ".join(
f"{directive} {value}" for directive, value in self.csp_directives.items()
)
# Add security headers — respect any already set by the route, so an
# endpoint can opt into less restrictive framing (e.g. user-content
# files served inline for the chat preview panel).
response.headers.setdefault("Content-Security-Policy", csp_value)
response.headers.setdefault("X-Content-Type-Options", "nosniff")
response.headers.setdefault("X-Frame-Options", "DENY")
response.headers.setdefault("X-XSS-Protection", "1; mode=block")
response.headers.setdefault("Referrer-Policy", "strict-origin-when-cross-origin")
response.headers["Permissions-Policy"] = (
"accelerometer=(), camera=(), geolocation=(), gyroscope=(), "
"magnetometer=(), microphone=(), payment=(), usb=()"
)
return response