diff --git a/management/internals/modules/agentnetwork/types/provider.go b/management/internals/modules/agentnetwork/types/provider.go index bf9401ceb..d604608f3 100644 --- a/management/internals/modules/agentnetwork/types/provider.go +++ b/management/internals/modules/agentnetwork/types/provider.go @@ -110,11 +110,7 @@ func (p *Provider) FromAPIRequest(req *api.AgentNetworkProviderRequest) { p.UpstreamURL = req.UpstreamUrl p.APIKeyProvided = req.ApiKey != nil if req.ApiKey != nil { - if strings.TrimSpace(*req.ApiKey) == "" { - p.APIKey = "" - } else { - p.APIKey = *req.ApiKey - } + p.APIKey = strings.TrimSpace(*req.ApiKey) } if req.ExtraValues != nil { // Replace the whole map (rather than merge) so unsetting a diff --git a/management/internals/modules/agentnetwork/types/provider_test.go b/management/internals/modules/agentnetwork/types/provider_test.go index 1678dc3b7..50b181d1a 100644 --- a/management/internals/modules/agentnetwork/types/provider_test.go +++ b/management/internals/modules/agentnetwork/types/provider_test.go @@ -100,6 +100,14 @@ func TestProvider_APIKeyPresenceAndResponse(t *testing.T) { assert.Equal(t, key, p.APIKey) assert.True(t, p.ToAPIResponse().HasApiKey) + paddedKey := " \tprotected-endpoint-token\n" + withPaddedKey := base() + withPaddedKey.ApiKey = &paddedKey + p.FromAPIRequest(withPaddedKey) + assert.True(t, p.APIKeyProvided) + assert.Equal(t, key, p.APIKey, "non-blank API keys must be normalized before storage") + assert.True(t, p.ToAPIResponse().HasApiKey, "the response must reflect the normalized stored key") + empty := "" clearKey := base() clearKey.ApiKey = &empty diff --git a/proxy/model_discovery_sync_test.go b/proxy/model_discovery_sync_test.go index 90b9aab35..4bc21754b 100644 --- a/proxy/model_discovery_sync_test.go +++ b/proxy/model_discovery_sync_test.go @@ -94,6 +94,17 @@ func (s *modelDiscoverySyncStream) Context() context.Context { return s.ctx } +func completeModelDiscoveryInitialSync(t *testing.T, stream *modelDiscoverySyncStream) { + t.Helper() + stream.recv <- &proto.SyncMappingsResponse{InitialSyncComplete: true} + select { + case sent := <-stream.sent: + require.NotNil(t, sent.GetAck()) + case <-time.After(time.Second): + t.Fatal("initial snapshot was not acknowledged") + } +} + func TestProxyCapabilitiesAdvertiseModelDiscovery(t *testing.T) { t.Parallel() @@ -178,10 +189,11 @@ func TestHandleSyncMappingsStreamRunsDiscoveryOutOfBand(t *testing.T) { stream.sendWait = 10 * time.Millisecond done := make(chan error, 1) - initialSyncDone := true + initialSyncDone := false go func() { done <- server.handleSyncMappingsStream(ctx, stream, &initialSyncDone, time.Time{}) }() + completeModelDiscoveryInitialSync(t, stream) stream.recv <- &proto.SyncMappingsResponse{ ModelDiscoveryRequest: &proto.ModelDiscoveryRequest{ @@ -242,10 +254,11 @@ func TestHandleSyncMappingsStreamBoundsConcurrentDiscovery(t *testing.T) { stream := newModelDiscoverySyncStream(ctx) done := make(chan error, 1) - initialSyncDone := true + initialSyncDone := false go func() { done <- server.handleSyncMappingsStream(ctx, stream, &initialSyncDone, time.Time{}) }() + completeModelDiscoveryInitialSync(t, stream) for i := range 5 { stream.recv <- &proto.SyncMappingsResponse{ @@ -298,7 +311,8 @@ func TestHandleSyncMappingsStreamRejectsMixedDiscoveryMessage(t *testing.T) { } stream := newModelDiscoverySyncStream(ctx) stream.recv <- &proto.SyncMappingsResponse{ - Mapping: []*proto.ProxyMapping{{Id: "mapping-1"}}, + Mapping: []*proto.ProxyMapping{{Id: "mapping-1"}}, + InitialSyncComplete: true, ModelDiscoveryRequest: &proto.ModelDiscoveryRequest{ RequestId: "request-1", }, @@ -307,5 +321,29 @@ func TestHandleSyncMappingsStreamRejectsMixedDiscoveryMessage(t *testing.T) { initialSyncDone := true err := server.handleSyncMappingsStream(ctx, stream, &initialSyncDone, time.Time{}) - require.EqualError(t, err, "model discovery message must not include mapping data") + require.EqualError(t, err, "model discovery message must not include mapping data or set initial_sync_complete") +} + +func TestHandleSyncMappingsStreamRejectsDiscoveryBeforeInitialSnapshot(t *testing.T) { + t.Parallel() + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + server := &Server{ + Logger: log.New(), + routerReady: closedChan(), + modelDiscoverer: &stubModelDiscoverer{}, + } + stream := newModelDiscoverySyncStream(ctx) + stream.recv <- &proto.SyncMappingsResponse{ + ModelDiscoveryRequest: &proto.ModelDiscoveryRequest{ + RequestId: "request-1", + }, + } + close(stream.recv) + + initialSyncDone := true + err := server.handleSyncMappingsStream(ctx, stream, &initialSyncDone, time.Time{}) + require.EqualError(t, err, "model discovery request received before initial sync completed") } diff --git a/proxy/server.go b/proxy/server.go index bf54a85b8..625f0a0fe 100644 --- a/proxy/server.go +++ b/proxy/server.go @@ -1372,6 +1372,7 @@ func (s *Server) handleSyncMappingsStream(ctx context.Context, stream proto.Prox } tracker := s.newSnapshotTracker(initialSyncDone, connectTime) + initialSnapshotComplete := false discoverer := s.modelDiscoverer if discoverer == nil { discoverer = modeldiscovery.New(s.Logger) @@ -1411,7 +1412,10 @@ func (s *Server) handleSyncMappingsStream(ctx context.Context, stream proto.Prox if discovery := msg.GetModelDiscoveryRequest(); discovery != nil { if len(msg.GetMapping()) != 0 || msg.GetInitialSyncComplete() { - return errors.New("model discovery message must not include mapping data") + return errors.New("model discovery message must not include mapping data or set initial_sync_complete") + } + if !initialSnapshotComplete { + return errors.New("model discovery request received before initial sync completed") } select { @@ -1455,7 +1459,11 @@ func (s *Server) handleSyncMappingsStream(ctx context.Context, stream proto.Prox return err } s.Logger.Debug("Processing mapping update completed") - tracker.recordBatch(ctx, s, msg.GetMapping(), msg.GetInitialSyncComplete(), batchStart) + syncComplete := msg.GetInitialSyncComplete() + tracker.recordBatch(ctx, s, msg.GetMapping(), syncComplete, batchStart) + if syncComplete { + initialSnapshotComplete = true + } if err := send(&proto.SyncMappingsRequest{ Msg: &proto.SyncMappingsRequest_Ack{Ack: &proto.SyncMappingsAck{}},