diff --git a/src/fast_agent/mcp/sampling.py b/src/fast_agent/mcp/sampling.py index 883563043..3f3cff78d 100644 --- a/src/fast_agent/mcp/sampling.py +++ b/src/fast_agent/mcp/sampling.py @@ -36,6 +36,9 @@ logger = get_logger(__name__) +# Stop reasons defined by the MCP sampling spec for a final (non-tool) result. +_SAMPLING_STOP_REASONS = (LlmStopReason.STOP_SEQUENCE, LlmStopReason.MAX_TOKENS) + @dataclass(frozen=True, slots=True) class _SamplingModelSelection: @@ -180,11 +183,14 @@ def _sampling_response( stop_reason="toolUse", ) + stop_reason = llm_response.stop_reason + if stop_reason not in _SAMPLING_STOP_REASONS: + stop_reason = LlmStopReason.END_TURN return CreateMessageResult( role=llm_response.role, content=TextContent(type="text", text=llm_response.first_text()), model=model, - stop_reason=LlmStopReason.END_TURN.value, + stop_reason=stop_reason.value, ) diff --git a/tests/unit/fast_agent/mcp/test_sampling.py b/tests/unit/fast_agent/mcp/test_sampling.py index c74030c5e..78a81458e 100644 --- a/tests/unit/fast_agent/mcp/test_sampling.py +++ b/tests/unit/fast_agent/mcp/test_sampling.py @@ -1,6 +1,11 @@ -from mcp_types import CreateMessageRequestParams, SamplingMessage, TextContent +import pytest +from mcp_types import CreateMessageRequestParams, CreateMessageResult, SamplingMessage, TextContent +from fast_agent.config import MCPSamplingSettings, MCPServerSettings +from fast_agent.mcp import sampling +from fast_agent.mcp.prompt import Prompt from fast_agent.mcp.sampling import sampling_agent_config +from fast_agent.types.llm_stop_reason import LlmStopReason def test_build_sampling_agent_config_with_system_prompt(): @@ -46,3 +51,38 @@ def test_build_sampling_agent_config_empty_system_prompt(): # Verify instruction is the empty string as received in params.systemPrompt assert config.instruction == "" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("llm_stop_reason", "expected"), + [ + (LlmStopReason.MAX_TOKENS, "maxTokens"), + (LlmStopReason.STOP_SEQUENCE, "stopSequence"), + (LlmStopReason.END_TURN, "endTurn"), + ], +) +async def test_sample_reports_why_generation_stopped( + monkeypatch: pytest.MonkeyPatch, llm_stop_reason: LlmStopReason, expected: str +) -> None: + class _StoppedLLM: + async def generate(self, *_args: object, **_kwargs: object): + return Prompt.assistant("partial", stop_reason=llm_stop_reason) + + monkeypatch.setattr(sampling, "create_sampling_llm", lambda *_args: _StoppedLLM()) + params = CreateMessageRequestParams( + max_tokens=5, + messages=[SamplingMessage(role="user", content=TextContent(type="text", text="Hello"))], + ) + + result = await sampling.sample( + params, + server_name="test", + server_config=MCPServerSettings(sampling=MCPSamplingSettings(model="passthrough")), + agent_model=None, + api_key=None, + app_context=None, + ) + + assert isinstance(result, CreateMessageResult) + assert result.stop_reason == expected