API-key normalization and model-discovery snapshot validation

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

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

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")
 }

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{}},