Capture raw provider traces for native-protocol (GenerateContent) agents so they reach llm_trajectory parity

[xdotli]

Problem

Native-protocol agents (those that speak a provider's own wire format, e.g. Google GenerateContent) talk to their provider directly and bypass the LiteLLM usage proxy entirely. Because the proxy is the sole producer of the raw provider request/response capture, these runs get no llm_trajectory.jsonl and report usage_source='unavailable' (no per-call token/cost usage). OpenAI/Anthropic-protocol agents that flow through the proxy get full raw capture; native agents do not. This is a cross-protocol trace-parity gap.

Data flow (proxy is the only raw-capture source):

• The raw per-call provider request/response is recorded by the LiteLLM CustomLogger callback (BenchFlowLiteLLMLogger.async_log_success_event / async_log_failure_event) which appends to the callback log. Source: src/benchflow/providers/litellm_logging.py:97 (class), :147 / :180 (success/failure events), :46 (callback_module_source). This runs only inside the proxy process.
• On proxy stop, that callback JSONL is parsed into a Trajectory of LLMExchange(request, response) records via trajectory_from_litellm_callback_log (src/benchflow/providers/litellm_logging.py:304) and stored on server.trajectory.
• The only producer of llm_trajectory.jsonl is RolloutSession._write_llm_trajectory (src/benchflow/rollout.py:3305): it reads getattr(getattr(usage_runtime, 'server', None), 'trajectory', None) and returns early if the trajectory is None or has no exchanges (:3309-3311). It is invoked at src/benchflow/rollout.py:2566, guarded by a non-None usage runtime.

Why native agents skip the proxy:

• _NATIVE_PROTOCOL_AGENTS = frozenset({"oracle", "gemini"}) (src/benchflow/providers/litellm_runtime.py:58).
• needs_litellm_runtime(agent, model) returns bool(model) and agent not in _NATIVE_PROTOCOL_AGENTS (src/benchflow/providers/litellm_runtime.py:269-271) → False for native agents.
• ensure_litellm_runtime short-circuits via _skip_litellm_runtime (src/benchflow/providers/litellm_runtime.py:914) and returns runtime=None; the agent keeps its raw provider key and talks to the provider directly.
• extract_usage(None) returns usage_unavailable() (src/benchflow/providers/litellm_runtime.py:1101).

Note: ACP-level traces (acp_trajectory.jsonl) ARE still captured for native agents — they are wired through the ACP session event stream independently of the proxy (src/benchflow/rollout.py:2064, :789). The gap is specifically the raw provider HTTP request/response and token/cost usage.

Empirical repro (against the release branch)

uv sync --extra dev --locked
uv run python -c "
import asyncio
from benchflow.providers.litellm_runtime import needs_litellm_runtime, ensure_litellm_runtime, extract_usage, _NATIVE_PROTOCOL_AGENTS
print('native set:', _NATIVE_PROTOCOL_AGENTS)
print('gemini needs proxy:', needs_litellm_runtime('gemini','gemini-2.5-pro'))
env, rt = asyncio.run(ensure_litellm_runtime(agent='gemini', agent_env={'GEMINI_API_KEY':'fake-key-xxx'}, model='gemini-2.5-pro', runtime=None, environment='host', usage_tracking='auto'))
print('runtime:', rt)
print('raw key kept:', env.get('GEMINI_API_KEY'))
print('base url injected:', env.get('BENCHFLOW_PROVIDER_BASE_URL'))
print('usage_source:', extract_usage(rt)['usage_source'])
"

Observed output:

native set: frozenset({'gemini', 'oracle'})
gemini needs proxy: False
runtime: None
raw key kept: fake-key-xxx
base url injected: None
usage_source: unavailable

With runtime=None, _write_llm_trajectory (rollout.py:3309-3311) short-circuits and writes no file; porting the method body against a stub runtime with one exchange writes the file, directly proving the native path produces no raw-capture file while the proxy path does.

Why it matters

Users comparing or debugging runs expect the same artifacts regardless of which provider protocol an agent speaks. Today a run on a native GenerateContent agent silently lacks llm_trajectory.jsonl and shows usage_source='unavailable', so there is no per-call raw request/response to diff and no token/cost accounting — while an equivalent OpenAI/Anthropic-protocol run has both. This breaks trace parity across protocols, undermines reproducibility and cost reporting, and makes native-protocol runs second-class for observability and evaluation.

Proposed approach

Provide a native-protocol capture path so these agents produce llm_trajectory.jsonl and real usage without forcing them through the OpenAI-surface proxy. Options:

1. Add a lightweight capture shim for native agents (e.g. a logging/recording proxy or interceptor in front of the native GenerateContent endpoint) that records raw request/response into the same Trajectory/LLMExchange structure consumed by _write_llm_trajectory, populating server.trajectory so the existing writer at rollout.py:3305 works unchanged.
2. Parse the native provider's response usage block (e.g. usageMetadata with promptTokenCount/candidatesTokenCount/totalTokenCount, which the proxy callback already special-cases at litellm_logging.py:213-222) into the usage metrics so usage_source is no longer unavailable for native runs.

Keep the OpenAI/Anthropic proxy path untouched; this is additive parity for the native branch. Mirror the existing Trajectory/LLMExchange shape so downstream consumers (trajectory_from_litellm_callback_log, _write_llm_trajectory) need no changes.

Acceptance criteria

• A run using a native-protocol agent (e.g. gemini) produces a non-empty trajectory/llm_trajectory.jsonl containing raw provider request/response exchanges.
• The same run reports a usage source other than unavailable (real token/cost usage parsed from the native provider response).
• OpenAI/Anthropic-protocol proxy capture is unchanged (no regression in existing llm_trajectory.jsonl content or usage).
• acp_trajectory.jsonl continues to be produced for native agents (no regression).
• Secret redaction is applied to the native capture identically to the proxy path (to_jsonl(redact_keys=True)).

Test plan

• Unit: feed a recorded native GenerateContent request/response pair through the new capture path and assert a Trajectory with one LLMExchange is produced and that _write_llm_trajectory writes the expected JSONL.
• Unit: assert usage parsing maps native usage fields to non-unavailable token/cost metrics.
• Integration: run a native-agent rollout (mock/recorded provider) and assert both llm_trajectory.jsonl and acp_trajectory.jsonl exist and that the usage-tracking metadata status is enabled, not unavailable.
• Regression: existing proxy-path trajectory and usage tests still pass.

Suggested labels: enhancement, providers, observability

[xdotli]

Part of a related set on Gemini/GenerateContent gateway parity: #671 (capture raw traces for native-protocol agents), #672 (configurable base URL for Gemini backends), #673 (faithful capture of gateway-translated payloads). #672 unblocks pointing Gemini at a custom/self-hosted endpoint; #671 and #673 close the trace-capture gaps (native path and cross-protocol gateway path respectively).

[bingran-you]

Codex triage/update on 2026-06-15: moved from triage to P1/status:ready/area:diagnostics+eval. This is higher priority than 672/673 because native-protocol agents need raw provider trace parity for healthy llm_trajectory artifacts. No linked implementation PR found in the current open PR set.

[bingran-you]

Automation validation (2026-06-16): current main still leaves this issue open. Gemini remains in the native-protocol path, needs_litellm_runtime("gemini", ...) is false, and native Gemini runs do not get a usage proxy runtime. _write_llm_trajectory() only persists usage_runtime.server.trajectory, so native GenerateContent traffic still does not produce raw provider request/response entries in trajectory/llm_trajectory.jsonl.\n\nRecent native ACP usage telemetry improves token accounting via usage_source="agent_native_acp", but that is not raw provider trace parity. None of the current open PRs (#784, #785, #788) address this. Keeping #671 as P1 / status:ready / project Todo; recommended next implementation is a native GenerateContent capture runtime/interceptor with redaction, usageMetadata parsing, and regression coverage for native llm_trajectory.jsonl.

[bingran-you]

Closing in the 2026-07-09 issue-maintenance pass: no follow-up from @bingran-you newer than 2026-07-07T06:41:02Z was found. Reopen if this still needs active BenchFlow repo work.