merge: async run queue + SSE architecture
This commit is contained in:
commit
81663e7fba
@ -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
|
||||
|
||||
@ -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: .
|
||||
|
||||
@ -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:
|
||||
|
||||
185
docs/latest/changelog/run-architecture-migration-v4.md
Normal file
185
docs/latest/changelog/run-architecture-migration-v4.md
Normal 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. 最终消息与工具调用正常入库。
|
||||
@ -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 # 代码最大行宽
|
||||
|
||||
@ -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)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# > === 模型管理分组 ===
|
||||
# =============================================================================
|
||||
|
||||
@ -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}")
|
||||
|
||||
# 预热 Redis(run 队列)
|
||||
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
5
server/worker_main.py
Normal file
@ -0,0 +1,5 @@
|
||||
"""ARQ worker entrypoint."""
|
||||
|
||||
from src.services.run_worker import WorkerSettings
|
||||
|
||||
__all__ = ["WorkerSettings"]
|
||||
@ -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:
|
||||
|
||||
102
src/repositories/agent_run_repository.py
Normal file
102
src/repositories/agent_run_repository.py
Normal 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()
|
||||
246
src/services/agent_run_service.py
Normal file
246
src/services/agent_run_service.py
Normal 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}
|
||||
241
src/services/run_queue_service.py
Normal file
241
src/services/run_queue_service.py
Normal 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
373
src/services/run_worker.py
Normal 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
|
||||
@ -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:
|
||||
|
||||
@ -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),
|
||||
}
|
||||
|
||||
308
test/test_agent_run_service.py
Normal file
308
test/test_agent_run_service.py
Normal 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 冲突"
|
||||
122
test/test_run_queue_service.py
Normal file
122
test/test_run_queue_service.py
Normal 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
154
test/test_run_worker.py
Normal 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"]
|
||||
@ -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
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@ -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) => {
|
||||
|
||||
@ -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,
|
||||
|
||||
Loading…
Reference in New Issue
Block a user