Skip to content

fix(mistral): drop eager response serialization, use allowlist metadata pass-through#604

Merged
Abhijeet Prasad (AbhiPrasad) merged 1 commit into
mainfrom
fix/mistral-drop-eager-serialization
Jul 21, 2026
Merged

fix(mistral): drop eager response serialization, use allowlist metadata pass-through#604
Abhijeet Prasad (AbhiPrasad) merged 1 commit into
mainfrom
fix/mistral-drop-eager-serialization

Conversation

@starfolkai

@starfolkai starfolkai Bot commented Jul 21, 2026

Copy link
Copy Markdown
Contributor

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_dumps 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

  • 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.
  • No cassette re-recording. HTTP behavior is unchanged; only in-process span shaping was touched.
  • cd py && nox -s "test_mistral(latest)" — 45/45 passed.
  • 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).
  • 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_failstest_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

Created by abhijeet

Slack thread

…ta pass-through

Align the Mistral integration with `.agents/skills/sdk-integrations/SKILL.md`
following the same pattern as #594 (litellm) and #583 (anthropic):

- **Response finalizers** no longer call `_normalized_mistral_dict(response)`
  to full-`model_dump` the response just to read a few fields. Read `choices`
  / `outputs` / `pages` / `usage` / `id` / `model` directly via the existing
  `_get_value` helper (dict-or-pydantic aware, handles Mistral's `Unset`
  sentinel). `bt_safe_deep_copy` in `logger.log_internal` already
  `model_dump`s pydantic values at log time.
- **`_build_request_metadata`** stops walk-and-dumping every metadata value
  (`temperature`, `top_p`, `tools`, `response_format`, `stop`, …). These
  don't carry attachments; pass through and let bt_json handle it.
- **Streaming aggregators** stop pre-dumping `usage` / `segments` into the
  intermediate dict; `_parse_usage_metrics` already accepts pydantic and
  the log-time deep-copy handles the output.
- **Tool-span walkers** (`_completion_tool_calls`, `_conversation_tool_outputs`,
  `_log_completion_tool_spans`, `_log_conversation_tool_spans`) switched to
  `_get_value` so they work on raw pydantic responses.

The multimodal walker is retained where attachment materialization matters:
the `_start_span` input path, `_append_delta_content` / `_merge_tool_calls`
streaming assembly, and per-field accumulators for conversation outputs.

Metadata is still allowlist-per-surface — no changes there.

Test updates (no new mocks, no cassette re-recording):
- `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 was a
  side-effect of `.model_dump()` failing on non-pydantic objects. The new
  behavior correctly reads usage via `getattr`; assert the positive.
- `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")`.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
@AbhiPrasad
Abhijeet Prasad (AbhiPrasad) merged commit a095270 into main Jul 21, 2026
82 checks passed
@AbhiPrasad
Abhijeet Prasad (AbhiPrasad) deleted the fix/mistral-drop-eager-serialization branch July 21, 2026 00:25
Abhijeet Prasad (AbhiPrasad) pushed a commit that referenced this pull request Jul 21, 2026
…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
…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>
Abhijeet Prasad (AbhiPrasad) pushed a commit that referenced this pull request Jul 21, 2026
…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
…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

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants