From 5728ddcc5cbae1798012231f64aabdee89f3309f Mon Sep 17 00:00:00 2001 From: Major Hayden Date: Tue, 7 Apr 2026 07:39:19 -0500 Subject: [PATCH 1/2] fix: add configuration guard to rlsapi_v1 /infer endpoint Call check_configuration_loaded() before processing requests so a missing config returns a clean HTTP 500 instead of an opaque crash. Matches the pattern used by every other endpoint. Ref: RSPEED-2817 Signed-off-by: Major Hayden --- src/app/endpoints/rlsapi_v1.py | 3 +++ tests/unit/app/endpoints/test_rlsapi_v1.py | 23 ++++++++++++++++++++++ 2 files changed, 26 insertions(+) diff --git a/src/app/endpoints/rlsapi_v1.py b/src/app/endpoints/rlsapi_v1.py index c78b00015..8c6bfd330 100644 --- a/src/app/endpoints/rlsapi_v1.py +++ b/src/app/endpoints/rlsapi_v1.py @@ -38,6 +38,7 @@ from models.rlsapi.requests import RlsapiV1InferRequest, RlsapiV1SystemInfo from models.rlsapi.responses import RlsapiV1InferData, RlsapiV1InferResponse from observability import InferenceEventData, build_inference_event, send_splunk_event +from utils.endpoints import check_configuration_loaded from utils.query import ( extract_provider_and_model_from_model_id, handle_known_apistatus_errors, @@ -595,6 +596,8 @@ async def infer_endpoint( # pylint: disable=R0914 # Authentication enforced by get_auth_dependency(), authorization by @authorize decorator. _ = auth + check_configuration_loaded(configuration) + request_id = get_suid() logger.info("Processing rlsapi v1 /infer request %s", request_id) diff --git a/tests/unit/app/endpoints/test_rlsapi_v1.py b/tests/unit/app/endpoints/test_rlsapi_v1.py index 208de1176..20bbbb31b 100644 --- a/tests/unit/app/endpoints/test_rlsapi_v1.py +++ b/tests/unit/app/endpoints/test_rlsapi_v1.py @@ -493,6 +493,29 @@ def test_get_rh_identity_context( # --- Test infer_endpoint --- +@pytest.mark.asyncio +async def test_infer_endpoint_configuration_not_loaded( + mocker: MockerFixture, + mock_auth_resolvers: None, + mock_request_factory: Callable[..., Any], + mock_background_tasks: Any, +) -> None: + """Test /infer returns HTTP 500 when configuration is not loaded.""" + mocker.patch.object(AppConfig(), "_configuration", None) + + infer_request = RlsapiV1InferRequest(question="How do I list files?") + mock_request = mock_request_factory() + + with pytest.raises(HTTPException) as exc_info: + await infer_endpoint( + infer_request=infer_request, + request=mock_request, + background_tasks=mock_background_tasks, + auth=MOCK_AUTH, + ) + assert exc_info.value.status_code == status.HTTP_500_INTERNAL_SERVER_ERROR + + @pytest.mark.asyncio async def test_infer_minimal_request( mocker: MockerFixture, From cb2db15aa31149c35bed8862046d17b25b7eb2c3 Mon Sep 17 00:00:00 2001 From: Major Hayden Date: Tue, 7 Apr 2026 08:17:52 -0500 Subject: [PATCH 2/2] feat: add rlsapi_v1 config section with quota enforcement Move allow_verbose_infer from the shared Customization config into a new dedicated rlsapi_v1 config section so CLA-specific settings are consolidated in one place and don't clutter shared configuration. Add configurable quota enforcement for /v1/infer via quota_subject, which selects the identity field (user_id, org_id, or system_id) used as the quota subject. Disabled by default (quota_subject: null). Startup validation rejects org_id/system_id quota_subject when the authentication module is not rh-identity, and warns when quota_subject is set but no quota limiters are configured. Falls back to user_id at runtime if a specific request lacks rh-identity data. No changes to the core quota system (utils/quota.py, quota limiters, factory, or scheduler). All existing endpoints are unaffected. Signed-off-by: Major Hayden --- examples/lightspeed-stack-rlsapi-cla.yaml | 11 + src/app/endpoints/rlsapi_v1.py | 71 ++++- src/configuration.py | 16 + src/models/config.py | 97 +++++- tests/unit/app/endpoints/test_rlsapi_v1.py | 300 ++++++++++++++++-- .../models/config/test_dump_configuration.py | 20 ++ .../config/test_rlsapi_v1_configuration.py | 172 ++++++++++ 7 files changed, 647 insertions(+), 40 deletions(-) create mode 100644 tests/unit/models/config/test_rlsapi_v1_configuration.py diff --git a/examples/lightspeed-stack-rlsapi-cla.yaml b/examples/lightspeed-stack-rlsapi-cla.yaml index 8dcedfa2b..fe00eebd5 100644 --- a/examples/lightspeed-stack-rlsapi-cla.yaml +++ b/examples/lightspeed-stack-rlsapi-cla.yaml @@ -20,6 +20,17 @@ inference: default_provider: google-vertex default_model: gemini-2.5-flash +# rlsapi v1 endpoint settings (CLA-specific) +rlsapi_v1: + # Quota enforcement: which identity field to track token usage by. + # Options: "user_id", "org_id", "system_id" + # Requires quota_handlers to be configured. Omit to disable quotas. + # quota_subject: "org_id" + + # Allow /v1/infer to return extended metadata when clients send + # "include_metadata": true. Should NOT be enabled in production. + # allow_verbose_infer: false + # Red Hat Identity authentication (typical for CLA deployments) authentication: module: "rh-identity" diff --git a/src/app/endpoints/rlsapi_v1.py b/src/app/endpoints/rlsapi_v1.py index 8c6bfd330..6d24962e9 100644 --- a/src/app/endpoints/rlsapi_v1.py +++ b/src/app/endpoints/rlsapi_v1.py @@ -40,9 +40,11 @@ from observability import InferenceEventData, build_inference_event, send_splunk_event from utils.endpoints import check_configuration_loaded from utils.query import ( + consume_query_tokens, extract_provider_and_model_from_model_id, handle_known_apistatus_errors, ) +from utils.quota import check_tokens_available from utils.responses import ( build_turn_summary, extract_text_from_response_items, @@ -456,12 +458,57 @@ def _is_verbose_enabled(infer_request: RlsapiV1InferRequest) -> bool: True if both server config and client request enable verbose mode. """ return ( - configuration.customization is not None - and configuration.customization.allow_verbose_infer - and infer_request.include_metadata + configuration.rlsapi_v1.allow_verbose_infer and infer_request.include_metadata ) +def _resolve_quota_subject(request: Request, auth: AuthTuple) -> str | None: + """Resolve the quota subject identifier based on rlsapi_v1 configuration. + + Returns None when quota enforcement is disabled (quota_subject not set), + signaling the caller to skip quota checks entirely. + + When the configured subject source (org_id or system_id) is unavailable + (e.g., rh-identity auth is not active), falls back to user_id from the + auth tuple so quota enforcement still applies. + + Args: + request: The FastAPI request object (for accessing rh-identity state). + auth: Authentication tuple from the configured auth provider. + + Returns: + The resolved subject identifier string, or None if quota is disabled. + """ + quota_subject = configuration.rlsapi_v1.quota_subject + if quota_subject is None: + return None + + user_id = auth[0] + + if quota_subject == "user_id": + return user_id + + org_id, system_id = _get_rh_identity_context(request) + + if quota_subject == "org_id": + if org_id == AUTH_DISABLED: + logger.warning( + "quota_subject is 'org_id' but rh-identity data is unavailable, " + "falling back to user_id" + ) + return user_id + return org_id + + # quota_subject == "system_id" + if system_id == AUTH_DISABLED: + logger.warning( + "quota_subject is 'system_id' but rh-identity data is unavailable, " + "falling back to user_id" + ) + return user_id + return system_id + + def _build_infer_response( response_text: str, request_id: str, @@ -594,10 +641,14 @@ async def infer_endpoint( # pylint: disable=R0914 HTTPException: 503 if the LLM service is unavailable. """ # Authentication enforced by get_auth_dependency(), authorization by @authorize decorator. - _ = auth - check_configuration_loaded(configuration) + # Quota enforcement: resolve subject and check availability before any work. + # No-op when quota_subject is not configured or no quota limiters exist. + quota_id = _resolve_quota_subject(request, auth) + if quota_id is not None: + check_tokens_available(configuration.quota_limiters, quota_id) + request_id = get_suid() logger.info("Processing rlsapi v1 /infer request %s", request_id) @@ -634,7 +685,7 @@ async def infer_endpoint( # pylint: disable=R0914 model_id=model_id, ) response_text = extract_text_from_response_items(response.output) - extract_token_usage(response.usage, model_id) + token_usage = extract_token_usage(response.usage, model_id) inference_time = time.monotonic() - start_time except _INFER_HANDLED_EXCEPTIONS as error: if response is not None: @@ -662,6 +713,14 @@ async def infer_endpoint( # pylint: disable=R0914 logger.warning("Empty response from LLM for request %s", request_id) response_text = constants.UNABLE_TO_PROCESS_RESPONSE + # Consume quota tokens after successful inference. + if quota_id is not None: + consume_query_tokens( + user_id=quota_id, + model_id=model_id, + token_usage=token_usage, + ) + _queue_splunk_event( background_tasks, infer_request, diff --git a/src/configuration.py b/src/configuration.py index e383d92a9..4eeca0460 100644 --- a/src/configuration.py +++ b/src/configuration.py @@ -27,6 +27,7 @@ OkpConfiguration, QuotaHandlersConfiguration, RagConfiguration, + RlsapiV1Configuration, ServiceConfiguration, SplunkConfiguration, UserDataCollection, @@ -286,6 +287,21 @@ def customization(self) -> Optional[Customization]: raise LogicError("logic error: configuration is not loaded") return self._configuration.customization + @property + def rlsapi_v1(self) -> RlsapiV1Configuration: + """Return rlsapi v1 endpoint configuration. + + Returns: + RlsapiV1Configuration: Configuration for the rlsapi v1 /infer + endpoint (CLA-specific settings). + + Raises: + LogicError: If the configuration has not been loaded. + """ + if self._configuration is None: + raise LogicError("logic error: configuration is not loaded") + return self._configuration.rlsapi_v1 + @property def inference(self) -> InferenceConfiguration: """Return inference configuration. diff --git a/src/models/config.py b/src/models/config.py index 906efe294..71c373361 100644 --- a/src/models/config.py +++ b/src/models/config.py @@ -1326,21 +1326,6 @@ class Customization(ConfigurationBase): agent_card_config: Optional[dict[str, Any]] = None custom_profile: Optional[CustomProfile] = Field(default=None, init=False) - # Debugging: Allow /v1/infer to return extended metadata - # WARNING: This should NOT be enabled in production environments. - # Setting this to True allows clients to request extended response data - # (tool_calls, rag_chunks, token_usage, etc.) from the /v1/infer endpoint - # by including "include_metadata": true in the request body. - # - # If this feature were wanted in production, consider RBAC-based access control instead: - # 1. Add Action.RLSAPI_V1_INFER_VERBOSE to models/config.py Action enum - # 2. Check authorization in infer_endpoint: - # if infer_request.include_metadata: - # if Action.RLSAPI_V1_INFER_VERBOSE not in request.state.authorized_actions: - # raise HTTPException(status_code=403, detail="Verbose infer not authorized") - # 3. Add the action to authorization rules for specific users/roles - allow_verbose_infer: bool = False - @model_validator(mode="after") def check_customization_model(self) -> Self: """ @@ -1381,6 +1366,35 @@ def check_customization_model(self) -> Self: return self +class RlsapiV1Configuration(ConfigurationBase): + """Configuration for the rlsapi v1 /infer endpoint. + + Settings specific to the RHEL Lightspeed Command Line Assistant (CLA) + stateless inference endpoint. Kept separate from shared configuration + sections so that CLA-specific options do not affect other endpoints. + """ + + allow_verbose_infer: bool = Field( + default=False, + title="Allow verbose infer", + description="Allow /v1/infer to return extended metadata " + "(tool_calls, rag_chunks, token_usage) when the client sends " + '"include_metadata": true. Should NOT be enabled in production. ' + "If production use is needed, consider RBAC-based access control " + "via an Action.RLSAPI_V1_INFER authorization rule.", + ) + + quota_subject: Optional[Literal["user_id", "org_id", "system_id"]] = Field( + default=None, + title="Quota subject", + description="Identity field used as the quota subject for /v1/infer. " + "When set, token quota enforcement is enabled for this endpoint. " + "Requires quota_handlers to be configured. " + '"org_id" and "system_id" require rh-identity authentication; ' + "falls back to user_id when rh-identity data is unavailable.", + ) + + class InferenceConfiguration(ConfigurationBase): """Inference configuration.""" @@ -1911,6 +1925,13 @@ class Configuration(ConfigurationBase): ) azure_entra_id: Optional[AzureEntraIdConfiguration] = None + rlsapi_v1: RlsapiV1Configuration = Field( + default_factory=RlsapiV1Configuration, + title="rlsapi v1 configuration", + description="Configuration for the rlsapi v1 /infer endpoint used by " + "the RHEL Lightspeed Command Line Assistant (CLA).", + ) + splunk: Optional[SplunkConfiguration] = Field( default=None, title="Splunk configuration", @@ -1993,6 +2014,52 @@ def validate_mcp_auth_headers(self) -> Self: self.mcp_servers = valid_mcp_servers return self + @model_validator(mode="after") + def validate_rlsapi_v1_quota_configuration(self) -> Self: + """Validate rlsapi_v1 quota settings against authentication and quota handlers. + + Enforces that quota_subject values requiring rh-identity data ("org_id", + "system_id") are only used when rh-identity authentication is active. + Warns when quota_subject is set but quota enforcement is not fully + configured (missing limiters or missing storage backend). + + Returns: + Self: The validated configuration instance. + + Raises: + ValueError: If quota_subject requires rh-identity but a different + authentication module is configured. + """ + quota_subject = self.rlsapi_v1.quota_subject # pylint: disable=no-member + if quota_subject is None: + return self + + auth_module = self.authentication.module # pylint: disable=no-member + + if quota_subject in ("org_id", "system_id") and ( + auth_module != constants.AUTH_MOD_RH_IDENTITY + ): + raise ValueError( + f"rlsapi_v1.quota_subject='{quota_subject}' requires " + f"authentication.module='{constants.AUTH_MOD_RH_IDENTITY}', " + f"but got '{auth_module}'. Use quota_subject='user_id' or " + f"switch authentication to '{constants.AUTH_MOD_RH_IDENTITY}'." + ) + + if not self.quota_handlers.limiters or ( # pylint: disable=no-member + self.quota_handlers.sqlite is None # pylint: disable=no-member + and self.quota_handlers.postgres is None # pylint: disable=no-member + ): + logger.warning( + "rlsapi_v1.quota_subject is '%s' but quota enforcement is not " + "fully configured. Token quota enforcement will not take effect " + "until at least one limiter and one quota storage backend " + "are configured.", + quota_subject, + ) + + return self + def dump(self, filename: str | Path = "configuration.json") -> None: """ Write the current Configuration model to a JSON file. diff --git a/tests/unit/app/endpoints/test_rlsapi_v1.py b/tests/unit/app/endpoints/test_rlsapi_v1.py index 20bbbb31b..dc2f284ce 100644 --- a/tests/unit/app/endpoints/test_rlsapi_v1.py +++ b/tests/unit/app/endpoints/test_rlsapi_v1.py @@ -25,6 +25,7 @@ _compile_prompt_template, _get_default_model_id, _get_rh_identity_context, + _resolve_quota_subject, infer_endpoint, retrieve_simple_response, ) @@ -60,8 +61,13 @@ def mock_custom_prompt_fixture(mocker: MockerFixture) -> Callable[[str], None]: def _set(prompt: str) -> None: mock_customization = mocker.Mock() mock_customization.system_prompt = prompt + mock_rlsapi_v1 = mocker.Mock() + mock_rlsapi_v1.allow_verbose_infer = False + mock_rlsapi_v1.quota_subject = None mock_config = mocker.Mock() mock_config.customization = mock_customization + mock_config.rlsapi_v1 = mock_rlsapi_v1 + mock_config.quota_limiters = [] mocker.patch("app.endpoints.rlsapi_v1.configuration", mock_config) return _set @@ -709,12 +715,14 @@ async def test_infer_include_metadata_respects_verbose_config( expect_metadata: bool, ) -> None: """Test /infer metadata inclusion controlled by dual opt-in (config + request).""" - custom_mock = mocker.Mock() - custom_mock.allow_verbose_infer = verbose_enabled - custom_mock.system_prompt = "You are a helpful assistant." + rlsapi_v1_mock = mocker.Mock() + rlsapi_v1_mock.allow_verbose_infer = verbose_enabled + rlsapi_v1_mock.quota_subject = None config_mock = mocker.Mock() config_mock.inference = mock_configuration.inference - config_mock.customization = custom_mock + config_mock.customization = mock_configuration.customization + config_mock.rlsapi_v1 = rlsapi_v1_mock + config_mock.quota_limiters = [] mocker.patch("app.endpoints.rlsapi_v1.configuration", config_mock) mock_response = mocker.Mock() @@ -763,12 +771,14 @@ def _setup_config_mock( verbose_enabled: bool, ) -> None: """Helper to set up configuration mock with verbose setting.""" - custom_mock = mocker.Mock() - custom_mock.allow_verbose_infer = verbose_enabled - custom_mock.system_prompt = "You are a helpful assistant." + rlsapi_v1_mock = mocker.Mock() + rlsapi_v1_mock.allow_verbose_infer = verbose_enabled + rlsapi_v1_mock.quota_subject = None config_mock = mocker.Mock() config_mock.inference = mock_configuration.inference - config_mock.customization = custom_mock + config_mock.customization = mock_configuration.customization + config_mock.rlsapi_v1 = rlsapi_v1_mock + config_mock.quota_limiters = [] mocker.patch("app.endpoints.rlsapi_v1.configuration", config_mock) @@ -862,30 +872,282 @@ async def test_infer_queues_splunk_event_on_success( assert call_args[0][2] == "infer_with_llm" +# --- Test _resolve_quota_subject --- + + +@pytest.mark.parametrize( + ("quota_subject", "rh_identity_setup", "expected"), + [ + pytest.param(None, None, None, id="disabled_no_identity"), + pytest.param( + None, + {"org_id": "org1", "user_id": "sys1"}, + None, + id="disabled_with_identity", + ), + pytest.param("user_id", None, "mock_user_id", id="user_id_no_identity"), + pytest.param( + "org_id", + {"org_id": "org123", "user_id": "sys456"}, + "org123", + id="org_id_with_identity", + ), + pytest.param( + "system_id", + {"org_id": "org123", "user_id": "sys456"}, + "sys456", + id="system_id_with_identity", + ), + pytest.param( + "org_id", + None, + "mock_user_id", + id="org_id_fallback_no_identity", + ), + pytest.param( + "system_id", + None, + "mock_user_id", + id="system_id_fallback_no_identity", + ), + pytest.param( + "org_id", + {"org_id": "", "user_id": "sys1"}, + "mock_user_id", + id="org_id_fallback_empty_org", + ), + pytest.param( + "system_id", + {"org_id": "org1", "user_id": ""}, + "mock_user_id", + id="system_id_fallback_empty_system", + ), + ], +) +def test_resolve_quota_subject( + mocker: MockerFixture, + mock_request_factory: Callable[..., Any], + quota_subject: str | None, + rh_identity_setup: dict[str, str] | None, + expected: str | None, +) -> None: + """Test _resolve_quota_subject resolves correct ID based on config and identity.""" + rlsapi_v1_mock = mocker.Mock() + rlsapi_v1_mock.quota_subject = quota_subject + config_mock = mocker.Mock() + config_mock.rlsapi_v1 = rlsapi_v1_mock + mocker.patch("app.endpoints.rlsapi_v1.configuration", config_mock) + + if rh_identity_setup is not None: + mock_rh_identity = mocker.Mock(spec=RHIdentityData) + mock_rh_identity.get_org_id.return_value = rh_identity_setup["org_id"] + mock_rh_identity.get_user_id.return_value = rh_identity_setup["user_id"] + mock_request = mock_request_factory(rh_identity=mock_rh_identity) + else: + mock_request = mock_request_factory() + + result = _resolve_quota_subject(mock_request, MOCK_AUTH) + assert result == expected + + +# --- Test quota enforcement in infer_endpoint --- + + +@pytest.fixture(name="mock_quota_config") +def mock_quota_config_fixture( + mocker: MockerFixture, + mock_configuration: AppConfig, +) -> Callable[[str], None]: + """Factory fixture that patches configuration with quota_subject enabled. + + Args: + mocker: The pytest mocker fixture. + mock_configuration: Base AppConfig to extend. + + Returns: + Callable that accepts a quota_subject value and patches configuration. + """ + + def _set(quota_subject: str) -> None: + rlsapi_v1_mock = mocker.Mock() + rlsapi_v1_mock.quota_subject = quota_subject + rlsapi_v1_mock.allow_verbose_infer = False + config_mock = mocker.Mock() + config_mock.inference = mock_configuration.inference + config_mock.customization = mock_configuration.customization + config_mock.rlsapi_v1 = rlsapi_v1_mock + config_mock.quota_limiters = [] + mocker.patch("app.endpoints.rlsapi_v1.configuration", config_mock) + + return _set + + @pytest.mark.asyncio -async def test_infer_queues_splunk_error_event_on_failure( +async def test_infer_quota_check_called_when_configured( + mocker: MockerFixture, + mock_quota_config: Callable[[str], None], + mock_llm_response: None, + mock_auth_resolvers: None, + mock_request_factory: Callable[..., Any], + mock_background_tasks: Any, +) -> None: + """Test /infer calls check_tokens_available when quota_subject is set.""" + mock_quota_config("user_id") + mock_check = mocker.patch("app.endpoints.rlsapi_v1.check_tokens_available") + mock_consume = mocker.patch("app.endpoints.rlsapi_v1.consume_query_tokens") + + response = await infer_endpoint( + infer_request=RlsapiV1InferRequest(question="How do I list files?"), + request=mock_request_factory(), + background_tasks=mock_background_tasks, + auth=MOCK_AUTH, + ) + + assert isinstance(response, RlsapiV1InferResponse) + mock_check.assert_called_once_with([], "mock_user_id") + mock_consume.assert_called_once() + + +@pytest.mark.asyncio +async def test_infer_quota_skipped_when_not_configured( mocker: MockerFixture, mock_configuration: AppConfig, - mock_api_connection_error: None, + mock_llm_response: None, mock_auth_resolvers: None, mock_request_factory: Callable[..., Any], mock_background_tasks: Any, ) -> None: - """Test that failed inference queues a Splunk error event.""" - infer_request = RlsapiV1InferRequest(question="Test question") - mock_request = mock_request_factory() + """Test /infer skips quota calls when quota_subject is None (default).""" + mock_check = mocker.patch("app.endpoints.rlsapi_v1.check_tokens_available") + mock_consume = mocker.patch("app.endpoints.rlsapi_v1.consume_query_tokens") + + await infer_endpoint( + infer_request=RlsapiV1InferRequest(question="How do I list files?"), + request=mock_request_factory(), + background_tasks=mock_background_tasks, + auth=MOCK_AUTH, + ) + + mock_check.assert_not_called() + mock_consume.assert_not_called() - with pytest.raises(HTTPException): + +@pytest.mark.asyncio +async def test_infer_quota_exceeded_returns_429( + mocker: MockerFixture, + mock_quota_config: Callable[[str], None], + mock_llm_response: None, + mock_auth_resolvers: None, + mock_request_factory: Callable[..., Any], + mock_background_tasks: Any, +) -> None: + """Test /infer returns HTTP 429 when quota is exceeded.""" + mock_quota_config("user_id") + mocker.patch( + "app.endpoints.rlsapi_v1.check_tokens_available", + side_effect=HTTPException(status_code=status.HTTP_429_TOO_MANY_REQUESTS), + ) + + with pytest.raises(HTTPException) as exc_info: await infer_endpoint( - infer_request=infer_request, - request=mock_request, + infer_request=RlsapiV1InferRequest(question="How do I list files?"), + request=mock_request_factory(), background_tasks=mock_background_tasks, auth=MOCK_AUTH, ) - mock_background_tasks.add_task.assert_called_once() - call_args = mock_background_tasks.add_task.call_args - assert call_args[0][2] == "infer_error" + assert exc_info.value.status_code == status.HTTP_429_TOO_MANY_REQUESTS + + +@pytest.mark.parametrize( + ("quota_subject", "rh_identity_setup", "expected_subject"), + [ + pytest.param( + "org_id", + {"org_id": "org123", "user_id": "sys456"}, + "org123", + id="org_id", + ), + pytest.param( + "system_id", + {"org_id": "org123", "user_id": "sys456"}, + "sys456", + id="system_id", + ), + ], +) +@pytest.mark.asyncio +async def test_infer_quota_with_rh_identity_subject( + mocker: MockerFixture, + mock_quota_config: Callable[[str], None], + mock_llm_response: None, + mock_auth_resolvers: None, + mock_request_factory: Callable[..., Any], + mock_background_tasks: Any, + quota_subject: str, + rh_identity_setup: dict[str, str], + expected_subject: str, +) -> None: + """Test /infer propagates org_id/system_id to quota check and consumption.""" + mock_quota_config(quota_subject) + + mock_rh_identity = mocker.Mock(spec=RHIdentityData) + mock_rh_identity.get_org_id.return_value = rh_identity_setup["org_id"] + mock_rh_identity.get_user_id.return_value = rh_identity_setup["user_id"] + + mock_check = mocker.patch("app.endpoints.rlsapi_v1.check_tokens_available") + mock_consume = mocker.patch("app.endpoints.rlsapi_v1.consume_query_tokens") + + await infer_endpoint( + infer_request=RlsapiV1InferRequest(question="How do I list files?"), + request=mock_request_factory(rh_identity=mock_rh_identity), + background_tasks=mock_background_tasks, + auth=MOCK_AUTH, + ) + + mock_check.assert_called_once_with([], expected_subject) + mock_consume.assert_called_once() + assert mock_consume.call_args.kwargs["user_id"] == expected_subject + + +@pytest.mark.asyncio +async def test_infer_quota_shield_blocked_does_not_consume_tokens( + mocker: MockerFixture, + mock_quota_config: Callable[[str], None], + mock_llm_response: None, + mock_auth_resolvers: None, + mock_request_factory: Callable[..., Any], + mock_background_tasks: Any, +) -> None: + """Test quota pre-check runs but tokens are NOT consumed when shield blocks.""" + mock_quota_config("user_id") + + blocked = ShieldModerationBlocked( + message="Blocked by moderation", + moderation_id="modr-test", + refusal_response=OpenAIResponseMessage( + role="assistant", + content="Blocked by moderation", + ), + ) + mocker.patch( + "app.endpoints.rlsapi_v1.run_shield_moderation", + new=mocker.AsyncMock(return_value=blocked), + ) + + mock_check = mocker.patch("app.endpoints.rlsapi_v1.check_tokens_available") + mock_consume = mocker.patch("app.endpoints.rlsapi_v1.consume_query_tokens") + + response = await infer_endpoint( + infer_request=RlsapiV1InferRequest(question="Bad question"), + request=mock_request_factory(), + background_tasks=mock_background_tasks, + auth=MOCK_AUTH, + ) + + assert response.data.text == "Blocked by moderation" + mock_check.assert_called_once_with([], "mock_user_id") + mock_consume.assert_not_called() # --- Test shield moderation --- diff --git a/tests/unit/models/config/test_dump_configuration.py b/tests/unit/models/config/test_dump_configuration.py index 3b531ab65..06a3ef08c 100644 --- a/tests/unit/models/config/test_dump_configuration.py +++ b/tests/unit/models/config/test_dump_configuration.py @@ -215,6 +215,10 @@ def test_dump_configuration(tmp_path: Path) -> None: "offline": True, "chunk_filter_query": None, }, + "rlsapi_v1": { + "allow_verbose_infer": False, + "quota_subject": None, + }, "splunk": None, "deployment_environment": "development", } @@ -568,6 +572,10 @@ def test_dump_configuration_with_quota_limiters(tmp_path: Path) -> None: "offline": True, "chunk_filter_query": None, }, + "rlsapi_v1": { + "allow_verbose_infer": False, + "quota_subject": None, + }, "splunk": None, "deployment_environment": "development", } @@ -798,6 +806,10 @@ def test_dump_configuration_with_quota_limiters_different_values( "offline": True, "chunk_filter_query": None, }, + "rlsapi_v1": { + "allow_verbose_infer": False, + "quota_subject": None, + }, "splunk": None, "deployment_environment": "development", } @@ -1003,6 +1015,10 @@ def test_dump_configuration_byok(tmp_path: Path) -> None: "offline": True, "chunk_filter_query": None, }, + "rlsapi_v1": { + "allow_verbose_infer": False, + "quota_subject": None, + }, "splunk": None, "deployment_environment": "development", } @@ -1193,6 +1209,10 @@ def test_dump_configuration_pg_namespace(tmp_path: Path) -> None: "offline": True, "chunk_filter_query": None, }, + "rlsapi_v1": { + "allow_verbose_infer": False, + "quota_subject": None, + }, "splunk": None, "deployment_environment": "development", } diff --git a/tests/unit/models/config/test_rlsapi_v1_configuration.py b/tests/unit/models/config/test_rlsapi_v1_configuration.py new file mode 100644 index 000000000..202fc4199 --- /dev/null +++ b/tests/unit/models/config/test_rlsapi_v1_configuration.py @@ -0,0 +1,172 @@ +"""Unit tests for RlsapiV1Configuration and related startup validators.""" + +import logging +from typing import Any + +import pytest +from pydantic import ValidationError + +from models.config import Configuration, RlsapiV1Configuration + +# --- Test RlsapiV1Configuration --- + + +def test_defaults() -> None: + """Test RlsapiV1Configuration defaults to disabled state.""" + config = RlsapiV1Configuration() + assert config.allow_verbose_infer is False + assert config.quota_subject is None + + +@pytest.mark.parametrize( + "quota_subject", + [ + pytest.param("user_id", id="user_id"), + pytest.param("org_id", id="org_id"), + pytest.param("system_id", id="system_id"), + ], +) +def test_valid_quota_subject_values(quota_subject: str) -> None: + """Test all valid quota_subject literal values are accepted.""" + config = RlsapiV1Configuration(quota_subject=quota_subject) + assert config.quota_subject == quota_subject + + +def test_invalid_quota_subject_rejected() -> None: + """Test invalid quota_subject value is rejected by Literal validation.""" + with pytest.raises(ValidationError, match="quota_subject"): + RlsapiV1Configuration( + quota_subject="invalid_value" # pyright: ignore[reportArgumentType] + ) + + +def test_rejects_unknown_fields() -> None: + """Test RlsapiV1Configuration rejects unknown fields (extra=forbid).""" + with pytest.raises(ValidationError, match="Extra inputs are not permitted"): + RlsapiV1Configuration( + unknown_field="should_fail" # pyright: ignore[reportCallIssue] + ) + + +# --- Test Configuration-level startup validators --- + + +def _build_config_dict(**overrides: Any) -> dict[str, Any]: + """Build a minimal Configuration dict with optional overrides. + + Args: + **overrides: Keys to override in the base config dict. + + Returns: + A dict suitable for Configuration(**dict). + """ + base: dict[str, Any] = { + "name": "test", + "service": {"host": "localhost", "port": 8080}, + "llama_stack": { + "api_key": "test-key", + "url": "http://test.com:1234", + "use_as_library_client": False, + }, + "user_data_collection": {}, + "authentication": {"module": "noop"}, + "authorization": {"access_rules": []}, + } + base.update(overrides) + return base + + +@pytest.mark.parametrize( + ("quota_subject", "auth_module"), + [ + pytest.param("org_id", "noop", id="org_id_noop"), + pytest.param("org_id", "k8s", id="org_id_k8s"), + pytest.param("system_id", "noop", id="system_id_noop"), + pytest.param("system_id", "k8s", id="system_id_k8s"), + ], +) +def test_identity_quota_subject_requires_rh_identity( + quota_subject: str, auth_module: str +) -> None: + """Test startup validation rejects org_id/system_id without rh-identity auth.""" + config_dict = _build_config_dict( + rlsapi_v1={"quota_subject": quota_subject}, + authentication={"module": auth_module}, + ) + with pytest.raises(ValidationError, match="rh-identity"): + Configuration(**config_dict) + + +def test_quota_subject_user_id_works_with_any_auth() -> None: + """Test user_id quota_subject is accepted with any authentication module.""" + config_dict = _build_config_dict( + rlsapi_v1={"quota_subject": "user_id"}, + authentication={"module": "noop"}, + ) + config = Configuration(**config_dict) + assert config.rlsapi_v1.quota_subject == "user_id" # pylint: disable=no-member + + +def test_quota_subject_org_id_accepted_with_rh_identity() -> None: + """Test org_id quota_subject is accepted when rh-identity auth is configured.""" + config_dict = _build_config_dict( + rlsapi_v1={"quota_subject": "org_id"}, + authentication={"module": "rh-identity", "rh_identity_config": {}}, + ) + config = Configuration(**config_dict) + assert config.rlsapi_v1.quota_subject == "org_id" # pylint: disable=no-member + + +def test_quota_subject_none_by_default() -> None: + """Test quota_subject defaults to None when not configured.""" + config_dict = _build_config_dict() + config = Configuration(**config_dict) + assert config.rlsapi_v1.quota_subject is None # pylint: disable=no-member + + +def test_quota_subject_warns_when_no_limiters(caplog: pytest.LogCaptureFixture) -> None: + """Test startup validation warns when quota_subject is set but no limiters exist.""" + config_dict = _build_config_dict( + rlsapi_v1={"quota_subject": "user_id"}, + authentication={"module": "noop"}, + quota_handlers={}, + ) + config_logger = logging.getLogger("models.config") + config_logger.propagate = True + try: + with caplog.at_level(logging.WARNING): + Configuration(**config_dict) + + assert "quota enforcement is not fully configured" in caplog.text + finally: + config_logger.propagate = False + + +def test_quota_subject_warns_when_no_storage_backend( + caplog: pytest.LogCaptureFixture, +) -> None: + """Test startup validation warns when limiters exist but no storage backend.""" + config_dict = _build_config_dict( + rlsapi_v1={"quota_subject": "user_id"}, + authentication={"module": "noop"}, + quota_handlers={ + "limiters": [ + { + "name": "test", + "type": "user_limiter", + "initial_quota": 1000, + "quota_increase": 0, + "period": "1 month", + } + ], + }, + ) + config_logger = logging.getLogger("models.config") + config_logger.propagate = True + try: + with caplog.at_level(logging.WARNING): + Configuration(**config_dict) + + assert "quota enforcement is not fully configured" in caplog.text + finally: + config_logger.propagate = False