diff --git a/tests/integration/spend/_partial_stream.py b/tests/integration/spend/_partial_stream.py new file mode 100644 index 00000000000..eda351f8337 --- /dev/null +++ b/tests/integration/spend/_partial_stream.py @@ -0,0 +1,165 @@ +import json +import threading +import uuid +from collections.abc import Callable, Iterable, Iterator +from hashlib import sha256 +from typing import Final + +import httpx +from integration._support.client import Gateway, Scenario, eventually, held +from integration._support.wire import Reply, Request, Wire +from openai.types.chat import ChatCompletionChunk +from pydantic import BaseModel, JsonValue, TypeAdapter + +JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +BACKEND: Final = "integration-partial-stream" +PIECE: Final = "partial stream billed tokens " +DELIVERED_PIECES: Final = 12 +_FAILURE_ROWS: Final = 'SELECT spend FROM "LiteLLM_SpendLogs" WHERE api_key=%s AND status=%s ORDER BY "startTime"' +ROWS: Final = 'SELECT status, spend FROM "LiteLLM_SpendLogs" WHERE api_key=%s ORDER BY "startTime"' + + +class _SpendRow(BaseModel): + spend: float + + +class StatusRow(BaseModel): + status: str + spend: float + + +class _BudgetError(BaseModel): + type: str + + +class _ErrorResponse(BaseModel): + error: _BudgetError + + +def _chunk(stream_id: str, delta: dict[str, JsonValue], finish_reason: str | None = None) -> bytes: + body: Final = { + "id": stream_id, + "object": "chat.completion.chunk", + "created": 1700000000, + "model": BACKEND, + "choices": [{"index": 0, "delta": delta, "finish_reason": finish_reason}], + } + return f"data: {json.dumps(body)}\n\n".encode() + + +def completion() -> bytes: + return json.dumps( + { + "id": f"chatcmpl-{uuid.uuid4().hex}", + "object": "chat.completion", + "created": 1700000000, + "model": BACKEND, + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + } + ).encode() + + +def peer(delivered: int | None, gate: threading.Event | None = None) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + if request.method != "POST": + return Reply(status=404) + if JSON_OBJECT.validate_json(request.body).get("stream") is True: + stream_id: Final = f"chatcmpl-{uuid.uuid4().hex}" + return Reply( + content_type="text/event-stream", + chunks=( + *(_chunk(stream_id, {"role": "assistant", "content": PIECE}) for _ in range(DELIVERED_PIECES)), + _chunk(stream_id, {}, "stop"), + b"data: [DONE]\n\n", + ), + abort_after=delivered, + pause_between_chunks=0.02, + gate_before_first=gate, + ) + return Reply(body=completion()) + + return respond + + +def priced_model(scenario: Scenario, wire: Wire) -> str: + return scenario.model( + model=f"openai/{BACKEND}", + api_base=f"{wire.url}/v1", + api_key="synthetic-openai-key", + input_cost_per_token=0.001, + output_cost_per_token=0.002, + ) + + +def stream(gateway: Gateway, model: str, key: str, **fields: JsonValue) -> httpx.Response: + return gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "max_tokens": 5, + "stream": True, + "messages": [{"role": "user", "content": "partial stream"}], + **fields, + }, + key=key, + ) + + +def broken_stream(gateway: Gateway, model: str, key: str, **fields: JsonValue) -> None: + response: Final = stream(gateway, model, key, **fields) + assert response.status_code == 200, response.text + events: Final = tuple( + line.removeprefix("data: ") + for line in response.text.splitlines() + if line.startswith("data: ") and line != "data: [DONE]" + ) + assert "error" in JSON_OBJECT.validate_json(events[-1]), response.text + delivered: Final = "".join(_streamed_text(ChatCompletionChunk.model_validate_json(event) for event in events[:-1])) + assert delivered == PIECE * DELIVERED_PIECES, response.text + + +def _streamed_text(chunks: Iterable[ChatCompletionChunk]) -> Iterator[str]: + for chunk in chunks: + yield from (choice.delta.content or "" for choice in chunk.choices) + + +def chat(gateway: Gateway, model: str, key: str, **fields: JsonValue) -> httpx.Response: + return gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "max_tokens": 5, + "messages": [{"role": "user", "content": f"after {uuid.uuid4().hex}"}], + **fields, + }, + key=key, + ) + + +def failure_spend(gateway: Gateway, key: str, rows: int = 1) -> tuple[float, ...]: + digest: Final = sha256(key.encode()).hexdigest() + eventually(lambda: gateway.rows(_FAILURE_ROWS, (digest, "failure")), lambda found: len(found) >= rows, seconds=70) + found: Final = held( + lambda: gateway.rows(_FAILURE_ROWS, (digest, "failure")), lambda latest: len(latest) == rows, holding=2 + ) + return tuple(_SpendRow.model_validate(row).spend for row in found) + + +def chats_sent(wire: Wire) -> int: + return sum(1 for request in wire.drain() if request.method == "POST") + + +def recovered_cost(gateway: Gateway, scenario: Scenario, model: str) -> float: + key: Final = scenario.key(models=[model], max_budget=100) + broken_stream(gateway, model, key) + (spend,) = failure_spend(gateway, key) + assert spend > 0 + return spend + + +def assert_budget_refusal(response: httpx.Response) -> None: + assert response.status_code == 422, response.text + assert _ErrorResponse.model_validate_json(response.content).error.type == "budget_exceeded", response.text diff --git a/tests/integration/spend/test_failure_spend_live_counters_chaos.py b/tests/integration/spend/test_failure_spend_live_counters_chaos.py new file mode 100644 index 00000000000..f0e9538c711 --- /dev/null +++ b/tests/integration/spend/test_failure_spend_live_counters_chaos.py @@ -0,0 +1,224 @@ +import json +import threading +import uuid +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path +from typing import Final + +import httpx +import pytest +import yaml +from integration._support.client import Gateway, eventually +from integration._support.process import owned_proxy_process +from integration._support.redis_process import owned_redis +from integration._support.timing import assert_not_delayed_by, timed +from integration._support.wire import Reply, Request, Wire, wire_server +from integration.spend._partial_stream import ( + DELIVERED_PIECES, + assert_budget_refusal, + broken_stream, + chat, + chats_sent, + failure_spend, + peer, + priced_model, + recovered_cost, +) +from pydantic import JsonValue + +from litellm.constants import REDIS_SOCKET_TIMEOUT + +_BURST: Final = 12 +_SPREAD: Final = 6 +_ATTACK_MARKER: Final = "synthetic-attack-marker" +_POLICY_ID: Final = "synthetic-priced-policy" +_POLICY_APPLY_PATH: Final = f"/guardrail/{_POLICY_ID}/version/DRAFT/apply" + + +def _proxy_config(directory: Path, name: str, sections: dict[str, JsonValue]) -> Path: + base: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + path: Final = directory / f"{name}.yaml" + path.write_text(yaml.safe_dump({**base, **sections})) + return path + + +def _redis_counter_config(directory: Path) -> Path: + base: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + return _proxy_config( + directory, + "redis-counters", + { + "router_settings": { + **base["router_settings"], + "redis_host": "os.environ/REDIS_HOST", + "redis_port": "os.environ/REDIS_PORT", + } + }, + ) + + +def _spread_requests(gateway: Gateway, model: str, key: str) -> tuple[httpx.Response, ...]: + return tuple( + gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "max_tokens": 5, "messages": [{"role": "user", "content": f"spread {uuid.uuid4().hex}"}]}, + key=key, + headers={"connection": "close"}, + ) + for _ in range(_SPREAD) + ) + + +@pytest.mark.timeout(600) +def test_broken_streams_through_a_redis_stall_log_once_and_the_budget_holds_on_every_worker( + gateway: Gateway, tmp_path: Path +) -> None: + release: Final = threading.Event() + with ( + owned_redis(tmp_path) as counters, + wire_server(peer(DELIVERED_PIECES)) as wire, + wire_server(peer(DELIVERED_PIECES, gate=release)) as held_wire, + owned_proxy_process( + gateway, + tmp_path, + {"REDIS_HOST": counters.host, "REDIS_PORT": str(counters.port)}, + config=_redis_counter_config(tmp_path), + workers=2, + ) as owned, + owned.gateway.scenario() as scenario, + ): + scenario.cleanups.callback(release.set) + candidate: Final = owned.gateway + model: Final = priced_model(scenario, wire) + stalled_model: Final = priced_model(scenario, held_wire) + cost: Final = recovered_cost(candidate, scenario, model) + control: Final = timed(lambda: broken_stream(candidate, model, scenario.key(models=[model]))) + burst_key: Final = scenario.key(models=[stalled_model]) + chats_sent(held_wire) + + with ThreadPoolExecutor(max_workers=_BURST) as pool: + pending: Final = tuple( + pool.submit(broken_stream, candidate, stalled_model, burst_key) for _ in range(_BURST) + ) + eventually(held_wire.received.qsize, lambda held: held >= _BURST, seconds=30) + with counters.paused(): + release.set() + drained: Final = timed(lambda: tuple(future.result() for future in pending)) + + assert_not_delayed_by((drained.seconds,), control=(control.seconds,), injected_delay=REDIS_SOCKET_TIMEOUT) + assert failure_spend(candidate, burst_key, rows=_BURST) == pytest.approx((cost,) * _BURST) + + budgeted: Final = scenario.key(models=[model], max_budget=cost / 2) + broken_stream(candidate, model, budgeted) + before_kill: Final = _spread_requests(candidate, model, budgeted) + owned.kill_worker(min(owned.worker_pids())) + after_kill: Final = _spread_requests(candidate, model, budgeted) + + for refused in (*before_kill, *after_kill): + assert_budget_refusal(refused) + + +def _policy(request: Request) -> Reply: + if request.method != "POST" or request.target != _POLICY_APPLY_PATH: + return Reply(status=404) + if _ATTACK_MARKER.encode() not in request.body: + return Reply(body=json.dumps({"action": "NONE", "outputs": [], "assessments": []}).encode()) + return Reply( + body=json.dumps( + { + "action": "GUARDRAIL_INTERVENED", + "outputs": [{"text": "blocked by synthetic policy"}], + "assessments": [ + {"contentPolicy": {"filters": [{"type": "VIOLENCE", "confidence": "HIGH", "action": "BLOCKED"}]}} + ], + "usage": {"contentPolicyUnits": 1}, + } + ).encode() + ) + + +def _policy_config(directory: Path, policy: Wire) -> Path: + return _proxy_config( + directory, + "priced-policy", + { + "guardrails": [ + { + "guardrail_name": "priced-bedrock-policy", + "litellm_params": { + "guardrail": "bedrock", + "mode": "pre_call", + "default_on": True, + "guardrailIdentifier": _POLICY_ID, + "guardrailVersion": "DRAFT", + "aws_region_name": "us-east-1", + "aws_access_key_id": "AKIASYNTHETICGUARDRAIL", + "aws_secret_access_key": "synthetic-guardrail-secret", + "aws_bedrock_runtime_endpoint": policy.url, + }, + } + ] + }, + ) + + +def _attack(gateway: Gateway, model: str, key: str) -> httpx.Response: + blocked: Final = gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": f"{_ATTACK_MARKER} {uuid.uuid4().hex}"}]}, + key=key, + ) + assert blocked.status_code == 400, blocked.text + return blocked + + +@pytest.mark.timeout(480) +def test_guardrail_cost_on_blocked_requests_counts_toward_the_budget(gateway: Gateway, tmp_path: Path) -> None: + with ( + wire_server(_policy) as policy, + wire_server(peer(None)) as wire, + owned_proxy_process(gateway, tmp_path, {}, config=_policy_config(tmp_path, policy), workers=2) as owned, + owned.gateway.scenario() as scenario, + ): + candidate: Final = owned.gateway + model: Final = priced_model(scenario, wire) + meter: Final = scenario.key(models=[model], max_budget=100) + _attack(candidate, model, meter) + (cost,) = failure_spend(candidate, meter) + assert cost > 0 + key: Final = scenario.key(models=[model], max_budget=cost * 1.5) + chats_sent(wire) + + _attack(candidate, model, key) + _attack(candidate, model, key) + after: Final = chat(candidate, model, key) + + assert_budget_refusal(after) + assert chats_sent(wire) == 0 + + +@pytest.mark.timeout(480) +def test_broken_streams_leave_the_budget_alone_while_spend_updates_are_disabled( + gateway: Gateway, tmp_path: Path +) -> None: + base: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config: Final = _proxy_config( + tmp_path, "spend-updates-off", {"general_settings": {**base["general_settings"], "disable_spend_updates": True}} + ) + with ( + wire_server(peer(DELIVERED_PIECES)) as wire, + gateway.scenario() as metering, + owned_proxy_process(gateway, tmp_path, {}, config=config) as owned, + owned.gateway.scenario() as scenario, + ): + cost: Final = recovered_cost(gateway, metering, priced_model(metering, wire)) + candidate: Final = owned.gateway + model: Final = priced_model(scenario, wire) + key: Final = scenario.key(models=[model], max_budget=cost / 2) + + broken_stream(candidate, model, key) + after: Final = chat(candidate, model, key) + + assert after.status_code == 200, after.text diff --git a/tests/integration/spend/test_partial_stream_live_budget.py b/tests/integration/spend/test_partial_stream_live_budget.py index fd7ca2de522..89fb66c9bbf 100644 --- a/tests/integration/spend/test_partial_stream_live_budget.py +++ b/tests/integration/spend/test_partial_stream_live_budget.py @@ -1,195 +1,350 @@ import json +import threading import uuid -from collections.abc import Callable, Iterable, Iterator +from collections.abc import Callable +from concurrent.futures import ThreadPoolExecutor from hashlib import sha256 from typing import Final import httpx import pytest -from integration._support.client import Gateway, Scenario, eventually, held -from integration._support.wire import Reply, Request, Wire, wire_server -from openai.types.chat import ChatCompletionChunk -from pydantic import BaseModel, JsonValue, TypeAdapter +from integration._support.client import Gateway, eventually, held +from integration._support.wire import Reply, Request, wire_server +from integration.spend._partial_stream import ( + BACKEND, + DELIVERED_PIECES, + JSON_OBJECT, + PIECE, + ROWS, + StatusRow, + assert_budget_refusal, + broken_stream, + chat, + chats_sent, + completion, + failure_spend, + peer, + priced_model, + recovered_cost, + stream, +) +from pydantic import JsonValue -_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) -_BACKEND: Final = "integration-partial-stream" -_PIECE: Final = "partial stream billed tokens " -_DELIVERED_PIECES: Final = 12 -_FAILURE_ROWS: Final = 'SELECT spend FROM "LiteLLM_SpendLogs" WHERE api_key=%s AND status=%s ORDER BY "startTime"' +_CONCURRENT_STREAMS: Final = 6 +_MESSAGES_BACKEND: Final = "claude-integration-partial-stream" +_UNCACHED: Final[dict[str, JsonValue]] = {"cache": {"no-cache": True}} -class _SpendRow(BaseModel): - spend: float +@pytest.mark.timeout(240) +def test_a_broken_stream_that_spends_past_the_budget_blocks_the_next_request(gateway: Gateway) -> None: + with wire_server(peer(DELIVERED_PIECES)) as wire, gateway.scenario() as scenario: + model: Final = priced_model(scenario, wire) + cost: Final = recovered_cost(gateway, scenario, model) + key: Final = scenario.key(models=[model], max_budget=cost / 2) + chats_sent(wire) + + broken_stream(gateway, model, key) + after: Final = chat(gateway, model, key) + + assert_budget_refusal(after) + assert failure_spend(gateway, key, rows=2) == pytest.approx((cost, 0.0)) + assert chats_sent(wire) == 1 -class _BudgetError(BaseModel): - type: str +@pytest.mark.timeout(240) +def test_broken_stream_spend_counts_once_toward_the_budget(gateway: Gateway) -> None: + with wire_server(peer(DELIVERED_PIECES)) as wire, gateway.scenario() as scenario: + model: Final = priced_model(scenario, wire) + cost: Final = recovered_cost(gateway, scenario, model) + key: Final = scenario.key(models=[model], max_budget=cost * 1.5) + + broken_stream(gateway, model, key) + broken_stream(gateway, model, key) + after: Final = chat(gateway, model, key) + + assert_budget_refusal(after) + assert failure_spend(gateway, key, rows=3) == pytest.approx((cost, cost, 0.0)) -class _ErrorResponse(BaseModel): - error: _BudgetError +@pytest.mark.timeout(240) +def test_a_stream_that_breaks_before_any_chunk_leaves_the_budget_untouched(gateway: Gateway) -> None: + with wire_server(peer(0)) as wire, gateway.scenario() as scenario: + model: Final = priced_model(scenario, wire) + key: Final = scenario.key(models=[model], max_budget=0.05) + + refused: Final = stream(gateway, model, key) + after: Final = chat(gateway, model, key) + + assert refused.status_code == 500, refused.text + + assert after.status_code == 200, after.text + assert failure_spend(gateway, key) == (0.0,) -def _chunk(delta: dict[str, JsonValue], finish_reason: str | None = None) -> bytes: - body: Final = { - "id": "chatcmpl-partial", - "object": "chat.completion.chunk", - "created": 1700000000, - "model": _BACKEND, - "choices": [{"index": 0, "delta": delta, "finish_reason": finish_reason}], - } - return f"data: {json.dumps(body)}\n\n".encode() +def _broken_stream_then_refusal( + gateway: Gateway, model: str, spender: str, other: str, **fields: JsonValue +) -> httpx.Response: + broken_stream(gateway, model, spender, **fields) + return chat(gateway, model, other, **fields) -def _completion() -> bytes: - return json.dumps( - { - "id": f"chatcmpl-{uuid.uuid4().hex}", - "object": "chat.completion", - "created": 1700000000, - "model": _BACKEND, - "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], - "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, - } - ).encode() +@pytest.mark.timeout(240) +def test_a_broken_stream_on_one_team_key_blocks_the_teams_other_keys(gateway: Gateway) -> None: + with wire_server(peer(DELIVERED_PIECES)) as wire, gateway.scenario() as scenario: + model: Final = priced_model(scenario, wire) + cost: Final = recovered_cost(gateway, scenario, model) + team: Final = scenario.team(models=[model], max_budget=cost / 2) + + after: Final = _broken_stream_then_refusal( + gateway, model, scenario.key(team_id=team, models=[model]), scenario.key(team_id=team, models=[model]) + ) + + assert_budget_refusal(after) -def _peer(delivered: int) -> Callable[[Request], Reply]: - frames: Final = ( - *(_chunk({"role": "assistant", "content": _PIECE}) for _ in range(_DELIVERED_PIECES)), - _chunk({}, "stop"), - b"data: [DONE]\n\n", - ) +@pytest.mark.timeout(240) +def test_a_broken_stream_on_one_of_a_users_keys_blocks_their_other_keys(gateway: Gateway) -> None: + with wire_server(peer(DELIVERED_PIECES)) as wire, gateway.scenario() as scenario: + model: Final = priced_model(scenario, wire) + cost: Final = recovered_cost(gateway, scenario, model) + user: Final = scenario.user(max_budget=cost / 2) + after: Final = _broken_stream_then_refusal( + gateway, model, scenario.key(user_id=user, models=[model]), scenario.key(user_id=user, models=[model]) + ) + + assert_budget_refusal(after) + + +@pytest.mark.timeout(240) +def test_a_broken_stream_for_an_end_user_blocks_their_next_request_on_any_key(gateway: Gateway) -> None: + end_user: Final = f"integration-end-user-{uuid.uuid4().hex}" + with wire_server(peer(DELIVERED_PIECES)) as wire, gateway.scenario() as scenario: + model: Final = priced_model(scenario, wire) + cost: Final = recovered_cost(gateway, scenario, model) + gateway.post("/customer/new", {"user_id": end_user, "max_budget": cost / 2}) + scenario.cleanups.callback(gateway.post, "/customer/delete", {"user_ids": [end_user]}) + + after: Final = _broken_stream_then_refusal( + gateway, model, scenario.key(models=[model]), scenario.key(models=[model]), user=end_user + ) + + assert_budget_refusal(after) + + +@pytest.mark.timeout(240) +def test_a_broken_stream_under_a_tag_blocks_the_next_request_with_that_tag(gateway: Gateway) -> None: + tag: Final = f"integration-partial-stream-{uuid.uuid4().hex}" + with wire_server(peer(DELIVERED_PIECES)) as wire, gateway.scenario() as scenario: + model: Final = priced_model(scenario, wire) + cost: Final = recovered_cost(gateway, scenario, model) + gateway.post("/tag/new", {"name": tag, "max_budget": cost / 2}) + scenario.cleanups.callback(gateway.post, "/tag/delete", {"name": tag}) + key: Final = scenario.key(models=[model]) + + after: Final = _broken_stream_then_refusal(gateway, model, key, key, metadata={"tags": [tag]}) + untagged: Final = chat(gateway, model, key) + + assert_budget_refusal(after) + assert untagged.status_code == 200, untagged.text + + +@pytest.mark.timeout(240) +def test_a_broken_stream_fills_the_keys_budget_window(gateway: Gateway) -> None: + with wire_server(peer(DELIVERED_PIECES)) as wire, gateway.scenario() as scenario: + model: Final = priced_model(scenario, wire) + cost: Final = recovered_cost(gateway, scenario, model) + key: Final = scenario.key(models=[model], budget_limits=[{"budget_duration": "1d", "max_budget": cost / 2}]) + + after: Final = _broken_stream_then_refusal(gateway, model, key, key) + + assert_budget_refusal(after) + + +@pytest.mark.timeout(240) +def test_concurrent_broken_streams_each_count_once_toward_the_budget(gateway: Gateway) -> None: + with wire_server(peer(DELIVERED_PIECES)) as wire, gateway.scenario() as scenario: + model: Final = priced_model(scenario, wire) + cost: Final = recovered_cost(gateway, scenario, model) + key: Final = scenario.key(models=[model], max_budget=cost * (_CONCURRENT_STREAMS + 0.5)) + + with ThreadPoolExecutor(max_workers=_CONCURRENT_STREAMS) as pool: + tuple(pool.map(lambda _: broken_stream(gateway, model, key), range(_CONCURRENT_STREAMS))) + within_budget: Final = chat(gateway, model, key) + broken_stream(gateway, model, key) + over_budget: Final = chat(gateway, model, key) + + assert within_budget.status_code == 200, within_budget.text + assert_budget_refusal(over_budget) + assert sorted(failure_spend(gateway, key, rows=_CONCURRENT_STREAMS + 2)) == pytest.approx( + [0.0] + [cost] * (_CONCURRENT_STREAMS + 1) + ) + + +@pytest.mark.timeout(240) +def test_completed_streams_count_once_toward_the_budget(gateway: Gateway) -> None: + with wire_server(peer(None)) as wire, gateway.scenario() as scenario: + model: Final = priced_model(scenario, wire) + meter: Final = scenario.key(models=[model], max_budget=100) + assert stream(gateway, model, meter, **_UNCACHED).status_code == 200 + (metered,) = eventually( + lambda: gateway.rows(ROWS, (sha256(meter.encode()).hexdigest(),)), + lambda found: len(found) == 1, + seconds=70, + ) + cost: Final = StatusRow.model_validate(metered).spend + key: Final = scenario.key(models=[model], max_budget=cost * 1.5) + + first: Final = stream(gateway, model, key, **_UNCACHED) + second: Final = stream(gateway, model, key, **_UNCACHED) + assert (first.status_code, second.status_code) == (200, 200), (first.text, second.text) + assert _rows(gateway, key, 2) == (StatusRow(status="success", spend=cost),) * 2 + third: Final = stream(gateway, model, key, **_UNCACHED) + + assert_budget_refusal(third) + + +def _hung_then_served(release: threading.Event, hung_model: str) -> Callable[[Request], Reply]: def respond(request: Request) -> Reply: if request.method != "POST": return Reply(status=404) - if _JSON_OBJECT.validate_json(request.body).get("stream") is True: - return Reply( - content_type="text/event-stream", chunks=frames, abort_after=delivered, pause_between_chunks=0.02 - ) - return Reply(body=_completion()) + if JSON_OBJECT.validate_json(request.body).get("model") == hung_model: + return Reply(chunks=(completion(),), gate_before_first=release) + return Reply(body=completion()) return respond -def _priced_model(scenario: Scenario, wire: Wire) -> str: - return scenario.model( - model=f"openai/{_BACKEND}", - api_base=f"{wire.url}/v1", - api_key="synthetic-openai-key", - input_cost_per_token=0.001, - output_cost_per_token=0.002, +def _timed_out_spend(gateway: Gateway, model: str, key: str) -> float: + timed_out: Final = chat(gateway, model, key, messages=[{"role": "user", "content": "hung"}]) + assert timed_out.status_code in (408, 500), timed_out.text + digest: Final = sha256(key.encode()).hexdigest() + landed: Final = eventually(lambda: gateway.rows(ROWS, (digest,)), lambda found: len(found) >= 3, seconds=70) + rows: Final = tuple( + StatusRow.model_validate(row) + for row in held(lambda: gateway.rows(ROWS, (digest,)), lambda latest: len(latest) == len(landed), holding=2) + ) + assert {row.status for row in rows} == {"failure"}, rows + return sum(row.spend for row in rows) + + +@pytest.mark.timeout(240) +def test_retries_the_provider_billed_before_a_timeout_count_toward_the_budget(gateway: Gateway) -> None: + hung_backend: Final = f"{BACKEND}-hung" + release: Final = threading.Event() + with wire_server(_hung_then_served(release, hung_backend)) as wire, gateway.scenario() as scenario: + scenario.cleanups.callback(release.set) + hung: Final = scenario.model( + model=f"openai/{hung_backend}", + api_base=f"{wire.url}/v1", + api_key="synthetic-openai-key", + input_cost_per_token=0.001, + output_cost_per_token=0.002, + timeout=1, + num_retries=2, + ) + served: Final = priced_model(scenario, wire) + billed: Final = _timed_out_spend(gateway, hung, scenario.key(models=[hung, served], max_budget=100)) + assert billed > 0 + key: Final = scenario.key(models=[hung, served], max_budget=billed / 2) + + assert _timed_out_spend(gateway, hung, key) == pytest.approx(billed) + after: Final = chat(gateway, served, key) + + assert_budget_refusal(after) + + +def _messages_event(kind: str, fields: dict[str, JsonValue]) -> bytes: + return f"event: {kind}\ndata: {json.dumps({'type': kind, **fields})}\n\n".encode() + + +def _messages_peer(request: Request) -> Reply: + if request.method != "POST" or request.target != "/v1/messages": + return Reply(status=404) + message: Final[dict[str, JsonValue]] = { + "id": f"msg_{uuid.uuid4().hex}", + "type": "message", + "role": "assistant", + "model": _MESSAGES_BACKEND, + "stop_reason": None, + "stop_sequence": None, + "usage": {"input_tokens": 10, "output_tokens": 1}, + } + if JSON_OBJECT.validate_json(request.body).get("stream") is not True: + return Reply( + body=json.dumps( + {**message, "content": [{"type": "text", "text": "ok"}], "stop_reason": "end_turn"} + ).encode() + ) + return Reply( + content_type="text/event-stream", + chunks=( + _messages_event("message_start", {"message": {**message, "content": []}}), + _messages_event("content_block_start", {"index": 0, "content_block": {"type": "text", "text": ""}}), + *( + _messages_event("content_block_delta", {"index": 0, "delta": {"type": "text_delta", "text": PIECE}}) + for _ in range(DELIVERED_PIECES) + ), + _messages_event("content_block_stop", {"index": 0}), + _messages_event( + "message_delta", + {"delta": {"stop_reason": "end_turn", "stop_sequence": None}, "usage": {"output_tokens": 60}}, + ), + _messages_event("message_stop", {}), + ), + abort_after=DELIVERED_PIECES + 2, + pause_between_chunks=0.02, ) -def _stream(gateway: Gateway, model: str, key: str) -> httpx.Response: +def _messages(gateway: Gateway, model: str, key: str, *, streamed: bool) -> httpx.Response: return gateway.request( "POST", - "/v1/chat/completions", + "/v1/messages", { "model": model, "max_tokens": 5, - "stream": True, - "messages": [{"role": "user", "content": "partial stream"}], + "stream": streamed, + "messages": [{"role": "user", "content": f"partial stream {uuid.uuid4().hex}"}], }, key=key, ) -def _broken_stream(gateway: Gateway, model: str, key: str) -> str: - response: Final = _stream(gateway, model, key) +def _broken_messages_stream(gateway: Gateway, model: str, key: str) -> None: + response: Final = _messages(gateway, model, key, streamed=True) assert response.status_code == 200, response.text - events: Final = tuple( - line.removeprefix("data: ") - for line in response.text.splitlines() - if line.startswith("data: ") and line != "data: [DONE]" - ) - assert "error" in _JSON_OBJECT.validate_json(events[-1]), response.text - return "".join(_streamed_text(ChatCompletionChunk.model_validate_json(event) for event in events[:-1])) + assert response.text.count(PIECE) == DELIVERED_PIECES, response.text + assert "message_stop" not in response.text, response.text -def _streamed_text(chunks: Iterable[ChatCompletionChunk]) -> Iterator[str]: - for chunk in chunks: - yield from (choice.delta.content or "" for choice in chunk.choices) - - -def _chat(gateway: Gateway, model: str, key: str) -> httpx.Response: - return gateway.request( - "POST", - "/v1/chat/completions", - {"model": model, "max_tokens": 5, "messages": [{"role": "user", "content": f"after {uuid.uuid4().hex}"}]}, - key=key, - ) - - -def _failure_spend(gateway: Gateway, key: str, rows: int = 1) -> tuple[float, ...]: +def _rows(gateway: Gateway, key: str, count: int) -> tuple[StatusRow, ...]: digest: Final = sha256(key.encode()).hexdigest() - eventually(lambda: gateway.rows(_FAILURE_ROWS, (digest, "failure")), lambda found: len(found) >= rows, seconds=70) - found: Final = held( - lambda: gateway.rows(_FAILURE_ROWS, (digest, "failure")), lambda latest: len(latest) == rows, holding=2 + eventually(lambda: gateway.rows(ROWS, (digest,)), lambda found: len(found) >= count, seconds=70) + return tuple( + StatusRow.model_validate(row) + for row in held(lambda: gateway.rows(ROWS, (digest,)), lambda latest: len(latest) == count, holding=2) ) - return tuple(_SpendRow.model_validate(row).spend for row in found) - - -def _chats_sent(wire: Wire) -> int: - return sum(1 for request in wire.drain() if request.method == "POST") - - -def _recovered_cost(gateway: Gateway, scenario: Scenario, model: str) -> float: - key: Final = scenario.key(models=[model], max_budget=100) - assert _broken_stream(gateway, model, key) == _PIECE * _DELIVERED_PIECES - (spend,) = _failure_spend(gateway, key) - assert spend > 0 - return spend - - -def _assert_budget_refusal(response: httpx.Response) -> None: - assert response.status_code == 422, response.text - assert _ErrorResponse.model_validate_json(response.content).error.type == "budget_exceeded", response.text @pytest.mark.timeout(240) -def test_a_broken_stream_that_spends_past_the_budget_blocks_the_next_request(gateway: Gateway) -> None: - with wire_server(_peer(_DELIVERED_PIECES)) as wire, gateway.scenario() as scenario: - model: Final = _priced_model(scenario, wire) - cost: Final = _recovered_cost(gateway, scenario, model) - key: Final = scenario.key(models=[model], max_budget=cost / 2) - _chats_sent(wire) +def test_a_broken_messages_stream_counts_toward_the_budget(gateway: Gateway) -> None: + with wire_server(_messages_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"anthropic/{_MESSAGES_BACKEND}", + api_base=wire.url, + api_key="synthetic-anthropic-key", + input_cost_per_token=0.001, + output_cost_per_token=0.002, + ) + meter: Final = scenario.key(models=[model], max_budget=100) + _broken_messages_stream(gateway, model, meter) + (metered,) = _rows(gateway, meter, 1) + assert metered.spend > 0, metered + key: Final = scenario.key(models=[model], max_budget=metered.spend / 2) - assert _broken_stream(gateway, model, key) == _PIECE * _DELIVERED_PIECES - after: Final = _chat(gateway, model, key) + _broken_messages_stream(gateway, model, key) + after: Final = _messages(gateway, model, key, streamed=False) - _assert_budget_refusal(after) - assert _failure_spend(gateway, key, rows=2) == pytest.approx((cost, 0.0)) - assert _chats_sent(wire) == 1 - - -@pytest.mark.timeout(240) -def test_broken_stream_spend_counts_once_toward_the_budget(gateway: Gateway) -> None: - with wire_server(_peer(_DELIVERED_PIECES)) as wire, gateway.scenario() as scenario: - model: Final = _priced_model(scenario, wire) - cost: Final = _recovered_cost(gateway, scenario, model) - key: Final = scenario.key(models=[model], max_budget=cost * 1.5) - - assert _broken_stream(gateway, model, key) == _PIECE * _DELIVERED_PIECES - assert _broken_stream(gateway, model, key) == _PIECE * _DELIVERED_PIECES - after: Final = _chat(gateway, model, key) - - _assert_budget_refusal(after) - assert _failure_spend(gateway, key, rows=3) == pytest.approx((cost, cost, 0.0)) - - -@pytest.mark.timeout(240) -def test_a_stream_that_breaks_before_any_chunk_leaves_the_budget_untouched(gateway: Gateway) -> None: - with wire_server(_peer(0)) as wire, gateway.scenario() as scenario: - model: Final = _priced_model(scenario, wire) - key: Final = scenario.key(models=[model], max_budget=0.05) - - refused: Final = _stream(gateway, model, key) - after: Final = _chat(gateway, model, key) - - assert refused.status_code == 500, refused.text - - assert after.status_code == 200, after.text - assert _failure_spend(gateway, key) == (0.0,) + assert_budget_refusal(after) + assert _rows(gateway, key, 2) == (metered, StatusRow(status="failure", spend=0.0))