codekingpro/portable-devtools

1"""Context-var storage utilities for Claude Agent SDK tracing.2 3This module stores the *parent run tree* — the root chain span opened by4``_traced_receive_response`` — for direct/default-session hook calls. Real5instrumented client sessions primarily parent hook spans via6``SessionState.root_run``.7 8A :class:`contextvars.ContextVar` is used so concurrent conversations on the9same thread (e.g. multiple ``ClaudeSDKClient`` instances driven by10``asyncio.gather``) each see their own parent run tree.11"""12 13from contextvars import ContextVar14from typing import Any, Optional15 16_parent_run_tree: ContextVar[Optional[Any]] = ContextVar(17    "langsmith_claude_agent_parent_run_tree", default=None18)19 20 21def set_parent_run_tree(run_tree: Any) -> Any:22    """Bind *run_tree* to the current context and return a reset token."""23    return _parent_run_tree.set(run_tree)24 25 26def clear_parent_run_tree(token: Any = None) -> None:27    """Reset the parent run tree in the current context.28 29    If a *token* from :func:`set_parent_run_tree` is provided, it is used30    to restore the previous value; otherwise the context is cleared.31    """32    if token is not None:33        try:34            _parent_run_tree.reset(token)35        except ValueError:36            _parent_run_tree.set(None)37    else:38        _parent_run_tree.set(None)39 40 41def get_parent_run_tree() -> Any:42    """Return the parent run tree bound to the current context."""43    return _parent_run_tree.get()44