"""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 typing import Any, Optional 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 ..task_progress import latest_run_progress 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: Optional[int] = None, before_idx: Optional[int] = None, after_idx: Optional[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, Message.artifact_refs, Message.attachment_refs, ) 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 progress_snapshot = latest_run_progress(s, tid) return { "has_more": has_more, "has_more_after": has_more_after, "progress_snapshot": progress_snapshot, "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), "artifact_refs": r.artifact_refs, "attachment_refs": r.attachment_refs, } 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 and not body.attachments: 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) # Omitted means a legacy client; explicit [] means a new message has no attachments. attachments_supplied = "attachments" in body.model_fields_set attachment_refs: Optional[list[dict]] = [] if attachments_supplied else None if body.attachments: from core.artifacts import ArtifactPathError from core.attachments import normalize_attachment_refs from core.paths import from_db_path root = load_user_root(user_id) with session_scope() as s: task_working_dir = s.execute( select(Task.working_dir).where( Task.task_id == tid, Task.user_id == user_id, ) ).scalar_one_or_none() if not task_working_dir: raise HTTPException(404, f"task not found: {tid}") working_dir = from_db_path(task_working_dir).resolve() try: attachment_refs = normalize_attachment_refs( body.attachments, working_dir=working_dir, user_root=root, ) except ArtifactPathError as e: raise HTTPException(400, str(e)) 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, attachment_refs=attachment_refs, ) 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"] if should_auto_title: from core.task_title import is_attachment_only_message should_auto_title = bool(content) and not is_attachment_only_message(content) # 快速入口只在首条消息时 pending=true。辅助标题与主 run 并行，不阻塞 # TTFT/SSE；纯附件消息保留 pending，等下一条自然语言再命名，避免标题暴露 # `[用户上传的文件]` 等内部标记。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", "run_id": str(claim.run_id), } @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, 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 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, "next_message_idx": 0, # 全删后 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: # Durable catch-up for events emitted before this subscriber existed. # Subscribing first keeps updates emitted during the DB read queued. with session_scope() as s: snapshot = latest_run_progress(s, tid) if snapshot is not None: yield sse_event("progress_snapshot", snapshot) 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", }, )