feat(tunnel): free finished tunnels at once and wait out 429s on --remote

A tunnel closed with code 1000 is freed without the 60s grace, a shared: true API key counts tunnels per X-Client-Id, and --remote waits out Retry-After and a full tunnel cap with a countdown; responses carry X-RateLimit-* headers.

docs/dev/rest-api/authentication.md

@@ -224,6 +224,19 @@ A limit comes from `serve.tiers.<name>.rate_limit` or from the API key's own
 [dashboard rate limits](dashboard.md#rate-limits) for the config and the counters, and
 [services configuration](../../reference/configuration/services.md) for the keys themselves.
 
+Every response to an API key that has a limit carries these headers, the `429` too. A
+WebSocket or event stream response sends its headers before the handler returns, so it has none.
+The CORS headers expose these headers and `Retry-After` to a browser client.
+
+| Header | Value |
+|---|---|
+| `X-RateLimit-Limit` | Requests per hour that the API key may make. |
+| `X-RateLimit-Remaining` | Requests left in the current window. |
+| `X-RateLimit-Reset` | Seconds until the window resets. |
+
+A `--remote` client that gets a `429` with `Retry-After` shows a countdown, waits that time
+(one hour at most), and then sends the request again. It stops after 3 waits for one request.
+
 !!! warning "`429` is not the `RATE_LIMITED` error body"
     `RATE_LIMITED` in [Errors](errors.md) is a *service* rate limit that a handler raises, and it
     arrives in the standard envelope with an `error_code`. The `429` here comes from the
@@ -311,6 +324,7 @@ With `--caddy`, unshackle also starts `caddy run` using the `Caddyfile` located
 | Missing-header body | `{"status": 401, "message": "Secret Key is Empty."}` |
 | Invalid-key body | `{"status": 401, "message": "Secret Key is Invalid."}` |
 | Over-limit body | `{"status": 429, "message": "Rate limit exceeded."}`, plus a `Retry-After` header |
+| Rate limit headers | `X-RateLimit-Limit`, `X-RateLimit-Remaining`, `X-RateLimit-Reset` on each response to a limited API key |
 | Rate limit config | `serve.tiers.<name>.rate_limit`, overridden by `serve.users.<key>.rate_limit` |
 | Disable auth | `--no-key` (all requests allowed) |
 | Config location | `serve.api_secret` and `serve.users` in `unshackle.yaml` |

docs/dev/rest-api/remote-sessions.md

@@ -448,7 +448,7 @@ The server mounts all these routes, even in `--remote-only` mode. Paths use the
 | Method | Path | Purpose |
 |---|---|---|
 | `POST` | `/api/session/create` | Build + authenticate a service instance; returns a session ID immediately |
-| `GET` | `/api/tunnel` | WebSocket that relays the connections of a `tunnel:<id>` proxy to the client; `?id=` reconnects |
+| `GET` | `/api/tunnel` | WebSocket that relays the connections of a `tunnel:<id>` proxy to the client; `?id=` reconnects, close code 1000 frees the tunnel at once |
 | `GET` | `/api/tunnel/check?id=` | Exit IP address and lookup time through the caller's tunnel, for `unshackle tunnel-test` |
 | `GET` | `/api/session/{id}/titles` | Fetch and cache the title list |
 | `POST` | `/api/session/{id}/tracks` | Resolve tracks, manifests, chapters, headers/cookies for a title |

docs/guide/proxies-and-vpn.md

@@ -336,8 +336,10 @@ A service that uses a server account does not work with `--local-proxy`, because
 logs in with its own account. If the tunnel drops, the client connects again, and a new request
 from the server waits up to 10 seconds for it. A request that was in progress when the tunnel
 dropped fails. The server closes a tunnel when its client has not connected for 60 seconds, and
-the client then stops with `The tunnel to the server closed`. An API key can have 4 tunnels
-open at once.
+the client then stops with `The tunnel to the server closed`. A client that exits closes its
+tunnel at once. An API key can have 4 tunnels open at once (4 for each person on a key with
+`shared: true`, see [services configuration](../reference/configuration/services.md)). A fifth run waits for a free
+tunnel and tries again every 15 seconds, for one hour at most.
 
 ## Basic (static proxies)
 

docs/reference/configuration/services.md

@@ -240,10 +240,17 @@ server vault with this grant, are the only source, and a missing content key sto
 
 `admin` is a boolean that lets the API key run the maintenance endpoints (clear-cache, clear-temp, refresh-services). It is `false` unless the entry sets it. Keys that have no `users` entry, such as `api_secret`, keep that access.
 
+`shared` is a boolean for an API key that a gateway in front of the server gives to many people.
+The gateway puts its own ID for each person in the `X-Client-Id` header, and the server then
+counts the `--local-proxy` tunnels of that API key for each person. Without `shared: true` the
+server ignores the header, because a client that connects to the server directly can set any
+value. Set it only when the gateway always replaces the header that the client sends.
+
 A `tier` names an entry under `serve.tiers`, which holds the settings that several API keys
 share. Today a tier carries `rate_limit`, the requests per hour that API key may make; an API
 key can also set its own `rate_limit`, which wins over its tier's. An API key with neither has
-no limit. Going over the limit answers `429` with a `Retry-After` header. The server keeps a
+no limit. Going over the limit answers `429` with a `Retry-After` header, and a `--remote`
+client waits that time before it tries again. The server keeps a
 fixed window rather than a sliding one: it opens on the first request and resets an hour
 later. A `tier` that names no entry under `serve.tiers`, or a `rate_limit` that is not a
 positive whole number, stops the server at startup, because the alternative is an API key that
@@ -341,6 +348,8 @@ serve:
     f3a4b5c6:
       services: [EXAMPLE1]
       server_vault: [EXAMPLE1]    # the client's device licenses; the server vault still serves keys
+    d7e8f9a0:
+      shared: true                # a gateway gives this key to many people; tunnels count per X-Client-Id
 ```
 
 !!! note "`dl` keys inside `serve`"

tests/remote/unit/test_dashboard.py

@@ -252,9 +252,13 @@ async def test_rate_limited_request_returns_429(aiohttp_client, monkeypatch) ->
     setup_routes(app, remote_only=True, dashboard=True)
     client = await aiohttp_client(app)
 
-    assert (await client.get("/api/services", headers={"X-Secret-Key": "tier-key"})).status == 200
+    resp = await client.get("/api/services", headers={"X-Secret-Key": "tier-key"})
+    assert resp.status == 200
+    assert (resp.headers["X-RateLimit-Limit"], resp.headers["X-RateLimit-Remaining"]) == ("1", "0")
+    assert 0 < int(resp.headers["X-RateLimit-Reset"]) <= 3600
     resp = await client.get("/api/services", headers={"X-Secret-Key": "tier-key"})
     assert resp.status == 429 and int(resp.headers["Retry-After"]) > 0
+    assert resp.headers["X-RateLimit-Remaining"] == "0"
     assert (await client.get("/api/health", headers={"X-Secret-Key": "tier-key"})).status == 200
     assert (await client.get("/api/health")).status == 401
     # The dashboard key leaves the middleware before the limit check, so it always answers.

tests/remote/unit/test_remote_client.py

@@ -4,6 +4,7 @@ from __future__ import annotations
 
 import json
 
+import click
 import pytest
 import responses
 
@@ -244,3 +245,93 @@ def test_proxy_error_page_logs_its_title(client: RemoteClient, caplog: pytest.Lo
         client.get("/api/session/s/titles")
     assert "HTTP 524: example.com | 524: A timeout occurred" in caplog.text
     assert "DOCTYPE" not in caplog.text
+
+
+@responses.activate
+def test_a_429_with_retry_after_waits_then_retries(client: RemoteClient, monkeypatch: pytest.MonkeyPatch) -> None:
+    waits: list[int] = []
+    monkeypatch.setattr("unshackle.core.remote_service.wait_retry_after", waits.append)
+    url = "http://srv:8786/api/services"
+    responses.add(responses.GET, url, status=429, json={"status": 429}, headers={"Retry-After": "7"})
+    responses.add(responses.GET, url, json={"status": "ok"})
+
+    assert client.get("/api/services") == {"status": "ok"}
+    assert waits == [7]
+
+
+@responses.activate
+def test_a_429_without_retry_after_fails_at_once(client: RemoteClient, monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr("unshackle.core.remote_service.wait_retry_after", pytest.fail)
+    responses.add(responses.GET, "http://srv:8786/api/services", status=429, json={"error_code": "RATE_LIMITED"})
+
+    with pytest.raises(SystemExit):
+        client.get("/api/services")
+
+
+def test_wait_retry_after_waits_out_its_own_deadline(monkeypatch: pytest.MonkeyPatch) -> None:
+    from unshackle.core import remote_service
+
+    clock = [0.0]
+    monkeypatch.setattr(remote_service.time, "monotonic", lambda: clock[0])
+    monkeypatch.setattr(remote_service.time, "sleep", lambda s: clock.__setitem__(0, clock[0] + s))
+    remote_service.wait_retry_after(3)
+    assert clock[0] == 3.0
+    remote_service.wait_retry_after(99999)
+    assert clock[0] == 3.0 + remote_service.RETRY_AFTER_MAX
+
+
+def test_a_rate_limited_tunnel_waits_then_opens(client: RemoteClient, monkeypatch: pytest.MonkeyPatch) -> None:
+    from unshackle.core import remote_service
+
+    attempts: list[object] = []
+
+    class Refused:
+        def __init__(self, *_: object) -> None:
+            self.retry_after = 5 if not attempts else None
+            self.at_limit = len(attempts) == 1
+            attempts.append(self)
+
+        def start(self) -> str:
+            if self.retry_after or self.at_limit:
+                raise ConnectionError("HTTP 429")
+            return "id"
+
+    waits: list[int] = []
+    monkeypatch.setattr(remote_service, "TunnelClient", Refused)
+    monkeypatch.setattr(remote_service, "wait_retry_after", lambda s, *_: waits.append(s))
+    assert remote_service.start_tunnel(client) is attempts[2]
+    assert waits == [5, remote_service.TUNNEL_SLOT_POLL]
+
+
+def test_a_tunnel_refused_for_another_reason_fails_at_once(
+    client: RemoteClient, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    from unshackle.core import remote_service
+
+    class Refused:
+        retry_after = None
+        at_limit = False
+
+        def __init__(self, *_: object) -> None:
+            pass
+
+        def start(self) -> str:
+            raise ConnectionError("HTTP 403")
+
+    monkeypatch.setattr(remote_service, "TunnelClient", Refused)
+    monkeypatch.setattr(remote_service, "wait_retry_after", pytest.fail)
+    with pytest.raises(click.ClickException, match="HTTP 403"):
+        remote_service.start_tunnel(client)
+
+
+@responses.activate
+def test_the_429_wait_gives_up_after_its_tries(client: RemoteClient, monkeypatch: pytest.MonkeyPatch) -> None:
+    from unshackle.core import remote_service
+
+    waits: list[int] = []
+    monkeypatch.setattr(remote_service, "wait_retry_after", waits.append)
+    responses.add(responses.GET, "http://srv:8786/api/services", status=429, json={}, headers={"Retry-After": "0"})
+
+    with pytest.raises(SystemExit):
+        client.get("/api/services")
+    assert waits == [1] * remote_service.RATE_LIMIT_TRIES

tests/remote/unit/test_tunnel.py

@@ -16,6 +16,7 @@ import pytest
 import requests
 import responses
 from aiohttp import web
+from multidict import CIMultiDict
 
 from unshackle.core import tunnel
 from unshackle.core.api.routes import tunnel as tunnel_route
@@ -356,6 +357,7 @@ def test_a_malformed_frame_still_leaves_the_tunnel_to_expire(server: str) -> Non
             async with http.ws_connect(server + ROUTE) as ws:
                 tunnel_id = (await ws.receive_json())["id"]
                 await ws.send_bytes(b"\x01")
+                await ws.close(code=aiohttp.WSCloseCode.GOING_AWAY)
         return str(tunnel_id)
 
     tunnel_id = asyncio.run(send_short_frame())
@@ -369,12 +371,112 @@ def test_an_api_key_has_a_tunnel_limit(server: str) -> None:
     clients = [TunnelClient(server, {"X-Secret-Key": "bob"}) for _ in range(tunnel.MAX_TUNNELS)]
     for client in clients:
         client.start()
-    with pytest.raises(ConnectionError, match="429"):
+    with pytest.raises(ConnectionError, match="429, the API key already has the maximum number of tunnels open"):
         TunnelClient(server, {"X-Secret-Key": "bob"}).start()
     for client in clients:
         client.stop()
 
 
+@pytest.mark.parametrize("shared", [True, False])
+def test_a_shared_api_key_counts_tunnels_for_each_client(
+    server: str, monkeypatch: pytest.MonkeyPatch, shared: bool
+) -> None:
+    from unshackle.core.config import config
+
+    monkeypatch.setitem(config.serve, "users", {"tier": {"shared": True} if shared else {}})
+    first = [TunnelClient(server, {"X-Secret-Key": "tier", "X-Client-Id": "u1"}) for _ in range(tunnel.MAX_TUNNELS)]
+    second = TunnelClient(server, {"X-Secret-Key": "tier", "X-Client-Id": "u2"})
+    try:
+        for client in first:
+            client.start()
+        if shared:
+            second.start()
+        else:
+            with pytest.raises(ConnectionError, match="maximum number of tunnels"):
+                second.start()
+        with pytest.raises(ConnectionError, match="maximum number of tunnels"):
+            TunnelClient(server, {"X-Secret-Key": "tier", "X-Client-Id": "u1"}).start()
+    finally:
+        for client in [*first, second]:
+            client.stop()
+
+
+def test_another_client_of_a_shared_api_key_cannot_reattach(server: str, monkeypatch: pytest.MonkeyPatch) -> None:
+    from unshackle.core.config import config
+
+    monkeypatch.setitem(config.serve, "users", {"tier": {"shared": True}})
+    owner = TunnelClient(server, {"X-Secret-Key": "tier", "X-Client-Id": "u1"})
+    tunnel_id = owner.start()
+    try:
+        resp = requests.get(
+            f"{server}{ROUTE}?id={tunnel_id}", headers={"X-Secret-Key": "tier", "X-Client-Id": "u2"}, timeout=10
+        )
+        assert resp.status_code == 404
+    finally:
+        owner.stop()
+
+
+def test_stopping_one_client_leaves_the_other_tunnels_of_its_api_key_open(
+    server: str, target: int, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    monkeypatch.setattr(tunnel, "is_public_ip", lambda address: True)
+    monkeypatch.setattr(tunnel, "dialled_itself", lambda writer: False)
+    clients = [TunnelClient(server, {"X-Secret-Key": "erin"}) for _ in range(3)]
+    stopped, *running = [client.start() for client in clients]
+    try:
+        clients[0].stop()
+        assert wait_until(lambda: stopped not in TUNNELS, 5)
+        for tunnel_id in running:
+            proxy = tunnel_proxy(tunnel_id, "erin")
+            assert proxy
+            resp = requests.get(f"http://127.0.0.1:{target}/", proxies={"all": proxy}, timeout=10)
+            assert resp.text == "hello through the tunnel"
+    finally:
+        for client in clients[1:]:
+            client.stop()
+
+
+@pytest.mark.parametrize(
+    ("value", "seconds"),
+    [
+        ("5", 5),
+        ("0", 1),
+        (None, None),
+        ("", None),
+        ("\u00b2", None),
+        ("9" * 5000, None),
+        ("Wed, 21 Oct 2026 07:28:00 GMT", None),
+    ],
+    ids=["seconds", "zero", "absent", "empty", "superscript", "huge", "http-date"],
+)
+def test_parse_retry_after(value: str | None, seconds: int | None) -> None:
+    assert tunnel.parse_retry_after(value) == seconds
+
+
+def test_a_client_that_exits_frees_its_tunnel_at_once(server: str) -> None:
+    async def open_then_leave() -> str:
+        async with aiohttp.ClientSession(headers={"X-Secret-Key": "dave"}) as http:
+            async with http.ws_connect(server + ROUTE) as ws:
+                return str((await ws.receive_json())["id"])
+
+    tunnel_id = asyncio.run(open_then_leave())
+    assert wait_until(lambda: tunnel_id not in TUNNELS, 5)
+
+
+@pytest.mark.parametrize(
+    ("status", "headers", "reason"),
+    [
+        (429, {"Retry-After": "120"}, "HTTP 429, the API key is over its hourly request limit, retry in 120s"),
+        (429, {"retry-after": "120"}, "HTTP 429, the API key is over its hourly request limit, retry in 120s"),
+        (429, {}, "HTTP 429, the API key already has the maximum number of tunnels open"),
+        (404, None, "HTTP 404"),
+    ],
+)
+def test_a_refused_tunnel_names_the_reason(status: int, headers: dict[str, str] | None, reason: str) -> None:
+    error = aiohttp.WSServerHandshakeError(None, (), status=status, headers=CIMultiDict(headers or {}))  # type: ignore[arg-type]
+    assert tunnel.refusal(error) == reason
+
+
 def test_a_connection_that_stops_reading_is_dropped() -> None:
     sent: list = []
 

unshackle/commands/serve.py

@@ -21,6 +21,7 @@ from unshackle.core.api.compression import compression_middleware
 from unshackle.core.api.events import publish_refresh_events, publish_service_event
 from unshackle.core.api.handlers import (
     DASHBOARD_PREFIX,
+    add_rate_limit_headers,
     api_key_authentication,
     dashboard_authentication,
     dashboard_key,
@@ -559,6 +560,8 @@ def serve(
                             f"https://git.gay/ready-dl/pyplayready serve v{pyplayready_version}"
                         )
 
+                    if request.path != "/api/health" and secret_key and known:
+                        add_rate_limit_headers(response, secret_key)
                     return response
 
                 return serve_authentication

unshackle/core/api/handlers.py

@@ -1354,6 +1354,7 @@ CORS_HEADERS = {
     "Access-Control-Allow-Methods": "GET, POST, PUT, DELETE, OPTIONS",
     "Access-Control-Allow-Headers": "Content-Type, X-Secret-Key, Authorization",
     "Access-Control-Max-Age": "3600",
+    "Access-Control-Expose-Headers": "Retry-After, X-RateLimit-Limit, X-RateLimit-Remaining, X-RateLimit-Reset",
 }
 
 
@@ -1408,10 +1409,19 @@ def rate_limit_response(secret_key: str) -> Optional[web.Response]:
     return web.json_response(
         {"status": 429, "message": "Rate limit exceeded."},
         status=429,
-        headers={"Retry-After": str(retry_after)},
+        headers={"Retry-After": str(retry_after), **stats.rate_limit_headers(secret_key)},
     )
 
 
+def add_rate_limit_headers(response: web.StreamResponse, secret_key: str) -> web.StreamResponse:
+    """Put *secret_key*'s rate-limit headers on *response*; a streamed response has sent its headers already."""
+    from unshackle.core.api.stats import stats
+
+    if not response.prepared:
+        response.headers.update(stats.rate_limit_headers(secret_key))
+    return response
+
+
 @web.middleware
 async def api_key_authentication(request: web.Request, handler: Any) -> web.StreamResponse:
     """Gate every non-dashboard route behind an API key in ``app["config"]["users"]``, then rate limit it.
@@ -1433,7 +1443,7 @@ async def api_key_authentication(request: web.Request, handler: Any) -> web.Stre
         return web.json_response({"status": 401, "message": "Secret Key is Invalid."}, status=401)
     if request.path == "/api/health":
         return await handler(request)
-    return rate_limit_response(secret_key) or await handler(request)
+    return rate_limit_response(secret_key) or add_rate_limit_headers(await handler(request), secret_key)
 
 
 def caller_key(request: Optional[web.Request] = None) -> str:
@@ -5649,19 +5659,30 @@ async def session_delete_handler(session_id: str, request: Optional[web.Request]
     return web.json_response(response)
 
 
+def shared_key_client(request: web.Request) -> Optional[str]:
+    """The person behind a ``shared: true`` API key, from the ``X-Client-Id`` its gateway sets; None for other keys.
+
+    Only a shared key reads the header, so a client that talks to the server directly cannot use it to pass a limit.
+    """
+    if serve_user_config(request_secret_key(request) or "").get("shared") is not True:
+        return None
+    return request.headers.get("X-Client-Id", "")[:64] or None
+
+
 async def tunnel_handler(request: web.Request) -> web.StreamResponse:
     """Open a ``--local-proxy`` tunnel, or reattach to the caller's tunnel named by ``?id=``."""
     owner = request_secret_key(request)
+    client = shared_key_client(request)
     tunnel_id = request.query.get("id")
     tunnel = TUNNELS.get(tunnel_id) if tunnel_id else None
-    if tunnel_id and (not tunnel or tunnel.owner != owner):
+    if tunnel_id and (not tunnel or tunnel.owner != owner or tunnel.client != client):
         raise APIError(APIErrorCode.NOT_FOUND, "No open tunnel has this id.")
-    if not tunnel and sum(t.owner == owner for t in TUNNELS.values()) >= MAX_TUNNELS:
+    if not tunnel and sum(t.owner == owner and t.client == client for t in TUNNELS.values()) >= MAX_TUNNELS:
         raise APIError(APIErrorCode.RATE_LIMITED, f"An API key can have {MAX_TUNNELS} tunnels open at once.")
     ws = web.WebSocketResponse(heartbeat=HEARTBEAT, max_msg_size=MAX_MESSAGE)
     await ws.prepare(request)  # a request that is not a WebSocket upgrade fails here, before a tunnel exists
     if not tunnel:
-        tunnel = Tunnel(owner)
+        tunnel = Tunnel(owner, client)
         await tunnel.start()
     await tunnel.serve(ws)
     return ws

unshackle/core/api/routes.py

@@ -2074,7 +2074,9 @@ async def tunnel(request: web.Request) -> web.StreamResponse:
     description: >-
       A WebSocket that a --remote client holds open. The server sends the requests of a remote session
       created with proxy tunnel:<id> through it, and the client makes them from its own connection.
-      The first message is {"id": <tunnel id>}.
+      The first message is {"id": <tunnel id>}. A client that has finished closes with code 1000,
+      which frees the tunnel at once; after any other close the tunnel waits 60 seconds for a reconnect.
+      An API key with shared: true counts its tunnels for each X-Client-Id value, which its gateway sets.
     parameters:
       - name: id
         in: query
@@ -2088,7 +2090,9 @@ async def tunnel(request: web.Request) -> web.StreamResponse:
       '404':
         description: No open tunnel with this id for this API key
       '429':
-        description: The API key has the maximum number of tunnels open
+        description: >-
+          The API key has the maximum number of tunnels open, or, with a Retry-After header,
+          the API key is over its hourly request limit
     """
     return await tunnel_handler(request)
 

unshackle/core/api/stats.py

@@ -183,6 +183,18 @@ class ServerStats:
         entry.window_used += 1
         return None
 
+    def rate_limit_headers(self, key: str) -> Dict[str, str]:
+        """The ``X-RateLimit-*`` headers for ``key``'s current window; none for an API key with no limit."""
+        limit = key_rate_limit(key)
+        if limit is None:
+            return {}
+        entry = self.key_stats(key)
+        return {
+            "X-RateLimit-Limit": str(limit),
+            "X-RateLimit-Remaining": str(max(0, limit - entry.window_used)),
+            "X-RateLimit-Reset": str(max(0, int(entry.window_start + RATE_LIMIT_WINDOW - time.time()))),
+        }
+
     def to_dict(self) -> Dict[str, Any]:
         return {
             "version": __version__,

unshackle/core/remote_service.py

@@ -13,6 +13,7 @@ import contextlib
 import hashlib
 import json
 import logging
+import math
 import random
 import re
 import shlex
@@ -35,9 +36,10 @@ from langcodes import Language
 from requests.adapters import HTTPAdapter, Retry
 from rich.padding import Padding
 from rich.rule import Rule
+from rich.spinner import Spinner
 
 from unshackle.core.config import config
-from unshackle.core.console import console, prompt_user, show_notice
+from unshackle.core.console import SyncLive, console, prompt_user, show_notice
 from unshackle.core.constants import AnyTrack
 from unshackle.core.credential import Credential
 from unshackle.core.drm import BlowfishStripe, cbcs, real_kids
@@ -52,7 +54,7 @@ from unshackle.core.titles.music import SONG_FIELDS, Album, Song
 from unshackle.core.tracks import Audio, Chapter, Chapters, Subtitle, Tracks, Video
 from unshackle.core.tracks.attachment import Attachment
 from unshackle.core.tracks.track import Track
-from unshackle.core.tunnel import TunnelClient
+from unshackle.core.tunnel import TunnelClient, parse_retry_after
 from unshackle.core.utils import run_proxy
 from unshackle.core.utils.collections import ci_get
 from unshackle.core.utils.redact import redact_path, redact_text, safe_display_url
@@ -69,6 +71,10 @@ DEFAULT_AUTH_HEADERS = ["X-Secret-Key", "X-Api-Key"]
 SESSION_LOST = ("SESSION_NOT_FOUND", "CONNECT", "TIMEOUT", "TUNNEL_CLOSED")
 # what a recovering call can meet on its way, so these log only once the recovery fails
 RECOVERABLE = (*SESSION_LOST, "SERVER_DRAINING")
+RETRY_AFTER_MAX = 3600
+RETRY_AFTER_LOCK = Lock()
+RATE_LIMIT_TRIES = 3
+TUNNEL_SLOT_POLL = 15
 REOPEN_WAIT = 1800.0  # seconds a lost session waits for its server to come back; covers a deploy's whole drain
 
 _recovering = local()
@@ -92,6 +98,23 @@ def redact_secrets(text: str, data: Optional[Dict[str, Any]] = None) -> str:
     return redact_text(text, secrets) or ""
 
 
+def wait_retry_after(seconds: int, reason: str = "Rate limited by the server") -> None:
+    """Count down a wait on the console, as ``--slow`` does, so the caller can retry; one hour at most, the rate-limit window.
+
+    Threads refused together share one countdown, each still waits out its own deadline.
+    """
+    deadline = time.monotonic() + min(seconds, RETRY_AFTER_MAX)
+    with RETRY_AFTER_LOCK:
+        if deadline <= time.monotonic():
+            return
+        log.info(f"{reason}, retrying in {math.ceil(deadline - time.monotonic())}s")
+        spinner = Spinner("dots")
+        with SyncLive(Padding(spinner, (0, 5)), console=console, refresh_per_second=12.5, transient=True):
+            while (left := deadline - time.monotonic()) > 0:
+                spinner.update(text=f"{reason}, retrying in {math.ceil(left)} seconds...")
+                time.sleep(min(1.0, left))
+
+
 class RemoteError(SystemExit):
     """A remote call that failed: logged, then exits a CLI run like ``SystemExit(1)``.
 
@@ -172,6 +195,7 @@ class RemoteClient:
             {"json": data} if body is None else {"data": body, "headers": {"Content-Type": "application/octet-stream"}}
         )
         gateway_retried = False
+        rate_waits = 0
         while True:
             try:
                 with self.in_flight(endpoint):
@@ -186,6 +210,14 @@ class RemoteClient:
             # servers differ on which header carries the key, so retry the rest before giving up
             if resp.status_code == 401 and self.next_auth_header():
                 continue
+            if (
+                resp.status_code == 429
+                and rate_waits < RATE_LIMIT_TRIES
+                and (retry_after := parse_retry_after(resp.headers.get("Retry-After")))
+            ):
+                rate_waits += 1
+                wait_retry_after(retry_after)
+                continue
             # a proxy HTML page, not the server's JSON, means a dropped tunnel stream, which one retry gets past
             if (
                 resp.status_code in (502, 503, 504)
@@ -831,11 +863,21 @@ def resolve_remote_proxy_arg(proxy_arg: str, providers: Optional[list] = None) -
 def start_tunnel(client: RemoteClient) -> TunnelClient:
     """Open a tunnel that sends the server's requests out through this machine's own connection."""
     headers = {k: str(v) for k, v in client.session.headers.items() if k.lower() != "connection"}
-    tunnel = TunnelClient(client.server_url, headers)
-    try:
-        tunnel.start()
-    except ConnectionError as e:
-        raise click.ClickException(f"--local-proxy could not open a tunnel to the server: {e}")
+    slot_deadline = time.monotonic() + RETRY_AFTER_MAX
+    rate_waits = 0
+    while True:
+        tunnel = TunnelClient(client.server_url, headers)
+        try:
+            tunnel.start()
+            break
+        except ConnectionError as e:
+            if tunnel.retry_after and rate_waits < RATE_LIMIT_TRIES:
+                rate_waits += 1
+                wait_retry_after(tunnel.retry_after)
+            elif tunnel.at_limit and time.monotonic() < slot_deadline:
+                wait_retry_after(TUNNEL_SLOT_POLL, "Every tunnel of this API key is in use")
+            else:
+                raise click.ClickException(f"--local-proxy could not open a tunnel to the server: {e}")
     client.tunnel = tunnel
     return tunnel
 

unshackle/core/tunnel.py

@@ -20,7 +20,8 @@ import threading
 from typing import Any, Optional
 
 import aiohttp
-from aiohttp import WSMsgType, web
+from aiohttp import WSCloseCode, WSMsgType, web
+from multidict import CIMultiDict
 
 log = logging.getLogger("tunnel")
 
@@ -144,10 +145,11 @@ TUNNELS: dict[str, "Tunnel"] = {}
 class Tunnel:
     """The server side of one client's tunnel: a loopback SOCKS5 listener and the client's current link."""
 
-    def __init__(self, owner: Optional[str]) -> None:
+    def __init__(self, owner: Optional[str], client: Optional[str] = None) -> None:
         self.id = secrets.token_urlsafe(16)
         self.secret = secrets.token_urlsafe(24)
         self.owner = owner
+        self.client = client
         self.link: Optional[Link] = None
         self.pending: dict[int, asyncio.Future[bool]] = {}
         self.sids = itertools.count(1)
@@ -290,7 +292,9 @@ class Tunnel:
                 self.link = None
                 self.linked.clear()
                 link.end_all()
-                if self.id in TUNNELS:
+                if ws.close_code == WSCloseCode.OK:
+                    self.expire()
+                elif self.id in TUNNELS:
                     self.expiry = asyncio.get_running_loop().call_later(GRACE, self.expire)
 
     def expire(self) -> None:
@@ -377,6 +381,27 @@ async def connect_public(host: str, port: int) -> tuple[asyncio.StreamReader, as
     raise error
 
 
+def parse_retry_after(value: Optional[str]) -> Optional[int]:
+    """Whole seconds, at least 1, from a delay-seconds Retry-After; None for an absent, HTTP-date or malformed one."""
+    if not value or len(value) > 9 or not (value.isascii() and value.isdigit()):
+        return None
+    return max(1, int(value))
+
+
+def retry_after(e: aiohttp.WSServerHandshakeError) -> Optional[int]:
+    """The Retry-After seconds of a rate-limit refusal, or None for any other refusal."""
+    return parse_retry_after(CIMultiDict(e.headers or {}).get("Retry-After")) if e.status == 429 else None
+
+
+def refusal(e: aiohttp.WSServerHandshakeError) -> str:
+    """Why the server refused a tunnel. aiohttp drops the response body, so a 429 is told apart by its headers."""
+    if e.status != 429:
+        return f"HTTP {e.status}"
+    if (wait := retry_after(e)) is not None:
+        return f"HTTP 429, the API key is over its hourly request limit, retry in {wait}s"
+    return "HTTP 429, the API key already has the maximum number of tunnels open"
+
+
 class TunnelClient:
     """The client side: a background thread that keeps the tunnel's WebSocket open and dials for the server."""
 
@@ -385,6 +410,8 @@ class TunnelClient:
         self.headers = headers
         self.id: Optional[str] = None
         self.error: Optional[BaseException] = None
+        self.retry_after: Optional[int] = None
+        self.at_limit = False
         self.lost = False
         self.ready = threading.Event()
         self.relays: set[asyncio.Task[None]] = set()
@@ -398,7 +425,7 @@ class TunnelClient:
         self.thread.start()
         if not self.ready.wait(timeout) or not self.id:
             self.stop()  # or the thread opens the tunnel later and nobody uses it
-            raise ConnectionError(f"Could not open a tunnel to the server: {self.error or 'timed out'}")
+            raise ConnectionError(str(self.error or "timed out"))
         return self.id
 
     def stop(self) -> None:
@@ -431,7 +458,9 @@ class TunnelClient:
                         backoff = 1
                         await self.serve(Link(ws))
                 except aiohttp.WSServerHandshakeError as e:
-                    self.error = e
+                    self.error = ConnectionError(refusal(e))
+                    self.retry_after = retry_after(e)
+                    self.at_limit = e.status == 429 and self.retry_after is None
                     # a 404 means the tunnel expired or the server has no tunnel route; a 5xx is a proxy in between
                     if not self.id or e.status < 500:
                         if self.id: