From 8243b821e3de24d6913e42004850916960d29598 Mon Sep 17 00:00:00 2001 From: dimakis Date: Sat, 4 Jul 2026 01:56:10 +0100 Subject: [PATCH] feat(observability): MLflow session tracking Each Mitzo session becomes an MLflow run in a "mitzo-sessions" experiment. Logs token usage, cost, turns, duration, and compactions as metrics. Fire-and-forget with graceful no-op when MLFLOW_TRACKING_URI is unset. Co-Authored-By: Claude Opus 4.6 --- .env.example | 3 + server/__tests__/mlflow.test.ts | 308 ++++++++++++++++++++++++++++++++ server/mlflow.ts | 225 +++++++++++++++++++++++ server/query-loop.ts | 41 +++++ 4 files changed, 577 insertions(+) create mode 100644 server/__tests__/mlflow.test.ts create mode 100644 server/mlflow.ts diff --git a/.env.example b/.env.example index 7c403046..09df8e5a 100644 --- a/.env.example +++ b/.env.example @@ -18,3 +18,6 @@ BASE_URL= # ContexGin Goal Registry (optional — tokens-to-goal tracking) CONTEXGIN_URL=http://localhost:8321 + +# MLflow experiment tracking (optional — session-level metrics) +MLFLOW_TRACKING_URI=http://localhost:5050 diff --git a/server/__tests__/mlflow.test.ts b/server/__tests__/mlflow.test.ts new file mode 100644 index 00000000..f226b49c --- /dev/null +++ b/server/__tests__/mlflow.test.ts @@ -0,0 +1,308 @@ +import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest'; + +// Stub fetch globally before importing the module +const mockFetch = vi.fn(); +vi.stubGlobal('fetch', mockFetch); + +// Must set env before importing +vi.stubEnv('MLFLOW_TRACKING_URI', 'http://localhost:5050'); + +describe('mlflow session tracking', () => { + let mlflow: typeof import('../mlflow.js'); + + beforeEach(async () => { + vi.resetModules(); + mockFetch.mockReset(); + vi.stubEnv('MLFLOW_TRACKING_URI', 'http://localhost:5050'); + mlflow = await import('../mlflow.js'); + mlflow._resetCache(); + }); + + afterEach(() => { + vi.unstubAllEnvs(); + }); + + it('isEnabled returns true when MLFLOW_TRACKING_URI is set', () => { + expect(mlflow.isEnabled()).toBe(true); + }); + + describe('createRun', () => { + it('creates experiment and run on first call', async () => { + // First call: search returns empty + mockFetch.mockResolvedValueOnce({ + ok: true, + json: () => Promise.resolve({ experiments: [] }), + }); + // Second call: create experiment + mockFetch.mockResolvedValueOnce({ + ok: true, + json: () => Promise.resolve({ experiment_id: '1' }), + }); + // Third call: create run + mockFetch.mockResolvedValueOnce({ + ok: true, + json: () => + Promise.resolve({ + run: { info: { run_id: 'run-abc' } }, + }), + }); + // Fourth call: log-batch params + mockFetch.mockResolvedValueOnce({ + ok: true, + json: () => Promise.resolve({}), + }); + + const runId = await mlflow.createRun({ + sessionId: 'sess-123', + mode: 'conversational', + model: 'opus', + cwd: '/repo', + branch: 'main', + }); + + expect(runId).toBe('run-abc'); + expect(mockFetch).toHaveBeenCalledTimes(4); + + // Verify experiment search + const searchCall = mockFetch.mock.calls[0]; + expect(searchCall[0]).toBe('http://localhost:5050/api/2.0/mlflow/experiments/search'); + + // Verify experiment create + const createExpCall = mockFetch.mock.calls[1]; + expect(createExpCall[0]).toBe('http://localhost:5050/api/2.0/mlflow/experiments/create'); + + // Verify run create + const createRunCall = mockFetch.mock.calls[2]; + expect(createRunCall[0]).toBe('http://localhost:5050/api/2.0/mlflow/runs/create'); + const runBody = JSON.parse(createRunCall[1].body); + expect(runBody.experiment_id).toBe('1'); + expect(runBody.run_name).toBe('sess-123'); + expect(runBody.tags).toEqual( + expect.arrayContaining([ + { key: 'session.id', value: 'sess-123' }, + { key: 'session.mode', value: 'conversational' }, + { key: 'session.model', value: 'opus' }, + ]), + ); + + // Verify params logged + const paramsCall = mockFetch.mock.calls[3]; + const paramsBody = JSON.parse(paramsCall[1].body); + expect(paramsBody.run_id).toBe('run-abc'); + expect(paramsBody.params).toEqual( + expect.arrayContaining([ + { key: 'cwd', value: '/repo' }, + { key: 'branch', value: 'main' }, + ]), + ); + }); + + it('reuses cached experiment ID on subsequent calls', async () => { + // First call: search finds existing experiment + mockFetch.mockResolvedValueOnce({ + ok: true, + json: () => + Promise.resolve({ + experiments: [{ experiment_id: '42', name: 'mitzo-sessions' }], + }), + }); + // Create run + mockFetch.mockResolvedValueOnce({ + ok: true, + json: () => + Promise.resolve({ + run: { info: { run_id: 'run-1' } }, + }), + }); + // Log params + mockFetch.mockResolvedValueOnce({ + ok: true, + json: () => Promise.resolve({}), + }); + + await mlflow.createRun({ sessionId: 'sess-1' }); + + // Second call should skip experiment search (cached) + mockFetch.mockResolvedValueOnce({ + ok: true, + json: () => + Promise.resolve({ + run: { info: { run_id: 'run-2' } }, + }), + }); + mockFetch.mockResolvedValueOnce({ + ok: true, + json: () => Promise.resolve({}), + }); + + const runId = await mlflow.createRun({ sessionId: 'sess-2' }); + expect(runId).toBe('run-2'); + // 3 calls for first run + 2 for second (no experiment search) + expect(mockFetch).toHaveBeenCalledTimes(5); + }); + + it('returns null on API error', async () => { + mockFetch.mockResolvedValueOnce({ + ok: false, + status: 500, + text: () => Promise.resolve('Internal Server Error'), + }); + + const runId = await mlflow.createRun({ sessionId: 'sess-fail' }); + expect(runId).toBeNull(); + }); + }); + + describe('endRun', () => { + it('logs metrics and updates run status', async () => { + // log-batch + mockFetch.mockResolvedValueOnce({ + ok: true, + json: () => Promise.resolve({}), + }); + // update + mockFetch.mockResolvedValueOnce({ + ok: true, + json: () => Promise.resolve({}), + }); + + await mlflow.endRun( + 'run-abc', + { + inputTokens: 1000, + outputTokens: 500, + cacheReadTokens: 200, + cacheCreationTokens: 100, + totalCostUsd: 0.05, + numTurns: 3, + durationMs: 30000, + durationApiMs: 25000, + numCompactions: 1, + }, + 'FINISHED', + ); + + expect(mockFetch).toHaveBeenCalledTimes(2); + + // Verify metrics batch + const batchCall = mockFetch.mock.calls[0]; + expect(batchCall[0]).toBe('http://localhost:5050/api/2.0/mlflow/runs/log-batch'); + const batchBody = JSON.parse(batchCall[1].body); + expect(batchBody.run_id).toBe('run-abc'); + expect(batchBody.metrics).toEqual( + expect.arrayContaining([ + expect.objectContaining({ key: 'tokens.input', value: 1000 }), + expect.objectContaining({ key: 'tokens.output', value: 500 }), + expect.objectContaining({ key: 'cost_usd', value: 0.05 }), + expect.objectContaining({ key: 'num_turns', value: 3 }), + expect.objectContaining({ key: 'num_compactions', value: 1 }), + ]), + ); + + // Verify run update + const updateCall = mockFetch.mock.calls[1]; + expect(updateCall[0]).toBe('http://localhost:5050/api/2.0/mlflow/runs/update'); + const updateBody = JSON.parse(updateCall[1].body); + expect(updateBody.run_id).toBe('run-abc'); + expect(updateBody.status).toBe('FINISHED'); + expect(updateBody.end_time).toBeGreaterThan(0); + }); + + it('skips compaction metric when undefined', async () => { + mockFetch.mockResolvedValueOnce({ + ok: true, + json: () => Promise.resolve({}), + }); + mockFetch.mockResolvedValueOnce({ + ok: true, + json: () => Promise.resolve({}), + }); + + await mlflow.endRun('run-abc', { + inputTokens: 0, + outputTokens: 0, + cacheReadTokens: 0, + cacheCreationTokens: 0, + totalCostUsd: 0, + numTurns: 0, + durationMs: 0, + durationApiMs: 0, + }); + + const batchBody = JSON.parse(mockFetch.mock.calls[0][1].body); + const keys = batchBody.metrics.map((m: { key: string }) => m.key); + expect(keys).not.toContain('num_compactions'); + }); + + it('is a no-op when runId is null', async () => { + await mlflow.endRun(null, { + inputTokens: 0, + outputTokens: 0, + cacheReadTokens: 0, + cacheCreationTokens: 0, + totalCostUsd: 0, + numTurns: 0, + durationMs: 0, + durationApiMs: 0, + }); + + expect(mockFetch).not.toHaveBeenCalled(); + }); + + it('sets FAILED status on error sessions', async () => { + mockFetch.mockResolvedValueOnce({ + ok: true, + json: () => Promise.resolve({}), + }); + mockFetch.mockResolvedValueOnce({ + ok: true, + json: () => Promise.resolve({}), + }); + + await mlflow.endRun( + 'run-err', + { + inputTokens: 0, + outputTokens: 0, + cacheReadTokens: 0, + cacheCreationTokens: 0, + totalCostUsd: 0, + numTurns: 0, + durationMs: 0, + durationApiMs: 0, + }, + 'FAILED', + ); + + const updateBody = JSON.parse(mockFetch.mock.calls[1][1].body); + expect(updateBody.status).toBe('FAILED'); + }); + }); +}); + +describe('mlflow disabled', () => { + it('is a no-op when MLFLOW_TRACKING_URI is not set', async () => { + vi.resetModules(); + mockFetch.mockReset(); + vi.stubEnv('MLFLOW_TRACKING_URI', ''); + + const mlflow = await import('../mlflow.js'); + expect(mlflow.isEnabled()).toBe(false); + + const runId = await mlflow.createRun({ sessionId: 'test' }); + expect(runId).toBeNull(); + expect(mockFetch).not.toHaveBeenCalled(); + + await mlflow.endRun('run-id', { + inputTokens: 0, + outputTokens: 0, + cacheReadTokens: 0, + cacheCreationTokens: 0, + totalCostUsd: 0, + numTurns: 0, + durationMs: 0, + durationApiMs: 0, + }); + expect(mockFetch).not.toHaveBeenCalled(); + }); +}); diff --git a/server/mlflow.ts b/server/mlflow.ts new file mode 100644 index 00000000..7f1ad63f --- /dev/null +++ b/server/mlflow.ts @@ -0,0 +1,225 @@ +/** + * MLflow session tracking — lightweight REST client. + * + * Each Mitzo session becomes an MLflow run in the "mitzo-sessions" experiment. + * All functions are no-ops when MLFLOW_TRACKING_URI is not set, so callers + * never need to guard. + * + * Endpoints used: + * POST /api/2.0/mlflow/experiments/search + * POST /api/2.0/mlflow/experiments/create + * POST /api/2.0/mlflow/runs/create + * POST /api/2.0/mlflow/runs/log-batch + * POST /api/2.0/mlflow/runs/update + */ + +import { createLogger } from './logger.js'; + +const log = createLogger('mlflow'); + +const TRACKING_URI = process.env.MLFLOW_TRACKING_URI ?? ''; +const EXPERIMENT_NAME = 'mitzo-sessions'; + +let cachedExperimentId: string | null = null; + +export function isEnabled(): boolean { + return TRACKING_URI.length > 0; +} + +// --------------------------------------------------------------------------- +// Internal HTTP helper +// --------------------------------------------------------------------------- + +async function post(path: string, body: Record): Promise { + const url = `${TRACKING_URI}${path}`; + const res = await fetch(url, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify(body), + signal: AbortSignal.timeout(5000), + }); + if (!res.ok) { + const text = await res.text().catch(() => ''); + throw new Error(`MLflow ${path} failed: ${res.status} ${text}`); + } + return res.json() as Promise; +} + +// --------------------------------------------------------------------------- +// Experiment management +// --------------------------------------------------------------------------- + +async function ensureExperiment(): Promise { + if (cachedExperimentId) return cachedExperimentId; + + // Search for existing experiment + interface SearchResult { + experiments?: Array<{ experiment_id: string; name: string }>; + } + const search = await post('/api/2.0/mlflow/experiments/search', { + max_results: 1, + filter: `name = '${EXPERIMENT_NAME}'`, + }); + if (search.experiments?.length) { + cachedExperimentId = search.experiments[0].experiment_id; + return cachedExperimentId; + } + + // Create new experiment + interface CreateResult { + experiment_id: string; + } + const created = await post('/api/2.0/mlflow/experiments/create', { + name: EXPERIMENT_NAME, + }); + cachedExperimentId = created.experiment_id; + log.info('created MLflow experiment', { name: EXPERIMENT_NAME, id: cachedExperimentId }); + return cachedExperimentId; +} + +// --------------------------------------------------------------------------- +// Public API +// --------------------------------------------------------------------------- + +export interface SessionParams { + sessionId: string; + mode?: string; + model?: string; + cwd?: string; + branch?: string; + clientType?: string; + agentDefHash?: string; + recipeHash?: string; +} + +export interface SessionMetrics { + inputTokens: number; + outputTokens: number; + cacheReadTokens: number; + cacheCreationTokens: number; + totalCostUsd: number; + numTurns: number; + durationMs: number; + durationApiMs: number; + numCompactions?: number; +} + +/** + * Create an MLflow run for a session. Returns the run_id, or null if disabled. + */ +export async function createRun(params: SessionParams): Promise { + if (!isEnabled()) return null; + + try { + const experimentId = await ensureExperiment(); + + interface RunResult { + run: { info: { run_id: string } }; + } + + const tags = [ + { key: 'session.id', value: params.sessionId }, + ...(params.mode ? [{ key: 'session.mode', value: params.mode }] : []), + ...(params.model ? [{ key: 'session.model', value: params.model }] : []), + ...(params.clientType ? [{ key: 'session.client_type', value: params.clientType }] : []), + ...(params.agentDefHash + ? [{ key: 'context.agent_def_hash', value: params.agentDefHash }] + : []), + ...(params.recipeHash ? [{ key: 'context.recipe_hash', value: params.recipeHash }] : []), + ]; + + const result = await post('/api/2.0/mlflow/runs/create', { + experiment_id: experimentId, + run_name: params.sessionId, + start_time: Date.now(), + tags, + }); + + const runId = result.run.info.run_id; + + // Log params (immutable session metadata) + const runParams = [ + { key: 'cwd', value: params.cwd ?? '' }, + { key: 'mode', value: params.mode ?? '' }, + { key: 'model', value: params.model ?? '' }, + { key: 'branch', value: params.branch ?? '' }, + ]; + + await post('/api/2.0/mlflow/runs/log-batch', { + run_id: runId, + params: runParams, + }); + + log.info('created MLflow run', { runId, sessionId: params.sessionId }); + return runId; + } catch (err) { + log.warn('failed to create MLflow run', { + error: err instanceof Error ? err.message : String(err), + }); + return null; + } +} + +/** + * End an MLflow run, logging final session metrics. + */ +export async function endRun( + runId: string | null, + metrics: SessionMetrics, + status: 'FINISHED' | 'FAILED' = 'FINISHED', +): Promise { + if (!isEnabled() || !runId) return; + + try { + const now = Date.now(); + + // Log all metrics in a single batch + const mlMetrics = [ + { key: 'tokens.input', value: metrics.inputTokens, timestamp: now, step: 0 }, + { key: 'tokens.output', value: metrics.outputTokens, timestamp: now, step: 0 }, + { key: 'tokens.cache_read', value: metrics.cacheReadTokens, timestamp: now, step: 0 }, + { + key: 'tokens.cache_creation', + value: metrics.cacheCreationTokens, + timestamp: now, + step: 0, + }, + { key: 'cost_usd', value: metrics.totalCostUsd, timestamp: now, step: 0 }, + { key: 'num_turns', value: metrics.numTurns, timestamp: now, step: 0 }, + { key: 'duration_ms', value: metrics.durationMs, timestamp: now, step: 0 }, + { key: 'duration_api_ms', value: metrics.durationApiMs, timestamp: now, step: 0 }, + ]; + + if (metrics.numCompactions !== undefined) { + mlMetrics.push({ + key: 'num_compactions', + value: metrics.numCompactions, + timestamp: now, + step: 0, + }); + } + + await post('/api/2.0/mlflow/runs/log-batch', { + run_id: runId, + metrics: mlMetrics, + }); + + await post('/api/2.0/mlflow/runs/update', { + run_id: runId, + status, + end_time: now, + }); + + log.info('ended MLflow run', { runId, status }); + } catch (err) { + log.warn('failed to end MLflow run', { + runId, + error: err instanceof Error ? err.message : String(err), + }); + } +} + +// For testing — reset cached experiment ID +export function _resetCache(): void { + cachedExperimentId = null; +} diff --git a/server/query-loop.ts b/server/query-loop.ts index 3b2f21ba..7bdc04e5 100644 --- a/server/query-loop.ts +++ b/server/query-loop.ts @@ -18,6 +18,7 @@ import { sendTurnCompleteNotification as apnsTurnComplete } from './apns.js'; import { extractSnippet } from './notification-helpers.js'; import { NOTIFY_SNIPPET_MAX_CHARS } from './constants.js'; import { createGoal, reportUsage, deriveGoalTitle } from './goal-client.js'; +import * as mlflow from './mlflow.js'; import { tracer } from './tracing.js'; import { context, trace, SpanStatusCode, type Span } from '@opentelemetry/api'; import { ProgressTracker } from './progress-tracker.js'; @@ -268,6 +269,7 @@ async function _runQueryLoopInner( let resolvedGoalId: string | undefined; let goalCreationPromise: Promise | undefined; let goalTitle: string | undefined; + let mlflowRunId: string | null = null; // Turn and tool span tracking let currentTurnSpan: Span | null = null; @@ -533,6 +535,22 @@ async function _runQueryLoopInner( return null; }); } + + // Create MLflow run for this session (fire-and-forget) + if (mlflow.isEnabled()) { + mlflow + .createRun({ + sessionId: resolvedSessionId, + mode: currentSession.mode, + model: currentSession.model, + cwd: currentSession.cwd, + branch: currentSession.branch, + }) + .then((id) => { + mlflowRunId = id; + }) + .catch(() => {}); + } } } else if (msg.type === 'result') { log.info('result received', { clientId, sessionId: msg.session_id }); @@ -1403,6 +1421,29 @@ async function _runQueryLoopInner( reason: caughtError ? 'error' : 'completed', }); } + // End MLflow run with final metrics (fire-and-forget). + // Uses the same usage data sources as recordUsage above. + if (mlflowRunId) { + const fallbackDurationMs = Date.now() - sessionStartedAt; + mlflow + .endRun( + mlflowRunId, + { + inputTokens: lastReportedUsage.inputTokens, + outputTokens: doneSent ? lastReportedUsage.outputTokens : cumulativeOutputTokens, + cacheReadTokens: lastReportedUsage.cacheReadTokens, + cacheCreationTokens: lastReportedUsage.cacheCreationTokens, + totalCostUsd: + lastReportedUsage.totalCostUsd || (finalSession?.cumulativeCostUsd ?? 0), + numTurns: doneSent ? lastReportedUsage.numTurns : turnIndex, + durationMs: doneSent ? lastReportedUsage.durationMs : fallbackDurationMs, + durationApiMs: lastReportedUsage.durationApiMs, + numCompactions, + }, + caughtError ? 'FAILED' : 'FINISHED', + ) + .catch(() => {}); + } // Clean up any open subagent spans (ERROR if catch was entered) const subagentCleanupStatus = caughtError ? SpanStatusCode.ERROR : SpanStatusCode.OK; for (const [, sub] of activeSubagents) {