fix(pi): repair MLflow SSE stream termination safely

[dgokeeffe]

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.

[usefactorapp]

Change

This PR establishes a unified AI agent framework within the ucode package,
integrating multiple agent implementations (Claude, OpenCode, Pi) with
Databricks infrastructure. It implements a comprehensive model discovery system
(Foundation Models, Bedrock, Gemini, OSS), a centralized configuration lifecycle
in src/ucode/agents/init.py, and a specialized proxy for MLflow tracing via
_mlflow_proxy.py.

Impact

The change introduces an extensive operational surface with 3,890 cross-service
boundary crossings, reflecting deep integration with Databricks workspace
management, authentication, and model provider services.

• Infrastructure: Depends on Databricks CLI and SCIM APIs for user identity
  (_scim_me) and workspace administration.
• Tracing: Implements an asynchronous server in _mlflow_proxy.py to relay
  Server-Sent Events (SSE) and handle POST requests, creating a dependency on
  the system's local network loopback for agent tracing.
• Agent Lifecycle: All agents utilize a common configuration and launch
  sequence governed by ensure_tool_binary_available, configure_tool, and
  launch.
• Upstream: 265 unique callers (fan-in) target the new logic, primarily
  through the agent initialization and model resolution interfaces.

Evidence

• Scale: 1,335 symbols were added or modified across 6 primary modules.
• Discovery Depth: The Databricks integration supports granular model
  categorization (Reasoning, OSS, GPT, Claude) and environment-specific
  constraints like token limits and context window parsing.
• Structural Branching: 43 downstream execution paths (fan-out) emerge from
  the core agent logic, covering model mapping, tool configuration, and
  authentication flows.

Validation

Deterministic validation is established for the following behaviors:

• Agent Pi Overlays: test_openai_provider_uses_openai_responses verifies the
  rendering of provider-specific responses.
• Model Constraints: test_gpt56_sol_model_entry_pins_1m_context validates that
  specific model configurations correctly apply 1M context suffixes.
• Lifecycle Entry Points: Validation evidence exists for the core launch,
  configure_tool, and ensure_tool_binary_available paths within the agent
  management layer.

Unverified

• No validation evidence was found for the _mlflow_proxy server's handling of
  specific SSE relay failures or verbatim response forwarding.
• One validation gap was identified in the discovery metrics for the
  src.ucode.agents namespace.
• Deterministic evidence of validation is missing for the Databricks budget
  resolution logic (resolve_current_budget_spend).

[usefactorapp]

If Factor’s analysis was useful to you and you’d like to use it across ucode, feel free to reach out at jishnu@tryfactor.app. Happy to discuss getting it set up for the repo.