Skip to content

Commit d735874

Browse files
tobiutobiu
committed
Implement Smart Lifecycle Management for ChromaDB processes #8284 
1 parent 8eb43be commit d735874

4 files changed

Lines changed: 110 additions & 4 deletions

File tree

‎ai/mcp/server/knowledge-base/openapi.yaml‎

Lines changed: 16 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -61,7 +61,14 @@ paths:
6161
summary: Start Database
6262
operationId: start_database
6363
description: |
64-
Starts the ChromaDB database instance as a background process.
64+
Starts the ChromaDB database instance as a background process, or connects to an existing one.
65+
66+
**Behavior:**
67+
- **Managed:** If no database is running on the configured port, this tool spawns a new process managed by this server. This process will be automatically cleaned up when the agent session ends.
68+
- **External:** If a database is already running, this tool simply confirms the connection. The server acts as a client and will NOT kill the database on exit.
69+
70+
**Multi-Agent / Swarm Recommendation:**
71+
For workflows involving multiple concurrent agents, it is **highly recommended** to start the database externally using `npm run ai:server` before starting the agents. This prevents unexpected disconnects for other agents when the "owner" agent exits.
6572
6673
**When to Use:**
6774
Use this tool if a `healthcheck` reveals that the database process is not running.
@@ -87,8 +94,15 @@ paths:
8794
description: |
8895
Stops the running ChromaDB database instance.
8996
97+
**Debug / Maintenance Tool:**
98+
This tool is generally **not required** for standard workflows. The server now implements automatic process cleanup (for managed instances) when the session concludes.
99+
100+
**Effect:**
101+
- **Managed:** Stops the process spawned by this server.
102+
- **External:** **No effect.** The server cannot stop a database it did not start.
103+
90104
**When to Use:**
91-
Use this tool to shut down the database process at the end of a session to free up resources.
105+
Use this only for debugging, forcing a restart, or freeing up resources in a long-running environment where the agent session itself does not exit.
92106
tags: ["Database Lifecycle"]
93107
responses:
94108
'200':

‎ai/mcp/server/knowledge-base/services/DatabaseLifecycleService.mjs‎

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -89,6 +89,13 @@ class DatabaseLifecycleService extends Base {
8989
spawnedProcess.on('spawn', () => {
9090
this.chromaProcess = spawnedProcess;
9191
logger.log(`ChromaDB (Knowledge Base) process started with PID: ${this.chromaProcess.pid}`);
92+
93+
// Register cleanup handlers
94+
this.cleanupHandler = this.cleanup.bind(this);
95+
process.on('exit', this.cleanupHandler);
96+
process.on('SIGINT', this.cleanupHandler);
97+
process.on('SIGTERM', this.cleanupHandler);
98+
9299
resolve();
93100
});
94101

@@ -108,6 +115,29 @@ class DatabaseLifecycleService extends Base {
108115
return result;
109116
}
110117

118+
/**
119+
* Handles process termination signals to ensure the child process is killed.
120+
* @param {string|number} signalOrCode
121+
*/
122+
async cleanup(signalOrCode) {
123+
if (this.chromaProcess) {
124+
logger.log(`[DatabaseLifecycleService] cleanup triggered by ${signalOrCode}`);
125+
try {
126+
// We use the synchronous kill here because async operations might not complete
127+
// reliably during the 'exit' event.
128+
process.kill(-this.chromaProcess.pid, 'SIGTERM');
129+
this.chromaProcess = null;
130+
} catch (e) {
131+
// Ignore errors if process is already gone
132+
}
133+
}
134+
135+
// If this was a signal (not a normal exit), we need to exit explicitly
136+
if (typeof signalOrCode === 'string') {
137+
process.exit(0);
138+
}
139+
}
140+
111141
/**
112142
* Waits for the ChromaDB server to respond to a heartbeat.
113143
* @returns {Promise<void>}
@@ -138,6 +168,15 @@ class DatabaseLifecycleService extends Base {
138168
this.chromaProcess.on('exit', () => {
139169
logger.log(`ChromaDB process with PID: ${pid} has been stopped.`);
140170
this.chromaProcess = null;
171+
172+
// Remove cleanup handlers
173+
if (this.cleanupHandler) {
174+
process.off('exit', this.cleanupHandler);
175+
process.off('SIGINT', this.cleanupHandler);
176+
process.off('SIGTERM', this.cleanupHandler);
177+
this.cleanupHandler = null;
178+
}
179+
141180
const result = { status: 'stopped' };
142181
this.fire('processStopped', { pid, managedByService: true });
143182
resolve(result);

‎ai/mcp/server/memory-core/openapi.yaml‎

Lines changed: 16 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -66,7 +66,14 @@ paths:
6666
summary: Start Database
6767
operationId: start_database
6868
description: |
69-
Starts the ChromaDB database instance for the Memory Core as a background process.
69+
Starts the ChromaDB database instance for the Memory Core as a background process, or connects to an existing one.
70+
71+
**Behavior:**
72+
- **Managed:** If no database is running on the configured port, this tool spawns a new process managed by this server. This process will be automatically cleaned up when the agent session ends.
73+
- **External:** If a database is already running (e.g., via `npm run ai:server-memory`), this tool simply confirms the connection. The server acts as a client and will NOT kill the database on exit.
74+
75+
**Multi-Agent / Swarm Recommendation:**
76+
For workflows involving multiple concurrent agents, it is **highly recommended** to start the database externally using `npm run ai:server-memory` before starting the agents. This prevents unexpected disconnects for other agents when the "owner" agent exits.
7077
7178
**When to Use:**
7279
Use this tool if a `healthcheck` reveals that the database process is not running.
@@ -92,8 +99,15 @@ paths:
9299
description: |
93100
Stops the running ChromaDB database instance for the Memory Core.
94101
102+
**Debug / Maintenance Tool:**
103+
This tool is generally **not required** for standard workflows. The server now implements automatic process cleanup (for managed instances) when the session concludes.
104+
105+
**Effect:**
106+
- **Managed:** Stops the process spawned by this server.
107+
- **External:** **No effect.** The server cannot stop a database it did not start.
108+
95109
**When to Use:**
96-
Use this tool to shut down the database process at the end of a session to free up resources.
110+
Use this only for debugging, forcing a restart, or freeing up resources in a long-running environment where the agent session itself does not exit.
97111
tags: ["Database Lifecycle"]
98112
responses:
99113
'200':

‎ai/mcp/server/memory-core/services/DatabaseLifecycleService.mjs‎

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -93,6 +93,13 @@ class DatabaseLifecycleService extends Base {
9393
spawnedProcess.on('spawn', () => {
9494
this.chromaProcess = spawnedProcess;
9595
logger.log(`ChromaDB (Memory Core) process started with PID: ${this.chromaProcess.pid}`);
96+
97+
// Register cleanup handlers
98+
this.cleanupHandler = this.cleanup.bind(this);
99+
process.on('exit', this.cleanupHandler);
100+
process.on('SIGINT', this.cleanupHandler);
101+
process.on('SIGTERM', this.cleanupHandler);
102+
96103
resolve();
97104
});
98105

@@ -120,6 +127,29 @@ class DatabaseLifecycleService extends Base {
120127
}
121128
}
122129

130+
/**
131+
* Handles process termination signals to ensure the child process is killed.
132+
* @param {string|number} signalOrCode
133+
*/
134+
async cleanup(signalOrCode) {
135+
if (this.chromaProcess) {
136+
logger.log(`[DatabaseLifecycleService] cleanup triggered by ${signalOrCode}`);
137+
try {
138+
// We use the synchronous kill here because async operations might not complete
139+
// reliably during the 'exit' event.
140+
process.kill(-this.chromaProcess.pid, 'SIGTERM');
141+
this.chromaProcess = null;
142+
} catch (e) {
143+
// Ignore errors if process is already gone
144+
}
145+
}
146+
147+
// If this was a signal (not a normal exit), we need to exit explicitly
148+
if (typeof signalOrCode === 'string') {
149+
process.exit(0);
150+
}
151+
}
152+
123153
/**
124154
* Waits for the ChromaDB server to respond to a heartbeat.
125155
* @returns {Promise<void>}
@@ -151,6 +181,15 @@ class DatabaseLifecycleService extends Base {
151181
this.chromaProcess.on('exit', () => {
152182
logger.log(`ChromaDB process with PID: ${pid} has been stopped.`);
153183
this.chromaProcess = null;
184+
185+
// Remove cleanup handlers
186+
if (this.cleanupHandler) {
187+
process.off('exit', this.cleanupHandler);
188+
process.off('SIGINT', this.cleanupHandler);
189+
process.off('SIGTERM', this.cleanupHandler);
190+
this.cleanupHandler = null;
191+
}
192+
154193
const result = { status: 'stopped' };
155194
this.fire('processStopped', { pid, managedByService: true });
156195
resolve(result);

0 commit comments

Comments
 (0)