"""SSE over POST with bounded workers, buffering and per-session exclusion. Run ONE process per Redis session namespace. A disconnected request retains its slot and session exclusion until the synchronous graph worker actually exits. """ import asyncio import json import logging import re import time from concurrent.futures import ThreadPoolExecutor from contextlib import asynccontextmanager from queue import Empty, Full, Queue from threading import Event, Lock from uuid import uuid4 from starlette.applications import Starlette from starlette.requests import ClientDisconnect from starlette.responses import JSONResponse, StreamingResponse from starlette.routing import Route from .config import load_settings from .runtime import Runtime from .company_classification import ClassificationService, classify_company from .clock import shanghai_now from .logging import configure_logging, log, safe_code, session_ref, summary_diagnostics, candidate_counts, validation_counts from .progress import STAGES, progress from .attachments import AttachmentError, resolve_arrays from .memory import Admission from src.step3_question_answer.timing import Timings, bind def sse(event, data): return f'event: {event}\ndata: {json.dumps(data, ensure_ascii=False)}\n\n' class Job: def __init__(self, thread_id, timeout, request_id=None, received_at=None, received_monotonic=None): self.thread_id = thread_id self.request_id = request_id or uuid4().hex self.received_at = received_at or shanghai_now() self.session_ref = session_ref(thread_id) self.outcome = 'running' self.events = Queue(maxsize=8) self.stopped = Event() self.finished = Event() self.started_at = time.monotonic() self.timings = Timings(self.started_at if received_monotonic is None else received_monotonic) self.deadline = self.started_at + timeout def emit(self, event, data): if self.stopped.is_set(): return if event == 'progress' and data.get('stage') in STAGES and data.get('status') in ('running', 'completed', 'waiting_input', 'failed'): info = progress(data['stage'], data['status']) self.log('stage', info.pop('message'), **info) elif event == 'error': self.outcome = 'failed' self.log('request_failed', '本轮处理失败', level=logging.ERROR, code=safe_code(data.get('code'))) elif event == 'result': recommendation = data.get('recommendation') or {} errors = data.get('errors') or [] analysis_errors = recommendation.get('analysis_errors') or [] summary_status = (recommendation.get('summary') or {}).get('status') highlight_errors = recommendation.get('support_highlight_errors') or [] degraded = bool(errors or analysis_errors or highlight_errors or recommendation.get('complete') is False or summary_status == 'fallback') if degraded: self.outcome = 'degraded' self.log('result_ready', '本轮结果已生成', level=logging.WARNING if degraded else logging.INFO, degraded=degraded, error_count=len(errors), analysis_error_count=len(analysis_errors), candidate_outcome_counts=candidate_counts(recommendation.get('candidate_diagnostics')), support_highlight_error_codes=[safe_code(e.get('code')) for e in highlight_errors[:16]], validation_error_counts=validation_counts(recommendation.get('candidate_diagnostics')), error_codes=[safe_code(c) for c in errors[:16]], analysis_error_codes=[safe_code(e.get('code')) for e in analysis_errors[:16]], summary_error_code=safe_code((recommendation.get('summary') or {}).get('error')) if (recommendation.get('summary') or {}).get('error') else None, summary_diagnostics=summary_diagnostics((recommendation.get('summary') or {}).get('diagnostics')), planning_error_code=safe_code((recommendation.get('search_plan') or {}).get('error')) if (recommendation.get('search_plan') or {}).get('error') else None, summary_status=summary_status if summary_status in ('generated', 'fallback', 'no_recommendations') else None) elif event == 'done': if data.get('status') == 'needs_input': self.outcome = 'needs_input' elif self.outcome == 'running': self.outcome = 'completed' frame = sse(event, {'thread_id': self.thread_id, 'request_id': self.request_id, 'data': data}) enqueue_started = self.timings.clock() while not self.stopped.is_set(): if time.monotonic() >= self.deadline: self.stopped.set() return try: self.events.put(frame, timeout=0.1) self.timings.add('event_queue', self.timings.clock() - enqueue_started) if event == 'item': self.timings.mark('first_card') elif event in ('answer_delta', 'answer', 'summary') and isinstance(data.get('text'), str) and data['text'].strip(): self.timings.mark('first_text') elif event == 'result': self.timings.mark('result_ready') elif event in ('done', 'error'): self.timings.mark('terminal_ready') return except Full: continue def log(self, event, message, **fields): log(event, message, request_id=self.request_id, session_ref=self.session_ref, elapsed_ms=round((time.monotonic() - self.started_at) * 1000), **fields) class Service: def __init__(self, runtime, settings): self.runtime, self.settings = runtime, settings self.admission = getattr(runtime, 'admission', None) or Admission(settings) self.pool = ThreadPoolExecutor(max_workers=settings.max_concurrent, thread_name_prefix='chat') self.lock = Lock() self.active = {} self.closing = False def submit(self, thread_id, question, request_id=None, received_at=None, pictures=None, files=None, received_monotonic=None): with self.lock: if thread_id in self.active: return None, 409, 'thread_busy' if self.closing or len(self.active) >= self.settings.max_concurrent: return None, 429, 'server_busy' lease, error = self.admission.acquire() if error: return None, 503 if error == 'memory_status_unavailable' else 429, error try: job = Job(thread_id, self.settings.stream_timeout_seconds, request_id, received_at, received_monotonic) job.memory_lease = lease self.active[thread_id] = job job.log('request_accepted', '请求已接纳', active_requests=len(self.active)) self.pool.submit(self._run, job, question, pictures or [], files or []) except Exception: self.active.pop(thread_id, None) lease.release() return None, 503, 'service_unavailable' return job, 200, None def _run(self, job, question, pictures=None, files=None): job.timings.add('worker_queue', job.timings.clock() - job.started_at) job.timings.mark('worker_started') job.log('worker_started', '开始处理问题') try: extra = {'pictures': pictures, 'files': files} if pictures or files else {} with bind(job.timings): self.runtime.run(job.thread_id, question, job.emit, job.stopped, request_id=job.request_id, received_at=job.received_at, **extra) except Exception: job.emit('error', {'code': 'processing_failed', 'message': '处理失败，请稍后重试或新建会话。'}) finally: # Never free the lock merely because the HTTP caller disconnected. with self.lock: self.active.pop(job.thread_id, None) job.memory_lease.release() job.log('request_performance', '本轮耗时统计', metrics=job.timings.snapshot(), thread_id=job.thread_id, outcome=job.outcome if job.outcome != 'running' else 'stopped' if job.stopped.is_set() else 'incomplete') job.log('worker_finished', '后台处理结束，会话锁已释放', outcome=job.outcome if job.outcome != 'running' else 'stopped' if job.stopped.is_set() else 'incomplete') job.finished.set() async def stream(self, job): heartbeat = time.monotonic() reason = 'client_disconnected' try: yield sse('accepted', {'thread_id': job.thread_id, 'request_id': job.request_id, 'data': {'message': '收到，我先了解一下您的问题。'}}) while True: if time.monotonic() >= job.deadline: reason = 'timeout' job.log('stream_timeout', '流响应超时', level=logging.WARNING, code='stream_timeout') yield sse('error', {'thread_id': job.thread_id, 'request_id': job.request_id, 'data': {'code': 'stream_timeout'}}) return try: frame = job.events.get_nowait() except Empty: if job.finished.is_set(): if reason != 'terminal_event': reason = 'worker_ended_without_terminal' return if time.monotonic() - heartbeat >= self.settings.heartbeat_seconds: yield sse('heartbeat', {'thread_id': job.thread_id, 'request_id': job.request_id, 'data': {'status': 'processing', 'message': '还在为您准备答复，请再稍等一会儿。', 'elapsed_seconds': int(time.monotonic() - job.started_at)}}) heartbeat = time.monotonic() await asyncio.sleep(0.02) else: if frame.startswith(('event: done\n', 'event: error\n')): reason = 'terminal_event' yield frame finally: job.stopped.set() job.log('stream_closed', '流连接结束', reason=reason, worker_finished=job.finished.is_set()) def close(self): with self.lock: self.closing = True for job in self.active.values(): job.stopped.set() self.pool.shutdown(wait=True) self.runtime.close() def create_app(settings=None, runtime_factory=Runtime): settings = settings or load_settings() @asynccontextmanager async def lifespan(app): configure_logging() log('service_starting', '正在初始化问答服务') try: runtime = await asyncio.to_thread(runtime_factory, settings) except Exception: log('service_startup_failed', '问答服务初始化失败', level=logging.ERROR) raise app.state.service = Service(runtime, settings) classify = getattr(runtime, 'classify_company', None) app.state.classification_service = ClassificationService( classify, settings.max_concurrent, app.state.service.admission) if classify else None app.state.classification_timeout = settings.stream_timeout_seconds log('service_ready', '问答服务已就绪', max_concurrent_requests=settings.max_concurrent) try: yield finally: log('service_stopping', '正在停止问答服务') if app.state.classification_service is not None: await asyncio.to_thread(app.state.classification_service.close) await asyncio.to_thread(app.state.service.close) log('service_stopped', '问答服务已停止') async def chat(request): request_id = uuid4().hex received_at = time.monotonic() received_time = shanghai_now() log('request_received', '收到问答请求', request_id=request_id, method='POST', path='/api/chat') def reject(code, status, thread_id=None, retry=False, detail=None): log('request_rejected', '请求未被接纳', level=logging.WARNING, request_id=request_id, session_ref=session_ref(thread_id) if thread_id is not None else None, status_code=status, code=code, elapsed_ms=round((time.monotonic() - received_at) * 1000)) headers = {'X-Request-ID': request_id} if retry: headers['Retry-After'] = '2' return JSONResponse({'error': code, **(detail or {})}, status_code=status, headers=headers) if request.headers.get('content-type', '').split(';')[0].strip().lower() != 'application/json': return reject('application_json_required', 415) try: body = bytearray() async with asyncio.timeout(15): async for chunk in request.stream(): if len(body) + len(chunk) > 65536: return reject('request_too_large', 413) body.extend(chunk) value = json.loads(body) if (not isinstance(value, dict) or not {'thread_id', 'question'} <= set(value) or set(value) - {'thread_id', 'question', 'pictures', 'files'}): raise ValueError() pictures, files = value.get('pictures', []), value.get('files', []) pictures, files = resolve_arrays(pictures, files, settings.attachment_base_url) thread_id, question = value['thread_id'], value['question'] if (not isinstance(thread_id, str) or not re.fullmatch(r'[A-Za-z0-9_-]{1,128}', thread_id) or not isinstance(question, str) or not question.strip() or len(question) > 8000): raise ValueError() except TimeoutError: return reject('request_timeout', 408) except AttachmentError as exc: kind, index = getattr(exc, 'kind', None), getattr(exc, 'index', None) if kind and index: title = ('图片' if kind == 'picture' else '附件') + str(index) return reject('invalid_request', 422, detail={ 'message': title + '的链接格式无效，请更换后重新提问。', 'attachment': {'kind': kind, 'index': index, 'title': title, 'error': 'invalid_attachment_url'}}) return reject('invalid_request', 422) except (ValueError, UnicodeError, ClientDisconnect): return reject('invalid_request', 422) job, status, error = request.app.state.service.submit(thread_id, question, request_id, received_time, pictures=pictures, files=files, received_monotonic=received_at) if error: return reject(error, status, thread_id, retry=True) return StreamingResponse(request.app.state.service.stream(job), media_type='text/event-stream', headers={'Cache-Control': 'no-cache, no-transform', 'X-Accel-Buffering': 'no', 'X-Request-ID': request_id}) async def health(request): return JSONResponse({'status': 'ready'}) return Starlette(routes=[Route('/api/chat', chat, methods=['POST']), Route('/api/company/classify', classify_company, methods=['POST']), Route('/health', health)], lifespan=lifespan)