Skip to content

feat(openai_agents): expose real usage, response_id, plumb previous_response_id, opt-in prompt_cache_key for stateful responses and prompt caching - #335

Merged
x merged 8 commits into
mainfrom
dpeticolas/prompt-cache-upstream
May 4, 2026
Merged

x merged 8 commits into
mainfrom
dpeticolas/prompt-cache-upstream

Conversation

@x

@x x commented Apr 30, 2026 •

Copy link
Copy Markdown
Contributor

Summary

Four changes to the streaming Responses API path in temporal_streaming_model.py.

My goal is to:

  • Make it easier to track prompt cache hits (usage)
  • Make it easier to get cache hits (prompt_cache_key)
  • Make it easier to use the responses API (response_id, previous_response_id)
  1. Capture real usage from ResponseCompletedEvent.response.usage and surface in the span's output dict.

    TemporalStreamingModel.get_response was constructing a zero-filled Usage and discarding event.response.usage. It now reads usage off the streaming protocol and falls back to zeros only on the error path (stream ends without a completed event).

    The len(''.join(reasoning_contents)) // 4 reasoning-tokens estimator is dropped — the real value arrives in the API response.

    The streaming_model_get_response span now carries {input_tokens, output_tokens, total_tokens, cached_input_tokens, reasoning_tokens} so traces show cache-hit rate.

  2. Return real response_id captured from ResponseCompletedEvent.response.id instead of the previously fabricated f"resp_{uuid.uuid4().hex[:8]}".

    OpenAI Agents SDK reads this field for previous_response_id chaining (run.py:145), exposes it as RunResult.last_response_id (result.py:108), and writes it into traces (tracing/span_data.py:164). A client-side UUID has never been issued by any server, so any caller that picks it up and tries to chain (the documented use case) gets a 400 "response not found".

    Latent since this file was added (commit 2f2a6ed7) because nothing in the codebase chains previous_response_id yet.

  3. Forward previous_response_id, conversation_id, and prompt from the SDK kwarg.

    The SDK's abstract Model.get_response has previous_response_id, conversation_id, prompt as required keyword-only params; the SDK threads previous_response_id down through _ServerConversationTracker when callers set it on Runner.run / RunConfig.

    Our implementation declared **kwargs # noqa: ARG002 and silently swallowed all three. Callers who used the official chaining API got no stateful behavior and no error. Replaced with explicit named params.

    This change is not necessarily required we need to pass previous_response_id but we could be doing that via extra_args. That said, the silent behavior is sketch. I could be convinced that the right thing to here is to raise an exception when we see them in kwargs and expect them in extra_args but this felt less dangerous and more intuitive to the caller.

  4. Plumb prompt_cache_key to responses.create as an opt-in parameter.

    Callers set it via model_settings.extra_args["prompt_cache_key"]. We do not auto-inject a default. I don't trust prompt_cache_key to be standard across all OpenAI-compatible endpoints.

    When unset, the parameter resolves to NOT_GIVEN and is omitted from the request body entirely.

Test plan

All 38 tests in test_streaming_model.py pass locally:

  • test_usage_captured_from_completed_event
  • test_usage_falls_back_when_no_completed_event
  • test_usage_emitted_in_span_output
  • test_response_id_captured_from_completed_event
  • test_response_id_is_none_when_no_completed_event — guards against the fake-UUID footgun
  • test_prompt_cache_key_not_sent_by_default — verifies NOT_GIVEN fallback for non-OpenAI compat
  • test_prompt_cache_key_forwarded_when_opted_in
  • test_previous_response_id_not_sent_by_default — verifies NOT_GIVEN fallback
  • test_previous_response_id_forwarded_via_sdk_kwarg — verifies the SDK's official chaining API now works
  • test_conversation_id_and_prompt_accepted_but_not_forwarded — verifies SDK contract compliance without surface-area expansion
  • All 28 pre-existing tests still pass (after stripping vestigial task_id kwargs)

Why these changes together

The motivating workstream (downstream of this PR) is a stateful-Responses-API migration for ST&S agents to chain via previous_response_id for 40–80% better cache utilization on reasoning models. That migration needs:

  • Real Usage to measure cache hit rate before/after.
  • Real response_id to actually pass back as previous_response_id.
  • previous_response_id actually reaching responses.create instead of being silently swallowed.
  • Optional prompt_cache_key for callers who want it without forcing it on everyone.

Once a release containing these changes lands and aimi-scale bumps its agentex-sdk pin, the runtime monkey-patch in st_s/*/agentex_usage_patch.py and the corresponding _apply_agentex_usage_patch() calls in each run_worker.py can be deleted.

Greptile Summary

Four correctness fixes to the Responses API streaming path: (1) real Usage is now read from ResponseCompletedEvent.response.usage and surfaced in the span; (2) response_id is the server-issued value instead of a fabricated UUID that would 400 on any chained turn; (3) previous_response_id, conversation_id, and prompt are explicitly accepted and forwarded rather than silently swallowed by **kwargs; (4) prompt_cache_key is opt-in via extra_args and resolves to NOT_GIVEN by default.

Confidence Score: 5/5

Safe to merge — all changes are correctness improvements with safe fallback paths and comprehensive test coverage.

No P0 or P1 findings. All four changes fix real latent bugs (fake UUID chaining, zero usage, silently dropped SDK params). Fallbacks are explicit (zeros, None, NOT_GIVEN). The removal of **kwargs is intentional and acknowledged. Tests cover every new code path.

No files require special attention.

Important Files Changed

Filename Overview
src/agentex/lib/core/temporal/plugins/openai_agents/models/temporal_streaming_model.py Captures real usage and response_id from ResponseCompletedEvent, forwards previous_response_id/conversation_id/prompt as explicit kwargs (replacing **kwargs swallowing), and plumbs opt-in prompt_cache_key — all correctness improvements with safe fallbacks.
src/agentex/lib/core/temporal/plugins/openai_agents/tests/test_streaming_model.py Adds 10 targeted tests for the new behaviour (usage capture, response_id, prompt_cache_key opt-in, previous_response_id chaining, conversation/prompt forwarding) and strips the now-vestigial task_id kwarg from 28 existing test calls.

Sequence Diagram

sequenceDiagram
    participant SDK as OpenAI Agents SDK
    participant TSM as TemporalStreamingModel
    participant OAI as OpenAI responses.create
    participant Redis as Redis (streaming)

    SDK->>TSM: get_response(..., previous_response_id, conversation_id, prompt)
    Note over TSM: pops prompt_cache_key from extra_args (NOT_GIVEN default)
    TSM->>OAI: responses.create(stream=True, previous_response_id, conversation, prompt, prompt_cache_key, ...)
    loop stream events
        OAI-->>TSM: ResponseTextDeltaEvent / etc.
        TSM->>Redis: stream partial content
    end
    OAI-->>TSM: ResponseCompletedEvent(response.id, response.usage)
    Note over TSM: captured_response_id = response.id<br/>captured_usage = response.usage
    TSM-->>SDK: ModelResponse(response_id=captured_response_id, usage=real_usage)
    Note over SDK: RunResult.last_response_id = captured_response_id<br/>(used for next turn's previous_response_id)
Loading

Reviews (7): Last reviewed commit: "Merge branch 'main' into dpeticolas/prom..." | Re-trigger Greptile

Loading
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.

3 participants