diff --git a/src/agents/run_internal/run_loop.py b/src/agents/run_internal/run_loop.py index 3d21d89fda..0b29719a02 100644 --- a/src/agents/run_internal/run_loop.py +++ b/src/agents/run_internal/run_loop.py @@ -230,6 +230,31 @@ ] +def _get_model_name(model: Any) -> str | None: + """Extract the string model name from a Model instance, if available. + + Most built-in model implementations (``OpenAIResponsesModel``, + ``OpenAIChatCompletionsModel``) expose a ``model`` attribute that contains + the underlying model name string. This helper retrieves it in a + forward-compatible, duck-typed way so that third-party model + implementations are handled gracefully. + + The function tries ``model`` first, then ``model_name`` for + implementations that prefer that name. Any exception raised while + resolving the attribute (for example, a custom descriptor or property + that throws) is swallowed so that usage accounting can never crash the + run, and ``None`` is returned when no string-valued name is available. + """ + for attr in ("model", "model_name"): + try: + value = getattr(model, attr, None) + except Exception: + continue + if isinstance(value, str) and value: + return value + return None + + async def _should_persist_stream_items( *, session: Session | None, @@ -1401,7 +1426,11 @@ async def rewind_model_request() -> None: ) if final_response is not None: - context_wrapper.usage.add(final_response.usage) + context_wrapper.usage.add( + final_response.usage, + agent_name=agent.name, + model_name=_get_model_name(model), + ) await asyncio.gather( ( agent.hooks.on_llm_end(context_wrapper, agent, final_response) @@ -1656,7 +1685,11 @@ async def rewind_model_request() -> None: # new deltas. server_conversation_tracker.mark_input_as_sent(filtered.input) - context_wrapper.usage.add(new_response.usage) + context_wrapper.usage.add( + new_response.usage, + agent_name=agent.name, + model_name=_get_model_name(model), + ) await asyncio.gather( ( diff --git a/src/agents/usage.py b/src/agents/usage.py index 28b723c872..7e6552b5a5 100644 --- a/src/agents/usage.py +++ b/src/agents/usage.py @@ -43,6 +43,8 @@ def deserialize_usage(usage_data: Mapping[str, Any]) -> Usage: entry.get("output_tokens_details") or {"reasoning_tokens": 0}, OutputTokensDetails(reasoning_tokens=0), ), + agent_name=entry.get("agent_name", None), + model_name=entry.get("model_name", None), ) ) @@ -76,6 +78,20 @@ class RequestUsage: output_tokens_details: OutputTokensDetails """Details about the output tokens for this individual request.""" + agent_name: str | None = None + """Name of the agent that made this request, if available. + + Populated automatically when an agent makes a model call so that callers can attribute + token usage and costs to specific agents in multi-agent workflows. + """ + + model_name: str | None = None + """Name of the model used for this request, if available. + + Populated automatically when an agent makes a model call so that callers can attribute + token usage and costs to specific models. + """ + def _normalize_input_tokens_details( v: InputTokensDetails | PromptTokensDetails | None, @@ -154,13 +170,23 @@ def __post_init__(self) -> None: if output_details_none or output_reasoning_none: self.output_tokens_details = OutputTokensDetails(reasoning_tokens=0) - def add(self, other: Usage) -> None: + def add( + self, + other: Usage, + *, + agent_name: str | None = None, + model_name: str | None = None, + ) -> None: """Add another Usage object to this one, aggregating all fields. This method automatically preserves request_usage_entries. Args: other: The Usage object to add to this one. + agent_name: Optional name of the agent making this request, used to annotate the + resulting ``RequestUsage`` entry for per-agent cost attribution. + model_name: Optional name of the model used for this request, used to annotate the + resulting ``RequestUsage`` entry for per-model cost attribution. """ self.requests += other.requests if other.requests else 0 self.input_tokens += other.input_tokens if other.input_tokens else 0 @@ -206,11 +232,26 @@ def add(self, other: Usage) -> None: total_tokens=other.total_tokens, input_tokens_details=input_details, output_tokens_details=output_details, + agent_name=agent_name, + model_name=model_name, ) self.request_usage_entries.append(request_usage) elif other.request_usage_entries: # If the other Usage already has individual request breakdowns, merge them. - self.request_usage_entries.extend(other.request_usage_entries) + # Apply agent_name/model_name to entries that don't already have them set, + # but copy each entry rather than mutating the original objects in place + # to avoid silent mis-attribution when the same Usage is added multiple times. + for entry in other.request_usage_entries: + annotated_entry = RequestUsage( + input_tokens=entry.input_tokens, + output_tokens=entry.output_tokens, + total_tokens=entry.total_tokens, + input_tokens_details=entry.input_tokens_details, + output_tokens_details=entry.output_tokens_details, + agent_name=agent_name if (agent_name is not None and entry.agent_name is None) else entry.agent_name, + model_name=model_name if (model_name is not None and entry.model_name is None) else entry.model_name, + ) + self.request_usage_entries.append(annotated_entry) def _serialize_usage_details(details: Any, default: dict[str, int]) -> dict[str, Any]: @@ -228,7 +269,7 @@ def serialize_usage(usage: Usage) -> dict[str, Any]: output_details = _serialize_usage_details(usage.output_tokens_details, {"reasoning_tokens": 0}) def _serialize_request_entry(entry: RequestUsage) -> dict[str, Any]: - return { + result: dict[str, Any] = { "input_tokens": entry.input_tokens, "output_tokens": entry.output_tokens, "total_tokens": entry.total_tokens, @@ -239,6 +280,11 @@ def _serialize_request_entry(entry: RequestUsage) -> dict[str, Any]: entry.output_tokens_details, {"reasoning_tokens": 0} ), } + if entry.agent_name is not None: + result["agent_name"] = entry.agent_name + if entry.model_name is not None: + result["model_name"] = entry.model_name + return result return { "requests": usage.requests, diff --git a/tests/test_usage.py b/tests/test_usage.py index 2a8fcaa6d0..9e02a6b8e6 100644 --- a/tests/test_usage.py +++ b/tests/test_usage.py @@ -377,3 +377,336 @@ def test_usage_normalizes_chat_completions_types(): assert isinstance(usage.output_tokens_details, OutputTokensDetails) assert usage.output_tokens_details.reasoning_tokens == 100 + + +# ============================================================================ +# Tests for agent_name and model_name on RequestUsage (issue #2100) +# ============================================================================ + + +def test_request_usage_default_agent_model_names_are_none(): + """Backward-compat: RequestUsage without agent_name/model_name defaults to None.""" + entry = RequestUsage( + input_tokens=10, + output_tokens=5, + total_tokens=15, + input_tokens_details=InputTokensDetails(cached_tokens=0), + output_tokens_details=OutputTokensDetails(reasoning_tokens=0), + ) + assert entry.agent_name is None + assert entry.model_name is None + + +def test_serialize_deserialize_roundtrip_preserves_agent_and_model_names(): + """JSON round-trip must preserve agent_name and model_name on each entry. + + This guards against a regression where serialize_usage drops the new + attribution fields, or deserialize_usage forgets to read them back. + Both branches of the conditional emit (entry-with-name and entry-without-name) + are exercised so the all-None fast path can't silently strip the keys. + """ + from agents.usage import deserialize_usage, serialize_usage + + named_entry = RequestUsage( + input_tokens=10, + output_tokens=5, + total_tokens=15, + input_tokens_details=InputTokensDetails(cached_tokens=0), + output_tokens_details=OutputTokensDetails(reasoning_tokens=0), + agent_name="Math Tutor", + model_name="gpt-4o", + ) + unnamed_entry = RequestUsage( + input_tokens=2, + output_tokens=1, + total_tokens=3, + input_tokens_details=InputTokensDetails(cached_tokens=0), + output_tokens_details=OutputTokensDetails(reasoning_tokens=0), + ) + original = Usage( + requests=2, + input_tokens=12, + output_tokens=6, + total_tokens=18, + request_usage_entries=[named_entry, unnamed_entry], + ) + + restored = deserialize_usage(serialize_usage(original)) + + assert len(restored.request_usage_entries) == 2 + restored_named = restored.request_usage_entries[0] + restored_unnamed = restored.request_usage_entries[1] + + assert restored_named.agent_name == "Math Tutor" + assert restored_named.model_name == "gpt-4o" + assert restored_named.input_tokens == 10 + assert restored_named.output_tokens == 5 + + assert restored_unnamed.agent_name is None + assert restored_unnamed.model_name is None + assert restored_unnamed.input_tokens == 2 + + +def test_request_usage_with_agent_and_model_names(): + """RequestUsage can be created with explicit agent_name and model_name.""" + entry = RequestUsage( + input_tokens=10, + output_tokens=5, + total_tokens=15, + input_tokens_details=InputTokensDetails(cached_tokens=0), + output_tokens_details=OutputTokensDetails(reasoning_tokens=0), + agent_name="Math Tutor", + model_name="gpt-4o", + ) + assert entry.agent_name == "Math Tutor" + assert entry.model_name == "gpt-4o" + + +def test_usage_add_propagates_agent_and_model_names(): + """Usage.add() with agent_name/model_name annotates the RequestUsage entry.""" + parent = Usage() + child = Usage( + requests=1, + input_tokens=65, + output_tokens=13, + total_tokens=78, + input_tokens_details=InputTokensDetails(cached_tokens=0), + output_tokens_details=OutputTokensDetails(reasoning_tokens=0), + ) + parent.add(child, agent_name="Code Reviewer", model_name="gpt-4o-mini") + + assert len(parent.request_usage_entries) == 1 + entry = parent.request_usage_entries[0] + assert entry.agent_name == "Code Reviewer" + assert entry.model_name == "gpt-4o-mini" + assert entry.input_tokens == 65 + assert entry.output_tokens == 13 + + +def test_usage_add_without_agent_model_names_stays_none(): + """Usage.add() without agent/model names leaves them as None (backward compat).""" + parent = Usage() + child = Usage( + requests=1, + input_tokens=20, + output_tokens=10, + total_tokens=30, + input_tokens_details=InputTokensDetails(cached_tokens=0), + output_tokens_details=OutputTokensDetails(reasoning_tokens=0), + ) + parent.add(child) + + assert len(parent.request_usage_entries) == 1 + entry = parent.request_usage_entries[0] + assert entry.agent_name is None + assert entry.model_name is None + + +def test_usage_add_merge_existing_entries_applies_agent_model_names(): + """When merging existing request_usage_entries, agent/model names are applied to unset ones.""" + # An existing entry without names + existing_entry = RequestUsage( + input_tokens=100, + output_tokens=50, + total_tokens=150, + input_tokens_details=InputTokensDetails(cached_tokens=0), + output_tokens_details=OutputTokensDetails(reasoning_tokens=0), + ) + parent = Usage() + child = Usage( + requests=2, # not 1, so it won't auto-create a new entry + input_tokens=100, + output_tokens=50, + total_tokens=150, + input_tokens_details=InputTokensDetails(cached_tokens=0), + output_tokens_details=OutputTokensDetails(reasoning_tokens=0), + request_usage_entries=[existing_entry], + ) + parent.add(child, agent_name="Triage Agent", model_name="gpt-4o") + + assert len(parent.request_usage_entries) == 1 + assert parent.request_usage_entries[0].agent_name == "Triage Agent" + assert parent.request_usage_entries[0].model_name == "gpt-4o" + + +def test_usage_add_merge_existing_entries_does_not_overwrite_names(): + """Existing agent/model names on entries are not overwritten during merge.""" + existing_entry = RequestUsage( + input_tokens=100, + output_tokens=50, + total_tokens=150, + input_tokens_details=InputTokensDetails(cached_tokens=0), + output_tokens_details=OutputTokensDetails(reasoning_tokens=0), + agent_name="Already Named Agent", + model_name="already-named-model", + ) + parent = Usage() + child = Usage( + requests=2, + input_tokens=100, + output_tokens=50, + total_tokens=150, + input_tokens_details=InputTokensDetails(cached_tokens=0), + output_tokens_details=OutputTokensDetails(reasoning_tokens=0), + request_usage_entries=[existing_entry], + ) + parent.add(child, agent_name="New Agent Name", model_name="new-model") + + # The existing names should NOT be overwritten + assert parent.request_usage_entries[0].agent_name == "Already Named Agent" + assert parent.request_usage_entries[0].model_name == "already-named-model" + + +@pytest.mark.asyncio +async def test_runner_run_populates_agent_name_in_request_usage(): + """Integration: Running an agent populates agent_name in RequestUsage entries.""" + from agents.usage import Usage as AgentUsage + + model_usage = AgentUsage( + requests=1, + input_tokens=42, + output_tokens=8, + total_tokens=50, + input_tokens_details=InputTokensDetails(cached_tokens=0), + output_tokens_details=OutputTokensDetails(reasoning_tokens=0), + ) + fake = FakeModel(initial_output=[get_text_message("hello")]) + fake.set_hardcoded_usage(model_usage) + agent = Agent(name="My Assistant", model=fake) + + result = await Runner.run(agent, input="hi") + + entries = result.context_wrapper.usage.request_usage_entries + assert len(entries) == 1 + assert entries[0].agent_name == "My Assistant" + + +@pytest.mark.asyncio +async def test_runner_run_populates_model_name_in_request_usage(): + """Integration: Running an agent populates model_name in RequestUsage entries.""" + from agents.usage import Usage as AgentUsage + + model_usage = AgentUsage( + requests=1, + input_tokens=30, + output_tokens=10, + total_tokens=40, + input_tokens_details=InputTokensDetails(cached_tokens=0), + output_tokens_details=OutputTokensDetails(reasoning_tokens=0), + ) + # FakeModel doesn't expose a `.model` attribute by default → model_name should be None + # We give it one to test that model_name is picked up. + fake = FakeModel(initial_output=[get_text_message("ok")]) + fake.model = "test-model-name" # type: ignore[attr-defined] + fake.set_hardcoded_usage(model_usage) + agent = Agent(name="Model-Aware Agent", model=fake) + + result = await Runner.run(agent, input="ping") + + entries = result.context_wrapper.usage.request_usage_entries + assert len(entries) == 1 + assert entries[0].model_name == "test-model-name" + + +@pytest.mark.asyncio +async def test_multi_agent_run_attributes_usage_to_correct_agents(): + """Multi-agent scenario: each RequestUsage entry has the right agent_name.""" + + from agents.usage import Usage as AgentUsage + from tests.test_responses import get_handoff_tool_call + + # Two separate models so we can track which agent's usage is which + triage_usage = AgentUsage( + requests=1, + input_tokens=100, + output_tokens=10, + total_tokens=110, + input_tokens_details=InputTokensDetails(cached_tokens=0), + output_tokens_details=OutputTokensDetails(reasoning_tokens=0), + ) + specialist_usage = AgentUsage( + requests=1, + input_tokens=200, + output_tokens=20, + total_tokens=220, + input_tokens_details=InputTokensDetails(cached_tokens=0), + output_tokens_details=OutputTokensDetails(reasoning_tokens=0), + ) + + specialist_model = FakeModel(initial_output=[get_text_message("specialist done")]) + specialist_model.model = "gpt-4o-specialist" # type: ignore[attr-defined] + specialist_model.set_hardcoded_usage(specialist_usage) + specialist_agent = Agent(name="Specialist Agent", model=specialist_model) + + triage_model = FakeModel() + triage_model.model = "gpt-4o-triage" # type: ignore[attr-defined] + triage_model.add_multiple_turn_outputs( + [ + [get_handoff_tool_call(specialist_agent)], + ] + ) + triage_model.set_hardcoded_usage(triage_usage) + triage_agent = Agent(name="Triage Agent", model=triage_model, handoffs=[specialist_agent]) + + result = await Runner.run(triage_agent, input="route me") + + all_entries = result.context_wrapper.usage.request_usage_entries + assert len(all_entries) == 2, f"Expected 2 request entries, got {len(all_entries)}" + + agent_names = [e.agent_name for e in all_entries] + assert "Triage Agent" in agent_names, f"Expected 'Triage Agent' in {agent_names}" + assert "Specialist Agent" in agent_names, f"Expected 'Specialist Agent' in {agent_names}" + + triage_entry = next(e for e in all_entries if e.agent_name == "Triage Agent") + assert triage_entry.input_tokens == 100 + assert triage_entry.model_name == "gpt-4o-triage", ( + f"Triage entry model_name should be 'gpt-4o-triage', got {triage_entry.model_name!r}" + ) + + specialist_entry = next(e for e in all_entries if e.agent_name == "Specialist Agent") + assert specialist_entry.input_tokens == 200 + assert specialist_entry.model_name == "gpt-4o-specialist", ( + "Specialist entry model_name should be 'gpt-4o-specialist', " + f"got {specialist_entry.model_name!r}" + ) + + +def test_add_does_not_mutate_other_entries() -> None: + """Adding a Usage with existing request_usage_entries must not mutate the original entries. + + Previously, the elif branch in Usage.add() called entry.agent_name = ... directly on + the objects inside other.request_usage_entries, causing silent mis-attribution when the + same Usage object was re-used or added to multiple aggregators. + """ + from openai.types.responses.response_usage import InputTokensDetails, OutputTokensDetails + + source_entry = RequestUsage( + input_tokens=50, + output_tokens=25, + total_tokens=75, + input_tokens_details=InputTokensDetails(cached_tokens=0), + output_tokens_details=OutputTokensDetails(reasoning_tokens=0), + agent_name=None, + model_name=None, + ) + + # Build a Usage that already has request_usage_entries (requests != 1 path) + other = Usage( + requests=2, + input_tokens=50, + output_tokens=25, + total_tokens=75, + request_usage_entries=[source_entry], + ) + + agg = Usage() + agg.add(other, agent_name="MyAgent", model_name="gpt-4o") + + # The aggregator should have a copy with the annotation applied + assert len(agg.request_usage_entries) == 1 + assert agg.request_usage_entries[0].agent_name == "MyAgent" + assert agg.request_usage_entries[0].model_name == "gpt-4o" + + # The original entry must NOT be mutated + assert source_entry.agent_name is None, "Original entry was mutated!" + assert source_entry.model_name is None, "Original entry was mutated!"