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, float]
 
     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
 
@@ -103,7 +109,11 @@ def test_user_timeout_limit_suppresses_default(self) -> None:
 
     def test_all_user_limits_suppress_all_defaults(self) -> None:
         agent = _make_agent(
-            middleware=[TokenLimitMiddleware(50_000), StepLimitMiddleware(10), TimeoutLimitMiddleware(30.0)]
+            middleware=[
+                TokenLimitMiddleware(50_000),
+                StepLimitMiddleware(10),
+                TimeoutLimitMiddleware(30.0),
+            ]
         )
         mw = list(agent.middleware or [])
         assert len([m for m in mw if isinstance(m, TokenLimitMiddleware)]) == 1
@@ -124,23 +134,34 @@ async def test_deadline_reset_on_each_invoke(self) -> None:
         mw = TimeoutLimitMiddleware(60.0)
         request = _make_agent_request()
 
-        await mw.agent_middleware(request, _noop_agent_handler)
-        first_deadline = mw._deadline  # pyright: ignore[reportPrivateUsage]
+        first_deadline: float | None = None
+        second_deadline: float | None = None
+
+        async def _first_agent_handler(_request: AgentRequest) -> AgentResponse[None]:
+            nonlocal first_deadline
+            first_deadline = mw._deadline_per_thread_id["foo"]  # pyright: ignore[reportPrivateUsage]
+            return AgentResponse(messages=[], structured_output=None)
+
+        async def _second_agent_handler(_request: AgentRequest) -> AgentResponse[None]:
+            nonlocal second_deadline
+            second_deadline = mw._deadline_per_thread_id["foo"]  # pyright: ignore[reportPrivateUsage]
+            return AgentResponse(messages=[], structured_output=None)
+
+        await mw.agent_middleware(request, _first_agent_handler)
 
-        await mw.agent_middleware(request, _noop_agent_handler)
-        second_deadline = mw._deadline  # pyright: ignore[reportPrivateUsage]
+        await mw.agent_middleware(request, _second_agent_handler)
 
         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.get("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)
@@ -153,17 +174,29 @@ class TestTokenLimitMiddleware(unittest.IsolatedAsyncioTestCase):
     async def test_raises_when_token_count_in_request_exceeds_limit(self) -> None:
         mw = TokenLimitMiddleware(200)
 
-        await mw.model_middleware(_make_model_request(token_count=100), _noop_model_handler)
-        await mw.model_middleware(_make_model_request(token_count=199), _noop_model_handler)
+        await mw.model_middleware(
+            _make_model_request(token_count=100), _noop_model_handler
+        )
+        await mw.model_middleware(
+            _make_model_request(token_count=199), _noop_model_handler
+        )
         with self.assertRaises(TokenLimitExceededException):
-            await mw.model_middleware(_make_model_request(token_count=200), _noop_model_handler)
+            await mw.model_middleware(
+                _make_model_request(token_count=200), _noop_model_handler
+            )
 
 
 class TestStepLimitMiddleware(unittest.IsolatedAsyncioTestCase):
     async def test_raises_when_steps_in_request_reach_limit(self) -> None:
         mw = StepLimitMiddleware(3)
 
-        await mw.model_middleware(_make_model_request(total_steps=1), _noop_model_handler)
-        await mw.model_middleware(_make_model_request(total_steps=2), _noop_model_handler)
+        await mw.model_middleware(
+            _make_model_request(total_steps=1), _noop_model_handler
+        )
+        await mw.model_middleware(
+            _make_model_request(total_steps=2), _noop_model_handler
+        )
         with self.assertRaises(StepsLimitExceededException):
-            await mw.model_middleware(_make_model_request(total_steps=3), _noop_model_handler)
+            await mw.model_middleware(
+                _make_model_request(total_steps=3), _noop_model_handler
+            )