fix(langchain): Use instance-level wrapping for AgentMiddleware hooks#4000
Open
MaxStenklyft wants to merge 3 commits intotraceloop:mainfrom
Open
fix(langchain): Use instance-level wrapping for AgentMiddleware hooks#4000MaxStenklyft wants to merge 3 commits intotraceloop:mainfrom
MaxStenklyft wants to merge 3 commits intotraceloop:mainfrom
Conversation
Replace class-level wrapt.wrap_function_wrapper on AgentMiddleware hook methods with instance-level wrapping via AgentMiddleware.__init__. The previous approach replaced base class methods with FunctionWrapper descriptors, which broke Python identity checks used by LangGraph's create_agent (e.g. m.__class__.before_agent is not AgentMiddleware.before_agent). This caused LangGraph to add graph nodes for every middleware hook regardless of whether it was overridden, scaling graph size linearly with the number of hooks per middleware and making it easy to exceed recursion limits. The fix wraps __init__ to patch hook methods on each instance after construction, preserving the base class methods untouched so identity checks work correctly.
|
|
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (1)
📝 WalkthroughWalkthroughThe instrumentation now wraps Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes 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)
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
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In
`@packages/opentelemetry-instrumentation-langchain/opentelemetry/instrumentation/langchain/__init__.py`:
- Around line 307-312: Update the comment/docstring near the
wrap_function_wrapper call for AgentMiddleware.__init__ (the block using
wrap_function_wrapper, AgentMiddleware.__init__, and _middleware_init_wrapper)
to explicitly acknowledge the tradeoff: explain that instance-level wrapping was
chosen to preserve LangGraph identity checks (e.g., m.__class__.before_agent
comparisons), that subclasses that override __init__ without calling super()
will not be wrapped, and that this is intentional and acceptable given the
alternative breakage; adjust the wording/severity in the docstring/tests to mark
this as an intentional limitation rather than a bug.
- Around line 288-305: The code currently mutates middleware instances in
_middleware_init_wrapper by setattr but never cleans them up; add a module-level
weakref.WeakSet (e.g. _patched_middleware_instances) and inside
_middleware_init_wrapper add each created instance to that set after patching;
update uninstrument() to iterate over _patched_middleware_instances and for each
instance remove the shadowed hook attributes (for each hook_name in
sync_wrappers and async_wrappers, if hasattr(instance, hook_name) then
delattr(instance, hook_name)), then clear the WeakSet and proceed to unwrap
AgentMiddleware.__init__ as before so no pre-existing instances keep emitting
spans.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 5eed5cf8-5477-44a2-a747-ba760caeaf1a
📒 Files selected for processing (3)
packages/opentelemetry-instrumentation-langchain/opentelemetry/instrumentation/langchain/__init__.pypackages/opentelemetry-instrumentation-langchain/tests/test_langgraph.pypackages/opentelemetry-instrumentation-langchain/tests/test_middleware_identity.py
Author
|
Noticed this when upgraded to langchain 1.2. We started hitting GRAPH_RECURSION_LIMIT due to no-op nodes introduced by this package. This change should eliminate the no-op nodes by fixing the identity check. |
Track patched middleware instances in a WeakSet and remove shadowed hook attributes during uninstrument(). This ensures pre-existing instances stop emitting spans after teardown. Also document the intentional tradeoff that subclasses overriding __init__ without calling super() will not be instrumented.
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
packages/opentelemetry-instrumentation-langchain/tests/test_middleware_identity.py (1)
32-35: Consider adding an opt-inConsoleSpanExporterin this fixture for hierarchy debugging.This test suite targets span behavior; a debug-only console exporter helps triage failures quickly without changing assertions.
🛠️ Suggested adjustment
+import os from opentelemetry.sdk.trace import TracerProvider -from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export import ConsoleSpanExporter, SimpleSpanProcessor from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter @@ def _instrument(): exporter = InMemorySpanExporter() provider = TracerProvider() provider.add_span_processor(SimpleSpanProcessor(exporter)) + if os.getenv("OTEL_LANGCHAIN_TEST_DEBUG_SPANS") == "1": + provider.add_span_processor(SimpleSpanProcessor(ConsoleSpanExporter()))🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@packages/opentelemetry-instrumentation-langchain/tests/test_middleware_identity.py` around lines 32 - 35, The test fixture currently registers only an InMemorySpanExporter with TracerProvider (exporter = InMemorySpanExporter(), provider = TracerProvider(), provider.add_span_processor(SimpleSpanProcessor(exporter))) which makes debugging span hierarchy harder; add an opt-in ConsoleSpanExporter (or SimpleSpanProcessor(ConsoleSpanExporter())) to the provider when a debug flag is set (e.g., env var like OTEL_SPAN_DEBUG or a pytest marker), or combine processors via a MultiSpanProcessor so the InMemorySpanExporter remains for assertions while ConsoleSpanExporter prints spans for triage; update the fixture to check the opt-in flag and register the ConsoleSpanExporter alongside the existing SimpleSpanProcessor(exporter).
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In
`@packages/opentelemetry-instrumentation-langchain/tests/test_middleware_identity.py`:
- Around line 69-73: Replace the unreliable identity check between bound methods
in the test by asserting the hook name exists in the instance's __dict__ (i.e.,
check hook_name in m.__dict__) to detect per-instance shadowing (replace the
current instance_method is not class_method.__get__(m, type(m)) assertion); also
add ConsoleSpanExporter to the test fixture (add
provider.add_span_processor(SimpleSpanProcessor(ConsoleSpanExporter())) as
commented optional debugging aid) so spans can be inspected during debugging;
references: m, hook_name, __dict__, test_uninstrument_removes_instance_patches,
provider.add_span_processor(SimpleSpanProcessor(ConsoleSpanExporter())).
---
Nitpick comments:
In
`@packages/opentelemetry-instrumentation-langchain/tests/test_middleware_identity.py`:
- Around line 32-35: The test fixture currently registers only an
InMemorySpanExporter with TracerProvider (exporter = InMemorySpanExporter(),
provider = TracerProvider(),
provider.add_span_processor(SimpleSpanProcessor(exporter))) which makes
debugging span hierarchy harder; add an opt-in ConsoleSpanExporter (or
SimpleSpanProcessor(ConsoleSpanExporter())) to the provider when a debug flag is
set (e.g., env var like OTEL_SPAN_DEBUG or a pytest marker), or combine
processors via a MultiSpanProcessor so the InMemorySpanExporter remains for
assertions while ConsoleSpanExporter prints spans for triage; update the fixture
to check the opt-in flag and register the ConsoleSpanExporter alongside the
existing SimpleSpanProcessor(exporter).
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: d75f61e6-b0c2-42e8-9546-76e24cc51e15
📒 Files selected for processing (2)
packages/opentelemetry-instrumentation-langchain/opentelemetry/instrumentation/langchain/__init__.pypackages/opentelemetry-instrumentation-langchain/tests/test_middleware_identity.py
🚧 Files skipped from review as they are similar to previous changes (1)
- packages/opentelemetry-instrumentation-langchain/opentelemetry/instrumentation/langchain/init.py
Comment on lines
+69
to
+73
| instance_method = getattr(m, hook_name) | ||
| class_method = getattr(m.__class__, hook_name) | ||
| # Instance attribute should shadow the class method (instrumented) | ||
| assert instance_method is not class_method.__get__(m, type(m)), ( | ||
| f"{hook_name} on instance should be instrumented (shadowed)" |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Verifies why `is not` on separately-bound methods is not a safe instrumentation signal.
uv run python - <<'PY'
class A:
def f(self): pass
a = A()
m1 = getattr(a, "f")
m2 = A.f.__get__(a, A)
print("m1 is m2 =", m1 is m2) # typically False even without patching
print("m1.__func__ is m2.__func__ =", m1.__func__ is m2.__func__) # typically True
PYRepository: traceloop/openllmetry
Length of output: 107
🏁 Script executed:
#!/bin/bash
# Test Python method binding behavior without uv
python3 <<'PY'
class A:
def f(self): pass
a = A()
m1 = getattr(a, "f")
m2 = A.f.__get__(a, A)
print("m1 is m2 =", m1 is m2)
print("m1.__func__ is m2.__func__ =", m1.__func__ is m2.__func__)
# Now test with instance attribute shadowing
a.f = lambda: "shadowed"
m3 = getattr(a, "f")
m4 = A.f.__get__(a, A)
print("\nAfter shadowing:")
print("m3 is m4 =", m3 is m4)
print("'f' in a.__dict__ =", 'f' in a.__dict__)
PYRepository: traceloop/openllmetry
Length of output: 175
🏁 Script executed:
# Read the actual test file to understand context
cat -n packages/opentelemetry-instrumentation-langchain/tests/test_middleware_identity.pyRepository: traceloop/openllmetry
Length of output: 4709
🏁 Script executed:
# Look for how hooks are instrumented in the codebase
rg -A 5 -B 5 "def _instrument" packages/opentelemetry-instrumentation-langchain/tests/test_middleware_identity.pyRepository: traceloop/openllmetry
Length of output: 356
The is not comparison on bound methods will always pass regardless of instrumentation.
Each call to getattr(instance, method) and class_method.__get__(instance, type(instance)) creates distinct bound method objects with different identities in Python, even when no patching occurs. This assertion cannot reliably detect whether hooks are actually shadowed on the instance.
The test should instead check for the presence of the hook in the instance's __dict__, which is the actual signal of instrumentation—this pattern is already used correctly in test_uninstrument_removes_instance_patches at line 83.
🔧 Proposed fix
def test_instance_hooks_are_instrumented(_instrument):
"""Instance-level hooks should be wrapped for tracing after construction."""
m = MyMiddleware()
for hook_name in ALL_HOOKS:
- instance_method = getattr(m, hook_name)
- class_method = getattr(m.__class__, hook_name)
- # Instance attribute should shadow the class method (instrumented)
- assert instance_method is not class_method.__get__(m, type(m)), (
- f"{hook_name} on instance should be instrumented (shadowed)"
- )
+ # Instrumentation shadows hooks on the instance.
+ assert hook_name in m.__dict__, (
+ f"{hook_name} should be in instance __dict__ after instrumentation"
+ )Additionally, per coding guidelines (**/*.py), add ConsoleSpanExporter to the fixture for debugging span hierarchy:
from opentelemetry.instrumentation.langchain import LangchainInstrumentor
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
+from opentelemetry.sdk.trace.export import ConsoleSpanExporter
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporterThen in the fixture, optionally add:
# Uncomment for debugging span hierarchy:
# provider.add_span_processor(SimpleSpanProcessor(ConsoleSpanExporter()))🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In
`@packages/opentelemetry-instrumentation-langchain/tests/test_middleware_identity.py`
around lines 69 - 73, Replace the unreliable identity check between bound
methods in the test by asserting the hook name exists in the instance's __dict__
(i.e., check hook_name in m.__dict__) to detect per-instance shadowing (replace
the current instance_method is not class_method.__get__(m, type(m)) assertion);
also add ConsoleSpanExporter to the test fixture (add
provider.add_span_processor(SimpleSpanProcessor(ConsoleSpanExporter())) as
commented optional debugging aid) so spans can be inspected during debugging;
references: m, hook_name, __dict__, test_uninstrument_removes_instance_patches,
provider.add_span_processor(SimpleSpanProcessor(ConsoleSpanExporter())).
Replace unreliable bound method identity comparison with __dict__ membership check. Python creates new bound method objects on every attribute access, so 'is not' always passes regardless of patching.
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.
Summary
_wrap_middleware_hooksuseswrapt.wrap_function_wrapperonAgentMiddlewarebase class methods (before_model,after_model, etc.). This replaces them withFunctionWrapperdescriptors, which return a newBoundFunctionWrapperon every attribute access. This breaks the identity checks in LangGraph'screate_agent(factory.py):The
is notcheck always evaluates toTrueafter wrapping, causing LangGraph to add graph nodes for every middleware hook regardless of whether it's actually overridden. This scales graph size linearly with the number of hooks per middleware, easily exceeding recursion limits.Fix
Replace class-level wrapping with instance-level wrapping via
AgentMiddleware.__init__. The wrapper patches hook methods on each instance after construction, leaving base class methods untouched so identity checks work correctly.Tests
test_middleware_identity.pywith two tests verifying base class identity is preserved and instance hooks are instrumentedtest_middleware_super_call_succeeds_despite_outer_failureto match new instance-level wrapping semanticsSummary by CodeRabbit
Bug Fixes
Tests