Enhance health_check tool with always-on counts, chunks check, and embedding reachability

src/tools/__tests__/health.test.ts

@@ -29,6 +29,7 @@ vi.mock('../../db/pg-client.js', () => ({
 
 vi.mock('../../services/embeddings.js', () => ({
 	isEmbeddingsConfigured: vi.fn(() => true),
+	generateEmbedding: vi.fn(async () => [0.1, 0.2, 0.3]),
 }));
 
 vi.mock('../../utils/health-helpers.js', async (importOriginal) => {
@@ -37,15 +38,16 @@ vi.mock('../../utils/health-helpers.js', async (importOriginal) => {
 		...orig,
 		serverStartTime: Date.now() - 60_000, // 1 minute ago
 		checkTable: vi.fn(async () => true),
+		estimateRowCount: vi.fn(async () => 0),
 	};
 });
 
 // --- Imports ---
 
-import { checkDatabaseConnection, isDatabaseConfigured } from '../../db/pg-client.js';
-import { isEmbeddingsConfigured } from '../../services/embeddings.js';
+import { checkDatabaseConnection, isDatabaseConfigured, queryCount } from '../../db/pg-client.js';
+import { generateEmbedding, isEmbeddingsConfigured } from '../../services/embeddings.js';
 import { config } from '../../utils/config.js';
-import { checkTable } from '../../utils/health-helpers.js';
+import { checkTable, estimateRowCount } from '../../utils/health-helpers.js';
 import { HealthCheckOutputSchema, registerHealthTool } from '../health.js';
 
 // --- Helpers ---
@@ -97,10 +99,13 @@ describe('health_check tool', () => {
 			expect.objectContaining({ ok: true }),
 		);
 		expect(result.structuredContent?.checks.embeddings).toEqual(
-			expect.objectContaining({ ok: true, model: 'text-embedding-3-small' }),
+			expect.objectContaining({ ok: true, model: 'text-embedding-3-small', provider: 'openai' }),
 		);
 		expect(result.structuredContent?.checks.documents).toEqual(
-			expect.objectContaining({ ok: true }),
+			expect.objectContaining({ ok: true, count: 0 }),
+		);
+		expect(result.structuredContent?.checks.chunks).toEqual(
+			expect.objectContaining({ ok: true, count: 0 }),
 		);
 	});
 
@@ -144,27 +149,34 @@ describe('health_check tool', () => {
 	it('checks memory tables when ENABLE_MEMORY is true', async () => {
 		(config as { ENABLE_MEMORY: boolean }).ENABLE_MEMORY = true;
 		const result = (await callHealth(false)) as {
-			structuredContent?: { checks: Record<string, { ok: boolean }> };
+			structuredContent?: {
+				checks: Record<string, { ok: boolean; entities?: number; observations?: number }>;
+			};
 		};
 		expect(result.structuredContent?.checks.memory).toBeDefined();
 		expect(result.structuredContent?.checks.memory.ok).toBe(true);
+		expect(result.structuredContent?.checks.memory.entities).toBe(0);
+		expect(result.structuredContent?.checks.memory.observations).toBe(0);
 	});
 
 	it('checks conversation tables when ENABLE_CONVERSATIONS is true', async () => {
 		(config as { ENABLE_CONVERSATIONS: boolean }).ENABLE_CONVERSATIONS = true;
 		const result = (await callHealth(false)) as {
-			structuredContent?: { checks: Record<string, { ok: boolean }> };
+			structuredContent?: { checks: Record<string, { ok: boolean; sessions?: number }> };
 		};
 		expect(result.structuredContent?.checks.conversations).toBeDefined();
 		expect(result.structuredContent?.checks.conversations.ok).toBe(true);
+		expect(result.structuredContent?.checks.conversations.sessions).toBe(0);
 	});
 
 	it('checks insight tables when ENABLE_INSIGHTS is true', async () => {
 		(config as { ENABLE_INSIGHTS: boolean }).ENABLE_INSIGHTS = true;
 		const result = (await callHealth(false)) as {
-			structuredContent?: { checks: Record<string, { ok: boolean }> };
+			structuredContent?: { checks: Record<string, { ok: boolean; count?: number }> };
 		};
 		expect(result.structuredContent?.checks.insights).toBeDefined();
+		expect(result.structuredContent?.checks.insights.ok).toBe(true);
+		expect(result.structuredContent?.checks.insights.count).toBe(0);
 		expect(result.structuredContent?.checks.insightQueue).toBeDefined();
 	});
 
@@ -187,6 +199,76 @@ describe('health_check tool', () => {
 		expect(result.structuredContent?.checks.embeddings.model).toBe('nomic-embed-text-v2-moe');
 	});
 
+	it('always includes counts even when verbose is false', async () => {
+		(config as Record<string, unknown>).ENABLE_MEMORY = true;
+		(config as Record<string, unknown>).ENABLE_CONVERSATIONS = true;
+		(config as Record<string, unknown>).ENABLE_INSIGHTS = true;
+
+		const result = (await callHealth(false)) as {
+			structuredContent?: { checks: Record<string, Record<string, unknown>> };
+		};
+		const checks = result.structuredContent?.checks;
+		expect(checks?.documents.count).toBe(0);
+		expect(checks?.chunks.count).toBe(0);
+		expect(checks?.memory.entities).toBe(0);
+		expect(checks?.memory.observations).toBe(0);
+		expect(checks?.conversations.sessions).toBe(0);
+		expect(checks?.insights.count).toBe(0);
+	});
+
+	it('returns degraded when embedding service is unreachable', async () => {
+		vi.mocked(generateEmbedding).mockRejectedValueOnce(new Error('Connection refused'));
+		const result = (await callHealth(false)) as {
+			structuredContent?: {
+				status: string;
+				checks: Record<string, { ok: boolean; error?: string }>;
+			};
+		};
+		expect(result.structuredContent?.status).toBe('degraded');
+		expect(result.structuredContent?.checks.embeddings.ok).toBe(false);
+		expect(result.structuredContent?.checks.embeddings.error).toBe('Connection refused');
+	});
+
+	it('insight queue failure sets status to degraded', async () => {
+		(config as Record<string, unknown>).ENABLE_INSIGHTS = true;
+		vi.mocked(checkTable).mockImplementation(async (table: string) => {
+			return table !== 'insight_queue';
+		});
+
+		const result = (await callHealth(false)) as {
+			structuredContent?: {
+				status: string;
+				checks: Record<string, { ok: boolean; error?: string }>;
+			};
+		};
+		expect(result.structuredContent?.status).toBe('degraded');
+		expect(result.structuredContent?.checks.insightQueue.ok).toBe(false);
+		expect(result.structuredContent?.checks.insightQueue.error).toContain('not accessible');
+	});
+
+	it('uses estimated counts by default and exact counts in verbose mode', async () => {
+		vi.mocked(estimateRowCount).mockResolvedValue(100);
+		vi.mocked(queryCount).mockResolvedValue(42);
+
+		const defaultResult = (await callHealth(false)) as {
+			structuredContent?: { checks: Record<string, { count?: number }> };
+		};
+		expect(defaultResult.structuredContent?.checks.documents.count).toBe(100);
+
+		const verboseResult = (await callHealth(true)) as {
+			structuredContent?: { checks: Record<string, { count?: number }> };
+		};
+		expect(verboseResult.structuredContent?.checks.documents.count).toBe(42);
+	});
+
+	it('includes provider in embeddings check', async () => {
+		(config as Record<string, unknown>).EMBEDDING_PROVIDER = 'google';
+		const result = (await callHealth(false)) as {
+			structuredContent?: { checks: Record<string, { provider?: string }> };
+		};
+		expect(result.structuredContent?.checks.embeddings.provider).toBe('google');
+	});
+
 	describe('output schema validation', () => {
 		const outputSchema = z.object(HealthCheckOutputSchema);
 

src/tools/health.ts

@@ -1,10 +1,21 @@
 import type { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js';
 import { z } from 'zod';
-import { checkDatabaseConnection, isDatabaseConfigured } from '../db/pg-client.js';
-import { isEmbeddingsConfigured } from '../services/embeddings.js';
+import {
+	checkDatabaseConnection,
+	isDatabaseConfigured,
+	queryCount,
+	queryOne,
+} from '../db/pg-client.js';
+import { generateEmbedding, isEmbeddingsConfigured } from '../services/embeddings.js';
 import { toolError, toolResponse } from '../utils/compact.js';
 import { config } from '../utils/config.js';
-import { checkTable, formatUptime, serverStartTime, timed } from '../utils/health-helpers.js';
+import {
+	checkTable,
+	estimateRowCount,
+	formatUptime,
+	serverStartTime,
+	timed,
+} from '../utils/health-helpers.js';
 import { logger } from '../utils/logger.js';
 
 import pkg from '../../package.json' with { type: 'json' };
@@ -16,8 +27,10 @@ const ComponentCheckSchema = z.object({
 	latencyMs: z.number().optional(),
 	error: z.string().optional(),
 	model: z.string().optional(),
+	provider: z.string().optional(),
 	count: z.number().optional(),
 	entities: z.number().optional(),
+	observations: z.number().optional(),
 	sessions: z.number().optional(),
 	pending: z.number().optional(),
 });
@@ -98,25 +111,43 @@ export function registerHealthTool(server: McpServer): void {
 					checks.database = { ok: false, error: 'DATABASE_URL not configured' };
 				}
 
-				// 2. Embeddings
-				const embeddingsOk = isEmbeddingsConfigured();
-				checks.embeddings = { ok: embeddingsOk, model: getEmbeddingModelName() };
-				if (!embeddingsOk) {
+				// 2. Embeddings (config + reachability)
+				const embeddingsConfigured = isEmbeddingsConfigured();
+				checks.embeddings = {
+					ok: embeddingsConfigured,
+					model: getEmbeddingModelName(),
+					provider: config.EMBEDDING_PROVIDER,
+				};
+				if (!embeddingsConfigured) {
 					checks.embeddings.error = `${config.EMBEDDING_PROVIDER} not configured`;
+					hasFailure = true;
+				} else {
+					try {
+						await generateEmbedding('health check');
+					} catch (err) {
+						checks.embeddings.ok = false;
+						checks.embeddings.error =
+							err instanceof Error ? err.message : 'Embedding service unreachable';
+						hasFailure = true;
+					}
 				}
 
 				// 3. Table checks (only if DB is connected)
 				if (dbOk) {
+					// Use estimated counts by default (pg_class catalog), exact count(*) in verbose
+					const getCount = verbose
+						? (table: string) => queryCount(`SELECT count(*) FROM ${table}`)
+						: estimateRowCount;
+
 					// Documents (always checked)
 					try {
 						const docsOk = await checkTable('documents');
 						checks.documents = { ok: docsOk };
 						if (!docsOk) {
 							checks.documents.error = "Table 'documents' not accessible";
 							hasFailure = true;
-						} else if (verbose) {
-							const { queryCount } = await import('../db/pg-client.js');
-							checks.documents.count = await queryCount('SELECT count(*) FROM documents');
+						} else {
+							checks.documents.count = await getCount('documents');
 						}
 					} catch (err) {
 						checks.documents = {
@@ -126,6 +157,24 @@ export function registerHealthTool(server: McpServer): void {
 						hasFailure = true;
 					}
 
+					// Chunks (always checked)
+					try {
+						const chunksOk = await checkTable('chunks');
+						checks.chunks = { ok: chunksOk };
+						if (!chunksOk) {
+							checks.chunks.error = "Table 'chunks' not accessible";
+							hasFailure = true;
+						} else {
+							checks.chunks.count = await getCount('chunks');
+						}
+					} catch (err) {
+						checks.chunks = {
+							ok: false,
+							error: err instanceof Error ? err.message : 'Check failed',
+						};
+						hasFailure = true;
+					}
+
 					// Memory (if enabled)
 					if (config.ENABLE_MEMORY) {
 						try {
@@ -134,9 +183,9 @@ export function registerHealthTool(server: McpServer): void {
 							if (!memOk) {
 								checks.memory.error = "Table 'memory_entities' not accessible";
 								hasFailure = true;
-							} else if (verbose) {
-								const { queryCount } = await import('../db/pg-client.js');
-								checks.memory.entities = await queryCount('SELECT count(*) FROM memory_entities');
+							} else {
+								checks.memory.entities = await getCount('memory_entities');
+								checks.memory.observations = await getCount('memory_observations');
 							}
 						} catch (err) {
 							checks.memory = {
@@ -155,11 +204,8 @@ export function registerHealthTool(server: McpServer): void {
 							if (!convOk) {
 								checks.conversations.error = "Table 'conversation_sessions' not accessible";
 								hasFailure = true;
-							} else if (verbose) {
-								const { queryCount } = await import('../db/pg-client.js');
-								checks.conversations.sessions = await queryCount(
-									'SELECT count(*) FROM conversation_sessions',
-								);
+							} else {
+								checks.conversations.sessions = await getCount('conversation_sessions');
 							}
 						} catch (err) {
 							checks.conversations = {
@@ -178,6 +224,8 @@ export function registerHealthTool(server: McpServer): void {
 							if (!insOk) {
 								checks.insights.error = "Table 'proactive_insights' not accessible";
 								hasFailure = true;
+							} else {
+								checks.insights.count = await getCount('proactive_insights');
 							}
 						} catch (err) {
 							checks.insights = {
@@ -193,8 +241,8 @@ export function registerHealthTool(server: McpServer): void {
 							checks.insightQueue = { ok: queueOk };
 							if (!queueOk) {
 								checks.insightQueue.error = "Table 'insight_queue' not accessible";
-							} else if (verbose) {
-								const { queryOne } = await import('../db/pg-client.js');
+								hasFailure = true;
+							} else {
 								const row = await queryOne<{ chunks_pending: number }>(
 									'SELECT chunks_pending FROM insight_queue WHERE id = 1',
 								);
@@ -207,6 +255,7 @@ export function registerHealthTool(server: McpServer): void {
 								ok: false,
 								error: err instanceof Error ? err.message : 'Check failed',
 							};
+							hasFailure = true;
 						}
 					}
 				}

src/utils/health-helpers.ts

@@ -33,3 +33,23 @@ export function formatUptime(seconds: number): string {
 	if (h > 0) return `${h}h ${m}m`;
 	return `${m}m ${seconds % 60}s`;
 }
+
+/** Estimated row count from pg_class (cheap, no table scan). */
+export async function estimateRowCount(tableName: string): Promise<number> {
+	if (!isDatabaseConfigured()) return -1;
+	if (!SAFE_IDENTIFIER.test(tableName)) return -1;
+	try {
+		const { rows } = await pgQuery<{ reltuples: string }>(
+			`SELECT c.reltuples
+			 FROM pg_class c
+			 JOIN pg_namespace n ON n.oid = c.relnamespace
+			 WHERE c.relname = $1 AND n.nspname = 'public'`,
+			[tableName],
+		);
+		if (!rows[0]) return -1;
+		const est = Number(rows[0].reltuples);
+		return est < 0 ? 0 : Math.round(est);
+	} catch {
+		return -1;
+	}
+}