Skip to content

Commit 8652244

Browse files
committed
fix: Add close() method to SQLAlchemySession for proper engine disposal
- Add _owns_engine flag to track engine ownership - from_url() sets _owns_engine=True (SDK creates engine) - Direct engine injection keeps _owns_engine=False (user manages) - close() disposes engine only when owned by session - Matches RedisSession pattern for resource lifecycle management Fixes resource leak when using from_url() in long-running applications
1 parent 9db9c01 commit 8652244

File tree

1 file changed

+14
-1
lines changed

1 file changed

+14
-1
lines changed

src/agents/extensions/memory/sqlalchemy_session.py

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -81,6 +81,7 @@ def __init__(
8181
self.session_id = session_id
8282
self._engine = engine
8383
self._lock = asyncio.Lock()
84+
self._owns_engine = False # Track if we own the engine
8485

8586
self._metadata = MetaData()
8687
self._sessions = Table(
@@ -159,7 +160,9 @@ def from_url(
159160
"""
160161
engine_kwargs = engine_kwargs or {}
161162
engine = create_async_engine(url, **engine_kwargs)
162-
return cls(session_id, engine=engine, **kwargs)
163+
session = cls(session_id, engine=engine, **kwargs)
164+
session._owns_engine = True # We created the engine, so we own it
165+
return session
163166

164167
async def _serialize_item(self, item: TResponseInputItem) -> str:
165168
"""Serialize an item to JSON string. Can be overridden by subclasses."""
@@ -319,3 +322,13 @@ async def clear_session(self) -> None:
319322
await sess.execute(
320323
delete(self._sessions).where(self._sessions.c.session_id == self.session_id)
321324
)
325+
326+
async def close(self) -> None:
327+
"""Close the database engine connection.
328+
329+
Only disposes the engine if this session owns it
330+
(i.e., created via from_url). If the engine was injected externally,
331+
the caller is responsible for managing its lifecycle.
332+
"""
333+
if self._owns_engine:
334+
await self._engine.dispose()

0 commit comments

Comments
 (0)