diff --git a/libs/arcade-mcp-server/arcade_mcp_server/managers/prompt.py b/libs/arcade-mcp-server/arcade_mcp_server/managers/prompt.py index b0c6d0241..e6734a548 100644 --- a/libs/arcade-mcp-server/arcade_mcp_server/managers/prompt.py +++ b/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,7 +96,10 @@ 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) @@ -80,7 +107,7 @@ async def get_prompt( 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: diff --git a/libs/arcade-mcp-server/arcade_mcp_server/mcp_app.py b/libs/arcade-mcp-server/arcade_mcp_server/mcp_app.py index d4da463ca..07c5ead46 100644 --- a/libs/arcade-mcp-server/arcade_mcp_server/mcp_app.py +++ b/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.") diff --git a/libs/arcade-mcp-server/arcade_mcp_server/server.py b/libs/arcade-mcp-server/arcade_mcp_server/server.py index de9a19d8d..fc8272d21 100644 --- a/libs/arcade-mcp-server/arcade_mcp_server/server.py +++ b/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: diff --git a/libs/tests/arcade_mcp_server/test_prompt.py b/libs/tests/arcade_mcp_server/test_prompt.py index cf2f69814..300365dcc 100644 --- a/libs/tests/arcade_mcp_server/test_prompt.py +++ b/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