#!/usr/bin/env python3 import json import unittest from validate_model_budget import ( ModelBudgetError, extract_response_usage, validate_provider_response, validate_request, validate_response, ) def packet( profile="feature", token_limit=4096, packet_bytes=1200, evidence_bytes=300, byte_limit=16384, evidence_byte_limit=None, ): budget = { "profile": profile, "token_limit": token_limit, "byte_limit": byte_limit, "packet_bytes": packet_bytes, "selected_evidence_bytes": evidence_bytes, } if evidence_byte_limit is not None: budget["evidence_byte_limit"] = evidence_byte_limit return { "budget": budget } class ModelBudgetTests(unittest.TestCase): def test_profile_reserves_output_and_safety_margin(self): result = validate_request(packet()) self.assertEqual(result["max_output_tokens"], 1024) self.assertEqual(result["reserved_total_tokens"], 2780) def test_non_packet_request_overhead_is_reserved(self): result = validate_request(packet(), overhead_tokens=700) self.assertEqual(result["overhead_tokens"], 700) self.assertEqual(result["reserved_total_tokens"], 3480) def test_request_overhead_can_exhaust_context_window(self): result = validate_request(packet(), context_window=4000, overhead_tokens=700) self.assertEqual(result["estimated_request_input_tokens"], 2200) def test_context_window_overflow_fails_before_request(self): with self.assertRaisesRegex(ModelBudgetError, "context budget exceeded"): validate_request(packet(), context_window=1700) def test_estimated_input_must_fit_declared_packet_budget(self): with self.assertRaisesRegex(ModelBudgetError, "input token budget exceeded"): validate_request(packet(token_limit=400), input_tokens=401) def test_byte_upper_bound_prevents_unsafe_heuristic_underestimate(self): with self.assertRaisesRegex(ModelBudgetError, "input token budget exceeded"): validate_request(packet(token_limit=1499)) def test_exact_tokenizer_count_can_be_supplied(self): result = validate_request(packet(token_limit=600), input_tokens=500) self.assertEqual(result["estimated_input_tokens"], 500) def test_explicit_output_budget_can_be_smaller(self): result = validate_request(packet(), context_window=5000, max_output_tokens=256) self.assertEqual(result["max_output_tokens"], 256) def test_output_budget_cannot_expand_profile(self): with self.assertRaisesRegex(ModelBudgetError, "output token budget exceeds profile"): validate_request(packet(), context_window=20000, max_output_tokens=2048) def test_input_budget_cannot_expand_profile(self): with self.assertRaisesRegex(ModelBudgetError, "packet token limit exceeds profile"): validate_request(packet(token_limit=5000), context_window=20000) def test_byte_budget_cannot_expand_profile(self): with self.assertRaisesRegex(ModelBudgetError, "byte limit exceeds profile"): validate_request(packet(byte_limit=16385), context_window=20000) def test_evidence_budget_is_checked_even_without_envelope(self): with self.assertRaisesRegex(ModelBudgetError, "selected evidence exceeds its evidence byte budget"): validate_request(packet(evidence_bytes=301, evidence_byte_limit=300)) def test_nested_accumulated_history_is_rejected(self): value = packet() value["scope"] = {"state": {"messages": ["old"]}} with self.assertRaisesRegex(ModelBudgetError, "forbidden accumulated context"): validate_request(value) def test_full_packet_rejects_unknown_top_level_context(self): value = packet() value.update({"context_id": "c1", "packet_id": "p1", "tool_output": "log"}) with self.assertRaisesRegex(ModelBudgetError, "unknown top-level context field"): validate_request(value) def test_packet_byte_overflow_is_rejected(self): with self.assertRaisesRegex(ModelBudgetError, "byte budget"): validate_request(packet(packet_bytes=16000, evidence_bytes=1000)) def test_full_packet_rejects_stale_serialized_size(self): value = packet() value.update({ "context_id": "c1", "packet_id": "p1", "evidence": {"selected": [], "excluded": []}, }) with self.assertRaisesRegex(ModelBudgetError, "does not match serialized packet"): validate_request(value) def test_packet_without_ids_cannot_lie_about_serialized_size(self): value = packet() value["task"] = "x" * 20000 with self.assertRaisesRegex(ModelBudgetError, "does not match serialized packet"): validate_request(value) def test_full_packet_rejects_inconsistent_evidence_totals(self): value = packet(packet_bytes=215, evidence_bytes=5) value.update({ "context_id": "c1", "packet_id": "p1", "evidence": {"selected": [{"path": "x", "bytes": 4}], "excluded": []}, }) with self.assertRaisesRegex(ModelBudgetError, "selected evidence byte metadata"): validate_request(value) def test_full_packet_rejects_accumulated_history(self): value = packet() value["history"] = ["old response"] with self.assertRaisesRegex(ModelBudgetError, "forbidden accumulated context"): validate_request(value) def test_full_packet_rejects_unknown_exclusion_reason(self): value = packet(packet_bytes=203, evidence_bytes=0) value.update({ "context_id": "c1", "packet_id": "p1", "evidence": {"selected": [], "excluded": [{"path": "x", "reason_code": "old"}]}, }) # Stabilize the serialized size after adding the evidence envelope. for _ in range(4): value["budget"]["packet_bytes"] = len( json.dumps(value, ensure_ascii=False, sort_keys=True, separators=(",", ":")).encode("utf-8") ) with self.assertRaisesRegex(ModelBudgetError, "unknown reason code"): validate_request(value) def test_response_overrun_is_rejected(self): with self.assertRaisesRegex(ModelBudgetError, "response exceeded output token budget"): validate_response(1025, max_output_tokens=1024) def test_response_length_finish_is_checkpointable(self): result = validate_response(1024, max_output_tokens=1024, finish_reason="length") self.assertTrue(result["truncated"]) def test_responses_api_incomplete_reason_is_checkpointable(self): result = validate_response( 1024, max_output_tokens=1024, finish_reason="max_output_tokens", ) self.assertTrue(result["truncated"]) def test_stream_envelope_usage_is_validated(self): event = { "type": "response.completed", "response": { "status": "completed", "usage": { "input_tokens": 300, "output_tokens": 40, "total_tokens": 340, }, }, } self.assertEqual(extract_response_usage(event)["output_tokens"], 40) result = validate_provider_response(event, max_output_tokens=64) self.assertEqual(result["status"], "completed") def test_provider_usage_total_cannot_be_inconsistent(self): response = { "status": "completed", "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 14}, } with self.assertRaisesRegex(ModelBudgetError, "total is smaller"): validate_provider_response(response, max_output_tokens=16) def test_failed_provider_response_is_not_reusable(self): response = { "status": "failed", "usage": {"input_tokens": 10, "output_tokens": 0, "total_tokens": 10}, } with self.assertRaisesRegex(ModelBudgetError, "not reusable"): validate_provider_response(response, max_output_tokens=16) def test_provider_response_without_completion_status_is_rejected(self): response = { "usage": {"input_tokens": 10, "output_tokens": 1, "total_tokens": 11}, } with self.assertRaisesRegex(ModelBudgetError, "missing completion status"): validate_provider_response(response, max_output_tokens=16) if __name__ == "__main__": unittest.main()