feat/add aws-bedrock-api support#247
Conversation
AmyTao
changed the title
![devin-ai-integration[bot]](https://avatars.githubusercontent.com/in/811515?s=60&v=4){kind=small}
This comment was marked as resolved.
This comment was marked as resolved.
Sorry, something went wrong.
Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com>
| def _bedrock_proxy_command( | ||
| *, | ||
| environment: str, | ||
| ) -> str: | ||
| if environment == "docker": | ||
| return "host.docker.internal" | ||
| return BEDROCK_PROXY_LOCAL_HOST |
There was a problem hiding this comment.
🔴 Bedrock proxy address 127.0.0.1 is unreachable from non-Docker (e.g. Daytona) sandboxes
The _bedrock_proxy_command function at src/benchflow/_provider_runtime.py:81-87 returns 127.0.0.1 for all non-Docker environments. However, the Bedrock proxy runs on the BenchFlow host machine (bound to 0.0.0.0), while Daytona sandboxes run on remote VMs. The agent inside the sandbox receives http://127.0.0.1:<port> as its API base URL and would connect to its own loopback—not the host running the proxy—resulting in a connection error. Daytona is widely used in this codebase (dozens of YAML configs and test fixtures default to environment: daytona).
Prompt for agents
The function _bedrock_proxy_command in src/benchflow/_provider_runtime.py returns 127.0.0.1 for all non-Docker environments. This is incorrect for Daytona, where the sandbox is a remote VM. The proxy runs on the BenchFlow host, but the agent gets 127.0.0.1 which points to the remote sandbox's own loopback.
Possible approaches:
1. For Daytona, the proxy would need to be accessible from the remote sandbox. This could involve tunneling or exposing the proxy on a public/reachable address. The function might need to raise an error for Daytona environments until proper support is added.
2. Alternatively, detect the Daytona case and raise a clear error like 'Bedrock proxy is not yet supported in Daytona environments. Use environment=docker or set up direct Bedrock access.'
3. For environments like Modal, the same issue applies.
The function is called from ensure_bedrock_proxy_runtime (line 118) with environment from the trial config.
Was this helpful? React with 👍 or 👎 to provide feedback.
| raise ValueError( | ||
| f"{required_key} required for model {model!r} but not set. " | ||
| "Pass it explicitly (for example via --agent-env/agent_env) " | ||
| "or define it in .env." | ||
| ) |
There was a problem hiding this comment.
🔴 Removal of subscription auth fallback breaks 4 existing tests in test_subscription_auth.py
The PR replaces the subscription-auth fallback (lines 341-345) with a plain raise ValueError(...), removing the check_subscription_auth() call that previously let users authenticated via CLI login (e.g. claude login, codex --login) bypass the API-key requirement when a model is specified. This breaks 4 existing tests in tests/test_subscription_auth.py that were NOT modified by this PR:
Failing tests and why they break
test_subscription_auth_detected(line 97-117): expectsresult["_BENCHFLOW_SUBSCRIPTION_AUTH"] == "1"but the new code raisesValueErrorinstead.test_no_auth_file_raises(line 119-136): expectsValueErrormatching"log in with the agent CLI"but the new message says"or define it in .env.".test_codex_subscription_auth(line 152-166): expects_BENCHFLOW_SUBSCRIPTION_AUTHmarker but getsValueError.test_gemini_subscription_auth(line 168-182): same as above.
This also violates the CLAUDE.md convention: "Don't rewrite passing tests to match new behavior." The corollary is that code changes should not silently break existing passing tests. These tests guard a user-facing feature (CLI login as API-key substitute) that the PR removes without updating the tests.
Prompt for agents
The PR removes the subscription auth fallback path in resolve_agent_env (around line 341) when a model is specified and the required API key is missing. The old code called check_subscription_auth() before raising ValueError, allowing users who had logged in via the agent CLI (e.g., claude login, codex --login) to bypass the API key requirement. Four tests in tests/test_subscription_auth.py depend on this behavior: test_subscription_auth_detected, test_no_auth_file_raises, test_codex_subscription_auth, and test_gemini_subscription_auth.
Either:
1. Restore the subscription auth fallback: before raising ValueError, check check_subscription_auth(agent, required_key) and set agent_env["_BENCHFLOW_SUBSCRIPTION_AUTH"] = "1" if it returns True. This preserves backward compatibility.
2. Or, if removing subscription auth for the model path is intentional, update the 4 broken tests in tests/test_subscription_auth.py to match the new behavior (the error message also changed from 'log in with the agent CLI' to 'or define it in .env').
The CLAUDE.md convention says 'Don't rewrite passing tests to match new behavior' — if option 2 is chosen, the PR description should explicitly justify this semantic change.
Was this helpful? React with 👍 or 👎 to provide feedback.
| agent_env = dict(agent_env or {}) | ||
| explicit_agent_env_keys = set(agent_env) | ||
| auto_inherit_env(agent_env) | ||
| auto_inherit_env(agent_env, source_env=load_dotenv_env()) |
There was a problem hiding this comment.
🔴 auto_inherit_env now reads from .env file instead of os.environ, silently dropping shell-exported API keys
The call to auto_inherit_env at src/benchflow/_agent_env.py:280 changed from auto_inherit_env(agent_env) (which defaulted to os.environ) to auto_inherit_env(agent_env, source_env=load_dotenv_env()). When load_dotenv_env() returns an empty dict (no .env file present), source_env is {} — not None — so the fallback to os.environ at line 101 never triggers. This means API keys set via export ANTHROPIC_API_KEY=sk-... in the user's shell are silently ignored. Users without a .env file who previously relied on shell-exported environment variables will get ValueError: ANTHROPIC_API_KEY required for model ... but not set errors with no indication that the inheritance source changed.
Prompt for agents
In resolve_agent_env at line 280, auto_inherit_env is called with source_env=load_dotenv_env(). When there is no .env file, load_dotenv_env() returns an empty dict {}, which is not None, so auto_inherit_env uses it as the source instead of falling back to os.environ. This breaks all users who rely on shell-exported API keys (export ANTHROPIC_API_KEY=sk-...).
The fix should either:
1. Merge .env and os.environ, with .env taking precedence: e.g. source = {**os.environ, **load_dotenv_env()} or chain them so .env wins but os.environ is still consulted.
2. Or fall back to os.environ when .env is empty/missing: pass None when load_dotenv_env returns empty, so the default os.environ path still works.
The intent appears to be that .env is a fallback source, not a replacement. The function's own docstring says 'Copy well-known agent env vars from a source mapping into agent_env' and its default behavior (source_env=None → os.environ) should be preserved when .env is absent.
Was this helpful? React with 👍 or 👎 to provide feedback.
| if self._env: | ||
| try: | ||
| await stop_provider_runtime(getattr(self, "_provider_runtime", None)) | ||
| self._provider_runtime = None | ||
| await self._env.stop(delete=True) | ||
| except Exception as e: | ||
| logger.warning(f"Cleanup failed: {e}") |
There was a problem hiding this comment.
🔴 Provider runtime stop failure prevents sandbox environment teardown
stop_provider_runtime() and self._env.stop(delete=True) are in the same try/except block. If stop_provider_runtime() raises an exception (e.g., the asyncio server is in an unexpected state), the except catches it and self._env.stop(delete=True) is never called. This leaves the sandbox environment (Docker container or Daytona workspace) running, creating a resource leak. Before this PR, self._env.stop() was the only operation in this try block and was always attempted (src/benchflow/trial.py:906-912).
| if self._env: | |
| try: | |
| await stop_provider_runtime(getattr(self, "_provider_runtime", None)) | |
| self._provider_runtime = None | |
| await self._env.stop(delete=True) | |
| except Exception as e: | |
| logger.warning(f"Cleanup failed: {e}") | |
| if self._env: | |
| try: | |
| await stop_provider_runtime(getattr(self, "_provider_runtime", None)) | |
| except Exception as e: | |
| logger.warning(f"Provider runtime cleanup failed: {e}") | |
| finally: | |
| self._provider_runtime = None | |
| try: | |
| await self._env.stop(delete=True) | |
| except Exception as e: | |
| logger.warning(f"Cleanup failed: {e}") |
Was this helpful? React with 👍 or 👎 to provide feedback.
| except Exception as exc: | ||
| logger.exception("Bedrock proxy request failed") | ||
| writer.write(_json_http_response(500, {"error": str(exc)})) | ||
| await writer.drain() |
There was a problem hiding this comment.
🟡 Corrupted HTTP response on mid-stream error in Bedrock proxy
When a streaming handler (_handle_messages or _handle_responses) successfully writes the HTTP/1.1 200 OK headers and some SSE chunks, then encounters an error (e.g., Bedrock stream failure or translation error), the exception propagates to _handle_connection's outer except block which writes a full HTTP 500 response (including status line and headers) into the already-started chunked response. This produces malformed output: the client sees the 200 headers, some valid chunks, then raw bytes of a 500 response embedded in the chunked stream. The client cannot parse this and will see a cryptic connection error rather than a meaningful error message.
Streaming error flow
- Client sends
POST /v1/messageswithstream: true - Handler writes
HTTP/1.1 200 OK+transfer-encoding: chunkedheaders (line 257-261) - Handler writes some SSE chunks successfully
- Exception occurs during
bedrock_stream_event_to_anthropic_sseorwriter.drain() - Exception propagates to
_handle_connectionexcept block (line 236-239) - Handler writes a full
HTTP/1.1 500 OKresponse into the already-started stream
Prompt for agents
The _handle_connection method in BedrockProxyServer has a top-level except handler that writes a full HTTP 500 response. This works for non-streaming requests (where no headers have been sent yet), but for streaming requests, headers have already been sent by the time an error occurs in the loop body of _handle_messages or _handle_responses.
Approach: Track whether response headers have been sent (e.g., with a boolean flag). In the except handler, only write the 500 response if headers haven't been sent yet. If headers were already sent (streaming case), either write an SSE error event into the chunked stream and terminate it properly (0\r\n\r\n), or just close the connection without writing anything more.
Alternatively, each streaming handler could have its own try/except around the streaming loop that sends a proper SSE error event and terminates the chunked stream before re-raising.
Was this helpful? React with 👍 or 👎 to provide feedback.
2 participants
Summary
Add initial AWS Bedrock API support to BenchFlow.
This PR introduces a host-side Bedrock proxy and wires
aws-bedrock/...models into the existing provider/agent flow so BenchFlow can route supported ACP agents through Bedrock instead of direct Anthropic endpoints.Current focus is the Sonnet 4.5 Claude path, with Bedrock provider plumbing and proxy infrastructure in place for broader expansion.
What Changed
Provider and Agent Plumbing
aws-bedrockas a registered providerBENCHFLOW_PROVIDER_*flow so agent launch remains consistent with other providersHost-Side Bedrock Proxy
invoke,invoke-with-response-stream, andcount-tokensTrial Integration
aws-bedrock/prefixConfig and Docs
.env.samplepyproject.tomlTests
Added and updated tests for:
Relevant test files:
Current Scope / Notes
ANTHROPIC_MODELin the Claude Bedrock launch path rather than relying on ClaudemodelOverrides.env.sampleincludesANTHROPIC_DEFAULT_SONNET_MODELas an example, but the current implementation path is driven by BenchFlow provider/runtime wiringExample Usage
Validation