- 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
parent
2c65c76ea3
commit
b09e48a395
|
|
@ -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 @@ class UnityInstanceMiddleware(Middleware):
|
|||
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 @@ class UnityInstanceMiddleware(Middleware):
|
|||
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)
|
||||
|
|
|
|||
Loading…
Reference in New Issue