Let installed extensions check chat requests

src/common/extensions.py

@@ -1,17 +1,19 @@
 """Extensions installed alongside InferiaLLM, found through the ``inferia.extensions`` entry point.
 
 Each entry point names a callable that returns an extension object. The gateway
-calls its ``mount(app)`` at startup and the migrate step its ``migrate(dsn)``; an
-extension may provide either. With nothing installed nothing changes, and an
-extension that fails is logged and skipped, never raised: InferiaLLM must start
-without it.
+calls its ``mount(app)`` at startup, the migrate step its ``migrate(dsn)``, and the
+inference app its ``before_request(...)`` on each chat request; an extension may
+provide any of them. With nothing installed nothing changes, and an extension
+that fails is logged and skipped, never raised: InferiaLLM must work without it.
 """
 from __future__ import annotations
 
 import logging
 from functools import lru_cache
 from importlib.metadata import entry_points
 
+from starlette.exceptions import HTTPException
+
 log = logging.getLogger(__name__)
 
 GROUP = "inferia.extensions"
@@ -46,3 +48,23 @@ async def migrate_all(dsn: str) -> None:
             await ext.migrate(dsn)
         except Exception:
             log.exception("extension %r failed to migrate; skipped", ext)
+
+
+async def before_request(**request) -> None:
+    """Let extensions check a request before the engine is called.
+
+    ``request`` holds request_id, user_id, org_id, deployment_id, model,
+    request_type and body (the JSON body, never file uploads). An extension
+    refuses by raising HTTPException, which reaches the client like core's own
+    refusals; any other error is logged and the request goes on.
+    """
+    for ext in installed():
+        hook = getattr(ext, "before_request", None)
+        if hook is None:
+            continue
+        try:
+            await hook(**request)
+        except HTTPException:
+            raise
+        except Exception:
+            log.exception("extension %r failed in before_request; skipped", ext)

src/inference/core/handlers/completion.py

@@ -11,6 +11,8 @@
 from fastapi import BackgroundTasks, HTTPException
 from fastapi.responses import StreamingResponse
 
+from common import extensions
+from common.http_client import request_id_ctx
 from inference.client import api_gateway_client
 from inference.config import settings
 from ..providers import resolve_upstream
@@ -79,6 +81,18 @@ async def handle(
         applied_policies.append("quota")
         await api_gateway_client.check_quota(user_context_id, model)
 
+        # 3b. Installed extensions; one may refuse with an HTTPException.
+        # /v1/messages comes through here too, already in this format.
+        await extensions.before_request(
+            request_id=request_id_ctx.get() or "",
+            user_id=user_context_id,
+            org_id=org_id,
+            deployment_id=str(deployment_id) if deployment_id else None,
+            model=model,
+            request_type="llm",
+            body=body,
+        )
+
         # 4. Prepare Provider Request
         endpoint_url = deployment.get("endpoint")
         if not endpoint_url:

src/tests/common/test_extensions.py

@@ -93,3 +93,38 @@ async def ext(dsn):
     monkeypatch.setattr(init, "migrate_all", ext)
     await init._apply_migrations("postgresql://x")
     assert order == [("core", "postgresql://x"), ("extensions", "postgresql://x")]
+
+
+class _Checker:
+    def __init__(self, error=None):
+        self.error, self.seen = error, []
+
+    async def before_request(self, **request):
+        self.seen.append(request)
+        if self.error:
+            raise self.error
+
+
+async def test_before_request_reaches_every_extension_that_has_it(install):
+    first, second = _Checker(), _Checker()
+    install(_EP("first", first), _EP("no_hook", _Ext()), _EP("second", second))
+    await extensions.before_request(request_id="r1", model="m", body={"messages": []})
+    assert first.seen == second.seen == [{"request_id": "r1", "model": "m", "body": {"messages": []}}]
+
+
+async def test_a_refusal_reaches_the_caller(install):
+    from fastapi import HTTPException
+
+    later = _Checker()
+    install(_EP("refuses", _Checker(HTTPException(status_code=400, detail="no"))), _EP("later", later))
+    with pytest.raises(HTTPException) as e:
+        await extensions.before_request(request_id="r1")
+    assert e.value.status_code == 400
+    assert later.seen == []  # refused requests stop at the first refusal
+
+
+async def test_a_broken_check_is_skipped_and_the_request_goes_on(install):
+    good = _Checker()
+    install(_EP("broken", _Checker(RuntimeError("bug"))), _EP("good", good))
+    await extensions.before_request(request_id="r1")
+    assert good.seen == [{"request_id": "r1"}]

src/tests/inference/test_extension_request_check.py

@@ -0,0 +1,77 @@
+"""Chat requests go past installed extensions after the quota check, before the engine."""
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+from fastapi import HTTPException
+
+from common.http_client import request_id_ctx
+from inference.core.handlers import completion as ch
+
+pytestmark = pytest.mark.asyncio
+
+BODY = {"model": "qwen", "messages": [{"role": "user", "content": "hi"}]}
+
+
+@pytest.fixture
+def handler(monkeypatch):
+    order = []
+    context = {
+        "deployment": {"id": "dep-1", "endpoint": "http://engine.test", "engine": "vllm"},
+        "user_id_context": "apikey:k1",
+        "org_id": "org-1",
+        "rate_limit_config": None,
+        "log_payloads": True,
+    }
+    monkeypatch.setattr(ch.GatewayService, "resolve_context", AsyncMock(return_value=context))
+
+    async def check_quota(user_id, model):
+        order.append("quota")
+
+    monkeypatch.setattr(ch.api_gateway_client, "check_quota", check_quota)
+
+    def resolve_upstream(*_a, **_k):
+        order.append("engine")
+        raise RuntimeError("stop before the engine call")
+
+    monkeypatch.setattr(ch, "resolve_upstream", resolve_upstream)
+    seen = []
+
+    async def before_request(**request):
+        order.append("extensions")
+        seen.append(request)
+
+    monkeypatch.setattr(ch.extensions, "before_request", before_request)
+    return order, seen, monkeypatch
+
+
+async def _handle():
+    token = request_id_ctx.set("trace-1")
+    try:
+        return await ch.CompletionHandler.handle("sk-key", dict(BODY), MagicMock())
+    finally:
+        request_id_ctx.reset(token)
+
+
+async def test_extensions_see_the_request_after_the_quota_and_before_the_engine(handler):
+    order, seen, _ = handler
+    with pytest.raises(RuntimeError, match="stop before the engine"):
+        await _handle()
+    assert order == ["quota", "extensions", "engine"]
+    assert seen == [{
+        "request_id": "trace-1", "user_id": "apikey:k1", "org_id": "org-1",
+        "deployment_id": "dep-1", "model": "qwen", "request_type": "llm", "body": BODY,
+    }]
+
+
+async def test_a_refusal_stops_the_request_before_the_engine(handler):
+    order, _, monkeypatch = handler
+
+    async def refuse(**request):
+        order.append("extensions")
+        raise HTTPException(status_code=400, detail="refused")
+
+    monkeypatch.setattr(ch.extensions, "before_request", refuse)
+    with pytest.raises(HTTPException) as e:
+        await _handle()
+    assert e.value.status_code == 400
+    assert order == ["quota", "extensions"]