()
| 66 | |
| 67 | |
| 68 | def main() -> None: |
| 69 | workspace = tempfile.mkdtemp(prefix="a3s-budget-") |
| 70 | agent = Agent.create(INLINE_CONFIG) |
| 71 | |
| 72 | # ----- Phase A: Deny ----- |
| 73 | guard = DenyingGuard() |
| 74 | opts = SessionOptions() |
| 75 | opts.permission_policy = PermissionPolicy(default_decision="allow") |
| 76 | opts.workspace_backend = LocalWorkspaceBackend(workspace) |
| 77 | opts.session_id = "budget-deny-test" |
| 78 | opts.budget_guard = guard |
| 79 | |
| 80 | # Also exercise the getter — read-back must match what we wrote. |
| 81 | assert opts.budget_guard is guard, "BudgetGuard getter must round-trip" |
| 82 | |
| 83 | session = agent.session(workspace, opts) |
| 84 | |
| 85 | try: |
| 86 | _ = session.send("hello") |
| 87 | except RuntimeError as exc: |
| 88 | msg = str(exc) |
| 89 | assert ( |
| 90 | "Budget exhausted" in msg or "llm_tokens" in msg |
| 91 | ), f"expected budget-exhausted error, got: {exc!r}" |
| 92 | else: |
| 93 | raise AssertionError("send() must raise when BudgetGuard denies") |
| 94 | |
| 95 | assert len(guard.llm_checks) == 1, ( |
| 96 | f"check_before_llm must be consulted exactly once, got {guard.llm_checks!r}" |
| 97 | ) |
| 98 | assert guard.llm_checks[0][0] == "budget-deny-test", ( |
| 99 | f"session_id must propagate, got {guard.llm_checks[0]!r}" |
| 100 | ) |
| 101 | assert len(guard.llm_records) == 0, ( |
| 102 | f"record_after_llm must not fire when call was denied, got {guard.llm_records!r}" |
| 103 | ) |
| 104 | |
| 105 | # ----- Phase B: Allow / no-op shape ----- |
| 106 | # A guard with only allow-style methods must not break send(). |
| 107 | # We can't actually send without provider credentials, so we just |
| 108 | # verify the SessionOptions roundtrip and that constructing a |
| 109 | # session succeeds. |
| 110 | allow_opts = SessionOptions() |
| 111 | allow_opts.permission_policy = PermissionPolicy(default_decision="allow") |
| 112 | allow_opts.workspace_backend = LocalWorkspaceBackend(workspace) |
| 113 | allow_opts.session_id = "budget-allow-test" |
| 114 | allow_opts.budget_guard = AllowingGuard() |
| 115 | _ = agent.session(workspace, allow_opts) |
| 116 | |
| 117 | # ----- Phase C: clear back to None ----- |
| 118 | opts.budget_guard = None |
| 119 | assert opts.budget_guard is None |
| 120 | |
| 121 | print("python sdk budget guard ok") |
| 122 | |
| 123 | |
| 124 | if __name__ == "__main__": |
no test coverage detected