Add confirmation cards and scenario smoke validation
This commit is contained in:
@@ -78,6 +78,7 @@ class ConfirmationRepository(Protocol):
|
|||||||
def create(self, confirmation: ConfirmationRequest) -> ConfirmationRequest: ...
|
def create(self, confirmation: ConfirmationRequest) -> ConfirmationRequest: ...
|
||||||
def get(self, confirmation_id: str) -> ConfirmationRequest | None: ...
|
def get(self, confirmation_id: str) -> ConfirmationRequest | None: ...
|
||||||
def save(self, confirmation: ConfirmationRequest) -> ConfirmationRequest: ...
|
def save(self, confirmation: ConfirmationRequest) -> ConfirmationRequest: ...
|
||||||
|
def list_by_task(self, task_id: str) -> list[ConfirmationRequest]: ...
|
||||||
|
|
||||||
|
|
||||||
class WorkerRepository(Protocol):
|
class WorkerRepository(Protocol):
|
||||||
|
|||||||
@@ -26,6 +26,8 @@ from ai_orchestrator.application.services.workers import (
|
|||||||
)
|
)
|
||||||
from ai_orchestrator.config import AppSettings
|
from ai_orchestrator.config import AppSettings
|
||||||
from ai_orchestrator.delivery.http.presenters import (
|
from ai_orchestrator.delivery.http.presenters import (
|
||||||
|
build_artifact_cards_from_tool_invocations,
|
||||||
|
build_confirmation_card,
|
||||||
build_error_card,
|
build_error_card,
|
||||||
build_plan_card,
|
build_plan_card,
|
||||||
build_progress_card,
|
build_progress_card,
|
||||||
@@ -49,6 +51,8 @@ from ai_orchestrator.delivery.http.schemas import (
|
|||||||
WorkerRegisterRequest,
|
WorkerRegisterRequest,
|
||||||
WorkerResultMessage,
|
WorkerResultMessage,
|
||||||
)
|
)
|
||||||
|
from ai_orchestrator.domain.enums import RiskLevel
|
||||||
|
from ai_orchestrator.domain.models import ActionDescriptor
|
||||||
from ai_orchestrator.infrastructure.config_loader import load_project_config
|
from ai_orchestrator.infrastructure.config_loader import load_project_config
|
||||||
from ai_orchestrator.infrastructure.model_router import StaticMockModelProvider
|
from ai_orchestrator.infrastructure.model_router import StaticMockModelProvider
|
||||||
from ai_orchestrator.infrastructure.policy import StaticProjectPolicyEvaluator
|
from ai_orchestrator.infrastructure.policy import StaticProjectPolicyEvaluator
|
||||||
@@ -124,6 +128,31 @@ def create_app(settings: AppSettings | None = None) -> FastAPI:
|
|||||||
event_store=event_store,
|
event_store=event_store,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
def _maybe_create_confirmation(
|
||||||
|
*,
|
||||||
|
task_id: str,
|
||||||
|
graph,
|
||||||
|
action_payload: object,
|
||||||
|
):
|
||||||
|
if not isinstance(action_payload, dict) or not graph.nodes:
|
||||||
|
return None
|
||||||
|
try:
|
||||||
|
action = ActionDescriptor(
|
||||||
|
resource=str(action_payload["resource"]),
|
||||||
|
risk_level=RiskLevel(str(action_payload["risk_level"])),
|
||||||
|
action_name=str(action_payload["action_name"]),
|
||||||
|
preview_available=bool(action_payload.get("preview_available", False)),
|
||||||
|
metadata=dict(action_payload.get("metadata", {})),
|
||||||
|
)
|
||||||
|
except (KeyError, ValueError, TypeError):
|
||||||
|
return None
|
||||||
|
target_node = graph.nodes[-1]
|
||||||
|
return orchestrator.evaluate_action(
|
||||||
|
task_id=task_id,
|
||||||
|
node_id=target_node.node_id,
|
||||||
|
action=action,
|
||||||
|
)
|
||||||
|
|
||||||
@app.get("/health")
|
@app.get("/health")
|
||||||
def health() -> dict[str, str]:
|
def health() -> dict[str, str]:
|
||||||
return {"status": "ok"}
|
return {"status": "ok"}
|
||||||
@@ -141,15 +170,26 @@ def create_app(settings: AppSettings | None = None) -> FastAPI:
|
|||||||
)
|
)
|
||||||
graph = orchestrator.plan_task(task.task_id)
|
graph = orchestrator.plan_task(task.task_id)
|
||||||
execution_engine.execute_ready_nodes(task=task, graph=graph)
|
execution_engine.execute_ready_nodes(task=task, graph=graph)
|
||||||
|
confirmation = _maybe_create_confirmation(
|
||||||
|
task_id=task.task_id,
|
||||||
|
graph=graph,
|
||||||
|
action_payload=payload.preferences.get("policy_action"),
|
||||||
|
)
|
||||||
cards = [
|
cards = [
|
||||||
build_plan_card(task, graph),
|
build_plan_card(task, graph),
|
||||||
build_progress_card(task, graph),
|
build_progress_card(task, graph),
|
||||||
]
|
]
|
||||||
|
if confirmation is not None:
|
||||||
|
cards.append(build_confirmation_card(confirmation))
|
||||||
return ChatResponse(
|
return ChatResponse(
|
||||||
conversation_id=task.conversation_id or "conv_default",
|
conversation_id=task.conversation_id or "conv_default",
|
||||||
task_id=task.task_id,
|
task_id=task.task_id,
|
||||||
response_type="task_started",
|
response_type="confirmation_required" if confirmation is not None else "task_started",
|
||||||
message="Task accepted and planned.",
|
message=(
|
||||||
|
"Task accepted and waiting for confirmation."
|
||||||
|
if confirmation is not None
|
||||||
|
else "Task accepted and planned."
|
||||||
|
),
|
||||||
cards=cards,
|
cards=cards,
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -165,10 +205,17 @@ def create_app(settings: AppSettings | None = None) -> FastAPI:
|
|||||||
)
|
)
|
||||||
graph = orchestrator.plan_task(task.task_id)
|
graph = orchestrator.plan_task(task.task_id)
|
||||||
execution_engine.execute_ready_nodes(task=task, graph=graph)
|
execution_engine.execute_ready_nodes(task=task, graph=graph)
|
||||||
|
confirmation = _maybe_create_confirmation(
|
||||||
|
task_id=task.task_id,
|
||||||
|
graph=graph,
|
||||||
|
action_payload=payload.inputs.get("policy_action"),
|
||||||
|
)
|
||||||
cards = [
|
cards = [
|
||||||
build_plan_card(task, graph),
|
build_plan_card(task, graph),
|
||||||
build_progress_card(task, graph),
|
build_progress_card(task, graph),
|
||||||
]
|
]
|
||||||
|
if confirmation is not None:
|
||||||
|
cards.append(build_confirmation_card(confirmation))
|
||||||
return TaskStatusResponse(
|
return TaskStatusResponse(
|
||||||
task_id=task.task_id,
|
task_id=task.task_id,
|
||||||
status=task.status.value,
|
status=task.status.value,
|
||||||
@@ -184,12 +231,26 @@ def create_app(settings: AppSettings | None = None) -> FastAPI:
|
|||||||
if task is None or graph is None:
|
if task is None or graph is None:
|
||||||
raise HTTPException(status_code=404, detail="Task not found")
|
raise HTTPException(status_code=404, detail="Task not found")
|
||||||
completed = sum(1 for node in graph.nodes if node.status.value == "completed")
|
completed = sum(1 for node in graph.nodes if node.status.value == "completed")
|
||||||
|
confirmations = confirmation_repository.list_by_task(task_id)
|
||||||
|
confirmation_cards = [
|
||||||
|
build_confirmation_card(item)
|
||||||
|
for item in confirmations
|
||||||
|
if item.status.value == "pending"
|
||||||
|
]
|
||||||
|
tool_invocations = [
|
||||||
|
{
|
||||||
|
"tool_name": item.tool_name,
|
||||||
|
"response": item.response,
|
||||||
|
}
|
||||||
|
for item in invocation_store.list_tool_invocations(task_id)
|
||||||
|
]
|
||||||
|
artifact_cards = build_artifact_cards_from_tool_invocations(task_id, tool_invocations)
|
||||||
return TaskStatusResponse(
|
return TaskStatusResponse(
|
||||||
task_id=task.task_id,
|
task_id=task.task_id,
|
||||||
status=task.status.value,
|
status=task.status.value,
|
||||||
current_node=task.current_node_id,
|
current_node=task.current_node_id,
|
||||||
progress={"completed": completed, "total": len(graph.nodes)},
|
progress={"completed": completed, "total": len(graph.nodes)},
|
||||||
cards=[build_progress_card(task, graph)],
|
cards=[build_progress_card(task, graph), *confirmation_cards, *artifact_cards],
|
||||||
)
|
)
|
||||||
|
|
||||||
@app.get("/tasks/{task_id}/events")
|
@app.get("/tasks/{task_id}/events")
|
||||||
|
|||||||
@@ -121,3 +121,33 @@ def build_task_event_dto(event: DomainEvent) -> dict[str, Any]:
|
|||||||
"payload": payload,
|
"payload": payload,
|
||||||
"card": card,
|
"card": card,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def build_artifact_cards_from_tool_invocations(
|
||||||
|
task_id: str,
|
||||||
|
tool_invocations: list[dict[str, Any]],
|
||||||
|
) -> list[dict[str, Any]]:
|
||||||
|
cards: list[dict[str, Any]] = []
|
||||||
|
for invocation in tool_invocations:
|
||||||
|
response = invocation.get("response", {})
|
||||||
|
artifacts = response.get("artifacts", []) or []
|
||||||
|
if artifacts:
|
||||||
|
for index, artifact in enumerate(artifacts, start=1):
|
||||||
|
cards.append(
|
||||||
|
build_artifact_card(
|
||||||
|
task_id=task_id,
|
||||||
|
artifact_type=str(artifact.get("type", "artifact")),
|
||||||
|
title=f"{invocation.get('tool_name', 'tool')} artifact {index}",
|
||||||
|
data=artifact,
|
||||||
|
)
|
||||||
|
)
|
||||||
|
elif response.get("content"):
|
||||||
|
cards.append(
|
||||||
|
build_artifact_card(
|
||||||
|
task_id=task_id,
|
||||||
|
artifact_type="tool_result",
|
||||||
|
title=f"{invocation.get('tool_name', 'tool')} result",
|
||||||
|
data=response["content"],
|
||||||
|
)
|
||||||
|
)
|
||||||
|
return cards
|
||||||
|
|||||||
@@ -59,6 +59,9 @@ class InMemoryConfirmationRepository(ConfirmationRepository):
|
|||||||
self.items[confirmation.confirmation_id] = confirmation
|
self.items[confirmation.confirmation_id] = confirmation
|
||||||
return confirmation
|
return confirmation
|
||||||
|
|
||||||
|
def list_by_task(self, task_id: str) -> list[ConfirmationRequest]:
|
||||||
|
return [item for item in self.items.values() if item.task_id == task_id]
|
||||||
|
|
||||||
|
|
||||||
@dataclass(slots=True)
|
@dataclass(slots=True)
|
||||||
class InMemoryWorkerRepository(WorkerRepository):
|
class InMemoryWorkerRepository(WorkerRepository):
|
||||||
|
|||||||
@@ -531,6 +531,13 @@ class SqlAlchemyConfirmationRepository(ConfirmationRepository):
|
|||||||
session.commit()
|
session.commit()
|
||||||
return confirmation
|
return confirmation
|
||||||
|
|
||||||
|
def list_by_task(self, task_id: str) -> list[ConfirmationRequest]:
|
||||||
|
with self.session_factory() as session:
|
||||||
|
records = session.scalars(
|
||||||
|
select(ConfirmationRecord).where(ConfirmationRecord.task_id == task_id)
|
||||||
|
).all()
|
||||||
|
return [_record_to_confirmation(record) for record in records]
|
||||||
|
|
||||||
|
|
||||||
@dataclass(slots=True)
|
@dataclass(slots=True)
|
||||||
class SqlAlchemyWorkerRepository(WorkerRepository):
|
class SqlAlchemyWorkerRepository(WorkerRepository):
|
||||||
|
|||||||
@@ -57,6 +57,33 @@ def test_chat_response_contains_ui_cards() -> None:
|
|||||||
assert any(card["type"] == "progress_card" for card in payload["cards"])
|
assert any(card["type"] == "progress_card" for card in payload["cards"])
|
||||||
|
|
||||||
|
|
||||||
|
def test_task_creation_with_policy_action_returns_confirmation_card() -> None:
|
||||||
|
client = TestClient(app)
|
||||||
|
|
||||||
|
response = client.post(
|
||||||
|
"/tasks",
|
||||||
|
json={
|
||||||
|
"project_id": "default",
|
||||||
|
"goal": "Need approval",
|
||||||
|
"inputs": {
|
||||||
|
"policy_action": {
|
||||||
|
"resource": "filesystem",
|
||||||
|
"risk_level": "write",
|
||||||
|
"action_name": "file.write",
|
||||||
|
"preview_available": True,
|
||||||
|
"metadata": {"path": "D:/Projects/test.txt"},
|
||||||
|
}
|
||||||
|
},
|
||||||
|
"execution_mode": "agent_graph",
|
||||||
|
},
|
||||||
|
)
|
||||||
|
|
||||||
|
payload = response.json()
|
||||||
|
|
||||||
|
assert response.status_code == 200
|
||||||
|
assert any(card["type"] == "confirmation_card" for card in payload["cards"])
|
||||||
|
|
||||||
|
|
||||||
def test_cancel_endpoint_marks_task_cancelled() -> None:
|
def test_cancel_endpoint_marks_task_cancelled() -> None:
|
||||||
client = TestClient(app)
|
client = TestClient(app)
|
||||||
|
|
||||||
@@ -166,6 +193,65 @@ def test_worker_websocket_heartbeat_and_poll() -> None:
|
|||||||
assert len(commands["commands"]) == 1
|
assert len(commands["commands"]) == 1
|
||||||
|
|
||||||
|
|
||||||
|
def test_task_get_returns_artifact_card_from_worker_result() -> None:
|
||||||
|
client = TestClient(app)
|
||||||
|
created = client.post(
|
||||||
|
"/tasks",
|
||||||
|
json={
|
||||||
|
"project_id": "default",
|
||||||
|
"goal": "Artifact check",
|
||||||
|
"inputs": {},
|
||||||
|
"execution_mode": "agent_graph",
|
||||||
|
},
|
||||||
|
).json()
|
||||||
|
registration = client.post(
|
||||||
|
"/workers/register",
|
||||||
|
json={
|
||||||
|
"worker_id": "worker_artifact_pc",
|
||||||
|
"name": "Artifact PC",
|
||||||
|
"capabilities": ["file.read"],
|
||||||
|
"version": "0.1.0",
|
||||||
|
"machine": "DESKTOP-3",
|
||||||
|
"os": "windows",
|
||||||
|
},
|
||||||
|
).json()
|
||||||
|
client.post(
|
||||||
|
f"/workers/{registration['session_id']}/commands",
|
||||||
|
json={
|
||||||
|
"task_id": created["task_id"],
|
||||||
|
"node_id": "node_artifact",
|
||||||
|
"command_name": "file.read",
|
||||||
|
"args": {"path": "D:/artifact.txt"},
|
||||||
|
},
|
||||||
|
)
|
||||||
|
with client.websocket_connect(f"/workers/ws/{registration['session_id']}") as websocket:
|
||||||
|
websocket.send_json({"event_type": "poll"})
|
||||||
|
commands = websocket.receive_json()["commands"]
|
||||||
|
websocket.send_json(
|
||||||
|
{
|
||||||
|
"event_type": "result",
|
||||||
|
"command_id": commands[0]["command_id"],
|
||||||
|
"task_id": created["task_id"],
|
||||||
|
"tool": "file.read",
|
||||||
|
"status": "success",
|
||||||
|
"started_at": "2026-07-03T10:00:00+00:00",
|
||||||
|
"finished_at": "2026-07-03T10:00:01+00:00",
|
||||||
|
"duration_ms": 1000,
|
||||||
|
"stdout": "",
|
||||||
|
"stderr": "",
|
||||||
|
"result": {"content": "ok"},
|
||||||
|
"artifacts": [{"type": "file", "path": "D:/artifact.txt"}],
|
||||||
|
"error": None,
|
||||||
|
}
|
||||||
|
)
|
||||||
|
websocket.receive_json()
|
||||||
|
|
||||||
|
response = client.get(f"/tasks/{created['task_id']}")
|
||||||
|
|
||||||
|
assert response.status_code == 200
|
||||||
|
assert any(card["type"] == "artifact_card" for card in response.json()["cards"])
|
||||||
|
|
||||||
|
|
||||||
def test_task_events_stream_returns_normalized_event_payload() -> None:
|
def test_task_events_stream_returns_normalized_event_payload() -> None:
|
||||||
client = TestClient(app)
|
client = TestClient(app)
|
||||||
created = client.post(
|
created = client.post(
|
||||||
|
|||||||
@@ -0,0 +1,226 @@
|
|||||||
|
from datetime import UTC, datetime
|
||||||
|
|
||||||
|
import httpx
|
||||||
|
|
||||||
|
from ai_orchestrator.application.services.execution import (
|
||||||
|
DefaultFinalizer,
|
||||||
|
FinalizerNodeRunner,
|
||||||
|
GraphExecutionEngine,
|
||||||
|
NoOpReviewer,
|
||||||
|
PlannerNodeRunner,
|
||||||
|
RunnerRegistry,
|
||||||
|
)
|
||||||
|
from ai_orchestrator.application.services.orchestrator import CreateTaskRequest, OrchestratorService
|
||||||
|
from ai_orchestrator.application.services.router import ConfigurableModelRouter, ProviderResponse
|
||||||
|
from ai_orchestrator.application.services.workers import (
|
||||||
|
CapabilityAwareWorkerGateway,
|
||||||
|
InMemoryWorkerConnectionManager,
|
||||||
|
RegisterWorkerRequest,
|
||||||
|
WorkerCommandResult,
|
||||||
|
WorkerService,
|
||||||
|
)
|
||||||
|
from ai_orchestrator.config import (
|
||||||
|
ExecutionConfig,
|
||||||
|
McpServerConfig,
|
||||||
|
ModelsConfig,
|
||||||
|
ModelSlotConfig,
|
||||||
|
PolicyConfig,
|
||||||
|
ProjectConfig,
|
||||||
|
ProjectMetadata,
|
||||||
|
)
|
||||||
|
from ai_orchestrator.infrastructure.mcp_client import McpHttpClient, McpToolGateway
|
||||||
|
from ai_orchestrator.infrastructure.model_router import StaticMockModelProvider
|
||||||
|
from ai_orchestrator.infrastructure.policy import StaticProjectPolicyEvaluator
|
||||||
|
from ai_orchestrator.infrastructure.storage.memory import (
|
||||||
|
InMemoryConfirmationRepository,
|
||||||
|
InMemoryEventStore,
|
||||||
|
InMemoryGraphRepository,
|
||||||
|
InMemoryInvocationStore,
|
||||||
|
InMemoryTaskRepository,
|
||||||
|
InMemoryWorkerRepository,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_task_can_plan_execute_and_finalize_over_two_scheduler_passes() -> None:
|
||||||
|
task_repository = InMemoryTaskRepository()
|
||||||
|
graph_repository = InMemoryGraphRepository()
|
||||||
|
confirmation_repository = InMemoryConfirmationRepository()
|
||||||
|
event_store = InMemoryEventStore()
|
||||||
|
orchestrator = OrchestratorService(
|
||||||
|
task_repository=task_repository,
|
||||||
|
graph_repository=graph_repository,
|
||||||
|
confirmation_repository=confirmation_repository,
|
||||||
|
event_store=event_store,
|
||||||
|
policy_evaluator=StaticProjectPolicyEvaluator(
|
||||||
|
projects={
|
||||||
|
"default": ProjectConfig(
|
||||||
|
project=ProjectMetadata(id="default", name="Default"),
|
||||||
|
models=ModelsConfig(
|
||||||
|
weak=ModelSlotConfig(provider="local", model="weak"),
|
||||||
|
strong=ModelSlotConfig(provider="local", model="strong"),
|
||||||
|
vision=ModelSlotConfig(provider="disabled"),
|
||||||
|
embedding=ModelSlotConfig(provider="local", model="embed"),
|
||||||
|
),
|
||||||
|
execution=ExecutionConfig(),
|
||||||
|
policy=PolicyConfig(),
|
||||||
|
)
|
||||||
|
}
|
||||||
|
),
|
||||||
|
)
|
||||||
|
registry = RunnerRegistry()
|
||||||
|
registry.register(PlannerNodeRunner())
|
||||||
|
registry.register(FinalizerNodeRunner(finalizer=DefaultFinalizer()))
|
||||||
|
engine = GraphExecutionEngine(
|
||||||
|
event_store=event_store,
|
||||||
|
runner_registry=registry,
|
||||||
|
reviewer=NoOpReviewer(),
|
||||||
|
)
|
||||||
|
|
||||||
|
task = orchestrator.create_task(
|
||||||
|
CreateTaskRequest(project_id="default", goal="Scenario", inputs={})
|
||||||
|
)
|
||||||
|
graph = orchestrator.plan_task(task.task_id)
|
||||||
|
engine.execute_ready_nodes(task=task, graph=graph)
|
||||||
|
engine.execute_ready_nodes(task=task, graph=graph)
|
||||||
|
|
||||||
|
assert graph.nodes[-1].output_data["status"] == "completed"
|
||||||
|
|
||||||
|
|
||||||
|
def test_worker_gateway_result_becomes_artifact_ready_invocation() -> None:
|
||||||
|
worker_repository = InMemoryWorkerRepository()
|
||||||
|
event_store = InMemoryEventStore()
|
||||||
|
invocation_store = InMemoryInvocationStore()
|
||||||
|
worker_service = WorkerService(worker_repository=worker_repository, event_store=event_store)
|
||||||
|
gateway = CapabilityAwareWorkerGateway(
|
||||||
|
worker_repository=worker_repository,
|
||||||
|
event_store=event_store,
|
||||||
|
invocation_store=invocation_store,
|
||||||
|
connection_manager=InMemoryWorkerConnectionManager(),
|
||||||
|
)
|
||||||
|
response = worker_service.register(
|
||||||
|
RegisterWorkerRequest(
|
||||||
|
worker_id="worker_home_pc",
|
||||||
|
name="Home PC",
|
||||||
|
machine="DESKTOP-1",
|
||||||
|
os="windows",
|
||||||
|
version="0.1.0",
|
||||||
|
capabilities=["file.read"],
|
||||||
|
)
|
||||||
|
)
|
||||||
|
queued = gateway.dispatch(
|
||||||
|
project_id="default",
|
||||||
|
worker_session_id=response.worker.session_id,
|
||||||
|
command_name="file.read",
|
||||||
|
args={"path": "D:/artifact.txt"},
|
||||||
|
task_context={"task_id": "task_1"},
|
||||||
|
)
|
||||||
|
polled = gateway.poll_commands(response.worker.session_id)
|
||||||
|
gateway.complete_command(
|
||||||
|
response.worker.session_id,
|
||||||
|
WorkerCommandResult(
|
||||||
|
command_id=polled[0].command_id,
|
||||||
|
task_id="task_1",
|
||||||
|
tool="file.read",
|
||||||
|
status="success",
|
||||||
|
started_at=polled[0].policy_context.get("started_at") or datetime.now(UTC),
|
||||||
|
finished_at=datetime.now(UTC),
|
||||||
|
duration_ms=100,
|
||||||
|
result={"content": "ok"},
|
||||||
|
artifacts=[{"type": "file", "path": "D:/artifact.txt"}],
|
||||||
|
),
|
||||||
|
)
|
||||||
|
|
||||||
|
assert queued.status == "queued"
|
||||||
|
assert (
|
||||||
|
invocation_store.list_tool_invocations("task_1")[-1].response["artifacts"][0]["path"]
|
||||||
|
== "D:/artifact.txt"
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_model_and_mcp_paths_can_run_in_same_runtime_context() -> None:
|
||||||
|
project = ProjectConfig(
|
||||||
|
project=ProjectMetadata(id="default", name="Default"),
|
||||||
|
models=ModelsConfig(
|
||||||
|
weak=ModelSlotConfig(provider="local", model="weak"),
|
||||||
|
strong=ModelSlotConfig(provider="external", model="strong"),
|
||||||
|
vision=ModelSlotConfig(provider="disabled"),
|
||||||
|
embedding=ModelSlotConfig(provider="local", model="embed"),
|
||||||
|
),
|
||||||
|
mcp_servers={"one_c": McpServerConfig(base_url="http://mcp.test")},
|
||||||
|
execution=ExecutionConfig(allow_external_models=True),
|
||||||
|
policy=PolicyConfig(),
|
||||||
|
)
|
||||||
|
invocation_store = InMemoryInvocationStore()
|
||||||
|
router = ConfigurableModelRouter(
|
||||||
|
project_configs={"default": project},
|
||||||
|
providers={
|
||||||
|
"local": StaticMockModelProvider(
|
||||||
|
provider_name="local",
|
||||||
|
responses=[
|
||||||
|
ProviderResponse(
|
||||||
|
status="success",
|
||||||
|
message={"content": ""},
|
||||||
|
tool_calls=[],
|
||||||
|
usage={},
|
||||||
|
),
|
||||||
|
ProviderResponse(
|
||||||
|
status="success",
|
||||||
|
message={"content": ""},
|
||||||
|
tool_calls=[],
|
||||||
|
usage={},
|
||||||
|
),
|
||||||
|
],
|
||||||
|
),
|
||||||
|
"external": StaticMockModelProvider(
|
||||||
|
provider_name="external",
|
||||||
|
responses=[
|
||||||
|
ProviderResponse(
|
||||||
|
status="success",
|
||||||
|
message={"content": "strong"},
|
||||||
|
tool_calls=[],
|
||||||
|
usage={"input_tokens": 1},
|
||||||
|
)
|
||||||
|
],
|
||||||
|
),
|
||||||
|
},
|
||||||
|
invocation_store=invocation_store,
|
||||||
|
)
|
||||||
|
|
||||||
|
def handler(request: httpx.Request) -> httpx.Response:
|
||||||
|
return httpx.Response(
|
||||||
|
200,
|
||||||
|
json={
|
||||||
|
"jsonrpc": "2.0",
|
||||||
|
"id": "1",
|
||||||
|
"result": {
|
||||||
|
"status": "success",
|
||||||
|
"content": {"rows": [{"id": 1}]},
|
||||||
|
"artifacts": [],
|
||||||
|
"logs": [],
|
||||||
|
"error": None,
|
||||||
|
},
|
||||||
|
},
|
||||||
|
)
|
||||||
|
|
||||||
|
gateway = McpToolGateway(
|
||||||
|
project_configs={"default": project},
|
||||||
|
client=McpHttpClient(transport=httpx.MockTransport(handler)),
|
||||||
|
invocation_store=invocation_store,
|
||||||
|
)
|
||||||
|
|
||||||
|
model_result = router.run(
|
||||||
|
project_id="default",
|
||||||
|
slot="weak",
|
||||||
|
messages=[{"role": "user", "content": "Hi"}],
|
||||||
|
task_context={"task_id": "task_1"},
|
||||||
|
)
|
||||||
|
tool_result = gateway.call(
|
||||||
|
project_id="default",
|
||||||
|
server_id="one_c",
|
||||||
|
tool_name="one_c.run_sql",
|
||||||
|
args={"query": "select 1"},
|
||||||
|
task_context={"task_id": "task_1"},
|
||||||
|
)
|
||||||
|
|
||||||
|
assert model_result.message["content"] == "strong"
|
||||||
|
assert tool_result.content["rows"][0]["id"] == 1
|
||||||
@@ -0,0 +1,321 @@
|
|||||||
|
import httpx
|
||||||
|
|
||||||
|
from ai_orchestrator.application.services.execution import (
|
||||||
|
DefaultFinalizer,
|
||||||
|
FinalizerNodeRunner,
|
||||||
|
GraphExecutionEngine,
|
||||||
|
NoOpReviewer,
|
||||||
|
PlannerNodeRunner,
|
||||||
|
RunnerRegistry,
|
||||||
|
)
|
||||||
|
from ai_orchestrator.application.services.orchestrator import CreateTaskRequest, OrchestratorService
|
||||||
|
from ai_orchestrator.application.services.router import ConfigurableModelRouter, ProviderResponse
|
||||||
|
from ai_orchestrator.application.services.workers import (
|
||||||
|
CapabilityAwareWorkerGateway,
|
||||||
|
InMemoryWorkerConnectionManager,
|
||||||
|
RegisterWorkerRequest,
|
||||||
|
WorkerCommandResult,
|
||||||
|
WorkerService,
|
||||||
|
)
|
||||||
|
from ai_orchestrator.config import (
|
||||||
|
ExecutionConfig,
|
||||||
|
McpServerConfig,
|
||||||
|
ModelsConfig,
|
||||||
|
ModelSlotConfig,
|
||||||
|
PolicyConfig,
|
||||||
|
ProjectConfig,
|
||||||
|
ProjectMetadata,
|
||||||
|
)
|
||||||
|
from ai_orchestrator.domain.enums import RiskLevel
|
||||||
|
from ai_orchestrator.domain.models import ActionDescriptor
|
||||||
|
from ai_orchestrator.infrastructure.mcp_client import McpHttpClient, McpToolGateway
|
||||||
|
from ai_orchestrator.infrastructure.model_router import StaticMockModelProvider
|
||||||
|
from ai_orchestrator.infrastructure.policy import StaticProjectPolicyEvaluator
|
||||||
|
from ai_orchestrator.infrastructure.storage.memory import (
|
||||||
|
InMemoryConfirmationRepository,
|
||||||
|
InMemoryEventStore,
|
||||||
|
InMemoryGraphRepository,
|
||||||
|
InMemoryInvocationStore,
|
||||||
|
InMemoryTaskRepository,
|
||||||
|
InMemoryWorkerRepository,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def _project_config(*, allow_external: bool = True, default_mode: str = "confirm") -> ProjectConfig:
|
||||||
|
return ProjectConfig(
|
||||||
|
project=ProjectMetadata(id="default", name="Default"),
|
||||||
|
models=ModelsConfig(
|
||||||
|
weak=ModelSlotConfig(provider="local", model="weak"),
|
||||||
|
strong=ModelSlotConfig(
|
||||||
|
provider="external" if allow_external else "disabled",
|
||||||
|
model="strong" if allow_external else None,
|
||||||
|
),
|
||||||
|
vision=ModelSlotConfig(provider="disabled"),
|
||||||
|
embedding=ModelSlotConfig(provider="local", model="embed"),
|
||||||
|
),
|
||||||
|
mcp_servers={"one_c": McpServerConfig(base_url="http://mcp.test")},
|
||||||
|
execution=ExecutionConfig(allow_external_models=allow_external),
|
||||||
|
policy=PolicyConfig(
|
||||||
|
default_mode=default_mode,
|
||||||
|
resources={
|
||||||
|
"filesystem": default_mode,
|
||||||
|
"shell": default_mode,
|
||||||
|
"sql": default_mode,
|
||||||
|
"mcp": default_mode,
|
||||||
|
"external_models": default_mode if allow_external else "disabled",
|
||||||
|
"desktop": default_mode,
|
||||||
|
"browser": default_mode,
|
||||||
|
"network": default_mode,
|
||||||
|
"cost": default_mode,
|
||||||
|
"system": default_mode,
|
||||||
|
},
|
||||||
|
),
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_smoke_simple_chat_and_graph_and_events() -> None:
|
||||||
|
task_repository = InMemoryTaskRepository()
|
||||||
|
graph_repository = InMemoryGraphRepository()
|
||||||
|
confirmation_repository = InMemoryConfirmationRepository()
|
||||||
|
event_store = InMemoryEventStore()
|
||||||
|
orchestrator = OrchestratorService(
|
||||||
|
task_repository=task_repository,
|
||||||
|
graph_repository=graph_repository,
|
||||||
|
confirmation_repository=confirmation_repository,
|
||||||
|
event_store=event_store,
|
||||||
|
policy_evaluator=StaticProjectPolicyEvaluator(projects={"default": _project_config()}),
|
||||||
|
)
|
||||||
|
registry = RunnerRegistry()
|
||||||
|
registry.register(PlannerNodeRunner())
|
||||||
|
registry.register(FinalizerNodeRunner(finalizer=DefaultFinalizer()))
|
||||||
|
engine = GraphExecutionEngine(
|
||||||
|
event_store=event_store,
|
||||||
|
runner_registry=registry,
|
||||||
|
reviewer=NoOpReviewer(),
|
||||||
|
)
|
||||||
|
|
||||||
|
task = orchestrator.create_task(
|
||||||
|
CreateTaskRequest(project_id="default", goal="hello", inputs={})
|
||||||
|
)
|
||||||
|
graph = orchestrator.plan_task(task.task_id)
|
||||||
|
engine.execute_ready_nodes(task=task, graph=graph)
|
||||||
|
|
||||||
|
assert len(graph.nodes) == 2
|
||||||
|
assert any(
|
||||||
|
event.event_type == "node_started" for event in event_store.list_by_task(task.task_id)
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_smoke_model_router_retry_and_fallback_and_external_disabled() -> None:
|
||||||
|
invocation_store = InMemoryInvocationStore()
|
||||||
|
project = _project_config(allow_external=True)
|
||||||
|
router = ConfigurableModelRouter(
|
||||||
|
project_configs={"default": project},
|
||||||
|
providers={
|
||||||
|
"local": StaticMockModelProvider(
|
||||||
|
provider_name="local",
|
||||||
|
responses=[
|
||||||
|
ProviderResponse(
|
||||||
|
status="success",
|
||||||
|
message={"content": ""},
|
||||||
|
tool_calls=[],
|
||||||
|
usage={},
|
||||||
|
),
|
||||||
|
ProviderResponse(
|
||||||
|
status="success",
|
||||||
|
message={"content": ""},
|
||||||
|
tool_calls=[],
|
||||||
|
usage={},
|
||||||
|
),
|
||||||
|
],
|
||||||
|
),
|
||||||
|
"external": StaticMockModelProvider(
|
||||||
|
provider_name="external",
|
||||||
|
responses=[
|
||||||
|
ProviderResponse(
|
||||||
|
status="success",
|
||||||
|
message={"content": "strong"},
|
||||||
|
tool_calls=[],
|
||||||
|
usage={},
|
||||||
|
)
|
||||||
|
],
|
||||||
|
),
|
||||||
|
},
|
||||||
|
invocation_store=invocation_store,
|
||||||
|
)
|
||||||
|
result = router.run(
|
||||||
|
project_id="default",
|
||||||
|
slot="weak",
|
||||||
|
messages=[{"role": "user", "content": "hi"}],
|
||||||
|
task_context={"task_id": "task_1"},
|
||||||
|
)
|
||||||
|
|
||||||
|
disabled_router = ConfigurableModelRouter(
|
||||||
|
project_configs={"default": _project_config(allow_external=False)},
|
||||||
|
providers={"local": StaticMockModelProvider(provider_name="local")},
|
||||||
|
invocation_store=InMemoryInvocationStore(),
|
||||||
|
)
|
||||||
|
disabled_result = disabled_router.run(
|
||||||
|
project_id="default",
|
||||||
|
slot="weak",
|
||||||
|
messages=[{"role": "user", "content": "hi"}],
|
||||||
|
task_context={"task_id": "task_2"},
|
||||||
|
)
|
||||||
|
|
||||||
|
assert result.status == "success"
|
||||||
|
assert result.message["content"] == "strong"
|
||||||
|
assert disabled_result.status in {"success", "degraded", "disabled"}
|
||||||
|
|
||||||
|
|
||||||
|
def test_smoke_policy_confirm_and_full_auto() -> None:
|
||||||
|
confirmation_repository = InMemoryConfirmationRepository()
|
||||||
|
event_store = InMemoryEventStore()
|
||||||
|
orchestrator_confirm = OrchestratorService(
|
||||||
|
task_repository=InMemoryTaskRepository(),
|
||||||
|
graph_repository=InMemoryGraphRepository(),
|
||||||
|
confirmation_repository=confirmation_repository,
|
||||||
|
event_store=event_store,
|
||||||
|
policy_evaluator=StaticProjectPolicyEvaluator(
|
||||||
|
projects={"default": _project_config(default_mode="confirm")}
|
||||||
|
),
|
||||||
|
)
|
||||||
|
task = orchestrator_confirm.create_task(
|
||||||
|
CreateTaskRequest(project_id="default", goal="confirm", inputs={})
|
||||||
|
)
|
||||||
|
graph = orchestrator_confirm.plan_task(task.task_id)
|
||||||
|
confirmation = orchestrator_confirm.evaluate_action(
|
||||||
|
task_id=task.task_id,
|
||||||
|
node_id=graph.nodes[-1].node_id,
|
||||||
|
action=ActionDescriptor(
|
||||||
|
resource="filesystem",
|
||||||
|
risk_level=RiskLevel.WRITE,
|
||||||
|
action_name="file.write",
|
||||||
|
preview_available=True,
|
||||||
|
),
|
||||||
|
)
|
||||||
|
|
||||||
|
orchestrator_auto = OrchestratorService(
|
||||||
|
task_repository=InMemoryTaskRepository(),
|
||||||
|
graph_repository=InMemoryGraphRepository(),
|
||||||
|
confirmation_repository=InMemoryConfirmationRepository(),
|
||||||
|
event_store=InMemoryEventStore(),
|
||||||
|
policy_evaluator=StaticProjectPolicyEvaluator(
|
||||||
|
projects={"default": _project_config(default_mode="full_auto")}
|
||||||
|
),
|
||||||
|
)
|
||||||
|
auto_task = orchestrator_auto.create_task(
|
||||||
|
CreateTaskRequest(project_id="default", goal="auto", inputs={})
|
||||||
|
)
|
||||||
|
auto_graph = orchestrator_auto.plan_task(auto_task.task_id)
|
||||||
|
auto_confirmation = orchestrator_auto.evaluate_action(
|
||||||
|
task_id=auto_task.task_id,
|
||||||
|
node_id=auto_graph.nodes[-1].node_id,
|
||||||
|
action=ActionDescriptor(
|
||||||
|
resource="filesystem",
|
||||||
|
risk_level=RiskLevel.WRITE,
|
||||||
|
action_name="file.write",
|
||||||
|
),
|
||||||
|
)
|
||||||
|
|
||||||
|
assert confirmation is not None
|
||||||
|
assert auto_confirmation is None
|
||||||
|
|
||||||
|
|
||||||
|
def test_smoke_mcp_and_worker_results_are_stored_and_finalizer_runs() -> None:
|
||||||
|
invocation_store = InMemoryInvocationStore()
|
||||||
|
|
||||||
|
def handler(request: httpx.Request) -> httpx.Response:
|
||||||
|
return httpx.Response(
|
||||||
|
200,
|
||||||
|
json={
|
||||||
|
"jsonrpc": "2.0",
|
||||||
|
"id": "1",
|
||||||
|
"result": {
|
||||||
|
"status": "success",
|
||||||
|
"content": {"rows": [1]},
|
||||||
|
"artifacts": [],
|
||||||
|
"logs": [],
|
||||||
|
"error": None,
|
||||||
|
},
|
||||||
|
},
|
||||||
|
)
|
||||||
|
|
||||||
|
mcp = McpToolGateway(
|
||||||
|
project_configs={"default": _project_config()},
|
||||||
|
client=McpHttpClient(transport=httpx.MockTransport(handler)),
|
||||||
|
invocation_store=invocation_store,
|
||||||
|
)
|
||||||
|
worker_repository = InMemoryWorkerRepository()
|
||||||
|
event_store = InMemoryEventStore()
|
||||||
|
worker_service = WorkerService(worker_repository=worker_repository, event_store=event_store)
|
||||||
|
gateway = CapabilityAwareWorkerGateway(
|
||||||
|
worker_repository=worker_repository,
|
||||||
|
event_store=event_store,
|
||||||
|
invocation_store=invocation_store,
|
||||||
|
connection_manager=InMemoryWorkerConnectionManager(),
|
||||||
|
)
|
||||||
|
worker = worker_service.register(
|
||||||
|
RegisterWorkerRequest(
|
||||||
|
worker_id="worker_home_pc",
|
||||||
|
name="Home PC",
|
||||||
|
machine="DESKTOP-1",
|
||||||
|
os="windows",
|
||||||
|
version="0.1.0",
|
||||||
|
capabilities=["file.read"],
|
||||||
|
)
|
||||||
|
).worker
|
||||||
|
queued = gateway.dispatch(
|
||||||
|
project_id="default",
|
||||||
|
worker_session_id=worker.session_id,
|
||||||
|
command_name="file.read",
|
||||||
|
args={"path": "D:/test.txt"},
|
||||||
|
task_context={"task_id": "task_1"},
|
||||||
|
)
|
||||||
|
command = gateway.poll_commands(worker.session_id)[0]
|
||||||
|
gateway.complete_command(
|
||||||
|
worker.session_id,
|
||||||
|
WorkerCommandResult(
|
||||||
|
command_id=command.command_id,
|
||||||
|
task_id="task_1",
|
||||||
|
tool="file.read",
|
||||||
|
status="success",
|
||||||
|
started_at=__import__("datetime").datetime.now(__import__("datetime").UTC),
|
||||||
|
finished_at=__import__("datetime").datetime.now(__import__("datetime").UTC),
|
||||||
|
duration_ms=1,
|
||||||
|
result={"content": "ok"},
|
||||||
|
),
|
||||||
|
)
|
||||||
|
tool_result = mcp.call(
|
||||||
|
project_id="default",
|
||||||
|
server_id="one_c",
|
||||||
|
tool_name="one_c.run_sql",
|
||||||
|
args={"query": "select 1"},
|
||||||
|
task_context={"task_id": "task_1"},
|
||||||
|
)
|
||||||
|
|
||||||
|
registry = RunnerRegistry()
|
||||||
|
registry.register(PlannerNodeRunner())
|
||||||
|
registry.register(FinalizerNodeRunner(finalizer=DefaultFinalizer()))
|
||||||
|
engine = GraphExecutionEngine(
|
||||||
|
event_store=event_store,
|
||||||
|
runner_registry=registry,
|
||||||
|
reviewer=NoOpReviewer(),
|
||||||
|
)
|
||||||
|
orchestrator = OrchestratorService(
|
||||||
|
task_repository=InMemoryTaskRepository(),
|
||||||
|
graph_repository=InMemoryGraphRepository(),
|
||||||
|
confirmation_repository=InMemoryConfirmationRepository(),
|
||||||
|
event_store=event_store,
|
||||||
|
policy_evaluator=StaticProjectPolicyEvaluator(projects={"default": _project_config()}),
|
||||||
|
)
|
||||||
|
task = orchestrator.create_task(
|
||||||
|
CreateTaskRequest(project_id="default", goal="finalize", inputs={})
|
||||||
|
)
|
||||||
|
graph = orchestrator.plan_task(task.task_id)
|
||||||
|
engine.execute_ready_nodes(task=task, graph=graph)
|
||||||
|
engine.execute_ready_nodes(task=task, graph=graph)
|
||||||
|
|
||||||
|
assert queued.status == "queued"
|
||||||
|
assert tool_result.status == "success"
|
||||||
|
assert len(invocation_store.list_tool_invocations("task_1")) >= 2
|
||||||
|
assert graph.nodes[-1].output_data["status"] == "completed"
|
||||||
Reference in New Issue
Block a user