Harden terminal WebSocket edge cases #37

Merged
jimmy merged 2 commits from dev into main 2026-03-15 23:22:09 +08:00
4 changed files with 47 additions and 19 deletions
+7
View File
@@ -28,6 +28,13 @@ ARG YQ_VERSION=v4.44.6
RUN wget -q -O /usr/local/bin/yq "https://github.com/mikefarah/yq/releases/download/${YQ_VERSION}/yq_linux_amd64" \ RUN wget -q -O /usr/local/bin/yq "https://github.com/mikefarah/yq/releases/download/${YQ_VERSION}/yq_linux_amd64" \
&& chmod +x /usr/local/bin/yq && chmod +x /usr/local/bin/yq
ARG TEA_VERSION=0.12.0
RUN wget -q -O /tmp/tea.tar.gz "https://dl.gitea.com/tea/${TEA_VERSION}/tea-${TEA_VERSION}-linux-amd64.tar.gz" \
&& tar -xzf /tmp/tea.tar.gz -C /tmp \
&& mv /tmp/tea-${TEA_VERSION}-linux-amd64/tea /usr/local/bin/tea \
&& chmod +x /usr/local/bin/tea \
&& rm -rf /tmp/tea.tar.gz /tmp/tea-${TEA_VERSION}-linux-amd64
RUN npm install -g @anthropic-ai/claude-code RUN npm install -g @anthropic-ai/claude-code
RUN python3 - <<'PY' RUN python3 - <<'PY'
+1
View File
@@ -14,4 +14,5 @@ zip
unzip unzip
make make
tmux tmux
tea
claude claude
+13
View File
@@ -190,18 +190,31 @@ async def ws_endpoint(websocket: WebSocket, host: str = Query("nas")):
rows = struct.unpack("!H", data[3:5])[0] rows = struct.unpack("!H", data[3:5])[0]
process.channel.change_terminal_size(cols, rows) process.channel.change_terminal_size(cols, rows)
else: else:
try:
process.stdin.write(data) process.stdin.write(data)
except (BrokenPipeError, OSError) as e:
logger.warning("Failed to write to stdin for %s: %s", host, e)
break
elif "text" in msg and msg["text"]: elif "text" in msg and msg["text"]:
if msg["text"] == "__ping__": if msg["text"] == "__ping__":
logger.debug("Received ping from %s, sending pong", host) logger.debug("Received ping from %s, sending pong", host)
await websocket.send_text("__pong__") await websocket.send_text("__pong__")
continue continue
try:
process.stdin.write(msg["text"].encode()) process.stdin.write(msg["text"].encode())
except (BrokenPipeError, OSError) as e:
logger.warning("Failed to write to stdin for %s: %s", host, e)
break
except Exception as e: except Exception as e:
logger.warning("WebSocket receive loop ended for %s: %s", host, e) logger.warning("WebSocket receive loop ended for %s: %s", host, e)
finally: finally:
read_task.cancel() read_task.cancel()
try:
await asyncio.wait_for(read_task, timeout=2.0)
except (asyncio.TimeoutError, asyncio.CancelledError):
pass
process.close() process.close()
finally: finally:
conn.close() conn.close()
await conn.wait_closed()
await _release_session(session_id) await _release_session(session_id)
+13 -6
View File
@@ -300,7 +300,9 @@
ws.send("__ping__"); ws.send("__ping__");
// Set timeout to detect missing pong // Set timeout to detect missing pong
pongTimeout = setTimeout(() => { pongTimeout = setTimeout(() => {
if (ws.readyState === 1) { // Check if still the same WebSocket instance and still open
const currentTab = tabData.get(tabId);
if (currentTab?.ws === ws && ws.readyState === 1) {
ws.close(1000, "keepalive ping timeout"); ws.close(1000, "keepalive ping timeout");
} }
}, PONG_TIMEOUT); }, PONG_TIMEOUT);
@@ -343,8 +345,8 @@
}; };
ws.onclose = async (e) => { ws.onclose = async (e) => {
const reason = e.reason || `Connection closed (code ${e.code})`; const reason = e.reason || `Connection closed (code ${e.code})`;
d = tabData.get(tabId); const tabState = tabData.get(tabId);
if (d) d.ws = null; if (tabState) tabState.ws = null;
// Handle auth failures (code 1008) // Handle auth failures (code 1008)
if (e.code === 1008) { if (e.code === 1008) {
@@ -359,8 +361,8 @@
// Try to refresh session // Try to refresh session
if (!closedManually && !authRecoveryInFlight) { if (!closedManually && !authRecoveryInFlight) {
authRecoveryInFlight = true; authRecoveryInFlight = true;
try {
const refreshed = await tryRefreshSession(); const refreshed = await tryRefreshSession();
authRecoveryInFlight = false;
if (refreshed && !closedManually) { if (refreshed && !closedManually) {
reconnecting = true; reconnecting = true;
updateTab(tabId, { updateTab(tabId, {
@@ -372,6 +374,11 @@
void connect(false); void connect(false);
return; return;
} }
} catch (err) {
console.error("Session refresh failed:", err);
} finally {
authRecoveryInFlight = false;
}
} }
handleAuthExpired("Session expired — please log in again"); handleAuthExpired("Session expired — please log in again");
return; return;
@@ -386,8 +393,8 @@
} }
}; };
ws.onerror = () => { ws.onerror = () => {
d = tabData.get(tabId); const tabState = tabData.get(tabId);
if (d) d.ws = null; if (tabState) tabState.ws = null;
if (!closedManually) scheduleReconnect("Connection failed"); if (!closedManually) scheduleReconnect("Connection failed");
else updateTab(tabId, { connected: false, error: "Connection failed" }); else updateTab(tabId, { connected: false, error: "Connection failed" });
}; };