154 lines
8.3 KiB
Python
154 lines
8.3 KiB
Python
"""ASGI streaming proxy. Preserve authentication, disconnects and SSE headers."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import logging
|
|
import re
|
|
import time
|
|
from urllib.parse import quote
|
|
|
|
import httpx
|
|
from starlette.requests import Request
|
|
from starlette.responses import JSONResponse, StreamingResponse
|
|
|
|
from .controller import PoolUnavailable
|
|
|
|
THREAD_PATH = re.compile(r"^/api/threads/([^/]+)(?:/.*)?$")
|
|
START_PATH = re.compile(r"^/api/threads/[^/]+/runs(?:/(?:stream|wait))?$")
|
|
HOP_HEADERS = {b"host", b"connection", b"keep-alive", b"transfer-encoding", b"upgrade", b"proxy-authenticate", b"proxy-authorization", b"te", b"trailer", b"content-length", b"x-deerflow-worker-token"}
|
|
|
|
|
|
def _main_has_persistent_chat_storage(app) -> bool:
|
|
checkpointer = getattr(app.state, "checkpointer", None)
|
|
store = getattr(app.state, "store", None)
|
|
if checkpointer is None or store is None:
|
|
return False
|
|
from langgraph.checkpoint.memory import InMemorySaver
|
|
from langgraph.store.memory import InMemoryStore
|
|
|
|
return not isinstance(checkpointer, InMemorySaver) and not isinstance(store, InMemoryStore)
|
|
|
|
|
|
async def send_upstream(client, method, url, headers, content):
|
|
# A keep-alive socket may close while a completed run is being queried.
|
|
# Retry only body-free reads, before downstream headers have been sent.
|
|
attempts = 2 if method in {"GET", "HEAD"} else 1
|
|
for attempt in range(attempts):
|
|
upstream = client.build_request(method, url, headers=headers, content=None if method in {"GET", "HEAD"} else content)
|
|
try:
|
|
return await client.send(upstream, stream=True)
|
|
except httpx.TransportError:
|
|
if attempt + 1 == attempts:
|
|
raise
|
|
await asyncio.sleep(0.05)
|
|
|
|
|
|
class WorkerProxy:
|
|
def __init__(self, app):
|
|
self.app = app
|
|
|
|
async def __call__(self, scope, receive, send):
|
|
if scope["type"] != "http":
|
|
return await self.app(scope, receive, send)
|
|
pool = getattr(scope["app"].state, "worker_pool", None)
|
|
match = THREAD_PATH.fullmatch(scope["path"])
|
|
if pool is None or match is None or match[1] in {"search", "count", "reference-debug-log"}:
|
|
return await self.app(scope, receive, send)
|
|
if pool.db is None and not pool.store.load().enabled:
|
|
return await self.app(scope, receive, send)
|
|
row = pool._ensure_db().execute("SELECT node FROM affinity WHERE thread=?", (match[1],)).fetchone()
|
|
creating = scope["method"] == "POST" and bool(START_PATH.fullmatch(scope["path"]))
|
|
if row is None and (not creating or not pool.store.load().enabled):
|
|
return await self.app(scope, receive, send)
|
|
# Authenticate before allocating durable routing state. The worker
|
|
# still performs the existing owner/permission checks for every route.
|
|
request = Request(scope, receive)
|
|
from fastapi import HTTPException
|
|
|
|
from app.gateway.deps import get_current_user_from_request, get_thread_store
|
|
|
|
try:
|
|
access_started = time.perf_counter()
|
|
user = getattr(request.state, "user", None) or await get_current_user_from_request(request)
|
|
allowed = await get_thread_store(request).check_access(match[1], str(user.id), require_existing=True)
|
|
if not allowed:
|
|
return await JSONResponse({"detail": "Thread not found"}, 404)(scope, receive, send)
|
|
if row is None:
|
|
from deerflow.persistence.thread_meta.base import is_excluded_system_thread
|
|
|
|
record = await get_thread_store(request).get(match[1], user_id=None)
|
|
if record and is_excluded_system_thread(record.get("metadata")):
|
|
return await self.app(scope, receive, send)
|
|
manager = getattr(scope["app"].state, "run_manager", None)
|
|
if creating and pool.local_threads.get(match[1], 0):
|
|
return await JSONResponse({"detail": "会话仍在主服务执行，请完成当前轮后再提交"}, 409)(scope, receive, send)
|
|
if creating and manager is not None and any(
|
|
record.thread_id == match[1]
|
|
and getattr(record.status, "value", record.status) in {"pending", "running"}
|
|
and (record.task is None or not record.task.done())
|
|
for record in manager._runs.values()
|
|
):
|
|
return await JSONResponse({"detail": "会话仍在执行，请完成当前轮后再提交"}, 409)(scope, receive, send)
|
|
node = await pool.reserve(match[1], create=creating)
|
|
scope.setdefault("state", {})["worker_access_ms"] = round((time.perf_counter() - access_started) * 1000, 2)
|
|
except HTTPException as exc:
|
|
return await JSONResponse({"detail": exc.detail}, exc.status_code)(scope, receive, send)
|
|
except PoolUnavailable as exc:
|
|
if creating and pool.local_threads.get(match[1], 0):
|
|
return await JSONResponse({"detail": str(exc)}, 409)(scope, receive, send)
|
|
return await JSONResponse({"detail": str(exc)}, 503, headers={"Retry-After": "3"})(scope, receive, send)
|
|
if node is None:
|
|
return await self.app(scope, receive, send)
|
|
if node.get("_fallback_to_main"):
|
|
fallback_worker = node.get("_worker")
|
|
if creating and fallback_worker is not None:
|
|
try:
|
|
active_response = await pool.http.get(
|
|
fallback_worker["url"] + "/internal/worker-thread-active/" + quote(match[1], safe=""),
|
|
headers={"x-deerflow-worker-token": pool.store.raw()["token"]},
|
|
timeout=2,
|
|
)
|
|
active_response.raise_for_status()
|
|
if active_response.json().get("active"):
|
|
return await JSONResponse({"detail": "该会话正在 Worker 执行，请等待当前回答完成"}, 409)(scope, receive, send)
|
|
except httpx.HTTPError:
|
|
return await JSONResponse({"detail": "无法确认该会话是否正在 Worker 执行，请稍后重试"}, 503, headers={"Retry-After": "2"})(scope, receive, send)
|
|
if not _main_has_persistent_chat_storage(scope["app"]):
|
|
return await JSONResponse(
|
|
{"detail": "主服务的持久化聊天存储当前不可用，无法安全接管 Worker 请求"},
|
|
503,
|
|
headers={"Retry-After": "3"},
|
|
)(scope, receive, send)
|
|
scope.setdefault("state", {})["worker_pool_fallback"] = True
|
|
scope["state"]["worker_pool_fallback_reason"] = node.get("_fallback_reason")
|
|
return await self.app(scope, receive, send)
|
|
response = None
|
|
scope.setdefault("state", {})["worker_node"] = node["name"]
|
|
try:
|
|
headers = [(k, v) for k, v in scope["headers"] if k.lower() not in HOP_HEADERS or k.lower() == b"host"]
|
|
headers.append((b"x-deerflow-worker-token", pool.store.raw()["token"].encode()))
|
|
url = node["url"] + scope["path"]
|
|
if scope.get("query_string"):
|
|
url += "?" + scope["query_string"].decode("ascii")
|
|
upstream_started = time.perf_counter()
|
|
response = await send_upstream(pool.http, scope["method"], url, headers, request.stream())
|
|
scope["state"]["worker_upstream_headers_ms"] = round((time.perf_counter() - upstream_started) * 1000, 2)
|
|
|
|
async def chunks():
|
|
async for chunk in response.aiter_raw():
|
|
yield chunk
|
|
|
|
downstream = StreamingResponse(chunks(), status_code=response.status_code)
|
|
downstream.raw_headers = [(k, v) for k, v in response.headers.raw if k.lower() not in HOP_HEADERS]
|
|
await downstream(scope, receive, send)
|
|
except httpx.HTTPError as exc:
|
|
logging.getLogger(__name__).warning("Worker transport failed method=%s node=%s error=%s", scope["method"], node["name"], type(exc).__name__)
|
|
if response is None:
|
|
await JSONResponse({"detail": "执行节点连接失败。任务可能已开始，请查看会话状态后再决定是否重试。"}, 503)(scope, receive, send)
|
|
# Once SSE headers are sent, close the stream so the SDK reconnects.
|
|
finally:
|
|
if response is not None:
|
|
await response.aclose()
|
|
pool.release(node["name"])
|