Skip to content

fix: capture init task ref + cancel on close for Neo4jDriver and FalkorDriver#1636

Open
nandanadileep wants to merge 3 commits into
getzep:mainfrom
nandanadileep:fix/orphan-init-task-drivers
Open

fix: capture init task ref + cancel on close for Neo4jDriver and FalkorDriver#1636
nandanadileep wants to merge 3 commits into
getzep:mainfrom
nandanadileep:fix/orphan-init-task-drivers

Conversation

@nandanadileep

Copy link
Copy Markdown

Summary

Both Neo4jDriver.__init__ and FalkorDriver.__init__ schedule build_indices_and_constraints() via loop.create_task() without capturing the task reference. This creates a fire-and-forget background task whose lifecycle is not managed.

Problem

  1. Orphan task: The task reference is discarded, so there's no way to cancel or await it during shutdown.
  2. Race condition: Callers that explicitly await driver.build_indices_and_constraints() race with the constructor-fired task, producing EquivalentSchemaRuleAlreadyExists errors.
  3. Unhandled exceptions: If the task fails (e.g., connection error), "Task exception was never retrieved" warnings appear.
  4. Mid-tx state: Under pytest-asyncio, when the test function returns and the event loop closes, the background task can be cancelled mid-transaction.

Fix

  • Capture the task in self._init_task: asyncio.Task | None
  • In close()\, cancel the task if still running and await it with suppress(asyncio.CancelledError)
  • Removed redundant comments and moved the import asyncio to the top of neo4j_driver.py

Testing

  • make lint passes
  • make typecheck passes
  • DISABLE_FALKORDB=1 DISABLE_KUZU=1 DISABLE_NEPTUNE=1 uv run pytest -m 'not integration' — 165 passed, 2 skipped (FalkorDB integration), 1 pre-existing Neo4j auth error (no local Neo4j)

Closes #1513

…orDriver

Both Neo4jDriver and FalkorDriver __init__ scheduled
build_indices_and_constraints() via loop.create_task() without
capturing the task reference. The orphan task could:

1. Race with explicit build_indices_and_constraints() calls
2. Produce 'Task exception was never retrieved' warnings on failure
3. Leave the connection in a mid-transaction state when the driver
   is closed concurrently

Now captures the task in self._init_task and cancels it in close(),
with suppress(CancelledError) to handle the cancellation cleanly.

Closes getzep#1513
@zep-cla-assistant

zep-cla-assistant Bot commented Jul 8, 2026

Copy link
Copy Markdown
Contributor

All contributors have signed the CLA ✍️ ✅
Posted by the CLA Assistant Lite bot.

@nandanadileep

Copy link
Copy Markdown
Author

I have read the CLA Document and I hereby sign the CLA behalf on myself, e-mail: nandanadileep29@gmail.com

@Naseem77

Copy link
Copy Markdown
Contributor

@nandanadileep Thanks for catching this! One gap: close() skips done() tasks, so a failed init task still leaks its exception. Calling task.exception() on done tasks too would complete it.

Previously, close() only cancelled and awaited tasks that were
still running, but skipped already-done tasks. If a done task
failed with an exception, the exception was never retrieved,
causing 'Task exception was never retrieved' warnings.

Now calls task.exception() on done tasks to consume any pending
exception and complete the task's lifecycle.
zep-cla-assistant Bot added a commit that referenced this pull request Jul 20, 2026

@prasmussen15 prasmussen15 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Please update close() to retrieve exceptions from completed init tasks too; skipping done() tasks still leaves a failed init task exception unobserved.

Task.exception() raises CancelledError if the task finished in a
cancelled state. If the init task was cancelled externally before
close() ran, the else branch would re-raise CancelledError out of
close(), breaking driver cleanup.

Skip exception retrieval for cancelled tasks; only call exception()
on tasks that completed normally or failed.
@nandanadileep

Copy link
Copy Markdown
Author

@prasmussen15 Thanks for the review! The requested change was already landed in d3fc190 (pushed Jul 20) — close() in both drivers now calls self._init_task.exception() on already-done tasks so a failed init task's exception is observed:

if self._init_task is not None:
    if not self._init_task.done():
        self._init_task.cancel()
        with suppress(asyncio.CancelledError):
            await self._init_task
    elif not self._init_task.cancelled():
        self._init_task.exception()

The review may have been against a stale diff. That said, it surfaced a related edge case: Task.exception() raises CancelledError if the done task was cancelled externally, which would escape close(). Commit 91124f5 guards against that with the cancelled() check above. Ready for re-review 🙏

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Neo4jDriver.__init__:57 schedules orphan create_task without cancel/await pair

3 participants