diff --git a/packages/agent/src/agent.ts b/packages/agent/src/agent.ts index edad2de46..3c88834f6 100644 --- a/packages/agent/src/agent.ts +++ b/packages/agent/src/agent.ts @@ -50,6 +50,8 @@ export default class Agent 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 extends FrameworkMounter * Restart the agent at runtime (remount routes). */ async restart(): Promise { - // 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; + } } /** diff --git a/packages/agent/test/agent.test.ts b/packages/agent/test/agent.test.ts index 7d41be47c..dce11a932 100644 --- a/packages/agent/test/agent.test.ts +++ b/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', () => {