diff --git a/studio/backend/cloudflare_tunnel.py b/studio/backend/cloudflare_tunnel.py index b1ddc74c32..78fce0c70a 100644 --- a/studio/backend/cloudflare_tunnel.py +++ b/studio/backend/cloudflare_tunnel.py @@ -20,6 +20,7 @@ import shutil import subprocess import sys import threading +import time from pathlib import Path from typing import Optional, Tuple @@ -40,6 +41,22 @@ _RELEASE_BASE = "https://github.com/cloudflare/cloudflared/releases/latest/downl _READY_TIMEOUT = 15.0 # seconds to wait for the URL + a registered edge connection _DOWNLOAD_TIMEOUT = 60 # urlopen timeout for the one-time binary download +# A registered edge connection does not mean the hostname resolves yet, so the +# URL is fetched once before it is advertised. +_PUBLIC_PROBE_PATH = "/api/health" +_PUBLIC_PROBE_MARKER = "Unsloth UI Backend" +# One deadline for DNS propagation + the health probe, bounding the startup stall. +_PUBLIC_PROBE_TIMEOUT = 45.0 +_PUBLIC_PROBE_ATTEMPT_TIMEOUT = 5.0 +_PUBLIC_PROBE_RETRY_DELAY = 1.0 + +# Wait for the hostname via DoH first: an early OS lookup negative-caches the +# NXDOMAIN for up to 30 min. +_DNS_POLL_DELAY = 2.0 +# Retry transient DoH failures, but give up fast when DoH is blocked outright. +_DNS_MAX_DOH_ERRORS = 3 +_DOH_URL = "https://cloudflare-dns.com/dns-query?name={host}&type=A" + def _windows_hidden_kwargs() -> dict: """Suppress a child console window on Windows; no-op elsewhere.""" @@ -191,6 +208,59 @@ def ensure_cloudflared() -> Optional[str]: return None +def _wait_for_dns(host: str, deadline: float) -> None: + import json + import urllib.request + + errors = 0 + while True: + answered = False + try: + req = urllib.request.Request( + _DOH_URL.format(host = host), + headers = {"Accept": "application/dns-json", "User-Agent": "unsloth-studio"}, + ) + with urllib.request.urlopen(req, timeout = 5) as response: + answered = bool(json.loads(response.read(65536)).get("Answer")) + errors = 0 + except Exception: + errors += 1 + if errors >= _DNS_MAX_DOH_ERRORS: + return + if answered: + return + remaining = deadline - time.monotonic() + if remaining <= 0: + return + time.sleep(min(_DNS_POLL_DELAY, remaining)) + + +def verify_public_url(url: str, timeout: float = _PUBLIC_PROBE_TIMEOUT) -> bool: + import json + import urllib.request + from urllib.parse import urlsplit + + deadline = time.monotonic() + timeout + host = urlsplit(url).hostname + if host: + _wait_for_dns(host, deadline) + + probe_url = f"{url.rstrip('/')}{_PUBLIC_PROBE_PATH}" + while True: + try: + req = urllib.request.Request(probe_url, headers = {"User-Agent": "unsloth-studio"}) + with urllib.request.urlopen(req, timeout = _PUBLIC_PROBE_ATTEMPT_TIMEOUT) as response: + body = response.read(4096) + if json.loads(body).get("service") == _PUBLIC_PROBE_MARKER: + return True + except Exception: + pass + remaining = deadline - time.monotonic() + if remaining <= 0: + return False + time.sleep(min(_PUBLIC_PROBE_RETRY_DELAY, remaining)) + + class CloudflareTunnel: """A cloudflared quick tunnel to http://localhost:. Best-effort throughout. @@ -322,11 +392,12 @@ def start_studio_tunnel(port: int, timeout: float = _READY_TIMEOUT) -> Optional[ """Start a quick tunnel and return its public URL once it is actually serving, or None (best-effort). - Waits for cloudflared to both mint the URL and register an edge connection - before returning, so the caller never advertises a URL that yields Cloudflare - error 1033 (HTTP 530). If a URL is minted but no connection registers within - the window (e.g. quic is blocked on this network), retries once forcing the - http2 protocol. On any failure the tunnel is stopped and None is returned. + Waits for cloudflared to both mint the URL and register an edge connection, + then fetches /api/health over the public URL, so the caller never advertises + a link that yields Cloudflare error 1033 (HTTP 530) or an unresolvable host. + If a URL is minted but no connection registers within the window (e.g. quic + is blocked on this network), retries once forcing the http2 protocol. On any + failure the tunnel is stopped and None is returned. """ global _active_tunnel, _shutdown_requested binary = ensure_cloudflared() @@ -349,9 +420,13 @@ def start_studio_tunnel(port: int, timeout: float = _READY_TIMEOUT) -> Optional[ prior, _active_tunnel = _active_tunnel, tunnel if prior is not None: prior.stop() + registered = False try: tunnel.start() url = tunnel.wait_for_ready(timeout) + registered = url is not None + if url and not verify_public_url(url): + url = None except Exception: url = None if url: @@ -371,6 +446,9 @@ def start_studio_tunnel(port: int, timeout: float = _READY_TIMEOUT) -> Optional[ # http2 will not help, so do not burn another window on it. if not saw_url: return None + # probe failure after registering is DNS propagation; http2 would not help + if registered: + return None return None diff --git a/studio/backend/tests/test_cloudflare_tunnel.py b/studio/backend/tests/test_cloudflare_tunnel.py index bb51cabf76..2094d15066 100644 --- a/studio/backend/tests/test_cloudflare_tunnel.py +++ b/studio/backend/tests/test_cloudflare_tunnel.py @@ -403,11 +403,234 @@ def test_reader_ignores_api_endpoint_failure_line(): assert t.error == "cloudflared exited before emitting a tunnel URL" +# ── public reachability probe ──────────────────────────────────────── + + +class _FakeResponse: + def __init__(self, body): + self._body = body + + def read(self, size = -1): + return self._body + + def __enter__(self): + return self + + def __exit__(self, *exc): + return False + + +def _patch_urlopen(monkeypatch, handler): + import urllib.request + monkeypatch.setattr(urllib.request, "urlopen", lambda req, timeout = None: handler(req)) + + +@pytest.fixture(autouse = True) +def _stub_dns_wait(monkeypatch, request): + if request.node.name.startswith("test_verify_public_url"): + monkeypatch.setattr(ct, "_wait_for_dns", lambda *a, **kw: None) + + +def test_wait_for_dns_polls_until_answer(monkeypatch): + calls = [] + + def handler(req): + calls.append(req.full_url) + if len(calls) < 3: + return _FakeResponse(b'{"Status":3}') + return _FakeResponse(b'{"Status":0,"Answer":[{"data":"104.16.0.1"}]}') + + _patch_urlopen(monkeypatch, handler) + monkeypatch.setattr(ct.time, "sleep", lambda _s: None) + ct._wait_for_dns("words.trycloudflare.com", ct.time.monotonic() + 5) + assert len(calls) == 3 + assert "name=words.trycloudflare.com" in calls[0] + + +def test_wait_for_dns_gives_up_at_deadline(monkeypatch): + _patch_urlopen(monkeypatch, lambda req: _FakeResponse(b'{"Status":3}')) + monkeypatch.setattr(ct.time, "sleep", lambda _s: None) + ct._wait_for_dns("words.trycloudflare.com", ct.time.monotonic() + 0.05) + + +def test_wait_for_dns_retries_transient_doh_error(monkeypatch): + calls = [] + + def handler(req): + calls.append(req.full_url) + if len(calls) < 3: + raise OSError("transient") + return _FakeResponse(b'{"Status":0,"Answer":[{"data":"104.16.0.1"}]}') + + _patch_urlopen(monkeypatch, handler) + monkeypatch.setattr(ct.time, "sleep", lambda _s: None) + ct._wait_for_dns("words.trycloudflare.com", ct.time.monotonic() + 5) + assert len(calls) == 3 + + +def test_wait_for_dns_bails_on_persistent_doh_errors(monkeypatch): + calls = [] + + def handler(req): + calls.append(req.full_url) + raise OSError("blocked") + + _patch_urlopen(monkeypatch, handler) + monkeypatch.setattr(ct.time, "sleep", lambda _s: None) + ct._wait_for_dns("words.trycloudflare.com", ct.time.monotonic() + 5) + assert len(calls) == ct._DNS_MAX_DOH_ERRORS + + +def test_verify_public_url_accepts_studio_marker(monkeypatch): + seen = {} + + def handler(req): + seen["url"] = req.full_url + return _FakeResponse(b'{"status":"healthy","service":"Unsloth UI Backend"}') + + _patch_urlopen(monkeypatch, handler) + assert ct.verify_public_url("https://words.trycloudflare.com") is True + assert seen["url"] == "https://words.trycloudflare.com/api/health" + + +def test_verify_public_url_waits_for_dns_first(monkeypatch): + order = [] + monkeypatch.setattr(ct, "_wait_for_dns", lambda host, deadline: order.append(("dns", host))) + + def handler(req): + order.append(("probe", req.full_url)) + return _FakeResponse(b'{"service":"Unsloth UI Backend"}') + + _patch_urlopen(monkeypatch, handler) + assert ct.verify_public_url("https://words.trycloudflare.com") is True + assert order[0] == ("dns", "words.trycloudflare.com") + assert order[1][0] == "probe" + + +def test_verify_public_url_dns_wait_and_probe_share_deadline(monkeypatch): + # An exhausted DNS wait leaves the probe a single attempt, not a fresh window. + calls = [] + monkeypatch.setattr(ct, "_wait_for_dns", lambda host, deadline: None) + + def handler(req): + calls.append(req.full_url) + raise OSError("unreachable") + + _patch_urlopen(monkeypatch, handler) + assert ct.verify_public_url("https://words.trycloudflare.com", timeout = 0) is False + assert len(calls) == 1 + + +def test_verify_public_url_retries_then_succeeds(monkeypatch): + calls = [] + + def handler(req): + calls.append(req.full_url) + if len(calls) < 3: + raise OSError("Name or service not known") + return _FakeResponse(b'{"service":"Unsloth UI Backend"}') + + _patch_urlopen(monkeypatch, handler) + monkeypatch.setattr(ct.time, "sleep", lambda _s: None) + assert ct.verify_public_url("https://words.trycloudflare.com") is True + assert len(calls) == 3 + + +def test_verify_public_url_rejects_unreachable_host(monkeypatch): + def handler(req): + raise OSError("Name or service not known") + + _patch_urlopen(monkeypatch, handler) + monkeypatch.setattr(ct.time, "sleep", lambda _s: None) + assert ct.verify_public_url("https://words.trycloudflare.com", timeout = 0.05) is False + + +def test_verify_public_url_rejects_foreign_responder(monkeypatch): + # e.g. a Cloudflare error page: no service marker in the body. + _patch_urlopen(monkeypatch, lambda req: _FakeResponse(b"error 1033")) + monkeypatch.setattr(ct.time, "sleep", lambda _s: None) + assert ct.verify_public_url("https://words.trycloudflare.com", timeout = 0.05) is False + + +@pytest.fixture(autouse = True) +def _stub_public_probe(monkeypatch, request): + # start_studio_tunnel tests use fake hostnames; keep them off the network. + if not request.node.name.startswith("test_start_studio_tunnel"): + return + monkeypatch.setattr(ct, "verify_public_url", lambda url, **kw: True) + + def test_start_studio_tunnel_no_binary(monkeypatch): monkeypatch.setattr(ct, "ensure_cloudflared", lambda: None) assert ct.start_studio_tunnel(8080) is None +def test_start_studio_tunnel_drops_url_that_is_not_publicly_reachable(monkeypatch): + attempts = [] + + class _Stub: + def __init__( + self, + port, + binary, + protocol = None, + ): + self.url = None + attempts.append(protocol) + + def start(self): + self.url = "https://words.trycloudflare.com" + + def wait_for_ready(self, timeout): + return self.url + + def stop(self): + pass + + monkeypatch.setattr(ct, "ensure_cloudflared", lambda: "/bin/cloudflared") + monkeypatch.setattr(ct, "CloudflareTunnel", _Stub) + monkeypatch.setattr(ct, "verify_public_url", lambda url, **kw: False) + assert ct.start_studio_tunnel(8080) is None + assert attempts == [None] + assert ct._active_tunnel is None + + +def test_start_studio_tunnel_returns_url_once_probe_passes(monkeypatch): + probed = [] + + class _Stub: + def __init__( + self, + port, + binary, + protocol = None, + ): + self.url = None + self.protocol = protocol + + def start(self): + self.url = "https://words.trycloudflare.com" + + def wait_for_ready(self, timeout): + return self.url + + def stop(self): + pass + + def _probe(url, **kw): + probed.append(url) + return True + + monkeypatch.setattr(ct, "ensure_cloudflared", lambda: "/bin/cloudflared") + monkeypatch.setattr(ct, "CloudflareTunnel", _Stub) + monkeypatch.setattr(ct, "verify_public_url", _probe) + try: + assert ct.start_studio_tunnel(8080) == "https://words.trycloudflare.com" + assert probed == ["https://words.trycloudflare.com"] + finally: + ct.stop_studio_tunnel() + + def test_start_studio_tunnel_registers_before_wait(monkeypatch): # The tunnel must be visible to stop_studio_tunnel() during the readiness # wait, else a shutdown in that window orphans cloudflared. diff --git a/studio/frontend/src/components/assistant-ui/thread.tsx b/studio/frontend/src/components/assistant-ui/thread.tsx index adab56582b..ee81ef6794 100644 --- a/studio/frontend/src/components/assistant-ui/thread.tsx +++ b/studio/frontend/src/components/assistant-ui/thread.tsx @@ -1570,6 +1570,18 @@ const Composer: FC<{ const t = setTimeout(() => writeComposerDraft(draftKey, composerText), 300); return () => clearTimeout(t); }, [composerText, draftKey]); + // Without this the restore effect above puts the sent text back when the + // runtime rebinds on the first message. + const draftKeyRef = useRef(draftKey); + useEffect(() => { + draftKeyRef.current = draftKey; + }, [draftKey]); + const clearStoredDraft = useCallback(() => { + const key = draftKeyRef.current; + if (key) { + writeComposerDraft(key, ""); + } + }, []); // react-textarea-autosize re-measures only on value change or window resize, // not on the width swap from expanding, so it keeps the taller height and // leaves a stray blank row. Nudge a resize whenever input width changes. @@ -1720,9 +1732,10 @@ const Composer: FC<{ setPendingSend(false); dismissWaitToast(); if (text.trim().length > 0 || attachments.length > 0) { + clearStoredDraft(); aui.composer().send(); } - }, [pendingSend, indexingActive, aui, dismissWaitToast]); + }, [pendingSend, indexingActive, aui, clearStoredDraft, dismissWaitToast]); // Drop any queued send + toast on unmount (e.g. thread switch). useEffect( @@ -1765,6 +1778,7 @@ const Composer: FC<{ flushResourcesSync(() => { aui.composer().setText(""); }); + clearStoredDraft(); startPromptQueue( [queuedPrompt], createPromptQueueTarget(), @@ -1798,6 +1812,7 @@ const Composer: FC<{ closeOverlay(); return; } + clearStoredDraft(); setImageToolsEnabled(true); setPendingImageEditReference({ threadId: overlay.threadId ?? referenceThreadId, @@ -1815,11 +1830,15 @@ const Composer: FC<{ ); }); closeOverlay(); + return; } + + clearStoredDraft(); }, [ aui, canQueueCurrentPrompt, + clearStoredDraft, closeOverlay, composerText, createPromptQueueTarget, @@ -1921,6 +1940,7 @@ const Composer: FC<{ flushResourcesSync(() => { aui.composer().setText(""); }); + clearStoredDraft(); startPromptQueue([queuedPrompt], createPromptQueueTarget(), true); }} onSendClick={interceptSend}