plumb context to start

Signed-off-by: Jeremy Drouillard <jeremy@stacklok.com>

Author: jerm-dro

cmd/vmcp/app/commands.go

@@ -311,6 +311,12 @@ func runServe(cmd *cobra.Command, _ []string) error {
 		if err != nil {
 			return fmt.Errorf("failed to create telemetry provider: %w", err)
 		}
+		defer func() {
+			err := telemetryProvider.Shutdown(ctx)
+			if err != nil {
+				logger.Errorf("failed to shutdown telemetry provider: %v", err)
+			}
+		}()
 	}
 
 	serverCfg := &vmcpserver.Config{
@@ -333,7 +339,7 @@ func runServe(cmd *cobra.Command, _ []string) error {
 	}
 
 	// Create server with discovery manager, backends, and workflow definitions
-	srv, err := vmcpserver.New(serverCfg, rtr, backendClient, discoveryMgr, backends, workflowDefs)
+	srv, err := vmcpserver.New(ctx, serverCfg, rtr, backendClient, discoveryMgr, backends, workflowDefs)
 	if err != nil {
 		return fmt.Errorf("failed to create Virtual MCP Server: %w", err)
 	}

pkg/vmcp/server/health_test.go

@@ -59,7 +59,10 @@ func createTestServer(t *testing.T) *server.Server {
 	// Mock Stop to be called during server shutdown
 	mockDiscoveryMgr.EXPECT().Stop().AnyTimes()
 
-	srv, err := server.New(&server.Config{
+	// Create context for server
+	ctx, cancel := context.WithCancel(t.Context())
+
+	srv, err := server.New(ctx, &server.Config{
 		Name:    "test-vmcp",
 		Version: "1.0.0",
 		Host:    "127.0.0.1",
@@ -68,9 +71,7 @@ func createTestServer(t *testing.T) *server.Server {
 	require.NoError(t, err)
 
 	// Start server in background
-	ctx, cancel := context.WithCancel(t.Context())
 	t.Cleanup(cancel)
-
 	errCh := make(chan error, 1)
 	go func() {
 		if err := srv.Start(ctx); err != nil {
@@ -175,7 +176,7 @@ func TestServer_SessionManager(t *testing.T) {
 		mockDiscoveryMgr := discoveryMocks.NewMockManager(ctrl)
 		rt := router.NewDefaultRouter()
 
-		srv, err := server.New(&server.Config{
+		srv, err := server.New(context.Background(), &server.Config{
 			Name:       "test-vmcp",
 			Version:    "1.0.0",
 			SessionTTL: 10 * time.Minute,
@@ -198,7 +199,7 @@ func TestServer_SessionManager(t *testing.T) {
 		rt := router.NewDefaultRouter()
 
 		customTTL := 15 * time.Minute
-		srv, err := server.New(&server.Config{
+		srv, err := server.New(context.Background(), &server.Config{
 			Name:       "test-vmcp",
 			Version:    "1.0.0",
 			SessionTTL: customTTL,

pkg/vmcp/server/integration_test.go

@@ -192,7 +192,7 @@ func TestIntegration_AggregatorToRouterToServer(t *testing.T) {
 	// Mock Stop to be called during server shutdown
 	mockDiscoveryMgr.EXPECT().Stop().Times(1)
 
-	srv, err := server.New(&server.Config{
+	srv, err := server.New(ctx, &server.Config{
 		Name:    "test-vmcp",
 		Version: "1.0.0",
 		Host:    "127.0.0.1",
@@ -322,7 +322,7 @@ func TestIntegration_HTTPRequestFlowWithRoutingTable(t *testing.T) {
 	}
 
 	// Create and start server
-	srv, err := server.New(&server.Config{
+	srv, err := server.New(ctx, &server.Config{
 		Name:           "test-vmcp",
 		Version:        "1.0.0",
 		Host:           "127.0.0.1",

pkg/vmcp/server/server.go

@@ -152,6 +152,7 @@ type Server struct {
 //
 //nolint:gocyclo // Complexity from hook logic is acceptable
 func New(
+	ctx context.Context,
 	cfg *Config,
 	rt router.Router,
 	backendClient vmcp.BackendClient,
@@ -205,7 +206,7 @@ func New(
 	if cfg.TelemetryProvider != nil {
 		var err error
 		backendClient, err = monitorBackends(
-			context.Background(),
+			ctx,
 			cfg.TelemetryProvider.MeterProvider(),
 			cfg.TelemetryProvider.TracerProvider(),
 			backends,

pkg/vmcp/server/server_test.go

@@ -76,7 +76,7 @@ func TestNew(t *testing.T) {
 			mockBackendClient := mocks.NewMockBackendClient(ctrl)
 			mockDiscoveryMgr := discoveryMocks.NewMockManager(ctrl)
 
-			s, err := server.New(tt.config, mockRouter, mockBackendClient, mockDiscoveryMgr, []vmcp.Backend{}, nil)
+			s, err := server.New(context.Background(), tt.config, mockRouter, mockBackendClient, mockDiscoveryMgr, []vmcp.Backend{}, nil)
 			require.NoError(t, err)
 			require.NotNil(t, s)
 
@@ -133,7 +133,7 @@ func TestServer_Address(t *testing.T) {
 			mockBackendClient := mocks.NewMockBackendClient(ctrl)
 			mockDiscoveryMgr := discoveryMocks.NewMockManager(ctrl)
 
-			s, err := server.New(tt.config, mockRouter, mockBackendClient, mockDiscoveryMgr, []vmcp.Backend{}, nil)
+			s, err := server.New(context.Background(), tt.config, mockRouter, mockBackendClient, mockDiscoveryMgr, []vmcp.Backend{}, nil)
 			require.NoError(t, err)
 			addr := s.Address()
 			assert.Equal(t, tt.expected, addr)
@@ -155,7 +155,7 @@ func TestServer_Stop(t *testing.T) {
 		mockDiscoveryMgr := discoveryMocks.NewMockManager(ctrl)
 		mockDiscoveryMgr.EXPECT().Stop().Times(1)
 
-		s, err := server.New(&server.Config{}, mockRouter, mockBackendClient, mockDiscoveryMgr, []vmcp.Backend{}, nil)
+		s, err := server.New(context.Background(), &server.Config{}, mockRouter, mockBackendClient, mockDiscoveryMgr, []vmcp.Backend{}, nil)
 		require.NoError(t, err)
 		err = s.Stop(context.Background())
 		require.NoError(t, err)

test/integration/vmcp/helpers/vmcp_server.go

@@ -165,7 +165,7 @@ func NewVMCPServer(
 	rtr := router.NewDefaultRouter()
 
 	// Create vMCP server with test-specific defaults
-	vmcpServer, err := vmcpserver.New(&vmcpserver.Config{
+	vmcpServer, err := vmcpserver.New(ctx, &vmcpserver.Config{
 		Name:              "test-vmcp",
 		Version:           "1.0.0",
 		Host:              "127.0.0.1",