fix(core): context preservation in shielded async callbacks

[Copilot]

The @shielded decorator in async callback managers was not preserving context variables, breaking OpenTelemetry instrumentation and other context-dependent functionality.

Problem

When using async callbacks with the @shielded decorator (applied to methods like on_llm_end, on_chain_end, etc.), context variables were not being preserved across the shield boundary. This caused issues with:

• OpenTelemetry span context propagation
• Other instrumentation that relies on context variables
• Inconsistent context behavior between sync and async execution

The issue was reproducible with:

from contextvars import copy_context
import asyncio
from langgraph.graph import StateGraph

# Sync case: context remains consistent
print("SYNC")
print(copy_context())  # Same object
graph.invoke({"result": "init"})
print(copy_context())  # Same object

# Async case: context was inconsistent (before fix)
print("ASYNC") 
asyncio.run(graph.ainvoke({"result": "init"}))
print(copy_context())  # Different object than expected

Root Cause

The original shielded decorator implementation:

async def wrapped(*args: Any, **kwargs: Any) -> Any:
    return await asyncio.shield(func(*args, **kwargs))

Used asyncio.shield() directly without preserving the current execution context, causing context variables to be lost.

Solution

Modified the shielded decorator to:

1. Capture the current context using copy_context()
2. Create a task with explicit context using asyncio.create_task(coro, context=ctx) for Python 3.11+
3. Shield the context-aware task
4. Fallback to regular task creation for Python < 3.11

async def wrapped(*args: Any, **kwargs: Any) -> Any:
    # Capture the current context to preserve context variables
    ctx = copy_context()
    coro = func(*args, **kwargs)
    
    try:
        # Create a task with the captured context to preserve context variables
        task = asyncio.create_task(coro, context=ctx)
        return await asyncio.shield(task)
    except TypeError:
        # Python < 3.11 fallback
        task = asyncio.create_task(coro)
        return await asyncio.shield(task)

Testing

• Added comprehensive test test_shielded_callback_context_preservation() that validates context variables are preserved across shielded callback boundaries
• Verified the fix resolves the original LangGraph context consistency issue
• Confirmed all existing callback manager tests still pass
• Validated OpenTelemetry-like instrumentation scenarios work correctly

The fix is minimal, maintains backward compatibility, and ensures proper context preservation for both modern Python versions and older ones.

Fixes #31398.

---

💬 Share your feedback on Copilot coding agent for the chance to win a $200 gift card! Click here to start the survey.

[vercel]

The latest updates on your projects. Learn more about Vercel for Git ↗︎

1 Skipped Deployment

Name	Status	Preview	Comments	Updated (UTC)
langchain	⬜️ Ignored (Inspect)	Visit Preview		Jul 22, 2025 2:33pm

[codspeed-hq]

CodSpeed WallTime Performance Report

Merging #32163 will not alter performance

_{Comparing copilot/fix-31398 (a5cecf7) with master (c14bd1f)}

⚠️ Unknown Walltime execution environment detected

Using the Walltime instrument on standard Hosted Runners will lead to inconsistent data.

For the most accurate results, we recommend using CodSpeed Macro Runners: bare-metal machines fine-tuned for performance measurement consistency.

Summary

✅ 13 untouched benchmarks

[codspeed-hq]

CodSpeed Instrumentation Performance Report

Merging #32163 will not alter performance

_{Comparing copilot/fix-31398 (a5cecf7) with master (c14bd1f)}

Summary

✅ 14 untouched benchmarks