roboco/tests/unit/gateway/test_work_session_auto_create.py

"""Wave C4 (2026-05-12): claim auto-creates a WorkSession row.

Smoke run 3 showed task.work_session_id null on every task. Pre-gateway
created the row at claim time so the panel/PR/merge subsystems could
track agent-per-task git activity (branch, commits, PR number/url,
merge status). Restoring that side-effect.

The choreographer's _claim_plan_start_run (and _resume_from_claimed)
calls TaskService.ensure_work_session(task_id, agent_id) after the
task reaches in_progress, which creates the WorkSession and stores
its id on the task.
"""

from __future__ import annotations

from unittest.mock import AsyncMock, MagicMock
from uuid import UUID, uuid4

import pytest
from roboco.services.gateway.choreographer import Choreographer, ChoreographerDeps

# #172: a developer fresh claim must carry a substantive step checklist.
# Inert on re-entry/error/non-dev paths, so safe to pass everywhere.
_STEPS = [
    {
        "title": "Implement the change",
        "description": (
            "edit the target file, add tests, run them, and stage the "
            "change for commit on the task branch"
        ),
    }
]
# Full parity: a fresh dev claim authors the same rich plan a PM does.
# These satisfy _dev_plan_gate (plan/approach >= 150 chars,
# technical_considerations, risks).
_GOOD_PLAN = (
    "Append the timestamp HTML comment to the very bottom of README.md without "
    "touching any other line, then commit it on the task branch and open a PR. "
    "Verify the diff is a single-line addition before submitting for QA."
)
_GOOD_TC = ["Use a trailing newline so the comment sits on its own line."]
_GOOD_RISKS = [
    {
        "risk": "An accidental reformat of README.md balloons the diff.",
        "mitigation": "Append only; assert the diff touches one line pre-commit.",
    }
]


def _make_task_svc(agent_id: UUID, task_id: UUID, *, status: str) -> AsyncMock:
    """Build a TaskService AsyncMock that completes the (claim, set_plan, start)
    sequence and returns a task with branch_name set (as the real service does
    after auto-creating the branch during claim side-effects).
    """
    in_progress_task = MagicMock(
        id=task_id,
        status="in_progress",
        plan={"text": "plan text"},
        assigned_to=agent_id,
        branch_name="feature/backend/abc",
        work_session_id=None,
        commits=[],
        pr_number=None,
        quick_context=None,
        team="backend",
        task_type="code",
        parent_task_id=None,
        sequence=0,
        project_id=uuid4(),
    )
    task_svc = AsyncMock()
    task_svc.get.return_value = MagicMock(
        id=task_id,
        status=status,
        plan=None,
        assigned_to=None,
        branch_name="feature/backend/abc",
        work_session_id=None,
        commits=[],
        pr_number=None,
        quick_context=None,
        team="backend",
        task_type="code",
        parent_task_id=None,
        sequence=0,
        project_id=uuid4(),
    )
    task_svc.agent_for.return_value = MagicMock(
        id=agent_id, role="developer", team="backend", slug=None
    )
    task_svc.list_in_progress_for_agent.return_value = []
    task_svc.list_paused_for_agent.return_value = []
    task_svc.get_subtasks.return_value = []
    task_svc.claim.return_value = MagicMock(
        id=task_id, status="claimed", plan=None, assigned_to=agent_id
    )
    task_svc.set_plan.return_value = MagicMock(
        id=task_id, status="claimed", plan={"text": "plan text"}, assigned_to=agent_id
    )
    task_svc.start.return_value = in_progress_task
    task_svc.ensure_work_session.return_value = None
    task_svc.session = MagicMock()
    task_svc.session.begin_nested = MagicMock(
        return_value=MagicMock(
            __aenter__=AsyncMock(return_value=None),
            __aexit__=AsyncMock(return_value=False),
        )
    )
    return task_svc


def _make_deps(task_svc: AsyncMock) -> ChoreographerDeps:
    evidence_repo = AsyncMock()
    for method in (
        "list_unread_a2a",
        "list_unread_mentions",
        "list_pending_notifications",
        "task_metadata_gaps",
        "recent_team_activity",
        "blockers_in_lane",
        "journal_highlights_for_task",
    ):
        getattr(evidence_repo, method).return_value = []
    return ChoreographerDeps(
        task=task_svc,
        work_session=AsyncMock(),
        git=AsyncMock(),
        a2a=AsyncMock(),
        journal=AsyncMock(),
        audit=AsyncMock(),
        evidence_repo=evidence_repo,
    )


@pytest.mark.asyncio
async def test_i_will_work_on_calls_ensure_work_session() -> None:
    """After a successful i_will_work_on, TaskService.ensure_work_session is
    called once with (task_id, agent_id) so a WorkSession row is created and
    task.work_session_id is populated.

    Wave C4 (2026-05-12): pre-gateway parity. Smoke run 3 showed
    task.work_session_id null on every in_progress task.
    """
    agent_id = uuid4()
    task_id = uuid4()
    task_svc = _make_task_svc(agent_id, task_id, status="pending")
    deps = _make_deps(task_svc)
    c = Choreographer(deps)

    env = await c.i_will_work_on(
        agent_id,
        task_id,
        plan=_GOOD_PLAN,
        steps=_STEPS,
        technical_considerations=_GOOD_TC,
        risks=_GOOD_RISKS,
    )

    assert env.error is None, f"Expected ok, got error={env.error} msg={env.message}"
    assert env.status == "in_progress"
    task_svc.ensure_work_session.assert_awaited_once_with(task_id, agent_id)


@pytest.mark.asyncio
async def test_i_will_plan_calls_ensure_work_session() -> None:
    """PMs also get a WorkSession via ensure_work_session (cell_pm role, planning
    task). Both i_will_work_on and i_will_plan share _claim_plan_start_run so
    the same hook fires for both.

    Wave C4 (2026-05-12): pre-gateway parity.
    """
    pm_agent_id = uuid4()
    task_id = uuid4()

    in_progress_task = MagicMock(
        id=task_id,
        status="in_progress",
        plan={"approach": "plan text", "sub_tasks": ["t1"]},
        assigned_to=pm_agent_id,
        branch_name="feature/main_pm/abc",
        work_session_id=None,
        commits=[],
        pr_number=None,
        quick_context=None,
        team="main_pm",
        task_type="planning",
        parent_task_id=None,
        sequence=0,
        project_id=uuid4(),
    )
    task_svc = AsyncMock()
    task_svc.get.return_value = MagicMock(
        id=task_id,
        status="pending",
        plan=None,
        assigned_to=None,
        branch_name="feature/main_pm/abc",
        work_session_id=None,
        commits=[],
        pr_number=None,
        quick_context=None,
        team="main_pm",
        task_type="planning",
        parent_task_id=None,
        sequence=0,
        project_id=uuid4(),
    )
    task_svc.agent_for.return_value = MagicMock(
        id=pm_agent_id, role="cell_pm", team="backend", slug=None
    )
    task_svc.list_in_progress_for_agent.return_value = []
    task_svc.list_paused_for_agent.return_value = []
    task_svc.get_subtasks.return_value = []
    task_svc.claim.return_value = MagicMock(
        id=task_id, status="claimed", plan=None, assigned_to=pm_agent_id
    )
    task_svc.set_plan.return_value = MagicMock(
        id=task_id,
        status="claimed",
        plan={"approach": "plan text"},
        assigned_to=pm_agent_id,
    )
    task_svc.start.return_value = in_progress_task
    task_svc.ensure_work_session.return_value = None
    task_svc.session = MagicMock()
    task_svc.session.begin_nested = MagicMock(
        return_value=MagicMock(
            __aenter__=AsyncMock(return_value=None),
            __aexit__=AsyncMock(return_value=False),
        )
    )
    evidence_repo = AsyncMock()
    for method in (
        "list_unread_a2a",
        "list_unread_mentions",
        "list_pending_notifications",
        "task_metadata_gaps",
        "recent_team_activity",
        "blockers_in_lane",
        "journal_highlights_for_task",
    ):
        getattr(evidence_repo, method).return_value = []
    deps = ChoreographerDeps(
        task=task_svc,
        work_session=AsyncMock(),
        git=AsyncMock(),
        a2a=AsyncMock(),
        journal=AsyncMock(),
        audit=AsyncMock(),
        evidence_repo=evidence_repo,
    )
    c = Choreographer(deps)

    rich_plan = {
        "approach": (
            "Single-cell decomposition: be-dev-1 owns the change end to end "
            "— branch, implement, test, open PR; QA reviews after the PR "
            "opens, documentation follows, then be-pm completes and submits "
            "up. No cross-cell dependencies for this planning task."
        ),
        "sub_tasks": [
            {
                "title": "t1",
                "description": (
                    "be-dev-1 implements the change with tests and opens the "
                    "leaf PR for QA review."
                ),
            }
        ],
    }
    env = await c.i_will_plan(
        pm_agent_id, task_id, plan="plan text", rich_plan=rich_plan
    )

    assert env.error is None, f"Expected ok, got error={env.error} msg={env.message}"
    assert env.status == "in_progress"
    task_svc.ensure_work_session.assert_awaited_once_with(task_id, pm_agent_id)


@pytest.mark.asyncio
async def test_ensure_work_session_not_called_when_start_fails() -> None:
    """If start() returns None (task in wrong state), ensure_work_session must
    NOT be called — the WorkSession must not be created for a failed transition.
    """
    agent_id = uuid4()
    task_id = uuid4()
    task_svc = _make_task_svc(agent_id, task_id, status="pending")
    task_svc.start.return_value = None  # start fails
    deps = _make_deps(task_svc)
    c = Choreographer(deps)

    env = await c.i_will_work_on(
        agent_id,
        task_id,
        plan=_GOOD_PLAN,
        steps=_STEPS,
        technical_considerations=_GOOD_TC,
        risks=_GOOD_RISKS,
    )

    assert env.error == "invalid_state"
    task_svc.ensure_work_session.assert_not_awaited()