feat: forward client request headers to upstream providers in bridge routes#214
Conversation
This stack of pull requests is managed by Graphite. Learn more about stacking. |
ssncferreira
force-pushed
the
ssncf/feat-forward-client-headers-mdlw
branch
3 times, most recently
from
a2da2f0 to
252f03a
Compare
ssncferreira
force-pushed
the
ssncf/feat-forward-client-headers-mdlw
branch
from
252f03a to
2d43f46
Compare
| // TODO(ssncferreira): inject actor headers directly in the client-header | ||
| // middleware instead of using SDK options. |
There was a problem hiding this comment.
Will handle this in a follow-up in order to not increase the scope of this PR.
pawbana
left a comment
pawbana
left a comment
There was a problem hiding this comment.
Overall LGTM although it is a bit hard to review new and changes to the tests 😅
provider/openai_test.go
Outdated
|
|
||
| var receivedHeaders http.Header | ||
|
|
||
| mockUpstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { |
There was a problem hiding this comment.
Would internal/integrationtest/mockupstream.go be useful here. Probably new method that creates upstreamResponse from string/bytes would be need.
I see now it may not have been a good idea to make it private 😞
There was a problem hiding this comment.
Good point, the mock upstream is similar in the provider tests. However, I think it makes sense to separate integration and provider unit tests, but there are definitely some overlaps where it would be cleaner to reuse some mocks like the upstream. Maybe this is something that should be moved to a shared package, wdyt?
provider/openai_test.go
Outdated
| func TestOpenAI_CreateInterceptor(t *testing.T) { | ||
| t.Parallel() | ||
|
|
||
| t.Run("ChatCompletions_ClientHeaders", func(t *testing.T) { |
There was a problem hiding this comment.
Sorry for begin pedantic about this but would it be possible to reformat this into table tests? 😅
IIUC both tests are doing basically the same but with different request path?
Maybe I'm just used to it but I feel table test format is much easier to understand and keep track of test cases. It usually also shortens the code.
From my experience AIs for some reason don't do write tests in table format first but when you ask for it directly in follow up prompt AI realizes it is possible or sometimes provides a reason why it is not worth it.
There was a problem hiding this comment.
Completely agree, this was to make it consistent with the other provider tests that use the same format. Addressed in c9298cc
And I can create a follow-up to update the other provider tests also to use a table driven format 👍
There was a problem hiding this comment.
I don't think we need yet another package for mocking upstream. I think in both integration and unit tests (I think those are a bit larger then unit test) it serves the same role.
ssncferreira
force-pushed
the
ssncf/feat-forward-client-headers-mdlw
branch
from
d005352 to
f1b1fc0
Compare
3 participants
Description
AI Bridge acts as a reverse proxy between clients and upstream AI providers. However, bridge routes were dropping client request headers as the interceptors construct new HTTP requests via provider SDKs, which set their own headers and discard the originals.
This PR changes bridge routes to forward client request headers to upstream providers, so the upstream sees the same headers the client sent, except headers related to auth, transport, or managed by aibridge.
Changes
intercept/client_headers.gowithSanitizeClientHeadersandBuildUpstreamHeaders. Client headers are sanitized by stripping hop-by-hop (RFC 2616), transport (Host,Accept-Encoding,Content-Length), and auth (Authorization,X-Api-Key) headers.X-AI-Bridge-Actor-*): injected by aibridge as per-request SDK options to identify the requesting user to the upstream.This is an intermediate step toward removing the SDK from the HTTP transport path. By forwarding client headers directly, aibridge is no longer dependent on SDK-managed headers, making future SDK removal simpler.
Follow-up
X-Forwarded-For,X-Forwarded-Host,X-Forwarded-Proto), while bridge routes do not. Since AI Bridge behaves as a reverse proxy, this handling should be consistent across all requests.ExtraHeadersmechanism for Anthropic and Copilot: now redundant since all client headers are forwarded.Tests
Tested manually with:
Closes: #192