diff --git a/aqueduct/gateway/decorators/__init__.py b/aqueduct/gateway/decorators/__init__.py index 66ddac62..ad86b430 100644 --- a/aqueduct/gateway/decorators/__init__.py +++ b/aqueduct/gateway/decorators/__init__.py @@ -10,6 +10,7 @@ mcp_transport_security, parse_jsonrpc_message, ) +from gateway.decorators.plugins import run_plugins from gateway.decorators.response_type import normalize_reasoning_fields, process_file_content from gateway.decorators.responses import check_tool_availability, validate_response_id @@ -30,6 +31,7 @@ "process_file_content", "require_files_api_client", "resolve_alias", + "run_plugins", "token_authenticated", "tos_accepted", "validate_response_id", diff --git a/aqueduct/gateway/decorators/auth.py b/aqueduct/gateway/decorators/auth.py index bca21e04..4d6ad7ca 100644 --- a/aqueduct/gateway/decorators/auth.py +++ b/aqueduct/gateway/decorators/auth.py @@ -14,6 +14,7 @@ from gateway.decorators.types import AsyncView, Decorator, ViewResult from gateway.response_error import error_response from management.models import Token +from management.plugins import resolve_active_plugins log = logging.getLogger("aqueduct") @@ -75,6 +76,7 @@ async def wrapper(request: ASGIRequest, *args: Any, **kwargs: Any) -> ViewResult log.error("Token not found during authentication") return unauthorized_response kwargs["token"] = token + request.active_plugins = await sync_to_async(resolve_active_plugins)() # type: ignore[attr-defined] return await view_func(request, *args, **kwargs) return wrapper diff --git a/aqueduct/gateway/decorators/plugins.py b/aqueduct/gateway/decorators/plugins.py new file mode 100644 index 00000000..92a90a86 --- /dev/null +++ b/aqueduct/gateway/decorators/plugins.py @@ -0,0 +1,33 @@ +from functools import wraps +from typing import Any + +from django.core.handlers.asgi import ASGIRequest + +from gateway.decorators.types import AsyncView, ViewResult +from gateway.response_error import error_response +from management.plugins import BlockedByPlugin, PluginSnippet, after_hook, before_hook + + +def run_plugins(view_func: AsyncView) -> AsyncView: + @wraps(view_func) + async def wrapper(request: ASGIRequest, *args: Any, **kwargs: Any) -> ViewResult: + plugins: list[PluginSnippet] = getattr(request, "active_plugins", []) + if not plugins: + return await view_func(request, *args, **kwargs) + + token = kwargs.get("token") + body = kwargs.get("pydantic_model") + + try: + final_body = before_hook(plugins, request, token, body) + except BlockedByPlugin as e: + return error_response(e.reason, status=e.status) + + if final_body is not None: + kwargs["pydantic_model"] = final_body + + response = await view_func(request, *args, **kwargs) + after_hook(plugins, request, token, response) + return response + + return wrapper diff --git a/aqueduct/gateway/tests/test_plugins_api.py b/aqueduct/gateway/tests/test_plugins_api.py new file mode 100644 index 00000000..e160262c --- /dev/null +++ b/aqueduct/gateway/tests/test_plugins_api.py @@ -0,0 +1,175 @@ +from typing import ClassVar + +from django.test import override_settings + +from gateway.tests.test_endpoints import ChatCompletionsBase +from management.models import Snippet, SnippetType +from management.plugins import _plugin_class +from mock_api.mock_configs import MockConfig + +AFTER_CALLS: list[bool] = [] +EVENTS: list[str] = [] +BODY_EXTRA_KEY = "custom_plugin_field" + + +BODY_MUTATING_PLUGIN = """\ +from gateway.tests.test_plugins_api import BODY_EXTRA_KEY +class Mutator(PluginSnippet): + def before_request(self, request, token, body): + body[BODY_EXTRA_KEY] = "plugin-added" + return body +""" + + +BLOCKING_PLUGIN = """\ +class Guard(PluginSnippet): + def before_request(self, request, token, body): + raise BlockedByPlugin("forbidden by guard", status=403) +""" + + +AFTER_PLUGIN = """\ +from gateway.tests.test_plugins_api import AFTER_CALLS +class Audit(PluginSnippet): + def after_response(self, request, token, response): + AFTER_CALLS.append(True) +""" + + +PLUGIN_RECORD_A = """\ +from gateway.tests.test_plugins_api import EVENTS +class A(PluginSnippet): + def before_request(self, request, token, body): + EVENTS.append("A") + return body +""" + + +PLUGIN_RECORD_B = """\ +from gateway.tests.test_plugins_api import EVENTS +class B(PluginSnippet): + def before_request(self, request, token, body): + EVENTS.append("B") + return body +""" + + +PLUGIN_SET_STEP = """\ +from gateway.tests.test_plugins_api import EVENTS +class SetStep(PluginSnippet): + def before_request(self, request, token, body): + body["step"] = "from-pipe" + return body +""" + + +PLUGIN_READ_STEP = """\ +from gateway.tests.test_plugins_api import EVENTS +class ReadStep(PluginSnippet): + def before_request(self, request, token, body): + EVENTS.append(body.get("step")) + return body +""" + + +PLUGIN_BOOM = """\ +class Boom(PluginSnippet): + def before_request(self, request, token, body): + raise RuntimeError("plugin exploded") +""" + + +def seed_plugin(name: str, code: str, order: int = 0) -> Snippet: + return Snippet.objects.create( + name=name, type=SnippetType.PLUGIN, active=True, order=order, code=code + ) + + +@override_settings(TIKA_SERVER_URL=None) +class PluginGatewayIntegrationTest(ChatCompletionsBase): + MESSAGES: ClassVar[list[dict[str, str]]] = [{"role": "user", "content": "hello"}] + + def setUp(self): + super().setUp() + _plugin_class.cache_clear() + AFTER_CALLS.clear() + EVENTS.clear() + + def tearDown(self): + _plugin_class.cache_clear() + super().tearDown() + + def test_no_plugins_keeps_response_unchanged(self): + resp = self._send_chat_completion(self.MESSAGES) + self.assertEqual(resp.status_code, 200) + self.assertFalse(AFTER_CALLS) + + def test_before_hook_transforms_body(self): + seed_plugin("mutate", BODY_MUTATING_PLUGIN) + resp = self._send_chat_completion(self.MESSAGES) + self.assertEqual(resp.status_code, 200) + + def test_before_hook_can_block(self): + seed_plugin("guard", BLOCKING_PLUGIN) + resp = self._send_chat_completion(self.MESSAGES) + self.assertEqual(resp.status_code, 403) + self.assertEqual(resp.json()["error"]["message"], "forbidden by guard") + + def test_after_hook_observes_response(self): + seed_plugin("audit", AFTER_PLUGIN) + resp = self._send_chat_completion(self.MESSAGES) + self.assertEqual(resp.status_code, 200) + self.assertEqual(AFTER_CALLS, [True]) + + def test_router_error_still_converted_by_catch(self): + seed_plugin("audit", AFTER_PLUGIN) + bad = MockConfig( + status_code=400, + response_data={"error": {"message": "upstream boom", "type": "invalid_request_error"}}, + ) + with self.mock_server.patch_external_api("chat/completions", bad): + resp = self._send_chat_completion(self.MESSAGES) + self.assertEqual(resp.status_code, 400) + self.assertIn("upstream boom", resp.json()["error"]["message"]) + + def test_plugins_execute_in_order(self): + seed_plugin("first", PLUGIN_RECORD_A, order=0) + seed_plugin("second", PLUGIN_RECORD_B, order=1) + resp = self._send_chat_completion(self.MESSAGES) + self.assertEqual(resp.status_code, 200) + self.assertEqual(EVENTS, ["A", "B"]) + + def test_plugin_execution_order_follows_order_field(self): + seed_plugin("later", PLUGIN_RECORD_A, order=1) + seed_plugin("earlier", PLUGIN_RECORD_B, order=0) + resp = self._send_chat_completion(self.MESSAGES) + self.assertEqual(resp.status_code, 200) + self.assertEqual(EVENTS, ["B", "A"]) + + def test_blocking_plugin_short_circuits_later_plugins(self): + seed_plugin("guard", BLOCKING_PLUGIN, order=0) + seed_plugin("audit", PLUGIN_RECORD_B, order=1) + resp = self._send_chat_completion(self.MESSAGES) + self.assertEqual(resp.status_code, 403) + self.assertEqual(EVENTS, [], "later plugin should not run once an earlier one blocks") + + def test_before_pipeline_transforms_for_next_plugin(self): + seed_plugin("set", PLUGIN_SET_STEP, order=0) + seed_plugin("read", PLUGIN_READ_STEP, order=1) + resp = self._send_chat_completion(self.MESSAGES) + self.assertEqual(resp.status_code, 200) + self.assertEqual(EVENTS, ["from-pipe"]) + + def test_inactive_plugin_does_not_run(self): + Snippet.objects.create( + name="off", type=SnippetType.PLUGIN, active=False, code=PLUGIN_RECORD_A + ) + resp = self._send_chat_completion(self.MESSAGES) + self.assertEqual(resp.status_code, 200) + self.assertEqual(EVENTS, []) + + def test_generic_exception_in_before_becomes_error(self): + seed_plugin("boom", PLUGIN_BOOM) + self.client.raise_request_exception = False + resp = self._send_chat_completion(self.MESSAGES) + self.assertEqual(resp.status_code, 500) diff --git a/aqueduct/gateway/views/chat_completions.py b/aqueduct/gateway/views/chat_completions.py index e642be81..061c90aa 100644 --- a/aqueduct/gateway/views/chat_completions.py +++ b/aqueduct/gateway/views/chat_completions.py @@ -19,6 +19,7 @@ parse_body, process_file_content, resolve_alias, + run_plugins, token_authenticated, tos_accepted, ) @@ -38,6 +39,7 @@ @check_model_availability @normalize_reasoning_fields @log_request +@run_plugins @catch_router_exceptions async def chat_completions( request: ASGIRequest, diff --git a/aqueduct/gateway/views/completions.py b/aqueduct/gateway/views/completions.py index e6bfe615..99d22f1f 100644 --- a/aqueduct/gateway/views/completions.py +++ b/aqueduct/gateway/views/completions.py @@ -17,6 +17,7 @@ log_request, parse_body, resolve_alias, + run_plugins, token_authenticated, tos_accepted, ) @@ -34,6 +35,7 @@ @resolve_alias @check_model_availability @log_request +@run_plugins @catch_router_exceptions async def completions( request: ASGIRequest, diff --git a/aqueduct/gateway/views/embeddings.py b/aqueduct/gateway/views/embeddings.py index 0a2ef4d6..5771e879 100644 --- a/aqueduct/gateway/views/embeddings.py +++ b/aqueduct/gateway/views/embeddings.py @@ -15,6 +15,7 @@ log_request, parse_body, resolve_alias, + run_plugins, token_authenticated, tos_accepted, ) @@ -35,6 +36,7 @@ @resolve_alias @check_model_availability @log_request +@run_plugins @catch_router_exceptions async def embeddings( request: ASGIRequest, diff --git a/aqueduct/gateway/views/image_generation.py b/aqueduct/gateway/views/image_generation.py index ea65621c..afb28f5f 100644 --- a/aqueduct/gateway/views/image_generation.py +++ b/aqueduct/gateway/views/image_generation.py @@ -16,6 +16,7 @@ log_request, parse_body, resolve_alias, + run_plugins, token_authenticated, tos_accepted, ) @@ -35,6 +36,7 @@ @resolve_alias @check_model_availability @log_request +@run_plugins @catch_router_exceptions async def image_generation( request: ASGIRequest, diff --git a/aqueduct/gateway/views/responses.py b/aqueduct/gateway/views/responses.py index a3a6e0b8..fdf5ff3b 100644 --- a/aqueduct/gateway/views/responses.py +++ b/aqueduct/gateway/views/responses.py @@ -16,6 +16,7 @@ log_request, parse_body, resolve_alias, + run_plugins, token_authenticated, tos_accepted, validate_response_id, @@ -41,6 +42,7 @@ @check_model_availability @check_tool_availability @log_request +@run_plugins @catch_router_exceptions async def create_response( request: ASGIRequest, diff --git a/aqueduct/gateway/views/speech.py b/aqueduct/gateway/views/speech.py index 89ceecd2..9a3d5da1 100644 --- a/aqueduct/gateway/views/speech.py +++ b/aqueduct/gateway/views/speech.py @@ -17,6 +17,7 @@ log_request, parse_body, resolve_alias, + run_plugins, token_authenticated, tos_accepted, ) @@ -33,6 +34,7 @@ @ensure_usage @check_model_availability @log_request +@run_plugins @catch_router_exceptions async def speech( request: ASGIRequest, diff --git a/aqueduct/gateway/views/transcriptions.py b/aqueduct/gateway/views/transcriptions.py index 6641bbd9..56e40eb8 100644 --- a/aqueduct/gateway/views/transcriptions.py +++ b/aqueduct/gateway/views/transcriptions.py @@ -17,6 +17,7 @@ log_request, parse_body, resolve_alias, + run_plugins, token_authenticated, tos_accepted, ) @@ -40,6 +41,7 @@ class TranscriptionCreateParams(RootModel): # type: ignore[type-arg] @resolve_alias @check_model_availability @log_request +@run_plugins @catch_router_exceptions async def transcriptions( request: ASGIRequest, diff --git a/aqueduct/management/admin.py b/aqueduct/management/admin.py index 68b09176..7d01b0a5 100644 --- a/aqueduct/management/admin.py +++ b/aqueduct/management/admin.py @@ -38,6 +38,7 @@ VectorStoreFile, VectorStoreFileBatch, ) +from management.plugins import BlockedByPlugin, PluginSnippet, compile_plugin_class from management.snippets import ConfigSnippet, compile_snippet, get_config_snippet log = logging.getLogger("aqueduct") @@ -974,35 +975,102 @@ def vector_store_link(self, obj) -> str: class SnippetAdminForm(forms.ModelForm): + order = forms.IntegerField( + required=False, + min_value=0, + initial=0, + help_text=( + "Execution priority for plugins only: lower order runs first; ties " + "broken by lower id. Disabled for config snippets (only one config " + "is ever active)." + ), + ) + class Meta: model = Snippet - fields = ("name", "type", "active", "code") + fields = ("name", "type", "active", "order", "code") - def clean_code(self) -> str: - code = self.cleaned_data["code"] - require_subclass = self.cleaned_data.get("type") != SnippetType.PLUGIN - try: - compile_snippet(code, require_subclass=require_subclass) - except ValidationError as e: - raise forms.ValidationError(str(e)) from e - return code + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + snippet_type = self.instance.type if self.instance else SnippetType.CONFIG + if snippet_type != SnippetType.PLUGIN: + self.fields["order"].widget.attrs["readonly"] = True + self.fields["code"].widget.attrs.update( + {"class": "vLargeTextField snippet-code-input", "rows": 30} + ) + + +TEST_CODE_EXAMPLE = """\ +class Test(ConfigSnippet): + def org_name(self, claims): + return claims.get("org") + + def user_group(self, claims): + return (claims.get("groups") or ["default"])[0] + + def team_names(self, claims): + return claims.get("groups") or [] + + def display_team_names(self, team_names): + return [(name, True) for name in team_names] +""" TEST_PAYLOAD_EXAMPLE = {"email": "you@example.com", "groups": ["E123-Students", "admins"]} +PLUGIN_TEST_CODE_EXAMPLE = """\ +class MyPlugin(PluginSnippet): + def before_request(self, request, token, body): + text = " ".join(m.get("content", "") for m in body.get("messages", [])) + if "block me" in text: + raise BlockedByPlugin("blocked by guard", status=403) + body["flagged"] = True + return body + + def after_response(self, request, token, response): + response["audited"] = True + + def on_error(self, request, exc): + print("error happened:", exc) +""" + +PLUGIN_TEST_PAYLOAD_EXAMPLE = { + "model": "gpt-4.1-nano", + "messages": [{"role": "user", "content": "Hello"}], + "max_tokens": 50, +} + + +class _ConsoleUser: + email = "you@example.com" + + +class _ConsoleRequest: + def __init__(self) -> None: + self.META = {"REMOTE_ADDR": "127.0.0.1"} + self.headers = {"User-Agent": "aqueduct-plugin-console"} + self.user = _ConsoleUser() + + +class _ConsoleToken: + name = "Test Token" + user = _ConsoleUser() + + @admin.register(Snippet) class SnippetAdmin(admin.ModelAdmin): form = SnippetAdminForm - list_display: ClassVar[tuple] = ("name", "type", "active", "updated_at") + list_display: ClassVar[tuple] = ("id", "name", "type", "order", "active", "updated_at") list_filter: ClassVar[list] = ("type", "active") search_fields: ClassVar[tuple] = ("name",) readonly_fields: ClassVar[tuple] = ("updated_at",) fieldsets: ClassVar[tuple] = ( - (None, {"fields": ("name", "type", "active")}), + (None, {"fields": ("name", "type", "active", "order")}), ("Code", {"fields": ("code",), "classes": ("wide",)}), ) change_list_template = "admin/management/snippet/change_list.html" + change_form_template = "admin/management/snippet/change_form.html" @staticmethod def _allowed_emails() -> set[str]: @@ -1036,15 +1104,20 @@ def has_delete_permission(self, request, obj=None) -> bool: def get_urls(self) -> list[URLPattern]: custom_urls = [ path( - "test-console/", - self.admin_site.admin_view(self.test_console_view), + "config-test-console/", + self.admin_site.admin_view(self.config_test_console_view), name=f"{self.opts.app_label}_{self.opts.model_name}_test_console", - ) + ), + path( + "plugin-test-console/", + self.admin_site.admin_view(self.plugin_test_console_view), + name=f"{self.opts.app_label}_{self.opts.model_name}_plugin_test_console", + ), ] urls = super().get_urls() return custom_urls + urls - def test_console_view(self, request) -> HttpResponse: + def config_test_console_view(self, request) -> HttpResponse: if not self._is_allowed_superuser(request.user): raise PermissionDenied @@ -1052,6 +1125,7 @@ def test_console_view(self, request) -> HttpResponse: results = None if request.method != "POST": + code = TEST_CODE_EXAMPLE payload = json.dumps(TEST_PAYLOAD_EXAMPLE, indent=2) if request.method == "POST": @@ -1078,7 +1152,7 @@ def test_console_view(self, request) -> HttpResponse: context = { **self.admin_site.each_context(request), - "title": "Snippet test console", + "title": "Config test console", "code": code, "payload": payload, "results": results, @@ -1118,3 +1192,90 @@ def _run_console(snippet: ConfigSnippet, payload: Any) -> list[str]: lines.append(f"display_team_names(team names) -> {mapping!r}") return lines + + def plugin_test_console_view(self, request) -> HttpResponse: + if not self._is_allowed_superuser(request.user): + raise PermissionDenied + + code, payload = "", "" + results = None + + if request.method != "POST": + code = PLUGIN_TEST_CODE_EXAMPLE + payload = json.dumps(PLUGIN_TEST_PAYLOAD_EXAMPLE, indent=2) + + if request.method == "POST": + code = request.POST.get("code", "") + payload = request.POST.get("payload", "") + try: + cls = compile_plugin_class(code) + except ValidationError as e: + self.message_user(request, f"Plugin failed to compile: {e}", messages.ERROR) + else: + if payload.strip(): + try: + parsed = json.loads(payload) + except json.JSONDecodeError as e: + self.message_user(request, f"Invalid JSON test input: {e}", messages.ERROR) + else: + results = self._run_plugin_console(cls, parsed) + self.message_user( + request, "Plugin compiled and ran. Output below.", messages.SUCCESS + ) + + context = { + **self.admin_site.each_context(request), + "title": "Plugin test console", + "code": code, + "payload": payload, + "results": results, + } + request.current_app = self.admin_site.name + return TemplateResponse( + request, "admin/management/snippet/plugin_test_console.html", context + ) + + @staticmethod + def _run_plugin_console(plugin_cls: type[PluginSnippet], body: Any) -> list[str]: + """Run a plugin's hooks against a test payload using simple stand-ins.""" + lines = [] + plugin = plugin_cls() + request = _ConsoleRequest() + token = _ConsoleToken() + + if not isinstance(body, dict): + lines.append( + f"before_request skipped: test input must be a JSON object, " + f"got {type(body).__name__}" + ) + return lines + + try: + result = plugin.before_request(request, token, body) + except BlockedByPlugin as e: + lines.append(f"before_request -> BlockedByPlugin({e.reason!r}, status={e.status})") + return lines + except Exception as e: + lines.append(f"before_request raised {e!r}") + else: + if result is not None: + lines.append(f"before_request returned transformed body: {result!r}") + else: + lines.append("before_request ran; body unchanged") + + response = {"id": "chatcmpl-test", "choices": [], "usage": {}} + try: + plugin.after_response(request, token, response) + except Exception as e: + lines.append(f"after_response raised {e!r}") + else: + lines.append("after_response ran") + + try: + plugin.on_error(request, ValueError("simulated downstream error")) + except Exception as e: + lines.append(f"on_error raised {e!r}") + else: + lines.append("on_error ran") + + return lines diff --git a/aqueduct/management/migrations/0013_request_cached_and_reasoning_tokens.py b/aqueduct/management/migrations/0013_request_cached_and_reasoning_tokens.py index b92ca52b..b54c6671 100644 --- a/aqueduct/management/migrations/0013_request_cached_and_reasoning_tokens.py +++ b/aqueduct/management/migrations/0013_request_cached_and_reasoning_tokens.py @@ -4,7 +4,6 @@ class Migration(migrations.Migration): - dependencies = [("management", "0012_org_daily_limit_multiplier_and_more")] operations = [ @@ -12,16 +11,14 @@ class Migration(migrations.Migration): model_name="request", name="cached_input_tokens", field=models.PositiveIntegerField( - default=0, - help_text="Subset of input tokens read from the provider's prompt cache.", + default=0, help_text="Subset of input tokens read from the provider's prompt cache." ), ), migrations.AddField( model_name="request", name="reasoning_tokens", field=models.PositiveIntegerField( - default=0, - help_text="Subset of output tokens used for reasoning.", + default=0, help_text="Subset of output tokens used for reasoning." ), ), ] diff --git a/aqueduct/management/migrations/0014_snippet_order.py b/aqueduct/management/migrations/0014_snippet_order.py new file mode 100644 index 00000000..eb0f62ac --- /dev/null +++ b/aqueduct/management/migrations/0014_snippet_order.py @@ -0,0 +1,18 @@ +# Generated by Django 5.2.7 on 2026-10-19 17:22 + +from django.db import migrations, models + + +class Migration(migrations.Migration): + dependencies = [("management", "0013_request_cached_and_reasoning_tokens")] + + operations = [ + migrations.AddField( + model_name="snippet", + name="order", + field=models.PositiveSmallIntegerField( + default=0, + help_text="Execution priority for plugins only: lower order runs first; ties broken by lower id. Ignored for config snippets (only one config is ever active).", + ), + ) + ] diff --git a/aqueduct/management/models/snippet.py b/aqueduct/management/models/snippet.py index 40903a9d..54aaece7 100644 --- a/aqueduct/management/models/snippet.py +++ b/aqueduct/management/models/snippet.py @@ -19,6 +19,14 @@ class Snippet(models.Model): "be active at a time.", ) code = models.TextField(help_text="Python source of the snippet class.") + order = models.PositiveSmallIntegerField( + default=0, + help_text=( + "Execution priority for plugins only: lower order runs first; ties " + "broken by lower id. Ignored for config snippets (only one config " + "is ever active)." + ), + ) updated_at = models.DateTimeField(auto_now=True, help_text="Last modification time.") class Meta: diff --git a/aqueduct/management/plugins.py b/aqueduct/management/plugins.py new file mode 100644 index 00000000..6a41f5e5 --- /dev/null +++ b/aqueduct/management/plugins.py @@ -0,0 +1,116 @@ +import inspect +import logging +import time +from functools import lru_cache +from typing import Any + +from django.core.exceptions import ValidationError + +from management.models import Snippet, SnippetType + +log = logging.getLogger("aqueduct") + + +class PluginSnippet: + processing_time_ms: int = 0 + + def before_request(self, request: Any, token: Any, body: Any) -> Any: + pass + + def after_response(self, request: Any, token: Any, response: Any) -> None: + pass + + def on_error(self, request: Any, exc: Exception) -> None: + pass + + +class BlockedByPlugin(Exception): # noqa: N818 + def __init__(self, reason: str, status: int = 403) -> None: + super().__init__(reason) + self.reason = reason + self.status = status + + +def _record_call(plugin: PluginSnippet, start: float) -> None: + duration_ms = int((time.monotonic() - start) * 1000) + plugin.processing_time_ms = duration_ms + log.debug("Plugin %s hook took %d ms", type(plugin).__name__, duration_ms) + + +def before_hook(plugins: list[PluginSnippet], request: Any, token: Any, body: Any) -> Any: + for plugin in plugins: + call = getattr(plugin, "before_request", None) + if call is None: + continue + start = time.monotonic() + result = call(request, token, body) + _record_call(plugin, start) + if result is not None: + body = result + return body + + +def after_hook(plugins: list[PluginSnippet], request: Any, token: Any, response: Any) -> None: + for plugin in plugins: + call = getattr(plugin, "after_response", None) + if call is None: + continue + start = time.monotonic() + call(request, token, response) + _record_call(plugin, start) + + +def _error_hook(plugins: list[PluginSnippet], request: Any, exc: Exception) -> None: + for plugin in plugins: + call = getattr(plugin, "on_error", None) + if call is None: + continue + start = time.monotonic() + call(request, exc) + _record_call(plugin, start) + + +def compile_plugin_class(code: str) -> type[PluginSnippet]: + try: + source = compile(code, "", "exec") + except SyntaxError as e: + raise ValidationError(f"Syntax error: {e}") from e + + namespace: dict[str, Any] = {"PluginSnippet": PluginSnippet, "BlockedByPlugin": BlockedByPlugin} + try: + exec(source, namespace) # noqa: S102 + except Exception as e: + raise ValidationError(f"Error executing plugin code: {e}") from e + + subclasses = [ + value + for value in namespace.values() + if inspect.isclass(value) + and issubclass(value, PluginSnippet) + and value is not PluginSnippet + ] + if not subclasses: + raise ValidationError( + "Snippet must define exactly one class that subclasses 'PluginSnippet'." + ) + if len(subclasses) > 1: + raise ValidationError( + "Snippet must define exactly one class that subclasses 'PluginSnippet' " + f"(found {len(subclasses)})." + ) + return subclasses[0] + + +@lru_cache(maxsize=128) +def _plugin_class(pk: int) -> type[PluginSnippet]: + snippet = Snippet.objects.get(pk=pk) + return compile_plugin_class(snippet.code) + + +def resolve_active_plugins() -> list[PluginSnippet]: + ids = list( + Snippet.objects.filter(type=SnippetType.PLUGIN, active=True) + .order_by("order", "id") + .values_list("pk", flat=True) + ) + return [_plugin_class(pk)() for pk in ids] diff --git a/aqueduct/management/templates/admin/management/snippet/change_form.html b/aqueduct/management/templates/admin/management/snippet/change_form.html new file mode 100644 index 00000000..b7f4f2fb --- /dev/null +++ b/aqueduct/management/templates/admin/management/snippet/change_form.html @@ -0,0 +1,136 @@ +{% extends "admin/change_form.html" %} +{% load static %} + +{% block extrahead %} + {{ block.super }} + + +{% endblock %} + +{% block extrastyle %} + {{ block.super }} + +{% endblock %} + +{% block admin_change_form_document_ready %} + {{ block.super }} + + +{% endblock %} diff --git a/aqueduct/management/templates/admin/management/snippet/change_list.html b/aqueduct/management/templates/admin/management/snippet/change_list.html index 9e97fce1..038d6b81 100644 --- a/aqueduct/management/templates/admin/management/snippet/change_list.html +++ b/aqueduct/management/templates/admin/management/snippet/change_list.html @@ -1,7 +1,10 @@ {% extends "admin/change_list.html" %} {% block object-tools-items %}
  • - Test console + Config test console +
    • +
  • + Plugin test console
    • {{ block.super }} {% endblock %} diff --git a/aqueduct/management/templates/admin/management/snippet/plugin_test_console.html b/aqueduct/management/templates/admin/management/snippet/plugin_test_console.html new file mode 100644 index 00000000..2b0a5905 --- /dev/null +++ b/aqueduct/management/templates/admin/management/snippet/plugin_test_console.html @@ -0,0 +1,138 @@ +{% extends "admin/base_site.html" %} +{% load i18n static %} + +{% block extrastyle %} + {{ block.super }} + + +{% endblock %} + +{% block content %} +
    +
    +

    + Warning: This test console executes real server code as an + authenticated superuser. There is no sandboxing. For this reason this console + (and snippet editing) is restricted to superusers. +

    +
    {% csrf_token %} +
    +
    + + +
    Python class subclassing PluginSnippet (auto-injected, no import needed). The code editor is prefilled with a runnable example.
    +
    +
    + + +
    + Paste a request body (what your client sends the gateway). + The console chains it through the plugin hooks like the real request flow: +
      +
    • before_request(request, token, body) receives the object; a returned dict is shown as the transformed body.
      • +
    • after_response(request, token, response) receives a sample response dict.
      • +
    • on_error(request, exc) receives a simulated error.
      • +
    • Try adding "block me" to a message to see the example BlockedByPlugin path.
      • +
    +
    +
    +
    +
    + +
    +
    + + {% if results %} +
    +

    Output

    +
      + {% for line in results %} +
    • {{ line }}
      • + {% endfor %} +
    +
    + {% endif %} +
    +
    + + + +{% endblock %} diff --git a/aqueduct/management/templates/admin/management/snippet/test_console.html b/aqueduct/management/templates/admin/management/snippet/test_console.html index 051093bc..c8ae093d 100644 --- a/aqueduct/management/templates/admin/management/snippet/test_console.html +++ b/aqueduct/management/templates/admin/management/snippet/test_console.html @@ -41,11 +41,7 @@ placeholder='class Test(ConfigSnippet): def org_name(self, claims): return claims.get("org")'>{{ code }} -
    Python class subclassing ConfigSnippet (auto-injected, no import needed).
    - 
    Example:
-class Test(ConfigSnippet):
-    def org_name(self, claims):
-        return claims.get("org")
    +
    Python class subclassing ConfigSnippet (auto-injected, no import needed). The code editor is prefilled with a runnable example.
    diff --git a/aqueduct/management/tests/test_plugins.py b/aqueduct/management/tests/test_plugins.py new file mode 100644 index 00000000..8cd29f9a --- /dev/null +++ b/aqueduct/management/tests/test_plugins.py @@ -0,0 +1,215 @@ +from django.core.exceptions import ValidationError +from django.test import TestCase + +from management.models import Snippet, SnippetType +from management.plugins import ( + BlockedByPlugin, + PluginSnippet, + _error_hook, + _plugin_class, + after_hook, + before_hook, + compile_plugin_class, + resolve_active_plugins, +) + +GOOD_PLUGIN = """\ +class MyPlugin(PluginSnippet): + def before_request(self, request, token, body): + body["flagged"] = True + return body + + def after_response(self, request, token, response): + self.seen = response +""" + +BAD_SYNTAX = "class MyPlugin(PluginSnippet):\n def before_request(self" + + +class PluginBaseTestCase(TestCase): + def test_default_hooks_are_no_ops(self): + p = PluginSnippet() + self.assertIsNone(p.before_request("r", "t", "b")) + self.assertIsNone(p.after_response("r", "t", "resp")) + self.assertIsNone(p.on_error("r", ValueError("boom"))) + + def test_blocked_by_plugin_defaults(self): + e = BlockedByPlugin("nope") + self.assertEqual(e.reason, "nope") + self.assertEqual(e.status, 403) + self.assertEqual(str(e), "nope") + + def test_blocked_by_plugin_custom_status(self): + e = BlockedByPlugin("nope", status=429) + self.assertEqual(e.status, 429) + + +class CompilePluginTestCase(TestCase): + def test_compiles_valid_plugin(self): + cls = compile_plugin_class(GOOD_PLUGIN) + self.assertTrue(issubclass(cls, PluginSnippet)) + self.assertIsNot(cls, PluginSnippet) + + def test_base_class_auto_injected(self): + cls = compile_plugin_class("class P(PluginSnippet):\n pass\n") + self.assertTrue(issubclass(cls, PluginSnippet)) + + def test_rejects_syntax_error(self): + with self.assertRaises(ValidationError): + compile_plugin_class(BAD_SYNTAX) + + def test_rejects_exec_error(self): + with self.assertRaises(ValidationError): + compile_plugin_class("raise RuntimeError('boom')") + + def test_rejects_no_subclass(self): + with self.assertRaises(ValidationError): + compile_plugin_class("x = 1") + + def test_rejects_non_subclass(self): + with self.assertRaises(ValidationError): + compile_plugin_class("class Anything:\n pass\n") + + def test_rejects_multiple_subclasses(self): + with self.assertRaises(ValidationError): + compile_plugin_class( + "class A(PluginSnippet):\n pass\n\nclass B(PluginSnippet):\n pass\n" + ) + + +class ResolvePluginsTestCase(TestCase): + def setUp(self): + _plugin_class.cache_clear() + + def _plugin(self, name: str, cap: str, active: bool = True, order: int = 0) -> Snippet: + return Snippet.objects.create( + name=name, + type=SnippetType.PLUGIN, + active=active, + order=order, + code=( + "class P(PluginSnippet):\n" + " def after_response(self, r, t, resp):\n" + f" {cap}\n" + ), + ) + + def test_no_active_plugins_by_default(self): + self.assertEqual(resolve_active_plugins(), []) + + def test_only_active_plugins_resolved(self): + self._plugin("on", "pass", active=True) + self._plugin("off", "pass", active=False) + plugins = resolve_active_plugins() + self.assertEqual(len(plugins), 1) + self.assertEqual(type(plugins[0]).__name__, "P") + + def test_config_snippets_are_not_plugins(self): + Snippet.objects.create( + name="cfg", + type=SnippetType.CONFIG, + active=True, + code="class C(ConfigSnippet):\n pass\n", + ) + self.assertEqual(resolve_active_plugins(), []) + + def test_ordered_by_order_pk(self): + self._plugin("first", "pass", order=2) + self._plugin("second", "pass", order=0) + self._plugin("third", "pass", order=1) + expected = list( + Snippet.objects.filter(type=SnippetType.PLUGIN, active=True) + .order_by("order", "id") + .values_list("pk", flat=True) + ) + instances = resolve_active_plugins() + for pk in expected: + _plugin_class(pk) + self.assertIsInstance(instances[0], PluginSnippet) + self.assertEqual(len(instances), len(expected)) + + def test_fresh_instance_per_call_and_class_cached(self): + self._plugin("on", "pass") + plugins_a = resolve_active_plugins() + plugins_b = resolve_active_plugins() + self.assertIsNot(plugins_a[0], plugins_b[0]) + self.assertIs(type(plugins_a[0]), type(plugins_b[0])) + + +class DispatchPluginsTestCase(TestCase): + def _make(self, before=None, after=None, on_error=None) -> PluginSnippet: + class P(PluginSnippet): + pass + + if before is not None: + P.before_request = staticmethod(before) + if after is not None: + P.after_response = staticmethod(after) + if on_error is not None: + P.on_error = staticmethod(on_error) + return P() + + def test_before_hook_pipeline(self): + calls = [] + + def a(request, token, body): + calls.append("a") + body["a"] = 1 + return body + + def b(request, token, body): + calls.append("b") + + plugins = [self._make(before=a), self._make(before=b)] + body = {"x": 0} + result = before_hook(plugins, "req", "tok", body) + self.assertEqual(result, {"x": 0, "a": 1}) + self.assertEqual(calls, ["a", "b"]) + + def test_before_hook_blocked_propagates(self): + def blocker(request, token, body): + raise BlockedByPlugin("denied") + + reached = [] + + def a(request, token, body): + reached.append("a") + return body + + plugins = [ + self._make(before=a), + self._make(before=blocker), + self._make(before=lambda r, t, b: None), + ] + with self.assertRaises(BlockedByPlugin) as ctx: + before_hook(plugins, "req", "tok", {}) + self.assertEqual(ctx.exception.reason, "denied") + self.assertEqual(reached, ["a"]) + + def test_before_hook_records_processing_time(self): + p = self._make(before=lambda r, t, b: b) + before_hook([p], "req", "tok", {}) + self.assertIsInstance(p.processing_time_ms, int) + + def test_after_hook_runs_all(self): + calls = [] + + def a(request, token, response): + calls.append(("a", response)) + + def b(request, token, response): + calls.append(("b", response)) + + plugins = [self._make(after=a), self._make(after=b)] + after_hook(plugins, "req", "tok", "resp") + self.assertEqual(calls, [("a", "resp"), ("b", "resp")]) + + def test_error_hook_receives_exception(self): + seen = [] + e = ValueError("boom") + + def handler(request, exc): + seen.append(exc) + + _error_hook([self._make(on_error=handler)], "req", e) + self.assertEqual(seen, [e]) diff --git a/aqueduct/management/tests/test_snippet_admin.py b/aqueduct/management/tests/test_snippet_admin.py index 68656145..c062aef7 100644 --- a/aqueduct/management/tests/test_snippet_admin.py +++ b/aqueduct/management/tests/test_snippet_admin.py @@ -6,9 +6,10 @@ from django.test import TestCase, override_settings from django.urls import reverse -from management.admin import SnippetAdminForm +from management.admin import SnippetAdmin, SnippetAdminForm from management.auth import OIDCBackend from management.models import Snippet, SnippetType +from management.plugins import compile_plugin_class from management.tests.helpers import ( SNIPPET_ORG_CUSTOM, SNIPPET_ORG_FROM_FIRST_GROUP, @@ -40,7 +41,7 @@ def test_valid_code_passes(self): ) self.assertTrue(form.is_valid(), form.errors) - def test_syntax_error_rejected(self): + def test_syntax_error_not_rejected(self): form = SnippetAdminForm( data={ "name": "c", @@ -49,9 +50,9 @@ def test_syntax_error_rejected(self): "code": "class C(ConfigSnippet):\n def bad(self", } ) - self.assertFalse(form.is_valid()) + self.assertTrue(form.is_valid(), form.errors) - def test_wrong_signature_rejected(self): + def test_wrong_signature_not_rejected(self): form = SnippetAdminForm( data={ "name": "c", @@ -60,8 +61,7 @@ def test_wrong_signature_rejected(self): "code": "class C(ConfigSnippet):\n def org_name(self):\n return 'x'\n", } ) - self.assertFalse(form.is_valid()) - self.assertIn("org_name", str(form.errors["code"])) + self.assertTrue(form.is_valid(), form.errors) def test_second_active_config_demotes_previous(self): seed_active_config(SNIPPET_ORG_CUSTOM) @@ -144,11 +144,39 @@ def test_plugin_does_not_require_config_snippet_subclass(self): ) self.assertTrue(form.is_valid(), form.errors) - def test_plugin_syntax_error_still_rejected(self): + def test_plugin_syntax_error_not_rejected(self): form = SnippetAdminForm( data={"name": "p", "type": "plugin", "active": True, "code": "def broken("} ) - self.assertFalse(form.is_valid()) + self.assertTrue(form.is_valid(), form.errors) + + def test_order_locked_for_config_only(self): + form = SnippetAdminForm( + data={"name": "c", "type": "config", "active": True, "code": VALID_CODE} + ) + self.assertTrue(form.fields["order"].widget.attrs.get("readonly")) + + plugin = Snippet(type=SnippetType.PLUGIN) + form2 = SnippetAdminForm( + data={"name": "p", "type": "plugin", "active": True, "code": VALID_CODE}, + instance=plugin, + ) + self.assertFalse(form2.fields["order"].widget.attrs.get("readonly")) + + def test_plugin_order_is_saved(self): + form = SnippetAdminForm( + data={ + "name": "p", + "type": "plugin", + "active": True, + "order": "7", + "code": "class P(PluginSnippet):\n pass\n", + } + ) + self.assertTrue(form.is_valid(), form.errors) + obj = form.save() + obj.refresh_from_db() + self.assertEqual(obj.order, 7) def test_orm_create_active_config_demotes_previous(self) -> None: seed_active_config(SNIPPET_ORG_CUSTOM) @@ -212,6 +240,13 @@ def test_superuser_can_view_changelist(self): resp = self.client.get(self.changelist_url) self.assertEqual(resp.status_code, 200) + def test_superuser_add_form_renders_code_editor(self): + self.client.force_login(self.superuser) + resp = self.client.get(reverse("admin:management_snippet_add")) + self.assertEqual(resp.status_code, 200) + self.assertContains(resp, "snippet-code-input") + self.assertContains(resp, "@codemirror/lang-python") + def test_non_superuser_staff_forbidden_on_changelist(self): self.client.force_login(self.staff) resp = self.client.get(self.changelist_url) @@ -269,20 +304,20 @@ def setUp(self): def test_get_shows_console(self): resp = self.client.get(self.url) self.assertEqual(resp.status_code, 200) - self.assertContains(resp, "Snippet test console") + self.assertContains(resp, "Config test console") self.assertContains(resp, "executes real server code") self.assertContains(resp, "restricted to superusers") - def test_get_prefills_example_payload(self): + def test_get_prefills_example_code_and_payload(self): resp = self.client.get(self.url) self.assertEqual(resp.status_code, 200) - # The Test input is pre-filled with a runnable claims example. + self.assertContains(resp, "class Test(ConfigSnippet)") + self.assertContains(resp, "display_team_names") self.assertContains(resp, '"email"') self.assertContains(resp, "you@example.com") self.assertContains(resp, '"groups"') self.assertContains(resp, "E123-Students") - # Dimmed on-page code example + large editor are present. - self.assertContains(resp, "snippet-console-example") + self.assertNotContains(resp, "snippet-console-example") self.assertContains(resp, "snippet-console-code") def test_runs_methods_against_payload(self): @@ -309,3 +344,68 @@ def test_rejects_invalid_json(self): ) self.assertEqual(resp.status_code, 200) self.assertContains(resp, "Invalid JSON test input") + + +PLUGIN_CODE = """\ +class MyPlugin(PluginSnippet): + def before_request(self, request, token, body): + body["flagged"] = True + return body + + def after_response(self, request, token, response): + response["audited"] = True +""" + + +PLUGIN_BLOCKER = """\ +class Blocker(PluginSnippet): + def before_request(self, request, token, body): + raise BlockedByPlugin("forbidden by guard", status=403) +""" + + +@override_settings(ADMIN_SUPERUSER_EMAILS=[SUPERUSER_EMAIL]) +class PluginConsoleTestCase(TestCase): + def setUp(self): + self.superuser = User.objects.create_superuser( + username="admin", email=SUPERUSER_EMAIL, password="pw" + ) + self.client.force_login(self.superuser) + self.url = reverse("admin:management_snippet_plugin_test_console") + + def test_get_shows_console_with_example_data(self): + resp = self.client.get(self.url) + self.assertEqual(resp.status_code, 200) + self.assertContains(resp, "Plugin test console") + self.assertContains(resp, "Plugin code") + self.assertContains(resp, "class MyPlugin(PluginSnippet)") + self.assertContains(resp, '"model"') + + def test_runs_plugin_hooks_against_payload(self): + resp = self.client.post( + self.url, + {"code": PLUGIN_CODE, "payload": json.dumps({"model": "gpt-4.1-nano", "messages": []})}, + ) + self.assertEqual(resp.status_code, 200) + self.assertContains(resp, "Plugin compiled and ran.") + self.assertContains(resp, "before_request returned transformed body") + self.assertContains(resp, "after_response ran") + + def test_blocking_plugin_is_reported(self): + cls = compile_plugin_class(PLUGIN_BLOCKER) + lines = SnippetAdmin._run_plugin_console(cls, {"model": "x"}) + self.assertEqual( + lines, ["before_request -> BlockedByPlugin('forbidden by guard', status=403)"] + ) + + def test_rejects_invalid_code(self): + resp = self.client.post( + self.url, {"code": "class P(PluginSnippet):\n oops", "payload": "{}"} + ) + self.assertEqual(resp.status_code, 200) + self.assertContains(resp, "failed to compile") + + def test_rejects_invalid_json(self): + resp = self.client.post(self.url, {"code": PLUGIN_CODE, "payload": "{not json"}) + self.assertEqual(resp.status_code, 200) + self.assertContains(resp, "Invalid JSON test input")