From ce2b263e52415af8b001d54bd04030d4495db4ed Mon Sep 17 00:00:00 2001 From: Mikhail Date: Fri, 3 Jul 2026 21:09:47 +0300 Subject: [PATCH] Add architecture package and application skeleton --- .gitignore | 22 ++ README.md | 109 +++++----- docs/adr/001-runtime-stack.md | 65 ++++++ docs/adr/002-storage-and-persistence.md | 46 +++++ docs/adr/003-event-stream-transport.md | 40 ++++ docs/adr/004-worker-transport.md | 33 +++ docs/adr/005-policy-evaluation-model.md | 55 +++++ docs/adr/006-event-taxonomy.md | 42 ++++ docs/adr/007-model-router-fallback.md | 32 +++ docs/architecture/04_domain_model.md | 112 ++++++++++ docs/architecture/05_state_machines.md | 66 ++++++ docs/architecture/06_internal_interfaces.md | 57 ++++++ docs/architecture/07_error_model.md | 34 ++++ docs/architecture/08_execution_model.md | 74 +++++++ docs/backlog/01_execution_backlog.md | 60 ++++++ docs/contracts/API_CONTRACTS.md | 22 +- docs/storage/01_schema.md | 158 ++++++++++++++ docs/testing/01_strategy.md | 36 ++++ pyproject.toml | 53 +++++ src/ai_orchestrator/__init__.py | 2 + src/ai_orchestrator/application/__init__.py | 2 + src/ai_orchestrator/application/ports.py | 102 ++++++++++ .../application/services/__init__.py | 2 + .../application/services/orchestrator.py | 134 ++++++++++++ .../application/services/workers.py | 47 +++++ src/ai_orchestrator/config.py | 82 ++++++++ src/ai_orchestrator/delivery/__init__.py | 2 + src/ai_orchestrator/delivery/http/__init__.py | 2 + src/ai_orchestrator/delivery/http/app.py | 190 +++++++++++++++++ src/ai_orchestrator/delivery/http/schemas.py | 66 ++++++ src/ai_orchestrator/domain/__init__.py | 2 + src/ai_orchestrator/domain/enums.py | 86 ++++++++ src/ai_orchestrator/domain/errors.py | 30 +++ src/ai_orchestrator/domain/events.py | 20 ++ src/ai_orchestrator/domain/models.py | 192 ++++++++++++++++++ .../infrastructure/__init__.py | 2 + .../infrastructure/config_loader.py | 15 ++ src/ai_orchestrator/infrastructure/logging.py | 14 ++ src/ai_orchestrator/infrastructure/policy.py | 34 ++++ .../infrastructure/storage/__init__.py | 2 + .../infrastructure/storage/memory.py | 85 ++++++++ src/ai_orchestrator/main.py | 6 + tests/integration/test_http_app.py | 53 +++++ tests/unit/test_graph.py | 21 ++ tests/unit/test_policy.py | 80 ++++++++ tests/unit/test_task_lifecycle.py | 15 ++ tests/unit/test_worker_service.py | 26 +++ 47 files changed, 2377 insertions(+), 53 deletions(-) create mode 100644 .gitignore create mode 100644 docs/adr/001-runtime-stack.md create mode 100644 docs/adr/002-storage-and-persistence.md create mode 100644 docs/adr/003-event-stream-transport.md create mode 100644 docs/adr/004-worker-transport.md create mode 100644 docs/adr/005-policy-evaluation-model.md create mode 100644 docs/adr/006-event-taxonomy.md create mode 100644 docs/adr/007-model-router-fallback.md create mode 100644 docs/architecture/04_domain_model.md create mode 100644 docs/architecture/05_state_machines.md create mode 100644 docs/architecture/06_internal_interfaces.md create mode 100644 docs/architecture/07_error_model.md create mode 100644 docs/architecture/08_execution_model.md create mode 100644 docs/backlog/01_execution_backlog.md create mode 100644 docs/storage/01_schema.md create mode 100644 docs/testing/01_strategy.md create mode 100644 pyproject.toml create mode 100644 src/ai_orchestrator/__init__.py create mode 100644 src/ai_orchestrator/application/__init__.py create mode 100644 src/ai_orchestrator/application/ports.py create mode 100644 src/ai_orchestrator/application/services/__init__.py create mode 100644 src/ai_orchestrator/application/services/orchestrator.py create mode 100644 src/ai_orchestrator/application/services/workers.py create mode 100644 src/ai_orchestrator/config.py create mode 100644 src/ai_orchestrator/delivery/__init__.py create mode 100644 src/ai_orchestrator/delivery/http/__init__.py create mode 100644 src/ai_orchestrator/delivery/http/app.py create mode 100644 src/ai_orchestrator/delivery/http/schemas.py create mode 100644 src/ai_orchestrator/domain/__init__.py create mode 100644 src/ai_orchestrator/domain/enums.py create mode 100644 src/ai_orchestrator/domain/errors.py create mode 100644 src/ai_orchestrator/domain/events.py create mode 100644 src/ai_orchestrator/domain/models.py create mode 100644 src/ai_orchestrator/infrastructure/__init__.py create mode 100644 src/ai_orchestrator/infrastructure/config_loader.py create mode 100644 src/ai_orchestrator/infrastructure/logging.py create mode 100644 src/ai_orchestrator/infrastructure/policy.py create mode 100644 src/ai_orchestrator/infrastructure/storage/__init__.py create mode 100644 src/ai_orchestrator/infrastructure/storage/memory.py create mode 100644 src/ai_orchestrator/main.py create mode 100644 tests/integration/test_http_app.py create mode 100644 tests/unit/test_graph.py create mode 100644 tests/unit/test_policy.py create mode 100644 tests/unit/test_task_lifecycle.py create mode 100644 tests/unit/test_worker_service.py diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..c081c78 --- /dev/null +++ b/.gitignore @@ -0,0 +1,22 @@ +.pytest_cache/ +.ruff_cache/ +.mypy_cache/ +.venv/ +__pycache__/ +*.pyc +*.pyo +*.pyd +.coverage +htmlcov/ +dist/ +build/ +site/ +*.egg-info/ +.idea/ +.vscode/ +.DS_Store +Thumbs.db +*.db +*.sqlite3 +runtime/ +artifacts/ diff --git a/README.md b/README.md index bc47a15..4f3ef26 100644 --- a/README.md +++ b/README.md @@ -1,68 +1,73 @@ # AI Orchestrator -Отдельный проект для серверного AI-orchestrator, который управляет задачами, моделями, MCP-tools и local workers. +Серверный runtime для AI-задач, моделей, MCP-tools и local workers. +Проект отделен от `Local LLM Platform` и строится как самостоятельное orchestration-ядро. -## Главная идея - -Этот проект не является частью `Local LLM Platform`. - -`Local LLM Platform` отвечает за: -- локальные модели; -- inference; -- registry; -- GPU profiles; -- model deployment; -- smoke tests моделей. +## Цели `AI Orchestrator` отвечает за: -- общение с пользователем; -- построение плана/графа выполнения; -- вызов моделей; -- вызов MCP-инструментов; -- работу с local worker; -- подтверждения действий; -- логи; + +- прием пользовательских задач; +- построение execution graph; +- вызов моделей через model router; +- вызов MCP tools; +- dispatch команд в local worker; +- policy-driven подтверждения; +- event log, audit и progress stream; - fallback между weak/strong/vision моделями. -## Ключевой принцип +## Принципы -В коде нет жестких запретов на действия. +- никаких hardcoded bans; +- все ограничения идут через project policy; +- worker не планирует задачу и не принимает интеллектуальных решений; +- 1С не попадает в ядро и подключается только через MCP; +- локальный Docker не используется как обязательный dev path. -Система должна работать через настраиваемую политику проекта: +## Архитектура -- `full_auto` — выполнять без подтверждений; -- `confirm` — спрашивать подтверждение для действий, отмеченных политикой; -- `manual` — всегда показывать план и ждать запуска; -- `disabled` — компонент отключен настройкой. +Проект разбит на 4 слоя: -Важно: `disabled` — это не hardcoded ban, а выбранная настройка проекта. +- `domain` — сущности, state machines, policy/value objects, события; +- `application` — use cases и orchestration services; +- `infrastructure` — storage, config, policy adapter, provider gateways; +- `delivery` — HTTP API, SSE, worker transport adapters. -## Базовая схема +Подробные решения зафиксированы в: -```text -User UI / API - ↓ -Conversation Manager - ↓ -Server Orchestrator - ↓ -Execution Graph - ├─ Planner - ├─ Workers - ├─ Reviewer - └─ Finalizer - ↓ -Model Router - ├─ weak model: local - ├─ strong model: local or external - └─ vision model: local or external - ↓ -Tools - ├─ MCP client - ├─ remote MCP servers - └─ local worker protocol +- [ADR](docs/adr/001-runtime-stack.md) +- [Domain model](docs/architecture/04_domain_model.md) +- [State machines](docs/architecture/05_state_machines.md) +- [Storage schema](docs/storage/01_schema.md) + +## Текущий статус + +В репозитории уже собраны: + +- архитектурный пакет с ADR; +- доменная модель и state machines; +- draft persistent storage schema; +- внутренние порты orchestrator; +- кодовый каркас `src/ai_orchestrator`; +- базовый FastAPI delivery layer; +- in-memory adapters для безопасного старта; +- unit tests на базовые состояния и policy. + +## Быстрый запуск + +```bash +python -m pip install -e .[dev] +python -m uvicorn ai_orchestrator.main:app --host 127.0.0.1 --port 8080 ``` -## Что должен сделать Codex +## Структура -Начать с документа `docs/codex/CODEX_TASK.md`. +```text +src/ai_orchestrator/ + domain/ + application/ + infrastructure/ + delivery/ +``` + +Исходные папки `src/orchestrator`, `src/policy` и другие пока остаются как legacy-концептуальные маркеры из исходного blueprint. diff --git a/docs/adr/001-runtime-stack.md b/docs/adr/001-runtime-stack.md new file mode 100644 index 0000000..b41877c --- /dev/null +++ b/docs/adr/001-runtime-stack.md @@ -0,0 +1,65 @@ +# ADR 001: Runtime Stack + +## Status + +Accepted + +## Context + +AI Orchestrator должен быть: + +- независимым от Local LLM Platform; +- пригодным для долгоживущих task runtime; +- удобным для строгой типизации контрактов; +- удобным для async I/O: models, MCP, worker connections, event streams; +- запускаемым без Docker на локальной машине; +- переносимым на удаленное окружение, включая `docker-test`. + +## Decision + +Основной стек: + +- язык: Python 3.13+; +- HTTP API: FastAPI; +- конфиги и схемы запроса/ответа: Pydantic v2; +- app settings: `pydantic-settings`; +- конфиги проекта: YAML; +- event streaming для UI: SSE как основной transport; +- worker transport: WebSocket; +- structured logging: JSON lines через стандартный logging layer и event envelopes; +- тесты: `pytest` + `pytest-asyncio`; +- runtime packaging: стандартный `pyproject.toml`. + +## Why This Stack + +Python хорошо подходит для orchestration-heavy систем, где важнее: + +- строгое моделирование состояния; +- быстрая интеграция с внешними AI/MCP endpoint; +- асинхронный I/O; +- прозрачные схемы данных; +- простая локальная разработка без контейнеров. + +FastAPI выбран как delivery adapter, а не как центр архитектуры. +Доменные модели и application services не должны зависеть от FastAPI. + +## Consequences + +Плюсы: + +- быстрый путь к строгим контрактам; +- хорошая ergonomics для async adapters; +- удобный локальный запуск без Docker; +- понятный переход к PostgreSQL и production deployment. + +Минусы: + +- нужен дисциплинированный layering, чтобы не “утонуть” в framework-driven code; +- CPU-heavy задачи должны оставаться вне основного request loop. + +## Non-Goals + +- локальный Docker как обязательный dev path; +- тяжелая зависимость от конкретной ORM на уровне domain; +- framework-first архитектура. + diff --git a/docs/adr/002-storage-and-persistence.md b/docs/adr/002-storage-and-persistence.md new file mode 100644 index 0000000..a210ee3 --- /dev/null +++ b/docs/adr/002-storage-and-persistence.md @@ -0,0 +1,46 @@ +# ADR 002: Storage And Persistence Model + +## Status + +Accepted + +## Context + +Проект не должен начинаться с чисто временного `in-memory` мышления, иначе потом придется переделывать: + +- task lifecycle; +- event log; +- confirmations; +- worker session tracking; +- replay и audit. + +## Decision + +Хранилище проектируется как persistent-first: + +- production target: PostgreSQL; +- local development fallback: SQLite без Docker; +- repository interfaces определяются в application/domain boundary; +- materialized current state хранится отдельно от append-only event log; +- критические переходы состояний должны логироваться как события. + +## Data Model Principles + +1. `tasks`, `task_nodes`, `confirmations`, `worker_sessions` хранят текущее состояние. +2. `task_events` хранит audit trail и feed для replay/debugging. +3. Вызовы моделей, MCP и worker сохраняются как отдельные invocation records. +4. Артефакты хранят metadata в БД, а payload может лежать во внешнем файловом хранилище. + +## Consequences + +Плюсы: + +- можно строить UI progress и audit без догадок; +- проще реализовать replay/resume; +- нет боли миграции с “простых dict” на реальную БД. + +Минусы: + +- немного более сложный старт; +- нужно заранее аккуратно продумать схему. + diff --git a/docs/adr/003-event-stream-transport.md b/docs/adr/003-event-stream-transport.md new file mode 100644 index 0000000..c276935 --- /dev/null +++ b/docs/adr/003-event-stream-transport.md @@ -0,0 +1,40 @@ +# ADR 003: User Event Stream Transport + +## Status + +Accepted + +## Context + +UI требует поток событий по task execution: + +- `task_started` +- `node_started` +- `node_completed` +- `confirmation_required` +- `task_completed` +- `task_failed` + +В исходных документах упомянуты `SSE/WebSocket`, но как основной transport нужно выбрать один. + +## Decision + +Основной transport для пользовательского event stream: Server-Sent Events. + +WebSocket не исключается, но считается вторичным adapter для будущих realtime-потребностей. + +## Rationale + +SSE проще для: + +- однонаправленного server-to-client прогресса; +- проксирования; +- reconnect semantics; +- дебага и совместимости с обычными HTTP-инструментами. + +WebSocket нужен не UI в первую очередь, а local worker transport. + +## Consequences + +`GET /tasks/{id}/events` фиксируется как SSE endpoint. + diff --git a/docs/adr/004-worker-transport.md b/docs/adr/004-worker-transport.md new file mode 100644 index 0000000..f9d9969 --- /dev/null +++ b/docs/adr/004-worker-transport.md @@ -0,0 +1,33 @@ +# ADR 004: Local Worker Transport + +## Status + +Accepted + +## Context + +Local Worker работает на машине пользователя и должен: + +- сам инициировать соединение; +- не требовать входящего порта на пользовательском ПК; +- поддерживать прогресс, heartbeat и dispatch commands. + +## Decision + +Основной transport для worker gateway: outbound WebSocket session от worker к серверу. + +## Required Semantics + +- registration handshake; +- capability advertisement; +- heartbeat; +- command dispatch; +- progress events; +- result envelope; +- cancellation; +- reconnect with new session id. + +## Consequences + +Worker transport остается отдельным delivery/infrastructure adapter и не врастает в orchestration core. + diff --git a/docs/adr/005-policy-evaluation-model.md b/docs/adr/005-policy-evaluation-model.md new file mode 100644 index 0000000..7dbf187 --- /dev/null +++ b/docs/adr/005-policy-evaluation-model.md @@ -0,0 +1,55 @@ +# ADR 005: Policy Evaluation Model + +## Status + +Accepted + +## Context + +Ключевой принцип проекта: без hardcoded bans. +Система должна принимать решения через policy engine, а не через встроенные запреты по имени tool или по содержимому команды. + +## Decision + +Policy engine работает в 4 шага: + +1. request/action нормализуется в `ActionDescriptor`; +2. descriptor классифицируется по `resource` и `risk_level`; +3. policy evaluator находит эффективное правило; +4. возвращается `PolicyDecision`. + +## Core Types + +`resource`: + +- `filesystem` +- `shell` +- `sql` +- `mcp` +- `external_models` +- `desktop` +- `browser` +- `network` +- `cost` +- `system` + +`risk_level`: + +- `safe` +- `write` +- `destructive` +- `system` +- `cost` + +`decision`: + +- `allow` +- `confirm` +- `manual` +- `disabled_by_config` + +## Important Rule + +Policy engine не получает “tool name only” как источник истины. +Он должен опираться на action classification и metadata от tool/model/worker adapters. + diff --git a/docs/adr/006-event-taxonomy.md b/docs/adr/006-event-taxonomy.md new file mode 100644 index 0000000..be711a8 --- /dev/null +++ b/docs/adr/006-event-taxonomy.md @@ -0,0 +1,42 @@ +# ADR 006: Event Taxonomy And Correlation + +## Status + +Accepted + +## Context + +Оркестратор должен быть отлаживаемым и воспроизводимым. +Для этого все значимые переходы должны иметь единый event model. + +## Decision + +Все события включают: + +- `event_id` +- `event_type` +- `occurred_at` +- `task_id` +- `conversation_id` when available +- `node_id` when available +- `correlation_id` +- `causation_id` +- `payload` + +## Event Families + +1. Task lifecycle +2. Node lifecycle +3. Policy decisions +4. Confirmation lifecycle +5. Model invocation lifecycle +6. Tool invocation lifecycle +7. Worker session lifecycle +8. Worker command lifecycle +9. Finalizer/result lifecycle +10. System warnings/errors + +## Consequences + +UI, audit log, replay tooling и debugging должны читать одну и ту же taxonomy, а не разрозненные лог-сообщения. + diff --git a/docs/adr/007-model-router-fallback.md b/docs/adr/007-model-router-fallback.md new file mode 100644 index 0000000..27883a9 --- /dev/null +++ b/docs/adr/007-model-router-fallback.md @@ -0,0 +1,32 @@ +# ADR 007: Model Router Fallback Semantics + +## Status + +Accepted + +## Context + +В проекте уже зафиксирована идея: + +- сначала weak; +- затем retry weak; +- затем strong fallback; +- external models не запрещаются кодом, а регулируются конфигом и policy. + +## Decision + +Fallback pipeline: + +1. slot resolution выбирает provider и model для requested slot; +2. weak model вызывается первой; +3. если ответ невалиден по output contract, выполняется ограниченный retry weak; +4. если weak исчерпан, router проверяет доступность strong; +5. если strong разрешен config+policy, выполняется strong fallback; +6. если strong недоступен, возвращается structured degraded result. + +## Notes + +- “невалиден” означает нарушение response schema, parser failure или явный quality rejection; +- fallback trail сохраняется в invocation log; +- внешняя strong model не вызывается, если `allow_external_models=false` или policy запрещает. + diff --git a/docs/architecture/04_domain_model.md b/docs/architecture/04_domain_model.md new file mode 100644 index 0000000..a1bae6c --- /dev/null +++ b/docs/architecture/04_domain_model.md @@ -0,0 +1,112 @@ +# Доменная модель + +## Слои + +### Domain + +Чистые сущности, value objects, state machines, классификация рисков, события. + +### Application + +Use cases: + +- create_task +- submit_chat_message +- plan_task +- execute_ready_nodes +- request_confirmation +- approve_confirmation +- reject_confirmation +- register_worker +- ingest_worker_result + +### Infrastructure + +- storage adapters +- config loader +- event publisher +- model providers +- MCP transport +- worker gateway +- artifact persistence + +### Delivery + +- HTTP API +- SSE task events +- worker WebSocket session endpoint + +## Главные сущности + +### ProjectConfig + +Описывает: + +- model slots +- policy modes +- execution limits +- enabled MCP servers +- worker access rules + +### Conversation + +Контейнер пользовательского диалога и связанных задач. + +### Task + +Единица выполнения пользовательской цели. + +Поля: + +- `id` +- `project_id` +- `conversation_id` +- `goal` +- `inputs` +- `status` +- `requested_mode` +- `effective_mode` +- `created_at` +- `updated_at` +- `result_summary` + +### ExecutionGraph + +Набор node definitions и execution dependencies для конкретной task. + +### Node + +Типы: + +- `planner` +- `model_call` +- `tool_call` +- `local_worker_call` +- `reviewer` +- `finalizer` +- `confirmation` + +### ConfirmationRequest + +Отдельная сущность для действий, требующих подтверждения. + +### WorkerSession + +Онлайн-сессия thin local worker. + +### ModelInvocation + +Record вызова model router/provider. + +### ToolInvocation + +Record вызова MCP tool или worker command. + +### Artifact + +Материализованный результат: patch, file diff, report, log bundle. + +### Event + +Audit и UI-событие в единой taxonomy. + diff --git a/docs/architecture/05_state_machines.md b/docs/architecture/05_state_machines.md new file mode 100644 index 0000000..e3c132d --- /dev/null +++ b/docs/architecture/05_state_machines.md @@ -0,0 +1,66 @@ +# State Machines + +## TaskStatus + +```text +created + -> planned + -> running + -> waiting_confirmation + -> waiting_manual + -> completed + -> failed + -> cancelled +``` + +Правила: + +- `created -> planned`: planner graph создан; +- `planned -> running`: есть готовые к исполнению nodes; +- `running -> waiting_confirmation`: execution остановлен на pending approval; +- `running -> waiting_manual`: policy/manual gate остановил автостарт; +- `running -> completed`: finalizer сформировал итог; +- `running -> failed`: unrecoverable error; +- `waiting_confirmation -> running`: approval получен; +- `waiting_manual -> running`: пользователь вручную продолжил выполнение. + +## NodeStatus + +```text +pending + -> ready + -> running + -> waiting_confirmation + -> completed + -> failed + -> skipped + -> cancelled +``` + +## ConfirmationStatus + +```text +pending + -> approved + -> rejected + -> expired + -> cancelled +``` + +## WorkerSessionStatus + +```text +connecting + -> online + -> busy + -> stale + -> disconnected +``` + +## Retry Principles + +- retry выполняется только для явно retryable failures; +- retry history фиксируется в node attempts и event log; +- confirmation-required state не считается failure; +- idempotency must be explicit for side-effecting nodes. + diff --git a/docs/architecture/06_internal_interfaces.md b/docs/architecture/06_internal_interfaces.md new file mode 100644 index 0000000..4eb5372 --- /dev/null +++ b/docs/architecture/06_internal_interfaces.md @@ -0,0 +1,57 @@ +# Internal Interfaces + +## Orchestrator Application Ports + +### TaskRepository + +- create task +- update task +- get task by id +- list active tasks + +### GraphRepository + +- save graph +- get graph +- update node state +- list ready nodes + +### ConfirmationRepository + +- create confirmation +- get confirmation +- resolve confirmation + +### WorkerRepository + +- register session +- update heartbeat +- assign command +- complete command + +### EventStore + +- append event +- list events by task + +### PolicyEvaluator + +- evaluate action descriptor against project policy + +### ModelRouter + +- execute model request and return structured invocation result + +### ToolGateway + +- call MCP tool and return normalized result + +### WorkerGateway + +- dispatch command to worker and receive normalized result + +## Design Rule + +Application services зависят только от этих портов. +Ни один use case не должен напрямую импортировать HTTP handlers, DB session objects или конкретный MCP/WebSocket client. + diff --git a/docs/architecture/07_error_model.md b/docs/architecture/07_error_model.md new file mode 100644 index 0000000..6275e12 --- /dev/null +++ b/docs/architecture/07_error_model.md @@ -0,0 +1,34 @@ +# Error Model + +## Error Families + +### ValidationError + +Некорректный вход, schema mismatch, invalid command envelope. + +### PolicyError + +Действие запрещено текущей конфигурацией или требует manual/confirm остановки. + +### RetryableInfrastructureError + +Временная ошибка транспорта, timeout, transient upstream failure. + +### NonRetryableInfrastructureError + +Постоянная ошибка конфигурации, unsupported capability, broken contract. + +### ExecutionError + +Ошибка бизнес-исполнения конкретного node. + +### CancellationError + +Task или node остановлены по explicit cancel. + +## User-Facing Behavior + +- validation и policy ошибки должны быть ясными и краткими; +- infra ошибки должны сохранять technical details в logs/events; +- task итог должен различать `failed` и `waiting for action`. + diff --git a/docs/architecture/08_execution_model.md b/docs/architecture/08_execution_model.md new file mode 100644 index 0000000..c507579 --- /dev/null +++ b/docs/architecture/08_execution_model.md @@ -0,0 +1,74 @@ +# Execution Model + +## Базовый цикл + +1. Пользователь создает `Task`. +2. Planner materializes `ExecutionGraph`. +3. Graph engine отмечает `ready` узлы. +4. Orchestrator выбирает runner для каждого ready node. +5. Перед side-effecting action выполняется policy evaluation. +6. Если нужен `confirm`, создается `ConfirmationRequest` и graph останавливается. +7. Если action разрешен, runner выполняет node. +8. Node результат сохраняется как output + event trail + optional artifacts. +9. Reviewer может валидировать промежуточный результат. +10. Finalizer формирует user-facing result и cards. + +## Node Families + +### Planner + +Создает или пересобирает graph. Не выполняет side effects. + +### Model Call + +Использует `ModelRouter`. +Может инициировать fallback, но только внутри router contract. + +### Tool Call + +Использует `ToolGateway` для MCP. + +### Local Worker Call + +Использует `WorkerGateway` для thin executor. + +### Reviewer + +Системный шаг quality gate: + +- schema validation +- semantic validation +- retry recommendation +- escalate to strong model + +### Finalizer + +Преобразует execution result в user result: + +- message text +- cards +- artifacts metadata +- summary + +## Execution Rules + +- workers не создают новые workers; +- model/router adapters не меняют graph напрямую; +- только orchestrator меняет task/node state; +- side effects всегда проходят через policy decision; +- event log пишется на каждом значимом переходе. + +## Manual And Resume + +`manual` означает, что orchestrator не стартует следующий шаг автоматически. +Task остается в `waiting_manual` до явного resume action. + +## Replay + +Replay не означает повтор side effects автоматически. +Replay должен уметь: + +- восстанавливать graph state из persisted state; +- переиздавать progress view; +- запускать retry только для разрешенных retryable nodes. + diff --git a/docs/backlog/01_execution_backlog.md b/docs/backlog/01_execution_backlog.md new file mode 100644 index 0000000..578bcab --- /dev/null +++ b/docs/backlog/01_execution_backlog.md @@ -0,0 +1,60 @@ +# Execution Backlog + +## Package A: Architecture Freeze + +- [x] runtime ADR +- [x] storage ADR +- [x] event stream ADR +- [x] worker transport ADR +- [x] policy evaluation ADR +- [x] event taxonomy ADR +- [x] model fallback ADR +- [x] domain model +- [x] state machines +- [x] storage schema draft +- [x] internal interfaces +- [x] error model +- [x] execution model + +## Package B: Core Hardening + +- [ ] graph scheduler service +- [ ] node runner abstraction +- [ ] reviewer contract +- [ ] finalizer contract +- [ ] resume/cancel use cases +- [ ] idempotency policy +- [ ] retry policy matrix + +## Package C: Persistence + +- [ ] SQLAlchemy models +- [ ] repositories for PostgreSQL/SQLite +- [ ] migrations +- [ ] artifact metadata persistence +- [ ] event store queries with pagination + +## Package D: Integrations + +- [ ] model router providers +- [ ] MCP transport adapter +- [ ] worker WebSocket gateway +- [ ] heartbeat monitor +- [ ] capability-aware dispatch + +## Package E: Delivery + +- [ ] typed worker registration schema +- [ ] task event DTO normalization +- [ ] confirmation cards +- [ ] progress cards +- [ ] artifact cards +- [ ] error cards + +## Package F: Validation + +- [ ] contract tests +- [ ] scenario tests +- [ ] smoke suite +- [ ] static checks in CI + diff --git a/docs/contracts/API_CONTRACTS.md b/docs/contracts/API_CONTRACTS.md index 082db51..1991f8d 100644 --- a/docs/contracts/API_CONTRACTS.md +++ b/docs/contracts/API_CONTRACTS.md @@ -59,7 +59,9 @@ ## GET /tasks/{task_id}/events -SSE/WebSocket stream событий: +Primary transport: SSE stream событий. + +WebSocket может быть добавлен как secondary adapter, но contract-first путь для UI: SSE. ```json { @@ -87,6 +89,24 @@ SSE/WebSocket stream событий: } ``` +## GET /workers + +Список известных worker sessions. + +```json +{ + "workers": [ + { + "session_id": "wrk_1", + "worker_id": "worker_home_pc", + "name": "Home PC", + "status": "online", + "capabilities": ["file.read", "command.run"] + } + ] +} +``` + ## POST /workers/register Local worker регистрируется на сервере. diff --git a/docs/storage/01_schema.md b/docs/storage/01_schema.md new file mode 100644 index 0000000..c1eee2a --- /dev/null +++ b/docs/storage/01_schema.md @@ -0,0 +1,158 @@ +# Persistent Storage Schema Draft + +## Database Target + +- production: PostgreSQL +- local development: SQLite + +## Tables + +### projects + +- `project_id` pk +- `name` +- `config_blob` +- `created_at` +- `updated_at` + +### conversations + +- `conversation_id` pk +- `project_id` +- `title` +- `created_at` +- `updated_at` + +### messages + +- `message_id` pk +- `conversation_id` +- `role` +- `content` +- `attachments_json` +- `created_at` + +### tasks + +- `task_id` pk +- `project_id` +- `conversation_id` +- `goal` +- `inputs_json` +- `status` +- `requested_mode` +- `effective_mode` +- `current_node_id` +- `result_summary_json` +- `created_at` +- `updated_at` + +### task_graphs + +- `task_id` pk +- `graph_version` +- `graph_metadata_json` +- `created_at` + +### task_nodes + +- `node_id` pk +- `task_id` +- `node_type` +- `status` +- `input_json` +- `output_json` +- `dependencies_json` +- `assigned_runner` +- `attempts` +- `retryable` +- `timeout_ms` +- `created_at` +- `updated_at` + +### task_events + +- `event_id` pk +- `task_id` +- `conversation_id` +- `node_id` +- `event_type` +- `correlation_id` +- `causation_id` +- `payload_json` +- `occurred_at` + +### confirmations + +- `confirmation_id` pk +- `task_id` +- `node_id` +- `status` +- `scope` +- `decision_json` +- `preview_json` +- `comment` +- `created_at` +- `resolved_at` + +### worker_sessions + +- `session_id` pk +- `worker_id` +- `name` +- `machine` +- `os` +- `version` +- `status` +- `capabilities_json` +- `last_heartbeat_at` +- `current_task_id` +- `connected_at` +- `disconnected_at` + +### model_invocations + +- `invocation_id` pk +- `task_id` +- `node_id` +- `slot` +- `provider` +- `model` +- `status` +- `request_json` +- `response_json` +- `usage_json` +- `fallback_from_invocation_id` +- `created_at` + +### tool_invocations + +- `invocation_id` pk +- `task_id` +- `node_id` +- `source_type` +- `source_id` +- `tool_name` +- `status` +- `request_json` +- `response_json` +- `created_at` + +### artifacts + +- `artifact_id` pk +- `task_id` +- `node_id` +- `artifact_type` +- `storage_uri` +- `metadata_json` +- `created_at` + +## Required Indexes + +- `tasks(status, updated_at)` +- `task_nodes(task_id, status)` +- `task_events(task_id, occurred_at)` +- `confirmations(status, created_at)` +- `worker_sessions(status, last_heartbeat_at)` + diff --git a/docs/testing/01_strategy.md b/docs/testing/01_strategy.md new file mode 100644 index 0000000..d181364 --- /dev/null +++ b/docs/testing/01_strategy.md @@ -0,0 +1,36 @@ +# Testing Strategy + +## Layers + +### Unit Tests + +- state transitions +- policy evaluation +- graph readiness +- fallback logic + +### Contract Tests + +- API schemas +- worker envelopes +- MCP result normalization +- model router output contracts + +### Integration Tests + +- repository implementations +- event streaming +- config loader +- storage adapters + +### Scenario Tests + +- create task -> plan -> execute -> finalize +- confirm flow +- weak retry then strong fallback +- worker result ingestion + +### Smoke Tests + +Обязательный список зафиксирован в [tests/smoke/SMOKE_TESTS.md](/Z:/codex/ai_orchestrator/tests/smoke/SMOKE_TESTS.md). + diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..73603dd --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,53 @@ +[build-system] +requires = ["setuptools>=69", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "ai-orchestrator" +version = "0.1.0" +description = "Server-side AI orchestrator/runtime for tasks, models, MCP tools, and local workers." +readme = "README.md" +requires-python = ">=3.13" +license = { text = "Proprietary" } +authors = [ + { name = "Mikhail" } +] +dependencies = [ + "fastapi>=0.115.0,<1.0.0", + "pydantic>=2.9.0,<3.0.0", + "pydantic-settings>=2.5.0,<3.0.0", + "pyyaml>=6.0.2,<7.0.0", + "sse-starlette>=2.1.3,<3.0.0", + "uvicorn>=0.32.0,<1.0.0", +] + +[project.optional-dependencies] +dev = [ + "pytest>=8.3.0,<9.0.0", + "pytest-asyncio>=0.24.0,<1.0.0", + "httpx>=0.27.0,<1.0.0", + "ruff>=0.6.0,<1.0.0", +] +db = [ + "sqlalchemy>=2.0.36,<3.0.0", + "alembic>=1.13.0,<2.0.0", + "asyncpg>=0.30.0,<1.0.0", +] + +[tool.setuptools] +package-dir = { "" = "src" } + +[tool.setuptools.packages.find] +where = ["src"] + +[tool.pytest.ini_options] +pythonpath = ["src"] +testpaths = ["tests"] +asyncio_default_fixture_loop_scope = "function" + +[tool.ruff] +line-length = 100 +target-version = "py313" + +[tool.ruff.lint] +select = ["E", "F", "I", "N", "UP", "B", "SIM"] diff --git a/src/ai_orchestrator/__init__.py b/src/ai_orchestrator/__init__.py new file mode 100644 index 0000000..68efff4 --- /dev/null +++ b/src/ai_orchestrator/__init__.py @@ -0,0 +1,2 @@ +"""AI Orchestrator package.""" + diff --git a/src/ai_orchestrator/application/__init__.py b/src/ai_orchestrator/application/__init__.py new file mode 100644 index 0000000..1063efd --- /dev/null +++ b/src/ai_orchestrator/application/__init__.py @@ -0,0 +1,2 @@ +"""Application layer.""" + diff --git a/src/ai_orchestrator/application/ports.py b/src/ai_orchestrator/application/ports.py new file mode 100644 index 0000000..da944a1 --- /dev/null +++ b/src/ai_orchestrator/application/ports.py @@ -0,0 +1,102 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, Protocol + +from ai_orchestrator.domain.events import DomainEvent +from ai_orchestrator.domain.models import ( + ActionDescriptor, + ConfirmationRequest, + ExecutionGraph, + PolicyDecision, + Task, + WorkerSession, +) + + +@dataclass(slots=True) +class ModelInvocationResult: + status: str + provider: str + model: str + message: dict[str, Any] + tool_calls: list[dict[str, Any]] = field(default_factory=list) + usage: dict[str, Any] = field(default_factory=dict) + error: dict[str, Any] | None = None + + +@dataclass(slots=True) +class ToolInvocationResult: + status: str + content: dict[str, Any] + artifacts: list[dict[str, Any]] = field(default_factory=list) + logs: list[str] = field(default_factory=list) + error: dict[str, Any] | None = None + + +class TaskRepository(Protocol): + def create(self, task: Task) -> Task: ... + def save(self, task: Task) -> Task: ... + def get(self, task_id: str) -> Task | None: ... + + +class GraphRepository(Protocol): + def save(self, graph: ExecutionGraph) -> ExecutionGraph: ... + def get(self, task_id: str) -> ExecutionGraph | None: ... + + +class ConfirmationRepository(Protocol): + def create(self, confirmation: ConfirmationRequest) -> ConfirmationRequest: ... + def get(self, confirmation_id: str) -> ConfirmationRequest | None: ... + def save(self, confirmation: ConfirmationRequest) -> ConfirmationRequest: ... + + +class WorkerRepository(Protocol): + def save(self, worker: WorkerSession) -> WorkerSession: ... + def get(self, session_id: str) -> WorkerSession | None: ... + def list_active(self) -> list[WorkerSession]: ... + + +class EventStore(Protocol): + def append(self, event: DomainEvent) -> DomainEvent: ... + def list_by_task(self, task_id: str) -> list[DomainEvent]: ... + + +class PolicyEvaluator(Protocol): + def evaluate(self, action: ActionDescriptor, project_id: str) -> PolicyDecision: ... + + +class ModelRouter(Protocol): + def run( + self, + *, + project_id: str, + slot: str, + messages: list[dict[str, Any]], + task_context: dict[str, Any], + ) -> ModelInvocationResult: ... + + +class ToolGateway(Protocol): + def call( + self, + *, + project_id: str, + server_id: str, + tool_name: str, + args: dict[str, Any], + task_context: dict[str, Any], + ) -> ToolInvocationResult: ... + + +class WorkerGateway(Protocol): + def dispatch( + self, + *, + project_id: str, + worker_session_id: str, + command_name: str, + args: dict[str, Any], + task_context: dict[str, Any], + ) -> ToolInvocationResult: ... + diff --git a/src/ai_orchestrator/application/services/__init__.py b/src/ai_orchestrator/application/services/__init__.py new file mode 100644 index 0000000..1add9ae --- /dev/null +++ b/src/ai_orchestrator/application/services/__init__.py @@ -0,0 +1,2 @@ +"""Application services.""" + diff --git a/src/ai_orchestrator/application/services/orchestrator.py b/src/ai_orchestrator/application/services/orchestrator.py new file mode 100644 index 0000000..bb8e597 --- /dev/null +++ b/src/ai_orchestrator/application/services/orchestrator.py @@ -0,0 +1,134 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +from ai_orchestrator.application.ports import ( + ConfirmationRepository, + EventStore, + GraphRepository, + PolicyEvaluator, + TaskRepository, +) +from ai_orchestrator.domain.enums import NodeType, PolicyDecisionType +from ai_orchestrator.domain.events import DomainEvent +from ai_orchestrator.domain.models import ( + ActionDescriptor, + ConfirmationRequest, + ExecutionGraph, + ExecutionNode, + Task, +) + + +@dataclass(slots=True) +class CreateTaskRequest: + project_id: str + goal: str + inputs: dict[str, Any] + conversation_id: str | None = None + requested_mode: str = "auto" + + +@dataclass(slots=True) +class OrchestratorService: + task_repository: TaskRepository + graph_repository: GraphRepository + confirmation_repository: ConfirmationRepository + event_store: EventStore + policy_evaluator: PolicyEvaluator + + def create_task(self, request: CreateTaskRequest) -> Task: + task = Task( + project_id=request.project_id, + goal=request.goal, + inputs=request.inputs, + conversation_id=request.conversation_id, + requested_mode=request.requested_mode, + effective_mode=request.requested_mode, + ) + self.task_repository.create(task) + self.event_store.append( + DomainEvent( + event_type="task_created", + task_id=task.task_id, + conversation_id=task.conversation_id, + payload={"goal": task.goal, "requested_mode": task.requested_mode}, + ) + ) + return task + + def plan_task(self, task_id: str) -> ExecutionGraph: + task = self._require_task(task_id) + planner = ExecutionNode( + task_id=task.task_id, + node_type=NodeType.PLANNER, + input_data={"goal": task.goal, "inputs": task.inputs}, + ) + finalizer = ExecutionNode( + task_id=task.task_id, + node_type=NodeType.FINALIZER, + input_data={"task_id": task.task_id}, + dependencies=[planner.node_id], + retryable=False, + ) + graph = ExecutionGraph(task_id=task.task_id, nodes=[planner, finalizer]) + self.graph_repository.save(graph) + task.mark_planned() + self.task_repository.save(task) + self.event_store.append( + DomainEvent( + event_type="task_planned", + task_id=task.task_id, + conversation_id=task.conversation_id, + payload={"node_count": len(graph.nodes)}, + ) + ) + return graph + + def evaluate_action( + self, + *, + task_id: str, + node_id: str, + action: ActionDescriptor, + ) -> ConfirmationRequest | None: + task = self._require_task(task_id) + decision = self.policy_evaluator.evaluate(action, task.project_id) + self.event_store.append( + DomainEvent( + event_type="policy_decision_made", + task_id=task.task_id, + conversation_id=task.conversation_id, + node_id=node_id, + payload={ + "action_name": action.action_name, + "resource": action.resource, + "risk_level": action.risk_level.value, + "decision": decision.decision.value, + }, + ) + ) + if decision.decision != PolicyDecisionType.CONFIRM: + return None + confirmation = ConfirmationRequest(task_id=task.task_id, node_id=node_id, decision=decision) + self.confirmation_repository.create(confirmation) + task.wait_for_confirmation(node_id=node_id) + self.task_repository.save(task) + self.event_store.append( + DomainEvent( + event_type="confirmation_requested", + task_id=task.task_id, + conversation_id=task.conversation_id, + node_id=node_id, + payload={"confirmation_id": confirmation.confirmation_id}, + ) + ) + return confirmation + + def _require_task(self, task_id: str) -> Task: + task = self.task_repository.get(task_id) + if task is None: + raise LookupError(f"Task not found: {task_id}") + return task + diff --git a/src/ai_orchestrator/application/services/workers.py b/src/ai_orchestrator/application/services/workers.py new file mode 100644 index 0000000..92de9ea --- /dev/null +++ b/src/ai_orchestrator/application/services/workers.py @@ -0,0 +1,47 @@ +from __future__ import annotations + +from dataclasses import dataclass + +from ai_orchestrator.application.ports import EventStore, WorkerRepository +from ai_orchestrator.domain.events import DomainEvent +from ai_orchestrator.domain.models import WorkerSession + + +@dataclass(slots=True) +class RegisterWorkerRequest: + worker_id: str + name: str + machine: str + os: str + version: str + capabilities: list[str] + + +@dataclass(slots=True) +class WorkerService: + worker_repository: WorkerRepository + event_store: EventStore + + def register(self, request: RegisterWorkerRequest) -> WorkerSession: + worker = WorkerSession( + worker_id=request.worker_id, + name=request.name, + machine=request.machine, + os=request.os, + version=request.version, + capabilities=request.capabilities, + ) + worker.mark_online() + self.worker_repository.save(worker) + self.event_store.append( + DomainEvent( + event_type="worker_registered", + task_id="system", + payload={ + "session_id": worker.session_id, + "worker_id": worker.worker_id, + "capabilities": worker.capabilities, + }, + ) + ) + return worker diff --git a/src/ai_orchestrator/config.py b/src/ai_orchestrator/config.py new file mode 100644 index 0000000..2e0f6eb --- /dev/null +++ b/src/ai_orchestrator/config.py @@ -0,0 +1,82 @@ +from __future__ import annotations + +from enum import StrEnum + +from pydantic import BaseModel, Field +from pydantic_settings import BaseSettings, SettingsConfigDict + + +class ProviderType(StrEnum): + LOCAL = "local" + EXTERNAL = "external" + DISABLED = "disabled" + + +class PolicyMode(StrEnum): + FULL_AUTO = "full_auto" + CONFIRM = "confirm" + MANUAL = "manual" + DISABLED = "disabled" + + +class ModelSlotConfig(BaseModel): + provider: ProviderType + model: str | None = None + base_url: str | None = None + api_key_env: str | None = None + + +class ModelsConfig(BaseModel): + weak: ModelSlotConfig + strong: ModelSlotConfig + vision: ModelSlotConfig + embedding: ModelSlotConfig + + +class ExecutionConfig(BaseModel): + max_graph_nodes: int = 20 + max_agent_steps: int = 8 + max_local_retries: int = 2 + allow_external_models: bool = False + allow_paid_fallback: bool = False + + +class PolicyResourceConfig(BaseModel): + filesystem: PolicyMode = PolicyMode.CONFIRM + shell: PolicyMode = PolicyMode.CONFIRM + sql: PolicyMode = PolicyMode.CONFIRM + mcp: PolicyMode = PolicyMode.CONFIRM + external_models: PolicyMode = PolicyMode.CONFIRM + desktop: PolicyMode = PolicyMode.CONFIRM + browser: PolicyMode = PolicyMode.CONFIRM + network: PolicyMode = PolicyMode.CONFIRM + cost: PolicyMode = PolicyMode.CONFIRM + system: PolicyMode = PolicyMode.CONFIRM + + +class PolicyConfig(BaseModel): + default_mode: PolicyMode = PolicyMode.CONFIRM + resources: PolicyResourceConfig = Field(default_factory=PolicyResourceConfig) + + +class ProjectMetadata(BaseModel): + id: str + name: str + + +class ProjectConfig(BaseModel): + project: ProjectMetadata + models: ModelsConfig | None = None + execution: ExecutionConfig = Field(default_factory=ExecutionConfig) + policy: PolicyConfig = Field(default_factory=PolicyConfig) + + +class AppSettings(BaseSettings): + model_config = SettingsConfigDict(env_prefix="AI_ORCH_", extra="ignore") + + app_name: str = "ai-orchestrator" + env: str = "dev" + host: str = "127.0.0.1" + port: int = 8080 + default_project_config: str = "configs/examples/local_only.yaml" + diff --git a/src/ai_orchestrator/delivery/__init__.py b/src/ai_orchestrator/delivery/__init__.py new file mode 100644 index 0000000..c3bcb5c --- /dev/null +++ b/src/ai_orchestrator/delivery/__init__.py @@ -0,0 +1,2 @@ +"""Delivery layer.""" + diff --git a/src/ai_orchestrator/delivery/http/__init__.py b/src/ai_orchestrator/delivery/http/__init__.py new file mode 100644 index 0000000..f3c91a9 --- /dev/null +++ b/src/ai_orchestrator/delivery/http/__init__.py @@ -0,0 +1,2 @@ +"""HTTP delivery adapters.""" + diff --git a/src/ai_orchestrator/delivery/http/app.py b/src/ai_orchestrator/delivery/http/app.py new file mode 100644 index 0000000..9981190 --- /dev/null +++ b/src/ai_orchestrator/delivery/http/app.py @@ -0,0 +1,190 @@ +from __future__ import annotations + +from fastapi import FastAPI, HTTPException +from sse_starlette.sse import EventSourceResponse + +from ai_orchestrator.application.services.orchestrator import CreateTaskRequest, OrchestratorService +from ai_orchestrator.application.services.workers import RegisterWorkerRequest, WorkerService +from ai_orchestrator.config import AppSettings +from ai_orchestrator.delivery.http.schemas import ( + ChatRequest, + ChatResponse, + ConfirmationApproveRequest, + ConfirmationRejectRequest, + CreateTaskRequestSchema, + TaskStatusResponse, + WorkerListItem, + WorkerListResponse, + WorkerRegisterRequest, +) +from ai_orchestrator.infrastructure.config_loader import load_project_config +from ai_orchestrator.infrastructure.policy import StaticProjectPolicyEvaluator +from ai_orchestrator.infrastructure.storage.memory import ( + InMemoryConfirmationRepository, + InMemoryEventStore, + InMemoryGraphRepository, + InMemoryTaskRepository, + InMemoryWorkerRepository, +) + + +def create_app(settings: AppSettings | None = None) -> FastAPI: + settings = settings or AppSettings() + project_config = load_project_config(settings.default_project_config) + task_repository = InMemoryTaskRepository() + graph_repository = InMemoryGraphRepository() + confirmation_repository = InMemoryConfirmationRepository() + worker_repository = InMemoryWorkerRepository() + event_store = InMemoryEventStore() + policy_evaluator = StaticProjectPolicyEvaluator( + projects={project_config.project.id: project_config} + ) + orchestrator = OrchestratorService( + task_repository=task_repository, + graph_repository=graph_repository, + confirmation_repository=confirmation_repository, + event_store=event_store, + policy_evaluator=policy_evaluator, + ) + + app = FastAPI(title="AI Orchestrator", version="0.1.0") + app.state.task_repository = task_repository + app.state.graph_repository = graph_repository + app.state.confirmation_repository = confirmation_repository + app.state.worker_repository = worker_repository + app.state.event_store = event_store + app.state.orchestrator = orchestrator + app.state.worker_service = WorkerService( + worker_repository=worker_repository, + event_store=event_store, + ) + + @app.get("/health") + def health() -> dict[str, str]: + return {"status": "ok"} + + @app.post("/chat", response_model=ChatResponse) + def post_chat(payload: ChatRequest) -> ChatResponse: + task = orchestrator.create_task( + CreateTaskRequest( + project_id=payload.project_id, + goal=payload.message, + inputs={"attachments": payload.attachments, "preferences": payload.preferences}, + conversation_id=payload.conversation_id or "conv_default", + requested_mode=payload.mode, + ) + ) + orchestrator.plan_task(task.task_id) + return ChatResponse( + conversation_id=task.conversation_id or "conv_default", + task_id=task.task_id, + response_type="task_started", + message="Task accepted and planned.", + cards=[{"type": "plan_card", "task_id": task.task_id}], + ) + + @app.post("/tasks", response_model=TaskStatusResponse) + def post_tasks(payload: CreateTaskRequestSchema) -> TaskStatusResponse: + task = orchestrator.create_task( + CreateTaskRequest( + project_id=payload.project_id, + goal=payload.goal, + inputs=payload.inputs, + requested_mode=payload.execution_mode, + ) + ) + graph = orchestrator.plan_task(task.task_id) + return TaskStatusResponse( + task_id=task.task_id, + status=task.status.value, + current_node=task.current_node_id, + progress={"completed": 0, "total": len(graph.nodes)}, + ) + + @app.get("/tasks/{task_id}", response_model=TaskStatusResponse) + def get_task(task_id: str) -> TaskStatusResponse: + task = task_repository.get(task_id) + graph = graph_repository.get(task_id) + if task is None or graph is None: + raise HTTPException(status_code=404, detail="Task not found") + completed = sum(1 for node in graph.nodes if node.status.value == "completed") + return TaskStatusResponse( + task_id=task.task_id, + status=task.status.value, + current_node=task.current_node_id, + progress={"completed": completed, "total": len(graph.nodes)}, + ) + + @app.get("/tasks/{task_id}/events") + async def get_task_events(task_id: str) -> EventSourceResponse: + events = event_store.list_by_task(task_id) + + async def iterator(): + for event in events: + yield { + "event": event.event_type, + "data": { + "event_id": event.event_id, + "task_id": event.task_id, + "node_id": event.node_id, + "timestamp": event.occurred_at.isoformat(), + "payload": event.payload, + }, + } + + return EventSourceResponse(iterator()) + + @app.post("/confirmations/{confirmation_id}/approve") + def approve_confirmation( + confirmation_id: str, + payload: ConfirmationApproveRequest, + ) -> dict[str, str]: + confirmation = confirmation_repository.get(confirmation_id) + if confirmation is None: + raise HTTPException(status_code=404, detail="Confirmation not found") + confirmation.approve(payload.comment) + confirmation_repository.save(confirmation) + return {"status": confirmation.status.value} + + @app.post("/confirmations/{confirmation_id}/reject") + def reject_confirmation( + confirmation_id: str, + payload: ConfirmationRejectRequest, + ) -> dict[str, str]: + confirmation = confirmation_repository.get(confirmation_id) + if confirmation is None: + raise HTTPException(status_code=404, detail="Confirmation not found") + confirmation.reject(payload.reason) + confirmation_repository.save(confirmation) + return {"status": confirmation.status.value} + + @app.get("/workers", response_model=WorkerListResponse) + def list_workers() -> WorkerListResponse: + return WorkerListResponse( + workers=[ + WorkerListItem( + session_id=worker.session_id, + worker_id=worker.worker_id, + name=worker.name, + status=worker.status.value, + capabilities=worker.capabilities, + ) + for worker in worker_repository.list_active() + ] + ) + + @app.post("/workers/register") + def register_worker(payload: WorkerRegisterRequest) -> dict[str, object]: + worker = app.state.worker_service.register( + RegisterWorkerRequest( + worker_id=payload.worker_id, + name=payload.name, + machine=payload.machine or payload.name, + os=payload.os or "unknown", + version=payload.version, + capabilities=payload.capabilities, + ) + ) + return {"session_id": worker.session_id, "status": worker.status.value} + + return app diff --git a/src/ai_orchestrator/delivery/http/schemas.py b/src/ai_orchestrator/delivery/http/schemas.py new file mode 100644 index 0000000..70af523 --- /dev/null +++ b/src/ai_orchestrator/delivery/http/schemas.py @@ -0,0 +1,66 @@ +from __future__ import annotations + +from typing import Any + +from pydantic import BaseModel, Field + + +class ChatRequest(BaseModel): + project_id: str = "default" + conversation_id: str | None = None + message: str + attachments: list[dict[str, Any]] = Field(default_factory=list) + mode: str = "auto" + preferences: dict[str, Any] = Field(default_factory=dict) + + +class ChatResponse(BaseModel): + conversation_id: str + task_id: str + response_type: str + message: str + cards: list[dict[str, Any]] = Field(default_factory=list) + + +class CreateTaskRequestSchema(BaseModel): + project_id: str + goal: str + inputs: dict[str, Any] = Field(default_factory=dict) + execution_mode: str = "agent_graph" + + +class TaskStatusResponse(BaseModel): + task_id: str + status: str + current_node: str | None + progress: dict[str, int] + + +class ConfirmationApproveRequest(BaseModel): + scope: str = "once" + comment: str | None = None + + +class ConfirmationRejectRequest(BaseModel): + reason: str + + +class WorkerRegisterRequest(BaseModel): + worker_id: str + name: str + capabilities: list[str] = Field(default_factory=list) + version: str + machine: str | None = None + os: str | None = None + + +class WorkerListItem(BaseModel): + session_id: str + worker_id: str + name: str + status: str + capabilities: list[str] = Field(default_factory=list) + + +class WorkerListResponse(BaseModel): + workers: list[WorkerListItem] = Field(default_factory=list) diff --git a/src/ai_orchestrator/domain/__init__.py b/src/ai_orchestrator/domain/__init__.py new file mode 100644 index 0000000..2568f61 --- /dev/null +++ b/src/ai_orchestrator/domain/__init__.py @@ -0,0 +1,2 @@ +"""Domain layer.""" + diff --git a/src/ai_orchestrator/domain/enums.py b/src/ai_orchestrator/domain/enums.py new file mode 100644 index 0000000..c3c74d8 --- /dev/null +++ b/src/ai_orchestrator/domain/enums.py @@ -0,0 +1,86 @@ +from __future__ import annotations + +from enum import StrEnum + + +class TaskStatus(StrEnum): + CREATED = "created" + PLANNED = "planned" + RUNNING = "running" + WAITING_CONFIRMATION = "waiting_confirmation" + WAITING_MANUAL = "waiting_manual" + COMPLETED = "completed" + FAILED = "failed" + CANCELLED = "cancelled" + + +class NodeStatus(StrEnum): + PENDING = "pending" + READY = "ready" + RUNNING = "running" + WAITING_CONFIRMATION = "waiting_confirmation" + COMPLETED = "completed" + FAILED = "failed" + SKIPPED = "skipped" + CANCELLED = "cancelled" + + +class NodeType(StrEnum): + PLANNER = "planner" + MODEL_CALL = "model_call" + TOOL_CALL = "tool_call" + LOCAL_WORKER_CALL = "local_worker_call" + REVIEWER = "reviewer" + FINALIZER = "finalizer" + CONFIRMATION = "confirmation" + + +class ConfirmationStatus(StrEnum): + PENDING = "pending" + APPROVED = "approved" + REJECTED = "rejected" + EXPIRED = "expired" + CANCELLED = "cancelled" + + +class WorkerSessionStatus(StrEnum): + CONNECTING = "connecting" + ONLINE = "online" + BUSY = "busy" + STALE = "stale" + DISCONNECTED = "disconnected" + + +class ResourceType(StrEnum): + FILESYSTEM = "filesystem" + SHELL = "shell" + SQL = "sql" + MCP = "mcp" + EXTERNAL_MODELS = "external_models" + DESKTOP = "desktop" + BROWSER = "browser" + NETWORK = "network" + COST = "cost" + SYSTEM = "system" + + +class RiskLevel(StrEnum): + SAFE = "safe" + WRITE = "write" + DESTRUCTIVE = "destructive" + SYSTEM = "system" + COST = "cost" + + +class PolicyDecisionType(StrEnum): + ALLOW = "allow" + CONFIRM = "confirm" + MANUAL = "manual" + DISABLED_BY_CONFIG = "disabled_by_config" + + +class ConfirmationScope(StrEnum): + ONCE = "once" + TASK = "task" + PROJECT = "project" + diff --git a/src/ai_orchestrator/domain/errors.py b/src/ai_orchestrator/domain/errors.py new file mode 100644 index 0000000..602f5e2 --- /dev/null +++ b/src/ai_orchestrator/domain/errors.py @@ -0,0 +1,30 @@ +from __future__ import annotations + + +class OrchestratorError(Exception): + """Base class for orchestrator errors.""" + + +class ValidationError(OrchestratorError): + """Raised when request or data validation fails.""" + + +class PolicyError(OrchestratorError): + """Raised when policy blocks or pauses execution.""" + + +class RetryableInfrastructureError(OrchestratorError): + """Raised for transient upstream or transport errors.""" + + +class NonRetryableInfrastructureError(OrchestratorError): + """Raised for permanent upstream or integration errors.""" + + +class ExecutionError(OrchestratorError): + """Raised when node execution fails semantically.""" + + +class CancellationError(OrchestratorError): + """Raised when task or node execution was cancelled.""" + diff --git a/src/ai_orchestrator/domain/events.py b/src/ai_orchestrator/domain/events.py new file mode 100644 index 0000000..b0d8472 --- /dev/null +++ b/src/ai_orchestrator/domain/events.py @@ -0,0 +1,20 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +from datetime import UTC, datetime +from typing import Any +from uuid import uuid4 + + +@dataclass(slots=True, frozen=True) +class DomainEvent: + event_type: str + task_id: str + payload: dict[str, Any] + conversation_id: str | None = None + node_id: str | None = None + correlation_id: str | None = None + causation_id: str | None = None + event_id: str = field(default_factory=lambda: f"evt_{uuid4().hex}") + occurred_at: datetime = field(default_factory=lambda: datetime.now(UTC)) + diff --git a/src/ai_orchestrator/domain/models.py b/src/ai_orchestrator/domain/models.py new file mode 100644 index 0000000..8dd979a --- /dev/null +++ b/src/ai_orchestrator/domain/models.py @@ -0,0 +1,192 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +from datetime import UTC, datetime +from typing import Any +from uuid import uuid4 + +from ai_orchestrator.domain.enums import ( + ConfirmationScope, + ConfirmationStatus, + NodeStatus, + NodeType, + PolicyDecisionType, + RiskLevel, + TaskStatus, + WorkerSessionStatus, +) + + +def _utcnow() -> datetime: + return datetime.now(UTC) + + +@dataclass(slots=True) +class Task: + project_id: str + goal: str + inputs: dict[str, Any] + conversation_id: str | None = None + requested_mode: str = "auto" + effective_mode: str = "auto" + task_id: str = field(default_factory=lambda: f"task_{uuid4().hex}") + status: TaskStatus = TaskStatus.CREATED + current_node_id: str | None = None + result_summary: dict[str, Any] = field(default_factory=dict) + created_at: datetime = field(default_factory=_utcnow) + updated_at: datetime = field(default_factory=_utcnow) + + def mark_planned(self) -> None: + self.status = TaskStatus.PLANNED + self.updated_at = _utcnow() + + def mark_running(self, node_id: str | None = None) -> None: + self.status = TaskStatus.RUNNING + self.current_node_id = node_id + self.updated_at = _utcnow() + + def wait_for_confirmation(self, node_id: str | None = None) -> None: + self.status = TaskStatus.WAITING_CONFIRMATION + self.current_node_id = node_id + self.updated_at = _utcnow() + + def wait_for_manual(self, node_id: str | None = None) -> None: + self.status = TaskStatus.WAITING_MANUAL + self.current_node_id = node_id + self.updated_at = _utcnow() + + def complete(self, summary: dict[str, Any]) -> None: + self.status = TaskStatus.COMPLETED + self.result_summary = summary + self.updated_at = _utcnow() + + def fail(self, summary: dict[str, Any]) -> None: + self.status = TaskStatus.FAILED + self.result_summary = summary + self.updated_at = _utcnow() + + +@dataclass(slots=True) +class ExecutionNode: + task_id: str + node_type: NodeType + input_data: dict[str, Any] + dependencies: list[str] = field(default_factory=list) + node_id: str = field(default_factory=lambda: f"node_{uuid4().hex}") + status: NodeStatus = NodeStatus.PENDING + output_data: dict[str, Any] = field(default_factory=dict) + assigned_runner: str | None = None + attempts: int = 0 + retryable: bool = True + timeout_ms: int | None = None + created_at: datetime = field(default_factory=_utcnow) + updated_at: datetime = field(default_factory=_utcnow) + + def is_ready(self, completed_node_ids: set[str]) -> bool: + return all(dep in completed_node_ids for dep in self.dependencies) + + def mark_ready(self) -> None: + self.status = NodeStatus.READY + self.updated_at = _utcnow() + + def mark_running(self) -> None: + self.status = NodeStatus.RUNNING + self.attempts += 1 + self.updated_at = _utcnow() + + def mark_waiting_confirmation(self) -> None: + self.status = NodeStatus.WAITING_CONFIRMATION + self.updated_at = _utcnow() + + def mark_completed(self, output_data: dict[str, Any]) -> None: + self.status = NodeStatus.COMPLETED + self.output_data = output_data + self.updated_at = _utcnow() + + def mark_failed(self, output_data: dict[str, Any] | None = None) -> None: + self.status = NodeStatus.FAILED + self.output_data = output_data or {} + self.updated_at = _utcnow() + + +@dataclass(slots=True) +class ExecutionGraph: + task_id: str + nodes: list[ExecutionNode] + + def completed_node_ids(self) -> set[str]: + return {node.node_id for node in self.nodes if node.status == NodeStatus.COMPLETED} + + def ready_nodes(self) -> list[ExecutionNode]: + completed = self.completed_node_ids() + ready: list[ExecutionNode] = [] + for node in self.nodes: + if node.status == NodeStatus.PENDING and node.is_ready(completed): + node.mark_ready() + if node.status == NodeStatus.READY: + ready.append(node) + return ready + + +@dataclass(slots=True) +class ActionDescriptor: + resource: str + risk_level: RiskLevel + action_name: str + preview_available: bool = False + metadata: dict[str, Any] = field(default_factory=dict) + + +@dataclass(slots=True) +class PolicyDecision: + decision: PolicyDecisionType + reason: str + requires_confirmation: bool = False + preview: dict[str, Any] = field(default_factory=dict) + + +@dataclass(slots=True) +class ConfirmationRequest: + task_id: str + node_id: str + decision: PolicyDecision + scope: ConfirmationScope = ConfirmationScope.ONCE + confirmation_id: str = field(default_factory=lambda: f"conf_{uuid4().hex}") + status: ConfirmationStatus = ConfirmationStatus.PENDING + comment: str | None = None + created_at: datetime = field(default_factory=_utcnow) + resolved_at: datetime | None = None + + def approve(self, comment: str | None = None) -> None: + self.status = ConfirmationStatus.APPROVED + self.comment = comment + self.resolved_at = _utcnow() + + def reject(self, comment: str | None = None) -> None: + self.status = ConfirmationStatus.REJECTED + self.comment = comment + self.resolved_at = _utcnow() + + +@dataclass(slots=True) +class WorkerSession: + worker_id: str + name: str + machine: str + os: str + version: str + capabilities: list[str] + session_id: str = field(default_factory=lambda: f"wrk_{uuid4().hex}") + status: WorkerSessionStatus = WorkerSessionStatus.CONNECTING + current_task_id: str | None = None + last_heartbeat_at: datetime = field(default_factory=_utcnow) + connected_at: datetime = field(default_factory=_utcnow) + disconnected_at: datetime | None = None + + def mark_online(self) -> None: + self.status = WorkerSessionStatus.ONLINE + self.last_heartbeat_at = _utcnow() + + def heartbeat(self) -> None: + self.last_heartbeat_at = _utcnow() + diff --git a/src/ai_orchestrator/infrastructure/__init__.py b/src/ai_orchestrator/infrastructure/__init__.py new file mode 100644 index 0000000..c82f3d5 --- /dev/null +++ b/src/ai_orchestrator/infrastructure/__init__.py @@ -0,0 +1,2 @@ +"""Infrastructure layer.""" + diff --git a/src/ai_orchestrator/infrastructure/config_loader.py b/src/ai_orchestrator/infrastructure/config_loader.py new file mode 100644 index 0000000..81f025c --- /dev/null +++ b/src/ai_orchestrator/infrastructure/config_loader.py @@ -0,0 +1,15 @@ +from __future__ import annotations + +from pathlib import Path + +import yaml + +from ai_orchestrator.config import ProjectConfig + + +def load_project_config(path: str | Path) -> ProjectConfig: + config_path = Path(path) + with config_path.open("r", encoding="utf-8") as fh: + raw = yaml.safe_load(fh) or {} + return ProjectConfig.model_validate(raw) + diff --git a/src/ai_orchestrator/infrastructure/logging.py b/src/ai_orchestrator/infrastructure/logging.py new file mode 100644 index 0000000..0692893 --- /dev/null +++ b/src/ai_orchestrator/infrastructure/logging.py @@ -0,0 +1,14 @@ +from __future__ import annotations + +import json +import logging +from typing import Any + + +def configure_logging(level: int = logging.INFO) -> None: + logging.basicConfig(level=level, format="%(message)s") + + +def log_structured(logger: logging.Logger, event_type: str, **payload: Any) -> None: + logger.info(json.dumps({"event_type": event_type, **payload}, ensure_ascii=True, default=str)) + diff --git a/src/ai_orchestrator/infrastructure/policy.py b/src/ai_orchestrator/infrastructure/policy.py new file mode 100644 index 0000000..6bf087f --- /dev/null +++ b/src/ai_orchestrator/infrastructure/policy.py @@ -0,0 +1,34 @@ +from __future__ import annotations + +from dataclasses import dataclass + +from ai_orchestrator.application.ports import PolicyEvaluator +from ai_orchestrator.config import PolicyMode, ProjectConfig +from ai_orchestrator.domain.enums import PolicyDecisionType +from ai_orchestrator.domain.models import ActionDescriptor, PolicyDecision + + +@dataclass(slots=True) +class StaticProjectPolicyEvaluator(PolicyEvaluator): + projects: dict[str, ProjectConfig] + + def evaluate(self, action: ActionDescriptor, project_id: str) -> PolicyDecision: + project = self.projects[project_id] + resource_mode = getattr( + project.policy.resources, + action.resource, + project.policy.default_mode, + ) + mapping = { + PolicyMode.FULL_AUTO: PolicyDecisionType.ALLOW, + PolicyMode.CONFIRM: PolicyDecisionType.CONFIRM, + PolicyMode.MANUAL: PolicyDecisionType.MANUAL, + PolicyMode.DISABLED: PolicyDecisionType.DISABLED_BY_CONFIG, + } + decision = mapping[resource_mode] + return PolicyDecision( + decision=decision, + reason=f"Policy mode for {action.resource}: {resource_mode.value}", + requires_confirmation=decision == PolicyDecisionType.CONFIRM, + preview={"preview_available": action.preview_available, **action.metadata}, + ) diff --git a/src/ai_orchestrator/infrastructure/storage/__init__.py b/src/ai_orchestrator/infrastructure/storage/__init__.py new file mode 100644 index 0000000..61b9049 --- /dev/null +++ b/src/ai_orchestrator/infrastructure/storage/__init__.py @@ -0,0 +1,2 @@ +"""Storage adapters.""" + diff --git a/src/ai_orchestrator/infrastructure/storage/memory.py b/src/ai_orchestrator/infrastructure/storage/memory.py new file mode 100644 index 0000000..75de0e8 --- /dev/null +++ b/src/ai_orchestrator/infrastructure/storage/memory.py @@ -0,0 +1,85 @@ +from __future__ import annotations + +from dataclasses import dataclass, field + +from ai_orchestrator.application.ports import ( + ConfirmationRepository, + EventStore, + GraphRepository, + TaskRepository, + WorkerRepository, +) +from ai_orchestrator.domain.events import DomainEvent +from ai_orchestrator.domain.models import ConfirmationRequest, ExecutionGraph, Task, WorkerSession + + +@dataclass(slots=True) +class InMemoryTaskRepository(TaskRepository): + items: dict[str, Task] = field(default_factory=dict) + + def create(self, task: Task) -> Task: + self.items[task.task_id] = task + return task + + def save(self, task: Task) -> Task: + self.items[task.task_id] = task + return task + + def get(self, task_id: str) -> Task | None: + return self.items.get(task_id) + + +@dataclass(slots=True) +class InMemoryGraphRepository(GraphRepository): + items: dict[str, ExecutionGraph] = field(default_factory=dict) + + def save(self, graph: ExecutionGraph) -> ExecutionGraph: + self.items[graph.task_id] = graph + return graph + + def get(self, task_id: str) -> ExecutionGraph | None: + return self.items.get(task_id) + + +@dataclass(slots=True) +class InMemoryConfirmationRepository(ConfirmationRepository): + items: dict[str, ConfirmationRequest] = field(default_factory=dict) + + def create(self, confirmation: ConfirmationRequest) -> ConfirmationRequest: + self.items[confirmation.confirmation_id] = confirmation + return confirmation + + def get(self, confirmation_id: str) -> ConfirmationRequest | None: + return self.items.get(confirmation_id) + + def save(self, confirmation: ConfirmationRequest) -> ConfirmationRequest: + self.items[confirmation.confirmation_id] = confirmation + return confirmation + + +@dataclass(slots=True) +class InMemoryWorkerRepository(WorkerRepository): + items: dict[str, WorkerSession] = field(default_factory=dict) + + def save(self, worker: WorkerSession) -> WorkerSession: + self.items[worker.session_id] = worker + return worker + + def get(self, session_id: str) -> WorkerSession | None: + return self.items.get(session_id) + + def list_active(self) -> list[WorkerSession]: + return list(self.items.values()) + + +@dataclass(slots=True) +class InMemoryEventStore(EventStore): + items: list[DomainEvent] = field(default_factory=list) + + def append(self, event: DomainEvent) -> DomainEvent: + self.items.append(event) + return event + + def list_by_task(self, task_id: str) -> list[DomainEvent]: + return [event for event in self.items if event.task_id == task_id] + diff --git a/src/ai_orchestrator/main.py b/src/ai_orchestrator/main.py new file mode 100644 index 0000000..3e92aba --- /dev/null +++ b/src/ai_orchestrator/main.py @@ -0,0 +1,6 @@ +from __future__ import annotations + +from ai_orchestrator.delivery.http.app import create_app + +app = create_app() + diff --git a/tests/integration/test_http_app.py b/tests/integration/test_http_app.py new file mode 100644 index 0000000..42486f2 --- /dev/null +++ b/tests/integration/test_http_app.py @@ -0,0 +1,53 @@ +from fastapi.testclient import TestClient + +from ai_orchestrator.main import app + + +def test_health_endpoint_returns_ok() -> None: + client = TestClient(app) + + response = client.get("/health") + + assert response.status_code == 200 + assert response.json() == {"status": "ok"} + + +def test_post_tasks_creates_planned_task() -> None: + client = TestClient(app) + + response = client.post( + "/tasks", + json={ + "project_id": "default", + "goal": "Prepare a plan", + "inputs": {"source": "test"}, + "execution_mode": "agent_graph", + }, + ) + + payload = response.json() + + assert response.status_code == 200 + assert payload["status"] == "planned" + assert payload["progress"]["total"] == 2 + + +def test_worker_registration_is_exposed_via_list_endpoint() -> None: + client = TestClient(app) + + registration = client.post( + "/workers/register", + json={ + "worker_id": "worker_home_pc", + "name": "Home PC", + "capabilities": ["file.read", "command.run"], + "version": "0.1.0", + "machine": "DESKTOP-1", + "os": "windows", + }, + ) + listing = client.get("/workers") + + assert registration.status_code == 200 + assert listing.status_code == 200 + assert any(worker["worker_id"] == "worker_home_pc" for worker in listing.json()["workers"]) diff --git a/tests/unit/test_graph.py b/tests/unit/test_graph.py new file mode 100644 index 0000000..7ba99c6 --- /dev/null +++ b/tests/unit/test_graph.py @@ -0,0 +1,21 @@ +from ai_orchestrator.domain.enums import NodeStatus, NodeType +from ai_orchestrator.domain.models import ExecutionGraph, ExecutionNode + + +def test_graph_marks_pending_node_ready_when_dependencies_are_completed() -> None: + first = ExecutionNode(task_id="task_1", node_type=NodeType.PLANNER, input_data={}) + second = ExecutionNode( + task_id="task_1", + node_type=NodeType.FINALIZER, + input_data={}, + dependencies=[first.node_id], + ) + first.mark_completed({"ok": True}) + graph = ExecutionGraph(task_id="task_1", nodes=[first, second]) + + ready = graph.ready_nodes() + + assert len(ready) == 1 + assert ready[0].node_id == second.node_id + assert second.status == NodeStatus.READY + diff --git a/tests/unit/test_policy.py b/tests/unit/test_policy.py new file mode 100644 index 0000000..97e5426 --- /dev/null +++ b/tests/unit/test_policy.py @@ -0,0 +1,80 @@ +from ai_orchestrator.config import ( + ExecutionConfig, + ModelsConfig, + ModelSlotConfig, + PolicyConfig, + PolicyMode, + ProjectConfig, + ProjectMetadata, +) +from ai_orchestrator.domain.enums import PolicyDecisionType, RiskLevel +from ai_orchestrator.domain.models import ActionDescriptor +from ai_orchestrator.infrastructure.policy import StaticProjectPolicyEvaluator + + +def test_policy_evaluator_returns_confirm_for_configured_resource() -> None: + project = 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(), + ) + evaluator = StaticProjectPolicyEvaluator(projects={"default": project}) + + decision = evaluator.evaluate( + ActionDescriptor( + resource="filesystem", + risk_level=RiskLevel.WRITE, + action_name="file.write", + preview_available=True, + ), + "default", + ) + + assert decision.decision == PolicyDecisionType.CONFIRM + assert decision.requires_confirmation is True + + +def test_policy_evaluator_returns_disabled_by_config() -> None: + project = 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( + default_mode=PolicyMode.CONFIRM, + resources={ + "filesystem": PolicyMode.CONFIRM, + "shell": PolicyMode.CONFIRM, + "sql": PolicyMode.CONFIRM, + "mcp": PolicyMode.CONFIRM, + "external_models": PolicyMode.DISABLED, + "desktop": PolicyMode.CONFIRM, + "browser": PolicyMode.CONFIRM, + "network": PolicyMode.CONFIRM, + "cost": PolicyMode.CONFIRM, + "system": PolicyMode.CONFIRM, + }, + ), + ) + evaluator = StaticProjectPolicyEvaluator(projects={"default": project}) + + decision = evaluator.evaluate( + ActionDescriptor( + resource="external_models", + risk_level=RiskLevel.COST, + action_name="model.external", + ), + "default", + ) + + assert decision.decision == PolicyDecisionType.DISABLED_BY_CONFIG diff --git a/tests/unit/test_task_lifecycle.py b/tests/unit/test_task_lifecycle.py new file mode 100644 index 0000000..f5d0e6e --- /dev/null +++ b/tests/unit/test_task_lifecycle.py @@ -0,0 +1,15 @@ +from ai_orchestrator.domain.enums import TaskStatus +from ai_orchestrator.domain.models import Task + + +def test_task_status_transitions() -> None: + task = Task(project_id="default", goal="Test", inputs={}) + + task.mark_planned() + task.mark_running() + task.wait_for_confirmation("node_1") + task.complete({"result": "ok"}) + + assert task.status == TaskStatus.COMPLETED + assert task.current_node_id == "node_1" + diff --git a/tests/unit/test_worker_service.py b/tests/unit/test_worker_service.py new file mode 100644 index 0000000..033fcb8 --- /dev/null +++ b/tests/unit/test_worker_service.py @@ -0,0 +1,26 @@ +from ai_orchestrator.application.services.workers import RegisterWorkerRequest, WorkerService +from ai_orchestrator.infrastructure.storage.memory import ( + InMemoryEventStore, + InMemoryWorkerRepository, +) + + +def test_worker_registration_creates_online_session_and_event() -> None: + worker_repository = InMemoryWorkerRepository() + event_store = InMemoryEventStore() + service = WorkerService(worker_repository=worker_repository, event_store=event_store) + + worker = service.register( + RegisterWorkerRequest( + worker_id="worker_home_pc", + name="Home PC", + machine="DESKTOP-1", + os="windows", + version="0.1.0", + capabilities=["file.read", "command.run"], + ) + ) + + assert worker.status.value == "online" + assert worker_repository.get(worker.session_id) is not None + assert event_store.items[-1].event_type == "worker_registered"