106 lines
4.3 KiB
Python
106 lines
4.3 KiB
Python
"""Frozen per-run budgets and fail-closed durable enforcement."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from datetime import UTC, datetime
|
|
from typing import Any
|
|
|
|
from deerflow.config.report_collaboration_config import ReportCollaborationConfig
|
|
from deerflow.persistence.report_collaboration.base import ReportCollaborationLimitError, ReportCollaborationStore
|
|
from deerflow.persistence.report_collaboration.codec import parse_lease_until
|
|
|
|
|
|
def freeze_run_budget(config: ReportCollaborationConfig) -> dict[str, Any]:
|
|
return {
|
|
"max_team_members": config.max_team_members,
|
|
"max_parallel_tasks": config.max_parallel_tasks,
|
|
"max_model_calls": config.max_model_calls,
|
|
"max_retrieval_calls": config.max_retrieval_calls,
|
|
"max_source_bytes": config.max_source_bytes,
|
|
"token_budget": config.token_budget,
|
|
"cost_budget": config.cost_budget,
|
|
"session_timeout_seconds": config.session_timeout_seconds,
|
|
"node_timeout_seconds": config.node_timeout_seconds,
|
|
"max_node_attempts": config.max_node_attempts,
|
|
"max_repair_rounds": config.max_repair_rounds,
|
|
}
|
|
|
|
|
|
class RunBudgetEnforcer:
|
|
def __init__(self, store: ReportCollaborationStore, run_id: str) -> None:
|
|
self._store = store
|
|
self.run_id = run_id
|
|
|
|
async def consume(self, deltas: dict[str, float | int], *, action: str, resource_id: str | None = None) -> dict[str, Any]:
|
|
try:
|
|
result = await self._store.consume_run_budget(self.run_id, deltas=deltas)
|
|
except ReportCollaborationLimitError as exc:
|
|
await self.fail(exc.code, str(exc), detail={"attempted_usage": exc.usage, "limits": exc.limits})
|
|
raise
|
|
run = await self._store.get_run_record(self.run_id)
|
|
if run is not None:
|
|
await self._store.append_audit(
|
|
session_id=str(run["session_id"]),
|
|
run_id=self.run_id,
|
|
category="budget",
|
|
action=action,
|
|
resource_type="run",
|
|
resource_id=resource_id or self.run_id,
|
|
detail={"delta": dict(deltas), "usage": result["usage"]},
|
|
)
|
|
return result
|
|
|
|
async def remaining_seconds(self) -> float | None:
|
|
run = await self._store.get_run_record(self.run_id)
|
|
if run is None:
|
|
return None
|
|
deadline = parse_lease_until(run.get("deadline_at"))
|
|
if deadline is None:
|
|
return None
|
|
return max(0.0, (deadline - datetime.now(UTC)).total_seconds())
|
|
|
|
async def assert_time_available(self) -> None:
|
|
remaining = await self.remaining_seconds()
|
|
if remaining is not None and remaining <= 0:
|
|
await self.fail("RUN_TIME_BUDGET_EXCEEDED", "运行总时长已超过预算")
|
|
raise ReportCollaborationLimitError("RUN_TIME_BUDGET_EXCEEDED", "运行总时长已超过预算")
|
|
|
|
async def fail(self, code: str, message: str, *, detail: dict[str, Any] | None = None) -> dict[str, Any]:
|
|
before = await self._store.get_run_record(self.run_id)
|
|
if before is None:
|
|
return {}
|
|
if str(before.get("status") or "") in {"completed", "cancelled"}:
|
|
return await self._store.get_run(self.run_id) or {}
|
|
already_failed = bool(before.get("budget_exhausted_reason"))
|
|
run = await self._store.fail_run_for_limit(self.run_id, reason=code)
|
|
if already_failed:
|
|
return run
|
|
session_id = str(before["session_id"])
|
|
await self._store.append_audit(
|
|
session_id=session_id,
|
|
run_id=self.run_id,
|
|
category="budget",
|
|
action="run.limit_exceeded",
|
|
resource_type="run",
|
|
resource_id=self.run_id,
|
|
outcome="denied",
|
|
reason=code,
|
|
detail={"message": message, **(detail or {})},
|
|
)
|
|
await self._store.append_event(
|
|
session_id=session_id,
|
|
run_id=self.run_id,
|
|
event_type="run.error",
|
|
data={"code": code, "message": message},
|
|
)
|
|
await self._store.append_event(
|
|
session_id=session_id,
|
|
run_id=self.run_id,
|
|
event_type="run.status.changed",
|
|
data={"status": "failed", "reason": code},
|
|
)
|
|
return run
|
|
|
|
|
|
__all__ = ["RunBudgetEnforcer", "freeze_run_budget"]
|