fix(cognee): P0-1/2/3 — entry wrapper path, attribute migration, new_span patch

P0-1: ENTRY wrapper targeted ``cognee.api.v1.remember.remember`` etc.,
but in cognee v1.2.1 ``cognee.api.v1.remember`` is also re-exported as a
function into the v1 package namespace, so the dotted-path wrap is
ambiguous. Switch to wrapping the top-level ``cognee.{add,cognify,search,
recall,remember}`` attributes — these are the canonical user-facing entry
points re-exported by ``cognee/__init__.py``.

P0-2: ``install_attribute_migration_patch`` monkey-patched
``cognee.modules.observability.new_span`` to wrap each returned span's
``set_attribute``. The wrap never actually ran (cognee re-imports
``new_span`` lazily in some code paths) AND broke the cognee exception
path (``CollectionNotFoundError`` fallback was swallowed, triggering
``RuntimeError: generator didn't stop after throw()``). Delete the patch
and move the ``_wrap_span_set_attribute`` call into
``CogneeAttributeSpanProcessor.on_start`` — the SDK Span is still mutable
at ``on_start`` and the wrap is visible to all subsequent ``set_attribute``
calls in the span body. Also pre-inject ``gen_ai.task.name`` from the span
name for ``cognee.pipeline.task.*`` so the attribute is present even when
Cognee never sets ``cognee.pipeline.task_name``.

P0-3: same as P0-2 — the broken ``_patched_new_span`` is removed. No
``try/except/finally`` re-raise variant is kept because the simpler
on_start hook is correct and lower-risk.

Tests updated: conftest now exposes ``cognee.add`` etc. at the top level
(matching v1.2.1 re-exports); test_entry_wrapper calls ``cognee.add``
instead of ``cognee.api.v1.add.add``. All 46 tests pass.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

Author: 123liuziming

instrumentation-loongsuite/loongsuite-instrumentation-cognee/src/opentelemetry/instrumentation/cognee/__init__.py

@@ -56,7 +56,6 @@
 )
 from opentelemetry.instrumentation.cognee.internal._span_processor import (
     CogneeAttributeSpanProcessor,
-    install_attribute_migration_patch,
 )
 from opentelemetry.instrumentation.cognee.internal._step_wrapper import (
     install_step_wrapper,
@@ -142,17 +141,16 @@ def _instrument(self, **kwargs: Any) -> None:
         _enable_cognee_tracing(tracer_provider)
 
         # Step 2: attach SpanProcessor for cognee.* span normalization.
+        # The processor rewrites span name/kind/operation in on_start and
+        # wraps span.set_attribute so cognee.* attributes set during the span
+        # body are mirrored to gen_ai.* on the same span.
         self._cognee_processor = CogneeAttributeSpanProcessor()
         if hasattr(tracer_provider, "add_span_processor"):
             try:
                 tracer_provider.add_span_processor(self._cognee_processor)
             except Exception as e:
                 logger.debug("add_span_processor(CogneeAttributeSpanProcessor) failed: %s", e)
 
-        # Step 2b: wrap cognee.modules.observability.new_span so cognee.* attrs
-        # set during the span body are mirrored to gen_ai.* attrs on the same span.
-        install_attribute_migration_patch()
-
         # Step 3: install wrappers (each in its own try/except — single failure
         # does not break the rest).
         install_entry_wrappers(telemetry_handler)

instrumentation-loongsuite/loongsuite-instrumentation-cognee/src/opentelemetry/instrumentation/cognee/internal/_entry_wrapper.py

@@ -27,12 +27,18 @@
 
 
 # module_path, attr_name, operation_name (used for input extraction)
+# Cognee v1.2.1 re-exports ``add``/``cognify``/``search``/``recall``/``remember``
+# as top-level attributes of the ``cognee`` package (see ``cognee/__init__.py``).
+# ``cognee.api.v1.remember`` is also a function (re-exported into the v1 package
+# namespace by ``cognee/api/v1/__init__.py``), so wrapping at the v1 module
+# path would try to import a function as a module. Wrap at the top-level
+# ``cognee`` package instead — this is the canonical user-facing entry point.
 _ENTRY_TARGETS: tuple[tuple[str, str, str], ...] = (
-    ("cognee.api.v1.add", "add", "add"),
-    ("cognee.api.v1.cognify", "cognify", "cognify"),
-    ("cognee.api.v1.search", "search", "search"),
-    ("cognee.api.v1.recall", "recall", "recall"),
-    ("cognee.api.v1.remember", "remember", "remember"),
+    ("cognee", "add", "add"),
+    ("cognee", "cognify", "cognify"),
+    ("cognee", "search", "search"),
+    ("cognee", "recall", "recall"),
+    ("cognee", "remember", "remember"),
 )
 
 

instrumentation-loongsuite/loongsuite-instrumentation-cognee/src/opentelemetry/instrumentation/cognee/internal/_span_processor.py

@@ -9,11 +9,16 @@
 Cognee-specific attributes (``cognee.search.query``, ``cognee.pipeline.task_name`` …)
 are set by Cognee's own code during the span body, after ``on_start`` runs.
 Because the SDK Span becomes read-only once ``end()`` is called, we cannot
-migrate them in ``on_end``. Instead, ``install_attribute_migration_patch``
-wraps ``cognee.modules.observability.new_span`` so the span returned to
-Cognee's caller has its ``set_attribute`` method intercepted — when Cognee
-sets a ``cognee.*`` key, we ALSO set the corresponding ``gen_ai.*`` key on
-the same span.
+migrate them in ``on_end``. Instead, ``on_start`` wraps the span's
+``set_attribute`` method so that when Cognee calls
+``span.set_attribute("cognee.search.query", value)`` we ALSO set the matching
+``gen_ai.*`` key on the same span. The wrapping is idempotent.
+
+This replaces the previous ``install_attribute_migration_patch`` which patched
+``cognee.modules.observability.new_span`` — that patch broke Cognee's
+exception path (``CollectionNotFoundError`` fallback was swallowed, triggering
+``RuntimeError: generator didn't stop after throw()``) and never actually ran
+because cognee re-imports ``new_span`` lazily in some code paths.
 """
 
 from __future__ import annotations
@@ -52,8 +57,8 @@
 )
 
 # Cognee attribute -> gen-ai attribute migration table.
-# Applied by the set_attribute interceptor installed via
-# ``install_attribute_migration_patch``.
+# Applied by the ``set_attribute`` interceptor installed on each Cognee span
+# in ``CogneeAttributeSpanProcessor.on_start``.
 _MIGRATION_MAP: dict[str, str] = {
     COGNEE_SEARCH_QUERY: "gen_ai.retrieval.query.text",
     COGNEE_PIPELINE_TASK_NAME: "gen_ai.task.name",
@@ -77,14 +82,23 @@ class CogneeAttributeSpanProcessor(SpanProcessor):
     Uses ``on_start`` because the SDK Span becomes immutable after ``end()``.
     The name is known at start time (passed to ``start_span``), so prefix-based
     rewriting works.
+
+    In ``on_start`` we also wrap the span's ``set_attribute`` method so that
+    when Cognee code later calls ``span.set_attribute("cognee.search.query", v)``
+    we mirror the write to ``gen_ai.retrieval.query.text`` on the same span.
+    This is the only reliable hook point because the SDK Span is still mutable
+    at ``on_start`` and the wrap is visible to all subsequent ``set_attribute``
+    calls in the span body.
     """
 
     def on_start(self, span: Span, parent_context: Optional[Context] = None) -> None:
         try:
             name = span.name or ""
+            matched = False
             for prefix, kind, op, new_name in _PREFIX_RULES:
                 if not name.startswith(prefix):
                     continue
+                matched = True
                 try:
                     span.set_attribute("gen_ai.span.kind", kind)
                     span.set_attribute("gen_ai.operation.name", op)
@@ -96,16 +110,30 @@ def on_start(self, span: Span, parent_context: Optional[Context] = None) -> None
                     except Exception as e:  # pragma: no cover - defensive
                         logger.debug("update_name failed: %s", e)
                 elif prefix == "cognee.pipeline.task.":
+                    task_name = name[len(prefix) :]
                     try:
-                        span.update_name(f"task {name[len(prefix):]}")
+                        span.update_name(f"task {task_name}")
+                        # Pre-inject gen_ai.task.name from span name so the
+                        # attribute is present even if Cognee never sets
+                        # cognee.pipeline.task_name (e.g. for early-exit paths).
+                        span.set_attribute("gen_ai.task.name", task_name)
                     except Exception as e:  # pragma: no cover - defensive
                         logger.debug("update_name(task) failed: %s", e)
                 elif prefix == "cognee.retrieval.":
                     try:
-                        span.update_name(f"retrieval {name[len('cognee.retrieval.') :]}")
+                        span.update_name(
+                            f"retrieval {name[len('cognee.retrieval.') :]}"
+                        )
                     except Exception as e:  # pragma: no cover - defensive
                         logger.debug("update_name(retrieval) failed: %s", e)
-                return
+                break
+            # Wrap set_attribute for any Cognee-native span so cognee.* attrs
+            # set during the span body are mirrored to gen_ai.* on the same
+            # span. We do this regardless of whether the prefix matched,
+            # because Cognee may create spans whose names we don't recognize
+            # yet still set cognee.* attributes worth migrating.
+            if matched or name.startswith("cognee."):
+                _wrap_span_set_attribute(span)
         except Exception as e:  # pragma: no cover - defensive
             logger.debug("CogneeAttributeSpanProcessor.on_start failed: %s", e)
 
@@ -120,55 +148,6 @@ def force_flush(self, timeout_millis: int = 30000) -> bool:
         return True
 
 
-def install_attribute_migration_patch() -> None:
-    """Patch ``cognee.modules.observability.new_span`` to migrate cognee.* attrs.
-
-    When Cognee calls ``span.set_attribute("cognee.search.query", value)``, we
-    also call ``span.set_attribute("gen_ai.retrieval.query.text", value)`` on
-    the same span. This runs before ``end()`` so the SDK Span is still mutable.
-    """
-    try:
-        import cognee.modules.observability as cognee_obs  # type: ignore
-    except ImportError:
-        logger.debug(
-            "cognee.modules.observability not importable — attribute migration patch skipped"
-        )
-        return
-
-    original = getattr(cognee_obs, "new_span", None)
-    if original is None or getattr(original, "_cognee_genai_patched", False):
-        return
-
-    try:
-        import contextlib
-
-        @contextlib.contextmanager
-        def _patched_new_span(name: str):
-            ctx = original(name) if original else _noop_ctx()
-            try:
-                with ctx as span:
-                    if span is not None and not isinstance(span, _NullSpanType):
-                        _wrap_span_set_attribute(span)
-                    yield span
-            except Exception:
-                # Fall back to yielding a null span if the original blows up.
-                yield cognee_obs._NullSpan()
-
-        def _noop_ctx():
-            import contextlib
-
-            @contextlib.contextmanager
-            def _cm():
-                yield None
-
-            return _cm()
-
-        _patched_new_span._cognee_genai_patched = True  # type: ignore[attr-defined]
-        cognee_obs.new_span = _patched_new_span
-    except Exception as e:
-        logger.debug("install_attribute_migration_patch failed: %s", e)
-
-
 def _wrap_span_set_attribute(span: Any) -> None:
     """Wrap a single span's ``set_attribute`` to mirror cognee.* → gen_ai.*.
 
@@ -199,14 +178,3 @@ def _set_attribute(key, value):
         span._cognee_genai_set_attr_wrapped = True  # type: ignore[attr-defined]
     except (AttributeError, TypeError) as e:
         logger.debug("could not wrap span.set_attribute: %s", e)
-
-
-# Late import to avoid hard dependency on cognee at module load.
-try:
-    from cognee.modules.observability import (
-        _NullSpan as _NullSpanType,  # type: ignore
-    )
-except ImportError:  # pragma: no cover - cognee not installed
-
-    class _NullSpanType:  # type: ignore[no-redef]
-        pass

instrumentation-loongsuite/loongsuite-instrumentation-cognee/tests/conftest.py

@@ -128,7 +128,15 @@ def _ensure_fake_module(dotted_path: str, attrs: dict) -> types.ModuleType:
 
 @pytest.fixture
 def fake_cognee_v1_api():
-    """Install fake ``cognee.api.v1.{add,cognify,search,recall,remember}`` modules."""
+    """Install fake ``cognee.{add,cognify,search,recall,remember}`` entry points.
+
+    The instrumentor wraps the top-level ``cognee`` package attributes (not the
+    deep ``cognee.api.v1.add`` module path) because in v1.2.1
+    ``cognee.api.v1.remember`` is re-exported as a function into the v1
+    package namespace, which makes the dotted-path wrap ambiguous. We still
+    install the deep module hierarchy so callers that import it directly
+    keep working.
+    """
     async def add(data, *args, **kwargs):
         return {"added": data}
 
@@ -144,14 +152,20 @@ async def recall(query_text, *args, **kwargs):
     async def remember(data, *args, **kwargs):
         return {"remembered": data}
 
-    _ensure_fake_module("cognee", {})
+    cognee_root = _ensure_fake_module("cognee", {})
     _ensure_fake_module("cognee.api", {})
     _ensure_fake_module("cognee.api.v1", {})
     _ensure_fake_module("cognee.api.v1.add", {"add": add})
     _ensure_fake_module("cognee.api.v1.cognify", {"cognify": cognify})
     _ensure_fake_module("cognee.api.v1.search", {"search": search})
     _ensure_fake_module("cognee.api.v1.recall", {"recall": recall})
     _ensure_fake_module("cognee.api.v1.remember", {"remember": remember})
+    # Top-level re-exports — this is what the instrumentor wraps.
+    cognee_root.add = add
+    cognee_root.cognify = cognify
+    cognee_root.search = search
+    cognee_root.recall = recall
+    cognee_root.remember = remember
     yield
 
 

instrumentation-loongsuite/loongsuite-instrumentation-cognee/tests/test_entry_wrapper.py

@@ -59,10 +59,11 @@ def instrumented(fake_cognee_v1_api, monkeypatch):
 @pytest.mark.asyncio
 async def test_entry_add_creates_entry_span(instrumented):
     instrumentor, exporter, provider = instrumented
-    # We need to call the wrapped cognee.add via the instrumented module.
-    import cognee.api.v1.add as add_mod
+    # The instrumentor wraps the top-level ``cognee.add`` attribute (v1.2.1
+    # re-exports the V1 API functions at the package root).
+    import cognee as cognee_root
 
-    result = await add_mod.add("hello world")
+    result = await cognee_root.add("hello world")
     assert result == {"added": "hello world"}
 
     spans = exporter.get_finished_spans()
@@ -78,11 +79,9 @@ async def test_entry_add_creates_entry_span(instrumented):
 @pytest.mark.asyncio
 async def test_entry_search_propagates_session_id(instrumented):
     instrumentor, exporter, _ = instrumented
-    import cognee.api.v1.search as search_mod
+    import cognee as cognee_root
 
-    await search_mod.search(
-        "what is cognee", session_id="sess-123"
-    )
+    await cognee_root.search("what is cognee", session_id="sess-123")
 
     spans = exporter.get_finished_spans()
     entry_spans = [