Skip to content

fix(agentcore): Convert SSE stream iterator to async for proper streaming support - #16293

Merged
4 commits merged into
BerriAI:mainfrom
aproorg:fix-agentcore
Nov 12, 2025
Merged

4 commits merged into
BerriAI:mainfrom
aproorg:fix-agentcore

Conversation

@busla

@busla busla commented Nov 5, 2025 •

Copy link
Copy Markdown
Contributor

Description

Fixes async/sync iterator mismatch in AgentCore SSE streaming that was causing the error:

me-8Uck0u4W11) Exception litellm.APIConnectionError:
litellm  | Traceback (most recent call last):
litellm  |   File "/usr/lib/python3.13/site-packages/litellm/main.py", line 595, in acompletion
litellm  |     response = await init_response
litellm  |                ^^^^^^^^^^^^^^^^^^^
litellm  |   File "/usr/lib/python3.13/site-packages/litellm/llms/custom_httpx/llm_http_handler.py", line 620, in acompletion_stream_function
litellm  |     return await provider_config.get_async_custom_stream_wrapper(
litellm  |            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
litellm  |     ...<10 lines>...
litellm  |     )
litellm  |     ^
litellm  |   File "/usr/lib/python3.13/site-packages/litellm/llms/base_llm/chat/transformation.py", line 406, in get_async_custom_stream_wrapper
litellm  |     raise NotImplementedError
litellm  | NotImplementedError

Problem

The AgentCoreSSEStreamIterator was implemented as a synchronous iterator (__iter__ and __next__) but was being used with async HTTP response streams from httpx. When httpx provides an async response in streaming mode, it requires async iteration.

Solution

Converted AgentCoreSSEStreamIterator from sync to async iterator:

  1. Changed __iter__() → __aiter__()
  2. Changed __next__() → async __anext__()
  3. Changed self.response.iter_lines() → self.response.aiter_lines()
  4. Changed for line in self.line_iterator: → async for line in self.line_iterator:
  5. Changed all StopIteration → StopAsyncIteration

The CustomStreamWrapper already has proper async handling (checks is_async_iterable() and uses async for), so it works correctly with the async iterator.

Additional Enhancement: Runtime User ID Support

Added support for tracking user identity across AgentCore requests:

  • Added _get_runtime_user_id() method to extract runtimeUserId from optional params
  • Updated transform_request() to set X-Amzn-Bedrock-AgentCore-Runtime-User-Id header when runtimeUserId is provided
  • Works alongside existing runtimeSessionId → X-Amzn-Bedrock-AgentCore-Runtime-Session-Id mapping

This allows AgentCore runtimes to track individual users within sessions for better personalization and analytics.

Changes

Files Modified:

  • litellm/llms/bedrock/chat/agentcore/sse_iterator.py - Fixed async streaming
  • litellm/llms/bedrock/chat/agentcore/transformation.py - Added user ID header support

Testing

  • Verified the async iterator pattern matches httpx async response requirements
  • Confirmed CustomStreamWrapper properly handles async iterables at lines 1811-1812
  • Confirmed both runtimeSessionId and runtimeUserId optional params are correctly mapped to headers

Impact

  • Fixes AgentCore streaming responses for async contexts
  • Adds user ID tracking capability for AgentCore
  • No breaking changes - CustomStreamWrapper already supports both sync and async iterators
  • Allows LibreChat and other clients to properly receive streaming responses from AgentCore

LibreChat example

# ...
  custom:
    - name: "litellm"
      apiKey: "${LITELLM_MASTER_KEY}"
      baseURL: "http://litellm:4000/v1"
      titleModel: "claude"
      titleConvo: true
      summarize: false
      forcePrompt: false
      models:
        default: ["agentcore"]
        fetch: true
      modelDisplayLabel: "LiteLLM (Federated)"
      addParams:
        runtimeSessionId: "{{LIBRECHAT_BODY_CONVERSATIONID}}"
        runtimeUserId: "{{LIBRECHAT_BODY_USERID}}"

LiteLLM example

model_list:
  - model_name: claude
    litellm_params:
      model: bedrock/eu.anthropic.claude-sonnet-4-5-20250929-v1:0
      aws_region_name: eu-central-1
  - model_name: agentcore
    litellm_params:
      model: bedrock/agentcore/arn:aws:bedrock-agentcore:eu-central-1:<masked>:runtime/<masked>
      aws_access_key_id: os.environ/AWS_ACCESS_KEY_ID
      aws_secret_access_key: os.environ/AWS_SECRET_ACCESS_KEY
      aws_region_name: eu-central-1
      stream: True
litellm_settings:
  drop_params: true
  set_verbose: true
  cache: false # Disable cache for testing
general_settings:
  forward_client_headers_to_llm_api: true
  master_key: "os.environ/LITELLM_MASTER_KEY"
  json_logs: true

@vercel

vercel Bot commented Nov 5, 2025

Copy link
Copy Markdown

@busla is attempting to deploy a commit to the CLERKIEAI Team on Vercel.

A member of the Team first needs to authorize it.

@ishaan-jaff ishaan-jaff 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.

please can you show the before and after of manual testing results

This worked for me when testing

@busla

busla commented Nov 5, 2025

Copy link
Copy Markdown
Contributor Author

please can you show the before and after of manual testing results

This worked for me when testing

updated the PR with the error before the change

@ishaan-jaff ishaan-jaff 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.

does this break existing functionality, can you run test_bedrock_agentcore.py with your changes + agents and validate it works

I think your changes would break existing sync completion

@busla

busla commented Nov 6, 2025

Copy link
Copy Markdown
Contributor Author

I tested this against my runtime which contains Strands agents with async stream, so the released LiteLLM feature breaks runtimes that stream responses.

Would be great if you could test my changes against your JSON response agentcore runtime.

AgentCore Runtime supports two response formats:

  1. text/event-stream - Server-Sent Events (SSE) streaming format
  2. application/json - Standard JSON response format

From the AWS documentation example:

if "text/event-stream" in response.get("contentType", ""):
    # Handle streaming response - uses iter_lines()
    for line in response["response"].iter_lines():
        # Process SSE chunks

elif response.get("contentType") == "application/json":
    # Handle standard JSON response
    content = json.loads(response_data)

Streaming Path (Affected by Changes)

File: transformation.py:405-420

  • Uses get_streaming_response() → returns AgentCoreSSEStreamIterator
  • This async iterator is used by CustomStreamWrapper
  • This is where my changes apply

Non-Streaming Path (NOT Affected)

File: transformation.py:557-623

  • Uses transform_response() → calls _get_parsed_response()
  • _get_parsed_response() line 319-346:
    • Checks content-type header
    • If application/json: calls _parse_json_response(response_json)
    • Otherwise: calls _parse_sse_stream(response_text)

Important:

def _parse_sse_stream(self, response_text: str) -> AgentCoreParsedResponse:
    # ...

This method takes a string parameter response_text which comes from:

response_text = raw_response.text  # Synchronous read of full response body

The async iterator (AgentCoreSSEStreamIterator) is ONLY used when:

  1. stream=True is passed in the request, AND
  2. LiteLLM calls get_streaming_response()

For synchronous responses:

  1. LiteLLM reads the complete response using raw_response.text (synchronous)
  2. Then parses the complete text string using _parse_sse_stream(response_text: str)
  3. Never touches the async iterator

Synchronous responses are NOT affected because they:

  • Read the full response body synchronously via raw_response.text
  • Parse the complete text string without using the async iterator
  • Work identically whether the response is JSON or SSE format

@busla

busla commented Nov 6, 2025

Copy link
Copy Markdown
Contributor Author

Updated the PR to map runtimeUserId -> X-Amzn-Bedrock-AgentCore-Runtime-User-Id for runtime oauth support.

@ghost

ghost commented Nov 11, 2025

Copy link
Copy Markdown

@busla i believe the ask was to just confirm if your pr works. Can you just confirm this PR works for you on streaming in Strands? (Which i assume uses bedrock runtime)

@busla

busla commented Nov 11, 2025

Copy link
Copy Markdown
Contributor Author

yes, it works on my agentcore agent that streams responses.

@ghost
ghost merged commit 63445f2 into BerriAI:main Nov 12, 2025
3 of 7 checks passed
fzowl pushed a commit to fzowl/litellm that referenced this pull request Jun 24, 2026
…ming support (BerriAI#16293)

* fix(agentcore): support async agentcore runtime streaming

* revert: CLAUDE.md

* revert: .gitignore

* fix: map runtimeUserId to X-Amzn-Bedrock-AgentCore-Runtime-User-Id header for runtime oauth support
This pull request was closed.
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