feat(bedrock): add async support for aioboto3

[dvirski]

Fixes #3510

Bedrock’s OpenTelemetry instrumentation worked for boto3 (sync) but produced zero spans when users called Bedrock through aioboto3. This PR adds full async support — all four Bedrock Runtime methods (converse, converse_stream, invoke_model, invoke_model_with_response_stream)
now emit OTel-semconv-compliant spans, prompts, completions, token counts, and finish reasons.

The fix is symmetric to the existing sync path: same downstream helpers, same span attributes, same event-emitter pipeline, same semconv compliance — async just adds the missing async hook + four async-aware wrappers.

Tests: Added tests/traces/test_aioboto3.py — 12 mock-based tests covering all 4 async methods × 3 attribute modes (legacy / events+content / events+no-content), matching the sync test_nova.py coverage. Uses mocks instead of VCR cassettes because vcrpy's aiohttp stub can't replay aiobotocore responses (vcrpy#927 (kevin1024/vcrpy#927)).

Example: Added packages/sample-app/sample_app/async_bedrock_example.py — a research-assistant @workflow chaining all 4 async methods (converse → converse_stream → invoke_model → invoke_model_with_response_stream), so the resulting Traceloop trace shows a real-life async Bedrock flow.

Stress test (local-only, not committed): Wrote async_bedrock_stress_paths.py exercising 4 edge-case paths the doc example doesn't cover — invalid model (exception recording), concurrent asyncio.gather (no state leakage between simultaneous calls), early break from streaming
(span still ends with partial content, no ERROR status per OTel guidance), and SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY (instrumentation correctly skipped). This caught a real bug in converse_stream where the span would leak on early break, which we then fixed.

Code review (adversarial subagent): Spawned a code-reviewer subagent with explicit instructions to find bugs in the new async code. It flagged ~17 concerns; we fixed the ones introduced by this PR and deferred the ones that pre-exist in the sync path. Fixed:
AsyncStreamingWrapper span leak on early break (try/finally + _done guard), BufferedAsyncBody buffer-mutation bug (immutable buffer + cursor), aiobotocore body never closed (now released before replacement), InstrumentedClientContext.aenter not exception-safe (wrapped
patching block so the inner ctx is properly exited on failure), _handle_async_call passing a fake response dict to emit_choice_events (now passes the real response), and _handle_async_call swallowing exceptions at debug log level (now error).

Summary by CodeRabbit

• New Features

  • Async Bedrock instrumentation for aioboto3/aiobotocore: tracing for invoke/converse (streaming and non-streaming), buffered async responses, and robust async stream handling to ensure metrics/events and connection cleanup.
  • Added an async sample workflow demonstrating classification, streaming answer generation, summarization, and follow-ups.

• Bug Fixes

  • Improved extraction of token usage from nested response fields and headers.

• Tests

  • Comprehensive async tests for legacy/event-stream modes, edge cases, and resource handling.

• Chores

  • Test config updated to support asyncio and async AWS test packages.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: e5aad96b-2282-451a-945f-af65fd7711ab

📥 Commits

Reviewing files that changed from the base of the PR and between 7039b7f and 7b2b084.

⛔ Files ignored due to path filters (2)

• packages/opentelemetry-instrumentation-bedrock/uv.lock is excluded by !**/*.lock
• packages/sample-app/uv.lock is excluded by !**/*.lock

📒 Files selected for processing (9)

• packages/opentelemetry-instrumentation-bedrock/opentelemetry/instrumentation/bedrock/__init__.py
• packages/opentelemetry-instrumentation-bedrock/opentelemetry/instrumentation/bedrock/reusable_streaming_body.py
• packages/opentelemetry-instrumentation-bedrock/opentelemetry/instrumentation/bedrock/span_utils.py
• packages/opentelemetry-instrumentation-bedrock/opentelemetry/instrumentation/bedrock/streaming_wrapper.py
• packages/opentelemetry-instrumentation-bedrock/pyproject.toml
• packages/opentelemetry-instrumentation-bedrock/tests/traces/test_aioboto3.py
• packages/opentelemetry-instrumentation-bedrock/tests/traces/test_nova.py
• packages/sample-app/pyproject.toml
• packages/sample-app/sample_app/async_bedrock_example.py

🚧 Files skipped from review as they are similar to previous changes (9)

• packages/sample-app/pyproject.toml
• packages/opentelemetry-instrumentation-bedrock/tests/traces/test_nova.py
• packages/opentelemetry-instrumentation-bedrock/opentelemetry/instrumentation/bedrock/reusable_streaming_body.py
• packages/opentelemetry-instrumentation-bedrock/pyproject.toml
• packages/opentelemetry-instrumentation-bedrock/opentelemetry/instrumentation/bedrock/streaming_wrapper.py
• packages/opentelemetry-instrumentation-bedrock/opentelemetry/instrumentation/bedrock/span_utils.py
• packages/opentelemetry-instrumentation-bedrock/opentelemetry/instrumentation/bedrock/init.py
• packages/sample-app/sample_app/async_bedrock_example.py
• packages/opentelemetry-instrumentation-bedrock/tests/traces/test_aioboto3.py

---

📝 Walkthrough

Walkthrough

Adds async aiobotocore support for Bedrock instrumentation: intercepts async client creation, wraps async Bedrock APIs (invoke_model*, converse*), buffers async responses, parses event streams with guaranteed span closure, adds tests and a sample async workflow, and hardens uninstrumentation.

Changes

Bedrock Async Instrumentation

Layer / File(s)	Summary
Async response body abstractions packages/opentelemetry-instrumentation-bedrock/opentelemetry/instrumentation/bedrock/reusable_streaming_body.py, packages/opentelemetry-instrumentation-bedrock/opentelemetry/instrumentation/bedrock/streaming_wrapper.py	BufferedAsyncBody provides in-memory async buffering; AsyncStreamingWrapper wraps async event streams, decodes JSON chunks by type, accumulates content/deltas, and ensures a single stream_done_callback invocation via a _done guard and try/finally.
Async instrumentation entry point and client interception packages/opentelemetry-instrumentation-bedrock/opentelemetry/instrumentation/bedrock/__init__.py	WRAPPED_METHODS adds aiobotocore.session.AioSession.create_client (async=True); _wrap_async_factory and _InstrumentedClientContext open async clients and monkey-patch invoke_model* / converse* with async wrappers after client creation.
Async call and stream handlers packages/opentelemetry-instrumentation-bedrock/opentelemetry/instrumentation/bedrock/__init__.py	Adds _handle_async_call (async buffer + close underlying body, replace with BufferedAsyncBody, emit metrics/events), _handle_async_stream_call (wrap with AsyncStreamingWrapper, finalize spans/attrs on stream end), and _handle_async_converse_stream with _ConverseStreamCloser to parse events, accumulate partial deltas, flush on early termination, and guarantee single span.end().
Async uninstrumentation robustness packages/opentelemetry-instrumentation-bedrock/opentelemetry/instrumentation/bedrock/__init__.py	_uninstrument now ignores ImportError, ModuleNotFoundError, and AttributeError during unpatch to avoid hard failures.
Amazon usage token extraction update packages/opentelemetry-instrumentation-bedrock/opentelemetry/instrumentation/bedrock/span_utils.py	Derives prompt/completion token counts from nested response_body["usage"] or response_body["metadata"]["usage"], with HTTP header fallbacks for missing values.
Test infrastructure and async API tests packages/opentelemetry-instrumentation-bedrock/tests/traces/test_aioboto3.py	Adds async tests with aioboto3.Session fixture, mock async bodies/streams and _patch_make_api_call, and tests for converse, invoke_model, invoke_model_with_response_stream, and converse_stream across legacy and event-based modes validating spans, GEN-AI attributes, streamed content aggregation, and log emission/content-omission rules plus edge cases and resource lifecycle assertions.
Dependencies and pytest configuration packages/opentelemetry-instrumentation-bedrock/pyproject.toml	Test dependency group adds aioboto3, aiobotocore, and pytest-asyncio; pytest configured with asyncio_mode = "auto".
Sample async research assistant workflow packages/sample-app/pyproject.toml, packages/sample-app/sample_app/async_bedrock_example.py	Sample app adds aioboto3 and the instrumentation dependency and provides research_assistant workflow demonstrating classify/generate/summarize/suggest steps using converse, converse_stream, invoke_model, and invoke_model_with_response_stream.

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~60 minutes

Suggested reviewers

• doronkopit5
• galzilber
• nina-kollman
• max-deygin-traceloop

Poem

🐰 A curious rabbit hops through async streams,
Buffering bytes and parsing JSON dreams,
Spans close once, events stitch the thread,
Tests mock the flow so no network's bled,
Now aioboto3 and Bedrock hum observability beams.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 20.69% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	Title clearly summarizes the main change: adding async support for aioboto3 in the Bedrock instrumentation, which is the primary objective of the PR.
Linked Issues check	✅ Passed	The PR comprehensively implements all objectives from issue #3510: adds async Bedrock instrumentation for aioboto3, emits bedrock span for converse, converse_stream, invoke_model, and invoke_model_with_response_stream, restores standard LLM/gen_ai attributes including token counts, and achieves parity with sync boto3 instrumentation.
Out of Scope Changes check	✅ Passed	All changes are within scope: async Bedrock instrumentation via aioboto3, supporting classes (BufferedAsyncBody, AsyncStreamingWrapper), comprehensive tests, documentation/examples, and a targeted span_utils.py fix for Nova token extraction discovered during review.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

📝 Generate docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch dr/feat(bedrock)-add-async-support-via-aioboto3-for-all-4-runtime-methods

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[CLAassistant]

All committers have signed the CLA.

[coderabbitai]

Actionable comments posted: 2

🧹 Nitpick comments (1)

packages/opentelemetry-instrumentation-bedrock/tests/traces/test_aioboto3.py (1)

393-417: ⚡ Quick win

Add explicit early-break stream-consumption regression tests.

PR notes mention fixing span leaks on early break; add tests that break after first chunk/event and still assert span closure + attributes/log integrity.

Also applies to: 498-514

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@packages/opentelemetry-instrumentation-bedrock/tests/traces/test_aioboto3.py`
around lines 393 - 417, Add explicit regression tests that consume only the
first event from the response stream and then break early to ensure spans are
closed and logs/attributes are exported; modify the existing
test_aioboto3_invoke_stream_legacy (and the analogous test at lines 498-514) to
iterate async for event in response["body"], read the first chunk, break
immediately, then ensure you exit the context (await ctx.__aexit__(None, None,
None)) and assert via span_exporter and log_exporter that the span is finished
and contains the expected attributes/log entries; keep the same client creation
helper (_patch_make_api_call and _fake_invoke_stream_response) and use
span_exporter, log_exporter assertions to validate no span leaks on early-break
consumption.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In
`@packages/opentelemetry-instrumentation-bedrock/tests/traces/test_aioboto3.py`:
- Around line 274-279: The test currently iterates the logs returned by
log_exporter.get_finished_logs() and asserts on any existing log bodies but can
false-pass if no logs are emitted; update each such block (the one using logs =
log_exporter.get_finished_logs(), the subsequent for log in logs loop and body =
dict(log.log_record.body)) to first assert a minimum expected log count (e.g.,
assert len(logs) > 0 or assert len([l for l in logs if l.log_record.body]) > 0)
before validating bodies, mirroring the with-content variants; apply the same
change to the other occurrences flagged (the similar blocks around the other
ranges) so the test fails when no logs are emitted.
- Around line 93-97: The mock _AsyncReadableBody.read currently returns
self._raw without consuming it; change read so that when amt is None it returns
the full buffer and sets self._raw to b'' (i.e., consume the bytes), and when
amt is provided slice and remove bytes from self._raw as the existing branch
does; update the mock implementation in the test helper where
_AsyncReadableBody.read is defined to match botocore/aiobotocore StreamingBody
semantics. Also strengthen the log assertions referenced (tests that inspect
no-content logs at the indicated assertions) by asserting that the logs list is
non-empty before checking their contents so the test fails if no logs were
emitted.

---

Nitpick comments:
In
`@packages/opentelemetry-instrumentation-bedrock/tests/traces/test_aioboto3.py`:
- Around line 393-417: Add explicit regression tests that consume only the first
event from the response stream and then break early to ensure spans are closed
and logs/attributes are exported; modify the existing
test_aioboto3_invoke_stream_legacy (and the analogous test at lines 498-514) to
iterate async for event in response["body"], read the first chunk, break
immediately, then ensure you exit the context (await ctx.__aexit__(None, None,
None)) and assert via span_exporter and log_exporter that the span is finished
and contains the expected attributes/log entries; keep the same client creation
helper (_patch_make_api_call and _fake_invoke_stream_response) and use
span_exporter, log_exporter assertions to validate no span leaks on early-break
consumption.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: 92e63e19-30c4-4168-9d8f-b67795b6c91a

📥 Commits

Reviewing files that changed from the base of the PR and between 6d3e696 and 3bb131b.

⛔ Files ignored due to path filters (2)

• packages/opentelemetry-instrumentation-bedrock/uv.lock is excluded by !**/*.lock
• packages/sample-app/uv.lock is excluded by !**/*.lock

📒 Files selected for processing (7)

• packages/opentelemetry-instrumentation-bedrock/opentelemetry/instrumentation/bedrock/__init__.py
• packages/opentelemetry-instrumentation-bedrock/opentelemetry/instrumentation/bedrock/reusable_streaming_body.py
• packages/opentelemetry-instrumentation-bedrock/opentelemetry/instrumentation/bedrock/streaming_wrapper.py
• packages/opentelemetry-instrumentation-bedrock/pyproject.toml
• packages/opentelemetry-instrumentation-bedrock/tests/traces/test_aioboto3.py
• packages/sample-app/pyproject.toml
• packages/sample-app/sample_app/async_bedrock_example.py

[doronkopit5]

1. Async exception logging is asymmetric with sync (__init__.py:683-689)

except Exception as e:
    logger.error(
        "OpenLLMetry failed to trace in _handle_async_call, error: %s", e, exc_info=True
    )

The sync path uses @dont_throw (utils.py:11-33), which logs at logger.debug. Promoting async-only to logger.error with exc_info=True means async users will see stack traces in their logs that sync users don't — an instrumentation bug becomes user-visible noise. Either bump both, or align async to logger.debug for parity.

2. Only 1 of 12 tests asserts token attributes (tests/traces/test_aioboto3.py)

test_aioboto3_converse_legacy (lines 232-233) is the only test asserting GEN_AI_USAGE_INPUT_TOKENS / GEN_AI_USAGE_OUTPUT_TOKENS. The other 11 (including test_aioboto3_invoke_model_legacy, test_aioboto3_invoke_stream_legacy, test_aioboto3_converse_stream_legacy) only assert span name + PROVIDER_NAME / OPERATION_NAME / REQUEST_MODEL. Issue #3510 explicitly calls out token attributes as missing — a regression in token extraction for invoke_model or streaming would not be caught. Add token-attribute assertions to all 12 tests (mock values are deterministic: input=11, output=16).

3. No committed tests for the claimed edge-case fixes

The PR description enumerates fixes for:

• Early break out of async for (AsyncStreamingWrapper + converse_stream span leak)
• Exception during iteration
• Concurrent asyncio.gather (no state leakage)
• SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY short-circuit

None of these are in test_aioboto3.py — the PR description notes they were tested in a local-only stress harness. Those are the very bugs the author identified and fixed; without committed tests they can silently regress. At minimum:

• test_aioboto3_converse_stream_early_break — break after first chunk, assert len(spans) == 1 with partial content.
• test_aioboto3_invoke_stream_exception — raise from inside the async for body, assert the span still ends.
• test_aioboto3_suppress_language_model_instrumentation — set the context key, assert len(spans) == 0.

7. Mocks bypass the real aiobotocore body lifecycle (test_aioboto3.py:87-98)

_AsyncReadableBody does not implement close(). The instrumentation (__init__.py:634-641) tries getattr(original_body, \"close\", None) — with the mock it returns None and the close-branch is silently skipped, so the test never exercises the close path. The mock also doesn't simulate aiobotocore's "body is consumed after first read" semantic, so the BufferedAsyncBody swap-in isn't truly proven necessary by these tests.

Add:

• close() (and aclose() if applicable) on _AsyncReadableBody that records invocation, and assert it gets called.
• A test that calls await response[\"body\"].read() twice and asserts the second call returns b\"\" — proves the buffered-body cursor semantic.

[dvirski]

@doronkopit5

Topics addressed:

#1 — Async exception logging asymmetric with sync

• bedrock/init.py:683-690 — replaced logger.error(..., exc_info=True) with logger.debug(..., traceback.format_exc()) to match @dont_throw.
• Added import traceback.

#2 — Only 1 of 12 tests asserts token attributes

• Added deterministic GEN_AI_USAGE_INPUT_TOKENS == 11 / GEN_AI_USAGE_OUTPUT_TOKENS == 16 assertions to all 12 existing async tests in test_aioboto3.py.
• Tightened the existing > 0 assertion in test_aioboto3_converse_legacy to == 11 / == 16.

#3 — No committed tests for the claimed edge-case fixes

• test_aioboto3_invoke_stream_early_break — drives iterator with anext() + aclose(), asserts span ends. Locks the try/finally + _done guard in AsyncStreamingWrapper.
• test_aioboto3_invoke_stream_exception_during_iteration — uses stream.athrow(RuntimeError("boom")) to inject mid-iteration exception, asserts span still ends.
• test_aioboto3_suppress_language_model_instrumentation — sets SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY, asserts len(spans) == 0.

#7 — Mocks bypass real aiobotocore body lifecycle

• _AsyncReadableBody — added async def close() with close_calls counter (exercises the await-if-coroutine branch in the instrumentation).
• test_aioboto3_invoke_model_closes_underlying_body — captures original body via factory wrapper, asserts close_calls == 1.
• test_aioboto3_invoke_model_body_is_rereadable — reads response["body"] twice, asserts both calls return identical bytes (proves BufferedAsyncBody swap-in is necessary).

---

Bonus — preexisting sync bug found while addressing #2

• bedrock/span_utils.py:_set_amazon_span_attributes — fixed Nova token extraction:
  • Non-streaming path was reading top-level inputTokens (always missing → 0 in prod). Now reads usage.inputTokens (where Nova actually puts it).
  • Streaming-accumulated path now also reads metadata.usage.{input,output}Tokens.
  • HTTP header (x-amzn-bedrock-input-token-count) used as fallback for both.
• Verified live against AWS Bedrock Nova-lite via probe script — confirmed response shape matches the cassette and the fix.

Bonus — sync regression locks for the above fix

• test_nova.py::test_nova_completion — added GEN_AI_USAGE_INPUT_TOKENS == 30 / OUTPUT_TOKENS == 51 (values from VCR cassette).
• test_nova.py::test_nova_invoke_stream — added > 0 token assertions to guard the streaming-branch fix.

Verification: 380/380 bedrock package tests pass, ruff clean.

[coderabbitai]

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In
`@packages/opentelemetry-instrumentation-bedrock/opentelemetry/instrumentation/bedrock/span_utils.py`:
- Around line 647-655: The token-extraction logic in span_utils.py incorrectly
uses truthy `or` fallbacks so empty dicts or zero counts get treated as missing;
update the logic in the block that computes `usage`, `total_prompt_tokens`, and
`total_completion_tokens` to check presence/None explicitly: set `usage =
response_body.get("usage") if "usage" in response_body else
response_body.get("metadata", {}).get("usage", {})` (or equivalent presence
check), and when extracting `inputTokens`/`outputTokens` use `if value is not
None` or `get(..., None)` checks before falling back to header values so that 0
is accepted and empty dicts from `usage` don't fall through. Ensure
`total_prompt_tokens` and `total_completion_tokens` use these explicit None
checks on `usage.get("inputTokens")` and `usage.get("outputTokens")` and only
then fallback to `headers.get(...)`.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: 14843c1a-b781-4425-8230-4a7be43aa42b

📥 Commits

Reviewing files that changed from the base of the PR and between 62af08f and 7039b7f.

⛔ Files ignored due to path filters (2)

• packages/opentelemetry-instrumentation-bedrock/uv.lock is excluded by !**/*.lock
• packages/sample-app/uv.lock is excluded by !**/*.lock

📒 Files selected for processing (9)

• packages/opentelemetry-instrumentation-bedrock/opentelemetry/instrumentation/bedrock/__init__.py
• packages/opentelemetry-instrumentation-bedrock/opentelemetry/instrumentation/bedrock/reusable_streaming_body.py
• packages/opentelemetry-instrumentation-bedrock/opentelemetry/instrumentation/bedrock/span_utils.py
• packages/opentelemetry-instrumentation-bedrock/opentelemetry/instrumentation/bedrock/streaming_wrapper.py
• packages/opentelemetry-instrumentation-bedrock/pyproject.toml
• packages/opentelemetry-instrumentation-bedrock/tests/traces/test_aioboto3.py
• packages/opentelemetry-instrumentation-bedrock/tests/traces/test_nova.py
• packages/sample-app/pyproject.toml
• packages/sample-app/sample_app/async_bedrock_example.py

✅ Files skipped from review due to trivial changes (1)

• packages/opentelemetry-instrumentation-bedrock/pyproject.toml

🚧 Files skipped from review as they are similar to previous changes (4)

• packages/opentelemetry-instrumentation-bedrock/opentelemetry/instrumentation/bedrock/reusable_streaming_body.py
• packages/sample-app/pyproject.toml
• packages/sample-app/sample_app/async_bedrock_example.py
• packages/opentelemetry-instrumentation-bedrock/opentelemetry/instrumentation/bedrock/init.py