From 12eb564eda2a947cb8e265dfcc4064cc0f42cade Mon Sep 17 00:00:00 2001
From: Daniel Chalef <131175+danielchalef@users.noreply.github.com>
Date: Wed, 29 Oct 2025 22:43:09 -0700
Subject: [PATCH] fix: Resolve database connection and episode processing
 errors

Fixed two critical runtime errors:

1. Database connection check for KuzuDB
   - KuzuDB session.run() returns None, causing async iteration error
   - Added special handling for KuzuDB (in-memory, no query needed)
   - Other databases (Neo4j, FalkorDB) still perform connection test

2. Episode processing parameter error
   - Changed 'episode_type' parameter to 'source' to match Graphiti API
   - Added required 'reference_time' parameter with current timestamp
   - Added datetime imports (UTC, datetime)

Errors fixed:
- 'async for' requires an object with __aiter__ method, got NoneType
- Graphiti.add_episode() got an unexpected keyword argument 'episode_type'
---
 mcp_server/src/graphiti_mcp_server.py    | 13 ++++++++++---
 mcp_server/src/services/queue_service.py |  5 +++--
 2 files changed, 13 insertions(+), 5 deletions(-)

diff --git a/mcp_server/src/graphiti_mcp_server.py b/mcp_server/src/graphiti_mcp_server.py
index 2d009af1..7ff18a06 100644
--- a/mcp_server/src/graphiti_mcp_server.py
+++ b/mcp_server/src/graphiti_mcp_server.py
@@ -657,12 +657,19 @@ async def get_status() -> StatusResponse:
     try:
         client = await graphiti_service.get_client()
 
-        # Test database connection with a simple query
-        # This works for all supported databases (Neo4j, FalkorDB, KuzuDB)
+        # For KuzuDB, just check if the client exists (it's in-memory)
+        if config.database.provider.lower() == 'kuzu':
+            provider_info = f'{config.database.provider} database (in-memory)'
+            return StatusResponse(
+                status='ok', message=f'Graphiti MCP server is running and connected to {provider_info}'
+            )
+
+        # For Neo4j and FalkorDB, test connection with a simple query
         async with client.driver.session() as session:
             result = await session.run('MATCH (n) RETURN count(n) as count')
             # Consume the result to verify query execution
-            _ = [record async for record in result]
+            if result:
+                _ = [record async for record in result]
 
         provider_info = f'{config.database.provider} database'
         return StatusResponse(
diff --git a/mcp_server/src/services/queue_service.py b/mcp_server/src/services/queue_service.py
index 93a0ba8e..502b988c 100644
--- a/mcp_server/src/services/queue_service.py
+++ b/mcp_server/src/services/queue_service.py
@@ -3,6 +3,7 @@
 import asyncio
 import logging
 from collections.abc import Awaitable, Callable
+from datetime import UTC, datetime
 from typing import Any
 
 logger = logging.getLogger(__name__)
@@ -134,9 +135,9 @@ class QueueService:
                     name=name,
                     episode_body=content,
                     source_description=source_description,
-                    episode_type=episode_type,
+                    source=episode_type,
                     group_id=group_id,
-                    reference_time=None,  # Let graphiti handle timing
+                    reference_time=datetime.now(UTC),
                     entity_types=entity_types,
                     uuid=uuid,
                 )