fix(tracing): use WeakKeyDictionary for per-loop client cache

Addresses Greptile review feedback on PR #362.  The original
`dict[int, AsyncSGPClient]` cache used `id(asyncio.get_running_loop())`
as the key.  In CPython `id()` returns a memory address, and once a
loop is garbage-collected its address can be assigned to a new loop —
a fresh loop hashing to a stale entry would receive a client whose
httpx.AsyncClient was bound to the dead loop, reintroducing the
"bound to a different event loop" error this PR was built to prevent.

Switching the cache to `weakref.WeakKeyDictionary` keyed on the loop
object itself fixes the bug: the entry is evicted automatically when
the loop is collected, so id() recycling can't cause stale-client
reuse.  Multi-loop caching benefit is preserved (better than the
single-slot pattern in TracingModule for agents that bounce between
loops).

Same fix applied to AgentexAsyncTracingProcessor.  Added a regression
test verifying the cache evicts a closed/dropped loop's entry after
gc.collect().

Author: smoreinis

src/agentex/lib/core/tracing/processors/agentex_tracing_processor.py

@@ -1,4 +1,5 @@
 import asyncio
+import weakref
 from typing import TYPE_CHECKING, Any, Dict, override
 
 from agentex import Agentex
@@ -74,8 +75,13 @@ def __init__(self, config: AgentexTracingProcessorConfig):  # noqa: ARG002
         # Per-event-loop client cache.  httpx.AsyncClient is bound to the
         # loop that created it, so in sync-ACP / streaming contexts (where
         # the active loop can change between requests) we keep one client
-        # per loop instead of disabling keepalive entirely.
-        self._clients_by_loop_id: dict[int, "AsyncAgentex"] = {}
+        # per loop instead of disabling keepalive entirely.  The cache is a
+        # WeakKeyDictionary so a GC'd loop and its client are evicted
+        # automatically — using id() as a key would reuse entries when
+        # CPython recycles a freed loop's memory address.
+        self._clients_by_loop: weakref.WeakKeyDictionary[
+            asyncio.AbstractEventLoop, "AsyncAgentex"
+        ] = weakref.WeakKeyDictionary()
 
     def _build_client(self) -> "AsyncAgentex":
         import httpx
@@ -91,13 +97,13 @@ def _build_client(self) -> "AsyncAgentex":
     @property
     def client(self) -> "AsyncAgentex":
         try:
-            loop_id = id(asyncio.get_running_loop())
+            loop = asyncio.get_running_loop()
         except RuntimeError:
             return self._build_client()
-        client = self._clients_by_loop_id.get(loop_id)
+        client = self._clients_by_loop.get(loop)
         if client is None:
             client = self._build_client()
-            self._clients_by_loop_id[loop_id] = client
+            self._clients_by_loop[loop] = client
         return client
 
     # TODO(AGX1-199): Add batch create/update endpoints to Agentex API and use

src/agentex/lib/core/tracing/processors/sgp_tracing_processor.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import asyncio
+import weakref
 from typing import cast, override
 
 import scale_gp_beta.lib.tracing as tracing
@@ -96,11 +97,14 @@ def __init__(self, config: SGPTracingProcessorConfig):
         self._config = config
         # Per-event-loop client cache.  httpx.AsyncClient ties its connection
         # pool to the loop it was created on; in sync-ACP / streaming contexts
-        # the active loop can change between requests.  Keying by loop id lets
+        # the active loop can change between requests.  Caching per loop lets
         # us keep keepalive on within each loop while staying safe across
-        # loops.  The construction can also happen at module import time when
-        # no loop is running, so we have to defer it until the first call.
-        self._clients_by_loop_id: dict[int, AsyncSGPClient] = {}
+        # loops.  The cache is a WeakKeyDictionary so a GC'd loop and its
+        # client are evicted automatically — using id() as a key would reuse
+        # entries when CPython recycles a freed loop's memory address.
+        self._clients_by_loop: weakref.WeakKeyDictionary[
+            asyncio.AbstractEventLoop, AsyncSGPClient
+        ] = weakref.WeakKeyDictionary()
         self.env_vars = EnvironmentVariables.refresh()
 
     def _build_client(self) -> AsyncSGPClient:
@@ -124,15 +128,15 @@ def _get_client(self) -> AsyncSGPClient | None:
         if self.disabled:
             return None
         try:
-            loop_id = id(asyncio.get_running_loop())
+            loop = asyncio.get_running_loop()
         except RuntimeError:
             # Called from outside an event loop — should not happen on the
             # hot path, but build a one-off client rather than crashing.
             return self._build_client()
-        client = self._clients_by_loop_id.get(loop_id)
+        client = self._clients_by_loop.get(loop)
         if client is None:
             client = self._build_client()
-            self._clients_by_loop_id[loop_id] = client
+            self._clients_by_loop[loop] = client
         return client
 
     @override

tests/lib/core/tracing/processors/test_sgp_tracing_processor.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import uuid
+import asyncio
 from datetime import UTC, datetime
 from unittest.mock import AsyncMock, MagicMock, patch
 
@@ -292,6 +293,44 @@ def capture_limits(*args, **kwargs):
             f"max_keepalive_connections={max_keepalive}"
         )
 
+    def test_cache_is_weakkeydict_and_evicts_dead_loops(self):
+        """Regression guard for the id()-reuse bug: the per-loop cache must
+        be a WeakKeyDictionary so a GC'd loop's entry is evicted.  Otherwise
+        a new loop landing at the same memory address would reuse the dead
+        loop's client, reintroducing the "bound to a different event loop"
+        error the per-loop cache was built to prevent.
+        """
+        import gc
+        import weakref
+
+        mock_env = MagicMock()
+        mock_env.refresh.return_value = MagicMock(ACP_TYPE=None, AGENT_NAME=None, AGENT_ID=None)
+
+        with patch(f"{MODULE}.EnvironmentVariables", mock_env), patch(f"{MODULE}.AsyncSGPClient"):
+            from agentex.lib.core.tracing.processors.sgp_tracing_processor import (
+                SGPAsyncTracingProcessor,
+            )
+
+            processor = SGPAsyncTracingProcessor(_make_config())
+
+        # Storage type itself: WeakKeyDictionary, not plain dict.
+        assert isinstance(processor._clients_by_loop, weakref.WeakKeyDictionary)
+
+        # End-to-end check: insert under a loop, drop the loop, the entry
+        # must vanish after GC.
+        loop = asyncio.new_event_loop()
+        try:
+            processor._clients_by_loop[loop] = MagicMock()
+            assert len(processor._clients_by_loop) == 1
+        finally:
+            loop.close()
+        del loop
+        gc.collect()
+        assert len(processor._clients_by_loop) == 0, (
+            "WeakKeyDictionary should have evicted the dead loop's entry; "
+            "remaining keys would cause stale-client reuse on id() recycling."
+        )
+
     async def test_disabled_processor_returns_none_client(self):
         """When config is missing api_key/account_id, _get_client must return
         None and no HTTP client must be constructed."""