feat: add context parameter to prompt handlers

libs/arcade-mcp-server/arcade_mcp_server/managers/prompt.py

@@ -6,8 +6,9 @@
 
 from __future__ import annotations
 
+import inspect
 import logging
-from typing import Callable
+from typing import Any, Callable
 
 from arcade_mcp_server.exceptions import NotFoundError, PromptError
 from arcade_mcp_server.managers.base import ComponentManager
@@ -22,10 +23,28 @@ class PromptHandler:
     def __init__(
         self,
         prompt: Prompt,
-        handler: Callable[[dict[str, str]], list[PromptMessage]] | None = None,
+        handler: Callable[..., list[PromptMessage]] | None = None,
     ) -> None:
         self.prompt = prompt
         self.handler = handler or self._default_handler
+        self._accepts_context = self._check_accepts_context(self.handler)
+
+    @staticmethod
+    def _check_accepts_context(handler: Callable[..., Any]) -> bool:
+        """Check if the handler accepts a context parameter (2-parameter signature)."""
+        try:
+            sig = inspect.signature(handler)
+            params = [
+                p
+                for p in sig.parameters.values()
+                if p.name != "self"
+                and p.kind
+                not in (p.VAR_POSITIONAL, p.VAR_KEYWORD, p.KEYWORD_ONLY)
+                and p.default is p.empty
+            ]
+            return len(params) >= 2
+        except (ValueError, TypeError):
+            return False
 
     def __eq__(self, other: object) -> bool:  # pragma: no cover - simple comparison
         if not isinstance(other, PromptHandler):
@@ -43,7 +62,9 @@ def _default_handler(self, arguments: dict[str, str]) -> list[PromptMessage]:
             )
         ]
 
-    async def get_messages(self, arguments: dict[str, str] | None = None) -> list[PromptMessage]:
+    async def get_messages(
+        self, arguments: dict[str, str] | None = None, context: Any | None = None
+    ) -> list[PromptMessage]:
         args = arguments or {}
 
         # Validate required arguments
@@ -52,7 +73,10 @@ async def get_messages(self, arguments: dict[str, str] | None = None) -> list[Pr
                 if arg.required and arg.name not in args:
                     raise PromptError(f"Required argument '{arg.name}' not provided")
 
-        result = self.handler(args)
+        if self._accepts_context:
+            result = self.handler(context, args)
+        else:
+            result = self.handler(args)
         if hasattr(result, "__await__"):
             result = await result
 
@@ -72,15 +96,18 @@ async def list_prompts(self) -> list[Prompt]:
         return [h.prompt for h in handlers]
 
     async def get_prompt(
-        self, name: str, arguments: dict[str, str] | None = None
+        self,
+        name: str,
+        arguments: dict[str, str] | None = None,
+        context: Any | None = None,
     ) -> GetPromptResult:
         try:
             handler = await self.registry.get(name)
         except KeyError:
             raise NotFoundError(f"Prompt '{name}' not found")
 
         try:
-            messages = await handler.get_messages(arguments)
+            messages = await handler.get_messages(arguments, context=context)
             return GetPromptResult(
                 description=handler.prompt.description,
                 messages=messages,
@@ -93,7 +120,7 @@ async def get_prompt(
     async def add_prompt(
         self,
         prompt: Prompt,
-        handler: Callable[[dict[str, str]], list[PromptMessage]] | None = None,
+        handler: Callable[..., list[PromptMessage]] | None = None,
     ) -> None:
         prompt_handler = PromptHandler(prompt, handler)
         await self.registry.upsert(prompt.name, prompt_handler)
@@ -109,7 +136,7 @@ async def update_prompt(
         self,
         name: str,
         prompt: Prompt,
-        handler: Callable[[dict[str, str]], list[PromptMessage]] | None = None,
+        handler: Callable[..., list[PromptMessage]] | None = None,
     ) -> Prompt:
         # Ensure exists
         try:

libs/arcade-mcp-server/arcade_mcp_server/mcp_app.py

@@ -774,7 +774,7 @@ def __init__(self, app: MCPApp) -> None:
         self._app = app
 
     async def add(
-        self, prompt: Prompt, handler: Callable[[dict[str, str]], list[PromptMessage]] | None = None
+        self, prompt: Prompt, handler: Callable[..., list[PromptMessage]] | None = None
     ) -> None:
         if self._app.server is None:
             raise ServerError("No server bound to app. Set app.server to use runtime prompts API.")

libs/arcade-mcp-server/arcade_mcp_server/server.py

@@ -2015,9 +2015,11 @@ async def _handle_get_prompt(
     ) -> JSONRPCResponse[GetPromptResult] | JSONRPCError:
         """Handle get prompt request."""
         try:
+            context = get_current_model_context()
             result = await self._prompt_manager.get_prompt(
                 message.params.name,
                 message.params.arguments if hasattr(message.params, "arguments") else None,
+                context=context,
             )
             return JSONRPCResponse(id=message.id, result=result)
         except NotFoundError:

libs/tests/arcade_mcp_server/test_prompt.py

@@ -1,10 +1,11 @@
 """Tests for Prompt Manager implementation."""
 
 import asyncio
+from unittest.mock import MagicMock
 
 import pytest
 from arcade_mcp_server.exceptions import NotFoundError, PromptError
-from arcade_mcp_server.managers.prompt import PromptManager
+from arcade_mcp_server.managers.prompt import PromptHandler, PromptManager
 from arcade_mcp_server.types import (
     GetPromptResult,
     Prompt,
@@ -239,3 +240,122 @@ async def error_prompt(args: dict[str, str]):
 
         with pytest.raises(PromptError):
             await manager.get_prompt("error_prompt", {})
+
+
+class TestPromptHandlerContext:
+    """Test context parameter support in prompt handlers."""
+
+    @pytest.mark.asyncio
+    async def test_handler_with_context_parameter(self):
+        """Test that a prompt handler with context receives the context object."""
+        received_context = None
+
+        async def handler_with_context(context, args: dict[str, str]) -> list[PromptMessage]:
+            nonlocal received_context
+            received_context = context
+            name = args.get("name", "World")
+            return [PromptMessage(role="user", content={"type": "text", "text": f"Hello {name}"})]
+
+        prompt = Prompt(
+            name="ctx_prompt",
+            description="A prompt with context",
+            arguments=[PromptArgument(name="name", required=False)],
+        )
+        manager = PromptManager()
+        await manager.add_prompt(prompt, handler_with_context)
+
+        mock_context = MagicMock()
+        result = await manager.get_prompt("ctx_prompt", {"name": "Alice"}, context=mock_context)
+
+        assert received_context is mock_context
+        assert len(result.messages) == 1
+        assert "Hello Alice" in result.messages[0].content["text"]
+
+    @pytest.mark.asyncio
+    async def test_handler_without_context_still_works(self):
+        """Test backward compatibility: old-style handlers without context still work."""
+
+        async def handler_no_context(args: dict[str, str]) -> list[PromptMessage]:
+            name = args.get("name", "World")
+            return [PromptMessage(role="user", content={"type": "text", "text": f"Hello {name}"})]
+
+        prompt = Prompt(
+            name="no_ctx_prompt",
+            description="A prompt without context",
+            arguments=[PromptArgument(name="name", required=False)],
+        )
+        manager = PromptManager()
+        await manager.add_prompt(prompt, handler_no_context)
+
+        mock_context = MagicMock()
+        result = await manager.get_prompt(
+            "no_ctx_prompt", {"name": "Bob"}, context=mock_context
+        )
+
+        assert len(result.messages) == 1
+        assert "Hello Bob" in result.messages[0].content["text"]
+
+    @pytest.mark.asyncio
+    async def test_handler_without_context_no_context_provided(self):
+        """Test old-style handler works when no context is provided at all."""
+
+        async def handler_no_context(args: dict[str, str]) -> list[PromptMessage]:
+            return [PromptMessage(role="user", content={"type": "text", "text": "Hi"})]
+
+        prompt = Prompt(name="simple", description="Simple prompt")
+        manager = PromptManager()
+        await manager.add_prompt(prompt, handler_no_context)
+
+        result = await manager.get_prompt("simple", {})
+        assert len(result.messages) == 1
+
+    @pytest.mark.asyncio
+    async def test_default_handler_works_without_context(self):
+        """Test that the default handler (no user handler) still works."""
+        prompt = Prompt(name="default_prompt", description="Default prompt")
+        manager = PromptManager()
+        await manager.add_prompt(prompt)
+
+        result = await manager.get_prompt("default_prompt", {}, context=MagicMock())
+        assert len(result.messages) == 1
+        assert "Default prompt" in result.messages[0].content["text"]
+
+    @pytest.mark.asyncio
+    async def test_sync_handler_with_context(self):
+        """Test that a synchronous handler with context works."""
+
+        def sync_handler_with_context(context, args: dict[str, str]) -> list[PromptMessage]:
+            return [
+                PromptMessage(
+                    role="user", content={"type": "text", "text": f"Sync: {args.get('name', '')}"}
+                )
+            ]
+
+        prompt = Prompt(
+            name="sync_ctx",
+            description="Sync with context",
+            arguments=[PromptArgument(name="name", required=False)],
+        )
+        manager = PromptManager()
+        await manager.add_prompt(prompt, sync_handler_with_context)
+
+        mock_context = MagicMock()
+        result = await manager.get_prompt("sync_ctx", {"name": "Test"}, context=mock_context)
+
+        assert "Sync: Test" in result.messages[0].content["text"]
+
+    def test_check_accepts_context_two_params(self):
+        """Test introspection detects 2-param handler as context-accepting."""
+
+        def handler(ctx, args):
+            pass
+
+        assert PromptHandler._check_accepts_context(handler) is True
+
+    def test_check_accepts_context_one_param(self):
+        """Test introspection detects 1-param handler as not context-accepting."""
+
+        def handler(args):
+            pass
+
+        assert PromptHandler._check_accepts_context(handler) is False