chore: sync with main and stabilize MCP server#59
Merged
visahak merged 7 commits intoFeb 23, 2026
Conversation
- Resolved merge conflicts with origin/main. - Implemented lazy initialization in MCP server to fix SSE race conditions. - Added environment variable support for SQLite path and data directory. - Simplified MCP E2E tests by switching to filesystem backend and ensuring isolation.
📝 WalkthroughWalkthroughIntroduces lazy, thread-safe Kaizen client and one-time namespace initialization in the MCP server; makes SQLiteManager accept an optional db_path with environment fallback; adds Milvus embedding-index/search and index creation; parameterizes E2E tests for Milvus/filesystem backends; minor dependency/TOML fixes and test mocks. Changes
Sequence Diagram(s)sequenceDiagram
participant ClientCaller as Client (API call)
participant MCP as MCP Server (mcp_server.get_client)
participant Kaizen as KaizenClient
participant Backend as Backend (Milvus / SQLite / FS)
ClientCaller->>MCP: call create_entity(...)
MCP->>MCP: acquire _client_init_lock
alt client not initialized
MCP->>Kaizen: KaizenClient(...)
Kaizen-->>MCP: client instance
MCP->>Backend: ensure namespace / create collection
Backend-->>MCP: namespace created
MCP-->>MCP: set _namespace_initialized = true
end
MCP->>Kaizen: use client to update_entities(...)
Kaizen->>Backend: persist entities / index (Milvus index/create)
Backend-->>Kaizen: ack (flush/reload if Milvus)
Kaizen-->>MCP: update result
MCP-->>ClientCaller: return created entity / error
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches
🧪 Generate unit tests (beta)
Tip Issue Planner is now in beta. Read the docs and try it out! Share your feedback on Discord. Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tests/e2e/test_mcp.py (1)
9-9: 🛠️ Refactor suggestion | 🟠 MajorUnused import:
milvus_client_settings.This import is a leftover from the previous Milvus-based test setup and is not referenced anywhere in the file. Ruff would flag this.
-from kaizen.config.milvus import milvus_client_settings🤖 Prompt for AI Agents
Before applying any fix, first verify the finding against the current code and decide whether a code change is actually needed. If the finding is not valid or no change is required, do not modify code for that item and briefly explain why it was skipped. In `@tests/e2e/test_mcp.py` at line 9, Remove the unused import milvus_client_settings from the top of the test file: locate the import statement that references milvus_client_settings and delete it, ensuring no other references to milvus_client_settings remain in the file so the test no longer contains an unused import.
🤖 Fix all issues with AI agents
Before applying any fix, first verify the finding against the current code and
decide whether a code change is actually needed. If the finding is not valid or
no change is required, do not modify code for that item and briefly explain why
it was skipped.
In `@kaizen/db/sqlite_manager.py`:
- Around line 28-30: The inline import of os inside SQLiteManager.__init__
should be moved to the module top-level with the other stdlib imports (datetime,
logging, sqlite3, threading) in sqlite_manager.py; remove the inline "import os"
from the __init__ method and add a single top-level "import os" near the other
imports so self.db_path = db_path or os.environ.get("KAIZEN_SQLITE_PATH",
"entities.sqlite.db") continues to work without per-instantiation imports.
In `@kaizen/frontend/mcp/mcp_server.py`:
- Line 10: Remove the unused import asynccontextmanager from the top of the
file: the import statement "from contextlib import asynccontextmanager" in
mcp_server.py is not referenced anywhere and should be deleted to satisfy
linting (Ruff) and avoid unused imports.
- Around line 200-205: In the exception handler inside create_entity (the except
Exception as e: block) remove the redundant traceback import and the
traceback.print_exc() call and rely only on logger.exception(...) to log the
error and traceback; keep the logger.exception(f"CRASH IN CREATE_ENTITY: {e}")
and the subsequent return json.dumps({"error": f"Server Error: {str(e)}"}) as
the error response.
- Around line 28-61: The broad except in get_client surrounding
_client.get_namespace_details(kaizen_config.namespace_id) should be narrowed and
the original error preserved: only catch exceptions that indicate the namespace
might not exist (e.g., NamespaceNotFoundException) or specific recoverable
errors you expect, and for any other Exception log the original exception with
full context (including the exception instance) and re-raise instead of blindly
calling _client.create_namespace; update the get_client function to log the
caught error (use the exception variable) before any retry/create attempt and
only attempt create_namespace when it makes sense.
- Around line 34-58: The lazy initialization of the global _client and
_namespace_initialized in get_client()/mcp_server.py is not
thread-safe—concurrent callers can race and create multiple KaizenClient
instances or duplicate namespace work; protect the initialization with a
threading.Lock (e.g., a module-level _init_lock) and use double-checked locking
around the blocks that create _client and set _namespace_initialized, wrapping
both the client creation (KaizenClient()) and namespace checks/creation
(get_namespace_details, create_namespace with kaizen_config.namespace_id) inside
the lock to ensure only one thread performs initialization, or alternatively
document the single-threaded assumption if you prefer not to add locking.
In `@tests/e2e/test_mcp.py`:
- Around line 15-46: The mcp fixture currently mutates os.environ directly
(lines setting 'KAIZEN_BACKEND', 'KAIZEN_NAMESPACE_ID', 'KAIZEN_DATA_DIR') and
never restores them; replace those direct environment assignments with pytest's
monkeypatch.setenv calls inside the mcp fixture (use the fixture parameter name
monkeypatch) so environment variables are automatically restored after the test,
keep the subsequent kaizen_config.__init__() and filesystem_settings.__init__()
calls, and leave the existing MCP reset logic (mcp_server_module._client and
_namespace_initialized) as-is.
- Around line 37-40: Replace the broad try/except around
kaizen_client.create_namespace('test') so it only catches the expected
NamespaceAlreadyExistsException; allow any other exceptions to propagate (or
re-raise them) so setup failures are not silently swallowed. Specifically,
change the except Exception: pass to an except NamespaceAlreadyExistsException:
pass (import or reference the exception from the same module that defines it)
and do not suppress other errors thrown by kaizen_client.create_namespace.
- Around line 106-119: Replace the manual sys.modules patching for
fs_backend.resolve_conflicts (where code assigns original_resolve, defines
mock_resolve_func, and sets fs_backend.resolve_conflicts) with pytest's
monkeypatch fixture: use monkeypatch.setattr to replace
kaizen.backend.filesystem.resolve_conflicts with the mock_resolve_func so
teardown is automatic; ensure the mock function signature and return
(EntityUpdate list using first_entity_id) remain the same and remove the manual
restore of original_resolve.
- Line 9: Remove the unused import milvus_client_settings from the top of the
test file: locate the import statement that references milvus_client_settings
and delete it, ensuring no other references to milvus_client_settings remain in
the file so the test no longer contains an unused import.
🧹 Nitpick comments (6)
🤖 Fix all nitpicks with AI agents
Before applying any fix, first verify the finding against the current code and decide whether a code change is actually needed. If the finding is not valid or no change is required, do not modify code for that item and briefly explain why it was skipped. In `@kaizen/db/sqlite_manager.py`: - Around line 28-30: The inline import of os inside SQLiteManager.__init__ should be moved to the module top-level with the other stdlib imports (datetime, logging, sqlite3, threading) in sqlite_manager.py; remove the inline "import os" from the __init__ method and add a single top-level "import os" near the other imports so self.db_path = db_path or os.environ.get("KAIZEN_SQLITE_PATH", "entities.sqlite.db") continues to work without per-instantiation imports. In `@kaizen/frontend/mcp/mcp_server.py`: - Around line 200-205: In the exception handler inside create_entity (the except Exception as e: block) remove the redundant traceback import and the traceback.print_exc() call and rely only on logger.exception(...) to log the error and traceback; keep the logger.exception(f"CRASH IN CREATE_ENTITY: {e}") and the subsequent return json.dumps({"error": f"Server Error: {str(e)}"}) as the error response. - Around line 28-61: The broad except in get_client surrounding _client.get_namespace_details(kaizen_config.namespace_id) should be narrowed and the original error preserved: only catch exceptions that indicate the namespace might not exist (e.g., NamespaceNotFoundException) or specific recoverable errors you expect, and for any other Exception log the original exception with full context (including the exception instance) and re-raise instead of blindly calling _client.create_namespace; update the get_client function to log the caught error (use the exception variable) before any retry/create attempt and only attempt create_namespace when it makes sense. - Around line 34-58: The lazy initialization of the global _client and _namespace_initialized in get_client()/mcp_server.py is not thread-safe—concurrent callers can race and create multiple KaizenClient instances or duplicate namespace work; protect the initialization with a threading.Lock (e.g., a module-level _init_lock) and use double-checked locking around the blocks that create _client and set _namespace_initialized, wrapping both the client creation (KaizenClient()) and namespace checks/creation (get_namespace_details, create_namespace with kaizen_config.namespace_id) inside the lock to ensure only one thread performs initialization, or alternatively document the single-threaded assumption if you prefer not to add locking. In `@tests/e2e/test_mcp.py`: - Around line 37-40: Replace the broad try/except around kaizen_client.create_namespace('test') so it only catches the expected NamespaceAlreadyExistsException; allow any other exceptions to propagate (or re-raise them) so setup failures are not silently swallowed. Specifically, change the except Exception: pass to an except NamespaceAlreadyExistsException: pass (import or reference the exception from the same module that defines it) and do not suppress other errors thrown by kaizen_client.create_namespace. - Around line 106-119: Replace the manual sys.modules patching for fs_backend.resolve_conflicts (where code assigns original_resolve, defines mock_resolve_func, and sets fs_backend.resolve_conflicts) with pytest's monkeypatch fixture: use monkeypatch.setattr to replace kaizen.backend.filesystem.resolve_conflicts with the mock_resolve_func so teardown is automatic; ensure the mock function signature and return (EntityUpdate list using first_entity_id) remain the same and remove the manual restore of original_resolve.tests/e2e/test_mcp.py (2)
37-40: Broadexcept Exception: passsilently swallows setup failures.If the namespace already exists, a
NamespaceAlreadyExistsExceptionis expected — but any other error (e.g., misconfigured backend, filesystem permissions) would be silently hidden, making test failures hard to diagnose.Proposed fix — catch only the expected exception
+ from kaizen.schema.exceptions import NamespaceAlreadyExistsException try: kaizen_client.create_namespace('test') - except Exception: + except NamespaceAlreadyExistsException: pass🤖 Prompt for AI Agents
Before applying any fix, first verify the finding against the current code and decide whether a code change is actually needed. If the finding is not valid or no change is required, do not modify code for that item and briefly explain why it was skipped. In `@tests/e2e/test_mcp.py` around lines 37 - 40, Replace the broad try/except around kaizen_client.create_namespace('test') so it only catches the expected NamespaceAlreadyExistsException; allow any other exceptions to propagate (or re-raise them) so setup failures are not silently swallowed. Specifically, change the except Exception: pass to an except NamespaceAlreadyExistsException: pass (import or reference the exception from the same module that defines it) and do not suppress other errors thrown by kaizen_client.create_namespace.
106-119: Directsys.modulesattribute patching is fragile — considermonkeypatch.setattr.The manual save-replace-restore pattern works (and the
finallyblock is good), but pytest'smonkeypatchfixture would handle the cleanup automatically and be more idiomatic. This is a minor suggestion since the current approach is functional and correctly wrapped intry/finally.🤖 Prompt for AI Agents
Before applying any fix, first verify the finding against the current code and decide whether a code change is actually needed. If the finding is not valid or no change is required, do not modify code for that item and briefly explain why it was skipped. In `@tests/e2e/test_mcp.py` around lines 106 - 119, Replace the manual sys.modules patching for fs_backend.resolve_conflicts (where code assigns original_resolve, defines mock_resolve_func, and sets fs_backend.resolve_conflicts) with pytest's monkeypatch fixture: use monkeypatch.setattr to replace kaizen.backend.filesystem.resolve_conflicts with the mock_resolve_func so teardown is automatic; ensure the mock function signature and return (EntityUpdate list using first_entity_id) remain the same and remove the manual restore of original_resolve.kaizen/frontend/mcp/mcp_server.py (3)
200-205: Redundant traceback logging —logger.exceptionalready includes the traceback.
traceback.print_exc()writes to stderr, andlogger.exception(...)on the next line also logs the full traceback. This produces duplicate output. Remove thetracebackimport andprint_exc()call, and rely solely onlogger.exception.Proposed fix
except Exception as e: - import traceback - traceback.print_exc() - logger.exception(f"CRASH IN CREATE_ENTITY: {e}") + logger.exception(f"Error in create_entity: {e}") return json.dumps({"error": f"Server Error: {str(e)}"})🤖 Prompt for AI Agents
Before applying any fix, first verify the finding against the current code and decide whether a code change is actually needed. If the finding is not valid or no change is required, do not modify code for that item and briefly explain why it was skipped. In `@kaizen/frontend/mcp/mcp_server.py` around lines 200 - 205, In the exception handler inside create_entity (the except Exception as e: block) remove the redundant traceback import and the traceback.print_exc() call and rely only on logger.exception(...) to log the error and traceback; keep the logger.exception(f"CRASH IN CREATE_ENTITY: {e}") and the subsequent return json.dumps({"error": f"Server Error: {str(e)}"}) as the error response.
28-61: Broad fallback inget_client()(lines 50–57) may mask real errors.The generic
except Exceptionon line 50 catches any failure fromget_namespace_details(connectivity errors, auth failures, serialization bugs, etc.) and blindly retries withcreate_namespace. This can hide the real root cause — if the backend is unreachable, the second call will also fail but with a potentially less informative error.Consider narrowing the catch to only exceptions that meaningfully indicate "namespace may not exist" or at minimum logging the original error before retrying:
Proposed fix
except NamespaceNotFoundException: logger.info(f"Creating namespace '{kaizen_config.namespace_id}'") _client.create_namespace(kaizen_config.namespace_id) logger.info(f"Namespace '{kaizen_config.namespace_id}' created successfully") - except Exception as e: - logger.error(f"Error ensuring namespace: {e}") - try: - _client.create_namespace(kaizen_config.namespace_id) - logger.info(f"Namespace '{kaizen_config.namespace_id}' created after error") - except Exception as create_error: - logger.error(f"Failed to create namespace: {create_error}") - raise + except Exception: + logger.exception(f"Unexpected error checking namespace '{kaizen_config.namespace_id}'") + raise🤖 Prompt for AI Agents
Before applying any fix, first verify the finding against the current code and decide whether a code change is actually needed. If the finding is not valid or no change is required, do not modify code for that item and briefly explain why it was skipped. In `@kaizen/frontend/mcp/mcp_server.py` around lines 28 - 61, The broad except in get_client surrounding _client.get_namespace_details(kaizen_config.namespace_id) should be narrowed and the original error preserved: only catch exceptions that indicate the namespace might not exist (e.g., NamespaceNotFoundException) or specific recoverable errors you expect, and for any other Exception log the original exception with full context (including the exception instance) and re-raise instead of blindly calling _client.create_namespace; update the get_client function to log the caught error (use the exception variable) before any retry/create attempt and only attempt create_namespace when it makes sense.
34-58: Lazy init with bare globals is not thread-safe.If
get_client()is called concurrently (e.g., two simultaneous SSE requests), both callers can see_client is Noneand create separateKaizenClientinstances, or race on_namespace_initialized. In a single-threaded async event loop this is likely fine since the tools are synchronous, but if the server ever handles requests on multiple threads this will break.Worth documenting the single-threaded assumption or adding a
threading.Lockguard if the intent is to be thread-safe.🤖 Prompt for AI Agents
Before applying any fix, first verify the finding against the current code and decide whether a code change is actually needed. If the finding is not valid or no change is required, do not modify code for that item and briefly explain why it was skipped. In `@kaizen/frontend/mcp/mcp_server.py` around lines 34 - 58, The lazy initialization of the global _client and _namespace_initialized in get_client()/mcp_server.py is not thread-safe—concurrent callers can race and create multiple KaizenClient instances or duplicate namespace work; protect the initialization with a threading.Lock (e.g., a module-level _init_lock) and use double-checked locking around the blocks that create _client and set _namespace_initialized, wrapping both the client creation (KaizenClient()) and namespace checks/creation (get_namespace_details, create_namespace with kaizen_config.namespace_id) inside the lock to ensure only one thread performs initialization, or alternatively document the single-threaded assumption if you prefer not to add locking.kaizen/db/sqlite_manager.py (1)
28-30: Moveimport osto the module top-level.
osis a standard-library module and there's no circular-import concern here. Placing it at the top of the file alongside the other stdlib imports (datetime,logging,sqlite3,threading) is more conventional and avoids a per-instantiation import lookup.Proposed fix
At the top of the file (after line 4):
import threading +import osThen in
__init__:def __init__(self, db_path: str | None = None): - import os self.db_path = db_path or os.environ.get("KAIZEN_SQLITE_PATH", "entities.sqlite.db")🤖 Prompt for AI Agents
Before applying any fix, first verify the finding against the current code and decide whether a code change is actually needed. If the finding is not valid or no change is required, do not modify code for that item and briefly explain why it was skipped. In `@kaizen/db/sqlite_manager.py` around lines 28 - 30, The inline import of os inside SQLiteManager.__init__ should be moved to the module top-level with the other stdlib imports (datetime, logging, sqlite3, threading) in sqlite_manager.py; remove the inline "import os" from the __init__ method and add a single top-level "import os" near the other imports so self.db_path = db_path or os.environ.get("KAIZEN_SQLITE_PATH", "entities.sqlite.db") continues to work without per-instantiation imports.
This comment was marked as outdated.
This comment was marked as outdated.
Sorry, something went wrong.
visahak
force-pushed
the
fix/mcp-sync-and-stabilization
branch
from
fcafc16 to
77d4a35
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI agents
Verify each finding against the current code and only fix it if needed.
In `@kaizen/backend/milvus.py`:
- Around line 197-226: The conditional checking for empty search results is
redundant: replace the combined check "if not results or len(results) == 0" with
a single falsy check "if not results" where the Milvus search response is
handled (around the code that calls self.milvus.search and assigns to results),
keeping the subsequent early return of an empty list and leaving the hit-parsing
logic that uses parse_milvus_entity intact.
- Around line 177-179: The current call pattern — invoking
self.milvus.flush(namespace_id) and self.milvus.load_collection(namespace_id) on
every update_entities call — will become a throughput bottleneck; change
update_entities in the Milvus client wrapper so it does not synchronously
flush+load on every write but instead batches or defers those operations: add a
short-lived buffer or queue keyed by namespace_id and trigger flush+load either
when the buffer reaches a configurable batch_size or on a periodic background
timer (or when a force_flush flag is passed), and expose configuration flags
(e.g., batch_size, flush_interval, force_flush) so callers can opt into
immediate consistency when required; keep the flush+load calls but move them to
the batching/worker path and ensure thread-safe enqueueing in update_entities.
- Around line 228-235: The except handler in delete_entity_by_id should
explicitly chain the original ValueError: change the except to "except
ValueError as e:" and re-raise the KaizenException with "from e" (i.e., raise
KaizenException(f\"Invalid entity ID: {entity_id}. Entity IDs must be
numeric.\") from e) so the original ValueError context is preserved for callers
and logs.
In `@kaizen/db/sqlite_manager.py`:
- Around line 29-32: Move the local import of os out of the __init__ method and
add it to the module-level imports alongside datetime and logging; then simplify
__init__ to use the already-imported os when resolving db_path (db_path or
os.environ.get("KAIZEN_SQLITE_PATH", "entities.sqlite.db")). Ensure no other
uses of os remain inside methods so the dependency is visible at the top of the
file.
In `@kaizen/frontend/mcp/mcp_server.py`:
- Around line 188-193: Remove the redundant traceback.print_exc() and the inline
import inside the except block in the create_entity error handler: move the
import traceback to the top-level imports of mcp_server.py (so the module is
imported once), delete the traceback.print_exc() line in the except Exception as
e block, and keep the existing logger.exception(f"CRASH IN CREATE_ENTITY: {e}")
and json.dumps({"error": ...}) return as-is so the full traceback is emitted via
logger.exception().
- Around line 49-56: The current broad "except Exception as e" around the
namespace check can mask real connection/back-end errors and then try to create
the namespace; change this to only attempt creation when the failure is a
namespace-missing error (catch the specific NamespaceNotFoundException from your
client) and otherwise log the original exception with full context and re-raise;
if you cannot import that specific exception type, at minimum replace
logger.error(f"Error ensuring namespace: {e}") with logger.exception(...) to
include traceback and then re-raise when the exception is not a
NamespaceNotFoundException; update the block that calls
_client.create_namespace(kaizen_config.namespace_id) so it runs only under the
namespace-not-found branch and keep the create_error handling unchanged.
In `@tests/e2e/test_mcp.py`:
- Around line 28-32: The milvus temp DB file is being created in CWD because
milvus_db_file is a relative filename; change creation to place the file inside
pytest's tmp_path (e.g., build the path using tmp_path /
f"test_{uuid.uuid4().hex[:8]}.db") and set milvus_client_settings.uri to the
stringified tmp_path path so Milvus uses the tmp-managed directory (refer to
backend_type check, milvus_db_file variable, and milvus_client_settings.uri);
this ensures pytest auto-cleans the DB file and simplifies teardown.
- Around line 165-172: The test reads os.environ["KAIZEN_BACKEND"] at runtime to
choose patch_target, which couples the test to an env-var side effect; change
the test to accept the backend explicitly from the test fixtures instead: have
the mcp fixture (or a new fixture) return the backend type (or a namedtuple
including backend) or access request.param in a parametrized fixture, then use
that backend value to set patch_target
("kaizen.backend.filesystem.resolve_conflicts" vs
"kaizen.backend.milvus.resolve_conflicts") and call patch(patch_target) — update
references to KAIZEN_BACKEND, patch_target, and the mcp fixture/request.param
accordingly.
🧹 Nitpick comments (7)
🤖 Fix all nitpicks with AI agents
Verify each finding against the current code and only fix it if needed. In `@kaizen/backend/milvus.py`: - Around line 197-226: The conditional checking for empty search results is redundant: replace the combined check "if not results or len(results) == 0" with a single falsy check "if not results" where the Milvus search response is handled (around the code that calls self.milvus.search and assigns to results), keeping the subsequent early return of an empty list and leaving the hit-parsing logic that uses parse_milvus_entity intact. - Around line 177-179: The current call pattern — invoking self.milvus.flush(namespace_id) and self.milvus.load_collection(namespace_id) on every update_entities call — will become a throughput bottleneck; change update_entities in the Milvus client wrapper so it does not synchronously flush+load on every write but instead batches or defers those operations: add a short-lived buffer or queue keyed by namespace_id and trigger flush+load either when the buffer reaches a configurable batch_size or on a periodic background timer (or when a force_flush flag is passed), and expose configuration flags (e.g., batch_size, flush_interval, force_flush) so callers can opt into immediate consistency when required; keep the flush+load calls but move them to the batching/worker path and ensure thread-safe enqueueing in update_entities. - Around line 228-235: The except handler in delete_entity_by_id should explicitly chain the original ValueError: change the except to "except ValueError as e:" and re-raise the KaizenException with "from e" (i.e., raise KaizenException(f\"Invalid entity ID: {entity_id}. Entity IDs must be numeric.\") from e) so the original ValueError context is preserved for callers and logs. In `@kaizen/db/sqlite_manager.py`: - Around line 29-32: Move the local import of os out of the __init__ method and add it to the module-level imports alongside datetime and logging; then simplify __init__ to use the already-imported os when resolving db_path (db_path or os.environ.get("KAIZEN_SQLITE_PATH", "entities.sqlite.db")). Ensure no other uses of os remain inside methods so the dependency is visible at the top of the file. In `@kaizen/frontend/mcp/mcp_server.py`: - Around line 188-193: Remove the redundant traceback.print_exc() and the inline import inside the except block in the create_entity error handler: move the import traceback to the top-level imports of mcp_server.py (so the module is imported once), delete the traceback.print_exc() line in the except Exception as e block, and keep the existing logger.exception(f"CRASH IN CREATE_ENTITY: {e}") and json.dumps({"error": ...}) return as-is so the full traceback is emitted via logger.exception(). - Around line 49-56: The current broad "except Exception as e" around the namespace check can mask real connection/back-end errors and then try to create the namespace; change this to only attempt creation when the failure is a namespace-missing error (catch the specific NamespaceNotFoundException from your client) and otherwise log the original exception with full context and re-raise; if you cannot import that specific exception type, at minimum replace logger.error(f"Error ensuring namespace: {e}") with logger.exception(...) to include traceback and then re-raise when the exception is not a NamespaceNotFoundException; update the block that calls _client.create_namespace(kaizen_config.namespace_id) so it runs only under the namespace-not-found branch and keep the create_error handling unchanged. In `@tests/e2e/test_mcp.py`: - Around line 165-172: The test reads os.environ["KAIZEN_BACKEND"] at runtime to choose patch_target, which couples the test to an env-var side effect; change the test to accept the backend explicitly from the test fixtures instead: have the mcp fixture (or a new fixture) return the backend type (or a namedtuple including backend) or access request.param in a parametrized fixture, then use that backend value to set patch_target ("kaizen.backend.filesystem.resolve_conflicts" vs "kaizen.backend.milvus.resolve_conflicts") and call patch(patch_target) — update references to KAIZEN_BACKEND, patch_target, and the mcp fixture/request.param accordingly.kaizen/backend/milvus.py (3)
197-226: Vector search path handles empty filters and varied hit formats well — one minor redundancy.Line 211:
if not results or len(results) == 0:—not resultsalready coversNone,[], and other falsy values, makinglen(results) == 0redundant. Minor nit.The hit-parsing logic at lines 217–224 correctly handles both
hit['entity']and flat-dict formats across PyMilvus versions.Minor simplification
- if not results or len(results) == 0: + if not results: return []🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@kaizen/backend/milvus.py` around lines 197 - 226, The conditional checking for empty search results is redundant: replace the combined check "if not results or len(results) == 0" with a single falsy check "if not results" where the Milvus search response is handled (around the code that calls self.milvus.search and assigns to results), keeping the subsequent early return of an empty list and leaving the hit-parsing logic that uses parse_milvus_entity intact.
177-179:flush+load_collectionafter everyupdate_entitiescall may become a performance bottleneck under high write throughput.This guarantees read-after-write consistency, which is correct for the current use case. If write frequency increases, consider batching or deferring these calls.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@kaizen/backend/milvus.py` around lines 177 - 179, The current call pattern — invoking self.milvus.flush(namespace_id) and self.milvus.load_collection(namespace_id) on every update_entities call — will become a throughput bottleneck; change update_entities in the Milvus client wrapper so it does not synchronously flush+load on every write but instead batches or defers those operations: add a short-lived buffer or queue keyed by namespace_id and trigger flush+load either when the buffer reaches a configurable batch_size or on a periodic background timer (or when a force_flush flag is passed), and expose configuration flags (e.g., batch_size, flush_interval, force_flush) so callers can opt into immediate consistency when required; keep the flush+load calls but move them to the batching/worker path and ensure thread-safe enqueueing in update_entities.
228-235: Numeric ID validation indelete_entity_by_id— consider chaining the originalValueError.The
raise KaizenException(...)inside theexcept ValueErrorblock will implicitly chain the exception. Using explicit chaining (from) is more intentional:Proposed fix
except ValueError: - raise KaizenException(f"Invalid entity ID: {entity_id}. Entity IDs must be numeric.") + raise KaizenException(f"Invalid entity ID: {entity_id}. Entity IDs must be numeric.") from NoneUsing
from Nonesuppresses the noisy implicit chain if theValueErrorcontext isn't useful to callers, or usefrom eif you want to preserve it explicitly.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@kaizen/backend/milvus.py` around lines 228 - 235, The except handler in delete_entity_by_id should explicitly chain the original ValueError: change the except to "except ValueError as e:" and re-raise the KaizenException with "from e" (i.e., raise KaizenException(f\"Invalid entity ID: {entity_id}. Entity IDs must be numeric.\") from e) so the original ValueError context is preserved for callers and logs.kaizen/frontend/mcp/mcp_server.py (2)
188-193: Redundanttraceback.print_exc()alongsidelogger.exception(); moveimport tracebackto the top of the file.
logger.exception()already emits the full traceback to the log. The additionaltraceback.print_exc()duplicates output to stderr and can interleave with structured logs. Also, importingtracebackinside the except block is unconventional for a stdlib module.Proposed fix
At the top of the file:
import json import logging +import traceback import uuidIn the except block:
except Exception as e: - import traceback - - traceback.print_exc() logger.exception(f"CRASH IN CREATE_ENTITY: {e}") return json.dumps({"error": f"Server Error: {str(e)}"})🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@kaizen/frontend/mcp/mcp_server.py` around lines 188 - 193, Remove the redundant traceback.print_exc() and the inline import inside the except block in the create_entity error handler: move the import traceback to the top-level imports of mcp_server.py (so the module is imported once), delete the traceback.print_exc() line in the except Exception as e block, and keep the existing logger.exception(f"CRASH IN CREATE_ENTITY: {e}") and json.dumps({"error": ...}) return as-is so the full traceback is emitted via logger.exception().
49-56: Genericexceptfallback may mask connection errors as namespace-creation attempts.When
get_namespace_detailsfails for a non-NamespaceNotFoundExceptionreason (e.g., connection timeout, backend unavailable), the code falls through to attemptcreate_namespace. This could produce confusing error messages or silently swallow the real issue. Consider logging the original exception type or narrowing the catch.Proposed fix — log the original error with full context
except Exception as e: - logger.error(f"Error ensuring namespace: {e}") + logger.error(f"Unexpected error checking namespace (type={type(e).__name__}): {e}") try: _client.create_namespace(kaizen_config.namespace_id) logger.info(f"Namespace '{kaizen_config.namespace_id}' created after error") except Exception as create_error: - logger.error(f"Failed to create namespace: {create_error}") + logger.error(f"Failed to create namespace after prior error: {create_error}") raise🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@kaizen/frontend/mcp/mcp_server.py` around lines 49 - 56, The current broad "except Exception as e" around the namespace check can mask real connection/back-end errors and then try to create the namespace; change this to only attempt creation when the failure is a namespace-missing error (catch the specific NamespaceNotFoundException from your client) and otherwise log the original exception with full context and re-raise; if you cannot import that specific exception type, at minimum replace logger.error(f"Error ensuring namespace: {e}") with logger.exception(...) to include traceback and then re-raise when the exception is not a NamespaceNotFoundException; update the block that calls _client.create_namespace(kaizen_config.namespace_id) so it runs only under the namespace-not-found branch and keep the create_error handling unchanged.kaizen/db/sqlite_manager.py (1)
29-32: Moveimport osto the top of the file.
osis a standard-library module with no risk of circular imports. Placing it inside__init__is unconventional and hides a dependency. Moving it to the top alongsidedatetime,logging, etc. is more idiomatic.Proposed fix
At the top of the file, add
osto the existing imports:import datetime import logging +import os import sqlite3 import threadingThen simplify
__init__:def __init__(self, db_path: str | None = None): - import os - self.db_path = db_path or os.environ.get("KAIZEN_SQLITE_PATH", "entities.sqlite.db")🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@kaizen/db/sqlite_manager.py` around lines 29 - 32, Move the local import of os out of the __init__ method and add it to the module-level imports alongside datetime and logging; then simplify __init__ to use the already-imported os when resolving db_path (db_path or os.environ.get("KAIZEN_SQLITE_PATH", "entities.sqlite.db")). Ensure no other uses of os remain inside methods so the dependency is visible at the top of the file.tests/e2e/test_mcp.py (1)
165-172: Dynamic patch target selection is correct but fragile.Reading
os.environ["KAIZEN_BACKEND"]inside the test body works with the current fixture, but couples the test to an env-var side effect. Consider accessingrequest.paramvia a separate fixture or returning the backend type from themcpfixture (e.g., as part of a named tuple) for explicitness.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/e2e/test_mcp.py` around lines 165 - 172, The test reads os.environ["KAIZEN_BACKEND"] at runtime to choose patch_target, which couples the test to an env-var side effect; change the test to accept the backend explicitly from the test fixtures instead: have the mcp fixture (or a new fixture) return the backend type (or a namedtuple including backend) or access request.param in a parametrized fixture, then use that backend value to set patch_target ("kaizen.backend.filesystem.resolve_conflicts" vs "kaizen.backend.milvus.resolve_conflicts") and call patch(patch_target) — update references to KAIZEN_BACKEND, patch_target, and the mcp fixture/request.param accordingly.
Contributor
Wow! 👏 |
vinodmut
approved these changes
Feb 18, 2026
There was a problem hiding this comment.
♻️ Duplicate comments (1)
tests/e2e/test_mcp.py (1)
15-105:⚠️ Potential issue | 🟠 MajorRestore env vars after the fixture to avoid cross‑test leakage.
Direct
os.environmutation isn’t restored on teardown, so later tests can inherit KAIZEN_* values. Usingpytest’smonkeypatchmakes teardown automatic.✅ Proposed fix (use monkeypatch)
`@pytest.fixture`(params=["milvus", "filesystem"]) -def mcp(request, tmp_path): +def mcp(request, tmp_path, monkeypatch): backend_type = request.param - os.environ["KAIZEN_NAMESPACE_ID"] = "test" - os.environ["KAIZEN_BACKEND"] = backend_type + monkeypatch.setenv("KAIZEN_NAMESPACE_ID", "test") + monkeypatch.setenv("KAIZEN_BACKEND", backend_type) # Common SQLite setup - os.environ["KAIZEN_SQLITE_PATH"] = str(tmp_path / f"test_{backend_type}.sqlite.db") + monkeypatch.setenv("KAIZEN_SQLITE_PATH", str(tmp_path / f"test_{backend_type}.sqlite.db")) # Backend-specific setup milvus_db_file = None original_milvus_uri = None if backend_type == "milvus": # Use a unique DB file inside tmp_path for each test to avoid socket/locking issues and ensure cleanup milvus_db_file = str(tmp_path / f"test_{uuid.uuid4().hex[:8]}.db") original_milvus_uri = milvus_client_settings.uri milvus_client_settings.uri = milvus_db_file # Note: currently milvus-lite creates a file, not a dir for the uri # We set KAIZEN_URI just in case, though settings init should pick it up if we did it before elif backend_type == "filesystem": - os.environ["KAIZEN_DATA_DIR"] = str(tmp_path) + monkeypatch.setenv("KAIZEN_DATA_DIR", str(tmp_path))🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/e2e/test_mcp.py` around lines 15 - 105, The fixture mcp mutates os.environ directly which leaks into other tests; change the fixture signature to accept pytest's monkeypatch (def mcp(request, tmp_path, monkeypatch):) and replace every os.environ[...] = ... with monkeypatch.setenv(...) for KAIZEN_NAMESPACE_ID, KAIZEN_BACKEND, KAIZEN_SQLITE_PATH and KAIZEN_DATA_DIR (only for the filesystem branch); leave the milvus_client_settings.uri override as-is but keep restoring original_milvus_uri at teardown, and remove any manual env cleanup since monkeypatch will automatically restore env vars after the fixture ends; ensure references to the fixture (mcp) and module resets (mcp_server_module._client, mcp_server_module._namespace_initialized) remain unchanged.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Duplicate comments:
In `@tests/e2e/test_mcp.py`:
- Around line 15-105: The fixture mcp mutates os.environ directly which leaks
into other tests; change the fixture signature to accept pytest's monkeypatch
(def mcp(request, tmp_path, monkeypatch):) and replace every os.environ[...] =
... with monkeypatch.setenv(...) for KAIZEN_NAMESPACE_ID, KAIZEN_BACKEND,
KAIZEN_SQLITE_PATH and KAIZEN_DATA_DIR (only for the filesystem branch); leave
the milvus_client_settings.uri override as-is but keep restoring
original_milvus_uri at teardown, and remove any manual env cleanup since
monkeypatch will automatically restore env vars after the fixture ends; ensure
references to the fixture (mcp) and module resets (mcp_server_module._client,
mcp_server_module._namespace_initialized) remain unchanged.
This was referenced Feb 23, 2026
Merged
9 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description:
Sync & Conflict Resolution: Merged latest changes from main and resolved conflicts in sqlite_manager.py, mcp_server.py, and pyproject.toml.
MCP Stability: Implemented lazy initialization for the KaizenClient in the MCP server to prevent race conditions with the SSE transport.
Test Isolation: Fixed data leakage in tests by correctly using tmp_path and the KAIZEN_DATA_DIR environment variable.
Code Quality: Formatted the codebase with ruff and fixed all mypy type checking errors (including installing types-requests).
Verification:
All 18 tests (MCP E2E and Milvus Unit) passed successfully.
ruff format and ruff check pass.
mypy passes with 0 errors.
Summary by CodeRabbit
New Features
Bug Fixes
Tests