mirror of
https://github.com/coder/coder.git
synced 2026-06-03 04:58:23 +00:00
Danny Kopping
5d40bac79f
### TL;DR Introduces an in-process `TransportFactory` for aibridge so that chatd (coder-agent LLM traffic) can route requests through the aibridged handler without crossing the HTTP route or requiring a license entitlement check. ### What changed? - Added a new `coderd/aibridge` package with a `TransportFactory` interface and a `Source` type for tagging the call site on request contexts. `SourceAgents` is defined as the constant for coder-agent traffic. - Implemented `NewTransportFactory` in `coderd/aibridged/transport.go`, which returns an `http.RoundTripper` that dispatches requests to the aibridged handler in-process. The response body is streamed through an `io.Pipe` so SSE/NDJSON/chunked responses propagate token-by-token. Handler panics are recovered and surfaced as 500 responses, and context cancellation closes the pipe with the appropriate error. - `RegisterInMemoryAIBridgedHTTPHandler` now also constructs a `TransportFactory` from the registered handler and stores it on `API.AIBridgeTransportFactory` (an `atomic.Pointer`), making it available to chatd without going through the license-gated HTTP route. - Added `API.AIBridgeTransportFactory` as a public `atomic.Pointer[aibridge.TransportFactory]` field on `coderd.API`. ### How to test? - `coderd/aibridged/transport_test.go` covers: transport creation, nil-handler errors, source attachment to context, header/status passthrough, streaming (SSE-style chunked writes visible before handler completion), context cancellation closing the body with an error, concurrent requests, handler panics producing 500s, and handlers that return without writing. - `coderd/aibridge_test.go` verifies that `AIBridgeTransportFactory` starts as nil on AGPL coderd, can be stored and loaded atomically, and that the stored factory correctly dispatches requests through the stub handler. ### Why make this change? Chatd needs to send LLM requests through aibridge in-process rather than via the external HTTP route, which is license-gated. The `TransportFactory` abstraction provides a clean seam: the entitlement check remains on the HTTP route for external callers, while in-process coder-agent traffic bypasses it through the factory. The `Source` type allows downstream handlers and logs to attribute traffic without gating behavior on the caller identity.
102 lines
3.0 KiB
Go
102 lines
3.0 KiB
Go
package coderd_test
|
|
|
|
import (
|
|
"io"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"testing"
|
|
|
|
"github.com/google/uuid"
|
|
"github.com/stretchr/testify/require"
|
|
|
|
"github.com/coder/coder/v2/coderd/aibridge"
|
|
"github.com/coder/coder/v2/coderd/coderdtest"
|
|
"github.com/coder/coder/v2/testutil"
|
|
)
|
|
|
|
// stubTransportFactory wires a deterministic handler through the
|
|
// AIBridgeTransportFactory hook so the AGPL side of the in-memory pipe can be
|
|
// exercised without pulling coderd/aibridged in.
|
|
type stubTransportFactory struct {
|
|
handler http.Handler
|
|
calls chan callRecord
|
|
}
|
|
|
|
type callRecord struct {
|
|
providerID uuid.UUID
|
|
source aibridge.Source
|
|
}
|
|
|
|
func (f *stubTransportFactory) TransportFor(providerID uuid.UUID, source aibridge.Source) (http.RoundTripper, error) {
|
|
f.calls <- callRecord{providerID: providerID, source: source}
|
|
return &handlerRoundTripper{handler: f.handler}, nil
|
|
}
|
|
|
|
// handlerRoundTripper is a minimal http.RoundTripper for the AGPL test. It
|
|
// does not stream; coderd/aibridged.transport_test.go already covers
|
|
// streaming semantics.
|
|
type handlerRoundTripper struct{ handler http.Handler }
|
|
|
|
func (h *handlerRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
|
|
rec := httptest.NewRecorder()
|
|
h.handler.ServeHTTP(rec, req)
|
|
resp := rec.Result()
|
|
resp.Request = req
|
|
return resp, nil
|
|
}
|
|
|
|
// Verify that a factory stored on coderd.API.AIBridgeTransportFactory is
|
|
// observable through the normal API lifecycle: cli/server.go registers it
|
|
// when the bridge daemon starts (see RegisterInMemoryAIBridgedHTTPHandler).
|
|
func TestAIBridgeTransportFactory_Registration(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
_, _, api := coderdtest.NewWithAPI(t, nil)
|
|
|
|
require.Nil(t, api.AIBridgeTransportFactory.Load(),
|
|
"AGPL coderd must not pre-populate the factory")
|
|
|
|
stub := &stubTransportFactory{
|
|
handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.Header().Set("Content-Type", "application/json")
|
|
w.WriteHeader(http.StatusOK)
|
|
_, _ = w.Write([]byte(`{"bridged":true}`))
|
|
}),
|
|
calls: make(chan callRecord, 4),
|
|
}
|
|
|
|
var asInterface aibridge.TransportFactory = stub
|
|
api.AIBridgeTransportFactory.Store(&asInterface)
|
|
|
|
loaded := api.AIBridgeTransportFactory.Load()
|
|
require.NotNil(t, loaded)
|
|
|
|
providerID := uuid.New()
|
|
rt, err := (*loaded).TransportFor(providerID, aibridge.SourceAgents)
|
|
require.NoError(t, err)
|
|
require.NotNil(t, rt)
|
|
|
|
select {
|
|
case got := <-stub.calls:
|
|
require.Equal(t, providerID, got.providerID)
|
|
require.Equal(t, aibridge.SourceAgents, got.source)
|
|
default:
|
|
t.Fatal("factory was not invoked")
|
|
}
|
|
|
|
ctx := testutil.Context(t, testutil.WaitShort)
|
|
req, err := http.NewRequestWithContext(ctx, http.MethodPost, "http://aibridge/v1/messages", nil)
|
|
require.NoError(t, err)
|
|
|
|
client := &http.Client{Transport: rt}
|
|
resp, err := client.Do(req)
|
|
require.NoError(t, err)
|
|
defer resp.Body.Close()
|
|
|
|
body, err := io.ReadAll(resp.Body)
|
|
require.NoError(t, err)
|
|
require.Equal(t, http.StatusOK, resp.StatusCode)
|
|
require.Equal(t, `{"bridged":true}`, string(body))
|
|
require.Equal(t, "application/json", resp.Header.Get("Content-Type"))
|
|
}
|