Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 22 additions & 8 deletions src/ucode/smart_routing/claude_pty.py
Original file line number Diff line number Diff line change
Expand Up @@ -187,8 +187,18 @@ def serve_first_prompt_socket(
stop: threading.Event,
*,
log: Callable[[str], None] = lambda _message: None,
ready: threading.Event | None = None,
) -> threading.Thread:
"""Serve the hook protocol, blocking exactly one non-command prompt."""
"""Serve the hook protocol, blocking exactly one non-command prompt.

Pass a ``threading.Event`` as *ready* to receive a reliable signal that the
socket is fully listening (i.e. after ``listen()``, not just ``bind()``).
On macOS the file-system path appears after ``bind()`` but connections are
only accepted after ``listen()``, so callers that poll ``path.exists()``
can connect before the server is ready. The *ready* event fires after
``listen()`` on success, or immediately on ``OSError`` so callers never
block forever on failure.
"""

def serve() -> None:
claimed = False
Expand All @@ -201,8 +211,12 @@ def serve() -> None:
server.settimeout(0.5)
except OSError as exc:
log(f"[ERR] first-prompt socket bind failed: {exc!r}")
if ready is not None:
ready.set() # unblock callers so they don't wait forever on failure
return
log(f"[READY] first-prompt socket {path}")
if ready is not None:
ready.set() # signal: listen() is done, connections are now accepted
try:
while not stop.is_set():
try:
Expand Down Expand Up @@ -317,14 +331,14 @@ def on_blocked_prompt(prompt: str, model: str) -> None:
pending["value"] = (prompt, model)
log(f"[ROUTE] first prompt -> {model!r}")

server_thread = serve_first_prompt_socket(
socket_path, route_prompt, on_blocked_prompt, stop, log=log
ready = threading.Event()
serve_first_prompt_socket(
socket_path, route_prompt, on_blocked_prompt, stop, log=log, ready=ready
)
socket_deadline = time.monotonic() + 2.0
while (
not socket_path.exists() and server_thread.is_alive() and time.monotonic() < socket_deadline
):
time.sleep(0.01)
# Wait for the socket to be listening (after listen(), not just bind()). On
# macOS bind() creates the file before listen() is called, so polling
# socket_path.exists() races. The ready event fires only after listen().
ready.wait(timeout=2.0)
if not socket_path.exists():
log("[ERR] first-prompt socket was not ready before Claude launch")
stop.set()
Expand Down
24 changes: 24 additions & 0 deletions tests/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
from __future__ import annotations

import os
import tempfile

import pytest

Expand Down Expand Up @@ -59,6 +60,29 @@ def reject_privileged_write(path, _desired_text):
databricks_mod.clear_model_services_cache()


@pytest.fixture()
def short_tmp_path():
"""A temporary directory with a short absolute path.

pytest's ``tmp_path`` fixture generates paths up to ~121 characters on
macOS (e.g. ``/private/var/folders/.../pytest-N/test_name0/``). Unix
domain sockets (``AF_UNIX``) on macOS have a hard path-length limit of
104 characters, so any ``.sock`` file placed inside ``tmp_path`` silently
raises ``OSError: AF_UNIX path too long``. This fixture uses
``tempfile.mkdtemp()`` which produces short paths like
``/var/folders/.../T/tmpXXXXXX`` (≤ 60 chars), safely inside the limit.
Use it in place of ``tmp_path`` whenever the test creates a Unix socket.
"""
import shutil
from pathlib import Path

d = tempfile.mkdtemp()
try:
yield Path(d)
finally:
shutil.rmtree(d, ignore_errors=True)


def _workspace() -> str:
ws = os.environ.get("UCODE_TEST_WORKSPACE", "").strip().rstrip("/")
return normalize_workspace_url(ws) if ws else ""
Expand Down
15 changes: 7 additions & 8 deletions tests/test_claude_smart_routing_v2.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,6 @@
import os
import sys
import threading
import time
from pathlib import Path

import pytest
Expand Down Expand Up @@ -122,22 +121,22 @@ def test_displays_catalog_name_while_retaining_routable_model(self):
assert "Selected Model : GLM 5.3 Flash" in result["reason"]
assert "anthropic-aigw-77df06ea" not in result["reason"]

def test_blocks_once_then_allows_replay(self, tmp_path):
socket_path = tmp_path / "first.sock"
def test_blocks_once_then_allows_replay(self, tmp_path, short_tmp_path):
socket_path = short_tmp_path / "first.sock"
blocked: list[tuple[str, str]] = []
stop = threading.Event()
ready = threading.Event()
claude_pty.serve_first_prompt_socket(
socket_path,
lambda _prompt: claude_pty.FirstPromptRoute(
model="sonnet", display_model="sonnet", rationale="Selected for a narrow task."
),
lambda prompt, model: blocked.append((prompt, model)),
stop,
ready=ready,
)
try:
deadline = time.monotonic() + 5
while not socket_path.exists() and time.monotonic() < deadline:
time.sleep(0.01)
ready.wait(timeout=5)
first = claude_pty.request_first_prompt_route(
socket_path, {"session_id": "s1", "prompt": "fix the parser"}
)
Expand Down Expand Up @@ -615,11 +614,11 @@ def is_alive():
socket_path=tmp_path / "missing.sock",
)

def test_direct_switch_restore_and_replay(self, tmp_path):
def test_direct_switch_restore_and_replay(self, tmp_path, short_tmp_path):
fake_claude = tmp_path / "fake_claude.py"
capture = tmp_path / "capture.json"
restored = tmp_path / "restored"
socket_path = tmp_path / "first.sock"
socket_path = short_tmp_path / "first.sock"
fake_claude.write_text(
"""
import json
Expand Down