Skip to content

fix(litellm): use underlying provider for metadata, drop eager serialization#594

Merged
Abhijeet Prasad (AbhiPrasad) merged 3 commits into
mainfrom
fix/litellm-provider-metadata-drop-eager-serialization
Jul 17, 2026
Merged

fix(litellm): use underlying provider for metadata, drop eager serialization#594
Abhijeet Prasad (AbhiPrasad) merged 3 commits into
mainfrom
fix/litellm-provider-metadata-drop-eager-serialization

Conversation

@starfolkai

@starfolkai starfolkai Bot commented Jul 17, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Summary

Align the LiteLLM integration with .agents/skills/sdk-integrations/SKILL.md. Audit turned up four issues:

  • Wrong metadata.provider. SKILL §Completion-style explicitly says LiteLLM must set metadata.provider to the underlying provider (openai / cohere / openrouter / …), not the framework name. Fixed by routing every span through _resolve_provider(model) which calls litellm.get_llm_provider(model); falls back to "litellm" only when the lookup fails. @lru_cache(maxsize=256) on the resolver so the hot path doesn't re-dispatch through litellm's routing table per call.
  • Excess serialization. Every non-streaming wrapper called _try_to_dict(response) on the full pydantic response then indexed into it — but bt_json already serializes at log time. Replaced with direct attribute reads via a new shared _get_field(obj, key, default) helper in integrations/utils.py that works on both dicts and SDK objects (openai / anthropic / cohere had all been open-coding this pattern). Streaming paths left intact.
  • Dead code. _is_litellm_patched was only used by crewai until fix(crewai): add provider metadata, route tools to metadata, drop eager serialization #587 moved crewai to the "wrapper never owns tokens" pattern; no remaining callers anywhere. Removed the helper and its two subprocess tests.
  • Excessive comments. Trimmed the multi-example wrap_litellm / _is_litellm_patched docstrings and the """wrapt wrapper for litellm.X.""" one-liner on every wrapper.

Also lifted the rerank output cap (results[:100]) into integrations/utils.RERANK_OUTPUT_MAX_RESULTS and documented the "rerank exception to no-silent-caps" rule in the SKILL. Cohere carries the same magic number and can adopt the shared constant in a follow-up.

Test plan

  • Updated existing VCR-backed tests to assert the underlying provider (openai, cohere, openrouter) instead of "litellm". No cassette re-recording — HTTP requests unchanged, span-shape assertions only.
  • Added positive metadata["provider"] == "openrouter" assertion to the existing OpenRouter cassette test as coverage for the new resolver.
  • Dropped the two _is_litellm_patched subprocess tests alongside the helper.
  • No new mocks/fakes, no cassette re-recording.
  • ruff check + ruff format --check on touched files — clean.
  • Reviewer: cd py && nox -s "test_litellm(latest)" and cd py && nox -s "test_litellm(1.74.0)" (author's box has a broken Python 3.14 install).

🤖 Generated with Claude Code

Created by abhijeet

Slack thread

@starfolkbot @claude
fix(litellm): use underlying provider for metadata, drop eager serial…
ba72da4 
…ization

Align the LiteLLM integration with .agents/skills/sdk-integrations/SKILL.md.
Audit against the SKILL surfaced four issues:

- **`metadata.provider` was the framework name.** SKILL §Completion-style
  explicitly says LiteLLM must set `metadata.provider` to the underlying
  provider (openai / cohere / openrouter / ...), not "litellm". Fixed by
  routing every span through `_resolve_provider(model)` which calls
  `litellm.get_llm_provider(model)`; falls back to `"litellm"` only if the
  lookup fails. `@lru_cache(maxsize=256)` on the resolver so the hot path
  doesn't re-dispatch through litellm's routing table per call.
- **Excess serialization.** Every non-streaming wrapper called
  `_try_to_dict(response)` on the full pydantic response then indexed
  into it — but `bt_json` already serializes at log time. Replaced with
  direct attribute reads via a shared `_get_field(obj, key, default)`
  helper that works on both dicts and SDK objects (openai / anthropic /
  cohere had all been open-coding this). Streaming paths left intact.
- **Dead code.** `_is_litellm_patched` was only used by crewai until #587
  moved crewai to the "wrapper never owns tokens" pattern; no remaining
  callers anywhere. Removed the helper and its two subprocess tests.
- **Excessive comments.** Trimmed the huge `wrap_litellm` /
  `_is_litellm_patched` docstrings and the "wrapt wrapper for litellm.X"
  one-liner on every wrapper.

Also lifted the rerank output cap (`results[:100]`) into
`integrations/utils.RERANK_OUTPUT_MAX_RESULTS` and documented the
"rerank exception to no-silent-caps" rule in the SKILL — cohere carries
the same magic number and can adopt the shared constant in a follow-up.

## Test plan

- [x] Updated existing VCR-backed tests to assert the underlying
      provider (`openai`, `cohere`, `openrouter`) instead of `"litellm"`.
      No cassette re-recording (HTTP requests unchanged; span-shape
      assertions only).
- [x] Added positive `metadata["provider"] == "openrouter"` assertion
      to the existing OpenRouter cassette test as coverage for the new
      resolver.
- [x] Dropped the two `_is_litellm_patched` subprocess tests alongside
      the helper.
- [x] **No new mocks/fakes, no cassette re-recording.**
- [x] `ruff check` + `ruff format --check` on touched files — clean.
- [ ] Reviewer to run: `cd py && nox -s "test_litellm(latest)"` and
      `cd py && nox -s "test_litellm(1.74.0)"` (author's box has a
      broken Python 3.14 install).

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
@starfolkai
starfolkai Bot force-pushed the fix/litellm-provider-metadata-drop-eager-serialization branch from 1bace7d to ba72da4 Compare July 17, 2026 20:40
starfolkbot and others added 2 commits July 17, 2026 20:54
@starfolkbot @claude
fix(litellm): update remaining provider assertions for CI
69e280f 
Follow-up to the provider-metadata change. CI surfaced two categories:

- Three ``auto_test_scripts/`` subprocess scripts (test_auto_litellm,
  test_patch_litellm_responses, test_patch_litellm_aresponses) still
  hardcoded ``provider == "litellm"``. All three call ``gpt-4o-mini`` →
  expect ``"openai"``.
- ``test_litellm_text_completion_metrics`` / ``test_litellm_atext_completion_metrics``
  use ``gpt-3.5-turbo-instruct``, for which
  ``litellm.get_llm_provider`` returns ``"text-completion-openai"`` (its
  legacy text-completion endpoint category, distinct from chat
  completions). Assertion updated to that string — it IS the provider
  identifier litellm routes to.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
@starfolkbot @claude
fix(litellm): normalize legacy /v1/completions endpoint providers
0fbfc86 
``litellm.get_llm_provider`` reports endpoint-category identifiers for
legacy ``/v1/completions`` routes (``text-completion-openai``,
``azure_text``, ``text-completion-codestral``, ``text-completion-inception``).
The SKILL says ``metadata.provider`` names whose pricing applies — same
billing entity as the chat variant — so filtering by ``openai`` in the
UI should include both chat and text-completion calls.

Added a small ``_PROVIDER_ALIASES`` table that folds those routing keys
back to the pricing provider before returning from ``_resolve_provider``.
List sourced from litellm ``constants.py``.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
@AbhiPrasad
Abhijeet Prasad (AbhiPrasad) merged commit 59c9d5b into main Jul 17, 2026
82 checks passed
@AbhiPrasad
Abhijeet Prasad (AbhiPrasad) deleted the fix/litellm-provider-metadata-drop-eager-serialization branch July 17, 2026 21:05
@starfolkai starfolkai Bot mentioned this pull request Jul 21, 2026
Merged
5 tasks
Abhijeet Prasad (AbhiPrasad) pushed a commit that referenced this pull request Jul 21, 2026
@starfolkai @starfolkbot @claude
fix(mistral): drop eager response serialization, use allowlist metada…
a095270 
…ta pass-through (#604)

## Summary

Aligns the Mistral integration with
`.agents/skills/sdk-integrations/SKILL.md`, following the same pattern
as #594 (litellm) and #583 (anthropic). Audit turned up two issues:

- **Excess serialization.** `_normalize_mistral_multimodal_value` (which
recursively calls `value.model_dump(mode="python", by_alias=True)`) was
invoked 25 times across `tracing.py`. Two sites did real waste:
- Every non-streaming response finalizer
(`_finalize_completion_response`, `_finalize_conversation_response`,
`_finalize_ocr_response`, `_finalize_embeddings_response`) full-dumped
the response just to read a handful of fields (`id`, `model`,
`choices`/`outputs`/`pages`, `usage`).
- `_build_request_metadata` walk-and-dumped every metadata value
(`temperature`, `top_p`, `tools`, `response_format`, `stop`, …) — none
of which carry attachments.

Both are redundant: `bt_safe_deep_copy` in `logger.log_internal` already
`model_dump`s pydantic values at log time. Replaced with attribute reads
through the existing `_get_value` helper (already dict-or-pydantic aware
and handles Mistral's `Unset` sentinel — no new helper needed).

- **Aggregators.** `_aggregate_completion_events` /
`_aggregate_conversation_events` / `_aggregate_transcription_events` /
`_aggregate_speech_events` pre-dumped `usage` / `segments` into the
intermediate dict. `_parse_usage_metrics` already accepts pydantic, and
`bt_safe_deep_copy` handles the output. Dropped.

- **Tool-span walkers.** `_completion_tool_calls`,
`_conversation_tool_outputs`, `_log_completion_tool_spans`,
`_log_conversation_tool_spans` switched to `_get_value` so they walk raw
pydantic responses instead of assuming pre-dumped dict shape.

Multimodal walker retained where attachment materialization matters
(`_start_span` input path, `_append_delta_content` / `_merge_tool_calls`
streaming assembly, per-field conversation-output accumulators).

Metadata is still allowlist-per-surface — no changes to what fields we
capture. Comments were not excessive; nothing to trim.

Net: −36 LOC in `tracing.py`, one dead helper file removed
(`_response_data_to_metadata`,
`_conversation_response_data_to_metadata`, `_conversation_outputs_data`,
`_ocr_output_data` collapsed into the pydantic-aware equivalents).

## Test plan

- [x] **No new mocks/fakes.** The two existing narrow fakes
(`_PlainResponse` for the plain-Python usage path, and lambda
fail-injectors for error-propagation tests) fall under the SKILL's
"narrow error injection / provider-independent helpers" carve-out. All
provider-behavior coverage stays on the 30+ `@pytest.mark.vcr`
cassette-backed tests.
- [x] **No cassette re-recording.** HTTP behavior is unchanged; only
in-process span shaping was touched.
- [x] `cd py && nox -s "test_mistral(latest)"` — 45/45 passed.
- [x] `cd py && nox -s "test_mistral(1.12.4)"` — 43/43 passed (2
pre-existing skips for speech API which doesn't exist on 1.12.4).
- [x] `ruff check` + `ruff format --check` on touched files — clean.

Two existing unit tests were updated because their assertions hinged on
the old eager-serialization behavior:
- `test_wrappers_ignore_usage_when_response_normalization_fails` →
`test_wrappers_read_usage_from_plain_python_responses`. The old test
asserted metrics were *dropped* for plain-Python responses — that only
happened because `.model_dump()` failed on non-pydantic objects and we
bailed out. The new code correctly reads `usage` via `getattr`;
assertion is now positive (metrics *are* present).
- `test_aggregate_completion_events_merges_tool_calls_and_content`:
`aggregated["usage"]` is now the raw pydantic `UsageInfo` (dumps at log
time, not in the aggregator); switched to `getattr(...,
"total_tokens")`.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

<!-- sfk:created-approved-by -->
Created by abhijeet

<!-- sfk:slack-thread -->
[Slack
thread](https://starfolkai.slack.com/archives/C0AQDETAVT3/p1784591709340249?thread_ts=1784591709.340249&cid=C0AQDETAVT3)

Co-authored-by: Starfolk <noreply@starfolk.ai>
Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com>
This was referenced Jul 21, 2026
Merged
Merged
Abhijeet Prasad (AbhiPrasad) pushed a commit that referenced this pull request Jul 21, 2026
@starfolkai @starfolkbot @claude
fix(openrouter): allowlist metadata pass-through, drop eager serializ…
de233f8 
…ation (#608)

## Summary

Aligns the OpenRouter integration with
`.agents/skills/sdk-integrations/SKILL.md`, following #604 (mistral),
#603 (livekit_agents), #594 (litellm), #593 (langchain), and #585
(claude_agent_sdk). Audit found two issues:

- **Excess serialization.** `sanitize_openrouter_logged_value` called
`bt_safe_deep_copy` on every span input, every kwargs dict, every
response (for both output and metadata extraction), and every usage
dict. Braintrust already runs `bt_safe_deep_copy` at log time in
`Span.log_internal` (`logger.py:4699`), so every one of those was a
duplicate walk. Dropped the wrapper; reads now go through a small
`_get_field(obj, key)` helper (dict-or-attribute) and pydantic responses
flow straight to `span.log(...)`.

- **Denylist metadata → per-surface allowlist.** `_OMITTED_KEYS` was a
small denylist (`execute`, `render`, `nextTurnParams`,
`requireApproval`); the old `_build_openrouter_metadata` then dumped
every *other* kwarg — including `messages`, `api_key`, `extra_headers`,
and unknown kwargs — into `metadata`. Replaced with three request
allowlists (`_CHAT_REQUEST_KEYS`, `_EMBEDDINGS_REQUEST_KEYS`,
`_RESPONSES_REQUEST_KEYS`) and three response allowlists
(`_CHAT_RESPONSE_KEYS`, `_EMBEDDINGS_RESPONSE_KEYS`,
`_RESPONSES_RESPONSE_KEYS`), mirroring `litellm/_SAFE_METADATA_KEYS` and
mistral's `_*_METADATA_KEYS` family. `is_byok` is still extracted from
`usage` (it's still nested there).

Also collapsed the three near-identical `_finalize_*_response` functions
into one parameterized `_finalize_response(..., *, output,
response_keys)`, and factored the shared model/provider stamping
(`_stamp_model_and_provider`) used by both request and response metadata
builders.

Net: 1 file changed, +262 / −178 (net −0 after adding the allowlist
tuples, but the tracing surface itself shrank meaningfully).

## Test plan

- [x] **No new mocks/fakes.** All coverage stays on the existing
`@pytest.mark.vcr` cassette-backed tests in
`integrations/openrouter/test_openrouter.py` (chat sync/stream/async,
embeddings, responses sync + async-stream, `setup_*_creates_spans`,
`setup_*_is_idempotent`, auto_instrument subprocess).
- [x] **No cassette re-recording.** HTTP behavior is unchanged; only
in-process span shaping was touched.
- [x] `cd py && nox -s "test_openrouter(latest)"` — 9/9 passed.
- [x] `cd py && nox -s "test_openrouter(0.6.0)"` — 9/9 passed.
- [x] `ruff check` + `ruff format --check` — clean.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

<!-- sfk:created-approved-by -->
Created by abhijeet

<!-- sfk:slack-thread -->
[Slack
thread](https://starfolkai.slack.com/archives/C0AQDETAVT3/p1784653510318699?thread_ts=1784653510.318699&cid=C0AQDETAVT3)

Co-authored-by: Starfolk <noreply@starfolk.ai>
Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com>
Abhijeet Prasad (AbhiPrasad) pushed a commit that referenced this pull request Jul 21, 2026
@starfolkai @starfolkbot @claude
fix(openai): allowlist metadata pass-through, drop eager response ser…
27c4c79 
…ialization (#606)

## Summary

Aligns the OpenAI integration with
`.agents/skills/sdk-integrations/SKILL.md`, following the same pattern
as #604 (mistral), #594 (litellm), #603 (livekit), and #583 (anthropic).
Audit turned up two issues:

- **Denylist metadata pass-through (real leak vector).** Every
`_parse_params` copy (`ChatCompletionWrapper`, `ResponseWrapper`,
`BaseWrapper` for Embedding/Moderation/Speech, `_AudioFileWrapper`,
`_ImageBaseWrapper`) did `metadata: {**params, "provider": "openai"}` —
every non-input request kwarg landed in span metadata.
`ResponseWrapper._parse_event_from_result` did the same on the response
side (`{k: v for k, v in result.items() if k not in ["output",
"usage"]}`). That flowed `extra_headers`, `extra_query`, `extra_body`,
user-provided Responses `metadata`, `user`, `safety_identifier`,
`prompt_cache_key`, and anything OpenAI adds in future SDK versions
straight into telemetry. Three tests actively ratified the header leak.

Replaced with per-surface allowlist tuples plus a shared
`_filter_metadata` helper that drops `NOT_GIVEN`/`Omit` sentinels and
serializes `response_format` inline. Mirrors Anthropic's
`METADATA_PARAMS`.

- **Eager response serialization.** `_try_to_dict(raw_response)` was
full-dumping every non-streaming Chat/Response API response just to read
`usage` + `choices`/`output` — `bt_safe_deep_copy` already dumps at log
time. Same fix pattern the mistral/litellm/anthropic PRs used: reads via
`getattr` through a small `_get_attr` helper, plus a narrow
`_choices_output` that dumps only `choices` (needed for the
audio-attachment walker in `_process_attachments_in_chat_output`).
`_parse_event_from_result` and `_extract_moderation_metadata` were made
dict-or-object aware.

Streaming paths still pre-dump chunks via `_try_to_dict(item)` because
the aggregator reads deeply nested delta fields — that's not the
"over-serialize" case the SKILL warns about.

## Note on `moderation`

Added to both chat + responses request allowlists and the responses
result allowlist. It's a Braintrust-proxy request parameter (`{"model":
"omni-moderation-latest"}`) that deep-merges with the response's
`moderation` object (`{"input": ..., "output": ...}`) at log time via
`merge_dicts`. The two existing `*_inline_moderation_metadata` tests
rely on this merge to see `logged_moderation["model"]`.

## Test plan

- [x] **No new mocks/fakes.** All provider-behavior coverage stays on
the existing `@pytest.mark.vcr` cassette-backed tests.
- [x] **No cassette re-recording.** HTTP behavior unchanged; only
in-process span shaping was touched.
- [x] `nox -s "test_openai(latest)"` — 76/76 pass across
`test_openai.py` (70), `test_openai_ddtrace.py` (3),
`test_openai_openrouter_gateway.py` (3)
- [x] `nox -s "test_openai(1.92.0)"`, `test_openai(1.77.0)`,
`test_openai(1.71.0)` — all green
- [x] `nox -s "test_openai_http2_streaming(latest)"` — 3/3
- [x] `nox -s "test_btx_openai(latest)"` — 5/5
- [x] `nox -s "test_openai_agents(latest)"` — 7/7
- [x] `ruff check` + `ruff format --check` on touched files — clean

Three tests updated: assertions that expected
`metadata["extra_headers"]["X-Stainless-Helper-Method"] ==
"chat.completions.stream"` flipped to `"extra_headers" not in
span["metadata"]` — pins the specific leak vector the fix closes.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

<!-- sfk:created-approved-by -->
Created by abhijeet

<!-- sfk:slack-thread -->
[Slack
thread](https://starfolkai.slack.com/archives/C0AQDETAVT3/p1784593895623669?thread_ts=1784593895.623669&cid=C0AQDETAVT3)

---------

Co-authored-by: Starfolk <noreply@starfolk.ai>
Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com>
This was referenced Jul 21, 2026
Merged
Merged
Abhijeet Prasad (AbhiPrasad) pushed a commit that referenced this pull request Jul 21, 2026
@starfolkai @starfolkbot @claude
fix(strands): allowlist metrics, add span_origin, fix agent metadata …
a1a1b11 
…reads (#612)

## Summary

Aligns the Strands integration with
`.agents/skills/sdk-integrations/SKILL.md`, following the pattern of
#606 (openai), #604 (mistral), #603 (livekit), #594 (litellm), #583
(anthropic). Audit turned up several issues:

- **Missing span_origin instrumentation name.** Every other integration
has `_INSTRUMENTATION = "<provider>-auto"` + a module-level `start_span`
wrapper; strands was the only one whose spans lacked
`context.span_origin.instrumentation.name`. Added the standard wrapper
and pass `internal={"instrumentation": _INSTRUMENTATION}` explicitly
through the nested `parent.start_span(...)` path, which doesn't route
through the module wrapper.

- **Denylist metrics extraction (real leak vector).**
`_end_model_invoke_span_wrapper` was doing `{k: v for k, v in
metrics.items() if _is_supported_metric_value(v)}` — every numeric key
strands emits (present + future) flowed into the span's top-level
`metrics` bag. Replaced with an explicit allowlist (`latencyMs`,
`timeToFirstByteMs`). Since neither key is spec-listed, they land under
`metadata.strands_metrics` rather than top-level `metrics`.

- **Silently-broken agent metadata extraction.**
`_agent_metrics_and_metadata` did `metrics = getattr(result, "metrics");
if not isinstance(metrics, dict): return {}, {}` — but
`AgentResult.metrics` is an `EventLoopMetrics` **dataclass**, so the
guard discarded everything. Renamed to `_agent_metadata_from_result` and
read `cycle_count` / `accumulated_usage` / `accumulated_metrics` off the
dataclass through a single `_pick_numeric(source, keys)` allowlist
helper. Moved `cycles` from `metrics` to `metadata` (not a spec metric
key).

- **Positional-arg reads.** `_start_agent_span_wrapper` /
`_start_model_invoke_span_wrapper` were `kwargs.get("tools")` /
`kwargs.get("system_prompt")` etc. Strands' signatures put those at
fixed positional slots — switched to `_arg(args, kwargs, N, name)` so
positional callers work.

- **Simplifications (via `/simplify`):** collapsed twin allowlist
helpers into `_pick_numeric`; collapsed the parent-vs-module `if/else`
in `_start_span_for_otel`; replaced the `_bt_parent_otel_span` metadata
back-channel with a direct `parent_otel_span=` kwarg; dropped the
now-orphan `metrics=` parameter from `_end_span_for_otel`; trimmed the
verbose `wrap_strands_tracer` and `_InvalidOtelSpanKey` docstrings.

## Known SKILL gaps NOT fixed here

Each deserves its own PR/discussion — flagging so they don't get lost:

1. **`metadata.provider` on LLM spans.** Strands' `Tracer` only receives
`model_id`, not the model class — provider inference would require
brittle id-pattern matching.
2. **Tokens under `metadata.strands_usage` rather than `metrics.tokens`
/ `prompt_tokens` / `completion_tokens`.** Fixing correctly means
deciding whether model-invoke should stay typed `llm` at all — Strands'
`OpenAIModel` / `AnthropicModel` delegate to the underlying provider
SDK, which per the SKILL's "framework llm-like spans that delegate" rule
should NOT be typed `llm`.
3. **No `test_auto_strands.py` subprocess auto-instrument test.**
Strands is wired into `auto.py` but lacks the subprocess parity test
other integrations have.

## Test plan

- [x] **No new mocks/fakes.** Existing cassette-backed VCR test
(`test_strands_openai_agent_traces_native_otel_lifecycle`) covers the
behavior; added a `span_origin.instrumentation.name == "strands-auto"`
assertion on every emitted span.
- [x] **No cassette re-recording.** HTTP behavior is unchanged; only
in-process span shaping was touched.
- [x] `cd py && nox -s "test_strands(latest)"` — 2/2 pass
- [x] `cd py && nox -s "test_strands(1.20.0)"` — 2/2 pass
- [x] `ruff check` + `ruff format --check` on touched files — clean
- [x] `nox -s pylint` — clean

🤖 Generated with [Claude Code](https://claude.com/claude-code)

<!-- sfk:created-approved-by -->
Created by abhijeet

<!-- sfk:slack-thread -->
[Slack
thread](https://starfolkai.slack.com/archives/C0AQDETAVT3/p1784659948425149?thread_ts=1784659948.425149&cid=C0AQDETAVT3)

Co-authored-by: Starfolk <noreply@starfolk.ai>
Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com>
Abhijeet Prasad (AbhiPrasad) pushed a commit that referenced this pull request Jul 21, 2026
@starfolkai @starfolkbot @claude
fix(pydantic_ai): allowlist kwargs pass-through, drop eager message r…
fff5bd2 
…eshape (#610)

## Summary

Aligns the pydantic_ai integration with
`.agents/skills/sdk-integrations/SKILL.md`, following the same pattern
as #606 (openai), #608 (openrouter), and #604/#603/#594
(anthropic/livekit/litellm/mistral). Audit turned up two issues:

- **Denylist metadata pass-through (real leak vector).**
`_build_agent_input_and_metadata` and
`_build_direct_model_input_and_metadata` did `for key, value in
kwargs.items(): if key not in [...]: input_data[key] = value` — every
non-`deps` / non-`{model,messages}` kwarg landed in span input. That
flowed `usage` (a mutable counter), `event_stream_handler` (a callable),
`infer_name`, `instrument`, and anything pydantic_ai adds in future SDK
versions straight into telemetry. `deps` (user's app-level container —
commonly holds API keys / DB connections) was the sole excluded key on
the agent path; on the direct-model path there was no exclusion at all.

Replaced with per-surface allowlist tuples (`_AGENT_INPUT_ALLOWLIST`,
`_DIRECT_MODEL_INPUT_ALLOWLIST`) covering the known kwargs across
`Agent.run` / `run_sync` / `run_stream` / `run_stream_events` /
`to_cli_sync` and `direct.model_request*`. Explicitly excludes `deps`,
`usage`, `event_stream_handler`, `infer_name`, `instrument`.

- **Eager message reshaping dropping fields.** `_shape_message` /
`_shape_content_part` / `_shape_model_response` reshaped every message
into shallow dicts via a hard-coded field allowlist (`_MESSAGE_FIELDS`,
`_PART_FIELDS`, `_RESPONSE_FIELDS`) even when no binary content needed
materialization. That dropped real pydantic_ai v2 fields —
`instructions`, `run_id`, `conversation_id`, `metadata`, `state`,
`provider_details`, `outcome`, ... — and burned CPU rebuilding dicts
that `bt_safe_deep_copy` (which already handles dataclasses + Pydantic
models) would produce anyway at log time.

Added `_has_binary_leaf` walker; the four shape helpers now
short-circuit to the raw object when no binary is present. Reshape (and
its field allowlist) only runs when a binary attachment actually needs
materialization.

Also:
- Extracted `_shape_toolset(ts)` from a nested 20-line block inside
`_build_agent_input_and_metadata`.
- Trimmed docstrings/comments that just restated the helper name. Kept
the ones explaining non-obvious *why* (portal-thread context
propagation, wrapper-vs-leaf token ownership, 1.78.0 ToolManager alias,
`_system_prompts` private-API access).
- Hoisted `import inspect` from inside `_shape_type` to module top.

## Test plan

- [x] **No new mocks/fakes.** All provider-behavior coverage stays on
the existing `@pytest.mark.vcr` cassette-backed tests.
- [x] **No cassette re-recording.** HTTP behavior unchanged; only
in-process input/shape logic was touched.
- [x] `nox -s "test_pydantic_ai_integration(latest)"` — 68/68 pass
(`pydantic-ai==2.13.0`)
- [x] `nox -s "test_pydantic_ai_integration(1.10.0)"` — 68/68 pass
(min-supported matrix version)
- [x] `nox -s "test_pydantic_ai_wrap_openai(latest)"` — green
- [x] `nox -s "test_pydantic_ai_logfire(latest)"` — 1/1 pass

Two new tests pin the security invariant:
- `test_agent_input_kwargs_are_allowlisted` — asserts `deps` / `usage` /
`event_stream_handler` / `infer_name` never land in span input;
allowlisted kwargs (`instructions`, `usage_limits`, `retries`,
`conversation_id`, `metadata`) do.
- `test_direct_model_request_kwargs_are_allowlisted` — asserts
`instrument` never leaks; `model_settings` / `model_request_parameters`
do.

One existing test replaced:
`test_v2_message_and_response_fields_are_shaped` →
`test_shape_message_and_response_pass_through_when_no_binary` — the old
test was locking in the reshape that was losing v2 fields.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

<!-- sfk:created-approved-by -->
Created by abhijeet

<!-- sfk:slack-thread -->
[Slack
thread](https://starfolkai.slack.com/archives/C0AQDETAVT3/p1784659060950489?thread_ts=1784659060.950489&cid=C0AQDETAVT3)

---------

Co-authored-by: Starfolk <noreply@starfolk.ai>
Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@AbhiPrasad Abhijeet Prasad (AbhiPrasad) AbhiPrasad approved these changes

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@AbhiPrasad @starfolkbot