206 lines
7.3 KiB
Python
206 lines
7.3 KiB
Python
import os
|
|
import subprocess
|
|
import json
|
|
import sqlite3
|
|
import pytest
|
|
import shutil
|
|
import yaml
|
|
import sys
|
|
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_stop_session_unvalidated_agent(mam_sandbox, mock_herdr, mock_agents):
|
|
"""
|
|
Test case 1: Unvalidated agent argument in stop_session.sh
|
|
Verify that calling stop_session.sh with an invalid agent name returns code 2
|
|
and exits with a clear error message to stderr.
|
|
"""
|
|
tmp_path = mam_sandbox
|
|
stop_script = tmp_path / "skills" / "multi-agent-mux-stop" / "scripts" / "stop_session.sh"
|
|
|
|
# 1. Register a session first
|
|
session_name = "invalid-agent-session-creator-claude"
|
|
mutation = """
|
|
d['herdr_sessions'] = [{
|
|
'name': 'invalid-agent-session-creator-claude',
|
|
'status': 'running',
|
|
'role': 'Creator',
|
|
'pane': {
|
|
'cwd': 'WS_PLACEHOLDER',
|
|
'pid': 8888,
|
|
'cmd': 'claude',
|
|
'cmd_full': 'claude'
|
|
}
|
|
}]
|
|
""".replace("WS_PLACEHOLDER", str(tmp_path))
|
|
res_mut = run_mutation(tmp_path, mutation)
|
|
assert res_mut.returncode == 0
|
|
|
|
# 2. Try stopping with --agent invalid_agent and --purge-conversation
|
|
cmd_stop = [
|
|
"bash", str(stop_script),
|
|
"--session", session_name,
|
|
"--agent", "invalid_agent",
|
|
"--purge-conversation",
|
|
"--yes"
|
|
]
|
|
res_stop = subprocess.run(cmd_stop, capture_output=True, text=True, cwd=str(tmp_path))
|
|
|
|
# It returns code 2
|
|
assert res_stop.returncode == 2
|
|
|
|
# And prints invalid agent type error to stderr
|
|
assert "ERROR: invalid agent type 'invalid_agent'" in res_stop.stderr
|
|
|
|
# The registry entry is NOT removed
|
|
yaml_path = tmp_path / ".mam" / "agent-sessions.yaml"
|
|
with open(yaml_path, 'r') as f:
|
|
reg = yaml.safe_load(f)
|
|
assert len(reg.get("herdr_sessions", [])) == 1
|
|
|
|
|
|
def test_special_character_workspace_slug_interpretation(mam_sandbox):
|
|
"""
|
|
Test case 2: Special character workspace slug interpretation
|
|
Verify that workspace paths with only special characters result in derived session names
|
|
that do NOT start with dashes, preventing option flags injection.
|
|
"""
|
|
tmp_path = mam_sandbox
|
|
lib_path = tmp_path / ".agents" / "skills" / "lib.sh"
|
|
|
|
# Workspace consisting of purely special characters
|
|
special_ws = tmp_path / "@#$*" / "@#$*"
|
|
special_ws.mkdir(parents=True, exist_ok=True)
|
|
|
|
cmd = f"source {lib_path} && derive_session_name '{special_ws}' 'claude'"
|
|
res = subprocess.run(["bash", "-c", cmd], capture_output=True, text=True, cwd=str(tmp_path))
|
|
assert res.returncode == 0
|
|
derived_name = res.stdout.strip()
|
|
|
|
# Confirm that the derived name does NOT start with a dash and is default "ws"
|
|
assert not derived_name.startswith("-")
|
|
assert derived_name == "ws-creator-claude"
|
|
|
|
|
|
def test_monitor_reconcile_race_on_purge(mam_sandbox, mock_herdr):
|
|
"""
|
|
Test case 3: Concurrency/Race condition during session purge
|
|
Verify that if a session is being purged, the presence of the purging lock file
|
|
prevents concurrent monitor checks from auto-registering it back.
|
|
"""
|
|
tmp_path = mam_sandbox
|
|
reconcile_script = tmp_path / "skills" / "multi-agent-mux-monitor" / "scripts" / "reconcile.sh"
|
|
|
|
# 1. Register a running session in herdr and YAML
|
|
session_name = "purged-race-creator-claude"
|
|
with open(mock_herdr, 'r') as f:
|
|
state = json.load(f)
|
|
state["agents"][session_name] = {
|
|
"status": "running",
|
|
"agent": "claude",
|
|
"cwd": str(tmp_path),
|
|
"pid": 9999,
|
|
"pane_id": "w1:p1",
|
|
"command": "claude"
|
|
}
|
|
with open(mock_herdr, 'w') as f:
|
|
json.dump(state, f, indent=2)
|
|
|
|
mutation = f"""
|
|
d['herdr_sessions'] = [{{
|
|
'name': '{session_name}',
|
|
'status': 'running',
|
|
'role': 'Creator',
|
|
'pane': {{
|
|
'cwd': 'WS_PLACEHOLDER',
|
|
'pid': 9999,
|
|
'cmd': 'claude',
|
|
'cmd_full': 'claude'
|
|
}}
|
|
}}]
|
|
""".replace("WS_PLACEHOLDER", str(tmp_path))
|
|
res_mut = run_mutation(tmp_path, mutation)
|
|
assert res_mut.returncode == 0
|
|
|
|
# 2. Simulate a purge: stop_session.sh deletes the registry row, but herdr session takes a moment to die.
|
|
# We remove it from the YAML registry, and create the purging lock file.
|
|
mutation_purge = f"""
|
|
d['herdr_sessions'] = [s for s in d.get('herdr_sessions', []) if s.get('name') != '{session_name}']
|
|
"""
|
|
res_purge_mut = run_mutation(tmp_path, mutation_purge)
|
|
assert res_purge_mut.returncode == 0
|
|
|
|
# Create the purging lock file
|
|
purging_file = tmp_path / ".mam" / f"purging-{session_name}"
|
|
purging_file.parent.mkdir(parents=True, exist_ok=True)
|
|
purging_file.touch()
|
|
|
|
# Confirm YAML registry has no sessions
|
|
yaml_path = tmp_path / ".mam" / "agent-sessions.yaml"
|
|
with open(yaml_path, 'r') as f:
|
|
reg = yaml.safe_load(f)
|
|
assert len(reg.get("herdr_sessions", [])) == 0
|
|
|
|
# 3. Run reconcile.sh once. It should NOT auto-register it back because the purging lock file exists!
|
|
cmd_reconcile = ["bash", str(reconcile_script), "--once", "--emit-diff"]
|
|
res_rec = subprocess.run(cmd_reconcile, capture_output=True, text=True, cwd=str(tmp_path))
|
|
assert res_rec.returncode == 0
|
|
|
|
# Verify that the session has NOT been auto-registered back in YAML
|
|
with open(yaml_path, 'r') as f:
|
|
reg_after = yaml.safe_load(f)
|
|
sessions_after = reg_after.get("herdr_sessions", [])
|
|
assert len(sessions_after) == 0
|
|
|
|
|
|
def test_mqtt_hmac_and_seq_validation(mam_sandbox):
|
|
"""
|
|
Test case 4: MQTT HMAC and sequence number validation
|
|
Verify that mqtt_common correctly verifies HMAC signatures and drops messages with
|
|
invalid signatures or out-of-order sequence numbers.
|
|
"""
|
|
# Import mqtt_common from sandboxed folder
|
|
sys.path.insert(0, str(mam_sandbox / ".agents" / "skills" / "multi-agent-mux-delegate-job" / "scripts"))
|
|
import mqtt_common
|
|
|
|
# 1. verify_hmac behaviour with no auth_token (PoC mode)
|
|
payload_poc = {"data": {"val": 123}}
|
|
assert mqtt_common.verify_hmac(payload_poc, None) is True
|
|
|
|
# 2. verify_hmac with auth_token and valid HMAC
|
|
auth_token = "mysecrettoken"
|
|
payload = {
|
|
"job_id": "job123",
|
|
"event": "started",
|
|
"seq": 1,
|
|
"data": {
|
|
"val": 123
|
|
}
|
|
}
|
|
|
|
# Compute valid signature
|
|
msg = json.dumps(payload, sort_keys=True, separators=(",", ":")).encode()
|
|
import hmac
|
|
import hashlib
|
|
sig = hmac.new(auth_token.encode(), msg, hashlib.sha256).hexdigest()
|
|
|
|
payload_signed = dict(payload)
|
|
payload_signed["data"] = dict(payload["data"])
|
|
payload_signed["data"]["hmac_sig"] = sig
|
|
|
|
assert mqtt_common.verify_hmac(payload_signed, auth_token) is True
|
|
|
|
# 3. verify_hmac with invalid signature
|
|
payload_signed["data"]["hmac_sig"] = "invalidsignature"
|
|
assert mqtt_common.verify_hmac(payload_signed, auth_token) is False
|