multi-agent-mux/tests/test_tier3_integration.py

import os
import subprocess
import json
import sqlite3
import pytest
import shutil
import yaml
from pathlib import Path

# Helper to run mutation on agent-sessions.yaml using atomic_dump_yaml in bash
def run_mutation(mam_sandbox, mutation_str, env=None):
    lib_path = mam_sandbox / ".agents" / "skills" / "lib.sh"
    yaml_path = mam_sandbox / ".mam" / "agent-sessions.yaml"
    cmd_str = f"source {lib_path} && atomic_dump_yaml {yaml_path}"
    run_env = dict(os.environ)
    if env:
        run_env.update(env)
    res = subprocess.run(["bash", "-c", cmd_str], input=mutation_str, capture_output=True, text=True, env=run_env)
    return res

def test_integration_create_options_combination(mam_sandbox, mock_herdr, mock_agents, monkeypatch):
    """
    Tier 3: Create Session Integration
    Covers pairwise combination: --onboard + --submit-job + --wrapper + HERDR_SERVER_NAME env override
    """
    tmp_path = mam_sandbox
    script_path = tmp_path / "skills" / "multi-agent-mux-create" / "scripts" / "create_session.sh"
    
    # Set HERDR_SERVER_NAME env var to test environment override combination
    monkeypatch.setenv("HERDR_SERVER_NAME", "custom_server")
    
    cmd = [
        "bash", str(script_path),
        "--workspace", str(tmp_path),
        "--agent", "claude",
        "--role", "Creator",
        "--onboard",
        "--submit-job", "Test onboard prompt",
        "--wrapper"
    ]
    
    res = subprocess.run(cmd, capture_output=True, text=True, cwd=str(tmp_path))
    if res.returncode != 0:
        with open(mock_herdr, 'r') as f:
            state = json.load(f)
        debug_info = (
            f"MOCK HERDR CALLS: {state.get('calls', [])}\n"
            f"MOCK HERDR AGENTS: {state.get('agents', {})}\n"
        )
        assert False, f"Stdout: {res.stdout}\nStderr: {res.stderr}\nDebug:\n{debug_info}"
        
    assert res.returncode == 0
    
    # Verify it is registered in yaml
    yaml_path = tmp_path / ".mam" / "agent-sessions.yaml"
    assert yaml_path.exists()
    
    with open(yaml_path, 'r') as f:
        reg = yaml.safe_load(f)
        
    sessions = reg.get("herdr_sessions", [])
    assert len(sessions) == 1
    session = sessions[0]
    
    assert session["role"] == "Creator"
    assert session["herdr_server"] == "custom_server"
    assert session["delegate_job_id"] is not None
    assert session["status"] == "running"
    
    # Verify delegate job JSON file exists and contains correct info
    job_id = session["delegate_job_id"]
    job_file = tmp_path / ".mam" / "jobs" / f"{job_id}.json"
    assert job_file.exists()
    
    with open(job_file, 'r') as jf:
        job_data = json.load(jf)
        
    assert job_data["job_id"] == job_id
    assert "Test onboard prompt" in job_data["prompt"]
    assert job_data["agent_session"] == f"herdr:{session['name']}"


def test_integration_stop_purge_combination(mam_sandbox, mock_herdr, mock_agents):
    """
    Tier 3: Stop Session Integration
    Covers pairwise combination: --purge-conversation + --yes + HOME override
    """
    tmp_path = mam_sandbox
    
    # 1. Run create_session to register a session with isolation enabled
    create_script = tmp_path / "skills" / "multi-agent-mux-create" / "scripts" / "create_session.sh"
    cmd_create = [
        "bash", str(create_script),
        "--workspace", str(tmp_path),
        "--agent", "claude",
        "--role", "Creator"
    ]
    res_create = subprocess.run(cmd_create, capture_output=True, text=True, cwd=str(tmp_path))
    assert res_create.returncode == 0, f"Stdout: {res_create.stdout}\nStderr: {res_create.stderr}"
    
    # Get session name and isolation uuid
    yaml_path = tmp_path / ".mam" / "agent-sessions.yaml"
    with open(yaml_path, 'r') as f:
        reg = yaml.safe_load(f)
    session = reg["herdr_sessions"][0]
    session_name = session["name"]
    iso_uuid = session["isolation"]["uuid"]
    iso_root = Path(session["isolation"]["root"])
    
    # Ensure isolation home root folder is provisioned
    assert iso_root.exists()
    
    # Locate dynamically generated conversation file in claude projects
    key = str(tmp_path).replace('/', '-').replace('_', '-')
    proj_dir = iso_root / "projects" / key
    assert proj_dir.exists(), f"Expected isolated projects directory {proj_dir} to exist"
    
    jsonls = list(proj_dir.glob("*.jsonl"))
    assert len(jsonls) == 1, f"Expected exactly 1 jsonl file in {proj_dir}, found {jsonls}"
    jsonl_file = jsonls[0]
    assert jsonl_file.exists()
    
    # Update mock herdr's state to match the running agent so the stop kill-chain works gracefully
    with open(mock_herdr, 'r') as f:
        state = json.load(f)
    state["agents"][session_name] = {
        "status": "running",
        "agent": "claude",
        "cwd": str(tmp_path),
        "pid": 12345,
        "pane_id": "w1:p1",
        "command": "claude",
        "buffer": "Anthropic Claude Ready"
    }
    with open(mock_herdr, 'w') as f:
        json.dump(state, f, indent=2)

    # 2. Calling stop_session without --yes fails for --purge-conversation
    stop_script = tmp_path / "skills" / "multi-agent-mux-stop" / "scripts" / "stop_session.sh"
    cmd_stop_no_yes = [
        "bash", str(stop_script),
        "--session", session_name,
        "--purge-conversation"
    ]
    res_stop_no_yes = subprocess.run(cmd_stop_no_yes, capture_output=True, text=True, cwd=str(tmp_path))
    assert res_stop_no_yes.returncode == 3
    assert "DANGER: --purge-conversation" in res_stop_no_yes.stdout
    
    # Ensure nothing was deleted yet
    assert iso_root.exists()
    assert jsonl_file.exists()
    
    # 3. Run stop_session with --yes
    cmd_stop_yes = [
        "bash", str(stop_script),
        "--session", session_name,
        "--purge-conversation",
        "--yes"
    ]
    res_stop_yes = subprocess.run(cmd_stop_yes, capture_output=True, text=True, cwd=str(tmp_path))
    assert res_stop_yes.returncode == 0, f"Stdout: {res_stop_yes.stdout}\nStderr: {res_stop_yes.stderr}"
    
    # Assertions:
    # - Isolation directory deleted
    assert not iso_root.exists()
    # - Conversation files deleted
    assert not jsonl_file.exists()
    # - Session completely removed from YAML/DB registry
    with open(yaml_path, 'r') as f:
        reg = yaml.safe_load(f)
    assert not any(s["name"] == session_name for s in reg.get("herdr_sessions", []))
    
    db_path = tmp_path / ".mam" / "agent-sessions.db"
    conn = sqlite3.connect(str(db_path))
    cursor = conn.cursor()
    cursor.execute("SELECT COUNT(*) FROM sessions WHERE name=?", (session_name,))
    assert cursor.fetchone()[0] == 0
    conn.close()


def test_integration_resume_fallbacks(mam_sandbox, mock_herdr, mock_agents):
    """
    Tier 3: Resume Session Fallbacks
    Covers boundary checks, invalid / missing UUIDs, and workspace disk scan fallback resolution.
    """
    tmp_path = mam_sandbox
    
    # 1. Calling resume_session.sh with missing arguments
    resume_script = tmp_path / "skills" / "multi-agent-mux-resume" / "scripts" / "resume_session.sh"
    cmd_missing = ["bash", str(resume_script), "--session", "test-session"]
    res_missing = subprocess.run(cmd_missing, capture_output=True, text=True, cwd=str(tmp_path))
    assert res_missing.returncode == 2
    assert "ERROR:" in res_missing.stderr
    
    # 2. Calling with non-existent session
    cmd_nonexistent = ["bash", str(resume_script), "--workspace", str(tmp_path), "--agent", "claude", "--session", "non-existent"]
    res_nonexistent = subprocess.run(cmd_nonexistent, capture_output=True, text=True, cwd=str(tmp_path))
    assert res_nonexistent.returncode == 1
    assert "ERROR: No saved session for" in res_nonexistent.stderr
    
    # 3. Fallback resolution: set up a stopped session in YAML with NO own ID
    session_name = "test-fallback-creator-claude"
    mutation = """
d['herdr_sessions'] = [{
    'name': 'test-fallback-creator-claude',
    'status': 'stopped',
    'role': 'Creator',
    'pane': {'cwd': 'WS_PLACEHOLDER'}
}]
""".replace("WS_PLACEHOLDER", str(tmp_path))
    res_mut = run_mutation(tmp_path, mutation)
    assert res_mut.returncode == 0
    
    # Seed a jsonl file representing a conversation under the workspace path
    key = str(tmp_path).replace('/', '-').replace('_', '-')
    proj_dir = tmp_path / ".claude" / "projects" / key
    proj_dir.mkdir(parents=True, exist_ok=True)
    
    scanned_uuid = "scanned-uuid-xyz-123"
    (proj_dir / f"{scanned_uuid}.jsonl").write_text(json.dumps({"sessionId": scanned_uuid}) + "\n")
    
    # Run resume_session.sh
    cmd_resume = ["bash", str(resume_script), "--workspace", str(tmp_path), "--agent", "claude", "--session", session_name]
    res_resume = subprocess.run(cmd_resume, capture_output=True, text=True, cwd=str(tmp_path))
    assert res_resume.returncode == 0, f"Stdout: {res_resume.stdout}\nStderr: {res_resume.stderr}"
    
    # Verify the session resumed using the resolved UUID
    with open(mock_herdr, 'r') as f:
        state = json.load(f)
    
    calls = state.get("calls", [])
    # Find the agent start/new-session call
    new_sess_call = None
    for call in calls:
        if "agent" in call and "start" in call and session_name in call:
            new_sess_call = call
            break
            
    assert new_sess_call is not None, f"Calls: {calls}"
    # Verify the argument contains the resolved scanned_uuid
    assert any(scanned_uuid in arg for arg in new_sess_call), f"Resume arguments: {new_sess_call}"


def test_integration_reconcile_diff_formats(mam_sandbox, mock_herdr):
    """
    Tier 3: Reconcile Drift Detection and Diff Formatting
    Covers outputs of reconcile.sh --once --emit-diff under dry-run and actual mutation.
    """
    tmp_path = mam_sandbox
    reconcile_script = tmp_path / "skills" / "multi-agent-mux-monitor" / "scripts" / "reconcile.sh"
    
    # Seed a running session in YAML/DB that is NOT in herdr (Drift class A)
    session_name = "drift-a-creator-claude"
    mutation = """
d['herdr_sessions'] = [{
    'name': 'drift-a-creator-claude',
    'status': 'running',
    'role': 'Creator',
    'pane': {
        'cwd': 'WS_PLACEHOLDER',
        'pid': 9876,
        'cmd': 'claude',
        'cmd_full': 'claude'
    }
}]
""".replace("WS_PLACEHOLDER", str(tmp_path))
    res_mut = run_mutation(tmp_path, mutation)
    assert res_mut.returncode == 0
    
    db_path = tmp_path / ".mam" / "agent-sessions.db"
    assert db_path.exists()
    
    # Verify DB content
    conn = sqlite3.connect(str(db_path))
    db_rows = conn.execute("SELECT * FROM sessions").fetchall()
    db_state = conn.execute("SELECT * FROM state").fetchall()
    conn.close()
    
    # Verify YAML content
    yaml_path = tmp_path / ".mam" / "agent-sessions.yaml"
    with open(yaml_path, 'r') as f:
        yaml_content = f.read()
        
    # Verify load_state_json output
    lib_path = tmp_path / ".agents" / "skills" / "lib.sh"
    res_load = subprocess.run(["bash", "-c", f"source {lib_path} && load_state_json"], capture_output=True, text=True)
    
    # Verify herdr ls -F output
    res_ls = subprocess.run(["herdr", "ls", "-F", "#{session_name}|#{session_created}"], capture_output=True, text=True)

    # 1. Run reconcile with --dry-run
    cmd_dry = ["bash", str(reconcile_script), "--once", "--emit-diff", "--dry-run"]
    res_dry = subprocess.run(cmd_dry, capture_output=True, text=True, cwd=str(tmp_path))
    assert res_dry.returncode == 0, f"Stderr: {res_dry.stderr}"
    
    # Verify stdout is valid JSON and reports class A drift
    try:
        data_dry = json.loads(res_dry.stdout)
    except Exception as e:
        assert False, f"Failed to parse JSON. stdout: {res_dry.stdout}, stderr: {res_dry.stderr}, error: {e}"
        
    assert "drifts" in data_dry, f"JSON: {data_dry}"
    drifts_dry = data_dry["drifts"]
    
    # Debug print on failure
    if not any(d["class"] == "A" and d["name"] == session_name for d in drifts_dry):
        debug_info = (
            f"DB_ROWS: {db_rows}\n"
            f"DB_STATE: {db_state}\n"
            f"YAML: {yaml_content}\n"
            f"LOAD_STATE_JSON STDOUT: {res_load.stdout}\n"
            f"LOAD_STATE_JSON STDERR: {res_load.stderr}\n"
            f"HERDR LS STDOUT: {res_ls.stdout}\n"
            f"HERDR LS STDERR: {res_ls.stderr}\n"
            f"AGENT_SESSIONS_YAML ENV: {os.environ.get('AGENT_SESSIONS_YAML')}\n"
        )
        # Execute Python sub-process trace
        py_code = """
import os, json, glob, subprocess, time, sqlite3
from datetime import datetime, timezone
import yaml

yaml_path = os.environ['YAML_PATH']
home = os.environ['HOME_DIR']

d = {}
ws_root = os.environ.get('WORKSPACE_ROOT')
print("WORKSPACE_ROOT:", ws_root)
script = f"source '{ws_root}/.agents/skills/lib.sh' && load_state_json"
out = subprocess.check_output(['bash', '-c', script])
print("LOAD STATE JSON OUT:", out.decode('utf-8'))
"""
        env = {
            "YAML_PATH": str(yaml_path),
            "HOME_DIR": str(tmp_path),
            "CLAUDE_PROJECT_DIR": str(tmp_path / ".claude" / "projects"),
            "LOCAL_BIN": str(tmp_path / ".local" / "bin"),
            "WORKSPACE_ROOT": str(tmp_path),
            "AGENT_SESSIONS_YAML": str(yaml_path),
            "PATH": os.environ.get("PATH", "")
        }
        res_test = subprocess.run(["python3", "-c", py_code], env=env, capture_output=True, text=True)
        debug_info += f"TEST_CODE_STDOUT: {res_test.stdout}\nTEST_CODE_STDERR: {res_test.stderr}\n"
        assert False, f"Drift class A not found in drifts. Output JSON: {json.dumps(data_dry, indent=2)}\nStderr: {res_dry.stderr}\nDebug:\n{debug_info}"
    
    # Verify dry-run made no state changes in SQLite
    conn = sqlite3.connect(str(db_path))
    row = conn.execute("SELECT status FROM sessions WHERE name=?", (session_name,)).fetchone()
    assert row[0] == "running"
    conn.close()
    
    # 2. Run reconcile without --dry-run
    cmd_real = ["bash", str(reconcile_script), "--once", "--emit-diff"]
    res_real = subprocess.run(cmd_real, capture_output=True, text=True, cwd=str(tmp_path))
    assert res_real.returncode == 0, f"Stderr: {res_real.stderr}"
    
    # Verify DB has been updated to terminated
    conn = sqlite3.connect(str(db_path))
    row = conn.execute("SELECT status FROM sessions WHERE name=?", (session_name,)).fetchone()
    assert row[0] == "terminated"
    conn.close()


def test_integration_invalid_arguments_exit_statuses(mam_sandbox, mock_herdr, mock_agents):
    """
    Tier 3: Invalid Arguments & Exit Status Checks
    Verifies that all scripts correctly handle invalid options and combinations with correct exit statuses.
    """
    tmp_path = mam_sandbox
    
    # 1. create_session.sh with missing arguments
    create_script = tmp_path / "skills" / "multi-agent-mux-create" / "scripts" / "create_session.sh"
    res = subprocess.run(["bash", str(create_script)], capture_output=True, text=True, cwd=str(tmp_path))
    assert res.returncode == 2
    assert "ERROR: --workspace required" in res.stderr
    
    res = subprocess.run(["bash", str(create_script), "--workspace", str(tmp_path)], capture_output=True, text=True, cwd=str(tmp_path))
    assert res.returncode == 2
    assert "ERROR: --agent required" in res.stderr
    
    # 2. stop_session.sh with missing arguments
    stop_script = tmp_path / "skills" / "multi-agent-mux-stop" / "scripts" / "stop_session.sh"
    res = subprocess.run(["bash", str(stop_script)], capture_output=True, text=True, cwd=str(tmp_path))
    assert res.returncode == 2
    assert "ERROR: --session required" in res.stderr
    
    # 3. stop_session.sh cannot infer agent from weird session name
    res = subprocess.run(["bash", str(stop_script), "--session", "bad-name"], capture_output=True, text=True, cwd=str(tmp_path))
    assert res.returncode == 2
    assert "ERROR: cannot infer agent" in res.stderr
    
    # 4. stop_session.sh on non-existent session
    res = subprocess.run(["bash", str(stop_script), "--session", "nonexistent-creator-claude"], capture_output=True, text=True, cwd=str(tmp_path))
    assert res.returncode == 1
    assert "ERROR: session 'nonexistent-creator-claude' not in" in res.stderr
    
    # 5. reconcile.sh with invalid flag
    reconcile_script = tmp_path / "skills" / "multi-agent-mux-monitor" / "scripts" / "reconcile.sh"
    res = subprocess.run(["bash", str(reconcile_script), "--invalid-flag"], capture_output=True, text=True, cwd=str(tmp_path))
    assert res.returncode == 2
    assert "ERROR: unknown arg" in res.stderr