merge: async run queue + SSE architecture

This commit is contained in:
肖泽涛 2026-02-25 16:29:54 +08:00
commit 81663e7fba
22 changed files with 4718 additions and 2066 deletions

View File

@ -1,5 +1,10 @@
MODEL_DIR=./models
SAVE_DIR=./saves
REDIS_URL=redis://redis:6379/0
RUN_EVENTS_STREAM_TTL_SECONDS=7200
RUN_CANCEL_KEY_TTL_SECONDS=1800
LANGGRAPH_CHECKPOINTER_BACKEND=postgres
VITE_USE_RUNS_API=false
# region model_provider
SILICONFLOW_API_KEY= # 推荐使用硅基流动免费服务 https://cloud.siliconflow.cn/i/Eo5yTHGJ
@ -33,4 +38,4 @@ TAVILY_API_KEY= # 获取搜索服务的 api key 请访问 https://app.tavily.co
# LightRag llm 并发限制
# MAX_ASYNC=5
# EMBEDDING_FUNC_MAX_ASYNC=8
# EMBEDDING_FUNC_MAX_ASYNC=8

View File

@ -17,6 +17,9 @@ services:
environment:
- HOST_IP=${HOST_IP:-}
- POSTGRES_URL=${POSTGRES_URL:-postgresql+asyncpg://${POSTGRES_USER:-postgres}:${POSTGRES_PASSWORD:-postgres}@postgres:5432/${POSTGRES_DB:-yuxi_know}}
- REDIS_URL=${REDIS_URL:-redis://redis:6379/0}
- RUN_EVENTS_STREAM_TTL_SECONDS=${RUN_EVENTS_STREAM_TTL_SECONDS:-7200}
- RUN_CANCEL_KEY_TTL_SECONDS=${RUN_CANCEL_KEY_TTL_SECONDS:-1800}
- NEO4J_URI=${NEO4J_URI:-bolt://graph:7687}
- NEO4J_USERNAME=${NEO4J_USERNAME:-neo4j}
- NEO4J_PASSWORD=${NEO4J_PASSWORD:-0123456789}
@ -40,6 +43,57 @@ services:
retries: 8
start_period: 180s
depends_on:
postgres:
condition: service_healthy
redis:
condition: service_healthy
milvus:
condition: service_healthy
minio:
condition: service_healthy
worker:
build:
context: .
dockerfile: docker/api.Dockerfile
image: yuxi-api:0.5.prod
container_name: worker-prod
working_dir: /app
networks:
- app-network
volumes:
- ./saves:/app/saves
extra_hosts:
- "host.docker.internal:host-gateway"
env_file:
- .env.prod
environment:
- HOST_IP=${HOST_IP:-}
- POSTGRES_URL=${POSTGRES_URL:-postgresql+asyncpg://${POSTGRES_USER:-postgres}:${POSTGRES_PASSWORD:-postgres}@postgres:5432/${POSTGRES_DB:-yuxi_know}}
- REDIS_URL=${REDIS_URL:-redis://redis:6379/0}
- RUN_EVENTS_STREAM_TTL_SECONDS=${RUN_EVENTS_STREAM_TTL_SECONDS:-7200}
- RUN_CANCEL_KEY_TTL_SECONDS=${RUN_CANCEL_KEY_TTL_SECONDS:-1800}
- NEO4J_URI=${NEO4J_URI:-bolt://graph:7687}
- NEO4J_USERNAME=${NEO4J_USERNAME:-neo4j}
- NEO4J_PASSWORD=${NEO4J_PASSWORD:-0123456789}
- MILVUS_URI=${MILVUS_URI:-http://milvus:19530}
- MILVUS_TOKEN=${MILVUS_TOKEN:-}
- MILVUS_DB_NAME=${MILVUS_DB_NAME:-default}
- MINERU_VL_SERVER=${MINERU_VL_SERVER:-http://mineru-vllm-server:30000}
- MINERU_API_URI=${MINERU_API_URI:-http://mineru-api:30001}
- PADDLEX_URI=${PADDLEX_URI:-http://paddlex:8080}
- MINIO_URI=${MINIO_URI:-http://milvus-minio:9000}
- MODEL_DIR_IN_DOCKER=/models
- RUNNING_IN_DOCKER=true
- NO_PROXY=localhost,127.0.0.1,milvus,graph,milvus-minio,milvus-etcd,etcd,minio,mineru,paddlex,api.siliconflow.cn
- no_proxy=localhost,127.0.0.1,milvus,graph,milvus-minio,milvus-etcd,etcd,minio,mineru,paddlex,api.siliconflow.cn
command: uv run --no-dev arq server.worker_main.WorkerSettings
restart: unless-stopped
depends_on:
postgres:
condition: service_healthy
redis:
condition: service_healthy
milvus:
condition: service_healthy
minio:
@ -61,6 +115,7 @@ services:
environment:
- NODE_ENV=production
- VITE_API_URL=http://api:5050
- VITE_USE_RUNS_API=${VITE_USE_RUNS_API:-false}
command: nginx -g "daemon off;"
restart: unless-stopped
@ -173,6 +228,21 @@ services:
- app-network
restart: unless-stopped
redis:
image: redis:7-alpine
container_name: redis
command: redis-server --appendonly yes
healthcheck:
test: ["CMD", "redis-cli", "ping"]
interval: 10s
timeout: 5s
retries: 10
volumes:
- ./docker/volumes/redis:/data
networks:
- app-network
restart: unless-stopped
mineru-vllm-server:
build:
context: .

View File

@ -28,6 +28,9 @@ services:
environment:
- HOST_IP=${HOST_IP:-}
- POSTGRES_URL=${POSTGRES_URL:-postgresql+asyncpg://${POSTGRES_USER:-postgres}:${POSTGRES_PASSWORD:-postgres}@postgres:5432/${POSTGRES_DB:-yuxi_know}}
- REDIS_URL=${REDIS_URL:-redis://redis:6379/0}
- RUN_EVENTS_STREAM_TTL_SECONDS=${RUN_EVENTS_STREAM_TTL_SECONDS:-7200}
- RUN_CANCEL_KEY_TTL_SECONDS=${RUN_CANCEL_KEY_TTL_SECONDS:-1800}
- NEO4J_URI=${NEO4J_URI:-bolt://graph:7687}
- NEO4J_USERNAME=${NEO4J_USERNAME:-neo4j}
- NEO4J_PASSWORD=${NEO4J_PASSWORD:-0123456789}
@ -54,6 +57,63 @@ services:
depends_on:
postgres:
condition: service_healthy
redis:
condition: service_healthy
milvus:
condition: service_healthy
minio:
condition: service_healthy
worker:
build:
context: .
dockerfile: docker/api.Dockerfile
image: yuxi-api:0.5.dev
container_name: worker-dev
working_dir: /app
volumes:
- ./server:/app/server
- ./src:/app/src
- ./saves:/app/saves
- ./test:/app/test
- ./scripts:/app/scripts
- ./.env:/app/.env
- ./pyproject.toml:/app/pyproject.toml
- ./uv.lock:/app/uv.lock
- ${MODEL_DIR:-./models}:/models
networks:
- app-network
extra_hosts:
- "host.docker.internal:host-gateway"
env_file:
- .env
environment:
- HOST_IP=${HOST_IP:-}
- POSTGRES_URL=${POSTGRES_URL:-postgresql+asyncpg://${POSTGRES_USER:-postgres}:${POSTGRES_PASSWORD:-postgres}@postgres:5432/${POSTGRES_DB:-yuxi_know}}
- REDIS_URL=${REDIS_URL:-redis://redis:6379/0}
- RUN_EVENTS_STREAM_TTL_SECONDS=${RUN_EVENTS_STREAM_TTL_SECONDS:-7200}
- RUN_CANCEL_KEY_TTL_SECONDS=${RUN_CANCEL_KEY_TTL_SECONDS:-1800}
- NEO4J_URI=${NEO4J_URI:-bolt://graph:7687}
- NEO4J_USERNAME=${NEO4J_USERNAME:-neo4j}
- NEO4J_PASSWORD=${NEO4J_PASSWORD:-0123456789}
- MILVUS_URI=${MILVUS_URI:-http://milvus:19530}
- MILVUS_DB_NAME=${MILVUS_DB_NAME:-default}
- MILVUS_TOKEN=${MILVUS_TOKEN:-}
- MINERU_VL_SERVER=${MINERU_VL_SERVER:-http://mineru-vllm-server:30000}
- MINERU_API_URI=${MINERU_API_URI:-http://mineru-api:30001}
- PADDLEX_URI=${PADDLEX_URI:-http://paddlex:8080}
- MINIO_URI=${MINIO_URI:-http://milvus-minio:9000}
- MODEL_DIR_IN_DOCKER=/models
- RUNNING_IN_DOCKER=true
- NO_PROXY=localhost,127.0.0.1,milvus,graph,milvus-minio,milvus-etcd-dev,etcd,minio,mineru,paddlex,api.siliconflow.cn
- no_proxy=localhost,127.0.0.1,milvus,graph,milvus-minio,milvus-etcd-dev,etcd,minio,mineru,paddlex,api.siliconflow.cn
command: uv run --no-dev arq server.worker_main.WorkerSettings
restart: unless-stopped
depends_on:
postgres:
condition: service_healthy
redis:
condition: service_healthy
milvus:
condition: service_healthy
minio:
@ -80,6 +140,7 @@ services:
environment:
- NODE_ENV=development
- VITE_API_URL=http://api:5050
- VITE_USE_RUNS_API=${VITE_USE_RUNS_API:-false}
command: pnpm run server
restart: unless-stopped
@ -203,6 +264,21 @@ services:
- app-network
restart: unless-stopped
redis:
image: redis:7-alpine
container_name: redis
command: redis-server --appendonly yes
healthcheck:
test: ["CMD", "redis-cli", "ping"]
interval: 10s
timeout: 5s
retries: 10
volumes:
- ./docker/volumes/redis:/data
networks:
- app-network
restart: unless-stopped
# lastest version: wget https://gcore.jsdelivr.net/gh/opendatalab/MinerU@master/docker/compose.yaml
mineru-vllm-server:
build:

View File

@ -0,0 +1,185 @@
# Run 流式架构改造说明
## 1. 改造目标
本次改造将对话输出从 **HTTP 直连流** 升级为 **异步任务执行 + SSE 增量拉取**,目标是:
1. 页面离开/刷新不影响后台执行。
2. 前端支持断线重连与续流。
3. 提升系统稳定性、并发能力与可观测性。
4. 控制存储成本:过程数据短期保存,结果数据长期保存。
---
## 2. 改造前后对比(仅与 HTTP 直连流对比)
| 维度 | 改造前HTTP 直连流) | 改造后Run + SSE |
|---|---|---|
| 触发方式 | `POST /agent/{id}` 后长连接直接流式输出 | `POST /runs` 创建任务worker 异步执行 |
| 任务生命周期 | 绑定前端连接 | 与前端连接解耦 |
| 页面离开/刷新 | 常导致任务中断或前端丢上下文 | 任务继续执行,前端可续流 |
| 前端消费方式 | 同一个请求内读取 chunk | `GET /runs/{id}/events?after_seq=...` 增量拉取 |
| 恢复能力 | 弱,重连后难恢复 | 强,依赖 seq 游标恢复 |
| 取消语义 | 中断连接即可能影响任务 | 仅显式 `cancel` 才取消任务 |
---
## 3. 架构方案
### 3.1 组件职责
1. **FastAPI**:负责创建 run、查询 run、SSE 输出、cancel 接口。
2. **ARQ Worker**:负责真正执行模型流式任务。
3. **Redis**
- ARQ 队列
- run 过程事件流Redis Stream
- 取消信号key + pub/sub
4. **Postgres**
- run 执行状态(`agent_runs`
- 最终业务消息(`messages/tool_calls`
- checkpointer会话运行状态
### 3.2 架构图
```mermaid
flowchart LR
FE["Frontend"] -->|"POST /runs"| API["FastAPI"]
API -->|"create run"| PG[("Postgres")]
API -->|"enqueue"| R[("Redis")]
W["ARQ Worker"] -->|"dequeue"| R
W -->|"update run status"| PG
W -->|"write stream events"| R
FE -->|"GET /runs/:id/events?after_seq=..."| API
API -->|"read incremental events"| R
API -->|"SSE events"| FE
FE -->|"POST /runs/:id/cancel"| API
API -->|"cancel mark"| PG
API -->|"publish cancel"| R
W -->|"persist messages/tool_calls"| PG
```
---
## 4. 端到端流程(事件流转)
```mermaid
sequenceDiagram
participant FE as Frontend
participant API as FastAPI
participant R as Redis
participant W as ARQ Worker
participant PG as Postgres
FE->>API: POST /api/chat/agent/{agent_id}/runs
API->>PG: create agent_runs(status=pending)
API->>R: enqueue process_agent_run(run_id)
API-->>FE: run_id
FE->>API: GET /api/chat/runs/{run_id}/events?after_seq=0
W->>R: dequeue run job
W->>PG: mark running
W->>R: append loading/tool/state events (stream)
API->>R: read events after_seq
API-->>FE: SSE incremental events
FE->>API: POST /api/chat/runs/{run_id}/cancel (optional)
API->>PG: mark cancel_requested
API->>R: publish cancel signal
W->>W: cancel current task
W->>PG: mark terminal status
W->>PG: persist messages/tool_calls
API-->>FE: close event
```
---
## 5. 接口与协议变更
### 5.1 对外路径(保持稳定)
1. `POST /api/chat/agent/{agent_id}/runs`
2. `GET /api/chat/runs/{run_id}`
3. `GET /api/chat/runs/{run_id}/events?after_seq=...`
4. `POST /api/chat/runs/{run_id}/cancel`
### 5.2 `after_seq` 语义
1. 主格式为字符串游标Redis Stream ID`1700000000000-3`)。
2. 兼容旧整数参数输入。
3. SSE 返回 `seq` 字段统一为字符串,前端按单调递增去重。
### 5.3 SSE 事件格式
```json
{
"run_id": "...",
"seq": "1700000000000-3",
"event_type": "loading",
"payload": {"items": [...]},
"ts": 1700000000000
}
```
控制事件:`heartbeat` / `error` / `close`
---
## 6. 前端行为变化
1. 本地记录活跃 run 快照:`active_run:{threadId}`。
2. 刷新/切回页面时按 `run_id + last_seq` 自动续流。
3. 接收事件先做 seq 去重,再更新 UI。
4. 保留打字机效果(`requestAnimationFrame + throttle`)。
5. 保留首条消息自动更新会话标题逻辑。
---
## 7. 稳定性设计
1. 幂等:`request_id` 避免重复创建 run。
2. 重试:仅可恢复错误触发 ARQ 重试(`max_tries=2`)。
3. 取消DB 状态 + Redis 信号双通道。
4. SSE 生命周期:心跳、超时、终态关闭、断线重连。
5. 状态单一真相:执行态在 `agent_runs`,业务态在 `messages/tool_calls + checkpointer`
---
## 8. 本次代码变更范围(未提交部分)
### 后端
1. `/Yuxi-Know/src/services/run_queue_service.py`
2. `/Yuxi-Know/src/services/run_worker.py`
3. `/Yuxi-Know/src/services/agent_run_service.py`
4. `/Yuxi-Know/src/repositories/agent_run_repository.py`
5. `/Yuxi-Know/server/routers/chat_router.py`
6. `/Yuxi-Know/server/worker_main.py`
7. `/Yuxi-Know/src/storage/postgres/manager.py`
8. `/Yuxi-Know/src/storage/postgres/models_business.py`
### 前端
1. `/Yuxi-Know/web/src/apis/agent_api.js`
2. `/Yuxi-Know/web/src/components/AgentChatComponent.vue`
### 测试
1. `/Yuxi-Know/test/test_run_queue_service.py`
2. `/Yuxi-Know/test/test_agent_run_service.py`
3. `/Yuxi-Know/test/test_run_worker.py`
### 配置
1. `/Yuxi-Know/docker-compose.yml`
2. `/Yuxi-Know/docker-compose.prod.yml`
3. `/Yuxi-Know/.env.template`
---
## 9. 验收标准
1. 发送消息后SSE 能持续收到增量事件。
2. 页面刷新后,可按 `after_seq` 恢复输出。
3. 页面离开不影响后台执行。
4. 取消后 run 状态正确收敛,输出停止。
5. 最终消息与工具调用正常入库。

View File

@ -22,6 +22,7 @@ dependencies = [
"langchain-text-splitters>=1.0",
"langgraph>=1.0.1",
"langgraph-checkpoint-sqlite>=3.0",
"langgraph-checkpoint-postgres>=2.0.0",
"langgraph-cli[inmem]>=0.4",
"langsmith>=0.4",
"lightrag-hku>=1.4.6",
@ -59,6 +60,7 @@ dependencies = [
"tomli-w",
"aiofiles>=24.1.0",
"aiohttp>=3.9.0",
"arq>=0.26.3",
"chardet>=5.0.0",
"deepagents>=0.2.5",
"json-repair>=0.54.0",
@ -67,6 +69,7 @@ dependencies = [
"docling>=2.68.0",
"loguru>=0.7.3",
"langfuse>=2.0.0",
"redis>=5.2.0",
]
[tool.ruff]
line-length = 120 # 代码最大行宽

View File

@ -3,7 +3,7 @@ import uuid
from fastapi import APIRouter, Body, Depends, HTTPException, Query, UploadFile, File
from fastapi.responses import StreamingResponse
from pydantic import BaseModel
from pydantic import BaseModel, Field
from sqlalchemy.ext.asyncio import AsyncSession
from src.storage.postgres.models_business import User
@ -13,6 +13,13 @@ from src import config as conf
from src.agents import agent_manager
from src.models import select_model
from src.services.chat_stream_service import get_agent_state_view, stream_agent_chat, stream_agent_resume
from src.services.agent_run_service import (
cancel_agent_run_view,
create_agent_run_view,
get_active_run_by_thread,
get_agent_run_view,
stream_agent_run_events,
)
from src.services.conversation_service import (
create_thread_view,
delete_thread_attachment_view,
@ -61,6 +68,12 @@ class AgentConfigUpdate(BaseModel):
config_json: dict | None = None
class AgentRunCreate(BaseModel):
query: str
config: dict = Field(default_factory=dict)
image_content: str | None = None
chat = APIRouter(prefix="/chat", tags=["chat"])
# =============================================================================
@ -394,6 +407,76 @@ async def chat_agent(
)
@chat.post("/agent/{agent_id}/runs")
async def create_agent_run(
agent_id: str,
payload: AgentRunCreate,
current_user: User = Depends(get_required_user),
db: AsyncSession = Depends(get_db),
):
"""创建异步 run 任务并入队(需要登录)"""
return await create_agent_run_view(
agent_id=agent_id,
query=payload.query,
config=payload.config or {},
image_content=payload.image_content,
current_user_id=str(current_user.id),
db=db,
)
@chat.get("/runs/{run_id}")
async def get_agent_run(
run_id: str,
current_user: User = Depends(get_required_user),
db: AsyncSession = Depends(get_db),
):
"""获取 run 状态(需要登录)"""
return await get_agent_run_view(run_id=run_id, current_user_id=str(current_user.id), db=db)
@chat.post("/runs/{run_id}/cancel")
async def cancel_agent_run(
run_id: str,
current_user: User = Depends(get_required_user),
db: AsyncSession = Depends(get_db),
):
"""取消 run需要登录"""
return await cancel_agent_run_view(run_id=run_id, current_user_id=str(current_user.id), db=db)
@chat.get("/runs/{run_id}/events")
async def stream_run_events(
run_id: str,
after_seq: str = Query("0"),
current_user: User = Depends(get_required_user),
):
"""SSE 拉取 run 事件(需要登录)"""
return StreamingResponse(
stream_agent_run_events(
run_id=run_id,
after_seq=after_seq,
current_user_id=str(current_user.id),
),
media_type="text/event-stream",
headers={
"Cache-Control": "no-cache",
"Connection": "keep-alive",
"X-Accel-Buffering": "no",
},
)
@chat.get("/thread/{thread_id}/active_run")
async def get_thread_active_run(
thread_id: str,
current_user: User = Depends(get_required_user),
db: AsyncSession = Depends(get_db),
):
"""获取当前会话活跃 run需要登录"""
return await get_active_run_by_thread(thread_id=thread_id, current_user_id=str(current_user.id), db=db)
# =============================================================================
# > === 模型管理分组 ===
# =============================================================================

View File

@ -5,6 +5,7 @@ from fastapi import FastAPI
from src.services.task_service import tasker
from src.services.mcp_service import init_mcp_servers
from src.services.skill_service import init_skills_cache
from src.services.run_queue_service import close_queue_clients, get_redis_client
from src.storage.postgres.manager import pg_manager
from src.knowledge import knowledge_base
from src.utils import logger
@ -40,7 +41,15 @@ async def lifespan(app: FastAPI):
except Exception as e:
logger.error(f"Failed to initialize knowledge base manager: {e}")
# 预热 Redisrun 队列)
try:
redis = await get_redis_client()
await redis.ping()
except Exception as e:
logger.warning(f"Run queue redis unavailable on startup: {e}")
await tasker.start()
yield
await tasker.shutdown()
await close_queue_clients()
await pg_manager.close()

5
server/worker_main.py Normal file
View File

@ -0,0 +1,5 @@
"""ARQ worker entrypoint."""
from src.services.run_worker import WorkerSettings
__all__ = ["WorkerSettings"]

View File

@ -181,19 +181,53 @@ class BaseAgent:
if self.checkpointer is not None:
return self.checkpointer
# 创建数据库连接并确保设置 checkpointer
checkpointer = None
backend = os.getenv("LANGGRAPH_CHECKPOINTER_BACKEND", "sqlite").strip().lower()
try:
checkpointer = AsyncSqliteSaver(await self.get_async_conn())
if backend == "postgres":
checkpointer = await self._create_postgres_checkpointer()
except Exception as e:
logger.error(f"构建 Graph 设置 checkpointer 时出错: {e}, 尝试使用内存存储")
checkpointer = InMemorySaver()
if checkpointer is None:
try:
checkpointer = AsyncSqliteSaver(await self.get_async_conn())
except Exception as e:
logger.error(f"构建 sqlite checkpointer 失败: {e}, 尝试使用内存存储")
checkpointer = InMemorySaver()
self.checkpointer = checkpointer
return self.checkpointer
async def _create_postgres_checkpointer(self):
postgres_url = os.getenv("POSTGRES_URL")
if not postgres_url:
logger.warning("POSTGRES_URL 未配置,无法启用 postgres checkpointer回退 sqlite")
return None
try:
from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver # type: ignore
except Exception as e:
logger.warning(f"langgraph postgres checkpointer 不可用,回退 sqlite: {e}")
return None
conn_str = postgres_url.replace("+asyncpg", "")
try:
saver_factory = getattr(AsyncPostgresSaver, "from_conn_string", None)
if callable(saver_factory):
saver = saver_factory(conn_str)
else:
saver = AsyncPostgresSaver(conn_str) # type: ignore[call-arg]
setup_fn = getattr(saver, "setup", None)
if callable(setup_fn):
result = setup_fn()
if hasattr(result, "__await__"):
await result
logger.info(f"{self.name} 使用 postgres checkpointer")
return saver
except Exception as e:
logger.warning(f"初始化 postgres checkpointer 失败,回退 sqlite: {e}")
return None
async def get_async_conn(self) -> aiosqlite.Connection:
"""获取异步数据库连接"""
if self._async_conn is not None:

View File

@ -0,0 +1,102 @@
"""Agent run repository."""
from __future__ import annotations
from sqlalchemy import and_, select
from sqlalchemy.ext.asyncio import AsyncSession
from src.storage.postgres.models_business import AgentRun
from src.utils.datetime_utils import utc_now_naive
TERMINAL_RUN_STATUSES = {"completed", "failed", "cancelled", "interrupted"}
class AgentRunRepository:
def __init__(self, db_session: AsyncSession):
self.db = db_session
async def get_run(self, run_id: str) -> AgentRun | None:
result = await self.db.execute(select(AgentRun).where(AgentRun.id == run_id))
return result.scalar_one_or_none()
async def get_run_by_request_id(self, request_id: str) -> AgentRun | None:
result = await self.db.execute(select(AgentRun).where(AgentRun.request_id == request_id))
return result.scalar_one_or_none()
async def get_run_for_user(self, run_id: str, user_id: str) -> AgentRun | None:
result = await self.db.execute(
select(AgentRun).where(and_(AgentRun.id == run_id, AgentRun.user_id == str(user_id)))
)
return result.scalar_one_or_none()
async def create_run(
self,
*,
run_id: str,
thread_id: str,
agent_id: str,
user_id: str,
request_id: str,
input_payload: dict,
) -> AgentRun:
run = AgentRun(
id=run_id,
thread_id=thread_id,
agent_id=agent_id,
user_id=str(user_id),
request_id=request_id,
input_payload=input_payload or {},
status="pending",
)
self.db.add(run)
await self.db.flush()
return run
async def mark_running(self, run_id: str) -> AgentRun | None:
run = await self._lock_run(run_id)
if not run:
return None
if run.status in TERMINAL_RUN_STATUSES:
return run
now = utc_now_naive()
run.status = "running"
run.started_at = run.started_at or now
run.updated_at = now
await self.db.flush()
return run
async def request_cancel(self, run_id: str) -> AgentRun | None:
run = await self._lock_run(run_id)
if not run:
return None
if run.status in TERMINAL_RUN_STATUSES:
return run
run.status = "cancel_requested"
run.updated_at = utc_now_naive()
await self.db.flush()
return run
async def set_terminal_status(
self,
run_id: str,
*,
status: str,
error_type: str | None = None,
error_message: str | None = None,
) -> AgentRun | None:
run = await self._lock_run(run_id)
if not run:
return None
if run.status in TERMINAL_RUN_STATUSES:
return run
run.status = status
run.error_type = error_type
run.error_message = error_message
run.finished_at = utc_now_naive()
run.updated_at = run.finished_at
await self.db.flush()
return run
async def _lock_run(self, run_id: str) -> AgentRun | None:
result = await self.db.execute(select(AgentRun).where(AgentRun.id == run_id).with_for_update())
return result.scalar_one_or_none()

View File

@ -0,0 +1,246 @@
"""Agent run service (run creation, polling stream, cancel)."""
from __future__ import annotations
import asyncio
import json
import os
import uuid
from collections.abc import AsyncIterator
from fastapi import HTTPException
from sqlalchemy.exc import IntegrityError
from sqlalchemy.ext.asyncio import AsyncSession
from src.agents import agent_manager
from src.repositories.agent_run_repository import AgentRunRepository, TERMINAL_RUN_STATUSES
from src.repositories.conversation_repository import ConversationRepository
from src.services.run_queue_service import (
get_arq_pool,
get_last_run_stream_seq,
list_run_stream_events,
normalize_after_seq,
publish_cancel_signal,
)
from src.storage.postgres.manager import pg_manager
from src.utils.datetime_utils import utc_now_naive
from src.utils.logging_config import logger
SSE_HEARTBEAT_SECONDS = int(os.getenv("RUN_SSE_HEARTBEAT_SECONDS", "15"))
SSE_MAX_CONNECTION_MINUTES = int(os.getenv("RUN_SSE_MAX_CONNECTION_MINUTES", "30"))
SSE_POLL_INTERVAL_SECONDS = float(os.getenv("RUN_SSE_POLL_INTERVAL_SECONDS", "1.0"))
def _build_run_response(run) -> dict:
return {
"run_id": run.id,
"thread_id": run.thread_id,
"status": run.status,
"request_id": run.request_id,
"stream_url": f"/api/chat/runs/{run.id}/events?after_seq=0",
}
def _format_sse(data: dict, event: str | None = None) -> str:
lines = []
if event:
lines.append(f"event: {event}")
lines.append(f"data: {json.dumps(data, ensure_ascii=False)}")
lines.append("")
return "\n".join(lines) + "\n"
async def create_agent_run_view(
*,
agent_id: str,
query: str,
config: dict,
image_content: str | None,
current_user_id: str,
db: AsyncSession,
) -> dict:
if not query:
raise HTTPException(status_code=422, detail="query 不能为空")
if not agent_manager.get_agent(agent_id):
raise HTTPException(status_code=404, detail=f"智能体 {agent_id} 不存在")
thread_id = (config or {}).get("thread_id")
if not thread_id:
raise HTTPException(status_code=422, detail="config.thread_id 不能为空")
conv_repo = ConversationRepository(db)
conversation = await conv_repo.get_conversation_by_thread_id(thread_id)
if not conversation or conversation.user_id != str(current_user_id) or conversation.status == "deleted":
raise HTTPException(status_code=404, detail="对话线程不存在")
request_id = str((config or {}).get("request_id") or uuid.uuid4())
run_repo = AgentRunRepository(db)
existing = await run_repo.get_run_by_request_id(request_id)
if existing and existing.user_id == str(current_user_id):
return _build_run_response(existing)
if existing and existing.user_id != str(current_user_id):
raise HTTPException(status_code=409, detail="request_id 冲突")
run_id = str(uuid.uuid4())
input_payload = {
"query": query,
"config": config or {},
"image_content": image_content,
"agent_id": agent_id,
"thread_id": thread_id,
"user_id": str(current_user_id),
"request_id": request_id,
"created_at": utc_now_naive().isoformat(),
}
try:
run = await run_repo.create_run(
run_id=run_id,
thread_id=thread_id,
agent_id=agent_id,
user_id=str(current_user_id),
request_id=request_id,
input_payload=input_payload,
)
await db.commit()
except IntegrityError:
await db.rollback()
existing = await run_repo.get_run_by_request_id(request_id)
if existing and existing.user_id == str(current_user_id):
return _build_run_response(existing)
raise HTTPException(status_code=409, detail="request_id 冲突")
queue = await get_arq_pool()
await queue.enqueue_job("process_agent_run", run.id, _job_id=f"run:{run.id}")
return _build_run_response(run)
async def get_agent_run_view(*, run_id: str, current_user_id: str, db: AsyncSession) -> dict:
repo = AgentRunRepository(db)
run = await repo.get_run_for_user(run_id, str(current_user_id))
if not run:
raise HTTPException(status_code=404, detail="运行任务不存在")
return {"run": run.to_dict()}
async def cancel_agent_run_view(*, run_id: str, current_user_id: str, db: AsyncSession) -> dict:
repo = AgentRunRepository(db)
run = await repo.get_run_for_user(run_id, str(current_user_id))
if not run:
raise HTTPException(status_code=404, detail="运行任务不存在")
run = await repo.request_cancel(run_id)
await publish_cancel_signal(run_id)
return {"run": run.to_dict() if run else None}
async def stream_agent_run_events(
*,
run_id: str,
after_seq: str | int,
current_user_id: str,
) -> AsyncIterator[str]:
started_at = utc_now_naive()
last_heartbeat_ts = started_at
last_seq = normalize_after_seq(after_seq)
try:
while True:
try:
async with pg_manager.get_async_session_context() as db:
repo = AgentRunRepository(db)
run = await repo.get_run_for_user(run_id, str(current_user_id))
if not run:
yield _format_sse({"run_id": run_id, "message": "运行任务不存在"}, event="error")
yield _format_sse({"run_id": run_id, "last_seq": last_seq}, event="close")
return
except asyncio.CancelledError:
raise
except Exception as e:
logger.warning(f"Run SSE DB error for run {run_id}: {e}")
yield _format_sse(
{
"run_id": run_id,
"message": "运行事件流暂时不可用,请重连",
"reason": "db_error",
},
event="error",
)
yield _format_sse({"run_id": run_id, "last_seq": last_seq}, event="close")
return
try:
events = await list_run_stream_events(run_id, after_seq=last_seq, limit=200)
except Exception as e:
logger.warning(f"Run SSE redis error for run {run_id}: {e}")
yield _format_sse(
{
"run_id": run_id,
"message": "运行事件流暂时不可用,请重连",
"reason": "redis_error",
},
event="error",
)
yield _format_sse({"run_id": run_id, "last_seq": last_seq}, event="close")
return
for event in events:
seq = str(event.get("seq") or "0-0")
last_seq = seq
yield _format_sse(
{
"run_id": run_id,
"seq": seq,
"event_type": event.get("event_type") or "message",
"payload": event.get("payload") or {},
"ts": event.get("ts"),
},
event=event.get("event_type") or "message",
)
if run.status in TERMINAL_RUN_STATUSES and not events:
terminal_seq = last_seq
if terminal_seq in {"", "0", "0-0"}:
terminal_seq = await get_last_run_stream_seq(run_id)
yield _format_sse(
{"run_id": run_id, "status": run.status, "last_seq": terminal_seq},
event="close",
)
return
now = utc_now_naive()
elapsed_seconds = (now - started_at).total_seconds()
heartbeat_elapsed = (now - last_heartbeat_ts).total_seconds()
if heartbeat_elapsed >= SSE_HEARTBEAT_SECONDS:
yield _format_sse({"run_id": run_id, "last_seq": last_seq}, event="heartbeat")
last_heartbeat_ts = now
if elapsed_seconds >= SSE_MAX_CONNECTION_MINUTES * 60:
yield _format_sse({"run_id": run_id, "last_seq": last_seq}, event="close")
return
await asyncio.sleep(SSE_POLL_INTERVAL_SECONDS)
except asyncio.CancelledError:
return
async def get_active_run_by_thread(*, thread_id: str, current_user_id: str, db: AsyncSession) -> dict:
from sqlalchemy import select
from src.storage.postgres.models_business import AgentRun
result = await db.execute(
select(AgentRun)
.where(
AgentRun.thread_id == thread_id,
AgentRun.user_id == str(current_user_id),
AgentRun.status.notin_(list(TERMINAL_RUN_STATUSES)),
)
.order_by(AgentRun.created_at.desc())
.limit(1)
)
run = result.scalar_one_or_none()
return {"run": run.to_dict() if run else None}

View File

@ -0,0 +1,241 @@
"""Run queue/redis helpers."""
from __future__ import annotations
import asyncio
import json
import os
from contextlib import asynccontextmanager
from datetime import UTC, datetime
from src.utils.logging_config import logger
REDIS_URL = os.getenv("REDIS_URL", "redis://redis:6379/0")
RUN_CANCEL_KEY_TTL_SECONDS = int(os.getenv("RUN_CANCEL_KEY_TTL_SECONDS", "1800"))
RUN_EVENTS_STREAM_TTL_SECONDS = int(os.getenv("RUN_EVENTS_STREAM_TTL_SECONDS", "7200"))
RUN_EVENTS_STREAM_MAXLEN = int(os.getenv("RUN_EVENTS_STREAM_MAXLEN", "0"))
RUN_CANCEL_CHANNEL = os.getenv("RUN_CANCEL_CHANNEL", "run:cancel:ch")
_redis_client = None
_arq_pool = None
def _redacted_redis_url(url: str) -> str:
if "@" in url:
return url.split("@", 1)[1]
return url
def _cancel_key(run_id: str) -> str:
return f"run:cancel:{run_id}"
def _event_stream_key(run_id: str) -> str:
return f"run:events:{run_id}"
def _is_valid_stream_seq(value: str) -> bool:
major, sep, minor = value.partition("-")
if sep != "-":
return False
return major.isdigit() and minor.isdigit()
def normalize_after_seq(after_seq: str | int | None) -> str:
"""Normalize after_seq cursor to redis stream id format."""
if after_seq is None:
return "0-0"
if isinstance(after_seq, int):
return "0-0"
text = str(after_seq).strip()
if not text:
return "0-0"
if _is_valid_stream_seq(text):
return text
return "0-0"
async def get_redis_client():
global _redis_client
if _redis_client is not None:
return _redis_client
try:
from redis.asyncio import Redis
except Exception as e:
raise RuntimeError("redis dependency is required for run queue") from e
redis = Redis.from_url(REDIS_URL, decode_responses=True)
try:
await redis.ping()
except Exception as e:
try:
await redis.aclose()
except Exception:
pass
raise RuntimeError(
f"Redis connection failed ({_redacted_redis_url(REDIS_URL)}): {e}"
) from e
_redis_client = redis
return _redis_client
async def get_arq_pool():
global _arq_pool
if _arq_pool is not None:
return _arq_pool
try:
from arq.connections import RedisSettings, create_pool
except Exception as e:
raise RuntimeError("arq dependency is required for run queue") from e
settings = RedisSettings.from_dsn(REDIS_URL)
_arq_pool = await create_pool(settings)
return _arq_pool
@asynccontextmanager
async def redis_pubsub(channel: str):
redis = await get_redis_client()
pubsub = redis.pubsub()
await pubsub.subscribe(channel)
try:
yield pubsub
finally:
try:
await pubsub.unsubscribe(channel)
finally:
await pubsub.close()
async def publish_cancel_signal(run_id: str) -> None:
redis = await get_redis_client()
key = _cancel_key(run_id)
try:
await redis.set(key, "1", ex=RUN_CANCEL_KEY_TTL_SECONDS)
await redis.publish(RUN_CANCEL_CHANNEL, run_id)
except Exception as e:
logger.warning(f"Failed to publish cancel signal for run {run_id}: {e}")
async def has_cancel_signal(run_id: str) -> bool:
redis = await get_redis_client()
key = _cancel_key(run_id)
try:
return bool(await redis.get(key))
except Exception as e:
logger.warning(f"Failed to read cancel signal for run {run_id}: {e}")
return False
async def wait_for_cancel_signal(run_id: str, poll_timeout_seconds: float = 1.0) -> bool:
if await has_cancel_signal(run_id):
return True
try:
async with redis_pubsub(RUN_CANCEL_CHANNEL) as pubsub:
while True:
msg = await pubsub.get_message(
ignore_subscribe_messages=True,
timeout=poll_timeout_seconds,
)
if msg and str(msg.get("data")) == run_id:
return True
if await has_cancel_signal(run_id):
return True
except asyncio.CancelledError:
raise
except Exception as e:
logger.warning(f"Failed to wait cancel signal for run {run_id}: {e}")
return False
async def clear_cancel_signal(run_id: str) -> None:
redis = await get_redis_client()
key = _cancel_key(run_id)
try:
await redis.delete(key)
except Exception as e:
logger.warning(f"Failed to clear cancel signal for run {run_id}: {e}")
async def append_run_stream_event(run_id: str, event_type: str, payload: dict) -> str:
redis = await get_redis_client()
key = _event_stream_key(run_id)
now_ms = int(datetime.now(tz=UTC).timestamp() * 1000)
fields = {
"event_type": event_type,
"payload": json.dumps(payload or {}, ensure_ascii=False),
"ts": str(now_ms),
}
kwargs = {}
if RUN_EVENTS_STREAM_MAXLEN > 0:
kwargs["maxlen"] = RUN_EVENTS_STREAM_MAXLEN
kwargs["approximate"] = True
event_id = await redis.xadd(key, fields, **kwargs)
await redis.expire(key, RUN_EVENTS_STREAM_TTL_SECONDS)
return str(event_id)
async def list_run_stream_events(
run_id: str,
*,
after_seq: str = "0-0",
limit: int = 200,
) -> list[dict]:
redis = await get_redis_client()
key = _event_stream_key(run_id)
start = "-" if after_seq in {"0", "0-0", ""} else f"({after_seq}"
rows = await redis.xrange(key, min=start, max="+", count=limit)
events = []
for event_id, fields in rows:
payload_raw = fields.get("payload") or "{}"
try:
payload = json.loads(payload_raw)
except Exception:
payload = {}
ts_value = fields.get("ts")
events.append(
{
"seq": str(event_id),
"event_type": fields.get("event_type") or "message",
"payload": payload,
"ts": int(ts_value) if ts_value else None,
}
)
return events
async def get_last_run_stream_seq(run_id: str) -> str:
redis = await get_redis_client()
key = _event_stream_key(run_id)
rows = await redis.xrevrange(key, max="+", min="-", count=1)
if not rows:
return "0-0"
event_id, _ = rows[0]
return str(event_id)
async def close_queue_clients() -> None:
global _redis_client, _arq_pool
if _arq_pool is not None:
try:
await _arq_pool.close()
except Exception:
pass
_arq_pool = None
if _redis_client is not None:
try:
await _redis_client.aclose()
except Exception:
pass
_redis_client = None

373
src/services/run_worker.py Normal file
View File

@ -0,0 +1,373 @@
"""ARQ worker for agent runs."""
from __future__ import annotations
import asyncio
import json
import os
import time
from dataclasses import dataclass, field
from sqlalchemy import select
from sqlalchemy.exc import OperationalError
from src.repositories.agent_run_repository import AgentRunRepository, TERMINAL_RUN_STATUSES
from src.services.chat_stream_service import stream_agent_chat
from src.services.run_queue_service import (
append_run_stream_event,
clear_cancel_signal,
has_cancel_signal,
wait_for_cancel_signal,
)
from src.storage.postgres.manager import pg_manager
from src.storage.postgres.models_business import User
from src.utils.logging_config import logger
LOADING_FLUSH_INTERVAL_MS = 100
LOADING_FLUSH_MAX_CHARS = 512
RUN_CANCEL_POLL_SECONDS = 0.2
REDIS_URL = os.getenv("REDIS_URL", "redis://redis:6379/0")
class RetryableRunError(Exception):
"""Error type that should trigger ARQ retry."""
class NonRetryableRunError(Exception):
"""Error type that should not trigger ARQ retry."""
@dataclass
class RunContext:
run_id: str
cancel_event: asyncio.Event = field(default_factory=asyncio.Event)
_watch_task: asyncio.Task | None = None
async def start(self) -> None:
if self._watch_task is None:
self._watch_task = asyncio.create_task(self._watch_cancel_signal())
async def close(self) -> None:
if self._watch_task:
self._watch_task.cancel()
await asyncio.gather(self._watch_task, return_exceptions=True)
self._watch_task = None
async def wait_cancelled(self) -> None:
await self.cancel_event.wait()
async def is_cancelled(self) -> bool:
if self.cancel_event.is_set():
return True
if await has_cancel_signal(self.run_id):
self.cancel_event.set()
return True
return False
async def _watch_cancel_signal(self) -> None:
while not self.cancel_event.is_set():
cancelled = await wait_for_cancel_signal(
self.run_id,
poll_timeout_seconds=RUN_CANCEL_POLL_SECONDS,
)
if cancelled:
self.cancel_event.set()
return
class ChunkedEventWriter:
def __init__(self, run_id: str, interval_ms: int = 100, max_chars: int = 512):
self.run_id = run_id
self.interval_seconds = interval_ms / 1000
self.max_chars = max_chars
self.buffer: list[dict] = []
self.buffer_chars = 0
self.last_flush = time.monotonic()
async def append(self, chunk: dict):
self.buffer.append(chunk)
content = chunk.get("response") or ""
self.buffer_chars += len(content) if isinstance(content, str) else 0
now = time.monotonic()
if (now - self.last_flush) >= self.interval_seconds or self.buffer_chars >= self.max_chars:
await self.flush()
async def flush(self):
if not self.buffer:
return
await append_run_event(self.run_id, "loading", {"items": self.buffer})
self.buffer = []
self.buffer_chars = 0
self.last_flush = time.monotonic()
async def _get_run(run_id: str):
async with pg_manager.get_async_session_context() as db:
repo = AgentRunRepository(db)
return await repo.get_run(run_id)
async def append_run_event(run_id: str, event_type: str, payload: dict):
await append_run_stream_event(run_id, event_type, payload)
async def mark_run_running(run_id: str):
async with pg_manager.get_async_session_context() as db:
repo = AgentRunRepository(db)
await repo.mark_running(run_id)
async def mark_run_terminal(run_id: str, status: str, error_type: str | None = None, error_message: str | None = None):
async with pg_manager.get_async_session_context() as db:
repo = AgentRunRepository(db)
await repo.set_terminal_status(run_id, status=status, error_type=error_type, error_message=error_message)
async def _load_user(user_id: str):
async with pg_manager.get_async_session_context() as db:
result = await db.execute(select(User).where(User.id == int(user_id)))
return result.scalar_one_or_none()
async def _is_cancel_requested(run_id: str) -> bool:
run = await _get_run(run_id)
return bool(run and run.status == "cancel_requested")
def _job_try(ctx) -> int:
if isinstance(ctx, dict):
try:
return int(ctx.get("job_try") or 1)
except Exception:
return 1
return 1
def _is_last_try(ctx) -> bool:
return _job_try(ctx) >= max(1, int(getattr(WorkerSettings, "max_tries", 1)))
def _is_retryable_exception(exc: Exception) -> bool:
if isinstance(exc, NonRetryableRunError):
return False
return isinstance(exc, (RetryableRunError, OperationalError, ConnectionError, TimeoutError, asyncio.TimeoutError))
def _iter_json_chunks(chunk_bytes: bytes) -> list[dict]:
text = chunk_bytes.decode("utf-8")
chunks: list[dict] = []
for line in text.splitlines():
line = line.strip()
if not line:
continue
try:
chunks.append(json.loads(line))
except Exception:
logger.warning(f"Failed to parse run stream chunk: {line[:200]}")
return chunks
async def _consume_stream_with_cancel(agen, run_ctx: RunContext):
while True:
next_task = asyncio.create_task(agen.__anext__())
cancel_task = asyncio.create_task(run_ctx.wait_cancelled())
done, _ = await asyncio.wait({next_task, cancel_task}, return_when=asyncio.FIRST_COMPLETED)
if cancel_task in done:
next_task.cancel()
await asyncio.gather(next_task, return_exceptions=True)
raise asyncio.CancelledError(f"run {run_ctx.run_id} cancelled")
cancel_task.cancel()
await asyncio.gather(cancel_task, return_exceptions=True)
try:
yield next_task.result()
except StopAsyncIteration:
return
async def process_agent_run(ctx, run_id: str):
run = await _get_run(run_id)
if not run:
logger.warning(f"Run not found: {run_id}")
return
if run.status in TERMINAL_RUN_STATUSES:
logger.info(f"Run already terminal, skip: {run_id}, status={run.status}")
return
payload = run.input_payload or {}
query = payload.get("query")
config = payload.get("config") or {}
agent_id = payload.get("agent_id")
image_content = payload.get("image_content")
user_id = payload.get("user_id")
request_id = payload.get("request_id")
user = await _load_user(user_id)
if not user:
await mark_run_terminal(run_id, "failed", "user_not_found", f"user {user_id} not found")
return
if not request_id:
request_id = run.request_id
meta = {
"request_id": request_id,
"query": query,
"agent_id": agent_id,
"server_model_name": config.get("model", agent_id),
"thread_id": config.get("thread_id"),
"user_id": user.id,
"has_image": bool(image_content),
}
await mark_run_running(run_id)
run_ctx = RunContext(run_id=run_id)
writer = ChunkedEventWriter(
run_id=run_id,
interval_ms=LOADING_FLUSH_INTERVAL_MS,
max_chars=LOADING_FLUSH_MAX_CHARS,
)
await run_ctx.start()
terminal_set = False
try:
async with pg_manager.get_async_session_context() as db:
stream = stream_agent_chat(
agent_id=agent_id,
query=query,
config=config,
meta=meta,
image_content=image_content,
current_user=user,
db=db,
)
async for chunk_bytes in _consume_stream_with_cancel(stream, run_ctx):
for chunk in _iter_json_chunks(chunk_bytes):
if chunk.get("status") == "loading":
await writer.append(chunk)
continue
await writer.flush()
status = chunk.get("status") or "event"
await append_run_event(run_id, status, {"chunk": chunk})
if status == "finished":
await mark_run_terminal(run_id, "completed")
terminal_set = True
elif status == "error":
await mark_run_terminal(
run_id,
"failed",
error_type=chunk.get("error_type") or "stream_error",
error_message=chunk.get("error_message") or chunk.get("message"),
)
terminal_set = True
elif status == "interrupted":
status_value = "cancelled" if await _is_cancel_requested(run_id) else "interrupted"
await mark_run_terminal(
run_id,
status_value,
error_type=status_value,
error_message=chunk.get("message"),
)
terminal_set = True
if await run_ctx.is_cancelled():
raise asyncio.CancelledError(f"run {run_id} cancelled")
await writer.flush()
if not terminal_set:
await mark_run_terminal(run_id, "completed")
await append_run_event(run_id, "finished", {"chunk": {"status": "finished", "request_id": request_id}})
except asyncio.CancelledError:
await writer.flush()
await append_run_event(
run_id,
"interrupted",
{"chunk": {"status": "interrupted", "message": "对话已取消", "request_id": request_id}},
)
await mark_run_terminal(run_id, "cancelled", error_type="cancelled", error_message="对话已取消")
logger.info(f"Run cancelled: {run_id}")
except Exception as e:
await writer.flush()
if _is_retryable_exception(e):
job_try = _job_try(ctx)
logger.warning(f"Run retryable failure {run_id} (try={job_try}): {e}")
await append_run_event(
run_id,
"error",
{
"chunk": {
"status": "error",
"error_type": "retryable_worker_error",
"error_message": str(e),
"request_id": request_id,
"retryable": True,
"job_try": job_try,
}
},
)
if _is_last_try(ctx):
await mark_run_terminal(
run_id,
"failed",
error_type="retryable_worker_error",
error_message=str(e),
)
logger.error(f"Run failed after retries exhausted {run_id}: {e}")
return
if isinstance(e, RetryableRunError):
raise
raise RetryableRunError(str(e)) from e
logger.error(f"Run failed {run_id}: {e}")
await append_run_event(
run_id,
"error",
{
"chunk": {
"status": "error",
"error_type": "worker_error",
"error_message": str(e),
"request_id": request_id,
"retryable": False,
}
},
)
await mark_run_terminal(run_id, "failed", error_type="worker_error", error_message=str(e))
return
finally:
await run_ctx.close()
await clear_cancel_signal(run_id)
async def _worker_startup(ctx):
pg_manager.initialize()
await pg_manager.create_business_tables()
await pg_manager.ensure_business_schema()
async def _worker_shutdown(ctx):
await pg_manager.close()
class WorkerSettings:
functions = [process_agent_run]
max_tries = 2
retry_jobs = True
job_timeout = 900
keep_result = 60
on_startup = _worker_startup
on_shutdown = _worker_shutdown
try:
from arq.connections import RedisSettings
redis_settings = RedisSettings.from_dsn(REDIS_URL)
except Exception:
redis_settings = None

View File

@ -172,6 +172,26 @@ class PostgresManager(metaclass=SingletonMeta):
"ALTER TABLE IF EXISTS skills ADD COLUMN IF NOT EXISTS mcp_dependencies JSONB DEFAULT '[]'::jsonb",
"ALTER TABLE IF EXISTS skills ADD COLUMN IF NOT EXISTS skill_dependencies JSONB DEFAULT '[]'::jsonb",
"ALTER TABLE IF EXISTS mcp_servers ADD COLUMN IF NOT EXISTS env JSONB",
"""
CREATE TABLE IF NOT EXISTS agent_runs (
id VARCHAR(64) PRIMARY KEY,
thread_id VARCHAR(64) NOT NULL,
agent_id VARCHAR(64) NOT NULL,
user_id VARCHAR(64) NOT NULL,
status VARCHAR(32) NOT NULL DEFAULT 'pending',
request_id VARCHAR(64) NOT NULL UNIQUE,
input_payload JSONB NOT NULL DEFAULT '{}'::jsonb,
error_type VARCHAR(64),
error_message TEXT,
started_at TIMESTAMPTZ,
finished_at TIMESTAMPTZ,
created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
)
""",
"CREATE INDEX IF NOT EXISTS idx_agent_runs_user_created ON agent_runs(user_id, created_at DESC)",
"CREATE INDEX IF NOT EXISTS idx_agent_runs_thread_created ON agent_runs(thread_id, created_at DESC)",
"CREATE INDEX IF NOT EXISTS idx_agent_runs_status_updated ON agent_runs(status, updated_at)",
]
async with self.async_engine.begin() as conn:
for stmt in stmts:

View File

@ -544,3 +544,46 @@ class TaskRecord(Base):
data.pop("payload", None)
data.pop("result", None)
return data
class AgentRun(Base):
"""AgentRun table - 运行任务表"""
__tablename__ = "agent_runs"
id = Column(String(64), primary_key=True, comment="Run ID (UUID)")
thread_id = Column(String(64), index=True, nullable=False, comment="Thread ID")
agent_id = Column(String(64), index=True, nullable=False, comment="Agent ID")
user_id = Column(String(64), index=True, nullable=False, comment="User ID")
status = Column(
String(32),
index=True,
nullable=False,
default="pending",
comment="Run status: pending/running/completed/failed/cancel_requested/cancelled/interrupted",
)
request_id = Column(String(64), unique=True, index=True, nullable=False, comment="Idempotency request ID")
input_payload = Column(JSON, nullable=False, default=dict, comment="Original input payload")
error_type = Column(String(64), nullable=True, comment="Error type")
error_message = Column(Text, nullable=True, comment="Error message")
started_at = Column(DateTime, nullable=True, comment="Start time")
finished_at = Column(DateTime, nullable=True, comment="Finish time")
created_at = Column(DateTime, default=utc_now_naive, comment="Creation time")
updated_at = Column(DateTime, default=utc_now_naive, onupdate=utc_now_naive, comment="Update time")
def to_dict(self) -> dict[str, Any]:
return {
"id": self.id,
"thread_id": self.thread_id,
"agent_id": self.agent_id,
"user_id": self.user_id,
"status": self.status,
"request_id": self.request_id,
"input_payload": self.input_payload or {},
"error_type": self.error_type,
"error_message": self.error_message,
"started_at": format_utc_datetime(self.started_at),
"finished_at": format_utc_datetime(self.finished_at),
"created_at": format_utc_datetime(self.created_at),
"updated_at": format_utc_datetime(self.updated_at),
}

View File

@ -0,0 +1,308 @@
from __future__ import annotations
from contextlib import asynccontextmanager
from types import SimpleNamespace
import pytest
from sqlalchemy.exc import IntegrityError
import src.services.agent_run_service as agent_run_service
@pytest.mark.asyncio
async def test_stream_agent_run_events_emits_error_and_close_on_db_error(monkeypatch: pytest.MonkeyPatch):
@asynccontextmanager
async def fake_session_ctx():
yield object()
class BrokenRepo:
def __init__(self, db):
self.db = db
async def get_run_for_user(self, run_id: str, user_id: str):
del run_id, user_id
raise RuntimeError("db down")
monkeypatch.setattr(agent_run_service.pg_manager, "get_async_session_context", fake_session_ctx)
monkeypatch.setattr(agent_run_service, "AgentRunRepository", BrokenRepo)
chunks = []
async for chunk in agent_run_service.stream_agent_run_events(
run_id="run-1",
after_seq="0",
current_user_id="1",
):
chunks.append(chunk)
assert len(chunks) == 2
assert chunks[0].startswith("event: error")
assert '"reason": "db_error"' in chunks[0]
assert chunks[1].startswith("event: close")
@pytest.mark.asyncio
async def test_stream_agent_run_events_reads_redis_and_close_terminal(monkeypatch: pytest.MonkeyPatch):
@asynccontextmanager
async def fake_session_ctx():
yield object()
class Repo:
def __init__(self, db):
self.db = db
async def get_run_for_user(self, run_id: str, user_id: str):
del run_id, user_id
return SimpleNamespace(status="completed")
calls = {"count": 0}
async def fake_list_events(run_id: str, *, after_seq: str, limit: int):
del run_id, after_seq, limit
calls["count"] += 1
if calls["count"] == 1:
return [
{
"seq": "1700000000000-0",
"event_type": "loading",
"payload": {"items": [{"status": "loading", "response": ""}]},
"ts": 1700000000000,
}
]
return []
async def fake_last_seq(run_id: str):
del run_id
return "1700000000000-0"
monkeypatch.setattr(agent_run_service.pg_manager, "get_async_session_context", fake_session_ctx)
monkeypatch.setattr(agent_run_service, "AgentRunRepository", Repo)
monkeypatch.setattr(agent_run_service, "list_run_stream_events", fake_list_events)
monkeypatch.setattr(agent_run_service, "get_last_run_stream_seq", fake_last_seq)
monkeypatch.setattr(agent_run_service, "SSE_POLL_INTERVAL_SECONDS", 0)
chunks = []
async for chunk in agent_run_service.stream_agent_run_events(
run_id="run-1",
after_seq="0",
current_user_id="1",
):
chunks.append(chunk)
assert any(item.startswith("event: loading") for item in chunks)
assert chunks[-1].startswith("event: close")
assert '"last_seq": "1700000000000-0"' in chunks[-1]
@pytest.mark.asyncio
async def test_create_agent_run_commits_before_enqueue(monkeypatch: pytest.MonkeyPatch):
class FakeDB:
def __init__(self):
self.order: list[str] = []
self.committed = False
async def commit(self):
self.order.append("commit")
self.committed = True
async def rollback(self):
raise AssertionError("rollback should not be called")
db = FakeDB()
created_run = SimpleNamespace(
id="run-1",
thread_id="thread-1",
status="pending",
request_id="req-1",
user_id="1",
)
class Repo:
def __init__(self, db_session):
self.db = db_session
async def get_run_by_request_id(self, request_id: str):
del request_id
return None
async def create_run(self, **kwargs):
assert kwargs["request_id"] == "req-1"
return created_run
class ConvRepo:
def __init__(self, db_session):
self.db = db_session
async def get_conversation_by_thread_id(self, thread_id: str):
del thread_id
return SimpleNamespace(user_id="1", status="active")
class Queue:
async def enqueue_job(self, job_name: str, run_id: str, _job_id: str):
assert job_name == "process_agent_run"
assert run_id == "run-1"
assert _job_id == "run:run-1"
db.order.append("enqueue")
assert db.committed is True
async def fake_get_arq_pool():
return Queue()
monkeypatch.setattr(agent_run_service.agent_manager, "get_agent", lambda agent_id: object())
monkeypatch.setattr(agent_run_service, "ConversationRepository", ConvRepo)
monkeypatch.setattr(agent_run_service, "AgentRunRepository", Repo)
monkeypatch.setattr(agent_run_service, "get_arq_pool", fake_get_arq_pool)
result = await agent_run_service.create_agent_run_view(
agent_id="ChatbotAgent",
query="hello",
config={"thread_id": "thread-1", "request_id": "req-1"},
image_content=None,
current_user_id="1",
db=db,
)
assert db.order == ["commit", "enqueue"]
assert result["run_id"] == "run-1"
assert result["request_id"] == "req-1"
@pytest.mark.asyncio
async def test_create_agent_run_handles_integrity_error_with_same_user_existing(monkeypatch: pytest.MonkeyPatch):
class FakeDB:
def __init__(self):
self.commit_called = 0
self.rollback_called = 0
async def commit(self):
self.commit_called += 1
raise IntegrityError("insert", {"request_id": "req-1"}, Exception("duplicate"))
async def rollback(self):
self.rollback_called += 1
db = FakeDB()
existing_run = SimpleNamespace(
id="run-existing",
thread_id="thread-1",
status="running",
request_id="req-1",
user_id="1",
)
state = {"lookup_count": 0}
class Repo:
def __init__(self, db_session):
self.db = db_session
async def get_run_by_request_id(self, request_id: str):
del request_id
state["lookup_count"] += 1
if state["lookup_count"] == 1:
return None
return existing_run
async def create_run(self, **kwargs):
del kwargs
return SimpleNamespace(
id="run-new",
thread_id="thread-1",
status="pending",
request_id="req-1",
user_id="1",
)
class ConvRepo:
def __init__(self, db_session):
self.db = db_session
async def get_conversation_by_thread_id(self, thread_id: str):
del thread_id
return SimpleNamespace(user_id="1", status="active")
async def fake_get_arq_pool():
raise AssertionError("should not enqueue on integrity fallback")
monkeypatch.setattr(agent_run_service.agent_manager, "get_agent", lambda agent_id: object())
monkeypatch.setattr(agent_run_service, "ConversationRepository", ConvRepo)
monkeypatch.setattr(agent_run_service, "AgentRunRepository", Repo)
monkeypatch.setattr(agent_run_service, "get_arq_pool", fake_get_arq_pool)
result = await agent_run_service.create_agent_run_view(
agent_id="ChatbotAgent",
query="hello",
config={"thread_id": "thread-1", "request_id": "req-1"},
image_content=None,
current_user_id="1",
db=db,
)
assert db.commit_called == 1
assert db.rollback_called == 1
assert result["run_id"] == "run-existing"
assert result["status"] == "running"
@pytest.mark.asyncio
async def test_create_agent_run_integrity_error_returns_409_for_other_user(monkeypatch: pytest.MonkeyPatch):
class FakeDB:
async def commit(self):
raise IntegrityError("insert", {"request_id": "req-1"}, Exception("duplicate"))
async def rollback(self):
return None
db = FakeDB()
existing_run = SimpleNamespace(
id="run-existing",
thread_id="thread-1",
status="pending",
request_id="req-1",
user_id="2",
)
state = {"lookup_count": 0}
class Repo:
def __init__(self, db_session):
self.db = db_session
async def get_run_by_request_id(self, request_id: str):
del request_id
state["lookup_count"] += 1
if state["lookup_count"] == 1:
return None
return existing_run
async def create_run(self, **kwargs):
del kwargs
return SimpleNamespace(
id="run-new",
thread_id="thread-1",
status="pending",
request_id="req-1",
user_id="1",
)
class ConvRepo:
def __init__(self, db_session):
self.db = db_session
async def get_conversation_by_thread_id(self, thread_id: str):
del thread_id
return SimpleNamespace(user_id="1", status="active")
monkeypatch.setattr(agent_run_service.agent_manager, "get_agent", lambda agent_id: object())
monkeypatch.setattr(agent_run_service, "ConversationRepository", ConvRepo)
monkeypatch.setattr(agent_run_service, "AgentRunRepository", Repo)
with pytest.raises(agent_run_service.HTTPException) as exc:
await agent_run_service.create_agent_run_view(
agent_id="ChatbotAgent",
query="hello",
config={"thread_id": "thread-1", "request_id": "req-1"},
image_content=None,
current_user_id="1",
db=db,
)
assert exc.value.status_code == 409
assert exc.value.detail == "request_id 冲突"

View File

@ -0,0 +1,122 @@
from __future__ import annotations
import pytest
import src.services.run_queue_service as run_queue_service
class _FakeRedisClient:
def __init__(self, *, ping_error: Exception | None = None):
self.ping_error = ping_error
self.ping_calls = 0
self.closed = False
async def ping(self):
self.ping_calls += 1
if self.ping_error is not None:
raise self.ping_error
async def aclose(self):
self.closed = True
class _FakeStreamRedis:
def __init__(self):
self.streams: dict[str, list[tuple[str, dict[str, str]]]] = {}
self.expire_calls: list[tuple[str, int]] = []
async def xadd(self, key: str, fields: dict[str, str], **kwargs):
del kwargs
stream = self.streams.setdefault(key, [])
event_id = f"{1700000000000 + len(stream)}-0"
stream.append((event_id, dict(fields)))
return event_id
async def expire(self, key: str, ttl: int):
self.expire_calls.append((key, ttl))
async def xrange(self, key: str, min: str, max: str, count: int):
del max
rows = list(self.streams.get(key, []))
if min.startswith("("):
cursor = min[1:]
rows = [(event_id, fields) for event_id, fields in rows if event_id > cursor]
elif min == "-":
rows = list(rows)
return rows[:count]
async def xrevrange(self, key: str, max: str, min: str, count: int):
del max, min
rows = list(reversed(self.streams.get(key, [])))
return rows[:count]
@pytest.mark.asyncio
async def test_get_redis_client_ping_once_and_cache(monkeypatch: pytest.MonkeyPatch):
redis_asyncio = pytest.importorskip("redis.asyncio")
fake_client = _FakeRedisClient()
monkeypatch.setattr(run_queue_service, "_redis_client", None)
monkeypatch.setattr(
redis_asyncio.Redis,
"from_url",
staticmethod(lambda *args, **kwargs: fake_client),
)
client_1 = await run_queue_service.get_redis_client()
client_2 = await run_queue_service.get_redis_client()
assert client_1 is fake_client
assert client_2 is fake_client
assert fake_client.ping_calls == 1
@pytest.mark.asyncio
async def test_get_redis_client_ping_fail_fast(monkeypatch: pytest.MonkeyPatch):
redis_asyncio = pytest.importorskip("redis.asyncio")
fake_client = _FakeRedisClient(ping_error=RuntimeError("redis unavailable"))
monkeypatch.setattr(run_queue_service, "_redis_client", None)
monkeypatch.setattr(
redis_asyncio.Redis,
"from_url",
staticmethod(lambda *args, **kwargs: fake_client),
)
with pytest.raises(RuntimeError, match="Redis connection failed"):
await run_queue_service.get_redis_client()
assert run_queue_service._redis_client is None
assert fake_client.closed is True
@pytest.mark.asyncio
async def test_run_stream_event_roundtrip(monkeypatch: pytest.MonkeyPatch):
fake_redis = _FakeStreamRedis()
monkeypatch.setattr(run_queue_service, "_redis_client", fake_redis)
run_id = "run-1"
seq1 = await run_queue_service.append_run_stream_event(run_id, "loading", {"items": [1]})
seq2 = await run_queue_service.append_run_stream_event(run_id, "finished", {"chunk": {"status": "finished"}})
assert seq1 < seq2
events = await run_queue_service.list_run_stream_events(run_id, after_seq="0-0", limit=100)
assert [item["event_type"] for item in events] == ["loading", "finished"]
assert events[0]["payload"] == {"items": [1]}
next_events = await run_queue_service.list_run_stream_events(run_id, after_seq=seq1, limit=100)
assert len(next_events) == 1
assert next_events[0]["seq"] == seq2
last_seq = await run_queue_service.get_last_run_stream_seq(run_id)
assert last_seq == seq2
def test_normalize_after_seq_stream_id_only():
assert run_queue_service.normalize_after_seq(None) == "0-0"
assert run_queue_service.normalize_after_seq(0) == "0-0"
assert run_queue_service.normalize_after_seq(5) == "0-0"
assert run_queue_service.normalize_after_seq("1700000000000-3") == "1700000000000-3"
assert run_queue_service.normalize_after_seq("12") == "0-0"
assert run_queue_service.normalize_after_seq("bad-value") == "0-0"

154
test/test_run_worker.py Normal file
View File

@ -0,0 +1,154 @@
from __future__ import annotations
from contextlib import asynccontextmanager
from types import SimpleNamespace
import pytest
import src.services.run_worker as run_worker
class _RaisingAsyncIter:
def __init__(self, exc: Exception):
self._exc = exc
def __aiter__(self):
return self
async def __anext__(self):
raise self._exc
class _BytesAsyncIter:
def __init__(self, values: list[bytes]):
self._values = list(values)
self._idx = 0
def __aiter__(self):
return self
async def __anext__(self):
if self._idx >= len(self._values):
raise StopAsyncIteration
value = self._values[self._idx]
self._idx += 1
return value
def _build_run() -> SimpleNamespace:
return SimpleNamespace(
status="pending",
request_id="req-1",
input_payload={
"query": "hello",
"config": {"thread_id": "thread-1"},
"agent_id": "ChatbotAgent",
"image_content": None,
"user_id": "1",
"request_id": "req-1",
},
)
def _patch_common(monkeypatch: pytest.MonkeyPatch, run_obj: SimpleNamespace):
@asynccontextmanager
async def fake_session_ctx():
yield object()
async def fake_noop(*args, **kwargs):
del args, kwargs
return None
async def fake_get_run(run_id: str):
del run_id
return run_obj
async def fake_load_user(user_id: str):
del user_id
return SimpleNamespace(id=1)
async def fake_not_cancelled(self):
del self
return False
monkeypatch.setattr(run_worker.pg_manager, "get_async_session_context", fake_session_ctx)
monkeypatch.setattr(run_worker, "_get_run", fake_get_run)
monkeypatch.setattr(run_worker, "_load_user", fake_load_user)
monkeypatch.setattr(run_worker, "mark_run_running", fake_noop)
monkeypatch.setattr(run_worker, "clear_cancel_signal", fake_noop)
monkeypatch.setattr(run_worker, "stream_agent_chat", lambda **kwargs: object())
monkeypatch.setattr(run_worker.RunContext, "start", fake_noop)
monkeypatch.setattr(run_worker.RunContext, "close", fake_noop)
monkeypatch.setattr(run_worker.RunContext, "is_cancelled", fake_not_cancelled)
@pytest.mark.asyncio
async def test_process_agent_run_non_retryable_error_marks_failed(monkeypatch: pytest.MonkeyPatch):
run_obj = _build_run()
_patch_common(monkeypatch, run_obj)
terminal_statuses: list[str] = []
events: list[str] = []
async def fake_append_event(run_id: str, event_type: str, payload: dict):
del run_id, payload
events.append(event_type)
async def fake_mark_terminal(run_id: str, status: str, error_type=None, error_message=None):
del run_id, error_type, error_message
terminal_statuses.append(status)
monkeypatch.setattr(run_worker, "append_run_event", fake_append_event)
monkeypatch.setattr(run_worker, "mark_run_terminal", fake_mark_terminal)
monkeypatch.setattr(
run_worker,
"_consume_stream_with_cancel",
lambda stream, run_ctx: _RaisingAsyncIter(RuntimeError("boom")),
)
await run_worker.process_agent_run({"job_try": 1}, "run-1")
assert "error" in events
assert terminal_statuses == ["failed"]
@pytest.mark.asyncio
async def test_process_agent_run_retryable_error_retries_then_completes(monkeypatch: pytest.MonkeyPatch):
run_obj = _build_run()
_patch_common(monkeypatch, run_obj)
terminal_statuses: list[str] = []
events: list[dict] = []
attempts = {"count": 0}
async def fake_append_event(run_id: str, event_type: str, payload: dict):
del run_id
events.append({"event_type": event_type, "payload": payload})
async def fake_mark_terminal(run_id: str, status: str, error_type=None, error_message=None):
del run_id, error_type, error_message
terminal_statuses.append(status)
def fake_consume(stream, run_ctx):
del stream, run_ctx
attempts["count"] += 1
if attempts["count"] == 1:
return _RaisingAsyncIter(run_worker.RetryableRunError("temporary failure"))
return _BytesAsyncIter([b'{"status":"finished","request_id":"req-1"}\n'])
monkeypatch.setattr(run_worker, "append_run_event", fake_append_event)
monkeypatch.setattr(run_worker, "mark_run_terminal", fake_mark_terminal)
monkeypatch.setattr(run_worker, "_consume_stream_with_cancel", fake_consume)
with pytest.raises(run_worker.RetryableRunError):
await run_worker.process_agent_run({"job_try": 1}, "run-1")
assert terminal_statuses == []
assert any(
item["event_type"] == "error"
and item["payload"]["chunk"].get("error_type") == "retryable_worker_error"
for item in events
)
await run_worker.process_agent_run({"job_try": 2}, "run-1")
assert terminal_statuses == ["completed"]

3991
uv.lock

File diff suppressed because it is too large Load Diff

View File

@ -195,6 +195,58 @@ export const agentApi = {
},
...restOptions
})
},
/**
* 创建异步运行任务Run
* @param {string} agentId - 智能体ID
* @param {Object} data - run 请求体
* @returns {Promise<Object>}
*/
createAgentRun: (agentId, data) =>
apiPost(`/api/chat/agent/${agentId}/runs`, {
query: data.query,
config: data.config || {},
image_content: data.image_content || null
}),
/**
* 获取 Run 状态
* @param {string} runId - run ID
* @returns {Promise<Object>}
*/
getAgentRun: (runId) => apiGet(`/api/chat/runs/${runId}`),
/**
* 取消 Run
* @param {string} runId - run ID
* @returns {Promise<Object>}
*/
cancelAgentRun: (runId) => apiPost(`/api/chat/runs/${runId}/cancel`, {}),
/**
* 获取线程活跃 Run
* @param {string} threadId - 线程ID
* @returns {Promise<Object>}
*/
getThreadActiveRun: (threadId) => apiGet(`/api/chat/thread/${threadId}/active_run`),
/**
* 打开 Run 事件 SSE 连接调用方负责关闭
* @param {string} runId - run ID
* @param {string|number} afterSeq - 起始 seq/cursor
* @param {Object} options - { signal }
* @returns {Promise<Response>}
*/
streamAgentRunEvents: (runId, afterSeq = '0', options = {}) => {
const { signal } = options
return fetch(`/api/chat/runs/${runId}/events?after_seq=${encodeURIComponent(String(afterSeq))}`, {
method: 'GET',
headers: {
...useUserStore().getAuthHeaders()
},
signal
})
}
}

View File

@ -242,6 +242,19 @@ const {
// ==================== LOCAL CHAT & UI STATE ====================
const userInput = ref('')
const useRunsApi =
import.meta.env.VITE_USE_RUNS_API === 'true' &&
localStorage.getItem('force_legacy_stream') !== 'true'
const ACTIVE_RUN_STORAGE_TTL_MS = 60 * 60 * 1000
const ACTIVE_RUN_CLIENT_ID = `${Date.now()}-${Math.random().toString(36).slice(2, 10)}`
const typingChunkQueue = []
let typingAnimationFrameId = null
let typingLastFrameTs = 0
let pendingTypingChars = 0
const MIN_TYPING_CPS = 32
const MAX_TYPING_CPS = 320
const TYPING_BACKLOG_HIGH_WATERMARK = 500
//
const exampleQuestions = computed(() => {
@ -347,21 +360,11 @@ const currentAgentState = computed(() => {
})
const countFiles = (files) => {
// dict StateBackend array
if (!files) return 0
if (typeof files === 'object' && !Array.isArray(files)) {
// dict : {"/attachments/file.md": {...}, ...}
return Object.keys(files).length
}
if (Array.isArray(files)) {
// array
let c = 0
for (const item of files) {
if (item && typeof item === 'object') c += Object.keys(item).length
}
return c
return files.reduce((c, item) => c + (item && typeof item === 'object' ? Object.keys(item).length : 0), 0)
}
return 0
return typeof files === 'object' ? Object.keys(files).length : 0
}
const hasAgentStateContent = computed(() => {
@ -508,6 +511,9 @@ onMounted(() => {
})
onUnmounted(() => {
stopTypingRenderLoop()
typingChunkQueue.length = 0
pendingTypingChars = 0
scrollController.cleanup()
// 线
resetOnGoingConv()
@ -521,6 +527,10 @@ const getThreadState = (threadId) => {
chatState.threadStates[threadId] = {
isStreaming: false,
streamAbortController: null,
runStreamAbortController: null,
activeRunId: null,
runLastSeq: '0',
lastRetryableJobTry: null,
onGoingConv: createOnGoingConvState(),
agentState: null // agentState
}
@ -533,30 +543,34 @@ const cleanupThreadState = (threadId) => {
if (!threadId) return
const threadState = chatState.threadStates[threadId]
if (threadState) {
clearTypingQueueForThread(threadId)
if (threadState.streamAbortController) {
threadState.streamAbortController.abort()
}
if (threadState.runStreamAbortController) {
threadState.runStreamAbortController.abort()
}
delete chatState.threadStates[threadId]
}
}
// ==================== STREAM HANDLING LOGIC ====================
const resetOnGoingConv = (threadId = null) => {
console.log(
`🔄 [RESET] Resetting on going conversation: ${new Date().toLocaleTimeString()}.${new Date().getMilliseconds()}`,
threadId
)
const targetThreadId = threadId || currentChatId.value
if (targetThreadId) {
// 线
const threadState = getThreadState(targetThreadId)
if (threadState) {
clearTypingQueueForThread(targetThreadId)
if (threadState.streamAbortController) {
threadState.streamAbortController.abort()
threadState.streamAbortController = null
}
if (threadState.runStreamAbortController) {
threadState.runStreamAbortController.abort()
threadState.runStreamAbortController = null
}
//
threadState.onGoingConv = createOnGoingConvState()
@ -665,10 +679,6 @@ const fetchThreadMessages = async ({ agentId, threadId, delay = 0 }) => {
try {
const response = await agentApi.getAgentHistory(agentId, threadId)
console.log(
`🔄 [FETCH] Thread messages: ${new Date().toLocaleTimeString()}.${new Date().getMilliseconds()}`,
response
)
threadMessages.value[threadId] = response.history || []
} catch (error) {
handleChatError(error, 'load')
@ -680,24 +690,11 @@ const fetchAgentState = async (agentId, threadId) => {
if (!agentId || !threadId) return
try {
const res = await agentApi.getAgentState(agentId, threadId)
// currentChatId state currentAgentState
// currentChatId null使 threadId
const targetChatId = currentChatId.value || threadId
console.log(
'[fetchAgentState] agentId:',
agentId,
'threadId:',
threadId,
'targetChatId:',
targetChatId,
'agent_state:',
JSON.stringify(res.agent_state || {})?.slice(0, 200)
)
const ts = getThreadState(targetChatId)
if (ts) {
ts.agentState = res.agent_state || null
} else {
// targetChatId state
const newTs = getThreadState(threadId)
if (newTs) newTs.agentState = res.agent_state || null
}
@ -706,6 +703,434 @@ const fetchAgentState = async (agentId, threadId) => {
}
}
const RUN_TERMINAL_STATUSES = new Set(['completed', 'failed', 'cancelled', 'interrupted'])
const getActiveRunStorageKey = (threadId) => `active_run:${threadId}`
const normalizeRunSeq = (value) => {
if (value === undefined || value === null) return '0'
const text = String(value).trim()
return text || '0'
}
const parseRunSeq = (value) => {
const text = normalizeRunSeq(value)
if (text.includes('-')) {
const [majorRaw, minorRaw] = text.split('-', 2)
let major = 0n
let minor = 0n
try {
major = BigInt(majorRaw || '0')
minor = BigInt(minorRaw || '0')
} catch {
return { kind: 'legacy', value: 0 }
}
return { kind: 'stream', major, minor }
}
const numberValue = Number.parseInt(text, 10)
if (!Number.isNaN(numberValue)) {
return { kind: 'legacy', value: numberValue }
}
return { kind: 'legacy', value: 0 }
}
const compareRunSeq = (incoming, current) => {
const left = parseRunSeq(incoming)
const right = parseRunSeq(current)
if (left.kind === 'stream' && right.kind === 'stream') {
if (left.major > right.major) return 1
if (left.major < right.major) return -1
if (left.minor > right.minor) return 1
if (left.minor < right.minor) return -1
return 0
}
if (left.kind === 'legacy' && right.kind === 'legacy') {
return left.value - right.value
}
if (left.kind === 'stream' && right.kind === 'legacy') return 1
return -1
}
const saveActiveRunSnapshot = (threadId, runId, lastSeq = '0') => {
if (!threadId || !runId) return
localStorage.setItem(
getActiveRunStorageKey(threadId),
JSON.stringify({
run_id: runId,
last_seq: normalizeRunSeq(lastSeq),
created_at: Date.now(),
client_id: ACTIVE_RUN_CLIENT_ID
})
)
}
const loadActiveRunSnapshot = (threadId) => {
if (!threadId) return null
try {
const raw = localStorage.getItem(getActiveRunStorageKey(threadId))
return raw ? JSON.parse(raw) : null
} catch {
return null
}
}
const clearActiveRunSnapshot = (threadId) => {
if (!threadId) return
localStorage.removeItem(getActiveRunStorageKey(threadId))
}
const splitChars = (text) => {
if (typeof text !== 'string' || !text) return []
return Array.from(text)
}
const calcTypingCps = () => {
if (pendingTypingChars <= 0) return MIN_TYPING_CPS
const ratio = Math.min(1, pendingTypingChars / TYPING_BACKLOG_HIGH_WATERMARK)
return Math.round(MIN_TYPING_CPS + ratio * (MAX_TYPING_CPS - MIN_TYPING_CPS))
}
const scheduleTypingRender = () => {
if (typingAnimationFrameId !== null) return
if (typeof window !== 'undefined' && typeof window.requestAnimationFrame === 'function') {
typingAnimationFrameId = window.requestAnimationFrame(drainTypingQueue)
} else {
typingAnimationFrameId = setTimeout(() => {
typingAnimationFrameId = null
drainTypingQueue(Date.now())
}, 16)
}
}
const stopTypingRenderLoop = () => {
if (typingAnimationFrameId === null) return
if (typeof window !== 'undefined' && typeof window.cancelAnimationFrame === 'function') {
window.cancelAnimationFrame(typingAnimationFrameId)
} else {
clearTimeout(typingAnimationFrameId)
}
typingAnimationFrameId = null
typingLastFrameTs = 0
}
const enqueueLoadingChunkForTyping = (threadId, chunk) => {
if (!threadId || !chunk) return
if (chunk.status !== 'loading') {
typingChunkQueue.push({ threadId, chunk, isChar: false })
scheduleTypingRender()
return
}
const msg = chunk.msg || {}
const contentChars = splitChars(msg.content)
if (contentChars.length === 0) {
typingChunkQueue.push({ threadId, chunk, isChar: false })
scheduleTypingRender()
return
}
for (const char of contentChars) {
typingChunkQueue.push({
threadId,
isChar: true,
chunk: {
...chunk,
msg: {
...msg,
content: char
}
}
})
}
pendingTypingChars += contentChars.length
scheduleTypingRender()
}
// threadId
const partitionTypingQueue = (threadId, onMatch = null) => {
const remaining = []
let charCount = 0
for (const item of typingChunkQueue) {
if (item.threadId === threadId) {
if (onMatch) onMatch(item)
if (item.isChar) charCount += 1
} else {
remaining.push(item)
}
}
typingChunkQueue.length = 0
typingChunkQueue.push(...remaining)
pendingTypingChars = Math.max(0, pendingTypingChars - charCount)
typingChunkQueue.length === 0 ? stopTypingRenderLoop() : scheduleTypingRender()
}
const clearTypingQueueForThread = (threadId) => {
if (!threadId || typingChunkQueue.length === 0) return
partitionTypingQueue(threadId)
}
const flushTypingQueueForThread = (threadId) => {
if (!threadId || typingChunkQueue.length === 0) return
partitionTypingQueue(threadId, (item) => handleStreamChunk(item.chunk, threadId))
}
function drainTypingQueue(frameTs = Date.now()) {
typingAnimationFrameId = null
if (typingChunkQueue.length === 0) {
typingLastFrameTs = 0
return
}
if (!typingLastFrameTs) {
typingLastFrameTs = frameTs
}
const elapsedSeconds = Math.max(0.001, (frameTs - typingLastFrameTs) / 1000)
typingLastFrameTs = frameTs
const cps = calcTypingCps()
let budget = Math.max(1, Math.floor(elapsedSeconds * cps))
while (budget > 0 && typingChunkQueue.length > 0) {
const item = typingChunkQueue.shift()
if (!item) break
handleStreamChunk(item.chunk, item.threadId)
if (item.isChar) {
pendingTypingChars = Math.max(0, pendingTypingChars - 1)
}
budget -= 1
}
if (typingChunkQueue.length > 0) {
scheduleTypingRender()
} else {
typingLastFrameTs = 0
}
}
const processRunSseResponse = async (response, onEvent) => {
if (!response || !response.body) return
const reader = response.body.getReader()
const decoder = new TextDecoder()
let buffer = ''
let eventType = 'message'
let dataLines = []
try {
while (true) {
const { done, value } = await reader.read()
if (done) break
buffer += decoder.decode(value, { stream: true })
const lines = buffer.split('\n')
buffer = lines.pop() || ''
for (const rawLine of lines) {
const line = rawLine.replace(/\r$/, '')
if (!line) {
if (dataLines.length > 0) {
const dataText = dataLines.join('\n')
try {
const parsed = JSON.parse(dataText)
onEvent(eventType, parsed)
} catch (e) {
console.warn('Failed to parse run SSE data:', e, dataText)
}
}
eventType = 'message'
dataLines = []
continue
}
if (line.startsWith('event:')) {
eventType = line.slice(6).trim() || 'message'
} else if (line.startsWith('data:')) {
dataLines.push(line.slice(5).trim())
}
}
}
if (dataLines.length > 0) {
const dataText = dataLines.join('\n')
try {
const parsed = JSON.parse(dataText)
onEvent(eventType, parsed)
} catch (e) {
console.warn('Failed to parse trailing run SSE data:', e, dataText)
}
}
} finally {
try {
reader.releaseLock()
} catch {
// ignore
}
}
}
const stopRunStreamSubscription = (threadId) => {
const ts = getThreadState(threadId)
if (!ts) return
if (ts.runStreamAbortController) {
ts.runStreamAbortController.abort()
ts.runStreamAbortController = null
}
}
const startRunStream = async (threadId, runId, afterSeq = '0') => {
if (!threadId || !runId || !useRunsApi) return
const ts = getThreadState(threadId)
if (!ts) return
stopRunStreamSubscription(threadId)
const runController = new AbortController()
ts.runStreamAbortController = runController
ts.activeRunId = runId
ts.runLastSeq = normalizeRunSeq(afterSeq)
ts.lastRetryableJobTry = null
ts.isStreaming = true
saveActiveRunSnapshot(threadId, runId, ts.runLastSeq)
try {
const response = await agentApi.streamAgentRunEvents(runId, ts.runLastSeq, {
signal: runController.signal
})
if (!response.ok) {
throw new Error(`SSE response not ok: ${response.status}`)
}
await processRunSseResponse(response, (event, data) => {
if (!data || ts.activeRunId !== runId) return
if (data.seq !== undefined && data.seq !== null) {
const incomingSeq = normalizeRunSeq(data.seq)
if (compareRunSeq(incomingSeq, ts.runLastSeq) <= 0) return
ts.runLastSeq = incomingSeq
saveActiveRunSnapshot(threadId, runId, incomingSeq)
}
if (event === 'heartbeat') return
const payload = data.payload || {}
const isRetryableError = event === 'error' && payload?.chunk?.retryable === true
if (isRetryableError) {
const parsedJobTry = Number.parseInt(payload?.chunk?.job_try, 10)
const retryJobTry = Number.isNaN(parsedJobTry) ? null : parsedJobTry
if (retryJobTry !== null && ts.lastRetryableJobTry === retryJobTry) {
return
}
ts.lastRetryableJobTry = retryJobTry
console.warn('Run encountered retryable error, waiting for worker retry', {
threadId,
runId,
retryJobTry,
errorType: payload?.chunk?.error_type
})
return
}
if (Array.isArray(payload.items)) {
payload.items.forEach((chunk) => {
enqueueLoadingChunkForTyping(threadId, chunk)
})
} else if (payload.chunk) {
enqueueLoadingChunkForTyping(threadId, payload.chunk)
}
if (event === 'close') {
flushTypingQueueForThread(threadId)
ts.isStreaming = false
if (RUN_TERMINAL_STATUSES.has(data.status)) {
ts.activeRunId = null
ts.lastRetryableJobTry = null
clearActiveRunSnapshot(threadId)
fetchThreadMessages({ agentId: currentAgentId.value, threadId, delay: 200 }).finally(() => {
fetchAgentState(currentAgentId.value, threadId)
})
} else if (ts.activeRunId === runId) {
window.setTimeout(() => {
if (ts.activeRunId === runId && !ts.runStreamAbortController) {
void startRunStream(threadId, runId, ts.runLastSeq)
}
}, 300)
}
}
if (event === 'finished' || event === 'error' || event === 'interrupted') {
flushTypingQueueForThread(threadId)
ts.isStreaming = false
ts.activeRunId = null
ts.lastRetryableJobTry = null
clearActiveRunSnapshot(threadId)
fetchThreadMessages({ agentId: currentAgentId.value, threadId, delay: 300 }).finally(() => {
resetOnGoingConv(threadId)
fetchAgentState(currentAgentId.value, threadId)
scrollController.scrollToBottom()
})
}
})
} catch (error) {
if (error?.name !== 'AbortError') {
console.error('Run SSE stream error:', error)
handleChatError(error, 'stream')
if (ts.activeRunId === runId) {
window.setTimeout(() => {
if (ts.activeRunId === runId && !ts.runStreamAbortController) {
void startRunStream(threadId, runId, ts.runLastSeq)
}
}, 500)
}
}
} finally {
if (ts.runStreamAbortController === runController) {
ts.runStreamAbortController = null
}
if (!ts.activeRunId) {
ts.isStreaming = false
}
}
}
const resumeActiveRunForThread = async (threadId) => {
if (!useRunsApi || !threadId) return
const ts = getThreadState(threadId)
if (!ts || ts.runStreamAbortController) return
const snapshot = loadActiveRunSnapshot(threadId)
if (snapshot?.run_id) {
if (Date.now() - Number(snapshot.created_at || 0) > ACTIVE_RUN_STORAGE_TTL_MS) {
clearActiveRunSnapshot(threadId)
} else {
try {
const runRes = await agentApi.getAgentRun(snapshot.run_id)
const run = runRes?.run
if (run && !RUN_TERMINAL_STATUSES.has(run.status)) {
await startRunStream(threadId, run.id, snapshot.last_seq || '0')
return
}
} catch {
// ignore
}
clearActiveRunSnapshot(threadId)
}
}
try {
const active = await agentApi.getThreadActiveRun(threadId)
const run = active?.run
if (run && !RUN_TERMINAL_STATUSES.has(run.status)) {
await startRunStream(threadId, run.id, 0)
return
}
} catch (e) {
console.warn('Failed to load active run for thread:', threadId, e)
}
ts.activeRunId = null
ts.runLastSeq = '0'
ts.isStreaming = false
clearActiveRunSnapshot(threadId)
}
const fetchMentionResources = async () => {
try {
const [dbsRes, mcpsRes] = await Promise.all([
@ -740,7 +1165,7 @@ const { approvalState, handleApproval, processApprovalInStream } = useApproval({
fetchThreadMessages
})
const { handleAgentResponse } = useAgentStreamHandler({
const { handleAgentResponse, handleStreamChunk } = useAgentStreamHandler({
getThreadState,
processApprovalInStream,
currentAgentId,
@ -866,6 +1291,8 @@ const selectChat = async (chatId) => {
previousThreadState.isStreaming = false
previousThreadState.streamAbortController = null
}
// run SSE
stopRunStreamSubscription(previousThreadId)
}
chatState.currentThreadId = chatId
@ -881,6 +1308,7 @@ const selectChat = async (chatId) => {
await nextTick()
scrollController.scrollToBottomStaticForce()
await fetchAgentState(currentAgentId.value, chatId)
await resumeActiveRunForThread(chatId)
}
const deleteChat = async (chatId) => {
@ -927,7 +1355,6 @@ const renameChat = async (data) => {
}
const handleSendMessage = async ({ image } = {}) => {
console.log('AgentChatComponent: handleSendMessage payload image:', image)
const text = userInput.value.trim()
if ((!text && !image) || !currentAgent.value || isProcessing.value) return
@ -948,18 +1375,42 @@ const handleSendMessage = async ({ image } = {}) => {
const threadState = getThreadState(threadId)
if (!threadState) return
if (useRunsApi) {
if ((threadMessages.value[threadId] || []).length === 0) {
const autoTitle = text.replace(/\s+/g, ' ').trim().slice(0, 255)
if (autoTitle) {
void updateThread(threadId, autoTitle).catch(() => {})
}
}
resetOnGoingConv(threadId)
threadState.isStreaming = true
try {
const runResp = await agentApi.createAgentRun(currentAgentId.value, {
query: text,
config: {
thread_id: threadId,
...(selectedAgentConfigId.value ? { agent_config_id: selectedAgentConfigId.value } : {})
},
image_content: image?.imageContent
})
const runId = runResp?.run_id
if (!runId) {
throw new Error('创建 run 失败:缺少 run_id')
}
await startRunStream(threadId, runId, 0)
} catch (error) {
threadState.isStreaming = false
handleChatError(error, 'send')
}
return
}
threadState.isStreaming = true
resetOnGoingConv(threadId)
threadState.streamAbortController = new AbortController()
try {
console.log('[AgentStateDebug][before_send]', {
threadId,
currentAgentId: currentAgentId.value,
capabilities: currentCapabilities.value,
supportsTodo: supportsTodo.value,
supportsFiles: supportsFiles.value
})
const response = await sendMessage({
agentId: currentAgentId.value,
threadId: threadId,
@ -978,22 +1429,13 @@ const handleSendMessage = async ({ image } = {}) => {
}
threadState.isStreaming = false
} finally {
console.log('[AgentStateDebug][send_finally_start]', {
threadId,
currentAgentId: currentAgentId.value,
supportsTodo: supportsTodo.value,
supportsFiles: supportsFiles.value
})
threadState.streamAbortController = null
//
fetchThreadMessages({ agentId: currentAgentId.value, threadId: threadId, delay: 500 }).finally(
() => {
console.log('[AgentStateDebug][history_refreshed_after_send]', {
threadId,
currentAgentId: currentAgentId.value
})
//
resetOnGoingConv(threadId)
fetchAgentState(currentAgentId.value, threadId)
scrollController.scrollToBottom()
}
)
@ -1004,26 +1446,37 @@ const handleSendMessage = async ({ image } = {}) => {
const handleSendOrStop = async (payload) => {
const threadId = currentChatId.value
const threadState = getThreadState(threadId)
if (isProcessing.value && threadState && threadState.streamAbortController) {
//
threadState.streamAbortController.abort()
//
try {
await fetchThreadMessages({ agentId: currentAgentId.value, threadId: threadId, delay: 500 })
message.info('已中断对话生成')
} catch (error) {
console.error('刷新消息历史失败:', error)
message.info('已中断对话生成')
if (isProcessing.value && threadState) {
if (useRunsApi && threadState.activeRunId) {
try {
await agentApi.cancelAgentRun(threadState.activeRunId)
message.info('已发送取消请求')
} catch (error) {
handleChatError(error, 'stop')
}
return
}
if (threadState.streamAbortController) {
//
threadState.streamAbortController.abort()
//
try {
await fetchThreadMessages({ agentId: currentAgentId.value, threadId: threadId, delay: 500 })
message.info('已中断对话生成')
} catch (error) {
console.error('刷新消息历史失败:', error)
message.info('已中断对话生成')
}
return
}
return
}
await handleSendMessage(payload)
}
// ==================== ====================
const handleApprovalWithStream = async (approved) => {
console.log('🔄 [STREAM] Starting resume stream processing')
const threadId = approvalState.threadId
if (!threadId) {
@ -1040,14 +1493,6 @@ const handleApprovalWithStream = async (approved) => {
}
try {
console.log('[AgentStateDebug][before_resume_send]', {
threadId,
currentAgentId: currentAgentId.value,
approved,
capabilities: currentCapabilities.value,
supportsTodo: supportsTodo.value,
supportsFiles: supportsFiles.value
})
// 使 composable
const response = await handleApproval(
approved,
@ -1057,28 +1502,13 @@ const handleApprovalWithStream = async (approved) => {
if (!response) return // handleApproval
console.log('🔄 [STREAM] Processing resume streaming response')
//
await handleAgentResponse(response, threadId, (chunk) => {
console.log('🔄 [STREAM] Processing chunk:', chunk)
})
console.log('🔄 [STREAM] Resume stream processing completed')
await handleAgentResponse(response, threadId)
} catch (error) {
console.error('❌ [STREAM] Resume stream failed:', error)
if (error.name !== 'AbortError') {
console.error('Resume approval error:', error)
// handleChatError useApproval
}
} finally {
console.log('🔄 [STREAM] Cleaning up streaming state')
console.log('[AgentStateDebug][resume_finally_start]', {
threadId,
currentAgentId: currentAgentId.value,
supportsTodo: supportsTodo.value,
supportsFiles: supportsFiles.value
})
if (threadState) {
threadState.isStreaming = false
threadState.streamAbortController = null
@ -1087,11 +1517,6 @@ const handleApprovalWithStream = async (approved) => {
//
fetchThreadMessages({ agentId: currentAgentId.value, threadId: threadId, delay: 500 }).finally(
() => {
console.log('[AgentStateDebug][history_refreshed_after_resume]', {
threadId,
currentAgentId: currentAgentId.value
})
//
resetOnGoingConv(threadId)
scrollController.scrollToBottom()
}
@ -1147,17 +1572,7 @@ const openAgentModal = () => emit('open-agent-modal')
const handleAgentStateRefresh = async (threadId = null) => {
if (!currentAgentId.value) return
// 使 threadId使 currentChatId
let chatId = threadId || currentChatId.value
console.log('[AgentStateDebug][manual_refresh_start]', {
inputThreadId: threadId,
currentChatId: currentChatId.value,
finalChatId: chatId,
currentAgentId: currentAgentId.value,
capabilities: currentCapabilities.value,
supportsTodo: supportsTodo.value,
supportsFiles: supportsFiles.value
})
const chatId = threadId || currentChatId.value
if (!chatId) return
await fetchAgentState(currentAgentId.value, chatId)
}
@ -1285,19 +1700,6 @@ onMounted(async () => {
scrollController.enableAutoScroll()
})
watch(
[currentAgentId, currentCapabilities, supportsTodo, supportsFiles],
([agentId, capabilities, todo, files]) => {
console.log('[AgentStateDebug][capabilities_changed]', {
currentAgentId: agentId,
capabilities,
supportsTodo: todo,
supportsFiles: files
})
},
{ immediate: true }
)
watch(
currentAgentId,
async (newAgentId, oldAgentId) => {

View File

@ -56,7 +56,7 @@ const processStreamResponse = async (response, onChunk) => {
} finally {
try {
reader.releaseLock()
} catch (e) {
} catch {
// Ignore errors on releasing lock
}
}
@ -128,7 +128,7 @@ export function useAgentStreamHandler({
todoCount: Array.isArray(chunk.agent_state?.todos) ? chunk.agent_state.todos.length : 0,
fileKeys: chunk.agent_state?.files ? Object.keys(chunk.agent_state.files) : []
})
if ((unref(supportsTodo) || unref(supportsFiles)) && chunk.agent_state) {
if (chunk.agent_state) {
console.log(`${debugPrefix}[agent_state_apply]`, {
threadId,
todos: chunk.agent_state?.todos || [],
@ -137,7 +137,7 @@ export function useAgentStreamHandler({
threadState.agentState = chunk.agent_state
} else {
console.warn(`${debugPrefix}[agent_state_skip]`, {
reason: 'capability_gate_or_empty_state',
reason: 'empty_state',
supportsTodo: unref(supportsTodo),
supportsFiles: unref(supportsFiles),
hasAgentState: !!chunk.agent_state,