Merge dev to main: Security improvements and comprehensive test infrastructure #39

Merged
jimmy merged 195 commits from dev into main 2026-04-08 01:42:27 +08:00
5 changed files with 333 additions and 56 deletions
Showing only changes of commit 5da2ae01a0 - Show all commits
+212 -1
View File
@@ -19,9 +19,19 @@ _pool: Optional[asyncpg.Pool] = None
async def get_pool() -> asyncpg.Pool: async def get_pool() -> asyncpg.Pool:
"""Get or create database connection pool""" """Get or create database connection pool with retry logic"""
global _pool global _pool
if _pool is None: if _pool is None:
import asyncio
import logging
logger = logging.getLogger(__name__)
max_retries = 3
retry_delay = 5
for attempt in range(max_retries):
try:
logger.info(f"Attempting to connect to OPC database (attempt {attempt + 1}/{max_retries})")
_pool = await asyncpg.create_pool( _pool = await asyncpg.create_pool(
host=DB_HOST, host=DB_HOST,
port=DB_PORT, port=DB_PORT,
@@ -31,6 +41,16 @@ async def get_pool() -> asyncpg.Pool:
min_size=2, min_size=2,
max_size=10, max_size=10,
) )
logger.info("Successfully connected to OPC database")
break
except Exception as e:
logger.error(f"Failed to connect to OPC database (attempt {attempt + 1}/{max_retries}): {e}")
if attempt < max_retries - 1:
logger.info(f"Retrying in {retry_delay} seconds...")
await asyncio.sleep(retry_delay)
else:
logger.error("All connection attempts failed")
raise Exception(f"Failed to connect to OPC database after {max_retries} attempts: {e}")
return _pool return _pool
@@ -606,3 +626,194 @@ async def get_agent_executions(
execution[json_field] = json.loads(execution[json_field]) execution[json_field] = json.loads(execution[json_field])
executions.append(execution) executions.append(execution)
return executions return executions
async def get_task_by_id(task_id: int) -> Optional[Dict[str, Any]]:
"""Get a single task by ID"""
pool = await get_pool()
async with pool.acquire() as conn:
row = await conn.fetchrow("SELECT * FROM tasks WHERE id = $1", task_id)
if not row:
return None
task = dict(row)
task["tags"] = json.loads(task["tags"]) if task["tags"] else []
return task
async def get_agent_execution(execution_id: int) -> Optional[Dict[str, Any]]:
"""Get a single agent execution by ID"""
pool = await get_pool()
async with pool.acquire() as conn:
row = await conn.fetchrow("SELECT * FROM agent_executions WHERE id = $1", execution_id)
if not row:
return None
execution = dict(row)
for json_field in ["input_context", "output_result", "actions_proposed", "actions_executed"]:
if execution[json_field]:
execution[json_field] = json.loads(execution[json_field])
return execution
async def approve_agent_execution(
execution_id: int,
approved: bool,
approved_by: str,
feedback: Optional[str] = None
) -> Dict[str, Any]:
"""Approve or reject an agent execution"""
pool = await get_pool()
async with pool.acquire() as conn:
# Update execution with approval status
row = await conn.fetchrow("""
UPDATE agent_executions
SET approved_by = $1,
approved_at = CURRENT_TIMESTAMP,
status = CASE WHEN $2 THEN 'approved' ELSE 'rejected' END
WHERE id = $3
RETURNING *
""", approved_by, approved, execution_id)
if not row:
raise ValueError(f"Execution {execution_id} not found")
execution = dict(row)
for json_field in ["input_context", "output_result", "actions_proposed", "actions_executed"]:
if execution[json_field]:
execution[json_field] = json.loads(execution[json_field])
# If feedback provided, add to output_result
if feedback:
if not execution["output_result"]:
execution["output_result"] = {}
execution["output_result"]["approval_feedback"] = feedback
await conn.execute("""
UPDATE agent_executions
SET output_result = $1
WHERE id = $2
""", json.dumps(execution["output_result"]), execution_id)
return execution
async def update_agent_config(
agent_id: str,
name: Optional[str] = None,
description: Optional[str] = None,
system_prompt: Optional[str] = None,
config: Optional[Dict[str, Any]] = None,
enabled: Optional[bool] = None
) -> Dict[str, Any]:
"""Update agent configuration"""
pool = await get_pool()
async with pool.acquire() as conn:
# Build update query dynamically
updates = []
params = []
param_count = 0
if name is not None:
param_count += 1
updates.append(f"name = ${param_count}")
params.append(name)
if description is not None:
param_count += 1
updates.append(f"description = ${param_count}")
params.append(description)
if system_prompt is not None:
param_count += 1
updates.append(f"system_prompt = ${param_count}")
params.append(system_prompt)
if config is not None:
param_count += 1
updates.append(f"config = ${param_count}")
params.append(json.dumps(config))
if enabled is not None:
param_count += 1
updates.append(f"enabled = ${param_count}")
params.append(enabled)
if not updates:
# No updates provided, just return current agent
return await get_agent(agent_id)
param_count += 1
params.append(agent_id)
query = f"UPDATE agents SET {', '.join(updates)} WHERE id = ${param_count} RETURNING *"
row = await conn.fetchrow(query, *params)
if not row:
raise ValueError(f"Agent {agent_id} not found")
agent = dict(row)
agent["capabilities"] = json.loads(agent["capabilities"]) if agent["capabilities"] else []
agent["config"] = json.loads(agent["config"]) if agent["config"] else {}
return agent
async def update_execution_status(
execution_id: int,
status: str,
output_result: Optional[Dict[str, Any]] = None,
actions_proposed: Optional[List[Dict[str, Any]]] = None,
actions_executed: Optional[List[Dict[str, Any]]] = None,
error_message: Optional[str] = None
) -> Dict[str, Any]:
"""Update agent execution status and results"""
pool = await get_pool()
async with pool.acquire() as conn:
# Check if started_at is already set
current = await conn.fetchrow("SELECT started_at FROM agent_executions WHERE id = $1", execution_id)
updates = ["status = $1"]
params = [status]
param_count = 1
# Only set started_at if transitioning to running and not already set
if status == "running" and (not current or not current["started_at"]):
updates.append("started_at = CURRENT_TIMESTAMP")
if status in ["completed", "failed", "rejected"]:
updates.append("completed_at = CURRENT_TIMESTAMP")
if output_result is not None:
param_count += 1
updates.append(f"output_result = ${param_count}")
params.append(json.dumps(output_result))
if actions_proposed is not None:
param_count += 1
updates.append(f"actions_proposed = ${param_count}")
params.append(json.dumps(actions_proposed))
if actions_executed is not None:
param_count += 1
updates.append(f"actions_executed = ${param_count}")
params.append(json.dumps(actions_executed))
if error_message is not None:
param_count += 1
updates.append(f"error_message = ${param_count}")
params.append(error_message)
param_count += 1
params.append(execution_id)
query = f"UPDATE agent_executions SET {', '.join(updates)} WHERE id = ${param_count} RETURNING *"
row = await conn.fetchrow(query, *params)
if not row:
raise ValueError(f"Execution {execution_id} not found")
execution = dict(row)
for json_field in ["input_context", "output_result", "actions_proposed", "actions_executed"]:
if execution[json_field]:
execution[json_field] = json.loads(execution[json_field])
return execution
+23 -11
View File
@@ -19,6 +19,7 @@ class AgentUpdate(BaseModel):
class ExecutionApproval(BaseModel): class ExecutionApproval(BaseModel):
approved: bool approved: bool
feedback: Optional[str] = None
@router.get("/agents") @router.get("/agents")
@@ -53,8 +54,17 @@ async def update_agent(
): ):
"""Update agent configuration (admin only)""" """Update agent configuration (admin only)"""
user: User = request.state.user user: User = request.state.user
# This would need an update_agent function in opc_db try:
raise HTTPException(status_code=501, detail="Not implemented yet") agent = await opc_db.update_agent_config(
agent_id=agent_id,
name=updates.name,
description=updates.description,
system_prompt=updates.system_prompt,
enabled=updates.enabled
)
return agent
except ValueError as e:
raise HTTPException(status_code=404, detail=str(e))
@router.post("/agents/{agent_id}/execute") @router.post("/agents/{agent_id}/execute")
@@ -99,8 +109,7 @@ async def get_execution(
): ):
"""Get execution details""" """Get execution details"""
user: User = request.state.user user: User = request.state.user
executions = await opc_db.get_agent_executions(limit=1000) execution = await opc_db.get_agent_execution(execution_id)
execution = next((e for e in executions if e["id"] == execution_id), None)
if not execution: if not execution:
raise HTTPException(status_code=404, detail="Execution not found") raise HTTPException(status_code=404, detail="Execution not found")
return execution return execution
@@ -114,10 +123,13 @@ async def approve_execution(
): ):
"""Approve or reject agent execution (for CxO level actions)""" """Approve or reject agent execution (for CxO level actions)"""
user: User = request.state.user user: User = request.state.user
# This would need an approve_execution function in opc_db try:
# For now, return placeholder execution = await opc_db.approve_agent_execution(
return { execution_id=execution_id,
"execution_id": execution_id, approved=approval.approved,
"approved": approval.approved, approved_by=user.username,
"approved_by": user.username feedback=approval.feedback
} )
return execution
except ValueError as e:
raise HTTPException(status_code=404, detail=str(e))
+20 -21
View File
@@ -2,7 +2,8 @@
OPC WebSocket Router - Real-time updates for tasks and agent executions OPC WebSocket Router - Real-time updates for tasks and agent executions
""" """
from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Depends from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Depends
from typing import Set from typing import Set, Any, Dict
from datetime import datetime
import json import json
import asyncio import asyncio
import logging import logging
@@ -16,6 +17,18 @@ router = APIRouter()
active_connections: Set[WebSocket] = set() active_connections: Set[WebSocket] = set()
def serialize_datetime(obj: Any) -> Any:
"""Recursively serialize datetime objects to ISO format strings"""
if isinstance(obj, datetime):
return obj.isoformat()
elif isinstance(obj, dict):
return {key: serialize_datetime(value) for key, value in obj.items()}
elif isinstance(obj, list):
return [serialize_datetime(item) for item in obj]
else:
return obj
class ConnectionManager: class ConnectionManager:
"""Manages WebSocket connections""" """Manages WebSocket connections"""
@@ -74,44 +87,30 @@ async def websocket_endpoint(websocket: WebSocket):
async def broadcast_task_update(task_id: int, action: str, task_data: dict): async def broadcast_task_update(task_id: int, action: str, task_data: dict):
"""Broadcast task update to all connected clients""" """Broadcast task update to all connected clients"""
# Serialize datetime objects to ISO format # Serialize all datetime objects recursively
serialized_data = task_data.copy() serialized_data = serialize_datetime(task_data)
for key in ["created_at", "updated_at", "completed_at", "due_date"]:
if serialized_data.get(key):
serialized_data[key] = serialized_data[key].isoformat() if hasattr(serialized_data[key], 'isoformat') else serialized_data[key]
timestamp = task_data.get("updated_at")
if timestamp and hasattr(timestamp, 'isoformat'):
timestamp = timestamp.isoformat()
message = { message = {
"type": "task_update", "type": "task_update",
"task_id": task_id, "task_id": task_id,
"action": action, # created, updated, moved, deleted "action": action, # created, updated, moved, deleted
"data": serialized_data, "data": serialized_data,
"timestamp": timestamp "timestamp": datetime.utcnow().isoformat()
} }
await manager.broadcast(message) await manager.broadcast(message)
async def broadcast_agent_execution(execution_id: int, status: str, execution_data: dict): async def broadcast_agent_execution(execution_id: int, status: str, execution_data: dict):
"""Broadcast agent execution update""" """Broadcast agent execution update"""
# Serialize datetime objects to ISO format # Serialize all datetime objects recursively
serialized_data = execution_data.copy() serialized_data = serialize_datetime(execution_data)
for key in ["started_at", "completed_at"]:
if serialized_data.get(key):
serialized_data[key] = serialized_data[key].isoformat() if hasattr(serialized_data[key], 'isoformat') else serialized_data[key]
timestamp = execution_data.get("started_at") or execution_data.get("completed_at")
if timestamp and hasattr(timestamp, 'isoformat'):
timestamp = timestamp.isoformat()
message = { message = {
"type": "agent_execution", "type": "agent_execution",
"execution_id": execution_id, "execution_id": execution_id,
"status": status, # pending, running, completed, failed, pending_approval "status": status, # pending, running, completed, failed, pending_approval
"data": serialized_data, "data": serialized_data,
"timestamp": timestamp "timestamp": datetime.utcnow().isoformat()
} }
await manager.broadcast(message) await manager.broadcast(message)
+54 -12
View File
@@ -65,12 +65,22 @@ class AgentExecutor:
async def _process_pending_executions(self): async def _process_pending_executions(self):
"""Process pending agent executions""" """Process pending agent executions"""
executions = await opc_db.get_agent_executions(status="pending", limit=10) # Process both pending and approved executions
print(f"[Executor] Found {len(executions)} pending executions") pending_executions = await opc_db.get_agent_executions(status="pending", limit=10)
approved_executions = await opc_db.get_agent_executions(status="approved", limit=10)
executions = pending_executions + approved_executions
print(f"[Executor] Found {len(pending_executions)} pending and {len(approved_executions)} approved executions")
for execution in executions: for execution in executions:
try: try:
print(f"[Executor] Processing execution {execution['id']} for agent {execution['agent_id']}") print(f"[Executor] Processing execution {execution['id']} for agent {execution['agent_id']} (status: {execution['status']})")
if execution['status'] == 'approved':
# Resume approved execution
await self._resume_approved_execution(execution)
else:
# Execute new pending execution
await self._execute_agent(execution) await self._execute_agent(execution)
except Exception as e: except Exception as e:
print(f"[Executor] Failed to execute agent {execution['agent_id']}: {e}") print(f"[Executor] Failed to execute agent {execution['agent_id']}: {e}")
@@ -91,8 +101,7 @@ class AgentExecutor:
await self._update_execution_status(execution["id"], "running") await self._update_execution_status(execution["id"], "running")
# Get task and context # Get task and context
tasks = await opc_db.get_tasks(limit=1000) task = await opc_db.get_task_by_id(task_id)
task = next((t for t in tasks if t["id"] == task_id), None)
if not task: if not task:
raise ValueError(f"Task {task_id} not found") raise ValueError(f"Task {task_id} not found")
@@ -118,6 +127,39 @@ class AgentExecutor:
# Send completion notification # Send completion notification
await self._send_completion_notification(agent, task, result) await self._send_completion_notification(agent, task, result)
async def _resume_approved_execution(self, execution: Dict[str, Any]):
"""Resume an approved execution and execute its actions"""
agent_id = execution["agent_id"]
task_id = execution["task_id"]
# Get agent
agent = self.agents.get(agent_id)
if not agent:
raise ValueError(f"Agent {agent_id} not found")
# Get task
task = await opc_db.get_task_by_id(task_id)
if not task:
raise ValueError(f"Task {task_id} not found")
# Mark as running
await self._update_execution_status(execution["id"], "running")
# Get proposed actions from execution
result = execution.get("output_result", {})
actions = execution.get("actions_proposed", [])
logger.info(f"Resuming approved execution for agent {agent.name} on task {task['title']}")
# Execute approved actions
await self._execute_actions(task, actions)
# Mark as completed
await self._mark_execution_completed(execution["id"], result)
# Send completion notification
await self._send_completion_notification(agent, task, result)
async def _build_context(self, task: Dict[str, Any]) -> Dict[str, Any]: async def _build_context(self, task: Dict[str, Any]) -> Dict[str, Any]:
"""Build context for agent execution""" """Build context for agent execution"""
context = { context = {
@@ -150,10 +192,16 @@ class AgentExecutor:
await self._action_send_notification(action) await self._action_send_notification(action)
elif action_type == "add_comment": elif action_type == "add_comment":
await self._action_add_comment(task, action) await self._action_add_comment(task, action)
elif action_type == "error":
# Handle error action - log and fail execution
error_msg = action.get("message", "Unknown error")
logger.error(f"Agent reported error: {error_msg}")
raise Exception(error_msg)
else: else:
logger.warning(f"Unknown action type: {action_type}") logger.warning(f"Unknown action type: {action_type}")
except Exception as e: except Exception as e:
logger.error(f"Failed to execute action {action_type}: {e}") logger.error(f"Failed to execute action {action_type}: {e}")
raise
async def _action_create_subtask(self, parent_task: Dict[str, Any], action: Dict[str, Any]): async def _action_create_subtask(self, parent_task: Dict[str, Any], action: Dict[str, Any]):
"""Create a subtask""" """Create a subtask"""
@@ -199,13 +247,7 @@ class AgentExecutor:
async def _update_execution_status(self, execution_id: int, status: str): async def _update_execution_status(self, execution_id: int, status: str):
"""Update execution status""" """Update execution status"""
pool = await opc_db.get_pool() await opc_db.update_execution_status(execution_id, status)
async with pool.acquire() as conn:
await conn.execute("""
UPDATE agent_executions
SET status = $1, started_at = CURRENT_TIMESTAMP
WHERE id = $2
""", status, execution_id)
async def _mark_execution_completed(self, execution_id: int, result: Dict[str, Any]): async def _mark_execution_completed(self, execution_id: int, result: Dict[str, Any]):
"""Mark execution as completed""" """Mark execution as completed"""
@@ -76,6 +76,9 @@ Only propose actions you can actually execute. Be specific and actionable."""
async def _call_llm(self, prompt: str) -> str: async def _call_llm(self, prompt: str) -> str:
"""Call LiteLLM proxy""" """Call LiteLLM proxy"""
import logging
logger = logging.getLogger(__name__)
model = self.config.get("model", "claude-sonnet-4-6") model = self.config.get("model", "claude-sonnet-4-6")
temperature = self.config.get("temperature", 0.7) temperature = self.config.get("temperature", 0.7)
@@ -85,7 +88,7 @@ Only propose actions you can actually execute. Be specific and actionable."""
if self.litellm_api_key and self.litellm_api_key.strip(): if self.litellm_api_key and self.litellm_api_key.strip():
headers["Authorization"] = f"Bearer {self.litellm_api_key}" headers["Authorization"] = f"Bearer {self.litellm_api_key}"
async with httpx.AsyncClient(timeout=60.0) as client: async with httpx.AsyncClient(timeout=30.0) as client:
response = await client.post( response = await client.post(
f"{self.litellm_url}/chat/completions", f"{self.litellm_url}/chat/completions",
headers=headers, headers=headers,
@@ -102,7 +105,17 @@ Only propose actions you can actually execute. Be specific and actionable."""
response.raise_for_status() response.raise_for_status()
data = response.json() data = response.json()
return data["choices"][0]["message"]["content"] return data["choices"][0]["message"]["content"]
except httpx.ConnectError as e:
logger.error(f"LiteLLM connection failed: {e}")
raise Exception(f"LiteLLM service unavailable: {str(e)}")
except httpx.TimeoutException as e:
logger.error(f"LiteLLM request timeout: {e}")
raise Exception(f"LiteLLM request timeout: {str(e)}")
except httpx.HTTPStatusError as e:
logger.error(f"LiteLLM HTTP error: {e.response.status_code} - {e.response.text}")
raise Exception(f"LiteLLM HTTP error {e.response.status_code}: {e.response.text}")
except Exception as e: except Exception as e:
logger.error(f"LiteLLM call failed: {e}")
raise Exception(f"LLM call failed: {str(e)}") raise Exception(f"LLM call failed: {str(e)}")
def _parse_response(self, response: str) -> Dict[str, Any]: def _parse_response(self, response: str) -> Dict[str, Any]: