refactor(cli,compiler): close litellm async clients inside the compile loop

#91 closed litellm's cached async clients from the CLI layer in a separate
asyncio.run(), i.e. after the loop that created them was already torn down,
with the error swallowed by a bare except. Close them in the same loop that
created them instead, by moving the cleanup into the compile coroutines.

- add compiler._close_async_llm_clients() (best-effort, logs at debug)
- call it from compile_short_doc / compile_long_doc finally blocks
- drop cli._close_litellm_async_clients() and the three try/finally wraps
- drop the no-op cleanup around index_long_document (indexer uses no litellm)
- revert the test_add_command assertion to assert_called_once()

Author: KylinMountain

openkb/agent/compiler.py

@@ -359,6 +359,23 @@ async def _llm_call_async(model: str, messages: list[dict], step_name: str, **kw
     return content.strip()
 
 
+async def _close_async_llm_clients() -> None:
+    """Close LiteLLM's cached async (aiohttp) clients for the current loop.
+
+    LiteLLM caches its async clients per event loop. ``add_single_file`` runs
+    each doc in its own ``asyncio.run`` loop, so without this the clients are
+    orphaned when the loop is torn down and their connections pile up in
+    CLOSE-WAIT, leaking sockets/FDs across a long ingest. Call this from a
+    ``finally`` inside the compile coroutines so the clients are closed in the
+    same loop that created them. Best-effort: never raises, so cleanup can't
+    mask a real compilation error or break ingest.
+    """
+    try:
+        await litellm.close_litellm_async_clients()
+    except Exception:
+        logger.debug("litellm async client cleanup failed", exc_info=True)
+
+
 def _warn_if_truncated(response, step_name: str, max_tokens: int | None) -> None:
     """Emit a warning when the LLM hit the max_tokens cap.
 
@@ -1981,11 +1998,16 @@ async def compile_short_doc(
         summary = summary_raw
 
     # --- Steps 2-4: Concept plan → generate/update → summary rewrite → index ---
-    await _compile_concepts(
-        wiki_dir, kb_dir, model, system_msg, doc_msg,
-        summary, doc_name, max_concurrency, doc_brief=doc_brief,
-        doc_type="short", rewrite_summary=True, entity_types=entity_types,
-    )
+    try:
+        await _compile_concepts(
+            wiki_dir, kb_dir, model, system_msg, doc_msg,
+            summary, doc_name, max_concurrency, doc_brief=doc_brief,
+            doc_type="short", rewrite_summary=True, entity_types=entity_types,
+        )
+    finally:
+        # Close per-loop litellm async clients before asyncio.run tears this
+        # loop down, to avoid the CLOSE-WAIT/FD leak across a long ingest.
+        await _close_async_llm_clients()
 
 
 async def compile_long_doc(
@@ -2026,8 +2048,13 @@ async def compile_long_doc(
     overview = _llm_call(model, [system_msg, doc_msg], "overview")
 
     # --- Steps 2-4: Concept plan → generate/update → index ---
-    await _compile_concepts(
-        wiki_dir, kb_dir, model, system_msg, doc_msg,
-        overview, doc_name, max_concurrency, doc_brief=doc_description,
-        doc_type="pageindex", entity_types=entity_types,
-    )
+    try:
+        await _compile_concepts(
+            wiki_dir, kb_dir, model, system_msg, doc_msg,
+            overview, doc_name, max_concurrency, doc_brief=doc_description,
+            doc_type="pageindex", entity_types=entity_types,
+        )
+    finally:
+        # Close per-loop litellm async clients before asyncio.run tears this
+        # loop down, to avoid the CLOSE-WAIT/FD leak across a long ingest.
+        await _close_async_llm_clients()

openkb/cli.py

@@ -259,20 +259,6 @@ def _clear_existing_skill_dir(kb_dir: Path, name: str) -> None:
         shutil.rmtree(target)
 
 
-def _close_litellm_async_clients() -> None:
-    """Best-effort cleanup of cached LiteLLM async clients.
-
-    LiteLLM caches aiohttp clients per event loop; with ``asyncio.run`` creating
-    a fresh loop per doc, the old clients' connections linger in CLOSE-WAIT and
-    accumulate sockets/FDs over a long ingest. Closing them after each
-    compile/index frees those connections. Cleanup must never break ingest.
-    """
-    try:
-        asyncio.run(litellm.close_litellm_async_clients())
-    except Exception:
-        pass
-
-
 def add_single_file(file_path: Path, kb_dir: Path) -> Literal["added", "skipped", "failed"]:
     """Convert, index, and compile a single document into the knowledge base.
 
@@ -321,10 +307,7 @@ def add_single_file(file_path: Path, kb_dir: Path) -> Literal["added", "skipped"
         click.echo(f"  Long document detected — indexing with PageIndex...")
         try:
             from openkb.indexer import index_long_document
-            try:
-                index_result = index_long_document(result.raw_path, kb_dir)
-            finally:
-                _close_litellm_async_clients()
+            index_result = index_long_document(result.raw_path, kb_dir)
         except Exception as exc:
             click.echo(f"  [ERROR] Indexing failed: {exc}")
             logger.debug("Indexing traceback:", exc_info=True)
@@ -334,13 +317,10 @@ def add_single_file(file_path: Path, kb_dir: Path) -> Literal["added", "skipped"
         click.echo(f"  Compiling long doc (doc_id={index_result.doc_id})...")
         for attempt in range(2):
             try:
-                try:
-                    asyncio.run(
-                        compile_long_doc(doc_name, summary_path, index_result.doc_id, kb_dir, model,
-                                         doc_description=index_result.description)
-                    )
-                finally:
-                    _close_litellm_async_clients()
+                asyncio.run(
+                    compile_long_doc(doc_name, summary_path, index_result.doc_id, kb_dir, model,
+                                     doc_description=index_result.description)
+                )
                 break
             except Exception as exc:
                 if attempt == 0:
@@ -354,10 +334,7 @@ def add_single_file(file_path: Path, kb_dir: Path) -> Literal["added", "skipped"
         click.echo(f"  Compiling short doc...")
         for attempt in range(2):
             try:
-                try:
-                    asyncio.run(compile_short_doc(doc_name, result.source_path, kb_dir, model))
-                finally:
-                    _close_litellm_async_clients()
+                asyncio.run(compile_short_doc(doc_name, result.source_path, kb_dir, model))
                 break
             except Exception as exc:
                 if attempt == 0:

tests/test_add_command.py

@@ -145,10 +145,5 @@ def test_add_short_doc_runs_compiler(self, tmp_path):
              patch("openkb.cli.convert_document", return_value=mock_result), \
              patch("openkb.cli.asyncio.run") as mock_arun:
             result = runner.invoke(cli, ["add", str(doc)])
-            # asyncio.run drives both the compile and the post-compile
-            # litellm async-client cleanup, so it is called more than once;
-            # assert the compiler coroutine itself was run.
-            assert mock_arun.called
-            ran = [c.args[0] for c in mock_arun.call_args_list if c.args]
-            assert any(getattr(co, "__name__", "") == "compile_short_doc" for co in ran)
+            mock_arun.assert_called_once()
             assert "OK" in result.output