feat: Extract shared database drivers for session, multimodal, and response stores (#357)

* feat: Extract shared database drivers for session, multimodal, and response stores

* feat: Refactor shared database drivers and retry logic

- Parameterize retry helper by exception types, allowing scoped retries for drivers.
- Migrate CosmosDB and Firestore drivers to shared `core/util/drivers/`.
- Consolidate `RedisLikeDriver` retry and health-check behavior across sessions and response stores.
- Standardize method names in CosmosDB to match DynamoDB drivers.
- Update consumer stores to explicit driver parameters and adjust related tests.

* feat: Refactor driver state management and connection handling for Redis and DynamoDB

* feat: Migrate thread stores to shared database drivers and consolidate connection handling

* fix: Correct line numbers in spec.md for thread store connection logic

* fix: Update `spec.md` to reflect package rename from `drivers` to `driver` across shared database driver sections

* feat: Add shared database connection drivers for Cosmos DB, DynamoDB, Firestore, Redis, and Valkey

* fix: Handle `None` type for `last_err` in retry logic and enforce minimum retry count

* feat: Introduce BaseDriver class for shared database connection logic and refactor drivers to inherit from it

Author: induwara-yaala

.agents/skills/ak-dev-architecture/SKILL.md

@@ -149,7 +149,7 @@ Provides image and file attachment support via a pluggable storage and PreHook a
 | Backend | Class | Module | Key traits |
 |---------|-------|--------|------------|
 | In-memory | `InMemoryAttachmentStore` | `storage/in_memory.py` | `ClassVar` dict, ephemeral, zero setup |
-| Redis | `RedisAttachmentStore` | `storage/redis.py` | Persistent, TTL, connection pooling |
+| Redis | `RedisAttachmentStore` | `storage/redis.py` | Persistent, TTL, shared `RedisDriver` (lazy connect, retry, ping/reconnect) |
 | DynamoDB | `DynamoDBAttachmentStore` | `storage/dynamodb.py` | Serverless/AWS, TTL via `expiry_time` |
 | Session cache | `SessionNonVolatileCacheAttachmentStore` | `storage/session_cache.py` | Legacy, stores in `nv_cache` (not recommended) |
 
@@ -239,6 +239,26 @@ Pluggable storage backends agents can read from and write to as tools:
 - **`KnowledgeBuilder`** (`knowledgebuilder.py`): Wraps one or more `KnowledgeBase` instances and `build()`s plain-function tools (`get_schemas`, `read_kb`, `write_kb`, `get_all_kb_descriptions`) for binding via a framework's `ToolBuilder`
 - **Backends**: `ChromaManager` (vector, `chroma.py`), `Neo4jManager` (graph, `neo4j.py`), `StarburstManager` (read-only SQL via Trino, `starburst.py`) — each behind an optional dependency extra (`chromadb`, `neo4j`, `trino`)
 
+## Shared Database Drivers (`ak-py/src/agentkernel/core/util/driver/`)
+
+The Session, Multimodal attachment, Response Store, and Thread backends share one set of
+connection drivers: `RedisDriver`, `ValkeyDriver` (both subclassing `_RedisLikeDriver`),
+`DynamoDBDriver`, `CosmosDBDriver`, and `FirestoreDriver`. Three rules govern the package:
+
+1. **Drivers never read `AKConfig`** — all connection parameters are explicit constructor
+   arguments; config reading and validation stay in the stores and factories
+2. **Drivers own the connection lifecycle** (lazy connect, 3-retry/2s back-off, Redis/Valkey
+   ping health-check with reconnect, `socket_connect_timeout=5`, TTL plumbing) plus a generic
+   command surface; key schemas, serialization, and data layouts stay in the store classes
+3. **Drivers expose their native handle** (`client` / `table` / `table_client` / `collection`)
+   for consumers whose data operations exceed the generic surface (e.g. the DynamoDB/Cosmos/
+   Firestore thread stores)
+
+`driver/__init__.py` has no eager imports — `redis`, `valkey`, `azure`, and `gcp` extras stay
+optional; consumers import the concrete module (`from agentkernel.core.util.driver.redis import
+RedisDriver`). Connect/reconnect is serialized by a per-instance `threading.Lock`, so drivers
+are safe to share across threads (e.g. response stores under `ECSOutputConsumer`).
+
 ## Directory Structure
 
 ```
@@ -256,6 +276,7 @@ ak-py/src/agentkernel/
 │   ├── chat_service.py      # ChatService, RequestBuilder, AgentHandler, ResponseBuilder
 │   ├── logger.py            # Logging setup
 │   ├── util/                # Shared utilities
+│   │   └── driver/          # Shared DB connection drivers (Redis, Valkey, DynamoDB, Cosmos DB, Firestore)
 │   └── session/             # Session store implementations
 │       ├── base.py           # SessionStore, SessionCache
 │       ├── serde.py          # Session (de)serialization helpers

.agents/skills/ak-dev-new-multimodal-storage/SKILL.md

@@ -21,7 +21,7 @@ This guide walks through adding a new attachment storage backend to Agent Kernel
 | Backend | Config value | Features | Extras |
 |---|---|---|---|
 | In-memory | `in_memory` | Ephemeral `ClassVar` dict, zero setup, single-process only | None |
-| Redis | `redis` | Persistent, TTL, connection pooling, distributed | `agentkernel[multimodal,redis]` |
+| Redis | `redis` | Persistent, TTL, shared `RedisDriver` (lazy connect, retry, ping/reconnect), distributed | `agentkernel[multimodal,redis]` |
 | DynamoDB | `dynamodb` | Serverless/AWS, TTL via `expiry_time`, fully managed | `agentkernel[multimodal,aws]` |
 | Session cache | `session_cache` | Legacy — stores in session `nv_cache` (causes bloat, not recommended) | None |
 
@@ -141,7 +141,7 @@ class <Backend>AttachmentStore(AttachmentStore):
 - **Max attachments pruning**: When `save()` is called and the session exceeds `max_attachments`, delete the oldest entry. Track order via timestamps or a per-session index
 - **Index consistency on delete**: When `delete()` is called, remove the attachment ID from the per-session index in addition to deleting the data entry. Skipping this step causes the index count to drift — the backend will think more attachments exist than actually do, breaking `max_attachments` enforcement on subsequent saves
 - **TTL**: If the backend supports time-based expiry (like Redis TTL or DynamoDB `expiry_time`), use it for automatic cleanup. Read the TTL value from the backend-specific config
-- **Connection management**: Use lazy initialization and connection pooling where possible. The store may be instantiated per-request (inside `AttachmentStorageManager.__init__`)
+- **Connection management**: Reuse the shared connection drivers in `agentkernel/core/util/driver/` (`RedisDriver`, `DynamoDBDriver`, etc.) — they provide lazy connect, 3-retry back-off, and (for Redis-like backends) ping/reconnect. The store may be instantiated per-request (inside `AttachmentStorageManager.__init__`), so keep the driver construction cheap (no eager connect)
 - **Serialization**: Attachment dicts must be JSON-serializable. The `data` field contains base64-encoded binary, so all values are strings, floats, or ints
 
 ### 3. Add Backend-Specific Configuration

.agents/skills/ak-dev-testing-conventions/SKILL.md

@@ -41,6 +41,11 @@ Tests live in `ak-py/tests/` and follow the naming convention `test_<module>.py`
 | `test_session.py` | Session state, caches, context vars |
 | `test_session_cache.py` | LRU SessionCache |
 | `test_sessions_in_memory.py` | InMemorySessionStore |
+| `test_sessions_redis.py` | RedisSessionStore missing-config error, shared RedisDriver retry exhaustion |
+| `test_sessions_valkey.py` | ValkeySessionStore round trips (fake client), shared ValkeyDriver retry exhaustion |
+| `test_sessions_dynamodb.py` | DynamoDBSessionStore Binary wrap/unwrap, missing-item skip (mocked driver) |
+| `test_shared_drivers.py` | Shared DB drivers (`core/util/driver/`): retry scope, ping/reconnect, command surface, DynamoDB item-dict semantics |
+| `test_multimodal_redis_store.py` | RedisAttachmentStore index TTL refresh, JSON round trip, pruning (mocked driver) |
 | `test_config.py` | AKConfig loading, env vars |
 | `test_test_config.py` | AKTestConfig (Test framework config) loading, defaults |
 | `test_tool.py` | ToolContext, cache |
@@ -64,7 +69,7 @@ Tests live in `ak-py/tests/` and follow the naming convention `test_<module>.py`
 | `test_lambda_router.py` | Lambda routing |
 | `test_sqs_handler.py` | AWS SQSHandler config, client, message sending |
 | `test_serverless_request_handle.py` | BaseRequest/BaseRunRequest parsing from serverless payloads |
-| `test_firestore_database_id.py` | FirestoreDriver named `database_id` configuration |
+| `test_firestore_database_id.py` | Shared `FirestoreDriver` (`core/util/driver/firestore.py`, explicit constructor params) named `database_id` configuration |
 | `test_ak_logger.py` | AKLogger level resolution, configuration |
 | `test_error_util.py` | `user_facing_error_message` error mapping |
 | `test_thread_runner.py` | ThreadRunner task validation, failure/shutdown semantics |

ak-py/src/agentkernel/core/builder.py

@@ -127,7 +127,7 @@ def build() -> SessionStore:
             or InMemorySessionStore (for all other types).
 
         :raises: Any exceptions raised by SessionStoreBuilder.Types.from_str(), AKConfig.get(),
-            RedisDriver(), RedisSessionStore(), or InMemorySessionStore() initialization.
+            RedisSessionStore(), or InMemorySessionStore() initialization.
         """
         session_store_type: SessionStoreBuilder.Types = SessionStoreBuilder.Types.from_str(AKConfig.get().session.type)
         Builder._log.info(f"Building {session_store_type} session store")

ak-py/src/agentkernel/core/config.py

@@ -175,13 +175,12 @@ class _GmailConfig(BaseModel):
     label_filter: str = Field(default="INBOX", description="Gmail label to monitor (e.g., INBOX, UNREAD)")
 
 
-class _MultimodalStorageRedisConfig(BaseModel):
-    url: str = Field(default="redis://localhost:6379", description="Redis connection URL")
+class _MultimodalStorageRedisConfig(_RedisConfig):
     ttl: int = Field(default=604800, description="Attachment TTL in seconds")
     prefix: str = Field(default="ak:attachments:", description="Key prefix for attachment keys")
 
 
-class _MultimodalStorageDynamoDBConfig(BaseModel):
+class _MultimodalStorageDynamoDBConfig(_DynamoDBConfig):
     table_name: str = Field(default="ak-attachments", description="DynamoDB table name for attachment storage")
     ttl: int = Field(default=604800, description="Attachment TTL in seconds (0 disables)")
 
@@ -212,27 +211,24 @@ class _MultimodalConfig(BaseModel):
     dynamodb: Optional[_MultimodalStorageDynamoDBConfig] = None
 
 
-class _ThreadRedisConfig(BaseModel):
-    url: str = Field(default="redis://localhost:6379", description="Redis connection URL. Use rediss:// for SSL")
+class _ThreadRedisConfig(_RedisConfig):
     ttl: int = Field(default=2592000, description="Thread TTL in seconds (0 disables)")
     prefix: str = Field(default="ak:thread:", description="Key prefix for Redis thread storage")
 
 
-class _ThreadDynamoDBConfig(BaseModel):
+class _ThreadDynamoDBConfig(_DynamoDBConfig):
     table_name: str = Field(
         default="ak-agent-threads",
         description="DynamoDB table name for thread storage. Table should have a partition key named 'session_id' (S) and a sort key named 'sk' (S)",
     )
     ttl: int = Field(default=0, description="DynamoDB item TTL in seconds (0 disables)")
 
 
-class _ThreadFirestoreConfig(BaseModel):
+class _ThreadFirestoreConfig(_FirestoreConfig):
     collection_name: str = Field(
         default="ak-agent-threads",
         description="Firestore collection name for thread storage. Each document ID is a session_id.",
     )
-    project_id: Optional[str] = Field(default=None, description="GCP project ID. If null, inferred from Application Default Credentials.")
-    database_id: Optional[str] = Field(default=None, description="Firestore database ID. If null, defaults to '(default)' database.")
     ttl: int = Field(default=0, description="Thread TTL in seconds (0 disables)")