diff --git a/src/api/status.ts b/src/api/status.ts index b5f47ee..e359dbd 100644 --- a/src/api/status.ts +++ b/src/api/status.ts @@ -61,6 +61,7 @@ const TOOL_DEFINITIONS: ToolDef[] = [ { name: 'update_document', group: 'Document' }, { name: 'add_note', group: 'Document' }, { name: 'get_stats', group: 'Stats' }, + { name: 'health_check', group: 'Stats' }, // Memory tools { diff --git a/src/tools/__tests__/health.test.ts b/src/tools/__tests__/health.test.ts index a8f0b6d..c0d01ce 100644 --- a/src/tools/__tests__/health.test.ts +++ b/src/tools/__tests__/health.test.ts @@ -1,17 +1,9 @@ import { beforeEach, describe, expect, it, vi } from 'vitest'; import { z } from 'zod'; -// --- Mocks (must be declared before imports) --- - vi.mock('../../utils/config.js', () => ({ config: { COMPACT_RESPONSES: false, - ENABLE_MEMORY: false, - ENABLE_CONVERSATIONS: false, - ENABLE_INSIGHTS: false, - EMBEDDING_PROVIDER: 'openai', - OLLAMA_MODEL: 'nomic-embed-text', - GOOGLE_EMBEDDING_MODEL: 'text-embedding-004', }, })); @@ -21,40 +13,15 @@ vi.mock('../../utils/logger.js', () => ({ vi.mock('../../db/pg-client.js', () => ({ isDatabaseConfigured: vi.fn(() => true), - checkDatabaseConnection: vi.fn(async () => true), pgQuery: vi.fn(), - queryCount: vi.fn(async () => 0), - queryOne: vi.fn(async () => null), -})); - -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) => { - const orig = await importOriginal(); - return { - ...orig, - serverStartTime: Date.now() - 60_000, // 1 minute ago - checkTable: vi.fn(async () => true), - estimateRowCount: vi.fn(async () => 0), - }; -}); - -// --- Imports --- - -import { checkDatabaseConnection, isDatabaseConfigured, queryCount } from '../../db/pg-client.js'; -import { generateEmbedding, isEmbeddingsConfigured } from '../../services/embeddings.js'; -import { config } from '../../utils/config.js'; -import { checkTable, estimateRowCount } from '../../utils/health-helpers.js'; +import { isDatabaseConfigured, pgQuery } from '../../db/pg-client.js'; import { HealthCheckOutputSchema, registerHealthTool } from '../health.js'; -// --- Helpers --- - -type ToolHandler = (args: { verbose: boolean }, extra: unknown) => Promise; +type ToolHandler = (args: Record, extra: unknown) => Promise; -function createHandler(): (verbose: boolean) => Promise { +function createHandler(): () => Promise { let handler: ToolHandler | undefined; const fakeServer = { @@ -64,250 +31,84 @@ function createHandler(): (verbose: boolean) => Promise { }; registerHealthTool(fakeServer as never); - if (!handler) throw new Error('registerHealthTool did not register a handler'); - const h = handler; - return (verbose: boolean) => h({ verbose }, {}); + const registeredHandler = handler; + return () => registeredHandler({}, {}); } -// --- Tests --- - describe('health_check tool', () => { let callHealth: ReturnType; beforeEach(() => { vi.clearAllMocks(); vi.mocked(isDatabaseConfigured).mockReturnValue(true); - vi.mocked(checkDatabaseConnection).mockResolvedValue(true); - vi.mocked(isEmbeddingsConfigured).mockReturnValue(true); - vi.mocked(checkTable).mockResolvedValue(true); - - (config as Record).ENABLE_MEMORY = false; - (config as Record).ENABLE_CONVERSATIONS = false; - (config as Record).ENABLE_INSIGHTS = false; - (config as Record).EMBEDDING_PROVIDER = 'openai'; - + vi.mocked(pgQuery).mockImplementation(async (sql: string) => { + if (sql === 'SELECT 1') return { rows: [{ '?column?': 1 }], rowCount: 1 }; + return { rows: [{ c: 3 }], rowCount: 1 }; + }); callHealth = createHandler(); }); it('returns healthy when all checks pass', async () => { - const result = (await callHealth(false)) as { + const result = (await callHealth()) as { structuredContent?: { status: string; checks: Record }; }; expect(result.structuredContent?.status).toBe('healthy'); expect(result.structuredContent?.checks.database).toEqual( expect.objectContaining({ ok: true }), ); - expect(result.structuredContent?.checks.embeddings).toEqual( - expect.objectContaining({ ok: true, model: 'text-embedding-3-small', provider: 'openai' }), - ); expect(result.structuredContent?.checks.documents).toEqual( - expect.objectContaining({ ok: true, count: 0 }), - ); - expect(result.structuredContent?.checks.chunks).toEqual( - expect.objectContaining({ ok: true, count: 0 }), + expect.objectContaining({ ok: true, count: 3 }), ); }); - it('returns unhealthy when database is not configured', async () => { + it('returns config error when database is not configured', async () => { vi.mocked(isDatabaseConfigured).mockReturnValue(false); - const result = (await callHealth(false)) as { - structuredContent?: { - status: string; - checks: Record; - }; - }; - expect(result.structuredContent?.status).toBe('unhealthy'); - expect(result.structuredContent?.checks.database.ok).toBe(false); - expect(result.structuredContent?.checks.database.error).toContain('not configured'); + const result = (await callHealth()) as { isError?: boolean; content?: Array<{ text: string }> }; + expect(result.isError).toBe(true); + expect(result.content?.[0]?.text).toContain('DATABASE_URL'); }); - it('returns unhealthy when database connection fails', async () => { - vi.mocked(checkDatabaseConnection).mockResolvedValue(false); - const result = (await callHealth(false)) as { - structuredContent?: { status: string; checks: Record }; - }; - expect(result.structuredContent?.status).toBe('unhealthy'); - expect(result.structuredContent?.checks.database.ok).toBe(false); - }); - - it('returns degraded when a table check fails', async () => { - vi.mocked(checkTable).mockImplementation(async (table: string) => { - return table !== 'documents'; + it('returns degraded when one subsystem fails', async () => { + vi.mocked(pgQuery).mockImplementation(async (sql: string) => { + if (sql.includes('conversation_sessions')) throw new Error('missing table'); + if (sql === 'SELECT 1') return { rows: [{ '?column?': 1 }], rowCount: 1 }; + return { rows: [{ c: 3 }], rowCount: 1 }; }); - - const result = (await callHealth(false)) as { + const result = (await callHealth()) as { structuredContent?: { status: string; checks: Record; }; }; expect(result.structuredContent?.status).toBe('degraded'); - expect(result.structuredContent?.checks.documents.ok).toBe(false); - }); - - 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; - }; - }; - 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); + expect(result.structuredContent?.checks.conversations.ok).toBe(false); + expect(result.structuredContent?.checks.conversations.error).toContain('missing table'); }); - 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 }; - }; - 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 }; - }; - 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(); - }); - - it('includes version and uptime', async () => { - const result = (await callHealth(false)) as { - structuredContent?: { version: string; uptime: string }; - }; - expect(result.structuredContent?.version).toBeDefined(); - expect(result.structuredContent?.uptime).toBeDefined(); - // Uptime should be a human-readable string - expect(result.structuredContent?.uptime).toMatch(/\d+[dhms]/); - }); - - it('reports correct embedding model for ollama', async () => { - (config as Record).EMBEDDING_PROVIDER = 'ollama'; - (config as Record).OLLAMA_MODEL = 'nomic-embed-text-v2-moe'; - const result = (await callHealth(false)) as { - structuredContent?: { checks: Record }; - }; - 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'; + it('handles empty tables gracefully', async () => { + vi.mocked(pgQuery).mockImplementation(async (sql: string) => { + if (sql === 'SELECT 1') return { rows: [{ '?column?': 1 }], rowCount: 1 }; + return { rows: [{ c: 0 }], rowCount: 1 }; }); - const result = (await callHealth(false)) as { + const result = (await callHealth()) as { structuredContent?: { status: string; - checks: Record; + 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); - }); + expect(result.structuredContent?.status).toBe('healthy'); + expect(result.structuredContent?.checks.documents.count).toBe(0); - 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'); + const outputSchema = z.object(HealthCheckOutputSchema); + expect(outputSchema.safeParse(result.structuredContent).success).toBe(true); }); - describe('output schema validation', () => { + it('output matches schema', async () => { const outputSchema = z.object(HealthCheckOutputSchema); - - it('healthy state passes schema validation', async () => { - const result = (await callHealth(false)) as { - structuredContent?: Record; - }; - const parsed = outputSchema.safeParse(result.structuredContent); - expect(parsed.success).toBe(true); - }); - - it('degraded state passes schema validation', async () => { - vi.mocked(checkTable).mockResolvedValue(false); - const result = (await callHealth(false)) as { - structuredContent?: Record; - }; - const parsed = outputSchema.safeParse(result.structuredContent); - expect(parsed.success).toBe(true); - }); - - it('unhealthy state passes schema validation', async () => { - vi.mocked(isDatabaseConfigured).mockReturnValue(false); - const result = (await callHealth(false)) as { - structuredContent?: Record; - }; - const parsed = outputSchema.safeParse(result.structuredContent); - expect(parsed.success).toBe(true); - }); - - it('verbose mode passes schema validation', async () => { - (config as { ENABLE_MEMORY: boolean }).ENABLE_MEMORY = true; - (config as { ENABLE_CONVERSATIONS: boolean }).ENABLE_CONVERSATIONS = true; - (config as { ENABLE_INSIGHTS: boolean }).ENABLE_INSIGHTS = true; - - const result = (await callHealth(true)) as { - structuredContent?: Record; - }; - const parsed = outputSchema.safeParse(result.structuredContent); - expect(parsed.success).toBe(true); - }); + const result = (await callHealth()) as { structuredContent?: Record }; + expect(outputSchema.safeParse(result.structuredContent).success).toBe(true); }); }); diff --git a/src/tools/health.ts b/src/tools/health.ts index a689087..eed442e 100644 --- a/src/tools/health.ts +++ b/src/tools/health.ts @@ -1,75 +1,117 @@ import type { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'; import { z } from 'zod'; -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, - estimateRowCount, - formatUptime, - serverStartTime, - timed, -} from '../utils/health-helpers.js'; +import { isDatabaseConfigured, pgQuery } from '../db/pg-client.js'; +import { configError, toolError, toolResponse } from '../utils/compact.js'; import { logger } from '../utils/logger.js'; -import pkg from '../../package.json' with { type: 'json' }; - -// --- Output Schema --- - const ComponentCheckSchema = z.object({ ok: z.boolean(), 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(), }); export const HealthCheckOutputSchema = { status: z.enum(['healthy', 'degraded', 'unhealthy']), checks: z.record(z.string(), ComponentCheckSchema), - version: z.string(), - uptime: z.string(), + timestamp: z.string(), }; -// --- Embedding model name lookup --- +export type TableName = + | 'documents' + | 'chunks' + | 'memory_entities' + | 'conversation_sessions' + | 'proactive_insights'; -const EMBEDDING_MODEL_MAP: Record = { - openai: 'text-embedding-3-small', -}; +async function countRows(table: TableName): Promise { + const result = await pgQuery<{ c: number }>(`SELECT COUNT(*)::int AS c FROM ${table}`); + return result.rows[0]?.c ?? 0; +} -function getEmbeddingModelName(): string { - if (config.EMBEDDING_PROVIDER === 'ollama') return config.OLLAMA_MODEL; - if (config.EMBEDDING_PROVIDER === 'google') return config.GOOGLE_EMBEDDING_MODEL; - return EMBEDDING_MODEL_MAP[config.EMBEDDING_PROVIDER] ?? 'unknown'; +async function handleHealthCheck() { + const checks: Record> = {}; + + const t0 = Date.now(); + try { + await pgQuery('SELECT 1'); + checks.database = { ok: true, latencyMs: Date.now() - t0 }; + } catch (error) { + checks.database = { + ok: false, + error: error instanceof Error ? error.message : String(error), + }; + } + + try { + checks.documents = { ok: true, count: await countRows('documents') }; + } catch (error) { + checks.documents = { + ok: false, + error: error instanceof Error ? error.message : String(error), + }; + } + + try { + checks.chunks = { ok: true, count: await countRows('chunks') }; + } catch (error) { + checks.chunks = { + ok: false, + error: error instanceof Error ? error.message : String(error), + }; + } + + try { + checks.memory = { ok: true, entities: await countRows('memory_entities') }; + } catch (error) { + checks.memory = { + ok: false, + error: error instanceof Error ? error.message : String(error), + }; + } + + try { + checks.conversations = { + ok: true, + sessions: await countRows('conversation_sessions'), + }; + } catch (error) { + checks.conversations = { + ok: false, + error: error instanceof Error ? error.message : String(error), + }; + } + + try { + checks.insights = { ok: true, count: await countRows('proactive_insights') }; + } catch (error) { + checks.insights = { + ok: false, + error: error instanceof Error ? error.message : String(error), + }; + } + + const values = Object.values(checks); + const allOk = values.every((check) => check.ok === true); + const anyOk = values.some((check) => check.ok === true); + + return { + status: allOk ? 'healthy' : anyOk ? 'degraded' : 'unhealthy', + checks, + timestamp: new Date().toISOString(), + }; } -/** - * Register the health_check tool — system diagnostics for agents. - */ export function registerHealthTool(server: McpServer): void { server.registerTool( 'health_check', { title: 'Health Check', description: - 'Check the health of the Textrawl server and all its subsystems. Returns pass/fail per component with an overall status. Use this as the first diagnostic step when something seems broken.', - inputSchema: { - verbose: z - .boolean() - .default(false) - .describe('Include latency measurements and row counts for each component'), - }, + 'Quick diagnostic check across all textrawl subsystems. Returns pass/fail per component with error details on failure. Use as the first call when other tools are failing.', + inputSchema: {}, outputSchema: HealthCheckOutputSchema, annotations: { readOnlyHint: true, @@ -78,206 +120,15 @@ export function registerHealthTool(server: McpServer): void { openWorldHint: false, }, }, - async ({ verbose }) => { - logger.info('health_check called', { verbose }); - - try { - const checks: Record> = {}; - let hasFailure = false; - let dbOk = false; - - // 1. Database connectivity - if (isDatabaseConfigured()) { - try { - const [connected, latencyMs] = await timed(checkDatabaseConnection); - dbOk = connected; - checks.database = { ok: connected }; - if (verbose) checks.database.latencyMs = latencyMs; - if (!connected) { - checks.database.error = 'Connection check failed'; - hasFailure = true; - } - } catch (err) { - dbOk = false; - hasFailure = true; - checks.database = { - ok: false, - error: err instanceof Error ? err.message : 'Connection failed', - }; - } - } else { - dbOk = false; - hasFailure = true; - checks.database = { ok: false, error: 'DATABASE_URL not configured' }; - } - - // 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; + async () => { + logger.info('health_check called'); - // 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 { - checks.documents.count = await getCount('documents'); - } - } catch (err) { - checks.documents = { - ok: false, - error: err instanceof Error ? err.message : 'Check failed', - }; - 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 { - const memOk = await checkTable('memory_entities'); - checks.memory = { ok: memOk }; - if (!memOk) { - checks.memory.error = "Table 'memory_entities' not accessible"; - hasFailure = true; - } else { - checks.memory.entities = await getCount('memory_entities'); - checks.memory.observations = await getCount('memory_observations'); - } - } catch (err) { - checks.memory = { - ok: false, - error: err instanceof Error ? err.message : 'Check failed', - }; - hasFailure = true; - } - } - - // Conversations (if enabled) - if (config.ENABLE_CONVERSATIONS) { - try { - const convOk = await checkTable('conversation_sessions'); - checks.conversations = { ok: convOk }; - if (!convOk) { - checks.conversations.error = "Table 'conversation_sessions' not accessible"; - hasFailure = true; - } else { - checks.conversations.sessions = await getCount('conversation_sessions'); - } - } catch (err) { - checks.conversations = { - ok: false, - error: err instanceof Error ? err.message : 'Check failed', - }; - hasFailure = true; - } - } - - // Insights (if enabled) - if (config.ENABLE_INSIGHTS) { - try { - const insOk = await checkTable('proactive_insights'); - checks.insights = { ok: insOk }; - if (!insOk) { - checks.insights.error = "Table 'proactive_insights' not accessible"; - hasFailure = true; - } else { - checks.insights.count = await getCount('proactive_insights'); - } - } catch (err) { - checks.insights = { - ok: false, - error: err instanceof Error ? err.message : 'Check failed', - }; - hasFailure = true; - } - - // Insight queue - try { - const queueOk = await checkTable('insight_queue'); - checks.insightQueue = { ok: queueOk }; - if (!queueOk) { - checks.insightQueue.error = "Table 'insight_queue' not accessible"; - hasFailure = true; - } else { - const row = await queryOne<{ chunks_pending: number }>( - 'SELECT chunks_pending FROM insight_queue WHERE id = 1', - ); - if (row) { - checks.insightQueue.pending = row.chunks_pending; - } - } - } catch (err) { - checks.insightQueue = { - ok: false, - error: err instanceof Error ? err.message : 'Check failed', - }; - hasFailure = true; - } - } - } - - // Overall status - let status: 'healthy' | 'degraded' | 'unhealthy'; - if (!dbOk) { - status = 'unhealthy'; - } else if (hasFailure) { - status = 'degraded'; - } else { - status = 'healthy'; - } - - const uptimeSeconds = Math.round((Date.now() - serverStartTime) / 1000); - const result = { - status, - checks, - version: pkg.version, - uptime: formatUptime(uptimeSeconds), - }; + if (!isDatabaseConfigured()) { + return configError('Database', 'Set DATABASE_URL'); + } + try { + const result = await handleHealthCheck(); return toolResponse({ compact: result, verbose: result,