217 lines
7.4 KiB
Python
217 lines
7.4 KiB
Python
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
import json
|
|
import sys
|
|
|
|
import pytest
|
|
|
|
ROOT=Path(__file__).resolve().parents[1]
|
|
APP=ROOT/"active-runtime/app"
|
|
sys.path.insert(0,str(APP))
|
|
|
|
from aios_online_agent_checkpoint import (
|
|
CheckpointBinding,
|
|
CheckpointError,
|
|
CheckpointStore,
|
|
)
|
|
import aios_online_agent_worker as worker
|
|
|
|
|
|
def binding(attempt: int = 1, *, tenant_head: str = "b"*40):
|
|
return CheckpointBinding(
|
|
task_id="task-p1-a",
|
|
attempt=attempt,
|
|
execution_digest="d"*64,
|
|
job_id="job-p1-a",
|
|
repository="aios/aios",
|
|
base_oid="a"*40,
|
|
expected_branch_oid=tenant_head,
|
|
target_branch="agent/task-p1-a",
|
|
generation=1,
|
|
)
|
|
|
|
|
|
def test_CHECKPOINT_DURABLE_WRITE(tmp_path):
|
|
store=CheckpointStore(tmp_path/"checkpoints.sqlite3")
|
|
first=store.append(binding(), "claimed", {"phase":"claimed","attempt":1})
|
|
second=store.append(binding(), "workspace_prepared", {"phase":"workspace_prepared","attempt":1})
|
|
assert first["sequence"] == 1
|
|
assert second["sequence"] == 2
|
|
latest=store.latest("task-p1-a")
|
|
assert latest["phase"]=="workspace_prepared"
|
|
assert latest["payload"]["attempt"]==1
|
|
|
|
|
|
def test_CHECKPOINT_SCOPE_ISOLATION(tmp_path):
|
|
store=CheckpointStore(tmp_path/"checkpoints.sqlite3")
|
|
store.append(binding(), "claimed", {"phase":"claimed"})
|
|
changed=CheckpointBinding(
|
|
task_id="task-p1-a",attempt=1,execution_digest="d"*64,job_id="job-p1-a",
|
|
repository="aios/other",base_oid="a"*40,expected_branch_oid="b"*40,
|
|
target_branch="agent/task-p1-a",generation=1,
|
|
)
|
|
with pytest.raises(CheckpointError,match="CHECKPOINT_SCOPE_MISMATCH_DENY"):
|
|
store.append(changed,"workspace_prepared",{"phase":"workspace_prepared"})
|
|
|
|
|
|
def test_RESUME_FROM_CHECKPOINT_across_attempt_same_binding(tmp_path):
|
|
store=CheckpointStore(tmp_path/"checkpoints.sqlite3")
|
|
store.append(binding(1), "source_capture_done", {
|
|
"phase":"source_capture_done",
|
|
"attempt":1,
|
|
"result":{"status":"SUCCEEDED","classification":"NEEDS_REWORK_SOURCE_CAPTURE"},
|
|
})
|
|
resumed=store.verify_latest_binding(binding(2))
|
|
assert resumed is not None
|
|
assert resumed["attempt"]==1
|
|
assert resumed["phase"]=="source_capture_done"
|
|
|
|
|
|
def test_checkpoint_new_attempt_can_start_from_claimed(tmp_path):
|
|
store=CheckpointStore(tmp_path/"checkpoints.sqlite3")
|
|
store.append(binding(1),"result_submitted",{"phase":"result_submitted"})
|
|
cp=store.append(binding(2),"claimed",{"phase":"claimed"})
|
|
assert cp["attempt"]==2
|
|
assert cp["sequence"]==1
|
|
|
|
|
|
def envelope():
|
|
return worker.ExecutionEnvelope(
|
|
task_id="task-p1-a",
|
|
job_id="job-p1-a",
|
|
repository="aios/aios",
|
|
base_oid="a"*40,
|
|
expected_branch_oid="b"*40,
|
|
target_branch="agent/task-p1-a",
|
|
branch_mode="CONTINUE_EXISTING",
|
|
generation=1,
|
|
idempotency_key="governed-p1-a",
|
|
action="source_capture",
|
|
adapter="source_capture",
|
|
capability_scope=("git.read","git.test"),
|
|
source_signals=(),
|
|
change_set=(),
|
|
verification=(),
|
|
governance={"task_id":"task-p1-a","job_id":"job-p1-a"},
|
|
workspace={},
|
|
digest="d"*64,
|
|
)
|
|
|
|
|
|
def test_RESULT_SUBMISSION_REPLAY_SAFE_attempt_aware_key():
|
|
env=envelope()
|
|
result={"status":"SUCCEEDED","execution_digest":env.digest}
|
|
one=worker.result_idempotency_key(env.task_id,1,env,result)
|
|
two=worker.result_idempotency_key(env.task_id,2,env,result)
|
|
assert one != two
|
|
assert "-a1-result-" in one
|
|
assert "-a2-result-" in two
|
|
assert one == worker.result_idempotency_key(env.task_id,1,env,result)
|
|
|
|
|
|
class RetryClient:
|
|
def __init__(self):
|
|
self.calls=[]
|
|
def result(self,task_id,lease_token,key,result):
|
|
self.calls.append((task_id,key,dict(result)))
|
|
return (599,{"error":"SERVICE_UNAVAILABLE"}) if len(self.calls)==1 else (201,{"status":"VERIFYING"})
|
|
def get(self,task_id):
|
|
return 500,{}
|
|
|
|
|
|
def test_RETRY_IDEMPOTENT_NO_DUPLICATE_SIDE_EFFECT_reuses_same_key(monkeypatch):
|
|
client=RetryClient()
|
|
monkeypatch.setattr(worker.time,"sleep",lambda _:None)
|
|
status,response=worker.submit_result_replay_safe(
|
|
client,"task-p1-a","lease-token-"+"x"*32,"stable-key",{"status":"SUCCEEDED"}
|
|
)
|
|
assert status==201
|
|
assert [item[1] for item in client.calls]==["stable-key","stable-key"]
|
|
|
|
|
|
class ConflictClient:
|
|
def __init__(self,result):
|
|
self.result_value=result
|
|
def result(self,*args):
|
|
return 409,{"error":"IDEMPOTENCY_CONFLICT"}
|
|
def get(self,task_id):
|
|
return 200,{"status":"VERIFYING","result":dict(self.result_value)}
|
|
|
|
|
|
def test_RESULT_SUBMISSION_REPLAY_SAFE_reconciles_exact_result():
|
|
result={"status":"SUCCEEDED","value":1}
|
|
status,response=worker.submit_result_replay_safe(
|
|
ConflictClient(result),"task-p1-a","lease-token-"+"x"*32,"stable-key",result
|
|
)
|
|
assert status==200
|
|
assert response["idempotent"] is True
|
|
assert response["reconciled"] is True
|
|
|
|
|
|
def test_RESULT_SUBMISSION_REPLAY_SAFE_rejects_conflicting_result():
|
|
status,response=worker.submit_result_replay_safe(
|
|
ConflictClient({"status":"SUCCEEDED","value":2}),
|
|
"task-p1-a","lease-token-"+"x"*32,"stable-key",
|
|
{"status":"SUCCEEDED","value":1},
|
|
)
|
|
assert status==409
|
|
assert response["error"]=="IDEMPOTENCY_CONFLICT"
|
|
|
|
|
|
def test_WORKSPACE_CLEANUP_RECOVERY(tmp_path,monkeypatch):
|
|
root=tmp_path/"workspaces"; root.mkdir()
|
|
stale=root/"task-old"; stale.mkdir(); (stale/"x").write_text("x")
|
|
monkeypatch.setattr(worker,"WORKSPACE_ROOT",root)
|
|
monkeypatch.setattr(worker,"STATE_ROOT",tmp_path/"state")
|
|
removed=worker.recover_stale_workspaces()
|
|
assert removed==1
|
|
assert not stale.exists()
|
|
|
|
|
|
class LeaseBridge:
|
|
def __init__(self):
|
|
self.calls=0
|
|
def call(self,route,body):
|
|
assert route=="/v2/writer-lease"
|
|
self.calls+=1
|
|
if self.calls==1:
|
|
return 599,{"error":"SERVICE_UNAVAILABLE"}
|
|
return 200,{
|
|
"lease_id":"lease-p1-a",
|
|
"job_id":"job-p1-a",
|
|
"task_id":"task-p1-a",
|
|
"repository":"aios/aios",
|
|
"base_oid":"a"*40,
|
|
"expected_branch_oid":"b"*40,
|
|
"branch_mode":"CONTINUE_EXISTING",
|
|
"generation":1,
|
|
"target_branch":"agent/task-p1-a",
|
|
"state":"ACTIVE",
|
|
}
|
|
|
|
|
|
class LeaseHeartbeatStub:
|
|
def ensure(self):
|
|
return None
|
|
|
|
|
|
def test_LEASE_RENEWAL_RECOVERY_retries_same_binding(monkeypatch):
|
|
env=envelope()
|
|
env=worker.ExecutionEnvelope(
|
|
task_id=env.task_id,job_id=env.job_id,repository=env.repository,
|
|
base_oid=env.base_oid,expected_branch_oid=env.expected_branch_oid,
|
|
target_branch=env.target_branch,branch_mode=env.branch_mode,generation=env.generation,
|
|
idempotency_key=env.idempotency_key,action=env.action,adapter=env.adapter,
|
|
capability_scope=env.capability_scope,source_signals=env.source_signals,
|
|
change_set=env.change_set,verification=env.verification,
|
|
governance={"task_id":env.task_id,"job_id":env.job_id,"lease_seconds":900},
|
|
workspace=env.workspace,digest=env.digest,
|
|
)
|
|
bridge=LeaseBridge()
|
|
monkeypatch.setattr(worker.time,"sleep",lambda _:None)
|
|
monkeypatch.setattr(worker,"STATE_ROOT",Path("/tmp")/"aios-p1-test-state")
|
|
response=worker.acquire_writer_lease_replay_safe(env,bridge,LeaseHeartbeatStub())
|
|
assert response["lease_id"]=="lease-p1-a"
|
|
assert bridge.calls==2
|