⚡️ Speed up method NativeCallbackHandler._resolve_parent_span_id by 35% in PR #12012 (LE-458)

[codeflash-ai]

⚡️ This pull request contains optimizations for PR #12012

If you approve this dependent PR, these changes will be merged into the original PR branch LE-458.

This PR will be automatically closed if the original PR is merged.

---

📄 35% (0.35x) speedup for NativeCallbackHandler._resolve_parent_span_id in src/backend/base/langflow/services/tracing/native_callback.py

⏱️ Runtime : 1.42 milliseconds → 1.05 milliseconds (best of 64 runs)

📝 Explanation and details

The optimization replaces two consecutive dictionary operations in _resolve_parent_span_id (a membership test parent_run_id in self._spans followed by a lookup inside _get_span_id) with a single dict.get() call and direct key access. Line profiler shows the hot path (3005 hits) dropped from 4.58 ms to 4.50 ms for the lookup plus 0.72 ms for the None-check, eliminating the redundant second lookup that previously cost 6.96 ms in the _get_span_id call. This reduces total function time by 42% (12.17 ms → 7.05 ms) because each parent_run_id resolution now touches the dictionary once instead of twice, with no behavioral changes or trade-offs.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 3012 Passed
⏪ Replay Tests	🔘 None Found
🔎 Concolic Coverage Tests	🔘 None Found
📊 Tests Coverage	100.0%

🌀 Click to see Generated Regression Tests

from datetime import datetime, timezone
from uuid import UUID, uuid4

# imports
import pytest  # used for our unit tests
from langflow.services.tracing.native_callback import NativeCallbackHandler


def test_returns_parent_span_when_parent_run_id_is_none():
    # Create a NativeCallbackHandler instance with a known parent_span_id.
    # We pass None for the tracer because the handler stores it but the method under test
    # does not interact with the tracer object.
    parent_span = UUID(int=1)  # deterministic UUID for test stability
    handler = NativeCallbackHandler(tracer=None, parent_span_id=parent_span)

    # When parent_run_id is None, the method should simply return the handler's parent_span_id.
    codeflash_output = handler._resolve_parent_span_id(None); result = codeflash_output


def test_returns_span_id_when_parent_run_id_present_in_spans():
    # Create handler with a fallback parent_span_id that should NOT be returned
    fallback_parent = UUID(int=999)
    handler = NativeCallbackHandler(tracer=None, parent_span_id=fallback_parent)

    # Prepare a run_id and a corresponding span_id and insert into the handler's _spans map.
    run_id = UUID(int=2)
    expected_span_id = UUID(int=123456)
    # The implementation expects a mapping; provide the typical keys recorded by _get_span_id.
    handler._spans[run_id] = {"span_id": expected_span_id, "start_time": datetime.now(timezone.utc)}

    # Because run_id is present in _spans, _resolve_parent_span_id should fetch the stored span_id.
    codeflash_output = handler._resolve_parent_span_id(run_id); result = codeflash_output


def test_returns_fallback_parent_span_when_parent_run_id_not_in_spans():
    # Handler with a known parent_span_id to fall back to.
    fallback_parent = UUID(int=42)
    handler = NativeCallbackHandler(tracer=None, parent_span_id=fallback_parent)

    # Use a run_id that is NOT in handler._spans.
    absent_run_id = UUID(int=9999)

    # Since the run_id is absent, the method should return the handler's parent_span_id.
    codeflash_output = handler._resolve_parent_span_id(absent_run_id); result = codeflash_output


def test_malformed_spans_entry_raises_key_error_when_span_id_missing():
    # If the _spans dict contains an entry for run_id but that entry
    # lacks the 'span_id' key, _resolve_parent_span_id will end up raising a KeyError
    # because _get_span_id returns self._spans[run_id]['span_id'] without creating one.
    handler = NativeCallbackHandler(tracer=None, parent_span_id=None)

    run_id = UUID(int=7)
    # Intentionally insert a malformed entry (no 'span_id' key).
    handler._spans[run_id] = {"start_time": datetime.now(timezone.utc)}

    # Expect a KeyError when resolving because 'span_id' key is missing.
    with pytest.raises(KeyError):
        handler._resolve_parent_span_id(run_id)


def test_non_uuid_parent_run_id_returns_fallback_parent():
    # Using a non-UUID, e.g. a string, should behave as "not present in _spans"
    # because _spans' keys are UUID instances. The method should then return
    # the handler's parent_span_id.
    fallback_parent = UUID(int=555)
    handler = NativeCallbackHandler(tracer=None, parent_span_id=fallback_parent)

    # A non-UUID truthy value
    non_uuid_key = "not-a-uuid"

    codeflash_output = handler._resolve_parent_span_id(non_uuid_key); result = codeflash_output


def test_resolve_parent_span_id_with_large_spans_map_returns_expected_for_many_entries():
    # Create a handler with a fallback parent span id.
    fallback = UUID(int=314159)
    handler = NativeCallbackHandler(tracer=None, parent_span_id=fallback)

    # Populate handler._spans with 1000 deterministic entries.
    n = 1000
    # Use deterministic UUIDs for reproducibility in tests: UUID(int=i) for keys and values.
    for i in range(1, n + 1):
        run_id = UUID(int=i)
        span_id = UUID(int=100000 + i)
        # record both span_id and start_time to mimic real entries
        handler._spans[run_id] = {"span_id": span_id, "start_time": datetime.now(timezone.utc)}

    # Verify resolution for each existing run_id returns the correct stored span_id.
    for i in range(1, n + 1):
        run_id = UUID(int=i)
        expected_span = UUID(int=100000 + i)
        codeflash_output = handler._resolve_parent_span_id(run_id); result = codeflash_output

    # Verify that for many absent run_ids the fallback parent is returned quickly and consistently.
    # We'll test 1000 absent keys (offset by n+1..2n) to ensure consistent behavior under load.
    for i in range(n + 1, 2 * n + 1):
        absent_run_id = UUID(int=2000000 + i)  # ensure these are not colliding with inserted keys
        codeflash_output = handler._resolve_parent_span_id(absent_run_id); result = codeflash_output


def test_repeated_calls_do_not_mutate_spans_for_absent_run_id():
    # Ensure calling _resolve_parent_span_id with a run_id that is absent does not create
    # a new entry in _spans (the method should not call _get_span_id in that case).
    fallback = UUID(int=777)
    handler = NativeCallbackHandler(tracer=None, parent_span_id=fallback)

    absent_run_id = UUID(int=888)

    # Call multiple times in a loop to simulate repeated resolution attempts.
    for _ in range(1000):
        codeflash_output = handler._resolve_parent_span_id(absent_run_id); result = codeflash_output
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
from datetime import datetime, timezone
from uuid import UUID, uuid4

# imports
import pytest
from langflow.services.tracing.native import NativeTracer
from langflow.services.tracing.native_callback import NativeCallbackHandler


# fixtures
@pytest.fixture
def mock_tracer():
    """Create a real NativeTracer instance for testing."""
    return NativeTracer()


@pytest.fixture
def callback_handler(mock_tracer):
    """Create a NativeCallbackHandler with a real tracer."""
    return NativeCallbackHandler(tracer=mock_tracer)


@pytest.fixture
def callback_handler_with_parent(mock_tracer):
    """Create a NativeCallbackHandler with an initial parent span ID."""
    parent_id = uuid4()
    return NativeCallbackHandler(tracer=mock_tracer, parent_span_id=parent_id)

To edit these changes git checkout codeflash/optimize-pr12012-2026-03-03T21.17.19 and push.

[github-actions]

Frontend Unit Test Coverage Report

Coverage Summary

Lines	Statements	Branches	Functions
	22.87% (7981/34890)	15.45% (4224/27334)	15.6% (1147/7348)

Unit Test Results

Tests	Skipped	Failures	Errors	Time
2611	0 💤	0 ❌	0 🔥	44.526s ⏱️

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 37.05%. Comparing base (b9c7e67) to head (fb49944).

❌ Your project status has failed because the head coverage (41.49%) is below the target coverage (60.00%). You can increase the head coverage or adjust the target coverage.

Additional details and impacted files

@@           Coverage Diff           @@
##           LE-458   #12019   +/-   ##
=======================================
  Coverage   37.05%   37.05%           
=======================================
  Files        1588     1588           
  Lines       77974    77967    -7     
  Branches    11803    11800    -3     
=======================================
+ Hits        28893    28894    +1     
+ Misses      47460    47454    -6     
+ Partials     1621     1619    -2     

Flag	Coverage Δ
backend	57.34% <100.00%> (+0.02%)	⬆️
frontend	20.49% <ø> (-0.01%)	⬇️
lfx	41.49% <ø> (-0.01%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

Files with missing lines	Coverage Δ
.../base/langflow/services/tracing/native_callback.py	91.22% <100.00%> (+0.10%)	⬆️

... and 11 files with indirect coverage changes

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.