zk-data-agent/tests/test_bash_bg.py

"""Tests for the bash run_in_background path (local fallback only).

Covers:
- BashBgStore: record_start / mark_completed / mark_auto_resumed race semantics.
- _run_bash_background local fallback: spawns a detached process with file
  redirects, registers it via the bg_register callback, output/pid/exit_code
  files materialize.
- _run_bash_status / _run_bash_kill: route through bg_status_query / bg_kill_request
  callbacks and format their result.

Remote (jupyter_runtime) path is intentionally not exercised here — that
requires a live wsh account and is covered by the e2e checklist in the plan.
"""

from __future__ import annotations

import time
from pathlib import Path

import pytest

from src.agent_tool_core import ToolExecutionContext
from src.agent_tools import (
    _run_bash_background,
    _run_bash_kill,
    _run_bash_status,
)
from src.agent_types import AgentPermissions
from src.bash_bg_store import BashBgStore, BgTaskSpec, BgTaskStatus


def _make_context(
    tmp_path: Path,
    *,
    bg_register=None,
    bg_status_query=None,
    bg_kill_request=None,
) -> ToolExecutionContext:
    return ToolExecutionContext(
        root=tmp_path,
        command_timeout_seconds=10.0,
        max_output_chars=4096,
        permissions=AgentPermissions(allow_shell_commands=True),
        scratchpad_directory=tmp_path,
        bg_register=bg_register,
        bg_status_query=bg_status_query,
        bg_kill_request=bg_kill_request,
        account_id='acct-test',
        session_id='sess-test',
        run_id='run-test',
    )


# ----- BashBgStore --------------------------------------------------------


def test_store_record_and_complete(tmp_path: Path) -> None:
    store = BashBgStore(tmp_path / 'bg.db')
    spec = BgTaskSpec(
        task_id='bg_abc',
        account_key='acct',
        session_id='sess',
        run_id='run',
        pid=12345,
        task_dir=str(tmp_path / 'bg_abc'),
        output_path=str(tmp_path / 'bg_abc' / 'output'),
        pid_path=str(tmp_path / 'bg_abc' / 'pid'),
        exit_code_path=str(tmp_path / 'bg_abc' / 'exit_code'),
        command='echo hi',
        started_at=time.time(),
        wait_for_completion=True,
    )
    store.record_start(spec)
    row = store.get('bg_abc')
    assert row is not None
    assert row['status'] == 'running'
    assert row['exit_code'] is None

    store.mark_completed('bg_abc', exit_code=0)
    row = store.get('bg_abc')
    assert row['status'] == 'completed'
    assert row['exit_code'] == 0
    assert row['finished_at'] is not None


def test_store_mark_auto_resumed_is_race_safe(tmp_path: Path) -> None:
    store = BashBgStore(tmp_path / 'bg.db')
    spec = BgTaskSpec(
        task_id='bg_race',
        account_key='acct',
        session_id='sess',
        run_id='run',
        pid=1,
        task_dir='/tmp/x',
        output_path='/tmp/x/output',
        pid_path='/tmp/x/pid',
        exit_code_path='/tmp/x/exit_code',
        command='true',
        started_at=time.time(),
        wait_for_completion=True,
    )
    store.record_start(spec)
    store.mark_completed('bg_race', exit_code=0)

    # First caller wins; subsequent callers always see False.
    assert store.mark_auto_resumed('bg_race') is True
    assert store.mark_auto_resumed('bg_race') is False
    assert store.mark_auto_resumed('bg_race') is False


def test_store_mark_killed_only_running(tmp_path: Path) -> None:
    store = BashBgStore(tmp_path / 'bg.db')
    spec = BgTaskSpec(
        task_id='bg_kill',
        account_key='acct',
        session_id='sess',
        run_id='run',
        pid=2,
        task_dir='/tmp/y',
        output_path='/tmp/y/output',
        pid_path='/tmp/y/pid',
        exit_code_path='/tmp/y/exit_code',
        command='sleep 100',
        started_at=time.time(),
        wait_for_completion=False,
    )
    store.record_start(spec)
    store.mark_killed('bg_kill')
    row = store.get('bg_kill')
    assert row['status'] == 'cancelled'

    # Re-killing a non-running task is a no-op (status stays cancelled).
    store.mark_killed('bg_kill')
    row = store.get('bg_kill')
    assert row['status'] == 'cancelled'


# ----- _run_bash_background local fallback --------------------------------


def test_run_bash_background_local_spawns_detached(tmp_path: Path) -> None:
    captured: list[BgTaskSpec] = []

    ctx = _make_context(tmp_path, bg_register=captured.append)
    content, metadata = _run_bash_background(
        {
            'command': 'echo hello-from-bg',
            'wait_for_completion': True,
        },
        ctx,
    )

    assert metadata['action'] == 'bash_bg'
    task_id = metadata['task_id']
    assert task_id.startswith('bg_')
    assert metadata['wait_for_completion'] is True
    assert 'pid=' in content

    # bg_register received exactly one spec, with the agent's session/account.
    assert len(captured) == 1
    spec = captured[0]
    assert spec.task_id == task_id
    assert spec.account_key == 'acct-test'
    assert spec.session_id == 'sess-test'
    assert spec.run_id == 'run-test'
    assert spec.command == 'echo hello-from-bg'
    assert spec.wait_for_completion is True

    # Wait for the detached subprocess to finish — bounded by exit_code file.
    deadline = time.monotonic() + 5.0
    exit_code_path = Path(spec.exit_code_path)
    while time.monotonic() < deadline and not exit_code_path.exists():
        time.sleep(0.05)
    assert exit_code_path.exists(), 'exit_code file did not appear in 5s'
    assert exit_code_path.read_text(encoding='utf-8').strip() == '0'

    output_text = Path(spec.output_path).read_text(encoding='utf-8')
    assert 'hello-from-bg' in output_text


def test_run_bash_background_failed_command_records_nonzero_exit(
    tmp_path: Path,
) -> None:
    captured: list[BgTaskSpec] = []
    ctx = _make_context(tmp_path, bg_register=captured.append)

    _content, metadata = _run_bash_background(
        {'command': 'exit 7'},
        ctx,
    )
    spec = captured[0]
    assert metadata['action'] == 'bash_bg'

    deadline = time.monotonic() + 5.0
    exit_code_path = Path(spec.exit_code_path)
    while time.monotonic() < deadline and not exit_code_path.exists():
        time.sleep(0.05)
    assert exit_code_path.exists()
    assert exit_code_path.read_text(encoding='utf-8').strip() == '7'


def test_run_bash_background_requires_shell_permission(tmp_path: Path) -> None:
    from src.agent_tool_core import ToolPermissionError

    ctx = ToolExecutionContext(
        root=tmp_path,
        command_timeout_seconds=10.0,
        max_output_chars=4096,
        permissions=AgentPermissions(allow_shell_commands=False),
        scratchpad_directory=tmp_path,
    )
    with pytest.raises(ToolPermissionError):
        _run_bash_background({'command': 'echo nope'}, ctx)


# ----- bash_status / bash_kill via callbacks ------------------------------


def test_bash_status_routes_through_callback(tmp_path: Path) -> None:
    fake_status = BgTaskStatus(
        task_id='bg_xyz',
        status='completed',
        pid=999,
        started_at=time.time() - 5.0,
        finished_at=time.time(),
        exit_code=0,
        output_path='/tmp/x/output',
        output_preview='all good\n',
        auto_resumed=False,
    )

    queries: list[str] = []

    def query(task_id: str) -> BgTaskStatus | None:
        queries.append(task_id)
        return fake_status

    ctx = _make_context(tmp_path, bg_status_query=query)
    content, metadata = _run_bash_status({'task_id': 'bg_xyz'}, ctx)

    assert queries == ['bg_xyz']
    assert metadata['action'] == 'bash_status'
    assert metadata['status'] == 'completed'
    assert metadata['exit_code'] == 0
    assert 'all good' in content
    assert 'task_id=bg_xyz' in content


def test_bash_status_unknown_task_raises(tmp_path: Path) -> None:
    from src.agent_tool_core import ToolExecutionError

    def query(_task_id: str):
        return None

    ctx = _make_context(tmp_path, bg_status_query=query)
    with pytest.raises(ToolExecutionError):
        _run_bash_status({'task_id': 'bg_missing'}, ctx)


def test_bash_status_without_callback_raises(tmp_path: Path) -> None:
    from src.agent_tool_core import ToolExecutionError

    ctx = _make_context(tmp_path)
    with pytest.raises(ToolExecutionError):
        _run_bash_status({'task_id': 'bg_anything'}, ctx)


def test_bash_kill_routes_through_callback(tmp_path: Path) -> None:
    killed: list[str] = []

    def kill(task_id: str) -> bool:
        killed.append(task_id)
        return True

    ctx = _make_context(tmp_path, bg_kill_request=kill)
    content, metadata = _run_bash_kill({'task_id': 'bg_kk'}, ctx)
    assert killed == ['bg_kk']
    assert metadata['killed'] is True
    assert 'bg_kk' in content


def test_bash_kill_failure_reports(tmp_path: Path) -> None:
    ctx = _make_context(tmp_path, bg_kill_request=lambda _t: False)
    content, metadata = _run_bash_kill({'task_id': 'bg_dead'}, ctx)
    assert metadata['killed'] is False
    assert '未能终止' in content


# ----- end-to-end: spawn + register + sync via real BashBgStore -----------


def test_local_bg_lifecycle_with_real_store(tmp_path: Path) -> None:
    """Plug a real BashBgStore into bg_register and check the row appears."""
    store = BashBgStore(tmp_path / 'bg.db')

    def register(spec: BgTaskSpec) -> None:
        store.record_start(spec)

    ctx = _make_context(tmp_path, bg_register=register)
    _content, metadata = _run_bash_background(
        {'command': 'echo lifecycle && exit 0'}, ctx,
    )
    task_id = metadata['task_id']
    row = store.get(task_id)
    assert row is not None
    assert row['status'] == 'running'

    # Wait for child to write exit_code (more reliable than kill -0, since
    # the unreaped child becomes a zombie and kill -0 still reports it alive).
    deadline = time.monotonic() + 5.0
    exit_code_path = Path(row['exit_code_path'])
    while time.monotonic() < deadline and not exit_code_path.exists():
        time.sleep(0.05)
    if not exit_code_path.exists():
        pytest.fail('background process did not exit within 5 seconds')

    # Simulate manager finalize step.
    exit_code = int(
        exit_code_path.read_text(encoding='utf-8').strip()
    )
    store.mark_completed(task_id, exit_code=exit_code)
    row = store.get(task_id)
    assert row['status'] == 'completed'
    assert row['exit_code'] == 0