Studio: fix stuck composer prompt on first send and unreachable --secure Cloudflare links (#7340)
* Studio: clear composer draft on send * Studio: verify the Cloudflare link is reachable before printing it * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: wait for tunnel DNS propagation before verifying the public URL * Studio: bound tunnel DNS wait and health probe by one deadline * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: keep composer draft when overlay send validation fails * Studio: retry transient DoH failures while waiting for tunnel DNS * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Nilay
GitHub
parent
d59c7bfd03
commit
bfb6b9600c
3 changed files with 327 additions and 6 deletions
|
|
@ -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:<port>. 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
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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"<html>error 1033</html>"))
|
||||
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.
|
||||
|
|
|
|||
|
|
@ -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}
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue