Fix feature_slug propagation to observe() spans

[Copilot]

Problem

feature_slug set via start_observe() was not propagating to child spans created with observe(), breaking feature tracking for API requests and nested operations.

with start_observe(name="handler", feature_slug="payments"):
    prompt = b.prompts.get_sync(slug="router")  # ❌ span missing feature_slug
    with observe(name="validation"):            # ❌ span missing feature_slug
        pass

Root Cause

start_observe() correctly stores feature_slug in OpenTelemetry context, and BasaltContextProcessor applies it to spans via on_start(). However, Observe class never reads from context - it only passes through explicit parameters.

Changes

Modified basalt/observability/api.py:

• Observe.__enter__(): Read feature_slug from OTel context before creating span
• Observe.__call__(): Consolidate context imports at method level, read feature_slug in both sync/async wrappers

# Inherit feature_slug from parent context if available
from .trace_context import FEATURE_SLUG_CONTEXT_KEY
current_feature_slug = otel_context.get_value(FEATURE_SLUG_CONTEXT_KEY)

# Pass to span creation
with _with_span_handle(..., feature_slug=current_feature_slug) as span:

Added tests in test_feature_slug_propagation.py:

• Context manager propagation (simulates API request spans)
• Decorator propagation (sync and async functions)

Impact

All spans within a start_observe() context now correctly inherit feature_slug:

• API request spans (trace_sync_request, trace_async_request)
• Manual observe() context managers
• @observe() decorated functions (sync/async)

This enables proper feature-level trace filtering and analytics in the Basalt UI.

Original prompt

Problem

The feature_slug attribute is not being propagated to spans created by Basalt API client requests (e.g., PromptsClient.get_sync()) even when they are nested within a start_observe() context that sets a feature_slug.

Current Behavior

When using the observability API like this:

from basalt import Basalt
from basalt.observability import start_observe, Identity

b = Basalt(api_key=API_KEY)

with start_observe(
    name="Trace A",
    feature_slug="categoriser",
    identity=Identity(
        user={"id": "87b000dj", "name": "Bobo"},
        organization={"id": "id-basalt", "name": "Basalt"},
    )
) as trace:
    prompt = b.prompts.get_sync(slug="router")  # ❌ feature_slug NOT propagated
    run_experiment("hello there")

The span created for the get_sync() API call does not inherit the basalt.span.feature_slug attribute from the parent trace context.

Root Cause

The functions trace_sync_request() and trace_async_request() in basalt/observability/request_tracing.py create spans using the observe() context manager. However:

1. These spans are created with observe() which creates child spans
2. The BasaltContextProcessor.on_start() method reads feature_slug from OpenTelemetry context via _apply_feature_slug_from_context()
3. BUT the observe() decorator does NOT explicitly read or propagate feature_slug from context

The feature_slug is stored in OpenTelemetry context via FEATURE_SLUG_CONTEXT_KEY, and BasaltContextProcessor is designed to apply it to all spans. However, the request tracing spans are not receiving it.

Expected Behavior

All spans created within a start_observe() context should automatically inherit the feature_slug attribute, including:

• Spans created by observe() decorators
• Spans created by API client requests (PromptsClient.get_sync(), etc.)
• Spans created by auto-instrumentation

This is critical for:

• Feature tracking and analytics
• Filtering traces by feature in the Basalt UI
• Associating API requests with specific features

Proposed Solution

Ensure that observe() calls within trace_sync_request() and trace_async_request() properly inherit feature_slug from the OpenTelemetry context by:

1. Option A (Recommended): Verify that BasaltContextProcessor is being invoked for these spans and debug why _apply_feature_slug_from_context() is not working
2. Option B: Explicitly read feature_slug from context in trace_sync_request() and trace_async_request() and pass it to the span

The fix should ensure that the feature_slug propagates correctly through the entire trace tree, including API request spans.

Files to Modify

• basalt/observability/request_tracing.py - Add explicit feature_slug propagation if needed
• Potentially basalt/observability/api.py - Ensure observe() reads feature_slug from context
• Add tests in tests/observability/test_feature_slug_propagation.py to verify API request spans inherit feature_slug

Test Case

Add a test that verifies:

def test_feature_slug_propagates_to_api_request_spans(setup_tracer):
    """Ensure feature_slug from start_observe reaches API request spans."""
    exporter = setup_tracer

    with start_observe(name="root", feature_slug="test-feature"):
        # Simulate an API request span
        with observe(name="api.prompts.get"):
            pass

    spans = exporter.get_finished_spans()
    for span in spans:
        assert span.attributes.get("basalt.span.feature_slug") == "test-feature", \
            f"Span {span.name} missing feature_slug"

Acceptance Criteria

• feature_slug is propagated to all spans created within a start_observe() context
• API request spans (created via trace_sync_request/trace_async_request) inherit feature_slug
• Existing tests pass
• New test added to verify the fix
• No breaking changes to the API

References

• Related processor: basalt/observability/processors.py - BasaltContextProcessor._apply_feature_slug_from_context()
• Context key: basalt/observability/trace_context.py - FEATURE_SLUG_CONTEXT_KEY
• Test file: tests/observability/test_feature_slug_propagation.py

This pull request was created from Copilot chat.

---

✨ Let Copilot coding agent set things up for you — coding agent works faster and does higher quality work when set up for your repo.