feat(tracing): support prompt linking via propagate_attributes

[hassiebp]

Summary

Adds a prompt parameter to propagate_attributes, enabling prompt↔trace linking for generations created by auto-instrumentation integrations (LiteLLM langfuse_otel, OpenAI Agents SDK, OpenInference, ...) where the user never creates the generation via the Langfuse SDK. Previously this required a hand-rolled OTel SpanProcessor per prompt (see langfuse/langfuse#9345).

prompt = langfuse.get_prompt("my-prompt")

with langfuse.propagate_attributes(prompt=prompt):
    # generations emitted by auto-instrumentation in here are linked
    completion = litellm.completion(model="gpt-4o", messages=prompt.compile(...))

Details

• prompt accepts a PromptClient, a dict, or any object exposing name (str) and version (int) — e.g. {"name": "my-prompt", "version": 3}. Digit-string versions are coerced to int; invalid values are dropped with a warning; fallback prompts are never linked (consistent with existing linking behavior).
• Propagates through the existing OTel context machinery as langfuse.observation.prompt.name / langfuse.observation.prompt.version span attributes, including as_baggage=True cross-service propagation (version restored as int on read).
• Precedence: an explicit prompt set at observation creation (start_observation(..., prompt=...)) wins over the propagated one (guard in LangfuseSpanProcessor.on_start).
• Composes with other propagated attributes: an inner propagate_attributes(prompt=...) preserves an outer session_id/user_id context; same-key nesting shadows within scope and restores on exit (covered by tests).

Server-side companion

Since propagation stamps every span in scope, the backend gates prompt linking to GENERATION observations: langfuse/langfuse#14962. This PR should be released after that change is deployed.

Tests

• 11 new unit tests in tests/unit/test_propagate_attributes.py (PromptClient/dict/duck-typed inputs, fallback skip, invalid name/version, explicit-over-propagated precedence, nested shadowing, baggage round-trip, composition with outer attributes).
• Full unit suite: 623 passed (pre-existing test_prompt.py fixture errors on Python 3.14 reproduce on a clean tree, unrelated).
• ruff check/format and mypy clean (enforced via pre-commit).

🤖 Generated with Claude Code

Greptile Summary

This PR adds a prompt parameter to propagate_attributes, letting users link a Langfuse prompt to all generations emitted by auto-instrumentation libraries (LiteLLM langfuse_otel, OpenAI Agents SDK, OpenInference) without touching the SDK's observation creation API.

• _extract_propagated_prompt validates the prompt input (PromptClient, dict, or duck-typed object), skips fallback prompts, coerces digit-string versions to int, and drops invalid inputs with a warning.
• Prompt attributes are propagated via the existing OTel context machinery and optionally as W3C baggage (version serialised to string and restored via isdigit() on read-back).
• In LangfuseSpanProcessor.on_start, an explicit prompt set at span-creation time takes precedence over the propagated one by filtering it from the propagated dict before set_attributes is called.

Confidence Score: 4/5

Safe to merge once the companion backend change (langfuse/langfuse#14962) is deployed; the SDK changes themselves introduce no regressions in the existing propagation path.

The core logic is sound and well-tested. Two small issues were found: the public type hint for the prompt parameter omits the duck-typed object case that is explicitly documented and tested, and the test helper _make_prompt_client places two imports inside the method body rather than at module level.

langfuse/_client/propagation.py — type annotation for the prompt parameter; tests/unit/test_propagate_attributes.py — inline imports in _make_prompt_client.

Sequence Diagram

%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
    participant User
    participant PA as propagate_attributes()
    participant EPP as _extract_propagated_prompt()
    participant SPA as _set_propagated_attribute()
    participant OTel as OTel Context / Baggage
    participant Span as Auto-instrumented Span
    participant LSP as LangfuseSpanProcessor.on_start()
    participant GPAC as _get_propagated_attributes_from_context()

    User->>PA: "propagate_attributes(prompt=prompt, as_baggage=False)"
    PA->>EPP: _extract_propagated_prompt(prompt)
    EPP-->>PA: (name, version) or None (fallback/invalid)

    alt prompt valid
        PA->>SPA: set prompt_name + prompt_version
        SPA->>OTel: set_value(context_key, value) [or set_baggage(str(version))]
    end

    User->>Span: auto-instrumentation creates span (e.g. LiteLLM)
    Span->>LSP: on_start(span, parent_context)
    LSP->>GPAC: _get_propagated_attributes_from_context(context)
    GPAC->>OTel: get context values / baggage entries
    OTel-->>GPAC: prompt_name, prompt_version (int restored via isdigit() if baggage)
    GPAC-->>LSP: propagated_attributes dict

    alt span has explicit OBSERVATION_PROMPT_NAME
        LSP->>LSP: drop prompt_name + prompt_version from propagated_attributes
    end

    LSP->>Span: span.set_attributes(propagated_attributes)
    Note over LSP,Span: Backend gates prompt linking to GENERATION observations

Loading

%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
sequenceDiagram
    participant User
    participant PA as propagate_attributes()
    participant EPP as _extract_propagated_prompt()
    participant SPA as _set_propagated_attribute()
    participant OTel as OTel Context / Baggage
    participant Span as Auto-instrumented Span
    participant LSP as LangfuseSpanProcessor.on_start()
    participant GPAC as _get_propagated_attributes_from_context()

    User->>PA: "propagate_attributes(prompt=prompt, as_baggage=False)"
    PA->>EPP: _extract_propagated_prompt(prompt)
    EPP-->>PA: (name, version) or None (fallback/invalid)

    alt prompt valid
        PA->>SPA: set prompt_name + prompt_version
        SPA->>OTel: set_value(context_key, value) [or set_baggage(str(version))]
    end

    User->>Span: auto-instrumentation creates span (e.g. LiteLLM)
    Span->>LSP: on_start(span, parent_context)
    LSP->>GPAC: _get_propagated_attributes_from_context(context)
    GPAC->>OTel: get context values / baggage entries
    OTel-->>GPAC: prompt_name, prompt_version (int restored via isdigit() if baggage)
    GPAC-->>LSP: propagated_attributes dict

    alt span has explicit OBSERVATION_PROMPT_NAME
        LSP->>LSP: drop prompt_name + prompt_version from propagated_attributes
    end

    LSP->>Span: span.set_attributes(propagated_attributes)
    Note over LSP,Span: Backend gates prompt linking to GENERATION observations

Loading

Prompt To Fix All With AI

Fix the following 2 code review issues. Work through them one at a time, proposing concise fixes.

---

### Issue 1 of 2
langfuse/_client/propagation.py:123-125
The public type hint `Union[PromptClient, Mapping[str, Any]]` doesn't cover pure duck-typed objects (objects with `name`/`version` attributes but not a `Mapping`), even though the docstring and test suite explicitly document and verify that use-case. Users who call `propagate_attributes(prompt=my_duck_typed_obj)` with mypy enabled will get a type error despite the call being fully functional at runtime.

```suggestion
    prompt: Optional[Union[PromptClient, Mapping[str, Any], Any]] = None,
    as_baggage: bool = False,
) -> _AgnosticContextManager[Any]:
```

### Issue 2 of 2
tests/unit/test_propagate_attributes.py:3377-3378
**Inline imports inside a method**
`from langfuse.api import Prompt_Text` and `from langfuse.model import TextPromptClient` are placed inside `_make_prompt_client` rather than at the top of the module. The team rule requires imports to be at module level.

_{Reviews (1): Last reviewed commit: "feat(tracing): support prompt linking vi..." | Re-trigger Greptile}

[github-actions]

@claude review

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: afcb25d42a

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[claude]

🟡 Version coercion in _extract_propagated_prompt (propagation.py:436-437) and the baggage-decode branch of _get_propagated_attributes_from_context (propagation.py:474-480) gate int(version) on str.isdigit(), which returns True for Unicode digit-like characters (superscripts '²', '¹', '⁷', etc.) that int() rejects with ValueError. Passing version='²' at context-manager entry — or receiving a crafted W3C baggage header like langfuse_prompt_version=² in LangfuseSpanProcessor.on_start — raises an uncaught ValueError instead of dropping the value with a warning per the docstring contract. Fix by replacing the gate with .isdecimal() plus an ASCII check (or re.fullmatch(r'-?[0-9]+', version)), or wrapping the int() call in try/except ValueError at both sites.

Extended reasoning...

The bug

The PR introduces two identical version-coercion sites that both use the same too-permissive gate:

Site 1 — _extract_propagated_prompt (propagation.py:436-437):

if isinstance(version, str) and version.isdigit():
    version = int(version)

Site 2 — _get_propagated_attributes_from_context baggage branch (propagation.py:474-480):

if (
    span_key == LangfuseOtelSpanAttributes.OBSERVATION_PROMPT_VERSION
    and isinstance(baggage_value, str)
    and baggage_value.isdigit()
):
    propagated_attributes[span_key] = int(baggage_value)
    continue

str.isdigit() is a strict superset of what int(str) accepts. int() accepts characters where .isdecimal() is True (ASCII 0-9 plus other decimal scripts like Arabic-Indic '٠'-'٩'), but rejects .isdigit()-only characters such as superscripts.

Step-by-step proof

Verified on Python 3.11:

>>> '²'.isdigit()
True
>>> int('²')
Traceback (most recent call last):
  ...
ValueError: invalid literal for int() with base 10: '²'

Same for '¹', '⁷', and other superscript/Unicode digit chars.

Scenario A (user input):

1. User writes with langfuse.propagate_attributes(prompt={"name": "x", "version": "²"}):
2. _propagate_attributes calls _extract_propagated_prompt({"name":"x","version":"²"}) (propagation.py:325).
3. Line 436: isinstance("²", str) and "²".isdigit() → True.
4. Line 437: int("²") raises ValueError.
5. The exception propagates out of the context-manager __enter__, so the with block never runs — the user's request crashes.

This contradicts the docstring contract stated on the same function: "Returns None (with a warning) if the value is invalid" — and the class docstring: "invalid values are dropped with a warning". The intended fall-through already exists at line 439-443; the gate simply lets the wrong things past it.

Scenario B (cross-service baggage):

1. Any incoming request carrying Baggage: langfuse_prompt_version=² populates OTel baggage on the request context.
2. When any span is created, LangfuseSpanProcessor.on_start runs (span_processor.py:145) and calls _get_propagated_attributes_from_context(context) with no try/except wrapper.
3. The baggage branch (line 474-480) hits baggage_value.isdigit() == True, then int("²") raises ValueError.
4. OTel's SynchronousMultiSpanProcessor.on_start dispatch does not catch processor exceptions, so the ValueError propagates up through Span.start → Tracer.start_span and aborts span creation for every span in the trace.

Why existing code doesn't prevent this

• The isinstance(version, int) fall-through at line 439-443 would gracefully drop an invalid version — but it's only reached if the int() conversion succeeds or is skipped. With the current gate, malformed input takes the raising branch instead.
• The type check isinstance(version, bool) right below correctly rejects True/False, showing the author was thinking about edge cases — the Unicode-digit case was just missed.
• On the baggage path there's no guard at all: the on_start call site is a bare method call with no exception handling.

Impact

• Scenario A: low probability trigger (users don't normally type superscripts), but a real contract violation. Any user who feeds a non-ASCII string through their prompt-version pipeline hits an uncaught exception where they were promised a warning.
• Scenario B: also low probability in practice — legitimate Langfuse→Langfuse baggage propagation uses str(int) on the sender side, so no legitimate producer emits '²'. But baggage is untrusted header data in general; a peer service or crafted request could inject arbitrary strings, and the resulting exception surfaces from SpanProcessor.on_start on every span in the affected request.

Fix

Either of these at both sites:

1. Tighten the gate:

if isinstance(version, str) and re.fullmatch(r"-?[0-9]+", version):
    version = int(version)

(or version.isdecimal() and version.isascii() — beware that plain .isdecimal() still accepts non-ASCII decimal digits like '٠', which int() does accept but which are almost certainly not what a Langfuse prompt version should look like.)

2. Or wrap in try/except and let the existing fall-through handle it:

if isinstance(version, str):
    try:
        version = int(version)
    except ValueError:
        pass

Option 1 is preferable at the baggage site because it keeps the code exception-free rather than relying on catch-and-continue in a hot path.

Severity

Marking as nit: the trigger requires unusual (superscript characters) or actively hostile (crafted baggage headers) input, and the fix is a two-line change confined to already-changed code. Worth cleaning up in this PR since both sites are new; not worth blocking merge over.

[claude]

🟡 The docstring on propagate_attributes promises "An explicit prompt passed to start_observation / update_current_generation takes precedence over the propagated one," but this is only enforced for child spans via the on_start guard in span_processor.py. When propagate_attributes(prompt=Y) is entered while a generation created with an explicit prompt=X is the current active span, _set_propagated_attribute unconditionally calls span.set_attribute on that current span (propagation.py:325-341 → 570-580), silently overwriting the explicit prompt link with the propagated one on export. Adding a symmetric guard in _propagate_attributes — skip the current-span writes for prompt_name/prompt_version when OBSERVATION_PROMPT_NAME is already present — would restore the documented precedence.

Extended reasoning...

What breaks

In _propagate_attributes (propagation.py:325-341), when a prompt is passed, _set_propagated_attribute is called with span=current_span for both prompt_name and prompt_version. That helper (propagation.py:570-580) checks only span is not None and span.is_recording() and then unconditionally calls span.set_attribute(key=span_key, value=value). OTel set_attribute overwrites any prior value under the same key, so an OBSERVATION_PROMPT_NAME/OBSERVATION_PROMPT_VERSION already set on the current span by create_generation_attributes (attributes.py:141-146) is silently replaced by the propagated prompt.

Why the existing precedence guard does not help

The on_start guard added in this PR (span_processor.py:147-158) inspects span.attributes at span-creation time and drops propagated prompt keys if the new child already has an explicit prompt. That guard only fires for new spans starting inside the propagation context. The outer generation was created before propagate_attributes is entered — its on_start already ran with empty propagated context, and its explicit prompt was written by set_attributes afterwards. The direct span.set_attribute inside _set_propagated_attribute never goes through on_start again, so there is nothing to enforce precedence on the current/parent span.

Other propagated attributes (user_id, session_id, tags, …) do not have this problem because they cannot be set at observation-construction time — only prompt has this dual write path.

Step-by-step proof

Given:

prompt_a = langfuse.get_prompt("prompt-a")   # version 1
prompt_b = langfuse.get_prompt("prompt-b")   # version 2

with langfuse.start_as_current_generation(name="outer", prompt=prompt_a) as gen:
    with langfuse.propagate_attributes(prompt=prompt_b):
        # ... auto-instrumented sub-call ...
        pass

1. start_as_current_generation(prompt=prompt_a) starts gen. on_start sees empty propagated context. Immediately after, create_generation_attributes sets OBSERVATION_PROMPT_NAME='prompt-a', OBSERVATION_PROMPT_VERSION=1 on gen (attributes.py:141-146).
2. propagate_attributes(prompt=prompt_b) calls _propagate_attributes, which extracts ('prompt-b', 2) and calls _set_propagated_attribute twice with span=current_span=gen.
3. Inside _set_propagated_attribute, gen.is_recording() is True, so it hits the else branch at line 580 and calls gen.set_attribute(OBSERVATION_PROMPT_NAME, 'prompt-b') — overwriting 'prompt-a' — then gen.set_attribute(OBSERVATION_PROMPT_VERSION, 2) — overwriting 1.
4. When gen ends and is exported, its prompt link points to prompt-b, not the explicitly-passed prompt-a.

The included test test_explicit_prompt_takes_precedence_over_propagated only covers the case where the explicit prompt is set on a child span created inside the propagation context — that path goes through on_start and is protected. The parent-span variant is not covered.

Suggested fix

Mirror the on_start guard inside _propagate_attributes — before writing prompt attributes to current_span, check whether current_span already carries OBSERVATION_PROMPT_NAME and, if so, skip the direct span.set_attribute for both prompt keys (still write them to the OTel context so child spans without an explicit prompt inherit them).