feat(console): add Artemis Admin & Trace Console backend and management plane
Introduces the comprehensive developer/administrator control plane and trace
inspection console backend for the ARTEMIS mobile automation platform.
Key Features & Backend Architecture:
- Core Server & Modular Routing:
* FastAPI modular architecture with dedicated routers for sessions, steps, tasks, media, stream, and replay.
* Multi-layer repository pattern with SQLite persistence for sessions, traces, and step data.
- Task Control & Queue Management:
* FIFO asynchronous task queue service with support for single/batch goal execution and profile selection (Flash / Pro).
* Session lifecycle controls (run, pause, resume, stop, clear queue).
* Startup orphan cleanup and automatic older replay archiving.
- Live Telemetry & IPC Communication:
* High-performance IPC server for streaming real-time execution logs, agent thoughts, and tool invocations.
* Server-Sent Events (SSE) /api/stream/{session_id} endpoint for live UI updates.
- Real-Time Device Streaming:
* Live ADB screen capture and MJPEG frame generator (/api/stream/device-live).
* Device connectivity and state inspection endpoint (/api/stream/device-state).
- Trace Inspection & Step Replay Engine:
* Hierarchical call tree inspection, pre/post step screenshots, and action overlay metadata.
* Isolated single-step replay manager (/api/replay) for interactive debugging and validation.
This commit is contained in:
Weilin Liu committed
1 parent
2fd70465e6
commit
c1803fcb86
29 files changed
+8889
No files matched your search
@@ -0,0 +1 @@
|
||||
# Copyright 2026 Google LLC
|
||||
@@ -0,0 +1,116 @@
|
||||
# Copyright 2026 Google LLC
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from pathlib import Path
|
||||
from fastapi import APIRouter, HTTPException
|
||||
from fastapi.responses import FileResponse, HTMLResponse
|
||||
|
||||
try:
|
||||
from admin_console.core.config import IMAGES_DIR, WORKSPACE_ROOT
|
||||
from admin_console.services.media_service import media_service
|
||||
except ImportError:
|
||||
from apps.admin_console.core.config import IMAGES_DIR, WORKSPACE_ROOT
|
||||
from apps.admin_console.services.media_service import media_service
|
||||
|
||||
router = APIRouter(tags=["media"])
|
||||
|
||||
|
||||
@router.get("/", response_class=HTMLResponse)
|
||||
@router.get("/admin", response_class=HTMLResponse)
|
||||
@router.get("/debug", response_class=HTMLResponse)
|
||||
async def get_index():
|
||||
admin_index = Path(__file__).resolve().parent.parent / "index.html"
|
||||
if admin_index.exists():
|
||||
return admin_index.read_text(encoding="utf-8")
|
||||
|
||||
admin_index_alt = WORKSPACE_ROOT / "apps" / "admin_console" / "index.html"
|
||||
if admin_index_alt.exists():
|
||||
return admin_index_alt.read_text(encoding="utf-8")
|
||||
|
||||
return "<h1>Admin Console index.html not found</h1>"
|
||||
|
||||
|
||||
@router.get("/images/{image_name}")
|
||||
@router.get("/api/images/{image_name}")
|
||||
async def get_image(image_name: str):
|
||||
if not image_name.endswith(".jpg"):
|
||||
image_name += ".jpg"
|
||||
image_path = IMAGES_DIR / image_name
|
||||
if not image_path.exists():
|
||||
raise HTTPException(status_code=404, detail="Image not found")
|
||||
|
||||
return FileResponse(image_path, media_type="image/jpeg")
|
||||
|
||||
|
||||
@router.get("/videos/{video_path:path}")
|
||||
async def get_video(video_path: str):
|
||||
path = Path(video_path)
|
||||
if not path.exists():
|
||||
if not path.is_absolute():
|
||||
candidate = WORKSPACE_ROOT / path
|
||||
if candidate.exists():
|
||||
path = candidate
|
||||
else:
|
||||
candidate_root = Path("/" + video_path.lstrip("/"))
|
||||
if candidate_root.exists():
|
||||
path = candidate_root
|
||||
else:
|
||||
path = candidate
|
||||
|
||||
if not path.exists():
|
||||
raise HTTPException(status_code=404, detail=f"Video not found at {path}")
|
||||
|
||||
suffix = path.suffix.lower()
|
||||
media_type = "video/mp4"
|
||||
if suffix == ".webm":
|
||||
media_type = "video/webm"
|
||||
elif suffix == ".mkv":
|
||||
media_type = "video/x-matroska"
|
||||
|
||||
return FileResponse(path, media_type=media_type)
|
||||
|
||||
|
||||
@router.get("/api/sessions/{session_id}/video")
|
||||
async def get_session_video(session_id: str):
|
||||
try:
|
||||
from admin_console.database.repositories.session_repository import session_repo
|
||||
except ImportError:
|
||||
from apps.admin_console.database.repositories.session_repository import session_repo
|
||||
|
||||
video_rec_map = session_repo.get_video_recordings_map()
|
||||
video_idx = media_service.build_video_index()
|
||||
row = session_repo.get_session_by_id(session_id)
|
||||
row_dict = dict(row) if row else {"session_id": session_id}
|
||||
v_url = media_service.resolve_video_url(row_dict, video_rec_map, video_idx)
|
||||
return {
|
||||
"session_id": session_id,
|
||||
"has_video": bool(v_url),
|
||||
"video_url": v_url,
|
||||
}
|
||||
|
||||
|
||||
@router.get("/local_file")
|
||||
async def get_local_file(path: str):
|
||||
p, media_type = media_service.get_safe_local_file(path)
|
||||
return FileResponse(p, media_type=media_type)
|
||||
|
||||
|
||||
@router.get("/api/sessions/{session_id}/plan")
|
||||
async def get_task_plan(session_id: str):
|
||||
return {"plan": media_service.get_task_plan_content(session_id)}
|
||||
|
||||
|
||||
@router.get("/api/sessions/{session_id}/notes")
|
||||
async def get_all_notes(session_id: str):
|
||||
return {"notes": media_service.get_session_notes_content(session_id)}
|
||||
@@ -0,0 +1,91 @@
|
||||
# Copyright 2026 Google LLC
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import traceback
|
||||
from fastapi import APIRouter, HTTPException
|
||||
|
||||
try:
|
||||
from admin_console.core.config import WORKSPACE_ROOT
|
||||
from admin_console.schemas.task_schema import ReplayRequest
|
||||
from admin_console.replay_manager import ReplayManager
|
||||
except ImportError:
|
||||
from apps.admin_console.core.config import WORKSPACE_ROOT
|
||||
from apps.admin_console.schemas.task_schema import ReplayRequest
|
||||
from apps.admin_console.replay_manager import ReplayManager
|
||||
|
||||
replay_manager = ReplayManager(WORKSPACE_ROOT)
|
||||
|
||||
|
||||
router = APIRouter(tags=["replay"])
|
||||
|
||||
|
||||
@router.get("/api/devices")
|
||||
async def list_devices():
|
||||
"""Dynamically queries the ADB server for connected Android devices."""
|
||||
return replay_manager.list_devices()
|
||||
|
||||
|
||||
@router.get("/api/replay/tools")
|
||||
async def get_replay_tools():
|
||||
"""Returns the list of registered tools available for replay."""
|
||||
return replay_manager.get_replay_tools()
|
||||
|
||||
|
||||
@router.get("/api/replay/config")
|
||||
async def get_replay_config(tool_name: str = "ask_explorer"):
|
||||
"""Returns the parameter list and types of the selected tool logic."""
|
||||
try:
|
||||
return replay_manager.get_replay_config(tool_name)
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.get("/api/sessions/{session_id}/replay_steps")
|
||||
async def get_replay_steps(session_id: str):
|
||||
"""Loads and formats metadata for all chunked steps in the session."""
|
||||
try:
|
||||
return replay_manager.get_replay_steps(session_id)
|
||||
except FileNotFoundError as e:
|
||||
raise HTTPException(status_code=404, detail=str(e))
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.post("/api/sessions/{session_id}/steps/{step_number}/replay")
|
||||
async def trigger_step_replay_endpoint(session_id: str, step_number: int, req: ReplayRequest):
|
||||
"""Triggers sandbox execution of step replay and returns the resulting traces trees."""
|
||||
try:
|
||||
return await replay_manager.run_step_replay(
|
||||
session_id=session_id,
|
||||
step_number=step_number,
|
||||
device_id=req.device_id,
|
||||
user_submits=req.user_submits,
|
||||
tool_name=req.tool_name,
|
||||
replay_id=req.replay_id,
|
||||
)
|
||||
except Exception as e:
|
||||
traceback.print_exc()
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.get("/api/sessions/{session_id}/steps/{step_number}/replay_traces")
|
||||
async def get_step_replay_traces_endpoint(
|
||||
session_id: str, step_number: int, tool_name: str = "ask_explorer"
|
||||
):
|
||||
"""Retrieves previously generated step replay traces if they exist."""
|
||||
try:
|
||||
return replay_manager.get_step_replay_traces(session_id, step_number, tool_name)
|
||||
except Exception as e:
|
||||
traceback.print_exc()
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
@@ -0,0 +1,127 @@
|
||||
# Copyright 2026 Google LLC
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import time
|
||||
from uuid import UUID
|
||||
from fastapi import APIRouter, HTTPException
|
||||
|
||||
try:
|
||||
from admin_console.core.config import DB_PATH, TRACES_PATH
|
||||
from admin_console.core.state import state
|
||||
from admin_console.database.repositories.session_repository import session_repo
|
||||
from admin_console.database.repositories.trace_repository import trace_repo
|
||||
from admin_console.services.media_service import media_service
|
||||
from admin_console.services.model_service import model_service
|
||||
except ImportError:
|
||||
from apps.admin_console.core.config import DB_PATH, TRACES_PATH
|
||||
from apps.admin_console.core.state import state
|
||||
from apps.admin_console.database.repositories.session_repository import session_repo
|
||||
from apps.admin_console.database.repositories.trace_repository import trace_repo
|
||||
from apps.admin_console.services.media_service import media_service
|
||||
from apps.admin_console.services.model_service import model_service
|
||||
|
||||
|
||||
router = APIRouter(tags=["sessions"])
|
||||
|
||||
|
||||
@router.get("/api/sessions")
|
||||
async def list_sessions():
|
||||
try:
|
||||
rows = session_repo.get_all_sessions()
|
||||
video_rec_map = session_repo.get_video_recordings_map()
|
||||
video_idx = media_service.build_video_index()
|
||||
|
||||
orphaned_ids = []
|
||||
result = []
|
||||
|
||||
for row_dict in rows:
|
||||
s_id = str(row_dict.get("session_id"))
|
||||
if row_dict.get("status") == "running":
|
||||
is_active = s_id in state.active_connections or (
|
||||
state.is_running and s_id == str(state.active_session_id)
|
||||
)
|
||||
if not is_active:
|
||||
row_dict["status"] = "failed"
|
||||
row_dict["end_time"] = time.time()
|
||||
orphaned_ids.append(s_id)
|
||||
|
||||
row_dict["video_url"] = media_service.resolve_video_url(
|
||||
row_dict, video_rec_map, video_idx
|
||||
)
|
||||
|
||||
llm_traces = session_repo.get_llm_traces_for_profile(s_id)
|
||||
agent_names = session_repo.get_agent_trace_names(s_id)
|
||||
sess_profile = model_service.resolve_session_profile(
|
||||
row_dict, llm_traces, state.current_profile, agent_names=agent_names
|
||||
)
|
||||
|
||||
row_dict["model_info"] = model_service.get_active_model_info(sess_profile)
|
||||
result.append(row_dict)
|
||||
|
||||
if orphaned_ids:
|
||||
try:
|
||||
harvested = session_repo.harvest_orphaned_sessions(orphaned_ids)
|
||||
print(f"[list_sessions] Auto-harvested {harvested} orphaned running session(s).")
|
||||
except Exception as e:
|
||||
print(f"[list_sessions] Failed to update orphaned sessions: {e}")
|
||||
|
||||
return result
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.get("/api/sessions/{session_id}/tree")
|
||||
async def get_tree(session_id: str):
|
||||
try:
|
||||
return trace_repo.get_trace_tree(session_id)
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.get("/api/sessions/{session_id}/background_tasks")
|
||||
async def get_session_background_tasks(session_id: str):
|
||||
try:
|
||||
return session_repo.get_background_tasks(session_id)
|
||||
except Exception:
|
||||
return []
|
||||
|
||||
|
||||
@router.post("/api/cleanup")
|
||||
async def cleanup_history_endpoint():
|
||||
try:
|
||||
from artemis.data_engine.storage import StorageManager
|
||||
|
||||
storage = StorageManager(DB_PATH, TRACES_PATH)
|
||||
storage.clear_all_data()
|
||||
return {
|
||||
"status": "success",
|
||||
"message": "History cleaned up successfully",
|
||||
}
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.post("/api/sessions/{session_id}/delete")
|
||||
async def delete_session_endpoint(session_id: str):
|
||||
try:
|
||||
from artemis.data_engine.storage import StorageManager
|
||||
|
||||
storage = StorageManager(DB_PATH, TRACES_PATH)
|
||||
storage.delete_session(UUID(session_id))
|
||||
return {
|
||||
"status": "success",
|
||||
"message": f"Session {session_id} deleted successfully",
|
||||
}
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
@@ -0,0 +1,108 @@
|
||||
# Copyright 2026 Google LLC
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import json
|
||||
import traceback
|
||||
from fastapi import APIRouter, HTTPException
|
||||
from fastapi.responses import PlainTextResponse
|
||||
|
||||
try:
|
||||
from admin_console.core.config import TEST_OUTPUTS_DIR
|
||||
from admin_console.database.repositories.step_repository import step_repo
|
||||
from admin_console.database.repositories.trace_repository import trace_repo
|
||||
from admin_console.services.media_service import media_service
|
||||
except ImportError:
|
||||
from apps.admin_console.core.config import TEST_OUTPUTS_DIR
|
||||
from apps.admin_console.database.repositories.step_repository import step_repo
|
||||
from apps.admin_console.database.repositories.trace_repository import trace_repo
|
||||
from apps.admin_console.services.media_service import media_service
|
||||
|
||||
|
||||
router = APIRouter(tags=["steps"])
|
||||
|
||||
|
||||
@router.get("/api/sessions/{session_id}/steps")
|
||||
async def get_session_steps(session_id: str, client: str | None = None):
|
||||
try:
|
||||
return step_repo.get_session_steps(session_id, client=client)
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.get("/api/steps/{step_id}/traces")
|
||||
async def get_step_traces_endpoint(step_id: str):
|
||||
try:
|
||||
session_id = step_repo.get_step_session_id(step_id)
|
||||
if not session_id:
|
||||
raise HTTPException(status_code=404, detail="Step not found")
|
||||
return trace_repo.get_step_traces_tree(session_id, step_id)
|
||||
except HTTPException as e:
|
||||
raise e
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
def _resolve_trace_db_path(session_id: str | None, step_number: int | None):
|
||||
if session_id and step_number is not None:
|
||||
sandbox_db_path = (
|
||||
TEST_OUTPUTS_DIR
|
||||
/ f"{session_id}_step_{step_number:02d}"
|
||||
/ "temp_traces"
|
||||
/ "data_engine.db"
|
||||
)
|
||||
if sandbox_db_path.exists():
|
||||
return sandbox_db_path
|
||||
return None
|
||||
|
||||
|
||||
@router.get("/api/traces/{trace_id}")
|
||||
async def get_trace(trace_id: str, session_id: str = None, step_number: int = None):
|
||||
try:
|
||||
db_path = _resolve_trace_db_path(session_id, step_number)
|
||||
trace_dict = trace_repo.get_trace_by_id(trace_id, db_path=db_path)
|
||||
if not trace_dict:
|
||||
raise HTTPException(status_code=404, detail="Trace not found")
|
||||
|
||||
if trace_dict.get("payload"):
|
||||
try:
|
||||
payload_obj = json.loads(trace_dict["payload"])
|
||||
trace_dict["payload"] = media_service.unwrap_payload(payload_obj)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
return trace_dict
|
||||
except HTTPException as e:
|
||||
raise e
|
||||
except Exception as e:
|
||||
traceback.print_exc()
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.get("/api/traces/{trace_id}/download")
|
||||
async def download_trace(trace_id: str, session_id: str = None, step_number: int = None):
|
||||
try:
|
||||
db_path = _resolve_trace_db_path(session_id, step_number)
|
||||
trace_dict = trace_repo.get_trace_by_id(trace_id, db_path=db_path)
|
||||
if not trace_dict or not trace_dict.get("payload"):
|
||||
raise HTTPException(status_code=404, detail="Payload not found")
|
||||
|
||||
return PlainTextResponse(
|
||||
trace_dict["payload"],
|
||||
headers={"Content-Disposition": f'attachment; filename="trace_{trace_id}.json"'},
|
||||
)
|
||||
except HTTPException as e:
|
||||
raise e
|
||||
except Exception as e:
|
||||
traceback.print_exc()
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
@@ -0,0 +1,56 @@
|
||||
# Copyright 2026 Google LLC
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""Device Screen Live Streaming Router.
|
||||
|
||||
Exposes real-time screen streaming endpoints for the Web UI.
|
||||
"""
|
||||
|
||||
from fastapi import APIRouter
|
||||
from fastapi.responses import JSONResponse, StreamingResponse
|
||||
|
||||
try:
|
||||
from admin_console.services.device_stream_service import device_stream_service
|
||||
except ImportError:
|
||||
from apps.admin_console.services.device_stream_service import device_stream_service
|
||||
|
||||
router = APIRouter(tags=["stream"])
|
||||
|
||||
|
||||
@router.get("/api/stream/device-live")
|
||||
async def stream_device_live():
|
||||
"""Streams live device screen frames as multipart MJPEG."""
|
||||
return StreamingResponse(
|
||||
device_stream_service.mjpeg_frame_generator(),
|
||||
media_type="multipart/x-mixed-replace; boundary=frame",
|
||||
headers={
|
||||
"Cache-Control": "no-cache, no-store, must-revalidate, max-age=0",
|
||||
"Pragma": "no-cache",
|
||||
"Expires": "0",
|
||||
"Connection": "keep-alive",
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
@router.get("/api/stream/device-state")
|
||||
async def get_device_stream_state():
|
||||
"""Returns whether an ADB device is connected and live streaming is available."""
|
||||
serial = await device_stream_service.get_device_serial()
|
||||
return JSONResponse(
|
||||
{
|
||||
"connected": serial is not None,
|
||||
"serial": serial,
|
||||
"live_stream_url": "/api/stream/device-live" if serial else None,
|
||||
}
|
||||
)
|
||||
@@ -0,0 +1,189 @@
|
||||
# Copyright 2026 Google LLC
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
from fastapi import APIRouter, HTTPException
|
||||
from fastapi.responses import StreamingResponse
|
||||
|
||||
try:
|
||||
from admin_console.core.state import state
|
||||
from admin_console.database.repositories.session_repository import session_repo
|
||||
from admin_console.schemas.task_schema import RunRequest
|
||||
from admin_console.services.ipc_service import ipc_service
|
||||
from admin_console.services.model_service import model_service
|
||||
from admin_console.services.task_queue_service import task_queue_service
|
||||
except ImportError:
|
||||
from apps.admin_console.core.state import state
|
||||
from apps.admin_console.database.repositories.session_repository import session_repo
|
||||
from apps.admin_console.schemas.task_schema import RunRequest
|
||||
from apps.admin_console.services.ipc_service import ipc_service
|
||||
from apps.admin_console.services.model_service import model_service
|
||||
from apps.admin_console.services.task_queue_service import task_queue_service
|
||||
|
||||
|
||||
router = APIRouter(tags=["tasks"])
|
||||
|
||||
|
||||
@router.post("/api/run")
|
||||
async def run_task(request: RunRequest):
|
||||
incoming_goals = []
|
||||
if request.goals:
|
||||
incoming_goals = request.goals
|
||||
elif request.goal:
|
||||
incoming_goals = [request.goal]
|
||||
|
||||
if not incoming_goals:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail="Either 'goal' or 'goals' list must be provided.",
|
||||
)
|
||||
|
||||
return await task_queue_service.enqueue_tasks(
|
||||
incoming_goals, profile=request.profile or "flash"
|
||||
)
|
||||
|
||||
|
||||
@router.post("/api/stop")
|
||||
async def stop_task(all: bool = False):
|
||||
stopped = task_queue_service.stop_tasks(clear_all=all)
|
||||
if stopped:
|
||||
return {"status": "stopped"}
|
||||
return {"status": "no_running_task"}
|
||||
|
||||
|
||||
@router.post("/api/resume")
|
||||
async def resume_task():
|
||||
resumed = task_queue_service.resume_task()
|
||||
if resumed:
|
||||
return {"status": "resumed"}
|
||||
return {"status": "not_paused"}
|
||||
|
||||
|
||||
@router.get("/api/status")
|
||||
async def get_status():
|
||||
latest_session = session_repo.get_latest_session()
|
||||
latest_session_id = latest_session.get("session_id") if latest_session else None
|
||||
bg_tasks = session_repo.get_background_tasks(latest_session_id) if latest_session_id else []
|
||||
|
||||
is_running = state.is_running
|
||||
running_sid = state.active_session_id or (
|
||||
session_repo.get_running_session_id() if is_running else None
|
||||
)
|
||||
active_profile = state.current_profile
|
||||
if not active_profile and (running_sid or latest_session_id):
|
||||
check_sid = running_sid or latest_session_id
|
||||
sess_row = session_repo.get_session_by_id(check_sid)
|
||||
if sess_row:
|
||||
llm_traces = session_repo.get_llm_traces_for_profile(check_sid)
|
||||
agent_names = session_repo.get_agent_trace_names(check_sid)
|
||||
active_profile = model_service.resolve_session_profile(
|
||||
sess_row, llm_traces, agent_names=agent_names
|
||||
)
|
||||
|
||||
model_info = model_service.get_active_model_info(active_profile)
|
||||
queue_data = state.queue_tasks
|
||||
|
||||
if state.current_process:
|
||||
if state.current_process.returncode is None:
|
||||
is_paused = state.is_paused
|
||||
return {
|
||||
"status": "paused" if is_paused else "running",
|
||||
"goal": state.current_goal,
|
||||
"pid": state.current_process.pid,
|
||||
"session_id": running_sid,
|
||||
"background_tasks": bg_tasks,
|
||||
"queue": queue_data,
|
||||
"model_info": model_info,
|
||||
}
|
||||
else:
|
||||
completed_status = {
|
||||
"status": "completed",
|
||||
"goal": state.current_goal,
|
||||
"returncode": state.current_process.returncode,
|
||||
"session_id": latest_session_id,
|
||||
"background_tasks": bg_tasks,
|
||||
"queue": queue_data,
|
||||
"model_info": model_info,
|
||||
}
|
||||
state.current_process = None
|
||||
return completed_status
|
||||
|
||||
if latest_session_id and str(latest_session_id) in state.active_connections:
|
||||
conn_info = state.active_connections[str(latest_session_id)]
|
||||
is_paused = state.is_paused
|
||||
conn_profile = conn_info.get("profile") or active_profile
|
||||
conn_model_info = model_service.get_active_model_info(conn_profile)
|
||||
return {
|
||||
"status": "paused" if is_paused else "running",
|
||||
"goal": conn_info.get("goal"),
|
||||
"pid": conn_info.get("pid"),
|
||||
"session_id": latest_session_id,
|
||||
"background_tasks": bg_tasks,
|
||||
"queue": queue_data,
|
||||
"model_info": conn_model_info,
|
||||
}
|
||||
|
||||
return {
|
||||
"status": "idle",
|
||||
"session_id": latest_session_id,
|
||||
"background_tasks": bg_tasks,
|
||||
"queue": queue_data,
|
||||
"model_info": model_info,
|
||||
}
|
||||
|
||||
|
||||
@router.get("/api/stream/{session_id}")
|
||||
async def stream_events(session_id: str, client: str | None = None):
|
||||
async def event_generator():
|
||||
queue = asyncio.Queue()
|
||||
|
||||
def callback(event_type, data):
|
||||
try:
|
||||
# Filter events by session_id when subscribed to a specific session
|
||||
if session_id and session_id not in ("all", "active"):
|
||||
evt_session_id = None
|
||||
if isinstance(data, dict):
|
||||
evt_session_id = data.get("session_id")
|
||||
|
||||
if evt_session_id and str(evt_session_id) != str(session_id):
|
||||
return
|
||||
|
||||
if (
|
||||
not evt_session_id
|
||||
and state.active_session_id
|
||||
and str(state.active_session_id) != str(session_id)
|
||||
):
|
||||
return
|
||||
|
||||
sanitized_data = ipc_service.sanitize_event_data(event_type, data)
|
||||
loop = asyncio.get_running_loop()
|
||||
loop.call_soon_threadsafe(queue.put_nowait, (event_type, sanitized_data))
|
||||
except RuntimeError:
|
||||
pass
|
||||
|
||||
state.add_subscriber(callback)
|
||||
yield f'event: info\ndata: {{"message": "Subscribed to session {session_id}"}}\n\n'
|
||||
|
||||
try:
|
||||
while True:
|
||||
try:
|
||||
event_type, data = await asyncio.wait_for(queue.get(), timeout=5.0)
|
||||
yield f"event: {event_type}\ndata: {json.dumps(data, default=str)}\n\n"
|
||||
except TimeoutError:
|
||||
yield "event: keep-alive\ndata: {}\n\n"
|
||||
except asyncio.CancelledError:
|
||||
state.remove_subscriber(callback)
|
||||
|
||||
return StreamingResponse(event_generator(), media_type="text/event-stream")
|
||||
Reference in new issue
Block a user