diff --git a/docs/backlog/01_execution_backlog.md b/docs/backlog/01_execution_backlog.md index ec383e9..b12a0ab 100644 --- a/docs/backlog/01_execution_backlog.md +++ b/docs/backlog/01_execution_backlog.md @@ -20,7 +20,7 @@ - [x] graph scheduler service - [x] node runner abstraction -- [ ] reviewer contract +- [x] reviewer contract - [x] finalizer contract - [x] resume/cancel use cases - [x] idempotency policy @@ -31,8 +31,8 @@ - [x] SQLAlchemy models - [x] repositories for PostgreSQL/SQLite - [x] migrations -- [ ] artifact metadata persistence -- [ ] event store queries with pagination +- [x] artifact metadata persistence +- [x] event store queries with pagination ## Package D: Integrations @@ -53,7 +53,7 @@ ## Package F: Validation -- [ ] contract tests -- [ ] scenario tests -- [ ] smoke suite +- [x] contract tests +- [x] scenario tests +- [x] smoke suite - [ ] static checks in CI diff --git a/docs/storage/01_schema.md b/docs/storage/01_schema.md index 8338987..689916d 100644 --- a/docs/storage/01_schema.md +++ b/docs/storage/01_schema.md @@ -169,6 +169,17 @@ Migration skeleton: - `confirmations(status, created_at)` - `worker_sessions(status, last_heartbeat_at)` +## Query Support + +Read-model adapters now support paginated queries for: + +- `task_events` +- `model_invocations` +- `tool_invocations` +- `artifacts` + +This is implemented in both in-memory and SQLAlchemy adapters to keep test/runtime behavior aligned. + ## Active Adapter Strategy - `memory` backend остается для быстрых unit/integration прогонов; diff --git a/migrations/versions/0001_initial_schema.py b/migrations/versions/0001_initial_schema.py index 290fc19..a853a27 100644 --- a/migrations/versions/0001_initial_schema.py +++ b/migrations/versions/0001_initial_schema.py @@ -153,8 +153,24 @@ def upgrade() -> None: op.create_index("ix_tool_invocations_task_id", "tool_invocations", ["task_id"]) op.create_index("ix_tool_invocations_created_at", "tool_invocations", ["created_at"]) + op.create_table( + "artifacts", + sa.Column("artifact_id", sa.String(length=64), primary_key=True), + sa.Column("task_id", sa.String(length=64), nullable=False), + sa.Column("node_id", sa.String(length=64), nullable=True), + sa.Column("artifact_type", sa.String(length=64), nullable=False), + sa.Column("storage_uri", sa.String(length=512), nullable=False), + sa.Column("metadata_json", sa.JSON(), nullable=False), + sa.Column("created_at", sa.DateTime(timezone=True), nullable=False), + ) + op.create_index("ix_artifacts_task_id", "artifacts", ["task_id"]) + op.create_index("ix_artifacts_created_at", "artifacts", ["created_at"]) + def downgrade() -> None: + op.drop_index("ix_artifacts_created_at", table_name="artifacts") + op.drop_index("ix_artifacts_task_id", table_name="artifacts") + op.drop_table("artifacts") op.drop_index("ix_tool_invocations_created_at", table_name="tool_invocations") op.drop_index("ix_tool_invocations_task_id", table_name="tool_invocations") op.drop_table("tool_invocations") diff --git a/src/ai_orchestrator/application/ports.py b/src/ai_orchestrator/application/ports.py index 48e6a41..4adccef 100644 --- a/src/ai_orchestrator/application/ports.py +++ b/src/ai_orchestrator/application/ports.py @@ -63,6 +63,24 @@ class ToolInvocationRecord: response: dict[str, Any] +@dataclass(slots=True) +class ArtifactRecord: + artifact_id: str + task_id: str + node_id: str | None + artifact_type: str + storage_uri: str + metadata: dict[str, Any] = field(default_factory=dict) + + +@dataclass(slots=True) +class Page: + items: list[Any] + limit: int + offset: int + total: int + + class TaskRepository(Protocol): def create(self, task: Task) -> Task: ... def save(self, task: Task) -> Task: ... @@ -90,6 +108,7 @@ class WorkerRepository(Protocol): class EventStore(Protocol): def append(self, event: DomainEvent) -> DomainEvent: ... def list_by_task(self, task_id: str) -> list[DomainEvent]: ... + def list_by_task_paginated(self, task_id: str, *, limit: int, offset: int) -> Page: ... class InvocationStore(Protocol): @@ -97,6 +116,18 @@ class InvocationStore(Protocol): def save_tool_invocation(self, record: ToolInvocationRecord) -> ToolInvocationRecord: ... def list_model_invocations(self, task_id: str) -> list[ModelInvocationRecord]: ... def list_tool_invocations(self, task_id: str) -> list[ToolInvocationRecord]: ... + def list_model_invocations_paginated( + self, task_id: str, *, limit: int, offset: int + ) -> Page: ... + def list_tool_invocations_paginated( + self, task_id: str, *, limit: int, offset: int + ) -> Page: ... + + +class ArtifactStore(Protocol): + def save_artifact(self, record: ArtifactRecord) -> ArtifactRecord: ... + def list_by_task(self, task_id: str) -> list[ArtifactRecord]: ... + def list_by_task_paginated(self, task_id: str, *, limit: int, offset: int) -> Page: ... class PolicyEvaluator(Protocol): diff --git a/src/ai_orchestrator/application/services/execution.py b/src/ai_orchestrator/application/services/execution.py index b0ea417..a2e35b4 100644 --- a/src/ai_orchestrator/application/services/execution.py +++ b/src/ai_orchestrator/application/services/execution.py @@ -1,6 +1,7 @@ from __future__ import annotations from dataclasses import dataclass, field +from typing import Any from ai_orchestrator.application.ports import EventStore, Finalizer, NodeRunner, Reviewer from ai_orchestrator.domain.enums import NodeStatus, NodeType @@ -39,6 +40,32 @@ class NoOpReviewer(Reviewer): } +@dataclass(slots=True) +class StructuredReviewer(Reviewer): + required_keys: tuple[str, ...] = ("status",) + + def review( + self, + *, + task: Task, + node: ExecutionNode, + output_data: dict[str, Any], + ) -> dict[str, Any]: + missing = [key for key in self.required_keys if key not in output_data] + reviewed = { + **output_data, + "task_id": task.task_id, + "node_id": node.node_id, + "review": { + "status": "accepted" if not missing else "rejected", + "missing_keys": missing, + }, + } + if missing: + reviewed["status"] = "review_failed" + return reviewed + + @dataclass(slots=True) class DefaultFinalizer(Finalizer): def finalize(self, *, task: Task, graph: ExecutionGraph) -> dict[str, object]: diff --git a/src/ai_orchestrator/application/services/workers.py b/src/ai_orchestrator/application/services/workers.py index 69c776a..954885c 100644 --- a/src/ai_orchestrator/application/services/workers.py +++ b/src/ai_orchestrator/application/services/workers.py @@ -5,6 +5,8 @@ from datetime import datetime from uuid import uuid4 from ai_orchestrator.application.ports import ( + ArtifactRecord, + ArtifactStore, EventStore, InvocationStore, ToolInvocationRecord, @@ -115,6 +117,7 @@ class CapabilityAwareWorkerGateway(WorkerGateway): worker_repository: WorkerRepository event_store: EventStore invocation_store: InvocationStore + artifact_store: ArtifactStore connection_manager: InMemoryWorkerConnectionManager def dispatch( @@ -244,6 +247,21 @@ class CapabilityAwareWorkerGateway(WorkerGateway): }, ) ) + for index, artifact in enumerate(result.artifacts or [], start=1): + self.artifact_store.save_artifact( + ArtifactRecord( + artifact_id=f"art_{uuid4().hex}", + task_id=result.task_id, + node_id=None, + artifact_type=str(artifact.get("type", "artifact")), + storage_uri=str( + artifact.get("path") + or artifact.get("uri") + or f"worker://{worker_session_id}/{result.command_id}/{index}" + ), + metadata=dict(artifact), + ) + ) return ToolInvocationResult( status=result.status, content=result.result or {}, diff --git a/src/ai_orchestrator/delivery/http/app.py b/src/ai_orchestrator/delivery/http/app.py index 2bc9e0c..4d7506b 100644 --- a/src/ai_orchestrator/delivery/http/app.py +++ b/src/ai_orchestrator/delivery/http/app.py @@ -9,9 +9,9 @@ from ai_orchestrator.application.services.execution import ( DefaultFinalizer, FinalizerNodeRunner, GraphExecutionEngine, - NoOpReviewer, PlannerNodeRunner, RunnerRegistry, + StructuredReviewer, ) from ai_orchestrator.application.services.lifecycle import TaskLifecycleService from ai_orchestrator.application.services.orchestrator import CreateTaskRequest, OrchestratorService @@ -69,6 +69,7 @@ def create_app(settings: AppSettings | None = None) -> FastAPI: worker_repository = storage.worker_repository event_store = storage.event_store invocation_store = storage.invocation_store + artifact_store = storage.artifact_store policy_evaluator = StaticProjectPolicyEvaluator( projects={project_config.project.id: project_config} ) @@ -77,6 +78,7 @@ def create_app(settings: AppSettings | None = None) -> FastAPI: worker_repository=worker_repository, event_store=event_store, invocation_store=invocation_store, + artifact_store=artifact_store, connection_manager=worker_connection_manager, ) model_router = ConfigurableModelRouter( @@ -93,7 +95,7 @@ def create_app(settings: AppSettings | None = None) -> FastAPI: execution_engine = GraphExecutionEngine( event_store=event_store, runner_registry=runner_registry, - reviewer=NoOpReviewer(), + reviewer=StructuredReviewer(), ) orchestrator = OrchestratorService( task_repository=task_repository, @@ -117,6 +119,7 @@ def create_app(settings: AppSettings | None = None) -> FastAPI: app.state.event_store = event_store app.state.storage = storage app.state.invocation_store = invocation_store + app.state.artifact_store = artifact_store app.state.model_router = model_router app.state.worker_gateway = worker_gateway app.state.worker_connection_manager = worker_connection_manager @@ -245,12 +248,27 @@ def create_app(settings: AppSettings | None = None) -> FastAPI: for item in invocation_store.list_tool_invocations(task_id) ] artifact_cards = build_artifact_cards_from_tool_invocations(task_id, tool_invocations) + persisted_artifact_cards = [ + { + "type": "artifact_card", + "task_id": item.task_id, + "artifact_type": item.artifact_type, + "title": item.metadata.get("title", item.artifact_type), + "data": {"storage_uri": item.storage_uri, **item.metadata}, + } + for item in artifact_store.list_by_task(task_id) + ] return TaskStatusResponse( task_id=task.task_id, status=task.status.value, current_node=task.current_node_id, progress={"completed": completed, "total": len(graph.nodes)}, - cards=[build_progress_card(task, graph), *confirmation_cards, *artifact_cards], + cards=[ + build_progress_card(task, graph), + *confirmation_cards, + *artifact_cards, + *persisted_artifact_cards, + ], ) @app.get("/tasks/{task_id}/events") diff --git a/src/ai_orchestrator/infrastructure/mcp_client.py b/src/ai_orchestrator/infrastructure/mcp_client.py index a25c033..c357589 100644 --- a/src/ai_orchestrator/infrastructure/mcp_client.py +++ b/src/ai_orchestrator/infrastructure/mcp_client.py @@ -7,6 +7,8 @@ from uuid import uuid4 import httpx from ai_orchestrator.application.ports import ( + ArtifactRecord, + ArtifactStore, InvocationStore, ToolGateway, ToolInvocationRecord, @@ -60,6 +62,7 @@ class McpToolGateway(ToolGateway): project_configs: dict[str, ProjectConfig] client: McpHttpClient invocation_store: InvocationStore + artifact_store: ArtifactStore def call( self, @@ -105,4 +108,19 @@ class McpToolGateway(ToolGateway): }, ) ) + for index, artifact in enumerate(result.artifacts, start=1): + self.artifact_store.save_artifact( + ArtifactRecord( + artifact_id=f"art_{uuid4().hex}", + task_id=str(task_context.get("task_id", project_id)), + node_id=str(task_context["node_id"]) if task_context.get("node_id") else None, + artifact_type=str(artifact.get("type", "artifact")), + storage_uri=str( + artifact.get("path") + or artifact.get("uri") + or f"mcp://{server_id}/{tool_name}/{index}" + ), + metadata=dict(artifact), + ) + ) return result diff --git a/src/ai_orchestrator/infrastructure/storage/factory.py b/src/ai_orchestrator/infrastructure/storage/factory.py index 28e5e52..e5594e8 100644 --- a/src/ai_orchestrator/infrastructure/storage/factory.py +++ b/src/ai_orchestrator/infrastructure/storage/factory.py @@ -4,6 +4,7 @@ from dataclasses import dataclass from ai_orchestrator.config import AppSettings from ai_orchestrator.infrastructure.storage.memory import ( + InMemoryArtifactStore, InMemoryConfirmationRepository, InMemoryEventStore, InMemoryGraphRepository, @@ -22,6 +23,7 @@ class StorageBundle: worker_repository: object event_store: object invocation_store: object + artifact_store: object engine: object | None = None session_factory: object | None = None @@ -35,6 +37,7 @@ def create_storage_bundle(settings: AppSettings) -> StorageBundle: worker_repository=InMemoryWorkerRepository(), event_store=InMemoryEventStore(), invocation_store=InMemoryInvocationStore(), + artifact_store=InMemoryArtifactStore(), ) if settings.storage_backend in {"sqlite", "postgres"}: bundle = create_sqlalchemy_storage(settings.database_url) @@ -45,6 +48,7 @@ def create_storage_bundle(settings: AppSettings) -> StorageBundle: worker_repository=bundle.worker_repository, event_store=bundle.event_store, invocation_store=bundle.invocation_store, + artifact_store=bundle.artifact_store, engine=bundle.engine, session_factory=bundle.session_factory, ) diff --git a/src/ai_orchestrator/infrastructure/storage/memory.py b/src/ai_orchestrator/infrastructure/storage/memory.py index 2b454b4..5719be7 100644 --- a/src/ai_orchestrator/infrastructure/storage/memory.py +++ b/src/ai_orchestrator/infrastructure/storage/memory.py @@ -3,11 +3,14 @@ from __future__ import annotations from dataclasses import dataclass, field from ai_orchestrator.application.ports import ( + ArtifactRecord, + ArtifactStore, ConfirmationRepository, EventStore, GraphRepository, InvocationStore, ModelInvocationRecord, + Page, TaskRepository, ToolInvocationRecord, WorkerRepository, @@ -89,6 +92,15 @@ class InMemoryEventStore(EventStore): def list_by_task(self, task_id: str) -> list[DomainEvent]: return [event for event in self.items if event.task_id == task_id] + def list_by_task_paginated(self, task_id: str, *, limit: int, offset: int) -> Page: + filtered = self.list_by_task(task_id) + return Page( + items=filtered[offset : offset + limit], + limit=limit, + offset=offset, + total=len(filtered), + ) + @dataclass(slots=True) class InMemoryInvocationStore(InvocationStore): @@ -108,3 +120,42 @@ class InMemoryInvocationStore(InvocationStore): def list_tool_invocations(self, task_id: str) -> list[ToolInvocationRecord]: return [record for record in self.tool_items if record.task_id == task_id] + + def list_model_invocations_paginated(self, task_id: str, *, limit: int, offset: int) -> Page: + filtered = self.list_model_invocations(task_id) + return Page( + items=filtered[offset : offset + limit], + limit=limit, + offset=offset, + total=len(filtered), + ) + + def list_tool_invocations_paginated(self, task_id: str, *, limit: int, offset: int) -> Page: + filtered = self.list_tool_invocations(task_id) + return Page( + items=filtered[offset : offset + limit], + limit=limit, + offset=offset, + total=len(filtered), + ) + + +@dataclass(slots=True) +class InMemoryArtifactStore(ArtifactStore): + items: list[ArtifactRecord] = field(default_factory=list) + + def save_artifact(self, record: ArtifactRecord) -> ArtifactRecord: + self.items.append(record) + return record + + def list_by_task(self, task_id: str) -> list[ArtifactRecord]: + return [record for record in self.items if record.task_id == task_id] + + def list_by_task_paginated(self, task_id: str, *, limit: int, offset: int) -> Page: + filtered = self.list_by_task(task_id) + return Page( + items=filtered[offset : offset + limit], + limit=limit, + offset=offset, + total=len(filtered), + ) diff --git a/src/ai_orchestrator/infrastructure/storage/sqlalchemy.py b/src/ai_orchestrator/infrastructure/storage/sqlalchemy.py index a045b3b..6985e1f 100644 --- a/src/ai_orchestrator/infrastructure/storage/sqlalchemy.py +++ b/src/ai_orchestrator/infrastructure/storage/sqlalchemy.py @@ -8,11 +8,14 @@ from sqlalchemy import JSON, Boolean, DateTime, Integer, String, Text, create_en from sqlalchemy.orm import DeclarativeBase, Mapped, Session, mapped_column, sessionmaker from ai_orchestrator.application.ports import ( + ArtifactRecord, + ArtifactStore, ConfirmationRepository, EventStore, GraphRepository, InvocationStore, ModelInvocationRecord, + Page, TaskRepository, ToolInvocationRecord, WorkerRepository, @@ -163,6 +166,18 @@ class ToolInvocationDbRecord(Base): created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), index=True) +class ArtifactDbRecord(Base): + __tablename__ = "artifacts" + + artifact_id: Mapped[str] = mapped_column(String(64), primary_key=True) + task_id: Mapped[str] = mapped_column(String(64), index=True) + node_id: Mapped[str | None] = mapped_column(String(64), nullable=True) + artifact_type: Mapped[str] = mapped_column(String(64)) + storage_uri: Mapped[str] = mapped_column(String(512)) + metadata_json: Mapped[dict] = mapped_column(JSON) + created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), index=True) + + def build_engine(database_url: str): if database_url.startswith("sqlite:///"): db_path = database_url.removeprefix("sqlite:///") @@ -416,6 +431,29 @@ def _record_to_tool_invocation(record: ToolInvocationDbRecord) -> ToolInvocation ) +def _artifact_to_record(record: ArtifactRecord) -> ArtifactDbRecord: + return ArtifactDbRecord( + artifact_id=record.artifact_id, + task_id=record.task_id, + node_id=record.node_id, + artifact_type=record.artifact_type, + storage_uri=record.storage_uri, + metadata_json=record.metadata, + created_at=datetime.now(UTC), + ) + + +def _record_to_artifact(record: ArtifactDbRecord) -> ArtifactRecord: + return ArtifactRecord( + artifact_id=record.artifact_id, + task_id=record.task_id, + node_id=record.node_id, + artifact_type=record.artifact_type, + storage_uri=record.storage_uri, + metadata=record.metadata_json or {}, + ) + + @dataclass(slots=True) class SqlAlchemyTaskRepository(TaskRepository): session_factory: sessionmaker[Session] @@ -595,6 +633,15 @@ class SqlAlchemyEventStore(EventStore): ).all() return [_record_to_event(record) for record in records] + def list_by_task_paginated(self, task_id: str, *, limit: int, offset: int) -> Page: + records = self.list_by_task(task_id) + return Page( + items=records[offset : offset + limit], + limit=limit, + offset=offset, + total=len(records), + ) + @dataclass(slots=True) class SqlAlchemyInvocationStore(InvocationStore): @@ -630,6 +677,53 @@ class SqlAlchemyInvocationStore(InvocationStore): ).all() return [_record_to_tool_invocation(record) for record in records] + def list_model_invocations_paginated(self, task_id: str, *, limit: int, offset: int) -> Page: + records = self.list_model_invocations(task_id) + return Page( + items=records[offset : offset + limit], + limit=limit, + offset=offset, + total=len(records), + ) + + def list_tool_invocations_paginated(self, task_id: str, *, limit: int, offset: int) -> Page: + records = self.list_tool_invocations(task_id) + return Page( + items=records[offset : offset + limit], + limit=limit, + offset=offset, + total=len(records), + ) + + +@dataclass(slots=True) +class SqlAlchemyArtifactStore(ArtifactStore): + session_factory: sessionmaker[Session] + + def save_artifact(self, record: ArtifactRecord) -> ArtifactRecord: + with self.session_factory() as session: + session.add(_artifact_to_record(record)) + session.commit() + return record + + def list_by_task(self, task_id: str) -> list[ArtifactRecord]: + with self.session_factory() as session: + records = session.scalars( + select(ArtifactDbRecord) + .where(ArtifactDbRecord.task_id == task_id) + .order_by(ArtifactDbRecord.created_at.asc()) + ).all() + return [_record_to_artifact(record) for record in records] + + def list_by_task_paginated(self, task_id: str, *, limit: int, offset: int) -> Page: + records = self.list_by_task(task_id) + return Page( + items=records[offset : offset + limit], + limit=limit, + offset=offset, + total=len(records), + ) + @dataclass(slots=True) class SqlAlchemyStorageBundle: @@ -641,6 +735,7 @@ class SqlAlchemyStorageBundle: worker_repository: SqlAlchemyWorkerRepository event_store: SqlAlchemyEventStore invocation_store: SqlAlchemyInvocationStore + artifact_store: SqlAlchemyArtifactStore def create_sqlalchemy_storage(database_url: str) -> SqlAlchemyStorageBundle: @@ -655,4 +750,5 @@ def create_sqlalchemy_storage(database_url: str) -> SqlAlchemyStorageBundle: worker_repository=SqlAlchemyWorkerRepository(session_factory), event_store=SqlAlchemyEventStore(session_factory), invocation_store=SqlAlchemyInvocationStore(session_factory), + artifact_store=SqlAlchemyArtifactStore(session_factory), ) diff --git a/tests/integration/test_storage_queries.py b/tests/integration/test_storage_queries.py new file mode 100644 index 0000000..5f2b062 --- /dev/null +++ b/tests/integration/test_storage_queries.py @@ -0,0 +1,86 @@ +from pathlib import Path + +from ai_orchestrator.application.ports import ( + ArtifactRecord, + ModelInvocationRecord, + ToolInvocationRecord, +) +from ai_orchestrator.config import AppSettings +from ai_orchestrator.domain.events import DomainEvent +from ai_orchestrator.infrastructure.storage.factory import create_storage_bundle + + +def test_sqlite_paginated_event_and_invocation_queries(tmp_path: Path) -> None: + db_path = tmp_path / "queries.sqlite3" + storage = create_storage_bundle( + AppSettings(storage_backend="sqlite", database_url=f"sqlite:///{db_path.as_posix()}") + ) + for index in range(5): + storage.event_store.append( + DomainEvent(event_type="test_event", task_id="task_1", payload={"index": index}) + ) + storage.invocation_store.save_model_invocation( + ModelInvocationRecord( + invocation_id=f"minv_{index}", + task_id="task_1", + node_id=None, + slot="weak", + provider="local", + model="mock", + status="success", + request={}, + response={"message": {"content": "ok"}}, + ) + ) + storage.invocation_store.save_tool_invocation( + ToolInvocationRecord( + invocation_id=f"tinv_{index}", + task_id="task_1", + node_id=None, + source_type="mcp", + source_id="one_c", + tool_name="one_c.run_sql", + status="success", + request={}, + response={"content": {"rows": []}}, + ) + ) + + events_page = storage.event_store.list_by_task_paginated("task_1", limit=2, offset=1) + model_page = storage.invocation_store.list_model_invocations_paginated( + "task_1", limit=2, offset=2 + ) + tool_page = storage.invocation_store.list_tool_invocations_paginated( + "task_1", limit=3, offset=1 + ) + + assert events_page.total == 5 + assert len(events_page.items) == 2 + assert model_page.total == 5 + assert len(model_page.items) == 2 + assert tool_page.total == 5 + assert len(tool_page.items) == 3 + + +def test_sqlite_artifact_store_persists_and_paginates(tmp_path: Path) -> None: + db_path = tmp_path / "artifacts.sqlite3" + storage = create_storage_bundle( + AppSettings(storage_backend="sqlite", database_url=f"sqlite:///{db_path.as_posix()}") + ) + for index in range(4): + storage.artifact_store.save_artifact( + ArtifactRecord( + artifact_id=f"art_{index}", + task_id="task_1", + node_id=None, + artifact_type="file", + storage_uri=f"file:///artifact/{index}", + metadata={"index": index}, + ) + ) + + page = storage.artifact_store.list_by_task_paginated("task_1", limit=2, offset=1) + + assert page.total == 4 + assert len(page.items) == 2 + assert page.items[0].storage_uri == "file:///artifact/1" diff --git a/tests/scenario/test_runtime_scenarios.py b/tests/scenario/test_runtime_scenarios.py index 0736591..140abd3 100644 --- a/tests/scenario/test_runtime_scenarios.py +++ b/tests/scenario/test_runtime_scenarios.py @@ -32,6 +32,7 @@ from ai_orchestrator.infrastructure.mcp_client import McpHttpClient, McpToolGate from ai_orchestrator.infrastructure.model_router import StaticMockModelProvider from ai_orchestrator.infrastructure.policy import StaticProjectPolicyEvaluator from ai_orchestrator.infrastructure.storage.memory import ( + InMemoryArtifactStore, InMemoryConfirmationRepository, InMemoryEventStore, InMemoryGraphRepository, @@ -90,11 +91,13 @@ def test_worker_gateway_result_becomes_artifact_ready_invocation() -> None: worker_repository = InMemoryWorkerRepository() event_store = InMemoryEventStore() invocation_store = InMemoryInvocationStore() + artifact_store = InMemoryArtifactStore() 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, + artifact_store=artifact_store, connection_manager=InMemoryWorkerConnectionManager(), ) response = worker_service.register( @@ -135,6 +138,7 @@ def test_worker_gateway_result_becomes_artifact_ready_invocation() -> None: invocation_store.list_tool_invocations("task_1")[-1].response["artifacts"][0]["path"] == "D:/artifact.txt" ) + assert artifact_store.list_by_task("task_1")[0].storage_uri == "D:/artifact.txt" def test_model_and_mcp_paths_can_run_in_same_runtime_context() -> None: @@ -206,6 +210,7 @@ def test_model_and_mcp_paths_can_run_in_same_runtime_context() -> None: project_configs={"default": project}, client=McpHttpClient(transport=httpx.MockTransport(handler)), invocation_store=invocation_store, + artifact_store=InMemoryArtifactStore(), ) model_result = router.run( diff --git a/tests/smoke/test_smoke_suite.py b/tests/smoke/test_smoke_suite.py index 0787b56..7bbb66b 100644 --- a/tests/smoke/test_smoke_suite.py +++ b/tests/smoke/test_smoke_suite.py @@ -1,3 +1,5 @@ +from datetime import UTC, datetime + import httpx from ai_orchestrator.application.services.execution import ( @@ -32,6 +34,7 @@ from ai_orchestrator.infrastructure.mcp_client import McpHttpClient, McpToolGate from ai_orchestrator.infrastructure.model_router import StaticMockModelProvider from ai_orchestrator.infrastructure.policy import StaticProjectPolicyEvaluator from ai_orchestrator.infrastructure.storage.memory import ( + InMemoryArtifactStore, InMemoryConfirmationRepository, InMemoryEventStore, InMemoryGraphRepository, @@ -223,6 +226,7 @@ def test_smoke_policy_confirm_and_full_auto() -> None: def test_smoke_mcp_and_worker_results_are_stored_and_finalizer_runs() -> None: invocation_store = InMemoryInvocationStore() + artifact_store = InMemoryArtifactStore() def handler(request: httpx.Request) -> httpx.Response: return httpx.Response( @@ -244,6 +248,7 @@ def test_smoke_mcp_and_worker_results_are_stored_and_finalizer_runs() -> None: project_configs={"default": _project_config()}, client=McpHttpClient(transport=httpx.MockTransport(handler)), invocation_store=invocation_store, + artifact_store=artifact_store, ) worker_repository = InMemoryWorkerRepository() event_store = InMemoryEventStore() @@ -252,6 +257,7 @@ def test_smoke_mcp_and_worker_results_are_stored_and_finalizer_runs() -> None: worker_repository=worker_repository, event_store=event_store, invocation_store=invocation_store, + artifact_store=artifact_store, connection_manager=InMemoryWorkerConnectionManager(), ) worker = worker_service.register( @@ -279,8 +285,8 @@ def test_smoke_mcp_and_worker_results_are_stored_and_finalizer_runs() -> None: 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), + started_at=datetime.now(UTC), + finished_at=datetime.now(UTC), duration_ms=1, result={"content": "ok"}, ), @@ -318,4 +324,5 @@ def test_smoke_mcp_and_worker_results_are_stored_and_finalizer_runs() -> None: assert queued.status == "queued" assert tool_result.status == "success" assert len(invocation_store.list_tool_invocations("task_1")) >= 2 + assert artifact_store.list_by_task("task_1") == [] assert graph.nodes[-1].output_data["status"] == "completed" diff --git a/tests/unit/test_mcp_client.py b/tests/unit/test_mcp_client.py index 6bf9e62..75c1bfa 100644 --- a/tests/unit/test_mcp_client.py +++ b/tests/unit/test_mcp_client.py @@ -10,7 +10,10 @@ from ai_orchestrator.config import ( ProjectMetadata, ) from ai_orchestrator.infrastructure.mcp_client import McpHttpClient, McpToolGateway -from ai_orchestrator.infrastructure.storage.memory import InMemoryInvocationStore +from ai_orchestrator.infrastructure.storage.memory import ( + InMemoryArtifactStore, + InMemoryInvocationStore, +) def test_mcp_tool_gateway_normalizes_result_and_persists_invocation() -> None: @@ -47,6 +50,7 @@ def test_mcp_tool_gateway_normalizes_result_and_persists_invocation() -> None: project_configs={"default": project}, client=McpHttpClient(transport=transport), invocation_store=invocation_store, + artifact_store=InMemoryArtifactStore(), ) result = gateway.call( diff --git a/tests/unit/test_reviewer.py b/tests/unit/test_reviewer.py new file mode 100644 index 0000000..1b8f736 --- /dev/null +++ b/tests/unit/test_reviewer.py @@ -0,0 +1,29 @@ +from ai_orchestrator.application.services.execution import StructuredReviewer +from ai_orchestrator.domain.enums import NodeType +from ai_orchestrator.domain.models import ExecutionNode, Task + + +def test_structured_reviewer_accepts_valid_output() -> None: + reviewer = StructuredReviewer() + task = Task(project_id="default", goal="Review", inputs={}) + node = ExecutionNode(task_id=task.task_id, node_type=NodeType.PLANNER, input_data={}) + + reviewed = reviewer.review( + task=task, + node=node, + output_data={"status": "planned", "goal": "Review"}, + ) + + assert reviewed["review"]["status"] == "accepted" + assert reviewed["task_id"] == task.task_id + + +def test_structured_reviewer_marks_missing_status_as_review_failed() -> None: + reviewer = StructuredReviewer() + task = Task(project_id="default", goal="Review", inputs={}) + node = ExecutionNode(task_id=task.task_id, node_type=NodeType.PLANNER, input_data={}) + + reviewed = reviewer.review(task=task, node=node, output_data={"goal": "Review"}) + + assert reviewed["status"] == "review_failed" + assert reviewed["review"]["status"] == "rejected" diff --git a/tests/unit/test_worker_gateway.py b/tests/unit/test_worker_gateway.py index 3b319de..6af6f89 100644 --- a/tests/unit/test_worker_gateway.py +++ b/tests/unit/test_worker_gateway.py @@ -10,6 +10,7 @@ from ai_orchestrator.application.services.workers import ( WorkerService, ) from ai_orchestrator.infrastructure.storage.memory import ( + InMemoryArtifactStore, InMemoryEventStore, InMemoryInvocationStore, InMemoryWorkerRepository, @@ -22,22 +23,25 @@ def _build_gateway() -> tuple[ InMemoryWorkerRepository, InMemoryEventStore, InMemoryInvocationStore, + InMemoryArtifactStore, ]: worker_repository = InMemoryWorkerRepository() event_store = InMemoryEventStore() invocation_store = InMemoryInvocationStore() + artifact_store = InMemoryArtifactStore() service = WorkerService(worker_repository=worker_repository, event_store=event_store) gateway = CapabilityAwareWorkerGateway( worker_repository=worker_repository, event_store=event_store, invocation_store=invocation_store, + artifact_store=artifact_store, connection_manager=InMemoryWorkerConnectionManager(), ) - return gateway, service, worker_repository, event_store, invocation_store + return gateway, service, worker_repository, event_store, invocation_store, artifact_store def test_worker_gateway_dispatches_and_completes_command() -> None: - gateway, service, worker_repository, event_store, invocation_store = _build_gateway() + gateway, service, worker_repository, event_store, invocation_store, _ = _build_gateway() response = service.register( RegisterWorkerRequest( worker_id="worker_home_pc", @@ -80,7 +84,7 @@ def test_worker_gateway_dispatches_and_completes_command() -> None: def test_worker_gateway_rejects_unsupported_capability() -> None: - gateway, service, _, _, _ = _build_gateway() + gateway, service, _, _, _, _ = _build_gateway() response = service.register( RegisterWorkerRequest( worker_id="worker_home_pc", @@ -103,7 +107,7 @@ def test_worker_gateway_rejects_unsupported_capability() -> None: def test_worker_service_heartbeat_and_stale_tracking() -> None: - _, service, worker_repository, event_store, _ = _build_gateway() + _, service, worker_repository, event_store, _, _ = _build_gateway() response = service.register( RegisterWorkerRequest( worker_id="worker_home_pc",