diff --git a/tests/integration/providers/test_responses_lifecycle_wire.py b/tests/integration/providers/test_responses_lifecycle_wire.py index 24e65256afa..02269885b44 100644 --- a/tests/integration/providers/test_responses_lifecycle_wire.py +++ b/tests/integration/providers/test_responses_lifecycle_wire.py @@ -4,16 +4,19 @@ import json import os import threading import uuid +from collections.abc import Callable +from types import MappingProxyType from typing import Final, Literal from urllib.parse import parse_qsl, urlsplit +import httpx import openai import pytest from integration._support.client import Gateway, Scenario, eventually, object_value, string_value from integration._support.responses_vendor import same_response from integration._support.wire import Reply, Request, Wire, wire_server from integration._support.worker_readiness import wait_for_worker_readiness -from openai.types.responses import Response, ResponseStreamEvent +from openai.types.responses import Response, ResponseCreatedEvent, ResponseStreamEvent from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter _MODEL: Final = "gpt-5" @@ -88,11 +91,14 @@ def _stream(identity: str) -> tuple[bytes, ...]: ) + (b"data: [DONE]\n\n",) -def _deployment(scenario: Scenario, wire: Wire, api_key: str = _API_KEY) -> str: +def _deployment( + scenario: Scenario, wire: Wire, api_key: str = _API_KEY, provider: str = "openai", **parameters: JsonValue +) -> str: model: Final = scenario.model( - model=f"openai/{_MODEL}", + model=f"{provider}/{_MODEL}", api_key=api_key, api_base=f"{wire.url}/v1", + **parameters, ) _wait_for_workers(scenario.gateway, model, 1) return model @@ -488,6 +494,164 @@ def test_a_client_leaving_a_resumed_stream_releases_the_provider_connection(gate assert urlsplit(left_from).path == f"/v1/responses/{raw_id}" +def _created_then(raw_id: str, read: Reply) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + if request.method == "GET" and request.target == "/v1/models": + return model_discovery_reply(_MODEL) + if request.method == "POST": + return Reply(body=_response(raw_id)) + return read + + return respond + + +def _event_types(answer: httpx.Response) -> list[JsonValue]: + return [ + _JSON_OBJECT.validate_json(line.removeprefix("data: "))["type"] + for line in answer.text.splitlines() + if line.startswith("data: {") + ] + + +_MALFORMED_RESUMES: Final = MappingProxyType( + { + "stream": {"stream": "maybe"}, + "starting_after": {"stream": "true", "starting_after": "-1"}, + } +) + + +@pytest.mark.parametrize("param", tuple(_MALFORMED_RESUMES)) +@pytest.mark.parametrize("prefix", _ALIASES) +def test_a_malformed_resume_query_is_refused_without_calling_the_provider( + gateway: Gateway, prefix: str, param: str +) -> None: + raw_id: Final = f"resp_malformed_{uuid.uuid4().hex}" + replay: Final = Reply(content_type="text/event-stream", chunks=_stream(raw_id)) + + with wire_server(_created_then(raw_id, replay)) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire) + created: Final = _client(gateway, gateway.key, prefix).responses.create(model=model, input="malformed resume") + refused: Final = gateway.request("GET", f"{prefix}/responses/{created.id}", params=_MALFORMED_RESUMES[param]) + requests: Final = drain_contract_requests(wire) + + assert refused.status_code == 400, refused.text + assert object_value(_JSON_OBJECT.validate_json(refused.content)["error"])["param"] == param, refused.text + assert [request.method for request in requests] == ["POST"], requests + + +@pytest.mark.parametrize("prefix", _ALIASES) +def test_a_malformed_resume_cursor_replays_the_whole_stream_when_drop_params_is_on( + gateway: Gateway, prefix: str +) -> None: + raw_id: Final = f"resp_malformed_dropped_{uuid.uuid4().hex}" + replay: Final = Reply(content_type="text/event-stream", chunks=_stream(raw_id)) + + with wire_server(_created_then(raw_id, replay)) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, drop_params=True) + created: Final = _client(gateway, gateway.key, prefix).responses.create(model=model, input="dropped cursor") + replayed: Final = gateway.request( + "GET", f"{prefix}/responses/{created.id}", params=_MALFORMED_RESUMES["starting_after"] + ) + requests: Final = drain_contract_requests(wire) + + assert replayed.status_code == 200, replayed.text + assert _event_types(replayed) == ["response.created", "response.completed"], replayed.text + assert [request.method for request in requests] == ["POST", "GET"], requests + assert _sorted_query(requests[1]) == (("stream", "true"),), requests[1] + + +@pytest.mark.parametrize("prefix", _ALIASES) +def test_resuming_on_a_provider_that_cannot_replay_is_refused_without_calling_it(gateway: Gateway, prefix: str) -> None: + raw_id: Final = f"resp_no_replay_{uuid.uuid4().hex}" + + with wire_server(_created_then(raw_id, Reply(body=_response(raw_id)))) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, provider="litellm_proxy") + created: Final = _client(gateway, gateway.key, prefix).responses.create(model=model, input="cannot replay") + refused: Final = gateway.request( + "GET", f"{prefix}/responses/{created.id}", params={"stream": "true", "starting_after": "3"} + ) + requests: Final = drain_contract_requests(wire) + + assert refused.status_code == 400, refused.text + message: Final = string_value(object_value(_JSON_OBJECT.validate_json(refused.content)["error"])["message"]) + assert "litellm_proxy does not support streaming a stored response" in message, refused.text + assert [request.method for request in requests] == ["POST"], requests + + +@pytest.mark.parametrize("prefix", _ALIASES) +def test_resuming_on_a_provider_that_cannot_replay_answers_the_stored_response_when_drop_params_is_on( + gateway: Gateway, prefix: str +) -> None: + raw_id: Final = f"resp_no_replay_dropped_{uuid.uuid4().hex}" + + with wire_server(_created_then(raw_id, Reply(body=_response(raw_id)))) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, provider="litellm_proxy", drop_params=True) + created: Final = _client(gateway, gateway.key, prefix).responses.create(model=model, input="dropped replay") + answered: Final = gateway.request( + "GET", f"{prefix}/responses/{created.id}", params={"stream": "true", "starting_after": "3"} + ) + requests: Final = drain_contract_requests(wire) + + assert answered.status_code == 200, answered.text + body: Final = _JSON_OBJECT.validate_json(answered.content) + assert (body["object"], body["status"], same_response(string_value(body["id"]), created.id)) == ( + "response", + "completed", + True, + ), answered.text + assert [request.method for request in requests] == ["POST", "GET"], requests + assert _sorted_query(requests[1]) == (), requests[1] + + +@pytest.mark.parametrize("prefix", _ALIASES) +def test_a_resumed_stream_waits_out_the_stream_timeout_not_the_request_timeout(gateway: Gateway, prefix: str) -> None: + raw_id: Final = f"resp_slow_replay_{uuid.uuid4().hex}" + slow: Final = Reply(content_type="text/event-stream", chunks=_stream(raw_id), pause_between_chunks=2) + + with wire_server(_created_then(raw_id, slow)) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, timeout=1, stream_timeout=30) + client: Final = _client(gateway, gateway.key, prefix) + created: Final = client.responses.create(model=model, input="slow replay") + events: Final = tuple(client.responses.retrieve(created.id, stream=True, starting_after=3)) + + assert [event.type for event in events] == ["response.created", "response.completed"], events + + +@pytest.mark.parametrize("prefix", _ALIASES) +def test_a_plain_retrieve_on_a_deployment_set_to_stream_answers_the_response_object( + gateway: Gateway, prefix: str +) -> None: + raw_id: Final = f"resp_streaming_deployment_{uuid.uuid4().hex}" + + def respond(request: Request) -> Reply: + if request.method == "GET" and request.target == "/v1/models": + return model_discovery_reply(_MODEL) + if request.method == "POST": + return Reply(content_type="text/event-stream", chunks=_stream(raw_id)) + return Reply(body=_response(raw_id)) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, stream=True) + events: Final = tuple( + _client(gateway, gateway.key, prefix).responses.create(model=model, input="forced stream", stream=True) + ) + created_id: Final = next(event.response.id for event in events if isinstance(event, ResponseCreatedEvent)) + retrieved: Final = gateway.request("GET", f"{prefix}/responses/{created_id}") + requests: Final = drain_contract_requests(wire) + + assert retrieved.status_code == 200, retrieved.text + assert retrieved.headers["content-type"].startswith("application/json"), retrieved.headers + body: Final = _JSON_OBJECT.validate_json(retrieved.content) + assert (body["object"], body["status"], same_response(string_value(body["id"]), created_id)) == ( + "response", + "completed", + True, + ), retrieved.text + assert [request.method for request in requests] == ["POST", "GET"], requests + assert _sorted_query(requests[1]) == (), requests[1] + + @pytest.mark.parametrize("prefix", _ALIASES) def test_delete_forwards_the_provider_id_once_and_reports_deletion( gateway: Gateway, diff --git a/tests/integration/spend/test_background_response_billed_once.py b/tests/integration/spend/test_background_response_billed_once.py index 5f948fac3ab..00ac917a92e 100644 --- a/tests/integration/spend/test_background_response_billed_once.py +++ b/tests/integration/spend/test_background_response_billed_once.py @@ -7,6 +7,7 @@ from collections.abc import Callable, Mapping from concurrent.futures import ThreadPoolExecutor from itertools import chain from typing import Final +from urllib.parse import parse_qsl, urlsplit import httpx import openai @@ -24,7 +25,7 @@ from integration._support.client import ( from integration._support.database import read_rows, write_rows from integration._support.otlp_sink import Span, decode_spans from integration._support.upstream import delete_scenario, register_scenario -from integration._support.wire import Request, wire_server +from integration._support.wire import Reply, Request, wire_server from integration.caching.response_cache_case import entry_key, redis_store from integration.security._callback_traffic import langfuse_sink from integration.security._sinks import Recorder @@ -67,6 +68,7 @@ from integration.spend._background_responses import ( scripted_routes, scripted_upstream, settling_reads, + sse_frame, streamed_reply, submissions, tag_spend, @@ -350,6 +352,45 @@ def test_a_streamed_background_response_read_later_bills_its_stream_alone(gatewa assert_charged_exactly(lambda: key_spend(key), PRICE) +def _finishes_only_when_resumed(raw_id: str) -> Callable[[Request], Reply]: + finished: Final = (sse_frame("response.completed", 9, response_body("completed", id=raw_id)) + "\n\n").encode() + queued: Final = json.dumps(response_body("queued", id=raw_id)).encode() + + def respond(request: Request) -> Reply: + if urlsplit(request.target).path == "/v1/models": + return Reply(body=b'{"object": "list", "data": []}') + if request.method == "GET" and ("stream", "true") in parse_qsl(urlsplit(request.target).query): + return Reply(content_type="text/event-stream", chunks=(finished, b"data: [DONE]\n\n")) + return Reply(body=queued) + + return respond + + +@pytest.mark.timeout(240) +def test_a_background_response_first_seen_finished_on_a_resumed_stream_bills_once(gateway: Gateway) -> None: + raw_id: Final = f"resp_resumed_{uuid.uuid4().hex}" + with wire_server(_finishes_only_when_resumed(raw_id)) as wire, gateway.scenario() as scenario: + key: Final = scenario.key() + model: Final = scenario.model( + model="openai/gpt-4o-mini", + api_key="sk-scripted-provider", + api_base=f"{wire.url}/v1", + input_cost_per_token=INPUT_COST_PER_TOKEN, + output_cost_per_token=OUTPUT_COST_PER_TOKEN, + ) + response_id: Final = create_background_response(scenario, key, model) + resumed: Final = tuple( + gateway.request( + "GET", f"/v1/responses/{response_id}", key=key, params={"stream": "true", "starting_after": "0"} + ) + for _ in range(READS) + ) + assert [(answer.status_code, "response.completed" in answer.text) for answer in resumed] == [ + (200, True) + ] * READS, [answer.text for answer in resumed] + assert_charged_exactly(lambda: key_spend(key), PRICE) + + def _finished(scenario: Scenario, key: str) -> tuple[Upstream, str]: upstream: Final = scripted_upstream(scenario, background_routes("completed")) return upstream, create_background_response(scenario, key, upstream.model)