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.core.errors import AppError
from app.domain.sandbox.models import SessionStatus, WorkspaceStatus
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 reuses the active session of the task's project. When there is
    none, it starts the project's workspace session itself (see ``_provision``):
    a fresh organization cannot start a sandbox through the API, because sandbox
    mutations are owner-only, and without a session every tool turn fails. When
    the session cannot be started 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.projects import ProjectRepository
        from app.infrastructure.repositories.sandbox import (
            SandboxSessionRepository,
            WorkspaceRepository,
        )

        if not task.project_id:
            return ""
        # A task only ever runs in its own organization's project, whatever the row says.
        own = await ProjectRepository(self.session).list_for_org(task.organization_id)
        if task.project_id not in {project.id for project in own}:
            return ""
        found = await WorkspaceRepository(self.session).list_for_project(task.project_id)
        workspaces = [w for w in found if w.organization_id == task.organization_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 await self._provision(task)

    async def _provision(self, task: Task) -> str:
        """Start the project's workspace session and return its id, or ``""``.

        A project keeps one workspace: the newest one that is not destroyed is
        reused, and only a project with none gets a new workspace, built from the
        platform template for the configured provider. A failure is recorded and
        never raised, so the turn runs without a container and the gates say so.
        """
        from app.application.container import build_sandbox_service
        from app.core.config import settings
        from app.infrastructure.error_log import record_error
        from app.infrastructure.repositories.sandbox import (
            SandboxTemplateRepository,
            WorkspaceRepository,
        )

        project_id = task.project_id
        if not project_id:
            return ""
        metadata = {"task_id": task.id, "project_id": project_id}
        service = build_sandbox_service(self.session)
        try:
            found = await WorkspaceRepository(self.session).list_for_project(project_id)
            workspaces = [w for w in found if w.organization_id == task.organization_id]
            workspace = next(
                (w for w in workspaces if w.status != WorkspaceStatus.DESTROYED), None
            )
            if workspace is None:
                templates = await SandboxTemplateRepository(self.session).list()
                if not templates:
                    await record_error(
                        kind="engine.sandbox_provision_failed",
                        message="no sandbox template is seeded",
                        organization_id=task.organization_id,
                        metadata=metadata,
                    )
                    return ""
                preferred = [
                    t for t in templates if str(t.provider) == settings.sandbox_provider
                ]
                workspace = await service.create_workspace(
                    organization_id=task.organization_id,
                    project_id=project_id,
                    template_id=(preferred or templates)[0].id,
                )
            started = await service.start_session(
                organization_id=task.organization_id, workspace_id=workspace.id
            )
        except AppError as exc:
            await record_error(
                kind="engine.sandbox_provision_failed",
                message=type(exc).__name__,
                exc=exc,
                organization_id=task.organization_id,
                metadata=metadata,
            )
            return ""
        return started.external_session_id or ""


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.billing import CreditWalletRepository
    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,
        wallets=CreditWalletRepository(session),
    )


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"]