diff --git a/litellm/proxy/hooks/responses_id_security.py b/litellm/proxy/hooks/responses_id_security.py index e554512ec91..c1384565a80 100644 --- a/litellm/proxy/hooks/responses_id_security.py +++ b/litellm/proxy/hooks/responses_id_security.py @@ -6,6 +6,7 @@ instead of writing immediately on each request. """ from collections.abc import AsyncGenerator, Callable, Mapping +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, cast from fastapi import HTTPException @@ -31,6 +32,16 @@ if TYPE_CHECKING: _RESPONSES_API_PROVIDER_PREFIX: Final = "/openai" _RESPONSES_API_CREATE_ROUTES: Final = frozenset({"/v1/responses", "/responses"}) +_ADDRESSED_RESPONSE_ID_FIELD_BY_CALL_TYPE: Final[Mapping[str, str]] = MappingProxyType( + { + "aresponses": "previous_response_id", + "acompact_responses": "previous_response_id", + "aget_responses": "response_id", + "adelete_responses": "response_id", + "acancel_responses": "response_id", + "alist_input_items": "response_id", + } +) _UNMANAGED_RESPONSE_ID_DETAIL: Final = ( "Forbidden. This response id was not issued by this proxy, so the proxy cannot tell who owns it. " @@ -120,17 +131,9 @@ class ResponsesIDSecurity(CustomLogger): data: dict, call_type: CallTypesLiteral, ) -> Exception | str | dict | None: - # MAP all the responses api response ids to the encrypted response ids - responses_api_call_types: Final = { - "aresponses", - "aget_responses", - "adelete_responses", - "acancel_responses", - "alist_input_items", - } - if call_type not in responses_api_call_types: + addressed_id_field: Final = _ADDRESSED_RESPONSE_ID_FIELD_BY_CALL_TYPE.get(call_type) + if addressed_id_field is None: return None - addressed_id_field: Final = "previous_response_id" if call_type == "aresponses" else "response_id" retained_id: Final = data.get(ADDRESSED_RESPONSE_ID_FIELD) addressed_id: Final = ( retained_id if isinstance(retained_id, str) and retained_id else data.get(addressed_id_field) diff --git a/tests/integration/_support/worker_readiness.py b/tests/integration/_support/worker_readiness.py new file mode 100644 index 00000000000..b77f12bb5ee --- /dev/null +++ b/tests/integration/_support/worker_readiness.py @@ -0,0 +1,64 @@ +from __future__ import annotations + +import time +from collections.abc import Iterator, Mapping +from itertools import accumulate +from types import MappingProxyType +from typing import Final + +import httpx +from integration._support.client import JSON_OBJECT, Gateway, object_value + +from litellm.constants import PROXY_CONFIG_RELOAD_INTERVAL_SECONDS + + +def _record_latest_readback( + latest_by_worker: Mapping[int, int], + observation: tuple[int, int], +) -> Mapping[int, int]: + worker_pid, count = observation + return MappingProxyType({**latest_by_worker, worker_pid: count}) + + +def wait_for_worker_readiness( + gateway: Gateway, + group: str, + deployments: int, + workers: int, +) -> None: + deadline: Final = time.monotonic() + (30 if workers == 1 else PROXY_CONFIG_RELOAD_INTERVAL_SECONDS * 2 + 30) + headers: Final = {"Authorization": f"Bearer {gateway.key}"} + + def readbacks() -> Iterator[tuple[int, int]]: + while time.monotonic() < deadline: + timeout: Final = deadline - time.monotonic() + with httpx.Client( + base_url=str(gateway.client.base_url), + timeout=timeout, + trust_env=False, + verify=False, + ) as client: + debug_response: Final = client.get("/debug/memory/summary", headers=headers) + response: Final = client.get( + "/model/info", + headers=headers, + ) + assert debug_response.status_code == 200, debug_response.text + assert response.status_code == 200, response.text + assert debug_response.extensions.get("network_stream") is not None + assert debug_response.extensions.get("network_stream") is response.extensions.get("network_stream") + debug_body: Final = JSON_OBJECT.validate_json(debug_response.content) + worker_pid: Final = debug_body.get("worker_pid") + assert isinstance(worker_pid, int), debug_body + body: Final = JSON_OBJECT.validate_json(response.content) + entries: Final = body.get("data") + assert isinstance(entries, list), body + count: Final = sum(object_value(entry).get("model_name") == group for entry in entries) + yield worker_pid, count + + initial_counts: Final[Mapping[int, int]] = MappingProxyType({}) + latest_counts: Final = accumulate(readbacks(), _record_latest_readback, initial=initial_counts) + for latest_by_worker in latest_counts: + if len(latest_by_worker) >= workers and all(count == deployments for count in latest_by_worker.values()): + return + raise AssertionError(f"Expected {deployments} deployments on {workers} workers before the readiness deadline") diff --git a/tests/integration/authorization/_responses_compact.py b/tests/integration/authorization/_responses_compact.py new file mode 100644 index 00000000000..8270d729667 --- /dev/null +++ b/tests/integration/authorization/_responses_compact.py @@ -0,0 +1,371 @@ +from __future__ import annotations + +import asyncio +import json +import os +import threading +import uuid +from collections import Counter +from collections.abc import Mapping +from dataclasses import dataclass, field +from hashlib import sha256 +from types import MappingProxyType +from typing import Final + +import httpx +import openai +from integration._support.client import Gateway, Scenario, eventually +from integration._support.database import read_rows +from integration._support.responses_vendor import same_response +from integration._support.wire import Reply, Request, Wire +from integration._support.worker_readiness import wait_for_worker_readiness +from openai.types.responses import CompactedResponse, ResponseInputParam +from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter + +MODEL: Final = "gpt-5" +API_KEY: Final = "synthetic-responses-key" +ALIASES: Final = ("/v1", "", "/openai/v1") +COMPACT_STATE: Final = "synthetic-compact-state" +COMPACT_USAGE: Final[Mapping[str, JsonValue]] = MappingProxyType( + { + "input_tokens": 7, + "input_tokens_details": {"cached_tokens": 0}, + "output_tokens": 3, + "output_tokens_details": {"reasoning_tokens": 0}, + "total_tokens": 10, + } +) +JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +PROXY_WORKERS: Final = max(1, int(os.environ.get("INTEGRATION_PROXY_WORKERS", "1"))) +SPEND_MARKER_HEADER: Final = "x-litellm-spend-logs-metadata" +_DISABLE_HINT: Final = ( + "To disable this security feature, set general_settings::disable_responses_id_security to True " + "in the config.yaml file." +) +OTHER_USER_DETAIL: Final = ( + f"Forbidden. The response id is not associated with the user, who this key belongs to. {_DISABLE_HINT}" +) +OTHER_TEAM_DETAIL: Final = ( + f"Forbidden. The response id is not associated with the team, who this key belongs to. {_DISABLE_HINT}" +) +UNMANAGED_DETAIL: Final = ( + "Forbidden. This response id was not issued by this proxy, so the proxy cannot tell who owns it. " + "To let keys address responses this proxy did not issue, set " + "general_settings::allow_unmanaged_response_ids to True in the config.yaml file." +) +SCRIPTED_FAILURES: Final[Mapping[str, int]] = MappingProxyType( + {"fail400": 400, "fail401": 401, "fail404": 404, "fail500": 500} +) + + +class _InputText(BaseModel): + model_config = ConfigDict(frozen=True) + + text: str + + +class _InputMessage(BaseModel): + model_config = ConfigDict(frozen=True) + + content: tuple[_InputText, ...] + + +_COMPACT_INPUT: Final = TypeAdapter(tuple[_InputMessage, ...]) + + +def marker(label: str) -> str: + return f"{label}_{uuid.uuid4().hex}" + + +def minted_id(created_marker: str) -> str: + return f"resp_minted_{created_marker}" + + +def compacted_id(compact_marker: str) -> str: + return f"resp_compacted_{compact_marker}" + + +def vendor_id() -> str: + return f"resp_vendor_{uuid.uuid4().hex}" + + +def compact_input(text: str) -> ResponseInputParam: + return [{"role": "user", "content": [{"type": "input_text", "text": text}]}] + + +def compact_input_json(text: str) -> list[JsonValue]: + return [{"role": "user", "content": [{"type": "input_text", "text": text}]}] + + +def compact_body(model: str, text: str, extra: Mapping[str, JsonValue]) -> dict[str, JsonValue]: + return {"model": model, "input": compact_input_json(text), **extra} + + +def upstream_compact_body(text: str, extra: Mapping[str, JsonValue]) -> dict[str, JsonValue]: + return {"model": MODEL, "input": compact_input_json(text), **extra} + + +def compact_text(body: Mapping[str, JsonValue]) -> str: + return _COMPACT_INPUT.validate_python(body["input"])[0].content[0].text + + +def _created_body(identity: str) -> bytes: + return json.dumps( + { + "id": identity, + "object": "response", + "created_at": 1, + "status": "completed", + "model": MODEL, + "output": [ + { + "id": f"msg_{identity}", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "created", "annotations": []}], + } + ], + "parallel_tool_calls": False, + "tool_choice": "auto", + "tools": [], + } + ).encode() + + +def _compacted_body(identity: str) -> bytes: + return json.dumps( + { + "id": identity, + "created_at": 1, + "object": "response.compaction", + "output": [ + { + "id": f"rs_{identity}", + "type": "reasoning", + "summary": [{"type": "summary_text", "text": "Compacted state"}], + "encrypted_content": COMPACT_STATE, + } + ], + "usage": dict(COMPACT_USAGE), + } + ).encode() + + +def _failure_body(status: int, text: str) -> bytes: + kind: Final = "server_error" if status >= 500 else "invalid_request_error" + return json.dumps( + {"error": {"message": f"scripted {status} for {text}", "type": kind, "param": None, "code": None}} + ).encode() + + +@dataclass(frozen=True, slots=True) +class CompactUpstream: + gates: Mapping[str, threading.Event] = field(default_factory=lambda: MappingProxyType({})) + + def respond(self, request: Request) -> Reply: + if request.method == "GET" and request.target == "/v1/models": + return Reply( + body=json.dumps( + {"object": "list", "data": [{"id": MODEL, "object": "model", "created": 1, "owned_by": "openai"}]} + ).encode() + ) + body: Final = JSON_OBJECT.validate_json(request.body) + if request.method == "POST" and request.target == "/v1/responses": + return Reply(body=_created_body(minted_id(str(body["input"])))) + assert request.method == "POST" and request.target == "/v1/responses/compact", request.target + text: Final = compact_text(body) + label: Final = text.split("_", 1)[0] + gate: Final = self.gates.get(label) + if gate is not None: + assert gate.wait(timeout=180), f"The held {label} compact was never released" + status: Final = SCRIPTED_FAILURES.get(label) + if status is not None: + return Reply(status=status, body=_failure_body(status, text)) + return Reply(body=_compacted_body(compacted_id(text))) + + +@dataclass(frozen=True, slots=True) +class Upstreamed: + target: str + body: Mapping[str, JsonValue] + + +def upstreamed(wire: Wire) -> tuple[Upstreamed, ...]: + return tuple( + Upstreamed(request.target, JSON_OBJECT.validate_json(request.body)) + for request in wire.drain() + if request.target != "/v1/models" + ) + + +def compacts(requests: tuple[Upstreamed, ...]) -> tuple[Mapping[str, JsonValue], ...]: + return tuple(request.body for request in requests if request.target == "/v1/responses/compact") + + +def compact_hits(requests: tuple[Upstreamed, ...], text: str) -> tuple[Mapping[str, JsonValue], ...]: + return tuple(body for body in compacts(requests) if compact_text(body) == text) + + +def leaks(requests: tuple[Upstreamed, ...], raw_id: str) -> bool: + return any(raw_id in json.dumps(dict(request.body)) for request in requests) + + +@dataclass(frozen=True, slots=True) +class Tenants: + model: str + owner: str + other_user: str + other_team: str + + +def deployment(scenario: Scenario, wire: Wire) -> str: + model: Final = scenario.model(model=f"openai/{MODEL}", api_key=API_KEY, api_base=f"{wire.url}/v1") + wait_for_worker_readiness(scenario.gateway, model, 1, PROXY_WORKERS) + return model + + +def tenants(scenario: Scenario, wire: Wire) -> Tenants: + model: Final = deployment(scenario, wire) + team_a: Final = scenario.team(models=[model]) + team_b: Final = scenario.team(models=[model]) + user_a: Final = scenario.member(team_a) + user_b: Final = scenario.member(team_b) + scenario.gateway.post("/team/member_add", {"team_id": team_b, "member": {"role": "user", "user_id": user_a}}) + return Tenants( + model=model, + owner=scenario.key(team_id=team_a, user_id=user_a, models=[model]), + other_user=scenario.key(team_id=team_b, user_id=user_b, models=[model]), + other_team=scenario.key(team_id=team_b, user_id=user_a, models=[model]), + ) + + +def base_url(gateway: Gateway) -> str: + return str(gateway.client.base_url).rstrip("/") + + +def spend_marker(text: str) -> dict[str, str]: + return {SPEND_MARKER_HEADER: json.dumps({"marker": text})} + + +def sdk(gateway: Gateway, key: str, prefix: str) -> openai.OpenAI: + return openai.OpenAI(base_url=f"{base_url(gateway)}{prefix}", api_key=key, max_retries=0) + + +def mint(gateway: Gateway, key: str, model: str, text: str) -> str: + created: Final = sdk(gateway, key, "/v1").responses.create( + model=model, input=text, extra_headers=spend_marker(text) + ) + assert created.id != minted_id(text) and same_response(created.id, minted_id(text)), created.model_dump_json() + return created.id + + +def sdk_compact(gateway: Gateway, key: str, prefix: str, model: str, text: str, previous: str) -> CompactedResponse: + return sdk(gateway, key, prefix).responses.compact( + model=model, + input=compact_input(text), + previous_response_id=previous, + extra_headers=spend_marker(text), + ) + + +async def _async_compact( + gateway: Gateway, key: str, prefix: str, model: str, text: str, previous: str +) -> CompactedResponse: + async with openai.AsyncOpenAI(base_url=f"{base_url(gateway)}{prefix}", api_key=key, max_retries=0) as client: + return await client.responses.compact( + model=model, + input=compact_input(text), + previous_response_id=previous, + extra_headers=spend_marker(text), + ) + + +def async_compact(gateway: Gateway, key: str, prefix: str, model: str, text: str, previous: str) -> CompactedResponse: + return asyncio.run(_async_compact(gateway, key, prefix, model, text, previous)) + + +def compact_post(gateway: Gateway, key: str, prefix: str, body: Mapping[str, JsonValue], text: str) -> httpx.Response: + return httpx.post( + f"{base_url(gateway)}{prefix}/responses/compact", + json=dict(body), + headers={"Authorization": f"Bearer {key}", **spend_marker(text)}, + timeout=120, + trust_env=False, + ) + + +def compact_raw(gateway: Gateway, key: str, content: bytes, text: str) -> httpx.Response: + return httpx.post( + f"{base_url(gateway)}/v1/responses/compact", + content=content, + headers={"Authorization": f"Bearer {key}", "content-type": "application/json", **spend_marker(text)}, + timeout=120, + trust_env=False, + ) + + +def assert_compacted(compacted: CompactedResponse, model: str, text: str) -> None: + assert compacted.id != compacted_id(text) and same_response(compacted.id, compacted_id(text)), ( + compacted.model_dump_json() + ) + assert compacted.model_dump(mode="json", exclude_none=True) == { + "id": compacted.id, + "created_at": 1, + "model": model, + "object": "response.compaction", + "output": [ + { + "id": f"rs_{compacted_id(text)}", + "type": "reasoning", + "summary": [{"type": "summary_text", "text": "Compacted state"}], + "encrypted_content": COMPACT_STATE, + } + ], + "usage": dict(COMPACT_USAGE), + }, compacted.model_dump_json() + + +def assert_compacted_reply(response: httpx.Response, model: str, text: str) -> str: + assert response.status_code == 200, response.text + compacted: Final = CompactedResponse.model_validate_json(response.content) + assert_compacted(compacted, model, text) + return compacted.id + + +def assert_refused(response: httpx.Response, detail: str) -> None: + assert response.status_code == 403, response.text + assert response.json() == { + "error": {"message": detail, "type": "permission_error", "param": None, "code": "403"} + }, response.text + + +def spend_rows(key: str) -> tuple[Mapping[str, JsonValue], ...]: + return tuple( + read_rows( + "SELECT request_id, status, call_type, metadata->'spend_logs_metadata'->>'marker' AS marker " + 'FROM "LiteLLM_SpendLogs" WHERE api_key = %s', + (sha256(key.encode()).hexdigest(),), + ) + ) + + +def _marker_counts(key: str, texts: frozenset[str]) -> Mapping[str, int]: + return Counter(str(row["marker"]) for row in spend_rows(key) if row["marker"] in texts) + + +def landed_once(key: str, texts: frozenset[str]) -> tuple[Mapping[str, JsonValue], ...]: + counts: Final = eventually(lambda: _marker_counts(key, texts), lambda found: frozenset(found) == texts, seconds=90) + assert dict(counts) == dict.fromkeys(texts, 1), counts + return tuple(row for row in spend_rows(key) if row["marker"] in texts) + + +def assert_compact_rows(key: str, texts: frozenset[str]) -> None: + rows: Final = landed_once(key, texts) + assert {(row["status"], row["call_type"]) for row in rows} == {("success", "acompact_responses")}, rows + + +def successful_markers(key: str, texts: frozenset[str]) -> frozenset[str]: + return frozenset( + str(row["marker"]) for row in spend_rows(key) if row["marker"] in texts and row["status"] == "success" + ) diff --git a/tests/integration/authorization/test_responses_compact_id_ownership.py b/tests/integration/authorization/test_responses_compact_id_ownership.py new file mode 100644 index 00000000000..e57118da129 --- /dev/null +++ b/tests/integration/authorization/test_responses_compact_id_ownership.py @@ -0,0 +1,720 @@ +from __future__ import annotations + +import json +from collections.abc import Callable, Mapping +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from functools import partial +from pathlib import Path +from typing import Final, Literal + +import httpx +import pytest +from integration._support.client import Gateway, eventually +from integration._support.process import owned_proxy +from integration._support.responses_vendor import same_response +from integration._support.wire import wire_server +from integration.authorization._responses_compact import ( + ALIASES, + MODEL, + OTHER_TEAM_DETAIL, + OTHER_USER_DETAIL, + SCRIPTED_FAILURES, + UNMANAGED_DETAIL, + CompactUpstream, + Tenants, + Upstreamed, + assert_compact_rows, + assert_compacted, + assert_compacted_reply, + assert_refused, + async_compact, + base_url, + compact_body, + compact_hits, + compact_post, + compact_raw, + compacted_id, + compacts, + leaks, + marker, + mint, + minted_id, + sdk, + sdk_compact, + successful_markers, + tenants, + upstream_compact_body, + upstreamed, + vendor_id, +) +from pydantic import JsonValue + +_CLIENTS: Final = ("sync", "async", "httpx") + + +def _compact_with( + client: Literal["sync", "async", "httpx"], + gateway: Gateway, + key: str, + prefix: str, + model: str, + text: str, + previous: str, +) -> None: + match client: + case "sync": + assert_compacted(sdk_compact(gateway, key, prefix, model, text, previous), model, text) + case "async": + assert_compacted(async_compact(gateway, key, prefix, model, text, previous), model, text) + case "httpx": + assert_compacted_reply( + compact_post(gateway, key, prefix, compact_body(model, text, {"previous_response_id": previous}), text), + model, + text, + ) + + +@pytest.mark.parametrize("prefix", ALIASES) +@pytest.mark.parametrize("client", _CLIENTS) +def test_owner_compact_forwards_the_raw_previous_response_id( + gateway: Gateway, client: Literal["sync", "async", "httpx"], prefix: str +) -> None: + upstream: Final = CompactUpstream() + with wire_server(upstream.respond) as wire, gateway.scenario() as scenario: + actors: Final = tenants(scenario, wire) + created: Final = marker("mint") + owned: Final = mint(gateway, actors.owner, actors.model, created) + assert len(upstreamed(wire)) == 1 + text: Final = marker("compact") + _compact_with(client, gateway, actors.owner, prefix, actors.model, text, owned) + assert compacts(upstreamed(wire)) == ( + upstream_compact_body(text, {"previous_response_id": minted_id(created)}), + ) + assert_compact_rows(actors.owner, frozenset({text})) + + +def test_admin_compact_decodes_owned_ids_and_forwards_vendor_ids(gateway: Gateway) -> None: + upstream: Final = CompactUpstream() + with wire_server(upstream.respond) as wire, gateway.scenario() as scenario: + actors: Final = tenants(scenario, wire) + created: Final = marker("mint") + owned: Final = mint(gateway, actors.owner, actors.model, created) + upstreamed(wire) + foreign: Final = vendor_id() + owned_text: Final = marker("admin_owned") + foreign_text: Final = marker("admin_vendor") + assert_compacted( + sdk_compact(gateway, gateway.key, "/v1", actors.model, owned_text, owned), actors.model, owned_text + ) + assert_compacted( + sdk_compact(gateway, gateway.key, "/v1", actors.model, foreign_text, foreign), actors.model, foreign_text + ) + assert compacts(upstreamed(wire)) == ( + upstream_compact_body(owned_text, {"previous_response_id": minted_id(created)}), + upstream_compact_body(foreign_text, {"previous_response_id": foreign}), + ) + + +def test_compacted_response_id_belongs_to_the_compacting_owner(gateway: Gateway) -> None: + upstream: Final = CompactUpstream() + with wire_server(upstream.respond) as wire, gateway.scenario() as scenario: + actors: Final = tenants(scenario, wire) + created: Final = marker("mint") + owned: Final = mint(gateway, actors.owner, actors.model, created) + first: Final = marker("compact") + compacted: Final = sdk_compact(gateway, actors.owner, "/v1", actors.model, first, owned) + assert_compacted(compacted, actors.model, first) + upstreamed(wire) + follow: Final = marker("follow") + followed: Final = sdk(gateway, actors.owner, "/v1").responses.create( + model=actors.model, input=follow, previous_response_id=compacted.id + ) + assert same_response(followed.id, minted_id(follow)), followed.model_dump_json() + assert [(request.target, request.body) for request in upstreamed(wire)] == [ + ("/v1/responses", {"model": MODEL, "input": follow, "previous_response_id": compacted_id(first)}) + ] + second: Final = marker("compact_again") + assert_compacted( + sdk_compact(gateway, actors.owner, "", actors.model, second, compacted.id), actors.model, second + ) + assert compacts(upstreamed(wire)) == ( + upstream_compact_body(second, {"previous_response_id": compacted_id(first)}), + ) + refused_user: Final = marker("refused_user") + refused_team: Final = marker("refused_team") + assert_refused( + compact_post( + gateway, + actors.other_user, + "/v1", + compact_body(actors.model, refused_user, {"previous_response_id": compacted.id}), + refused_user, + ), + OTHER_USER_DETAIL, + ) + assert_refused( + compact_post( + gateway, + actors.other_team, + "/openai/v1", + compact_body(actors.model, refused_team, {"previous_response_id": compacted.id}), + refused_team, + ), + OTHER_TEAM_DETAIL, + ) + assert upstreamed(wire) == () + + +def test_compact_without_previous_response_id_is_forwarded_unchanged(gateway: Gateway) -> None: + upstream: Final = CompactUpstream() + with wire_server(upstream.respond) as wire, gateway.scenario() as scenario: + actors: Final = tenants(scenario, wire) + texts: Final = tuple(marker(f"bare{index}") for index in range(3)) + for key, prefix, text in zip((actors.owner, actors.other_user, gateway.key), ALIASES, texts, strict=True): + assert_compacted_reply( + compact_post(gateway, key, prefix, compact_body(actors.model, text, {}), text), actors.model, text + ) + assert compacts(upstreamed(wire)) == tuple(upstream_compact_body(text, {}) for text in texts) + assert_compact_rows(actors.owner, frozenset({texts[0]})) + + +@pytest.mark.parametrize("prefix", ALIASES) +def test_compact_refuses_other_tenants_and_unmanaged_ids_before_the_provider(gateway: Gateway, prefix: str) -> None: + upstream: Final = CompactUpstream() + with wire_server(upstream.respond) as wire, gateway.scenario() as scenario: + actors: Final = tenants(scenario, wire) + created: Final = marker("mint") + owned: Final = mint(gateway, actors.owner, actors.model, created) + upstreamed(wire) + refusals: Final = ( + (actors.other_user, owned, OTHER_USER_DETAIL, marker("other_user")), + (actors.other_team, owned, OTHER_TEAM_DETAIL, marker("other_team")), + (actors.owner, vendor_id(), UNMANAGED_DETAIL, marker("unmanaged")), + ) + for key, previous, detail, text in refusals: + assert_refused( + compact_post( + gateway, key, prefix, compact_body(actors.model, text, {"previous_response_id": previous}), text + ), + detail, + ) + assert upstreamed(wire) == () + barrier: Final = marker("barrier") + assert_compacted_reply( + compact_post( + gateway, + actors.owner, + prefix, + compact_body(actors.model, barrier, {"previous_response_id": owned}), + barrier, + ), + actors.model, + barrier, + ) + assert compacts(upstreamed(wire)) == ( + upstream_compact_body(barrier, {"previous_response_id": minted_id(created)}), + ) + assert_compact_rows(actors.owner, frozenset({barrier})) + refused_texts: Final = frozenset(text for _, _, _, text in refusals) + assert successful_markers(actors.owner, refused_texts) == frozenset() + assert successful_markers(actors.other_user, refused_texts) == frozenset() + assert successful_markers(actors.other_team, refused_texts) == frozenset() + + +_MALFORMED: Final[Mapping[str, Callable[[str], JsonValue]]] = { + "int": lambda owned: 0, + "list": lambda owned: [owned], + "object": lambda owned: {"id": owned}, + "empty": lambda owned: "", + "null": lambda owned: None, +} + + +@pytest.mark.parametrize("shape", tuple(_MALFORMED)) +def test_non_string_previous_response_id_never_reaches_the_provider_decoded(gateway: Gateway, shape: str) -> None: + upstream: Final = CompactUpstream() + with wire_server(upstream.respond) as wire, gateway.scenario() as scenario: + actors: Final = tenants(scenario, wire) + created: Final = marker("mint") + owned: Final = mint(gateway, actors.owner, actors.model, created) + upstreamed(wire) + text: Final = marker(f"malformed_{shape}") + value: Final = _MALFORMED[shape](owned) + response: Final = compact_post( + gateway, actors.other_user, "/v1", compact_body(actors.model, text, {"previous_response_id": value}), text + ) + received: Final = upstreamed(wire) + assert not leaks(received, minted_id(created)), received + assert (response.status_code, compacts(received)) == _MALFORMED_OUTCOME[shape](text, value), response.text + assert gateway.request("GET", "/health/liveliness").status_code == 200 + + +def _forwarded(text: str, value: JsonValue) -> tuple[int, tuple[Mapping[str, JsonValue], ...]]: + return 200, (upstream_compact_body(text, {"previous_response_id": value}),) + + +def _dropped(text: str, value: JsonValue) -> tuple[int, tuple[Mapping[str, JsonValue], ...]]: + return 200, (upstream_compact_body(text, {}),) + + +_MALFORMED_OUTCOME: Final[Mapping[str, Callable[[str, JsonValue], tuple[int, tuple[Mapping[str, JsonValue], ...]]]]] = { + "int": _forwarded, + "list": _forwarded, + "object": _forwarded, + "empty": _forwarded, + "null": _dropped, +} + + +def _tampered(owned: str) -> str: + middle: Final = len(owned) // 2 + swapped: Final = "B" if owned[middle] == "A" else "A" + return f"{owned[:middle]}{swapped}{owned[middle + 1 :]}" + + +_DEGENERATE: Final[Mapping[str, Callable[[str], str]]] = { + "oversized": lambda owned: "resp_" + "a" * 5120, + "oversized_unprefixed": lambda owned: "x" * 5120, + "bare_prefix": lambda owned: "resp_", + "tampered": _tampered, + "doubled_prefix": lambda owned: f"resp_{owned}", +} + + +@pytest.mark.parametrize("shape", tuple(_DEGENERATE)) +def test_degenerate_previous_response_ids_are_refused_as_unmanaged(gateway: Gateway, shape: str) -> None: + upstream: Final = CompactUpstream() + with wire_server(upstream.respond) as wire, gateway.scenario() as scenario: + actors: Final = tenants(scenario, wire) + created: Final = marker("mint") + owned: Final = mint(gateway, actors.owner, actors.model, created) + upstreamed(wire) + text: Final = marker(f"degenerate_{shape}") + assert_refused( + compact_post( + gateway, + actors.owner, + "/v1", + compact_body(actors.model, text, {"previous_response_id": _DEGENERATE[shape](owned)}), + text, + ), + UNMANAGED_DETAIL, + ) + assert upstreamed(wire) == () + assert gateway.request("GET", "/health/liveliness").status_code == 200 + + +def _duplicated(model: str, text: str, first: str, last: str) -> bytes: + return ( + f'{{"model": {json.dumps(model)}, "input": [{{"role": "user", "content": [{{"type": "input_text", ' + f'"text": {json.dumps(text)}}}]}}], "previous_response_id": {json.dumps(first)}, ' + f'"previous_response_id": {json.dumps(last)}}}' + ).encode() + + +def test_duplicated_previous_response_id_is_checked_on_the_value_the_provider_gets(gateway: Gateway) -> None: + upstream: Final = CompactUpstream() + with wire_server(upstream.respond) as wire, gateway.scenario() as scenario: + actors: Final = tenants(scenario, wire) + victim_mint: Final = marker("victim") + attacker_mint: Final = marker("attacker") + victim: Final = mint(gateway, actors.owner, actors.model, victim_mint) + attacker: Final = mint(gateway, actors.other_user, actors.model, attacker_mint) + upstreamed(wire) + same: Final = marker("dup_same") + assert_compacted_reply( + compact_raw(gateway, actors.owner, _duplicated(actors.model, same, victim, victim), same), + actors.model, + same, + ) + victim_last: Final = marker("dup_victim_last") + assert_refused( + compact_raw( + gateway, actors.other_user, _duplicated(actors.model, victim_last, attacker, victim), victim_last + ), + OTHER_USER_DETAIL, + ) + own_last: Final = marker("dup_own_last") + assert_compacted_reply( + compact_raw(gateway, actors.other_user, _duplicated(actors.model, own_last, victim, attacker), own_last), + actors.model, + own_last, + ) + received: Final = upstreamed(wire) + assert compacts(received) == ( + upstream_compact_body(same, {"previous_response_id": minted_id(victim_mint)}), + upstream_compact_body(own_last, {"previous_response_id": minted_id(attacker_mint)}), + ) + + +def test_client_sent_retained_response_id_cannot_address_another_owner(gateway: Gateway) -> None: + upstream: Final = CompactUpstream() + with wire_server(upstream.respond) as wire, gateway.scenario() as scenario: + actors: Final = tenants(scenario, wire) + victim_mint: Final = marker("victim") + attacker_mint: Final = marker("attacker") + victim: Final = mint(gateway, actors.owner, actors.model, victim_mint) + attacker: Final = mint(gateway, actors.other_user, actors.model, attacker_mint) + upstreamed(wire) + retained_victim: Final = marker("retained_victim") + assert_refused( + compact_post( + gateway, + actors.other_user, + "/v1", + compact_body( + actors.model, + retained_victim, + {"previous_response_id": attacker, "_litellm_addressed_response_id": victim}, + ), + retained_victim, + ), + OTHER_USER_DETAIL, + ) + retained_number: Final = marker("retained_number") + assert_refused( + compact_post( + gateway, + actors.other_user, + "/v1", + compact_body( + actors.model, retained_number, {"previous_response_id": victim, "_litellm_addressed_response_id": 7} + ), + retained_number, + ), + OTHER_USER_DETAIL, + ) + assert upstreamed(wire) == () + retained_own: Final = marker("retained_own") + assert_compacted_reply( + compact_post( + gateway, + actors.other_user, + "/v1", + compact_body( + actors.model, + retained_own, + {"previous_response_id": victim, "_litellm_addressed_response_id": attacker}, + ), + retained_own, + ), + actors.model, + retained_own, + ) + received: Final = upstreamed(wire) + assert compacts(received) == ( + upstream_compact_body(retained_own, {"previous_response_id": minted_id(attacker_mint)}), + ) + assert not leaks(received, minted_id(victim_mint)), received + + +@pytest.mark.parametrize("status", tuple(SCRIPTED_FAILURES.values())) +def test_provider_errors_on_an_owned_compact_reach_the_caller(gateway: Gateway, status: int) -> None: + upstream: Final = CompactUpstream() + with wire_server(upstream.respond) as wire, gateway.scenario() as scenario: + actors: Final = tenants(scenario, wire) + created: Final = marker("mint") + owned: Final = mint(gateway, actors.owner, actors.model, created) + upstreamed(wire) + failing: Final = marker(f"fail{status}") + response: Final = compact_post( + gateway, actors.owner, "/v1", compact_body(actors.model, failing, {"previous_response_id": owned}), failing + ) + assert response.status_code == status, response.text + assert f"scripted {status} for {failing}" in response.json()["error"]["message"], response.text + assert compact_hits(upstreamed(wire), failing) == ( + upstream_compact_body(failing, {"previous_response_id": minted_id(created)}), + ) + after: Final = marker("after_failure") + assert_compacted_reply( + compact_post( + gateway, actors.owner, "/v1", compact_body(actors.model, after, {"previous_response_id": owned}), after + ), + actors.model, + after, + ) + assert compacts(upstreamed(wire)) == ( + upstream_compact_body(after, {"previous_response_id": minted_id(created)}), + ) + assert_compact_rows(actors.owner, frozenset({after})) + + +def test_response_minted_by_a_key_without_owner_is_addressable_by_any_key(gateway: Gateway) -> None: + upstream: Final = CompactUpstream() + with wire_server(upstream.respond) as wire, gateway.scenario() as scenario: + actors: Final = tenants(scenario, wire) + ownerless: Final = scenario.key(models=[actors.model]) + created: Final = marker("ownerless") + minted: Final = mint(gateway, ownerless, actors.model, created) + upstreamed(wire) + text: Final = marker("compact_ownerless") + assert_compacted(sdk_compact(gateway, actors.other_user, "/v1", actors.model, text, minted), actors.model, text) + assert compacts(upstreamed(wire)) == ( + upstream_compact_body(text, {"previous_response_id": minted_id(created)}), + ) + + +def test_identical_compacts_each_reach_the_provider_with_the_cache_on(gateway: Gateway) -> None: + upstream: Final = CompactUpstream() + with wire_server(upstream.respond) as wire, gateway.scenario() as scenario: + actors: Final = tenants(scenario, wire) + created: Final = marker("mint") + owned: Final = mint(gateway, actors.owner, actors.model, created) + upstreamed(wire) + text: Final = marker("twin") + body: Final = compact_body(actors.model, text, {"previous_response_id": owned}) + rows: Final = (marker("twin_row"), marker("twin_row")) + for row in rows: + assert_compacted_reply(compact_post(gateway, actors.owner, "/v1", body, row), actors.model, text) + assert ( + compacts(upstreamed(wire)) + == (upstream_compact_body(text, {"previous_response_id": minted_id(created)}),) * 2 + ) + assert_compact_rows(actors.owner, frozenset(rows)) + assert_refused(compact_post(gateway, actors.other_user, "/v1", body, marker("twin_refused")), OTHER_USER_DETAIL) + assert upstreamed(wire) == () + + +def test_cached_create_still_checks_the_previous_response_id_owner(gateway: Gateway) -> None: + upstream: Final = CompactUpstream() + with wire_server(upstream.respond) as wire, gateway.scenario() as scenario: + actors: Final = tenants(scenario, wire) + created: Final = marker("mint") + owned: Final = mint(gateway, actors.owner, actors.model, created) + upstreamed(wire) + follow: Final = marker("cached_follow") + owner_client: Final = sdk(gateway, actors.owner, "/v1") + sent: Final = {"model": MODEL, "input": follow, "previous_response_id": minted_id(created)} + + def create() -> tuple[str, tuple[Upstreamed, ...]]: + reply: Final = owner_client.responses.create(model=actors.model, input=follow, previous_response_id=owned) + return reply.id, upstreamed(wire) + + first_id, first_received = create() + assert [(request.target, request.body) for request in first_received] == [("/v1/responses", sent)] + hit_id, hit_received = eventually(create, lambda attempt: not attempt[1], seconds=60) + assert same_response(first_id, minted_id(follow)) and same_response(hit_id, minted_id(follow)), hit_id + assert hit_received == () + refused: Final = httpx.post( + f"{base_url(gateway)}/v1/responses", + json={"model": actors.model, "input": follow, "previous_response_id": owned}, + headers={"Authorization": f"Bearer {actors.other_user}"}, + timeout=120, + trust_env=False, + ) + assert_refused(refused, OTHER_USER_DETAIL) + assert upstreamed(wire) == () + + +@dataclass(frozen=True, slots=True) +class _BurstCall: + kind: Literal["owner", "other_user", "other_team", "unmanaged", "admin_vendor", "bare"] + prefix: str + text: str + + +_BURST_KINDS: Final = ("owner", "other_user", "other_team", "unmanaged", "admin_vendor", "bare") +_REFUSALS: Final[Mapping[str, str]] = { + "other_user": OTHER_USER_DETAIL, + "other_team": OTHER_TEAM_DETAIL, + "unmanaged": UNMANAGED_DETAIL, +} + + +def _burst_call(gateway: Gateway, actors: Tenants, owned: str, foreign: str, call: _BurstCall) -> httpx.Response: + plans: Final[Mapping[str, tuple[str, Mapping[str, JsonValue]]]] = { + "owner": (actors.owner, {"previous_response_id": owned}), + "other_user": (actors.other_user, {"previous_response_id": owned}), + "other_team": (actors.other_team, {"previous_response_id": owned}), + "unmanaged": (actors.owner, {"previous_response_id": foreign}), + "admin_vendor": (gateway.key, {"previous_response_id": foreign}), + "bare": (actors.owner, {}), + } + key, extra = plans[call.kind] + return compact_post(gateway, key, call.prefix, compact_body(actors.model, call.text, extra), call.text) + + +def test_concurrent_mixed_compact_burst_gives_every_caller_its_own_verdict(gateway: Gateway) -> None: + upstream: Final = CompactUpstream() + with wire_server(upstream.respond) as wire, gateway.scenario() as scenario: + actors: Final = tenants(scenario, wire) + created: Final = marker("mint") + owned: Final = mint(gateway, actors.owner, actors.model, created) + upstreamed(wire) + foreign: Final = vendor_id() + calls: Final = tuple( + _BurstCall(_BURST_KINDS[index % 6], ALIASES[index % 3], marker(f"burst{index}")) for index in range(36) + ) + with ThreadPoolExecutor(max_workers=len(calls)) as pool: + responses: Final = tuple(pool.map(partial(_burst_call, gateway, actors, owned, foreign), calls)) + for call, response in zip(calls, responses, strict=True): + if call.kind in _REFUSALS: + assert_refused(response, _REFUSALS[call.kind]) + else: + assert_compacted_reply(response, actors.model, call.text) + received: Final = upstreamed(wire) + expected: Final[Mapping[str, Mapping[str, JsonValue]]] = { + "owner": {"previous_response_id": minted_id(created)}, + "admin_vendor": {"previous_response_id": foreign}, + "bare": {}, + } + assert {call.text: compact_hits(received, call.text) for call in calls} == { + call.text: (upstream_compact_body(call.text, expected[call.kind]),) if call.kind in expected else () + for call in calls + } + owner_texts: Final = frozenset(call.text for call in calls if call.kind in ("owner", "bare")) + assert_compact_rows(actors.owner, owner_texts) + refused_texts: Final = frozenset(call.text for call in calls if call.kind in _REFUSALS) + assert successful_markers(actors.owner, refused_texts) == frozenset() + assert successful_markers(actors.other_user, refused_texts) == frozenset() + assert successful_markers(actors.other_team, refused_texts) == frozenset() + + +def _settings_proxy_config(tmp_path: Path, name: str, settings: str) -> Path: + config: Final = tmp_path / f"{name}.yaml" + config.write_text(f"general_settings:\n{settings}") + return config + + +def test_allow_unmanaged_response_ids_still_scopes_proxy_issued_ids_on_compact( + gateway: Gateway, tmp_path: Path +) -> None: + config: Final = _settings_proxy_config(tmp_path, "allow_unmanaged", " allow_unmanaged_response_ids: true\n") + upstream: Final = CompactUpstream() + with wire_server(upstream.respond) as wire, gateway.scenario() as scenario: + actors: Final = tenants(scenario, wire) + with owned_proxy(gateway, tmp_path, {"INTEGRATION_PROXY_WORKERS": "1"}, config=config) as allowed: + created: Final = marker("mint") + owned: Final = mint(allowed, actors.owner, actors.model, created) + upstreamed(wire) + foreign: Final = vendor_id() + vendor_text: Final = marker("allowed_vendor") + assert_compacted_reply( + compact_post( + allowed, + actors.owner, + "/v1", + compact_body(actors.model, vendor_text, {"previous_response_id": foreign}), + vendor_text, + ), + actors.model, + vendor_text, + ) + other_text: Final = marker("allowed_other_user") + assert_refused( + compact_post( + allowed, + actors.other_user, + "/v1", + compact_body(actors.model, other_text, {"previous_response_id": owned}), + other_text, + ), + OTHER_USER_DETAIL, + ) + owner_text: Final = marker("allowed_owner") + assert_compacted_reply( + compact_post( + allowed, + actors.owner, + "/v1", + compact_body(actors.model, owner_text, {"previous_response_id": owned}), + owner_text, + ), + actors.model, + owner_text, + ) + received: Final = upstreamed(wire) + assert compacts(received) == ( + upstream_compact_body(vendor_text, {"previous_response_id": foreign}), + upstream_compact_body(owner_text, {"previous_response_id": minted_id(created)}), + ) + + +def test_disable_responses_id_security_lets_any_key_compact_and_forwards_raw_ids( + gateway: Gateway, tmp_path: Path +) -> None: + config: Final = _settings_proxy_config(tmp_path, "disable_security", " disable_responses_id_security: true\n") + upstream: Final = CompactUpstream() + with wire_server(upstream.respond) as wire, gateway.scenario() as scenario: + actors: Final = tenants(scenario, wire) + sealed_mint: Final = marker("sealed") + sealed: Final = mint(gateway, actors.owner, actors.model, sealed_mint) + with owned_proxy(gateway, tmp_path, {"INTEGRATION_PROXY_WORKERS": "1"}, config=config) as disabled: + open_mint: Final = marker("open") + opened: Final = mint(disabled, actors.owner, actors.model, open_mint) + upstreamed(wire) + foreign: Final = vendor_id() + plans: Final = ( + (marker("disabled_open"), opened, minted_id(open_mint)), + (marker("disabled_sealed"), sealed, minted_id(sealed_mint)), + (marker("disabled_vendor"), foreign, foreign), + ) + for text, previous, _ in plans: + assert_compacted_reply( + compact_post( + disabled, + actors.other_user, + "/v1", + compact_body(actors.model, text, {"previous_response_id": previous}), + text, + ), + actors.model, + text, + ) + received: Final = upstreamed(wire) + assert compacts(received) == tuple( + upstream_compact_body(text, {"previous_response_id": raw}) for text, _, raw in plans + ) + + +def test_null_and_empty_id_security_settings_keep_the_compact_checks(gateway: Gateway, tmp_path: Path) -> None: + config: Final = _settings_proxy_config( + tmp_path, + "null_and_empty", + " allow_unmanaged_response_ids: null\n disable_responses_id_security: ''\n", + ) + upstream: Final = CompactUpstream() + with wire_server(upstream.respond) as wire, gateway.scenario() as scenario: + actors: Final = tenants(scenario, wire) + with owned_proxy(gateway, tmp_path, {"INTEGRATION_PROXY_WORKERS": "1"}, config=config) as defaulted: + created: Final = marker("mint") + owned: Final = mint(defaulted, actors.owner, actors.model, created) + upstreamed(wire) + vendor_text: Final = marker("defaulted_vendor") + assert_refused( + compact_post( + defaulted, + actors.owner, + "/v1", + compact_body(actors.model, vendor_text, {"previous_response_id": vendor_id()}), + vendor_text, + ), + UNMANAGED_DETAIL, + ) + other_text: Final = marker("defaulted_other_user") + assert_refused( + compact_post( + defaulted, + actors.other_user, + "/v1", + compact_body(actors.model, other_text, {"previous_response_id": owned}), + other_text, + ), + OTHER_USER_DETAIL, + ) + assert upstreamed(wire) == () + owner_text: Final = marker("defaulted_owner") + assert_compacted_reply( + compact_post( + defaulted, + actors.owner, + "/v1", + compact_body(actors.model, owner_text, {"previous_response_id": owned}), + owner_text, + ), + actors.model, + owner_text, + ) + received: Final = upstreamed(wire) + assert compacts(received) == (upstream_compact_body(owner_text, {"previous_response_id": minted_id(created)}),) diff --git a/tests/integration/authorization/test_responses_compact_id_ownership_chaos.py b/tests/integration/authorization/test_responses_compact_id_ownership_chaos.py new file mode 100644 index 00000000000..d74838da9cb --- /dev/null +++ b/tests/integration/authorization/test_responses_compact_id_ownership_chaos.py @@ -0,0 +1,200 @@ +from __future__ import annotations + +import threading +from collections.abc import Mapping +from concurrent.futures import Future, ThreadPoolExecutor +from dataclasses import dataclass +from types import MappingProxyType +from typing import Final, Literal +from urllib.parse import urlsplit + +import httpx +from integration._support.client import Gateway, eventually +from integration._support.wire import wire_server +from integration.authorization._responses_compact import ( + ALIASES, + OTHER_TEAM_DETAIL, + OTHER_USER_DETAIL, + UNMANAGED_DETAIL, + CompactUpstream, + Tenants, + Upstreamed, + assert_compact_rows, + assert_compacted_reply, + assert_refused, + compact_body, + compact_hits, + compact_post, + marker, + mint, + minted_id, + successful_markers, + tenants, + upstream_compact_body, + upstreamed, + vendor_id, +) + +_Kind = Literal["owner", "other_user", "other_team", "unmanaged"] +_KINDS: Final[tuple[_Kind, ...]] = ("owner", "other_user", "other_team", "unmanaged") +_REFUSING: Final[tuple[_Kind, ...]] = ("other_user", "other_team", "unmanaged") +_OWNER: Final[tuple[_Kind, ...]] = ("owner",) +_REFUSALS: Final[Mapping[str, str]] = { + "other_user": OTHER_USER_DETAIL, + "other_team": OTHER_TEAM_DETAIL, + "unmanaged": UNMANAGED_DETAIL, +} + + +@dataclass(frozen=True, slots=True) +class _Call: + kind: _Kind + prefix: str + text: str + + +@dataclass(frozen=True, slots=True) +class _Outcome: + call: _Call + response: httpx.Response | None + + +@dataclass(frozen=True, slots=True) +class _Target: + gateway: Gateway + actors: Tenants + owned: str + foreign: str + + +def _calls(label: str, count: int, kinds: tuple[_Kind, ...] = _KINDS) -> tuple[_Call, ...]: + return tuple( + _Call(kinds[index % len(kinds)], ALIASES[index % len(ALIASES)], marker(label)) for index in range(count) + ) + + +def _send(target: _Target, call: _Call) -> _Outcome: + plans: Final[Mapping[str, tuple[str, str]]] = { + "owner": (target.actors.owner, target.owned), + "other_user": (target.actors.other_user, target.owned), + "other_team": (target.actors.other_team, target.owned), + "unmanaged": (target.actors.owner, target.foreign), + } + key, previous = plans[call.kind] + body: Final = compact_body(target.actors.model, call.text, {"previous_response_id": previous}) + try: + return _Outcome(call, compact_post(target.gateway, key, call.prefix, body, call.text)) + except httpx.TransportError: + return _Outcome(call, None) + + +def _submit(pool: ThreadPoolExecutor, target: _Target, calls: tuple[_Call, ...]) -> tuple[Future[_Outcome], ...]: + return tuple(pool.submit(_send, target, call) for call in calls) + + +def _results(futures: tuple[Future[_Outcome], ...]) -> tuple[_Outcome, ...]: + return tuple(future.result(timeout=180) for future in futures) + + +def _send_all(target: _Target, calls: tuple[_Call, ...]) -> tuple[_Outcome, ...]: + with ThreadPoolExecutor(max_workers=len(calls)) as pool: + return _results(_submit(pool, target, calls)) + + +def _done(futures: tuple[Future[_Outcome], ...]) -> int: + return sum(future.done() for future in futures) + + +def _assert_verdict(outcome: _Outcome, model: str) -> None: + assert outcome.response is not None, outcome.call + if outcome.call.kind in _REFUSALS: + assert_refused(outcome.response, _REFUSALS[outcome.call.kind]) + return + assert_compacted_reply(outcome.response, model, outcome.call.text) + + +def _texts(outcomes: tuple[_Outcome, ...], kinds: tuple[_Kind, ...]) -> frozenset[str]: + return frozenset(outcome.call.text for outcome in outcomes if outcome.call.kind in kinds) + + +def _assert_upstream( + received: tuple[Upstreamed, ...], owner_texts: frozenset[str], refused_texts: frozenset[str], created: str +) -> None: + assert {text: compact_hits(received, text) for text in owner_texts} == { + text: (upstream_compact_body(text, {"previous_response_id": minted_id(created)}),) for text in owner_texts + } + assert {text: compact_hits(received, text) for text in refused_texts} == dict.fromkeys(refused_texts, ()) + + +def _settled(target: _Target) -> bool: + text: Final = marker("settle") + body: Final = compact_body(target.actors.model, text, {}) + try: + return compact_post(target.gateway, target.actors.owner, "/v1", body, text).status_code == 200 + except httpx.TransportError: + return False + + +def test_upstream_outage_mid_burst_keeps_refusals_and_lands_each_owner_compact_once(gateway: Gateway) -> None: + upstream: Final = CompactUpstream() + created: Final = marker("mint") + with gateway.scenario() as scenario: + with wire_server(upstream.respond) as first: + port: Final = urlsplit(first.url).port + actors: Final = tenants(scenario, first) + target: Final = _Target(gateway, actors, mint(gateway, actors.owner, actors.model, created), vendor_id()) + upstreamed(first) + before: Final = _send_all(target, _calls("before", 16)) + before_received: Final = upstreamed(first) + during: Final = _send_all(target, _calls("during", 16)) + alive: Final = gateway.request("GET", "/health/liveliness") + assert port is not None + with wire_server(upstream.respond, port=port) as second: + eventually(lambda: _settled(target), lambda settled: settled, seconds=60) + after: Final = _send_all(target, _calls("after", 16)) + after_received: Final = upstreamed(second) + assert alive.status_code == 200, alive.text + for outcome in (*before, *after): + _assert_verdict(outcome, actors.model) + for outcome in during: + assert outcome.response is not None, outcome.call + if outcome.call.kind in _REFUSALS: + assert_refused(outcome.response, _REFUSALS[outcome.call.kind]) + else: + assert outcome.response.status_code >= 500, outcome.response.text + _assert_upstream(before_received, _texts(before, _OWNER), _texts(before, _REFUSING), created) + _assert_upstream(after_received, _texts(after, _OWNER), _texts((*during, *after), _REFUSING), created) + assert {text: compact_hits(after_received, text) for text in _texts(during, _OWNER)} == dict.fromkeys( + _texts(during, _OWNER), () + ) + assert_compact_rows(actors.owner, _texts((*before, *after), _OWNER)) + refused: Final = _texts((*before, *during, *after), _REFUSING) + assert successful_markers(actors.other_user, refused) == frozenset() + assert successful_markers(actors.other_team, refused) == frozenset() + assert successful_markers(actors.owner, refused | _texts(during, _OWNER)) == frozenset() + + +def test_held_owner_compacts_do_not_hold_up_refusals(gateway: Gateway) -> None: + release: Final = threading.Event() + upstream: Final = CompactUpstream(MappingProxyType({"held": release})) + created: Final = marker("mint") + with wire_server(upstream.respond) as wire, gateway.scenario() as scenario: + actors: Final = tenants(scenario, wire) + target: Final = _Target(gateway, actors, mint(gateway, actors.owner, actors.model, created), vendor_id()) + upstreamed(wire) + held: Final = _calls("held", 8, _OWNER) + with ThreadPoolExecutor(max_workers=len(held)) as pool: + try: + pending: Final = _submit(pool, target, held) + eventually(lambda: wire.received.qsize(), lambda size: size >= len(held), seconds=90) + refused: Final = _send_all(target, _calls("refusal", 12, _REFUSING)) + done_while_held: Final = _done(pending) + finally: + release.set() + served: Final = _results(pending) + received: Final = upstreamed(wire) + assert done_while_held == 0 + for outcome in (*refused, *served): + _assert_verdict(outcome, actors.model) + _assert_upstream(received, _texts(served, _OWNER), _texts(refused, _REFUSING), created) + assert_compact_rows(actors.owner, _texts(served, _OWNER)) diff --git a/tests/unit/test_responses_id_security.py b/tests/unit/test_responses_id_security.py index 704a52fc202..c7fd7fe247f 100644 --- a/tests/unit/test_responses_id_security.py +++ b/tests/unit/test_responses_id_security.py @@ -805,6 +805,7 @@ _FABRICATED_UNMANAGED_ID = "resp_fabricatedunmanagedidbbbbbbbbbbbbbbbb" _UNIT_TEST_SALT_KEY = "lit6837-unit-test-salt-key" _ADDRESSED_ID_FIELD_BY_CALL_TYPE = { "aresponses": "previous_response_id", + "acompact_responses": "previous_response_id", "aget_responses": "response_id", "adelete_responses": "response_id", "acancel_responses": "response_id", @@ -865,6 +866,52 @@ class TestUnrecognizedResponseIdIsRejected: assert "allow_unmanaged_response_ids" in exc_info.value.detail assert data[field] == _FABRICATED_UNMANAGED_ID + @pytest.mark.asyncio + @pytest.mark.parametrize("call_type", sorted(_ADDRESSED_ID_FIELD_BY_CALL_TYPE)) + async def test_owner_gets_the_provider_id_on_every_route(self, mock_cache, salt_key_env, call_type): + field = _ADDRESSED_ID_FIELD_BY_CALL_TYPE[call_type] + hook = _hook() + owner = _auth() + data = {field: _issue_managed_id(hook, owner)} + + result = await hook.async_pre_call_hook( + user_api_key_dict=owner, + cache=mock_cache, + data=data, + call_type=call_type, + ) + + assert result[field] == _FABRICATED_PROVIDER_RESPONSE_ID + + @pytest.mark.asyncio + @pytest.mark.parametrize("call_type", sorted(_ADDRESSED_ID_FIELD_BY_CALL_TYPE)) + @pytest.mark.parametrize( + ("caller_user_id", "caller_team_id", "detail"), + [ + ("stranger-user", "owner-team", "not associated with the user"), + ("owner-user", "stranger-team", "not associated with the team"), + ], + ) + async def test_other_user_or_team_is_rejected_on_every_route( + self, mock_cache, salt_key_env, call_type, caller_user_id, caller_team_id, detail + ): + field = _ADDRESSED_ID_FIELD_BY_CALL_TYPE[call_type] + hook = _hook() + issued_id = _issue_managed_id(hook, _auth()) + data = {field: issued_id} + + with pytest.raises(HTTPException) as exc_info: + await hook.async_pre_call_hook( + user_api_key_dict=_auth(user_id=caller_user_id, team_id=caller_team_id), + cache=mock_cache, + data=data, + call_type=call_type, + ) + + assert exc_info.value.status_code == 403 + assert detail in exc_info.value.detail + assert data[field] == issued_id + @pytest.mark.asyncio async def test_owner_can_still_address_the_id_the_proxy_issued_it(self, mock_cache, salt_key_env): hook = _hook() @@ -937,31 +984,37 @@ class TestUnmanagedResponseIdEscapeHatches: "general_settings", [{"allow_unmanaged_response_ids": True}, {"disable_responses_id_security": True}], ) - async def test_opted_in_settings_forward_the_id_untouched(self, mock_cache, salt_key_env, general_settings): - data = {"response_id": _FABRICATED_UNMANAGED_ID} + @pytest.mark.parametrize("call_type", sorted(_ADDRESSED_ID_FIELD_BY_CALL_TYPE)) + async def test_opted_in_settings_forward_the_id_untouched( + self, mock_cache, salt_key_env, general_settings, call_type + ): + field = _ADDRESSED_ID_FIELD_BY_CALL_TYPE[call_type] + data = {field: _FABRICATED_UNMANAGED_ID} result = await _hook(general_settings=general_settings).async_pre_call_hook( user_api_key_dict=_auth(), cache=mock_cache, data=data, - call_type="aget_responses", + call_type=call_type, ) - assert result["response_id"] == _FABRICATED_UNMANAGED_ID + assert result[field] == _FABRICATED_UNMANAGED_ID @pytest.mark.asyncio - async def test_proxy_without_a_signing_key_forwards_the_id_untouched(self, mock_cache, monkeypatch): + @pytest.mark.parametrize("call_type", sorted(_ADDRESSED_ID_FIELD_BY_CALL_TYPE)) + async def test_proxy_without_a_signing_key_forwards_the_id_untouched(self, mock_cache, monkeypatch, call_type): monkeypatch.delenv("LITELLM_SALT_KEY", raising=False) - data = {"response_id": _FABRICATED_UNMANAGED_ID} + field = _ADDRESSED_ID_FIELD_BY_CALL_TYPE[call_type] + data = {field: _FABRICATED_UNMANAGED_ID} result = await _hook(signing_key=None).async_pre_call_hook( user_api_key_dict=_auth(), cache=mock_cache, data=data, - call_type="aget_responses", + call_type=call_type, ) - assert result["response_id"] == _FABRICATED_UNMANAGED_ID + assert result[field] == _FABRICATED_UNMANAGED_ID @pytest.mark.asyncio async def test_proxy_admin_may_address_an_unmanaged_id(self, mock_cache, salt_key_env):