"""Messages 路由:历史/目录/发消息起 run/取消/清空/润色/bg proc/SSE。 drain 背压与 Web 特有模型门控在这里；run 抢占/调度收口在 web/run_lifecycle.py，BG worker 本体在 web/runs.py。 """ from __future__ import annotations import asyncio import time from uuid import UUID from fastapi import Depends, HTTPException from fastapi.responses import StreamingResponse from sqlalchemy import select, update from core.storage import session_scope from core.storage.models import Message, Task, User from ..broker import broker from ..common import ( assert_owns_task, iso, outline_snippet, sse_event, ) from ..model_gate import ( FALLBACK_MODEL_PROFILE, list_image_variants, list_video_variants, model_allowed_for_user, resolve_image_model, resolve_model_profile, resolve_video_model, skill_pinned_profiles, ) from ..run_lifecycle import ( RunScheduleError, RunTaskBusy, RunTaskNotFound, claim_run_with_message, schedule_claimed_run, ) from ..schemas import MessageRequest, OptimizePromptRequest from ..userfiles import load_user_root def _proc_view(m: dict) -> dict: """core.procs.list_* 的条目 → API 形态。elapsed:running 现算,finished 用 exit_code mtime 定格(前端秒数展示与工具卡同体验)。""" from core import procs as _procs d = m.get("_dir") st, ec = m.get("_status"), m.get("_exit_code") created = float(m.get("created_ts") or 0) or None now = time.time() if st == "finished" and d is not None: fin = _procs.finished_at(d) or created or now elapsed = max(0.0, fin - (created or fin)) else: elapsed = max(0.0, now - created) if created else 0.0 return { "task_id": m.get("task_id"), "proc_id": m.get("proc_id"), "kind": m.get("kind"), "command": m.get("command") or "", "status": st, "exit_code": ec, "killed": bool(m.get("killed")), "created_at": m.get("created_at"), "elapsed_s": int(elapsed), "timeout_s": m.get("timeout_s"), } def register_message_routes(app, *, require_user) -> None: @app.get("/v1/tasks/{task_id}/messages", tags=["messages"]) def list_messages( task_id: str, limit: int = None, before_idx: int = None, after_idx: int = None, user_id: UUID = Depends(require_user), ): """task 历史消息(idx 升序);LiteLLM 原 payload 透传给前端,自行渲染。 分页(双向窗口): - 不传 limit → 升序全量返回(向后兼容旧前端),两个 has_more 都 false。 - 传 limit(默认)→ 取**尾部**最近 limit 条(idx desc + limit 再 reverse 回升序)。 - 传 before_idx → 只取 idx < before_idx 的更早部分(向上翻页)。 - 传 after_idx → 只取 idx > after_idx 的更新部分(向下翻页;从目录跳到旧消息后用)。 响应恒含 has_more(窗口之前是否还有更早)+ has_more_after(窗口之后是否还有更新)。 """ try: tid = UUID(task_id) except ValueError: raise HTTPException(404, f"invalid task id: {task_id!r}") with session_scope() as s: assert_owns_task(s, tid, user_id) cols = ( Message.idx, Message.payload, Message.tokens_in, Message.tokens_out, Message.model_profile, Message.created_at, ) if limit is None: # 旧行为:升序全量 rows = s.execute( select(*cols).where(Message.task_id == tid).order_by(Message.idx) ).all() elif after_idx is not None: # 向下窗口:正序取 idx > after_idx 的最早 limit 条 rows = list(s.execute( select(*cols) .where(Message.task_id == tid, Message.idx > after_idx) .order_by(Message.idx).limit(limit) ).all()) else: # 尾部 / 向上窗口:倒序取 limit 条,再翻回升序 q = select(*cols).where(Message.task_id == tid) if before_idx is not None: q = q.where(Message.idx < before_idx) rows = list(s.execute(q.order_by(Message.idx.desc()).limit(limit)).all()) rows.reverse() # 窗口两端外是否还有(供前端顶/底 sentinel 决定要不要继续补) has_more = False has_more_after = False if rows: first_idx = rows[0].idx last_idx = rows[-1].idx has_more = s.execute( select(Message.idx) .where(Message.task_id == tid, Message.idx < first_idx) .limit(1) ).first() is not None has_more_after = s.execute( select(Message.idx) .where(Message.task_id == tid, Message.idx > last_idx) .limit(1) ).first() is not None return { "has_more": has_more, "has_more_after": has_more_after, "messages": [ { "idx": r.idx, "payload": dict(r.payload), "tokens_in": r.tokens_in, "tokens_out": r.tokens_out, "model_profile": r.model_profile, # 0006:assistant 行非空,标产生该 msg 的模型 "created_at": iso(r.created_at), } for r in rows ] } @app.get("/v1/tasks/{task_id}/outline", tags=["messages"]) def task_outline(task_id: str, user_id: UUID = Depends(require_user)): """消息目录:全部 user 轮次的 {idx, snippet}(idx 升序),供右侧圆点轨道导航。 只取 role=user 的 idx + content 首行片段,不回传整 payload(轻量,长任务也快); 走 (task_id, idx) 索引按 task 收窄,role 过滤为残余条件。前端点圆点 → 已加载则 scrollIntoView,未加载则用 before_idx 拉居中窗口再定位。 """ try: tid = UUID(task_id) except ValueError: raise HTTPException(404, f"invalid task id: {task_id!r}") with session_scope() as s: assert_owns_task(s, tid, user_id) rows = s.execute( select(Message.idx, Message.payload["content"].astext) .where( Message.task_id == tid, Message.payload["role"].astext == "user", ) .order_by(Message.idx) ).all() return { "items": [ {"idx": r[0], "snippet": outline_snippet(r[1])} for r in rows ] } @app.post("/v1/tasks/{task_id}/messages", status_code=202, tags=["messages"]) async def post_message( task_id: str, body: MessageRequest, user_id: UUID = Depends(require_user), ): """发消息 + 起 BG run。返 `{events_url}`,客户端立刻订阅 SSE 拿流式。 单活 run:`SELECT … FOR UPDATE` 锁 task 行 + 活跃状态检查 + 标 running, 全收进一个事务挡住"用户连点 send 两条消息"导致两个 BG 线程争 `messages.idx`。 tasks.run_status in ('running','cancelling') → 409;'error' 走起新 run 时清掉 (跟 ok / cancelled 一样视为可重启)。 """ try: tid = UUID(task_id) except ValueError: raise HTTPException(404, f"invalid task id: {task_id!r}") # 关停 drain 期:拒新 run,带 Retry-After 让客户端退避重试(部署窗口背压)。 if getattr(app.state, "draining", None) is not None and app.state.draining.is_set(): raise HTTPException( 503, "server is restarting; retry shortly", headers={"Retry-After": "3"}, ) content = (body.content or "").strip() if not content: raise HTTPException(400, "empty content") # 快捷指令展开(与渠道入口共用 core/shortcuts.py):整条精确命中触发词 → 换成完整 # 指令。在起 run 之前、落库之前展开,模型看到的就是完整指令(不进上下文、不问模型)。 from core.agent_builder import resolve_workspace as _resolve_ws from core import shortcuts as _shortcuts _ws = await asyncio.to_thread(_resolve_ws, None) content, _sc_hit = await asyncio.to_thread(_shortcuts.expand, _ws, user_id, content) if _sc_hit: print(f"[shortcut] {str(user_id)[:8]} '{_sc_hit}' expanded") # 所有可能返回 4xx 的 variant 校验必须发生在事务写入之前。否则 task 已提交 # running 后才发现参数非法，会留下一个实际上没有 worker 的假活跃任务。 image_variant = resolve_image_model(body.image_model, user_id=user_id) video_variant = resolve_video_model(body.video_model, user_id=user_id) def _prepare_claim(s, task): values: dict = {} # 档位门控:存量 task 的模型已不在用户档位内(如管理员下调了档位)→ 本次起 # 持久落回 flash(基线必含),UI 下拉随之显示 flash。不报错、不打断会话历史, # 符合"老 task 下次发消息直接切 flash"。当前 task 模型仍在档内则原样不动。 # plan/role 用同一 session 读(避免在 FOR UPDATE 事务里再开嵌套 session)。 cur_profile = task.model_profile or "" if cur_profile: from core.model_access import is_allowed urow = s.execute( select(User.plan, User.role).where(User.user_id == user_id) ).first() plan = (urow.plan if urow else "") or "" role = (urow.role if urow else "user") or "user" # skill 定向模型豁免:档外但属内置 skill 定向(建 task 默认 / load_skill # 热切写入)→ 不降,产品决策放行。开关关闭时集合为空 → 恢复降级,存量 # 定向 task 自然落回 flash。 if not is_allowed(cur_profile, plan, role) and cur_profile not in skill_pinned_profiles(): fb_profile, fb_model_id = resolve_model_profile(FALLBACK_MODEL_PROFILE) values["model_profile"] = fb_profile values["model"] = fb_model_id return values, { "title_profile": values.get("model_profile", cur_profile), "should_auto_title": bool(task.auto_title_pending), } try: claim = claim_run_with_message( tid, user_id, content, prepare=_prepare_claim, ) except RunTaskNotFound: raise HTTPException(404, f"task not found: {tid}") except RunTaskBusy as e: raise HTTPException( 409, f"task already has an active run (status={e.status}); " f"wait for it to finish or cancel", ) try: schedule_claimed_run( app, tid, user_id, content, image_variant=image_variant, video_variant=video_variant, ) except RunScheduleError: raise HTTPException( 500, "message persisted, but background scheduling failed; retry this task", ) title_profile = claim.metadata["title_profile"] should_auto_title = claim.metadata["should_auto_title"] # 快速入口只在首条消息时 pending=true。辅助标题与主 run 并行，不阻塞 # TTFT/SSE；generate_task_title 内再次查闸并用条件 UPDATE 防人工改名竞态。 if should_auto_title: from core.task_title import generate_task_title_safe title_task = asyncio.create_task(asyncio.to_thread( generate_task_title_safe, task_id=tid, user_id=user_id, user_message=content, model_profile=title_profile, )) app.state.aux_tasks.add(title_task) title_task.add_done_callback(app.state.aux_tasks.discard) return {"events_url": f"/v1/tasks/{tid}/events"} @app.post("/v1/tasks/{task_id}/cancel", status_code=202, tags=["tasks"]) def cancel_task( task_id: str, user_id: UUID = Depends(require_user), ): """向当前 task 的活跃 run 发协作式 cancel 信号。 - 单活 run 形态下"取消当前活动"语义无歧义;客户端只需 task_id - 校验 task 归属 user;否则 404 - tasks.run_status 不是 `running` → 409(idle / cancelling / error 都不能 cancel) - 标 `cancelling`(过渡态),BG 线程 loop 在 stream chunk 间 + 工具调用之间 poll 看见即退; 退出后 finally 写终态(正常→idle,异常→error) - LLM 走 streaming,cancel 延迟 ~ 单 chunk 间隔(100ms 级) """ try: tid = UUID(task_id) except ValueError: raise HTTPException(404, f"invalid task id: {task_id!r}") with session_scope() as s: row = s.execute( select(Task.run_status, Task.title_source) .where(Task.task_id == tid, Task.user_id == user_id) .with_for_update() ).first() if row is None: raise HTTPException(404, f"task not found: {tid}") if row.run_status != "running": raise HTTPException( 409, f"task not running (run_status={row.run_status}); cannot cancel", ) s.execute( update(Task).where(Task.task_id == tid).values(run_status="cancelling") ) broker.request_cancel(tid) return {"ok": True, "task_id": str(tid), "run_status": "cancelling"} # ───────────── Background procs(bg proc,DESIGN §8.12)───────────── @app.get("/v1/procs", tags=["tasks"]) async def list_user_procs(user_id: UUID = Depends(require_user)): """当前用户全部 task 的后台进程(shell/run_python background=true 启动)。 用户级而非 task 级:前端全局轮询,切到别的 task 也能收到"后台任务完成"提示。 纯文件系统读取(user_root/.zcbot_procs),无 DB;docker backend 的 running 探测走 docker inspect,放 to_thread 防塞 event loop。 """ anchor = load_user_root(user_id) from core import procs as _procs items = await asyncio.to_thread(_procs.list_all_procs, anchor) return {"procs": [_proc_view(m) for m in items]} @app.post("/v1/tasks/{task_id}/procs/{proc_id}/kill", tags=["tasks"]) async def kill_task_proc( task_id: str, proc_id: str, user_id: UUID = Depends(require_user), ): """强制终止一个后台进程(host 杀进程树 / docker rm -f 容器)。幂等:已结束返 ok。""" try: tid = UUID(task_id) except ValueError: raise HTTPException(404, f"invalid task id: {task_id!r}") with session_scope() as s: assert_owns_task(s, tid, user_id) anchor = load_user_root(user_id) from core import procs as _procs d = _procs.proc_dir(anchor, str(tid), proc_id) if d is None or not d.is_dir(): raise HTTPException(404, f"background proc not found: {proc_id!r}") meta = _procs.read_meta(d) if meta is None: raise HTTPException(404, f"background proc metadata missing: {proc_id!r}") msg = await asyncio.to_thread(_procs.kill_proc, meta, d) return {"ok": True, "proc_id": proc_id, "message": msg} @app.post("/v1/tasks/{task_id}/clear", tags=["messages"]) def clear_messages(task_id: str, user_id: UUID = Depends(require_user)): """清空当前 task 全部 messages,token 累计 / cost / run_error 归零。 同 working_dir 下的 FS 文件不动(沿用 task delete 的"FS 视图可重生"心智 — 中间产物保留,模型重起对话时可继续基于已有素材推进)。 usage_events 不动:那是用户级账户级用量记账,不该被对话清理影响。 - 活跃 run(running / cancelling)期间拒绝:409(先 cancel) - error 状态可清:顺手 run_status='idle' + run_error=None - 跨 user → 404 """ try: tid = UUID(task_id) except ValueError: raise HTTPException(404, f"invalid task id: {task_id!r}") from sqlalchemy import delete as _delete from ..common import task_dict with session_scope() as s: row = s.execute( select(Task.run_status) .where(Task.task_id == tid, Task.user_id == user_id) .with_for_update() ).first() if row is None: raise HTTPException(404, f"task not found: {tid}") if row.run_status in ("running", "cancelling"): raise HTTPException( 409, f"task has an active run (status={row.run_status}); " f"cancel it first", ) s.execute(_delete(Message).where(Message.task_id == tid)) reset_values: dict[str, Any] = { "tokens_prompt": 0, "tokens_completion": 0, "cost_cny": 0, "run_status": "idle", "run_error": None, # 全删后 idx 从 0 重起,base 必须归零否则 load 窗口起点悬空(0019); # 摘要一并清,否则清空后的对话还会被注入旧前情摘要(0021) "context_base_idx": 0, "context_summary": None, } # 自动标题描述的是当前对话主题：清空后回到草稿名，下一条消息重新命名。 # manual/fixed 分别代表用户显式命名与渠道/调度固定名，均不得覆盖。 if row.title_source == "auto": reset_values.update( name="新对话", auto_title_pending=True, auto_title_version=Task.auto_title_version + 1, ) s.execute( update(Task).where(Task.task_id == tid).values(**reset_values) ) task_row = s.execute(select(Task).where(Task.task_id == tid)).scalar_one() d = task_dict(task_row, n_messages=0) return d @app.post("/v1/tasks/{task_id}/optimize_prompt", tags=["messages"]) def optimize_prompt( task_id: str, body: OptimizePromptRequest, user_id: UUID = Depends(require_user), ): """用 task 当前 model 润色用户草稿 prompt;返回优化后的文本。 - 同步调用(短文本,3-5s),非 stream - 不入 messages 表;**不**累计到 tasks.tokens_prompt/completion(顶栏数字保持 只反映主对话)。usage_events 单独写一行 kind="prompt_optimize",方便对账 + 按 kind GROUP BY 评估"这个按钮值不值" - 不与主对话 run 互斥(它不写 messages,无 idx 竞争)— 用户在 LLM 流式 回复期间也可润色下一条草稿 - image_model 影响 meta-prompt 里给 LLM 的下游 tool 提示;不动 DB """ from decimal import Decimal from core.agent_builder import load_config from core.capabilities import ModelCapabilities from core.llm import LLM from core.paths import ROOT from core.storage.models import UsageEvent from core.storage.usage import USD_TO_CNY try: tid = UUID(task_id) except ValueError: raise HTTPException(404, f"invalid task id: {task_id!r}") text = (body.text or "").strip() if not text: raise HTTPException(400, "empty text") if len(text) > 4000: raise HTTPException(400, "text too long (>4000 chars)") with session_scope() as s: row = s.execute( select(Task.model_profile) .where(Task.task_id == tid, Task.user_id == user_id) ).first() if row is None: raise HTTPException(404, f"task not found: {tid}") task_model_profile = row.model_profile or "" cfg = load_config() chosen_profile = task_model_profile or cfg["default_model"] # 档位门控:task 存量模型已不在用户档位内 → 润色也落回 flash(与 send 路径一致, # 不持久改 task,仅本次润色调用降级)。 if chosen_profile and not model_allowed_for_user(chosen_profile, user_id): chosen_profile = FALLBACK_MODEL_PROFILE try: caps = ModelCapabilities.load(chosen_profile, ROOT / cfg["models_dir"]) except (FileNotFoundError, ValueError) as e: raise HTTPException(500, f"invalid task model_profile {chosen_profile!r}: {e}") # 收集下游 tool 上下文:对话模型 display_name + 当前选中 image/video variant 元数据 chat_model_display = caps.display_name or chosen_profile image_variant_hint = "" img_variant = (body.image_model or "").strip() if img_variant: for k, v in list_image_variants(): if k == img_variant: name = v.get("display_name") or k sz = v.get("default_size") size_seg = f"默认尺寸 {sz}" if sz else "输出尺寸固定不可调" image_variant_hint = ( f"\n下游生图工具:{name}({size_seg},支持中英文 prompt," f"擅长写实/插画/构图描述)。若用户意图涉及画面/封面/插图," f"润色后的文本要给出适合该模型的画面细节(主体/风格/光线/构图)。" ) video_variant_hint = "" vid_variant = (body.video_model or "").strip() if vid_variant: for k, v in list_video_variants(): if k == vid_variant: name = v.get("display_name") or k res = v.get("default_resolution") or "720p" dur = v.get("default_duration") or 5 video_variant_hint = ( f"\n下游生视频工具:{name}(默认 {res} / {dur}s / 16:9)。" f"若用户意图涉及视频/动画/动起来,润色后的文本要补全 " f"主体在做什么(运动)+ 镜头怎么动 + 场景 + 风格,而非静态画面描述。" ) meta_prompt = ( f"你的任务是润色用户输入的草稿,使之成为一个清晰、完整、可执行的 prompt。\n" f"当前对话模型:{chat_model_display}。{image_variant_hint}{video_variant_hint}\n\n" f"规则:\n" f"1. 只输出润色后的文本本身,不要任何解释、前后缀、引号、markdown 代码块包裹\n" f"2. 保留用户原始语言(中文/英文)\n" f"3. 补全模糊点(主体、目标、风格、约束),但不要无中生有改变用户意图\n" f"4. 长度合理 — 简短诉求润色后也应当简洁,不要堆砌\n\n" f"用户草稿:\n{text}" ) llm = LLM(caps) try: response = llm.chat( messages=[{"role": "user", "content": meta_prompt}], tools=None, ) except Exception as e: raise HTTPException(502, f"llm call failed: {type(e).__name__}: {e}") try: optimized = (response.choices[0].message.content or "").strip() except Exception: raise HTTPException(502, "llm response missing content") if not optimized: raise HTTPException(502, "llm returned empty optimization") usage = getattr(response, "usage", None) prompt_tokens = int(getattr(usage, "prompt_tokens", 0) or 0) completion_tokens = int(getattr(usage, "completion_tokens", 0) or 0) try: from litellm import completion_cost cost_usd_raw = completion_cost(completion_response=response) cost_usd = Decimal(str(cost_usd_raw)) if cost_usd_raw else Decimal("0") except Exception: cost_usd = Decimal("0") cost_cny = (cost_usd * USD_TO_CNY).quantize(Decimal("0.000001")) try: with session_scope() as s: s.add(UsageEvent( user_id=user_id, task_id=tid, message_id=None, kind="prompt_optimize", model_profile=chosen_profile, units={ "tokens_in": prompt_tokens, "tokens_out": completion_tokens, "usd_to_cny": float(USD_TO_CNY), "image_model_hint": img_variant or "", "video_model_hint": vid_variant or "", }, cost_cny=cost_cny, )) except Exception as e: # 记账失败不阻塞返结果 — 用户拿到润色文本要紧,事后人工补 print(f"[optimize_prompt] usage record failed: {type(e).__name__}: {e}", flush=True) return { "optimized": optimized, "model_profile": chosen_profile, "tokens_in": prompt_tokens, "tokens_out": completion_tokens, "cost_cny": float(cost_cny), } @app.get("/v1/tasks/{task_id}/events", tags=["tasks"]) async def stream_events( task_id: str, user_id: UUID = Depends(require_user), ): """SSE 流。订阅当前 task 的活动 event(单活 run 形态下无歧义)。 事件类型:run_start / llm_start / text / reasoning / tool_call / tool_result / llm_end / cancelled / error / done。data 是 JSON dict (已剔除 `type` 字段,移到 event 名)。 """ try: tid = UUID(task_id) except ValueError: raise HTTPException(404, f"invalid task id: {task_id!r}") with session_scope() as s: assert_owns_task(s, tid, user_id) run_status = s.execute( select(Task.run_status).where(Task.task_id == tid) ).scalar_one() # 重连保护:若 task 不在活跃态(进程重启 / reaper 已收尾 / 自然结束), # 直接吐 done 关流。否则 broker 进程内队列空,客户端会无限挂在 ping 上。 is_active = run_status in ("running", "cancelling") async def gen(): yield b": connected\nretry: 3000\n\n" if not is_active: yield sse_event("done", {}) return q = broker.subscribe(tid) try: while True: try: ev = await asyncio.wait_for(q.get(), timeout=30.0) except asyncio.TimeoutError: yield b": ping\n\n" continue ev_type = ev.get("type", "msg") payload = {k: v for k, v in ev.items() if k != "type"} yield sse_event(ev_type, payload) if ev_type in ("done", "error"): break except asyncio.CancelledError: pass # 客户端断开,静默退 finally: broker.unsubscribe(tid, q) return StreamingResponse( gen(), media_type="text/event-stream", headers={ "Cache-Control": "no-cache", "Connection": "keep-alive", "X-Accel-Buffering": "no", }, )