From 161bd46aac6557a42a6d63f59dc35d44db31a1e6 Mon Sep 17 00:00:00 2001 From: MyNameisPI Date: Fri, 5 Dec 2025 15:27:41 +0800 Subject: [PATCH] Fix: Add middleware support for resource context injection(#431) - Extract common logic from on_call_tool to _inject_unity_instance method. - Add on_read_resource to inject unity_instance into resource contexts. - Fix resources unable to get unity_instance from context after set_active_instance. --- Server/src/transport/unity_instance_middleware.py | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/Server/src/transport/unity_instance_middleware.py b/Server/src/transport/unity_instance_middleware.py index 8f3a0209e..1df1984e1 100644 --- a/Server/src/transport/unity_instance_middleware.py +++ b/Server/src/transport/unity_instance_middleware.py @@ -42,7 +42,7 @@ class UnityInstanceMiddleware(Middleware): Middleware that manages per-session Unity instance selection. Stores active instance per session_id and injects it into request state - for all tool calls. + for all tool and resource calls. """ def __init__(self): @@ -83,8 +83,8 @@ def clear_active_instance(self, ctx) -> None: with self._lock: self._active_by_key.pop(key, None) - async def on_call_tool(self, context: MiddlewareContext, call_next): - """Inject active Unity instance into tool context if available.""" + async def _inject_unity_instance(self, context: MiddlewareContext) -> None: + """Inject active Unity instance into context if available.""" ctx = context.fastmcp_context active_instance = self.get_active_instance(ctx) @@ -129,4 +129,13 @@ async def on_call_tool(self, context: MiddlewareContext, call_next): ctx.set_state("unity_instance", active_instance) if session_id is not None: ctx.set_state("unity_session_id", session_id) + + async def on_call_tool(self, context: MiddlewareContext, call_next): + """Inject active Unity instance into tool context if available.""" + await self._inject_unity_instance(context) + return await call_next(context) + + async def on_read_resource(self, context: MiddlewareContext, call_next): + """Inject active Unity instance into resource context if available.""" + await self._inject_unity_instance(context) return await call_next(context)