File size: 4,488 Bytes
18716c3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 | from types import SimpleNamespace
import pytest
from agent.core.agent_loop import Handlers
from agent.core.session import Event
from agent.core.usage_thresholds import (
USAGE_THRESHOLD_TOOL_NAME,
next_usage_warning_threshold,
)
class FakeUsageApprovalSession:
def __init__(self, *, continuation="continue_agent"):
self.pending_approval = {
"kind": USAGE_THRESHOLD_TOOL_NAME,
"tool_call_id": "usage-threshold-1",
"threshold_usd": 5.0,
"current_spend_usd": 12.25,
"next_threshold_usd": 10.0,
"billing_source": "app_telemetry_session",
"continuation": continuation,
"history_size": 3,
"final_response": "done",
}
self.context_manager = SimpleNamespace(items=[])
self.usage_warning_next_threshold_usd = 5.0
self.events: list[Event] = []
self.turn_count = 0
self.auto_saved = False
async def send_event(self, event: Event):
self.events.append(event)
def increment_turn(self):
self.turn_count += 1
async def auto_save_if_needed(self):
self.auto_saved = True
def test_next_usage_warning_threshold_advances_past_current_spend():
assert next_usage_warning_threshold(4.99, 5.0) == 5.0
assert next_usage_warning_threshold(5.0, 5.0) == 10.0
assert next_usage_warning_threshold(12.25, 5.0) == 20.0
assert next_usage_warning_threshold(40.0, 20.0) == 80.0
@pytest.mark.asyncio
async def test_usage_threshold_approval_resumes_agent(monkeypatch):
session = FakeUsageApprovalSession(continuation="continue_agent")
resumed = False
async def fake_run_agent(run_session, text):
nonlocal resumed
assert run_session is session
assert text == ""
resumed = True
monkeypatch.setattr(Handlers, "run_agent", fake_run_agent)
await Handlers.exec_approval(
session,
[{"tool_call_id": "usage-threshold-1", "approved": True}],
)
assert session.pending_approval is None
assert resumed is True
assert session.usage_warning_next_threshold_usd == 20.0
assert [event.event_type for event in session.events] == [
"tool_state_change",
"tool_output",
]
assert session.events[-1].data["success"] is True
@pytest.mark.asyncio
async def test_usage_threshold_approval_completes_finished_turn(monkeypatch):
session = FakeUsageApprovalSession(continuation="complete_turn")
async def fail_run_agent(*args, **kwargs):
raise AssertionError("complete_turn must not call run_agent")
monkeypatch.setattr(Handlers, "run_agent", fail_run_agent)
await Handlers.exec_approval(
session,
[{"tool_call_id": "usage-threshold-1", "approved": True}],
)
assert session.pending_approval is None
assert [event.event_type for event in session.events] == [
"tool_state_change",
"tool_output",
"turn_complete",
]
assert session.events[-1].data == {
"history_size": 3,
"final_response": "done",
}
assert session.turn_count == 1
assert session.auto_saved is True
@pytest.mark.asyncio
async def test_usage_threshold_rejection_stops_turn(monkeypatch):
session = FakeUsageApprovalSession()
async def fail_run_agent(*args, **kwargs):
raise AssertionError("rejection must not call run_agent")
monkeypatch.setattr(Handlers, "run_agent", fail_run_agent)
await Handlers.exec_approval(
session,
[{"tool_call_id": "usage-threshold-1", "approved": False}],
)
assert session.pending_approval is None
assert [event.event_type for event in session.events] == [
"tool_state_change",
"tool_output",
"interrupted",
]
assert session.events[1].data["success"] is False
assert session.turn_count == 1
assert session.auto_saved is True
@pytest.mark.asyncio
async def test_abandon_complete_turn_usage_threshold_completes_prior_turn():
session = FakeUsageApprovalSession(continuation="complete_turn")
await Handlers._abandon_pending_approval(session)
assert session.pending_approval is None
assert [event.event_type for event in session.events] == [
"tool_state_change",
"turn_complete",
]
assert session.events[-1].data == {
"history_size": 3,
"final_response": "done",
}
assert session.turn_count == 1
assert session.auto_saved is True
|