592 lines
21 KiB
Python
592 lines
21 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
import logging
|
|
from typing import Any
|
|
|
|
from app.core.contracts import (
|
|
CriticScore,
|
|
ExecutionDirective,
|
|
PermissionDecision,
|
|
PermissionRequest,
|
|
RuntimeEvent,
|
|
SecretRequest,
|
|
ToolCall,
|
|
UserTask,
|
|
)
|
|
from app.core.execution_scheduler import ExecutionScheduler
|
|
from app.events.event_bus import EventBus
|
|
from app.events.event_types import (
|
|
CRITIC_CALLED,
|
|
CRITIC_RESULT,
|
|
PERMISSION_REQUESTED,
|
|
PERMISSION_RESOLVED,
|
|
PLAN_FAILED,
|
|
PLAN_STARTED,
|
|
SECRET_REQUESTED,
|
|
STEP_STARTED,
|
|
STEPPED_COMPLETED,
|
|
TOOL_CALLED,
|
|
TOOL_COMPLETED,
|
|
)
|
|
from app.models.async_adapters import AsyncCriticAdapter, AsyncCoderAdapter
|
|
from app.memory.write_policy import MemoryWritePolicy
|
|
from app.memory.interface import MemoryInterface
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class ExecutionEngine:
|
|
def __init__(
|
|
self,
|
|
event_bus: EventBus,
|
|
tool_registry,
|
|
permission_service,
|
|
scheduler: ExecutionScheduler | None = None,
|
|
critic: AsyncCriticAdapter | None = None,
|
|
memory_policy: MemoryWritePolicy | None = None,
|
|
memory_interface: MemoryInterface | None = None,
|
|
prompts: dict[str, str] | None = None,
|
|
) -> None:
|
|
self._event_bus = event_bus
|
|
self._tool_registry = tool_registry
|
|
self._permission_service = permission_service
|
|
self._scheduler = scheduler or ExecutionScheduler()
|
|
self._critic = critic
|
|
self._coder: AsyncCoderAdapter | None = None
|
|
self._memory_policy = memory_policy
|
|
self._memory_interface = memory_interface
|
|
self._prompts = prompts or {}
|
|
|
|
def set_critic(self, critic: AsyncCriticAdapter) -> None:
|
|
self._critic = critic
|
|
|
|
def set_coder(self, coder: AsyncCoderAdapter) -> None:
|
|
self._coder = coder
|
|
|
|
def set_memory_policy(self, policy: MemoryWritePolicy) -> None:
|
|
self._memory_policy = policy
|
|
|
|
def execute(
|
|
self,
|
|
task: UserTask,
|
|
directive: ExecutionDirective,
|
|
permission_override: PermissionDecision | None = None,
|
|
secret_override: str | None = None,
|
|
password_override: str | None = None,
|
|
) -> dict[str, Any]:
|
|
scheduled = self._scheduler.next_directive(directive)
|
|
self._publish(task, STEP_STARTED, {"directive_type": scheduled.type})
|
|
|
|
if scheduled.type == "plan":
|
|
return self._execute_plan(
|
|
task=task,
|
|
directive=scheduled,
|
|
permission_override=permission_override,
|
|
secret_override=secret_override,
|
|
password_override=password_override,
|
|
)
|
|
|
|
if scheduled.type == "tool":
|
|
return self._execute_tool(
|
|
task=task,
|
|
directive=scheduled,
|
|
permission_override=permission_override,
|
|
secret_override=secret_override,
|
|
password_override=password_override,
|
|
)
|
|
|
|
if scheduled.type == "respond":
|
|
return {
|
|
"status": "completed",
|
|
"result": {
|
|
"message": f"Runtime accepted task: {task.input}",
|
|
"mode": scheduled.payload.get("mode", "direct_response"),
|
|
},
|
|
}
|
|
|
|
if scheduled.type == "coder":
|
|
return self._execute_coder(
|
|
task=task,
|
|
directive=scheduled,
|
|
)
|
|
|
|
if scheduled.type == "fail":
|
|
return {
|
|
"status": "failed",
|
|
"result": {"error": scheduled.reason or "Execution failed."},
|
|
}
|
|
|
|
return {
|
|
"status": "completed",
|
|
"result": {
|
|
"message": "Directive accepted.",
|
|
"directive_type": scheduled.type,
|
|
},
|
|
}
|
|
|
|
def _execute_plan(
|
|
self,
|
|
task: UserTask,
|
|
directive: ExecutionDirective,
|
|
permission_override: PermissionDecision | None = None,
|
|
secret_override: str | None = None,
|
|
password_override: str | None = None,
|
|
) -> dict[str, Any]:
|
|
# Unified format: {"type": "plan", "payload": {"steps": [...]}}
|
|
# Need to extract steps from nested payload
|
|
import json
|
|
|
|
payload = directive.payload
|
|
steps_data = []
|
|
|
|
# If payload has "steps" directly, use them
|
|
if "steps" in payload:
|
|
steps_data = payload.get("steps", [])
|
|
# If payload is a string (JSON), parse it
|
|
elif isinstance(payload, str) and payload.strip().startswith("{"):
|
|
try:
|
|
parsed = json.loads(payload)
|
|
steps_data = parsed.get("payload", {}).get("steps", [])
|
|
except:
|
|
steps_data = []
|
|
|
|
if steps_data:
|
|
plan_json = json.dumps({"type": "plan", "payload": {"steps": steps_data}})
|
|
else:
|
|
plan_json = json.dumps(payload)
|
|
|
|
plan_steps = self._scheduler.parse_plan_steps(plan_json, task.task_id)
|
|
|
|
if not plan_steps:
|
|
return {
|
|
"status": "failed",
|
|
"result": {"error": "Failed to parse plan steps from directive"},
|
|
}
|
|
|
|
if not self._scheduler.validate_no_cycles(plan_steps):
|
|
self._publish(task, PLAN_FAILED, {"error": "Cycle detected in plan"})
|
|
return {
|
|
"status": "failed",
|
|
"result": {"error": "Cycle detected in plan"},
|
|
}
|
|
|
|
graph = self._scheduler.build_task_graph(plan_steps)
|
|
self._publish(task, PLAN_STARTED, {"steps": len(plan_steps)})
|
|
|
|
completed_steps: set[str] = set()
|
|
step_results: list[dict[str, Any]] = []
|
|
|
|
ready_steps = self._get_ready_steps(graph, completed_steps)
|
|
|
|
while ready_steps:
|
|
step = ready_steps.pop(0)
|
|
|
|
# Handle respond kind directly without tool execution
|
|
if step.kind == "respond":
|
|
result = {
|
|
"status": "completed",
|
|
"result": {
|
|
"message": step.args.get("text", step.description),
|
|
},
|
|
}
|
|
else:
|
|
step_directive = ExecutionDirective(
|
|
type=step.kind,
|
|
payload={
|
|
"tool": step.tool,
|
|
"args": step.args,
|
|
},
|
|
requires_permission=step.requires_confirmation,
|
|
reason=step.description,
|
|
)
|
|
|
|
result = self._execute_tool(
|
|
task=task,
|
|
directive=step_directive,
|
|
permission_override=permission_override,
|
|
secret_override=secret_override,
|
|
password_override=password_override,
|
|
)
|
|
|
|
# If tool needs permission - return immediately, don't continue execution
|
|
if result.get("status") == "awaiting_permission":
|
|
return {
|
|
"status": "awaiting_permission",
|
|
"result": result.get("result", {}),
|
|
"step_results": step_results,
|
|
}
|
|
|
|
step_results.append({
|
|
"step_id": step.id,
|
|
"result": result,
|
|
})
|
|
|
|
completed_steps.add(step.id)
|
|
self._publish(task, STEPPED_COMPLETED, {
|
|
"step_id": step.id,
|
|
"status": result.get("status"),
|
|
})
|
|
|
|
# If tool needs permission or failed - return immediately, don't continue execution
|
|
if result.get("status") == "failed":
|
|
return {
|
|
"status": "failed",
|
|
"result": {
|
|
"error": f"Step {step.id} failed",
|
|
"failed_step": step.id,
|
|
"step_results": step_results,
|
|
},
|
|
}
|
|
|
|
requires_execution = directive.payload.get("requires_execution", True)
|
|
if requires_execution and self._critic:
|
|
critic_result = self._evaluate_with_critic(
|
|
task, step, result
|
|
)
|
|
if critic_result:
|
|
# Convert to dict for JSON serialization
|
|
result["critic_score"] = critic_result.model_dump(mode="json") if hasattr(critic_result, 'model_dump') else dict(critic_result)
|
|
self._save_critique_to_memory(task, step, critic_result)
|
|
|
|
ready_steps = self._get_ready_steps(graph, completed_steps)
|
|
|
|
return {
|
|
"status": "completed",
|
|
"result": {
|
|
"message": f"Plan executed: {len(completed_steps)} steps completed",
|
|
"step_results": step_results,
|
|
},
|
|
}
|
|
|
|
def _get_ready_steps(
|
|
self,
|
|
graph: dict[str, Any],
|
|
completed: set[str],
|
|
) -> list:
|
|
if not graph or not graph.get("nodes"):
|
|
return []
|
|
|
|
step_map: dict = graph.get("step_map", {})
|
|
ready = []
|
|
|
|
for node in graph["nodes"]:
|
|
node_id = node["id"]
|
|
if node_id in completed:
|
|
continue
|
|
|
|
deps = node.get("depends_on", [])
|
|
if all(dep in completed for dep in deps):
|
|
step = step_map.get(node_id)
|
|
if step:
|
|
ready.append(step)
|
|
|
|
return ready
|
|
|
|
def _evaluate_with_critic(
|
|
self,
|
|
task: UserTask,
|
|
step,
|
|
result: dict[str, Any],
|
|
) -> CriticScore | None:
|
|
if not self._critic:
|
|
return None
|
|
|
|
critic_prompt = self._build_critic_prompt(step, result)
|
|
|
|
self._publish(task, CRITIC_CALLED, {"step_id": step.id})
|
|
|
|
try:
|
|
critic_output = asyncio.run(self._critic.generate(critic_prompt))
|
|
score = self._parse_critic_score(critic_output)
|
|
|
|
self._publish(task, CRITIC_RESULT, {
|
|
"step_id": step.id,
|
|
"score": score.model_dump(mode="json") if score else None,
|
|
})
|
|
|
|
if score:
|
|
result["critic_score"] = {
|
|
"correctness": score.correctness,
|
|
"usefulness": score.usefulness,
|
|
"safety": score.safety,
|
|
"memory_store": score.memory_store,
|
|
"weight": score.weight,
|
|
"explanation": score.explanation,
|
|
}
|
|
|
|
return score
|
|
|
|
except Exception as e:
|
|
logger.warning(f"Critic evaluation failed: {e}")
|
|
self._publish(task, CRITIC_RESULT, {
|
|
"step_id": step.id,
|
|
"error": str(e),
|
|
})
|
|
return None
|
|
|
|
def _save_critique_to_memory(
|
|
self,
|
|
task: UserTask,
|
|
step,
|
|
score: CriticScore,
|
|
) -> None:
|
|
"""Save critic evaluation as critique entry in memory."""
|
|
if not self._memory_interface:
|
|
return
|
|
|
|
try:
|
|
tool_name = step.tool
|
|
tool_args = step.args or {}
|
|
args_str = ", ".join([f"{k}={v}" for k, v in tool_args.items()])
|
|
|
|
critique_text = f"Tool: {tool_name}({args_str}) | Task: {task.input[:100]} | Scores: correctness={score.correctness}, usefulness={score.usefulness}, safety={score.safety} | {score.explanation}"
|
|
|
|
metadata = {
|
|
"task_input": task.input,
|
|
"tool": tool_name,
|
|
"args": tool_args,
|
|
"step_id": step.id,
|
|
"scores": {
|
|
"correctness": score.correctness,
|
|
"usefulness": score.usefulness,
|
|
"safety": score.safety,
|
|
},
|
|
}
|
|
|
|
self._memory_interface.insert(
|
|
text=critique_text,
|
|
kind="critique",
|
|
source="critic",
|
|
task_id=task.task_id,
|
|
session_id=task.session_id,
|
|
weight=score.weight,
|
|
metadata=metadata,
|
|
)
|
|
logger.info(f"Saved critique to memory: {tool_name} task_id={task.task_id}")
|
|
|
|
except Exception as e:
|
|
logger.warning(f"Failed to save critique to memory: {e}")
|
|
|
|
def _build_critic_prompt(self, step, result: dict[str, Any]) -> str:
|
|
base_prompt = self._prompts.get("critic", "")
|
|
tool_result = result.get("result", {})
|
|
|
|
return f"""{base_prompt}
|
|
|
|
Step: {step.description}
|
|
Tool: {step.tool}
|
|
Args: {step.args}
|
|
|
|
Result:
|
|
{json.dumps(tool_result, indent=2)}
|
|
|
|
Evaluate and respond with JSON:
|
|
{{"correctness": 0.0-1.0, "usefulness": 0.0-1.0, "safety": 0.0-1.0, "memory_store": true|false, "weight": 0.0-1.0, "explanation": "..."}}"""
|
|
|
|
def _parse_critic_score(self, output: str) -> CriticScore | None:
|
|
try:
|
|
json_start = output.find("{")
|
|
json_end = output.rfind("}") + 1
|
|
if json_start < 0:
|
|
return None
|
|
|
|
json_str = output[json_start:json_end]
|
|
data = json.loads(json_str)
|
|
|
|
return CriticScore(
|
|
correctness=data.get("correctness", 0.5),
|
|
usefulness=data.get("usefulness", 0.5),
|
|
safety=data.get("safety", 1.0),
|
|
memory_store=data.get("memory_store", False),
|
|
weight=data.get("weight", 0.5),
|
|
explanation=data.get("explanation", ""),
|
|
)
|
|
|
|
except (json.JSONDecodeError, ValueError, TypeError) as e:
|
|
logger.warning(f"Critic score parsing failed: {e}")
|
|
return None
|
|
|
|
def _execute_coder(
|
|
self,
|
|
task: UserTask,
|
|
directive: ExecutionDirective,
|
|
) -> dict[str, Any]:
|
|
if not self._coder:
|
|
return {"status": "failed", "result": {"error": "Coder model not available"}}
|
|
|
|
coder_task = directive.payload.get("task", "")
|
|
if not coder_task:
|
|
return {"status": "failed", "result": {"error": "Missing task for coder"}}
|
|
|
|
try:
|
|
output = asyncio.run(self._coder.generate(coder_task))
|
|
|
|
return {
|
|
"status": "completed",
|
|
"result": {"code": output},
|
|
}
|
|
except Exception as e:
|
|
logger.warning(f"Coder execution failed: {e}")
|
|
return {"status": "failed", "result": {"error": str(e)}}
|
|
|
|
def _execute_tool(
|
|
self,
|
|
task: UserTask,
|
|
directive: ExecutionDirective,
|
|
permission_override: PermissionDecision | None = None,
|
|
secret_override: str | None = None,
|
|
password_override: str | None = None,
|
|
) -> dict[str, Any]:
|
|
tool_name = str(directive.payload.get("tool", "")).strip()
|
|
tool_args = dict(directive.payload.get("args", {}))
|
|
|
|
if password_override:
|
|
tool_args["password"] = password_override
|
|
|
|
if not tool_name:
|
|
return {"status": "failed", "result": {"error": "Missing tool name"}}
|
|
|
|
# Tool-first: validate tool exists in registry
|
|
available_tools = self._tool_registry.list_names()
|
|
if tool_name not in available_tools:
|
|
return {"status": "failed", "result": {"error": f"Unknown tool: {tool_name}. Available tools: {available_tools}"}}
|
|
|
|
permission_result = None
|
|
|
|
# Check permission for shell_exec and file_write
|
|
if tool_name == "shell_exec":
|
|
permission_result = self._permission_service.check_shell_command(
|
|
task_id=task.task_id,
|
|
session_id=task.session_id,
|
|
command=str(tool_args.get("command", "")),
|
|
)
|
|
elif tool_name == "file_write":
|
|
# Allow writing to runtime data directory without permission check
|
|
write_path = str(tool_args.get("path", ""))
|
|
if "allowed_commands.json" in write_path or "/data/runtime" in write_path:
|
|
# Internal system write - allow without permission
|
|
permission_result = {"decision": "allowed", "path": write_path}
|
|
else:
|
|
permission_result = self._permission_service.check_write_path(
|
|
task_id=task.task_id,
|
|
session_id=task.session_id,
|
|
path=write_path,
|
|
)
|
|
|
|
# Handle permission result
|
|
if permission_result:
|
|
decision = permission_result.get("decision", "unknown")
|
|
|
|
# Hard stop - deny execution
|
|
if decision == "hard_stop":
|
|
self._publish(task, PERMISSION_REQUESTED, permission_result)
|
|
return {
|
|
"status": "failed",
|
|
"result": {
|
|
"error": f"Command blocked: {permission_result.get('reason', 'Hard stop command')}",
|
|
"command": permission_result.get("command", ""),
|
|
},
|
|
}
|
|
|
|
# Cached - already allowed
|
|
if decision in ("allowed_always", "allowed") or permission_result.get("cached"):
|
|
self._publish(task, PERMISSION_RESOLVED, permission_result)
|
|
|
|
# Need user confirmation - return immediately, don't continue execution
|
|
elif decision == "prompt":
|
|
self._publish(task, PERMISSION_REQUESTED, permission_result)
|
|
return {
|
|
"status": "awaiting_permission",
|
|
"result": {
|
|
"error": "Permission required before execution.",
|
|
"permission_request": permission_result,
|
|
},
|
|
}
|
|
|
|
# Hard stop - return immediately
|
|
elif decision == "deny":
|
|
self._publish(task, PERMISSION_RESOLVED, permission_result)
|
|
return {
|
|
"status": "failed",
|
|
"result": {
|
|
"error": "Permission denied",
|
|
"command": permission_result.get("command", ""),
|
|
},
|
|
}
|
|
|
|
# Deny
|
|
elif decision == "deny":
|
|
self._publish(task, PERMISSION_RESOLVED, permission_result)
|
|
return {
|
|
"status": "failed",
|
|
"result": {
|
|
"error": "Permission denied",
|
|
"command": permission_result.get("command", ""),
|
|
},
|
|
}
|
|
|
|
if tool_name == "shell_exec":
|
|
command = str(tool_args.get("command", ""))
|
|
if command.startswith("sudo ") and secret_override is None:
|
|
secret_request = SecretRequest(
|
|
task_id=task.task_id,
|
|
session_id=task.session_id,
|
|
kind="sudo_password",
|
|
prompt="Sudo password required",
|
|
command=command,
|
|
)
|
|
self._publish(task, SECRET_REQUESTED, secret_request.model_dump(mode="json"))
|
|
return {
|
|
"status": "awaiting_input",
|
|
"result": {
|
|
"error": "Secret required",
|
|
"secret_request": secret_request.model_dump(mode="json"),
|
|
},
|
|
}
|
|
if command.startswith("sudo ") and secret_override is not None:
|
|
tool_args["command"] = f"sudo -S -p '' {command[len('sudo '):]}"
|
|
tool_args["stdin_secret"] = f"{secret_override}\n"
|
|
|
|
tool_call = ToolCall(
|
|
tool=tool_name,
|
|
args=tool_args,
|
|
task_id=task.task_id,
|
|
step_id="step-1",
|
|
)
|
|
self._publish(task, TOOL_CALLED, tool_call.model_dump(mode="json"))
|
|
tool_result = self._tool_registry.get(tool_name).execute(task=task, args=tool_args)
|
|
self._publish(task, TOOL_COMPLETED, tool_result.model_dump(mode="json"))
|
|
|
|
needs_sudo = tool_result.metadata.get("needs_sudo", False) if tool_result.metadata else False
|
|
|
|
if not tool_result.ok and needs_sudo:
|
|
return {
|
|
"status": "awaiting_password",
|
|
"result": {
|
|
"task_id": task.task_id,
|
|
"needs_sudo": True,
|
|
"command": tool_args.get("command", ""),
|
|
"error": tool_result.error or "Permission denied",
|
|
"tool_result": tool_result.model_dump(mode="json"),
|
|
},
|
|
}
|
|
|
|
return {
|
|
"status": "completed" if tool_result.ok else "failed",
|
|
"result": tool_result.model_dump(mode="json"),
|
|
}
|
|
|
|
def _publish(self, task: UserTask, event_type: str, payload: dict[str, Any]) -> None:
|
|
if not self._event_bus:
|
|
return
|
|
event = RuntimeEvent(
|
|
task_id=task.task_id,
|
|
session_id=task.session_id,
|
|
sequence=self._event_bus.next_sequence(task.task_id),
|
|
type=event_type,
|
|
payload=payload,
|
|
)
|
|
self._event_bus.publish(event)
|