增加 Validator 工具与材料预算

cyber_agent/core/resource_budget.py

@@ -27,13 +27,17 @@ MAX_TOTAL_TOKENS_ENV = "AGENT_MAX_TOTAL_TOKENS"
 MAX_TOTAL_COST_USD_ENV = "AGENT_MAX_TOTAL_COST_USD"
 MAX_DURATION_SECONDS_ENV = "AGENT_MAX_DURATION_SECONDS"
 RESERVED_FINAL_CALLS_ENV = "AGENT_RESERVED_FINAL_CALLS"
+MAX_VALIDATION_TOOL_CALLS_ENV = "AGENT_MAX_VALIDATION_TOOL_CALLS"
+MAX_VALIDATION_MATERIAL_CHARS_ENV = "AGENT_MAX_VALIDATION_MATERIAL_CHARS"
 
 DEFAULT_MAX_TOTAL_AGENTS = 50
 DEFAULT_MAX_LLM_CALLS = 150
 DEFAULT_MAX_TOTAL_TOKENS = 1_500_000
 DEFAULT_MAX_TOTAL_COST_USD = 15.0
 DEFAULT_MAX_DURATION_SECONDS = 3_600
-DEFAULT_RESERVED_FINAL_CALLS = 1
+DEFAULT_RESERVED_FINAL_CALLS = 8
+DEFAULT_MAX_VALIDATION_TOOL_CALLS = 300
+DEFAULT_MAX_VALIDATION_MATERIAL_CHARS = 1_000_000
 
 BudgetDimension = Literal[
     "agents",
@@ -41,6 +45,8 @@ BudgetDimension = Literal[
     "tokens",
     "cost_usd",
     "duration_seconds",
+    "validation_tool_calls",
+    "validation_material_chars",
 ]
 LLMCallPurpose = Literal["ordinary", "root_validator"]
 
@@ -105,6 +111,8 @@ class ResourceBudget:
     max_total_cost_usd: float = DEFAULT_MAX_TOTAL_COST_USD
     max_duration_seconds: int = DEFAULT_MAX_DURATION_SECONDS
     reserved_final_calls: int = DEFAULT_RESERVED_FINAL_CALLS
+    max_validation_tool_calls: int = DEFAULT_MAX_VALIDATION_TOOL_CALLS
+    max_validation_material_chars: int = DEFAULT_MAX_VALIDATION_MATERIAL_CHARS
 
     def __post_init__(self) -> None:
         if not isinstance(self.enabled, bool):
@@ -115,6 +123,8 @@ class ResourceBudget:
             "max_total_tokens": self.max_total_tokens,
             "max_duration_seconds": self.max_duration_seconds,
             "reserved_final_calls": self.reserved_final_calls,
+            "max_validation_tool_calls": self.max_validation_tool_calls,
+            "max_validation_material_chars": self.max_validation_material_chars,
         }
         for name, value in integer_limits.items():
             if isinstance(value, bool) or not isinstance(value, int) or value <= 0:
@@ -159,6 +169,16 @@ class ResourceBudget:
             reserved_final_calls=_positive_int(
                 values, RESERVED_FINAL_CALLS_ENV, DEFAULT_RESERVED_FINAL_CALLS
             ),
+            max_validation_tool_calls=_positive_int(
+                values,
+                MAX_VALIDATION_TOOL_CALLS_ENV,
+                DEFAULT_MAX_VALIDATION_TOOL_CALLS,
+            ),
+            max_validation_material_chars=_positive_int(
+                values,
+                MAX_VALIDATION_MATERIAL_CHARS_ENV,
+                DEFAULT_MAX_VALIDATION_MATERIAL_CHARS,
+            ),
         )
 
     @classmethod
@@ -197,6 +217,8 @@ class ResourceUsage:
     completion_tokens: int
     total_tokens: int
     total_cost_usd: float
+    validation_tool_calls: int
+    validation_material_chars: int
     started_at: str
     last_denial: dict[str, Any] | None = None
     exhausted_reason: str | None = None
@@ -212,6 +234,8 @@ class ResourceUsage:
             completion_tokens=0,
             total_tokens=0,
             total_cost_usd=0.0,
+            validation_tool_calls=0,
+            validation_material_chars=0,
             started_at=(started_at or _utc_now()).isoformat(),
         )
 
@@ -241,6 +265,8 @@ class ResourceUsage:
             "prompt_tokens": self.prompt_tokens,
             "completion_tokens": self.completion_tokens,
             "total_tokens": self.total_tokens,
+            "validation_tool_calls": self.validation_tool_calls,
+            "validation_material_chars": self.validation_material_chars,
         }
         for name, value in counters.items():
             if isinstance(value, bool) or not isinstance(value, int) or value < 0:
@@ -266,6 +292,8 @@ class ResourceUsage:
             "tokens",
             "cost_usd",
             "duration_seconds",
+            "validation_tool_calls",
+            "validation_material_chars",
         }
         if self.last_denial is not None:
             if set(self.last_denial) != {"dimension", "detail", "denied_at"}:
@@ -531,6 +559,82 @@ class ResourceBudgetController:
                 ),
             )
 
+    async def record_validation_usage(
+        self,
+        root_trace_id: str,
+        budget: ResourceBudget,
+        *,
+        tool_calls: int = 0,
+        material_chars: int = 0,
+        provider_cost_usd: float = 0.0,
+    ) -> ResourceUsage:
+        """原子预留 Validator 工具调用或登记可信材料字符。
+
+        搜索/打开在真实外部调用前用 ``tool_calls=1`` 预留；网页正文、
+        Artifact字符和 Provider 报告的成本事后登记，超限先落实际消耗。
+        """
+        for name, value in {
+            "tool_calls": tool_calls,
+            "material_chars": material_chars,
+        }.items():
+            if isinstance(value, bool) or not isinstance(value, int) or value < 0:
+                raise ValueError(f"{name} must be a non-negative integer")
+        if (
+            isinstance(provider_cost_usd, bool)
+            or not isinstance(provider_cost_usd, (int, float))
+            or not math.isfinite(float(provider_cost_usd))
+            or provider_cost_usd < 0
+        ):
+            raise ValueError("provider_cost_usd must be a non-negative number")
+        if not tool_calls and not material_chars and not provider_cost_usd:
+            return await self.get_usage(root_trace_id)
+        async with self._lock(root_trace_id):
+            usage = await self.get_usage(root_trace_id)
+            self._raise_if_terminal(usage, budget)
+            await self._check_time_locked(root_trace_id, usage, budget)
+            proposed_calls = usage.validation_tool_calls + tool_calls
+            if budget.enabled and proposed_calls > budget.max_validation_tool_calls:
+                usage = await self._deny_locked(
+                    root_trace_id,
+                    usage,
+                    "validation_tool_calls",
+                    f"Validator tool calls would reach {proposed_calls}",
+                )
+                raise self._exceeded("validation_tool_calls", usage, budget)
+            usage.validation_tool_calls = proposed_calls
+            usage.validation_material_chars += material_chars
+            usage.total_cost_usd += float(provider_cost_usd)
+            exceeded: BudgetDimension | None = None
+            if (
+                budget.enabled
+                and usage.validation_material_chars
+                > budget.max_validation_material_chars
+            ):
+                exceeded = "validation_material_chars"
+            elif budget.enabled and usage.total_cost_usd > budget.max_total_cost_usd:
+                exceeded = "cost_usd"
+            if exceeded:
+                usage = await self._deny_locked(
+                    root_trace_id,
+                    usage,
+                    exceeded,
+                    (
+                        "Validator material characters reached "
+                        f"{usage.validation_material_chars}"
+                        if exceeded == "validation_material_chars"
+                        else "Validator provider cost reached "
+                        f"{usage.total_cost_usd}"
+                    ),
+                )
+            else:
+                await self.trace_store.replace_resource_usage(
+                    root_trace_id,
+                    usage.to_dict(),
+                )
+            if exceeded:
+                raise self._exceeded(exceeded, usage, budget)
+            return usage
+
     @staticmethod
     def _validate_reported_usage(
         prompt_tokens: int,
@@ -665,6 +769,8 @@ class ResourceBudgetController:
             "tokens": "tokens",
             "cost_usd": "cost_usd",
             "duration_seconds": "duration_seconds",
+            "validation_tool_calls": "validation_tool_calls",
+            "validation_material_chars": "validation_material_chars",
         }
         if dimension in terminal_dimensions:
             typed_dimension = terminal_dimensions[dimension]

tests/test_recursive_resource_budget.py

@@ -32,7 +32,9 @@ class ResourceBudgetModelTest(unittest.TestCase):
         self.assertEqual(1_500_000, budget.max_total_tokens)
         self.assertEqual(15.0, budget.max_total_cost_usd)
         self.assertEqual(3_600, budget.max_duration_seconds)
-        self.assertEqual(1, budget.reserved_final_calls)
+        self.assertEqual(8, budget.reserved_final_calls)
+        self.assertEqual(300, budget.max_validation_tool_calls)
+        self.assertEqual(1_000_000, budget.max_validation_material_chars)
 
     def test_environment_is_strict_and_snapshot_round_trips(self):
         environ = {
@@ -43,9 +45,13 @@ class ResourceBudgetModelTest(unittest.TestCase):
             "AGENT_MAX_TOTAL_COST_USD": "1.25",
             "AGENT_MAX_DURATION_SECONDS": "90",
             "AGENT_RESERVED_FINAL_CALLS": "2",
+            "AGENT_MAX_VALIDATION_TOOL_CALLS": "18",
+            "AGENT_MAX_VALIDATION_MATERIAL_CHARS": "9000",
         }
         budget = ResourceBudget.from_environment(environ)
         self.assertEqual(budget, ResourceBudget.from_dict(budget.to_dict()))
+        self.assertEqual(18, budget.max_validation_tool_calls)
+        self.assertEqual(9000, budget.max_validation_material_chars)
 
         invalid = dict(environ, AGENT_RESOURCE_BUDGET_ENABLED="yes")
         with self.assertRaisesRegex(ValueError, "must be 'true' or 'false'"):
@@ -57,6 +63,8 @@ class ResourceBudgetModelTest(unittest.TestCase):
             ("AGENT_MAX_TOTAL_COST_USD", "nan"),
             ("AGENT_MAX_DURATION_SECONDS", "none"),
             ("AGENT_RESERVED_FINAL_CALLS", "0"),
+            ("AGENT_MAX_VALIDATION_TOOL_CALLS", "0"),
+            ("AGENT_MAX_VALIDATION_MATERIAL_CHARS", "-1"),
         ):
             with self.subTest(name=name):
                 invalid = dict(environ, **{name: value})
@@ -266,6 +274,7 @@ class ResourceBudgetControllerTest(unittest.IsolatedAsyncioTestCase):
             max_total_tokens=1,
             max_total_cost_usd=0.01,
             max_duration_seconds=1,
+            reserved_final_calls=1,
         )
         await self.controller.initialize(self.root_trace_id, budget)
         await self.controller.reserve_agents(self.root_trace_id, budget, 5)
@@ -285,6 +294,98 @@ class ResourceBudgetControllerTest(unittest.IsolatedAsyncioTestCase):
         self.assertEqual(200, usage.total_tokens)
         self.assertIsNone(usage.exhausted_reason)
 
+    async def test_validation_tool_calls_are_atomically_limited(self):
+        budget = ResourceBudget(max_validation_tool_calls=3)
+        await self.controller.initialize(self.root_trace_id, budget)
+        results = await asyncio.gather(
+            *(
+                self.controller.record_validation_usage(
+                    self.root_trace_id,
+                    budget,
+                    tool_calls=1,
+                )
+                for _ in range(6)
+            ),
+            return_exceptions=True,
+        )
+        self.assertEqual(3, sum(isinstance(item, ResourceUsage) for item in results))
+        self.assertEqual(
+            3,
+            (await self.controller.get_usage(self.root_trace_id)).validation_tool_calls,
+        )
+        self.assertTrue(any(
+            isinstance(item, ResourceBudgetExceeded)
+            and item.dimension == "validation_tool_calls"
+            for item in results
+        ))
+
+    async def test_validation_material_overage_is_recorded_before_denial(self):
+        budget = ResourceBudget(max_validation_material_chars=10)
+        await self.controller.initialize(self.root_trace_id, budget)
+        await self.controller.record_validation_usage(
+            self.root_trace_id,
+            budget,
+            material_chars=8,
+        )
+        with self.assertRaises(ResourceBudgetExceeded) as caught:
+            await self.controller.record_validation_usage(
+                self.root_trace_id,
+                budget,
+                material_chars=5,
+            )
+        self.assertEqual("validation_material_chars", caught.exception.dimension)
+        usage = await self.controller.get_usage(self.root_trace_id)
+        self.assertEqual(13, usage.validation_material_chars)
+        self.assertEqual(
+            "budget_exhausted:validation_material_chars",
+            usage.exhausted_reason,
+        )
+
+    async def test_validation_counters_do_not_consume_llm_or_agent_counts(self):
+        budget = ResourceBudget()
+        await self.controller.initialize(self.root_trace_id, budget)
+        usage = await self.controller.record_validation_usage(
+            self.root_trace_id,
+            budget,
+            tool_calls=2,
+            material_chars=500,
+        )
+        self.assertEqual(1, usage.total_agents)
+        self.assertEqual(0, usage.llm_calls)
+        self.assertEqual(2, usage.validation_tool_calls)
+        self.assertEqual(500, usage.validation_material_chars)
+
+    async def test_provider_cost_uses_the_tree_cost_budget(self):
+        budget = ResourceBudget(max_total_cost_usd=0.02)
+        await self.controller.initialize(self.root_trace_id, budget)
+        with self.assertRaises(ResourceBudgetExceeded) as caught:
+            await self.controller.record_validation_usage(
+                self.root_trace_id,
+                budget,
+                provider_cost_usd=0.03,
+            )
+        self.assertEqual("cost_usd", caught.exception.dimension)
+        usage = await self.controller.get_usage(self.root_trace_id)
+        self.assertAlmostEqual(0.03, usage.total_cost_usd)
+        self.assertEqual("budget_exhausted:cost_usd", usage.exhausted_reason)
+
+    async def test_disabled_budget_still_observes_validation_usage(self):
+        budget = ResourceBudget(
+            enabled=False,
+            max_validation_tool_calls=1,
+            max_validation_material_chars=1,
+        )
+        await self.controller.initialize(self.root_trace_id, budget)
+        usage = await self.controller.record_validation_usage(
+            self.root_trace_id,
+            budget,
+            tool_calls=3,
+            material_chars=100,
+        )
+        self.assertEqual(3, usage.validation_tool_calls)
+        self.assertEqual(100, usage.validation_material_chars)
+        self.assertIsNone(usage.exhausted_reason)
+
 
 if __name__ == "__main__":
     unittest.main()