Use thread_ids in TimeoutLimitMiddleware

Author: mateusz834

splunklib/ai/hooks.py

@@ -182,30 +182,34 @@ class TimeoutLimitMiddleware(AgentMiddleware):
     """
 
     _seconds: float
-    _deadline: float | None
+    _deadline_per_thread_id: dict[str, Any]
 
     def __init__(self, seconds: float) -> None:
         self._seconds = seconds
-        self._deadline = None
+        self._deadline_per_thread_id = {}
 
     @override
     async def agent_middleware(
         self,
         request: AgentRequest,
         handler: AgentMiddlewareHandler,
     ) -> AgentResponse[Any | None]:
-        # WARN: this might not work with agents handling
-        # different threads at the same time.
-        self._deadline = monotonic() + self._seconds
-        return await handler(request)
+        try:
+            # Agent loop starting.
+            self._deadline_per_thread_id[request.thread_id] = (
+                monotonic() + self._seconds
+            )
+            return await handler(request)
+        finally:
+            del self._deadline_per_thread_id[request.thread_id]  # don't leak memory
 
     @override
     async def model_middleware(
         self,
         request: ModelRequest,
         handler: ModelMiddlewareHandler,
     ) -> ModelResponse:
-        if self._deadline is not None and monotonic() >= self._deadline:
+        if monotonic() >= self._deadline_per_thread_id[request.state.thread_id]:
             raise TimeoutExceededException(timeout_seconds=self._seconds)
         return await handler(request)
 

tests/unit/ai/test_default_limits.py

@@ -28,7 +28,13 @@
     TokenLimitMiddleware,
 )
 from splunklib.ai.messages import AIMessage, AgentResponse
-from splunklib.ai.middleware import AgentMiddleware, AgentRequest, AgentState, ModelRequest, ModelResponse
+from splunklib.ai.middleware import (
+    AgentMiddleware,
+    AgentRequest,
+    AgentState,
+    ModelRequest,
+    ModelResponse,
+)
 from splunklib.ai.model import OpenAIModel
 from splunklib.client import Service
 
@@ -125,22 +131,22 @@ async def test_deadline_reset_on_each_invoke(self) -> None:
         request = _make_agent_request()
 
         await mw.agent_middleware(request, _noop_agent_handler)
-        first_deadline = mw._deadline  # pyright: ignore[reportPrivateUsage]
+        first_deadline = mw._deadline_per_thread_id["foo"]  # pyright: ignore[reportPrivateUsage]
 
         await mw.agent_middleware(request, _noop_agent_handler)
-        second_deadline = mw._deadline  # pyright: ignore[reportPrivateUsage]
+        second_deadline = mw._deadline_per_thread_id["foo"]  # pyright: ignore[reportPrivateUsage]
 
         assert first_deadline is not None
         assert second_deadline is not None
         assert second_deadline >= first_deadline
 
     async def test_deadline_is_none_before_first_invoke(self) -> None:
         mw = TimeoutLimitMiddleware(60.0)
-        assert mw._deadline is None  # pyright: ignore[reportPrivateUsage]
+        assert mw._deadline_per_thread_id["foo"] is None  # pyright: ignore[reportPrivateUsage]
 
     async def test_timeout_fires_when_deadline_exceeded(self) -> None:
         mw = TimeoutLimitMiddleware(60.0)
-        mw._deadline = monotonic() - 1.0  # pyright: ignore[reportPrivateUsage]  # already in the past
+        mw._deadline_per_thread_id["foo"] = monotonic() - 1.0  # pyright: ignore[reportPrivateUsage]  # already in the past
 
         state = AgentState(messages=[], total_steps=0, token_count=0, thread_id="foo")
         request = ModelRequest(system_message="", state=state)