[management] Close the review's nitpicks on live model discovery

Five smaller points from the same review as the four already fixed.

**The no-redirect policy had no test.** Every Fetch test injects an
HTTPClient, which bypasses httpClient() and therefore the policy entirely,
so nothing asserted that a 302 is refused — and the policy is a security
control: a redirect moves the request to a host checkPublicHost never
resolved. TestRedirectsAreNotFollowed drives the real constructor against an
httptest server that redirects to the cloud metadata address, and asserts
exactly one request leaves the client.

**An injected client silently lost that policy.** Production is safe today
only because NewManager passes a nil HTTPClient; any future non-test
injection would have dropped the guarantee with no signal. An injected
client that states no policy now inherits ours.

Implemented by copying the client rather than assigning into it. Writing
c.HTTPClient.CheckRedirect from httpClient() would mutate a struct shared by
every request goroutine for the process's lifetime — a data race, and the
exact hazard the same review's next point warns about. The copy shares the
Transport, which is safe for concurrent use by design.

**Two comments described things that were not true.** The manager's doc
claimed reading a stored credential "is permission-checked as one"; there is
no per-record Read check, just the single Create check, which covers it
because Create is stronger and the lookup is account-scoped. Said that
instead. The modelDiscovery field now records that it is shared across
requests and must stay read-only after construction.

**The handler test asserted neither the label contract nor the upstream.**
The response omits label entirely when a vendor supplies none, and the
dashboard falls back to the id on absence — an empty string would render a
blank row. The fixture had no label-less model to prove it with (the review
described one, but both existing entries carry labels), so this adds one.
The upstream assertion matters because Bedrock's region is read back out of
it. Also fixes the "blank catalog provide" subtest name.
This commit is contained in:
mlsmaycon
2026-08-22 13:20:27 +00:00
committed by GitHub
parent ca59f9f7a8
commit 67b1373dfe
4 changed files with 107 additions and 15 deletions

View File

@@ -60,6 +60,9 @@ func TestDiscoverModelsReturnsTheVendorList(t *testing.T) {
stub := &discoveryManagerStub{models: []modeldiscovery.Model{
{ID: "eu.anthropic.claude-haiku-4-5-20251001-v1:0", Label: "EU Claude Haiku 4.5", PricingKnown: true},
{ID: "global.cohere.embed-v4:0", Label: "Global Cohere Embed v4"},
// A vendor that supplies no display name at all. Bedrock does for
// every profile, but the OpenAI listing carries none.
{ID: "gpt-4o-mini", PricingKnown: true},
}}
rec := postDiscovery(t, stub, `{
@@ -71,7 +74,7 @@ func TestDiscoverModelsReturnsTheVendorList(t *testing.T) {
var out api.AgentNetworkModelDiscoveryResponse
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &out))
require.Len(t, out.Models, 2)
require.Len(t, out.Models, 3)
assert.Equal(t, "eu.anthropic.claude-haiku-4-5-20251001-v1:0", out.Models[0].Id)
assert.True(t, out.Models[0].PricingKnown)
@@ -79,8 +82,19 @@ func TestDiscoverModelsReturnsTheVendorList(t *testing.T) {
// from a priced one: registering it silently would meter at zero.
assert.False(t, out.Models[1].PricingKnown)
require.NotNil(t, out.Models[0].Label, "the vendor supplied a display name")
assert.Equal(t, "EU Claude Haiku 4.5", *out.Models[0].Label)
// A vendor that supplies no name must omit the key rather than send an
// empty string: the dashboard falls back to the id on absence, and would
// render a blank row for "".
assert.Nil(t, out.Models[2].Label, "an absent label must not serialize")
assert.NotContains(t, rec.Body.String(), `"label":""`)
assert.Equal(t, "bedrock_api", stub.gotReq.CatalogID)
assert.Equal(t, "aws-bearer", stub.gotReq.APIKey)
// The upstream is what the region is read back out of for Bedrock, so
// losing it here would break discovery for every regional provider.
assert.Equal(t, "https://bedrock-runtime.eu-central-1.amazonaws.com", stub.gotReq.UpstreamURL)
assert.Empty(t, stub.gotRecordID)
}
@@ -151,9 +165,9 @@ func TestDiscoverModelsReportsCallerInputAsBadRequest(t *testing.T) {
func TestDiscoverModelsRejectsMalformedRequests(t *testing.T) {
for name, body := range map[string]string{
"not json": `{`,
"no catalog provider": `{"api_key":"sk"}`,
"blank catalog provide": `{"catalog_provider_id":" ","api_key":"sk"}`,
"not json": `{`,
"no catalog provider": `{"api_key":"sk"}`,
"blank catalog provider": `{"catalog_provider_id":" ","api_key":"sk"}`,
} {
t.Run(name, func(t *testing.T) {
stub := &discoveryManagerStub{}

View File

@@ -128,6 +128,10 @@ type managerImpl struct {
// modelDiscovery queries vendors for the models a credential can reach.
// A field rather than a package call so tests can drive it without
// reaching the network.
//
// One instance serves every request for the process's lifetime, so its
// fields must stay read-only after construction: lazy initialisation
// inside Fetch or httpClient would race across request goroutines.
modelDiscovery *modeldiscovery.Client
// reconcileCache holds the last set of synthesised proxy mappings
@@ -182,12 +186,13 @@ func (m *managerImpl) GetProvider(ctx context.Context, accountID, userID, provid
//
// recordID, when set, names an existing provider whose stored credential and
// upstream are used instead of the ones in req — so the dashboard can refresh
// the list without ever holding the key. Reading a stored credential is a read
// of that provider, and is permission-checked as one.
// the list without ever holding the key.
//
// Gated on Create rather than Read: this spends the operator's credential
// against a third party, which is not something a read-only role should be
// able to make the server do.
// able to make the server do. That one check also covers reading the stored
// record — Create is strictly stronger than Read here, and the lookup is
// scoped to accountID, so another account's record is never reachable.
func (m *managerImpl) DiscoverProviderModels(ctx context.Context, accountID, userID string, req modeldiscovery.Request, recordID string) ([]modeldiscovery.Model, error) {
if err := m.requirePermission(ctx, accountID, userID, modules.AgentNetworkProviders, operations.Create); err != nil {
return nil, err

View File

@@ -357,22 +357,37 @@ func decorate(entry catalog.Provider, ids []listedModel) []Model {
return out
}
// refuseRedirect is the redirect policy every discovery request runs under. A
// redirect is a way to move the request to a host checkPublicHost never saw,
// so none are followed.
func refuseRedirect(*http.Request, []*http.Request) error {
return http.ErrUseLastResponse
}
func (c *Client) httpClient() *http.Client {
if c.HTTPClient != nil {
return c.HTTPClient
if c.HTTPClient.CheckRedirect != nil {
return c.HTTPClient
}
// An injected client that states no policy still gets ours: the
// no-redirect guarantee should not depend on the caller remembering it.
//
// Copied rather than assigned into: one Client is shared by every
// request for the process's lifetime, so writing to its fields here
// would race across request goroutines. The copy shares the Transport,
// which is safe for concurrent use by design.
clone := *c.HTTPClient
clone.CheckRedirect = refuseRedirect
return &clone
}
transport := guardedTransport
if c.AllowPrivateHosts {
transport = http.DefaultTransport
}
return &http.Client{
Timeout: fetchTimeout,
Transport: transport,
// A redirect is a way to move the request to a host the guard above
// never checked, so none are followed.
CheckRedirect: func(*http.Request, []*http.Request) error {
return http.ErrUseLastResponse
},
Timeout: fetchTimeout,
Transport: transport,
CheckRedirect: refuseRedirect,
}
}

View File

@@ -4,6 +4,7 @@ import (
"context"
"io"
"net/http"
"net/http/httptest"
"net/netip"
"strings"
"testing"
@@ -255,6 +256,63 @@ func TestHostGuardResolvesAndRejectsLocalhost(t *testing.T) {
assert.Contains(t, err.Error(), "non-public")
}
// TestRedirectsAreNotFollowed covers a gap the other tests leave open: they all
// inject an HTTPClient, which bypasses httpClient() and therefore the redirect
// policy entirely. The policy is a security control — a 302 moves the request
// to a host checkPublicHost never resolved — so it needs a test that goes
// through the constructor the manager actually uses.
func TestRedirectsAreNotFollowed(t *testing.T) {
var hits int
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
hits++
http.Redirect(w, r, "http://169.254.169.254/latest/meta-data/", http.StatusFound)
}))
t.Cleanup(srv.Close)
for name, cl := range map[string]*Client{
// The production shape: no injected client at all.
"default client": {AllowPrivateHosts: true},
// An injected client that states no policy must inherit ours rather
// than silently chasing the redirect.
"injected client with no policy": {
AllowPrivateHosts: true,
HTTPClient: &http.Client{},
},
} {
t.Run(name, func(t *testing.T) {
hits = 0
req, err := http.NewRequest(http.MethodGet, srv.URL, nil)
require.NoError(t, err)
resp, err := cl.httpClient().Do(req)
require.NoError(t, err)
t.Cleanup(func() { _ = resp.Body.Close() })
assert.Equal(t, http.StatusFound, resp.StatusCode,
"the redirect must be surfaced, not followed to an unchecked host")
assert.Equal(t, 1, hits, "exactly one request must leave the client")
})
}
}
// TestInjectedClientKeepsItsOwnRedirectPolicy pins that the default above is a
// default, not an override, and that supplying it does not mutate the caller's
// client — one Client is shared across every request, so a write here would
// race.
func TestInjectedClientKeepsItsOwnRedirectPolicy(t *testing.T) {
own := func(*http.Request, []*http.Request) error { return nil }
injected := &http.Client{CheckRedirect: own}
cl := &Client{HTTPClient: injected}
assert.Same(t, injected, cl.httpClient(),
"a client that states a policy must be handed back untouched")
bare := &http.Client{}
cl = &Client{HTTPClient: bare}
require.NotSame(t, bare, cl.httpClient(), "the policy must be applied to a copy")
assert.Nil(t, bare.CheckRedirect, "the caller's client must not be written to")
}
// TestDialGuardRejectsRebindingToANonPublicAddress covers the window between
// the two DNS lookups. checkPublicHost resolves the host, then the transport
// resolves it again to dial; a name whose owner answers the first with a public