165 lines
4.9 KiB
Python
165 lines
4.9 KiB
Python
from datetime import timedelta
|
|
|
|
from celery.exceptions import SoftTimeLimitExceeded
|
|
from django.utils import timezone
|
|
|
|
from apps.agents.models import AgentRun
|
|
from apps.agents.services import build_config_snapshot, execute_agent_run
|
|
from apps.agents.tasks import cleanup_expired_agent_runs
|
|
from apps.agents.tasks import execute_agent_run as execute_agent_run_task
|
|
from apps.audit.models import OperationAuditLog
|
|
from apps.modelhub.services import CompletionResult
|
|
|
|
from .helpers import create_agent
|
|
|
|
|
|
def create_run(version):
|
|
return AgentRun.objects.create(
|
|
definition=version.definition,
|
|
agent_version=version,
|
|
sandbox=version.status == "draft",
|
|
input_payload={"problem": "无法登录"},
|
|
config_snapshot=build_config_snapshot(version),
|
|
total_steps=version.steps.count(),
|
|
)
|
|
|
|
|
|
def test_ordered_runtime_records_steps_usage_and_final_output(db, monkeypatch):
|
|
_, version = create_agent()
|
|
outputs = iter(
|
|
[
|
|
CompletionResult('{"summary":"账号异常"}', 10, 5),
|
|
CompletionResult('{"reply":"请重置密码"}', 12, 6),
|
|
]
|
|
)
|
|
monkeypatch.setattr(
|
|
"apps.modelhub.services._request_completion",
|
|
lambda *args, **kwargs: next(outputs),
|
|
)
|
|
run = create_run(version)
|
|
|
|
execute_agent_run(run.id)
|
|
|
|
run.refresh_from_db()
|
|
assert run.status == "succeeded"
|
|
assert run.final_output == {"reply": "请重置密码"}
|
|
assert run.input_tokens == 22
|
|
assert run.output_tokens == 11
|
|
assert list(run.steps.values_list("code", flat=True)) == ["structure", "reply"]
|
|
|
|
|
|
def test_invalid_structure_is_repaired_once(db, monkeypatch):
|
|
_, version = create_agent(two_steps=False)
|
|
outputs = iter(
|
|
[
|
|
CompletionResult("not-json", 4, 2),
|
|
CompletionResult('{"summary":"已修复"}', 5, 3),
|
|
]
|
|
)
|
|
monkeypatch.setattr(
|
|
"apps.modelhub.services._request_completion",
|
|
lambda *args, **kwargs: next(outputs),
|
|
)
|
|
run = create_run(version)
|
|
|
|
execute_agent_run(run.id)
|
|
|
|
run.refresh_from_db()
|
|
step = run.steps.get()
|
|
assert run.status == "succeeded"
|
|
assert step.attempts.count() == 2
|
|
assert list(step.attempts.values_list("attempt_type", flat=True)) == [
|
|
"initial",
|
|
"structure_repair",
|
|
]
|
|
|
|
|
|
def test_second_invalid_structure_fails_and_stops(db, monkeypatch):
|
|
_, version = create_agent()
|
|
outputs = iter(
|
|
[
|
|
CompletionResult("not-json", 4, 2),
|
|
CompletionResult("still-not-json", 5, 3),
|
|
]
|
|
)
|
|
monkeypatch.setattr(
|
|
"apps.modelhub.services._request_completion",
|
|
lambda *args, **kwargs: next(outputs),
|
|
)
|
|
run = create_run(version)
|
|
|
|
execute_agent_run(run.id)
|
|
|
|
run.refresh_from_db()
|
|
assert run.status == "failed"
|
|
assert run.steps.count() == 1
|
|
assert run.steps.get().status == "failed"
|
|
|
|
|
|
def test_cancellation_before_first_step_stops_run(db):
|
|
_, version = create_agent()
|
|
run = create_run(version)
|
|
run.cancellation_requested_at = timezone.now()
|
|
run.save(update_fields=["cancellation_requested_at"])
|
|
|
|
execute_agent_run(run.id)
|
|
|
|
run.refresh_from_db()
|
|
assert run.status == "cancelled"
|
|
assert run.steps.count() == 0
|
|
|
|
|
|
def test_cancellation_after_model_response_stops_following_steps(db, monkeypatch):
|
|
_, version = create_agent()
|
|
run = create_run(version)
|
|
|
|
def request(*args, **kwargs):
|
|
AgentRun.objects.filter(id=run.id).update(cancellation_requested_at=timezone.now())
|
|
return CompletionResult('{"summary":"已完成当前请求"}', 3, 2)
|
|
|
|
monkeypatch.setattr("apps.modelhub.services._request_completion", request)
|
|
|
|
execute_agent_run(run.id)
|
|
|
|
run.refresh_from_db()
|
|
assert run.status == "cancelled"
|
|
assert list(run.steps.values_list("code", flat=True)) == ["structure"]
|
|
assert run.steps.get().status == "succeeded"
|
|
assert run.final_output is None
|
|
|
|
|
|
def test_soft_time_limit_marks_run_timed_out_and_audits(db, monkeypatch):
|
|
_, version = create_agent()
|
|
run = create_run(version)
|
|
monkeypatch.setattr(
|
|
"apps.agents.tasks.execute_agent_run_service",
|
|
lambda *_: (_ for _ in ()).throw(SoftTimeLimitExceeded()),
|
|
)
|
|
|
|
execute_agent_run_task.run(run.id)
|
|
|
|
run.refresh_from_db()
|
|
assert run.status == "timed_out"
|
|
assert run.error_code == "run_timeout"
|
|
assert OperationAuditLog.objects.filter(
|
|
target_id=str(run.public_id),
|
|
action="run_finished",
|
|
result="timed_out",
|
|
).exists()
|
|
|
|
|
|
def test_cleanup_removes_expired_runs_but_keeps_audit(db, settings):
|
|
_, version = create_agent()
|
|
run = create_run(version)
|
|
AgentRun.objects.filter(id=run.id).update(created_at=timezone.now() - timedelta(days=91))
|
|
|
|
deleted = cleanup_expired_agent_runs()
|
|
|
|
assert deleted == 1
|
|
assert not AgentRun.objects.filter(id=run.id).exists()
|
|
assert OperationAuditLog.objects.filter(
|
|
module="agents",
|
|
action="cleanup_runs",
|
|
request_snapshot__deleted_count=1,
|
|
).exists()
|