diff --git a/src/tools/__tests__/health.test.ts b/src/tools/__tests__/health.test.ts index 4368f45..a8f0b6d 100644 --- a/src/tools/__tests__/health.test.ts +++ b/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 }; + structuredContent?: { + checks: Record; + }; }; 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 }; + structuredContent?: { checks: Record }; }; 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 }; + structuredContent?: { checks: Record }; }; 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).ENABLE_MEMORY = true; + (config as Record).ENABLE_CONVERSATIONS = true; + (config as Record).ENABLE_INSIGHTS = true; + + const result = (await callHealth(false)) as { + structuredContent?: { checks: Record> }; + }; + 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; + }; + }; + 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).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; + }; + }; + 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 }; + }; + expect(defaultResult.structuredContent?.checks.documents.count).toBe(100); + + const verboseResult = (await callHealth(true)) as { + structuredContent?: { checks: Record }; + }; + expect(verboseResult.structuredContent?.checks.documents.count).toBe(42); + }); + + it('includes provider in embeddings check', async () => { + (config as Record).EMBEDDING_PROVIDER = 'google'; + const result = (await callHealth(false)) as { + structuredContent?: { checks: Record }; + }; + expect(result.structuredContent?.checks.embeddings.provider).toBe('google'); + }); + describe('output schema validation', () => { const outputSchema = z.object(HealthCheckOutputSchema); diff --git a/src/tools/health.ts b/src/tools/health.ts index 0e91f28..a689087 100644 --- a/src/tools/health.ts +++ b/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,15 +111,34 @@ 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'); @@ -114,9 +146,8 @@ export function registerHealthTool(server: McpServer): void { 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; } } } diff --git a/src/utils/health-helpers.ts b/src/utils/health-helpers.ts index 500f39c..27268f3 100644 --- a/src/utils/health-helpers.ts +++ b/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 { + 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; + } +}