118 lines
4.9 KiB
Python
118 lines
4.9 KiB
Python
from pathlib import Path
|
|
import hashlib, json, sqlite3, sys
|
|
|
|
import pytest
|
|
|
|
ROOT=Path(__file__).resolve().parents[1]
|
|
sys.path.insert(0,str(ROOT/"active-runtime/app"))
|
|
from aios_agent_task_retry_extension import RetryExtensionRequest, RetryExtensionStore, RetryExtensionError
|
|
|
|
AUTHORIZATION_ID="owner-scheme-a-20260817"
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _authorization_id(monkeypatch):
|
|
monkeypatch.setenv("AIOS_AGENT_TASK_AUTHORIZATION_ID", AUTHORIZATION_ID)
|
|
|
|
|
|
def canon(v):
|
|
return hashlib.sha256(json.dumps(v,sort_keys=True,separators=(",",":")).encode()).hexdigest()
|
|
|
|
|
|
def db(tmp_path):
|
|
p=tmp_path/"state.sqlite3"
|
|
c=sqlite3.connect(p)
|
|
c.executescript("""
|
|
CREATE TABLE tasks(
|
|
task_id TEXT PRIMARY KEY,status TEXT,attempt INTEGER,max_attempts INTEGER,revision INTEGER,
|
|
claimed_by TEXT,lease_token_hash TEXT,lease_expires_at TEXT,verification_json TEXT,result_json TEXT
|
|
);
|
|
CREATE TABLE task_events(
|
|
event_id INTEGER PRIMARY KEY AUTOINCREMENT,task_id TEXT,event_type TEXT,agent_id TEXT,
|
|
payload_json TEXT,created_at TEXT
|
|
);
|
|
""")
|
|
verification={"decision":"NEEDS_REWORK","reason":"BASE_OID_BINDING_INVALID"}
|
|
c.execute("INSERT INTO tasks VALUES (?,?,?,?,?,?,?,?,?,?)",
|
|
("task-a","FAILED",10,10,41,None,None,None,json.dumps(verification),json.dumps({"status":"FAILED"})))
|
|
c.commit(); c.close()
|
|
return p,verification
|
|
|
|
|
|
def request(v):
|
|
return RetryExtensionRequest.from_mapping({
|
|
"expected_status":"FAILED","expected_revision":41,"expected_attempt":10,"expected_max_attempts":10,
|
|
"next_max_attempts":11,"previous_verification_sha256":canon(v),
|
|
"authorization_id":AUTHORIZATION_ID,"idempotency_key":"retry-task-a-attempt11",
|
|
})
|
|
|
|
|
|
def test_retry_extension_cas_and_audit(tmp_path):
|
|
p,v=db(tmp_path)
|
|
out=RetryExtensionStore(p).extend("task-a",request(v))
|
|
assert out["status"]=="QUEUED" and out["max_attempts"]==11 and out["revision"]==42
|
|
c=sqlite3.connect(p)
|
|
assert c.execute("select status,attempt,max_attempts,revision from tasks").fetchone()==("QUEUED",10,11,42)
|
|
event=c.execute("select event_type,payload_json from task_events").fetchone()
|
|
assert event[0]=="TASK_RETRY_EXTENSION_AUTHORIZED"
|
|
assert json.loads(event[1])["authorization_id"]=="owner-scheme-a-20260817"
|
|
|
|
|
|
def test_retry_extension_idempotent_replay(tmp_path):
|
|
p,v=db(tmp_path); store=RetryExtensionStore(p); req=request(v)
|
|
first=store.extend("task-a",req); second=store.extend("task-a",req)
|
|
assert first["idempotent"] is False and second["idempotent"] is True
|
|
c=sqlite3.connect(p)
|
|
assert c.execute("select count(*) from task_retry_extensions").fetchone()[0]==1
|
|
|
|
|
|
def test_retry_extension_normalizes_stale_terminal_claim_without_active_lease(tmp_path):
|
|
p,v=db(tmp_path)
|
|
c=sqlite3.connect(p)
|
|
c.execute("update tasks set claimed_by='previous-worker'")
|
|
c.commit(); c.close()
|
|
out=RetryExtensionStore(p).extend("task-a",request(v))
|
|
assert out["status"]=="QUEUED"
|
|
c=sqlite3.connect(p)
|
|
assert c.execute(
|
|
"select claimed_by,lease_token_hash,lease_expires_at from tasks"
|
|
).fetchone()==(None,None,None)
|
|
payload=json.loads(c.execute("select payload_json from task_events").fetchone()[0])
|
|
assert payload["previous_claimed_by"]=="previous-worker"
|
|
assert payload["stale_terminal_claim_normalized"] is True
|
|
assert c.execute("select count(*) from task_events").fetchone()[0]==1
|
|
|
|
|
|
def test_retry_extension_verification_drift_denied(tmp_path):
|
|
p,v=db(tmp_path)
|
|
bad=request({"decision":"NEEDS_REWORK","reason":"other"})
|
|
with pytest.raises(RetryExtensionError,match="RETRY_EXTENSION_VERIFICATION_DRIFT"):
|
|
RetryExtensionStore(p).extend("task-a",bad)
|
|
|
|
|
|
def test_retry_extension_authorization_id_mismatch_rejected(tmp_path, monkeypatch):
|
|
p,v=db(tmp_path)
|
|
monkeypatch.setenv("AIOS_AGENT_TASK_AUTHORIZATION_ID", "other-owner-id")
|
|
with pytest.raises(RetryExtensionError,match="RETRY_EXTENSION_AUTHORIZATION_ID_MISMATCH"):
|
|
RetryExtensionStore(p).extend("task-a",request(v))
|
|
|
|
|
|
def test_retry_extension_bounded_one_step():
|
|
v={"decision":"NEEDS_REWORK"}
|
|
data={
|
|
"expected_status":"FAILED","expected_revision":41,"expected_attempt":10,"expected_max_attempts":10,
|
|
"next_max_attempts":12,"previous_verification_sha256":canon(v),
|
|
"authorization_id":"owner-scheme-a-20260817","idempotency_key":"retry-task-a-attempt11",
|
|
}
|
|
with pytest.raises(RetryExtensionError,match="RETRY_EXTENSION_MUST_BE_BOUNDED_ONE_STEP"):
|
|
RetryExtensionRequest.from_mapping(data)
|
|
|
|
|
|
def test_retry_extension_unit_is_unprivileged_and_requires_owner_binding():
|
|
unit=(ROOT/"deploy/systemd/aios-agent-task-retry-extension.service").read_text()
|
|
assert "User=aios-agent-task" in unit
|
|
assert "Group=aios-agent-task" in unit
|
|
assert "LoadCredential=agent_task_control_token:" in unit
|
|
assert "EnvironmentFile=/etc/aios-gitops/agent-task-retry-extension.env" in unit
|
|
assert "User=root" not in unit
|