Skip to content

fix(pi): repair MLflow SSE stream termination safely - #335

Open
dgokeeffe wants to merge 10 commits into
databricks:mainfrom
dgokeeffe:review/mlflow-sse-proxy
Open

fix(pi): repair MLflow SSE stream termination safely#335
dgokeeffe wants to merge 10 commits into
databricks:mainfrom
dgokeeffe:review/mlflow-sse-proxy

Conversation

@dgokeeffe

Copy link
Copy Markdown

Issue and stack

Closes #215.

Depends on #334 (and its prerequisite stack). Review only the four-file MLflow proxy range:

dgokeeffe/ucode@347f81d...8e78937

Do not merge before #334. The incremental range contains no capability-discovery or cache-diagnostic changes.


Unit 06 — MLflow SSE repair proxy

Objective and user-visible behavior

Let Pi consume otherwise valid MLflow SSE streams that omit terminal finish_reason, without converting upstream errors into success. Recover stale loopback configuration, bind a fresh port per launch, restore direct config before shutdown, and preserve primary failures/interrupt behavior.

Candidate: review/mlflow-sse-proxy (8e78937), mandatory PR target unit 05 (347f81d).

Exact scope

Production: new src/ucode/agents/_mlflow_proxy.py and lifecycle wiring in src/ucode/agents/pi.py.
Tests: new tests/test_mlflow_proxy.py and lifecycle coverage in tests/test_agent_pi.py.
Non-goals: model discovery/metadata, OpenCode, native providers, cache behavior, general-purpose HTTP proxying.

Before / after reproduction

git show 347f81d:src/ucode/agents/pi.py | rg '_start_oss_proxy|_mlflow_proxy' || echo '<absent>'
# <absent>

git show review/mlflow-sse-proxy:src/ucode/agents/pi.py | rg '_start_oss_proxy|_mlflow_proxy'
# loopback proxy import, startup, and lifecycle calls are present

Focused validation:

uv run --frozen pytest tests/test_mlflow_proxy.py tests/test_agent_pi.py -q
# 98 passed
uv run --frozen ruff check .
# All checks passed
uv run --frozen ruff format --check \
  src/ucode/agents/_mlflow_proxy.py src/ucode/agents/pi.py \
  tests/test_mlflow_proxy.py tests/test_agent_pi.py
# 4 files already formatted

Full suite:

uv run --frozen pytest -q
# 2 failed, 1793 passed, 36 skipped

The installed Claude capture and preferred-Pi-model state assertion reproduce on integrated dev and are outside this four-file incremental range.

Stream/lifecycle evidence

Loopback integration tests cover healthy finish passthrough, injected finish before [DONE], clean EOF recovery, transport-truncation non-success, multi-line SSE events, event errors, JSON error chunks, malformed SSE data, upstream HTTP error bytes/status, non-streaming passthrough, redirects, invalid targets/content lengths, disconnects, and no-body responses. Pi lifecycle tests cover no-OSS bypass, stale-port replacement, recursive-loopback prevention, first-write ordering, refresh-preserves-loopback behavior, serialized in-flight refresh/direct restoration, normal exit, SIGINT, setup failure, restore failure, server shutdown, and direct-config restoration. No request bodies, authorization headers, tokens, or live identifiers are logged.

Impact map

  • Agent: Pi only.
  • Provider: only databricks-mlflow; Claude/OpenAI/Gemini remain direct.
  • Config: MLflow base URL is session-temporarily loopback and restored to direct gateway before listener release.
  • Network/security: binds 127.0.0.1 on an ephemeral port, rejects absolute/foreign targets and redirects, strips hop-by-hop headers, and never logs credentials/bodies.
  • Managed config/MCP/skills: unchanged.

Rollback and residual risk

Revert the two unit-06 commits (2a73144, 8e78937); unit 05's model metadata/discovery remains. Models whose gateway streams omit terminal finish reasons may again fail in Pi. Residual risk is maintenance of a local HTTP boundary; tests intentionally bound accepted path, response repair conditions, shutdown, and error passthrough.

Hygiene

git diff --check review/oss-capability-discovery..review/mlflow-sse-proxy passes. Four-file range; no capability-discovery changes, cache scripts, generated files, credentials, uv.lock, .pi-subagents/, or goal.md; no merge markers or unresolved index entries.

dgokeeffe added 10 commits July 26, 2026 10:36
…directing HOME

Pi honors the PI_CODING_AGENT_DIR env var to resolve its config
directory (~/.pi/agent), so redirecting /Users/david.okeeffe to APP_DIR/pi-home was
unnecessary. The HOME redirect broke macOS keychain default resolution
under ucode: the Security framework looks for the login keychain under
the redirected HOME, finds none, and security default-keychain returns
'A default keychain could not be found'. As a result gh auth, the git
credential helper, and any keychain-backed tool failed inside pi.

Setting PI_CODING_AGENT_DIR to the existing PI_CONFIG_DIR preserves
config isolation (models.json/settings.json/sessions still land under
APP_DIR/pi-home/.pi/agent) while leaving /Users/david.okeeffe as the user's real home,
so the login keychain stays discoverable.

Tests: the two pi e2e sites monkeypatched PI_UCODE_HOME/PI_CONFIG_PATH
to redirect pi at a tmp home. They now also patch PI_CONFIG_DIR (read by
build_runtime_env) and PI_SETTINGS_PATH/PI_SETTINGS_BACKUP_PATH (previously
masked because the HOME redirect made pi read settings from the un-patched
real APP_DIR path).
Expose the GLM and Kimi coding-model cohort through Pi and OpenCode with shared token limits and reasoning metadata. Keep unsupported chat models out of discovery, including Inkling until gateway issue databricks#215 is fixed, and retain the GPT-OSS Responses API routing guard.
Centralize Claude family/version parsing so Pi metadata, adaptive-thinking compatibility, and Claude Code's [1m] selector cannot drift. Cover Sonnet 4.5, Opus 4.6, future major versions, Fable fallback, and prefixed model IDs.
Configure the Databricks OpenAI Responses provider alongside the validated GLM/Kimi provider, and fall back to foundation-model serving endpoints when UC model services are unavailable.
…/consolidated-pi-opencode-upstream

# Conflicts:
#	tests/test_e2e.py
`_pi_gpt_model_entry` declared `reasoning: True` without an off-state, so for
the thinking-off case Pi's Responses builder fell back to
`reasoning: {effort: "none"}` (pi-ai openai-responses.js, the
`thinkingLevelMap?.off !== null` branch). `"none"` is only valid on gpt-5.1+,
so every request to gpt-5, gpt-5-mini, gpt-5-nano and gpt-5-5-pro was rejected:

  BAD_REQUEST: Unsupported value: 'none' is not supported with the 'gpt-5'
  model. Supported values are: 'minimal', 'low', 'medium', and 'high'.

Setting `thinkingLevelMap: {"off": None}` makes Pi omit `reasoning` entirely,
which the gateway accepts for all 14 codex ids. Verified against
/ai-gateway/codex/v1/responses: effort="none" 400s on gpt-5/-mini/-nano/-5-5-pro
and 200s on gpt-5-1..-5-6; omitting `reasoning` is 200 everywhere.
`{"off": "minimal"}` was rejected as an alternative because gpt-5-5-pro 400s on
it too. Same pattern already used for the Gemini 3.x entries.

The rest of Pi's Responses payload was bisected against the gateway and is
fine: store:false, prompt_cache_key, prompt_cache_retention:"24h",
prompt_cache_options, include:["reasoning.encrypted_content"], developer role,
flat tool schemas, and the session_id / x-client-request-id affinity headers.

Regression was hard to spot because the gateway returns
{"error_code","message"} rather than OpenAI's {"error":...}, so Pi's
error-body.js recovery no-ops and every 400 renders as
"OpenAI API error (400): 400 status code (no body)". Reported upstream as
earendil-works/pi#7748.

Refs databricks#286
…i-opencode-upstream

# Conflicts:
#	src/ucode/agents/pi.py
#	src/ucode/databricks.py
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

AI Gateway MLflow chat-completions omits finish_reason on natural stop (breaks inkling in Pi)

1 participant