smart-org/apps/api/tests/test_engine_flow.py

"""End-to-end: the scheduler walks a task from a queued idea to merged code.

This is the 4.6 flow test. The repositories, the scheduler, the manager and the
turn runner are all real; only the model (scripted), the sandbox (scripted) and
git (in-memory ``FakeRepoWorkspace``) are fakes. What it defends:

* a task advances ``analysis → implementation → test → final review → merge``
  one stage per tick, and the coarse ``tasks.status`` follows the stage
* the platform — never the model — runs git: ``commit``/``merge_into``/``push``
  were called, the commit author carries the platform address, and the model
  was never asked to run a command
* every gate leaves a ``gate_results`` row (diff guard, no-progress, tests,
  green)
* red tests and a forbidden diff send the work *back* with cycle + 1 instead of
  merging it
* ``analysis_only`` never creates a worktree
* an overdrawn wallet pauses the task before a single turn is spent
* a transient provider failure is retried, then handed to the manager
"""

from __future__ import annotations

from collections.abc import AsyncIterator
from dataclasses import dataclass

import pytest_asyncio

from app.application.engine.events import RunFeed
from app.application.engine.gates import GateRunner
from app.application.engine.manager import ManagerTurns
from app.application.engine.runner import TurnRunner
from app.application.engine.scheduler import BUDGET_PAUSED_MESSAGE, ENGINE_COMMIT_EMAIL, Scheduler
from app.application.ports import CommandResult
from app.application.ports.delivery_hooks import NullDeliveryHooks
from app.application.ports.repo_workspace import DiffResult, FakeRepoWorkspace, MergeOutcome
from app.core.errors import ProviderUnavailableError
from app.domain.base import new_id
from app.domain.workforce.engine import EmployeeProfile
from app.domain.workforce.models import Employee, RoleKey, TaskStatus
from app.domain.workforce.pipeline import PipelineKey, Stage, TaskKind, TaskPipelineState
from app.infrastructure.db.models import (
    AiModelRow,
    AiProviderRow,
    OrganizationRow,
    ProjectRow,
    TaskRow,
    UserRow,
)
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.engine import (
    EmployeeProfileRepository,
    EngineTurnRepository,
    GateResultRepository,
    PipelineStateRepository,
)
from app.infrastructure.repositories.workforce import (
    EmployeeRepository,
    MessageRepository,
    RunRepository,
    RunStepRepository,
    TaskRepository,
    ThreadRepository,
    ToolCallRepository,
)
from tests.engine_fakes import (
    FakeGrants,
    FakeSandbox,
    RecordingMemory,
    ScriptedModel,
    ScriptedReply,
)


@dataclass
class _World:
    organization_id: str
    user_id: str
    project_id: str
    strong_id: str


class _OverdrawnWallet:
    """A wallet below zero: the org is overdrawn, so the engine stops before a turn."""

    balance = -1.0


class _ZeroWallet:
    """A wallet at exactly zero: the normal starting state of a fresh, post-paid org."""

    balance = 0.0


class _Wallets:
    def __init__(self, wallet: _OverdrawnWallet | _ZeroWallet | None) -> None:
        self._wallet = wallet

    async def get_for_org(self, organization_id: str):
        return self._wallet


class _OverdrawnWalletRepo(_Wallets):
    def __init__(self) -> None:
        super().__init__(_OverdrawnWallet())


@pytest_asyncio.fixture
async def world(session_factory, seeded) -> AsyncIterator[_World]:
    async with session_factory() as session:
        user = UserRow(email=f"flow-{new_id()[:8]}@example.com", password_hash="x")
        session.add(user)
        await session.flush()
        org = OrganizationRow(
            name="سازمان جریان", slug=f"flow-{user.id[:6]}", owner_user_id=user.id
        )
        session.add(org)
        await session.flush()
        project = ProjectRow(
            id=new_id(),
            organization_id=org.id,
            name="app",
            slug="app",
            test_command="pytest -q",
        )
        session.add(project)
        provider_id = (await session.execute(AiProviderRow.__table__.select().limit(1))).first().id
        strong = await AiModelRepository(session).create(
            AiModelRow(
                provider_id=provider_id,
                internal_name=f"strong-{new_id()[:4]}",
                display_name_for_admin="مدل قوی",
                capabilities=["text", "code", "reasoning"],
                input_price_per_1m=0.5,
                output_price_per_1m=3.0,
                is_active=True,
            )
        )
        await session.commit()
        yield _World(
            organization_id=org.id,
            user_id=user.id,
            project_id=project.id,
            strong_id=strong.id,
        )


async def _employee(session, world: _World, name: str, role: RoleKey) -> Employee:
    employee = await EmployeeRepository(session).create(
        Employee(organization_id=world.organization_id, name=name, role_key=role)
    )
    await EmployeeProfileRepository(session).upsert(
        EmployeeProfile(employee_id=employee.id, ai_model_id=world.strong_id, roles=[role.value])
    )
    return employee


async def _team(session, world: _World) -> dict[RoleKey, Employee]:
    return {
        RoleKey.MANAGER: await _employee(session, world, "مدیر", RoleKey.MANAGER),
        RoleKey.PLANNER: await _employee(session, world, "برنامه‌ریز", RoleKey.PLANNER),
        RoleKey.DEVELOPER: await _employee(session, world, "توسعه‌دهنده", RoleKey.DEVELOPER),
        RoleKey.REVIEWER: await _employee(session, world, "بازبین", RoleKey.REVIEWER),
        RoleKey.TESTER: await _employee(session, world, "آزمون‌گر", RoleKey.TESTER),
    }


async def _task(session, world: _World, *, stage: Stage, pipeline=PipelineKey.PLAN_CODE_TEST,
                kind=TaskKind.FEATURE) -> TaskRow:
    row = TaskRow(
        organization_id=world.organization_id,
        title="کار جریان",
        project_id=world.project_id,
        created_by_user_id=world.user_id,
    )
    session.add(row)
    await session.flush()
    await PipelineStateRepository(session).create(
        TaskPipelineState(
            task_id=row.id,
            organization_id=world.organization_id,
            project_id=world.project_id,
            pipeline=pipeline,
            kind=kind,
            stage=stage,
        )
    )
    await session.flush()
    return row


def _scheduler(
    session,
    world: _World,
    *,
    scripted,
    sandbox: FakeSandbox,
    workspace: FakeRepoWorkspace,
    memory: RecordingMemory,
    delivery,
    wallets=None,
) -> Scheduler:
    grants = FakeGrants({world.strong_id: "strong"})
    runner = TurnRunner(
        provider=scripted,
        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=RunFeed(),
    )
    manager = ManagerTurns(
        provider=scripted,
        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,
        feed=RunFeed(),
    )
    gate_runner = GateRunner(
        results=GateResultRepository(session),
        feed=RunFeed(),
        profiles=AppProfileRepository(session),
    )

    async def _session_for(_task) -> str:
        return "sess-1"

    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=workspace,
        sandbox=sandbox,
        memory=memory,
        delivery=delivery,
        wallets=wallets,
        sandbox_session_for=_session_for,
    )


async def _gates(session, task_id: str) -> dict[str, bool]:
    rows = await GateResultRepository(session).list_for_task(task_id)
    return {row.gate.value: row.ok for row in rows}


def _calls(workspace: FakeRepoWorkspace, name: str) -> list[dict]:
    return [kwargs for called, kwargs in workspace.calls if called == name]


async def _run(scheduler: Scheduler, session, ticks: int) -> None:
    for _ in range(ticks):
        await scheduler.tick()
        await session.flush()


# ---------------------------------------------------------------------------
# the whole pipeline
# ---------------------------------------------------------------------------


async def test_a_task_runs_from_triage_to_merged_code(world, session_factory):
    scripted = ScriptedModel(
        [
            ScriptedReply(content='[[ASSIGN employee="" pipeline="plan_code_review_test"/]]'),
            ScriptedReply(content="[[PLAN_SPEC]]پلان کار[[/PLAN_SPEC]]"),
            ScriptedReply(content="[[WORK_REPORT]]پیاده‌سازی شد[[/WORK_REPORT]]"),
            ScriptedReply(content="[[REVIEW_PASS]]قابل قبول است[[/REVIEW_PASS]]"),
            ScriptedReply(content="[[TEST_PASS]]سبز است[[/TEST_PASS]]"),
            ScriptedReply(content="[[COMPLETE/]]"),
        ]
    )
    sandbox = FakeSandbox()
    workspace = FakeRepoWorkspace()
    memory = RecordingMemory()
    delivery = NullDeliveryHooks()

    async with session_factory() as session:
        await _team(session, world)
        row = await _task(session, world, stage=Stage.PENDING_MANAGER_REVIEW)
        workspace.diffs[row.id] = DiffResult(
            changed_files=["app/main.py"], text="--- a\n+++ b\n@@\n+x\n"
        )
        await session.commit()
        scheduler = _scheduler(
            session,
            world,
            scripted=scripted,
            sandbox=sandbox,
            workspace=workspace,
            memory=memory,
            delivery=delivery,
        )
        await _run(scheduler, session, 7)
        state = await PipelineStateRepository(session).get(row.id)
        task = await TaskRepository(session).get(row.id)
        gates = await _gates(session, row.id)

    assert state is not None and state.stage is Stage.COMPLETED
    assert state.branch == f"task/{row.id}"
    assert task is not None and task.status is TaskStatus.DONE
    # Five worker turns plus one triage and one final verdict; the merge is the
    # platform's, so it costs no model call.
    assert scripted.calls == 6
    assert gates["diff_guard"] is True
    assert gates["no_progress"] is True
    assert gates["tests"] is True
    assert gates["green"] is True

    commits = _calls(workspace, "commit")
    assert len(commits) == 1 and ENGINE_COMMIT_EMAIL in commits[0]["author"]
    assert _calls(workspace, "merge_into")[0]["target"] == "dev"
    assert _calls(workspace, "push")[0]["branch"] == "dev"
    assert len(_calls(workspace, "remove_worktree")) == 1
    assert delivery.pushed and delivery.pushed[0]["source"] == "engine"
    assert len(memory.finished) == 1

    # The model was never asked to run git: no shell command reached the sandbox
    # for git, and the only command it saw was the test command.
    assert all("git" not in call["command"] for call in sandbox.executed)


async def test_the_scheduler_never_runs_git_through_the_model(world, session_factory):
    scripted = ScriptedModel([ScriptedReply(content="[[WORK_REPORT]]x[[/WORK_REPORT]]")])
    sandbox = FakeSandbox()
    workspace = FakeRepoWorkspace()
    async with session_factory() as session:
        await _team(session, world)
        row = await _task(session, world, stage=Stage.PENDING_IMPLEMENTATION)
        workspace.diffs[row.id] = DiffResult(changed_files=["a.py"], text="d")
        await session.commit()
        scheduler = _scheduler(
            session,
            world,
            scripted=scripted,
            sandbox=sandbox,
            workspace=workspace,
            memory=RecordingMemory(),
            delivery=NullDeliveryHooks(),
        )
        await _run(scheduler, session, 1)
    assert scripted.all_text()  # a prompt was built
    assert all("git" not in call["command"] for call in sandbox.executed)


# ---------------------------------------------------------------------------
# gates send work back
# ---------------------------------------------------------------------------


async def test_red_tests_send_the_work_back_then_green_lets_it_through(
    world, session_factory
):
    scripted = ScriptedModel(
        [
            ScriptedReply(content="[[WORK_REPORT]]تلاش یک[[/WORK_REPORT]]"),
            ScriptedReply(content="[[WORK_REPORT]]تلاش دو[[/WORK_REPORT]]"),
        ]
    )
    sandbox = FakeSandbox(
        commands={
            "pytest": [
                CommandResult(exit_code=1, stdout="", stderr="AssertionError"),
                CommandResult(exit_code=0, stdout="ok"),
            ]
        }
    )
    workspace = FakeRepoWorkspace()
    async with session_factory() as session:
        await _team(session, world)
        row = await _task(session, world, stage=Stage.PENDING_IMPLEMENTATION)
        workspace.diffs[row.id] = DiffResult(changed_files=["a.py"], text="diff-1")
        await session.commit()
        scheduler = _scheduler(
            session,
            world,
            scripted=scripted,
            sandbox=sandbox,
            workspace=workspace,
            memory=RecordingMemory(),
            delivery=NullDeliveryHooks(),
        )
        await _run(scheduler, session, 1)
        after_red = await PipelineStateRepository(session).get(row.id)
        gates_red = await _gates(session, row.id)

        # A different diff on the second attempt, or the no-progress gate fires.
        workspace.diffs[row.id] = DiffResult(changed_files=["a.py", "b.py"], text="diff-2")
        await _run(scheduler, session, 1)
        after_green = await PipelineStateRepository(session).get(row.id)
        gates_green = await _gates(session, row.id)

    assert after_red is not None and after_red.stage is Stage.PENDING_IMPLEMENTATION
    assert after_red.cycle == 1
    assert gates_red["tests"] is False
    assert after_green is not None and after_green.stage is Stage.PENDING_TEST
    assert gates_green["tests"] is True


async def test_a_forbidden_diff_never_reaches_git(world, session_factory):
    scripted = ScriptedModel([ScriptedReply(content="[[WORK_REPORT]]x[[/WORK_REPORT]]")])
    sandbox = FakeSandbox()
    workspace = FakeRepoWorkspace()
    async with session_factory() as session:
        await _team(session, world)
        row = await _task(session, world, stage=Stage.PENDING_IMPLEMENTATION)
        workspace.diffs[row.id] = DiffResult(
            changed_files=[".github/workflows/ci.yml"], text="ci"
        )
        await session.commit()
        scheduler = _scheduler(
            session,
            world,
            scripted=scripted,
            sandbox=sandbox,
            workspace=workspace,
            memory=RecordingMemory(),
            delivery=NullDeliveryHooks(),
        )
        await _run(scheduler, session, 1)
        state = await PipelineStateRepository(session).get(row.id)
        gates = await _gates(session, row.id)

    assert state is not None and state.stage is Stage.PENDING_IMPLEMENTATION
    assert state.cycle == 1
    assert gates["diff_guard"] is False
    assert _calls(workspace, "commit") == []


async def test_a_tester_failure_sends_the_work_back(world, session_factory):
    scripted = ScriptedModel([ScriptedReply(content="[[TEST_FAIL]]قرمز[[/TEST_FAIL]]")])
    workspace = FakeRepoWorkspace()
    async with session_factory() as session:
        await _team(session, world)
        row = await _task(session, world, stage=Stage.PENDING_TEST)
        await session.commit()
        scheduler = _scheduler(
            session,
            world,
            scripted=scripted,
            sandbox=FakeSandbox(),
            workspace=workspace,
            memory=RecordingMemory(),
            delivery=NullDeliveryHooks(),
        )
        await _run(scheduler, session, 1)
        state = await PipelineStateRepository(session).get(row.id)
    assert state is not None and state.stage is Stage.PENDING_IMPLEMENTATION
    assert state.cycle == 1


async def test_a_test_pass_without_green_tests_is_not_believed(world, session_factory):
    scripted = ScriptedModel([ScriptedReply(content="[[TEST_PASS]]سبز[[/TEST_PASS]]")])
    workspace = FakeRepoWorkspace()
    async with session_factory() as session:
        await _team(session, world)
        row = await _task(session, world, stage=Stage.PENDING_TEST)
        await session.commit()
        scheduler = _scheduler(
            session,
            world,
            scripted=scripted,
            sandbox=FakeSandbox(),
            workspace=workspace,
            memory=RecordingMemory(),
            delivery=NullDeliveryHooks(),
        )
        await _run(scheduler, session, 1)
        state = await PipelineStateRepository(session).get(row.id)
        gates = await _gates(session, row.id)
    assert state is not None and state.stage is Stage.PENDING_IMPLEMENTATION
    assert gates["green"] is False


# ---------------------------------------------------------------------------
# analysis, budget, dependencies, failures
# ---------------------------------------------------------------------------


async def test_analysis_only_finishes_without_a_worktree(world, session_factory):
    scripted = ScriptedModel([ScriptedReply(content="[[PLAN_SPEC]]گزارش[[/PLAN_SPEC]]")])
    workspace = FakeRepoWorkspace()
    async with session_factory() as session:
        await _team(session, world)
        row = await _task(
            session,
            world,
            stage=Stage.PENDING_ANALYSIS,
            pipeline=PipelineKey.ANALYSIS_ONLY,
            kind=TaskKind.ANALYSIS,
        )
        await session.commit()
        scheduler = _scheduler(
            session,
            world,
            scripted=scripted,
            sandbox=FakeSandbox(),
            workspace=workspace,
            memory=RecordingMemory(),
            delivery=NullDeliveryHooks(),
        )
        await _run(scheduler, session, 1)
        state = await PipelineStateRepository(session).get(row.id)
        task = await TaskRepository(session).get(row.id)

    assert state is not None and state.stage is Stage.COMPLETED
    assert task is not None and task.status is TaskStatus.DONE
    assert _calls(workspace, "create_task_worktree") == []
    assert _calls(workspace, "commit") == []


async def test_an_overdrawn_wallet_pauses_the_task_before_a_turn_is_spent(
    world, session_factory
):
    scripted = ScriptedModel([ScriptedReply(content="[[ASSIGN/]]")])
    workspace = FakeRepoWorkspace()
    async with session_factory() as session:
        await _team(session, world)
        row = await _task(session, world, stage=Stage.PENDING_MANAGER_REVIEW)
        await session.commit()
        scheduler = _scheduler(
            session,
            world,
            scripted=scripted,
            sandbox=FakeSandbox(),
            workspace=workspace,
            memory=RecordingMemory(),
            delivery=NullDeliveryHooks(),
            wallets=_OverdrawnWalletRepo(),
        )
        await _run(scheduler, session, 1)
        state = await PipelineStateRepository(session).get(row.id)
        task = await TaskRepository(session).get(row.id)

    assert state is not None and state.stage is Stage.PAUSED
    assert task is not None and task.status is TaskStatus.WAITING_APPROVAL
    assert scripted.calls == 0
    assert BUDGET_PAUSED_MESSAGE


async def test_a_zero_wallet_does_not_pause_a_fresh_org(world, session_factory):
    """A 0.0 balance is every org's starting state, so it must not pause the task.

    This is an intentional behaviour change. The scheduler used to pause on
    ``balance <= 0``, which froze every new org before its first turn: usage is
    post-paid and nothing debits a wallet. Only an overdrawn wallet (below zero)
    pauses now. The old test that encoded the previous rule was renamed and
    updated, as PLAN-v2 §7.4 requires for a deliberate behaviour change.
    """
    scripted = ScriptedModel([ScriptedReply(content="[[ASSIGN/]]")])
    workspace = FakeRepoWorkspace()
    async with session_factory() as session:
        await _team(session, world)
        row = await _task(session, world, stage=Stage.PENDING_MANAGER_REVIEW)
        await session.commit()
        scheduler = _scheduler(
            session,
            world,
            scripted=scripted,
            sandbox=FakeSandbox(),
            workspace=workspace,
            memory=RecordingMemory(),
            delivery=NullDeliveryHooks(),
            wallets=_Wallets(_ZeroWallet()),
        )
        await _run(scheduler, session, 1)
        state = await PipelineStateRepository(session).get(row.id)

    assert state is not None and state.stage is not Stage.PAUSED
    assert scripted.calls == 1


async def test_a_task_waits_for_its_dependencies(world, session_factory):
    scripted = ScriptedModel(
        [
            ScriptedReply(content='[[TEST_PASS waive="مشکل شناخته‌شده"]]سبز[[/TEST_PASS]]'),
            ScriptedReply(content="[[WORK_REPORT]]done[[/WORK_REPORT]]"),
        ]
    )
    workspace = FakeRepoWorkspace()
    async with session_factory() as session:
        await _team(session, world)
        first = await _task(session, world, stage=Stage.PENDING_TEST)
        second = await _task(session, world, stage=Stage.PENDING_IMPLEMENTATION)
        state = await PipelineStateRepository(session).get(second.id)
        assert state is not None
        await PipelineStateRepository(session).patch(
            task_id=second.id, expected_version=state.version, depends_on=[first.id]
        )
        workspace.diffs[second.id] = DiffResult(changed_files=["b.py"], text="diff-b")
        await session.commit()
        scheduler = _scheduler(
            session,
            world,
            scripted=scripted,
            sandbox=FakeSandbox(),
            workspace=workspace,
            memory=RecordingMemory(),
            delivery=NullDeliveryHooks(),
        )
        await _run(scheduler, session, 1)
        blocked = await PipelineStateRepository(session).get(second.id)

        # The dependency finished, so the next tick lets the work run.
        dependency = await PipelineStateRepository(session).get(first.id)
        assert dependency is not None
        await PipelineStateRepository(session).patch(
            task_id=first.id,
            expected_version=dependency.version,
            stage=str(Stage.COMPLETED),
        )
        await _run(scheduler, session, 1)
        unblocked = await PipelineStateRepository(session).get(second.id)

    assert blocked is not None and blocked.stage is Stage.PENDING_IMPLEMENTATION
    assert unblocked is not None and unblocked.stage is Stage.PENDING_TEST


async def test_a_transient_failure_is_retried_then_handed_to_the_manager(
    world, session_factory
):
    scripted = ScriptedModel(
        [
            ScriptedReply(raises=ProviderUnavailableError("429 rate limit")),
            ScriptedReply(raises=ProviderUnavailableError("429 rate limit")),
            ScriptedReply(raises=ProviderUnavailableError("429 rate limit")),
        ]
    )
    workspace = FakeRepoWorkspace()
    async with session_factory() as session:
        await _team(session, world)
        row = await _task(session, world, stage=Stage.PENDING_IMPLEMENTATION)
        workspace.diffs[row.id] = DiffResult(changed_files=["a.py"], text="d")
        await session.commit()
        scheduler = _scheduler(
            session,
            world,
            scripted=scripted,
            sandbox=FakeSandbox(),
            workspace=workspace,
            memory=RecordingMemory(),
            delivery=NullDeliveryHooks(),
        )
        await _run(scheduler, session, 3)
        state = await PipelineStateRepository(session).get(row.id)

    assert state is not None and state.stage is Stage.ERROR_REVIEW
    assert state.retry_count == 2


async def test_a_merge_conflict_sends_the_work_back_for_another_pass(
    world, session_factory
):
    scripted = ScriptedModel()
    workspace = FakeRepoWorkspace()
    workspace.sync_outcomes = [MergeOutcome(ok=False, conflict_files=["app/main.py"])]
    async with session_factory() as session:
        await _team(session, world)
        row = await _task(session, world, stage=Stage.MERGE_REQUEST)
        workspace.diffs[row.id] = DiffResult(changed_files=["app/main.py"], text="d")
        await session.commit()
        scheduler = _scheduler(
            session,
            world,
            scripted=scripted,
            sandbox=FakeSandbox(),
            workspace=workspace,
            memory=RecordingMemory(),
            delivery=NullDeliveryHooks(),
        )
        await _run(scheduler, session, 1)
        state = await PipelineStateRepository(session).get(row.id)

    assert state is not None and state.stage is Stage.PENDING_IMPLEMENTATION
    assert state.retry_count == 1 and state.cycle == 1