|
1 | | -from typing import Optional |
2 | | -import asyncio |
3 | | -import json |
4 | | -from fastapi import APIRouter, HTTPException, Query, BackgroundTasks |
5 | | -from fastapi.responses import JSONResponse, StreamingResponse |
| 1 | +"""Legacy re-export of `pure_auto_codeql.api.analysis_routes`.""" |
6 | 2 |
|
7 | | -from api.models import ( |
8 | | - AnalysisRequest, |
9 | | - AnalysisTaskInfo, |
10 | | - AnalysisResult, |
11 | | - TaskListResponse, |
12 | | - TaskStatus, |
13 | | - ErrorResponse |
14 | | -) |
15 | | -from api.config import get_config |
16 | | -from api.task_manager import get_task_manager |
17 | | -from pure_auto_codeql.application import ( |
18 | | - AnalysisValidationError, |
19 | | - validate_analysis_case, |
20 | | -) |
| 3 | +from importlib import import_module |
| 4 | +import sys as _sys |
21 | 5 |
|
22 | | - |
23 | | -router = APIRouter(prefix="/analysis", tags=["analysis"]) |
24 | | - |
25 | | -# 启动任务 |
26 | | -@router.post("/start", response_model=AnalysisTaskInfo, status_code=202) |
27 | | -async def start_analysis( |
28 | | - request: AnalysisRequest, |
29 | | - background_tasks: BackgroundTasks |
30 | | -) -> AnalysisTaskInfo: |
31 | | - |
32 | | - try: |
33 | | - validate_analysis_case(request.case_id, projects_dir=get_config().projects_dir) |
34 | | - except AnalysisValidationError as e: |
35 | | - raise HTTPException( |
36 | | - status_code=e.status_code, |
37 | | - detail=str(e), |
38 | | - ) from e |
39 | | - |
40 | | - task_manager = get_task_manager() |
41 | | - task_id = task_manager.create_task(request.case_id) |
42 | | - |
43 | | - config = { |
44 | | - 'language': request.language, |
45 | | - 'max_rounds': request.max_rounds, |
46 | | - 'enable_cve_analysis': request.enable_cve_analysis, |
47 | | - 'enable_sink_analysis': request.enable_sink_analysis, |
48 | | - 'show_thinking': False, |
49 | | - 'refresh_intel': False |
50 | | - } |
51 | | - |
52 | | - background_tasks.add_task(task_manager.start_task, task_id, config) |
53 | | - |
54 | | - task_info = task_manager.get_task_status(task_id) |
55 | | - if not task_info: |
56 | | - raise HTTPException( |
57 | | - status_code=500, |
58 | | - detail="任务创建失败" |
59 | | - ) |
60 | | - |
61 | | - return task_info |
62 | | - |
63 | | - |
64 | | -@router.get("/{task_id}/status", response_model=AnalysisTaskInfo) |
65 | | -async def get_task_status(task_id: str) -> AnalysisTaskInfo: |
66 | | - task_manager = get_task_manager() |
67 | | - task_info = task_manager.get_task_status(task_id) |
68 | | - |
69 | | - if not task_info: |
70 | | - raise HTTPException( |
71 | | - status_code=404, |
72 | | - detail=f"任务 '{task_id}' 不存在" |
73 | | - ) |
74 | | - |
75 | | - return task_info |
76 | | - |
77 | | - |
78 | | -@router.get("/{task_id}/result", response_model=AnalysisResult) |
79 | | -async def get_task_result(task_id: str) -> AnalysisResult: |
80 | | - task_manager = get_task_manager() |
81 | | - |
82 | | - task_info = task_manager.get_task_status(task_id) |
83 | | - if not task_info: |
84 | | - raise HTTPException( |
85 | | - status_code=404, |
86 | | - detail=f"任务 '{task_id}' 不存在" |
87 | | - ) |
88 | | - |
89 | | - if task_info.status == TaskStatus.PENDING: |
90 | | - raise HTTPException( |
91 | | - status_code=409, |
92 | | - detail="任务尚未开始执行" |
93 | | - ) |
94 | | - |
95 | | - if task_info.status == TaskStatus.RUNNING: |
96 | | - raise HTTPException( |
97 | | - status_code=409, |
98 | | - detail="任务正在执行中,请稍后查询" |
99 | | - ) |
100 | | - |
101 | | - if task_info.status == TaskStatus.CANCELLED: |
102 | | - raise HTTPException( |
103 | | - status_code=410, |
104 | | - detail="任务已被取消" |
105 | | - ) |
106 | | - |
107 | | - if task_info.status == TaskStatus.FAILED: |
108 | | - raise HTTPException( |
109 | | - status_code=500, |
110 | | - detail=f"任务执行失败: {task_info.error}" |
111 | | - ) |
112 | | - |
113 | | - result = task_manager.get_task_result(task_id) |
114 | | - if not result: |
115 | | - raise HTTPException( |
116 | | - status_code=404, |
117 | | - detail="任务结果不存在" |
118 | | - ) |
119 | | - |
120 | | - return result |
121 | | - |
122 | | - |
123 | | -@router.delete("/{task_id}", status_code=200) |
124 | | -async def cancel_task(task_id: str) -> dict: |
125 | | - task_manager = get_task_manager() |
126 | | - |
127 | | - task_info = task_manager.get_task_status(task_id) |
128 | | - if not task_info: |
129 | | - raise HTTPException( |
130 | | - status_code=404, |
131 | | - detail=f"任务 '{task_id}' 不存在" |
132 | | - ) |
133 | | - |
134 | | - if task_info.status not in [TaskStatus.PENDING, TaskStatus.RUNNING]: |
135 | | - raise HTTPException( |
136 | | - status_code=400, |
137 | | - detail=f"无法取消状态为 '{task_info.status.value}' 的任务" |
138 | | - ) |
139 | | - |
140 | | - success = await task_manager.cancel_task(task_id) |
141 | | - if not success: |
142 | | - raise HTTPException( |
143 | | - status_code=500, |
144 | | - detail="任务取消失败" |
145 | | - ) |
146 | | - |
147 | | - return { |
148 | | - "message": "任务已成功取消", |
149 | | - "task_id": task_id |
150 | | - } |
151 | | - |
152 | | - |
153 | | -@router.get("/tasks", response_model=TaskListResponse) |
154 | | -async def list_tasks( |
155 | | - status: Optional[TaskStatus] = Query(None, description="按状态过滤"), |
156 | | - page: int = Query(1, ge=1, description="页码"), |
157 | | - page_size: int = Query(20, ge=1, le=100, description="每页大小") |
158 | | -) -> TaskListResponse: |
159 | | - task_manager = get_task_manager() |
160 | | - |
161 | | - offset = (page - 1) * page_size |
162 | | - |
163 | | - tasks, total = task_manager.list_tasks( |
164 | | - status_filter=status, |
165 | | - limit=page_size, |
166 | | - offset=offset |
167 | | - ) |
168 | | - |
169 | | - return TaskListResponse( |
170 | | - tasks=tasks, |
171 | | - total=total, |
172 | | - page=page, |
173 | | - page_size=page_size |
174 | | - ) |
175 | | - |
176 | | - |
177 | | -@router.post("/cleanup", status_code=200) |
178 | | -async def cleanup_old_tasks( |
179 | | - max_age_hours: int = Query(24, ge=1, le=168, description="最大保留时间(小时)") |
180 | | -) -> dict: |
181 | | - task_manager = get_task_manager() |
182 | | - task_manager.cleanup_old_tasks(max_age_hours) |
183 | | - |
184 | | - return { |
185 | | - "message": f"已清理超过 {max_age_hours} 小时的旧任务" |
186 | | - } |
187 | | - |
188 | | - |
189 | | -@router.get("/{task_id}/stream") |
190 | | -async def stream_task_output(task_id: str): |
191 | | - """ |
192 | | - 通过 Server-Sent Events (SSE) 流式输出任务的实时事件 |
193 | | - |
194 | | - Args: |
195 | | - task_id: 任务ID |
196 | | - |
197 | | - Returns: |
198 | | - StreamingResponse: SSE 格式的事件流 |
199 | | - |
200 | | - Raises: |
201 | | - HTTPException 404: 任务不存在或事件队列未创建 |
202 | | - HTTPException 410: 任务已结束且事件队列已清理 |
203 | | - """ |
204 | | - task_manager = get_task_manager() |
205 | | - |
206 | | - task_info = task_manager.get_task_status(task_id) |
207 | | - if not task_info: |
208 | | - raise HTTPException( |
209 | | - status_code=404, |
210 | | - detail=f"任务 '{task_id}' 不存在" |
211 | | - ) |
212 | | - |
213 | | - if task_id not in task_manager._event_queues: |
214 | | - if task_info.status in [TaskStatus.COMPLETED, TaskStatus.FAILED, TaskStatus.CANCELLED]: |
215 | | - raise HTTPException( |
216 | | - status_code=410, |
217 | | - detail=f"任务已结束,事件流不再可用。请使用 /api/analysis/{task_id}/result 获取最终结果" |
218 | | - ) |
219 | | - else: |
220 | | - raise HTTPException( |
221 | | - status_code=404, |
222 | | - detail="任务事件队列未创建,任务可能尚未启动" |
223 | | - ) |
224 | | - |
225 | | - # 事件生成器 |
226 | | - async def event_generator(): |
227 | | - """从队列读取事件并格式化为 SSE""" |
228 | | - queue = task_manager._event_queues.get(task_id) |
229 | | - if not queue: |
230 | | - return |
231 | | - |
232 | | - try: |
233 | | - while True: |
234 | | - try: |
235 | | - # 等待事件,超时时间 30 秒(发送心跳) |
236 | | - event = await asyncio.wait_for(queue.get(), timeout=30.0) |
237 | | - except asyncio.TimeoutError: |
238 | | - # 超时时发送心跳 |
239 | | - yield f": heartbeat\n\n" |
240 | | - continue |
241 | | - |
242 | | - # SSE 格式: event: <type>\ndata: <json>\n\n |
243 | | - event_type = event.get('type', 'message') |
244 | | - |
245 | | - event_data = { |
246 | | - 'type': str(event_type) if hasattr(event_type, 'value') else event_type, |
247 | | - 'timestamp': event.get('timestamp'), |
248 | | - 'step_name': event.get('step_name'), |
249 | | - 'message': event.get('message'), |
250 | | - 'data': event.get('data', {}) |
251 | | - } |
252 | | - |
253 | | - event_json = json.dumps(event_data, ensure_ascii=False) |
254 | | - |
255 | | - yield f"event: {event_data['type']}\n" |
256 | | - yield f"data: {event_json}\n\n" |
257 | | - |
258 | | - if event_data['type'] in ['completed', 'error']: |
259 | | - if event.get('data', {}).get('task_id') == task_id: |
260 | | - break |
261 | | - |
262 | | - except asyncio.CancelledError: |
263 | | - pass |
264 | | - except Exception as e: |
265 | | - # 发送错误事件 |
266 | | - error_event = { |
267 | | - 'type': 'error', |
268 | | - 'timestamp': None, |
269 | | - 'step_name': 'stream', |
270 | | - 'message': f'流式传输错误: {str(e)}', |
271 | | - 'data': {'error': str(e)} |
272 | | - } |
273 | | - yield f"event: error\n" |
274 | | - yield f"data: {json.dumps(error_event, ensure_ascii=False)}\n\n" |
275 | | - |
276 | | - return StreamingResponse( |
277 | | - event_generator(), |
278 | | - media_type="text/event-stream", |
279 | | - headers={ |
280 | | - "Cache-Control": "no-cache", |
281 | | - "Connection": "keep-alive", |
282 | | - "X-Accel-Buffering": "no", # 禁用 nginx 缓冲 |
283 | | - } |
284 | | - ) |
| 6 | +_sys.modules[__name__] = import_module("pure_auto_codeql.api.analysis_routes") |
0 commit comments