diff --git a/go/adk/pkg/a2a/executor.go b/go/adk/pkg/a2a/executor.go index f4ba0ab5d..08937a085 100644 --- a/go/adk/pkg/a2a/executor.go +++ b/go/adk/pkg/a2a/executor.go @@ -270,7 +270,11 @@ func (e *KAgentExecutor) Execute(ctx context.Context, reqCtx *a2asrv.RequestCont lastNonPartialParts a2atype.ContentParts hitlParts a2atype.ContentParts runErr error + usage turnUsage ) + // Resumed tasks (HITL cycles, follow-up messages) carry the previously + // persisted total; seed it so kagent_usage_total stays a task-lifetime sum. + usage.seedFromTask(reqCtx.StoredTask) for adkEvent, adkErr := range r.Run(ctx, userID, sessionID, content, runConfig) { if adkErr != nil { @@ -287,6 +291,9 @@ func (e *KAgentExecutor) Execute(ctx context.Context, reqCtx *a2asrv.RequestCont invocationSpan.SetAttributes(attribute.String("gcp.vertex.agent.invocation_id", invocationID)) } + // Aggregate token usage for the terminal status update. + usage.add(adkEvent) + // Build per-event metadata (inherits baseMeta + adds invocation_id, usage etc.). eventMeta := buildEventMeta(baseMeta, adkEvent) @@ -299,6 +306,7 @@ func (e *KAgentExecutor) Execute(ctx context.Context, reqCtx *a2asrv.RequestCont a2atype.TextPart{Text: fmt.Sprintf("LLM error: %s %s", adkEvent.ErrorCode, adkEvent.ErrorMessage)}) failed := a2atype.NewStatusUpdateEvent(reqCtx, a2atype.TaskStateFailed, errMsg) failed.Final = true + usage.stamp(eventMeta) failed.Metadata = eventMeta return queue.Write(ctx, failed) } @@ -311,6 +319,7 @@ func (e *KAgentExecutor) Execute(ctx context.Context, reqCtx *a2asrv.RequestCont a2atype.TextPart{Text: fmt.Sprintf("LLM error: %s %s", adkEvent.ErrorCode, adkEvent.ErrorMessage)}) failed := a2atype.NewStatusUpdateEvent(reqCtx, a2atype.TaskStateFailed, errMsg) failed.Final = true + usage.stamp(eventMeta) failed.Metadata = eventMeta return queue.Write(ctx, failed) } @@ -383,6 +392,7 @@ func (e *KAgentExecutor) Execute(ctx context.Context, reqCtx *a2asrv.RequestCont if invocationID != "" { finalMeta[adka2a.ToA2AMetaKey("invocation_id")] = invocationID } + usage.stamp(finalMeta) if runErr != nil { errMsg := newAgentMessage(reqCtx, a2atype.TextPart{Text: runErr.Error()}) diff --git a/go/adk/pkg/a2a/usage.go b/go/adk/pkg/a2a/usage.go new file mode 100644 index 000000000..301131bfd --- /dev/null +++ b/go/adk/pkg/a2a/usage.go @@ -0,0 +1,110 @@ +package a2a + +import ( + "math" + + a2atype "github.com/a2aproject/a2a-go/a2a" + adksession "google.golang.org/adk/v2/session" + "google.golang.org/genai" +) + +// turnUsage accumulates token usage across the ADK events of one execution so +// the aggregated total can be emitted on terminal status updates. Partial +// (streaming chunk) events are skipped: each LLM call reports its usage on the +// final non-partial event, so summing partials would double-count. +type turnUsage struct { + promptTokens int64 + completionTokens int64 + thoughtsTokens int64 + cachedContentTokens int64 + totalTokens int64 + modelVersion string +} + +func (u *turnUsage) add(event *adksession.Event) { + if event == nil || event.Partial || event.UsageMetadata == nil { + return + } + u.promptTokens += int64(event.UsageMetadata.PromptTokenCount) + u.completionTokens += int64(event.UsageMetadata.CandidatesTokenCount) + u.thoughtsTokens += int64(event.UsageMetadata.ThoughtsTokenCount) + u.cachedContentTokens += int64(event.UsageMetadata.CachedContentTokenCount) + u.totalTokens += int64(event.UsageMetadata.TotalTokenCount) + if event.ModelVersion != "" { + u.modelVersion = event.ModelVersion + } +} + +// seedFromTask primes the accumulator with the total already persisted on a +// resumed task, so tasks spanning multiple executions (HITL input-required +// cycles, follow-up messages) report a task-lifetime total instead of the last +// segment only. +func (u *turnUsage) seedFromTask(task *a2atype.Task) { + if task == nil || task.Metadata == nil { + return + } + prior, ok := task.Metadata[GetKAgentMetadataKey("usage_total")].(map[string]any) + if !ok { + return + } + u.promptTokens += metadataTokenCount(prior["promptTokenCount"]) + u.completionTokens += metadataTokenCount(prior["candidatesTokenCount"]) + u.thoughtsTokens += metadataTokenCount(prior["thoughtsTokenCount"]) + u.cachedContentTokens += metadataTokenCount(prior["cachedContentTokenCount"]) + u.totalTokens += metadataTokenCount(prior["totalTokenCount"]) + if modelVersion, ok := prior["modelVersion"].(string); ok && modelVersion != "" { + u.modelVersion = modelVersion + } +} + +// metadataTokenCount reads a numeric token count from stored task metadata. +// Counts are float64 after a JSON round-trip but keep an integer type with +// in-memory task stores. +func metadataTokenCount(value any) int64 { + switch n := value.(type) { + case float64: + return int64(n) + case int64: + return n + case int32: + return int64(n) + case int: + return int64(n) + default: + return 0 + } +} + +func (u *turnUsage) empty() bool { + return u.promptTokens == 0 && u.completionTokens == 0 && u.totalTokens == 0 +} + +// stamp attaches the aggregate to meta under kagent_usage_total. The value is +// serialized exactly like the per-event adk_usage_metadata (same genai type, +// same JSON mapping) plus modelVersion, so consumers can share one parser. +func (u *turnUsage) stamp(meta map[string]any) { + if u.empty() { + return + } + total, err := toA2AMetadataMap(&genai.GenerateContentResponseUsageMetadata{ + PromptTokenCount: clampInt32(u.promptTokens), + CandidatesTokenCount: clampInt32(u.completionTokens), + ThoughtsTokenCount: clampInt32(u.thoughtsTokens), + CachedContentTokenCount: clampInt32(u.cachedContentTokens), + TotalTokenCount: clampInt32(u.totalTokens), + }) + if err != nil || total == nil { + return + } + if u.modelVersion != "" { + total["modelVersion"] = u.modelVersion + } + meta[GetKAgentMetadataKey("usage_total")] = total +} + +func clampInt32(v int64) int32 { + if v > math.MaxInt32 { + return math.MaxInt32 + } + return int32(v) +} diff --git a/go/adk/pkg/a2a/usage_test.go b/go/adk/pkg/a2a/usage_test.go new file mode 100644 index 000000000..056997465 --- /dev/null +++ b/go/adk/pkg/a2a/usage_test.go @@ -0,0 +1,122 @@ +package a2a + +import ( + "math" + "testing" + + a2atype "github.com/a2aproject/a2a-go/a2a" + "github.com/stretchr/testify/require" + "google.golang.org/adk/v2/model" + "google.golang.org/adk/v2/server/adka2a" //nolint:staticcheck // kagent still uses a2a-go v1; this ADK package is the compatibility adapter. + adksession "google.golang.org/adk/v2/session" + "google.golang.org/genai" +) + +func usageEvent(prompt, completion, total int32, modelVersion string, partial bool) *adksession.Event { + return &adksession.Event{ + LLMResponse: model.LLMResponse{ + UsageMetadata: &genai.GenerateContentResponseUsageMetadata{ + PromptTokenCount: prompt, + CandidatesTokenCount: completion, + TotalTokenCount: total, + }, + ModelVersion: modelVersion, + Partial: partial, + }, + } +} + +func stampedTotal(t *testing.T, usage *turnUsage) map[string]any { + t.Helper() + meta := map[string]any{} + usage.stamp(meta) + total, ok := meta[GetKAgentMetadataKey("usage_total")].(map[string]any) + require.True(t, ok, "kagent_usage_total must be stamped") + return total +} + +func TestTurnUsageAggregatesNonPartialEvents(t *testing.T) { + var usage turnUsage + require.True(t, usage.empty()) + + usage.add(usageEvent(100, 20, 120, "model-a", false)) + usage.add(usageEvent(200, 30, 230, "model-b", false)) + + require.False(t, usage.empty()) + require.Equal(t, map[string]any{ + "promptTokenCount": float64(300), + "candidatesTokenCount": float64(50), + "totalTokenCount": float64(350), + "modelVersion": "model-b", + }, stampedTotal(t, &usage)) +} + +func TestTurnUsageSkipsPartialAndEmptyEvents(t *testing.T) { + var usage turnUsage + + usage.add(nil) + usage.add(&adksession.Event{}) + usage.add(usageEvent(999, 999, 999, "chunk-model", true)) + + require.True(t, usage.empty()) + meta := map[string]any{} + usage.stamp(meta) + require.NotContains(t, meta, GetKAgentMetadataKey("usage_total"), + "empty usage must not stamp the key") + + usage.add(usageEvent(10, 5, 15, "", false)) + require.Equal(t, map[string]any{ + "promptTokenCount": float64(10), + "candidatesTokenCount": float64(5), + "totalTokenCount": float64(15), + }, stampedTotal(t, &usage), "modelVersion must be omitted when no event carried one") +} + +func TestTurnUsageShapeMatchesPerEventUsageMetadata(t *testing.T) { + event := usageEvent(10, 5, 15, "", false) + + var usage turnUsage + usage.add(event) + + perEventMeta := buildEventMeta(map[string]any{}, event) + require.Equal(t, + perEventMeta[adka2a.ToA2AMetaKey("usage_metadata")], + stampedTotal(t, &usage), + "kagent_usage_total must serialize identically to adk_usage_metadata") +} + +func TestTurnUsageSeedFromTaskAccumulatesAcrossExecutions(t *testing.T) { + var usage turnUsage + // float64 values mimic a JSON round-trip through the task store. + usage.seedFromTask(&a2atype.Task{Metadata: map[string]any{ + GetKAgentMetadataKey("usage_total"): map[string]any{ + "promptTokenCount": float64(100), + "candidatesTokenCount": float64(20), + "totalTokenCount": float64(120), + "modelVersion": "model-a", + }, + }}) + usage.add(usageEvent(200, 30, 230, "", false)) + + require.Equal(t, map[string]any{ + "promptTokenCount": float64(300), + "candidatesTokenCount": float64(50), + "totalTokenCount": float64(350), + "modelVersion": "model-a", + }, stampedTotal(t, &usage)) +} + +func TestTurnUsageSeedFromTaskIgnoresMissingOrMalformed(t *testing.T) { + var usage turnUsage + usage.seedFromTask(nil) + usage.seedFromTask(&a2atype.Task{}) + usage.seedFromTask(&a2atype.Task{Metadata: map[string]any{ + GetKAgentMetadataKey("usage_total"): "not-a-map", + }}) + require.True(t, usage.empty()) +} + +func TestClampInt32(t *testing.T) { + require.Equal(t, int32(42), clampInt32(42)) + require.Equal(t, int32(math.MaxInt32), clampInt32(math.MaxInt32+1)) +} diff --git a/python/packages/kagent-adk/src/kagent/adk/_agent_executor.py b/python/packages/kagent-adk/src/kagent/adk/_agent_executor.py index 0ed10a317..c666719b2 100644 --- a/python/packages/kagent-adk/src/kagent/adk/_agent_executor.py +++ b/python/packages/kagent-adk/src/kagent/adk/_agent_executor.py @@ -53,6 +53,7 @@ from ._mcp_toolset import is_anyio_cross_task_cancel_scope_error from ._remote_a2a_tool import SubagentSessionProvider +from ._turn_usage import TurnUsage from .converters.event_converter import convert_event_to_a2a_events, serialize_metadata_value from .converters.part_converter import convert_a2a_part_to_genai_part, convert_genai_part_to_a2a_part from .converters.request_converter import convert_a2a_request_to_adk_run_args @@ -241,12 +242,15 @@ async def _execute_impl( # Handle the request and publish updates to the event queue runner = await self._resolve_runner() + # Shared with _handle_request so failed terminals still carry the + # metadata accumulated before the error (invocation_id, usage total). + run_metadata: dict[str, Any] = {} try: - await self._handle_request(context, event_queue, runner, run_args) + await self._handle_request(context, event_queue, runner, run_args, run_metadata) except asyncio.CancelledError as e: logger.error("A2A request execution was cancelled", exc_info=True) error_message = str(e) or "A2A request execution was cancelled." - await self._publish_failed_status_event(context, event_queue, error_message) + await self._publish_failed_status_event(context, event_queue, error_message, metadata=run_metadata) except Exception as e: logger.error("Error handling A2A request: %s", e, exc_info=True) @@ -267,7 +271,7 @@ async def _execute_impl( "2. Use a model that supports function calling (e.g., OpenAI, Anthropic, or Gemini models)." ) # Publish failure event - await self._publish_failed_status_event(context, event_queue, error_message) + await self._publish_failed_status_event(context, event_queue, error_message, metadata=run_metadata) finally: clear_kagent_span_attributes(context_token) # close the runner which cleans up the mcptoolsets @@ -321,6 +325,7 @@ async def _publish_failed_status_event( context: RequestContext, event_queue: EventQueue, error_message: str, + metadata: Optional[dict[str, Any]] = None, ) -> None: try: await event_queue.enqueue_event( @@ -337,6 +342,7 @@ async def _publish_failed_status_event( ), context_id=context.context_id, final=True, + metadata=metadata or None, ) ) except BaseException as enqueue_error: @@ -526,6 +532,7 @@ async def _handle_request( event_queue: EventQueue, runner: Runner, run_args: dict[str, Any], + run_metadata: dict[str, Any], ): # ensure the session exists session = await self._prepare_session(context, run_args, runner) @@ -561,11 +568,13 @@ async def _handle_request( ) # Base metadata for events (invocation_id will be updated once we see it from ADK) - run_metadata = { - get_kagent_metadata_key("app_name"): runner.app_name, - get_kagent_metadata_key("user_id"): run_args["user_id"], - get_kagent_metadata_key("session_id"): run_args["session_id"], - } + run_metadata.update( + { + get_kagent_metadata_key("app_name"): runner.app_name, + get_kagent_metadata_key("user_id"): run_args["user_id"], + get_kagent_metadata_key("session_id"): run_args["session_id"], + } + ) # publish the task working event await event_queue.enqueue_event( @@ -587,6 +596,12 @@ async def _handle_request( real_invocation_id: str | None = None last_usage_metadata = None + # Aggregate token usage across the run for the terminal status update. + # Resumed tasks (HITL cycles, follow-up messages) carry the previously + # persisted total; seed it so kagent_usage_total stays a task-lifetime sum. + turn_usage = TurnUsage() + turn_usage.seed_from_task(context.current_task) + # Build a mapping of tool name -> subagent session ID once so the # event converter can stamp it onto function_call DataParts. subagent_session_ids: dict[str, str] = {} @@ -595,36 +610,43 @@ async def _handle_request( subagent_session_ids[tool.name] = tool.subagent_session_id task_result_aggregator = TaskResultAggregator() - async with Aclosing(runner.run_async(**run_args)) as agen: - async for adk_event in agen: - # Capture the real invocation_id from the first ADK event that has one - event_inv_id = getattr(adk_event, "invocation_id", None) - if event_inv_id and not real_invocation_id: - real_invocation_id = event_inv_id - run_metadata[get_kagent_metadata_key("invocation_id")] = real_invocation_id - - # Track the last usage_metadata so it can be included in the final - # event's run_metadata. The A2A task_manager merges run_metadata into - # task.metadata, making it available to callers (e.g. KAgentRemoteA2ATool). - if getattr(adk_event, "usage_metadata", None) is not None: - last_usage_metadata = adk_event.usage_metadata - - for a2a_event in convert_event_to_a2a_events( - adk_event, - invocation_context, - context.task_id, - context.context_id, - subagent_session_ids=subagent_session_ids or None, - ): - # Only aggregate non-partial events to avoid duplicates from streaming chunks - # Partial events are sent to frontend for display but not accumulated - if not adk_event.partial: - task_result_aggregator.process_event(a2a_event) - await event_queue.enqueue_event(a2a_event) - - # Break on confirmation events that use long running tools - if getattr(adk_event, "long_running_tool_ids", None): - break + try: + async with Aclosing(runner.run_async(**run_args)) as agen: + async for adk_event in agen: + # Capture the real invocation_id from the first ADK event that has one + event_inv_id = getattr(adk_event, "invocation_id", None) + if event_inv_id and not real_invocation_id: + real_invocation_id = event_inv_id + run_metadata[get_kagent_metadata_key("invocation_id")] = real_invocation_id + + # Track the last usage_metadata so it can be included in the final + # event's run_metadata. The A2A task_manager merges run_metadata into + # task.metadata, making it available to callers (e.g. KAgentRemoteA2ATool). + if getattr(adk_event, "usage_metadata", None) is not None: + last_usage_metadata = adk_event.usage_metadata + + turn_usage.add(adk_event) + + for a2a_event in convert_event_to_a2a_events( + adk_event, + invocation_context, + context.task_id, + context.context_id, + subagent_session_ids=subagent_session_ids or None, + ): + # Only aggregate non-partial events to avoid duplicates from streaming chunks + # Partial events are sent to frontend for display but not accumulated + if not adk_event.partial: + task_result_aggregator.process_event(a2a_event) + await event_queue.enqueue_event(a2a_event) + + # Break on confirmation events that use long running tools + if getattr(adk_event, "long_running_tool_ids", None): + break + finally: + # Stamp on every exit so failed terminals published by the caller + # carry the tokens burned before the error too. + turn_usage.stamp(run_metadata) # Attach the last LLM usage to run_metadata so the A2A task_manager # merges it into task.metadata on the completed Task object. diff --git a/python/packages/kagent-adk/src/kagent/adk/_turn_usage.py b/python/packages/kagent-adk/src/kagent/adk/_turn_usage.py new file mode 100644 index 000000000..ee20397e3 --- /dev/null +++ b/python/packages/kagent-adk/src/kagent/adk/_turn_usage.py @@ -0,0 +1,98 @@ +from __future__ import annotations + +from typing import Any, Optional + +from a2a.types import Task +from google.adk.events import Event +from google.genai import types as genai_types +from kagent.core.a2a import get_kagent_metadata_key + +from .converters.event_converter import serialize_metadata_value + +USAGE_TOTAL_KEY = get_kagent_metadata_key("usage_total") + + +class TurnUsage: + """Accumulates token usage across the ADK events of one execution so the + aggregated total can be emitted on terminal status updates. + + Partial (streaming chunk) events are skipped: each LLM call reports its + usage on the final non-partial event, so summing partials would + double-count. + """ + + def __init__(self) -> None: + self.prompt_tokens = 0 + self.completion_tokens = 0 + self.thoughts_tokens = 0 + self.cached_content_tokens = 0 + self.total_tokens = 0 + self.model_version: Optional[str] = None + + def add(self, event: Optional[Event]) -> None: + if event is None or event.partial: + return + usage = event.usage_metadata + if usage is None: + return + self.prompt_tokens += usage.prompt_token_count or 0 + self.completion_tokens += usage.candidates_token_count or 0 + self.thoughts_tokens += usage.thoughts_token_count or 0 + self.cached_content_tokens += usage.cached_content_token_count or 0 + self.total_tokens += usage.total_token_count or 0 + if event.model_version: + self.model_version = event.model_version + + def seed_from_task(self, task: Optional[Task]) -> None: + """Prime the accumulator with the total already persisted on a resumed + task, so tasks spanning multiple executions (HITL input-required + cycles, follow-up messages) report a task-lifetime total instead of the + last segment only.""" + if task is None or not task.metadata: + return + prior = task.metadata.get(USAGE_TOTAL_KEY) + if not isinstance(prior, dict): + return + self.prompt_tokens += _token_count(prior.get("promptTokenCount")) + self.completion_tokens += _token_count(prior.get("candidatesTokenCount")) + self.thoughts_tokens += _token_count(prior.get("thoughtsTokenCount")) + self.cached_content_tokens += _token_count(prior.get("cachedContentTokenCount")) + self.total_tokens += _token_count(prior.get("totalTokenCount")) + model_version = prior.get("modelVersion") + if isinstance(model_version, str) and model_version: + self.model_version = model_version + + def empty(self) -> bool: + return self.prompt_tokens == 0 and self.completion_tokens == 0 and self.total_tokens == 0 + + def stamp(self, metadata: dict[str, Any]) -> None: + """Attach the aggregate to metadata under kagent_usage_total. The value + is serialized exactly like the per-event kagent_usage_metadata (same + genai type, same serializer) plus modelVersion, so consumers can share + one parser.""" + if self.empty(): + return + total = serialize_metadata_value( + genai_types.GenerateContentResponseUsageMetadata( + prompt_token_count=self.prompt_tokens or None, + candidates_token_count=self.completion_tokens or None, + thoughts_token_count=self.thoughts_tokens or None, + cached_content_token_count=self.cached_content_tokens or None, + total_token_count=self.total_tokens or None, + ) + ) + if not isinstance(total, dict): + return + if self.model_version: + total["modelVersion"] = self.model_version + metadata[USAGE_TOTAL_KEY] = total + + +def _token_count(value: Any) -> int: + """Read a numeric token count from stored task metadata; counts may be int + or float depending on the task store's JSON round-trip.""" + if isinstance(value, bool): + return 0 + if isinstance(value, (int, float)): + return int(value) + return 0 diff --git a/python/packages/kagent-adk/tests/unittests/test_turn_usage.py b/python/packages/kagent-adk/tests/unittests/test_turn_usage.py new file mode 100644 index 000000000..cc9086f45 --- /dev/null +++ b/python/packages/kagent-adk/tests/unittests/test_turn_usage.py @@ -0,0 +1,120 @@ +from a2a.types import Task, TaskState, TaskStatus +from google.adk.events import Event +from google.genai import types as genai_types + +from kagent.adk._turn_usage import USAGE_TOTAL_KEY, TurnUsage +from kagent.adk.converters.event_converter import serialize_metadata_value + + +def usage_event(prompt: int, completion: int, total: int, model_version: str | None, partial: bool) -> Event: + return Event( + author="agent", + partial=partial, + model_version=model_version, + usage_metadata=genai_types.GenerateContentResponseUsageMetadata( + prompt_token_count=prompt, + candidates_token_count=completion, + total_token_count=total, + ), + ) + + +def stamped_total(usage: TurnUsage) -> dict: + metadata: dict = {} + usage.stamp(metadata) + assert USAGE_TOTAL_KEY in metadata, "kagent_usage_total must be stamped" + return metadata[USAGE_TOTAL_KEY] + + +def task_with_total(total: dict) -> Task: + return Task( + id="task-1", + context_id="ctx-1", + status=TaskStatus(state=TaskState.input_required), + metadata={USAGE_TOTAL_KEY: total}, + ) + + +def test_aggregates_non_partial_events(): + usage = TurnUsage() + assert usage.empty() + + usage.add(usage_event(100, 20, 120, "model-a", partial=False)) + usage.add(usage_event(200, 30, 230, "model-b", partial=False)) + + assert not usage.empty() + assert stamped_total(usage) == { + "promptTokenCount": 300, + "candidatesTokenCount": 50, + "totalTokenCount": 350, + "modelVersion": "model-b", + } + + +def test_skips_partial_and_empty_events(): + usage = TurnUsage() + + usage.add(None) + usage.add(Event(author="agent")) + usage.add(usage_event(999, 999, 999, "chunk-model", partial=True)) + + assert usage.empty() + metadata: dict = {} + usage.stamp(metadata) + assert USAGE_TOTAL_KEY not in metadata, "empty usage must not stamp the key" + + usage.add(usage_event(10, 5, 15, None, partial=False)) + assert stamped_total(usage) == { + "promptTokenCount": 10, + "candidatesTokenCount": 5, + "totalTokenCount": 15, + }, "modelVersion must be omitted when no event carried one" + + +def test_shape_matches_per_event_usage_metadata(): + event = usage_event(10, 5, 15, None, partial=False) + + usage = TurnUsage() + usage.add(event) + + assert stamped_total(usage) == serialize_metadata_value(event.usage_metadata), ( + "kagent_usage_total must serialize identically to kagent_usage_metadata" + ) + + +def test_seed_from_task_accumulates_across_executions(): + usage = TurnUsage() + usage.seed_from_task( + task_with_total( + { + "promptTokenCount": 100, + "candidatesTokenCount": 20, + "totalTokenCount": 120, + "modelVersion": "model-a", + } + ) + ) + usage.add(usage_event(200, 30, 230, None, partial=False)) + + assert stamped_total(usage) == { + "promptTokenCount": 300, + "candidatesTokenCount": 50, + "totalTokenCount": 350, + "modelVersion": "model-a", + } + + +def test_seed_from_task_handles_float_counts(): + usage = TurnUsage() + usage.seed_from_task(task_with_total({"promptTokenCount": 100.0, "totalTokenCount": 120.0})) + + assert usage.prompt_tokens == 100 + assert usage.total_tokens == 120 + + +def test_seed_from_task_ignores_missing_or_malformed(): + usage = TurnUsage() + usage.seed_from_task(None) + usage.seed_from_task(Task(id="t", context_id="c", status=TaskStatus(state=TaskState.completed))) + usage.seed_from_task(task_with_total("not-a-dict")) # type: ignore[arg-type] + assert usage.empty()