Skip to content

AI-472 Add replay-safe Google ADK metrics sample - #355

Open
xumaple wants to merge 3 commits into
mainfrom
maplexu/AI-472-adk-metrics
Open

AI-472 Add replay-safe Google ADK metrics sample#355
xumaple wants to merge 3 commits into
mainfrom
maplexu/AI-472-adk-metrics

Conversation

@xumaple

@xumaple xumaple commented Aug 25, 2026

Copy link
Copy Markdown
Contributor

This PR adds a Google ADK sample that exports ADK OpenTelemetry metrics through ReplaySafeMeterProvider. A deterministic local model keeps the sample API-key-free, and the Workflow includes a post-metrics replay boundary so duplicate observations are detectable.

Note: The sample pins the SDK revision that introduces ReplaySafeMeterProvider until that API is available in a release.

@xumaple
xumaple requested review from a team as code owners August 25, 2026 19:30

@DABH DABH left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Reviewed at 0b92811. The core thesis holds and I verified it rather than taking the README's word for it — I probed _skip_recording directly and got 6 recordings live, 6 suppressed under Replayer, counts_after == counts_before — and the happy path works end to end against a real dev server. Switching this commit to Replayer and relative equality was the right response to the macOS failures; it's one indentation away from green.

One item that isn't tied to a line in this diff: metrics is missing from the suite scenario table. google_adk_agents/README.md:39-46 lists all six siblings and nothing in the repo links to google_adk_agents/metrics/README.md, so the sample is unreachable from any index.

| [metrics](./metrics/README.md) | Google ADK's OpenTelemetry metrics exported to a local Prometheus endpoint, with `ReplaySafeMeterProvider` keeping replay from double-counting observations. |

That row is only accurate once the nested pyproject.toml goes away — the suite's generic run block at :48-53 doesn't work for this scenario today. See the pyproject.toml:8 comment.

id=f"google-adk-agents-metrics-{uuid.uuid4()}",
task_queue=task_queue,
)
result = await handle.result()

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is why CI is red on all eight jobs. The async with await WorkflowEnvironment.start_time_skipping() opened at :23 exits at :41, which shuts the ephemeral server down, and handle.fetch_history() at :50 then RPCs a dead address — RPCError: tcp connect error after nine gRPC retries against 127.0.0.1:43563 in job 98639038931. Deterministic on every platform, not a macOS timing thing this time.

Fetch the history while the server is still up. I applied exactly this and the test goes green in 0.85s:

Suggested change
result = await handle.result()
result = await handle.result()
history = await handle.fetch_history()

Upstream does the same — sdk-python/tests/contrib/google_adk_agents/test_replay_metrics.py:296 fetches history as the last statement inside async with Worker(...). Note this disappears on its own if you take the test-placement suggestion at :10, since the session-scoped env fixture outlives the test.

assert counts_before_replay["gen_ai.client.token.usage"] > 0

await Replayer(workflows=[MetricsWorkflow], plugins=[plugin]).replay_workflow(
await handle.fetch_history()

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Companion to the :40 suggestion — use the history fetched inside the environment block.

Suggested change
await handle.fetch_history()
history

ADK_METER_SCOPE = "gcp.vertex.agent"


async def test_metrics_are_not_inflated_by_replay() -> None:

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is the only real test in the repo outside tests/find . \( -name 'test_*.py' -o -name '*_test.py' \) | grep -v '^./tests/' returns two paths and the other (polling/test_service.py) has no test functions, it just matches pytest's glob. All six ADK siblings live at tests/google_adk_agents/<scenario>_test.py with the signature async def test_basic(client: Client, monkeypatch: pytest.MonkeyPatch).

The cost isn't just tidiness. tests/conftest.py doesn't apply out here, so this can't use the session-scoped env/client fixtures at :40-57, it ignores --workflow-environment (:18-23) so both CI passes run an identical path and each start a redundant server, and it hand-rolls the environment whose premature teardown is the failure at :50. It also collects first among the ADK tests purely because it sorts before tests/ — in the a058b20 macOS log it starts at 19:35:18.26 against 19:35:30.83+ for the siblings, so it alone paid ADK's cold in-workflow import anthropic (google/adk/flows/llm_flows/contents.py:62), which is what blew the 2s deadlock budget on 3/3 macOS runners.

Moving it to tests/google_adk_agents/metrics_test.py taking client: Client, dropping the inline WorkflowEnvironment, and registering LocalMetricsModel through monkeypatch the way tests/openai_agents/_mock_model.py:59 does would fix the blocker and the macOS flakiness together.

Comment on lines +23 to +27
async with await WorkflowEnvironment.start_time_skipping() as environment:
plugin = GoogleAdkPlugin()
config = environment.client.config()
config["plugins"] = [*config["plugins"], plugin]
client = type(environment.client)(**config)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

start_time_skipping takes plugins directly (temporalio/testing/_workflow.py:241), so the client rebuild isn't needed. Also type(environment.client)(**config) where all six siblings just write Client(**config).

Suggested change
async with await WorkflowEnvironment.start_time_skipping() as environment:
plugin = GoogleAdkPlugin()
config = environment.client.config()
config["plugins"] = [*config["plugins"], plugin]
client = type(environment.client)(**config)
async with await WorkflowEnvironment.start_time_skipping(
plugins=[GoogleAdkPlugin()]
) as environment:
client = environment.client
task_queue = f"google-adk-agents-metrics-{uuid.uuid4()}"

If you take this, plugin is no longer in scope for the Replayer(...) call at :49 — construct it once above the async with and pass the same instance to both.

continue
for metric in scope_metrics.metrics:
counts[metric.name] = sum(
getattr(point, "count", 1) for point in metric.data.data_points

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

getattr(point, "count", 1) reads as 1 per attribute set for a NumberDataPoint, which has value and no count — so a Counter that had doubled would still compare equal. Dead branch today since every ADK instrument here is a histogram, but it's silently defeating the assertion this test exists for. Better to fail loudly on anything unexpected:

Suggested change
getattr(point, "count", 1) for point in metric.data.data_points
point.count for point in metric.data.data_points

for part in event.content.parts:
if part.text:
final_text = part.text
await workflow.sleep(timedelta(milliseconds=1))

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Undocumented scaffolding — no comment in a file with no comments, no README mention, no sibling precedent; the rationale ("a post-metrics replay boundary so duplicate observations are detectable") lives only in the PR description. At head it changes nothing measurable either way (6/6), and it only ever mattered under the max_cached_workflows=0 that's being removed.

Suggested change
await workflow.sleep(timedelta(milliseconds=1))
return final_text

timedelta on :1 becomes unused with it. Note root ruff runs only --select I, so an orphan import won't be flagged.

Comment on lines +13 to +21
```shell
uv run --project google_adk_agents/metrics python -m google_adk_agents.metrics.run_worker
```

Then run the Workflow:

```shell
uv run --project google_adk_agents/metrics python -m google_adk_agents.metrics.run_metrics_workflow
```

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Companion to the pyproject.toml:8 comment — once the nested project is gone, --project has nothing to point at, and the suite README already documents the house form at :48-53.

Suggested change
```shell
uv run --project google_adk_agents/metrics python -m google_adk_agents.metrics.run_worker
```
Then run the Workflow:
```shell
uv run --project google_adk_agents/metrics python -m google_adk_agents.metrics.run_metrics_workflow
```
```shell
uv run python -m google_adk_agents.metrics.run_worker
```
Then run the Workflow:
```shell
uv run python -m google_adk_agents.metrics.run_metrics_workflow
```

The starter prints `Replay-safe metrics are ready.` Inspect the metrics exposed by the worker:

```shell
curl http://127.0.0.1:9464/metrics | rg 'gen_ai'

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

rg is the only ripgrep invocation in any markdown file in the repo (grep -rln '| rg ' --include='*.md' . gives one hit) and it isn't a documented prerequisite.

Suggested change
curl http://127.0.0.1:9464/metrics | rg 'gen_ai'
curl -s http://127.0.0.1:9464/metrics | grep gen_ai

curl http://127.0.0.1:9464/metrics | rg 'gen_ai'
```

The output includes `gen_ai.invoke_agent` metrics, `gen_ai.client.operation.duration`, and `gen_ai.client.token.usage` from instrumentation scope `gcp.vertex.agent`. The worker sets `max_cached_workflows=0`, forcing replay between Workflow tasks. `ReplaySafeMeterProvider` drops observations made while replaying, so replay does not multiply the recorded counts.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Two corrections here.

The no-inflation sentence is true — I verified it — but it's the README's only statement about recording semantics and it omits the half the SDK spells out at _meter_provider.py:221-223: recordings are first-execution-only, and a Workflow task retry re-executes live and can record again. The guard is in_workflow() and is_replaying_history_events(), and a failed WFT appends WorkflowTaskFailed without advancing the replay boundary — which is exactly how a058b20 produced assert 2 == 1 on macOS after a TMPRL1101 deadlock. For a sample whose entire subject is metric accuracy, a reader will otherwise take these as exactly-once.

Second, the instrumentation scope isn't observable in the output — the Prometheus exporter emits no otel_scope_* labels, and the real line is the munged gen_ai_invoke_agent_duration_seconds_count{gen_ai_agent_name="metrics_agent"} 1.0. And the max_cached_workflows sentence goes away with the run_worker.py:28 change.

Suggested change
The output includes `gen_ai.invoke_agent` metrics, `gen_ai.client.operation.duration`, and `gen_ai.client.token.usage` from instrumentation scope `gcp.vertex.agent`. The worker sets `max_cached_workflows=0`, forcing replay between Workflow tasks. `ReplaySafeMeterProvider` drops observations made while replaying, so replay does not multiply the recorded counts.
The output includes `gen_ai.invoke_agent` metrics, `gen_ai.client.operation.duration`, and `gen_ai.client.token.usage`, exported with dots munged to underscores — for example `gen_ai_invoke_agent_duration_seconds_count{gen_ai_agent_name="metrics_agent"} 1.0`. `ReplaySafeMeterProvider` drops observations made while replaying, so replay does not multiply the recorded counts.
Recordings are first-execution-only rather than exactly-once: replay is suppressed, but a Workflow task retry re-executes live and can record again. Treat these counters as at-least-once — aggregate with rates and percentiles rather than relying on exact counts.

Comment on lines +16 to +18
async def generate_content_async(
self, llm_request: LlmRequest, stream: bool = False
) -> AsyncGenerator[LlmResponse, None]:

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

stream=True silently yields a single non-partial response rather than streaming, so a reader who points the streaming scenario at this model gets quietly wrong behavior instead of an error.

Suggested change
async def generate_content_async(
self, llm_request: LlmRequest, stream: bool = False
) -> AsyncGenerator[LlmResponse, None]:
async def generate_content_async(
self, llm_request: LlmRequest, stream: bool = False
) -> AsyncGenerator[LlmResponse, None]:
if stream:
raise NotImplementedError(
"LocalMetricsModel does not implement streaming responses."
)

Minor, separately: metrics_workflow.py:17 retypes "local-metrics-model" instead of importing MODEL_NAME from :8 here.

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.

2 participants