Skip to content

feat(agent): record where each step's time went in the trace - #4067

Merged
kovtcharov-amd merged 7 commits into
amd:mainfrom
kovtcharov:feat/step-timing-trace
Sep 26, 2026
Merged

kovtcharov-amd merged 7 commits into
amd:mainfrom
kovtcharov:feat/step-timing-trace

Conversation

@kovtcharov

Copy link
Copy Markdown
Contributor

On a coding benchmark GAIA took 2,916 s against Claude Code's 557 s, and GAIA's own trace couldn't say why. Each step recorded token counts but no timing, so the answer (long reasoning bursts at ~50 tok/s) had to be rebuilt from Lemonade's server log. Now every step's stats record says how long the model, the tools and everything else took, with a time to first token, a finish reason, reasoning tokens and a decode rate. The result dict also gains a run summary that names the three slowest steps and why they were slow.

  • Times come from the agent's own clock and tokens from each response's own usage, not Lemonade's /stats, so the numbers are correct for cloud models too.
  • A streamed cloud call used to report no tokens at all, so its step got no stats record. It now asks for usage and gets one.
🔍 Technical details

Real non-streamed step 1 against fireworks.deepseek-v4-flash-0731. The existing keys (type, step, performance_stats) are unchanged:

{"type": "stats", "step": 1,
 "performance_stats": {"prompt_tokens": 472, "completion_tokens": 69, "total_tokens": 541, "reasoning_tokens": 9},
 "step_seconds": 1.1254, "llm_seconds": 1.1238, "tool_seconds": 0.0002, "overhead_seconds": 0.0014,
 "ttft_seconds": null, "finish_reason": "tool_calls", "reasoning_tokens": 9,
 "output_tok_per_s": 61.4, "llm_calls": 1, "tools": [{"name": "add", "seconds": 0.0002}]}
  • Each step runs from the top of its loop iteration to the top of the next one. overhead_seconds is whatever is left over, so the three parts always add up to step_seconds. The record is filled in when the step closes, because the step keeps running after the record is appended.
  • AgentSDK times every request, including failed and cancelled ones, and passes it to a per-turn sink. Only calls made on the loop thread are counted. An LLM call made inside a tool body is already part of that tool's time.
  • output_tok_per_s: for a streamed call it is the tokens after the first divided by the time after the first token. A non-streamed call has no first-token time, so the rate is taken over the whole request and is a lower bound.
  • ttft_seconds is recorded by the provider at the first content, reasoning or tool-call fragment. Tool-call-only streams yield nothing until the end, so timing the first chunk the SDK receives would report the whole request as ttft.
  • get_last_finish_reason() has the same name and meaning as in the cut-off-continuation branch, which keeps the merge simple.
  • A local stream still reports /stats in performance_stats, because that carries the ttft the UI shows. stream_options.include_usage is sent only to cloud models, so the request to local backends (llama.cpp, FLM) is unchanged.
  • No eval run: no prompt, tool schema, parsing or error classification changed.

Test plan

  • python -m pytest tests/unit/agents/test_step_timing.py -q: 17 tests covering field types, the parts summing to the step (fake clock and a real loop), ordering of the slowest steps, SDK sink contents, Lemonade usage, ttft and finish reason, and the cloud-only include_usage request
  • python -m pytest tests/unit tests/test_sdk.py -q shows no new failures against main
  • Against a live Lemonade with a cloud model, run one agent turn with streaming=False and one with streaming=True. Both should show the new fields on each stats record and timing_summary in the result

A benchmark run took 2,916 s against a reference agent's 557 s, and the trace
could not say why: each step's stats record carried token counts but no
timing, so the breakdown had to be rebuilt from Lemonade's server log.

Each step's stats record now also carries step_seconds, llm_seconds,
tool_seconds and overhead_seconds (the remainder, so the three always add up
to the step), ttft_seconds for streamed calls, finish_reason,
reasoning_tokens, output_tok_per_s, the number of model calls the step made
(retries and continuations), and per-tool seconds. process_query's result
gains timing_summary: run totals plus the three slowest steps and why.

Times come from the agent's own clock and tokens from each response's own
usage, never /stats, so the figures hold for cloud models. A cloud stream now
asks for usage, so a streamed cloud step gets a stats record at all.
@github-actions github-actions Bot added documentation Documentation changes chat Chat SDK changes llm LLM backend changes tests Test changes performance Performance-critical changes agents labels Sep 21, 2026
@github-actions

Copy link
Copy Markdown
Contributor

Approve with suggestions

This adds a per-step time breakdown to the agent trace — how long the model took, how long the tools took, and what's left over — plus a run summary naming the three slowest steps. The design is sound: overhead is computed as the remainder so the parts always reconcile against the step's wall clock, and tokens come from each response's own usage rather than Lemonade's server-global /stats, which is what makes the numbers trustworthy for cloud models.

The one thing worth settling before merge: to get token counts out of a streamed cloud call, the PR starts sending a new field on the outgoing request to Lemonade. That's a contract with a service outside this repo, and it's currently only verified against a mocked SDK — the test proves we send the field, not that Lemonade accepts it. The real trace output in the description is from the non-streamed path. Streaming is the default path for chat and the UI, so a rejected request there would break cloud chat outright rather than just lose a metric. One live run against a cloud model with streaming on, pasted into the description, closes it. This is the same shape as the #1655 regression the project docs call out, which is why it's worth the extra minute rather than a shrug.

Everything else is small: one timing call sits outside the error guard its neighbour uses, and the eval's per-step duration still reports zero even though the real number is now sitting on the same record.

Real-world evidence

No evidence-bundle.md was produced for this PR, so nothing was exercised on my side — and pytest isn't installed on this runner, so I could not run the new suite either. My verdict rests on static review plus what the PR itself shows.

The description does carry genuine evidence for part of the change: a real step-1 stats record captured against a live cloud model (fireworks.deepseek-v4-flash-0731) through the non-streamed path, with the new fields populated and the pre-existing keys unchanged. That covers the non-streamed surface well. The streamed surface — which is where the new outgoing request field and the time-to-first-token stamping live — has unit tests only, and the test plan's live-Lemonade item is unchecked. That gap is the 🟡 above.

🔍 Technical details

🟡 Important

New outgoing request field verified only against a mock (src/gaia/llm/lemonade_client.py:2228)

stream_options={"include_usage": True} is a new field on the request to Lemonade's cloud-proxy path. tests/unit/agents/test_step_timing.py:434 asserts it is sent; nothing asserts Lemonade accepts it. There's no prior use of stream_options anywhere in the tree, so there's no established precedent that the proxy forwards it rather than rejecting the body. CLAUDE.md → "Mocks prove 'we called it,' not 'the call is valid.'" — #1655 is the canonical case (a recipe= field Lemonade 400s, green under a mocked HTTP layer).

Blast radius if the field is rejected: LemonadeClient.chat_completions(stream=True) raises for every cloud model, i.e. all streaming cloud chat, not just the new metric.

Ask: one live streamed turn against a cloud model, with the resulting stats record pasted into the description alongside the non-streamed one already there. (Local streaming is unaffected — the field is gated on cloud_model_provider(model) and test_local_stream_request_is_unchanged pins that.)

Worth noting the downstream shape change is handled correctly: the usage-bearing final chunk arrives with choices: [], and LemonadeProvider._handle_stream guards with if "choices" in chunk and chunk["choices"]. The only other stream consumer is the __main__ demo in lemonade_client.py:5254, which would IndexError on an empty choices — but it runs DEFAULT_MODEL_NAME, so the field is never sent there. Not worth changing.

🟢 Minor

1. Step timing call sits outside the "never displace a tool error" guard (src/gaia/agents/base/agent.py:4073)

The recorder.record_tool call immediately below is wrapped in try/except with an explicit # noqa: BLE001 - never displace a tool error. The new step_timer.record_tool on the line above is not, so a raise there would replace the tool's own result or exception. It's arithmetic and an append, so the odds are low — but the invariant is stated right there, and this is a finally block.

            if step_timer is not None:
                try:
                    step_timer.record_tool(tool_name or "<unnamed>", elapsed, waited)
                except Exception as e:  # noqa: BLE001 - never displace a tool error
                    logger.warning("could not record step tool timing: %s", e)

2. The eval's per-step duration is still hardcoded to zero (src/gaia/eval/performance.py:276)

# /stats has no per-step "duration"; tolerated as 0. — true before this PR, no longer true after it. step_seconds now lives on the same stats record extract_step_stats is already walking, so StepResult.duration_ms could carry the real number for free. Not a bug, but the comment now documents a limitation the PR removed.

                        duration_ms=int(
                            (content.get("step_seconds") or stats.get("duration", 0) or 0)
                            * 1000
                        ),

3. "always add up" is a rounding-strength claim (docs/spec/agent-base.mdx)

Each of the four values is independently round(..., 4), so the three parts can differ from step_seconds by up to ~1e-4. The PR's own agent-loop test uses abs=1e-3 rather than exact equality, which is the honest reading. Consider "add up to step_seconds" without "always", or round the remainder as step_seconds - llm - tools after rounding.

Strengths

  • Overhead as the remainder, not a separate measurement. llm + tools + overhead == step_seconds holds by construction rather than by everyone remembering to instrument every branch — which is exactly the property that makes the numbers usable for the "where did 2,916s go" question that motivated this.
  • The thread-ownership guard in StepTimer (step_timing.py:336) is the non-obvious correctness bit: an LLM call made inside a tool body is already inside that tool's seconds, and counting it as model time too would push the parts past the step. Good call, and it's tested (test_calls_from_another_thread_are_not_counted).
  • Fail-loud on a call with no wall time (record_llm_call raising ValueError, attach raising RuntimeError) rather than quietly writing a zero — matches the project's no-silent-fallbacks rule, and a silently-zero timing record is worse than no record.
  • The local-stream /stats carve-out is deliberate and pinned. _last_streamed keeps local streaming on /stats because that's what carries the time-to-first-token the UI renders; test_stream_stamps_ttft_on_a_tool_call_fragment_and_keeps_trailing_usage locks it in. Easy thing to have regressed by accident.
  • Stamping ttft on the first tool-call fragment rather than the first yielded chunk is the right choice — a tool-call-only stream yields nothing until the end, so the naive version would have reported the whole request as ttft.
  • Test coverage is genuinely good for a metrics change: fake clock for exact arithmetic, a real agent loop for the integration shape, and an assertion that the sink is detached when the turn ends so a later turn can't inherit it.

…real step durations

Three loose ends in the step-timing trace. A failure inside the new timing
call could replace a tool's own result or error, because it sat outside the
try/except its neighbour uses for exactly that reason. Eval scorecards kept
reporting every step as 0ms even though the agent now measures the real wall
time on the same record. And both the spec and the module docstring promised
the three parts "always" sum to the step total, which independent rounding of
each value cannot guarantee.
@github-actions github-actions Bot added the eval Evaluation framework changes label Sep 21, 2026
The trace's three parts were each rounded independently, so reading them
off a step could leave the reader to explain a ~1e-4 discrepancy against
the step total. Deriving overhead from the already-rounded parts makes
the reported numbers reconcile exactly, which is what the spec claimed
in the first place; the claim is restored rather than softened, and the
loop test asserts equality instead of a 1e-3 tolerance.

The run summary is built the same way and gets the same treatment, so a
reader is not told one rule that holds per step and another per run.
@kovtcharov-amd

Copy link
Copy Markdown
Collaborator

Two of the four findings were already fixed on this branch before I got here — the timer no longer displaces a tool's result, and eval scorecards now carry the real step duration instead of reporting every step as 0ms. I verified both against the current head rather than re-applying them.

The rounding finding I took further than the suggestion. The branch had softened the docs to "give or take rounding"; instead I made the claim true. Overhead is now the remainder of the already-rounded parts, so the three numbers a reader sees add back to the step total exactly, and the test asserts equality rather than a tolerance. Same treatment for the run summary, so there aren't two different rules to remember.

Still outstanding: the live streamed turn against a cloud model. No Lemonade server is reachable from here, so I did not run it and have not faked evidence for it.

What I can say without the live run is that the new field is standard and correctly placed. stream_options: {"include_usage": true} is a documented OpenAI Chat Completions field, and the one constraint attached to it — only valid when streaming — is satisfied here, because it is set on the streaming path only and gated to cloud models. The residual risk the live test would actually retire is narrower than "does this field work": it's whether the specific provider behind Lemonade's proxy accepts it. That varies in practice, so it is worth one real turn before merge.

🔍 Technical details

Already fixed on head (51180f921), verified not re-applied:

  • agent.py ~4073 — step_timer.record_tool is wrapped in try/except Exception with the # noqa: BLE001 - never displace a tool error comment, matching its recorder.record_tool neighbour.
  • performance.py ~276 — duration_ms now reads content.get("step_seconds") or stats.get("duration", 0) or 0, and the stale "/stats has no per-step duration; tolerated as 0" comment was rewritten to say step_seconds is the measurement and /stats the fallback. tests/unit/eval/test_performance_extractor.py covers it.

Rounding — made exact instead of reworded. In step_timing.py, _close() and finish() both did round(wall - llm - tools, 4) against independently rounded siblings. Now:

r_wall = round(wall, _ROUND)
r_llm = round(llm, _ROUND)
r_tools = round(tools, _ROUND)
...
"overhead_seconds": round(r_wall - r_llm - r_tools, _ROUND),

The module docstring and docs/spec/agent-base.mdx both state the parts add up exactly; the _close() inline comment that still said "always sum to the step" is no longer contradicting the softened docstring above it. tests/unit/agents/test_step_timing.py now asserts round(parts, 4) == rec["step_seconds"] and the equivalent for the summary, replacing pytest.approx(..., abs=1e-3).

On stream_options (lemonade_client.py ~2228). It is set inside _stream_chat_chunks, where "stream": True is hardcoded, and only when self.cloud_model_provider(model) is truthy — so the "only valid with stream: true" constraint cannot be violated from this call site, and local models never see the field. Documented by OpenAI as an optional stream_options object on POST /v1/chat/completions; effect is one extra SSE chunk before [DONE] with empty choices and a populated usage, which is exactly what this code wants.

The reason a live turn still matters: acceptance is not uniform across OpenAI-compatible backends. vLLM, llama.cpp's llama-server, LiteLLM and Ollama all support include_usage, but several providers 400 on the field outright as unrecognised (xAI/Grok: Argument not supported: stream_options; Fireworks: Extra inputs are not permitted, field: 'stream_options'; Azure OpenAI in some configurations). So the question the live test answers is whether the cloud provider on the other side of Lemonade's proxy is in the permissive set — a mock cannot answer that, and neither can I from here.

Tests run: tests/unit/agents/test_step_timing.py tests/unit/eval/test_performance_extractor.py — 50 passed. The wider tests/unit/eval run has 8 failures in test_release_scorecard.py, all UnicodeEncodeError: 'charmap' codec can't encode character 'Σ' — a Windows cp1252 console-encoding issue in files this branch does not touch.

# Conflicts:
#	src/gaia/agents/base/agent.py
@kovtcharov-amd

Copy link
Copy Markdown
Collaborator

Merged main in. The only conflict was two unrelated imports landing on the same line of the base agent — kept both.

🔍 Technical details

src/gaia/agents/base/agent.py: this branch adds StepTimer, main added resolve_project_root. Union, alphabetical. Both symbols verified in use in the merged file, so neither is a dangling import.

pytest tests/unit/agents/ → 1183 passed, 50 skipped. The 2 errors in test_large_tool_result_truncation.py are a pre-existing Windows-only limit (the parametrised test ID overflows PYTEST_CURRENT_TEST's 32767-char env-var cap) — reproduced identically on main's copy of that file, so unrelated to this branch.

# Conflicts:
#	src/gaia/llm/providers/lemonade.py
@kovtcharov-amd

Copy link
Copy Markdown
Collaborator

Merged main — the conflict was one import line, and nothing from either side was lost.

main picked up two commits that land in the same three files this branch touches: the context-overflow fix (#4095) and the Graph error-shape fix (#4109). Only the Lemonade provider's import block actually collided — both sides had added an import to it. Everything else auto-merged into regions the other side never visited.

I checked the result against both parents rather than trusting the auto-merge, and it comes out exact in both directions: the branch still adds precisely its 11 files at +1069/-49, and main's two commits still land byte-for-byte.

🔍 Technical details

The conflict: src/gaia/llm/providers/lemonade.py imports. The branch added import time (for perf_counter TTFT sampling); main added import re (for classify_lemonade_exception's message matching). Both are live in the merged file, so the resolution keeps both in isort order. No other hunk in the file conflicted — main's work sits in the module-level classifiers, the branch's in LemonadeProvider.chat/_handle_stream/_capture_usage.

Auto-merged files verified by hand, both directions:

  • src/gaia/agents/base/agent.py — main rewrote two exception helpers around line 4910 to call classify_lemonade_exception directly instead of importing gaia.ui._chat_helpers (which drops a except Exception: return None swallow). The branch's changes are StepTimer wiring in _attach_step_timer, _execute_tool_timed, _process_query_impl and the two result dicts. Disjoint regions; both present.
  • src/gaia/llm/lemonade_client.py — main adds the module-level active_profile_ctx_size(); the branch adds stream_options={"include_usage": True} and the usage passthrough inside the streaming method. Disjoint; both present.
  • src/gaia/cli.py, src/gaia/ui/_chat_helpers.py, src/gaia/agents/tools/_email/graph.py — main-only, taken verbatim.

_classify_chat_exception survives as an alias in _chat_helpers.py, so the branch has no dangling reference to it.

rerere is enabled locally but recorded a preimage rather than replaying a cached resolution, so nothing was resolved without review.

Tests — 396 passed across test_step_timing, test_performance_extractor, test_ctx_overflow_device_profile_2884, test_lemonade_error_classification, test_turn_metrics{,_wire}, test_claude_provider, test_lemonade_{cloud_models,ctx_ceiling,client_ctx_override,model_loading}, test_email_tools.

Four test_cli_prompt_* cases in test_ctx_overflow_device_profile_2884.py fail, but they fail identically on plain origin/main — the Windows ProactorEventLoop self-pipe opens a local socket and the unit-test network guard rejects it (#4112). Not introduced here.

Kalin Ovtcharov and others added 2 commits September 23, 2026 16:31
# Conflicts:
#	src/gaia/llm/providers/claude.py
main shipped this branch's one contested change while it was open. amd#3739 added
`stream_options: {"include_usage": true}` to every stream, not just cloud ones,
which retires the review's blocking ask: the field is already in production and
already exercised against a live local backend.

Two things the auto-merge got wrong without conflicting:

- **`_capture_usage` was defined twice in `LemonadeProvider`.** main's
  `(usage, timings)` version silently shadowed the branch's `(response)` one,
  so the branch's two call sites were passing the wrong shape. main's is the
  richer of the two (it adds `cached_tokens`), so it stays — but it read
  reasoning tokens only from `completion_tokens_details`, dropping the
  `output_tokens_details` fallback the branch's `_reasoning_tokens` helper
  handled. The helper is wired back in, so neither shape is lost.
- **`stream_options` was set twice** on the streaming request — main's
  unconditional one, then the branch's cloud-only one overwriting it with the
  same value. The dead cloud gate is gone.

`test_local_stream_request_is_unchanged` asserted local streams never carry the
field. amd#3739 deliberately decided otherwise, so the test now pins that contract
instead of the branch's older cloud-only one.

`step_timer.finish()` moved below main's new step-cap completion call, so that
call's seconds are counted rather than falling outside the run summary.
@kovtcharov-amd

Copy link
Copy Markdown
Collaborator

The blocking ask on this PR — "prove Lemonade accepts stream_options" — is now moot: main already ships it. #3739 added stream_options: {"include_usage": true} to every stream, not just cloud ones, so the field has been in production and exercised against a live local backend since before this review. The risk the live run was meant to retire has already been taken.

Merging main in was not clean, though, and neither problem raised a conflict marker. LemonadeProvider ended up with two _capture_usage definitions, main's silently shadowing this branch's — so both of this branch's call sites were passing the wrong argument shape. And stream_options was being set twice on the same request.

🔍 Technical details

The shadowed method. main's _capture_usage(usage, timings) is defined after this branch's _capture_usage(response), so Python kept main's and the branch's self._capture_usage(response) / self._capture_usage(chunk) calls were passing a whole response where a usage dict was expected. main's is the richer version (it adds cached_tokens), so it stays — but it read reasoning tokens only from completion_tokens_details, dropping the output_tokens_details fallback this branch's _reasoning_tokens helper handled. The helper is wired back into main's version, so neither shape is lost and it is no longer dead code. Both call sites moved to the (usage, timings) signature.

Found by AST-walking every class in the tree for duplicate method names, not by reading the diff.

The duplicate field. main sets stream_options unconditionally inside the request_params literal; this branch then re-set it, cloud-gated, to the same value. Same result, but the gate read as meaningful when it no longer was. Removed.

test_local_stream_request_is_unchanged asserted a local stream never carries the field. #3739 deliberately decided the opposite — "invisible locally, because llama.cpp answers the /stats poll" — so the test now pins that contract instead, renamed test_local_stream_also_asks_for_usage.

step_timer.finish() moved below main's new step-cap completion call (the extra tools-withheld call when a turn runs out of steps). It increments steps_taken, so it is a real step with real seconds; closing the timer above it dropped that time from the run summary.

Tests: test_step_timing.py + test_performance_extractor.py — 55 passed. Wider run across tests/unit/eval, tests/unit/agents, test_lemonade_error_classification, test_turn_metrics — 2392 passed, 57 skipped, one failure in tests/unit/eval/test_runner.py::test_rewritten_interpreter_is_logged, which this PR doesn't touch and which I reproduced on main's own copies of both runner.py and test_runner.py.

@kovtcharov-amd
kovtcharov-amd added this pull request to the merge queue Sep 26, 2026
Merged via the queue into amd:main with commit a02a233 Sep 26, 2026
50 of 54 checks passed
@kovtcharov-amd
kovtcharov-amd deleted the feat/step-timing-trace branch September 26, 2026 04:10
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

agents chat Chat SDK changes documentation Documentation changes eval Evaluation framework changes llm LLM backend changes performance Performance-critical changes tests Test changes

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants