fix(agent): avoid concurent restart lead to an empty schema sending

packages/agent/src/agent.ts

@@ -50,6 +50,8 @@ export default class Agent<S extends TSchema = TSchema> extends FrameworkMounter
   private mcpEnabled = false;
   private mcpEnabledTools?: ToolName[];
 
+  private isRestarting = false;
+
   /**
    * Create a new Agent Builder.
    * If any options are missing, the default will be applied:
@@ -111,11 +113,25 @@ export default class Agent<S extends TSchema = TSchema> extends FrameworkMounter
    * Restart the agent at runtime (remount routes).
    */
   async restart(): Promise<void> {
-    // We force sending schema when restarting
-    const { router, mcpHttpCallback } = await this.buildRouterAndSendSchema();
+    if (this.isRestarting) {
+      this.options.logger('Debug', 'Agent is already restarting. Do nothing.');
+
+      return;
+    }
+
+    this.options.logger('Info', `Agent is restarting...`);
+
+    this.isRestarting = true;
 
-    this.setMcpCallback(mcpHttpCallback ?? null);
-    await this.remount(router);
+    try {
+      // We force sending schema when restarting
+      const { router, mcpHttpCallback } = await this.buildRouterAndSendSchema();
+
+      this.setMcpCallback(mcpHttpCallback ?? null);
+      await this.remount(router);
+    } finally {
+      this.isRestarting = false;
+    }
   }
 
   /**

packages/agent/test/agent.test.ts

@@ -223,6 +223,68 @@ describe('Agent', () => {
 
       expect(DataSourceCustomizer.prototype.use).toHaveBeenCalledTimes(2);
     });
+
+    describe('when a restart is already in progress', () => {
+      test('should ignore concurrent restart calls and not rebuild routes again', async () => {
+        const mockLogger = jest.fn();
+        const agent = new Agent({ ...options, logger: mockLogger });
+
+        await agent.start();
+        jest.clearAllMocks();
+
+        // Make the next getDataSource call hang so the first restart is still
+        // in progress when the second one is triggered.
+        let releaseFirstRestart: () => void = () => {};
+
+        jest.mocked(DataSourceCustomizer.prototype.getDataSource).mockImplementationOnce(
+          () =>
+            new Promise(resolve => {
+              releaseFirstRestart = () => resolve(factories.dataSource.build());
+            }),
+        );
+
+        const firstRestart = agent.restart();
+        const secondRestart = agent.restart();
+
+        releaseFirstRestart();
+        await Promise.all([firstRestart, secondRestart]);
+
+        expect(DataSourceCustomizer.prototype.getDataSource).toHaveBeenCalledTimes(1);
+        expect(mockSetupRoute).toHaveBeenCalledTimes(1);
+        expect(mockBootstrap).toHaveBeenCalledTimes(1);
+        expect(mockMakeRoutes).toHaveBeenCalledTimes(1);
+        expect(mockLogger).toHaveBeenCalledWith(
+          'Debug',
+          'Agent is already restarting. Do nothing.',
+        );
+      });
+
+      test('should allow a new restart once the previous one has finished', async () => {
+        const agent = new Agent(options);
+
+        await agent.start();
+        jest.clearAllMocks();
+
+        await agent.restart();
+        await agent.restart();
+
+        expect(DataSourceCustomizer.prototype.getDataSource).toHaveBeenCalledTimes(2);
+        expect(mockSetupRoute).toHaveBeenCalledTimes(2);
+        expect(mockBootstrap).toHaveBeenCalledTimes(2);
+      });
+
+      test('should log that the agent is restarting on each accepted restart', async () => {
+        const mockLogger = jest.fn();
+        const agent = new Agent({ ...options, logger: mockLogger });
+
+        await agent.start();
+        mockLogger.mockClear();
+
+        await agent.restart();
+
+        expect(mockLogger).toHaveBeenCalledWith('Info', 'Agent is restarting...');
+      });
+    });
   });
 
   describe('Production', () => {