Get sampling and elicitation working with new reverse proxy #470
Conversation
|
Caution Review failedThe pull request is closed. WalkthroughAdds nested sampling and elicitation example servers and integrates them into asyncio and Temporal example clients/servers. Threads optional runtime Context through client/session creation. Refactors client session to upstream-pass-through with local SamplingHandler fallback and elicitation forwarding. Adjusts Temporal executor/task semantics and notification behavior. Normalizes gateway URLs and revises client proxy timeout handling. Changes
Sequence Diagram(s)sequenceDiagram
autonumber
actor User
participant Client as Example Client
participant App as MCPApp (Asyncio/Temporal)
participant Nested as Nested Sampling Server
participant Upstream as Upstream MCP (optional)
rect rgba(230,240,255,0.4)
note over Client,App: sampling_demo flow
User->>Client: Invoke sampling_demo(topic)
Client->>App: tool:sampling_demo(topic)
App->>App: Configure MCPServerSettings (nested_sampling)
App->>Nested: gen_client.connect/init
App->>Nested: tools/call get_haiku(topic)
Nested-->>App: Tool result (TextContent)
App-->>Client: Extracted text (or "")
end
rect rgba(240,255,230,0.4)
note over Client,Upstream: Direct sampling path
Client->>App: sampling/createMessage(params)
alt Upstream available
App->>Upstream: Forward sampling request
Upstream-->>App: CreateMessageResult
else No upstream
App->>App: SamplingHandler.handle_sampling(params)
App-->>App: Local LLM/human-in-loop (if configured)
end
App-->>Client: CreateMessageResult or ErrorData
end
sequenceDiagram
autonumber
actor User
participant Client as Example Client
participant App as MCPApp
participant Nested as Nested Elicitation Server
participant Upstream as Upstream MCP (optional)
rect rgba(255,240,230,0.4)
note over Client,App: elicitation_demo flow
User->>Client: Invoke elicitation_demo(action)
Client->>App: tool:elicitation_demo(action)
App->>App: Configure MCPServerSettings (nested_elicitation)
App->>Nested: gen_client.connect/init
App->>Nested: tools/call confirm_action(action)
Nested-->>App: Confirmation result text
App-->>Client: Extracted text (or "")
end
rect rgba(240,230,255,0.4)
note over Client,Upstream: Elicitation request
Client->>App: elicitation/create(params)
alt Upstream available
App->>Upstream: Forward via ElicitRequest
Upstream-->>App: Result
else No upstream
App->>App: Local elicitation handler (console callback)
end
App-->>Client: Result or Error
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Suggested reviewers
Poem
Pre-merge checks and finishing touches❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
📜 Recent review detailsConfiguration used: CodeRabbit UI Review profile: CHILL Plan: Pro ⛔ Files ignored due to path filters (1)
📒 Files selected for processing (14)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (4)
src/mcp_agent/server/app_server.py (2)
405-460: Guard “latest upstream session” fallback to prevent cross‑client leakageFalling back to app.context.upstream_session risks routing notifications/requests to the wrong client in multi‑client scenarios. Gate this behind a flag (env/config) and/or verify session_id/correlation before using the fallback. Otherwise restrict to non‑sensitive methods only.
Also applies to: 566-676, 795-826, 910-929
1902-1911: Run cleanup may never execute for workflow_id‑only queriesunstable: When run_id is None, awaiting _unregister_session(run_id) silently no‑ops under the except and leaks mappings. Guard and attempt derivation from status.
Apply:
- try: - state = str(status.get("status", "")).lower() - if state in ("completed", "error", "cancelled"): - try: - await _unregister_session(run_id) - except Exception: - pass - except Exception: - pass + try: + state = str(status.get("status", "")).lower() + if state in ("completed", "error", "cancelled"): + exec_run_id = run_id + if exec_run_id is None: + exec_run_id = status.get("run_id") or (status.get("state") or {}).get("run_id") + if exec_run_id: + await _unregister_session(exec_run_id) + except Exception: + # best-effort cleanup + passsrc/mcp_agent/mcp/mcp_server_registry.py (1)
107-117: start_server: wrong timedelta construction (interpreted as days, not seconds)timedelta(config.read_timeout_seconds) sets days. Should be seconds=... This inflates timeouts by 86400x.
Apply:
- read_timeout_seconds = ( - timedelta(config.read_timeout_seconds) + read_timeout = ( + timedelta(seconds=config.read_timeout_seconds) if config.read_timeout_seconds else None )And pass read_timeout below (stdio/http/sse/ws blocks):
- read_timeout_seconds, + read_timeout,Do this for all four client_session_factory(...) calls in this file.
src/mcp_agent/mcp/gen_client.py (1)
47-56: Action required: fix unusedcontextparameter in connect()connect(..., context: Optional[Context] = None) is unused — the call at src/mcp_agent/mcp/gen_client.py:67 does not pass it, and MCPConnectionManager.get_server does not accept a context arg (src/mcp_agent/mcp/mcp_connection_manager.py:519). Choose one:
- Thread context through persistent path: add
context: Optional[Context] = Noneto MCPConnectionManager.get_server signature and propagate it through launch_server/ServerConnection/transport factory so initialize_server receives it; then update gen_client.connect to call get_server(..., context=context). (Files: src/mcp_agent/mcp/mcp_connection_manager.py, src/mcp_agent/mcp/gen_client.py)- Or remove the unused
contextparameter from gen_client.connect (and update callers).
🧹 Nitpick comments (21)
src/mcp_agent/server/app_server.py (3)
396-403: Idempotency set can grow unbounded_IDEMPOTENCY_KEYS_SEEN never prunes entries. Add TTL/cleanup when a run completes (e.g., in _unregister_session) to avoid memory growth under long‑lived processes.
530-532: Avoid returning raw exception strings to clientsReturning str(e_mapped) may leak internals. Return a generic message and log details server‑side.
Apply:
- return JSONResponse( - {"ok": False, "error": str(e_mapped)}, status_code=500 - ) + logger.exception("Forwarding error (mapped session)", exc_info=e_mapped) + return JSONResponse({"ok": False, "error": "internal_error"}, status_code=500)
868-869: Human prompt logging: include minimal metadata onlyConsider excluding high‑entropy ids (e.g., full idempotency keys) from info logs to reduce log PII/sensitivity. Keep them at debug if needed.
examples/mcp_agent_server/shared/nested_elicitation_server.py (1)
12-23: Happy path only—handle explicit rejection/error casesYou already treat non‑Accepted as “declined.” Consider surfacing a specific message when the elicitation is rejected vs. cancelled/error to aid UX, but not blocking.
src/mcp_agent/mcp/gen_client.py (1)
16-26: Type of client_session_factory mismatches actual call (now passes context=)The factory is annotated for 3 params, but start paths pass context=. Update the type to accept context to avoid mypy/pyright errors.
Apply:
-from typing import AsyncGenerator, Callable, Optional +from typing import AsyncGenerator, Callable, Optional, Protocol +class ClientSessionFactory(Protocol): + def __call__( + self, + read_stream: MemoryObjectReceiveStream, + write_stream: MemoryObjectSendStream, + read_timeout: timedelta | None, + *, + context: Optional[Context] = None, + ) -> ClientSession: ...Then:
- client_session_factory: Callable[ - [MemoryObjectReceiveStream, MemoryObjectSendStream, timedelta | None], - ClientSession, - ] = MCPAgentClientSession, + client_session_factory: ClientSessionFactory = MCPAgentClientSession,examples/mcp_agent_server/shared/nested_sampling_server.py (1)
7-33: Sampling tool: solid integrationGood use of model preferences and content checks. Consider trimming max_tokens to 80 to align with example Temporal workflow for consistent demos, but optional.
src/mcp_agent/mcp/mcp_server_registry.py (3)
291-302: initialize_server: init_hook typingParameter is typed InitHookCallable but default None. Use Optional for correctness.
Apply:
- init_hook: InitHookCallable = None, + init_hook: Optional[InitHookCallable] = None,
334-343: Typo: intialization_callback → initialization_callbackMinor rename for readability.
Apply:
- intialization_callback = ( + initialization_callback = ( init_hook if init_hook is not None else self.init_hooks.get(server_name) ) - - if intialization_callback: + if initialization_callback: logger.info(f"{server_name}: Executing init hook") - intialization_callback(session, config.auth) + initialization_callback(session, config.auth)
111-115: client_session_factory type should accept context kwLike gen_client, update the Callable typing or introduce a Protocol to reflect context kwarg; otherwise type checkers will flag calls.
Apply:
-from typing import Callable, Dict, AsyncGenerator, Optional, TYPE_CHECKING +from typing import Callable, Dict, AsyncGenerator, Optional, TYPE_CHECKING, Protocol +class ClientSessionFactory(Protocol): + def __call__( + self, + read_stream: MemoryObjectReceiveStream, + write_stream: MemoryObjectSendStream, + read_timeout: timedelta | None, + *, + context: Optional["Context"] = None, + ) -> ClientSession: ...And:
- client_session_factory: Callable[ - [MemoryObjectReceiveStream, MemoryObjectSendStream, timedelta | None], - ClientSession, - ] = ClientSession, + client_session_factory: ClientSessionFactory = ClientSession,…and same in initialize_server.
Also applies to: 295-299
examples/mcp_agent_server/asyncio/main.py (4)
128-141: Harden nested server registration: restrict tool surface and avoid repeated mutation.
- Security: limit exposed tools from the nested server.
- Idempotency: avoid re-assigning server settings on every call.
Apply:
- _app.context.config.mcp.servers[nested_name] = MCPServerSettings( + if nested_name not in _app.context.config.mcp.servers: + _app.context.config.mcp.servers[nested_name] = MCPServerSettings( name=nested_name, command="uv", args=["run", nested_path], description="Nested server providing a haiku generator using sampling", - ) + allowed_tools={"get_haiku"}, + )
142-155: Handle nested tool call failures explicitly.Wrap the call to prevent unbound/local exceptions from bubbling to the server, and log context.
- async with gen_client( - nested_name, _app.context.server_registry, context=_app.context - ) as client: - result = await client.call_tool("get_haiku", {"topic": topic}) + try: + async with gen_client( + nested_name, _app.context.server_registry, context=_app.context + ) as client: + result = await client.call_tool("get_haiku", {"topic": topic}) + except Exception as e: + _app.logger.warning("sampling_demo: nested call failed", data={"error": str(e), "server": nested_name}) + return ""
169-181: Apply the same hardening to the elicitation nested server.Limit tools and avoid redefining on every call.
- _app.context.config.mcp.servers[nested_name] = MCPServerSettings( + if nested_name not in _app.context.config.mcp.servers: + _app.context.config.mcp.servers[nested_name] = MCPServerSettings( name=nested_name, command="uv", args=["run", nested_path], description="Nested server demonstrating elicitation", - ) + allowed_tools={"confirm_action"}, + )
182-194: Add defensive error handling for elicitation nested call.Prevents tool failures from crashing the tool invocation.
- async with gen_client( - nested_name, _app.context.server_registry, context=_app.context - ) as client: - # The nested server will call context.session.elicit() internally - result = await client.call_tool("confirm_action", {"action": action}) + try: + async with gen_client( + nested_name, _app.context.server_registry, context=_app.context + ) as client: + # The nested server will call context.session.elicit() internally + result = await client.call_tool("confirm_action", {"action": action}) + except Exception as e: + _app.logger.warning("elicitation_demo: nested call failed", data={"error": str(e), "server": nested_name}) + return ""src/mcp_agent/mcp/sampling_handler.py (2)
159-167: Ensuregenerate_strreceives a string; robustly flatten message content.
generate_strtypically expectsstr. Passing a list may rely on implicit handling. Also handle non-text content deterministically.- messages = [] + messages: list[str] = [] for m in params.messages: if hasattr(m.content, "text") and m.content.text: messages.append(m.content.text) elif hasattr(m.content, "data") and m.content.data: messages.append(str(m.content.data)) else: messages.append(str(m.content)) @@ - text = await llm.generate_str(message=messages, request_params=req_params) + text = await llm.generate_str(message="\n\n".join(messages), request_params=req_params)Also applies to: 178-184
64-79: Use a distinct error code for user rejections.Mapping user-driven rejections to JSON‑RPC
-32603(internal error) makes monitoring harder. Consider a custom application error or-32000range for server-defined errors.- return ErrorData( - code=-32603, message=f"Sampling request rejected by user: {reason}" - ) + return ErrorData(code=-32001, message=f"Sampling request rejected by user: {reason}") @@ - return ErrorData( - code=-32603, message=f"Response rejected by user: {reason}" - ) + return ErrorData(code=-32002, message=f"Response rejected by user: {reason}")examples/mcp_agent_server/temporal/basic_agent_server.py (4)
169-195: Constrain nested sampling server and handle failures.
- Restrict exposed tools for least privilege.
- Avoid re-registering server settings on every run.
- Wrap nested call with error handling.
- app_ctx.config.mcp.servers[nested_name] = MCPServerSettings( + if nested_name not in app_ctx.config.mcp.servers: + app_ctx.config.mcp.servers[nested_name] = MCPServerSettings( name=nested_name, command="uv", args=["run", nested_path], description="Nested server providing a haiku generator using sampling", - ) + allowed_tools={"get_haiku"}, + ) @@ - async with gen_client( - nested_name, app_ctx.server_registry, context=app_ctx - ) as client: - result = await client.call_tool("get_haiku", {"topic": topic}) + try: + async with gen_client( + nested_name, app_ctx.server_registry, context=app_ctx + ) as client: + result = await client.call_tool("get_haiku", {"topic": topic}) + except Exception as e: + app_ctx.app.logger.warning("call_nested_sampling failed", data={"error": str(e), "server": nested_name}) + return ""
197-223: Apply the same constraints and error handling to nested elicitation.- app_ctx.config.mcp.servers[nested_name] = MCPServerSettings( + if nested_name not in app_ctx.config.mcp.servers: + app_ctx.config.mcp.servers[nested_name] = MCPServerSettings( name=nested_name, command="uv", args=["run", nested_path], description="Nested server demonstrating elicitation", - ) + allowed_tools={"confirm_action"}, + ) @@ - async with gen_client( - nested_name, app_ctx.server_registry, context=app_ctx - ) as client: - result = await client.call_tool("confirm_action", {"action": action}) + try: + async with gen_client( + nested_name, app_ctx.server_registry, context=app_ctx + ) as client: + result = await client.call_tool("confirm_action", {"action": action}) + except Exception as e: + app_ctx.app.logger.warning("call_nested_elicitation failed", data={"error": str(e), "server": nested_name}) + return ""
235-253: Guard against missing upstream session in workflow context.If
app.context.upstream_sessionis None, this will raise. Add a clear error or fallback.- direct = await app.context.upstream_session.create_message( + if not app.context.upstream_session: + app.logger.error("SamplingWorkflow: upstream_session not configured; cannot perform direct sampling") + direct_text = "" + else: + direct = await app.context.upstream_session.create_message( messages=[ SamplingMessage( role="user", content=TextContent( type="text", text=f"Write a haiku about {input}." ), ) ], system_prompt="You are a poet.", max_tokens=80, model_preferences=ModelPreferences( hints=[ModelHint(name="gpt-4o-mini")], costPriority=0.1, speedPriority=0.8, intelligencePriority=0.1, ), - ) + ) + try: + direct_text = ( + direct.content.text if isinstance(direct.content, TextContent) else "" + ) + except Exception: + direct_text = ""
286-295: Same guard for elicitation: upstream may be absent.- direct = await app.context.upstream_session.elicit( + if not app.context.upstream_session: + app.logger.error("ElicitationWorkflow: upstream_session not configured; cannot perform direct elicitation") + direct_text = "accepted=" + else: + direct = await app.context.upstream_session.elicit( message=f"Do you want to {input}?", requestedSchema=schema, - ) - direct_text = f"accepted={getattr(direct, 'action', '')}" + ) + direct_text = f"accepted={getattr(direct, 'action', '')}"src/mcp_agent/mcp/mcp_agent_client_session.py (2)
339-340: Prefer structured logging for consistency.Align with the logger’s structured style used elsewhere.
- logger.debug(f"Handling sampling request: {params}") + logger.debug("Handling sampling request", data=params.model_dump())
341-356: Consider logging upstream pass-through failures before returning ErrorData.A warning helps trace intermittent upstream issues.
except Exception as e: - return ErrorData(code=-32603, message=str(e)) + logger.warning("Upstream sampling forwarding failed", data={"error": str(e)}) + return ErrorData(code=-32603, message=str(e))
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (10)
examples/mcp_agent_server/asyncio/main.py(3 hunks)examples/mcp_agent_server/shared/nested_elicitation_server.py(1 hunks)examples/mcp_agent_server/shared/nested_sampling_server.py(1 hunks)examples/mcp_agent_server/temporal/basic_agent_server.py(2 hunks)src/mcp_agent/cli/core/utils.py(1 hunks)src/mcp_agent/mcp/gen_client.py(4 hunks)src/mcp_agent/mcp/mcp_agent_client_session.py(6 hunks)src/mcp_agent/mcp/mcp_server_registry.py(9 hunks)src/mcp_agent/mcp/sampling_handler.py(1 hunks)src/mcp_agent/server/app_server.py(8 hunks)
🧰 Additional context used
🧠 Learnings (2)
📚 Learning: 2025-07-22T18:59:49.368Z
Learnt from: CR
PR: lastmile-ai/mcp-agent#0
File: examples/usecases/reliable_conversation/CLAUDE.md:0-0
Timestamp: 2025-07-22T18:59:49.368Z
Learning: Applies to examples/usecases/reliable_conversation/examples/reliable_conversation/src/**/*.py : Use mcp-agent's Agent abstraction for ALL LLM interactions, including quality evaluation, to ensure consistent tool access, logging, and error handling.
Applied to files:
src/mcp_agent/mcp/sampling_handler.pysrc/mcp_agent/mcp/mcp_server_registry.pyexamples/mcp_agent_server/asyncio/main.py
📚 Learning: 2025-08-28T15:07:10.015Z
Learnt from: saqadri
PR: lastmile-ai/mcp-agent#386
File: src/mcp_agent/mcp/mcp_server_registry.py:110-116
Timestamp: 2025-08-28T15:07:10.015Z
Learning: In MCP server registry methods, when client_session_factory parameters are updated to accept additional context parameters, ensure the type hints match what is actually passed (Context instance vs ServerSession) and that the default factory (MCPAgentClientSession) can handle the number of arguments being passed to avoid TypeError at runtime.
Applied to files:
src/mcp_agent/mcp/mcp_server_registry.pysrc/mcp_agent/mcp/gen_client.py
🧬 Code graph analysis (9)
src/mcp_agent/cli/core/utils.py (2)
src/mcp_agent/config.py (1)
Settings(583-690)src/mcp_agent/cli/mcp_app/api_client.py (1)
MCPApp(22-29)
examples/mcp_agent_server/shared/nested_sampling_server.py (3)
src/mcp_agent/core/context.py (1)
mcp(102-103)src/mcp_agent/executor/temporal/session_proxy.py (1)
create_message(235-271)examples/mcp_agent_server/temporal/basic_agent_server.py (4)
run(49-87)run(140-166)run(230-267)run(275-298)
src/mcp_agent/mcp/sampling_handler.py (6)
src/mcp_agent/core/context.py (2)
mcp(102-103)Context(57-103)src/mcp_agent/core/context_dependent.py (1)
ContextDependent(9-48)src/mcp_agent/workflows/llm/augmented_llm.py (1)
RequestParams(124-174)src/mcp_agent/workflows/llm/llm_selector.py (2)
ModelSelector(98-449)select_best_model(136-284)src/mcp_agent/human_input/types.py (1)
HumanInputRequest(7-29)src/mcp_agent/workflows/factory.py (1)
create_llm(68-86)
src/mcp_agent/mcp/mcp_server_registry.py (2)
src/mcp_agent/app.py (1)
context(149-154)src/mcp_agent/core/context.py (1)
Context(57-103)
examples/mcp_agent_server/shared/nested_elicitation_server.py (3)
src/mcp_agent/core/context.py (1)
mcp(102-103)examples/mcp_agent_server/shared/nested_sampling_server.py (1)
main(35-36)examples/mcp_agent_server/temporal/basic_agent_server.py (5)
main(301-313)run(49-87)run(140-166)run(230-267)run(275-298)
src/mcp_agent/mcp/gen_client.py (4)
src/mcp_agent/core/context.py (2)
mcp(102-103)Context(57-103)src/mcp_agent/logging/logger.py (1)
get_logger(513-541)src/mcp_agent/mcp/mcp_server_registry.py (1)
ServerRegistry(54-395)src/mcp_agent/mcp/mcp_agent_client_session.py (1)
MCPAgentClientSession(73-428)
src/mcp_agent/mcp/mcp_agent_client_session.py (3)
src/mcp_agent/mcp/sampling_handler.py (2)
SamplingHandler(33-184)handle_sampling(39-58)src/mcp_agent/app.py (4)
context(149-154)logger(193-210)upstream_session(173-174)upstream_session(177-178)src/mcp_agent/logging/logger.py (2)
debug(261-269)warning(281-289)
examples/mcp_agent_server/asyncio/main.py (5)
src/mcp_agent/human_input/handler.py (1)
console_input_callback(68-78)src/mcp_agent/elicitation/handler.py (1)
console_elicitation_callback(138-167)src/mcp_agent/mcp/gen_client.py (1)
gen_client(17-44)src/mcp_agent/app.py (4)
config(157-158)tool(715-767)context(149-154)server_registry(161-162)src/mcp_agent/config.py (1)
MCPServerSettings(52-113)
examples/mcp_agent_server/temporal/basic_agent_server.py (5)
src/mcp_agent/human_input/handler.py (1)
console_input_callback(68-78)src/mcp_agent/elicitation/handler.py (1)
console_elicitation_callback(138-167)src/mcp_agent/mcp/gen_client.py (1)
gen_client(17-44)src/mcp_agent/app.py (8)
config(157-158)logger(193-210)MCPApp(42-965)server_registry(161-162)context(149-154)run(397-420)upstream_session(173-174)upstream_session(177-178)src/mcp_agent/config.py (1)
MCPServerSettings(52-113)
🔇 Additional comments (14)
src/mcp_agent/server/app_server.py (3)
608-611: Model dumps: good aliasing and null‑exclusionUsing model_dump(by_alias=True, exclude_none=True) is correct and keeps wire shapes stable.
Also applies to: 632-635, 651-654, 670-673
685-685: 503 for session_not_available is appropriateService‑unavailable fits when no target client session is bound.
937-939: Consistent error envelopeGood: consistent {"error": "..."} with 503 for missing session. Ensure client code treats this as retryable.
Do downstream callers retry 503 on internal routes? If not, add a lightweight retry where applicable.
src/mcp_agent/cli/core/utils.py (2)
30-33: Signature reformat only; OKNo behavioral change; consistent with project style.
41-43: Docstring matches behaviorArgs and resolution order are accurate.
examples/mcp_agent_server/shared/nested_elicitation_server.py (1)
26-31: Entrypoint: OKSimple run guard looks good.
src/mcp_agent/mcp/mcp_server_registry.py (3)
12-13: Imports: Optional used; TYPE_CHECKING guard is correctGood defensive typing without runtime import of Context.
321-328: Propagate context into start_server: goodThe context threading matches session factory signature changes.
107-117: Resolved — pass renamed read_timeout_seconds consistently
All client_session_factory invocations in src/mcp_agent/mcp/mcp_server_registry.py pass read_timeout_seconds (calls at ~lines 154, 208, 248, 273).examples/mcp_agent_server/asyncio/main.py (1)
43-45: Good wiring of local human-input and elicitation callbacks.This correctly enables console-based approval/elicitation when no upstream is present and allows upstream pass-through when connected.
examples/mcp_agent_server/temporal/basic_agent_server.py (1)
32-38: Nice: console callbacks wired for Temporal example as well.This keeps behavior consistent between asyncio and Temporal demos.
src/mcp_agent/mcp/mcp_agent_client_session.py (3)
120-121: Good: centralize local sampling via SamplingHandler.This keeps upstream pass-through and local handling consistent with the rest of the stack.
370-384: Nice fallback strategy for elicitation.Prefer upstream; on failure, degrade locally with a warning—good UX.
391-399: Don't rely on injecting extras without verifying the SDK version — prefer passing params unchanged.
Elicitation models set model_config = ConfigDict(extra="allow") so model_copy(update={"server_name": ...}) is allowed, but a known issue can cause top-level RequestParams to ignore extras in some SDK versions; confirm your installed MCP Python SDK version or replace with:
elicitation_response = await self.context.elicitation_handler(params) (or log server_name).
| llm = create_llm( | ||
| agent_name="sampling", | ||
| server_names=[], | ||
| instruction=None, | ||
| provider=model_info.provider, | ||
| model=model_info.name, | ||
| request_params=None, | ||
| context=None, | ||
| ) |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Pass app Context to LLM factory to preserve tracing, logging, and settings.
Using context=None breaks the preference to run all LLM interactions under the app’s context (tracing, token counting, credentials). This can lead to inconsistent behavior vs. the rest of the system.
- llm = create_llm(
+ llm = create_llm(
agent_name="sampling",
server_names=[],
instruction=None,
provider=model_info.provider,
model=model_info.name,
request_params=None,
- context=None,
+ context=self.context,
)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| llm = create_llm( | |
| agent_name="sampling", | |
| server_names=[], | |
| instruction=None, | |
| provider=model_info.provider, | |
| model=model_info.name, | |
| request_params=None, | |
| context=None, | |
| ) | |
| llm = create_llm( | |
| agent_name="sampling", | |
| server_names=[], | |
| instruction=None, | |
| provider=model_info.provider, | |
| model=model_info.name, | |
| request_params=None, | |
| context=self.context, | |
| ) |
🤖 Prompt for AI Agents
In src/mcp_agent/mcp/sampling_handler.py around lines 149-157 the LLM is created
with context=None which breaks app-level tracing, logging, token counting and
credential propagation; change the call to pass the active application context
object (the same app/request context used elsewhere in this module — e.g. pass
the existing `app_context`/`ctx`/`request.app` variable) instead of None, and if
that variable isn’t in scope, obtain or accept the current app context and
thread it into create_llm so all LLM interactions inherit tracing, logging and
settings.
There was a problem hiding this comment.
Actionable comments posted: 0
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (3)
src/mcp_agent/mcp/gen_client.py (1)
47-71: Propagate Context through persistent-connection path & update client_session_factory typing
- connect() accepts a context but does not forward it — update src/mcp_agent/mcp/gen_client.py:connect to pass context to connection_manager.get_server(..., context=context).
- MCPConnectionManager must carry/forward context: change get_server / launch_server signatures and ServerConnection to accept Optional[Context] (store it on the ServerConnection) and ensure create_session calls the factory with the context (currently it calls _client_session_factory(read_stream, send_stream, read_timeout) — include context).
- Update client_session_factory typing everywhere (gen_client, mcp_connection_manager, mcp_server_registry) to include Optional[Context] as the final parameter (e.g. Callable[[MemoryObjectReceiveStream, MemoryObjectSendStream, timedelta | None, Optional[Context]], ClientSession] or use a Protocol/ParamSpec).
- MCPAgentClientSession.init already accepts context (src/mcp_agent/mcp/mcp_agent_client_session.py), so the changes above are necessary for proper propagation and to keep types accurate.
- Update all get_server call sites (e.g., src/mcp_agent/mcp/mcp_aggregator.py, src/mcp_agent/agents/agent.py) to forward context or document intended behavior if context should not propagate.
src/mcp_agent/mcp/mcp_server_registry.py (2)
135-139: Bug: read_timeout_seconds interpreted as days, not secondstimedelta(config.read_timeout_seconds) treats the value as days. This will massively increase timeouts and cause hangs.
Apply this fix:
- read_timeout_seconds = ( - timedelta(config.read_timeout_seconds) - if config.read_timeout_seconds - else None - ) + read_timeout_seconds = ( + timedelta(seconds=config.read_timeout_seconds) + if config.read_timeout_seconds + else None + )
111-117: Default factory + type hint mismatch will raise TypeError when passing contextstart_server passes context= to client_session_factory, but the default is ClientSession which does not accept it. Same pattern in initialize_server. This will crash when callers don’t override the factory.
Two safe options (prefer A):
A) Make MCPAgentClientSession the default and widen the type to accept kwargs.
- client_session_factory: Callable[ - [MemoryObjectReceiveStream, MemoryObjectSendStream, timedelta | None], - ClientSession, - ] = ClientSession, + client_session_factory: Callable[..., ClientSession] = MCPAgentClientSession,Do the same in initialize_server (lines 295-299).
B) Add a small helper that tries with context= and falls back if unsupported, then use it at all call sites:
+ def _build_session(factory, rs, ws, rt, ctx): + try: + return factory(rs, ws, rt, context=ctx) + except TypeError: + return factory(rs, ws, rt)And replace each client_session_factory(..., context=context) with _build_session(client_session_factory, ..., context).
🧹 Nitpick comments (18)
src/mcp_agent/server/app_server.py (1)
685-685: Inconsistent error envelope across endpointsSome endpoints return {"ok": False, "error": ...} while others return {"error": ...} only. Consider standardizing the error shape for clients.
Apply where feasible:
- return JSONResponse({"error": "session_not_available"}, status_code=503) + return JSONResponse({"ok": False, "error": "session_not_available"}, status_code=503)Also applies to: 937-939
examples/mcp_agent_server/asyncio/main.py (2)
118-141: Initialize servers dict before assignment to avoid NoneType subscripting_config.mcp or _config.mcp.servers may be None in some runs. Guard before writing nested server entries.
-from mcp_agent.mcp.gen_client import gen_client -from mcp_agent.config import MCPServerSettings +from mcp_agent.mcp.gen_client import gen_client +from mcp_agent.config import MCPServerSettings, MCPSettings @@ - _app.context.config.mcp.servers[nested_name] = MCPServerSettings( + if _app.context.config.mcp is None: + _app.context.config.mcp = MCPSettings() + if _app.context.config.mcp.servers is None: + _app.context.config.mcp.servers = {} + _app.context.config.mcp.servers[nested_name] = MCPServerSettings( @@ - _app.context.config.mcp.servers[nested_name] = MCPServerSettings( + if _app.context.config.mcp is None: + _app.context.config.mcp = MCPSettings() + if _app.context.config.mcp.servers is None: + _app.context.config.mcp.servers = {} + _app.context.config.mcp.servers[nested_name] = MCPServerSettings(Also applies to: 157-181
149-155: Defensive content extractionCallToolResult.content items may not always be text; guard attribute access to avoid AttributeError.
- try: - if result.content and len(result.content) > 0: - return result.content[0].text or "" + try: + if result.content and len(result.content) > 0: + first = result.content[0] + return getattr(first, "text", "") or "" except Exception: pass return ""Also applies to: 188-194
src/mcp_agent/mcp/mcp_server_registry.py (6)
195-203: Inconsistent timeout units for SSE vs Streamable HTTPFor Streamable HTTP, sse_read_timeout is a timedelta; for SSE, it’s an int. If the SSE client expects the same unit as Streamable HTTP, this inconsistency can cause subtle bugs.
If SSE also accepts timedelta, normalize:
- if config.read_timeout_seconds: - kwargs["sse_read_timeout"] = config.read_timeout_seconds + if config.read_timeout_seconds: + kwargs["sse_read_timeout"] = timedelta(seconds=config.read_timeout_seconds)If not, leave as-is.
Also applies to: 237-242
158-159: All transport paths pass context=; ties to the factory issueThese calls assume the factory accepts context=. They’ll fail if Option A (default factory) or Option B (helper) above isn’t applied consistently.
After fixing the factory mismatch, no change needed here.
Also applies to: 212-213, 252-253, 278-278
301-302: initialize_server: mirror factory/typing fix and forward contextGood forwarding of context to start_server. Ensure the factory default/typing fix is also applied here to avoid the same TypeError.
Also applies to: 323-329
334-343: Typo: intialization_callback variableSpelling nit; not functional, but worth fixing for readability (initialization_callback).
- intialization_callback = ( + initialization_callback = ( init_hook if init_hook is not None else self.init_hooks.get(server_name) ) - if intialization_callback: + if initialization_callback: logger.info(f"{server_name}: Executing init hook") - intialization_callback(session, config.auth) + initialization_callback(session, config.auth)
299-300: Optional type for init_hookinit_hook is annotated as InitHookCallable but defaulted to None. Make it Optional[InitHookCallable] for accuracy.
- init_hook: InitHookCallable = None, + init_hook: Optional[InitHookCallable] = None,
68-90: Config mutability assumptionServerRegistry stores config.mcp.servers by reference. If a different config dict is later assigned (not mutated), registry won’t see it. If your examples mutate the same dict (they do), you’re fine; otherwise document/guard this.
src/mcp_agent/mcp/mcp_agent_client_session.py (2)
369-399: Elicitation: injecting server_name into MCPElicitRequestParams may violate schemaparams.model_copy(update={"server_name": server_name}) will fail if ElicitRequestParams forbids extra fields (common in pydantic models). Consider attaching this info via a supported field (e.g., params.metadata if present) or by passing server_name separately to the local handler.
Simplest safe change (fallback to metadata when server_name isn’t a valid field):
- elicitation_request = params.model_copy(update={"server_name": server_name}) + elicitation_request = params + if server_name: + try: + elicitation_request = params.model_copy(update={"server_name": server_name}) + except Exception: + # Fallback: stuff it into a generic metadata map if available + meta = getattr(params, "metadata", None) + if isinstance(meta, dict): + meta = {**meta, "server_name": server_name} + elicitation_request = params.model_copy(update={"metadata": meta})If ElicitRequestParams has no metadata, consider omitting server_name or logging it instead.
291-295: Notification logging at info level may be noisy_received_notification logs at info. Consider debug to reduce log volume in busy sessions.
examples/mcp_agent_server/temporal/basic_agent_server.py (5)
178-187: Ensure nested server registration is visible to ServerRegistryYou add to app_ctx.config.mcp.servers, but gen_client looks up the server via app_ctx.server_registry. If the registry captured a different dict at construction, this may raise “Server 'nested_sampling' not found”.
Safer to update both:
app_ctx.config.mcp.servers[nested_name] = MCPServerSettings( name=nested_name, command="uv", args=["run", nested_path], description="Nested server providing a haiku generator using sampling", ) + # Ensure registry sees the new server immediately + if app_ctx.server_registry and nested_name not in app_ctx.server_registry.registry: + app_ctx.server_registry.registry[nested_name] = app_ctx.config.mcp.servers[nested_name]Repeat similarly for nested_elicitation.
185-194: Tool result parsing: guard content typingResult content can be a list of union types. Current try/except is fine; optionally check isinstance(item, TextContent) for clarity.
235-267: SamplingWorkflow assumes upstream_session is setDirect call to app.context.upstream_session without None check will fail outside the proxy environment. Add a guard to degrade gracefully in local runs.
- direct = await app.context.upstream_session.create_message( + if not app.context.upstream_session: + app.logger.warning("No upstream_session; skipping direct sampling.") + direct_text = "" + else: + direct = await app.context.upstream_session.create_message( ... - try: - direct_text = ( - direct.content.text if isinstance(direct.content, TextContent) else "" - ) - except Exception: - direct_text = "" + try: + direct_text = ( + direct.content.text if isinstance(direct.content, TextContent) else "" + ) + except Exception: + direct_text = ""
281-289: ElicitationWorkflow assumes upstream_session is setMirror the guard used for sampling to avoid crashes when upstream isn’t configured.
180-181: Operational note: depends on 'uv' being availableThe nested servers are launched via uv. Document this dependency in the example README, or detect and emit a helpful error if missing.
Also applies to: 208-209
src/mcp_agent/mcp/sampling_handler.py (2)
140-146: Require modelPreferences for local samplingReasonable constraint to avoid recursive/ambiguous model selection. Consider a clearer error code (e.g., -32602 Invalid params) rather than -32603 Internal error.
- raise ToolError("Model preferences must be provided for sampling requests") + raise ToolError("Model preferences must be provided for sampling requests") # consider mapping to -32602Or map the caught exception to ErrorData(code=-32602, message=...).
159-167: Messages extraction: handle None and non-text contentIf params.messages is None or empty, llm.generate_str([]) may produce weak outputs. Consider defaulting to [""] or requiring at least one user message.
- messages = [] - for m in params.messages: + messages = [] + for m in (params.messages or []): ... + if not messages: + messages = [""]
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (10)
examples/mcp_agent_server/asyncio/main.py(3 hunks)examples/mcp_agent_server/shared/nested_elicitation_server.py(1 hunks)examples/mcp_agent_server/shared/nested_sampling_server.py(1 hunks)examples/mcp_agent_server/temporal/basic_agent_server.py(2 hunks)src/mcp_agent/cli/core/utils.py(1 hunks)src/mcp_agent/mcp/gen_client.py(4 hunks)src/mcp_agent/mcp/mcp_agent_client_session.py(6 hunks)src/mcp_agent/mcp/mcp_server_registry.py(9 hunks)src/mcp_agent/mcp/sampling_handler.py(1 hunks)src/mcp_agent/server/app_server.py(8 hunks)
🧰 Additional context used
🧠 Learnings (2)
📚 Learning: 2025-08-28T15:07:10.015Z
Learnt from: saqadri
PR: lastmile-ai/mcp-agent#386
File: src/mcp_agent/mcp/mcp_server_registry.py:110-116
Timestamp: 2025-08-28T15:07:10.015Z
Learning: In MCP server registry methods, when client_session_factory parameters are updated to accept additional context parameters, ensure the type hints match what is actually passed (Context instance vs ServerSession) and that the default factory (MCPAgentClientSession) can handle the number of arguments being passed to avoid TypeError at runtime.
Applied to files:
src/mcp_agent/mcp/gen_client.pysrc/mcp_agent/mcp/mcp_server_registry.pysrc/mcp_agent/mcp/mcp_agent_client_session.py
📚 Learning: 2025-07-22T18:59:49.368Z
Learnt from: CR
PR: lastmile-ai/mcp-agent#0
File: examples/usecases/reliable_conversation/CLAUDE.md:0-0
Timestamp: 2025-07-22T18:59:49.368Z
Learning: Applies to examples/usecases/reliable_conversation/examples/reliable_conversation/src/**/*.py : Use mcp-agent's Agent abstraction for ALL LLM interactions, including quality evaluation, to ensure consistent tool access, logging, and error handling.
Applied to files:
src/mcp_agent/server/app_server.pyexamples/mcp_agent_server/asyncio/main.pysrc/mcp_agent/mcp/mcp_server_registry.py
🧬 Code graph analysis (9)
src/mcp_agent/cli/core/utils.py (2)
src/mcp_agent/config.py (1)
Settings(583-690)src/mcp_agent/cli/mcp_app/api_client.py (1)
MCPApp(22-29)
examples/mcp_agent_server/shared/nested_elicitation_server.py (3)
src/mcp_agent/core/context.py (1)
mcp(102-103)examples/mcp_agent_server/shared/nested_sampling_server.py (1)
main(35-36)examples/mcp_agent_server/temporal/basic_agent_server.py (5)
main(301-313)run(49-87)run(140-166)run(230-267)run(275-298)
src/mcp_agent/mcp/sampling_handler.py (7)
src/mcp_agent/core/context.py (2)
mcp(102-103)Context(57-103)src/mcp_agent/core/context_dependent.py (1)
ContextDependent(9-48)src/mcp_agent/workflows/llm/augmented_llm.py (1)
RequestParams(124-174)src/mcp_agent/workflows/llm/llm_selector.py (2)
ModelSelector(98-449)select_best_model(136-284)src/mcp_agent/mcp/mcp_agent_client_session.py (1)
send_request(149-231)src/mcp_agent/human_input/types.py (1)
HumanInputRequest(7-29)src/mcp_agent/workflows/factory.py (1)
create_llm(68-86)
src/mcp_agent/mcp/gen_client.py (3)
src/mcp_agent/core/context.py (2)
mcp(102-103)Context(57-103)src/mcp_agent/mcp/mcp_server_registry.py (1)
ServerRegistry(54-395)src/mcp_agent/mcp/mcp_agent_client_session.py (1)
MCPAgentClientSession(73-428)
examples/mcp_agent_server/shared/nested_sampling_server.py (4)
src/mcp_agent/core/context.py (1)
mcp(102-103)src/mcp_agent/executor/temporal/session_proxy.py (1)
create_message(235-271)examples/mcp_agent_server/shared/nested_elicitation_server.py (1)
main(26-27)examples/mcp_agent_server/temporal/basic_agent_server.py (5)
main(301-313)run(49-87)run(140-166)run(230-267)run(275-298)
examples/mcp_agent_server/asyncio/main.py (4)
src/mcp_agent/human_input/handler.py (1)
console_input_callback(68-78)src/mcp_agent/elicitation/handler.py (1)
console_elicitation_callback(138-167)src/mcp_agent/mcp/gen_client.py (1)
gen_client(17-44)src/mcp_agent/config.py (1)
MCPServerSettings(52-113)
src/mcp_agent/mcp/mcp_server_registry.py (2)
src/mcp_agent/app.py (1)
context(149-154)src/mcp_agent/core/context.py (1)
Context(57-103)
examples/mcp_agent_server/temporal/basic_agent_server.py (6)
src/mcp_agent/human_input/handler.py (1)
console_input_callback(68-78)src/mcp_agent/elicitation/handler.py (1)
console_elicitation_callback(138-167)src/mcp_agent/core/context.py (2)
mcp(102-103)Context(57-103)src/mcp_agent/mcp/gen_client.py (1)
gen_client(17-44)src/mcp_agent/config.py (1)
MCPServerSettings(52-113)src/mcp_agent/executor/temporal/session_proxy.py (2)
create_message(235-271)elicit(273-286)
src/mcp_agent/mcp/mcp_agent_client_session.py (2)
src/mcp_agent/mcp/sampling_handler.py (2)
SamplingHandler(33-184)handle_sampling(39-58)src/mcp_agent/app.py (4)
context(149-154)logger(193-210)upstream_session(173-174)upstream_session(177-178)
🔇 Additional comments (10)
src/mcp_agent/cli/core/utils.py (1)
30-33: Signature reformatted only — LGTMNo behavioral change; callers remain unaffected.
src/mcp_agent/server/app_server.py (1)
608-611: JSONResponse wrapping style changes — OKMulti-line model_dump calls and error response formatting are cosmetic; behavior unchanged.
Also applies to: 632-635, 651-654, 670-673, 530-532
examples/mcp_agent_server/shared/nested_elicitation_server.py (1)
12-23: Elicitation flow is soundUses elicit_with_validation with a Pydantic schema and cleanly handles acceptance/decline. Good demo quality.
examples/mcp_agent_server/shared/nested_sampling_server.py (1)
7-33: Sampling tool implementation looks correctAPI usage matches CreateMessageResult/TextContent patterns; reasonable defaults (tokens/temperature) and model preferences.
src/mcp_agent/mcp/mcp_agent_client_session.py (3)
119-121: Initialize SamplingHandler with session contextLooks good; ensures local sampling has access to runtime Context. No issues.
339-359: Sampling: upstream pass-through with local fallback via SamplingHandlerSolid approach. Minimal surface and clear behavior.
Please confirm that self.context is always available here; ContextDependent will fall back to global context if not explicitly provided. Given registry now threads context, we should be safe.
149-232: Tracing propagation and attribute recordingGood span attributes and trace context injection into params.meta. This will help correlate downstream requests.
examples/mcp_agent_server/temporal/basic_agent_server.py (1)
32-38: Thread local callbacks into appNice: wiring human_input and elicitation callbacks enables local flows during development.
src/mcp_agent/mcp/sampling_handler.py (2)
39-59: Upstream pass-through with error captureGood: forwards CreateMessageRequest to upstream and wraps exceptions as ErrorData. Clear behavior.
168-176: Local LLM call: good minimal request shapingSolid use of RequestParams and model selection; avoids history and tool call recursion. LGTM.
Also applies to: 178-184
1 participant
With this change we should get all server -> client MCP capabilities fully supported in MCP-agent, including in distributed systems environments (i.e. Temporal). The communication will work in workflow context, in activity context, and outside workflow context.
Summary by CodeRabbit
New Features
Improvements
Documentation