Spaces:
Running
Running
| """Budget gateway checks. All HTTP requests target loopback fake servers only.""" | |
| import concurrent.futures | |
| import contextlib | |
| import copy | |
| import http.client | |
| from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer | |
| import importlib.util | |
| import io | |
| import json | |
| from pathlib import Path | |
| import tempfile | |
| import threading | |
| import time | |
| import unittest | |
| from unittest import mock | |
| spec = importlib.util.spec_from_file_location("budget_proxy", Path(__file__).parents[1] / "budget_proxy.py") | |
| proxy = importlib.util.module_from_spec(spec) | |
| spec.loader.exec_module(proxy) | |
| def configuration(): | |
| return { | |
| "global_api_token_cap": 20_000_000, | |
| "category_caps": {"benchmark": 18_000_000, "validation": 2_000_000}, | |
| "daily_headroom_tokens": 20_000_000, | |
| "max_inflight_requests": 2, | |
| "runs": [ | |
| {"run_id": "run-a", "api_token_cap": 8_000_000, "category": "benchmark", "task": "jq", "condition": "single"}, | |
| {"run_id": "run-b", "api_token_cap": 8_000_000, "category": "benchmark", "task": "xsv", "condition": "peers"}, | |
| {"run_id": "validation-a", "api_token_cap": 2_000_000, "category": "validation"}, | |
| ], | |
| } | |
| class LedgerTests(unittest.TestCase): | |
| def setUp(self): | |
| self.temp = tempfile.TemporaryDirectory() | |
| self.path = Path(self.temp.name) / "ledger.sqlite" | |
| self.config = configuration() | |
| def tearDown(self): | |
| self.temp.cleanup() | |
| def test_concurrent_reservations_cannot_overspend(self): | |
| unit = proxy.CONTEXT_LIMIT + 100 | |
| self.config["global_api_token_cap"] = 3 * unit | |
| ledger = proxy.Ledger(self.path, self.config) | |
| barrier = threading.Barrier(12) | |
| def reserve(index): | |
| barrier.wait() | |
| try: | |
| return ledger.reserve("run-a", f"agent-{index}", "hash", 100) | |
| except proxy.BudgetExceeded: | |
| return None | |
| with concurrent.futures.ThreadPoolExecutor(max_workers=12) as pool: | |
| results = list(pool.map(reserve, range(12))) | |
| admitted = [item for item in results if item is not None] | |
| self.assertEqual(len(admitted), 3) | |
| self.assertEqual(ledger.totals()["api_tokens_reserved"], 3 * unit) | |
| def test_per_run_category_and_global_caps_are_independent(self): | |
| unit = proxy.CONTEXT_LIMIT + 100 | |
| for field, expected in [("run", "episode_api_budget"), ("category", "category_api_budget"), ("global", "global_api_budget")]: | |
| with self.subTest(field=field): | |
| config = configuration() | |
| if field == "run": | |
| config["runs"][0]["api_token_cap"] = unit | |
| elif field == "category": | |
| config["category_caps"]["benchmark"] = unit | |
| else: | |
| config["global_api_token_cap"] = unit | |
| ledger = proxy.Ledger(Path(self.temp.name) / f"{field}.sqlite", config) | |
| ledger.reserve("run-a", "agent-0", "hash", 100) | |
| next_run = "run-a" if field == "run" else "run-b" | |
| with self.assertRaisesRegex(proxy.BudgetExceeded, expected): | |
| ledger.reserve(next_run, "agent-1", "hash", 100) | |
| if field != "global": | |
| ledger.reserve("validation-a", "agent-0", "hash", 100) | |
| def test_unknown_failure_keeps_full_charge_and_restart_never_refunds(self): | |
| ledger = proxy.Ledger(self.path, self.config) | |
| first, amount = ledger.reserve("run-a", "agent-0", "one", 100) | |
| ledger.settle(first, None, 1, 502) | |
| second, second_amount = ledger.reserve("run-b", "agent-1", "two", 100) | |
| restored = proxy.Ledger(self.path, self.config) | |
| totals = restored.totals() | |
| self.assertEqual(totals["api_tokens_charged"], amount + second_amount) | |
| self.assertEqual(totals["api_tokens_reserved"], 0) | |
| self.assertEqual(totals["unknown_usage_requests"], 2) | |
| with restored.connect() as db: | |
| self.assertEqual(db.execute("SELECT status FROM requests WHERE id=?", (second,)).fetchone()[0], "unknown_after_restart") | |
| self.assertEqual(proxy.Ledger(self.path, self.config).totals(), totals) | |
| def test_usage_refunds_unused_reservation_without_double_counting_reasoning(self): | |
| ledger = proxy.Ledger(self.path, self.config) | |
| ident, _ = ledger.reserve("run-a", "agent-0", "hash", 100) | |
| row = ledger.settle(ident, {"prompt_tokens": 1000, "completion_tokens": 90, | |
| "total_tokens": 1090, "completion_tokens_details": {"reasoning_tokens": 70}}, 0.2, 200) | |
| self.assertEqual(row["charged"], 1090) | |
| self.assertEqual(row["reserved"], 0) | |
| self.assertEqual(ledger.totals()["api_output_tokens"], 90) | |
| with self.assertRaises(ValueError): | |
| ledger.settle(ident, {"prompt_tokens": 1000, "completion_tokens": 90}, 0.2, 200) | |
| def test_known_rejection_is_zero_charge_and_daily_headroom_stops_new_calls(self): | |
| ledger = proxy.Ledger(self.path, self.config) | |
| ident, _ = ledger.reserve("run-a", "agent-0", "hash", 100) | |
| row = ledger.settle(ident, None, 0.2, 429, remaining_daily=19_000_000, rejected=True) | |
| self.assertEqual(row["charged"], 0) | |
| with self.assertRaisesRegex(proxy.BudgetExceeded, "daily_quota_headroom"): | |
| ledger.reserve("run-b", "agent-0", "hash", 100) | |
| def test_mutating_budget_on_restart_is_rejected(self): | |
| proxy.Ledger(self.path, self.config) | |
| changed = copy.deepcopy(self.config) | |
| changed["global_api_token_cap"] += 1 | |
| with self.assertRaises(ValueError): | |
| proxy.Ledger(self.path, changed) | |
| def test_repeated_history_is_not_recounted_and_compaction_starts_new_window(self): | |
| self.config["runs"][0].update(context_token_cap=1_000_000, agent_context_token_cap=1_000_000) | |
| ledger = proxy.Ledger(self.path, self.config) | |
| body = {"messages": [{"role": "user", "content": "Implement this task"}]} | |
| ident, _ = ledger.reserve("run-a", "agent-0", "one", 100, body) | |
| ledger.settle(ident, {"prompt_tokens": 1000, "completion_tokens": 90}, 0.2, 200) | |
| first = ledger.totals()["context_tokens_estimated"] | |
| self.assertEqual(first, 1090) | |
| body["messages"].extend([{"role": "assistant", "content": "already counted output"}, | |
| {"role": "tool", "content": "new tool result"}]) | |
| ident, _ = ledger.reserve("run-a", "agent-0", "two", 100, body) | |
| ledger.settle(ident, {"prompt_tokens": 10_000, "completion_tokens": 90}, 0.2, 200) | |
| second = ledger.totals()["context_tokens_estimated"] | |
| self.assertGreater(second, first + 90) | |
| self.assertLess(second, first + 300) | |
| compacted = {"messages": [{"role": "user", "content": "Summary of the preceding conversation"}]} | |
| ident, _ = ledger.reserve("run-a", "agent-0", "three", 100, compacted) | |
| ledger.settle(ident, {"prompt_tokens": 500, "completion_tokens": 90}, 0.2, 200) | |
| self.assertEqual(ledger.totals()["context_tokens_estimated"] - second, 590) | |
| def test_agent_and_shared_context_limits_block_before_dispatch(self): | |
| for scope in ("agent", "episode"): | |
| with self.subTest(scope=scope): | |
| config = configuration() | |
| config["runs"][0].update(context_token_cap=2000 if scope == "episode" else 1_000_000, | |
| agent_context_token_cap=1500 if scope == "agent" else 1_000_000) | |
| ledger = proxy.Ledger(Path(self.temp.name) / f"context-{scope}.sqlite", config) | |
| body = {"messages": [{"role": "user", "content": "Implement this task"}]} | |
| ident, _ = ledger.reserve("run-a", "agent-0", "one", 100, body) | |
| ledger.settle(ident, {"prompt_tokens": 1000, "completion_tokens": 100}, 0.2, 200) | |
| with self.assertRaisesRegex(proxy.BudgetExceeded, f"{scope}_context_budget_estimated"): | |
| ledger.reserve("run-a", "agent-0" if scope == "agent" else "agent-1", "two", 500, body) | |
| class BackoffTests(unittest.TestCase): | |
| def test_exponential_jitter_and_provider_delay_are_bounded(self): | |
| with mock.patch.object(proxy.secrets, "randbelow", return_value=500): | |
| self.assertEqual(proxy.rate_limit_delay({}, 0, 0), 2.5) | |
| self.assertEqual(proxy.rate_limit_delay({}, 4, 0), 32.5) | |
| self.assertEqual(proxy.rate_limit_delay({"Retry-After": "60"}, 0, 0), 60) | |
| self.assertEqual(proxy.rate_limit_delay({"Retry-After": "nonsense"}, 0, 0), 2.5) | |
| self.assertIsNone(proxy.rate_limit_delay({"Retry-After": "60"}, 0, 61)) | |
| self.assertIsNone(proxy.rate_limit_delay({"x-should-retry": "false"}, 0, 0)) | |
| with mock.patch.object(proxy.time, "time", return_value=0): | |
| self.assertEqual(proxy.rate_limit_delay({"Retry-After": "Thu, 01 Jan 1970 00:01:00 GMT"}, 0, 0), 60) | |
| class FakeProvider(BaseHTTPRequestHandler): | |
| def log_message(self, *_): | |
| pass | |
| def do_POST(self): | |
| body = json.loads(self.rfile.read(int(self.headers["Content-Length"]))) | |
| with self.server.lock: | |
| self.server.calls.append({"body": body, "authorization": self.headers.get("Authorization")}) | |
| mode = self.server.modes.pop(0) if self.server.modes else self.server.mode | |
| self.server.inflight += 1 | |
| self.server.peak = max(self.server.peak, self.server.inflight) | |
| try: | |
| if mode in ("http_500", "http_429"): | |
| self.send_response(429 if mode == "http_429" else 500) | |
| if mode == "http_429": | |
| self.send_header("Retry-After", self.server.retry_after) | |
| self.send_header("x-ratelimit-remaining-tokens-day", self.server.remaining_tokens) | |
| self.send_header("x-should-retry", self.server.should_retry) | |
| details = b"private-provider-error-detail" | |
| self.send_header("Content-Length", str(len(details))) | |
| self.end_headers() | |
| self.wfile.write(details) | |
| return | |
| self.send_response(200) | |
| self.send_header("Content-Type", "text/event-stream") | |
| self.send_header("x-ratelimit-remaining-tokens-day", "149000000") | |
| self.end_headers() | |
| time.sleep(0.02) | |
| events = [{"choices": [{"delta": {"reasoning_content": "fake reasoning"}}]}, | |
| {"choices": [{"delta": {"content": "fake answer"}}]}] | |
| if mode != "missing_usage": | |
| events.append({"choices": [], "usage": {"prompt_tokens": 1000, "completion_tokens": 90, | |
| "total_tokens": 1090, "completion_tokens_details": {"reasoning_tokens": 70}}}) | |
| for event in events: | |
| self.wfile.write(b"data: " + json.dumps(event).encode() + b"\n\n") | |
| self.wfile.write(b"data: [DONE]\n\n") | |
| finally: | |
| with self.server.lock: | |
| self.server.inflight -= 1 | |
| class GatewayTests(unittest.TestCase): | |
| def setUp(self): | |
| self.temp = tempfile.TemporaryDirectory() | |
| self.provider = ThreadingHTTPServer(("127.0.0.1", 0), FakeProvider) | |
| self.provider.mode = "success" | |
| self.provider.modes = [] | |
| self.provider.retry_after = "0" | |
| self.provider.remaining_tokens = "149000000" | |
| self.provider.should_retry = "true" | |
| self.provider.calls = [] | |
| self.provider.lock = threading.Lock() | |
| self.provider.inflight = self.provider.peak = 0 | |
| self.provider_thread = threading.Thread(target=self.provider.serve_forever, daemon=True) | |
| self.provider_thread.start() | |
| config = configuration() | |
| config["runs"][0].update(context_token_cap=1_000_000, agent_context_token_cap=1_000_000) | |
| self.gateway = proxy.Gateway(("127.0.0.1", 0), config, | |
| Path(self.temp.name) / "ledger.sqlite", Path(self.temp.name) / "events.jsonl", | |
| "synthetic-provider-secret", "synthetic-local-secret", | |
| endpoint=f"http://127.0.0.1:{self.provider.server_port}/chat/completions") | |
| self.gateway_thread = threading.Thread(target=self.gateway.serve_forever, daemon=True) | |
| self.gateway_thread.start() | |
| def tearDown(self): | |
| self.gateway.shutdown() | |
| self.gateway.server_close() | |
| self.provider.shutdown() | |
| self.provider.server_close() | |
| self.gateway_thread.join(timeout=2) | |
| self.provider_thread.join(timeout=2) | |
| self.temp.cleanup() | |
| def call(self, message="test prompt", authorization="Bearer synthetic-local-secret", path=None, headers=None): | |
| conn = http.client.HTTPConnection("127.0.0.1", self.gateway.server_port, timeout=5) | |
| body = {"model": "Beam-501B-A23B", "messages": [{"role": "user", "content": message}], | |
| "stream": True, "max_tokens": 100, "temperature": 0, "reasoning_effort": "low"} | |
| conn.request("POST", path or "/runs/run-a/agents/agent-0/v1/chat/completions", | |
| json.dumps(body), {"Authorization": authorization, "Content-Type": "application/json", **(headers or {})}) | |
| response = conn.getresponse() | |
| status, data = response.status, response.read() | |
| conn.close() | |
| return status, data | |
| def wait_settled(self): | |
| deadline = time.monotonic() + 3 | |
| while self.gateway.ledger.totals()["api_tokens_reserved"] and time.monotonic() < deadline: | |
| time.sleep(0.01) | |
| def test_stream_usage_auth_payload_and_redaction(self): | |
| output = io.StringIO() | |
| with contextlib.redirect_stdout(output), contextlib.redirect_stderr(output): | |
| status, data = self.call("private-test-prompt") | |
| self.wait_settled() | |
| self.assertEqual(status, 200) | |
| self.assertIn(b"[DONE]", data) | |
| self.assertEqual(self.gateway.ledger.totals()["api_tokens_charged"], 1090) | |
| sent = self.provider.calls[0] | |
| self.assertEqual(sent["authorization"], "Bearer synthetic-provider-secret") | |
| self.assertEqual(sent["body"]["stream_options"], {"include_usage": True}) | |
| self.assertEqual(sent["body"]["reasoning_effort"], "medium") | |
| self.assertEqual(sent["body"]["temperature"], 0.7) | |
| logs = output.getvalue() + (Path(self.temp.name) / "events.jsonl").read_text() | |
| with self.gateway.ledger.connect() as db: | |
| logs += str([dict(row) for row in db.execute("SELECT * FROM requests")]) | |
| for forbidden in ["synthetic-provider-secret", "synthetic-local-secret", "private-test-prompt", "fake reasoning"]: | |
| self.assertNotIn(forbidden, logs) | |
| def test_rate_limit_rejection_is_retried_and_every_attempt_is_accounted(self): | |
| self.provider.modes = ["http_429", "success"] | |
| with mock.patch.object(proxy, "rate_limit_delay", return_value=0): | |
| status, data = self.call(headers={"X-Study-Request-ID": "model-call-1", "X-Study-Purpose": "generation"}) | |
| self.wait_settled() | |
| self.assertEqual(status, 200) | |
| self.assertIn(b"[DONE]", data) | |
| self.assertEqual(len(self.provider.calls), 2) | |
| with self.gateway.ledger.connect() as db: | |
| rows = list(db.execute("SELECT * FROM requests ORDER BY started_at")) | |
| self.assertEqual([r["status"] for r in rows], ["rejected", "complete"]) | |
| self.assertEqual([r["charged"] for r in rows], [0, 1090]) | |
| self.assertEqual([r["context_charged"] for r in rows], [0, 1090]) | |
| self.assertEqual({r["run_id"] for r in rows}, {"run-a"}) | |
| self.assertEqual({r["client_request_id"] for r in rows}, {"model-call-1"}) | |
| self.assertEqual({r["purpose"] for r in rows}, {"generation"}) | |
| self.assertEqual({r["request_sha256"] for r in rows}, {rows[0]["request_sha256"]}) | |
| def test_retry_events_pass_real_tracking_sanitizer_and_advance_outbox_cursor(self): | |
| tracking_spec = importlib.util.spec_from_file_location("gateway_test_tracking", Path(__file__).parents[1] / "tracking.py") | |
| tracking = importlib.util.module_from_spec(tracking_spec) | |
| tracking_spec.loader.exec_module(tracking) | |
| self.provider.modes = ["http_429", "success"] | |
| with mock.patch.object(proxy, "rate_limit_delay", return_value=0): | |
| status, _ = self.call() | |
| self.assertEqual(status, 200) | |
| events = [json.loads(line) for line in self.gateway.events.read_text().splitlines()] | |
| sanitized = [tracking.sanitize(event) for event in events] | |
| self.assertEqual(sum(event["metrics"].get("provider_retry", 0) for event in sanitized), 1) | |
| tracker = tracking.Tracker(Path(self.temp.name) / "tracking") | |
| try: | |
| tracker.ingest(self.gateway.events) | |
| self.assertEqual(tracker.pending(), len(events)) | |
| self.assertEqual(tracker.db.execute("SELECT offset FROM cursor").fetchone()[0], self.gateway.events.stat().st_size) | |
| tracker.ingest(self.gateway.events) | |
| self.assertEqual(tracker.pending(), len(events)) | |
| finally: | |
| tracker.close() | |
| def test_concurrent_agents_recover_without_bypassing_concurrency_or_accounting(self): | |
| self.provider.modes = ["http_429"] * 4 | |
| with mock.patch.object(proxy, "rate_limit_delay", return_value=0): | |
| with concurrent.futures.ThreadPoolExecutor(max_workers=5) as pool: | |
| results = list(pool.map(lambda index: self.call(path=f"/runs/run-a/agents/agent-{index}/v1/chat/completions"), range(5))) | |
| self.wait_settled() | |
| self.assertTrue(all(status == 200 for status, _ in results)) | |
| self.assertEqual(len(self.provider.calls), 9) | |
| self.assertLessEqual(self.provider.peak, 2) | |
| self.assertEqual(self.gateway.ledger.totals()["api_tokens_charged"], 5 * 1090) | |
| self.assertEqual(self.gateway.ledger.totals()["context_tokens_estimated"], 5 * 1090) | |
| def test_persistent_rate_limits_stop_after_bounded_attempts_without_charge(self): | |
| self.provider.mode = "http_429" | |
| with mock.patch.object(proxy, "rate_limit_delay", return_value=0): | |
| status, data = self.call() | |
| self.wait_settled() | |
| self.assertEqual(status, 429) | |
| self.assertIn(b"Reflection HTTP 429", data) | |
| self.assertNotIn(b"private-provider-error-detail", data) | |
| self.assertEqual(len(self.provider.calls), proxy.MAX_PROVIDER_ATTEMPTS) | |
| self.assertEqual(self.gateway.ledger.totals()["api_tokens_charged"], 0) | |
| events = [json.loads(line) for line in self.gateway.events.read_text().splitlines()] | |
| self.assertEqual(len([e for e in events if e.get("metrics", {}).get("provider_retry") == 1]), proxy.MAX_PROVIDER_ATTEMPTS - 1) | |
| for secret in ("synthetic-provider-secret", "synthetic-local-secret", "test prompt", "private-provider-error-detail"): | |
| self.assertNotIn(secret, self.gateway.events.read_text()) | |
| def test_provider_no_retry_header_stops_retries(self): | |
| self.provider.mode = "http_429" | |
| self.provider.should_retry = "false" | |
| status, _ = self.call() | |
| self.assertEqual(status, 429) | |
| self.assertEqual(len(self.provider.calls), 1) | |
| def test_rate_limit_daily_quota_header_blocks_retry(self): | |
| self.provider.mode = "http_429" | |
| self.provider.remaining_tokens = "19000000" | |
| with mock.patch.object(proxy, "rate_limit_delay", return_value=0): | |
| status, data = self.call() | |
| self.assertEqual(status, 402) | |
| self.assertIn(b"daily_quota_headroom", data) | |
| self.assertEqual(len(self.provider.calls), 1) | |
| self.assertEqual(self.gateway.ledger.totals()["api_tokens_charged"], 0) | |
| def test_every_retry_rechecks_original_episode_budget(self): | |
| self.provider.modes = ["http_429", "success"] | |
| reserve = self.gateway.ledger.reserve | |
| with mock.patch.object(self.gateway.ledger, "reserve", side_effect=[reserve("run-a", "agent-0", "hash", 100), proxy.BudgetExceeded("episode_api_budget")]) as calls: | |
| with mock.patch.object(proxy, "rate_limit_delay", return_value=0): | |
| status, data = self.call() | |
| self.assertEqual(status, 402) | |
| self.assertIn(b"episode_api_budget", data) | |
| self.assertEqual(calls.call_count, 2) | |
| self.assertEqual(len(self.provider.calls), 1) | |
| self.assertEqual(self.gateway.ledger.totals()["api_tokens_charged"], 0) | |
| def test_provider_retry_after_beyond_deadline_does_not_retry_early(self): | |
| self.provider.mode = "http_429" | |
| self.provider.retry_after = "600" | |
| status, _ = self.call() | |
| self.assertEqual(status, 429) | |
| self.assertEqual(len(self.provider.calls), 1) | |
| def test_disconnected_client_does_not_dispatch_after_backoff(self): | |
| self.provider.mode = "http_429" | |
| self.provider.retry_after = "5" | |
| conn = http.client.HTTPConnection("127.0.0.1", self.gateway.server_port, timeout=5) | |
| body = {"model": "Beam-501B-A23B", "messages": [], "stream": True, "max_tokens": 100} | |
| conn.request("POST", "/runs/run-a/agents/agent-0/v1/chat/completions", json.dumps(body), | |
| {"Authorization": "Bearer synthetic-local-secret", "Content-Type": "application/json"}) | |
| deadline = time.monotonic() + 3 | |
| while not self.provider.calls and time.monotonic() < deadline: | |
| time.sleep(0.01) | |
| self.assertEqual(len(self.provider.calls), 1) | |
| conn.close() | |
| self.wait_settled() | |
| # Acquiring every slot proves the waiting request has released its slot. | |
| for _ in range(2): | |
| self.assertTrue(self.gateway.slots.acquire(timeout=2)) | |
| self.gateway.slots.release() | |
| self.gateway.slots.release() | |
| self.assertEqual(len(self.provider.calls), 1) | |
| def test_unauthorized_request_never_reaches_provider_or_reserves_budget(self): | |
| status, _ = self.call(authorization="Bearer wrong-key") | |
| self.assertEqual(status, 401) | |
| self.assertEqual(self.provider.calls, []) | |
| self.assertEqual(self.gateway.ledger.totals()["api_requests"], 0) | |
| def test_unregistered_run_is_rejected_without_dispatch(self): | |
| status, data = self.call(path="/runs/unknown/agents/agent-0/v1/chat/completions") | |
| self.assertEqual(status, 402) | |
| self.assertIn(b"unregistered_run", data) | |
| self.assertEqual(self.provider.calls, []) | |
| def test_ledger_failure_releases_concurrency_slot_without_dispatch(self): | |
| with mock.patch.object(self.gateway.ledger, "reserve", side_effect=RuntimeError("sensitive detail")): | |
| status, data = self.call() | |
| self.assertEqual(status, 500) | |
| self.assertNotIn(b"sensitive detail", data) | |
| self.assertEqual(self.provider.calls, []) | |
| status, _ = self.call() | |
| self.assertEqual(status, 200) | |
| self.wait_settled() | |
| def test_missing_usage_and_transport_failure_retain_full_reservation(self): | |
| for mode, expected_status in [("missing_usage", 200), ("http_500", 500)]: | |
| with self.subTest(mode=mode): | |
| self.provider.mode = mode | |
| before = self.gateway.ledger.totals()["api_tokens_charged"] | |
| status, _ = self.call() | |
| self.wait_settled() | |
| self.assertEqual(status, expected_status) | |
| self.assertEqual(self.gateway.ledger.totals()["api_tokens_charged"] - before, proxy.CONTEXT_LIMIT + 100) | |
| def test_global_concurrency_limit_covers_all_agents_and_summary_requests(self): | |
| def request(index): | |
| # Compaction is another completion on this same registered provider route. | |
| prompt = "Summarize the prior conversation" if index % 2 else "Continue implementation" | |
| return self.call(prompt, path=f"/runs/run-a/agents/agent-{index}/v1/chat/completions") | |
| with concurrent.futures.ThreadPoolExecutor(max_workers=6) as pool: | |
| results = list(pool.map(request, range(6))) | |
| self.wait_settled() | |
| self.assertTrue(all(status == 200 for status, _ in results)) | |
| self.assertEqual(len(self.provider.calls), 6) | |
| self.assertLessEqual(self.provider.peak, 2) | |
| self.assertEqual(self.gateway.ledger.totals()["api_tokens_charged"], 6 * 1090) | |
| if __name__ == "__main__": | |
| unittest.main() | |