Fix: Add middleware support for resource context injection(#431) (#432)

- 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.

Co-authored-by: MyNameisPI <eqjbc@qq.com>
main
MyNameisPI 2025-12-06 02:33:40 +08:00 committed by GitHub
parent 2c65c76ea3
commit b09e48a395
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 12 additions and 3 deletions

View File

@ -42,7 +42,7 @@ class UnityInstanceMiddleware(Middleware):
Middleware that manages per-session Unity instance selection. Middleware that manages per-session Unity instance selection.
Stores active instance per session_id and injects it into request state 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): def __init__(self):
@ -83,8 +83,8 @@ class UnityInstanceMiddleware(Middleware):
with self._lock: with self._lock:
self._active_by_key.pop(key, None) self._active_by_key.pop(key, None)
async def on_call_tool(self, context: MiddlewareContext, call_next): async def _inject_unity_instance(self, context: MiddlewareContext) -> None:
"""Inject active Unity instance into tool context if available.""" """Inject active Unity instance into context if available."""
ctx = context.fastmcp_context ctx = context.fastmcp_context
active_instance = self.get_active_instance(ctx) active_instance = self.get_active_instance(ctx)
@ -129,4 +129,13 @@ class UnityInstanceMiddleware(Middleware):
ctx.set_state("unity_instance", active_instance) ctx.set_state("unity_instance", active_instance)
if session_id is not None: if session_id is not None:
ctx.set_state("unity_session_id", session_id) 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) return await call_next(context)