smart-org/apps/api/app/application/container_engine.py

"""Wiring for the AI-workforce workstream (team bootstrap).

The org wizard calls ``build_workforce_bootstrap(session)`` after it commits a
new organization, then ``bootstrap_team`` creates the starting team. The team
is built on the *caller's* session, so the employees and the "waiting for a
model" note land in the same transaction as the organization.
"""

from __future__ import annotations

from typing import TYPE_CHECKING

from sqlalchemy.ext.asyncio import AsyncSession

from app.application.ports.workforce_bootstrap import WorkforceBootstrapPort
from app.domain.sandbox.models import SessionStatus
from app.domain.workforce.models import Task

if TYPE_CHECKING:
    from app.application.engine.manager import ManagerTurns
    from app.application.engine.runner import TurnRunner
    from app.application.engine.scheduler import Scheduler


class _EngineWorkforceBootstrap:
    """The real bootstrap: one manager, planner, developer and tester.

    Everything is built here, lazily, from the session: ``container.py`` is a
    frozen composition root and the engine must not widen it, so the repos and
    ports the manager needs are assembled where they are used.
    """

    def __init__(self, session: AsyncSession) -> None:
        self.session = session

    async def bootstrap_team(self, *, organization_id: str, project_id: str) -> None:
        from app.application.container import get_model_provider
        from app.application.container_org import build_memory, build_model_grants
        from app.application.engine.manager import ManagerTurns
        from app.infrastructure.repositories.ai_control import (
            AiModelRepository,
            ModelCredentialRepository,
            RateCardRepository,
        )
        from app.infrastructure.repositories.audit import UsageLedgerRepository
        from app.infrastructure.repositories.engine import (
            EmployeeProfileRepository,
            PipelineStateRepository,
        )
        from app.infrastructure.repositories.workforce import (
            EmployeeRepository,
            MessageRepository,
            TaskRepository,
            ThreadRepository,
        )

        session = self.session
        manager = ManagerTurns(
            provider=get_model_provider(),
            employees=EmployeeRepository(session),
            messages=MessageRepository(session),
            tasks=TaskRepository(session),
            threads=ThreadRepository(session),
            models=AiModelRepository(session),
            credentials=ModelCredentialRepository(session),
            rate_cards=RateCardRepository(session),
            ledger=UsageLedgerRepository(session),
            grants=build_model_grants(session),
            memory=build_memory(session),
            profiles=EmployeeProfileRepository(session),
            states=PipelineStateRepository(session),
        )
        await manager.bootstrap_team(
            organization_id=organization_id, project_id=project_id
        )


def build_workforce_bootstrap(session: AsyncSession) -> WorkforceBootstrapPort:
    return _EngineWorkforceBootstrap(session)


# ---------------------------------------------------------------------------
# the scheduler
# ---------------------------------------------------------------------------


class _EngineSandboxSessions:
    """The sandbox session the engine's turns run in.

    The engine does not *create* sandboxes: a workspace and a running session
    belong to the project and are started through the sandbox flow. The engine
    only reuses the active session of the task's project. When there is none it
    returns ``""``, and the tests gate treats that as red and sends the work
    back — never "merge untested because there was no container".
    """

    #: A session in one of these states can still run a command.
    ACTIVE = frozenset(
        {
            SessionStatus.STARTING,
            SessionStatus.RUNNING,
            SessionStatus.PAUSED,
        }
    )

    def __init__(self, session: AsyncSession) -> None:
        self.session = session

    async def for_task(self, task: Task) -> str:
        from app.infrastructure.repositories.sandbox import (
            SandboxSessionRepository,
            WorkspaceRepository,
        )

        if not task.project_id:
            return ""
        workspaces = await WorkspaceRepository(self.session).list_for_project(
            task.project_id
        )
        sessions = SandboxSessionRepository(self.session)
        for workspace in workspaces:
            for sandbox_session in await sessions.list_for_workspace(workspace.id):
                if (
                    sandbox_session.status in self.ACTIVE
                    and sandbox_session.external_session_id
                ):
                    return sandbox_session.external_session_id
        return ""


def _build_manager(
    session: AsyncSession,
    *,
    runner: TurnRunner,
    provider,
    grants,
    memory,
) -> ManagerTurns:
    from app.application.engine.manager import ManagerTurns
    from app.infrastructure.repositories.ai_control import (
        AiModelRepository,
        ModelCredentialRepository,
        RateCardRepository,
    )
    from app.infrastructure.repositories.audit import UsageLedgerRepository
    from app.infrastructure.repositories.engine import (
        EmployeeProfileRepository,
        PipelineStateRepository,
    )
    from app.infrastructure.repositories.workforce import (
        EmployeeRepository,
        MessageRepository,
        TaskRepository,
        ThreadRepository,
    )

    return ManagerTurns(
        provider=provider,
        employees=EmployeeRepository(session),
        messages=MessageRepository(session),
        tasks=TaskRepository(session),
        threads=ThreadRepository(session),
        models=AiModelRepository(session),
        credentials=ModelCredentialRepository(session),
        rate_cards=RateCardRepository(session),
        ledger=UsageLedgerRepository(session),
        grants=grants,
        memory=memory,
        profiles=EmployeeProfileRepository(session),
        states=PipelineStateRepository(session),
        runner=runner,
    )


def build_scheduler(session: AsyncSession) -> Scheduler:
    """The engine's driver, assembled on the caller's session.

    Same rule as the bootstrap: ``container.py`` is frozen, so the engine wires
    itself here. Every port the scheduler touches is built from the session or
    from the existing composition helpers, so the worker and the API cannot
    disagree about which adapter backs a port.
    """
    from app.application.container import get_model_provider
    from app.application.container_delivery import (
        build_delivery_hooks,
        build_repo_workspace,
    )
    from app.application.container_org import build_memory, build_model_grants
    from app.application.engine.events import RunFeed
    from app.application.engine.gates import GateRunner
    from app.application.engine.runner import TurnRunner
    from app.application.engine.scheduler import Scheduler
    from app.infrastructure.adapters.registry import get_sandbox_provider
    from app.infrastructure.repositories.ai_control import (
        AiModelRepository,
        ModelCredentialRepository,
        RateCardRepository,
    )
    from app.infrastructure.repositories.app_profile import AppProfileRepository
    from app.infrastructure.repositories.audit import UsageLedgerRepository
    from app.infrastructure.repositories.billing import CreditWalletRepository
    from app.infrastructure.repositories.engine import (
        EmployeeProfileRepository,
        EngineTurnRepository,
        GateResultRepository,
        PipelineStateRepository,
        TaskSpendRepository,
    )
    from app.infrastructure.repositories.workforce import (
        EmployeeRepository,
        RunRepository,
        RunStepRepository,
        TaskRepository,
        ToolCallRepository,
    )

    provider = get_model_provider()
    sandbox = get_sandbox_provider()
    grants = build_model_grants(session)
    memory = build_memory(session)
    feed = RunFeed()

    runner = TurnRunner(
        provider=provider,
        models=AiModelRepository(session),
        credentials=ModelCredentialRepository(session),
        rate_cards=RateCardRepository(session),
        ledger=UsageLedgerRepository(session),
        grants=grants,
        sandbox=sandbox,
        memory=memory,
        turns=EngineTurnRepository(session),
        run_steps=RunStepRepository(session),
        tool_calls=ToolCallRepository(session),
        runs=RunRepository(session),
        states=PipelineStateRepository(session),
        feed=feed,
    )
    manager = _build_manager(
        session, runner=runner, provider=provider, grants=grants, memory=memory
    )
    gate_runner = GateRunner(
        results=GateResultRepository(session),
        feed=feed,
        profiles=AppProfileRepository(session),
    )
    sandbox_sessions = _EngineSandboxSessions(session)

    return Scheduler(
        states=PipelineStateRepository(session),
        tasks=TaskRepository(session),
        employees=EmployeeRepository(session),
        profiles=EmployeeProfileRepository(session),
        runs=RunRepository(session),
        turns=EngineTurnRepository(session),
        gates=GateResultRepository(session),
        manager=manager,
        runner=runner,
        gate_runner=gate_runner,
        repo_workspace=build_repo_workspace(session),
        sandbox=sandbox,
        memory=memory,
        delivery=build_delivery_hooks(session),
        wallets=CreditWalletRepository(session),
        spend=TaskSpendRepository(session),
        sandbox_session_for=sandbox_sessions.for_task,
    )


__all__ = ["build_scheduler", "build_workforce_bootstrap"]