Agent/agent/trace/run_api.py

"""
Trace 控制 API — 新建 / 运行 / 停止 / 反思

提供 POST 端点触发 Agent 执行和控制。需要通过 set_runner() 注入 AgentRunner 实例。
执行在后台异步进行，客户端通过 WebSocket (/api/traces/{trace_id}/watch) 监听实时更新。

端点：
  POST /api/traces              — 新建 Trace 并执行
  POST /api/traces/{id}/run     — 运行（统一续跑 + 回溯）
  POST /api/traces/{id}/stop    — 停止运行中的 Trace
  POST /api/traces/{id}/reflect — 反思，在 trace 末尾追加反思 prompt 运行，结果追加到 experiences 文件
  GET  /api/traces/running      — 列出正在运行的 Trace
  GET  /api/experiences         — 读取经验文件内容
"""

import asyncio
import logging
import os
from datetime import datetime
from typing import Any, Dict, List, Optional

from fastapi import APIRouter, HTTPException
from pydantic import BaseModel, Field

logger = logging.getLogger(__name__)

router = APIRouter(prefix="/api/traces", tags=["run"])

# 经验 API 使用独立 prefix
experiences_router = APIRouter(prefix="/api", tags=["experiences"])


# ===== 全局 Runner（由 api_server.py 注入）=====

_runner = None


def set_runner(runner):
    """注入 AgentRunner 实例"""
    global _runner
    _runner = runner


def _get_runner():
    if _runner is None:
        raise HTTPException(
            status_code=503,
            detail="AgentRunner not configured. Server is in read-only mode.",
        )
    return _runner


# ===== Request / Response 模型 =====


class CreateRequest(BaseModel):
    """新建执行"""
    messages: List[Dict[str, Any]] = Field(..., description="OpenAI SDK 格式的输入消息")
    model: str = Field("gpt-4o", description="模型名称")
    temperature: float = Field(0.3)
    max_iterations: int = Field(200)
    system_prompt: Optional[str] = Field(None, description="自定义 system prompt（None = 从 skills 自动构建）")
    tools: Optional[List[str]] = Field(None, description="工具白名单（None = 全部）")
    name: Optional[str] = Field(None, description="任务名称（None = 自动生成）")
    uid: Optional[str] = Field(None)


class TraceRunRequest(BaseModel):
    """运行（统一续跑 + 回溯）"""
    messages: List[Dict[str, Any]] = Field(
        default_factory=list,
        description="追加的新消息（可为空，用于重新生成场景）",
    )
    insert_after: Optional[int] = Field(
        None,
        description="回溯插入点的 message sequence。None = 从末尾续跑，int = 回溯到该 sequence 后运行",
    )


class ReflectRequest(BaseModel):
    """反思请求"""
    focus: Optional[str] = Field(None, description="反思重点（可选）")


class RunResponse(BaseModel):
    """操作响应（立即返回，后台执行）"""
    trace_id: str
    status: str = "started"
    message: str = ""


class StopResponse(BaseModel):
    """停止响应"""
    trace_id: str
    status: str  # "stopping" | "not_running"


class ReflectResponse(BaseModel):
    """反思响应"""
    trace_id: str
    reflection: str


# ===== 后台执行 =====

_running_tasks: Dict[str, asyncio.Task] = {}


async def _run_in_background(trace_id: str, messages: List[Dict], config):
    """后台执行 agent，消费 run() 的所有 yield"""
    runner = _get_runner()
    try:
        async for _item in runner.run(messages=messages, config=config):
            pass  # WebSocket 广播由 runner 内部的 store 事件驱动
    except Exception as e:
        logger.error(f"Background run failed for {trace_id}: {e}")
    finally:
        _running_tasks.pop(trace_id, None)


async def _run_with_trace_signal(
    messages: List[Dict], config, trace_id_future: asyncio.Future,
):
    """后台执行 agent，通过 Future 将 trace_id 传回给等待的 endpoint"""
    from agent.trace.models import Trace

    runner = _get_runner()
    trace_id: Optional[str] = None
    try:
        async for item in runner.run(messages=messages, config=config):
            if isinstance(item, Trace) and not trace_id_future.done():
                trace_id = item.trace_id
                trace_id_future.set_result(trace_id)
    except Exception as e:
        if not trace_id_future.done():
            trace_id_future.set_exception(e)
        logger.error(f"Background run failed: {e}")
    finally:
        if trace_id:
            _running_tasks.pop(trace_id, None)


# ===== 路由 =====


@router.post("", response_model=RunResponse)
async def create_and_run(req: CreateRequest):
    """
    新建 Trace 并开始执行

    立即返回 trace_id，后台异步执行。
    通过 WebSocket /api/traces/{trace_id}/watch 监听实时更新。
    """
    from agent.core.runner import RunConfig

    _get_runner()  # 验证 Runner 已配置

    config = RunConfig(
        model=req.model,
        temperature=req.temperature,
        max_iterations=req.max_iterations,
        system_prompt=req.system_prompt,
        tools=req.tools,
        name=req.name,
        uid=req.uid,
    )

    # 启动后台执行，通过 Future 等待 trace_id（Phase 1 完成后即返回）
    trace_id_future: asyncio.Future[str] = asyncio.get_running_loop().create_future()
    task = asyncio.create_task(
        _run_with_trace_signal(req.messages, config, trace_id_future)
    )

    trace_id = await trace_id_future
    _running_tasks[trace_id] = task

    return RunResponse(
        trace_id=trace_id,
        status="started",
        message=f"Execution started. Watch via WebSocket: /api/traces/{trace_id}/watch",
    )


@router.post("/{trace_id}/run", response_model=RunResponse)
async def run_trace(trace_id: str, req: TraceRunRequest):
    """
    运行已有 Trace（统一续跑 + 回溯）

    - insert_after 为 null（或省略）：从末尾续跑
    - insert_after 为 int：回溯到该 sequence 后运行
    - messages 为空 + insert_after 为 int：重新生成（从该位置重跑，不插入新消息）

    insert_after 的值是 message 的 sequence 号。如果指定的 sequence 是一条带
    tool_calls 的 assistant 消息，系统会自动扩展截断点到其所有 tool response 之后。
    """
    from agent.core.runner import RunConfig

    runner = _get_runner()

    # 验证 trace 存在
    if runner.trace_store:
        trace = await runner.trace_store.get_trace(trace_id)
        if not trace:
            raise HTTPException(status_code=404, detail=f"Trace not found: {trace_id}")

    # 检查是否已在运行
    if trace_id in _running_tasks and not _running_tasks[trace_id].done():
        raise HTTPException(status_code=409, detail="Trace is already running")

    config = RunConfig(trace_id=trace_id, insert_after=req.insert_after)
    task = asyncio.create_task(_run_in_background(trace_id, req.messages, config))
    _running_tasks[trace_id] = task

    mode = "rewind" if req.insert_after is not None else "continue"
    return RunResponse(
        trace_id=trace_id,
        status="started",
        message=f"Run ({mode}) started. Watch via WebSocket: /api/traces/{trace_id}/watch",
    )


@router.post("/{trace_id}/stop", response_model=StopResponse)
async def stop_trace(trace_id: str):
    """
    停止运行中的 Trace

    设置取消信号，agent loop 在下一个 LLM 调用前检查并退出。
    Trace 状态置为 "stopped"。
    """
    runner = _get_runner()

    # 通过 runner 的 stop 方法设置取消信号
    stopped = await runner.stop(trace_id)

    if not stopped:
        # 检查是否在 _running_tasks 但 runner 不知道（可能已完成）
        if trace_id in _running_tasks:
            task = _running_tasks[trace_id]
            if not task.done():
                task.cancel()
                _running_tasks.pop(trace_id, None)
                return StopResponse(trace_id=trace_id, status="stopping")
        return StopResponse(trace_id=trace_id, status="not_running")

    return StopResponse(trace_id=trace_id, status="stopping")


@router.post("/{trace_id}/reflect", response_model=ReflectResponse)
async def reflect_trace(trace_id: str, req: ReflectRequest):
    """
    触发反思

    在 trace 末尾追加一条包含反思 prompt 的 user message，运行 agent 获取反思结果，
    将结果追加到 experiences 文件（默认 ./cache/experiences.md）。

    反思消息作为侧枝（side branch）：运行前保存 head_sequence，运行后恢复。
    这样反思消息不会出现在主对话路径上。
    """
    from agent.core.runner import RunConfig
    from agent.trace.compaction import build_reflect_prompt

    runner = _get_runner()

    if not runner.trace_store:
        raise HTTPException(status_code=503, detail="TraceStore not configured")

    # 验证 trace 存在
    trace = await runner.trace_store.get_trace(trace_id)
    if not trace:
        raise HTTPException(status_code=404, detail=f"Trace not found: {trace_id}")

    # 检查是否仍在运行
    if trace_id in _running_tasks and not _running_tasks[trace_id].done():
        raise HTTPException(status_code=409, detail="Cannot reflect on a running trace. Stop it first.")

    # 保存当前 head_sequence（反思完成后恢复，使反思消息成为侧枝）
    saved_head_sequence = trace.head_sequence

    # 构建反思 prompt
    prompt = build_reflect_prompt()
    if req.focus:
        prompt += f"\n\n请特别关注：{req.focus}"

    # 以续跑方式运行：追加 user message，agent 回复反思内容
    config = RunConfig(trace_id=trace_id)
    result = await runner.run_result(
        messages=[{"role": "user", "content": prompt}],
        config=config,
    )

    reflection_text = result.get("summary", "")

    # 恢复 head_sequence（反思消息成为侧枝，不影响主路径）
    await runner.trace_store.update_trace(trace_id, head_sequence=saved_head_sequence)

    # 追加到 experiences 文件
    if reflection_text:
        experiences_path = getattr(runner, "experiences_path", "./cache/experiences.md")
        if experiences_path:
            os.makedirs(os.path.dirname(experiences_path), exist_ok=True)
            header = f"\n\n---\n\n## {trace_id} ({datetime.now().strftime('%Y-%m-%d %H:%M')})\n\n"
            with open(experiences_path, "a", encoding="utf-8") as f:
                f.write(header + reflection_text + "\n")
            logger.info(f"Reflection appended to {experiences_path}")

    return ReflectResponse(
        trace_id=trace_id,
        reflection=reflection_text,
    )


@router.get("/running", tags=["run"])
async def list_running():
    """列出正在运行的 Trace"""
    running = []
    for tid, task in list(_running_tasks.items()):
        if task.done():
            _running_tasks.pop(tid, None)
        else:
            running.append(tid)
    return {"running": running}


# ===== 经验 API =====


@experiences_router.get("/experiences")
async def list_experiences():
    """读取经验文件内容"""
    runner = _get_runner()
    experiences_path = getattr(runner, "experiences_path", "./cache/experiences.md")

    if not experiences_path or not os.path.exists(experiences_path):
        return {"content": "", "path": experiences_path}

    with open(experiences_path, "r", encoding="utf-8") as f:
        content = f.read()

    return {"content": content, "path": experiences_path}