feat(observability): MLflow session experiment tracking by dimakis · Pull Request #425 · dimakis/mitzo

dimakis · 2026-07-02T00:08:43Z

Summary

Each Mitzo session becomes an MLflow run in the mitzo-sessions experiment
Run tags capture session metadata (mode, model, branch, session ID)
Final metrics logged at session end: tokens (input/output/cache), cost, turns, duration, compactions
All MLflow calls are fire-and-forget no-ops when MLFLOW_TRACKING_URI is unset
Handles both normal completion and fallback paths (disconnect/abort/timeout)

Files

server/mlflow.ts — new REST client (experiment management, run create/end)
server/query-loop.ts — wired into session ID resolution + finally block
server/__tests__/mlflow.test.ts — 9 tests covering create, end, caching, errors, disabled mode
.env.example — added MLFLOW_TRACKING_URI

Design

Spec: mgmt/local_features/mlflow-session-tracking/spec.md

Every commit to an agent definition or boot context recipe becomes an independent variable in a continuous experiment, with token spend/cost/turns as dependent variables. Phase 1 tracks session-level metrics; Phase 2 adds per-turn stepping and context consumed artifacts.

Test plan

npx vitest run server/__tests__/mlflow.test.ts — 9/9 passing
Deploy with MLFLOW_TRACKING_URI=http://localhost:5050, run a session, verify run appears in MLflow UI

🤖 Generated with Claude Code

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 <noreply@anthropic.com>

dimakis

Centaur Review

Found 5 issue(s) (1 warning).

`server/query-loop.ts`

Clean, well-tested observability module. The main concern is a race between fire-and-forget createRun and synchronous mlflowRunId check in endRun — store and await the promise like the adjacent goalCreationPromise pattern.

🟡 bugs (L540): Race condition: createRun is fire-and-forget (promise not stored), but endRun in the finally block checks mlflowRunId synchronously. If the session ends before createRun resolves (fast abort, immediate error, slow MLflow server), mlflowRunId is still null and endRun never fires — orphaning the MLflow run in RUNNING state. The goalCreationPromise pattern 10 lines above shows the fix: store the promise and await it before endRun. This is especially relevant since createRun makes 3–4 HTTP requests with 5s timeouts each. [fixable]
🔵 bugs (L1437): totalCostUsd: lastReportedUsage.totalCostUsd || (finalSession?.cumulativeCostUsd ?? 0) uses || which treats 0 as falsy. A zero-cost session (legitimate SDK result with $0.00) falls through to the alternative source. Should use ?? for consistency with the existing fallback at line 1411 which uses ??. In practice both sources converge to the same value when cost is 0, so this won't corrupt data, but it's semantically incorrect. [fixable]
🔵 unsafe_assumptions (L553): .catch(() => {}) silently swallows errors from createRun at the call site. While mlflow.ts logs internally, a reader of query-loop.ts has no indication errors are handled. The goalCreationPromise pattern above uses .catch((err) => { log.warn(...); return null; }) — consider matching that pattern for consistency and debuggability. [fixable]

`server/mlflow.ts`

Clean, well-tested observability module. The main concern is a race between fire-and-forget createRun and synchronous mlflowRunId check in endRun — store and await the promise like the adjacent goalCreationPromise pattern.

🔵 style (L1): 14-line module-level doc comment violates the project convention: "Never write multi-paragraph docstrings or multi-line comment blocks — one short line max." The endpoints list is especially verbose — callers don't need to know the REST paths. Consider collapsing to a single-line summary or removing entirely since the module name and exports are self-documenting. [fixable]

`server/tests/mlflow.test.ts`

Clean, well-tested observability module. The main concern is a race between fire-and-forget createRun and synchronous mlflowRunId check in endRun — store and await the promise like the adjacent goalCreationPromise pattern.

🔵 missing_tests: No test for the isEnabled() === false path of endRun — the test covers runId === null (line 230) but not the case where isEnabled() is false AND runId is non-null (the early return at the top of endRun checks both). Minor gap since the implementation clearly short-circuits, but it would complete the disabled-state coverage. [fixable]

dimakis · 2026-07-04T01:00:54Z

            }
+
+            // Create MLflow run for this session (fire-and-forget)
+            if (mlflow.isEnabled()) {


🟡 bugs: Race condition: createRun is fire-and-forget (promise not stored), but endRun in the finally block checks mlflowRunId synchronously. If the session ends before createRun resolves (fast abort, immediate error, slow MLflow server), mlflowRunId is still null and endRun never fires — orphaning the MLflow run in RUNNING state. The goalCreationPromise pattern 10 lines above shows the fix: store the promise and await it before endRun. This is especially relevant since createRun makes 3–4 HTTP requests with 5s timeouts each. [fixable]

dimakis · 2026-07-04T01:00:54Z

+              cacheReadTokens: lastReportedUsage.cacheReadTokens,
+              cacheCreationTokens: lastReportedUsage.cacheCreationTokens,
+              totalCostUsd:
+                lastReportedUsage.totalCostUsd || (finalSession?.cumulativeCostUsd ?? 0),


🔵 bugs: totalCostUsd: lastReportedUsage.totalCostUsd || (finalSession?.cumulativeCostUsd ?? 0) uses || which treats 0 as falsy. A zero-cost session (legitimate SDK result with $0.00) falls through to the alternative source. Should use ?? for consistency with the existing fallback at line 1411 which uses ??. In practice both sources converge to the same value when cost is 0, so this won't corrupt data, but it's semantically incorrect. [fixable]

dimakis · 2026-07-04T01:00:54Z

+                  mlflowRunId = id;
+                })
+                .catch(() => {});
+            }


🔵 unsafe_assumptions: .catch(() => {}) silently swallows errors from createRun at the call site. While mlflow.ts logs internally, a reader of query-loop.ts has no indication errors are handled. The goalCreationPromise pattern above uses .catch((err) => { log.warn(...); return null; }) — consider matching that pattern for consistency and debuggability. [fixable]

dimakis · 2026-07-04T01:00:54Z

@@ -0,0 +1,225 @@
+/**


🔵 style: 14-line module-level doc comment violates the project convention: "Never write multi-paragraph docstrings or multi-line comment blocks — one short line max." The endpoints list is especially verbose — callers don't need to know the REST paths. Consider collapsing to a single-line summary or removing entirely since the module name and exports are self-documenting. [fixable]

dimakis

Centaur Review

Found 4 issue(s) (1 warning).

`server/query-loop.ts`

Clean MLflow integration following existing patterns (goal tracking, notification fire-and-forget). One real race condition: the fire-and-forget createRun can leave orphaned MLflow runs if the session ends before the promise resolves — store the promise and await it like goalCreationPromise.

🟡 bugs (L535): Race condition: createRun is fire-and-forget (.then() assigns mlflowRunId), but endRun in the finally block guards on if (mlflowRunId). If the session ends before createRun resolves (fast error, slow MLflow — up to 4 network hops × 5s timeout each), mlflowRunId is still null, so endRun is silently skipped and the MLflow run is orphaned in RUNNING state. The existing goalCreationPromise pattern solves this correctly: store the promise, then await it before consuming the result. Fix: let mlflowRunPromise: Promise<string | null> | undefined; → store the promise → await mlflowRunPromise in the finally block before calling endRun. [fixable]
🔵 unsafe_assumptions (L1438): totalCostUsd: lastReportedUsage.totalCostUsd || (finalSession?.cumulativeCostUsd ?? 0) uses ||, which treats a legitimate zero cost as falsy and falls through to the session accumulator. In practice both paths yield zero so the result is correct, but ?? would express the intent more precisely ("use fallback only when the value is null/undefined, not when it's zero"). [fixable]
🔵 style (L548): The .catch(() => {}) on createRun silently swallows errors. While createRun has internal try/catch that logs warnings and returns null, the outer catch would only fire for truly unexpected failures (e.g., in the .then() callback). The adjacent goalCreationPromise logs in its catch. Minor inconsistency — consider .catch(() => { /* logged internally */ }) or matching the goal pattern. [fixable]

`server/mlflow.ts`

Clean MLflow integration following existing patterns (goal tracking, notification fire-and-forget). One real race condition: the fire-and-forget createRun can leave orphaned MLflow runs if the session ends before the promise resolves — store the promise and await it like goalCreationPromise.

🔵 style (L1): The 14-line module-level JSDoc comment lists endpoints and describes behavior. CLAUDE.md says "default to writing no comments" and "never write multi-line comment blocks." The endpoint list is useful as a quick reference, but could be trimmed to a one-liner or moved to the CLAUDE.md architecture section where it's already partially documented. [fixable]

dimakis · 2026-07-04T09:41:54Z

@@ -533,6 +535,22 @@ async function _runQueryLoopInner(
                return null;


🟡 bugs: Race condition: createRun is fire-and-forget (.then() assigns mlflowRunId), but endRun in the finally block guards on if (mlflowRunId). If the session ends before createRun resolves (fast error, slow MLflow — up to 4 network hops × 5s timeout each), mlflowRunId is still null, so endRun is silently skipped and the MLflow run is orphaned in RUNNING state. The existing goalCreationPromise pattern solves this correctly: store the promise, then await it before consuming the result. Fix: let mlflowRunPromise: Promise<string | null> | undefined; → store the promise → await mlflowRunPromise in the finally block before calling endRun. [fixable]

dimakis · 2026-07-04T09:41:54Z

+              cacheCreationTokens: lastReportedUsage.cacheCreationTokens,
+              totalCostUsd:
+                lastReportedUsage.totalCostUsd || (finalSession?.cumulativeCostUsd ?? 0),
+              numTurns: doneSent ? lastReportedUsage.numTurns : turnIndex,


🔵 unsafe_assumptions: totalCostUsd: lastReportedUsage.totalCostUsd || (finalSession?.cumulativeCostUsd ?? 0) uses ||, which treats a legitimate zero cost as falsy and falls through to the session accumulator. In practice both paths yield zero so the result is correct, but ?? would express the intent more precisely ("use fallback only when the value is null/undefined, not when it's zero"). [fixable]

dimakis · 2026-07-04T09:41:54Z

+                  model: currentSession.model,
+                  cwd: currentSession.cwd,
+                  branch: currentSession.branch,
+                })


🔵 style: The .catch(() => {}) on createRun silently swallows errors. While createRun has internal try/catch that logs warnings and returns null, the outer catch would only fire for truly unexpected failures (e.g., in the .then() callback). The adjacent goalCreationPromise logs in its catch. Minor inconsistency — consider .catch(() => { /* logged internally */ }) or matching the goal pattern. [fixable]

dimakis · 2026-07-04T09:41:54Z

@@ -0,0 +1,225 @@
+/**


🔵 style: The 14-line module-level JSDoc comment lists endpoints and describes behavior. CLAUDE.md says "default to writing no comments" and "never write multi-line comment blocks." The endpoint list is useful as a quick reference, but could be trimmed to a one-liner or moved to the CLAUDE.md architecture section where it's already partially documented. [fixable]

dimakis

Centaur Review

Found 5 issue(s) (1 warning).

`server/query-loop.ts`

Clean, well-tested MLflow integration. The main concern is a race condition where createRun is fire-and-forget but endRun depends on its result — short sessions can orphan MLflow runs. Follow the existing goalCreationPromise pattern to store and await the promise.

🟡 bugs (L539): Race condition: createRun is fire-and-forget — mlflowRunId is set inside a .then() callback, but endRun at line 1426 checks if (mlflowRunId) synchronously. If the session ends before createRun resolves (fast error, immediate disconnect), mlflowRunId is still null, the endRun is skipped, and the MLflow run is orphaned in RUNNING state forever. The codebase already has a pattern for this: goalCreationPromise (line 529) stores the promise and awaits it at session end (line 602). Store the createRun promise similarly and await it before the endRun block. [fixable]
🔵 bugs (L1436): lastReportedUsage.totalCostUsd || (finalSession?.cumulativeCostUsd ?? 0) uses ||, which treats a legitimate 0 cost the same as "no data reported." A zero-cost session (e.g. cached response) would fall through to the cumulativeCostUsd fallback. Use ?? instead of || for numeric fallback to distinguish zero from absent. [fixable]

`server/mlflow.ts`

Clean, well-tested MLflow integration. The main concern is a race condition where createRun is fire-and-forget but endRun depends on its result — short sessions can orphan MLflow runs. Follow the existing goalCreationPromise pattern to store and await the promise.

🔵 unsafe_assumptions (L20): TRACKING_URI is captured once at module load time. This is consistent with other env-based config in the codebase (e.g. OTEL_EXPORTER_OTLP_ENDPOINT), so it's fine — but worth noting that isEnabled() can't change at runtime without a server restart.
🔵 style (L1): The module-level doc block is 14 lines long. Per project conventions: "default to writing no comments" and "never write multi-paragraph docstrings or multi-line comment blocks — one short line max." Consider trimming to one line or removing entirely since the function names and types are self-documenting. [fixable]

`server/tests/mlflow.test.ts`

Clean, well-tested MLflow integration. The main concern is a race condition where createRun is fire-and-forget but endRun depends on its result — short sessions can orphan MLflow runs. Follow the existing goalCreationPromise pattern to store and await the promise.

🔵 missing_tests: No test for endRun when the MLflow API call fails (e.g., log-batch returns 500). The createRun error path is tested, but endRun's try/catch (which swallows errors and logs a warning) is not exercised. Adding a test would confirm the error is gracefully handled and doesn't throw. [fixable]

dimakis · 2026-07-04T11:25:53Z

              });
            }
+
+            // Create MLflow run for this session (fire-and-forget)


🟡 bugs: Race condition: createRun is fire-and-forget — mlflowRunId is set inside a .then() callback, but endRun at line 1426 checks if (mlflowRunId) synchronously. If the session ends before createRun resolves (fast error, immediate disconnect), mlflowRunId is still null, the endRun is skipped, and the MLflow run is orphaned in RUNNING state forever. The codebase already has a pattern for this: goalCreationPromise (line 529) stores the promise and awaits it at session end (line 602). Store the createRun promise similarly and await it before the endRun block. [fixable]

dimakis · 2026-07-04T11:25:53Z

+              outputTokens: doneSent ? lastReportedUsage.outputTokens : cumulativeOutputTokens,
+              cacheReadTokens: lastReportedUsage.cacheReadTokens,
+              cacheCreationTokens: lastReportedUsage.cacheCreationTokens,
+              totalCostUsd:


🔵 bugs: lastReportedUsage.totalCostUsd || (finalSession?.cumulativeCostUsd ?? 0) uses ||, which treats a legitimate 0 cost the same as "no data reported." A zero-cost session (e.g. cached response) would fall through to the cumulativeCostUsd fallback. Use ?? instead of || for numeric fallback to distinguish zero from absent. [fixable]

dimakis · 2026-07-04T11:25:53Z

+
+const log = createLogger('mlflow');
+
+const TRACKING_URI = process.env.MLFLOW_TRACKING_URI ?? '';


🔵 unsafe_assumptions: TRACKING_URI is captured once at module load time. This is consistent with other env-based config in the codebase (e.g. OTEL_EXPORTER_OTLP_ENDPOINT), so it's fine — but worth noting that isEnabled() can't change at runtime without a server restart.

dimakis · 2026-07-04T11:25:53Z

@@ -0,0 +1,225 @@
+/**


🔵 style: The module-level doc block is 14 lines long. Per project conventions: "default to writing no comments" and "never write multi-paragraph docstrings or multi-line comment blocks — one short line max." Consider trimming to one line or removing entirely since the function names and types are self-documenting. [fixable]

dimakis force-pushed the feat/mlflow-session-tracking branch from 8a10522 to 8243b82 Compare July 4, 2026 00:56

dimakis commented Jul 4, 2026

View reviewed changes

		@@ -533,6 +535,22 @@ async function _runQueryLoopInner(
		return null;


		const log = createLogger('mlflow');

		const TRACKING_URI = process.env.MLFLOW_TRACKING_URI ?? '';

Conversation

dimakis commented Jul 2, 2026

Summary

Files

Design

Test plan

Uh oh!

dimakis left a comment

Choose a reason for hiding this comment

Centaur Review

server/query-loop.ts

server/mlflow.ts

server/__tests__/mlflow.test.ts

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

dimakis left a comment

Choose a reason for hiding this comment

Centaur Review

server/query-loop.ts

server/mlflow.ts

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

dimakis left a comment

Choose a reason for hiding this comment

Centaur Review

server/query-loop.ts

server/mlflow.ts

server/__tests__/mlflow.test.ts

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

`server/query-loop.ts`

`server/mlflow.ts`

`server/tests/mlflow.test.ts`

`server/query-loop.ts`

`server/mlflow.ts`

`server/query-loop.ts`

`server/mlflow.ts`

`server/tests/mlflow.test.ts`