mirror of
https://github.com/lintsinghua/DeepAudit.git
synced 2026-10-25 10:03:19 +08:00
Improve audit reliability, security, and application performance
This commit is contained in:
@@ -26,7 +26,10 @@ on:
|
||||
default: true
|
||||
|
||||
jobs:
|
||||
quality:
|
||||
uses: ./.github/workflows/quality.yml
|
||||
build-and-push:
|
||||
needs: quality
|
||||
name: 构建并推送镜像
|
||||
runs-on: ubuntu-latest
|
||||
|
||||
@@ -42,7 +45,7 @@ jobs:
|
||||
if: ${{ github.event.inputs.build_frontend == 'true' }}
|
||||
uses: actions/setup-node@v4
|
||||
with:
|
||||
node-version: '20'
|
||||
node-version: '22.14.0'
|
||||
|
||||
- name: 安装 pnpm
|
||||
if: ${{ github.event.inputs.build_frontend == 'true' }}
|
||||
|
||||
@@ -0,0 +1,68 @@
|
||||
name: Quality checks
|
||||
|
||||
on:
|
||||
pull_request:
|
||||
push:
|
||||
branches: [main]
|
||||
workflow_call:
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
|
||||
jobs:
|
||||
security-regressions:
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 15
|
||||
services:
|
||||
postgres:
|
||||
image: postgres:15-alpine
|
||||
env:
|
||||
POSTGRES_USER: postgres
|
||||
POSTGRES_PASSWORD: ci-test-password
|
||||
POSTGRES_DB: deepaudit_test
|
||||
ports: ['5432:5432']
|
||||
options: >-
|
||||
--health-cmd "pg_isready -U postgres" --health-interval 5s
|
||||
--health-timeout 5s --health-retries 10
|
||||
env:
|
||||
SECRET_KEY: ci-test-secret-for-deepaudit-at-least-32-chars
|
||||
ENVIRONMENT: test
|
||||
DATABASE_URL: postgresql+asyncpg://postgres:ci-test-password@localhost:5432/deepaudit_test
|
||||
TEST_DATABASE_URL: postgresql+asyncpg://postgres:ci-test-password@localhost:5432/deepaudit_test
|
||||
defaults:
|
||||
run:
|
||||
working-directory: backend
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: astral-sh/setup-uv@v5
|
||||
with:
|
||||
version: '0.11.7'
|
||||
enable-cache: true
|
||||
- run: uv sync --frozen --group dev --python 3.12
|
||||
- name: Migrate disposable database
|
||||
run: uv run --frozen --python 3.12 alembic upgrade head
|
||||
# PostgreSQL integration tests run without external LLM calls or a sandbox.
|
||||
- name: Backend tests including security regressions
|
||||
run: uv run --frozen --python 3.12 pytest tests
|
||||
|
||||
frontend:
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 15
|
||||
defaults:
|
||||
run:
|
||||
working-directory: frontend
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: pnpm/action-setup@v4
|
||||
with:
|
||||
version: 9
|
||||
- uses: actions/setup-node@v4
|
||||
with:
|
||||
node-version: '22.14.0'
|
||||
cache: pnpm
|
||||
cache-dependency-path: frontend/pnpm-lock.yaml
|
||||
- run: pnpm install --frozen-lockfile
|
||||
- run: pnpm type-check
|
||||
- run: pnpm lint
|
||||
- run: pnpm test
|
||||
- run: pnpm build
|
||||
@@ -21,7 +21,10 @@ on:
|
||||
- 'v*.*.*'
|
||||
|
||||
jobs:
|
||||
quality:
|
||||
uses: ./.github/workflows/quality.yml
|
||||
build-and-release:
|
||||
needs: quality
|
||||
name: 构建并发布
|
||||
runs-on: ubuntu-latest
|
||||
|
||||
@@ -40,7 +43,7 @@ jobs:
|
||||
- name: 设置 Node.js
|
||||
uses: actions/setup-node@v4
|
||||
with:
|
||||
node-version: '20'
|
||||
node-version: '22.14.0'
|
||||
|
||||
# 3. 安装 pnpm
|
||||
- name: 安装 pnpm
|
||||
|
||||
@@ -10,7 +10,10 @@ on:
|
||||
workflow_dispatch:
|
||||
|
||||
jobs:
|
||||
quality:
|
||||
uses: ./.github/workflows/quality.yml
|
||||
check-and-release:
|
||||
needs: quality
|
||||
name: 检查并发布
|
||||
runs-on: ubuntu-latest
|
||||
|
||||
|
||||
@@ -245,6 +245,9 @@ DeepAudit/
|
||||
|
||||
### 方式一:一行命令部署(推荐)
|
||||
|
||||
> 当前源码已收紧默认安全配置:部署前须配置至少 32 字符的随机 `SECRET_KEY`,
|
||||
> 默认不再创建演示管理员。已有部署升级前请阅读[安全配置与迁移说明](docs/CONFIGURATION.md#后端配置)。
|
||||
|
||||
使用预构建的 Docker 镜像,无需克隆代码,一行命令即可启动:
|
||||
|
||||
```bash
|
||||
@@ -321,7 +324,7 @@ git clone https://github.com/lintsinghua/DeepAudit.git && cd DeepAudit
|
||||
|
||||
# 2. 配置环境变量
|
||||
cp backend/env.example backend/.env
|
||||
# 编辑 backend/.env 填入你的 LLM API Key
|
||||
# 编辑 backend/.env 填入随机 SECRET_KEY(用 openssl rand -hex 32 生成)及 LLM API Key
|
||||
|
||||
# 3. 一键启动
|
||||
docker compose up -d
|
||||
@@ -363,6 +366,8 @@ source .venv/bin/activate
|
||||
|
||||
# 启动 API 服务
|
||||
uvicorn app.main:app --reload
|
||||
# 另开终端,同目录启动审计 worker
|
||||
python -m app.worker
|
||||
```
|
||||
|
||||
### 3. 前端启动
|
||||
|
||||
@@ -30,7 +30,7 @@ target_metadata = Base.metadata
|
||||
# can be acquired:
|
||||
# my_important_option = config.get_main_option("my_important_option")
|
||||
# ... etc.
|
||||
config.set_main_option("sqlalchemy.url", settings.DATABASE_URL)
|
||||
config.set_main_option("sqlalchemy.url", settings.DATABASE_URL.replace("%", "%%"))
|
||||
|
||||
|
||||
def run_migrations_offline() -> None:
|
||||
@@ -92,4 +92,3 @@ else:
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,45 @@
|
||||
"""Durable audit delivery and bounded list query indexes."""
|
||||
from alembic import op
|
||||
import sqlalchemy as sa
|
||||
|
||||
revision = "009_durable_audit_jobs"
|
||||
down_revision = "008_add_files_with_findings"
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
def upgrade():
|
||||
op.create_table(
|
||||
"audit_jobs",
|
||||
sa.Column("id", sa.String(36), primary_key=True),
|
||||
sa.Column("kind", sa.String(20), nullable=False),
|
||||
sa.Column("status", sa.String(20), nullable=False, server_default="queued"),
|
||||
sa.Column("payload", sa.JSON(), nullable=False, server_default="{}"),
|
||||
sa.Column("checkpoint", sa.JSON(), nullable=False, server_default="{}"),
|
||||
sa.Column("attempts", sa.Integer(), nullable=False, server_default="0"),
|
||||
sa.Column("cancel_requested", sa.Boolean(), nullable=False, server_default=sa.false()),
|
||||
sa.Column("worker_id", sa.String(100)),
|
||||
sa.Column("lease_until", sa.DateTime(timezone=True)),
|
||||
sa.Column("error", sa.Text()),
|
||||
sa.Column("created_at", sa.DateTime(timezone=True), nullable=False, server_default=sa.func.now()),
|
||||
sa.Column("updated_at", sa.DateTime(timezone=True), nullable=False, server_default=sa.func.now()),
|
||||
)
|
||||
op.create_index("ix_audit_jobs_delivery", "audit_jobs", ["status", "lease_until", "created_at"])
|
||||
op.create_index("ix_tasks_project_created", "audit_tasks", ["project_id", "created_at", "id"])
|
||||
op.create_index("ix_issues_task_created", "audit_issues", ["task_id", "created_at", "id"])
|
||||
op.create_index("ix_events_task_sequence", "agent_events", ["task_id", "sequence"])
|
||||
# Pre-queue releases may have left tasks pending/running during an upgrade.
|
||||
# Delivery reuses their persisted scan configuration and project data.
|
||||
op.execute("""INSERT INTO audit_jobs (id, kind)
|
||||
SELECT id, 'agent' FROM agent_tasks
|
||||
WHERE status NOT IN ('completed', 'failed', 'cancelled', 'paused')""")
|
||||
op.execute("""INSERT INTO audit_jobs (id, kind)
|
||||
SELECT id, CASE WHEN task_type = 'zip_upload' THEN 'zip' ELSE 'repository' END
|
||||
FROM audit_tasks WHERE status IN ('pending', 'running')""")
|
||||
|
||||
|
||||
def downgrade():
|
||||
op.drop_index("ix_events_task_sequence", table_name="agent_events")
|
||||
op.drop_index("ix_issues_task_created", table_name="audit_issues")
|
||||
op.drop_index("ix_tasks_project_created", table_name="audit_tasks")
|
||||
op.drop_table("audit_jobs")
|
||||
@@ -0,0 +1,110 @@
|
||||
"""Agent Events HTTP endpoints."""
|
||||
|
||||
import logging
|
||||
from typing import Any, List
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException, Query, Request
|
||||
from fastapi.responses import StreamingResponse
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
from sqlalchemy.future import select
|
||||
|
||||
from app.api import deps
|
||||
from app.db.session import get_db
|
||||
from app.models.agent_task import (
|
||||
AgentEvent,
|
||||
AgentTask,
|
||||
)
|
||||
from app.models.project import Project
|
||||
from app.models.user import User
|
||||
from app.services.agent_audit.events import stream_events
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
from app.services.agent_audit.schemas import AgentEventResponse
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
|
||||
async def authorize_task(task_id, db, current_user):
|
||||
task = await db.get(AgentTask, task_id)
|
||||
if not task:
|
||||
raise HTTPException(status_code=404, detail="任务不存在")
|
||||
project = await db.get(Project, task.project_id)
|
||||
if not project or project.owner_id != current_user.id:
|
||||
raise HTTPException(status_code=403, detail="无权访问此任务")
|
||||
|
||||
|
||||
@router.get("/{task_id}/events")
|
||||
@router.get("/{task_id}/stream")
|
||||
async def stream_agent_events(
|
||||
task_id: str,
|
||||
request: Request,
|
||||
after_sequence: int = Query(0, ge=0),
|
||||
include_thinking: bool = True,
|
||||
include_tool_calls: bool = True,
|
||||
db: AsyncSession = Depends(get_db),
|
||||
current_user: User = Depends(deps.get_current_user),
|
||||
):
|
||||
await authorize_task(task_id, db, current_user)
|
||||
# Release the authorization transaction before a potentially long-lived stream.
|
||||
await db.rollback()
|
||||
previous = request.headers.get("last-event-id", "")
|
||||
if previous.isdigit():
|
||||
after_sequence = max(after_sequence, int(previous))
|
||||
return StreamingResponse(
|
||||
stream_events(
|
||||
task_id,
|
||||
request,
|
||||
after_sequence,
|
||||
include_thinking=include_thinking,
|
||||
include_tool_calls=include_tool_calls,
|
||||
),
|
||||
media_type="text/event-stream",
|
||||
headers={
|
||||
"Cache-Control": "no-cache",
|
||||
"X-Accel-Buffering": "no",
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
@router.get("/{task_id}/events/list", response_model=List[AgentEventResponse])
|
||||
async def list_agent_events(
|
||||
task_id: str,
|
||||
after_sequence: int = Query(0, ge=0),
|
||||
limit: int = Query(100, ge=1, le=500),
|
||||
db: AsyncSession = Depends(get_db),
|
||||
current_user: User = Depends(deps.get_current_user),
|
||||
) -> Any:
|
||||
"""
|
||||
获取 Agent 事件列表
|
||||
"""
|
||||
task = await db.get(AgentTask, task_id)
|
||||
if not task:
|
||||
raise HTTPException(status_code=404, detail="任务不存在")
|
||||
|
||||
project = await db.get(Project, task.project_id)
|
||||
if not project or project.owner_id != current_user.id:
|
||||
raise HTTPException(status_code=403, detail="无权访问此任务")
|
||||
|
||||
result = await db.execute(
|
||||
select(AgentEvent)
|
||||
.where(AgentEvent.task_id == task_id)
|
||||
.where(AgentEvent.sequence > after_sequence)
|
||||
.order_by(AgentEvent.sequence)
|
||||
.limit(limit)
|
||||
)
|
||||
events = result.scalars().all()
|
||||
|
||||
# 🔥 Debug logging
|
||||
logger.debug(
|
||||
f"[EventsList] Task {task_id}: returning {len(events)} events (after_sequence={after_sequence})"
|
||||
)
|
||||
if events:
|
||||
logger.debug(
|
||||
f"[EventsList] First event: type={events[0].event_type}, seq={events[0].sequence}"
|
||||
)
|
||||
if len(events) > 1:
|
||||
logger.debug(
|
||||
f"[EventsList] Last event: type={events[-1].event_type}, seq={events[-1].sequence}"
|
||||
)
|
||||
|
||||
return events
|
||||
@@ -0,0 +1,242 @@
|
||||
"""Agent Inspection HTTP endpoints."""
|
||||
|
||||
import json
|
||||
import logging
|
||||
from typing import Any, List, Optional
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException, Query
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
from sqlalchemy.future import select
|
||||
|
||||
from app.api import deps
|
||||
from app.db.session import get_db
|
||||
from app.models.agent_task import (
|
||||
AgentTask,
|
||||
)
|
||||
from app.models.project import Project
|
||||
from app.models.user import User
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
from app.services.agent_audit.schemas import (
|
||||
AgentTreeNodeResponse,
|
||||
AgentTreeResponse,
|
||||
CheckpointResponse,
|
||||
)
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
|
||||
@router.get("/{task_id}/agent-tree", response_model=AgentTreeResponse)
|
||||
async def get_agent_tree(
|
||||
task_id: str,
|
||||
db: AsyncSession = Depends(get_db),
|
||||
current_user: User = Depends(deps.get_current_user),
|
||||
) -> Any:
|
||||
"""
|
||||
获取任务的 Agent 树结构
|
||||
|
||||
返回动态 Agent 树的完整结构,包括:
|
||||
- 所有 Agent 节点
|
||||
- 父子关系
|
||||
- 执行状态
|
||||
- 发现统计
|
||||
"""
|
||||
task = await db.get(AgentTask, task_id)
|
||||
if not task:
|
||||
raise HTTPException(status_code=404, detail="任务不存在")
|
||||
|
||||
project = await db.get(Project, task.project_id)
|
||||
if not project or project.owner_id != current_user.id:
|
||||
raise HTTPException(status_code=403, detail="无权访问此任务")
|
||||
|
||||
from app.models.audit_job import AuditJob
|
||||
|
||||
job = await db.get(AuditJob, task_id)
|
||||
if task.status not in {"completed", "failed", "cancelled"} and job:
|
||||
snapshot = (job.checkpoint or {}).get("live_tree")
|
||||
if snapshot:
|
||||
return AgentTreeResponse(**snapshot)
|
||||
|
||||
# 从数据库获取(已完成的任务)
|
||||
from app.models.agent_task import AgentTreeNode
|
||||
|
||||
result = await db.execute(
|
||||
select(AgentTreeNode)
|
||||
.where(AgentTreeNode.task_id == task_id)
|
||||
.order_by(AgentTreeNode.depth, AgentTreeNode.created_at)
|
||||
)
|
||||
db_nodes = result.scalars().all()
|
||||
|
||||
if not db_nodes:
|
||||
return AgentTreeResponse(
|
||||
task_id=task_id,
|
||||
nodes=[],
|
||||
)
|
||||
|
||||
# 构建响应
|
||||
nodes = []
|
||||
root_id = None
|
||||
running = 0
|
||||
completed = 0
|
||||
failed = 0
|
||||
|
||||
for node in db_nodes:
|
||||
if node.parent_agent_id is None:
|
||||
root_id = node.agent_id
|
||||
|
||||
if node.status == "running":
|
||||
running += 1
|
||||
elif node.status == "completed":
|
||||
completed += 1
|
||||
elif node.status == "failed":
|
||||
failed += 1
|
||||
|
||||
# 🔥 FIX: 对于 Orchestrator (root agent),使用 task 的 findings_count
|
||||
# 这确保了正确显示聚合的 findings 总数
|
||||
if node.parent_agent_id is None:
|
||||
# Root agent uses task's total findings
|
||||
node_findings_count = task.findings_count or 0
|
||||
else:
|
||||
node_findings_count = node.findings_count or 0
|
||||
|
||||
nodes.append(
|
||||
AgentTreeNodeResponse(
|
||||
id=node.id,
|
||||
agent_id=node.agent_id,
|
||||
agent_name=node.agent_name,
|
||||
agent_type=node.agent_type,
|
||||
parent_agent_id=node.parent_agent_id,
|
||||
depth=node.depth,
|
||||
task_description=node.task_description,
|
||||
knowledge_modules=node.knowledge_modules,
|
||||
status=node.status,
|
||||
result_summary=node.result_summary,
|
||||
findings_count=node_findings_count,
|
||||
iterations=node.iterations or 0,
|
||||
tokens_used=node.tokens_used or 0,
|
||||
tool_calls=node.tool_calls or 0,
|
||||
duration_ms=node.duration_ms,
|
||||
children=[],
|
||||
)
|
||||
)
|
||||
|
||||
# 🔥 使用 task.findings_count 作为 total_findings,确保一致性
|
||||
return AgentTreeResponse(
|
||||
task_id=task_id,
|
||||
root_agent_id=root_id,
|
||||
total_agents=len(nodes),
|
||||
running_agents=running,
|
||||
completed_agents=completed,
|
||||
failed_agents=failed,
|
||||
total_findings=task.findings_count or 0,
|
||||
nodes=nodes,
|
||||
)
|
||||
|
||||
|
||||
@router.get("/{task_id}/checkpoints", response_model=List[CheckpointResponse])
|
||||
async def list_checkpoints(
|
||||
task_id: str,
|
||||
agent_id: Optional[str] = None,
|
||||
limit: int = Query(20, ge=1, le=100),
|
||||
db: AsyncSession = Depends(get_db),
|
||||
current_user: User = Depends(deps.get_current_user),
|
||||
) -> Any:
|
||||
"""
|
||||
获取任务的检查点列表
|
||||
|
||||
用于:
|
||||
- 查看执行历史
|
||||
- 状态恢复
|
||||
- 调试分析
|
||||
"""
|
||||
task = await db.get(AgentTask, task_id)
|
||||
if not task:
|
||||
raise HTTPException(status_code=404, detail="任务不存在")
|
||||
|
||||
project = await db.get(Project, task.project_id)
|
||||
if not project or project.owner_id != current_user.id:
|
||||
raise HTTPException(status_code=403, detail="无权访问此任务")
|
||||
|
||||
from app.models.agent_task import AgentCheckpoint
|
||||
|
||||
query = select(AgentCheckpoint).where(AgentCheckpoint.task_id == task_id)
|
||||
|
||||
if agent_id:
|
||||
query = query.where(AgentCheckpoint.agent_id == agent_id)
|
||||
|
||||
query = query.order_by(AgentCheckpoint.created_at.desc()).limit(limit)
|
||||
|
||||
result = await db.execute(query)
|
||||
checkpoints = result.scalars().all()
|
||||
|
||||
return [
|
||||
CheckpointResponse(
|
||||
id=cp.id,
|
||||
agent_id=cp.agent_id,
|
||||
agent_name=cp.agent_name,
|
||||
agent_type=cp.agent_type,
|
||||
iteration=cp.iteration,
|
||||
status=cp.status,
|
||||
total_tokens=cp.total_tokens or 0,
|
||||
tool_calls=cp.tool_calls or 0,
|
||||
findings_count=cp.findings_count or 0,
|
||||
checkpoint_type=cp.checkpoint_type or "auto",
|
||||
checkpoint_name=cp.checkpoint_name,
|
||||
created_at=cp.created_at.isoformat() if cp.created_at else None,
|
||||
)
|
||||
for cp in checkpoints
|
||||
]
|
||||
|
||||
|
||||
@router.get("/{task_id}/checkpoints/{checkpoint_id}")
|
||||
async def get_checkpoint_detail(
|
||||
task_id: str,
|
||||
checkpoint_id: str,
|
||||
db: AsyncSession = Depends(get_db),
|
||||
current_user: User = Depends(deps.get_current_user),
|
||||
) -> Any:
|
||||
"""
|
||||
获取检查点详情
|
||||
|
||||
返回完整的 Agent 状态数据
|
||||
"""
|
||||
task = await db.get(AgentTask, task_id)
|
||||
if not task:
|
||||
raise HTTPException(status_code=404, detail="任务不存在")
|
||||
|
||||
project = await db.get(Project, task.project_id)
|
||||
if not project or project.owner_id != current_user.id:
|
||||
raise HTTPException(status_code=403, detail="无权访问此任务")
|
||||
|
||||
from app.models.agent_task import AgentCheckpoint
|
||||
|
||||
checkpoint = await db.get(AgentCheckpoint, checkpoint_id)
|
||||
if not checkpoint or checkpoint.task_id != task_id:
|
||||
raise HTTPException(status_code=404, detail="检查点不存在")
|
||||
|
||||
# 解析状态数据
|
||||
state_data = {}
|
||||
if checkpoint.state_data:
|
||||
try:
|
||||
state_data = json.loads(checkpoint.state_data)
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
|
||||
return {
|
||||
"id": checkpoint.id,
|
||||
"task_id": checkpoint.task_id,
|
||||
"agent_id": checkpoint.agent_id,
|
||||
"agent_name": checkpoint.agent_name,
|
||||
"agent_type": checkpoint.agent_type,
|
||||
"parent_agent_id": checkpoint.parent_agent_id,
|
||||
"iteration": checkpoint.iteration,
|
||||
"status": checkpoint.status,
|
||||
"total_tokens": checkpoint.total_tokens,
|
||||
"tool_calls": checkpoint.tool_calls,
|
||||
"findings_count": checkpoint.findings_count,
|
||||
"checkpoint_type": checkpoint.checkpoint_type,
|
||||
"checkpoint_name": checkpoint.checkpoint_name,
|
||||
"state_data": state_data,
|
||||
"metadata": checkpoint.checkpoint_metadata,
|
||||
"created_at": checkpoint.created_at.isoformat() if checkpoint.created_at else None,
|
||||
}
|
||||
@@ -0,0 +1,457 @@
|
||||
"""Agent Reports HTTP endpoints."""
|
||||
|
||||
import logging
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException, Query
|
||||
from sqlalchemy import case
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
from sqlalchemy.future import select
|
||||
|
||||
from app.api import deps
|
||||
from app.db.session import get_db
|
||||
from app.models.agent_task import (
|
||||
AgentFinding,
|
||||
AgentTask,
|
||||
)
|
||||
from app.models.project import Project
|
||||
from app.models.user import User
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
|
||||
@router.get("/{task_id}/report")
|
||||
async def generate_audit_report(
|
||||
task_id: str,
|
||||
format: str = Query("markdown", regex="^(markdown|json)$"),
|
||||
db: AsyncSession = Depends(get_db),
|
||||
current_user: User = Depends(deps.get_current_user),
|
||||
):
|
||||
"""
|
||||
生成审计报告
|
||||
|
||||
支持 Markdown 和 JSON 格式
|
||||
"""
|
||||
task = await db.get(AgentTask, task_id)
|
||||
if not task:
|
||||
raise HTTPException(status_code=404, detail="任务不存在")
|
||||
|
||||
project = await db.get(Project, task.project_id)
|
||||
if not project or project.owner_id != current_user.id:
|
||||
raise HTTPException(status_code=403, detail="无权访问此任务")
|
||||
|
||||
# 获取此任务的所有发现
|
||||
findings = await db.execute(
|
||||
select(AgentFinding)
|
||||
.where(AgentFinding.task_id == task_id)
|
||||
.order_by(
|
||||
case(
|
||||
(AgentFinding.severity == "critical", 1),
|
||||
(AgentFinding.severity == "high", 2),
|
||||
(AgentFinding.severity == "medium", 3),
|
||||
(AgentFinding.severity == "low", 4),
|
||||
else_=5,
|
||||
),
|
||||
AgentFinding.created_at.desc(),
|
||||
)
|
||||
)
|
||||
findings = findings.scalars().all()
|
||||
|
||||
# 🔥 Helper function to normalize severity for comparison (case-insensitive)
|
||||
def normalize_severity(sev: str) -> str:
|
||||
return str(sev).lower().strip() if sev else ""
|
||||
|
||||
# Log findings for debugging
|
||||
logger.info(f"[Report] Task {task_id}: Found {len(findings)} findings from database")
|
||||
if findings:
|
||||
for i, f in enumerate(findings[:3]): # Log first 3
|
||||
logger.debug(
|
||||
f"[Report] Finding {i + 1}: severity='{f.severity}', title='{f.title[:50] if f.title else 'N/A'}'"
|
||||
)
|
||||
|
||||
if format == "json":
|
||||
# Enhanced JSON report with full metadata
|
||||
return {
|
||||
"report_metadata": {
|
||||
"task_id": task.id,
|
||||
"project_id": task.project_id,
|
||||
"project_name": project.name,
|
||||
"generated_at": datetime.now(timezone.utc).isoformat(),
|
||||
"task_status": task.status,
|
||||
"duration_seconds": int((task.completed_at - task.started_at).total_seconds())
|
||||
if task.completed_at and task.started_at
|
||||
else None,
|
||||
},
|
||||
"summary": {
|
||||
"security_score": task.security_score,
|
||||
"total_files_analyzed": task.analyzed_files,
|
||||
"total_findings": len(findings),
|
||||
"verified_findings": sum(1 for f in findings if f.is_verified),
|
||||
"severity_distribution": {
|
||||
"critical": sum(
|
||||
1 for f in findings if normalize_severity(f.severity) == "critical"
|
||||
),
|
||||
"high": sum(1 for f in findings if normalize_severity(f.severity) == "high"),
|
||||
"medium": sum(
|
||||
1 for f in findings if normalize_severity(f.severity) == "medium"
|
||||
),
|
||||
"low": sum(1 for f in findings if normalize_severity(f.severity) == "low"),
|
||||
},
|
||||
"agent_metrics": {
|
||||
"total_iterations": task.total_iterations,
|
||||
"tool_calls": task.tool_calls_count,
|
||||
"tokens_used": task.tokens_used,
|
||||
},
|
||||
},
|
||||
"findings": [
|
||||
{
|
||||
"id": f.id,
|
||||
"title": f.title,
|
||||
"severity": f.severity,
|
||||
"vulnerability_type": f.vulnerability_type,
|
||||
"description": f.description,
|
||||
"file_path": f.file_path,
|
||||
"line_start": f.line_start,
|
||||
"line_end": f.line_end,
|
||||
"code_snippet": f.code_snippet,
|
||||
"is_verified": f.is_verified,
|
||||
"has_poc": f.has_poc,
|
||||
"poc_code": f.poc_code,
|
||||
"poc_description": f.poc_description,
|
||||
"poc_steps": f.poc_steps,
|
||||
"confidence": f.ai_confidence,
|
||||
"suggestion": f.suggestion,
|
||||
"fix_code": f.fix_code,
|
||||
"created_at": f.created_at.isoformat() if f.created_at else None,
|
||||
}
|
||||
for f in findings
|
||||
],
|
||||
}
|
||||
|
||||
# Generate Enhanced Markdown Report
|
||||
timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
|
||||
|
||||
# Calculate statistics
|
||||
total = len(findings)
|
||||
critical = sum(1 for f in findings if normalize_severity(f.severity) == "critical")
|
||||
high = sum(1 for f in findings if normalize_severity(f.severity) == "high")
|
||||
medium = sum(1 for f in findings if normalize_severity(f.severity) == "medium")
|
||||
low = sum(1 for f in findings if normalize_severity(f.severity) == "low")
|
||||
verified = sum(1 for f in findings if f.is_verified)
|
||||
with_poc = sum(1 for f in findings if f.has_poc)
|
||||
|
||||
# Calculate duration
|
||||
duration_str = "N/A"
|
||||
if task.completed_at and task.started_at:
|
||||
duration = (task.completed_at - task.started_at).total_seconds()
|
||||
if duration >= 3600:
|
||||
duration_str = f"{duration / 3600:.1f} 小时"
|
||||
elif duration >= 60:
|
||||
duration_str = f"{duration / 60:.1f} 分钟"
|
||||
else:
|
||||
duration_str = f"{int(duration)} 秒"
|
||||
|
||||
md_lines = []
|
||||
|
||||
# Header
|
||||
md_lines.append("# DeepAudit 安全审计报告")
|
||||
md_lines.append("")
|
||||
md_lines.append("---")
|
||||
md_lines.append("")
|
||||
|
||||
# Report Info
|
||||
md_lines.append("## 报告信息")
|
||||
md_lines.append("")
|
||||
md_lines.append(f"| 属性 | 内容 |")
|
||||
md_lines.append(f"|----------|-------|")
|
||||
md_lines.append(f"| **项目名称** | {project.name} |")
|
||||
md_lines.append(f"| **任务 ID** | `{task.id[:8]}...` |")
|
||||
md_lines.append(f"| **生成时间** | {timestamp} |")
|
||||
md_lines.append(f"| **任务状态** | {task.status.upper()} |")
|
||||
md_lines.append(f"| **耗时** | {duration_str} |")
|
||||
md_lines.append("")
|
||||
|
||||
# Executive Summary
|
||||
md_lines.append("## 执行摘要")
|
||||
md_lines.append("")
|
||||
|
||||
score = task.security_score
|
||||
if score is not None:
|
||||
if score >= 80:
|
||||
score_assessment = "良好 - 建议进行少量优化"
|
||||
score_icon = "通过"
|
||||
elif score >= 60:
|
||||
score_assessment = "中等 - 存在若干问题需要关注"
|
||||
score_icon = "警告"
|
||||
else:
|
||||
score_assessment = "严重 - 需要立即进行修复"
|
||||
score_icon = "未通过"
|
||||
md_lines.append(f"**安全评分: {int(score)}/100** [{score_icon}]")
|
||||
md_lines.append(f"*{score_assessment}*")
|
||||
else:
|
||||
md_lines.append("**安全评分:** 未计算")
|
||||
md_lines.append("")
|
||||
|
||||
# Findings Summary
|
||||
md_lines.append("### 漏洞发现概览")
|
||||
md_lines.append("")
|
||||
md_lines.append(f"| 严重程度 | 数量 | 已验证 |")
|
||||
md_lines.append(f"|----------|-------|----------|")
|
||||
if critical > 0:
|
||||
md_lines.append(
|
||||
f"| **严重 (CRITICAL)** | {critical} | {sum(1 for f in findings if normalize_severity(f.severity) == 'critical' and f.is_verified)} |"
|
||||
)
|
||||
if high > 0:
|
||||
md_lines.append(
|
||||
f"| **高危 (HIGH)** | {high} | {sum(1 for f in findings if normalize_severity(f.severity) == 'high' and f.is_verified)} |"
|
||||
)
|
||||
if medium > 0:
|
||||
md_lines.append(
|
||||
f"| **中危 (MEDIUM)** | {medium} | {sum(1 for f in findings if normalize_severity(f.severity) == 'medium' and f.is_verified)} |"
|
||||
)
|
||||
if low > 0:
|
||||
md_lines.append(
|
||||
f"| **低危 (LOW)** | {low} | {sum(1 for f in findings if normalize_severity(f.severity) == 'low' and f.is_verified)} |"
|
||||
)
|
||||
md_lines.append(f"| **总计** | {total} | {verified} |")
|
||||
md_lines.append("")
|
||||
|
||||
# Audit Metrics
|
||||
md_lines.append("### 审计指标")
|
||||
md_lines.append("")
|
||||
md_lines.append(f"- **分析文件数:** {task.analyzed_files} / {task.total_files}")
|
||||
md_lines.append(f"- **Agent 迭代次数:** {task.total_iterations}")
|
||||
md_lines.append(f"- **工具调用次数:** {task.tool_calls_count}")
|
||||
md_lines.append(f"- **Token 消耗:** {task.tokens_used:,}")
|
||||
if with_poc > 0:
|
||||
md_lines.append(f"- **生成的 PoC:** {with_poc}")
|
||||
md_lines.append("")
|
||||
|
||||
# Detailed Findings
|
||||
if not findings:
|
||||
md_lines.append("## 漏洞详情")
|
||||
md_lines.append("")
|
||||
md_lines.append("*本次审计未发现安全漏洞。*")
|
||||
md_lines.append("")
|
||||
else:
|
||||
# Group findings by severity
|
||||
severity_map = {
|
||||
"critical": "严重 (Critical)",
|
||||
"high": "高危 (High)",
|
||||
"medium": "中危 (Medium)",
|
||||
"low": "低危 (Low)",
|
||||
}
|
||||
|
||||
for severity_level, severity_name in severity_map.items():
|
||||
severity_findings = [
|
||||
f for f in findings if normalize_severity(f.severity) == severity_level
|
||||
]
|
||||
if not severity_findings:
|
||||
continue
|
||||
|
||||
md_lines.append(f"## {severity_name} 漏洞")
|
||||
md_lines.append("")
|
||||
|
||||
for i, f in enumerate(severity_findings, 1):
|
||||
verified_badge = "[已验证]" if f.is_verified else "[未验证]"
|
||||
poc_badge = " [含 PoC]" if f.has_poc else ""
|
||||
|
||||
md_lines.append(f"### {severity_level.upper()}-{i}: {f.title}")
|
||||
md_lines.append("")
|
||||
md_lines.append(f"**{verified_badge}**{poc_badge} | 类型: `{f.vulnerability_type}`")
|
||||
md_lines.append("")
|
||||
|
||||
if f.file_path:
|
||||
location = f"`{f.file_path}"
|
||||
if f.line_start:
|
||||
location += f":{f.line_start}"
|
||||
if f.line_end and f.line_end != f.line_start:
|
||||
location += f"-{f.line_end}"
|
||||
location += "`"
|
||||
md_lines.append(f"**位置:** {location}")
|
||||
md_lines.append("")
|
||||
|
||||
if f.ai_confidence:
|
||||
md_lines.append(f"**AI 置信度:** {int(f.ai_confidence * 100)}%")
|
||||
md_lines.append("")
|
||||
|
||||
if f.description:
|
||||
md_lines.append("**漏洞描述:**")
|
||||
md_lines.append("")
|
||||
md_lines.append(f.description)
|
||||
md_lines.append("")
|
||||
|
||||
if f.code_snippet:
|
||||
# 🔥 v2.1: 增强语言检测,避免默认 python 标记错误
|
||||
lang = "text" # 默认使用 text 而非 python
|
||||
if f.file_path:
|
||||
ext = f.file_path.split(".")[-1].lower()
|
||||
lang_map = {
|
||||
# Python
|
||||
"py": "python",
|
||||
"pyw": "python",
|
||||
"pyi": "python",
|
||||
# JavaScript/TypeScript
|
||||
"js": "javascript",
|
||||
"mjs": "javascript",
|
||||
"cjs": "javascript",
|
||||
"ts": "typescript",
|
||||
"mts": "typescript",
|
||||
"jsx": "jsx",
|
||||
"tsx": "tsx",
|
||||
# Web
|
||||
"html": "html",
|
||||
"htm": "html",
|
||||
"css": "css",
|
||||
"scss": "scss",
|
||||
"sass": "sass",
|
||||
"less": "less",
|
||||
"vue": "vue",
|
||||
"svelte": "svelte",
|
||||
# Backend
|
||||
"java": "java",
|
||||
"kt": "kotlin",
|
||||
"kts": "kotlin",
|
||||
"go": "go",
|
||||
"rs": "rust",
|
||||
"rb": "ruby",
|
||||
"erb": "erb",
|
||||
"php": "php",
|
||||
"phtml": "php",
|
||||
# C-family
|
||||
"c": "c",
|
||||
"h": "c",
|
||||
"cpp": "cpp",
|
||||
"cc": "cpp",
|
||||
"cxx": "cpp",
|
||||
"hpp": "cpp",
|
||||
"cs": "csharp",
|
||||
# Shell/Script
|
||||
"sh": "bash",
|
||||
"bash": "bash",
|
||||
"zsh": "zsh",
|
||||
"ps1": "powershell",
|
||||
"psm1": "powershell",
|
||||
# Config
|
||||
"json": "json",
|
||||
"yaml": "yaml",
|
||||
"yml": "yaml",
|
||||
"toml": "toml",
|
||||
"ini": "ini",
|
||||
"cfg": "ini",
|
||||
"xml": "xml",
|
||||
"xhtml": "xml",
|
||||
# Database
|
||||
"sql": "sql",
|
||||
# Other
|
||||
"md": "markdown",
|
||||
"markdown": "markdown",
|
||||
"sol": "solidity",
|
||||
"swift": "swift",
|
||||
"r": "r",
|
||||
"R": "r",
|
||||
"lua": "lua",
|
||||
"pl": "perl",
|
||||
"pm": "perl",
|
||||
"ex": "elixir",
|
||||
"exs": "elixir",
|
||||
"erl": "erlang",
|
||||
"hs": "haskell",
|
||||
"scala": "scala",
|
||||
"sc": "scala",
|
||||
"clj": "clojure",
|
||||
"cljs": "clojure",
|
||||
"dart": "dart",
|
||||
"groovy": "groovy",
|
||||
"gradle": "groovy",
|
||||
}
|
||||
lang = lang_map.get(ext, "text")
|
||||
md_lines.append("**漏洞代码:**")
|
||||
md_lines.append("")
|
||||
md_lines.append(f"```{lang}")
|
||||
md_lines.append(f.code_snippet.strip())
|
||||
md_lines.append("```")
|
||||
md_lines.append("")
|
||||
|
||||
if f.suggestion:
|
||||
md_lines.append("**修复建议:**")
|
||||
md_lines.append("")
|
||||
md_lines.append(f.suggestion)
|
||||
md_lines.append("")
|
||||
|
||||
if f.fix_code:
|
||||
md_lines.append("**参考修复代码:**")
|
||||
md_lines.append("")
|
||||
md_lines.append(f"```{lang if f.file_path else 'text'}")
|
||||
md_lines.append(f.fix_code.strip())
|
||||
md_lines.append("```")
|
||||
md_lines.append("")
|
||||
|
||||
# 🔥 添加 PoC 详情
|
||||
if f.has_poc:
|
||||
md_lines.append("**概念验证 (PoC):**")
|
||||
md_lines.append("")
|
||||
|
||||
if f.poc_description:
|
||||
md_lines.append(f"*{f.poc_description}*")
|
||||
md_lines.append("")
|
||||
|
||||
if f.poc_steps:
|
||||
md_lines.append("**复现步骤:**")
|
||||
md_lines.append("")
|
||||
for step_idx, step in enumerate(f.poc_steps, 1):
|
||||
md_lines.append(f"{step_idx}. {step}")
|
||||
md_lines.append("")
|
||||
|
||||
if f.poc_code:
|
||||
md_lines.append("**PoC 代码:**")
|
||||
md_lines.append("")
|
||||
md_lines.append("```")
|
||||
md_lines.append(f.poc_code.strip())
|
||||
md_lines.append("```")
|
||||
md_lines.append("")
|
||||
|
||||
md_lines.append("---")
|
||||
md_lines.append("")
|
||||
|
||||
# Remediation Priority
|
||||
if critical > 0 or high > 0:
|
||||
md_lines.append("## 修复优先级建议")
|
||||
md_lines.append("")
|
||||
md_lines.append("基于已发现的漏洞,我们建议按以下优先级进行修复:")
|
||||
md_lines.append("")
|
||||
priority_idx = 1
|
||||
if critical > 0:
|
||||
md_lines.append(
|
||||
f"{priority_idx}. **立即修复:** 处理 {critical} 个严重漏洞 - 可能造成严重影响"
|
||||
)
|
||||
priority_idx += 1
|
||||
if high > 0:
|
||||
md_lines.append(f"{priority_idx}. **高优先级:** 在 1 周内修复 {high} 个高危漏洞")
|
||||
priority_idx += 1
|
||||
if medium > 0:
|
||||
md_lines.append(f"{priority_idx}. **中优先级:** 在 2-4 周内修复 {medium} 个中危漏洞")
|
||||
priority_idx += 1
|
||||
if low > 0:
|
||||
md_lines.append(f"{priority_idx}. **低优先级:** 在日常维护中处理 {low} 个低危漏洞")
|
||||
priority_idx += 1
|
||||
md_lines.append("")
|
||||
|
||||
# Footer
|
||||
md_lines.append("---")
|
||||
md_lines.append("")
|
||||
md_lines.append("*本报告由 DeepAudit - AI 驱动的安全分析系统生成*")
|
||||
md_lines.append("")
|
||||
content = "\n".join(md_lines)
|
||||
|
||||
filename = f"audit_report_{task.id[:8]}_{datetime.now().strftime('%Y%m%d')}.md"
|
||||
|
||||
from fastapi.responses import Response
|
||||
|
||||
return Response(
|
||||
content=content,
|
||||
media_type="text/markdown",
|
||||
headers={"Content-Disposition": f"attachment; filename={filename}"},
|
||||
)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -24,7 +24,7 @@ SENSITIVE_LLM_FIELDS = [
|
||||
'qwenApiKey', 'deepseekApiKey', 'zhipuApiKey', 'moonshotApiKey',
|
||||
'baiduApiKey', 'minimaxApiKey', 'doubaoApiKey'
|
||||
]
|
||||
SENSITIVE_OTHER_FIELDS = ['githubToken', 'gitlabToken']
|
||||
SENSITIVE_OTHER_FIELDS = ['githubToken', 'gitlabToken', 'giteaToken']
|
||||
|
||||
|
||||
def encrypt_config(config: dict, sensitive_fields: list) -> dict:
|
||||
|
||||
@@ -1,3 +1,5 @@
|
||||
from app.services.audit_queue import enqueue
|
||||
from app.services.archive import save_upload
|
||||
from typing import Any, List, Optional
|
||||
from fastapi import APIRouter, Depends, HTTPException, BackgroundTasks, UploadFile, File
|
||||
from fastapi.responses import FileResponse
|
||||
@@ -82,6 +84,7 @@ class StatsResponse(BaseModel):
|
||||
total_issues: int
|
||||
resolved_issues: int
|
||||
avg_quality_score: float = 0.0
|
||||
issue_types: dict[str, int] = {}
|
||||
|
||||
@router.post("/", response_model=ProjectResponse)
|
||||
async def create_project(
|
||||
@@ -157,67 +160,8 @@ async def get_stats(
|
||||
"""
|
||||
Get statistics for current user.
|
||||
"""
|
||||
# 只统计当前用户的项目
|
||||
projects_result = await db.execute(
|
||||
select(Project).where(Project.owner_id == current_user.id)
|
||||
)
|
||||
projects = projects_result.scalars().all()
|
||||
project_ids = [p.id for p in projects]
|
||||
|
||||
# 统计旧的 AuditTask
|
||||
tasks_result = await db.execute(
|
||||
select(AuditTask).where(AuditTask.project_id.in_(project_ids)) if project_ids else select(AuditTask).where(False)
|
||||
)
|
||||
tasks = tasks_result.scalars().all()
|
||||
task_ids = [t.id for t in tasks]
|
||||
|
||||
# 统计旧的 AuditIssue
|
||||
issues_result = await db.execute(
|
||||
select(AuditIssue).where(AuditIssue.task_id.in_(task_ids)) if task_ids else select(AuditIssue).where(False)
|
||||
)
|
||||
issues = issues_result.scalars().all()
|
||||
|
||||
# 🔥 同时统计新的 AgentTask
|
||||
agent_tasks_result = await db.execute(
|
||||
select(AgentTask).where(AgentTask.project_id.in_(project_ids)) if project_ids else select(AgentTask).where(False)
|
||||
)
|
||||
agent_tasks = agent_tasks_result.scalars().all()
|
||||
agent_task_ids = [t.id for t in agent_tasks]
|
||||
|
||||
# 🔥 统计 AgentFinding
|
||||
agent_findings_result = await db.execute(
|
||||
select(AgentFinding).where(AgentFinding.task_id.in_(agent_task_ids)) if agent_task_ids else select(AgentFinding).where(False)
|
||||
)
|
||||
agent_findings = agent_findings_result.scalars().all()
|
||||
|
||||
# 合并统计(旧任务 + 新 Agent 任务)
|
||||
total_tasks = len(tasks) + len(agent_tasks)
|
||||
completed_tasks = (
|
||||
len([t for t in tasks if t.status == "completed"]) +
|
||||
len([t for t in agent_tasks if t.status == AgentTaskStatus.COMPLETED])
|
||||
)
|
||||
total_issues = len(issues) + len(agent_findings)
|
||||
resolved_issues = (
|
||||
len([i for i in issues if i.status == "resolved"]) +
|
||||
len([f for f in agent_findings if f.status in ("fixed", "wont_fix", "false_positive")])
|
||||
)
|
||||
|
||||
# 计算平均质量分(只统计已完成且有质量分的任务)
|
||||
quality_scores = (
|
||||
[t.quality_score for t in tasks if t.status == "completed" and t.quality_score and t.quality_score > 0] +
|
||||
[t.quality_score for t in agent_tasks if t.status == AgentTaskStatus.COMPLETED and t.quality_score and t.quality_score > 0]
|
||||
)
|
||||
avg_quality_score = sum(quality_scores) / len(quality_scores) if quality_scores else 0.0
|
||||
|
||||
return {
|
||||
"total_projects": len(projects),
|
||||
"active_projects": len([p for p in projects if p.is_active]),
|
||||
"total_tasks": total_tasks,
|
||||
"completed_tasks": completed_tasks,
|
||||
"total_issues": total_issues,
|
||||
"resolved_issues": resolved_issues,
|
||||
"avg_quality_score": avg_quality_score,
|
||||
}
|
||||
from app.services.project_stats import project_stats
|
||||
return await project_stats(db, current_user.id)
|
||||
|
||||
@router.get("/{id}", response_model=ProjectResponse)
|
||||
async def read_project(
|
||||
@@ -508,6 +452,9 @@ async def scan_project(
|
||||
if not project:
|
||||
raise HTTPException(status_code=404, detail="项目不存在")
|
||||
|
||||
if project.owner_id != current_user.id:
|
||||
raise HTTPException(status_code=403, detail="无权操作此项目")
|
||||
|
||||
# 获取分支和排除模式
|
||||
branch_name = scan_request.branch_name if scan_request else None
|
||||
exclude_patterns = scan_request.exclude_patterns if scan_request else None
|
||||
@@ -523,51 +470,10 @@ async def scan_project(
|
||||
scan_config=json.dumps(scan_request.dict()) if scan_request else "{}"
|
||||
)
|
||||
db.add(task)
|
||||
await enqueue(db, task, "repository")
|
||||
await db.commit()
|
||||
await db.refresh(task)
|
||||
|
||||
# 获取用户配置(包含解密敏感字段)
|
||||
from app.core.encryption import decrypt_sensitive_data
|
||||
|
||||
# 需要解密的敏感字段列表
|
||||
SENSITIVE_LLM_FIELDS = [
|
||||
'llmApiKey', 'geminiApiKey', 'openaiApiKey', 'claudeApiKey',
|
||||
'qwenApiKey', 'deepseekApiKey', 'zhipuApiKey', 'moonshotApiKey',
|
||||
'baiduApiKey', 'minimaxApiKey', 'doubaoApiKey'
|
||||
]
|
||||
SENSITIVE_OTHER_FIELDS = ['githubToken', 'gitlabToken']
|
||||
|
||||
def decrypt_config(config_dict: dict, sensitive_fields: list) -> dict:
|
||||
"""解密配置中的敏感字段"""
|
||||
decrypted = config_dict.copy()
|
||||
for field in sensitive_fields:
|
||||
if field in decrypted and decrypted[field]:
|
||||
decrypted[field] = decrypt_sensitive_data(decrypted[field])
|
||||
return decrypted
|
||||
|
||||
result = await db.execute(
|
||||
select(UserConfig).where(UserConfig.user_id == current_user.id)
|
||||
)
|
||||
config = result.scalar_one_or_none()
|
||||
user_config = {}
|
||||
if config:
|
||||
llm_config = json.loads(config.llm_config) if config.llm_config else {}
|
||||
other_config = json.loads(config.other_config) if config.other_config else {}
|
||||
# 解密敏感字段
|
||||
llm_config = decrypt_config(llm_config, SENSITIVE_LLM_FIELDS)
|
||||
other_config = decrypt_config(other_config, SENSITIVE_OTHER_FIELDS)
|
||||
user_config = {
|
||||
'llmConfig': llm_config,
|
||||
'otherConfig': other_config,
|
||||
}
|
||||
|
||||
# 将扫描配置注入到 user_config 中,以便 scan_repo_task 使用
|
||||
if scan_request and scan_request.file_paths:
|
||||
user_config['scan_config'] = {'file_paths': scan_request.file_paths}
|
||||
|
||||
# Trigger Background Task
|
||||
background_tasks.add_task(scan_repo_task, task.id, AsyncSessionLocal, user_config)
|
||||
|
||||
return {"task_id": task.id, "status": "started"}
|
||||
|
||||
|
||||
@@ -642,14 +548,8 @@ async def upload_project_zip(
|
||||
temp_file_path = f"/tmp/{temp_file_id}.zip"
|
||||
|
||||
try:
|
||||
with open(temp_file_path, "wb") as buffer:
|
||||
shutil.copyfileobj(file.file, buffer)
|
||||
|
||||
# 检查文件大小
|
||||
file_size = os.path.getsize(temp_file_path)
|
||||
if file_size > 500 * 1024 * 1024: # 500MB limit
|
||||
raise HTTPException(status_code=400, detail="文件大小不能超过500MB")
|
||||
|
||||
await save_upload(file, temp_file_path)
|
||||
|
||||
# 保存到持久化存储
|
||||
meta = await save_project_zip(id, temp_file_path, file.filename)
|
||||
|
||||
|
||||
@@ -1,3 +1,8 @@
|
||||
from app.services.user_configuration import get_user_config_dict
|
||||
from pathlib import Path
|
||||
from app.services.audit_queue import enqueue
|
||||
from app.services.zip_scan import process_zip_task, normalize_path
|
||||
from app.services.archive import extract_zip, save_upload
|
||||
from fastapi import APIRouter, UploadFile, File, Form, Depends, BackgroundTasks, HTTPException
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
from sqlalchemy.future import select
|
||||
@@ -20,212 +25,47 @@ from app.models.project import Project
|
||||
from app.models.analysis import InstantAnalysis
|
||||
from app.models.user_config import UserConfig
|
||||
from app.services.llm.service import LLMService
|
||||
from app.services.scanner import task_control, is_text_file, should_exclude, get_language_from_path, get_analysis_config
|
||||
from app.services.scanner import (
|
||||
task_control,
|
||||
is_text_file,
|
||||
should_exclude,
|
||||
get_language_from_path,
|
||||
get_analysis_config,
|
||||
)
|
||||
from app.services.zip_storage import load_project_zip, save_project_zip, has_project_zip
|
||||
from app.core.config import settings
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
|
||||
def normalize_path(path: str) -> str:
|
||||
"""
|
||||
统一路径分隔符为正斜杠,确保跨平台兼容性
|
||||
Windows 使用反斜杠 (\),Unix/Mac 使用正斜杠 (/)
|
||||
统一转换为正斜杠以保证一致性
|
||||
"""
|
||||
return path.replace("\\", "/")
|
||||
|
||||
|
||||
# 支持的文件扩展名
|
||||
TEXT_EXTENSIONS = [
|
||||
".js", ".ts", ".tsx", ".jsx", ".py", ".java", ".go", ".rs",
|
||||
".cpp", ".c", ".h", ".cc", ".hh", ".cs", ".php", ".rb",
|
||||
".kt", ".swift", ".sql", ".sh", ".json", ".yml", ".yaml"
|
||||
".js",
|
||||
".ts",
|
||||
".tsx",
|
||||
".jsx",
|
||||
".py",
|
||||
".java",
|
||||
".go",
|
||||
".rs",
|
||||
".cpp",
|
||||
".c",
|
||||
".h",
|
||||
".cc",
|
||||
".hh",
|
||||
".cs",
|
||||
".php",
|
||||
".rb",
|
||||
".kt",
|
||||
".swift",
|
||||
".sql",
|
||||
".sh",
|
||||
".json",
|
||||
".yml",
|
||||
".yaml",
|
||||
]
|
||||
|
||||
|
||||
async def process_zip_task(task_id: str, file_path: str, db_session_factory, user_config: dict = None):
|
||||
"""后台ZIP文件处理任务"""
|
||||
async with db_session_factory() as db:
|
||||
task = await db.get(AuditTask, task_id)
|
||||
if not task:
|
||||
return
|
||||
|
||||
try:
|
||||
task.status = "running"
|
||||
task.started_at = datetime.now(timezone.utc)
|
||||
await db.commit()
|
||||
|
||||
# 创建使用用户配置的LLM服务实例
|
||||
llm_service = LLMService(user_config=user_config or {})
|
||||
|
||||
# Extract ZIP
|
||||
extract_dir = Path(f"/tmp/{task_id}")
|
||||
extract_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
with zipfile.ZipFile(file_path, 'r') as zip_ref:
|
||||
zip_ref.extractall(extract_dir)
|
||||
|
||||
# 获取用户自定义排除模式
|
||||
scan_config = (user_config or {}).get('scan_config', {})
|
||||
custom_exclude_patterns = scan_config.get('exclude_patterns', [])
|
||||
|
||||
# Find files
|
||||
files_to_scan = []
|
||||
for root, dirs, files in os.walk(extract_dir):
|
||||
# 排除常见非代码目录
|
||||
dirs[:] = [d for d in dirs if d not in ['node_modules', '__pycache__', '.git', 'dist', 'build', 'vendor']]
|
||||
|
||||
for file in files:
|
||||
full_path = Path(root) / file
|
||||
# 统一使用正斜杠,确保跨平台兼容性
|
||||
rel_path = normalize_path(str(full_path.relative_to(extract_dir)))
|
||||
|
||||
# 检查文件类型和排除规则(包含用户自定义排除模式)
|
||||
if is_text_file(rel_path) and not should_exclude(rel_path, custom_exclude_patterns):
|
||||
try:
|
||||
content = full_path.read_text(errors='ignore')
|
||||
if len(content) <= settings.MAX_FILE_SIZE_BYTES:
|
||||
files_to_scan.append({
|
||||
"path": rel_path,
|
||||
"content": content
|
||||
})
|
||||
except:
|
||||
pass
|
||||
|
||||
# 获取分析配置(优先使用用户配置)
|
||||
analysis_config = get_analysis_config(user_config)
|
||||
max_analyze_files = analysis_config['max_analyze_files']
|
||||
llm_gap_ms = analysis_config['llm_gap_ms']
|
||||
|
||||
# 限制文件数量
|
||||
# 如果指定了特定文件,则只分析这些文件
|
||||
target_files = scan_config.get('file_paths', [])
|
||||
if target_files:
|
||||
# 统一目标文件路径的分隔符,确保匹配一致性
|
||||
normalized_targets = {normalize_path(p) for p in target_files}
|
||||
print(f"🎯 ZIP任务: 指定分析 {len(normalized_targets)} 个文件")
|
||||
files_to_scan = [f for f in files_to_scan if f['path'] in normalized_targets]
|
||||
elif max_analyze_files > 0:
|
||||
files_to_scan = files_to_scan[:max_analyze_files]
|
||||
|
||||
task.total_files = len(files_to_scan)
|
||||
await db.commit()
|
||||
|
||||
print(f"📊 ZIP任务 {task_id}: 找到 {len(files_to_scan)} 个文件 (最大文件数: {max_analyze_files}, 请求间隔: {llm_gap_ms}ms)")
|
||||
|
||||
total_issues = 0
|
||||
total_lines = 0
|
||||
quality_scores = []
|
||||
scanned_files = 0
|
||||
failed_files = 0
|
||||
|
||||
for file_info in files_to_scan:
|
||||
# 检查是否取消
|
||||
if task_control.is_cancelled(task_id):
|
||||
print(f"🛑 ZIP任务 {task_id} 已被取消")
|
||||
task.status = "cancelled"
|
||||
task.completed_at = datetime.now(timezone.utc)
|
||||
await db.commit()
|
||||
task_control.cleanup_task(task_id)
|
||||
return
|
||||
|
||||
try:
|
||||
content = file_info['content']
|
||||
total_lines += content.count('\n') + 1
|
||||
language = get_language_from_path(file_info['path'])
|
||||
|
||||
# 获取规则集和提示词模板ID
|
||||
scan_config = (user_config or {}).get('scan_config', {})
|
||||
rule_set_id = scan_config.get('rule_set_id')
|
||||
prompt_template_id = scan_config.get('prompt_template_id')
|
||||
|
||||
# 使用规则集和提示词模板进行分析
|
||||
if rule_set_id or prompt_template_id:
|
||||
result = await llm_service.analyze_code_with_rules(
|
||||
content, language,
|
||||
rule_set_id=rule_set_id,
|
||||
prompt_template_id=prompt_template_id,
|
||||
db_session=db
|
||||
)
|
||||
else:
|
||||
result = await llm_service.analyze_code(content, language)
|
||||
|
||||
issues = result.get("issues", [])
|
||||
for i in issues:
|
||||
issue = AuditIssue(
|
||||
task_id=task.id,
|
||||
file_path=file_info['path'],
|
||||
line_number=i.get('line', 1),
|
||||
column_number=i.get('column'),
|
||||
issue_type=i.get('type', 'maintainability'),
|
||||
severity=i.get('severity', 'low'),
|
||||
title=i.get('title', 'Issue'),
|
||||
message=i.get('title', 'Issue'),
|
||||
description=i.get('description'),
|
||||
suggestion=i.get('suggestion'),
|
||||
code_snippet=i.get('code_snippet'),
|
||||
ai_explanation=json.dumps(i.get('xai')) if i.get('xai') else None,
|
||||
status="open"
|
||||
)
|
||||
db.add(issue)
|
||||
total_issues += 1
|
||||
|
||||
if "quality_score" in result:
|
||||
quality_scores.append(result["quality_score"])
|
||||
|
||||
scanned_files += 1
|
||||
task.scanned_files = scanned_files
|
||||
task.total_lines = total_lines
|
||||
task.issues_count = total_issues
|
||||
await db.commit()
|
||||
|
||||
print(f"📈 ZIP任务 {task_id}: 进度 {scanned_files}/{len(files_to_scan)}")
|
||||
|
||||
# 请求间隔
|
||||
await asyncio.sleep(llm_gap_ms / 1000)
|
||||
|
||||
except Exception as file_error:
|
||||
failed_files += 1
|
||||
print(f"❌ ZIP任务分析文件失败 ({file_info['path']}): {file_error}")
|
||||
await asyncio.sleep(llm_gap_ms / 1000)
|
||||
|
||||
# 完成任务
|
||||
avg_quality_score = sum(quality_scores) / len(quality_scores) if quality_scores else 100.0
|
||||
|
||||
# 如果有文件需要分析但全部失败,标记为失败
|
||||
if len(files_to_scan) > 0 and scanned_files == 0:
|
||||
task.status = "failed"
|
||||
task.completed_at = datetime.now(timezone.utc)
|
||||
task.scanned_files = 0
|
||||
task.total_lines = total_lines
|
||||
task.issues_count = 0
|
||||
task.quality_score = 0
|
||||
await db.commit()
|
||||
print(f"❌ ZIP任务 {task_id} 失败: 所有 {len(files_to_scan)} 个文件分析均失败,请检查 LLM API 配置")
|
||||
else:
|
||||
task.status = "completed"
|
||||
task.completed_at = datetime.now(timezone.utc)
|
||||
task.scanned_files = scanned_files
|
||||
task.total_lines = total_lines
|
||||
task.issues_count = total_issues
|
||||
task.quality_score = avg_quality_score
|
||||
await db.commit()
|
||||
print(f"✅ ZIP任务 {task_id} 完成: 扫描 {scanned_files} 个文件, 发现 {total_issues} 个问题")
|
||||
task_control.cleanup_task(task_id)
|
||||
|
||||
except Exception as e:
|
||||
print(f"❌ ZIP扫描失败: {e}")
|
||||
task.status = "failed"
|
||||
task.completed_at = datetime.now(timezone.utc)
|
||||
await db.commit()
|
||||
task_control.cleanup_task(task_id)
|
||||
finally:
|
||||
# Cleanup - 只清理解压目录,不删除源ZIP文件(已持久化存储)
|
||||
if extract_dir.exists():
|
||||
shutil.rmtree(extract_dir)
|
||||
|
||||
|
||||
@router.post("/upload-zip")
|
||||
async def scan_zip(
|
||||
background_tasks: BackgroundTasks,
|
||||
@@ -243,65 +83,47 @@ async def scan_zip(
|
||||
project = await db.get(Project, project_id)
|
||||
if not project:
|
||||
raise HTTPException(status_code=404, detail="项目不存在")
|
||||
|
||||
|
||||
# 检查权限:只有项目所有者可以上传
|
||||
if project.owner_id != current_user.id:
|
||||
raise HTTPException(status_code=403, detail="无权操作此项目")
|
||||
|
||||
|
||||
# Validate file
|
||||
if not file.filename.lower().endswith('.zip'):
|
||||
if not file.filename.lower().endswith(".zip"):
|
||||
raise HTTPException(status_code=400, detail="请上传ZIP格式文件")
|
||||
|
||||
|
||||
# Save Uploaded File to temp
|
||||
file_id = str(uuid.uuid4())
|
||||
file_path = f"/tmp/{file_id}.zip"
|
||||
with open(file_path, "wb") as buffer:
|
||||
shutil.copyfileobj(file.file, buffer)
|
||||
|
||||
# Check file size
|
||||
file_size = os.path.getsize(file_path)
|
||||
if file_size > 500 * 1024 * 1024: # 500MB limit
|
||||
os.remove(file_path)
|
||||
raise HTTPException(status_code=400, detail="文件大小不能超过500MB")
|
||||
|
||||
# 保存ZIP文件到持久化存储
|
||||
await save_project_zip(project_id, file_path, file.filename)
|
||||
|
||||
# Parse scan_config if provided
|
||||
parsed_scan_config = {}
|
||||
if scan_config:
|
||||
try:
|
||||
parsed_scan_config = json.loads(scan_config)
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
await save_upload(file, file_path)
|
||||
|
||||
# Create Task
|
||||
task = AuditTask(
|
||||
project_id=project_id,
|
||||
created_by=current_user.id,
|
||||
task_type="zip_upload",
|
||||
status="pending",
|
||||
scan_config=scan_config if scan_config else "{}"
|
||||
)
|
||||
db.add(task)
|
||||
await db.commit()
|
||||
await db.refresh(task)
|
||||
try:
|
||||
# 保存ZIP文件到持久化存储
|
||||
await save_project_zip(project_id, file_path, file.filename)
|
||||
|
||||
# 获取用户配置
|
||||
user_config = await get_user_config_dict(db, current_user.id)
|
||||
|
||||
# 将扫描配置注入到 user_config 中(包括规则集、提示词模板和排除模式)
|
||||
if parsed_scan_config:
|
||||
user_config['scan_config'] = {
|
||||
'file_paths': parsed_scan_config.get('file_paths', []),
|
||||
'exclude_patterns': parsed_scan_config.get('exclude_patterns', []),
|
||||
'rule_set_id': parsed_scan_config.get('rule_set_id'),
|
||||
'prompt_template_id': parsed_scan_config.get('prompt_template_id'),
|
||||
}
|
||||
if scan_config:
|
||||
try:
|
||||
parsed_scan_config = json.loads(scan_config)
|
||||
if not isinstance(parsed_scan_config, dict):
|
||||
raise ValueError("scan_config must be an object")
|
||||
except (ValueError, TypeError) as exc:
|
||||
raise HTTPException(status_code=400, detail="扫描配置必须是有效 JSON 对象") from exc
|
||||
|
||||
# Trigger Background Task - 使用持久化存储的文件路径
|
||||
stored_zip_path = await load_project_zip(project_id)
|
||||
background_tasks.add_task(process_zip_task, task.id, stored_zip_path or file_path, AsyncSessionLocal, user_config)
|
||||
# Create Task
|
||||
task = AuditTask(
|
||||
project_id=project_id,
|
||||
created_by=current_user.id,
|
||||
task_type="zip_upload",
|
||||
status="pending",
|
||||
scan_config=scan_config if scan_config else "{}",
|
||||
)
|
||||
db.add(task)
|
||||
stored_zip_path = await load_project_zip(project_id)
|
||||
await enqueue(db, task, "zip", zip_path=file_path)
|
||||
await db.commit()
|
||||
await db.refresh(task)
|
||||
finally:
|
||||
Path(file_path).unlink(missing_ok=True)
|
||||
|
||||
return {"task_id": task.id, "status": "queued"}
|
||||
|
||||
@@ -329,43 +151,29 @@ async def scan_stored_zip(
|
||||
project = await db.get(Project, project_id)
|
||||
if not project:
|
||||
raise HTTPException(status_code=404, detail="项目不存在")
|
||||
|
||||
|
||||
# 检查权限:只有项目所有者可以扫描
|
||||
if project.owner_id != current_user.id:
|
||||
raise HTTPException(status_code=403, detail="无权操作此项目")
|
||||
|
||||
|
||||
# 检查是否有存储的ZIP文件
|
||||
stored_zip_path = await load_project_zip(project_id)
|
||||
if not stored_zip_path:
|
||||
raise HTTPException(status_code=400, detail="项目没有已存储的ZIP文件,请先上传")
|
||||
|
||||
|
||||
# Create Task
|
||||
task = AuditTask(
|
||||
project_id=project_id,
|
||||
created_by=current_user.id,
|
||||
task_type="zip_upload",
|
||||
status="pending",
|
||||
scan_config=json.dumps(scan_request.dict()) if scan_request else "{}"
|
||||
scan_config=json.dumps(scan_request.dict()) if scan_request else "{}",
|
||||
)
|
||||
db.add(task)
|
||||
await enqueue(db, task, "zip", zip_path=stored_zip_path)
|
||||
await db.commit()
|
||||
await db.refresh(task)
|
||||
|
||||
# 获取用户配置
|
||||
user_config = await get_user_config_dict(db, current_user.id)
|
||||
|
||||
# 将扫描配置注入到 user_config 中(包括规则集、提示词模板和排除模式)
|
||||
if scan_request:
|
||||
user_config['scan_config'] = {
|
||||
'file_paths': scan_request.file_paths or [],
|
||||
'exclude_patterns': scan_request.exclude_patterns or [],
|
||||
'rule_set_id': scan_request.rule_set_id,
|
||||
'prompt_template_id': scan_request.prompt_template_id,
|
||||
}
|
||||
|
||||
# Trigger Background Task
|
||||
background_tasks.add_task(process_zip_task, task.id, stored_zip_path, AsyncSessionLocal, user_config)
|
||||
|
||||
return {"task_id": task.id, "status": "queued"}
|
||||
|
||||
|
||||
@@ -389,82 +197,39 @@ class InstantAnalysisResponse(BaseModel):
|
||||
from_attributes = True
|
||||
|
||||
|
||||
async def get_user_config_dict(db: AsyncSession, user_id: str) -> dict:
|
||||
"""获取用户配置字典(包含解密敏感字段)"""
|
||||
from app.core.encryption import decrypt_sensitive_data
|
||||
|
||||
# 需要解密的敏感字段列表(与 config.py 保持一致)
|
||||
SENSITIVE_LLM_FIELDS = [
|
||||
'llmApiKey', 'geminiApiKey', 'openaiApiKey', 'claudeApiKey',
|
||||
'qwenApiKey', 'deepseekApiKey', 'zhipuApiKey', 'moonshotApiKey',
|
||||
'baiduApiKey', 'minimaxApiKey', 'doubaoApiKey'
|
||||
]
|
||||
SENSITIVE_OTHER_FIELDS = ['githubToken', 'gitlabToken']
|
||||
|
||||
def decrypt_config(config: dict, sensitive_fields: list) -> dict:
|
||||
"""解密配置中的敏感字段"""
|
||||
decrypted = config.copy()
|
||||
for field in sensitive_fields:
|
||||
if field in decrypted and decrypted[field]:
|
||||
decrypted[field] = decrypt_sensitive_data(decrypted[field])
|
||||
return decrypted
|
||||
|
||||
result = await db.execute(
|
||||
select(UserConfig).where(UserConfig.user_id == user_id)
|
||||
)
|
||||
config = result.scalar_one_or_none()
|
||||
if not config:
|
||||
return {}
|
||||
|
||||
# 解析配置
|
||||
llm_config = json.loads(config.llm_config) if config.llm_config else {}
|
||||
other_config = json.loads(config.other_config) if config.other_config else {}
|
||||
|
||||
# 解密敏感字段
|
||||
llm_config = decrypt_config(llm_config, SENSITIVE_LLM_FIELDS)
|
||||
other_config = decrypt_config(other_config, SENSITIVE_OTHER_FIELDS)
|
||||
|
||||
return {
|
||||
'llmConfig': llm_config,
|
||||
'otherConfig': other_config,
|
||||
}
|
||||
|
||||
|
||||
@router.post("/instant")
|
||||
async def instant_analysis(
|
||||
req: InstantAnalysisRequest,
|
||||
db: AsyncSession = Depends(get_db),
|
||||
current_user: User = Depends(deps.get_current_user),
|
||||
current_user: User = Depends(deps.get_current_user),
|
||||
) -> Any:
|
||||
"""
|
||||
Perform instant code analysis.
|
||||
"""
|
||||
# 获取用户配置
|
||||
user_config = await get_user_config_dict(db, current_user.id)
|
||||
|
||||
|
||||
# 创建使用用户配置的LLM服务实例
|
||||
llm_service = LLMService(user_config=user_config)
|
||||
|
||||
|
||||
start_time = datetime.now(timezone.utc)
|
||||
|
||||
|
||||
try:
|
||||
# 如果指定了提示词模板,使用自定义分析
|
||||
# 统一使用 analyze_code_with_rules,会自动使用默认模板
|
||||
result = await llm_service.analyze_code_with_rules(
|
||||
req.code, req.language,
|
||||
req.code,
|
||||
req.language,
|
||||
prompt_template_id=req.prompt_template_id,
|
||||
db_session=db,
|
||||
use_default_template=True # 没有指定模板时使用数据库中的默认模板
|
||||
use_default_template=True, # 没有指定模板时使用数据库中的默认模板
|
||||
)
|
||||
except Exception as e:
|
||||
# 分析失败,返回错误信息
|
||||
error_msg = str(e)
|
||||
print(f"❌ 即时分析失败: {error_msg}")
|
||||
raise HTTPException(
|
||||
status_code=500,
|
||||
detail=f"代码分析失败: {error_msg}"
|
||||
)
|
||||
|
||||
raise HTTPException(status_code=500, detail=f"代码分析失败: {error_msg}")
|
||||
|
||||
end_time = datetime.now(timezone.utc)
|
||||
duration = (end_time - start_time).total_seconds()
|
||||
|
||||
@@ -476,18 +241,14 @@ async def instant_analysis(
|
||||
analysis_result=json.dumps(result),
|
||||
issues_count=len(result.get("issues", [])),
|
||||
quality_score=result.get("quality_score", 0),
|
||||
analysis_time=duration
|
||||
analysis_time=duration,
|
||||
)
|
||||
db.add(analysis)
|
||||
await db.commit()
|
||||
await db.refresh(analysis)
|
||||
|
||||
|
||||
# Return result with analysis ID for export functionality
|
||||
return {
|
||||
**result,
|
||||
"analysis_id": analysis.id,
|
||||
"analysis_time": duration
|
||||
}
|
||||
return {**result, "analysis_id": analysis.id, "analysis_time": duration}
|
||||
|
||||
|
||||
@router.get("/instant/history", response_model=List[InstantAnalysisResponse])
|
||||
@@ -523,13 +284,13 @@ async def delete_instant_analysis(
|
||||
.where(InstantAnalysis.user_id == current_user.id)
|
||||
)
|
||||
analysis = result.scalar_one_or_none()
|
||||
|
||||
|
||||
if not analysis:
|
||||
raise HTTPException(status_code=404, detail="分析记录不存在")
|
||||
|
||||
|
||||
await db.delete(analysis)
|
||||
await db.commit()
|
||||
|
||||
|
||||
return {"message": "删除成功"}
|
||||
|
||||
|
||||
@@ -542,12 +303,10 @@ async def delete_all_instant_analyses(
|
||||
Delete all instant analysis records for current user.
|
||||
"""
|
||||
from sqlalchemy import delete
|
||||
|
||||
await db.execute(
|
||||
delete(InstantAnalysis).where(InstantAnalysis.user_id == current_user.id)
|
||||
)
|
||||
|
||||
await db.execute(delete(InstantAnalysis).where(InstantAnalysis.user_id == current_user.id))
|
||||
await db.commit()
|
||||
|
||||
|
||||
return {"message": "已清空所有历史记录"}
|
||||
|
||||
|
||||
@@ -562,7 +321,7 @@ async def export_instant_report_pdf(
|
||||
"""
|
||||
from fastapi.responses import Response
|
||||
from app.services.report_generator import ReportGenerator
|
||||
|
||||
|
||||
# 获取即时分析记录
|
||||
result = await db.execute(
|
||||
select(InstantAnalysis)
|
||||
@@ -570,29 +329,25 @@ async def export_instant_report_pdf(
|
||||
.where(InstantAnalysis.user_id == current_user.id)
|
||||
)
|
||||
analysis = result.scalar_one_or_none()
|
||||
|
||||
|
||||
if not analysis:
|
||||
raise HTTPException(status_code=404, detail="分析记录不存在")
|
||||
|
||||
|
||||
# 解析分析结果
|
||||
try:
|
||||
analysis_result = json.loads(analysis.analysis_result) if analysis.analysis_result else {}
|
||||
except json.JSONDecodeError:
|
||||
analysis_result = {}
|
||||
|
||||
|
||||
# 生成 PDF
|
||||
pdf_bytes = ReportGenerator.generate_instant_report(
|
||||
analysis_result,
|
||||
analysis.language,
|
||||
analysis.analysis_time
|
||||
analysis_result, analysis.language, analysis.analysis_time
|
||||
)
|
||||
|
||||
|
||||
# 返回 PDF 文件
|
||||
filename = f"instant-analysis-{analysis.language}-{analysis.id[:8]}.pdf"
|
||||
return Response(
|
||||
content=pdf_bytes,
|
||||
media_type="application/pdf",
|
||||
headers={
|
||||
"Content-Disposition": f'attachment; filename="{filename}"'
|
||||
}
|
||||
headers={"Content-Disposition": f'attachment; filename="{filename}"'},
|
||||
)
|
||||
|
||||
@@ -1,3 +1,7 @@
|
||||
from typing import Annotated
|
||||
from fastapi import Query, Response
|
||||
from sqlalchemy import func, case
|
||||
from app.services.audit_queue import request_cancel
|
||||
from typing import Any, List, Optional
|
||||
from fastapi import APIRouter, Depends, HTTPException
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
@@ -87,6 +91,11 @@ class AuditTaskSchema(BaseModel):
|
||||
|
||||
@router.get("/", response_model=List[AuditTaskSchema])
|
||||
async def list_tasks(
|
||||
response: Response = None,
|
||||
skip: Annotated[int, Query(ge=0)] = 0,
|
||||
limit: Annotated[int, Query(ge=1, le=200)] = 50,
|
||||
status: Optional[str] = None,
|
||||
search: Annotated[str, Query(max_length=200)] = "",
|
||||
project_id: Optional[str] = None,
|
||||
db: AsyncSession = Depends(get_db),
|
||||
current_user: User = Depends(deps.get_current_user),
|
||||
@@ -105,7 +114,16 @@ async def list_tasks(
|
||||
query = query.where(AuditTask.project_id.in_(user_project_ids)) if user_project_ids else query.where(False)
|
||||
if project_id:
|
||||
query = query.where(AuditTask.project_id == project_id)
|
||||
query = query.order_by(AuditTask.created_at.desc())
|
||||
if status:
|
||||
query = query.where(AuditTask.status == status)
|
||||
if search:
|
||||
pattern = "%" + search.replace("%", r"\%").replace("_", r"\_") + "%"
|
||||
query = query.where(AuditTask.project.has(Project.name.ilike(pattern)) | AuditTask.task_type.ilike(pattern))
|
||||
if response is not None:
|
||||
response.headers["X-Total-Count"] = str(await db.scalar(
|
||||
select(func.count()).select_from(query.order_by(None).subquery())
|
||||
))
|
||||
query = query.order_by(AuditTask.created_at.desc(), AuditTask.id.desc()).offset(skip).limit(limit)
|
||||
result = await db.execute(query)
|
||||
return result.scalars().all()
|
||||
|
||||
@@ -157,6 +175,7 @@ async def cancel_task(
|
||||
raise HTTPException(status_code=400, detail="只能取消待处理或运行中的任务")
|
||||
|
||||
# 标记任务为取消
|
||||
await request_cancel(db, id)
|
||||
task_control.cancel_task(id)
|
||||
|
||||
# 更新数据库状态
|
||||
@@ -170,6 +189,11 @@ async def cancel_task(
|
||||
@router.get("/{id}/issues", response_model=List[AuditIssueSchema])
|
||||
async def read_task_issues(
|
||||
id: str,
|
||||
response: Response = None,
|
||||
skip: Annotated[int, Query(ge=0)] = 0,
|
||||
limit: Annotated[int, Query(ge=1, le=200)] = 50,
|
||||
severity: Optional[str] = None,
|
||||
status: Optional[str] = None,
|
||||
db: AsyncSession = Depends(get_db),
|
||||
current_user: User = Depends(deps.get_current_user),
|
||||
) -> Any:
|
||||
@@ -188,15 +212,20 @@ async def read_task_issues(
|
||||
if task.created_by != current_user.id:
|
||||
raise HTTPException(status_code=403, detail="无权查看此任务的问题")
|
||||
|
||||
result = await db.execute(
|
||||
select(AuditIssue)
|
||||
.where(AuditIssue.task_id == id)
|
||||
.order_by(
|
||||
# 按严重程度排序
|
||||
AuditIssue.severity.desc(),
|
||||
AuditIssue.created_at.desc()
|
||||
)
|
||||
)
|
||||
query = select(AuditIssue).where(AuditIssue.task_id == id)
|
||||
if severity:
|
||||
query = query.where(AuditIssue.severity == severity)
|
||||
if status:
|
||||
query = query.where(AuditIssue.status == status)
|
||||
if response is not None:
|
||||
response.headers["X-Total-Count"] = str(await db.scalar(
|
||||
select(func.count()).select_from(query.subquery())
|
||||
))
|
||||
severity_order = case({"critical": 0, "high": 1, "medium": 2, "low": 3},
|
||||
value=AuditIssue.severity, else_=4)
|
||||
result = await db.execute(query.order_by(
|
||||
severity_order, AuditIssue.created_at.desc(), AuditIssue.id.desc(),
|
||||
).offset(skip).limit(limit))
|
||||
return result.scalars().all()
|
||||
|
||||
|
||||
@@ -211,6 +240,12 @@ async def update_issue(
|
||||
"""
|
||||
Update issue status (e.g., resolve, mark as false positive).
|
||||
"""
|
||||
task = await db.get(AuditTask, task_id)
|
||||
if not task:
|
||||
raise HTTPException(status_code=404, detail="任务不存在")
|
||||
if task.created_by != current_user.id:
|
||||
raise HTTPException(status_code=403, detail="无权修改此任务的问题")
|
||||
|
||||
result = await db.execute(
|
||||
select(AuditIssue)
|
||||
.where(AuditIssue.id == issue_id, AuditIssue.task_id == task_id)
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
from typing import List, Union, Optional
|
||||
from pydantic import AnyHttpUrl, validator
|
||||
from typing import List, Union, Optional, Literal
|
||||
from pydantic import AnyHttpUrl, Field, field_validator, model_validator, validator
|
||||
from pydantic_settings import BaseSettings
|
||||
|
||||
|
||||
@@ -8,13 +8,31 @@ class Settings(BaseSettings):
|
||||
API_V1_STR: str = "/api/v1"
|
||||
|
||||
# SECURITY
|
||||
SECRET_KEY: str = "changethis_in_production_to_a_long_random_string"
|
||||
ENVIRONMENT: Literal["development", "test", "production"] = "production"
|
||||
SECRET_KEY: str = Field(min_length=32)
|
||||
DEMO_ENABLED: bool = False
|
||||
ALGORITHM: str = "HS256"
|
||||
ACCESS_TOKEN_EXPIRE_MINUTES: int = 60 * 24 * 8 # 8 days
|
||||
|
||||
# CORS
|
||||
BACKEND_CORS_ORIGINS: List[AnyHttpUrl] = []
|
||||
|
||||
@field_validator("SECRET_KEY")
|
||||
@classmethod
|
||||
def reject_placeholder_secret(cls, value: str) -> str:
|
||||
if value in {
|
||||
"changethis_in_production_to_a_long_random_string",
|
||||
"your-super-secret-key-change-this-in-production",
|
||||
} or len(value.strip()) < 32:
|
||||
raise ValueError("SECRET_KEY must be a random secret of at least 32 characters")
|
||||
return value
|
||||
|
||||
@model_validator(mode="after")
|
||||
def validate_demo_environment(self):
|
||||
if self.ENVIRONMENT == "production" and self.DEMO_ENABLED:
|
||||
raise ValueError("DEMO_ENABLED is not allowed in production")
|
||||
return self
|
||||
|
||||
@validator("BACKEND_CORS_ORIGINS", pre=True)
|
||||
def assemble_cors_origins(cls, v: Union[str, List[str]]) -> Union[List[str], str]:
|
||||
if isinstance(v, str) and not v.startswith("["):
|
||||
@@ -34,7 +52,10 @@ class Settings(BaseSettings):
|
||||
def assemble_db_connection(cls, v: str | None, values: dict[str, any]) -> str:
|
||||
if isinstance(v, str):
|
||||
return v
|
||||
return str(f"postgresql+asyncpg://{values.get('POSTGRES_USER')}:{values.get('POSTGRES_PASSWORD')}@{values.get('POSTGRES_SERVER')}/{values.get('POSTGRES_DB')}")
|
||||
from sqlalchemy.engine import URL
|
||||
return URL.create("postgresql+asyncpg", username=values.get("POSTGRES_USER"),
|
||||
password=values.get("POSTGRES_PASSWORD"), host=values.get("POSTGRES_SERVER"),
|
||||
database=values.get("POSTGRES_DB")).render_as_string(hide_password=False)
|
||||
|
||||
# LLM配置
|
||||
LLM_PROVIDER: str = "openai" # gemini, openai, claude, qwen, deepseek, zhipu, moonshot, baidu, minimax, doubao, ollama
|
||||
@@ -82,6 +103,18 @@ class Settings(BaseSettings):
|
||||
|
||||
# ZIP文件存储配置
|
||||
ZIP_STORAGE_PATH: str = "./uploads/zip_files" # ZIP文件存储目录
|
||||
ZIP_MAX_UPLOAD_BYTES: int = Field(default=500 * 1024 * 1024, gt=0)
|
||||
ZIP_MAX_EXTRACTED_BYTES: int = Field(default=2 * 1024 * 1024 * 1024, gt=0)
|
||||
ZIP_MAX_FILE_BYTES: int = Field(default=50 * 1024 * 1024, gt=0)
|
||||
ZIP_MAX_ENTRIES: int = Field(default=50000, gt=0)
|
||||
ZIP_MAX_COMPRESSION_RATIO: int = Field(default=200, gt=0)
|
||||
AUDIT_WORKSPACE_PATH: str = "./uploads/audit_workspaces"
|
||||
WORKER_POLL_SECONDS: float = Field(default=1.0, gt=0)
|
||||
WORKER_HEARTBEAT_SECONDS: float = Field(default=2.0, gt=0)
|
||||
WORKER_LEASE_SECONDS: int = Field(default=30, ge=10)
|
||||
WORKER_MAX_ATTEMPTS: int = Field(default=3, ge=1)
|
||||
WORKER_TASK_TIMEOUT_SECONDS: int = Field(default=3600, ge=30)
|
||||
WORKER_RETENTION_DAYS: int = Field(default=7, ge=1)
|
||||
|
||||
# 输出语言配置 - 支持 zh-CN(中文)和 en-US(英文)
|
||||
OUTPUT_LANGUAGE: str = "zh-CN"
|
||||
@@ -109,6 +142,7 @@ class Settings(BaseSettings):
|
||||
AGENT_TIMEOUT_SECONDS: int = 1800 # Agent 超时时间(30分钟)
|
||||
|
||||
# 沙箱配置(必须)
|
||||
SANDBOX_ENABLED: bool = True
|
||||
SANDBOX_IMAGE: str = "deepaudit/sandbox:latest" # 沙箱 Docker 镜像
|
||||
SANDBOX_MEMORY_LIMIT: str = "512m" # 沙箱内存限制
|
||||
SANDBOX_CPU_LIMIT: float = 1.0 # 沙箱 CPU 限制
|
||||
@@ -122,6 +156,12 @@ class Settings(BaseSettings):
|
||||
RAG_CHUNK_OVERLAP: int = 50 # 代码块重叠(Token)
|
||||
RAG_TOP_K: int = 10 # 检索返回数量
|
||||
|
||||
@model_validator(mode="after")
|
||||
def validate_worker_timing(self):
|
||||
if self.WORKER_HEARTBEAT_SECONDS >= self.WORKER_LEASE_SECONDS:
|
||||
raise ValueError("WORKER_HEARTBEAT_SECONDS must be less than WORKER_LEASE_SECONDS")
|
||||
return self
|
||||
|
||||
class Config:
|
||||
case_sensitive = True
|
||||
env_file = ".env"
|
||||
|
||||
@@ -0,0 +1,52 @@
|
||||
"""Bound ZIP request bodies before Starlette spools multipart uploads to disk."""
|
||||
|
||||
from starlette.formparsers import MultiPartException
|
||||
from starlette.responses import JSONResponse
|
||||
from app.core.config import settings
|
||||
|
||||
|
||||
class UploadLimitMiddleware:
|
||||
def __init__(self, app):
|
||||
self.app = app
|
||||
|
||||
async def __call__(self, scope, receive, send):
|
||||
path = scope.get("path", "")
|
||||
is_upload = path == settings.API_V1_STR + "/scan/upload-zip" or (
|
||||
path.startswith(settings.API_V1_STR + "/projects/") and path.endswith("/zip")
|
||||
)
|
||||
if scope["type"] != "http" or scope.get("method") != "POST" or not is_upload:
|
||||
return await self.app(scope, receive, send)
|
||||
# Multipart boundaries and small form fields have a separate 1 MiB allowance.
|
||||
limit = settings.ZIP_MAX_UPLOAD_BYTES + 1024 * 1024
|
||||
headers = dict(scope.get("headers", []))
|
||||
try:
|
||||
declared = int(headers.get(b"content-length", b"0"))
|
||||
except ValueError:
|
||||
declared = 0
|
||||
rejection = JSONResponse(
|
||||
{"detail": "ZIP upload exceeds the request quota"}, status_code=413
|
||||
)
|
||||
if declared > limit:
|
||||
return await rejection(scope, receive, send)
|
||||
total, exceeded = 0, False
|
||||
|
||||
async def bounded_receive():
|
||||
nonlocal total, exceeded
|
||||
message = await receive()
|
||||
if message["type"] == "http.request":
|
||||
total += len(message.get("body", b""))
|
||||
if total > limit:
|
||||
exceeded = True
|
||||
# Starlette closes partially spooled files on this exception.
|
||||
raise MultiPartException("ZIP upload exceeds the request quota")
|
||||
return message
|
||||
|
||||
async def bounded_send(message):
|
||||
if exceeded and message["type"] == "http.response.start":
|
||||
message = {**message, "status": 413}
|
||||
await send(message)
|
||||
|
||||
try:
|
||||
await self.app(scope, bounded_receive, bounded_send)
|
||||
except MultiPartException:
|
||||
await rejection(scope, receive, send)
|
||||
@@ -8,7 +8,8 @@ from datetime import datetime, timedelta, timezone
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
from sqlalchemy.future import select
|
||||
|
||||
from app.core.security import get_password_hash
|
||||
from app.core.config import settings
|
||||
from app.core.security import get_password_hash, verify_password
|
||||
from app.models.user import User
|
||||
from app.models.project import Project, ProjectMember
|
||||
from app.models.audit import AuditTask, AuditIssue
|
||||
@@ -265,12 +266,19 @@ async def init_db(db: AsyncSession) -> None:
|
||||
"""
|
||||
logger.info("开始初始化数据库...")
|
||||
|
||||
# 创建演示用户
|
||||
demo_user = await create_demo_user(db)
|
||||
|
||||
# 创建演示数据
|
||||
if demo_user:
|
||||
await create_demo_data(db, demo_user)
|
||||
if settings.DEMO_ENABLED:
|
||||
demo_user = await create_demo_user(db)
|
||||
if demo_user:
|
||||
await create_demo_data(db, demo_user)
|
||||
else:
|
||||
# Retire the known default credential on upgrades without deleting user data.
|
||||
result = await db.execute(select(User).where(User.email == DEFAULT_DEMO_EMAIL))
|
||||
demo_user = result.scalars().first()
|
||||
if demo_user and demo_user.is_active and verify_password(
|
||||
DEFAULT_DEMO_PASSWORD, demo_user.hashed_password
|
||||
):
|
||||
demo_user.is_active = False
|
||||
logger.warning("Disabled legacy demo account using the default password")
|
||||
|
||||
await db.commit()
|
||||
|
||||
|
||||
+23
-21
@@ -64,6 +64,9 @@ async def lifespan(app: FastAPI):
|
||||
await init_db(db)
|
||||
logger.info(" - 数据库初始化完成")
|
||||
except Exception as e:
|
||||
if settings.ENVIRONMENT == "production":
|
||||
# Bootstrap also disables legacy default credentials: fail closed.
|
||||
raise
|
||||
# 表不存在时静默跳过,等待用户运行数据库迁移
|
||||
error_msg = str(e)
|
||||
if "does not exist" in error_msg or "UndefinedTableError" in error_msg:
|
||||
@@ -71,25 +74,11 @@ async def lifespan(app: FastAPI):
|
||||
else:
|
||||
logger.warning(f"数据库初始化跳过: {e}")
|
||||
|
||||
# 检查 Agent 服务
|
||||
logger.info("检查 Agent 核心服务...")
|
||||
issues = await check_agent_services()
|
||||
if issues:
|
||||
logger.warning("=" * 50)
|
||||
logger.warning("Agent 服务检查发现问题:")
|
||||
for issue in issues:
|
||||
logger.warning(f" - {issue}")
|
||||
logger.warning("部分功能可能不可用,请检查配置")
|
||||
logger.warning("=" * 50)
|
||||
else:
|
||||
logger.info(" - Agent 核心服务检查通过")
|
||||
logger.info("审计任务通过持久化队列交由独立 worker 执行")
|
||||
|
||||
logger.info("=" * 50)
|
||||
logger.info("DeepAudit 后端服务已启动")
|
||||
logger.info(f"API 文档: http://localhost:8000/docs")
|
||||
logger.info("=" * 50)
|
||||
logger.info("演示账户: demo@example.com / demo123")
|
||||
logger.info("=" * 50)
|
||||
|
||||
yield
|
||||
|
||||
@@ -102,13 +91,17 @@ app = FastAPI(
|
||||
lifespan=lifespan
|
||||
)
|
||||
|
||||
# Configure CORS - Allow all origins in development
|
||||
from app.core.upload_limit import UploadLimitMiddleware
|
||||
app.add_middleware(UploadLimitMiddleware)
|
||||
|
||||
# Same-origin deployments need no CORS entries; cross-origin clients are explicit.
|
||||
app.add_middleware(
|
||||
CORSMiddleware,
|
||||
allow_origins=["*"], # In production, replace with specific frontend URL
|
||||
allow_origins=[str(origin).rstrip("/") for origin in settings.BACKEND_CORS_ORIGINS],
|
||||
allow_credentials=True,
|
||||
allow_methods=["*"],
|
||||
allow_headers=["*"],
|
||||
expose_headers=["X-Total-Count"],
|
||||
)
|
||||
|
||||
app.include_router(api_router, prefix=settings.API_V1_STR)
|
||||
@@ -124,8 +117,17 @@ async def root():
|
||||
return {
|
||||
"message": "Welcome to DeepAudit API",
|
||||
"docs": "/docs",
|
||||
"demo_account": {
|
||||
"email": "demo@example.com",
|
||||
"password": "demo123"
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@app.get("/ready")
|
||||
async def readiness():
|
||||
from fastapi.responses import JSONResponse
|
||||
from sqlalchemy import select
|
||||
from app.models.audit_job import AuditJob
|
||||
try:
|
||||
async with AsyncSessionLocal() as db:
|
||||
await db.execute(select(AuditJob.id).limit(1))
|
||||
return {"status": "ready"}
|
||||
except Exception:
|
||||
return JSONResponse(status_code=503, content={"status": "unavailable"})
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
from .user import User
|
||||
from .audit_job import AuditJob
|
||||
from .user_config import UserConfig
|
||||
from .project import Project, ProjectMember
|
||||
from .audit import AuditTask, AuditIssue
|
||||
@@ -12,4 +13,3 @@ from .agent_task import (
|
||||
)
|
||||
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,21 @@
|
||||
"""Durable delivery and recovery metadata for audit execution."""
|
||||
from sqlalchemy import Boolean, Column, DateTime, Integer, JSON, String, Text, Index
|
||||
from sqlalchemy.sql import func
|
||||
from app.db.base import Base
|
||||
|
||||
|
||||
class AuditJob(Base):
|
||||
__tablename__ = "audit_jobs"
|
||||
id = Column(String(36), primary_key=True) # Same id as the owning audit task.
|
||||
kind = Column(String(20), nullable=False)
|
||||
status = Column(String(20), nullable=False, default="queued")
|
||||
payload = Column(JSON, nullable=False, default=dict)
|
||||
checkpoint = Column(JSON, nullable=False, default=dict)
|
||||
attempts = Column(Integer, nullable=False, default=0)
|
||||
cancel_requested = Column(Boolean, nullable=False, default=False)
|
||||
worker_id = Column(String(100), nullable=True)
|
||||
lease_until = Column(DateTime(timezone=True), nullable=True)
|
||||
error = Column(Text, nullable=True)
|
||||
created_at = Column(DateTime(timezone=True), server_default=func.now(), nullable=False)
|
||||
updated_at = Column(DateTime(timezone=True), server_default=func.now(), nullable=False)
|
||||
__table_args__ = (Index("ix_audit_jobs_delivery", "status", "lease_until", "created_at"),)
|
||||
@@ -297,6 +297,7 @@ class BaseAgent(ABC):
|
||||
self._total_tokens = 0
|
||||
self._tool_calls = 0
|
||||
self._cancelled = False
|
||||
self._cancel_callback = None
|
||||
|
||||
# 获取超时配置
|
||||
self._timeout_config = self._get_timeout_config()
|
||||
@@ -512,9 +513,6 @@ class BaseAgent(ABC):
|
||||
self._cancelled = True
|
||||
logger.info(f"[{self.name}] Cancel requested")
|
||||
|
||||
# 🔥 外部取消检查回调
|
||||
self._cancel_callback = None
|
||||
|
||||
def set_cancel_callback(self, callback) -> None:
|
||||
"""设置外部取消检查回调"""
|
||||
self._cancel_callback = callback
|
||||
|
||||
@@ -19,6 +19,7 @@ from datetime import datetime, timezone
|
||||
from .state import AgentState, AgentStatus
|
||||
from .registry import agent_registry
|
||||
from .message import message_bus, MessageType
|
||||
from ..config import get_agent_config
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -88,7 +89,7 @@ class DynamicAgentExecutor:
|
||||
tools: Dict[str, Any],
|
||||
event_emitter=None,
|
||||
max_parallel: int = 5,
|
||||
default_timeout: int = 600,
|
||||
default_timeout: Optional[int] = None,
|
||||
):
|
||||
"""
|
||||
初始化执行器
|
||||
@@ -104,7 +105,10 @@ class DynamicAgentExecutor:
|
||||
self.tools = tools
|
||||
self.event_emitter = event_emitter
|
||||
self.max_parallel = max_parallel
|
||||
self.default_timeout = default_timeout
|
||||
self.default_timeout = (
|
||||
default_timeout if default_timeout is not None
|
||||
else get_agent_config().sub_agent_timeout_seconds
|
||||
)
|
||||
|
||||
# 执行状态
|
||||
self._tasks: Dict[str, ExecutionTask] = {}
|
||||
|
||||
@@ -84,6 +84,33 @@ class AgentGraphController:
|
||||
"note": "Agent将在当前迭代完成后停止",
|
||||
}
|
||||
|
||||
def stop_agent_tree(self, root_agent_id: str) -> Dict[str, Any]:
|
||||
"""Stop one task's root and descendants without touching other task trees."""
|
||||
tree = agent_registry.get_agent_tree()
|
||||
nodes = tree["nodes"]
|
||||
if root_agent_id not in nodes:
|
||||
return {"success": False, "stopped": [], "error": "Agent root not found"}
|
||||
|
||||
pending = [root_agent_id]
|
||||
children: Dict[Optional[str], List[str]] = {}
|
||||
for agent_id, node in nodes.items():
|
||||
children.setdefault(node.get("parent_id"), []).append(agent_id)
|
||||
visited = set()
|
||||
stopped = []
|
||||
failed = []
|
||||
while pending:
|
||||
agent_id = pending.pop()
|
||||
if agent_id in visited:
|
||||
continue
|
||||
visited.add(agent_id)
|
||||
pending.extend(children.get(agent_id, []))
|
||||
result = self.stop_agent(agent_id)
|
||||
if result.get("success"):
|
||||
stopped.append(agent_id)
|
||||
else:
|
||||
failed.append({"id": agent_id, "error": result.get("error")})
|
||||
return {"success": not failed, "stopped": stopped, "failed": failed}
|
||||
|
||||
def stop_all_agents(self, exclude_root: bool = True) -> Dict[str, Any]:
|
||||
"""
|
||||
停止所有Agent
|
||||
|
||||
@@ -4,6 +4,8 @@ Agent 事件管理器
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import time
|
||||
from .config import get_agent_config
|
||||
import json
|
||||
import logging
|
||||
from typing import Optional, Dict, Any, List, AsyncGenerator, Callable
|
||||
@@ -54,17 +56,16 @@ class AgentEventEmitter:
|
||||
self.event_manager = event_manager
|
||||
self._sequence = 0
|
||||
self._current_phase = None
|
||||
self._emit_lock = asyncio.Lock()
|
||||
|
||||
async def emit(self, event_data: AgentEventData):
|
||||
"""发射事件"""
|
||||
self._sequence += 1
|
||||
event_data.phase = event_data.phase or self._current_phase
|
||||
|
||||
await self.event_manager.add_event(
|
||||
task_id=self.task_id,
|
||||
sequence=self._sequence,
|
||||
**event_data.to_dict()
|
||||
)
|
||||
async with self._emit_lock:
|
||||
self._sequence += 1
|
||||
event_data.phase = event_data.phase or self._current_phase
|
||||
await self.event_manager.add_event(
|
||||
task_id=self.task_id, sequence=self._sequence, **event_data.to_dict()
|
||||
)
|
||||
|
||||
async def emit_phase_start(self, phase: str, message: Optional[str] = None):
|
||||
"""发射阶段开始事件"""
|
||||
@@ -263,6 +264,8 @@ class EventManager:
|
||||
self.db_session_factory = db_session_factory
|
||||
self._event_queues: Dict[str, asyncio.Queue] = {}
|
||||
self._event_callbacks: Dict[str, List[Callable]] = {}
|
||||
self._token_buffers = {}
|
||||
self._persist_lock = asyncio.Lock()
|
||||
|
||||
async def add_event(
|
||||
self,
|
||||
@@ -300,13 +303,27 @@ class EventManager:
|
||||
"timestamp": timestamp.isoformat(),
|
||||
}
|
||||
|
||||
# 保存到数据库(跳过高频事件如 thinking_token)
|
||||
skip_db_events = {"thinking_token"}
|
||||
if self.db_session_factory and event_type not in skip_db_events:
|
||||
try:
|
||||
await self._save_event_to_db(event_data)
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to save event to database: {e}")
|
||||
# Persist ordinary events immediately and batch adjacent tokens from the same agent.
|
||||
if self.db_session_factory:
|
||||
async with self._persist_lock:
|
||||
if event_type == "thinking_token":
|
||||
buffered, since = self._token_buffers.get(task_id, (None, time.monotonic()))
|
||||
if buffered and any(
|
||||
(buffered.get("metadata") or {}).get(key) != (metadata or {}).get(key)
|
||||
for key in ("agent_id", "agent_name")
|
||||
):
|
||||
await self._flush_tokens_unlocked(task_id)
|
||||
buffered, since = None, time.monotonic()
|
||||
token = (event_data.get("metadata") or {}).get("token", "")
|
||||
if buffered:
|
||||
token = buffered["metadata"].get("token", "") + token
|
||||
buffered = {**event_data, "metadata": {**(metadata or {}), "token": token}}
|
||||
self._token_buffers[task_id] = (buffered, since)
|
||||
if len(token) >= 1024 or time.monotonic() - since >= 0.2:
|
||||
await self._flush_tokens_unlocked(task_id)
|
||||
else:
|
||||
await self._flush_tokens_unlocked(task_id)
|
||||
await self._save_event_to_db(event_data)
|
||||
|
||||
# 推送到队列(非阻塞)
|
||||
if task_id in self._event_queues:
|
||||
@@ -335,6 +352,17 @@ class EventManager:
|
||||
|
||||
return event_id
|
||||
|
||||
async def _flush_tokens_unlocked(self, task_id):
|
||||
buffered = self._token_buffers.get(task_id)
|
||||
if buffered:
|
||||
await self._save_event_to_db(buffered[0])
|
||||
self._token_buffers.pop(task_id, None)
|
||||
|
||||
async def flush_tokens(self):
|
||||
async with self._persist_lock:
|
||||
for task_id in list(self._token_buffers):
|
||||
await self._flush_tokens_unlocked(task_id)
|
||||
|
||||
async def _save_event_to_db(self, event_data: Dict):
|
||||
"""保存事件到数据库"""
|
||||
from app.models.agent_task import AgentEvent
|
||||
@@ -493,11 +521,12 @@ class EventManager:
|
||||
logger.info(f"[StreamEvents] Task {task_id}: Entering real-time loop, queue size: {queue.qsize()}")
|
||||
|
||||
# 然后实时推送新事件
|
||||
heartbeat_interval = get_agent_config().sse_heartbeat_interval_seconds
|
||||
try:
|
||||
while True:
|
||||
try:
|
||||
logger.debug(f"[StreamEvents] Task {task_id}: Waiting for next event from queue...")
|
||||
event = await asyncio.wait_for(queue.get(), timeout=30)
|
||||
event = await asyncio.wait_for(queue.get(), timeout=heartbeat_interval)
|
||||
logger.debug(f"[StreamEvents] Task {task_id}: Got event from queue: {event.get('event_type')}")
|
||||
|
||||
# 🔥 过滤掉序列号 <= after_sequence 的事件
|
||||
@@ -544,4 +573,3 @@ class EventManager:
|
||||
self._event_callbacks.clear()
|
||||
|
||||
logger.debug("EventManager closed")
|
||||
|
||||
|
||||
@@ -65,9 +65,9 @@ class SandboxManager:
|
||||
try:
|
||||
import docker
|
||||
logger.info(f"🔄 Attempting to connect to Docker... (lib: {docker.__file__})")
|
||||
self._docker_client = docker.from_env()
|
||||
self._docker_client = await asyncio.to_thread(docker.from_env)
|
||||
# 测试连接
|
||||
self._docker_client.ping()
|
||||
await asyncio.to_thread(self._docker_client.ping)
|
||||
self._initialized = True
|
||||
self._init_error = None
|
||||
logger.info("✅ Docker sandbox manager initialized successfully")
|
||||
@@ -92,6 +92,24 @@ class SandboxManager:
|
||||
if self.is_available:
|
||||
return "Docker Service Available"
|
||||
return f"Docker Service Unavailable. Error: {self._init_error or 'Not initialized'}"
|
||||
|
||||
def _daemon_path(self, path: str) -> str:
|
||||
"""Resolve a worker-container path to the Docker daemon's volume source."""
|
||||
import socket
|
||||
from pathlib import Path
|
||||
resolved = Path(path).resolve()
|
||||
if not Path("/.dockerenv").exists():
|
||||
return str(resolved)
|
||||
container = self._docker_client.containers.get(socket.gethostname())
|
||||
mounts = sorted(container.attrs.get("Mounts", []),
|
||||
key=lambda mount: len(mount["Destination"]), reverse=True)
|
||||
for mount in mounts:
|
||||
try:
|
||||
relative = resolved.relative_to(mount["Destination"])
|
||||
return str(Path(mount["Source"]) / relative)
|
||||
except ValueError:
|
||||
continue
|
||||
raise ValueError("Sandbox project must be located on a shared worker volume")
|
||||
|
||||
async def execute_command(
|
||||
self,
|
||||
@@ -149,10 +167,9 @@ class SandboxManager:
|
||||
"network_mode": self.config.network_mode,
|
||||
"user": self.config.user,
|
||||
"read_only": self.config.read_only,
|
||||
"volumes": {
|
||||
temp_dir: {"bind": "/workspace", "mode": "rw"},
|
||||
},
|
||||
"labels": {"deepaudit.job_id": os.environ.get("DEEPAUDIT_JOB_ID", "local")},
|
||||
"tmpfs": {
|
||||
"/workspace": "rw,size=512m,mode=1777",
|
||||
"/home/sandbox": "rw,size=100m,mode=1777",
|
||||
"/tmp": "rw,size=100m,mode=1777"
|
||||
},
|
||||
@@ -280,13 +297,14 @@ class SandboxManager:
|
||||
"user": self.config.user,
|
||||
"read_only": self.config.read_only,
|
||||
"volumes": {
|
||||
host_workdir: {"bind": "/workspace", "mode": "ro"}, # 只读挂载项目代码
|
||||
self._daemon_path(host_workdir): {"bind": "/workspace", "mode": "ro"},
|
||||
},
|
||||
"tmpfs": {
|
||||
"/home/sandbox": "rw,size=100m,mode=1777",
|
||||
"/tmp": "rw,size=100m,mode=1777" # 添加 /tmp 目录供工具写入临时文件
|
||||
},
|
||||
"working_dir": "/workspace",
|
||||
"labels": {"deepaudit.job_id": os.environ.get("DEEPAUDIT_JOB_ID", "local")},
|
||||
"environment": container_env,
|
||||
}
|
||||
|
||||
@@ -1237,4 +1255,4 @@ sys.argv = ['script.py', '{test_command}']
|
||||
escaped_code = full_code.replace("'", "'\"'\"'")
|
||||
command = f"python3 -c '{escaped_code}'"
|
||||
|
||||
return await self.sandbox_manager.execute_command(command, timeout=30)
|
||||
return await self.sandbox_manager.execute_command(command, timeout=30)
|
||||
|
||||
@@ -0,0 +1 @@
|
||||
"""Agent audit application services, independent of HTTP routing."""
|
||||
@@ -0,0 +1,82 @@
|
||||
"""Replayable event stream shared by every API replica."""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import time
|
||||
|
||||
from sqlalchemy import select
|
||||
|
||||
from app.db.session import async_session_factory
|
||||
from app.models.agent_task import AgentEvent, AgentTask
|
||||
|
||||
|
||||
def event_payload(event):
|
||||
return {
|
||||
"id": event.id,
|
||||
"type": event.event_type,
|
||||
"event_type": event.event_type,
|
||||
"phase": event.phase,
|
||||
"message": event.message,
|
||||
"sequence": event.sequence,
|
||||
"timestamp": event.created_at.isoformat() if event.created_at else None,
|
||||
"tool_name": event.tool_name,
|
||||
"tool_input": event.tool_input,
|
||||
"tool_output": event.tool_output,
|
||||
"tool_duration_ms": event.tool_duration_ms,
|
||||
"finding_id": event.finding_id,
|
||||
"tokens_used": event.tokens_used,
|
||||
"metadata": event.event_metadata,
|
||||
"tool": {
|
||||
"name": event.tool_name,
|
||||
"input": event.tool_input,
|
||||
"output": event.tool_output,
|
||||
"duration_ms": event.tool_duration_ms,
|
||||
}
|
||||
if event.tool_name
|
||||
else None,
|
||||
}
|
||||
|
||||
|
||||
def encode_sse(payload):
|
||||
sequence = payload.get("sequence")
|
||||
event_id = f"id: {sequence}\n" if sequence is not None else ""
|
||||
return (
|
||||
f"{event_id}event: {payload['type']}\ndata: {json.dumps(payload, ensure_ascii=False)}\n\n"
|
||||
)
|
||||
|
||||
|
||||
async def stream_events(
|
||||
task_id, request, after_sequence=0, *, include_thinking=True, include_tool_calls=True
|
||||
):
|
||||
cursor = after_sequence
|
||||
last_heartbeat = time.monotonic()
|
||||
while not await request.is_disconnected():
|
||||
async with async_session_factory() as db:
|
||||
events = list(
|
||||
await db.scalars(
|
||||
select(AgentEvent)
|
||||
.where(
|
||||
AgentEvent.task_id == task_id,
|
||||
AgentEvent.sequence > cursor,
|
||||
)
|
||||
.order_by(AgentEvent.sequence)
|
||||
.limit(200)
|
||||
)
|
||||
)
|
||||
task = await db.get(AgentTask, task_id)
|
||||
status = task.status if task else "deleted"
|
||||
for event in events:
|
||||
cursor = event.sequence
|
||||
if not include_thinking and str(event.event_type).startswith("thinking"):
|
||||
continue
|
||||
if not include_tool_calls and str(event.event_type).startswith("tool"):
|
||||
continue
|
||||
yield encode_sse(event_payload(event))
|
||||
if not events and status in {"completed", "failed", "cancelled", "deleted"}:
|
||||
yield encode_sse({"type": "task_end", "status": status, "sequence": cursor})
|
||||
return
|
||||
if time.monotonic() - last_heartbeat >= 15:
|
||||
yield encode_sse({"type": "heartbeat"})
|
||||
last_heartbeat = time.monotonic()
|
||||
if len(events) < 200:
|
||||
await asyncio.sleep(0.25)
|
||||
@@ -0,0 +1,520 @@
|
||||
"""Execution services for agent audits."""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import os
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from sqlalchemy import func, select
|
||||
from sqlalchemy.orm import selectinload
|
||||
|
||||
from app.core.encryption import decrypt_sensitive_data
|
||||
from app.db.session import async_session_factory
|
||||
from app.models.agent_task import (
|
||||
AgentEvent,
|
||||
AgentTask,
|
||||
AgentTaskPhase,
|
||||
AgentTaskStatus,
|
||||
)
|
||||
from app.services.agent.agents.base import AgentResult
|
||||
from app.services.audit_queue import load_checkpoint, save_checkpoint
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
from .results import _calculate_security_score, _get_user_config, _save_agent_tree, _save_findings
|
||||
from .runtime import (
|
||||
_cancelled_tasks,
|
||||
_running_asyncio_tasks,
|
||||
_running_event_managers,
|
||||
_running_orchestrators,
|
||||
_running_tasks,
|
||||
is_task_cancelled,
|
||||
)
|
||||
from .tools import _collect_project_info, _initialize_tools
|
||||
from .workspace import _get_project_root
|
||||
|
||||
|
||||
async def _execute_agent_task(task_id: str, user_config_override=None):
|
||||
"""
|
||||
在后台执行 Agent 任务 - 使用动态 Agent 树架构
|
||||
|
||||
架构:OrchestratorAgent 作为大脑,动态调度子 Agent
|
||||
"""
|
||||
import time
|
||||
|
||||
from app.core.config import settings
|
||||
from app.services.agent.agents import (
|
||||
AnalysisAgent,
|
||||
OrchestratorAgent,
|
||||
ReconAgent,
|
||||
VerificationAgent,
|
||||
)
|
||||
from app.services.agent.core import agent_registry
|
||||
from app.services.agent.event_manager import AgentEventEmitter, EventManager
|
||||
from app.services.agent.tools import SandboxManager
|
||||
from app.services.llm.service import LLMService
|
||||
|
||||
# 🔥 在任务最开始就初始化 Docker 沙箱管理器
|
||||
# 这样可以确保整个任务生命周期内使用同一个管理器,并且尽早发现 Docker 问题
|
||||
logger.info(f"🚀 Starting execution for task {task_id}")
|
||||
sandbox_manager = SandboxManager()
|
||||
await sandbox_manager.initialize()
|
||||
logger.info(
|
||||
f"🐳 Global Sandbox Manager initialized (Available: {sandbox_manager.is_available})"
|
||||
)
|
||||
|
||||
# 🔥 提前创建事件管理器,以便在克隆仓库和索引时发送实时日志
|
||||
from app.services.agent.event_manager import AgentEventEmitter, EventManager
|
||||
|
||||
event_manager = EventManager(db_session_factory=async_session_factory)
|
||||
event_manager.create_queue(task_id)
|
||||
event_emitter = AgentEventEmitter(task_id, event_manager)
|
||||
async with async_session_factory() as event_db:
|
||||
event_emitter._sequence = (
|
||||
await event_db.scalar(
|
||||
select(func.max(AgentEvent.sequence)).where(AgentEvent.task_id == task_id)
|
||||
)
|
||||
) or 0
|
||||
_running_event_managers[task_id] = event_manager
|
||||
|
||||
async with async_session_factory() as db:
|
||||
orchestrator = None
|
||||
start_time = time.time()
|
||||
|
||||
try:
|
||||
# 获取任务
|
||||
task = await db.get(AgentTask, task_id, options=[selectinload(AgentTask.project)])
|
||||
if not task:
|
||||
logger.error(f"Task {task_id} not found")
|
||||
return
|
||||
|
||||
# 获取项目
|
||||
project = task.project
|
||||
if not project:
|
||||
logger.error(f"Project not found for task {task_id}")
|
||||
return
|
||||
|
||||
# 🔥 发送任务开始事件 - 使用 phase_start 让前端知道进入准备阶段
|
||||
await event_emitter.emit_phase_start("preparation", f"🚀 任务开始执行: {project.name}")
|
||||
|
||||
# 更新任务阶段为准备中
|
||||
task.status = AgentTaskStatus.RUNNING
|
||||
task.started_at = datetime.now(timezone.utc)
|
||||
task.current_phase = AgentTaskPhase.PLANNING # preparation 对应 PLANNING
|
||||
await db.commit()
|
||||
|
||||
# 获取用户配置(需要在获取项目根目录之前,以便传递 token)
|
||||
user_config = (
|
||||
user_config_override
|
||||
if user_config_override is not None
|
||||
else await _get_user_config(db, task.created_by)
|
||||
)
|
||||
|
||||
# 从用户配置中提取 token和SSH密钥(用于私有仓库克隆)
|
||||
other_config = (user_config or {}).get("otherConfig", {})
|
||||
github_token = other_config.get("githubToken") or settings.GITHUB_TOKEN
|
||||
gitlab_token = other_config.get("gitlabToken") or settings.GITLAB_TOKEN
|
||||
gitea_token = other_config.get("giteaToken") or settings.GITEA_TOKEN
|
||||
|
||||
# 解密SSH私钥
|
||||
ssh_private_key = None
|
||||
if "sshPrivateKey" in other_config:
|
||||
try:
|
||||
encrypted_key = other_config["sshPrivateKey"]
|
||||
ssh_private_key = decrypt_sensitive_data(encrypted_key)
|
||||
logger.info("成功解密SSH私钥")
|
||||
except Exception as e:
|
||||
logger.warning(f"解密SSH私钥失败: {e}")
|
||||
|
||||
# 获取项目根目录(传递任务指定的分支和认证 token/SSH密钥)
|
||||
# 🔥 传递 event_emitter 以发送克隆进度
|
||||
project_root = await _get_project_root(
|
||||
project,
|
||||
task_id,
|
||||
task.branch_name,
|
||||
github_token=github_token,
|
||||
gitlab_token=gitlab_token,
|
||||
gitea_token=gitea_token, # 🔥 新增
|
||||
ssh_private_key=ssh_private_key, # 🔥 新增SSH密钥
|
||||
event_emitter=event_emitter, # 🔥 新增
|
||||
)
|
||||
|
||||
# 🔥 自动修正 target_files 路径
|
||||
# 如果发生了目录调整(例如 ZIP 解压后只有一层目录,root 被下移),
|
||||
# 原有的 target_files (如 "Prefix/file.php") 可能无法匹配。
|
||||
# 我们需要检测并移除这些无效的前缀。
|
||||
if task.target_files and len(task.target_files) > 0:
|
||||
# 1. 检查是否存在不匹配的文件
|
||||
all_exist = True
|
||||
for tf in task.target_files:
|
||||
if not os.path.exists(os.path.join(project_root, tf)):
|
||||
all_exist = False
|
||||
break
|
||||
|
||||
if not all_exist:
|
||||
logger.info(f"Target files path mismatch detected in {project_root}")
|
||||
# 尝试通过路径匹配来修复
|
||||
# 获取当前根目录的名称
|
||||
root_name = os.path.basename(project_root)
|
||||
|
||||
new_target_files = []
|
||||
fixed_count = 0
|
||||
|
||||
for tf in task.target_files:
|
||||
# 检查文件是否以 root_name 开头(例如 "PHP-Project/index.php" 而 root 是 ".../PHP-Project")
|
||||
if tf.startswith(root_name + "/"):
|
||||
fixed_path = tf[len(root_name) + 1 :]
|
||||
if os.path.exists(os.path.join(project_root, fixed_path)):
|
||||
new_target_files.append(fixed_path)
|
||||
fixed_count += 1
|
||||
continue
|
||||
|
||||
# 如果上面的没匹配,尝试暴力搜索(只针对未找到的文件)
|
||||
# 这种情况比较少见,先保留原样或标记为丢失
|
||||
if os.path.exists(os.path.join(project_root, tf)):
|
||||
new_target_files.append(tf)
|
||||
else:
|
||||
# 尝试查看 tf 的 basename 是否在根目录直接存在(针对常见的最简情况)
|
||||
basename = os.path.basename(tf)
|
||||
if os.path.exists(os.path.join(project_root, basename)):
|
||||
new_target_files.append(basename)
|
||||
fixed_count += 1
|
||||
else:
|
||||
# 实在找不到,保留原样,让后续流程报错或忽略
|
||||
new_target_files.append(tf)
|
||||
|
||||
if fixed_count > 0:
|
||||
logger.info(f"🔧 Auto-fixed {fixed_count} target file paths")
|
||||
await event_emitter.emit_info(
|
||||
f"🔧 自动修正了 {fixed_count} 个目标文件的路径"
|
||||
)
|
||||
task.target_files = new_target_files
|
||||
|
||||
# 🔥 重新验证修正后的文件
|
||||
valid_target_files = []
|
||||
if task.target_files:
|
||||
for tf in task.target_files:
|
||||
if os.path.exists(os.path.join(project_root, tf)):
|
||||
valid_target_files.append(tf)
|
||||
else:
|
||||
logger.warning(f"⚠️ Target file not found: {tf}")
|
||||
|
||||
if not valid_target_files:
|
||||
raise ValueError("找不到指定的扫描文件,请检查文件路径后重试")
|
||||
elif len(valid_target_files) < len(task.target_files):
|
||||
logger.warning(
|
||||
f"⚠️ Partial target files missing. Found {len(valid_target_files)}/{len(task.target_files)}"
|
||||
)
|
||||
task.target_files = valid_target_files
|
||||
|
||||
logger.info(f"🚀 Task {task_id} started with Dynamic Agent Tree architecture")
|
||||
|
||||
# 🔥 获取项目根目录后检查取消
|
||||
if is_task_cancelled(task_id):
|
||||
logger.info(f"[Cancel] Task {task_id} cancelled after project preparation")
|
||||
raise asyncio.CancelledError("任务已取消")
|
||||
|
||||
checkpoint = await load_checkpoint(task_id)
|
||||
if checkpoint.get("result"):
|
||||
result = AgentResult(**checkpoint["result"])
|
||||
sub_agent_stats = checkpoint.get("sub_agent_stats", [])
|
||||
await event_emitter.emit_info("♻️ 恢复已完成的分析结果")
|
||||
else:
|
||||
# 创建 LLM 服务
|
||||
llm_service = LLMService(user_config=user_config)
|
||||
|
||||
# 初始化工具集 - 传递排除模式和目标文件以及预初始化的 sandbox_manager
|
||||
# 🔥 传递 event_emitter 以发送索引进度,传递 task_id 以支持取消
|
||||
tools = await _initialize_tools(
|
||||
project_root,
|
||||
llm_service,
|
||||
user_config,
|
||||
sandbox_manager=sandbox_manager,
|
||||
exclude_patterns=task.exclude_patterns,
|
||||
target_files=task.target_files,
|
||||
project_id=str(project.id), # 🔥 传递 project_id 用于 RAG
|
||||
event_emitter=event_emitter, # 🔥 新增
|
||||
task_id=task_id, # 🔥 新增:用于取消检查
|
||||
)
|
||||
|
||||
# 🔥 初始化工具后检查取消
|
||||
if is_task_cancelled(task_id):
|
||||
logger.info(f"[Cancel] Task {task_id} cancelled after tools initialization")
|
||||
raise asyncio.CancelledError("任务已取消")
|
||||
|
||||
# 创建子 Agent
|
||||
recon_agent = ReconAgent(
|
||||
llm_service=llm_service,
|
||||
tools=tools.get("recon", {}),
|
||||
event_emitter=event_emitter,
|
||||
)
|
||||
|
||||
analysis_agent = AnalysisAgent(
|
||||
llm_service=llm_service,
|
||||
tools=tools.get("analysis", {}),
|
||||
event_emitter=event_emitter,
|
||||
)
|
||||
|
||||
verification_agent = VerificationAgent(
|
||||
llm_service=llm_service,
|
||||
tools=tools.get("verification", {}),
|
||||
event_emitter=event_emitter,
|
||||
)
|
||||
|
||||
# 创建 Orchestrator Agent
|
||||
orchestrator = OrchestratorAgent(
|
||||
llm_service=llm_service,
|
||||
tools=tools.get("orchestrator", {}),
|
||||
event_emitter=event_emitter,
|
||||
sub_agents={
|
||||
"recon": recon_agent,
|
||||
"analysis": analysis_agent,
|
||||
"verification": verification_agent,
|
||||
},
|
||||
)
|
||||
|
||||
# 🔥 设置外部取消检查回调
|
||||
# 这确保即使 runner.cancel() 失败,Agent 也能通过 checking 全局标志感知取消
|
||||
def check_global_cancel():
|
||||
return is_task_cancelled(task_id)
|
||||
|
||||
orchestrator.set_cancel_callback(check_global_cancel)
|
||||
# 同时也为子 Agent 设置(虽然 Orchestrator 会传播)
|
||||
recon_agent.set_cancel_callback(check_global_cancel)
|
||||
analysis_agent.set_cancel_callback(check_global_cancel)
|
||||
verification_agent.set_cancel_callback(check_global_cancel)
|
||||
|
||||
# 注册到全局
|
||||
_running_orchestrators[task_id] = orchestrator
|
||||
_running_tasks[task_id] = orchestrator # 兼容旧的取消逻辑
|
||||
_running_event_managers[task_id] = event_manager # 用于 SSE 流
|
||||
|
||||
# 🔥 清理旧的 Agent 注册表,避免显示多个树
|
||||
from app.services.agent.core import agent_registry
|
||||
|
||||
agent_registry.clear()
|
||||
|
||||
# 注册 Orchestrator 到 Agent Registry(使用其内置方法)
|
||||
orchestrator._register_to_registry(task="Root orchestrator for security audit")
|
||||
|
||||
await event_emitter.emit_info("🧠 动态 Agent 树架构启动")
|
||||
await event_emitter.emit_info(f"📁 项目路径: {project_root}")
|
||||
|
||||
# 收集项目信息 - 传递排除模式和目标文件
|
||||
project_info = await _collect_project_info(
|
||||
project_root,
|
||||
project.name,
|
||||
exclude_patterns=task.exclude_patterns,
|
||||
target_files=task.target_files,
|
||||
)
|
||||
|
||||
# 更新任务文件统计
|
||||
task.total_files = project_info.get("file_count", 0)
|
||||
await db.commit()
|
||||
|
||||
# 构建输入数据
|
||||
input_data = {
|
||||
"project_info": project_info,
|
||||
"config": {
|
||||
"target_vulnerabilities": task.target_vulnerabilities or [],
|
||||
"verification_level": task.verification_level or "sandbox",
|
||||
"exclude_patterns": task.exclude_patterns or [],
|
||||
"target_files": task.target_files or [],
|
||||
"max_iterations": task.max_iterations or 50,
|
||||
},
|
||||
"project_root": project_root,
|
||||
"task_id": task_id,
|
||||
}
|
||||
|
||||
# 执行 Orchestrator
|
||||
await event_emitter.emit_phase_start(
|
||||
"orchestration", "🎯 Orchestrator 开始编排审计流程"
|
||||
)
|
||||
task.current_phase = AgentTaskPhase.ANALYSIS
|
||||
await db.commit()
|
||||
|
||||
# 🔥 将 orchestrator.run() 包装在 asyncio.Task 中,以便可以强制取消
|
||||
run_task = asyncio.create_task(orchestrator.run(input_data))
|
||||
_running_asyncio_tasks[task_id] = run_task
|
||||
|
||||
try:
|
||||
result = await run_task
|
||||
finally:
|
||||
_running_asyncio_tasks.pop(task_id, None)
|
||||
|
||||
sub_agent_stats = [
|
||||
agent.get_stats()
|
||||
for agent in orchestrator.sub_agents.values()
|
||||
if hasattr(agent, "get_stats")
|
||||
]
|
||||
if result.success:
|
||||
await save_checkpoint(
|
||||
task_id, result=result.to_dict(), sub_agent_stats=sub_agent_stats
|
||||
)
|
||||
|
||||
# 处理结果
|
||||
duration_ms = int((time.time() - start_time) * 1000)
|
||||
|
||||
await db.refresh(task)
|
||||
|
||||
if result.success:
|
||||
# 🔥 CRITICAL FIX: Log and save findings with detailed debugging
|
||||
findings = result.data.get("findings", [])
|
||||
logger.info(
|
||||
f"[AgentTask] Task {task_id} completed with {len(findings)} findings from Orchestrator"
|
||||
)
|
||||
|
||||
# 🔥 Debug: Log each finding for verification
|
||||
for i, f in enumerate(findings[:5]): # Log first 5
|
||||
if isinstance(f, dict):
|
||||
logger.debug(
|
||||
f"[AgentTask] Finding {i + 1}: {f.get('title', 'N/A')[:50]} - {f.get('severity', 'N/A')}"
|
||||
)
|
||||
|
||||
# 🔥 v2.1: 传递 project_root 用于文件路径验证
|
||||
saved_count = await _save_findings(db, task_id, findings, project_root=project_root)
|
||||
logger.info(
|
||||
f"[AgentTask] Saved {saved_count}/{len(findings)} findings (filtered {len(findings) - saved_count} hallucinations)"
|
||||
)
|
||||
|
||||
# 更新任务统计
|
||||
# 🔥 CRITICAL FIX: 在设置完成前再次检查取消状态
|
||||
# 避免 "取消后后端继续运行并最终标记为完成" 的问题
|
||||
if is_task_cancelled(task_id):
|
||||
logger.info(
|
||||
f"[AgentTask] Task {task_id} was cancelled, overriding success result"
|
||||
)
|
||||
task.status = AgentTaskStatus.CANCELLED
|
||||
else:
|
||||
task.status = AgentTaskStatus.COMPLETED
|
||||
task.completed_at = datetime.now(timezone.utc)
|
||||
task.current_phase = AgentTaskPhase.REPORTING
|
||||
task.findings_count = saved_count # 🔥 v2.1: 使用实际保存的数量(排除幻觉)
|
||||
|
||||
# 🔥 CRITICAL FIX: 累加所有子 Agent 的统计,而不仅仅是 Orchestrator 的
|
||||
total_iterations = result.iterations
|
||||
tool_calls_count = result.tool_calls
|
||||
tokens_used = result.tokens_used
|
||||
|
||||
for sub_stats in sub_agent_stats:
|
||||
total_iterations += sub_stats.get("iterations", 0)
|
||||
tool_calls_count += sub_stats.get("tool_calls", 0)
|
||||
tokens_used += sub_stats.get("tokens_used", 0)
|
||||
|
||||
task.total_iterations = total_iterations
|
||||
task.tool_calls_count = tool_calls_count
|
||||
task.tokens_used = tokens_used
|
||||
|
||||
# 🔥 统计文件数量
|
||||
# analyzed_files = 实际扫描过的文件数(任务完成时等于 total_files)
|
||||
# files_with_findings = 有漏洞发现的唯一文件数
|
||||
task.analyzed_files = task.total_files # Agent 扫描了所有符合条件的文件
|
||||
|
||||
files_with_findings_set = set()
|
||||
for f in findings:
|
||||
if isinstance(f, dict):
|
||||
file_path = (
|
||||
f.get("file_path")
|
||||
or f.get("file")
|
||||
or f.get("location", "").split(":")[0]
|
||||
)
|
||||
if file_path:
|
||||
files_with_findings_set.add(file_path)
|
||||
task.files_with_findings = len(files_with_findings_set)
|
||||
|
||||
# 统计严重程度和验证状态
|
||||
verified_count = 0
|
||||
task.critical_count = task.high_count = task.medium_count = task.low_count = 0
|
||||
for f in findings:
|
||||
if isinstance(f, dict):
|
||||
sev = str(f.get("severity", "low")).lower()
|
||||
if sev == "critical":
|
||||
task.critical_count += 1
|
||||
elif sev == "high":
|
||||
task.high_count += 1
|
||||
elif sev == "medium":
|
||||
task.medium_count += 1
|
||||
elif sev == "low":
|
||||
task.low_count += 1
|
||||
# 🔥 统计已验证的发现
|
||||
if f.get("is_verified") or f.get("verdict") == "confirmed":
|
||||
verified_count += 1
|
||||
task.verified_count = verified_count
|
||||
|
||||
# 计算安全评分
|
||||
task.security_score = _calculate_security_score(findings)
|
||||
task.quality_score = _calculate_security_score(findings)
|
||||
# 🔥 注意: progress_percentage 是计算属性,不需要手动设置
|
||||
# 当 status = COMPLETED 时会自动返回 100.0
|
||||
|
||||
await db.commit()
|
||||
|
||||
await event_emitter.emit_task_complete(
|
||||
findings_count=len(findings),
|
||||
duration_ms=duration_ms,
|
||||
)
|
||||
|
||||
logger.info(
|
||||
f"✅ Task {task_id} completed: {len(findings)} findings, {duration_ms}ms"
|
||||
)
|
||||
else:
|
||||
# 🔥 检查是否是取消导致的失败
|
||||
if result.error == "任务已取消":
|
||||
# 状态可能已经被 cancel API 更新,只需确保一致性
|
||||
if task.status != AgentTaskStatus.CANCELLED:
|
||||
task.status = AgentTaskStatus.CANCELLED
|
||||
task.completed_at = datetime.now(timezone.utc)
|
||||
await db.commit()
|
||||
logger.info(f"🛑 Task {task_id} cancelled")
|
||||
else:
|
||||
task.status = AgentTaskStatus.FAILED
|
||||
task.error_message = result.error or "Unknown error"
|
||||
task.completed_at = datetime.now(timezone.utc)
|
||||
await db.commit()
|
||||
|
||||
await event_emitter.emit_error(result.error or "Unknown error")
|
||||
logger.error(f"❌ Task {task_id} failed: {result.error}")
|
||||
|
||||
except asyncio.CancelledError:
|
||||
logger.info(f"Task {task_id} cancelled")
|
||||
try:
|
||||
task = await db.get(AgentTask, task_id)
|
||||
if task:
|
||||
task.status = AgentTaskStatus.CANCELLED
|
||||
task.completed_at = datetime.now(timezone.utc)
|
||||
await db.commit()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Task {task_id} failed: {e}", exc_info=True)
|
||||
|
||||
try:
|
||||
task = await db.get(AgentTask, task_id)
|
||||
if task:
|
||||
task.status = AgentTaskStatus.FAILED
|
||||
task.error_message = str(e)[:1000]
|
||||
task.completed_at = datetime.now(timezone.utc)
|
||||
await db.commit()
|
||||
except Exception as db_error:
|
||||
logger.error(f"Failed to update task status: {db_error}")
|
||||
|
||||
finally:
|
||||
# 🔥 在清理之前保存 Agent 树到数据库
|
||||
try:
|
||||
async with async_session_factory() as save_db:
|
||||
await _save_agent_tree(save_db, task_id)
|
||||
except Exception as save_error:
|
||||
logger.error(f"Failed to save agent tree: {save_error}")
|
||||
|
||||
# 清理
|
||||
_running_orchestrators.pop(task_id, None)
|
||||
_running_tasks.pop(task_id, None)
|
||||
_running_event_managers.pop(task_id, None)
|
||||
_running_asyncio_tasks.pop(task_id, None) # 🔥 清理 asyncio task
|
||||
await event_manager.flush_tokens()
|
||||
_cancelled_tasks.discard(task_id) # 🔥 清理取消标志
|
||||
|
||||
# 🔥 清理整个 Agent 注册表(包括所有子 Agent)
|
||||
agent_registry.clear()
|
||||
|
||||
logger.debug(f"Task {task_id} cleaned up")
|
||||
@@ -0,0 +1,69 @@
|
||||
"""Publish live worker state for API replicas without sharing process memory."""
|
||||
|
||||
from sqlalchemy import update
|
||||
|
||||
from app.db.session import AsyncSessionLocal
|
||||
from app.models.agent_task import AgentTask
|
||||
from app.models.audit_job import AuditJob
|
||||
|
||||
from .runtime import _running_orchestrators
|
||||
|
||||
|
||||
async def publish_progress(task_id):
|
||||
orchestrator = _running_orchestrators.get(task_id)
|
||||
if not orchestrator:
|
||||
return
|
||||
from app.services.agent.core import agent_registry
|
||||
|
||||
tree = agent_registry.get_agent_tree()
|
||||
totals = {"iterations": 0, "tool_calls": 0, "tokens_used": 0}
|
||||
nodes = []
|
||||
for agent_id, node in tree.get("nodes", {}).items():
|
||||
instance = agent_registry.get_agent(agent_id)
|
||||
stats = instance.get_stats() if instance and hasattr(instance, "get_stats") else {}
|
||||
for key in totals:
|
||||
totals[key] += stats.get(key, 0)
|
||||
result = node.get("result") or {}
|
||||
nodes.append(
|
||||
dict(
|
||||
id=agent_id,
|
||||
agent_id=agent_id,
|
||||
agent_name=node.get("name", "Unknown"),
|
||||
agent_type=node.get("type", "unknown"),
|
||||
parent_agent_id=node.get("parent_id"),
|
||||
task_description=node.get("task"),
|
||||
knowledge_modules=node.get("knowledge_modules") or [],
|
||||
status=node.get("status", "unknown"),
|
||||
findings_count=len(result.get("findings", [])),
|
||||
iterations=stats.get("iterations", 0),
|
||||
tool_calls=stats.get("tool_calls", 0),
|
||||
tokens_used=stats.get("tokens_used", 0),
|
||||
children=[],
|
||||
)
|
||||
)
|
||||
async with AsyncSessionLocal() as db:
|
||||
job = await db.get(AuditJob, task_id, with_for_update=True)
|
||||
if not job:
|
||||
return
|
||||
task = await db.get(AgentTask, task_id)
|
||||
snapshot = dict(
|
||||
task_id=task_id,
|
||||
root_agent_id=tree.get("root_agent_id"),
|
||||
nodes=nodes,
|
||||
total_agents=len(nodes),
|
||||
total_findings=task.findings_count or 0,
|
||||
running_agents=sum(n["status"] == "running" for n in nodes),
|
||||
completed_agents=sum(n["status"] == "completed" for n in nodes),
|
||||
failed_agents=sum(n["status"] == "failed" for n in nodes),
|
||||
)
|
||||
job.checkpoint = {**(job.checkpoint or {}), "live_tree": snapshot}
|
||||
await db.execute(
|
||||
update(AgentTask)
|
||||
.where(AgentTask.id == task_id)
|
||||
.values(
|
||||
total_iterations=totals["iterations"],
|
||||
tool_calls_count=totals["tool_calls"],
|
||||
tokens_used=totals["tokens_used"],
|
||||
)
|
||||
)
|
||||
await db.commit()
|
||||
@@ -0,0 +1,446 @@
|
||||
"""Results services for agent audits."""
|
||||
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
from typing import Any, Dict, List, Optional
|
||||
from uuid import UUID, uuid4, uuid5
|
||||
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
from sqlalchemy.future import select
|
||||
|
||||
from app.models.agent_task import (
|
||||
AgentFinding,
|
||||
FindingStatus,
|
||||
VulnerabilitySeverity,
|
||||
)
|
||||
from app.models.user_config import UserConfig
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
async def _get_user_config(db: AsyncSession, user_id: Optional[str]) -> Optional[Dict[str, Any]]:
|
||||
"""获取用户配置"""
|
||||
if not user_id:
|
||||
return None
|
||||
|
||||
try:
|
||||
from app.api.v1.endpoints.config import (
|
||||
SENSITIVE_LLM_FIELDS,
|
||||
SENSITIVE_OTHER_FIELDS,
|
||||
decrypt_config,
|
||||
)
|
||||
|
||||
result = await db.execute(select(UserConfig).where(UserConfig.user_id == user_id))
|
||||
config = result.scalar_one_or_none()
|
||||
|
||||
if config and config.llm_config:
|
||||
user_llm_config = json.loads(config.llm_config) if config.llm_config else {}
|
||||
user_other_config = json.loads(config.other_config) if config.other_config else {}
|
||||
|
||||
user_llm_config = decrypt_config(user_llm_config, SENSITIVE_LLM_FIELDS)
|
||||
user_other_config = decrypt_config(user_other_config, SENSITIVE_OTHER_FIELDS)
|
||||
|
||||
return {
|
||||
"llmConfig": user_llm_config,
|
||||
"otherConfig": user_other_config,
|
||||
}
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to get user config: {e}")
|
||||
|
||||
return None
|
||||
|
||||
|
||||
async def _save_findings(
|
||||
db: AsyncSession,
|
||||
task_id: str,
|
||||
findings: List[Dict],
|
||||
project_root: Optional[str] = None,
|
||||
) -> int:
|
||||
"""
|
||||
保存发现到数据库
|
||||
|
||||
🔥 增强版:支持多种 Agent 输出格式,健壮的字段映射
|
||||
🔥 v2.1: 添加文件路径验证,过滤幻觉发现
|
||||
|
||||
Args:
|
||||
db: 数据库会话
|
||||
task_id: 任务ID
|
||||
findings: 发现列表
|
||||
project_root: 项目根目录(用于验证文件路径)
|
||||
|
||||
Returns:
|
||||
int: 实际保存的发现数量
|
||||
"""
|
||||
from app.models.agent_task import VulnerabilityType
|
||||
|
||||
logger.info(f"[SaveFindings] Starting to save {len(findings)} findings for task {task_id}")
|
||||
|
||||
if not findings:
|
||||
logger.warning(f"[SaveFindings] No findings to save for task {task_id}")
|
||||
return 0
|
||||
|
||||
# 🔥 Case-insensitive mapping preparation
|
||||
severity_map = {
|
||||
"critical": VulnerabilitySeverity.CRITICAL,
|
||||
"high": VulnerabilitySeverity.HIGH,
|
||||
"medium": VulnerabilitySeverity.MEDIUM,
|
||||
"low": VulnerabilitySeverity.LOW,
|
||||
"info": VulnerabilitySeverity.INFO,
|
||||
}
|
||||
|
||||
type_map = {
|
||||
"sql_injection": VulnerabilityType.SQL_INJECTION,
|
||||
"nosql_injection": VulnerabilityType.NOSQL_INJECTION,
|
||||
"xss": VulnerabilityType.XSS,
|
||||
"command_injection": VulnerabilityType.COMMAND_INJECTION,
|
||||
"code_injection": VulnerabilityType.CODE_INJECTION,
|
||||
"path_traversal": VulnerabilityType.PATH_TRAVERSAL,
|
||||
"ssrf": VulnerabilityType.SSRF,
|
||||
"xxe": VulnerabilityType.XXE,
|
||||
"auth_bypass": VulnerabilityType.AUTH_BYPASS,
|
||||
"idor": VulnerabilityType.IDOR,
|
||||
"sensitive_data_exposure": VulnerabilityType.SENSITIVE_DATA_EXPOSURE,
|
||||
"hardcoded_secret": VulnerabilityType.HARDCODED_SECRET,
|
||||
"deserialization": VulnerabilityType.DESERIALIZATION,
|
||||
"weak_crypto": VulnerabilityType.WEAK_CRYPTO,
|
||||
"file_inclusion": VulnerabilityType.FILE_INCLUSION,
|
||||
"race_condition": VulnerabilityType.RACE_CONDITION,
|
||||
"business_logic": VulnerabilityType.BUSINESS_LOGIC,
|
||||
"memory_corruption": VulnerabilityType.MEMORY_CORRUPTION,
|
||||
}
|
||||
|
||||
existing_ids = set(
|
||||
await db.scalars(select(AgentFinding.id).where(AgentFinding.task_id == task_id))
|
||||
)
|
||||
seen_ids = set()
|
||||
saved_count = 0
|
||||
logger.info(f"Saving {len(findings)} findings for task {task_id}")
|
||||
|
||||
for finding in findings:
|
||||
if not isinstance(finding, dict):
|
||||
logger.debug(f"[SaveFindings] Skipping non-dict finding: {type(finding)}")
|
||||
continue
|
||||
|
||||
try:
|
||||
# 🔥 Handle severity (case-insensitive, support multiple field names)
|
||||
raw_severity = (
|
||||
str(finding.get("severity") or finding.get("risk") or "medium").lower().strip()
|
||||
)
|
||||
severity_enum = severity_map.get(raw_severity, VulnerabilitySeverity.MEDIUM)
|
||||
|
||||
# 🔥 Handle vulnerability type (case-insensitive & snake_case normalization)
|
||||
# Support multiple field names: vulnerability_type, type, vuln_type
|
||||
raw_type = (
|
||||
str(
|
||||
finding.get("vulnerability_type")
|
||||
or finding.get("type")
|
||||
or finding.get("vuln_type")
|
||||
or "other"
|
||||
)
|
||||
.lower()
|
||||
.strip()
|
||||
.replace(" ", "_")
|
||||
.replace("-", "_")
|
||||
)
|
||||
|
||||
type_enum = type_map.get(raw_type, VulnerabilityType.OTHER)
|
||||
|
||||
# 🔥 Additional fallback for common Agent output variations
|
||||
if "sqli" in raw_type or "sql" in raw_type:
|
||||
type_enum = VulnerabilityType.SQL_INJECTION
|
||||
if "xss" in raw_type:
|
||||
type_enum = VulnerabilityType.XSS
|
||||
if "rce" in raw_type or "command" in raw_type or "cmd" in raw_type:
|
||||
type_enum = VulnerabilityType.COMMAND_INJECTION
|
||||
if "traversal" in raw_type or "lfi" in raw_type or "rfi" in raw_type:
|
||||
type_enum = VulnerabilityType.PATH_TRAVERSAL
|
||||
if "ssrf" in raw_type:
|
||||
type_enum = VulnerabilityType.SSRF
|
||||
if "xxe" in raw_type:
|
||||
type_enum = VulnerabilityType.XXE
|
||||
if "auth" in raw_type:
|
||||
type_enum = VulnerabilityType.AUTH_BYPASS
|
||||
if "secret" in raw_type or "credential" in raw_type or "password" in raw_type:
|
||||
type_enum = VulnerabilityType.HARDCODED_SECRET
|
||||
if "deserial" in raw_type:
|
||||
type_enum = VulnerabilityType.DESERIALIZATION
|
||||
|
||||
# 🔥 Handle file path (support multiple field names)
|
||||
file_path = (
|
||||
finding.get("file_path")
|
||||
or finding.get("file")
|
||||
or finding.get("location", "").split(":")[0]
|
||||
if ":" in finding.get("location", "")
|
||||
else finding.get("location")
|
||||
)
|
||||
|
||||
# 🔥 v2.1: 文件路径验证 - 过滤幻觉发现
|
||||
if project_root and file_path:
|
||||
# 清理路径(移除可能的行号)
|
||||
clean_path = (
|
||||
file_path.split(":")[0].strip() if ":" in file_path else file_path.strip()
|
||||
)
|
||||
full_path = os.path.join(project_root, clean_path)
|
||||
|
||||
if not os.path.isfile(full_path):
|
||||
# 尝试作为绝对路径
|
||||
if not (os.path.isabs(clean_path) and os.path.isfile(clean_path)):
|
||||
logger.warning(
|
||||
f"[SaveFindings] 🚫 跳过幻觉发现: 文件不存在 '{file_path}' "
|
||||
f"(title: {finding.get('title', 'N/A')[:50]})"
|
||||
)
|
||||
continue # 跳过这个发现
|
||||
|
||||
# 🔥 Handle line numbers (support multiple formats)
|
||||
line_start = finding.get("line_start") or finding.get("line")
|
||||
if not line_start and ":" in finding.get("location", ""):
|
||||
try:
|
||||
line_start = int(finding.get("location", "").split(":")[1])
|
||||
except (ValueError, IndexError):
|
||||
line_start = None
|
||||
|
||||
line_end = finding.get("line_end") or line_start
|
||||
|
||||
# 🔥 Handle code snippet (support multiple field names)
|
||||
code_snippet = (
|
||||
finding.get("code_snippet") or finding.get("code") or finding.get("vulnerable_code")
|
||||
)
|
||||
|
||||
# 🔥 Handle title (generate from type if not provided)
|
||||
title = finding.get("title")
|
||||
if not title:
|
||||
# Generate title from vulnerability type and file
|
||||
type_display = raw_type.replace("_", " ").title()
|
||||
if file_path:
|
||||
title = f"{type_display} in {os.path.basename(file_path)}"
|
||||
else:
|
||||
title = f"{type_display} Vulnerability"
|
||||
|
||||
# 🔥 Handle description (support multiple field names)
|
||||
description = (
|
||||
finding.get("description")
|
||||
or finding.get("details")
|
||||
or finding.get("explanation")
|
||||
or finding.get("impact")
|
||||
or ""
|
||||
)
|
||||
|
||||
# 🔥 Handle suggestion/recommendation
|
||||
suggestion = (
|
||||
finding.get("suggestion")
|
||||
or finding.get("recommendation")
|
||||
or finding.get("remediation")
|
||||
or finding.get("fix")
|
||||
)
|
||||
|
||||
# 🔥 Handle confidence (map to ai_confidence field in model)
|
||||
confidence = finding.get("confidence") or finding.get("ai_confidence") or 0.5
|
||||
if isinstance(confidence, str):
|
||||
try:
|
||||
confidence = float(confidence)
|
||||
except ValueError:
|
||||
confidence = 0.5
|
||||
|
||||
# 🔥 Handle verification status
|
||||
is_verified = finding.get("is_verified", False)
|
||||
if finding.get("verdict") == "confirmed":
|
||||
is_verified = True
|
||||
|
||||
# 🔥 Handle PoC information
|
||||
poc_data = finding.get("poc", {})
|
||||
has_poc = bool(poc_data)
|
||||
poc_code = None
|
||||
poc_description = None
|
||||
poc_steps = None
|
||||
|
||||
if isinstance(poc_data, dict):
|
||||
poc_description = poc_data.get("description")
|
||||
poc_steps = poc_data.get("steps")
|
||||
poc_code = poc_data.get("payload") or poc_data.get("code")
|
||||
elif isinstance(poc_data, str):
|
||||
poc_description = poc_data
|
||||
|
||||
# 🔥 Handle verification details
|
||||
verification_method = finding.get("verification_method")
|
||||
verification_result = None
|
||||
if finding.get("verification_details"):
|
||||
verification_result = {"details": finding.get("verification_details")}
|
||||
|
||||
# 🔥 Handle CWE and CVSS
|
||||
cwe_id = finding.get("cwe_id") or finding.get("cwe")
|
||||
cvss_score = finding.get("cvss_score") or finding.get("cvss")
|
||||
if isinstance(cvss_score, str):
|
||||
try:
|
||||
cvss_score = float(cvss_score)
|
||||
except ValueError:
|
||||
cvss_score = None
|
||||
|
||||
identity = json.dumps([file_path, line_start, title, type_enum], ensure_ascii=False)
|
||||
finding_id = str(uuid5(UUID(task_id), identity))
|
||||
if finding_id in seen_ids:
|
||||
continue
|
||||
seen_ids.add(finding_id)
|
||||
if finding_id in existing_ids:
|
||||
saved_count += 1
|
||||
continue
|
||||
db_finding = AgentFinding(
|
||||
id=finding_id,
|
||||
task_id=task_id,
|
||||
vulnerability_type=type_enum,
|
||||
severity=severity_enum,
|
||||
title=title[:500] if title else "Unknown Vulnerability",
|
||||
description=description[:5000] if description else "",
|
||||
file_path=file_path[:500] if file_path else None,
|
||||
line_start=line_start,
|
||||
line_end=line_end,
|
||||
code_snippet=code_snippet[:10000] if code_snippet else None,
|
||||
suggestion=suggestion[:5000] if suggestion else None,
|
||||
is_verified=is_verified,
|
||||
ai_confidence=confidence, # 🔥 FIX: Use ai_confidence, not confidence
|
||||
status=FindingStatus.VERIFIED if is_verified else FindingStatus.NEW,
|
||||
# 🔥 Additional fields
|
||||
has_poc=has_poc,
|
||||
poc_code=poc_code,
|
||||
poc_description=poc_description,
|
||||
poc_steps=poc_steps,
|
||||
verification_method=verification_method,
|
||||
verification_result=verification_result,
|
||||
cvss_score=cvss_score,
|
||||
# References for CWE
|
||||
references=[{"cwe": cwe_id}] if cwe_id else None,
|
||||
)
|
||||
db.add(db_finding)
|
||||
saved_count += 1
|
||||
logger.debug(f"[SaveFindings] Prepared finding: {title[:50]}... ({severity_enum})")
|
||||
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to save finding: {e}, data: {finding}")
|
||||
import traceback
|
||||
|
||||
logger.debug(f"[SaveFindings] Traceback: {traceback.format_exc()}")
|
||||
|
||||
logger.info(f"Successfully prepared {saved_count} findings for commit")
|
||||
|
||||
try:
|
||||
await db.commit()
|
||||
logger.info(f"[SaveFindings] Successfully committed {saved_count} findings to database")
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to commit findings: {e}")
|
||||
await db.rollback()
|
||||
raise
|
||||
|
||||
return saved_count
|
||||
|
||||
|
||||
def _calculate_security_score(findings: List[Dict]) -> float:
|
||||
"""计算安全评分"""
|
||||
if not findings:
|
||||
return 100.0
|
||||
|
||||
# 基于发现的严重程度计算扣分
|
||||
deductions = {
|
||||
"critical": 25,
|
||||
"high": 15,
|
||||
"medium": 8,
|
||||
"low": 3,
|
||||
"info": 1,
|
||||
}
|
||||
|
||||
total_deduction = 0
|
||||
for f in findings:
|
||||
if isinstance(f, dict):
|
||||
sev = f.get("severity", "low")
|
||||
total_deduction += deductions.get(sev, 3)
|
||||
|
||||
score = max(0, 100 - total_deduction)
|
||||
return float(score)
|
||||
|
||||
|
||||
async def _save_agent_tree(db: AsyncSession, task_id: str) -> None:
|
||||
"""
|
||||
保存 Agent 树到数据库
|
||||
|
||||
🔥 在任务完成前调用,将内存中的 Agent 树持久化到数据库
|
||||
"""
|
||||
from app.models.agent_task import AgentTreeNode
|
||||
from app.services.agent.core import agent_registry
|
||||
|
||||
try:
|
||||
tree = agent_registry.get_agent_tree()
|
||||
nodes = tree.get("nodes", {})
|
||||
|
||||
if not nodes:
|
||||
logger.warning(f"[SaveAgentTree] No agent nodes to save for task {task_id}")
|
||||
return
|
||||
|
||||
logger.info(f"[SaveAgentTree] Saving {len(nodes)} agent nodes for task {task_id}")
|
||||
|
||||
# 计算每个节点的深度
|
||||
def get_depth(agent_id: str, visited: set = None) -> int:
|
||||
if visited is None:
|
||||
visited = set()
|
||||
if agent_id in visited:
|
||||
return 0
|
||||
visited.add(agent_id)
|
||||
node = nodes.get(agent_id)
|
||||
if not node:
|
||||
return 0
|
||||
parent_id = node.get("parent_id")
|
||||
if not parent_id:
|
||||
return 0
|
||||
return 1 + get_depth(parent_id, visited)
|
||||
|
||||
from sqlalchemy import delete
|
||||
|
||||
await db.execute(delete(AgentTreeNode).where(AgentTreeNode.task_id == task_id))
|
||||
saved_count = 0
|
||||
for agent_id, node_data in nodes.items():
|
||||
# 获取 Agent 实例的统计数据
|
||||
agent_instance = agent_registry.get_agent(agent_id)
|
||||
iterations = 0
|
||||
tool_calls = 0
|
||||
tokens_used = 0
|
||||
|
||||
if agent_instance and hasattr(agent_instance, "get_stats"):
|
||||
stats = agent_instance.get_stats()
|
||||
iterations = stats.get("iterations", 0)
|
||||
tool_calls = stats.get("tool_calls", 0)
|
||||
tokens_used = stats.get("tokens_used", 0)
|
||||
|
||||
# 从结果中获取发现数量
|
||||
findings_count = 0
|
||||
result_summary = None
|
||||
if node_data.get("result"):
|
||||
result = node_data.get("result", {})
|
||||
if isinstance(result, dict):
|
||||
findings_count = len(result.get("findings", []))
|
||||
if result.get("summary"):
|
||||
result_summary = str(result.get("summary"))[:2000]
|
||||
|
||||
tree_node = AgentTreeNode(
|
||||
id=str(uuid4()),
|
||||
task_id=task_id,
|
||||
agent_id=agent_id,
|
||||
agent_name=node_data.get("name", "Unknown"),
|
||||
agent_type=node_data.get("type", "unknown"),
|
||||
parent_agent_id=node_data.get("parent_id"),
|
||||
depth=get_depth(agent_id),
|
||||
task_description=node_data.get("task"),
|
||||
knowledge_modules=node_data.get("knowledge_modules"),
|
||||
status=node_data.get("status", "unknown"),
|
||||
result_summary=result_summary,
|
||||
findings_count=findings_count,
|
||||
iterations=iterations,
|
||||
tool_calls=tool_calls,
|
||||
tokens_used=tokens_used,
|
||||
)
|
||||
db.add(tree_node)
|
||||
saved_count += 1
|
||||
|
||||
await db.commit()
|
||||
logger.info(f"[SaveAgentTree] Successfully saved {saved_count} agent nodes to database")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"[SaveAgentTree] Failed to save agent tree: {e}", exc_info=True)
|
||||
await db.rollback()
|
||||
@@ -0,0 +1,11 @@
|
||||
"""State local to one isolated audit worker process."""
|
||||
|
||||
_running_tasks = {}
|
||||
_running_asyncio_tasks = {}
|
||||
_running_orchestrators = {}
|
||||
_running_event_managers = {}
|
||||
_cancelled_tasks = set()
|
||||
|
||||
|
||||
def is_task_cancelled(task_id):
|
||||
return task_id in _cancelled_tasks
|
||||
@@ -0,0 +1,238 @@
|
||||
import logging
|
||||
from datetime import datetime
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class AgentTaskCreate(BaseModel):
|
||||
"""创建 Agent 任务请求"""
|
||||
|
||||
project_id: str = Field(..., description="项目 ID")
|
||||
name: Optional[str] = Field(None, description="任务名称")
|
||||
description: Optional[str] = Field(None, description="任务描述")
|
||||
|
||||
# 审计配置
|
||||
audit_scope: Optional[dict] = Field(None, description="审计范围")
|
||||
target_vulnerabilities: Optional[List[str]] = Field(
|
||||
default=["sql_injection", "xss", "command_injection", "path_traversal", "ssrf"],
|
||||
description="目标漏洞类型",
|
||||
)
|
||||
verification_level: str = Field(
|
||||
"sandbox", description="验证级别: analysis_only, sandbox, generate_poc"
|
||||
)
|
||||
|
||||
# 分支
|
||||
branch_name: Optional[str] = Field(None, description="分支名称")
|
||||
|
||||
# 排除模式
|
||||
exclude_patterns: Optional[List[str]] = Field(
|
||||
default=["node_modules", "__pycache__", ".git", "*.min.js"], description="排除模式"
|
||||
)
|
||||
|
||||
# 文件范围
|
||||
target_files: Optional[List[str]] = Field(None, description="指定扫描的文件")
|
||||
|
||||
# Agent 配置
|
||||
max_iterations: int = Field(50, ge=1, le=200, description="最大迭代次数")
|
||||
timeout_seconds: int = Field(1800, ge=60, le=7200, description="超时时间(秒)")
|
||||
|
||||
|
||||
class AgentTaskResponse(BaseModel):
|
||||
"""Agent 任务响应 - 包含所有前端需要的字段"""
|
||||
|
||||
id: str
|
||||
project_id: str
|
||||
name: Optional[str]
|
||||
description: Optional[str]
|
||||
task_type: str = "agent_audit"
|
||||
status: str
|
||||
current_phase: Optional[str]
|
||||
current_step: Optional[str] = None
|
||||
|
||||
# 进度统计
|
||||
total_files: int = 0
|
||||
indexed_files: int = 0
|
||||
analyzed_files: int = 0
|
||||
total_chunks: int = 0
|
||||
|
||||
# Agent 统计
|
||||
total_iterations: int = 0
|
||||
tool_calls_count: int = 0
|
||||
tokens_used: int = 0
|
||||
|
||||
# 发现统计(兼容两种命名)
|
||||
findings_count: int = 0
|
||||
total_findings: int = 0 # 兼容字段
|
||||
verified_count: int = 0
|
||||
verified_findings: int = 0 # 兼容字段
|
||||
false_positive_count: int = 0
|
||||
|
||||
# 严重程度统计
|
||||
critical_count: int = 0
|
||||
high_count: int = 0
|
||||
medium_count: int = 0
|
||||
low_count: int = 0
|
||||
|
||||
# 评分
|
||||
quality_score: float = 0.0
|
||||
security_score: Optional[float] = None
|
||||
|
||||
# 进度百分比
|
||||
progress_percentage: float = 0.0
|
||||
|
||||
# 时间
|
||||
created_at: datetime
|
||||
started_at: Optional[datetime] = None
|
||||
completed_at: Optional[datetime] = None
|
||||
|
||||
# 配置
|
||||
audit_scope: Optional[dict] = None
|
||||
target_vulnerabilities: Optional[List[str]] = None
|
||||
verification_level: Optional[str] = None
|
||||
exclude_patterns: Optional[List[str]] = None
|
||||
target_files: Optional[List[str]] = None
|
||||
|
||||
# 错误信息
|
||||
error_message: Optional[str] = None
|
||||
|
||||
class Config:
|
||||
from_attributes = True
|
||||
|
||||
|
||||
class AgentEventResponse(BaseModel):
|
||||
"""Agent 事件响应"""
|
||||
|
||||
id: str
|
||||
task_id: str
|
||||
event_type: str
|
||||
phase: Optional[str]
|
||||
message: Optional[str] = None
|
||||
sequence: int
|
||||
# 🔥 ORM 字段名是 created_at,序列化为 timestamp
|
||||
created_at: datetime = Field(serialization_alias="timestamp")
|
||||
|
||||
# 工具相关字段
|
||||
tool_name: Optional[str] = None
|
||||
tool_input: Optional[Dict[str, Any]] = None
|
||||
tool_output: Optional[Dict[str, Any]] = None
|
||||
tool_duration_ms: Optional[int] = None
|
||||
|
||||
# 其他字段
|
||||
progress_percent: Optional[float] = None
|
||||
finding_id: Optional[str] = None
|
||||
tokens_used: Optional[int] = None
|
||||
# 🔥 ORM 字段名是 event_metadata,序列化为 metadata
|
||||
event_metadata: Optional[Dict[str, Any]] = Field(default=None, serialization_alias="metadata")
|
||||
|
||||
model_config = {
|
||||
"from_attributes": True,
|
||||
"populate_by_name": True,
|
||||
"by_alias": True, # 🔥 关键:确保序列化时使用别名
|
||||
}
|
||||
|
||||
|
||||
class AgentFindingResponse(BaseModel):
|
||||
"""Agent 发现响应"""
|
||||
|
||||
id: str
|
||||
task_id: str
|
||||
vulnerability_type: str
|
||||
severity: str
|
||||
title: str
|
||||
description: Optional[str]
|
||||
file_path: Optional[str]
|
||||
line_start: Optional[int]
|
||||
line_end: Optional[int]
|
||||
code_snippet: Optional[str]
|
||||
|
||||
is_verified: bool
|
||||
# 🔥 FIX: Map from ai_confidence in ORM, make Optional with default
|
||||
confidence: Optional[float] = Field(default=0.5, validation_alias="ai_confidence")
|
||||
status: str
|
||||
|
||||
suggestion: Optional[str] = None
|
||||
poc: Optional[dict] = None
|
||||
|
||||
created_at: datetime
|
||||
|
||||
model_config = {
|
||||
"from_attributes": True,
|
||||
"populate_by_name": True, # Allow both 'confidence' and 'ai_confidence'
|
||||
}
|
||||
|
||||
|
||||
class TaskSummaryResponse(BaseModel):
|
||||
"""任务摘要响应"""
|
||||
|
||||
task_id: str
|
||||
status: str
|
||||
security_score: Optional[int]
|
||||
|
||||
total_findings: int
|
||||
verified_findings: int
|
||||
|
||||
severity_distribution: Dict[str, int]
|
||||
vulnerability_types: Dict[str, int]
|
||||
|
||||
duration_seconds: Optional[int]
|
||||
phases_completed: List[str]
|
||||
|
||||
|
||||
class AgentTreeNodeResponse(BaseModel):
|
||||
"""Agent 树节点响应"""
|
||||
|
||||
id: str
|
||||
agent_id: str
|
||||
agent_name: str
|
||||
agent_type: str
|
||||
parent_agent_id: Optional[str] = None
|
||||
depth: int = 0
|
||||
task_description: Optional[str] = None
|
||||
knowledge_modules: Optional[List[str]] = None
|
||||
status: str = "created"
|
||||
result_summary: Optional[str] = None
|
||||
findings_count: int = 0
|
||||
iterations: int = 0
|
||||
tokens_used: int = 0
|
||||
tool_calls: int = 0
|
||||
duration_ms: Optional[int] = None
|
||||
children: List["AgentTreeNodeResponse"] = []
|
||||
|
||||
class Config:
|
||||
from_attributes = True
|
||||
|
||||
|
||||
class AgentTreeResponse(BaseModel):
|
||||
"""Agent 树响应"""
|
||||
|
||||
task_id: str
|
||||
root_agent_id: Optional[str] = None
|
||||
total_agents: int = 0
|
||||
running_agents: int = 0
|
||||
completed_agents: int = 0
|
||||
failed_agents: int = 0
|
||||
total_findings: int = 0
|
||||
nodes: List[AgentTreeNodeResponse] = []
|
||||
|
||||
|
||||
class CheckpointResponse(BaseModel):
|
||||
"""检查点响应"""
|
||||
|
||||
id: str
|
||||
agent_id: str
|
||||
agent_name: str
|
||||
agent_type: str
|
||||
iteration: int
|
||||
status: str
|
||||
total_tokens: int = 0
|
||||
tool_calls: int = 0
|
||||
findings_count: int = 0
|
||||
checkpoint_type: str = "auto"
|
||||
checkpoint_name: Optional[str] = None
|
||||
created_at: Optional[str] = None
|
||||
|
||||
class Config:
|
||||
from_attributes = True
|
||||
@@ -0,0 +1,512 @@
|
||||
"""Tools services for agent audits."""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import os
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from app.services.audit_queue import workspace_for
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
from .runtime import is_task_cancelled
|
||||
|
||||
|
||||
async def _initialize_tools(
|
||||
project_root: str,
|
||||
llm_service,
|
||||
user_config: Optional[Dict[str, Any]],
|
||||
sandbox_manager: Any, # 传递预初始化的 SandboxManager
|
||||
exclude_patterns: Optional[List[str]] = None,
|
||||
target_files: Optional[List[str]] = None,
|
||||
project_id: Optional[str] = None, # 🔥 用于 RAG collection_name
|
||||
event_emitter: Optional[Any] = None, # 🔥 新增:用于发送实时日志
|
||||
task_id: Optional[str] = None, # 🔥 新增:用于取消检查
|
||||
) -> Dict[str, Dict[str, Any]]:
|
||||
"""初始化工具集
|
||||
|
||||
Args:
|
||||
project_root: 项目根目录
|
||||
llm_service: LLM 服务
|
||||
user_config: 用户配置
|
||||
sandbox_manager: 沙箱管理器
|
||||
exclude_patterns: 排除模式列表
|
||||
target_files: 目标文件列表
|
||||
project_id: 项目 ID(用于 RAG collection_name)
|
||||
event_emitter: 事件发送器(用于发送实时日志)
|
||||
task_id: 任务 ID(用于取消检查)
|
||||
"""
|
||||
from app.core.config import settings
|
||||
from app.services.agent.knowledge import (
|
||||
GetVulnerabilityKnowledgeTool,
|
||||
SecurityKnowledgeQueryTool,
|
||||
)
|
||||
from app.services.agent.tools import (
|
||||
BanditTool,
|
||||
CreateVulnerabilityReportTool,
|
||||
DataFlowAnalysisTool,
|
||||
FileReadTool,
|
||||
FileSearchTool,
|
||||
FunctionContextTool,
|
||||
GitleaksTool,
|
||||
ListFilesTool,
|
||||
NpmAuditTool, # 🔥 Added missing tools
|
||||
OSVScannerTool,
|
||||
PatternMatchTool,
|
||||
RAGQueryTool,
|
||||
ReflectTool,
|
||||
SafetyTool,
|
||||
SecurityCodeSearchTool,
|
||||
SemgrepTool,
|
||||
ThinkTool,
|
||||
TruffleHogTool,
|
||||
)
|
||||
|
||||
# 🔥 RAG 相关导入
|
||||
from app.services.rag import CodeIndexer, CodeRetriever, EmbeddingService, IndexUpdateMode
|
||||
|
||||
# 辅助函数:发送事件
|
||||
async def emit(message: str, level: str = "info"):
|
||||
if event_emitter:
|
||||
logger.debug(f"[EMIT-TOOLS] Sending {level}: {message[:60]}...")
|
||||
if level == "info":
|
||||
await event_emitter.emit_info(message)
|
||||
elif level == "warning":
|
||||
await event_emitter.emit_warning(message)
|
||||
elif level == "error":
|
||||
await event_emitter.emit_error(message)
|
||||
else:
|
||||
logger.warning(f"[EMIT-TOOLS] No event_emitter, skipping: {message[:60]}...")
|
||||
|
||||
# ============ 🔥 初始化 RAG 系统 ============
|
||||
retriever = None
|
||||
try:
|
||||
await emit(f"🔍 正在初始化 RAG 系统...")
|
||||
|
||||
# 从用户配置中获取 embedding 配置
|
||||
user_llm_config = (user_config or {}).get("llmConfig", {})
|
||||
user_other_config = (user_config or {}).get("otherConfig", {})
|
||||
user_embedding_config = user_other_config.get("embedding_config", {})
|
||||
|
||||
# Embedding Provider 优先级:用户嵌入配置 > 环境变量
|
||||
embedding_provider = user_embedding_config.get("provider") or getattr(
|
||||
settings, "EMBEDDING_PROVIDER", "openai"
|
||||
)
|
||||
|
||||
# Embedding Model 优先级:用户嵌入配置 > 环境变量
|
||||
embedding_model = user_embedding_config.get("model") or getattr(
|
||||
settings, "EMBEDDING_MODEL", "text-embedding-3-small"
|
||||
)
|
||||
|
||||
# API Key 优先级:用户嵌入配置 > 环境变量 EMBEDDING_API_KEY > 用户 LLM 配置 > 环境变量 LLM_API_KEY
|
||||
# 注意:API Key 可以共享,因为很多用户使用同一个 OpenAI Key 做 LLM 和 Embedding
|
||||
embedding_api_key = (
|
||||
user_embedding_config.get("api_key")
|
||||
or getattr(settings, "EMBEDDING_API_KEY", None)
|
||||
or user_llm_config.get("llmApiKey")
|
||||
or getattr(settings, "LLM_API_KEY", "")
|
||||
or ""
|
||||
)
|
||||
|
||||
# Base URL 优先级:用户嵌入配置 > 环境变量 EMBEDDING_BASE_URL > None(使用提供商默认地址)
|
||||
# 🔥 重要:Base URL 不应该回退到 LLM 的 base_url,因为 Embedding 和 LLM 可能使用完全不同的服务
|
||||
# 例如:LLM 使用 SiliconFlow,但 Embedding 使用 HuggingFace
|
||||
embedding_base_url = (
|
||||
user_embedding_config.get("base_url")
|
||||
or getattr(settings, "EMBEDDING_BASE_URL", None)
|
||||
or None
|
||||
)
|
||||
|
||||
logger.info(
|
||||
f"RAG 配置: provider={embedding_provider}, model={embedding_model}, base_url={embedding_base_url or '(使用默认)'}"
|
||||
)
|
||||
await emit(f"📊 Embedding 配置: {embedding_provider}/{embedding_model}")
|
||||
|
||||
# 创建 Embedding 服务
|
||||
embedding_service = EmbeddingService(
|
||||
provider=embedding_provider,
|
||||
model=embedding_model,
|
||||
api_key=embedding_api_key,
|
||||
base_url=embedding_base_url,
|
||||
)
|
||||
# 使用用户配置的 batch_size
|
||||
embedding_service.batch_size = user_embedding_config.get("batch_size", 100)
|
||||
|
||||
# 创建 collection_name(基于 project_id)
|
||||
collection_name = f"project_{project_id}" if project_id else "default_project"
|
||||
|
||||
# 🔥 v2.0: 创建 CodeIndexer 并进行智能索引
|
||||
# 智能索引会自动:
|
||||
# - 检测 embedding 模型变更,如需要则自动重建
|
||||
# - 对比文件 hash,只更新变化的文件(增量更新)
|
||||
indexer = CodeIndexer(
|
||||
collection_name=collection_name,
|
||||
embedding_service=embedding_service,
|
||||
persist_directory=str(workspace_for(task_id) / "vectors")
|
||||
if task_id
|
||||
else settings.VECTOR_DB_PATH,
|
||||
)
|
||||
|
||||
logger.info(f"📝 开始智能索引项目: {project_root}")
|
||||
await emit(f"📝 正在构建代码向量索引...")
|
||||
|
||||
index_progress = None
|
||||
last_progress_update = 0
|
||||
last_embedding_progress = [0] # 使用列表以便在闭包中修改
|
||||
embedding_total = [0] # 记录总数
|
||||
|
||||
# 🔥 嵌入进度回调函数(同步,但会调度异步任务)
|
||||
def on_embedding_progress(processed: int, total: int):
|
||||
embedding_total[0] = total
|
||||
# 每处理 50 个或完成时更新
|
||||
if processed - last_embedding_progress[0] >= 50 or processed == total:
|
||||
last_embedding_progress[0] = processed
|
||||
percentage = (processed / total * 100) if total > 0 else 0
|
||||
msg = f"🔢 嵌入进度: {processed}/{total} ({percentage:.0f}%)"
|
||||
logger.info(msg)
|
||||
# 使用 asyncio.create_task 调度异步 emit
|
||||
try:
|
||||
loop = asyncio.get_running_loop()
|
||||
loop.create_task(emit(msg))
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to emit embedding progress: {e}")
|
||||
|
||||
# 🔥 创建取消检查函数,用于在嵌入批处理中检查取消状态
|
||||
def check_cancelled() -> bool:
|
||||
return task_id is not None and is_task_cancelled(task_id)
|
||||
|
||||
async for progress in indexer.smart_index_directory(
|
||||
directory=project_root,
|
||||
exclude_patterns=exclude_patterns or [],
|
||||
include_patterns=target_files, # 🔥 传递 target_files 限制索引范围
|
||||
update_mode=IndexUpdateMode.SMART,
|
||||
embedding_progress_callback=on_embedding_progress,
|
||||
cancel_check=check_cancelled, # 🔥 传递取消检查函数
|
||||
):
|
||||
# 🔥 在索引过程中检查取消状态
|
||||
if check_cancelled():
|
||||
logger.info(f"[Cancel] RAG indexing cancelled for task {task_id}")
|
||||
raise asyncio.CancelledError("任务已取消")
|
||||
|
||||
index_progress = progress
|
||||
# 每处理 10 个文件或有重要变化时发送进度更新
|
||||
if (
|
||||
progress.processed_files - last_progress_update >= 10
|
||||
or progress.processed_files == progress.total_files
|
||||
):
|
||||
if progress.total_files > 0:
|
||||
await emit(
|
||||
f"📝 索引进度: {progress.processed_files}/{progress.total_files} 文件 "
|
||||
f"({progress.progress_percentage:.0f}%)"
|
||||
)
|
||||
last_progress_update = progress.processed_files
|
||||
|
||||
# 🔥 发送状态消息(如嵌入向量生成进度)
|
||||
if progress.status_message:
|
||||
await emit(progress.status_message)
|
||||
progress.status_message = "" # 清空已发送的消息
|
||||
|
||||
if index_progress:
|
||||
summary = (
|
||||
f"✅ 索引完成: 模式={index_progress.update_mode}, "
|
||||
f"新增={index_progress.added_files}, "
|
||||
f"更新={index_progress.updated_files}, "
|
||||
f"删除={index_progress.deleted_files}, "
|
||||
f"代码块={index_progress.indexed_chunks}"
|
||||
)
|
||||
logger.info(summary)
|
||||
await emit(summary)
|
||||
|
||||
# 创建 CodeRetriever(用于搜索)
|
||||
# 🔥 传递 api_key,用于自动适配 collection 的 embedding 配置
|
||||
retriever = CodeRetriever(
|
||||
collection_name=collection_name,
|
||||
embedding_service=embedding_service,
|
||||
persist_directory=str(workspace_for(task_id) / "vectors")
|
||||
if task_id
|
||||
else settings.VECTOR_DB_PATH,
|
||||
api_key=embedding_api_key, # 🔥 传递 api_key 以支持自动切换 embedding
|
||||
)
|
||||
|
||||
logger.info(f"✅ RAG 系统初始化成功: collection={collection_name}")
|
||||
await emit(f"✅ RAG 系统初始化成功")
|
||||
|
||||
except Exception as e:
|
||||
logger.warning(f"⚠️ RAG 系统初始化失败: {e}")
|
||||
await emit(f"⚠️ RAG 系统初始化失败: {e}", "warning")
|
||||
import traceback
|
||||
|
||||
logger.debug(f"RAG 初始化异常详情:\n{traceback.format_exc()}")
|
||||
retriever = None
|
||||
|
||||
# 基础工具 - 传递排除模式和目标文件
|
||||
base_tools = {
|
||||
"read_file": FileReadTool(project_root, exclude_patterns, target_files),
|
||||
"list_files": ListFilesTool(project_root, exclude_patterns, target_files),
|
||||
"search_code": FileSearchTool(project_root, exclude_patterns, target_files),
|
||||
"think": ThinkTool(),
|
||||
"reflect": ReflectTool(),
|
||||
}
|
||||
|
||||
# Recon 工具
|
||||
recon_tools = {
|
||||
**base_tools,
|
||||
# 🔥 外部侦察工具 (Recon 阶段也需要使用这些工具来收集初步信息)
|
||||
"semgrep_scan": SemgrepTool(project_root, sandbox_manager),
|
||||
"bandit_scan": BanditTool(project_root, sandbox_manager),
|
||||
"gitleaks_scan": GitleaksTool(project_root, sandbox_manager),
|
||||
"npm_audit": NpmAuditTool(project_root, sandbox_manager),
|
||||
"safety_scan": SafetyTool(project_root, sandbox_manager),
|
||||
"trufflehog_scan": TruffleHogTool(project_root, sandbox_manager),
|
||||
"osv_scan": OSVScannerTool(project_root, sandbox_manager),
|
||||
}
|
||||
|
||||
# 🔥 注册 RAG 工具到 Recon Agent
|
||||
if retriever:
|
||||
recon_tools["rag_query"] = RAGQueryTool(retriever)
|
||||
logger.info("✅ RAG 工具 (rag_query) 已注册到 Recon Agent")
|
||||
|
||||
# Analysis 工具
|
||||
# 🔥 导入智能扫描工具
|
||||
from app.services.agent.tools import QuickAuditTool, SmartScanTool
|
||||
|
||||
analysis_tools = {
|
||||
**base_tools,
|
||||
# 🔥 智能扫描工具(推荐首先使用)
|
||||
"smart_scan": SmartScanTool(project_root),
|
||||
"quick_audit": QuickAuditTool(project_root),
|
||||
# 模式匹配工具(增强版)
|
||||
"pattern_match": PatternMatchTool(project_root),
|
||||
# 数据流分析
|
||||
"dataflow_analysis": DataFlowAnalysisTool(llm_service),
|
||||
# 外部安全工具 (传入共享的 sandbox_manager)
|
||||
"semgrep_scan": SemgrepTool(project_root, sandbox_manager),
|
||||
"bandit_scan": BanditTool(project_root, sandbox_manager),
|
||||
"gitleaks_scan": GitleaksTool(project_root, sandbox_manager),
|
||||
"npm_audit": NpmAuditTool(project_root, sandbox_manager),
|
||||
"safety_scan": SafetyTool(project_root, sandbox_manager),
|
||||
"trufflehog_scan": TruffleHogTool(project_root, sandbox_manager),
|
||||
"osv_scan": OSVScannerTool(project_root, sandbox_manager),
|
||||
# 安全知识查询
|
||||
"query_security_knowledge": SecurityKnowledgeQueryTool(),
|
||||
"get_vulnerability_knowledge": GetVulnerabilityKnowledgeTool(),
|
||||
}
|
||||
|
||||
# 🔥 注册 RAG 工具到 Analysis Agent
|
||||
if retriever:
|
||||
analysis_tools["rag_query"] = RAGQueryTool(retriever)
|
||||
analysis_tools["security_search"] = SecurityCodeSearchTool(retriever)
|
||||
analysis_tools["function_context"] = FunctionContextTool(retriever)
|
||||
logger.info(
|
||||
"✅ RAG 工具 (rag_query, security_search, function_context) 已注册到 Analysis Agent"
|
||||
)
|
||||
else:
|
||||
logger.warning("⚠️ RAG 未初始化,rag_query/security_search/function_context 工具不可用")
|
||||
|
||||
# Verification 工具
|
||||
# 🔥 导入沙箱工具
|
||||
from app.services.agent.tools import (
|
||||
# 漏洞验证专用工具
|
||||
CommandInjectionTestTool,
|
||||
DeserializationTestTool,
|
||||
ExtractFunctionTool,
|
||||
GoTestTool,
|
||||
JavaScriptTestTool,
|
||||
JavaTestTool,
|
||||
PathTraversalTestTool,
|
||||
# 多语言代码测试工具
|
||||
PhpTestTool,
|
||||
PythonTestTool,
|
||||
RubyTestTool,
|
||||
# 🔥 新增:通用代码执行工具 (LLM 驱动的 Fuzzing Harness)
|
||||
RunCodeTool,
|
||||
SandboxHttpTool,
|
||||
SandboxTool,
|
||||
ShellTestTool,
|
||||
SqlInjectionTestTool,
|
||||
SstiTestTool,
|
||||
UniversalCodeTestTool,
|
||||
UniversalVulnTestTool,
|
||||
VulnerabilityVerifyTool,
|
||||
XssTestTool,
|
||||
)
|
||||
|
||||
verification_tools = {
|
||||
**base_tools,
|
||||
# 🔥 沙箱验证工具
|
||||
"sandbox_exec": SandboxTool(sandbox_manager),
|
||||
"sandbox_http": SandboxHttpTool(sandbox_manager),
|
||||
"verify_vulnerability": VulnerabilityVerifyTool(sandbox_manager),
|
||||
# 🔥 多语言代码测试工具
|
||||
"php_test": PhpTestTool(sandbox_manager, project_root),
|
||||
"python_test": PythonTestTool(sandbox_manager, project_root),
|
||||
"javascript_test": JavaScriptTestTool(sandbox_manager, project_root),
|
||||
"java_test": JavaTestTool(sandbox_manager, project_root),
|
||||
"go_test": GoTestTool(sandbox_manager, project_root),
|
||||
"ruby_test": RubyTestTool(sandbox_manager, project_root),
|
||||
"shell_test": ShellTestTool(sandbox_manager, project_root),
|
||||
"universal_code_test": UniversalCodeTestTool(sandbox_manager, project_root),
|
||||
# 🔥 漏洞验证专用工具
|
||||
"test_command_injection": CommandInjectionTestTool(sandbox_manager, project_root),
|
||||
"test_sql_injection": SqlInjectionTestTool(sandbox_manager, project_root),
|
||||
"test_xss": XssTestTool(sandbox_manager, project_root),
|
||||
"test_path_traversal": PathTraversalTestTool(sandbox_manager, project_root),
|
||||
"test_ssti": SstiTestTool(sandbox_manager, project_root),
|
||||
"test_deserialization": DeserializationTestTool(sandbox_manager, project_root),
|
||||
"universal_vuln_test": UniversalVulnTestTool(sandbox_manager, project_root),
|
||||
# 🔥 新增:通用代码执行工具 (LLM 驱动的 Fuzzing Harness)
|
||||
"run_code": RunCodeTool(sandbox_manager, project_root),
|
||||
"extract_function": ExtractFunctionTool(project_root),
|
||||
# 报告工具 - 🔥 v2.1: 传递 project_root 用于文件验证
|
||||
"create_vulnerability_report": CreateVulnerabilityReportTool(project_root),
|
||||
}
|
||||
|
||||
# Orchestrator 工具(主要是思考工具)
|
||||
orchestrator_tools = {
|
||||
"think": ThinkTool(),
|
||||
"reflect": ReflectTool(),
|
||||
}
|
||||
|
||||
return {
|
||||
"recon": recon_tools,
|
||||
"analysis": analysis_tools,
|
||||
"verification": verification_tools,
|
||||
"orchestrator": orchestrator_tools,
|
||||
}
|
||||
|
||||
|
||||
async def _collect_project_info(
|
||||
project_root: str,
|
||||
project_name: str,
|
||||
exclude_patterns: Optional[List[str]] = None,
|
||||
target_files: Optional[List[str]] = None,
|
||||
) -> Dict[str, Any]:
|
||||
"""收集项目信息
|
||||
|
||||
Args:
|
||||
project_root: 项目根目录
|
||||
project_name: 项目名称
|
||||
exclude_patterns: 排除模式列表
|
||||
target_files: 目标文件列表
|
||||
|
||||
🔥 重要:当指定了 target_files 时,返回的项目结构应该只包含目标文件相关的信息,
|
||||
以确保 Orchestrator 和子 Agent 看到的是一致的、过滤后的视图。
|
||||
"""
|
||||
import fnmatch
|
||||
|
||||
info = {
|
||||
"name": project_name,
|
||||
"root": project_root,
|
||||
"languages": [],
|
||||
"file_count": 0,
|
||||
"structure": {},
|
||||
}
|
||||
|
||||
try:
|
||||
# 默认排除目录
|
||||
exclude_dirs = {
|
||||
"node_modules",
|
||||
"__pycache__",
|
||||
".git",
|
||||
"venv",
|
||||
".venv",
|
||||
"build",
|
||||
"dist",
|
||||
"target",
|
||||
".idea",
|
||||
".vscode",
|
||||
}
|
||||
|
||||
# 从用户配置的排除模式中提取目录
|
||||
if exclude_patterns:
|
||||
for pattern in exclude_patterns:
|
||||
if pattern.endswith("/**"):
|
||||
exclude_dirs.add(pattern[:-3])
|
||||
elif "/" not in pattern and "*" not in pattern:
|
||||
exclude_dirs.add(pattern)
|
||||
|
||||
# 目标文件集合
|
||||
target_files_set = set(target_files) if target_files else None
|
||||
|
||||
lang_map = {
|
||||
".py": "Python",
|
||||
".js": "JavaScript",
|
||||
".ts": "TypeScript",
|
||||
".java": "Java",
|
||||
".go": "Go",
|
||||
".php": "PHP",
|
||||
".rb": "Ruby",
|
||||
".rs": "Rust",
|
||||
".c": "C",
|
||||
".cpp": "C++",
|
||||
}
|
||||
|
||||
# 🔥 收集过滤后的文件列表
|
||||
filtered_files = []
|
||||
filtered_dirs = set()
|
||||
|
||||
for root, dirs, files in os.walk(project_root):
|
||||
dirs[:] = [d for d in dirs if d not in exclude_dirs]
|
||||
|
||||
for f in files:
|
||||
relative_path = os.path.relpath(os.path.join(root, f), project_root)
|
||||
|
||||
# 检查是否在目标文件列表中
|
||||
if target_files_set and relative_path not in target_files_set:
|
||||
continue
|
||||
|
||||
# 检查排除模式
|
||||
should_skip = False
|
||||
if exclude_patterns:
|
||||
for pattern in exclude_patterns:
|
||||
if fnmatch.fnmatch(relative_path, pattern) or fnmatch.fnmatch(f, pattern):
|
||||
should_skip = True
|
||||
break
|
||||
if should_skip:
|
||||
continue
|
||||
|
||||
info["file_count"] += 1
|
||||
filtered_files.append(relative_path)
|
||||
|
||||
# 🔥 收集文件所在的目录
|
||||
dir_path = os.path.dirname(relative_path)
|
||||
if dir_path:
|
||||
# 添加目录及其父目录
|
||||
parts = dir_path.split(os.sep)
|
||||
for i in range(len(parts)):
|
||||
filtered_dirs.add(os.sep.join(parts[: i + 1]))
|
||||
|
||||
ext = os.path.splitext(f)[1].lower()
|
||||
if ext in lang_map and lang_map[ext] not in info["languages"]:
|
||||
info["languages"].append(lang_map[ext])
|
||||
|
||||
# 🔥 根据是否有目标文件限制,生成不同的结构信息
|
||||
if target_files_set:
|
||||
# 当指定了目标文件时,只显示目标文件和相关目录
|
||||
info["structure"] = {
|
||||
"directories": sorted(list(filtered_dirs))[:20],
|
||||
"files": filtered_files[:30],
|
||||
"scope_limited": True, # 🔥 标记这是限定范围的视图
|
||||
"scope_message": f"审计范围限定为 {len(filtered_files)} 个指定文件",
|
||||
}
|
||||
else:
|
||||
# 全项目审计时,显示顶层目录结构
|
||||
try:
|
||||
top_items = os.listdir(project_root)
|
||||
info["structure"] = {
|
||||
"directories": [
|
||||
d
|
||||
for d in top_items
|
||||
if os.path.isdir(os.path.join(project_root, d)) and d not in exclude_dirs
|
||||
],
|
||||
"files": [
|
||||
f for f in top_items if os.path.isfile(os.path.join(project_root, f))
|
||||
][:20],
|
||||
"scope_limited": False,
|
||||
}
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to collect project info: {e}")
|
||||
|
||||
return info
|
||||
@@ -0,0 +1,658 @@
|
||||
"""Workspace services for agent audits."""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
import zipfile
|
||||
from typing import Any, Optional
|
||||
|
||||
from app.core.config import settings
|
||||
from app.db.session import AsyncSessionLocal
|
||||
from app.models.audit_job import AuditJob
|
||||
from app.models.project import Project
|
||||
from app.services.audit_queue import load_checkpoint, save_checkpoint, workspace_for
|
||||
from app.services.git_ssh_service import GitSSHOperations
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
from .runtime import is_task_cancelled
|
||||
|
||||
|
||||
def validate_git_url(url: str) -> bool:
|
||||
"""
|
||||
验证 Git URL 是否安全
|
||||
|
||||
Args:
|
||||
url: Git URL
|
||||
|
||||
Returns:
|
||||
bool: URL 是否安全
|
||||
"""
|
||||
if not url:
|
||||
return False
|
||||
|
||||
from urllib.parse import urlparse
|
||||
|
||||
parsed = urlparse(url)
|
||||
|
||||
# 只允许 http, https, git, ssh 协议
|
||||
allowed_schemes = {"http", "https", "git", "ssh", "git@"}
|
||||
if parsed.scheme and parsed.scheme not in allowed_schemes:
|
||||
return False
|
||||
|
||||
# 检查是否包含可疑的命令注入字符
|
||||
dangerous_patterns = [";", "|", "&", "$(", "`", "\n", "\r", "\t"]
|
||||
for pattern in dangerous_patterns:
|
||||
if pattern in url:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
|
||||
def validate_branch_name(branch: str) -> bool:
|
||||
"""
|
||||
验证 Git 分支名称是否安全
|
||||
|
||||
Args:
|
||||
branch: 分支名称
|
||||
|
||||
Returns:
|
||||
bool: 分支名称是否安全
|
||||
"""
|
||||
if not branch:
|
||||
return False
|
||||
|
||||
# Git 分支名称规则:只允许字母、数字、下划线、连字符、点、斜杠
|
||||
# 参考: https://git-scm.com/docs/git-check-ref-format
|
||||
pattern = r"^[a-zA-Z0-9_\-\.\/]+$"
|
||||
if not re.match(pattern, branch):
|
||||
return False
|
||||
|
||||
# 防止路径遍历
|
||||
if ".." in branch or branch.startswith("/") or branch.endswith("/"):
|
||||
return False
|
||||
|
||||
# 限制长度
|
||||
if len(branch) > 256:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
|
||||
def is_path_safe(base_path: str, target_path: str) -> bool:
|
||||
"""
|
||||
检查目标路径是否在基础目录内(防止路径遍历)
|
||||
|
||||
Args:
|
||||
base_path: 基础目录
|
||||
target_path: 目标路径
|
||||
|
||||
Returns:
|
||||
bool: 路径是否安全
|
||||
"""
|
||||
# 规范化路径
|
||||
abs_base = os.path.abspath(base_path)
|
||||
abs_target = os.path.abspath(os.path.join(base_path, target_path))
|
||||
|
||||
# 检查目标路径是否在基础目录内
|
||||
return abs_target.startswith(abs_base + os.sep) or abs_target == abs_base
|
||||
|
||||
|
||||
def safe_extract_zip(zip_ref: zipfile.ZipFile, extract_dir: str, task_id: str) -> None:
|
||||
"""
|
||||
安全解压 ZIP 文件,防止 Zip Slip 攻击
|
||||
|
||||
Args:
|
||||
zip_ref: ZipFile 对象
|
||||
extract_dir: 解压目标目录
|
||||
task_id: 任务 ID(用于取消检查)
|
||||
"""
|
||||
from app.services.archive import extract_zip
|
||||
|
||||
extract_zip(
|
||||
zip_ref, extract_dir, strip_root=True, cancel_check=lambda: is_task_cancelled(task_id)
|
||||
)
|
||||
|
||||
|
||||
async def _get_project_root(
|
||||
project: Project,
|
||||
task_id: str,
|
||||
branch_name: Optional[str] = None,
|
||||
github_token: Optional[str] = None,
|
||||
gitlab_token: Optional[str] = None,
|
||||
gitea_token: Optional[str] = None, # 🔥 新增
|
||||
ssh_private_key: Optional[str] = None, # 🔥 新增:SSH私钥(用于SSH认证)
|
||||
event_emitter: Optional[Any] = None, # 🔥 新增:用于发送实时日志
|
||||
) -> str:
|
||||
"""
|
||||
获取项目根目录
|
||||
|
||||
支持两种项目类型:
|
||||
- ZIP 项目:解压 ZIP 文件到临时目录
|
||||
- 仓库项目:克隆仓库到临时目录
|
||||
|
||||
Args:
|
||||
project: 项目对象
|
||||
task_id: 任务ID
|
||||
branch_name: 分支名称(仓库项目使用,优先于 project.default_branch)
|
||||
github_token: GitHub 访问令牌(用于私有仓库)
|
||||
gitlab_token: GitLab 访问令牌(用于私有仓库)
|
||||
gitea_token: Gitea 访问令牌(用于私有仓库)
|
||||
ssh_private_key: SSH私钥(用于SSH认证)
|
||||
event_emitter: 事件发送器(用于发送实时日志)
|
||||
|
||||
Returns:
|
||||
项目根目录路径
|
||||
|
||||
Raises:
|
||||
RuntimeError: 当项目文件获取失败时
|
||||
"""
|
||||
import shutil
|
||||
import subprocess
|
||||
import zipfile
|
||||
from urllib.parse import urlparse, urlunparse
|
||||
|
||||
# 辅助函数:发送事件
|
||||
async def emit(message: str, level: str = "info"):
|
||||
if event_emitter:
|
||||
if level == "info":
|
||||
await event_emitter.emit_info(message)
|
||||
elif level == "warning":
|
||||
await event_emitter.emit_warning(message)
|
||||
elif level == "error":
|
||||
await event_emitter.emit_error(message)
|
||||
|
||||
# 🔥 辅助函数:检查取消状态
|
||||
def check_cancelled():
|
||||
if is_task_cancelled(task_id):
|
||||
raise asyncio.CancelledError("任务已取消")
|
||||
|
||||
workspace = workspace_for(task_id)
|
||||
from app.services.audit_queue import project_snapshot
|
||||
|
||||
project = await project_snapshot(task_id, project)
|
||||
checkpoint = await load_checkpoint(task_id)
|
||||
cached_root = checkpoint.get("project_root")
|
||||
if (
|
||||
cached_root
|
||||
and os.path.isdir(cached_root)
|
||||
and os.path.commonpath([str(workspace), os.path.realpath(cached_root)]) == str(workspace)
|
||||
):
|
||||
await emit("♻️ 恢复已准备的项目快照")
|
||||
return cached_root
|
||||
base_path = str(workspace / "source")
|
||||
|
||||
# 确保目录存在且为空
|
||||
if os.path.exists(base_path):
|
||||
shutil.rmtree(base_path)
|
||||
os.makedirs(base_path, exist_ok=True)
|
||||
|
||||
# 🔥 在开始任何操作前检查取消
|
||||
check_cancelled()
|
||||
|
||||
# 根据项目类型处理
|
||||
if project.source_type == "zip":
|
||||
# 🔥 ZIP 项目:解压 ZIP 文件
|
||||
check_cancelled() # 🔥 解压前检查
|
||||
await emit(f"📦 正在解压项目文件...")
|
||||
from app.services.zip_storage import load_project_zip
|
||||
|
||||
async with AsyncSessionLocal() as snapshot_db:
|
||||
job = await snapshot_db.get(AuditJob, task_id)
|
||||
zip_path = (job.payload or {}).get("source_zip") if job else None
|
||||
zip_path = zip_path or await load_project_zip(project.id)
|
||||
|
||||
if zip_path and os.path.exists(zip_path):
|
||||
try:
|
||||
check_cancelled() # 🔥 解压前再次检查
|
||||
with zipfile.ZipFile(zip_path, "r") as zip_ref:
|
||||
# 🔥 逐个文件解压,支持取消检查
|
||||
# 🔥 Security Fix: 使用 safe_extract_zip 替代 extract,防止 Zip Slip 和软链接攻击
|
||||
await asyncio.to_thread(safe_extract_zip, zip_ref, base_path, task_id)
|
||||
logger.info(f"✅ Extracted ZIP project {project.id} to {base_path}")
|
||||
await emit(f"✅ ZIP 文件解压完成")
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to extract ZIP {zip_path}: {e}")
|
||||
await emit(f"❌ 解压失败: {e}", "error")
|
||||
raise RuntimeError(f"无法解压项目文件: {e}")
|
||||
else:
|
||||
logger.warning(f"⚠️ ZIP file not found for project {project.id}")
|
||||
await emit(f"❌ ZIP 文件不存在", "error")
|
||||
raise RuntimeError(f"项目 ZIP 文件不存在: {project.id}")
|
||||
|
||||
elif project.source_type == "repository" and project.repository_url:
|
||||
# 🔥 仓库项目:优先使用 ZIP 下载(更快更稳定),git clone 作为回退
|
||||
repo_url = project.repository_url
|
||||
repo_type = project.repository_type or "other"
|
||||
|
||||
# 🔥 安全验证:检查 Git URL 是否安全
|
||||
if not validate_git_url(repo_url):
|
||||
logger.error(f"❌ 无效的 Git URL: {repo_url}")
|
||||
await emit(f"❌ 无效的仓库 URL", "error")
|
||||
raise RuntimeError(f"无效的仓库 URL: {repo_url}")
|
||||
|
||||
await emit(f"🔄 正在获取仓库: {repo_url}")
|
||||
|
||||
# 检测是否为SSH URL(SSH链接不支持ZIP下载)
|
||||
is_ssh_url = GitSSHOperations.is_ssh_url(repo_url)
|
||||
|
||||
# 解析仓库 URL 获取 owner/repo
|
||||
parsed = urlparse(repo_url)
|
||||
path_parts = parsed.path.strip("/").replace(".git", "").split("/")
|
||||
if len(path_parts) >= 2:
|
||||
owner, repo = path_parts[0], path_parts[1]
|
||||
else:
|
||||
owner, repo = None, None
|
||||
|
||||
# 构建分支尝试顺序
|
||||
branches_to_try = []
|
||||
if branch_name:
|
||||
# 🔥 安全验证:检查分支名称是否安全
|
||||
if not validate_branch_name(branch_name):
|
||||
logger.error(f"❌ 无效的分支名称: {branch_name}")
|
||||
await emit(f"❌ 无效的分支名称", "error")
|
||||
raise RuntimeError(f"无效的分支名称: {branch_name}")
|
||||
branches_to_try.append(branch_name)
|
||||
if not branch_name:
|
||||
if project.default_branch and project.default_branch not in branches_to_try:
|
||||
# 🔥 安全验证:检查默认分支名称是否安全
|
||||
if validate_branch_name(project.default_branch):
|
||||
branches_to_try.append(project.default_branch)
|
||||
for common_branch in ["main", "master"]:
|
||||
if common_branch not in branches_to_try:
|
||||
branches_to_try.append(common_branch)
|
||||
|
||||
download_success = False
|
||||
last_error = ""
|
||||
|
||||
# ============ 方案1: 优先使用 ZIP 下载(更快更稳定)============
|
||||
# SSH链接直接跳过ZIP下载,使用git clone
|
||||
if is_ssh_url:
|
||||
logger.info(f"检测到SSH URL,跳过ZIP下载,直接使用Git克隆")
|
||||
await emit(f"🔑 检测到SSH认证,使用Git克隆...")
|
||||
|
||||
if owner and repo and not is_ssh_url:
|
||||
import httpx
|
||||
|
||||
for branch in branches_to_try:
|
||||
check_cancelled()
|
||||
|
||||
# 清理目录
|
||||
if os.path.exists(base_path) and os.listdir(base_path):
|
||||
shutil.rmtree(base_path)
|
||||
os.makedirs(base_path, exist_ok=True)
|
||||
|
||||
# 构建 ZIP 下载 URL
|
||||
if repo_type == "github" or "github.com" in repo_url:
|
||||
# GitHub ZIP 下载 URL
|
||||
zip_url = f"https://github.com/{owner}/{repo}/archive/refs/heads/{branch}.zip"
|
||||
headers = {}
|
||||
if github_token:
|
||||
headers["Authorization"] = f"token {github_token}"
|
||||
elif repo_type == "gitlab" or "gitlab" in repo_url:
|
||||
# GitLab ZIP 下载 URL(需要对 owner/repo 进行 URL 编码)
|
||||
import urllib.parse
|
||||
|
||||
project_path = urllib.parse.quote(f"{owner}/{repo}", safe="")
|
||||
gitlab_host = parsed.netloc
|
||||
zip_url = f"https://{gitlab_host}/api/v4/projects/{project_path}/repository/archive.zip?sha={branch}"
|
||||
headers = {}
|
||||
if gitlab_token:
|
||||
headers["PRIVATE-TOKEN"] = gitlab_token
|
||||
else:
|
||||
# 其他平台,跳过 ZIP 下载
|
||||
break
|
||||
|
||||
logger.info(f"📦 尝试下载 ZIP 归档 (分支: {branch})...")
|
||||
await emit(f"📦 尝试下载 ZIP 归档 (分支: {branch})")
|
||||
|
||||
try:
|
||||
zip_temp_path = str(workspace / "download.zip")
|
||||
|
||||
async def download_zip():
|
||||
async with httpx.AsyncClient(timeout=60.0, follow_redirects=True) as client:
|
||||
async with client.stream("GET", zip_url, headers=headers) as resp:
|
||||
if resp.status_code != 200:
|
||||
return False, f"HTTP {resp.status_code}"
|
||||
import aiofiles
|
||||
|
||||
total = 0
|
||||
async with aiofiles.open(zip_temp_path, "wb") as output:
|
||||
async for chunk in resp.aiter_bytes(64 * 1024):
|
||||
check_cancelled()
|
||||
total += len(chunk)
|
||||
if total > settings.ZIP_MAX_UPLOAD_BYTES:
|
||||
raise ValueError(
|
||||
"Repository ZIP exceeds the upload quota"
|
||||
)
|
||||
await output.write(chunk)
|
||||
return True, None
|
||||
|
||||
# 使用取消检查循环
|
||||
download_task = asyncio.create_task(download_zip())
|
||||
while not download_task.done():
|
||||
check_cancelled()
|
||||
try:
|
||||
success, error = await asyncio.wait_for(
|
||||
asyncio.shield(download_task), timeout=1.0
|
||||
)
|
||||
break
|
||||
except asyncio.TimeoutError:
|
||||
continue
|
||||
|
||||
if download_task.done():
|
||||
success, error = download_task.result()
|
||||
|
||||
if success and os.path.exists(zip_temp_path):
|
||||
# 解压 ZIP
|
||||
check_cancelled()
|
||||
with zipfile.ZipFile(zip_temp_path, "r") as zip_ref:
|
||||
# 🔥 使用安全解压函数,防止 Zip Slip 攻击
|
||||
await asyncio.to_thread(safe_extract_zip, zip_ref, base_path, task_id)
|
||||
|
||||
# 清理临时文件
|
||||
os.remove(zip_temp_path)
|
||||
logger.info(f"✅ ZIP 下载成功 (分支: {branch})")
|
||||
await emit(f"✅ 仓库获取成功 (ZIP下载, 分支: {branch})")
|
||||
download_success = True
|
||||
break
|
||||
else:
|
||||
last_error = error or "下载失败"
|
||||
logger.warning(f"ZIP 下载失败 (分支 {branch}): {last_error}")
|
||||
await emit(f"⚠️ ZIP 下载失败,尝试其他分支...", "warning")
|
||||
# 清理临时文件
|
||||
if os.path.exists(zip_temp_path):
|
||||
os.remove(zip_temp_path)
|
||||
|
||||
except asyncio.CancelledError:
|
||||
logger.info(f"[Cancel] ZIP download cancelled for task {task_id}")
|
||||
raise
|
||||
except Exception as e:
|
||||
last_error = str(e)
|
||||
logger.warning(f"ZIP 下载异常 (分支 {branch}): {e}")
|
||||
await emit(f"⚠️ ZIP 下载异常: {str(e)[:50]}...", "warning")
|
||||
|
||||
# ============ 方案2: 回退到 git clone ============
|
||||
if not download_success:
|
||||
if is_ssh_url:
|
||||
# SSH链接直接使用git clone,不是"失败"
|
||||
pass # 已在上面输出提示
|
||||
else:
|
||||
await emit(f"🔄 ZIP 下载失败,回退到 Git 克隆...")
|
||||
logger.info("ZIP download failed, falling back to git clone")
|
||||
|
||||
# 检查 git 是否可用
|
||||
try:
|
||||
git_check = subprocess.run(
|
||||
["git", "--version"], capture_output=True, text=True, timeout=10
|
||||
)
|
||||
if git_check.returncode != 0:
|
||||
await emit(f"❌ Git 未安装", "error")
|
||||
raise RuntimeError("Git 未安装,无法克隆仓库。")
|
||||
except FileNotFoundError:
|
||||
await emit(f"❌ Git 未安装", "error")
|
||||
raise RuntimeError("Git 未安装,无法克隆仓库。")
|
||||
except subprocess.TimeoutExpired:
|
||||
await emit(f"❌ Git 检测超时", "error")
|
||||
raise RuntimeError("Git 检测超时")
|
||||
|
||||
# 构建带认证的 URL
|
||||
auth_url = repo_url
|
||||
if repo_type == "github" and github_token:
|
||||
auth_url = urlunparse(
|
||||
(
|
||||
parsed.scheme,
|
||||
f"{github_token}@{parsed.netloc}",
|
||||
parsed.path,
|
||||
parsed.params,
|
||||
parsed.query,
|
||||
parsed.fragment,
|
||||
)
|
||||
)
|
||||
await emit(f"🔐 使用 GitHub Token 认证")
|
||||
elif repo_type == "gitlab" and gitlab_token:
|
||||
auth_url = urlunparse(
|
||||
(
|
||||
parsed.scheme,
|
||||
f"oauth2:{gitlab_token}@{parsed.netloc}",
|
||||
parsed.path,
|
||||
parsed.params,
|
||||
parsed.query,
|
||||
parsed.fragment,
|
||||
)
|
||||
)
|
||||
await emit(f"🔐 使用 GitLab Token 认证")
|
||||
elif repo_type == "gitea" and gitea_token:
|
||||
auth_url = urlunparse(
|
||||
(
|
||||
parsed.scheme,
|
||||
f"{gitea_token}@{parsed.netloc}",
|
||||
parsed.path,
|
||||
parsed.params,
|
||||
parsed.query,
|
||||
parsed.fragment,
|
||||
)
|
||||
)
|
||||
await emit(f"🔐 使用 Gitea Token 认证")
|
||||
elif is_ssh_url and ssh_private_key:
|
||||
await emit(f"🔐 使用 SSH Key 认证")
|
||||
|
||||
for branch in branches_to_try:
|
||||
check_cancelled()
|
||||
|
||||
if os.path.exists(base_path) and os.listdir(base_path):
|
||||
shutil.rmtree(base_path)
|
||||
os.makedirs(base_path, exist_ok=True)
|
||||
|
||||
logger.info(f"🔄 尝试克隆分支: {branch}")
|
||||
await emit(f"🔄 尝试克隆分支: {branch}")
|
||||
|
||||
try:
|
||||
# SSH URL使用GitSSHOperations(支持SSH密钥认证)
|
||||
if is_ssh_url and ssh_private_key:
|
||||
|
||||
async def run_ssh_clone():
|
||||
return await asyncio.to_thread(
|
||||
GitSSHOperations.clone_repo_with_ssh,
|
||||
repo_url,
|
||||
ssh_private_key,
|
||||
base_path,
|
||||
branch,
|
||||
)
|
||||
|
||||
clone_task = asyncio.create_task(run_ssh_clone())
|
||||
while not clone_task.done():
|
||||
check_cancelled()
|
||||
try:
|
||||
result = await asyncio.wait_for(
|
||||
asyncio.shield(clone_task), timeout=1.0
|
||||
)
|
||||
break
|
||||
except asyncio.TimeoutError:
|
||||
continue
|
||||
|
||||
if clone_task.done():
|
||||
result = clone_task.result()
|
||||
|
||||
# GitSSHOperations返回字典格式
|
||||
if result.get("success"):
|
||||
logger.info(f"✅ Git 克隆成功 (SSH, 分支: {branch})")
|
||||
await emit(f"✅ 仓库获取成功 (SSH克隆, 分支: {branch})")
|
||||
download_success = True
|
||||
break
|
||||
else:
|
||||
last_error = result.get("message", "未知错误")
|
||||
logger.warning(f"SSH克隆失败 (分支 {branch}): {last_error[:200]}")
|
||||
await emit(f"⚠️ 分支 {branch} SSH克隆失败...", "warning")
|
||||
else:
|
||||
# HTTPS URL使用标准git clone
|
||||
async def run_clone():
|
||||
return await asyncio.to_thread(
|
||||
subprocess.run,
|
||||
[
|
||||
"git",
|
||||
"clone",
|
||||
"--depth",
|
||||
"1",
|
||||
"--branch",
|
||||
branch,
|
||||
auth_url,
|
||||
base_path,
|
||||
],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=120,
|
||||
)
|
||||
|
||||
clone_task = asyncio.create_task(run_clone())
|
||||
while not clone_task.done():
|
||||
check_cancelled()
|
||||
try:
|
||||
result = await asyncio.wait_for(
|
||||
asyncio.shield(clone_task), timeout=1.0
|
||||
)
|
||||
break
|
||||
except asyncio.TimeoutError:
|
||||
continue
|
||||
|
||||
if clone_task.done():
|
||||
result = clone_task.result()
|
||||
|
||||
if result.returncode == 0:
|
||||
logger.info(f"✅ Git 克隆成功 (分支: {branch})")
|
||||
await emit(f"✅ 仓库获取成功 (Git克隆, 分支: {branch})")
|
||||
download_success = True
|
||||
break
|
||||
else:
|
||||
last_error = result.stderr
|
||||
logger.warning(f"克隆失败 (分支 {branch}): {last_error[:200]}")
|
||||
await emit(f"⚠️ 分支 {branch} 克隆失败...", "warning")
|
||||
except subprocess.TimeoutExpired:
|
||||
last_error = f"克隆分支 {branch} 超时"
|
||||
logger.warning(last_error)
|
||||
await emit(f"⚠️ 分支 {branch} 克隆超时...", "warning")
|
||||
except asyncio.CancelledError:
|
||||
logger.info(f"[Cancel] Git clone cancelled for task {task_id}")
|
||||
raise
|
||||
|
||||
# 尝试默认分支
|
||||
if not download_success:
|
||||
check_cancelled()
|
||||
await emit(f"🔄 尝试使用仓库默认分支...")
|
||||
|
||||
if os.path.exists(base_path) and os.listdir(base_path):
|
||||
shutil.rmtree(base_path)
|
||||
os.makedirs(base_path, exist_ok=True)
|
||||
|
||||
try:
|
||||
# SSH URL使用GitSSHOperations(不指定分支)
|
||||
if is_ssh_url and ssh_private_key:
|
||||
|
||||
async def run_default_ssh_clone():
|
||||
return await asyncio.to_thread(
|
||||
GitSSHOperations.clone_repo_with_ssh,
|
||||
repo_url,
|
||||
ssh_private_key,
|
||||
base_path,
|
||||
branch,
|
||||
)
|
||||
|
||||
clone_task = asyncio.create_task(run_default_ssh_clone())
|
||||
while not clone_task.done():
|
||||
check_cancelled()
|
||||
try:
|
||||
result = await asyncio.wait_for(
|
||||
asyncio.shield(clone_task), timeout=1.0
|
||||
)
|
||||
break
|
||||
except asyncio.TimeoutError:
|
||||
continue
|
||||
|
||||
if clone_task.done():
|
||||
result = clone_task.result()
|
||||
|
||||
if result.get("success"):
|
||||
logger.info(f"✅ Git 克隆成功 (SSH, 默认分支)")
|
||||
await emit(f"✅ 仓库获取成功 (SSH克隆, 默认分支)")
|
||||
download_success = True
|
||||
else:
|
||||
last_error = result.get("message", "未知错误")
|
||||
else:
|
||||
# HTTPS URL使用标准git clone
|
||||
async def run_default_clone():
|
||||
return await asyncio.to_thread(
|
||||
subprocess.run,
|
||||
["git", "clone", "--depth", "1", auth_url, base_path],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=120,
|
||||
)
|
||||
|
||||
clone_task = asyncio.create_task(run_default_clone())
|
||||
while not clone_task.done():
|
||||
check_cancelled()
|
||||
try:
|
||||
result = await asyncio.wait_for(
|
||||
asyncio.shield(clone_task), timeout=1.0
|
||||
)
|
||||
break
|
||||
except asyncio.TimeoutError:
|
||||
continue
|
||||
|
||||
if clone_task.done():
|
||||
result = clone_task.result()
|
||||
|
||||
if result.returncode == 0:
|
||||
logger.info(f"✅ Git 克隆成功 (默认分支)")
|
||||
await emit(f"✅ 仓库获取成功 (Git克隆, 默认分支)")
|
||||
download_success = True
|
||||
else:
|
||||
last_error = result.stderr
|
||||
except subprocess.TimeoutExpired:
|
||||
last_error = "克隆超时"
|
||||
except asyncio.CancelledError:
|
||||
logger.info(f"[Cancel] Git clone cancelled for task {task_id}")
|
||||
raise
|
||||
|
||||
if not download_success:
|
||||
# 分析错误原因
|
||||
error_msg = "克隆仓库失败"
|
||||
if "Authentication failed" in last_error or "401" in last_error:
|
||||
error_msg = "认证失败,请检查 GitHub/GitLab Token 配置"
|
||||
elif "not found" in last_error.lower() or "404" in last_error:
|
||||
error_msg = "仓库不存在或无访问权限"
|
||||
elif "Could not resolve host" in last_error:
|
||||
error_msg = "无法解析主机名,请检查网络连接"
|
||||
elif "Permission denied" in last_error or "403" in last_error:
|
||||
error_msg = "无访问权限,请检查仓库权限或 Token"
|
||||
else:
|
||||
error_msg = f"克隆仓库失败: {last_error[:200]}"
|
||||
|
||||
logger.error(f"❌ {error_msg}")
|
||||
await emit(f"❌ {error_msg}", "error")
|
||||
raise RuntimeError(error_msg)
|
||||
|
||||
# 验证目录不为空
|
||||
if not os.listdir(base_path):
|
||||
await emit(f"❌ 项目目录为空", "error")
|
||||
raise RuntimeError(f"项目目录为空,可能是克隆/解压失败: {base_path}")
|
||||
|
||||
# 🔥 智能检测:如果解压后只有一个子目录(常见于 ZIP 文件),
|
||||
# 则使用那个子目录作为真正的项目根目录
|
||||
# 例如:/tmp/deepaudit/UUID/PHP-Project/ -> 返回 /tmp/deepaudit/UUID/PHP-Project
|
||||
items = os.listdir(base_path)
|
||||
# 过滤掉 macOS 产生的 __MACOSX 目录和隐藏文件
|
||||
real_items = [item for item in items if not item.startswith("__") and not item.startswith(".")]
|
||||
|
||||
if len(real_items) == 1:
|
||||
single_item_path = os.path.join(base_path, real_items[0])
|
||||
if os.path.isdir(single_item_path):
|
||||
logger.info(
|
||||
f"🔍 检测到单层嵌套目录,自动调整项目根目录: {base_path} -> {single_item_path}"
|
||||
)
|
||||
await emit(f"🔍 检测到嵌套目录,自动调整为: {real_items[0]}")
|
||||
base_path = single_item_path
|
||||
|
||||
await emit(f"📁 项目准备完成: {base_path}")
|
||||
await save_checkpoint(task_id, project_root=base_path)
|
||||
return base_path
|
||||
@@ -0,0 +1,48 @@
|
||||
"""Reuse completed file analyses after worker interruption; commit files atomically."""
|
||||
|
||||
import asyncio
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
from app.services.audit_queue import workspace_for
|
||||
|
||||
|
||||
async def analyze_file(
|
||||
task_id,
|
||||
file_path,
|
||||
content,
|
||||
service,
|
||||
language,
|
||||
*,
|
||||
rule_set_id=None,
|
||||
prompt_template_id=None,
|
||||
db=None,
|
||||
):
|
||||
key = hashlib.sha256((file_path + "\0" + content).encode()).hexdigest()
|
||||
cache = workspace_for(task_id) / "analysis" / f"{key}.json"
|
||||
if cache.exists():
|
||||
try:
|
||||
return json.loads(await asyncio.to_thread(cache.read_text))
|
||||
except (ValueError, OSError):
|
||||
pass
|
||||
if rule_set_id or prompt_template_id:
|
||||
result = await service.analyze_code_with_rules(
|
||||
content,
|
||||
language,
|
||||
rule_set_id=rule_set_id,
|
||||
prompt_template_id=prompt_template_id,
|
||||
db_session=db,
|
||||
)
|
||||
else:
|
||||
result = await service.analyze_code(content, language)
|
||||
|
||||
def persist():
|
||||
cache.parent.mkdir(parents=True, exist_ok=True)
|
||||
temporary = cache.with_suffix(".tmp")
|
||||
with open(temporary, "w", encoding="utf-8") as output:
|
||||
os.chmod(temporary, 0o600)
|
||||
json.dump(result, output, ensure_ascii=False)
|
||||
os.replace(temporary, cache)
|
||||
|
||||
await asyncio.to_thread(persist)
|
||||
return result
|
||||
@@ -0,0 +1,134 @@
|
||||
"""Bounded ZIP ingestion shared by uploads, repository downloads and workers."""
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import stat
|
||||
import tempfile
|
||||
import zipfile
|
||||
from pathlib import Path, PurePosixPath
|
||||
from typing import Callable
|
||||
|
||||
from app.core.config import settings
|
||||
|
||||
|
||||
class ArchiveError(ValueError):
|
||||
pass
|
||||
|
||||
|
||||
def validate_zip(archive: zipfile.ZipFile) -> list[zipfile.ZipInfo]:
|
||||
entries = archive.infolist()
|
||||
if len(entries) > settings.ZIP_MAX_ENTRIES:
|
||||
raise ArchiveError("ZIP contains too many entries")
|
||||
total = 0
|
||||
names = set()
|
||||
for entry in entries:
|
||||
name = entry.filename.replace("\\", "/")
|
||||
path = PurePosixPath(name)
|
||||
if (
|
||||
path.is_absolute()
|
||||
or ".." in path.parts
|
||||
or ":" in name
|
||||
or "\x00" in name
|
||||
or not path.parts
|
||||
):
|
||||
raise ArchiveError("ZIP contains an unsafe path")
|
||||
mode = entry.external_attr >> 16
|
||||
if stat.S_ISLNK(mode) or (
|
||||
stat.S_IFMT(mode) and not (stat.S_ISREG(mode) or stat.S_ISDIR(mode))
|
||||
):
|
||||
raise ArchiveError("ZIP links and special files are not allowed")
|
||||
canonical = str(path).casefold()
|
||||
if canonical in names:
|
||||
raise ArchiveError("ZIP contains duplicate paths")
|
||||
names.add(canonical)
|
||||
if entry.flag_bits & 1:
|
||||
raise ArchiveError("Encrypted ZIP files are not supported")
|
||||
if entry.file_size > settings.ZIP_MAX_FILE_BYTES:
|
||||
raise ArchiveError("ZIP entry exceeds the per-file quota")
|
||||
total += entry.file_size
|
||||
if total > settings.ZIP_MAX_EXTRACTED_BYTES:
|
||||
raise ArchiveError("ZIP exceeds the total extraction quota")
|
||||
if entry.file_size > max(entry.compress_size, 1) * settings.ZIP_MAX_COMPRESSION_RATIO:
|
||||
raise ArchiveError("ZIP compression ratio exceeds the quota")
|
||||
return entries
|
||||
|
||||
|
||||
def extract_zip(
|
||||
archive: zipfile.ZipFile,
|
||||
destination: str,
|
||||
*,
|
||||
strip_root: bool = False,
|
||||
cancel_check: Callable[[], bool] | None = None,
|
||||
) -> None:
|
||||
entries = validate_zip(archive)
|
||||
root = Path(destination).resolve()
|
||||
root.mkdir(parents=True, exist_ok=True)
|
||||
if any(root.iterdir()):
|
||||
raise ArchiveError("ZIP destination must be empty")
|
||||
useful = [e for e in entries if not e.filename.startswith("__MACOSX/")]
|
||||
parts = [PurePosixPath(e.filename.replace("\\", "/")).parts for e in useful]
|
||||
prefix = None
|
||||
if strip_root and parts and len({p[0] for p in parts}) == 1:
|
||||
first = parts[0][0]
|
||||
if all(len(p) > 1 or e.is_dir() for p, e in zip(parts, useful)):
|
||||
prefix = first
|
||||
# Stage writes so a corrupt archive, quota violation or cancellation leaves no partial tree.
|
||||
with tempfile.TemporaryDirectory(prefix=".extract-", dir=root.parent) as staging:
|
||||
stage = Path(staging)
|
||||
total = 0
|
||||
for entry in useful:
|
||||
if cancel_check and cancel_check():
|
||||
raise asyncio.CancelledError("ZIP extraction cancelled")
|
||||
path = PurePosixPath(entry.filename.replace("\\", "/"))
|
||||
if prefix:
|
||||
path = PurePosixPath(*path.parts[1:])
|
||||
if str(path) == ".":
|
||||
continue
|
||||
target = stage.joinpath(*path.parts)
|
||||
if entry.is_dir():
|
||||
target.mkdir(parents=True, exist_ok=True)
|
||||
continue
|
||||
target.parent.mkdir(parents=True, exist_ok=True)
|
||||
written = 0
|
||||
with archive.open(entry) as source, target.open("xb") as output:
|
||||
while chunk := source.read(64 * 1024):
|
||||
if cancel_check and cancel_check():
|
||||
raise asyncio.CancelledError("ZIP extraction cancelled")
|
||||
written += len(chunk)
|
||||
total += len(chunk)
|
||||
if (
|
||||
written > settings.ZIP_MAX_FILE_BYTES
|
||||
or total > settings.ZIP_MAX_EXTRACTED_BYTES
|
||||
):
|
||||
raise ArchiveError("ZIP extraction quota exceeded")
|
||||
output.write(chunk)
|
||||
for child in stage.iterdir():
|
||||
os.replace(child, root / child.name)
|
||||
|
||||
|
||||
async def save_upload(upload, destination: str) -> None:
|
||||
"""Apply the quota while reading, before the whole request reaches disk."""
|
||||
import aiofiles
|
||||
|
||||
try:
|
||||
total = 0
|
||||
async with aiofiles.open(destination, "wb") as output:
|
||||
while chunk := await upload.read(64 * 1024):
|
||||
total += len(chunk)
|
||||
if total > settings.ZIP_MAX_UPLOAD_BYTES:
|
||||
raise ArchiveError("ZIP upload exceeds the quota")
|
||||
await output.write(chunk)
|
||||
|
||||
def validate():
|
||||
with zipfile.ZipFile(destination) as archive:
|
||||
validate_zip(archive)
|
||||
|
||||
await asyncio.to_thread(validate)
|
||||
except (ArchiveError, zipfile.BadZipFile) as exc:
|
||||
Path(destination).unlink(missing_ok=True)
|
||||
from fastapi import HTTPException
|
||||
|
||||
raise HTTPException(status_code=400, detail=str(exc)) from exc
|
||||
except BaseException:
|
||||
Path(destination).unlink(missing_ok=True)
|
||||
raise
|
||||
@@ -0,0 +1,105 @@
|
||||
"""Transactional queue operations. Task creation and enqueue share one commit."""
|
||||
|
||||
import asyncio
|
||||
import hashlib
|
||||
import shutil
|
||||
from pathlib import Path
|
||||
from datetime import datetime, timezone
|
||||
from sqlalchemy import update
|
||||
|
||||
from app.core.config import settings
|
||||
from app.db.session import AsyncSessionLocal
|
||||
from app.models.audit_job import AuditJob
|
||||
|
||||
|
||||
def workspace_for(task_id: str) -> Path:
|
||||
from uuid import UUID
|
||||
|
||||
UUID(task_id) # Paths never accept arbitrary request text.
|
||||
return Path(settings.AUDIT_WORKSPACE_PATH).resolve() / task_id
|
||||
|
||||
|
||||
async def enqueue(db, task, kind: str, *, zip_path: str | None = None):
|
||||
await db.flush()
|
||||
import json
|
||||
from app.core.encryption import encrypt_sensitive_data
|
||||
from app.models.project import Project
|
||||
from app.services.user_configuration import get_user_config_dict
|
||||
|
||||
configuration = await get_user_config_dict(db, task.created_by)
|
||||
payload = {"user_config": encrypt_sensitive_data(json.dumps(configuration))}
|
||||
project = await db.get(Project, task.project_id)
|
||||
if project:
|
||||
payload["project"] = {
|
||||
name: getattr(project, name)
|
||||
for name in (
|
||||
"id",
|
||||
"name",
|
||||
"source_type",
|
||||
"repository_url",
|
||||
"repository_type",
|
||||
"default_branch",
|
||||
)
|
||||
}
|
||||
if zip_path:
|
||||
directory = workspace_for(task.id)
|
||||
directory.mkdir(parents=True, exist_ok=True)
|
||||
destination = directory / "source.zip"
|
||||
await asyncio.to_thread(shutil.copyfile, zip_path, destination)
|
||||
payload["source_zip"] = str(destination)
|
||||
db.add(AuditJob(id=task.id, kind=kind, payload=payload))
|
||||
|
||||
|
||||
async def project_snapshot(task_id: str, fallback):
|
||||
from types import SimpleNamespace
|
||||
|
||||
async with AsyncSessionLocal() as db:
|
||||
job = await db.get(AuditJob, task_id)
|
||||
project = (job.payload or {}).get("project") if job else None
|
||||
return SimpleNamespace(**project) if project else fallback
|
||||
|
||||
|
||||
async def request_cancel(db, task_id: str):
|
||||
await db.execute(
|
||||
update(AuditJob)
|
||||
.where(AuditJob.id == task_id)
|
||||
.values(
|
||||
cancel_requested=True,
|
||||
updated_at=datetime.now(timezone.utc),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
async def load_checkpoint(task_id: str) -> dict:
|
||||
async with AsyncSessionLocal() as db:
|
||||
job = await db.get(AuditJob, task_id)
|
||||
return dict(job.checkpoint or {}) if job else {}
|
||||
|
||||
|
||||
async def save_checkpoint(task_id: str, **changes):
|
||||
async with AsyncSessionLocal() as db:
|
||||
job = await db.get(AuditJob, task_id, with_for_update=True)
|
||||
if job:
|
||||
job.checkpoint = {**(job.checkpoint or {}), **changes}
|
||||
if job.kind == "agent":
|
||||
import json
|
||||
from app.models.agent_task import AgentCheckpoint
|
||||
|
||||
stage = "result" if "result" in changes else "workspace"
|
||||
db.add(
|
||||
AgentCheckpoint(
|
||||
task_id=task_id,
|
||||
agent_id=task_id,
|
||||
agent_name="Audit worker",
|
||||
agent_type="worker",
|
||||
status="running",
|
||||
checkpoint_type="auto",
|
||||
checkpoint_name=stage,
|
||||
state_data=json.dumps(changes, ensure_ascii=False),
|
||||
)
|
||||
)
|
||||
await db.commit()
|
||||
|
||||
|
||||
def advisory_key(task_id: str) -> int:
|
||||
return int.from_bytes(hashlib.sha256(task_id.encode()).digest()[:8], "big", signed=True)
|
||||
@@ -0,0 +1,73 @@
|
||||
"""Database-side aggregates; dashboards never load every task and issue."""
|
||||
|
||||
from sqlalchemy import select, func, union_all
|
||||
from app.models.project import Project
|
||||
from app.models.audit import AuditTask, AuditIssue
|
||||
from app.models.agent_task import AgentTask, AgentFinding
|
||||
|
||||
|
||||
async def project_stats(db, user_id):
|
||||
projects = select(Project.id).where(Project.owner_id == user_id)
|
||||
total_projects, active_projects = (
|
||||
await db.execute(
|
||||
select(
|
||||
func.count(),
|
||||
func.count().filter(Project.is_active.is_(True)),
|
||||
)
|
||||
.select_from(Project)
|
||||
.where(Project.owner_id == user_id)
|
||||
)
|
||||
).one()
|
||||
tasks = union_all(
|
||||
*[
|
||||
select(model.status.label("status"), model.quality_score.label("score")).where(
|
||||
model.project_id.in_(projects)
|
||||
)
|
||||
for model in (AuditTask, AgentTask)
|
||||
]
|
||||
).subquery()
|
||||
total_tasks, completed_tasks, average = (
|
||||
await db.execute(
|
||||
select(
|
||||
func.count(),
|
||||
func.count().filter(tasks.c.status == "completed"),
|
||||
func.avg(tasks.c.score).filter(
|
||||
(tasks.c.status == "completed") & (tasks.c.score > 0)
|
||||
),
|
||||
).select_from(tasks)
|
||||
)
|
||||
).one()
|
||||
issues = union_all(
|
||||
select(
|
||||
AuditIssue.issue_type.label("kind"), (AuditIssue.status == "resolved").label("resolved")
|
||||
).where(
|
||||
AuditIssue.task_id.in_(select(AuditTask.id).where(AuditTask.project_id.in_(projects)))
|
||||
),
|
||||
select(
|
||||
AgentFinding.vulnerability_type.label("kind"),
|
||||
AgentFinding.status.in_(["fixed", "wont_fix", "false_positive"]).label("resolved"),
|
||||
).where(
|
||||
AgentFinding.task_id.in_(select(AgentTask.id).where(AgentTask.project_id.in_(projects)))
|
||||
),
|
||||
).subquery()
|
||||
total_issues, resolved_issues = (
|
||||
await db.execute(
|
||||
select(
|
||||
func.count(),
|
||||
func.count().filter(issues.c.resolved.is_(True)),
|
||||
).select_from(issues)
|
||||
)
|
||||
).one()
|
||||
issue_types = dict(
|
||||
(await db.execute(select(issues.c.kind, func.count()).group_by(issues.c.kind))).all()
|
||||
)
|
||||
return dict(
|
||||
total_projects=total_projects,
|
||||
active_projects=active_projects,
|
||||
total_tasks=total_tasks,
|
||||
completed_tasks=completed_tasks,
|
||||
total_issues=total_issues,
|
||||
resolved_issues=resolved_issues,
|
||||
avg_quality_score=float(average or 0),
|
||||
issue_types=issue_types,
|
||||
)
|
||||
@@ -15,8 +15,6 @@ import base64
|
||||
if sys.platform == 'darwin':
|
||||
os.environ['DYLD_FALLBACK_LIBRARY_PATH'] = '/opt/homebrew/lib:' + os.environ.get('DYLD_FALLBACK_LIBRARY_PATH', '')
|
||||
|
||||
from weasyprint import HTML, CSS
|
||||
from weasyprint.text.fonts import FontConfiguration
|
||||
from jinja2 import Template
|
||||
|
||||
class ReportGenerator:
|
||||
@@ -452,6 +450,10 @@ class ReportGenerator:
|
||||
|
||||
@classmethod
|
||||
def _render_pdf(cls, context: Dict[str, Any]) -> bytes:
|
||||
# Load native PDF dependencies only when rendering a PDF.
|
||||
from weasyprint import HTML
|
||||
from weasyprint.text.fonts import FontConfiguration
|
||||
|
||||
# 注入 Logo
|
||||
context['logo_b64'] = cls._get_logo_base64()
|
||||
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
from app.services.analysis_checkpoint import analyze_file
|
||||
"""
|
||||
仓库扫描服务 - 支持GitHub, GitLab 和 Gitea 仓库扫描
|
||||
"""
|
||||
@@ -319,6 +320,8 @@ async def scan_repo_task(task_id: str, db_session_factory, user_config: dict = N
|
||||
|
||||
# 2. 获取项目信息
|
||||
project = await db.get(Project, task.project_id)
|
||||
from app.services.audit_queue import project_snapshot
|
||||
project = await project_snapshot(task_id, project)
|
||||
if not project:
|
||||
raise Exception("项目不存在")
|
||||
|
||||
@@ -528,15 +531,10 @@ async def scan_repo_task(task_id: str, db_session_factory, user_config: dict = N
|
||||
rule_set_id = scan_config.get('rule_set_id')
|
||||
prompt_template_id = scan_config.get('prompt_template_id')
|
||||
|
||||
if rule_set_id or prompt_template_id:
|
||||
analysis = await llm_service.analyze_code_with_rules(
|
||||
content, language,
|
||||
rule_set_id=rule_set_id,
|
||||
prompt_template_id=prompt_template_id,
|
||||
db_session=db
|
||||
)
|
||||
else:
|
||||
analysis = await llm_service.analyze_code(content, language)
|
||||
analysis = await analyze_file(
|
||||
task_id, file_info['path'], content, llm_service, language,
|
||||
rule_set_id=rule_set_id, prompt_template_id=prompt_template_id, db=db,
|
||||
)
|
||||
print(f"✅ LLM 分析完成: {file_info['path']}")
|
||||
|
||||
# 再次检查是否取消(LLM分析后)
|
||||
|
||||
@@ -0,0 +1,51 @@
|
||||
import json
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
from app.models.user_config import UserConfig
|
||||
|
||||
|
||||
async def get_user_config_dict(db: AsyncSession, user_id: str) -> dict:
|
||||
"""获取用户配置字典(包含解密敏感字段)"""
|
||||
from app.core.encryption import decrypt_sensitive_data
|
||||
|
||||
# 需要解密的敏感字段列表(与 config.py 保持一致)
|
||||
SENSITIVE_LLM_FIELDS = [
|
||||
"llmApiKey",
|
||||
"geminiApiKey",
|
||||
"openaiApiKey",
|
||||
"claudeApiKey",
|
||||
"qwenApiKey",
|
||||
"deepseekApiKey",
|
||||
"zhipuApiKey",
|
||||
"moonshotApiKey",
|
||||
"baiduApiKey",
|
||||
"minimaxApiKey",
|
||||
"doubaoApiKey",
|
||||
]
|
||||
SENSITIVE_OTHER_FIELDS = ["githubToken", "gitlabToken", "giteaToken"]
|
||||
|
||||
def decrypt_config(config: dict, sensitive_fields: list) -> dict:
|
||||
"""解密配置中的敏感字段"""
|
||||
decrypted = config.copy()
|
||||
for field in sensitive_fields:
|
||||
if field in decrypted and decrypted[field]:
|
||||
decrypted[field] = decrypt_sensitive_data(decrypted[field])
|
||||
return decrypted
|
||||
|
||||
result = await db.execute(select(UserConfig).where(UserConfig.user_id == user_id))
|
||||
config = result.scalar_one_or_none()
|
||||
if not config:
|
||||
return {}
|
||||
|
||||
# 解析配置
|
||||
llm_config = json.loads(config.llm_config) if config.llm_config else {}
|
||||
other_config = json.loads(config.other_config) if config.other_config else {}
|
||||
|
||||
# 解密敏感字段
|
||||
llm_config = decrypt_config(llm_config, SENSITIVE_LLM_FIELDS)
|
||||
other_config = decrypt_config(other_config, SENSITIVE_OTHER_FIELDS)
|
||||
|
||||
return {
|
||||
"llmConfig": llm_config,
|
||||
"otherConfig": other_config,
|
||||
}
|
||||
@@ -0,0 +1,247 @@
|
||||
from app.services.analysis_checkpoint import analyze_file
|
||||
from app.services.archive import extract_zip
|
||||
from datetime import datetime, timezone
|
||||
import shutil
|
||||
import os
|
||||
import json
|
||||
from pathlib import Path
|
||||
import zipfile
|
||||
import asyncio
|
||||
|
||||
from app.models.audit import AuditTask, AuditIssue
|
||||
from app.services.llm.service import LLMService
|
||||
from app.services.scanner import (
|
||||
task_control,
|
||||
is_text_file,
|
||||
should_exclude,
|
||||
get_language_from_path,
|
||||
get_analysis_config,
|
||||
)
|
||||
from app.core.config import settings
|
||||
|
||||
|
||||
def normalize_path(path: str) -> str:
|
||||
"""
|
||||
统一路径分隔符为正斜杠,确保跨平台兼容性
|
||||
Windows 使用反斜杠 (\),Unix/Mac 使用正斜杠 (/)
|
||||
统一转换为正斜杠以保证一致性
|
||||
"""
|
||||
return path.replace("\\", "/")
|
||||
|
||||
|
||||
async def process_zip_task(
|
||||
task_id: str,
|
||||
file_path: str,
|
||||
db_session_factory,
|
||||
user_config: dict = None,
|
||||
*,
|
||||
project_root: str | None = None,
|
||||
):
|
||||
"""后台ZIP文件处理任务"""
|
||||
async with db_session_factory() as db:
|
||||
task = await db.get(AuditTask, task_id)
|
||||
if not task:
|
||||
return
|
||||
|
||||
extract_dir = Path(project_root) if project_root else Path(f"/tmp/{task_id}")
|
||||
try:
|
||||
task.status = "running"
|
||||
task.started_at = datetime.now(timezone.utc)
|
||||
await db.commit()
|
||||
|
||||
# 创建使用用户配置的LLM服务实例
|
||||
llm_service = LLMService(user_config=user_config or {})
|
||||
|
||||
if not project_root:
|
||||
if extract_dir.exists():
|
||||
await asyncio.to_thread(shutil.rmtree, extract_dir)
|
||||
extract_dir.mkdir(parents=True, exist_ok=True)
|
||||
with zipfile.ZipFile(file_path, "r") as zip_ref:
|
||||
await asyncio.to_thread(
|
||||
extract_zip,
|
||||
zip_ref,
|
||||
str(extract_dir),
|
||||
cancel_check=lambda: task_control.is_cancelled(task_id),
|
||||
)
|
||||
|
||||
# 获取用户自定义排除模式
|
||||
scan_config = (user_config or {}).get("scan_config", {})
|
||||
custom_exclude_patterns = scan_config.get("exclude_patterns", [])
|
||||
|
||||
# Find files
|
||||
files_to_scan = []
|
||||
for root, dirs, files in os.walk(extract_dir):
|
||||
dirs.sort()
|
||||
files.sort()
|
||||
# 排除常见非代码目录
|
||||
dirs[:] = [
|
||||
d
|
||||
for d in dirs
|
||||
if d not in ["node_modules", "__pycache__", ".git", "dist", "build", "vendor"]
|
||||
]
|
||||
|
||||
for file in files:
|
||||
full_path = Path(root) / file
|
||||
# 统一使用正斜杠,确保跨平台兼容性
|
||||
rel_path = normalize_path(str(full_path.relative_to(extract_dir)))
|
||||
|
||||
# 检查文件类型和排除规则(包含用户自定义排除模式)
|
||||
if (
|
||||
not full_path.is_symlink()
|
||||
and is_text_file(rel_path)
|
||||
and not should_exclude(rel_path, custom_exclude_patterns)
|
||||
):
|
||||
try:
|
||||
if full_path.stat().st_size > settings.MAX_FILE_SIZE_BYTES:
|
||||
continue
|
||||
if full_path.stat().st_size <= settings.MAX_FILE_SIZE_BYTES:
|
||||
files_to_scan.append(
|
||||
{
|
||||
"path": rel_path,
|
||||
}
|
||||
)
|
||||
except:
|
||||
pass
|
||||
|
||||
# 获取分析配置(优先使用用户配置)
|
||||
analysis_config = get_analysis_config(user_config)
|
||||
max_analyze_files = analysis_config["max_analyze_files"]
|
||||
llm_gap_ms = analysis_config["llm_gap_ms"]
|
||||
|
||||
# 限制文件数量
|
||||
# 如果指定了特定文件,则只分析这些文件
|
||||
target_files = scan_config.get("file_paths", [])
|
||||
if target_files:
|
||||
# 统一目标文件路径的分隔符,确保匹配一致性
|
||||
normalized_targets = {normalize_path(p) for p in target_files}
|
||||
print(f"🎯 ZIP任务: 指定分析 {len(normalized_targets)} 个文件")
|
||||
files_to_scan = [f for f in files_to_scan if f["path"] in normalized_targets]
|
||||
elif max_analyze_files > 0:
|
||||
files_to_scan = files_to_scan[:max_analyze_files]
|
||||
|
||||
task.total_files = len(files_to_scan)
|
||||
await db.commit()
|
||||
|
||||
print(
|
||||
f"📊 ZIP任务 {task_id}: 找到 {len(files_to_scan)} 个文件 (最大文件数: {max_analyze_files}, 请求间隔: {llm_gap_ms}ms)"
|
||||
)
|
||||
|
||||
total_issues = 0
|
||||
total_lines = 0
|
||||
quality_scores = []
|
||||
scanned_files = 0
|
||||
failed_files = 0
|
||||
|
||||
for file_info in files_to_scan:
|
||||
# 检查是否取消
|
||||
if task_control.is_cancelled(task_id):
|
||||
print(f"🛑 ZIP任务 {task_id} 已被取消")
|
||||
task.status = "cancelled"
|
||||
task.completed_at = datetime.now(timezone.utc)
|
||||
await db.commit()
|
||||
task_control.cleanup_task(task_id)
|
||||
return
|
||||
|
||||
try:
|
||||
content = await asyncio.to_thread(
|
||||
(extract_dir / file_info["path"]).read_text, errors="ignore"
|
||||
)
|
||||
total_lines += content.count("\n") + 1
|
||||
language = get_language_from_path(file_info["path"])
|
||||
|
||||
# 获取规则集和提示词模板ID
|
||||
scan_config = (user_config or {}).get("scan_config", {})
|
||||
rule_set_id = scan_config.get("rule_set_id")
|
||||
prompt_template_id = scan_config.get("prompt_template_id")
|
||||
|
||||
# 使用规则集和提示词模板进行分析
|
||||
result = await analyze_file(
|
||||
task_id,
|
||||
file_info["path"],
|
||||
content,
|
||||
llm_service,
|
||||
language,
|
||||
rule_set_id=rule_set_id,
|
||||
prompt_template_id=prompt_template_id,
|
||||
db=db,
|
||||
)
|
||||
|
||||
issues = result.get("issues", [])
|
||||
for i in issues:
|
||||
issue = AuditIssue(
|
||||
task_id=task.id,
|
||||
file_path=file_info["path"],
|
||||
line_number=i.get("line", 1),
|
||||
column_number=i.get("column"),
|
||||
issue_type=i.get("type", "maintainability"),
|
||||
severity=i.get("severity", "low"),
|
||||
title=i.get("title", "Issue"),
|
||||
message=i.get("title", "Issue"),
|
||||
description=i.get("description"),
|
||||
suggestion=i.get("suggestion"),
|
||||
code_snippet=i.get("code_snippet"),
|
||||
ai_explanation=json.dumps(i.get("xai")) if i.get("xai") else None,
|
||||
status="open",
|
||||
)
|
||||
db.add(issue)
|
||||
total_issues += 1
|
||||
|
||||
if "quality_score" in result:
|
||||
quality_scores.append(result["quality_score"])
|
||||
|
||||
scanned_files += 1
|
||||
task.scanned_files = scanned_files
|
||||
task.total_lines = total_lines
|
||||
task.issues_count = total_issues
|
||||
await db.commit()
|
||||
|
||||
print(f"📈 ZIP任务 {task_id}: 进度 {scanned_files}/{len(files_to_scan)}")
|
||||
|
||||
# 请求间隔
|
||||
await asyncio.sleep(llm_gap_ms / 1000)
|
||||
|
||||
except Exception as file_error:
|
||||
failed_files += 1
|
||||
print(f"❌ ZIP任务分析文件失败 ({file_info['path']}): {file_error}")
|
||||
await asyncio.sleep(llm_gap_ms / 1000)
|
||||
|
||||
# 完成任务
|
||||
avg_quality_score = (
|
||||
sum(quality_scores) / len(quality_scores) if quality_scores else 100.0
|
||||
)
|
||||
|
||||
# 如果有文件需要分析但全部失败,标记为失败
|
||||
if len(files_to_scan) > 0 and scanned_files == 0:
|
||||
task.status = "failed"
|
||||
task.completed_at = datetime.now(timezone.utc)
|
||||
task.scanned_files = 0
|
||||
task.total_lines = total_lines
|
||||
task.issues_count = 0
|
||||
task.quality_score = 0
|
||||
await db.commit()
|
||||
print(
|
||||
f"❌ ZIP任务 {task_id} 失败: 所有 {len(files_to_scan)} 个文件分析均失败,请检查 LLM API 配置"
|
||||
)
|
||||
else:
|
||||
task.status = "completed"
|
||||
task.completed_at = datetime.now(timezone.utc)
|
||||
task.scanned_files = scanned_files
|
||||
task.total_lines = total_lines
|
||||
task.issues_count = total_issues
|
||||
task.quality_score = avg_quality_score
|
||||
await db.commit()
|
||||
print(
|
||||
f"✅ ZIP任务 {task_id} 完成: 扫描 {scanned_files} 个文件, 发现 {total_issues} 个问题"
|
||||
)
|
||||
task_control.cleanup_task(task_id)
|
||||
|
||||
except Exception as e:
|
||||
print(f"❌ ZIP扫描失败: {e}")
|
||||
task.status = "failed"
|
||||
task.completed_at = datetime.now(timezone.utc)
|
||||
await db.commit()
|
||||
task_control.cleanup_task(task_id)
|
||||
finally:
|
||||
# Cleanup - 只清理解压目录,不删除源ZIP文件(已持久化存储)
|
||||
if not project_root and extract_dir.exists():
|
||||
await asyncio.to_thread(shutil.rmtree, extract_dir)
|
||||
@@ -1,3 +1,5 @@
|
||||
import asyncio
|
||||
import tempfile
|
||||
"""
|
||||
ZIP文件存储服务
|
||||
用于管理项目的ZIP文件持久化存储
|
||||
@@ -45,7 +47,16 @@ async def save_project_zip(project_id: str, file_path: str, original_filename: s
|
||||
meta_path = get_project_zip_meta_path(project_id)
|
||||
|
||||
# 复制文件到存储目录
|
||||
shutil.copy2(file_path, target_path)
|
||||
def atomic_copy():
|
||||
with tempfile.NamedTemporaryFile(dir=target_path.parent, delete=False) as temporary:
|
||||
temporary_path = Path(temporary.name)
|
||||
try:
|
||||
shutil.copy2(file_path, temporary_path)
|
||||
os.chmod(temporary_path, 0o600)
|
||||
os.replace(temporary_path, target_path)
|
||||
finally:
|
||||
temporary_path.unlink(missing_ok=True)
|
||||
await asyncio.to_thread(atomic_copy)
|
||||
|
||||
# 获取文件大小
|
||||
file_size = os.path.getsize(target_path)
|
||||
|
||||
@@ -0,0 +1,360 @@
|
||||
"""PostgreSQL-backed worker. Each job runs in its own process and holds a DB lock.
|
||||
|
||||
Run `python -m app.worker`; scale by starting more workers. No API process executes jobs.
|
||||
Session advisory locks prevent simultaneous execution, including after a lease expires.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import contextlib
|
||||
import logging
|
||||
import os
|
||||
import signal
|
||||
import socket
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
from datetime import datetime, timedelta, timezone
|
||||
|
||||
from sqlalchemy import select, text, update, or_, delete
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from app.core.config import settings
|
||||
from app.db.session import engine, AsyncSessionLocal
|
||||
from app.models.audit_job import AuditJob
|
||||
from app.models.audit import AuditTask, AuditIssue
|
||||
from app.models.agent_task import AgentTask
|
||||
from app.services.audit_queue import advisory_key
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
TERMINAL = {"completed", "failed", "cancelled"}
|
||||
|
||||
|
||||
def heartbeat_file():
|
||||
return Path(settings.AUDIT_WORKSPACE_PATH) / ".workers" / socket.gethostname()
|
||||
|
||||
|
||||
def touch_heartbeat():
|
||||
path = heartbeat_file()
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
path.touch()
|
||||
|
||||
|
||||
async def cleanup_artifacts():
|
||||
from app.services.audit_queue import workspace_for
|
||||
import shutil
|
||||
|
||||
cutoff = datetime.now(timezone.utc) - timedelta(days=settings.WORKER_RETENTION_DAYS)
|
||||
cursor = ""
|
||||
while True:
|
||||
async with AsyncSessionLocal() as db:
|
||||
ids = list(
|
||||
await db.scalars(
|
||||
select(AuditJob.id)
|
||||
.where(
|
||||
AuditJob.status.in_(TERMINAL),
|
||||
AuditJob.updated_at < cutoff,
|
||||
AuditJob.id > cursor,
|
||||
)
|
||||
.order_by(AuditJob.id)
|
||||
.limit(500)
|
||||
)
|
||||
)
|
||||
if not ids:
|
||||
return
|
||||
for task_id in ids:
|
||||
await asyncio.to_thread(shutil.rmtree, workspace_for(task_id), True)
|
||||
cursor = ids[-1]
|
||||
|
||||
|
||||
def cleanup_sandboxes(job_id):
|
||||
try:
|
||||
import docker
|
||||
|
||||
with docker.from_env() as client:
|
||||
for container in client.containers.list(
|
||||
all=True, filters={"label": f"deepaudit.job_id={job_id}"}
|
||||
):
|
||||
container.remove(force=True)
|
||||
except Exception:
|
||||
logger.debug("Sandbox cleanup unavailable", exc_info=True)
|
||||
|
||||
|
||||
async def candidate_ids() -> list[str]:
|
||||
async with AsyncSessionLocal() as db:
|
||||
result = await db.scalars(
|
||||
select(AuditJob.id)
|
||||
.where(
|
||||
or_(
|
||||
AuditJob.status == "queued",
|
||||
(AuditJob.status == "running")
|
||||
& (AuditJob.lease_until < datetime.now(timezone.utc)),
|
||||
)
|
||||
)
|
||||
.order_by(AuditJob.created_at, AuditJob.id)
|
||||
.limit(100)
|
||||
)
|
||||
return list(result)
|
||||
|
||||
|
||||
async def execute_payload(job: AuditJob):
|
||||
from app.core.encryption import decrypt_sensitive_data
|
||||
from app.services.user_configuration import get_user_config_dict
|
||||
import json
|
||||
|
||||
model = AgentTask if job.kind == "agent" else AuditTask
|
||||
async with AsyncSessionLocal() as db:
|
||||
task = await db.get(model, job.id)
|
||||
encoded_config = (job.payload or {}).get("user_config")
|
||||
config = (
|
||||
json.loads(decrypt_sensitive_data(encoded_config))
|
||||
if encoded_config
|
||||
else await get_user_config_dict(db, task.created_by)
|
||||
)
|
||||
if job.kind == "agent":
|
||||
from app.services.agent_audit.execution import _execute_agent_task
|
||||
|
||||
await _execute_agent_task(job.id, user_config_override=config)
|
||||
return
|
||||
from app.services.zip_storage import load_project_zip
|
||||
|
||||
async with AsyncSessionLocal() as db:
|
||||
task = await db.get(AuditTask, job.id)
|
||||
config["scan_config"] = json.loads(task.scan_config or "{}")
|
||||
# Per-file LLM outputs are cached durably; replay replaces derived rows.
|
||||
await db.execute(delete(AuditIssue).where(AuditIssue.task_id == task.id))
|
||||
task.scanned_files = task.total_lines = task.issues_count = 0
|
||||
await db.commit()
|
||||
project_id = task.project_id
|
||||
if job.kind == "zip":
|
||||
from app.services.zip_scan import process_zip_task
|
||||
|
||||
source = (job.payload or {}).get("source_zip") or await load_project_zip(project_id)
|
||||
if not source:
|
||||
raise ValueError("ZIP source is no longer available")
|
||||
await process_zip_task(job.id, source, AsyncSessionLocal, config)
|
||||
elif job.kind == "repository":
|
||||
from app.models.project import Project
|
||||
from app.services.agent_audit.workspace import _get_project_root
|
||||
from app.services.zip_scan import process_zip_task
|
||||
|
||||
async with AsyncSessionLocal() as db:
|
||||
project = await db.get(Project, project_id)
|
||||
other = config.get("otherConfig", {})
|
||||
ssh_key = other.get("sshPrivateKey")
|
||||
root = await _get_project_root(
|
||||
project,
|
||||
job.id,
|
||||
branch_name=task.branch_name,
|
||||
github_token=other.get("githubToken") or settings.GITHUB_TOKEN,
|
||||
gitlab_token=other.get("gitlabToken") or settings.GITLAB_TOKEN,
|
||||
gitea_token=other.get("giteaToken") or settings.GITEA_TOKEN,
|
||||
ssh_private_key=decrypt_sensitive_data(ssh_key) if ssh_key else None,
|
||||
)
|
||||
await process_zip_task(job.id, None, AsyncSessionLocal, config, project_root=root)
|
||||
else:
|
||||
raise ValueError(f"Unknown job kind: {job.kind}")
|
||||
|
||||
|
||||
async def run_job(job_id: str) -> bool:
|
||||
"""Return False if another process owns the job. Lock lasts until child exit."""
|
||||
async with engine.connect() as connection:
|
||||
acquired = await connection.scalar(
|
||||
text("SELECT pg_try_advisory_lock(:key)"), {"key": advisory_key(job_id)}
|
||||
)
|
||||
await connection.commit()
|
||||
if not acquired:
|
||||
return False
|
||||
try:
|
||||
async with AsyncSession(bind=connection, expire_on_commit=False) as db:
|
||||
job = await db.get(AuditJob, job_id, with_for_update=True)
|
||||
if not job or job.status in TERMINAL:
|
||||
return True
|
||||
model = AgentTask if job.kind == "agent" else AuditTask
|
||||
task = await db.get(model, job_id)
|
||||
if not task:
|
||||
job.status = "cancelled"
|
||||
await db.commit()
|
||||
return True
|
||||
if task.status in TERMINAL:
|
||||
job.status = task.status
|
||||
await db.commit()
|
||||
return True
|
||||
if job.cancel_requested or job.attempts >= settings.WORKER_MAX_ATTEMPTS:
|
||||
job.status = task.status = "cancelled" if job.cancel_requested else "failed"
|
||||
task.completed_at = datetime.now(timezone.utc)
|
||||
job.error = (
|
||||
"Cancelled" if job.cancel_requested else "Worker recovery limit exceeded"
|
||||
)
|
||||
await db.commit()
|
||||
return True
|
||||
job.attempts += 1
|
||||
job.status = "running"
|
||||
job.worker_id = f"{socket.gethostname()}:{os.getpid()}"
|
||||
job.lease_until = datetime.now(timezone.utc) + timedelta(
|
||||
seconds=settings.WORKER_LEASE_SECONDS
|
||||
)
|
||||
job.updated_at = datetime.now(timezone.utc)
|
||||
await db.commit()
|
||||
timeout = (
|
||||
getattr(task, "timeout_seconds", None) or settings.WORKER_TASK_TIMEOUT_SECONDS
|
||||
)
|
||||
|
||||
os.environ["DEEPAUDIT_JOB_ID"] = job_id
|
||||
if settings.SANDBOX_ENABLED:
|
||||
await asyncio.to_thread(cleanup_sandboxes, job_id)
|
||||
touch_heartbeat()
|
||||
work = asyncio.create_task(execute_payload(job))
|
||||
monitor_error = None
|
||||
|
||||
async def heartbeat():
|
||||
nonlocal monitor_error
|
||||
while not work.done():
|
||||
await asyncio.sleep(settings.WORKER_HEARTBEAT_SECONDS)
|
||||
try:
|
||||
result = await connection.execute(
|
||||
update(AuditJob)
|
||||
.where(AuditJob.id == job_id)
|
||||
.values(
|
||||
lease_until=datetime.now(timezone.utc)
|
||||
+ timedelta(seconds=settings.WORKER_LEASE_SECONDS),
|
||||
updated_at=datetime.now(timezone.utc),
|
||||
)
|
||||
.returning(AuditJob.cancel_requested)
|
||||
)
|
||||
cancelled = result.scalar_one()
|
||||
await connection.commit()
|
||||
touch_heartbeat()
|
||||
except Exception:
|
||||
# The session owns the lock. Do not let execution outlive a lost session.
|
||||
logger.exception("Lost worker lock connection; terminating job process")
|
||||
os._exit(75)
|
||||
from app.services.agent_audit.runtime import _running_event_managers
|
||||
|
||||
try:
|
||||
from app.services.agent_audit.progress import publish_progress
|
||||
|
||||
await publish_progress(job_id)
|
||||
for event_manager in list(_running_event_managers.values()):
|
||||
await event_manager.flush_tokens()
|
||||
except Exception as exc:
|
||||
monitor_error = f"Event persistence failed: {type(exc).__name__}"
|
||||
work.cancel()
|
||||
return
|
||||
if cancelled:
|
||||
from app.services.agent_audit.runtime import _cancelled_tasks
|
||||
from app.services.scanner import task_control
|
||||
|
||||
_cancelled_tasks.add(job_id)
|
||||
task_control.cancel_task(job_id)
|
||||
work.cancel()
|
||||
return
|
||||
|
||||
monitor = asyncio.create_task(heartbeat())
|
||||
error = None
|
||||
try:
|
||||
done, _ = await asyncio.wait({work}, timeout=timeout)
|
||||
if not done:
|
||||
error = f"Task timeout after {timeout} seconds"
|
||||
work.cancel()
|
||||
await asyncio.wait_for(work, timeout=10 if not done else None)
|
||||
except asyncio.CancelledError:
|
||||
error = error or "Execution interrupted"
|
||||
except Exception as exc:
|
||||
error = f"{type(exc).__name__}: {str(exc)[:1000]}"
|
||||
logger.exception("Job %s failed", job_id)
|
||||
finally:
|
||||
monitor.cancel()
|
||||
with contextlib.suppress(asyncio.CancelledError):
|
||||
await monitor
|
||||
error = monitor_error or error
|
||||
await connection.rollback()
|
||||
async with AsyncSession(bind=connection, expire_on_commit=False) as db:
|
||||
current = await db.get(AuditJob, job_id, with_for_update=True)
|
||||
task = await db.get(model, job_id)
|
||||
if current.cancel_requested:
|
||||
current.status = "cancelled"
|
||||
elif error:
|
||||
current.status = "failed"
|
||||
else:
|
||||
current.status = task.status if task and task.status in TERMINAL else "failed"
|
||||
current.error = error
|
||||
current.lease_until = None
|
||||
current.updated_at = datetime.now(timezone.utc)
|
||||
if task:
|
||||
task.status = current.status
|
||||
task.completed_at = datetime.now(timezone.utc)
|
||||
if error and hasattr(task, "error_message"):
|
||||
task.error_message = error
|
||||
await db.commit()
|
||||
return True
|
||||
finally:
|
||||
with contextlib.suppress(Exception):
|
||||
await connection.execute(
|
||||
text("SELECT pg_advisory_unlock(:key)"), {"key": advisory_key(job_id)}
|
||||
)
|
||||
await connection.commit()
|
||||
|
||||
|
||||
async def serve():
|
||||
stopping = asyncio.Event()
|
||||
loop = asyncio.get_running_loop()
|
||||
for sig in (signal.SIGTERM, signal.SIGINT):
|
||||
loop.add_signal_handler(sig, stopping.set)
|
||||
last_cleanup = 0.0
|
||||
while not stopping.is_set():
|
||||
try:
|
||||
touch_heartbeat()
|
||||
if time.monotonic() - last_cleanup >= 3600:
|
||||
await cleanup_artifacts()
|
||||
last_cleanup = time.monotonic()
|
||||
candidates = await candidate_ids()
|
||||
for job_id in candidates:
|
||||
if stopping.is_set():
|
||||
break
|
||||
process = await asyncio.create_subprocess_exec(
|
||||
sys.executable,
|
||||
"-m",
|
||||
"app.worker",
|
||||
"--job",
|
||||
job_id,
|
||||
start_new_session=True,
|
||||
)
|
||||
done = asyncio.create_task(process.wait())
|
||||
stop = asyncio.create_task(stopping.wait())
|
||||
finished, _ = await asyncio.wait(
|
||||
{done, stop},
|
||||
return_when=asyncio.FIRST_COMPLETED,
|
||||
timeout=settings.WORKER_TASK_TIMEOUT_SECONDS + 60,
|
||||
)
|
||||
if not finished:
|
||||
os.killpg(process.pid, signal.SIGKILL)
|
||||
stop.cancel()
|
||||
if stopping.is_set() and process.returncode is None:
|
||||
# Leave status running for recovery; never mark deployment shutdown as user cancellation.
|
||||
os.killpg(process.pid, signal.SIGTERM)
|
||||
try:
|
||||
await asyncio.wait_for(done, timeout=10)
|
||||
except asyncio.TimeoutError:
|
||||
os.killpg(process.pid, signal.SIGKILL)
|
||||
await done
|
||||
else:
|
||||
await done
|
||||
except Exception:
|
||||
logger.exception("Worker polling failed")
|
||||
try:
|
||||
await asyncio.wait_for(stopping.wait(), timeout=settings.WORKER_POLL_SECONDS)
|
||||
except asyncio.TimeoutError:
|
||||
pass
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
if sys.argv[1:] == ["--health"]:
|
||||
try:
|
||||
sys.exit(0 if time.time() - heartbeat_file().stat().st_mtime < 60 else 1)
|
||||
except OSError:
|
||||
sys.exit(1)
|
||||
elif len(sys.argv) == 3 and sys.argv[1] == "--job":
|
||||
asyncio.run(run_job(sys.argv[2]))
|
||||
else:
|
||||
asyncio.run(serve())
|
||||
+26
-2
@@ -20,9 +20,17 @@ POSTGRES_DB=deepaudit
|
||||
# =============================================
|
||||
# 安全配置
|
||||
# =============================================
|
||||
# JWT 签名密钥 - 生产环境必须修改为随机字符串!
|
||||
# 运行环境:production / development / test
|
||||
ENVIRONMENT=production
|
||||
# 演示账户仅允许在 development / test 中显式启用
|
||||
DEMO_ENABLED=false
|
||||
# JWT 签名及配置加密密钥,必须填写至少 32 字符的随机字符串。
|
||||
# 已有部署请保留安全的原密钥;修改密钥会使现有 Token 和加密配置失效。
|
||||
# 建议使用: openssl rand -hex 32
|
||||
SECRET_KEY=your-super-secret-key-change-this-in-production
|
||||
SECRET_KEY=
|
||||
|
||||
# 同域反向代理可保持空数组;跨域开发按实际前端地址配置。
|
||||
BACKEND_CORS_ORIGINS=["http://localhost:5173","http://localhost:3000"]
|
||||
|
||||
# JWT 加密算法
|
||||
ALGORITHM=HS256
|
||||
@@ -216,3 +224,19 @@ ZIP_STORAGE_PATH=./uploads/zip_files
|
||||
# zh-CN: 中文
|
||||
# en-US: 英文
|
||||
OUTPUT_LANGUAGE=zh-CN
|
||||
|
||||
# 独立审计 worker 与共享任务目录(详见 docs/OPERATIONS.md)
|
||||
AUDIT_WORKSPACE_PATH=./uploads/audit_workspaces
|
||||
WORKER_POLL_SECONDS=1
|
||||
WORKER_HEARTBEAT_SECONDS=2
|
||||
WORKER_LEASE_SECONDS=30
|
||||
WORKER_MAX_ATTEMPTS=3
|
||||
WORKER_TASK_TIMEOUT_SECONDS=3600
|
||||
WORKER_RETENTION_DAYS=7
|
||||
|
||||
# ZIP 上传、仓库归档下载与解压配额,单位为字节
|
||||
ZIP_MAX_UPLOAD_BYTES=524288000
|
||||
ZIP_MAX_EXTRACTED_BYTES=2147483648
|
||||
ZIP_MAX_FILE_BYTES=52428800
|
||||
ZIP_MAX_ENTRIES=50000
|
||||
ZIP_MAX_COMPRESSION_RATIO=200
|
||||
|
||||
@@ -5,8 +5,14 @@ Shared test fixtures for all backend tests.
|
||||
import pytest
|
||||
import tempfile
|
||||
import shutil
|
||||
import os
|
||||
from unittest.mock import patch, MagicMock
|
||||
|
||||
# Configure imports before the application constructs its settings singleton.
|
||||
os.environ["SECRET_KEY"] = "test_secret_key_for_testing_1234567890"
|
||||
os.environ["ENVIRONMENT"] = "test"
|
||||
os.environ["DEMO_ENABLED"] = "false"
|
||||
|
||||
from app.core.config import Settings
|
||||
|
||||
|
||||
|
||||
@@ -373,6 +373,7 @@ class TestUpdateIssue:
|
||||
"""Resolve an issue and set resolved_by/resolved_at."""
|
||||
issue = _make_issue(status="open")
|
||||
db = AsyncMock()
|
||||
db.get.return_value = _make_task()
|
||||
|
||||
issue_result = MagicMock()
|
||||
issue_scalars = MagicMock()
|
||||
@@ -395,6 +396,7 @@ class TestUpdateIssue:
|
||||
async def test_update_issue_not_found(self):
|
||||
"""Raises 404 when issue does not exist."""
|
||||
db = AsyncMock()
|
||||
db.get.return_value = _make_task()
|
||||
|
||||
issue_result = MagicMock()
|
||||
issue_scalars = MagicMock()
|
||||
@@ -416,6 +418,7 @@ class TestUpdateIssue:
|
||||
"""Update issue status to false_positive (not resolved, so no resolved_by)."""
|
||||
issue = _make_issue(status="open")
|
||||
db = AsyncMock()
|
||||
db.get.return_value = _make_task()
|
||||
|
||||
issue_result = MagicMock()
|
||||
issue_scalars = MagicMock()
|
||||
@@ -433,3 +436,31 @@ class TestUpdateIssue:
|
||||
assert issue.status == "false_positive"
|
||||
# resolved_by should NOT be set since status is not "resolved"
|
||||
assert issue.resolved_by is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_other_user_cannot_modify_issue(self):
|
||||
db = AsyncMock()
|
||||
db.get.return_value = _make_task(created_by=OTHER_USER_ID)
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await update_issue(
|
||||
task_id=TASK_ID, issue_id="issue-001",
|
||||
issue_update=IssueUpdateSchema(status="resolved"),
|
||||
db=db, current_user=_make_user(),
|
||||
)
|
||||
assert exc.value.status_code == 403
|
||||
db.execute.assert_not_awaited()
|
||||
db.commit.assert_not_awaited()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_missing_task_cannot_modify_issue(self):
|
||||
db = AsyncMock()
|
||||
db.get.return_value = None
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await update_issue(
|
||||
task_id=TASK_ID, issue_id="issue-001",
|
||||
issue_update=IssueUpdateSchema(status="resolved"),
|
||||
db=db, current_user=_make_user(),
|
||||
)
|
||||
assert exc.value.status_code == 404
|
||||
db.execute.assert_not_awaited()
|
||||
db.commit.assert_not_awaited()
|
||||
|
||||
@@ -0,0 +1,127 @@
|
||||
import asyncio
|
||||
import io
|
||||
import stat
|
||||
import zipfile
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
from fastapi import HTTPException
|
||||
from app.services import archive
|
||||
|
||||
|
||||
def make_zip(entries, compression=zipfile.ZIP_STORED):
|
||||
buffer = io.BytesIO()
|
||||
with zipfile.ZipFile(buffer, "w", compression=compression) as output:
|
||||
for name, content in entries:
|
||||
output.writestr(name, content)
|
||||
buffer.seek(0)
|
||||
return buffer
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"name", ["../escape.py", "/escape.py", "a/../../escape.py", "a\\..\\escape.py", "C:/escape.py"]
|
||||
)
|
||||
def test_rejects_unsafe_paths_before_writing(tmp_path, name):
|
||||
with zipfile.ZipFile(make_zip([("good.py", "ok"), (name, "bad")])) as source:
|
||||
with pytest.raises(archive.ArchiveError):
|
||||
archive.extract_zip(source, str(tmp_path / "out"))
|
||||
assert not list(tmp_path.rglob("*.py"))
|
||||
|
||||
|
||||
def test_rejects_links_and_case_collisions():
|
||||
link = zipfile.ZipInfo("link")
|
||||
link.external_attr = (stat.S_IFLNK | 0o777) << 16
|
||||
for entries in [[(link, "../../outside")], [("A.py", "a"), ("a.py", "b")]]:
|
||||
with zipfile.ZipFile(make_zip(entries)) as source:
|
||||
with pytest.raises(archive.ArchiveError):
|
||||
archive.validate_zip(source)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"setting,value,entries",
|
||||
[
|
||||
("ZIP_MAX_ENTRIES", 1, [("a", "a"), ("b", "b")]),
|
||||
("ZIP_MAX_FILE_BYTES", 2, [("a", "abc")]),
|
||||
("ZIP_MAX_EXTRACTED_BYTES", 3, [("a", "ab"), ("b", "cd")]),
|
||||
("ZIP_MAX_COMPRESSION_RATIO", 2, [("a", "a" * 10000)]),
|
||||
],
|
||||
)
|
||||
def test_each_resource_quota(monkeypatch, setting, value, entries):
|
||||
monkeypatch.setattr(archive.settings, setting, value)
|
||||
with zipfile.ZipFile(make_zip(entries, zipfile.ZIP_DEFLATED)) as source:
|
||||
with pytest.raises(archive.ArchiveError):
|
||||
archive.validate_zip(source)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"entries,expected",
|
||||
[
|
||||
([("root/a.py", "a"), ("root/b.py", "b")], ["a.py", "b.py"]),
|
||||
([("README.md", "readme"), ("src/a.py", "a")], ["README.md", "src/a.py"]),
|
||||
],
|
||||
)
|
||||
def test_strip_common_directory_preserves_root_files(tmp_path, entries, expected):
|
||||
with zipfile.ZipFile(make_zip(entries)) as source:
|
||||
archive.extract_zip(source, str(tmp_path), strip_root=True)
|
||||
assert (
|
||||
sorted(str(p.relative_to(tmp_path)) for p in tmp_path.rglob("*") if p.is_file()) == expected
|
||||
)
|
||||
|
||||
|
||||
def test_cancellation_discards_staged_files(tmp_path):
|
||||
checks = iter([False, False, True])
|
||||
with zipfile.ZipFile(make_zip([("first", "ok"), ("second", "no")])) as source:
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
archive.extract_zip(source, str(tmp_path / "out"), cancel_check=lambda: next(checks))
|
||||
assert list((tmp_path / "out").iterdir()) == []
|
||||
assert list(tmp_path.iterdir()) == [tmp_path / "out"]
|
||||
|
||||
|
||||
async def test_upload_stops_at_quota_and_removes_partial_file(tmp_path, monkeypatch):
|
||||
monkeypatch.setattr(archive.settings, "ZIP_MAX_UPLOAD_BYTES", 3)
|
||||
upload = AsyncMock()
|
||||
upload.read.side_effect = [b"abc", b"d", b"should not read"]
|
||||
destination = tmp_path / "upload.zip"
|
||||
with pytest.raises(HTTPException) as error:
|
||||
await archive.save_upload(upload, str(destination))
|
||||
assert error.value.status_code == 400
|
||||
assert upload.read.await_count == 2
|
||||
assert not destination.exists()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("declare_size", [True, False])
|
||||
async def test_multipart_limit_runs_before_spooling_full_request(monkeypatch, declare_size):
|
||||
import httpx
|
||||
from fastapi import FastAPI, File, UploadFile
|
||||
from app.core.upload_limit import UploadLimitMiddleware, settings as upload_settings
|
||||
|
||||
monkeypatch.setattr(upload_settings, "ZIP_MAX_UPLOAD_BYTES", 8)
|
||||
monkeypatch.setattr(archive.settings, "ZIP_MAX_UPLOAD_BYTES", 8)
|
||||
app = FastAPI()
|
||||
app.add_middleware(UploadLimitMiddleware)
|
||||
called = []
|
||||
|
||||
@app.post("/api/v1/scan/upload-zip")
|
||||
async def upload(file: UploadFile = File(...)):
|
||||
called.append(True)
|
||||
return {}
|
||||
|
||||
consumed = []
|
||||
|
||||
async def body():
|
||||
yield b'--test\r\nContent-Disposition: form-data; name="file"; filename="a.zip"\r\nContent-Type: application/zip\r\n\r\n'
|
||||
for index in range(10):
|
||||
consumed.append(index)
|
||||
yield b"x" * (512 * 1024)
|
||||
yield b"\r\n--test--\r\n"
|
||||
|
||||
headers = {"Content-Type": "multipart/form-data; boundary=test"}
|
||||
if declare_size:
|
||||
headers["Content-Length"] = str(6 * 1024 * 1024)
|
||||
async with httpx.AsyncClient(
|
||||
transport=httpx.ASGITransport(app=app), base_url="http://test"
|
||||
) as client:
|
||||
response = await client.post("/api/v1/scan/upload-zip", content=body(), headers=headers)
|
||||
assert response.status_code == 413
|
||||
assert not called
|
||||
assert len(consumed) < 10
|
||||
@@ -0,0 +1,473 @@
|
||||
"""Real PostgreSQL regressions. TEST_DATABASE_URL must name a disposable migrated DB.
|
||||
Each test creates and removes only its own rows; no schema or database is truncated.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from types import SimpleNamespace
|
||||
from uuid import uuid4
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
from sqlalchemy import delete, select, update, func
|
||||
from sqlalchemy.ext.asyncio import create_async_engine, async_sessionmaker
|
||||
from sqlalchemy.pool import NullPool
|
||||
|
||||
from app import worker
|
||||
from app.db import session
|
||||
from app.models import User, Project
|
||||
from app.models.audit import AuditTask, AuditIssue
|
||||
from app.models.audit_job import AuditJob
|
||||
from app.models.agent_task import AgentTask, AgentEvent, AgentFinding
|
||||
from app.services import audit_queue
|
||||
from app.services.agent_audit import events
|
||||
|
||||
pytestmark = pytest.mark.skipif(
|
||||
not os.getenv("TEST_DATABASE_URL"), reason="disposable PostgreSQL not configured"
|
||||
)
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def queue_db(monkeypatch, tmp_path):
|
||||
engine = create_async_engine(os.environ["TEST_DATABASE_URL"], poolclass=NullPool)
|
||||
factory = async_sessionmaker(engine, expire_on_commit=False)
|
||||
for module in [worker, session, audit_queue]:
|
||||
monkeypatch.setattr(module, "AsyncSessionLocal", factory)
|
||||
monkeypatch.setattr(worker, "engine", engine)
|
||||
monkeypatch.setattr(events, "async_session_factory", factory)
|
||||
monkeypatch.setattr(worker.settings, "AUDIT_WORKSPACE_PATH", str(tmp_path))
|
||||
monkeypatch.setattr(worker.settings, "SANDBOX_ENABLED", False)
|
||||
monkeypatch.setattr(worker.settings, "WORKER_HEARTBEAT_SECONDS", 0.05)
|
||||
user = User(email=f"{uuid4()}@example.com", hashed_password="test")
|
||||
async with factory() as db:
|
||||
db.add(user)
|
||||
await db.flush()
|
||||
project = Project(
|
||||
name="Queue regression",
|
||||
owner_id=user.id,
|
||||
source_type="repository",
|
||||
repository_url="https://example.com/source.git",
|
||||
)
|
||||
db.add(project)
|
||||
await db.commit()
|
||||
ids = []
|
||||
|
||||
async def create(kind="repository", **kwargs):
|
||||
async with factory() as db:
|
||||
model = AgentTask if kind == "agent" else AuditTask
|
||||
task = model(
|
||||
project_id=project.id,
|
||||
created_by=user.id,
|
||||
task_type="agent_audit" if kind == "agent" else "repository",
|
||||
**kwargs,
|
||||
)
|
||||
db.add(task)
|
||||
await audit_queue.enqueue(db, task, kind)
|
||||
await db.commit()
|
||||
ids.append(task.id)
|
||||
return task
|
||||
|
||||
yield SimpleNamespace(factory=factory, create=create, project=project, user=user, ids=ids)
|
||||
async with factory() as db:
|
||||
await db.execute(delete(AuditJob).where(AuditJob.id.in_(ids)))
|
||||
await db.execute(delete(AuditIssue).where(AuditIssue.task_id.in_(ids)))
|
||||
await db.execute(delete(AuditTask).where(AuditTask.id.in_(ids)))
|
||||
await db.execute(delete(Project).where(Project.id == project.id))
|
||||
await db.execute(delete(User).where(User.id == user.id))
|
||||
await db.commit()
|
||||
await engine.dispose()
|
||||
|
||||
|
||||
async def finish(factory, task_id, kind="repository"):
|
||||
async with factory() as db:
|
||||
model = AgentTask if kind == "agent" else AuditTask
|
||||
await db.execute(update(model).where(model.id == task_id).values(status="completed"))
|
||||
await db.commit()
|
||||
|
||||
|
||||
async def test_enqueue_rolls_back_with_task(queue_db):
|
||||
async with queue_db.factory() as db:
|
||||
task = AuditTask(
|
||||
project_id=queue_db.project.id, created_by=queue_db.user.id, task_type="repository"
|
||||
)
|
||||
db.add(task)
|
||||
await audit_queue.enqueue(db, task, "repository")
|
||||
task_id = task.id
|
||||
await db.rollback()
|
||||
async with queue_db.factory() as db:
|
||||
assert await db.get(AuditJob, task_id) is None
|
||||
assert await db.get(AuditTask, task_id) is None
|
||||
|
||||
|
||||
async def test_snapshot_encrypted_and_independent_of_project_changes(queue_db):
|
||||
task = await queue_db.create()
|
||||
async with queue_db.factory() as db:
|
||||
job = await db.get(AuditJob, task.id)
|
||||
assert job.payload["user_config"] != "{}"
|
||||
await db.execute(
|
||||
update(Project)
|
||||
.where(Project.id == task.project_id)
|
||||
.values(repository_url="https://example.com/changed.git")
|
||||
)
|
||||
await db.commit()
|
||||
snapshot = await audit_queue.project_snapshot(task.id, None)
|
||||
assert snapshot.repository_url == "https://example.com/source.git"
|
||||
|
||||
|
||||
async def test_expired_lease_does_not_allow_concurrent_execution(queue_db, monkeypatch):
|
||||
task = await queue_db.create()
|
||||
started, release = asyncio.Event(), asyncio.Event()
|
||||
calls = []
|
||||
|
||||
async def execute(job):
|
||||
calls.append(job.id)
|
||||
started.set()
|
||||
await release.wait()
|
||||
await finish(queue_db.factory, job.id)
|
||||
|
||||
monkeypatch.setattr(worker, "execute_payload", execute)
|
||||
running = asyncio.create_task(worker.run_job(task.id))
|
||||
await asyncio.wait_for(started.wait(), 5)
|
||||
try:
|
||||
async with queue_db.factory() as db:
|
||||
await db.execute(
|
||||
update(AuditJob)
|
||||
.where(AuditJob.id == task.id)
|
||||
.values(lease_until=datetime.now(timezone.utc) - timedelta(seconds=100))
|
||||
)
|
||||
await db.commit()
|
||||
assert await worker.run_job(task.id) is False
|
||||
finally:
|
||||
release.set()
|
||||
await asyncio.wait_for(running, 5)
|
||||
assert calls == [task.id]
|
||||
async with queue_db.factory() as db:
|
||||
job = await db.get(AuditJob, task.id)
|
||||
assert (job.status, job.attempts) == ("completed", 1)
|
||||
|
||||
|
||||
async def test_recovery_after_lock_owner_process_dies(queue_db, monkeypatch):
|
||||
task = await queue_db.create()
|
||||
code = """import asyncio, asyncpg, sys
|
||||
async def main():
|
||||
c = await asyncpg.connect(sys.argv[1])
|
||||
await c.execute('SELECT pg_advisory_lock($1)', int(sys.argv[3]))
|
||||
await c.execute("UPDATE audit_jobs SET status='running', attempts=1, lease_until=now()-interval '1 minute' WHERE id=$1", sys.argv[2])
|
||||
print('locked', flush=True)
|
||||
await asyncio.Event().wait()
|
||||
asyncio.run(main())"""
|
||||
process = await asyncio.create_subprocess_exec(
|
||||
os.sys.executable,
|
||||
"-c",
|
||||
code,
|
||||
os.environ["TEST_DATABASE_URL"].replace("postgresql+asyncpg", "postgresql"),
|
||||
task.id,
|
||||
str(audit_queue.advisory_key(task.id)),
|
||||
stdout=asyncio.subprocess.PIPE,
|
||||
)
|
||||
try:
|
||||
assert await asyncio.wait_for(process.stdout.readline(), 10) == b"locked\n"
|
||||
assert await worker.run_job(task.id) is False
|
||||
finally:
|
||||
process.kill()
|
||||
await process.wait()
|
||||
|
||||
async def execute(job):
|
||||
await finish(queue_db.factory, job.id)
|
||||
|
||||
monkeypatch.setattr(worker, "execute_payload", execute)
|
||||
assert task.id in await worker.candidate_ids()
|
||||
assert await worker.run_job(task.id)
|
||||
async with queue_db.factory() as db:
|
||||
job = await db.get(AuditJob, task.id)
|
||||
assert (job.status, job.attempts) == ("completed", 2)
|
||||
|
||||
|
||||
async def test_durable_cancel_stops_running_work(queue_db, monkeypatch):
|
||||
task = await queue_db.create()
|
||||
started, cancelled = asyncio.Event(), asyncio.Event()
|
||||
|
||||
async def execute(job):
|
||||
started.set()
|
||||
try:
|
||||
await asyncio.Event().wait()
|
||||
finally:
|
||||
cancelled.set()
|
||||
|
||||
monkeypatch.setattr(worker, "execute_payload", execute)
|
||||
running = asyncio.create_task(worker.run_job(task.id))
|
||||
await asyncio.wait_for(started.wait(), 5)
|
||||
async with queue_db.factory() as db:
|
||||
await audit_queue.request_cancel(db, task.id)
|
||||
await db.commit()
|
||||
await asyncio.wait_for(running, 5)
|
||||
assert cancelled.is_set()
|
||||
async with queue_db.factory() as db:
|
||||
assert (await db.get(AuditJob, task.id)).status == "cancelled"
|
||||
assert (await db.get(AuditTask, task.id)).status == "cancelled"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"state,attempts,expected", [("completed", 1, "completed"), ("pending", 3, "failed")]
|
||||
)
|
||||
async def test_completed_result_and_recovery_limit_never_reexecute(
|
||||
queue_db, monkeypatch, state, attempts, expected
|
||||
):
|
||||
task = await queue_db.create(status=state)
|
||||
async with queue_db.factory() as db:
|
||||
await db.execute(
|
||||
update(AuditJob)
|
||||
.where(AuditJob.id == task.id)
|
||||
.values(status="running", attempts=attempts)
|
||||
)
|
||||
await db.commit()
|
||||
execute = AsyncMock()
|
||||
monkeypatch.setattr(worker, "execute_payload", execute)
|
||||
await worker.run_job(task.id)
|
||||
execute.assert_not_called()
|
||||
async with queue_db.factory() as db:
|
||||
assert (await db.get(AuditJob, task.id)).status == expected
|
||||
|
||||
|
||||
async def test_checkpoint_and_findings_replay_are_idempotent(queue_db):
|
||||
from app.services.agent_audit.results import _save_findings
|
||||
|
||||
task = await queue_db.create("agent")
|
||||
await audit_queue.save_checkpoint(task.id, project_root="/snapshot")
|
||||
await audit_queue.save_checkpoint(task.id, result={"success": True})
|
||||
assert await audit_queue.load_checkpoint(task.id) == {
|
||||
"project_root": "/snapshot",
|
||||
"result": {"success": True},
|
||||
}
|
||||
findings = [
|
||||
{
|
||||
"file_path": "app.py",
|
||||
"line": 5,
|
||||
"title": "Untrusted query",
|
||||
"severity": "high",
|
||||
"vulnerability_type": "sql_injection",
|
||||
"description": "unsafe query",
|
||||
}
|
||||
]
|
||||
async with queue_db.factory() as db:
|
||||
await _save_findings(db, task.id, findings + findings)
|
||||
await _save_findings(db, task.id, findings)
|
||||
assert (
|
||||
await db.scalar(
|
||||
select(func.count())
|
||||
.select_from(AgentFinding)
|
||||
.where(AgentFinding.task_id == task.id)
|
||||
)
|
||||
== 1
|
||||
)
|
||||
|
||||
|
||||
async def test_stream_replay_broadcast_and_token_batching(queue_db):
|
||||
from app.services.agent.event_manager import EventManager, AgentEventEmitter, AgentEventData
|
||||
|
||||
task = await queue_db.create("agent")
|
||||
manager = EventManager(db_session_factory=queue_db.factory)
|
||||
emitter = AgentEventEmitter(task.id, manager)
|
||||
await emitter.emit(AgentEventData(event_type="thinking_token", metadata={"token": "hello "}))
|
||||
await emitter.emit(AgentEventData(event_type="thinking_token", metadata={"token": "world"}))
|
||||
await emitter.emit(AgentEventData(event_type="task_complete", message="done"))
|
||||
await finish(queue_db.factory, task.id, "agent")
|
||||
request = SimpleNamespace(is_disconnected=AsyncMock(return_value=False))
|
||||
|
||||
async def collect(cursor=0):
|
||||
return [item async for item in events.stream_events(task.id, request, cursor)]
|
||||
|
||||
first, second = await asyncio.gather(collect(), collect())
|
||||
assert first == second
|
||||
assert len(first) == 3
|
||||
assert "hello world" in first[0] and "id: 2" in first[0]
|
||||
assert "id: 3" in first[1]
|
||||
assert await collect(2) == first[1:]
|
||||
|
||||
|
||||
async def test_sql_pagination_total_and_owner_isolation(queue_db):
|
||||
from fastapi import Response
|
||||
from app.api.v1.endpoints.tasks import list_tasks
|
||||
from app.services.project_stats import project_stats
|
||||
|
||||
tasks = [await queue_db.create(status="completed") for _ in range(3)]
|
||||
async with queue_db.factory() as db:
|
||||
response = Response()
|
||||
page = await list_tasks(
|
||||
response=response,
|
||||
skip=1,
|
||||
limit=1,
|
||||
status="completed",
|
||||
search="Queue",
|
||||
db=db,
|
||||
current_user=queue_db.user,
|
||||
)
|
||||
assert len(page) == 1 and response.headers["X-Total-Count"] == "3"
|
||||
assert page[0].id in {t.id for t in tasks}
|
||||
assert (
|
||||
await list_tasks(skip=0, limit=10, db=db, current_user=SimpleNamespace(id=str(uuid4())))
|
||||
== []
|
||||
)
|
||||
stats = await project_stats(db, queue_db.user.id)
|
||||
assert stats["total_tasks"] == stats["completed_tasks"] == 3
|
||||
assert stats["total_projects"] == 1
|
||||
|
||||
|
||||
async def test_live_tree_published_for_another_api_process(queue_db, monkeypatch):
|
||||
from app.services.agent_audit import progress
|
||||
from app.services.agent_audit.runtime import _running_orchestrators
|
||||
from app.services.agent.core import agent_registry
|
||||
from app.api.v1.endpoints.agent_inspection import get_agent_tree
|
||||
|
||||
task = await queue_db.create("agent", status="running")
|
||||
monkeypatch.setattr(progress, "AsyncSessionLocal", queue_db.factory)
|
||||
monkeypatch.setitem(_running_orchestrators, task.id, object())
|
||||
monkeypatch.setattr(
|
||||
agent_registry,
|
||||
"get_agent_tree",
|
||||
lambda: {
|
||||
"root_agent_id": "root",
|
||||
"nodes": {
|
||||
"root": {
|
||||
"id": "root",
|
||||
"name": "Orchestrator",
|
||||
"type": "orchestrator",
|
||||
"status": "running",
|
||||
},
|
||||
},
|
||||
},
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
agent_registry,
|
||||
"get_agent",
|
||||
lambda _: SimpleNamespace(
|
||||
get_stats=lambda: {"iterations": 2, "tool_calls": 3, "tokens_used": 50}
|
||||
),
|
||||
)
|
||||
await progress.publish_progress(task.id)
|
||||
async with queue_db.factory() as db:
|
||||
persisted = await db.get(AgentTask, task.id)
|
||||
assert (persisted.total_iterations, persisted.tool_calls_count, persisted.tokens_used) == (
|
||||
2,
|
||||
3,
|
||||
50,
|
||||
)
|
||||
tree = await get_agent_tree(task.id, db=db, current_user=queue_db.user)
|
||||
assert tree.root_agent_id == "root" and tree.running_agents == 1
|
||||
assert tree.nodes[0].agent_name == "Orchestrator"
|
||||
|
||||
|
||||
async def test_token_batches_keep_agents_separate(queue_db):
|
||||
from app.services.agent.event_manager import EventManager
|
||||
|
||||
task = await queue_db.create("agent")
|
||||
manager = EventManager(db_session_factory=queue_db.factory)
|
||||
for sequence, name in enumerate(["first", "second"], 1):
|
||||
await manager.add_event(
|
||||
task.id,
|
||||
"thinking_token",
|
||||
sequence=sequence,
|
||||
metadata={"token": name, "agent_name": name},
|
||||
)
|
||||
await manager.flush_tokens()
|
||||
async with queue_db.factory() as db:
|
||||
rows = list(
|
||||
await db.scalars(
|
||||
select(AgentEvent)
|
||||
.where(AgentEvent.task_id == task.id)
|
||||
.order_by(AgentEvent.sequence)
|
||||
)
|
||||
)
|
||||
assert [row.event_metadata for row in rows] == [
|
||||
{"token": "first", "agent_name": "first"},
|
||||
{"token": "second", "agent_name": "second"},
|
||||
]
|
||||
|
||||
|
||||
async def test_zip_worker_runs_real_pipeline_and_reuses_file_outputs(
|
||||
queue_db, monkeypatch, tmp_path
|
||||
):
|
||||
import zipfile
|
||||
from app.services import zip_scan
|
||||
|
||||
task = await queue_db.create("zip")
|
||||
source = tmp_path / "original.zip"
|
||||
with zipfile.ZipFile(source, "w") as output:
|
||||
output.writestr("app.py", 'print("hello")')
|
||||
async with queue_db.factory() as db:
|
||||
job = await db.get(AuditJob, task.id)
|
||||
job.payload = {**job.payload, "source_zip": str(source)}
|
||||
await db.commit()
|
||||
llm = SimpleNamespace(
|
||||
analyze_code=AsyncMock(
|
||||
return_value={
|
||||
"quality_score": 95,
|
||||
"issues": [
|
||||
{"line": 1, "type": "security", "severity": "high", "title": "Test issue"},
|
||||
],
|
||||
}
|
||||
)
|
||||
)
|
||||
monkeypatch.setattr(zip_scan, "LLMService", lambda **_: llm)
|
||||
monkeypatch.setattr(
|
||||
zip_scan, "get_analysis_config", lambda _: {"max_analyze_files": 0, "llm_gap_ms": 0}
|
||||
)
|
||||
for attempt in range(2):
|
||||
assert await worker.run_job(task.id)
|
||||
async with queue_db.factory() as db:
|
||||
persisted = await db.get(AuditTask, task.id)
|
||||
assert persisted.status == "completed"
|
||||
assert (persisted.scanned_files, persisted.issues_count, persisted.quality_score) == (
|
||||
1,
|
||||
1,
|
||||
95,
|
||||
)
|
||||
assert (
|
||||
await db.scalar(
|
||||
select(func.count())
|
||||
.select_from(AuditIssue)
|
||||
.where(AuditIssue.task_id == task.id)
|
||||
)
|
||||
== 1
|
||||
)
|
||||
if attempt == 0:
|
||||
persisted.status = "running"
|
||||
job = await db.get(AuditJob, task.id)
|
||||
job.status = "running"
|
||||
await db.commit()
|
||||
assert llm.analyze_code.await_count == 1
|
||||
|
||||
|
||||
async def test_completed_agent_checkpoint_finishes_without_model_calls(queue_db, monkeypatch, tmp_path):
|
||||
from app.services.agent_audit import execution
|
||||
from app.services.agent.tools import SandboxManager
|
||||
from app.services.agent.agents.base import AgentResult
|
||||
task = await queue_db.create('agent')
|
||||
root = tmp_path / 'snapshot'
|
||||
root.mkdir()
|
||||
(root / 'app.py').write_text('print("hello")')
|
||||
await audit_queue.save_checkpoint(task.id, result=AgentResult(success=True, data={'findings': []}, tokens_used=15).to_dict())
|
||||
monkeypatch.setattr(execution, '_get_project_root', AsyncMock(return_value=str(root)))
|
||||
initialize_tools = AsyncMock(side_effect=AssertionError('must reuse completed result'))
|
||||
monkeypatch.setattr(execution, '_initialize_tools', initialize_tools)
|
||||
monkeypatch.setattr(SandboxManager, 'initialize', AsyncMock())
|
||||
await worker.run_job(task.id)
|
||||
initialize_tools.assert_not_called()
|
||||
async with queue_db.factory() as db:
|
||||
persisted = await db.get(AgentTask, task.id)
|
||||
assert persisted.status == 'completed'
|
||||
assert persisted.tokens_used == 15
|
||||
assert (await db.get(AuditJob, task.id)).status == 'completed'
|
||||
|
||||
|
||||
async def test_stream_authorization_releases_database_connection(queue_db):
|
||||
from app.api.v1.endpoints.agent_events import stream_agent_events
|
||||
task = await queue_db.create('agent', status='completed')
|
||||
request = SimpleNamespace(headers={'last-event-id': '7'}, is_disconnected=AsyncMock(return_value=False))
|
||||
async with queue_db.factory() as db:
|
||||
response = await stream_agent_events(task.id, request, after_sequence=0, db=db, current_user=queue_db.user)
|
||||
assert not db.in_transaction()
|
||||
chunks = [chunk async for chunk in response.body_iterator]
|
||||
assert len(chunks) == 1 and 'id: 7' in chunks[0]
|
||||
@@ -0,0 +1,60 @@
|
||||
from types import SimpleNamespace
|
||||
from uuid import uuid4
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
from app.services import audit_queue
|
||||
from app.services.analysis_checkpoint import analyze_file
|
||||
from app.services.agent.tools.sandbox_tool import SandboxManager
|
||||
|
||||
|
||||
async def test_file_cache_reuses_result_and_invalidates_changed_source(tmp_path, monkeypatch):
|
||||
monkeypatch.setattr(audit_queue.settings, "AUDIT_WORKSPACE_PATH", str(tmp_path))
|
||||
task_id = str(uuid4())
|
||||
service = SimpleNamespace(
|
||||
analyze_code=AsyncMock(return_value={"issues": [], "quality_score": 90})
|
||||
)
|
||||
first = await analyze_file(task_id, "app.py", "old", service, "python")
|
||||
assert await analyze_file(task_id, "app.py", "old", service, "python") == first
|
||||
assert service.analyze_code.await_count == 1
|
||||
await analyze_file(task_id, "app.py", "new", service, "python")
|
||||
assert service.analyze_code.await_count == 2
|
||||
for file in tmp_path.rglob("*.json"):
|
||||
file.write_text("corrupt")
|
||||
await analyze_file(task_id, "app.py", "new", service, "python")
|
||||
assert service.analyze_code.await_count == 3
|
||||
|
||||
|
||||
async def test_failed_analysis_is_not_cached(tmp_path, monkeypatch):
|
||||
monkeypatch.setattr(audit_queue.settings, "AUDIT_WORKSPACE_PATH", str(tmp_path))
|
||||
service = SimpleNamespace(
|
||||
analyze_code=AsyncMock(side_effect=[RuntimeError("unavailable"), {"issues": []}])
|
||||
)
|
||||
task_id = str(uuid4())
|
||||
with pytest.raises(RuntimeError):
|
||||
await analyze_file(task_id, "a.py", "source", service, "python")
|
||||
assert not list(tmp_path.rglob("*.json"))
|
||||
assert await analyze_file(task_id, "a.py", "source", service, "python") == {"issues": []}
|
||||
|
||||
|
||||
def test_sandbox_maps_shared_volume_to_daemon_path(monkeypatch):
|
||||
from pathlib import Path
|
||||
|
||||
manager = SandboxManager()
|
||||
manager._docker_client = MagicMock()
|
||||
manager._docker_client.containers.get.return_value.attrs = {
|
||||
"Mounts": [
|
||||
{"Destination": "/app", "Source": "/daemon/app"},
|
||||
{"Destination": "/app/uploads", "Source": "/daemon/volumes/uploads/_data"},
|
||||
]
|
||||
}
|
||||
original = Path.exists
|
||||
monkeypatch.setattr(
|
||||
Path, "exists", lambda path: True if str(path) == "/.dockerenv" else original(path)
|
||||
)
|
||||
assert (
|
||||
manager._daemon_path("/app/uploads/task/source")
|
||||
== "/daemon/volumes/uploads/_data/task/source"
|
||||
)
|
||||
with pytest.raises(ValueError, match="shared worker volume"):
|
||||
manager._daemon_path("/tmp/not-shared")
|
||||
@@ -0,0 +1,117 @@
|
||||
"""Regression coverage for deployment credentials, demo bootstrap, and CORS."""
|
||||
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from pydantic import ValidationError
|
||||
|
||||
from app.core.config import Settings
|
||||
from app.db import init_db as bootstrap
|
||||
|
||||
|
||||
def test_secret_is_required(monkeypatch):
|
||||
monkeypatch.delenv("SECRET_KEY")
|
||||
with pytest.raises(ValidationError, match="SECRET_KEY"):
|
||||
Settings(_env_file=None)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("secret", [
|
||||
"short", " " * 32,
|
||||
"changethis_in_production_to_a_long_random_string",
|
||||
"your-super-secret-key-change-this-in-production",
|
||||
])
|
||||
def test_unsafe_secrets_are_rejected(secret):
|
||||
with pytest.raises(ValidationError, match="SECRET_KEY"):
|
||||
Settings(_env_file=None, SECRET_KEY=secret)
|
||||
|
||||
|
||||
def test_production_rejects_demo_mode():
|
||||
with pytest.raises(ValidationError, match="DEMO_ENABLED"):
|
||||
Settings(_env_file=None, ENVIRONMENT="production", DEMO_ENABLED=True)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("default_password", [True, False])
|
||||
async def test_existing_demo_is_disabled_only_with_default_password(default_password):
|
||||
user = SimpleNamespace(is_active=True, hashed_password="stored-hash")
|
||||
db = AsyncMock()
|
||||
result = MagicMock()
|
||||
result.scalars.return_value.first.return_value = user
|
||||
db.execute.return_value = result
|
||||
with (
|
||||
patch.object(bootstrap, "settings", Settings(_env_file=None, DEMO_ENABLED=False)),
|
||||
patch.object(bootstrap, "verify_password", return_value=default_password),
|
||||
patch.object(bootstrap, "create_demo_user", new_callable=AsyncMock) as create,
|
||||
patch("app.services.init_templates.init_templates_and_rules", new_callable=AsyncMock),
|
||||
):
|
||||
await bootstrap.init_db(db)
|
||||
assert user.is_active is (not default_password)
|
||||
create.assert_not_awaited()
|
||||
db.commit.assert_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fresh_install_does_not_create_demo():
|
||||
db = AsyncMock()
|
||||
result = MagicMock()
|
||||
result.scalars.return_value.first.return_value = None
|
||||
db.execute.return_value = result
|
||||
with (
|
||||
patch.object(bootstrap, "settings", Settings(_env_file=None, DEMO_ENABLED=False)),
|
||||
patch.object(bootstrap, "create_demo_user", new_callable=AsyncMock) as create,
|
||||
patch("app.services.init_templates.init_templates_and_rules", new_callable=AsyncMock),
|
||||
):
|
||||
await bootstrap.init_db(db)
|
||||
create.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_demo_can_be_explicitly_enabled_in_development():
|
||||
db = AsyncMock()
|
||||
with (
|
||||
patch.object(bootstrap, "settings", Settings(
|
||||
_env_file=None, ENVIRONMENT="development", DEMO_ENABLED=True,
|
||||
)),
|
||||
patch.object(bootstrap, "create_demo_user", new_callable=AsyncMock) as create,
|
||||
patch.object(bootstrap, "create_demo_data", new_callable=AsyncMock) as data,
|
||||
patch("app.services.init_templates.init_templates_and_rules", new_callable=AsyncMock),
|
||||
):
|
||||
await bootstrap.init_db(db)
|
||||
create.assert_awaited_once_with(db)
|
||||
data.assert_awaited_once_with(db, create.return_value)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_root_does_not_publish_credentials_and_cors_rejects_unknown_origin():
|
||||
from app.main import app
|
||||
|
||||
async with httpx.AsyncClient(
|
||||
transport=httpx.ASGITransport(app=app), base_url="http://test",
|
||||
) as client:
|
||||
response = await client.get("/")
|
||||
assert "demo_account" not in response.json()
|
||||
assert "demo123" not in response.text
|
||||
preflight = await client.options("/api/v1/tasks/", headers={
|
||||
"Origin": "https://untrusted.example",
|
||||
"Access-Control-Request-Method": "GET",
|
||||
})
|
||||
assert preflight.status_code == 400
|
||||
assert "access-control-allow-origin" not in preflight.headers
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_production_does_not_start_when_security_bootstrap_fails():
|
||||
from app import main
|
||||
|
||||
with (
|
||||
patch.object(main, "settings", Settings(_env_file=None, ENVIRONMENT="production")),
|
||||
patch.object(main, "AsyncSessionLocal", return_value=AsyncMock()),
|
||||
patch.object(main, "init_db", side_effect=RuntimeError("bootstrap failed")),
|
||||
patch.object(main, "check_agent_services", new_callable=AsyncMock) as check,
|
||||
):
|
||||
with pytest.raises(RuntimeError, match="bootstrap failed"):
|
||||
async with main.lifespan(main.app):
|
||||
pytest.fail("Production startup must fail closed")
|
||||
check.assert_not_awaited()
|
||||
@@ -0,0 +1,72 @@
|
||||
"""Cancelling one audit must never stop another audit's agents."""
|
||||
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from app.services.agent.core import graph_controller
|
||||
from app.services.agent.core.registry import AgentRegistry
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def task_trees(monkeypatch):
|
||||
registry = AgentRegistry()
|
||||
agents = {}
|
||||
for agent_id, parent in [
|
||||
("a", None), ("a-child", "a"), ("a-grandchild", "a-child"),
|
||||
("b", None), ("b-child", "b"),
|
||||
]:
|
||||
agents[agent_id] = MagicMock()
|
||||
registry.register_agent(
|
||||
agent_id, agent_id, "analysis", "audit", parent_id=parent,
|
||||
agent_instance=agents[agent_id],
|
||||
)
|
||||
monkeypatch.setattr(graph_controller, "agent_registry", registry)
|
||||
return registry, agents
|
||||
|
||||
|
||||
def test_only_selected_tree_is_stopped(task_trees):
|
||||
registry, agents = task_trees
|
||||
result = graph_controller.agent_graph_controller.stop_agent_tree("a")
|
||||
assert result["success"]
|
||||
assert set(result["stopped"]) == {"a", "a-child", "a-grandchild"}
|
||||
for agent_id in ("a", "a-child", "a-grandchild"):
|
||||
agents[agent_id].cancel.assert_called_once()
|
||||
for agent_id in ("b", "b-child"):
|
||||
agents[agent_id].cancel.assert_not_called()
|
||||
assert registry.get_agent_node(agent_id)["status"] == "running"
|
||||
|
||||
|
||||
def test_unknown_root_never_falls_back_to_global_stop(task_trees):
|
||||
_, agents = task_trees
|
||||
result = graph_controller.agent_graph_controller.stop_agent_tree("missing")
|
||||
assert not result["success"]
|
||||
for agent in agents.values():
|
||||
agent.cancel.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("runner_present", [True, False])
|
||||
async def test_cancel_endpoint_scopes_stop_to_task(task_trees, monkeypatch, runner_present):
|
||||
from app.api.v1.endpoints import agent_tasks
|
||||
|
||||
_, agents = task_trees
|
||||
task = SimpleNamespace(
|
||||
id="task-a", project_id="project-a", status=agent_tasks.AgentTaskStatus.RUNNING,
|
||||
)
|
||||
db = AsyncMock()
|
||||
db.get.side_effect = [task, SimpleNamespace(owner_id="owner-a")]
|
||||
runner = SimpleNamespace(agent_id="a", cancel=MagicMock())
|
||||
monkeypatch.setattr(agent_tasks, "_running_tasks", {"task-a": runner} if runner_present else {})
|
||||
monkeypatch.setattr(agent_tasks, "_running_asyncio_tasks", {})
|
||||
monkeypatch.setattr(agent_tasks, "_cancelled_tasks", set())
|
||||
await agent_tasks.cancel_agent_task("task-a", db, SimpleNamespace(id="owner-a"))
|
||||
assert task.status == agent_tasks.AgentTaskStatus.CANCELLED
|
||||
assert "task-a" in agent_tasks._cancelled_tasks
|
||||
db.commit.assert_awaited_once()
|
||||
if runner_present:
|
||||
runner.cancel.assert_called_once()
|
||||
agents["a-grandchild"].cancel.assert_called_once()
|
||||
agents["b"].cancel.assert_not_called()
|
||||
agents["b-child"].cancel.assert_not_called()
|
||||
@@ -15,7 +15,7 @@ services:
|
||||
- postgres_data:/var/lib/postgresql/data
|
||||
environment:
|
||||
- POSTGRES_USER=postgres
|
||||
- POSTGRES_PASSWORD=postgres
|
||||
- POSTGRES_PASSWORD=${POSTGRES_PASSWORD:?Set POSTGRES_PASSWORD}
|
||||
- POSTGRES_DB=deepaudit
|
||||
healthcheck:
|
||||
test: ["CMD-SHELL", "pg_isready -U postgres"]
|
||||
@@ -43,13 +43,20 @@ services:
|
||||
restart: unless-stopped
|
||||
volumes:
|
||||
- backend_uploads:/app/uploads
|
||||
- /var/run/docker.sock:/var/run/docker.sock
|
||||
ports:
|
||||
- "8000:8000"
|
||||
environment:
|
||||
- DATABASE_URL=postgresql+asyncpg://postgres:postgres@db:5432/deepaudit
|
||||
environment: &audit_environment
|
||||
- POSTGRES_SERVER=db
|
||||
- POSTGRES_USER=postgres
|
||||
- POSTGRES_DB=deepaudit
|
||||
- POSTGRES_PASSWORD=${POSTGRES_PASSWORD:?Set POSTGRES_PASSWORD}
|
||||
- SECRET_KEY=${SECRET_KEY:?Set SECRET_KEY to a random secret of at least 32 characters}
|
||||
- ENVIRONMENT=production
|
||||
- DEMO_ENABLED=false
|
||||
- BACKEND_CORS_ORIGINS=${BACKEND_CORS_ORIGINS:-[]}
|
||||
- REDIS_URL=redis://redis:6379/0
|
||||
- AGENT_ENABLED=true
|
||||
- AUDIT_WORKSPACE_PATH=/app/uploads/audit_workspaces
|
||||
- SANDBOX_ENABLED=true
|
||||
- SANDBOX_IMAGE=ghcr.nju.edu.cn/lintsinghua/deepaudit-sandbox:latest
|
||||
# LLM 配置 - 请根据需要修改
|
||||
@@ -70,13 +77,46 @@ services:
|
||||
condition: service_completed_successfully
|
||||
networks:
|
||||
- deepaudit-network
|
||||
healthcheck:
|
||||
test: ["CMD", ".venv/bin/python", "-c", "import urllib.request; urllib.request.urlopen('http://localhost:8000/ready')"]
|
||||
interval: 10s
|
||||
timeout: 5s
|
||||
retries: 5
|
||||
start_period: 30s
|
||||
|
||||
|
||||
worker:
|
||||
image: ghcr.nju.edu.cn/lintsinghua/deepaudit-backend:latest
|
||||
restart: unless-stopped
|
||||
init: true
|
||||
command: [".venv/bin/python", "-m", "app.worker"]
|
||||
environment: *audit_environment
|
||||
volumes:
|
||||
- backend_uploads:/app/uploads
|
||||
- /var/run/docker.sock:/var/run/docker.sock
|
||||
depends_on:
|
||||
backend:
|
||||
condition: service_healthy
|
||||
healthcheck:
|
||||
test: ["CMD", ".venv/bin/python", "-m", "app.worker", "--health"]
|
||||
interval: 15s
|
||||
timeout: 5s
|
||||
retries: 3
|
||||
start_period: 30s
|
||||
networks:
|
||||
- deepaudit-network
|
||||
|
||||
# 数据库迁移服务 - 在后端启动前自动执行
|
||||
db-migrate:
|
||||
image: ghcr.nju.edu.cn/lintsinghua/deepaudit-backend:latest
|
||||
restart: "no"
|
||||
environment:
|
||||
- DATABASE_URL=postgresql+asyncpg://postgres:postgres@db:5432/deepaudit
|
||||
- POSTGRES_SERVER=db
|
||||
- POSTGRES_USER=postgres
|
||||
- POSTGRES_DB=deepaudit
|
||||
- POSTGRES_PASSWORD=${POSTGRES_PASSWORD:?Set POSTGRES_PASSWORD}
|
||||
- SECRET_KEY=${SECRET_KEY:?Set SECRET_KEY to a random secret of at least 32 characters}
|
||||
- ENVIRONMENT=production
|
||||
command: [".venv/bin/alembic", "upgrade", "head"]
|
||||
depends_on:
|
||||
db:
|
||||
|
||||
+45
-5
@@ -12,7 +12,7 @@ services:
|
||||
- postgres_data:/var/lib/postgresql/data
|
||||
environment:
|
||||
- POSTGRES_USER=postgres
|
||||
- POSTGRES_PASSWORD=postgres
|
||||
- POSTGRES_PASSWORD=${POSTGRES_PASSWORD:?Set POSTGRES_PASSWORD}
|
||||
- POSTGRES_DB=deepaudit
|
||||
healthcheck:
|
||||
test: ["CMD-SHELL", "pg_isready -U postgres"]
|
||||
@@ -40,13 +40,20 @@ services:
|
||||
restart: unless-stopped
|
||||
volumes:
|
||||
- backend_uploads:/app/uploads
|
||||
- /var/run/docker.sock:/var/run/docker.sock
|
||||
ports:
|
||||
- "8000:8000"
|
||||
environment:
|
||||
- DATABASE_URL=postgresql+asyncpg://postgres:postgres@db:5432/deepaudit
|
||||
environment: &audit_environment
|
||||
- POSTGRES_SERVER=db
|
||||
- POSTGRES_USER=postgres
|
||||
- POSTGRES_DB=deepaudit
|
||||
- POSTGRES_PASSWORD=${POSTGRES_PASSWORD:?Set POSTGRES_PASSWORD}
|
||||
- SECRET_KEY=${SECRET_KEY:?Set SECRET_KEY to a random secret of at least 32 characters}
|
||||
- ENVIRONMENT=production
|
||||
- DEMO_ENABLED=false
|
||||
- BACKEND_CORS_ORIGINS=${BACKEND_CORS_ORIGINS:-[]}
|
||||
- REDIS_URL=redis://redis:6379/0
|
||||
- AGENT_ENABLED=true
|
||||
- AUDIT_WORKSPACE_PATH=/app/uploads/audit_workspaces
|
||||
- SANDBOX_ENABLED=true
|
||||
- SANDBOX_IMAGE=ghcr.io/lintsinghua/deepaudit-sandbox:latest
|
||||
# LLM 配置 - 请根据需要修改
|
||||
@@ -72,13 +79,46 @@ services:
|
||||
condition: service_completed_successfully
|
||||
networks:
|
||||
- deepaudit-network
|
||||
healthcheck:
|
||||
test: ["CMD", ".venv/bin/python", "-c", "import urllib.request; urllib.request.urlopen('http://localhost:8000/ready')"]
|
||||
interval: 10s
|
||||
timeout: 5s
|
||||
retries: 5
|
||||
start_period: 30s
|
||||
|
||||
|
||||
worker:
|
||||
image: ghcr.io/lintsinghua/deepaudit-backend:latest
|
||||
restart: unless-stopped
|
||||
init: true
|
||||
command: [".venv/bin/python", "-m", "app.worker"]
|
||||
environment: *audit_environment
|
||||
volumes:
|
||||
- backend_uploads:/app/uploads
|
||||
- /var/run/docker.sock:/var/run/docker.sock
|
||||
depends_on:
|
||||
backend:
|
||||
condition: service_healthy
|
||||
healthcheck:
|
||||
test: ["CMD", ".venv/bin/python", "-m", "app.worker", "--health"]
|
||||
interval: 15s
|
||||
timeout: 5s
|
||||
retries: 3
|
||||
start_period: 30s
|
||||
networks:
|
||||
- deepaudit-network
|
||||
|
||||
# 数据库迁移服务 - 在后端启动前自动执行
|
||||
db-migrate:
|
||||
image: ghcr.io/lintsinghua/deepaudit-backend:latest
|
||||
restart: "no"
|
||||
environment:
|
||||
- DATABASE_URL=postgresql+asyncpg://postgres:postgres@db:5432/deepaudit
|
||||
- POSTGRES_SERVER=db
|
||||
- POSTGRES_USER=postgres
|
||||
- POSTGRES_DB=deepaudit
|
||||
- POSTGRES_PASSWORD=${POSTGRES_PASSWORD:?Set POSTGRES_PASSWORD}
|
||||
- SECRET_KEY=${SECRET_KEY:?Set SECRET_KEY to a random secret of at least 32 characters}
|
||||
- ENVIRONMENT=production
|
||||
command: [".venv/bin/alembic", "upgrade", "head"]
|
||||
depends_on:
|
||||
db:
|
||||
|
||||
+33
-2
@@ -43,15 +43,15 @@ services:
|
||||
volumes:
|
||||
# - ./backend/app:/app/app:ro # 挂载代码目录,修改后自动生效
|
||||
- backend_uploads:/app/uploads
|
||||
- /var/run/docker.sock:/var/run/docker.sock # 沙箱执行必须
|
||||
ports:
|
||||
- "8000:8000"
|
||||
env_file:
|
||||
- ./backend/.env
|
||||
environment:
|
||||
environment: &audit_environment
|
||||
- DATABASE_URL=postgresql+asyncpg://postgres:postgres@db:5432/deepaudit
|
||||
- REDIS_URL=redis://redis:6379/0
|
||||
- AGENT_ENABLED=true
|
||||
- AUDIT_WORKSPACE_PATH=/app/uploads/audit_workspaces
|
||||
- SANDBOX_ENABLED=true
|
||||
- SANDBOX_IMAGE=deepaudit/sandbox:latest # 使用本地构建的沙箱镜像
|
||||
# 禁用代理设置,防止容器内无法连接外部 API
|
||||
@@ -69,6 +69,37 @@ services:
|
||||
command: sh -c ".venv/bin/alembic upgrade head && .venv/bin/uvicorn app.main:app --host 0.0.0.0 --port 8000 --reload"
|
||||
networks:
|
||||
- deepaudit-network
|
||||
healthcheck:
|
||||
test: ["CMD", ".venv/bin/python", "-c", "import urllib.request; urllib.request.urlopen('http://localhost:8000/ready')"]
|
||||
interval: 10s
|
||||
timeout: 5s
|
||||
retries: 5
|
||||
start_period: 30s
|
||||
|
||||
|
||||
worker:
|
||||
build:
|
||||
context: ./backend
|
||||
restart: unless-stopped
|
||||
init: true
|
||||
command: [".venv/bin/python", "-m", "app.worker"]
|
||||
environment: *audit_environment
|
||||
env_file:
|
||||
- ./backend/.env
|
||||
volumes:
|
||||
- backend_uploads:/app/uploads
|
||||
- /var/run/docker.sock:/var/run/docker.sock
|
||||
depends_on:
|
||||
backend:
|
||||
condition: service_healthy
|
||||
healthcheck:
|
||||
test: ["CMD", ".venv/bin/python", "-m", "app.worker", "--health"]
|
||||
interval: 15s
|
||||
timeout: 5s
|
||||
retries: 3
|
||||
start_period: 30s
|
||||
networks:
|
||||
- deepaudit-network
|
||||
|
||||
frontend:
|
||||
build:
|
||||
|
||||
+24
-1
@@ -1,5 +1,7 @@
|
||||
# 配置说明
|
||||
|
||||
> 本次升级新增独立审计 worker,需要运行数据库迁移并同步更新 API、worker 与前端。详见[运行与升级说明](OPERATIONS.md)。
|
||||
|
||||
本文档详细介绍 DeepAudit 的所有配置选项,包括后端环境变量、前端配置和运行时配置。
|
||||
|
||||
## 目录
|
||||
@@ -34,6 +36,24 @@ DeepAudit 采用前后端分离架构,数据存储在后端 PostgreSQL 数据
|
||||
cp backend/env.example backend/.env
|
||||
```
|
||||
|
||||
启动前,使用 `openssl rand -hex 32` 生成随机密钥,填写到 `backend/.env` 的
|
||||
`SECRET_KEY` 中。密钥不能为空、不能使用旧的示例值,且至少需要 32 个字符。
|
||||
使用生产 Compose 文件时,在项目根目录 `.env` 中设置相同变量,或通过
|
||||
`docker compose --env-file backend/.env -f docker-compose.prod.yml up -d` 显式传入。
|
||||
国内镜像配置 `docker-compose.prod.cn.yml` 使用相同方式。
|
||||
|
||||
默认 `ENVIRONMENT=production`、`DEMO_ENABLED=false`,不再自动创建演示管理员。
|
||||
如需本地演示,必须同时设置 `ENVIRONMENT=development` 和 `DEMO_ENABLED=true`;
|
||||
生产环境会拒绝启用演示账户。已有部署升级后,仍使用默认密码的原演示账户会被禁用,
|
||||
其数据保留;升级前请创建独立管理员或修改原演示账户密码。
|
||||
|
||||
`SECRET_KEY` 也用于加密用户保存的 API Key 等配置。已有部署请保留安全的原密钥;
|
||||
如果原来使用默认密钥,则必须更换,之后需要重新登录并重新保存加密配置。
|
||||
生产环境数据库初始化失败时,应用会停止启动,请先完成数据库迁移。
|
||||
|
||||
同域反向代理部署可将 `BACKEND_CORS_ORIGINS` 设为 `[]`。前后端跨域时,
|
||||
填写实际前端来源,例如 `["http://localhost:5173"]`,不支持通配符。
|
||||
|
||||
### 完整配置参考
|
||||
|
||||
```env
|
||||
@@ -49,7 +69,10 @@ POSTGRES_DB=deepaudit # 数据库名称
|
||||
# DATABASE_URL= # 完整数据库连接字符串(可选,会覆盖上述配置)
|
||||
|
||||
# ========== 安全配置 ==========
|
||||
SECRET_KEY=your-super-secret-key # JWT 签名密钥(生产环境必须修改!)
|
||||
SECRET_KEY= # 必填,至少 32 字符的随机密钥
|
||||
ENVIRONMENT=production
|
||||
DEMO_ENABLED=false
|
||||
BACKEND_CORS_ORIGINS=[] # 跨域部署时填写实际前端来源
|
||||
ALGORITHM=HS256 # JWT 加密算法
|
||||
ACCESS_TOKEN_EXPIRE_MINUTES=11520 # Token 过期时间(分钟),默认 8 天
|
||||
|
||||
|
||||
+11
-3
@@ -1,5 +1,7 @@
|
||||
# 部署指南
|
||||
|
||||
> 本次升级新增独立审计 worker,需要运行数据库迁移并同步更新 API、worker 与前端。详见[运行与升级说明](OPERATIONS.md)。
|
||||
|
||||
本文档详细介绍 DeepAudit v3.0.0 的各种部署方式,包括 Docker Compose 一键部署、Agent 审计模式部署和本地开发环境搭建。
|
||||
|
||||
## 目录
|
||||
@@ -36,12 +38,15 @@ docker compose up -d
|
||||
|
||||
### 演示账户
|
||||
|
||||
系统启动时会自动创建演示账户,包含示例项目和审计数据,可直接体验完整功能:
|
||||
默认不创建演示账户。仅在本地开发中同时设置 `ENVIRONMENT=development` 和
|
||||
`DEMO_ENABLED=true` 后,系统才会创建以下演示账户及示例数据:
|
||||
|
||||
- 📧 邮箱:`demo@example.com`
|
||||
- 🔑 密码:`demo123`
|
||||
|
||||
> ⚠️ **安全提示**: 生产环境部署后,请删除演示账户或修改密码。
|
||||
> 生产环境禁止启用演示账户,并要求配置至少 32 字符的随机 `SECRET_KEY`。
|
||||
> 升级会禁用仍使用默认密码的旧演示账户。升级前请创建独立管理员或修改密码,
|
||||
> 详情见[安全配置与迁移说明](CONFIGURATION.md#后端配置)。
|
||||
|
||||
---
|
||||
|
||||
@@ -79,7 +84,8 @@ POSTGRES_PASSWORD=postgres
|
||||
POSTGRES_DB=deepaudit
|
||||
|
||||
# 安全配置(生产环境请修改)
|
||||
SECRET_KEY=your-super-secret-key-change-this-in-production
|
||||
# 使用 openssl rand -hex 32 生成随机值;已有部署保留原有安全密钥
|
||||
SECRET_KEY=<填入随机密钥>
|
||||
|
||||
# LLM 配置(必填)
|
||||
LLM_PROVIDER=openai
|
||||
@@ -332,6 +338,8 @@ uv run alembic upgrade head
|
||||
|
||||
# 6. 启动后端服务(开发模式,支持热重载)
|
||||
uv run uvicorn app.main:app --reload --port 8000
|
||||
# 另开终端,同目录启动审计 worker
|
||||
uv run python -m app.worker
|
||||
```
|
||||
|
||||
### 前端启动
|
||||
|
||||
@@ -0,0 +1,95 @@
|
||||
# 审计队列、资源限制与升级
|
||||
|
||||
## 运行方式
|
||||
|
||||
审计 API 只创建任务和数据库队列记录。独立 worker 从 PostgreSQL 领取任务,每个任务运行在单独的进程中。普通仓库、ZIP 和 Agent 审计均使用该队列;即时单段代码分析仍由请求直接处理。
|
||||
|
||||
源码部署需要同时启动 API 和 worker,工作目录均为 `backend`:
|
||||
|
||||
```bash
|
||||
uv sync --frozen --group dev
|
||||
uv run alembic upgrade head
|
||||
uv run uvicorn app.main:app --host 0.0.0.0 --port 8000
|
||||
# 在另一个终端中启动,使用相同 .env 和工作目录
|
||||
uv run python -m app.worker
|
||||
```
|
||||
|
||||
仓库中的三个 Compose 配置均已包含 worker。API、worker、数据库迁移必须使用包含此次改动的同一版本代码;旧的已发布镜像可能没有 worker 入口。测试本地改动时使用源码构建:
|
||||
|
||||
```bash
|
||||
# 先配置 backend/.env 中的随机 SECRET_KEY、数据库和模型参数
|
||||
# 首次部署的账户通过注册页面创建;默认不再生成演示管理员。
|
||||
docker compose up -d --build
|
||||
```
|
||||
|
||||
worker 默认每次处理一个任务。增加吞吐量可启动多个 worker,例如 `docker compose up -d --scale worker=3`。每个 worker 都必须访问同一数据库和共享上传目录,使用相同的 `SECRET_KEY`。跨主机部署需要共享文件存储,并确保沙箱所在的 Docker 主机能访问该存储。
|
||||
|
||||
## 升级步骤
|
||||
|
||||
1. 备份 PostgreSQL、上传目录和原有密钥;停止旧版本 API 与后台执行进程,防止旧任务与新 worker 同时运行。
|
||||
2. 检查安全配置。生产环境禁止启用演示账户,`SECRET_KEY` 必须是至少 32 字符的随机值。保留原有安全密钥,否则旧的加密配置和队列配置快照无法解密。生产 Compose 还要求明确设置 `POSTGRES_PASSWORD`;对已有数据库,修改环境变量不会自动修改数据库用户密码,应与数据库实际密码保持一致。
|
||||
3. 执行 `alembic upgrade head`,创建 `audit_jobs` 表和查询索引。迁移会把旧版本遗留的待执行、执行中任务加入队列,暂停任务保持暂停。
|
||||
4. 启动新版 API、worker 和前端。检查 API `/ready` 与 worker 健康状态,然后提交一个小项目验证部署。
|
||||
|
||||
使用生产 Compose 时,在根目录 `.env` 或进程环境设置 `SECRET_KEY`、`POSTGRES_PASSWORD` 和模型配置。普通源码 Compose 从 `backend/.env` 读取应用配置,数据库仍采用其本地开发默认配置。
|
||||
|
||||
## 恢复与取消
|
||||
|
||||
任务创建与入队在同一数据库事务内提交。worker 持有 PostgreSQL 会话锁,租约过期也不会让第二个 worker 同时执行同一任务。进程退出或部署重启释放锁后,租约到期的任务可被重新领取;最多执行 3 次。
|
||||
|
||||
任务保留源码快照、加密的用户配置快照、逐文件分析缓存和已完成的 Agent 结果。恢复时复用已完成结果,未完成的分析步骤会重新执行,因此中断中的模型调用可能重复计费。尚未完成下载的仓库会重新下载;快照准备完成后,重试使用该快照。这里不提供模型生成到一半的逐 token 恢复。
|
||||
|
||||
取消请求写入数据库,worker 在心跳时检查并中止当前任务。部署停止与用户取消分别处理:部署中断的任务可以恢复。启动任务前,worker 会清理带有该任务标签的遗留沙箱容器。
|
||||
|
||||
完成、失败或取消的任务工作目录默认保留 7 天,worker 每小时清理一次到期目录。数据库中的报告和事件仍保留,项目原始 ZIP 存储独立保留。不要在运行中手动删除共享工作目录。
|
||||
|
||||
## 资源配置
|
||||
|
||||
以下参数放在 `backend/.env`;自定义生产 Compose 环境时应同时传给 API 和 worker。
|
||||
|
||||
| 参数 | 默认值 | 用途 |
|
||||
| --- | --- | --- |
|
||||
| `AUDIT_WORKSPACE_PATH` | `./uploads/audit_workspaces` | 源码快照、逐文件缓存、向量数据;Compose 使用 `/app/uploads/audit_workspaces` |
|
||||
| `WORKER_POLL_SECONDS` | 1 | 队列轮询间隔 |
|
||||
| `WORKER_HEARTBEAT_SECONDS` | 2 | 租约刷新与取消检查间隔 |
|
||||
| `WORKER_LEASE_SECONDS` | 30 | 失联任务重新领取的等待时间;应大于心跳间隔 |
|
||||
| `WORKER_MAX_ATTEMPTS` | 3 | 单任务最大执行次数 |
|
||||
| `WORKER_TASK_TIMEOUT_SECONDS` | 3600 | 普通任务执行上限;Agent 使用任务自己的超时设置 |
|
||||
| `WORKER_RETENTION_DAYS` | 7 | 终态任务工作目录保留天数 |
|
||||
| `ZIP_MAX_UPLOAD_BYTES` | 524288000 | 压缩上传或下载大小上限,500 MiB |
|
||||
| `ZIP_MAX_EXTRACTED_BYTES` | 2147483648 | ZIP 解压总量上限,2 GiB |
|
||||
| `ZIP_MAX_FILE_BYTES` | 52428800 | 单个 ZIP 条目大小上限,50 MiB |
|
||||
| `ZIP_MAX_ENTRIES` | 50000 | ZIP 最大条目数 |
|
||||
| `ZIP_MAX_COMPRESSION_RATIO` | 200 | 单条目最大压缩比 |
|
||||
|
||||
上传和仓库归档下载按块读取。ZIP 在解压前校验所有条目,禁止路径穿越、符号链接、加密条目和重复路径;实际解压时再次计量。取消或校验失败时丢弃暂存内容。普通扫描逐个读取文件,另受 `MAX_FILE_SIZE_BYTES` 和扫描文件数量配置限制。
|
||||
|
||||
Docker socket 仅挂载给 worker。沙箱以任务标签标识,并将 worker 内的共享卷路径映射到 Docker 主机路径。API 无需访问 Docker socket。
|
||||
|
||||
## 事件、分页与健康检查
|
||||
|
||||
事件先写入数据库,再由 API 按序号读取。`/agent-tasks/{id}/events` 与 `/stream` 都支持 `after_sequence` 和 `Last-Event-ID`,多个客户端各自读取。高频思考 token 合并成块,最后一块由 worker 心跳或后续事件刷盘;尚未刷盘时进程崩溃可能丢失这一小段思考文本,不影响已保存的报告结果。
|
||||
|
||||
任务、问题与发现列表返回 `X-Total-Count`,支持 `skip`、`limit`,分页上限由各端点约束。任务页面按服务端状态和搜索条件分页;导出显式读取全部分页。统计面板由数据库聚合,避免逐任务拉取所有问题。页面通过按路由加载减少首屏下载量。
|
||||
|
||||
- `/health`:API 进程存活。
|
||||
- `/ready`:数据库可访问,队列表已迁移。
|
||||
- `python -m app.worker --health`:worker 最近 60 秒有活动心跳。容器健康检查持续失败时查看 worker 日志及数据库连接。
|
||||
|
||||
## 验证
|
||||
|
||||
```bash
|
||||
cd backend
|
||||
uv run --frozen pytest tests
|
||||
# 设置独立测试数据库并先迁移,运行真实队列恢复、锁与事件测试
|
||||
DATABASE_URL="$TEST_DATABASE_URL" uv run --frozen alembic upgrade head
|
||||
TEST_DATABASE_URL=postgresql+asyncpg://USER:PASSWORD@localhost/TEST_DB uv run --frozen pytest tests/test_durable_execution.py
|
||||
|
||||
cd ../frontend
|
||||
pnpm type-check
|
||||
pnpm lint
|
||||
pnpm test
|
||||
pnpm build
|
||||
```
|
||||
|
||||
CI 自动启动独立 PostgreSQL、执行迁移、后端测试,以及前端类型检查、静态检查、测试和构建。发布工作流依赖这些检查。测试不调用外部模型;真实模型与 Docker 沙箱的端到端验证需要部署环境。
|
||||
+3
-3
@@ -3,7 +3,7 @@
|
||||
# =============================================
|
||||
# 使用 Nginx 提供静态文件和反向代理 (支持 SSE 流式传输)
|
||||
|
||||
FROM node:20-slim AS builder
|
||||
FROM node:22.14.0-slim AS builder
|
||||
|
||||
WORKDIR /app
|
||||
|
||||
@@ -19,7 +19,7 @@ ENV NO_PROXY="*"
|
||||
|
||||
# 配置国内 npm 镜像源
|
||||
RUN npm config set registry https://registry.npmmirror.com && \
|
||||
npm install -g pnpm && \
|
||||
npm install -g pnpm@9 && \
|
||||
pnpm config set registry https://registry.npmmirror.com
|
||||
|
||||
# 复制依赖文件
|
||||
@@ -28,7 +28,7 @@ COPY package.json pnpm-lock.yaml ./
|
||||
# 增加网络超时设置和并发数限制,防止 ARM 架构构建卡死
|
||||
RUN pnpm config set network-timeout 300000 && \
|
||||
pnpm config set fetch-retries 5 && \
|
||||
pnpm install --no-frozen-lockfile --network-concurrency 1
|
||||
pnpm install --frozen-lockfile --network-concurrency 1
|
||||
|
||||
# 复制源代码
|
||||
COPY . .
|
||||
|
||||
@@ -0,0 +1,7 @@
|
||||
{
|
||||
"$schema": "./node_modules/@biomejs/biome/configuration_schema.json",
|
||||
"files": {
|
||||
"includes": ["**", "!!dist"],
|
||||
"experimentalScannerIgnores": [".git", "dist"]
|
||||
}
|
||||
}
|
||||
@@ -6,9 +6,9 @@
|
||||
"dev": "vite",
|
||||
"build": "vite build",
|
||||
"preview": "vite preview --strictPort --port 5173",
|
||||
"lint": "tsgo -p tsconfig.check.json; biome lint --only=correctness/noUndeclaredDependencies; ast-grep scan",
|
||||
"lint": "tsgo -p tsconfig.check.json && biome lint src --only=correctness/noUndeclaredDependencies && ast-grep scan",
|
||||
"lint:fix": "biome check --apply .",
|
||||
"type-check": "tsc --noEmit",
|
||||
"type-check": "tsc --noEmit -p tsconfig.check.json",
|
||||
"format": "biome format --write .",
|
||||
"clean": "rm -rf dist node_modules/.vite",
|
||||
"setup": "node scripts/setup.js || ./scripts/setup.sh",
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
import { useState } from "react";
|
||||
import { useState, Suspense } from "react";
|
||||
import { BrowserRouter, Routes, Route, Outlet } from "react-router-dom";
|
||||
import { Toaster } from "sonner";
|
||||
import Sidebar from "@/components/layout/Sidebar";
|
||||
@@ -19,7 +19,9 @@ function AppLayout() {
|
||||
className={`transition-all duration-300 min-h-screen ${collapsed ? "md:ml-20" : "md:ml-64"
|
||||
}`}
|
||||
>
|
||||
<Outlet />
|
||||
<Suspense fallback={<div role="status" className="p-8 text-muted-foreground">加载中…</div>}>
|
||||
<Outlet />
|
||||
</Suspense>
|
||||
</main>
|
||||
</div>
|
||||
);
|
||||
|
||||
+13
-12
@@ -1,15 +1,16 @@
|
||||
import Dashboard from "@/pages/Dashboard";
|
||||
import Projects from "@/pages/Projects";
|
||||
import ProjectDetail from "@/pages/ProjectDetail";
|
||||
import RecycleBin from "@/pages/RecycleBin";
|
||||
import InstantAnalysis from "@/pages/InstantAnalysis";
|
||||
import AuditTasks from "@/pages/AuditTasks";
|
||||
import TaskDetail from "@/pages/TaskDetail";
|
||||
import AgentAudit from "@/pages/AgentAudit";
|
||||
import AdminDashboard from "@/pages/AdminDashboard";
|
||||
import Account from "@/pages/Account";
|
||||
import AuditRules from "@/pages/AuditRules";
|
||||
import PromptManager from "@/pages/PromptManager";
|
||||
import { lazy } from "react";
|
||||
const Dashboard = lazy(() => import("@/pages/Dashboard"));
|
||||
const Projects = lazy(() => import("@/pages/Projects"));
|
||||
const ProjectDetail = lazy(() => import("@/pages/ProjectDetail"));
|
||||
const RecycleBin = lazy(() => import("@/pages/RecycleBin"));
|
||||
const InstantAnalysis = lazy(() => import("@/pages/InstantAnalysis"));
|
||||
const AuditTasks = lazy(() => import("@/pages/AuditTasks"));
|
||||
const TaskDetail = lazy(() => import("@/pages/TaskDetail"));
|
||||
const AgentAudit = lazy(() => import("@/pages/AgentAudit"));
|
||||
const AdminDashboard = lazy(() => import("@/pages/AdminDashboard"));
|
||||
const Account = lazy(() => import("@/pages/Account"));
|
||||
const AuditRules = lazy(() => import("@/pages/AuditRules"));
|
||||
const PromptManager = lazy(() => import("@/pages/PromptManager"));
|
||||
import type { ReactNode } from 'react';
|
||||
|
||||
export interface RouteConfig {
|
||||
|
||||
@@ -0,0 +1,38 @@
|
||||
import { Button } from "@/components/ui/button";
|
||||
|
||||
export function Pagination({
|
||||
page,
|
||||
pageSize,
|
||||
total,
|
||||
onChange,
|
||||
disabled = false,
|
||||
}: {
|
||||
page: number;
|
||||
pageSize: number;
|
||||
total: number;
|
||||
onChange: (page: number) => void;
|
||||
disabled?: boolean;
|
||||
}) {
|
||||
const pages = Math.max(1, Math.ceil(total / pageSize));
|
||||
return (
|
||||
<nav aria-label="分页" className="flex items-center justify-end gap-3 py-4">
|
||||
<span className="text-sm text-muted-foreground">
|
||||
共 {total} 条 · 第 {page + 1} / {pages} 页
|
||||
</span>
|
||||
<Button
|
||||
variant="outline"
|
||||
disabled={disabled || page === 0}
|
||||
onClick={() => onChange(page - 1)}
|
||||
>
|
||||
上一页
|
||||
</Button>
|
||||
<Button
|
||||
variant="outline"
|
||||
disabled={disabled || page + 1 >= pages}
|
||||
onClick={() => onChange(page + 1)}
|
||||
>
|
||||
下一页
|
||||
</Button>
|
||||
</nav>
|
||||
);
|
||||
}
|
||||
@@ -3,7 +3,7 @@
|
||||
* 显示当前使用的数据库模式
|
||||
*/
|
||||
|
||||
import { Database, Cloud, Eye, Server } from 'lucide-react';
|
||||
import { Server } from 'lucide-react';
|
||||
import { Badge } from '@/components/ui/badge';
|
||||
import { dbMode } from '@/shared/config/database';
|
||||
|
||||
@@ -17,34 +17,6 @@ export function DatabaseStatus() {
|
||||
variant: 'default' as const,
|
||||
description: '数据存储在后端 PostgreSQL 数据库'
|
||||
};
|
||||
case 'local':
|
||||
return {
|
||||
icon: Database,
|
||||
label: '本地数据库',
|
||||
variant: 'default' as const,
|
||||
description: '数据存储在浏览器本地'
|
||||
};
|
||||
case 'supabase':
|
||||
return {
|
||||
icon: Cloud,
|
||||
label: 'Supabase 云端',
|
||||
variant: 'secondary' as const,
|
||||
description: '数据存储在云端(已废弃)'
|
||||
};
|
||||
case 'demo':
|
||||
return {
|
||||
icon: Eye,
|
||||
label: '演示模式',
|
||||
variant: 'outline' as const,
|
||||
description: '使用演示数据,不会持久化'
|
||||
};
|
||||
default:
|
||||
return {
|
||||
icon: Database,
|
||||
label: '未知模式',
|
||||
variant: 'destructive' as const,
|
||||
description: ''
|
||||
};
|
||||
}
|
||||
};
|
||||
|
||||
@@ -70,38 +42,6 @@ export function DatabaseStatusDetail() {
|
||||
description: '数据存储在后端 PostgreSQL 数据库中,通过 REST API 访问。支持多用户、多设备同步。',
|
||||
tips: '提示:所有数据操作都通过后端 API 进行,确保网络连接正常。'
|
||||
};
|
||||
case 'local':
|
||||
return {
|
||||
icon: Database,
|
||||
label: '本地数据库模式',
|
||||
variant: 'default' as const,
|
||||
description: '数据存储在浏览器 IndexedDB 中,完全本地化,隐私安全。',
|
||||
tips: '提示:定期导出数据以防丢失。'
|
||||
};
|
||||
case 'supabase':
|
||||
return {
|
||||
icon: Cloud,
|
||||
label: 'Supabase 云端模式(已废弃)',
|
||||
variant: 'secondary' as const,
|
||||
description: '此模式已不再使用,请使用后端数据库模式。',
|
||||
tips: '提示:已迁移到后端 PostgreSQL 数据库。'
|
||||
};
|
||||
case 'demo':
|
||||
return {
|
||||
icon: Eye,
|
||||
label: '演示模式',
|
||||
variant: 'outline' as const,
|
||||
description: '使用内置演示数据,所有操作不会持久化保存。',
|
||||
tips: '提示:配置数据库以保存您的数据。'
|
||||
};
|
||||
default:
|
||||
return {
|
||||
icon: Database,
|
||||
label: '未知模式',
|
||||
variant: 'destructive' as const,
|
||||
description: '数据库配置异常',
|
||||
tips: ''
|
||||
};
|
||||
}
|
||||
};
|
||||
|
||||
|
||||
@@ -0,0 +1,36 @@
|
||||
import { act, renderHook } from '@testing-library/react';
|
||||
import { describe, expect, it } from 'vitest';
|
||||
import { useAgentAuditState } from '../hooks/useAgentAuditState';
|
||||
|
||||
describe('streamed findings', () => {
|
||||
it('normalizes a partial event before it reaches the findings UI', () => {
|
||||
const { result } = renderHook(() => useAgentAuditState());
|
||||
act(() => result.current.dispatch({
|
||||
type: 'ADD_FINDING',
|
||||
payload: { id: 'finding-1', task_id: 'task-1', title: 'SQL injection' },
|
||||
}));
|
||||
expect(result.current.findings[0]).toMatchObject({
|
||||
id: 'finding-1', task_id: 'task-1', title: 'SQL injection',
|
||||
status: 'new', is_verified: false, has_poc: false,
|
||||
description: null, file_path: null, line_start: null,
|
||||
});
|
||||
expect(Object.values(result.current.findings[0])).not.toContain(undefined);
|
||||
});
|
||||
|
||||
it('preserves supplied values and deduplicates repeated events', () => {
|
||||
const { result } = renderHook(() => useAgentAuditState());
|
||||
act(() => {
|
||||
const payload = {
|
||||
id: 'finding-1', ai_confidence: 0, is_verified: true,
|
||||
has_poc: true, poc_code: 'proof', severity: 'critical',
|
||||
};
|
||||
result.current.dispatch({ type: 'ADD_FINDING', payload });
|
||||
result.current.dispatch({ type: 'ADD_FINDING', payload });
|
||||
});
|
||||
expect(result.current.findings).toHaveLength(1);
|
||||
expect(result.current.findings[0]).toMatchObject({
|
||||
ai_confidence: 0, is_verified: true, has_poc: true,
|
||||
poc_code: 'proof', severity: 'critical',
|
||||
});
|
||||
});
|
||||
});
|
||||
@@ -44,7 +44,28 @@ function agentAuditReducer(state: AgentAuditState, action: AgentAuditAction): Ag
|
||||
|
||||
case 'ADD_FINDING': {
|
||||
// 🔥 添加单个 finding,避免重复
|
||||
const newFinding = action.payload;
|
||||
const partial = action.payload;
|
||||
const newFinding: AgentFinding = {
|
||||
id: partial.id,
|
||||
task_id: partial.task_id ?? state.task?.id ?? '',
|
||||
vulnerability_type: partial.vulnerability_type ?? 'unknown',
|
||||
severity: partial.severity ?? 'medium',
|
||||
title: partial.title ?? 'Vulnerability found',
|
||||
description: partial.description ?? null,
|
||||
file_path: partial.file_path ?? null,
|
||||
line_start: partial.line_start ?? null,
|
||||
line_end: partial.line_end ?? null,
|
||||
code_snippet: partial.code_snippet ?? null,
|
||||
status: partial.status ?? 'new',
|
||||
is_verified: partial.is_verified ?? false,
|
||||
has_poc: partial.has_poc ?? false,
|
||||
poc_code: partial.poc_code ?? null,
|
||||
suggestion: partial.suggestion ?? null,
|
||||
fix_code: partial.fix_code ?? null,
|
||||
ai_explanation: partial.ai_explanation ?? null,
|
||||
ai_confidence: partial.ai_confidence ?? null,
|
||||
created_at: partial.created_at ?? '',
|
||||
};
|
||||
const existingIds = new Set(state.findings.map(f => f.id));
|
||||
if (newFinding.id && existingIds.has(newFinding.id)) {
|
||||
return state; // 已存在,不添加
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
import { parseSSE as parseEventStream } from "@/shared/api/sse";
|
||||
/**
|
||||
* Resilient Stream Hook
|
||||
* Enhanced stream connection with automatic reconnection, heartbeat monitoring,
|
||||
@@ -80,6 +81,10 @@ export function useResilientStream(
|
||||
const heartbeatTimeoutRef = useRef<ReturnType<typeof setTimeout> | null>(null);
|
||||
const streamOptionsRef = useRef(streamOptions);
|
||||
const isDisconnectingRef = useRef(false);
|
||||
const cursorRef = useRef({ taskId, sequence: streamOptions.afterSequence ?? 0 });
|
||||
if (cursorRef.current.taskId !== taskId) {
|
||||
cursorRef.current = { taskId, sequence: streamOptions.afterSequence ?? 0 };
|
||||
}
|
||||
|
||||
// Update refs when options change
|
||||
streamOptionsRef.current = streamOptions;
|
||||
@@ -153,45 +158,16 @@ export function useResilientStream(
|
||||
// ============ SSE Parsing ============
|
||||
|
||||
const parseSSE = useCallback((buffer: string): { parsed: StreamEventData[]; remaining: string } => {
|
||||
const parsed: StreamEventData[] = [];
|
||||
const lines = buffer.split('\n');
|
||||
let remaining = '';
|
||||
let currentEvent: Partial<StreamEventData> = {};
|
||||
|
||||
for (let i = 0; i < lines.length; i++) {
|
||||
const line = lines[i];
|
||||
|
||||
if (line === '') {
|
||||
if (currentEvent.type) {
|
||||
parsed.push(currentEvent as StreamEventData);
|
||||
currentEvent = {};
|
||||
}
|
||||
continue;
|
||||
}
|
||||
|
||||
if (i === lines.length - 1 && !buffer.endsWith('\n')) {
|
||||
remaining = line;
|
||||
break;
|
||||
}
|
||||
|
||||
if (line.startsWith('event:')) {
|
||||
currentEvent.type = line.slice(6).trim() as StreamEventData['type'];
|
||||
} else if (line.startsWith('data:')) {
|
||||
try {
|
||||
const data = JSON.parse(line.slice(5).trim());
|
||||
currentEvent = { ...currentEvent, ...data };
|
||||
} catch {
|
||||
// Ignore parse errors
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return { parsed, remaining };
|
||||
return parseEventStream(buffer);
|
||||
}, []);
|
||||
|
||||
// ============ Event Handling ============
|
||||
|
||||
const handleEvent = useCallback((event: StreamEventData) => {
|
||||
if (typeof event.sequence === 'number' && event.type !== 'task_end') {
|
||||
if (event.sequence <= cursorRef.current.sequence) return;
|
||||
cursorRef.current.sequence = event.sequence;
|
||||
}
|
||||
const opts = streamOptionsRef.current;
|
||||
|
||||
// Extract agent_name from metadata
|
||||
@@ -310,7 +286,7 @@ export function useResilientStream(
|
||||
const params = new URLSearchParams({
|
||||
include_thinking: String(streamOptionsRef.current.includeThinking ?? true),
|
||||
include_tool_calls: String(streamOptionsRef.current.includeToolCalls ?? true),
|
||||
after_sequence: String(streamOptionsRef.current.afterSequence ?? 0),
|
||||
after_sequence: String(Math.max(cursorRef.current.sequence, streamOptionsRef.current.afterSequence ?? 0)),
|
||||
});
|
||||
|
||||
const url = `/api/v1/agent-tasks/${taskId}/stream?${params}`;
|
||||
@@ -346,7 +322,10 @@ export function useResilientStream(
|
||||
if (isDisconnectingRef.current) break;
|
||||
|
||||
const { done, value } = await readerRef.current.read();
|
||||
if (done) break;
|
||||
if (done) {
|
||||
if (!isDisconnectingRef.current) throw new Error('事件连接已中断');
|
||||
break;
|
||||
}
|
||||
|
||||
buffer += decoder.decode(value, { stream: true });
|
||||
const events = parseSSE(buffer);
|
||||
|
||||
@@ -1,43 +1,44 @@
|
||||
import { Pagination } from "@/components/common/Pagination";
|
||||
import { fetchPage } from "@/shared/api/pagination";
|
||||
import { useDebounce } from "@/shared/hooks/useDebounce";
|
||||
/**
|
||||
* Audit Tasks Page
|
||||
* Cyberpunk Terminal Aesthetic
|
||||
* 支持普通审计任务和Agent审计任务
|
||||
*/
|
||||
|
||||
import { useState, useEffect, useRef } from "react";
|
||||
import { useEffect,useRef,useState } from "react";
|
||||
|
||||
import { Button } from "@/components/ui/button";
|
||||
import { Badge } from "@/components/ui/badge";
|
||||
import { Input } from "@/components/ui/input";
|
||||
import { Progress } from "@/components/ui/progress";
|
||||
import {
|
||||
Activity,
|
||||
AlertTriangle,
|
||||
CheckCircle,
|
||||
Clock,
|
||||
Search,
|
||||
FileText,
|
||||
Calendar,
|
||||
Plus,
|
||||
XCircle,
|
||||
ArrowUpRight,
|
||||
Shield,
|
||||
Terminal,
|
||||
Bot,
|
||||
Zap,
|
||||
Download
|
||||
} from "lucide-react";
|
||||
import { api } from "@/shared/config/database";
|
||||
import { apiClient } from "@/shared/api/serverClient";
|
||||
import type { AuditTask } from "@/shared/types";
|
||||
import { Link, useNavigate } from "react-router-dom";
|
||||
import { toast } from "sonner";
|
||||
import CreateTaskDialog from "@/components/audit/CreateTaskDialog";
|
||||
import TerminalProgressDialog from "@/components/audit/TerminalProgressDialog";
|
||||
import ExportReportDialog from "@/components/reports/ExportReportDialog";
|
||||
import { calculateTaskProgress } from "@/shared/utils/utils";
|
||||
import { getAgentTasks, cancelAgentTask, getAgentFindings, type AgentTask, type AgentFinding } from "@/shared/api/agentTasks";
|
||||
import { Badge } from "@/components/ui/badge";
|
||||
import { Button } from "@/components/ui/button";
|
||||
import { Input } from "@/components/ui/input";
|
||||
import { Progress } from "@/components/ui/progress";
|
||||
import ReportExportDialog from "@/pages/AgentAudit/components/ReportExportDialog";
|
||||
import { cancelAgentTask,getAgentFindings,type AgentFinding,type AgentTask } from "@/shared/api/agentTasks";
|
||||
import { api } from "@/shared/config/database";
|
||||
import type { AuditTask } from "@/shared/types";
|
||||
import { calculateTaskProgress } from "@/shared/utils/utils";
|
||||
import {
|
||||
Activity,
|
||||
AlertTriangle,
|
||||
ArrowUpRight,
|
||||
Bot,
|
||||
Calendar,
|
||||
CheckCircle,
|
||||
Clock,
|
||||
Download,
|
||||
FileText,
|
||||
Plus,
|
||||
Search,
|
||||
Terminal,
|
||||
XCircle,
|
||||
Zap
|
||||
} from "lucide-react";
|
||||
import { Link,useNavigate } from "react-router-dom";
|
||||
import { toast } from "sonner";
|
||||
|
||||
// Zombie task detection config
|
||||
const ZOMBIE_TIMEOUT = 180000; // 3 minutes without progress is potentially stuck
|
||||
@@ -54,6 +55,16 @@ export default function AuditTasks() {
|
||||
const [loading, setLoading] = useState(true);
|
||||
const [searchTerm, setSearchTerm] = useState("");
|
||||
const [statusFilter, setStatusFilter] = useState<string>("all");
|
||||
const [page, setPage] = useState(0);
|
||||
const [taskTotal, setTaskTotal] = useState(0);
|
||||
const [agentTotal, setAgentTotal] = useState(0);
|
||||
const pageSize = 20;
|
||||
const debouncedSearch = useDebounce(searchTerm, 300);
|
||||
const query = { skip: page * pageSize, limit: pageSize, search: debouncedSearch,
|
||||
...(statusFilter === 'all' ? {} : { status: statusFilter }) };
|
||||
const taskRequest = useRef(0);
|
||||
const agentRequest = useRef(0);
|
||||
|
||||
const [showCreateDialog, setShowCreateDialog] = useState(false);
|
||||
const [cancellingTaskId, setCancellingTaskId] = useState<string | null>(null);
|
||||
const [showTerminal, setShowTerminal] = useState(false);
|
||||
@@ -78,23 +89,28 @@ export default function AuditTasks() {
|
||||
useEffect(() => {
|
||||
loadTasks();
|
||||
loadAgentTasks();
|
||||
}, []);
|
||||
}, [page, debouncedSearch, statusFilter, activeTab]);
|
||||
|
||||
useEffect(() => { setPage(0); }, [debouncedSearch, statusFilter, activeTab]);
|
||||
|
||||
// 加载Agent任务(支持静默更新,不触发 loading 状态)
|
||||
const loadAgentTasks = async (silent = false) => {
|
||||
const ticket = ++agentRequest.current;
|
||||
try {
|
||||
if (!silent) {
|
||||
setAgentLoading(true);
|
||||
}
|
||||
const data = await getAgentTasks();
|
||||
setAgentTasks(data);
|
||||
const data = await fetchPage<AgentTask>('/agent-tasks/', query);
|
||||
if (ticket !== agentRequest.current) return;
|
||||
setAgentTasks(data.items);
|
||||
setAgentTotal(data.total);
|
||||
} catch (error) {
|
||||
console.error('Failed to load agent tasks:', error);
|
||||
if (!silent) {
|
||||
toast.error("加载Agent任务失败");
|
||||
}
|
||||
} finally {
|
||||
if (!silent) {
|
||||
if (ticket === agentRequest.current && !silent) {
|
||||
setAgentLoading(false);
|
||||
}
|
||||
}
|
||||
@@ -111,10 +127,12 @@ export default function AuditTasks() {
|
||||
return;
|
||||
}
|
||||
|
||||
let disposed = false;
|
||||
const intervalId = setInterval(async () => {
|
||||
try {
|
||||
const updatedData = await api.getAuditTasks();
|
||||
const updatedData = (await api.getAuditTaskPage(query)).items;
|
||||
|
||||
if (disposed) return;
|
||||
setTasks(prevTasks => {
|
||||
return prevTasks.map(prevTask => {
|
||||
const updated = updatedData.find(t => t.id === prevTask.id);
|
||||
@@ -165,8 +183,8 @@ export default function AuditTasks() {
|
||||
}
|
||||
}, 3000);
|
||||
|
||||
return () => clearInterval(intervalId);
|
||||
}, [tasks.map(t => t.id + t.status).join(',')]);
|
||||
return () => { disposed = true; clearInterval(intervalId); };
|
||||
}, [tasks.map(t => t.id + t.status).join(','), page, debouncedSearch, statusFilter]);
|
||||
|
||||
// 自动刷新Agent任务(静默更新,不显示 loading)
|
||||
useEffect(() => {
|
||||
@@ -178,7 +196,7 @@ export default function AuditTasks() {
|
||||
|
||||
const intervalId = setInterval(() => loadAgentTasks(true), 5000);
|
||||
return () => clearInterval(intervalId);
|
||||
}, [agentTasks.map(t => t.id + t.status).join(',')]);
|
||||
}, [agentTasks.map(t => t.id + t.status).join(','), page, debouncedSearch, statusFilter, activeTab]);
|
||||
|
||||
const handleCancelTask = async (taskId: string) => {
|
||||
if (cancellingTaskId) return;
|
||||
@@ -218,7 +236,7 @@ export default function AuditTasks() {
|
||||
try {
|
||||
setExportingTaskId(task.id);
|
||||
// 获取任务的问题列表
|
||||
const issuesResponse = await apiClient.get(`/tasks/${task.id}/issues`);
|
||||
const issuesResponse = { data: await api.getAuditIssues(task.id) };
|
||||
setExportTask(task);
|
||||
setExportIssues(issuesResponse.data || []);
|
||||
setShowExportDialog(true);
|
||||
@@ -248,10 +266,13 @@ export default function AuditTasks() {
|
||||
};
|
||||
|
||||
const loadTasks = async () => {
|
||||
const ticket = ++taskRequest.current;
|
||||
try {
|
||||
setLoading(true);
|
||||
const data = await api.getAuditTasks();
|
||||
setTasks(data);
|
||||
const data = await api.getAuditTaskPage(query);
|
||||
if (ticket !== taskRequest.current) return;
|
||||
setTasks(data.items);
|
||||
setTaskTotal(data.total);
|
||||
} catch (error) {
|
||||
console.error('Failed to load tasks:', error);
|
||||
toast.error("加载任务失败");
|
||||
@@ -300,19 +321,9 @@ export default function AuditTasks() {
|
||||
});
|
||||
};
|
||||
|
||||
const filteredTasks = tasks.filter(task => {
|
||||
const matchesSearch = task.project?.name.toLowerCase().includes(searchTerm.toLowerCase()) ||
|
||||
task.task_type.toLowerCase().includes(searchTerm.toLowerCase());
|
||||
const matchesStatus = statusFilter === "all" || task.status === statusFilter;
|
||||
return matchesSearch && matchesStatus;
|
||||
});
|
||||
const filteredTasks = tasks;
|
||||
const filteredAgentTasks = agentTasks;
|
||||
|
||||
const filteredAgentTasks = agentTasks.filter(task => {
|
||||
const matchesSearch = (task.name || '').toLowerCase().includes(searchTerm.toLowerCase()) ||
|
||||
task.task_type.toLowerCase().includes(searchTerm.toLowerCase());
|
||||
const matchesStatus = statusFilter === "all" || task.status === statusFilter;
|
||||
return matchesSearch && matchesStatus;
|
||||
});
|
||||
|
||||
// 统计数据
|
||||
const regularStats = {
|
||||
@@ -945,6 +956,9 @@ export default function AuditTasks() {
|
||||
</>
|
||||
)}
|
||||
|
||||
<Pagination page={page} pageSize={pageSize} total={activeTab === 'agent' ? agentTotal : taskTotal}
|
||||
onChange={setPage} disabled={activeTab === 'agent' ? agentLoading : loading} />
|
||||
|
||||
{/* Create Task Dialog */}
|
||||
<CreateTaskDialog
|
||||
open={showCreateDialog}
|
||||
|
||||
@@ -3,26 +3,42 @@
|
||||
* Cyberpunk Terminal Aesthetic
|
||||
*/
|
||||
|
||||
import { useState, useEffect } from "react";
|
||||
import { Button } from "@/components/ui/button";
|
||||
import { Badge } from "@/components/ui/badge";
|
||||
import {
|
||||
LineChart, Line, PieChart, Pie, Cell,
|
||||
XAxis, YAxis, CartesianGrid, Tooltip, ResponsiveContainer
|
||||
} from "recharts";
|
||||
import {
|
||||
Activity, AlertTriangle, Clock, Code,
|
||||
FileText, GitBranch, Shield, TrendingUp, Zap,
|
||||
BarChart3, Target, ArrowUpRight, Calendar,
|
||||
MessageSquare, Bot, Cpu, Terminal
|
||||
} from "lucide-react";
|
||||
import { api, dbMode, isDemoMode } from "@/shared/config/database";
|
||||
import type { Project, AuditTask, ProjectStats, UnifiedTask } from "@/shared/types";
|
||||
import { Link } from "react-router-dom";
|
||||
import { toast } from "sonner";
|
||||
import { getRuleSets } from "@/shared/api/rules";
|
||||
import { Button } from "@/components/ui/button";
|
||||
import { getAgentTasks,type AgentTask } from "@/shared/api/agentTasks";
|
||||
import { getPromptTemplates } from "@/shared/api/prompts";
|
||||
import { getAgentTasks, type AgentTask } from "@/shared/api/agentTasks";
|
||||
import { getRuleSets } from "@/shared/api/rules";
|
||||
import { api,dbMode,isDemoMode } from "@/shared/config/database";
|
||||
import type { AuditTask,Project,ProjectStats,UnifiedTask } from "@/shared/types";
|
||||
import {
|
||||
Activity,AlertTriangle,
|
||||
ArrowUpRight,
|
||||
BarChart3,
|
||||
Bot,
|
||||
Calendar,
|
||||
Clock,Code,
|
||||
Cpu,
|
||||
FileText,GitBranch,
|
||||
MessageSquare,
|
||||
Shield,
|
||||
Target,
|
||||
Terminal,
|
||||
TrendingUp,Zap
|
||||
} from "lucide-react";
|
||||
import { useEffect,useState } from "react";
|
||||
import { Link } from "react-router-dom";
|
||||
import {
|
||||
CartesianGrid,
|
||||
Cell,
|
||||
Line,
|
||||
LineChart,
|
||||
Pie,
|
||||
PieChart,
|
||||
ResponsiveContainer,
|
||||
Tooltip,
|
||||
XAxis,YAxis
|
||||
} from "recharts";
|
||||
import { toast } from "sonner";
|
||||
|
||||
export default function Dashboard() {
|
||||
const [stats, setStats] = useState<ProjectStats | null>(null);
|
||||
@@ -45,7 +61,7 @@ export default function Dashboard() {
|
||||
const results = await Promise.allSettled([
|
||||
api.getProjectStats(),
|
||||
api.getProjects(),
|
||||
api.getAuditTasks(),
|
||||
api.getAuditTaskPage({ limit: 10 }).then(page => page.items),
|
||||
getAgentTasks({ limit: 10 })
|
||||
]);
|
||||
|
||||
@@ -108,17 +124,8 @@ export default function Dashboard() {
|
||||
}
|
||||
|
||||
try {
|
||||
const allIssues = await Promise.all(
|
||||
tasks.map(task => api.getAuditIssues(task.id).catch(() => []))
|
||||
);
|
||||
const flatIssues = allIssues.flat();
|
||||
|
||||
if (flatIssues.length > 0) {
|
||||
const typeCount: Record<string, number> = {};
|
||||
flatIssues.forEach(issue => {
|
||||
typeCount[issue.issue_type] = (typeCount[issue.issue_type] || 0) + 1;
|
||||
});
|
||||
|
||||
const typeCount = results[0].status === 'fulfilled' ? (results[0].value.issue_types || {}) : {};
|
||||
if (Object.keys(typeCount).length > 0) {
|
||||
const typeMap: Record<string, { name: string; color: string }> = {
|
||||
security: { name: '安全问题', color: '#f43f5e' },
|
||||
bug: { name: '潜在Bug', color: '#f97316' },
|
||||
|
||||
+576
-742
File diff suppressed because it is too large
Load Diff
@@ -3,7 +3,7 @@ import { useEffect, useMemo, useState } from "react";
|
||||
import { api } from "@/shared/config/database";
|
||||
import type { Project, AuditTask } from "@/shared/types";
|
||||
import type { AgentTask } from "@/shared/api/agentTasks";
|
||||
import { getAgentTasks } from "@/shared/api/agentTasks";
|
||||
import { fetchAllPages } from "@/shared/api/pagination";
|
||||
import type { UnifiedTask } from "@/shared/types";
|
||||
|
||||
export type ProjectDetailCombinedStats = {
|
||||
@@ -28,7 +28,7 @@ export function useProjectDetailData(projectId: string | undefined) {
|
||||
const [projectRes, auditTasksRes, agentTasksRes] = await Promise.allSettled([
|
||||
api.getProjectById(projectId),
|
||||
api.getAuditTasks(projectId),
|
||||
getAgentTasks({ project_id: projectId }),
|
||||
fetchAllPages<AgentTask>("/agent-tasks/", { project_id: projectId }),
|
||||
]);
|
||||
|
||||
if (projectRes.status === "fulfilled") setProject(projectRes.value);
|
||||
|
||||
@@ -1,7 +1,8 @@
|
||||
import { fetchAllPages } from "@/shared/api/pagination";
|
||||
import { useEffect, useMemo, useState } from "react";
|
||||
|
||||
import { apiClient } from "@/shared/api/serverClient";
|
||||
import type { AuditTask } from "@/shared/types";
|
||||
import type { AuditTask, AuditIssue } from "@/shared/types";
|
||||
import type { AgentTask, AgentFinding } from "@/shared/api/agentTasks";
|
||||
import type {
|
||||
AggregatedAuditIssue,
|
||||
@@ -53,22 +54,12 @@ async function mapWithConcurrency<T, R>(
|
||||
return results;
|
||||
}
|
||||
|
||||
async function fetchAuditIssues(taskId: string) {
|
||||
const res = await withTimeout(
|
||||
apiClient.get(`/tasks/${taskId}/issues`),
|
||||
PROJECT_DETAIL_REQUEST_TIMEOUT_MS,
|
||||
`GET /tasks/${taskId}/issues`
|
||||
);
|
||||
return res.data;
|
||||
async function fetchAuditIssues(taskId: string): Promise<AuditIssue[]> {
|
||||
return fetchAllPages<AuditIssue>(`/tasks/${taskId}/issues`);
|
||||
}
|
||||
|
||||
async function fetchAgentFindings(taskId: string): Promise<AgentFinding[]> {
|
||||
const res = await withTimeout(
|
||||
apiClient.get(`/agent-tasks/${taskId}/findings`),
|
||||
PROJECT_DETAIL_REQUEST_TIMEOUT_MS,
|
||||
`GET /agent-tasks/${taskId}/findings`
|
||||
);
|
||||
return res.data;
|
||||
return fetchAllPages<AgentFinding>(`/agent-tasks/${taskId}/findings`);
|
||||
}
|
||||
|
||||
function normalizeSeverity(severity: unknown): LatestProblem["severity"] {
|
||||
|
||||
@@ -0,0 +1,446 @@
|
||||
import { Button } from "@/components/ui/button";
|
||||
|
||||
import { Badge } from "@/components/ui/badge";
|
||||
|
||||
import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs";
|
||||
|
||||
import {
|
||||
AlertTriangle,
|
||||
CheckCircle,
|
||||
ChevronDown,
|
||||
Code,
|
||||
FileText,
|
||||
Info,
|
||||
Lightbulb,
|
||||
Shield,
|
||||
Zap,
|
||||
} from "lucide-react";
|
||||
|
||||
import {
|
||||
DropdownMenu,
|
||||
DropdownMenuContent,
|
||||
DropdownMenuItem,
|
||||
DropdownMenuTrigger,
|
||||
} from "@/components/ui/dropdown-menu";
|
||||
|
||||
import type { AuditIssue } from "@/shared/types";
|
||||
|
||||
function parseAIExplanation(aiExplanation: string) {
|
||||
try {
|
||||
const parsed = JSON.parse(aiExplanation);
|
||||
if (parsed.xai) {
|
||||
return parsed.xai;
|
||||
}
|
||||
if (parsed.what || parsed.why || parsed.how) {
|
||||
return parsed;
|
||||
}
|
||||
return null;
|
||||
} catch (error) {
|
||||
return null;
|
||||
}
|
||||
}
|
||||
|
||||
// Issues List Component
|
||||
export function IssuesList({
|
||||
issues,
|
||||
onStatusChange,
|
||||
}: {
|
||||
issues: AuditIssue[];
|
||||
onStatusChange?: (issue: AuditIssue, newStatus: string) => void;
|
||||
}) {
|
||||
const getSeverityClasses = (severity: string) => {
|
||||
switch (severity) {
|
||||
case "critical":
|
||||
return "severity-critical";
|
||||
case "high":
|
||||
return "severity-high";
|
||||
case "medium":
|
||||
return "severity-medium";
|
||||
case "low":
|
||||
return "severity-low";
|
||||
default:
|
||||
return "severity-info";
|
||||
}
|
||||
};
|
||||
|
||||
const getTypeIcon = (type: string) => {
|
||||
switch (type) {
|
||||
case "security":
|
||||
return <Shield className="w-4 h-4" />;
|
||||
case "bug":
|
||||
return <AlertTriangle className="w-4 h-4" />;
|
||||
case "performance":
|
||||
return <Zap className="w-4 h-4" />;
|
||||
case "style":
|
||||
return <Code className="w-4 h-4" />;
|
||||
case "maintainability":
|
||||
return <FileText className="w-4 h-4" />;
|
||||
default:
|
||||
return <Info className="w-4 h-4" />;
|
||||
}
|
||||
};
|
||||
|
||||
const criticalIssues = issues.filter(
|
||||
(issue) => issue.severity === "critical",
|
||||
);
|
||||
const highIssues = issues.filter((issue) => issue.severity === "high");
|
||||
const mediumIssues = issues.filter((issue) => issue.severity === "medium");
|
||||
const lowIssues = issues.filter((issue) => issue.severity === "low");
|
||||
|
||||
const renderIssue = (issue: AuditIssue, index: number) => (
|
||||
<div
|
||||
key={issue.id || index}
|
||||
className="cyber-card p-4 hover:border-border transition-all group"
|
||||
>
|
||||
<div className="flex items-start justify-between mb-3">
|
||||
<div className="flex items-start space-x-3">
|
||||
<div
|
||||
className={`w-10 h-10 rounded-lg flex items-center justify-center ${
|
||||
issue.severity === "critical"
|
||||
? "bg-rose-500/20 text-rose-400"
|
||||
: issue.severity === "high"
|
||||
? "bg-orange-500/20 text-orange-400"
|
||||
: issue.severity === "medium"
|
||||
? "bg-amber-500/20 text-amber-400"
|
||||
: "bg-sky-500/20 text-sky-400"
|
||||
}`}
|
||||
>
|
||||
{getTypeIcon(issue.issue_type)}
|
||||
</div>
|
||||
<div className="flex-1">
|
||||
<h4 className="font-bold text-base text-foreground mb-1 group-hover:text-primary transition-colors uppercase">
|
||||
{issue.title}
|
||||
</h4>
|
||||
<div className="flex items-center space-x-1 text-xs text-muted-foreground font-mono">
|
||||
<FileText className="w-3 h-3" />
|
||||
<span className="bg-muted px-2 py-0.5 rounded border border-border">
|
||||
{issue.file_path}
|
||||
</span>
|
||||
</div>
|
||||
{issue.line_number && (
|
||||
<div className="flex items-center space-x-1 text-xs text-muted-foreground mt-1 font-mono">
|
||||
<span className="text-primary">></span>
|
||||
<span>LINE: {issue.line_number}</span>
|
||||
{issue.column_number && (
|
||||
<span>, COL: {issue.column_number}</span>
|
||||
)}
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
</div>
|
||||
<div className="flex items-center gap-2">
|
||||
{onStatusChange && (
|
||||
<DropdownMenu>
|
||||
<DropdownMenuTrigger asChild>
|
||||
<Button
|
||||
variant="outline"
|
||||
size="sm"
|
||||
className="text-xs font-mono"
|
||||
>
|
||||
{issue.status === "resolved"
|
||||
? "已解决"
|
||||
: issue.status === "false_positive"
|
||||
? "误报"
|
||||
: "待处理"}
|
||||
<ChevronDown className="w-3 h-3 ml-1" />
|
||||
</Button>
|
||||
</DropdownMenuTrigger>
|
||||
<DropdownMenuContent align="end">
|
||||
<DropdownMenuItem
|
||||
onClick={() => onStatusChange(issue, "resolved")}
|
||||
>
|
||||
已解决
|
||||
</DropdownMenuItem>
|
||||
<DropdownMenuItem
|
||||
onClick={() => onStatusChange(issue, "false_positive")}
|
||||
>
|
||||
误报
|
||||
</DropdownMenuItem>
|
||||
<DropdownMenuItem onClick={() => onStatusChange(issue, "open")}>
|
||||
恢复
|
||||
</DropdownMenuItem>
|
||||
</DropdownMenuContent>
|
||||
</DropdownMenu>
|
||||
)}
|
||||
<Badge
|
||||
className={`${getSeverityClasses(issue.severity)} font-bold uppercase px-2 py-1 rounded text-xs`}
|
||||
>
|
||||
{issue.severity === "critical"
|
||||
? "严重"
|
||||
: issue.severity === "high"
|
||||
? "高"
|
||||
: issue.severity === "medium"
|
||||
? "中等"
|
||||
: "低"}
|
||||
</Badge>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
{issue.description && (
|
||||
<div className="bg-muted border border-border p-3 mb-3 rounded font-mono">
|
||||
<div className="flex items-center mb-1 border-b border-border pb-1">
|
||||
<Info className="w-3 h-3 text-muted-foreground mr-1" />
|
||||
<span className="font-bold text-muted-foreground text-xs uppercase">
|
||||
问题详情
|
||||
</span>
|
||||
</div>
|
||||
<p className="text-foreground text-xs leading-relaxed mt-1">
|
||||
{issue.description}
|
||||
</p>
|
||||
</div>
|
||||
)}
|
||||
|
||||
{issue.code_snippet && (
|
||||
<div className="cyber-bg-elevated p-3 mb-3 border border-border rounded">
|
||||
<div className="flex items-center justify-between mb-2 border-b border-border pb-1">
|
||||
<div className="flex items-center space-x-1">
|
||||
<div className="w-4 h-4 bg-primary rounded flex items-center justify-center">
|
||||
<Code className="w-2 h-2 text-foreground" />
|
||||
</div>
|
||||
<span className="text-emerald-600 dark:text-emerald-400 text-xs font-bold font-mono uppercase">
|
||||
CODE_SNIPPET
|
||||
</span>
|
||||
</div>
|
||||
{issue.line_number && (
|
||||
<span className="text-muted-foreground text-xs font-mono">
|
||||
LINE: {issue.line_number}
|
||||
</span>
|
||||
)}
|
||||
</div>
|
||||
<div className="bg-slate-100 dark:bg-black/40 p-2 border border-border rounded">
|
||||
<pre className="text-xs text-emerald-700 dark:text-emerald-400 font-mono overflow-x-auto">
|
||||
<code>{issue.code_snippet}</code>
|
||||
</pre>
|
||||
</div>
|
||||
</div>
|
||||
)}
|
||||
|
||||
<div className="space-y-3">
|
||||
{issue.suggestion && (
|
||||
<div className="bg-sky-500/10 border border-sky-500/30 p-3 rounded">
|
||||
<div className="flex items-center mb-2 border-b border-sky-500/20 pb-1">
|
||||
<div className="w-5 h-5 bg-sky-500/20 border border-sky-500/40 rounded flex items-center justify-center mr-2">
|
||||
<Lightbulb className="w-3 h-3 text-sky-600 dark:text-sky-400" />
|
||||
</div>
|
||||
<span className="font-bold text-sky-700 dark:text-sky-300 text-sm uppercase">
|
||||
修复建议
|
||||
</span>
|
||||
</div>
|
||||
<p className="text-sky-800 dark:text-sky-200/80 text-xs leading-relaxed font-mono">
|
||||
{issue.suggestion}
|
||||
</p>
|
||||
</div>
|
||||
)}
|
||||
|
||||
{issue.ai_explanation &&
|
||||
(() => {
|
||||
const parsedExplanation = parseAIExplanation(issue.ai_explanation);
|
||||
|
||||
if (parsedExplanation) {
|
||||
return (
|
||||
<div className="bg-violet-500/10 border border-violet-500/30 p-3 rounded">
|
||||
<div className="flex items-center mb-2 border-b border-violet-500/20 pb-1">
|
||||
<div className="w-5 h-5 bg-violet-500/20 border border-violet-500/40 rounded flex items-center justify-center mr-2">
|
||||
<Zap className="w-3 h-3 text-violet-600 dark:text-violet-400" />
|
||||
</div>
|
||||
<span className="font-bold text-violet-700 dark:text-violet-300 text-sm uppercase">
|
||||
AI 解释
|
||||
</span>
|
||||
</div>
|
||||
|
||||
<div className="space-y-2 text-xs font-mono">
|
||||
{parsedExplanation.what && (
|
||||
<div className="border-l-2 border-rose-500 pl-2">
|
||||
<span className="font-bold text-rose-600 dark:text-rose-400 uppercase">
|
||||
问题:
|
||||
</span>
|
||||
<span className="text-foreground ml-1">
|
||||
{parsedExplanation.what}
|
||||
</span>
|
||||
</div>
|
||||
)}
|
||||
|
||||
{parsedExplanation.why && (
|
||||
<div className="border-l-2 border-amber-500 pl-2">
|
||||
<span className="font-bold text-amber-600 dark:text-amber-400 uppercase">
|
||||
原因:
|
||||
</span>
|
||||
<span className="text-foreground ml-1">
|
||||
{parsedExplanation.why}
|
||||
</span>
|
||||
</div>
|
||||
)}
|
||||
|
||||
{parsedExplanation.how && (
|
||||
<div className="border-l-2 border-emerald-500 pl-2">
|
||||
<span className="font-bold text-emerald-600 dark:text-emerald-400 uppercase">
|
||||
方案:
|
||||
</span>
|
||||
<span className="text-foreground ml-1">
|
||||
{parsedExplanation.how}
|
||||
</span>
|
||||
</div>
|
||||
)}
|
||||
|
||||
{parsedExplanation.learn_more && (
|
||||
<div className="border-l-2 border-sky-500 pl-2">
|
||||
<span className="font-bold text-sky-600 dark:text-sky-400 uppercase">
|
||||
链接:
|
||||
</span>
|
||||
<a
|
||||
href={parsedExplanation.learn_more}
|
||||
target="_blank"
|
||||
rel="noopener noreferrer"
|
||||
className="text-sky-600 dark:text-sky-400 hover:text-sky-500 dark:hover:text-sky-300 hover:underline ml-1 font-bold"
|
||||
>
|
||||
{parsedExplanation.learn_more}
|
||||
</a>
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
</div>
|
||||
);
|
||||
} else {
|
||||
return (
|
||||
<div className="bg-violet-500/10 border border-violet-500/30 p-3 rounded">
|
||||
<div className="flex items-center mb-2 border-b border-violet-500/20 pb-1">
|
||||
<Zap className="w-4 h-4 text-violet-600 dark:text-violet-400 mr-2" />
|
||||
<span className="font-bold text-violet-700 dark:text-violet-300 text-sm uppercase">
|
||||
AI 解释
|
||||
</span>
|
||||
</div>
|
||||
<p className="text-foreground text-xs leading-relaxed font-mono">
|
||||
{issue.ai_explanation}
|
||||
</p>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
})()}
|
||||
</div>
|
||||
</div>
|
||||
);
|
||||
|
||||
if (issues.length === 0) {
|
||||
return (
|
||||
<div className="cyber-card p-16 text-center border-dashed">
|
||||
<CheckCircle className="w-16 h-16 text-emerald-600 dark:text-emerald-400 mx-auto mb-4" />
|
||||
<h3 className="text-xl font-bold text-emerald-700 dark:text-emerald-300 mb-2 uppercase">
|
||||
代码质量优秀!
|
||||
</h3>
|
||||
<p className="text-emerald-600 dark:text-emerald-400/80 mb-4 font-mono">
|
||||
恭喜!没有发现任何问题
|
||||
</p>
|
||||
<div className="bg-emerald-500/10 border border-emerald-500/30 p-4 max-w-md mx-auto rounded">
|
||||
<p className="text-emerald-700 dark:text-emerald-300/80 text-sm font-mono">
|
||||
您的代码通过了所有质量检查,包括安全性、性能、可维护性等各个方面的评估。
|
||||
</p>
|
||||
</div>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
||||
return (
|
||||
<Tabs defaultValue="all" className="w-full">
|
||||
<TabsList className="grid w-full grid-cols-5 bg-muted border border-border p-1 h-auto gap-1 rounded">
|
||||
<TabsTrigger
|
||||
value="all"
|
||||
className="data-[state=active]:bg-primary data-[state=active]:text-foreground font-mono font-bold uppercase py-2 text-muted-foreground transition-all rounded-sm text-xs"
|
||||
>
|
||||
全部 ({issues.length})
|
||||
</TabsTrigger>
|
||||
<TabsTrigger
|
||||
value="critical"
|
||||
className="data-[state=active]:bg-rose-500 data-[state=active]:text-foreground font-mono font-bold uppercase py-2 text-muted-foreground transition-all rounded-sm text-xs"
|
||||
>
|
||||
严重 ({criticalIssues.length})
|
||||
</TabsTrigger>
|
||||
<TabsTrigger
|
||||
value="high"
|
||||
className="data-[state=active]:bg-orange-500 data-[state=active]:text-foreground font-mono font-bold uppercase py-2 text-muted-foreground transition-all rounded-sm text-xs"
|
||||
>
|
||||
高 ({highIssues.length})
|
||||
</TabsTrigger>
|
||||
<TabsTrigger
|
||||
value="medium"
|
||||
className="data-[state=active]:bg-amber-500 data-[state=active]:text-background font-mono font-bold uppercase py-2 text-muted-foreground transition-all rounded-sm text-xs"
|
||||
>
|
||||
中等 ({mediumIssues.length})
|
||||
</TabsTrigger>
|
||||
<TabsTrigger
|
||||
value="low"
|
||||
className="data-[state=active]:bg-sky-500 data-[state=active]:text-foreground font-mono font-bold uppercase py-2 text-muted-foreground transition-all rounded-sm text-xs"
|
||||
>
|
||||
低 ({lowIssues.length})
|
||||
</TabsTrigger>
|
||||
</TabsList>
|
||||
|
||||
<TabsContent value="all" className="space-y-4 mt-6">
|
||||
{issues.map((issue, index) => renderIssue(issue, index))}
|
||||
</TabsContent>
|
||||
|
||||
<TabsContent value="critical" className="space-y-4 mt-6">
|
||||
{criticalIssues.length > 0 ? (
|
||||
criticalIssues.map((issue, index) => renderIssue(issue, index))
|
||||
) : (
|
||||
<div className="cyber-card p-12 text-center border-dashed">
|
||||
<CheckCircle className="w-16 h-16 text-emerald-400 mx-auto mb-4" />
|
||||
<h3 className="text-lg font-bold text-foreground uppercase mb-2">
|
||||
没有发现严重问题
|
||||
</h3>
|
||||
<p className="text-muted-foreground font-mono">
|
||||
代码在严重级别的检查中表现良好
|
||||
</p>
|
||||
</div>
|
||||
)}
|
||||
</TabsContent>
|
||||
|
||||
<TabsContent value="high" className="space-y-4 mt-6">
|
||||
{highIssues.length > 0 ? (
|
||||
highIssues.map((issue, index) => renderIssue(issue, index))
|
||||
) : (
|
||||
<div className="cyber-card p-12 text-center border-dashed">
|
||||
<CheckCircle className="w-16 h-16 text-emerald-400 mx-auto mb-4" />
|
||||
<h3 className="text-lg font-bold text-foreground uppercase mb-2">
|
||||
没有发现高优先级问题
|
||||
</h3>
|
||||
<p className="text-muted-foreground font-mono">
|
||||
代码在高优先级检查中表现良好
|
||||
</p>
|
||||
</div>
|
||||
)}
|
||||
</TabsContent>
|
||||
|
||||
<TabsContent value="medium" className="space-y-4 mt-6">
|
||||
{mediumIssues.length > 0 ? (
|
||||
mediumIssues.map((issue, index) => renderIssue(issue, index))
|
||||
) : (
|
||||
<div className="cyber-card p-12 text-center border-dashed">
|
||||
<CheckCircle className="w-16 h-16 text-emerald-400 mx-auto mb-4" />
|
||||
<h3 className="text-lg font-bold text-foreground uppercase mb-2">
|
||||
没有发现中等优先级问题
|
||||
</h3>
|
||||
<p className="text-muted-foreground font-mono">
|
||||
代码在中等优先级检查中表现良好
|
||||
</p>
|
||||
</div>
|
||||
)}
|
||||
</TabsContent>
|
||||
|
||||
<TabsContent value="low" className="space-y-4 mt-6">
|
||||
{lowIssues.length > 0 ? (
|
||||
lowIssues.map((issue, index) => renderIssue(issue, index))
|
||||
) : (
|
||||
<div className="cyber-card p-12 text-center border-dashed">
|
||||
<CheckCircle className="w-16 h-16 text-emerald-400 mx-auto mb-4" />
|
||||
<h3 className="text-lg font-bold text-foreground uppercase mb-2">
|
||||
没有发现低优先级问题
|
||||
</h3>
|
||||
<p className="text-muted-foreground font-mono">
|
||||
代码在低优先级检查中表现良好
|
||||
</p>
|
||||
</div>
|
||||
)}
|
||||
</TabsContent>
|
||||
</Tabs>
|
||||
);
|
||||
}
|
||||
@@ -187,7 +187,7 @@ describe("agentTasks API", () => {
|
||||
|
||||
expect(apiClient.get).toHaveBeenCalledWith(
|
||||
"/agent-tasks/task-1/findings",
|
||||
{ params: { severity: "critical", is_verified: true } },
|
||||
{ params: { severity: "critical", is_verified: true, skip: 0, limit: 100 } },
|
||||
);
|
||||
expect(result).toEqual(mockFindings);
|
||||
});
|
||||
|
||||
@@ -0,0 +1,22 @@
|
||||
import { beforeEach, describe, expect, it, vi } from 'vitest';
|
||||
import { apiClient } from '../serverClient';
|
||||
import { fetchAllPages, fetchPage } from '../pagination';
|
||||
vi.mock('../serverClient', () => ({ apiClient: { get: vi.fn() } }));
|
||||
beforeEach(() => vi.resetAllMocks());
|
||||
describe('paged API', () => {
|
||||
it('uses server totals and forwards filters', async () => {
|
||||
vi.mocked(apiClient.get).mockResolvedValue({ data: [1], headers: {'x-total-count': '205'} });
|
||||
expect(await fetchPage('/tasks/', {skip: 20, limit: 20, status: 'running'})).toEqual({items: [1], total: 205});
|
||||
expect(apiClient.get).toHaveBeenCalledWith('/tasks/', {params: {skip: 20, limit: 20, status: 'running'}});
|
||||
});
|
||||
it('exports every page without truncating at the default limit', async () => {
|
||||
const first = Array.from({length: 100}, (_, i) => i);
|
||||
vi.mocked(apiClient.get).mockResolvedValueOnce({data: first}).mockResolvedValueOnce({data: [100, 101]});
|
||||
expect(await fetchAllPages('/findings', {severity: 'high'})).toEqual([...first, 100, 101]);
|
||||
expect(apiClient.get).toHaveBeenLastCalledWith('/findings', {params: {severity: 'high', skip: 100, limit: 100}});
|
||||
});
|
||||
it('does not publish a partial export after a later page fails', async () => {
|
||||
vi.mocked(apiClient.get).mockResolvedValueOnce({data: Array(100).fill(1)}).mockRejectedValueOnce(new Error('offline'));
|
||||
await expect(fetchAllPages('/issues')).rejects.toThrow('offline');
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,28 @@
|
||||
import { describe, expect, it, vi } from 'vitest';
|
||||
import { parseSSE } from '../sse';
|
||||
import { AgentStreamHandler } from '../agentStream';
|
||||
|
||||
describe('event stream recovery', () => {
|
||||
it('preserves events at every possible network split', () => {
|
||||
const input = 'id: 7\r\nevent: thinking_token\r\ndata: {"sequence":7,"metadata":{"token":"中文"}}\r\n\r\n';
|
||||
for (let i = 0; i <= input.length; i++) {
|
||||
const first = parseSSE(input.slice(0, i));
|
||||
const second = parseSSE(first.remaining + input.slice(i));
|
||||
expect([...first.parsed, ...second.parsed]).toEqual([{type: 'thinking_token', sequence: 7, metadata: {token: '中文'}}]);
|
||||
expect(second.remaining).toBe('');
|
||||
}
|
||||
});
|
||||
it('ignores malformed complete records and keeps subsequent records', () => {
|
||||
expect(parseSSE('data: broken\n\nevent: heartbeat\ndata: {}\n\n').parsed).toEqual([{type: 'heartbeat'}]);
|
||||
});
|
||||
it('advances reconnect cursor and suppresses duplicate replay, preserving task_end', () => {
|
||||
const onEvent = vi.fn();
|
||||
const handler = new AgentStreamHandler('task', {onEvent});
|
||||
const internal = handler as unknown as {handleEvent: (event: unknown) => void; options: {afterSequence: number}};
|
||||
internal.handleEvent({type: 'info', sequence: 8});
|
||||
internal.handleEvent({type: 'info', sequence: 8});
|
||||
internal.handleEvent({type: 'task_end', sequence: 8, status: 'completed'});
|
||||
expect(onEvent).toHaveBeenCalledTimes(2);
|
||||
expect(internal.options.afterSequence).toBe(8);
|
||||
});
|
||||
});
|
||||
@@ -1,3 +1,4 @@
|
||||
import { parseSSE } from "./sse";
|
||||
/**
|
||||
* Agent 流式事件处理
|
||||
*
|
||||
@@ -207,7 +208,7 @@ export class AgentStreamHandler {
|
||||
const { done, value } = await this.reader.read();
|
||||
|
||||
if (done) {
|
||||
console.log('[AgentStream] Reader done, stream ended');
|
||||
if (!this.isDisconnecting) throw new Error('事件连接已中断');
|
||||
break;
|
||||
}
|
||||
|
||||
@@ -275,51 +276,17 @@ export class AgentStreamHandler {
|
||||
* 解析 SSE 格式
|
||||
*/
|
||||
private parseSSE(buffer: string): { parsed: StreamEventData[]; remaining: string } {
|
||||
const parsed: StreamEventData[] = [];
|
||||
const lines = buffer.split('\n');
|
||||
let remaining = '';
|
||||
let currentEvent: Partial<StreamEventData> = {};
|
||||
|
||||
for (let i = 0; i < lines.length; i++) {
|
||||
const line = lines[i];
|
||||
|
||||
// 空行表示事件结束
|
||||
if (line === '') {
|
||||
if (currentEvent.type) {
|
||||
parsed.push(currentEvent as StreamEventData);
|
||||
currentEvent = {};
|
||||
}
|
||||
continue;
|
||||
}
|
||||
|
||||
// 检查是否是最后一行(可能不完整)
|
||||
if (i === lines.length - 1 && !buffer.endsWith('\n')) {
|
||||
remaining = line;
|
||||
break;
|
||||
}
|
||||
|
||||
// 解析 event: 行
|
||||
if (line.startsWith('event:')) {
|
||||
currentEvent.type = line.slice(6).trim() as StreamEventType;
|
||||
}
|
||||
// 解析 data: 行
|
||||
else if (line.startsWith('data:')) {
|
||||
try {
|
||||
const data = JSON.parse(line.slice(5).trim());
|
||||
currentEvent = { ...currentEvent, ...data };
|
||||
} catch {
|
||||
// 忽略解析错误
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return { parsed, remaining };
|
||||
return parseSSE(buffer);
|
||||
}
|
||||
|
||||
/**
|
||||
* 处理事件
|
||||
*/
|
||||
private handleEvent(event: StreamEventData): void {
|
||||
if (typeof event.sequence === 'number' && event.type !== 'task_end') {
|
||||
if (event.sequence <= (this.options.afterSequence ?? 0)) return;
|
||||
this.options.afterSequence = event.sequence;
|
||||
}
|
||||
// Extract agent_name from metadata if present
|
||||
if (event.metadata?.agent_name && !event.agent_name) {
|
||||
event.agent_name = event.metadata.agent_name as string;
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
import { fetchAllPages } from "./pagination";
|
||||
/**
|
||||
* Agent Tasks API
|
||||
* Agent 审计任务相关的 API 调用
|
||||
@@ -214,8 +215,10 @@ export async function getAgentFindings(
|
||||
is_verified?: boolean;
|
||||
}
|
||||
): Promise<AgentFinding[]> {
|
||||
const response = await apiClient.get(`/agent-tasks/${taskId}/findings`, { params });
|
||||
return response.data;
|
||||
const { is_verified, ...filters } = params || {};
|
||||
return fetchAllPages<AgentFinding>(`/agent-tasks/${taskId}/findings`, {
|
||||
...filters, ...(is_verified === undefined ? {} : {is_verified}),
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
import { fetchPage, fetchAllPages, type Page, type PageOptions } from "./pagination";
|
||||
import { apiClient } from "./serverClient";
|
||||
import type {
|
||||
Profile,
|
||||
@@ -158,8 +159,15 @@ export const api = {
|
||||
|
||||
async getAuditTasks(projectId?: string): Promise<AuditTask[]> {
|
||||
const params = projectId ? { project_id: projectId } : {};
|
||||
const res = await apiClient.get('/tasks/', { params });
|
||||
return res.data;
|
||||
return fetchAllPages<AuditTask>('/tasks/', params);
|
||||
},
|
||||
|
||||
async getAuditTaskPage(options: PageOptions = {}, projectId?: string): Promise<Page<AuditTask>> {
|
||||
return fetchPage<AuditTask>('/tasks/', { ...options, ...(projectId ? {project_id: projectId} : {}) });
|
||||
},
|
||||
|
||||
async getAuditIssuePage(taskId: string, options: PageOptions = {}): Promise<Page<AuditIssue>> {
|
||||
return fetchPage<AuditIssue>(`/tasks/${taskId}/issues`, { ...options });
|
||||
},
|
||||
|
||||
async getAuditTaskById(id: string): Promise<AuditTask | null> {
|
||||
@@ -198,8 +206,7 @@ export const api = {
|
||||
// ==================== AuditIssue 相关方法 ====================
|
||||
|
||||
async getAuditIssues(taskId: string): Promise<AuditIssue[]> {
|
||||
const res = await apiClient.get(`/tasks/${taskId}/issues`);
|
||||
return res.data;
|
||||
return fetchAllPages<AuditIssue>(`/tasks/${taskId}/issues`);
|
||||
},
|
||||
|
||||
async createAuditIssue(_issue: Omit<AuditIssue, 'id' | 'created_at' | 'task' | 'resolver'>): Promise<AuditIssue> {
|
||||
@@ -269,6 +276,7 @@ export const api = {
|
||||
total_issues: number;
|
||||
resolved_issues: number;
|
||||
avg_quality_score: number;
|
||||
issue_types?: Record<string, number>;
|
||||
}> {
|
||||
try {
|
||||
const res = await apiClient.get('/projects/stats');
|
||||
|
||||
@@ -0,0 +1,38 @@
|
||||
import { apiClient } from "./serverClient";
|
||||
|
||||
export interface Page<T> {
|
||||
items: T[];
|
||||
total: number;
|
||||
}
|
||||
export interface PageOptions {
|
||||
skip?: number;
|
||||
limit?: number;
|
||||
status?: string;
|
||||
search?: string;
|
||||
severity?: string;
|
||||
}
|
||||
|
||||
export async function fetchPage<T>(
|
||||
url: string,
|
||||
params: Record<string, unknown> = {},
|
||||
): Promise<Page<T>> {
|
||||
const response = await apiClient.get<T[]>(url, { params });
|
||||
return {
|
||||
items: response.data,
|
||||
total: Number(response.headers?.["x-total-count"] ?? response.data.length),
|
||||
};
|
||||
}
|
||||
|
||||
/** Explicitly used by exports and legacy aggregate views that require every row. */
|
||||
export async function fetchAllPages<T>(
|
||||
url: string,
|
||||
params: Record<string, unknown> = {},
|
||||
): Promise<T[]> {
|
||||
const items: T[] = [];
|
||||
const limit = 100;
|
||||
for (let skip = 0; ; skip += limit) {
|
||||
const page = await fetchPage<T>(url, { ...params, skip, limit });
|
||||
items.push(...page.items);
|
||||
if (page.items.length < limit) return items;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,35 @@
|
||||
import type { StreamEventData } from "./agentStream";
|
||||
|
||||
/** Preserve the entire unfinished event across arbitrary network chunk boundaries. */
|
||||
export function parseSSE(buffer: string): {
|
||||
parsed: StreamEventData[];
|
||||
remaining: string;
|
||||
} {
|
||||
const parsed: StreamEventData[] = [];
|
||||
const delimiter = /\r?\n\r?\n/g;
|
||||
let start = 0;
|
||||
for (
|
||||
let match = delimiter.exec(buffer);
|
||||
match;
|
||||
match = delimiter.exec(buffer)
|
||||
) {
|
||||
const block = buffer.slice(start, match.index);
|
||||
start = match.index + match[0].length;
|
||||
let type = "";
|
||||
const data: string[] = [];
|
||||
for (const line of block.split(/\r?\n/)) {
|
||||
if (line.startsWith("event:")) type = line.slice(6).trim();
|
||||
if (line.startsWith("data:")) data.push(line.slice(5).trimStart());
|
||||
}
|
||||
if (!data.length) continue;
|
||||
try {
|
||||
const payload = JSON.parse(data.join("\n"));
|
||||
if (typeof payload !== "object" || payload === null) continue;
|
||||
const event = { type, ...payload } as StreamEventData;
|
||||
if (event.type) parsed.push(event);
|
||||
} catch {
|
||||
/* Ignore malformed complete events; preserve subsequent events. */
|
||||
}
|
||||
}
|
||||
return { parsed, remaining: buffer.slice(start) };
|
||||
}
|
||||
@@ -14,6 +14,7 @@ export enum LogLevel {
|
||||
export enum LogCategory {
|
||||
USER_ACTION = 'USER_ACTION', // 用户操作
|
||||
API_CALL = 'API_CALL', // API调用
|
||||
PERFORMANCE = 'PERFORMANCE', // 请求性能
|
||||
SYSTEM = 'SYSTEM', // 系统事件
|
||||
CONSOLE_ERROR = 'CONSOLE_ERROR', // 控制台错误
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user