diff --git a/.agents/skills/verify-diff-lifecycle/covered.lock b/.agents/skills/verify-diff-lifecycle/covered.lock index d7c002f12..dcea0caf5 100644 --- a/.agents/skills/verify-diff-lifecycle/covered.lock +++ b/.agents/skills/verify-diff-lifecycle/covered.lock @@ -139,8 +139,8 @@ file lua/claudecode/logger.lua 17b404dda41c3591 == lua/claudecode/selection.lua file lua/claudecode/selection.lua 05663913bc425259 == lua/claudecode/server/client.lua -file lua/claudecode/server/client.lua 5789c6ff18045f9f -block create_client 647b7523dea26af6 +file lua/claudecode/server/client.lua 4304e6496b5c4f32 +block create_client 24fa2a266b76e6ad == lua/claudecode/server/frame.lua file lua/claudecode/server/frame.lua 3c50766875878d56 == lua/claudecode/server/handshake.lua @@ -161,7 +161,7 @@ site lua/claudecode/server/init.lua: if not _G.claude_deferred_responses then @ == lua/claudecode/server/mock.lua file lua/claudecode/server/mock.lua def1248d2793d559 == lua/claudecode/server/tcp.lua -file lua/claudecode/server/tcp.lua f1f9833b4322b1ae +file lua/claudecode/server/tcp.lua 63cbb8922056279a == lua/claudecode/server/utils.lua file lua/claudecode/server/utils.lua 2915ae87b48f04ca == lua/claudecode/terminal.lua diff --git a/lua/claudecode/server/README.md b/lua/claudecode/server/README.md index 8622c5aa4..c11f3a2b0 100644 --- a/lua/claudecode/server/README.md +++ b/lua/claudecode/server/README.md @@ -35,13 +35,13 @@ Layers, top to bottom: - `M.create_server(config, callbacks, auth_token) -> server|nil, err`: tries ports in the configured range until one binds and listens [server/tcp.lua:M.create_server] - `M.find_available_port(min_port, max_port)`: probe-only port check [server/tcp.lua:M.find_available_port] - `M.send_to_client`, `M.broadcast`, `M.get_client_count`, and `M.get_clients_info` [server/tcp.lua:M.send_to_client] -- `M.stop_server(server)`: sends a Close frame with code 1001 "Server shutting down" to clients that finished the handshake or whose 101 reply is still being written, closes the TCP handle of the others, then closes the listener [server/tcp.lua:M.stop_server] [server/client.lua:M.close_client] -- `M.start_ping_timer(server, interval)`: default interval is 30000 ms [server/tcp.lua:M.start_ping_timer] +- `M.stop_server(server)`: closes the check and idle handles of a stall confirmation that has not run, sends a Close frame with code 1001 "Server shutting down" to clients that finished the handshake or whose 101 reply is still being written, closes the TCP handle of the others, then closes the listener [server/tcp.lua:M.stop_server] [server/client.lua:M.close_client] +- `M.start_ping_timer(server, interval)`: default interval is 30000 ms. It sets `server.keepalive_timeout` to `interval * 2`, and clients accepted from then on carry it [server/tcp.lua:M.start_ping_timer] - Internal: `M._bind_and_listen`, `M._handle_new_connection`, `M._disconnect_client`, and `M._remove_client` [server/tcp.lua:M._bind_and_listen] ### Connection, handshake, and frames -- `client.create_client(tcp_handle)`, `process_data(...)`, `send_message`, `send_ping`, `close_client`, `is_client_alive(client, timeout=30000)`, and `get_client_info` [server/client.lua:M.process_data] +- `client.create_client(tcp_handle, keepalive_timeout?)`, `process_data(...)`, `send_message`, `send_ping`, `close_client`, `abandon_close_frame`, `is_client_alive(client, timeout=30000)`, `sample_write_progress(client)`, `is_client_reading(client)`, and `get_client_info` [server/client.lua:M.process_data] - `handshake.process_handshake(request, expected_auth_token) -> ok, response, headers?` [server/handshake.lua:M.process_handshake] - `handshake.validate_upgrade_request`, `create_handshake_response`, `parse_request_line`, `is_websocket_endpoint`, `create_error_response`, and `extract_http_request` [server/handshake.lua:M.validate_upgrade_request] - `frame.OPCODE`: `CONTINUATION 0x0`, `TEXT 0x1`, `BINARY 0x2`, `CLOSE 0x8`, `PING 0x9`, `PONG 0xA` [server/frame.lua:M.OPCODE] @@ -77,7 +77,9 @@ Layers, top to bottom: - When the JSON decodes to a non-table value (for example a number or a boolean), `_handle_message` replies -32600 with `"id": null`. It does not raise [server/init.lua:M._handle_message] - A message with a truthy `id` is a request. Any other message is a notification. A valid notification never gets a response. The -32700 and -32600 checks run before the `id` check, so a malformed notification gets an error reply [server/init.lua:M._handle_message] - A deferred response is keyed by `tostring(coroutine)` in `_G.claude_deferred_responses`. A result with neither `content` nor `error` becomes -32603 [server/init.lua:M._setup_deferred_response] -- The ping timer closes a connected client when two ticks in a row find no pong within `interval * 2`. The first such tick pings again and sets `client.pong_overdue`, because libuv runs timers before it reads sockets, so a tick that Neovim runs late after a block finds the pong of the last ping still unread. The next tick finds that pong, or drops the client. A pong clears `pong_overdue` at the next tick. The Close frame of a dropped client carries 1011 "Keepalive timeout" (RFC 6455 reserves 1006, so no Close frame carries it), `on_disconnect` gets 1006 "Connection timeout", and the timer logs `Closed client_: no reply for s` at info level. A peer that stopped reading does not get that Close frame. Once its receive window is full, every write to it (replies, pings, and the Close frame) waits in libuv's queue, and closing the socket drops the queue. Until the timeout, the queued replies stay in memory. A tick gap above `interval * 1.5` counts as a wake from sleep and resets `last_pong` of each connected client. The 101 write callback sets `last_pong` again, so the keepalive window starts when the client is connected, not at the accept. A main loop that blocks during the handshake then does not count against the client [server/tcp.lua:M.start_ping_timer] [server/client.lua:M.process_data] +- The ping timer closes a connected client when two ticks in a row find no pong within `interval * 2`. The first such tick pings again and sets `client.pong_overdue`, because a tick that Neovim runs late after a block can come before the loop has read the pong of the last ping. The next tick finds that pong, or drops the client. A tick that finds a pong clears `pong_overdue`. The Close frame of a dropped client carries 1011 "Keepalive timeout" (RFC 6455 reserves 1006, so no Close frame carries it), `on_disconnect` gets 1006 "Connection timeout", and the timer logs `Closed client_: no reply for s` at info level. A peer that stopped reading does not get that Close frame: once its receive window is full, every write to it (replies, pings, and the Close frame) waits in libuv's queue, and closing the socket drops the queue [server/tcp.lua:M.start_ping_timer] [server/tcp.lua:drop_client] +- A Claude that answers each ping but never reads would keep that queue, which is Neovim's memory, without limit, so the server also closes a connected client that answers and whose write queue made no progress at two ticks in a row. Each tick calls `is_client_reading` for each connected client, which calls `sample_write_progress`: it compares `queued_bytes - get_write_queue_size()` with its value at the last progress, and an empty queue counts as progress. A tick that finds progress clears `client.write_stalled`. The first tick that finds none sets it. The second asks for a confirmation, and a confirmation that still finds no progress drops the client: the Close frame carries 1011 "Not reading", `on_disconnect` gets 1006 "Write queue stalled", and the server logs `Closed client_: not reading for s, KB queued` at info level. The confirmation runs right after the next poll. libuv reports that a socket is writable again only in a poll phase, so a tick that Neovim runs after a main-loop block inside an I/O callback can find no progress although Claude read during the block, whether the tick is late or on time. `after_next_poll` arms a check handle, which libuv runs right after each poll phase in every version, and an idle handle, which keeps that poll from waiting for a timer or a socket. The check callback closes both handles, removes them from `server.pending_polls`, and judges the queue once more: progress clears `write_stalled`, and no progress drops the client. A client that left or started to leave in the meantime, because the poll read its EOF, its Close frame or a protocol error, is skipped without a log line. `stop_server` closes the handles of a confirmation that has not run, which `server.pending_polls` lists. If libuv cannot create a handle, the tick leaves the client for the next tick. The look is a sample at one moment, not proof. libuv takes its snapshot of the ready sockets when a poll starts, so a callback that blocks inside that poll, after the snapshot, hides what Claude reads during the block from the check callback that follows. A client that was stalled for the whole window and resumes during such a block is dropped. That is a deadline race, as with any timeout: the client was stalled when the window ended, and what it read afterwards is not reported yet. Verifier probes closed a reader that resumed 0.1 s into a 0.5 s to 1.5 s block inside the confirming poll, on Neovim 0.12.3 and 0.8.0, at 3 s, 5 s and 30 s intervals. A block after the tick and before that poll, or in a scheduled callback, keeps the reader. The bound is two to three ticks after the last progress, which is 60 s to 90 s at the default interval, plus the delay of ticks that Neovim runs late by less than 1.5 intervals. A tick gap above `interval * 1.5` counts as a wake from sleep. A wake tick starts a fresh stall window: it clears the count and judges nothing, so it delays the drop by one tick when the count was clear, and by two ticks when the wake tick takes the place of the confirming tick. The next two ticks always count, wake ticks or not, because wake ticks that come again and again are a loop that stays blocked, not sleeps. So whatever the pattern of late ticks, a stalled client is closed at the second tick after a fresh window started. A wake tick three or more ticks after the last fresh start starts a new window. The bound is never the size of the queue, so a Claude that reads slowly, or is sent one message of tens of MB, stays connected while libuv writes any of it. libuv writes more only after the kernel reports the socket writable again, which took 1.2 to 1.5 MB of reading on a Linux loopback socket with default buffers (measured), so a Claude that reads less than about 1.4 MB per window counts as stalled while a backlog exists, and one that reads just faster than that still lets the queue grow as long as Neovim queues more than it reads. `write` queues a frame of more than 1 MiB in pieces of 1 MiB, because libuv on Windows lowers the queue only when a whole write request is done, so a large frame shows its progress per piece there too, and the callback of the last piece reports the frame. [server/tcp.lua:M.start_ping_timer] [server/tcp.lua:confirm_stall] [server/tcp.lua:after_next_poll] [server/tcp.lua:M.stop_server] [server/client.lua:M.is_client_reading] [server/client.lua:M.sample_write_progress] [server/client.lua:write] +- On one tick the pong rule runs first. A client without a pong gets the second chance or the keepalive drop whatever its write queue shows, and only a client that answered is confirmed for a stall. A pong that arrives while the queue is stalled does not save the client: the next tick clears `pong_overdue`, and then asks for the confirmation. Every connected client is sampled on every tick except one that starts a fresh stall window, answering or not, so "two ticks in a row" means consecutive ticks. A tick gap above `interval * 1.5` counts as a wake from sleep and resets `last_pong` of each connected client, and its own alive check finds the fresh pong and clears `pong_overdue`. A wake tick that starts a fresh stall window also resets `last_drain`, clears `write_stalled`, and does not sample the queue. `last_drain` feeds only the `not reading for s` of the log line, so the reset keeps the sleep out of it. A wake tick within two ticks after the start of a fresh window resets `last_pong` only, and counts like any tick. The 101 write callback sets `last_pong` again, so the keepalive window starts when the client is connected, not at the accept. A main loop that blocks during the handshake then does not count against the client [server/tcp.lua:M.start_ping_timer] [server/client.lua:M.process_data] - `_disconnect_client` is idempotent: `on_disconnect` fires once per client id. A client that already left the server (it stopped before a scheduled `on_close` or an EOF ran) gets its socket closed there, unless the write callback of its Close frame will close it [server/tcp.lua:M._disconnect_client] - A client's read error or protocol error reaches the server only as its `on_disconnect` reason (`Client read error: …` or `Client error: …`), which `init.lua` logs at debug level. `on_error` carries only the server's own failures (listen, accept, and creating a client handle or the ping timer), which `init.lua` logs at error level [server/tcp.lua:M._handle_new_connection] [server/init.lua:M.start] - A rejected handshake logs `Rejected client_: ` at warn level, with the reason from the body of the error response. Once the error response is written, `on_error` removes the client from the server and closes its socket [server/client.lua:M.process_data] @@ -85,7 +87,7 @@ Layers, top to bottom: - The server logs each error, warning, and info line from a libuv callback through `logger.one_line`, so the line never stops Neovim at "Press ENTER" [server/init.lua:M.start] [server/client.lua:M.process_data] [server/tcp.lua:M.start_ping_timer] - Each connection gets the id `client_`, where `n` counts the connections of this Neovim process in `_G.claudecode_clients_created`. An id never repeats, even after a reload of `claudecode.server.*`. A Claude CLI that reconnects gets a new id, so it does not own the diffs it opened before [server/client.lua:M.create_client] - When the state is already `closing` or `closed`, `close_client` sends nothing [server/client.lua:M.close_client] -- A Close frame queued by `close_client` is followed by the close of the socket, from its write callback, so `client.close_pending` is set until then and `_remove_client` leaves the socket open. After a protocol error `process_data` calls `on_error` at once, so the client leaves `server.clients` and `on_disconnect` runs while the frame can still wait in libuv's queue behind the replies of a Claude that does not read. A Claude that never reads keeps its socket and the bytes queued for it until it reads, hangs up, or Neovim exits. The ping timer clears `close_pending`, so a timed-out client's socket closes at once [server/client.lua:M.close_client] [server/client.lua:fail_connection] [server/tcp.lua:release_socket] [server/tcp.lua:M.start_ping_timer] +- A Close frame queued by `close_client` is followed by the close of the socket, from its write callback, so `client.close_pending` is set until then and `_remove_client` leaves the socket open. After a protocol error `process_data` calls `on_error` at once, so the client leaves `server.clients` and `on_disconnect` runs while the frame can still wait in libuv's queue behind the replies of a Claude that does not read. The ping timer cannot see that client, because it left `server.clients`, so `close_client` starts a deadline for the frame (`client.close_timer`, a repeating timer of `client.keepalive_timeout`). A socket that wrote no byte of its queue during a keepalive window is closed, and the queue with it, while a Claude that reads keeps its socket. The window starts with the frame, and one window without progress is enough here: the client is leaving anyway, so a close that comes too early after a blocked loop costs only its Close frame. The write callback and the deadline each close the socket if it is still open and end the other, so the first one to run wins, and a deadline never touches a socket that closed. A server without a ping timer has no keepalive window and its clients get no deadline, but `M.start` starts the timer in the same call that listens, before any client can connect. `drop_client` calls `abandon_close_frame`, so a client that the ping timer drops, for a keepalive timeout or a stalled write queue, has its socket closed at once [server/client.lua:M.close_client] [server/client.lua:fail_connection] [server/client.lua:start_close_timer] [server/tcp.lua:release_socket] [server/tcp.lua:drop_client] - `sha1` rejects non-string input and input over 10 MiB [server/utils.lua:M.sha1] ## Side Effects @@ -99,7 +101,9 @@ Layers, top to bottom: - `client.create_client` increments `_G.claudecode_clients_created` [server/client.lua:M.create_client] - `tcp.lua` seeds `math.randomseed` once per process. `init.lua` calls `math.random` at load. Both change the global PRNG [server/tcp.lua:ensure_rng_seeded] - `client.process_data` writes handshake, PONG, and Close frames to the socket. It runs `on_message` and `on_close` through `vim.schedule`. It calls `on_error` directly after a failed handshake write, and after it queues the Close frame for a protocol error [server/client.lua:M.process_data] [server/client.lua:M.close_client] -- The ping timer writes ping frames and closes timed-out clients [server/tcp.lua:M.start_ping_timer] +- The ping timer writes ping frames and closes clients that timed out or stopped reading [server/tcp.lua:M.start_ping_timer] +- A tick that finds the second stalled tick in a row creates a `vim.loop` check handle and an idle handle for one confirmation. The check callback closes them, and `stop_server` closes those that did not run [server/tcp.lua:after_next_poll] [server/tcp.lua:M.stop_server] +- `close_client` starts a repeating `vim.loop` timer for the Close frame of a client that has `keepalive_timeout`. It runs until the frame is written, the socket closes, or the deadline closes the socket [server/client.lua:M.close_client] [server/client.lua:start_close_timer] - `handshake.lua`, `frame.lua`, and `utils.lua` do no I/O [server/handshake.lua:M] - `mock.lua` opens no sockets and keeps its clients and messages in its own `M.state` tables. Its `mcp.tool.invoke` handler calls the real `tools.handle_invoke`, so a simulated tool call has that tool's real side effects (for example `saveDocument` writes the file) [server/mock.lua:M.register_handlers] [server/mock.lua:M.state] diff --git a/lua/claudecode/server/client.lua b/lua/claudecode/server/client.lua index 742612882..a168d4193 100644 --- a/lua/claudecode/server/client.lua +++ b/lua/claudecode/server/client.lua @@ -27,15 +27,22 @@ local CLOSE_REASONS = { ---@field pending table|nil The reads of an incomplete frame, with their total length in `size` and the frame's in `needed` ---@field fragments table|nil Joined pieces of the text message whose last fragment has not arrived, and their total `size` ---@field handshake_complete boolean Whether WebSocket handshake is complete ----@field close_pending boolean|nil True while a Close frame waits to be written: its write callback closes the socket +---@field close_pending boolean|nil True while a Close frame waits to be written: its write callback closes the socket, or its deadline does +---@field close_timer table|nil The deadline of the Close frame: a repeating timer that runs while `close_pending` is set +---@field keepalive_timeout number|nil The window in milliseconds for a pong and for write progress, from the server's ping interval ---@field last_ping number Timestamp of last ping sent ---@field last_pong number Timestamp of last pong received +---@field queued_bytes integer The bytes passed to the socket's write queue since the connection began +---@field drained_bytes integer The bytes libuv had written to the socket when `sample_write_progress` last saw progress +---@field last_drain number Timestamp of the last `sample_write_progress` check that found the write queue empty or smaller +---@field write_stalled boolean|nil True after a ping tick found no write progress. The next such tick asks for a look after the next poll, and only that look drops the client ---@field pong_overdue boolean|nil True after a ping timer tick found no pong for two intervals. The next such tick drops the client ---Create a new WebSocket client ---@param tcp_handle table The vim.loop TCP handle +---@param keepalive_timeout number|nil The window in milliseconds that the server's ping timer allows for a pong and for write progress. Without it the client has no deadline for its Close frame ---@return WebSocketClient client The client object -function M.create_client(tcp_handle) +function M.create_client(tcp_handle, keepalive_timeout) _G.claudecode_clients_created = (_G.claudecode_clients_created or 0) + 1 local client_id = "client_" .. _G.claudecode_clients_created @@ -47,11 +54,88 @@ function M.create_client(tcp_handle) handshake_complete = false, last_ping = 0, last_pong = vim.loop.now(), + queued_bytes = 0, + drained_bytes = 0, + last_drain = vim.loop.now(), + keepalive_timeout = keepalive_timeout, } return client end +-- The size of one write request for a large frame. libuv lowers the size of a socket's write queue by the bytes it wrote +-- on Linux and macOS, but on Windows only by a whole request when it is done. Queued in pieces, a large frame shows its +-- progress there too, and `is_client_reading` does not take a Claude that reads it for a stalled one. +local WRITE_PIECE_BYTES = 1024 * 1024 + +---Write to the socket of a client. Every write goes through here, so `queued_bytes` counts them all: the bytes still in +---libuv's write queue are `queued_bytes - drained_bytes`, and `is_client_reading` tells a stalled queue from a slow one. +---A frame of more than `WRITE_PIECE_BYTES` is queued in pieces. libuv finishes the requests of a socket in order, so the +---callback of the last piece says that the whole frame is written, or that a piece failed. +---@param client WebSocketClient +---@param data string +---@param callback function? Optional callback: function(err) +local function write(client, data, callback) + client.queued_bytes = client.queued_bytes + #data + if #data <= WRITE_PIECE_BYTES then + client.tcp_handle:write(data, callback) + return + end + for first = 1, #data, WRITE_PIECE_BYTES do + local last_piece = first + WRITE_PIECE_BYTES > #data + client.tcp_handle:write(data:sub(first, first + WRITE_PIECE_BYTES - 1), last_piece and callback or nil) + end +end + +---Stop the deadline of the Close frame. Safe to call when none runs. +---@param client WebSocketClient +local function stop_close_timer(client) + local timer = client.close_timer + if timer then + client.close_timer = nil + timer:stop() + timer:close() + end +end + +---Start the deadline of a Close frame that was just queued. The ping timer sees only the clients in the server's list, +---and a client that broke the protocol, or that the server stopped, has left it with its frame still queued. A socket +---that wrote no byte for a keepalive window is closed, and the queue with it, as for the ping timer. The window starts +---with the frame, and a Claude that reads keeps the socket open. One window without progress is enough here: the client +---is leaving already, so a close that comes too early after a blocked loop costs only its Close frame. +---@param client WebSocketClient +local function start_close_timer(client) + local timeout = client.keepalive_timeout + -- Not when the write callback already ran: the frame is written and the socket is closed. + if not timeout or not client.close_pending then + return + end + local timer = vim.loop.new_timer() + if not timer then + return + end + stop_close_timer(client) + client.close_timer = timer + client.drained_bytes = client.queued_bytes - client.tcp_handle:get_write_queue_size() + client.last_drain = vim.loop.now() + timer:start(timeout, timeout, function() + -- The write callback, a drop by the ping timer, or a close of the handle ended the wait: stop, and touch nothing. + -- Whichever of this callback and the write callback runs first closes the socket, and the other finds it closed. + if client.close_timer ~= timer then + return + end + local handle = client.tcp_handle + if not client.close_pending or handle:is_closing() then + stop_close_timer(client) + elseif not M.sample_write_progress(client) then + stop_close_timer(client) + client.close_pending = nil + client.state = "closed" + handle:close() + end + end) +end + ---Adds a data frame to the message the client is sending. A text message arrives as one TEXT frame with FIN set, or ---as fragments: a TEXT frame with FIN clear, then CONTINUATION frames up to one with FIN set. A fragment may end ---inside a code point, so only the whole text is checked for UTF-8. @@ -174,7 +258,7 @@ function M.process_data(client, data, on_message, on_close, on_error, auth_token client.state = "closing" end - client.tcp_handle:write(response_from_handshake, function(err) + write(client, response_from_handshake, function(err) if err or not success then -- on_error removes the client from the server and closes its socket. A server that stopped while the reply -- was written has already dropped the client, so the socket is closed here too. @@ -268,7 +352,7 @@ function M.process_data(client, data, on_message, on_close, on_error, auth_token if client.state == "connected" then local close_frame = frame.create_close_frame(code, reason) - client.tcp_handle:write(close_frame) + write(client, close_frame) client.state = "closing" end @@ -283,7 +367,7 @@ function M.process_data(client, data, on_message, on_close, on_error, auth_token break elseif opcode == frame.OPCODE.PING then local pong_frame = frame.create_pong_frame(parsed_frame.payload) - client.tcp_handle:write(pong_frame) + write(client, pong_frame) elseif opcode == frame.OPCODE.PONG then client.last_pong = vim.loop.now() end @@ -310,7 +394,7 @@ function M.send_message(client, message, callback) -- Every outbound message passes here. A buffer name or path in it can hold any bytes, and the client drops the -- connection on a text frame that is not valid UTF-8. local text_frame = frame.create_text_frame(utils.valid_utf8(message)) - client.tcp_handle:write(text_frame, callback) + write(client, text_frame, callback) end ---Send a ping to a client @@ -322,7 +406,7 @@ function M.send_ping(client, data) end local ping_frame = frame.create_ping_frame(data or "") - client.tcp_handle:write(ping_frame) + write(client, ping_frame) client.last_ping = vim.loop.now() end @@ -353,16 +437,19 @@ function M.close_client(client, code, reason) -- An upgrading client gets the Close frame too: libuv writes it after the 101 reply that is still queued. if client.handshake_complete or client.state == "upgrading" then local close_frame = frame.create_close_frame(code, reason) - -- The write callback closes the socket, so nothing may close it first: that would cancel the frame. libuv always - -- invokes the callback (on success or error), so the socket is never left open. + -- The write callback closes the socket, so nothing may close it first: that would cancel the frame. libuv invokes + -- the callback on success, on error and when the socket closes, so the socket is never left open. A Claude that + -- does not read can keep the frame in the queue for good: the deadline closes the socket then. client.close_pending = true - client.tcp_handle:write(close_frame, function() + write(client, close_frame, function() client.close_pending = nil + stop_close_timer(client) client.state = "closed" if not client.tcp_handle:is_closing() then client.tcp_handle:close() end end) + start_close_timer(client) -- Mark as "closing" while the Close frame write/teardown is in flight. The -- write callback transitions to "closed". Do not clobber the "closed" state -- set synchronously on the not-handshake-complete branch below. @@ -373,6 +460,13 @@ function M.close_client(client, code, reason) end end +---Stop waiting for the write of a Close frame, because the caller closes the socket at once. +---@param client WebSocketClient The client object +function M.abandon_close_frame(client) + client.close_pending = nil + stop_close_timer(client) +end + ---Check if a client connection is alive ---@param client WebSocketClient The client object ---@param timeout number Timeout in milliseconds (default: 30000) @@ -388,6 +482,40 @@ function M.is_client_alive(client, timeout) return (now - client.last_pong) < timeout end +---Check whether libuv wrote a byte of the write queue of a client since the last check that saw progress, or the queue +---is empty. A client that reads slowly makes progress, and so does one that is sent a single large message, however big. +---@param client WebSocketClient The client object +---@return boolean progress False when the queue is not empty and nothing of it was written since the last progress +function M.sample_write_progress(client) + local queue_size = client.tcp_handle:get_write_queue_size() + local drained = client.queued_bytes - queue_size + if queue_size == 0 or drained > client.drained_bytes then + client.drained_bytes, client.last_drain = drained, vim.loop.now() + return true + end + return false +end + +---Check that a client still reads what the server writes to it. The keepalive cannot tell: a Claude whose reader is +---stuck can still answer each ping, while replies and pings pile up in libuv's write queue, in Neovim's memory. +---Call it once per ping tick. A tick can run before libuv reported the bytes that Claude read during a main-loop block, +---so a tick never drops a client. A tick that finds progress, or an empty queue, clears the count. The first that finds +---none sets it, and the second returns false. `tcp.lua` then looks once more right after the next poll, and only that +---look drops the client. +---@param client WebSocketClient The client object +---@return boolean reading False when this tick, and the one before it, found no write progress +function M.is_client_reading(client) + if M.sample_write_progress(client) then + client.write_stalled = nil + return true + end + if client.write_stalled then + return false + end + client.write_stalled = true + return true +end + ---Get client info for debugging ---@param client WebSocketClient The client object ---@return table info Client information diff --git a/lua/claudecode/server/tcp.lua b/lua/claudecode/server/tcp.lua index c4fc11573..639d4d93c 100644 --- a/lua/claudecode/server/tcp.lua +++ b/lua/claudecode/server/tcp.lua @@ -13,6 +13,8 @@ local M = {} ---@field on_open function Callback for a client whose WebSocket handshake completed ---@field on_disconnect function Callback for client disconnections ---@field on_error fun(err_msg: string) Callback for the server's own errors, not a client's +---@field keepalive_timeout number|nil The window in milliseconds for a pong and for write progress: twice the ping interval, set by `start_ping_timer` +---@field pending_polls table|nil The check and idle handles of each stall judgement that waits for the next poll: closed by `stop_server` -- Seed Lua's PRNG exactly once per process. #282 removed the implicit seeding -- that used to happen via utils.shuffle_array (math.randomseed(os.time())), which @@ -197,7 +199,7 @@ function M._handle_new_connection(server) end -- Create WebSocket client wrapper - local client = client_manager.create_client(client_tcp) + local client = client_manager.create_client(client_tcp, server.keepalive_timeout) server.clients[client.id] = client -- Set up data handler @@ -287,6 +289,20 @@ function M._remove_client(server, client) end end +---Drop a client that the ping timer gave up on, and report it at once. RFC 6455 reserves 1006, so the Close frame says +---1011, as other servers do for a keepalive timeout. The connection still ends without the client's Close frame, which +---is what 1006 means here. The socket does not wait for that frame: a Claude that stopped answering or reading may +---never read it, and the frame waits behind everything queued for it. Closing now drops the frame and the queue. +---@param server TCPServer The server object +---@param client WebSocketClient The client to drop +---@param close_reason string The reason in the Close frame +---@param disconnect_reason string The reason that `on_disconnect` gets, with the code 1006 +local function drop_client(server, client, close_reason, disconnect_reason) + client_manager.close_client(client, 1011, close_reason) + client_manager.abandon_close_frame(client) + M._disconnect_client(server, client, 1006, disconnect_reason) +end + ---Send a message to a specific client ---@param server TCPServer The server object ---@param client_id string The client ID @@ -364,9 +380,87 @@ function M.get_clients_info(server) return clients end +---Run `fn` once, right after libuv polls its sockets the next time. libuv reports that a socket is writable again +---only in a poll phase, and it runs a check handle right after each poll phase, in every libuv version. A ping tick +---that Neovim runs after a main-loop block inside an I/O callback can run before that poll, so it cannot see the bytes +---that Claude read during the block: what runs here can. An idle handle keeps that poll from waiting for a timer or a +---socket. Both handles are closed once `fn` ran. `stop_server` closes those that did not run yet. +---@param server TCPServer The server object +---@param fn function +---@return boolean armed False when libuv could not create a handle, and `fn` will not run +local function after_next_poll(server, fn) + local check = vim.loop.new_check() + local idle = vim.loop.new_idle() + if not check or not idle then + if check then + check:close() + end + if idle then + idle:close() + end + return false + end + local handles = { check, idle } + server.pending_polls = server.pending_polls or {} + server.pending_polls[handles] = true + idle:start(function() end) + check:start(function() + server.pending_polls[handles] = nil + for _, handle in ipairs(handles) do + handle:stop() + handle:close() + end + fn() + end) + return true +end + +---Close the handles of the polls that `fn` did not run after yet. +---@param server TCPServer The server object +local function close_pending_polls(server) + local pending = server.pending_polls + server.pending_polls = nil + for handles in pairs(pending or {}) do + for _, handle in ipairs(handles) do + handle:stop() + if not handle:is_closing() then + handle:close() + end + end + end +end + +---Drop a client that answered its ping and whose write queue made no progress at two ticks in a row, unless the next +---poll shows progress. A Claude whose reader is stuck can still answer each ping, so only its write queue shows it. +---The queue is Neovim's memory: it goes with the client. The tick can run before libuv reported the bytes that Claude +---read during a block, so the stall is judged once more right after the next poll, and only that judgement drops. +---@param server TCPServer The server object +---@param client WebSocketClient The client whose queue made no progress at two ticks in a row +local function confirm_stall(server, client) + after_next_poll(server, function() + if server.clients[client.id] ~= client or client.state ~= "connected" then + return + end + if client_manager.sample_write_progress(client) then + client.write_stalled = nil + return + end + local time_since_drain = math.floor((vim.loop.now() - client.last_drain) / 1000) + local queued_kb = math.floor(client.tcp_handle:get_write_queue_size() / 1024) + require("claudecode.logger").one_line( + "info", + "server", + string.format("Closed %s: not reading for %ds, %d KB queued", client.id, time_since_drain, queued_kb) + ) + drop_client(server, client, "Not reading", "Write queue stalled") + end) +end + ---Stop the TCP server ---@param server TCPServer The server object function M.stop_server(server) + close_pending_polls(server) + -- Close all clients for _, client in pairs(server.clients) do client_manager.close_client(client, 1001, "Server shutting down") @@ -388,6 +482,8 @@ end function M.start_ping_timer(server, interval) interval = interval or 30000 -- 30 seconds local last_run = vim.loop.now() + -- The ticks since the last one that started a fresh stall window: see `fresh_start` below. + local ticks_since_fresh = math.huge local timer = vim.loop.new_timer() if not timer then @@ -395,6 +491,11 @@ function M.start_ping_timer(server, interval) return nil end + -- A client has this long to answer a ping, and to let libuv write some of its queue. A Close frame gets the same + -- window (see `client.close_client`), and the clients accepted from now on carry it. + local timeout = interval * 2 + server.keepalive_timeout = timeout + timer:start(interval, interval, function() local now = vim.loop.now() local elapsed = now - last_run @@ -402,9 +503,18 @@ function M.start_ping_timer(server, interval) -- Detect potential system sleep: timer interval was significantly exceeded -- Allow 50% grace period (e.g., 45s instead of 30s) to account for system load local is_wake_from_sleep = elapsed > (interval * 1.5) + -- A wake tick gives each client a fresh stall window: Claude had no chance yet to read after the sleep, so the stall + -- count starts again and this tick judges nothing. The next two ticks always count, wake ticks or not, because wake + -- ticks that come again and again are a loop that stays blocked, not sleeps. A stalled client is then closed at the + -- second of them, and no pattern of late ticks keeps it. + ticks_since_fresh = ticks_since_fresh + 1 + local fresh_start = is_wake_from_sleep and ticks_since_fresh > 2 + if fresh_start then + ticks_since_fresh = 0 + end if is_wake_from_sleep then - -- After system sleep/wake, reset all client pong timestamps to prevent false timeouts + -- After system sleep/wake, reset all client pong and write-progress timestamps to prevent false timeouts -- This gives clients a fresh keepalive window since the time jump isn't their fault require("claudecode.logger").debug( "server", @@ -415,21 +525,34 @@ function M.start_ping_timer(server, interval) ) for _, client in pairs(server.clients) do if client.state == "connected" then + -- The alive check below finds the fresh pong and clears `pong_overdue`. client.last_pong = now + if fresh_start then + client.last_drain = now + client.write_stalled = nil + end end end end for _, client in pairs(server.clients) do if client.state == "connected" then + -- Every tick samples the write queue, so that "two ticks in a row" are consecutive ticks, whether the client + -- answered or not. A tick that starts a fresh stall window does not. + local reading = fresh_start or client_manager.is_client_reading(client) -- Check if client is alive (local connections, so use standard timeout) - if client_manager.is_client_alive(client, interval * 2) then + if client_manager.is_client_alive(client, timeout) then client.pong_overdue = nil client_manager.send_ping(client, "ping") + -- Only a client that answered is dropped for a stall, and only after the next poll confirms it. + if not reading then + confirm_stall(server, client) + end elseif not client.pong_overdue then - -- libuv runs timers before it reads sockets, so a tick that Neovim runs late after a block finds the pong of - -- the last ping still unread. The first tick that finds no pong for two intervals pings again. The next one - -- finds the pong that was read in between, or drops the client. + -- A tick that Neovim runs late after a block can come before the loop has read the pong of the last ping. + -- The first tick that finds no pong for two intervals pings again. The next one finds the pong that was read + -- in between, or drops the client. A client without a pong gets this second chance whatever its write queue + -- shows: the stall rule above judges only a client that answered. client.pong_overdue = true client_manager.send_ping(client, "ping") else @@ -440,13 +563,7 @@ function M.start_ping_timer(server, interval) "server", string.format("Closed %s: no reply for %ds", client.id, time_since_pong) ) - -- RFC 6455 reserves 1006, so the Close frame says 1011, as other servers do for a keepalive timeout. The - -- connection still ends without the client's Close frame, which is what 1006 means here. - client_manager.close_client(client, 1011, "Keepalive timeout") - -- The socket does not wait for the Close frame: a Claude that stopped answering may never read it. Closing - -- now drops the frame when Claude's window is full. - client.close_pending = nil - M._disconnect_client(server, client, 1006, "Connection timeout") + drop_client(server, client, "Keepalive timeout", "Connection timeout") end end end diff --git a/tests/mocks/vim.lua b/tests/mocks/vim.lua index 54734b91f..dda6e3267 100644 --- a/tests/mocks/vim.lua +++ b/tests/mocks/vim.lua @@ -1244,6 +1244,9 @@ local vim = { end return true end, + get_write_queue_size = function(self) + return 0 + end, close = function(self) return true end, @@ -1265,6 +1268,40 @@ local vim = { end, } end, + -- Check handles run a callback right after the poll phase of a loop iteration, and idle handles keep that poll from + -- waiting. The mock never runs them. + new_idle = function() + return { + start = function(self, callback) + return true + end, + stop = function(self) + return true + end, + close = function(self) + return true + end, + is_closing = function(self) + return false + end, + } + end, + new_check = function() + return { + start = function(self, callback) + return true + end, + stop = function(self) + return true + end, + close = function(self) + return true + end, + is_closing = function(self) + return false + end, + } + end, now = function() return os.time() * 1000 end, diff --git a/tests/unit/server/connection_lifecycle_spec.lua b/tests/unit/server/connection_lifecycle_spec.lua index 33f59cd9f..b7f6996f3 100644 --- a/tests/unit/server/connection_lifecycle_spec.lua +++ b/tests/unit/server/connection_lifecycle_spec.lua @@ -5,8 +5,8 @@ require("tests.busted_setup") describe("The lifecycle of a WebSocket connection", function() - local tcp, frame, server, callbacks, scheduled - local clock, saved + local tcp, frame, utils, server, callbacks, scheduled + local clock, saved, timers, checks, idles local token = "test-token-0123456789" local interval = 30000 @@ -22,8 +22,64 @@ describe("The lifecycle of a WebSocket connection", function() end tcp = require("claudecode.server.tcp") frame = require("claudecode.server.frame") - scheduled, clock = nil, 1000000 - saved = { schedule = vim.schedule, now = vim.loop.now, new_timer = vim.loop.new_timer } + scheduled, clock, timers, checks, idles = nil, 1000000, {}, {}, {} + utils = require("claudecode.utils") + saved = { + schedule = vim.schedule, + now = vim.loop.now, + new_timer = vim.loop.new_timer, + new_check = vim.loop.new_check, + new_idle = vim.loop.new_idle, + valid_utf8 = utils.valid_utf8, + } + -- `send_message` scans each outbound text with `utils.valid_utf8`, byte by byte in Lua. Under luacov every byte costs + -- a line hook, so the multi-MiB messages of these tests took minutes in CI. They check how a message is queued and + -- when its client is dropped, not its encoding: `send_utf8_spec.lua` tests that. + utils.valid_utf8 = function(text) + return text + end + -- Each timer waits for the test to run it, and records how it ended. The ping timer is one of them. + vim.loop.new_timer = function() + local timer = { stopped = false, closed = false } + function timer:start(delay, repeat_ms, callback) + self.delay, self.repeat_ms, self.callback = delay, repeat_ms, callback + end + function timer:stop() + self.stopped = true + end + function timer:close() + assert(not self.closed, "handle is already closing") + self.closed = true + end + table.insert(timers, timer) + return timer + end + -- libuv runs a check callback right after the poll phase of a loop iteration, and an active idle handle keeps that + -- poll from waiting. These handles run when the test calls `poll()`, and record how they ended. + local function fake_handle(list) + return function() + local handle = { active = false, closed = false } + function handle:start(callback) + self.active, self.callback = true, callback + return true + end + function handle:stop() + self.active = false + return true + end + function handle:close() + assert(not self.closed, "handle is already closing") + self.closed, self.active = true, false + end + function handle:is_closing() + return self.closed + end + table.insert(list, handle) + return handle + end + end + vim.loop.new_check = fake_handle(checks) + vim.loop.new_idle = fake_handle(idles) -- While `scheduled` is a list, scheduled callbacks wait in it, as they wait for the main loop in Neovim. vim.schedule = function(fn) if scheduled then @@ -48,11 +104,14 @@ describe("The lifecycle of a WebSocket connection", function() after_each(function() vim.schedule, vim.loop.now, vim.loop.new_timer = saved.schedule, saved.now, saved.new_timer + vim.loop.new_check, vim.loop.new_idle = saved.new_check, saved.new_idle + utils.valid_utf8 = saved.valid_utf8 end) - ---Accept a connection whose socket works like libuv's. + ---Accept a connection whose socket works like libuv's. `handle:get_write_queue_size()` counts the bytes of the writes + ---that wait, and `handle.drain(n)` lets the kernel accept `n` bytes of them, the oldest write first. ---@return table client - ---@return table writes Each write: { data = string, done = function(err) } + ---@return table writes Each write: { data = string, left = integer, done = function(err) } ---@return table handle local function connect() tcp._handle_new_connection(server) @@ -61,7 +120,30 @@ describe("The lifecycle of a WebSocket connection", function() local closed = false function handle:write(data, callback) assert(not closed, "write after close") - table.insert(writes, { data = data, done = callback or function() end }) + local write = { data = data, left = #data } + function write.done(err) + write.left = 0 + if callback then + callback(err) + end + end + table.insert(writes, write) + end + function handle:get_write_queue_size() + local size = 0 + for _, write in ipairs(writes) do + size = size + write.left + end + return size + end + function handle.drain(bytes) + for _, write in ipairs(writes) do + local taken = math.min(write.left, bytes) + write.left, bytes = write.left - taken, bytes - taken + if taken > 0 and write.left == 0 then + write.done() + end + end end function handle:close() assert(not closed, "handle is already closing") @@ -91,20 +173,50 @@ describe("The lifecycle of a WebSocket connection", function() end ---Start the ping timer on the server, and return the function that runs on each tick. - local function start_ping_timer() - local tick - vim.loop.new_timer = function() - return { - start = function(_, _, _, callback) - tick = callback - end, - stop = function() end, - close = function() end, - } + ---One poll phase of the loop: libuv polls its sockets (`during` runs, and reports what Claude read), then the armed + ---check callbacks run. + ---@param during function|nil What libuv reports while it polls + local function poll(during) + if during then + during() + end + for _, handle in ipairs(checks) do + if handle.active then + handle.callback() + end end + end + + ---Start the ping timer. The first result runs one loop iteration: the timer callback, then a poll phase. The second + ---result runs the timer callback alone, as a tick that libuv runs before it polls. + local function start_ping_timer() + local before = #timers tcp.start_ping_timer(server, interval) - vim.loop.new_timer = saved.new_timer - return tick + assert.equals(before + 1, #timers) + local ping_timer = timers[#timers] + return function() + ping_timer.callback() + poll() + end, ping_timer.callback + end + + ---The timer that waits for the Close frame of a client. It repeats every keepalive window, two ping intervals, and + ---the ping timer every interval. + local function close_deadline() + local deadlines = {} + for _, timer in ipairs(timers) do + if timer.repeat_ms == 2 * interval then + table.insert(deadlines, timer) + end + end + assert.equals(1, #deadlines, "no deadline waits for the Close frame") + return deadlines[1] + end + + ---One keepalive window passes and the timer runs. + local function run_after_window(timer) + clock = clock + timer.repeat_ms + timer.callback() end local function parse(write) @@ -168,10 +280,15 @@ describe("The lifecycle of a WebSocket connection", function() local tick = start_ping_timer() local client, writes, handle = connect_and_upgrade() local pong = frame.create_frame(frame.OPCODE.PONG, "ping", true, true) + -- A Claude answers a ping after it read it, so libuv has written the pings queued for it by then. + local function claude_pongs() + handle.drain(math.huge) + handle._read_cb(nil, pong) + end clock = clock + interval tick() - handle._read_cb(nil, pong) + claude_pongs() clock = clock + interval tick() -- Neovim blocks right after this ping, before it reads the pong, for 1.3 intervals clock = clock + 1.3 * interval @@ -179,7 +296,7 @@ describe("The lifecycle of a WebSocket connection", function() assert.is_not_nil(server.clients[client.id], "the client was dropped by a late tick") assert.spy(callbacks.on_disconnect).was_not_called() - handle._read_cb(nil, pong) -- the loop reads the pong after the late tick + claude_pongs() -- the loop reads the pong after the late tick clock = clock + interval tick() clock = clock + 1.3 * interval -- a second late tick, again with the pong unread @@ -293,4 +410,955 @@ describe("The lifecycle of a WebSocket connection", function() assert.is_true(handle:is_closing(), "the socket waits for a Claude that does not read") end) end) + + -- A PONG keeps the keepalive window open, so only the write queue can show that Claude stopped reading. The queue + -- counts as stalled when two observations in a row, one per ping tick, found that libuv wrote no byte of it. + describe("when Claude answers each ping but stops reading", function() + local reply = string.rep("x", 262144) + + ---Claude's PONG reaches the server, then one interval passes and the ping timer runs. + local function pong_and_tick(tick, handle) + handle._read_cb(nil, frame.create_frame(frame.OPCODE.PONG, "", true, true)) + clock = clock + interval + tick() + end + + it("closes the client once the write queue stalls for two intervals, and frees the queue", function() + local tick = start_ping_timer() + local client, writes, handle = connect_and_upgrade() + + -- Every interval Neovim queues another reply, and Claude reads none of them. + for _ = 1, 2 do + tcp.send_to_client(server, client.id, reply) + pong_and_tick(tick, handle) + end + assert.is_not_nil(server.clients[client.id], "the client was dropped before the window ended") + assert.is_true(handle:get_write_queue_size() > 2 * #reply) + + tcp.send_to_client(server, client.id, reply) + pong_and_tick(tick, handle) + + assert.is_nil(server.clients[client.id]) + assert.spy(callbacks.on_disconnect).was_called(1) + assert.spy(callbacks.on_disconnect).was_called_with(client, 1006, "Write queue stalled") + assert.equals(1011, close_status(writes[#writes])) + assert.is_true(handle:is_closing(), "the socket still holds the queue") + end) + + it("reports the stall as one info line", function() + local logger = require("claudecode.logger") + local lines = {} + local saved_one_line = logger.one_line + logger.one_line = function(level, component, text) + table.insert(lines, { level = level, component = component, text = text }) + end + local tick = start_ping_timer() + local client, _, handle = connect_and_upgrade() + tcp.send_to_client(server, client.id, reply) + + for _ = 1, 3 do + pong_and_tick(tick, handle) + end + logger.one_line = saved_one_line + + assert.same({ + { + level = "info", + component = "server", + text = "Closed " .. client.id .. ": not reading for 60s, 256 KB queued", + }, + }, lines) + end) + + it("keeps a client that reads slowly, although its queue keeps growing", function() + local tick = start_ping_timer() + local client, _, handle = connect_and_upgrade() + + -- Neovim queues a reply every interval and the kernel takes 1 KB of the queue: Claude reads, but slowly. + for _ = 1, 20 do + tcp.send_to_client(server, client.id, reply) + handle.drain(1024) + pong_and_tick(tick, handle) + end + + assert.is_not_nil(server.clients[client.id], "a Claude that reads was dropped") + assert.spy(callbacks.on_disconnect).was_not_called() + assert.is_true(handle:get_write_queue_size() > 19 * #reply) + end) + + it("keeps a client that is sent one large message, however long it takes to read", function() + local tick = start_ping_timer() + local client, _, handle = connect_and_upgrade() + local large = string.rep("x", 16 * 1024 * 1024) + tcp.send_to_client(server, client.id, large) + + -- Claude reads 1 MB per interval, so the message needs 16 intervals, many times the window. + for _ = 1, 20 do + handle.drain(1024 * 1024) + pong_and_tick(tick, handle) + end + + assert.is_not_nil(server.clients[client.id], "a Claude that reads one large message was dropped") + assert.spy(callbacks.on_disconnect).was_not_called() + assert.is_true(handle:get_write_queue_size() < 1024, "only the pings of the last intervals are left") + end) + + it("counts the window from the last byte Claude read, not from the connection", function() + local tick = start_ping_timer() + local client, _, handle = connect_and_upgrade() + tcp.send_to_client(server, client.id, reply) + + -- Claude reads for five intervals. Then it stops. + for _ = 1, 5 do + handle.drain(1024) + pong_and_tick(tick, handle) + end + assert.is_not_nil(server.clients[client.id]) + pong_and_tick(tick, handle) -- no progress for one interval + assert.is_not_nil(server.clients[client.id], "the window ended after one interval") + pong_and_tick(tick, handle) -- no progress for two + + assert.is_nil(server.clients[client.id]) + end) + + it("keeps a client that reads every ping, so that its queue is empty at each tick", function() + local tick = start_ping_timer() + local client, _, handle = connect_and_upgrade() + + for _ = 1, 10 do + handle.drain(math.huge) + pong_and_tick(tick, handle) + end + + assert.is_not_nil(server.clients[client.id]) + assert.spy(callbacks.on_disconnect).was_not_called() + end) + + it("gives the client a fresh window after Neovim slept, as it does for the keepalive", function() + local tick = start_ping_timer() + local client, _, handle = connect_and_upgrade() + tcp.send_to_client(server, client.id, reply) + pong_and_tick(tick, handle) + + -- The machine sleeps for ten minutes. The timer runs first when it wakes, before Claude could read. + clock = clock + 600000 + tick() + assert.is_not_nil(server.clients[client.id], "the sleep counted against the client") + + pong_and_tick(tick, handle) + assert.is_not_nil(server.clients[client.id], "the window did not restart at the wake") + pong_and_tick(tick, handle) + + assert.is_nil(server.clients[client.id]) + end) + end) + + -- Two rules judge a connected client on each ping tick. The pong rule gives a client without a pong one more tick + -- (`pong_overdue`) and then drops it. The stall rule drops a client that answers, whose write queue made no progress at + -- two observations in a row. The pong rule runs first, so only a client that answers is judged for a stall. + describe("when a client misses a pong and its write queue stalls", function() + local reply = string.rep("x", 262144) + + local function pong(handle) + handle._read_cb(nil, frame.create_frame(frame.OPCODE.PONG, "", true, true)) + end + + ---Start a client whose queue never drains, which answers the first tick and the third, and not the second. + local function stalled_client_that_misses_a_pong() + local tick = start_ping_timer() + local client, writes, handle = connect_and_upgrade() + tcp.send_to_client(server, client.id, reply) + pong(handle) + clock = clock + interval + tick() -- the 101 reply was written: that is progress, and the last one + clock = clock + interval + tick() -- no pong for two intervals: the second chance + return tick, client, writes, handle + end + + it("gives a silent client its second chance although its queue stalled too, then drops it for the pong", function() + local client_manager = require("claudecode.server.client") + local tick = start_ping_timer() + local client, writes, handle = connect_and_upgrade() + tcp.send_to_client(server, client.id, reply) + -- No pong for two intervals and no byte written for two ticks: both rules find this client dead at the first tick. + client.last_pong = clock - 2 * interval + client.last_drain = clock - 3 * interval + client.drained_bytes = client.queued_bytes - handle:get_write_queue_size() + client.write_stalled = true -- the tick before found no progress either + assert.is_false(client_manager.is_client_alive(client, 2 * interval)) + assert.is_false(client_manager.is_client_reading(client)) + + tick() + + assert.is_not_nil(server.clients[client.id], "the stall skipped the second chance of the pong rule") + assert.is_true(client.pong_overdue) + assert.equals(frame.OPCODE.PING, parse(writes[#writes]).opcode) + + clock = clock + interval + tick() + + assert.is_nil(server.clients[client.id]) + assert.spy(callbacks.on_disconnect).was_called(1) + assert.spy(callbacks.on_disconnect).was_called_with(client, 1006, "Connection timeout") + assert.equals(1011, close_status(writes[#writes])) + assert.equals("Keepalive timeout", parse(writes[#writes]).payload:sub(3)) + assert.is_true(handle:is_closing()) + end) + + it("drops a client for the stall when its pong arrives, because the pong does not save a stalled queue", function() + local tick, client, writes, handle = stalled_client_that_misses_a_pong() + assert.is_true(client.pong_overdue) + + pong(handle) -- the late pong is read, but nothing of the queue is written + clock = clock + interval + tick() + + assert.is_nil(client.pong_overdue, "the tick that found the pong kept the flag") + assert.is_nil(server.clients[client.id]) + assert.spy(callbacks.on_disconnect).was_called(1) + assert.spy(callbacks.on_disconnect).was_called_with(client, 1006, "Write queue stalled") + assert.equals("Not reading", parse(writes[#writes]).payload:sub(3)) + end) + + it("keeps a client whose pong arrives while its queue drains, and clears the flag", function() + local tick = start_ping_timer() + local client, _, handle = connect_and_upgrade() + tcp.send_to_client(server, client.id, reply) + pong(handle) + clock = clock + interval + tick() + clock = clock + interval + tick() -- no pong: the second chance + assert.is_true(client.pong_overdue) + + handle.drain(1024) -- Claude reads a little, and answers + pong(handle) + clock = clock + interval + tick() + + assert.is_nil(client.pong_overdue) + assert.is_not_nil(server.clients[client.id]) + assert.spy(callbacks.on_disconnect).was_not_called() + end) + + it("clears both flags and restarts the stall window when the machine wakes from sleep", function() + local tick, client, _, handle = stalled_client_that_misses_a_pong() + assert.is_true(client.pong_overdue) + assert.is_true(client.write_stalled) + + clock = clock + 600000 -- the machine slept for ten minutes, and the timer runs first when it wakes + tick() + assert.is_nil(client.pong_overdue, "the wake kept the flag") + assert.is_nil(client.write_stalled, "the wake kept the stall count") + assert.is_not_nil(server.clients[client.id], "the sleep counted against the client") + + pong(handle) + clock = clock + interval + tick() + assert.is_not_nil(server.clients[client.id], "the stall window did not restart at the wake") + pong(handle) + clock = clock + interval + tick() + + assert.is_nil(server.clients[client.id]) + assert.spy(callbacks.on_disconnect).was_called_with(client, 1006, "Write queue stalled") + end) + end) + + -- libuv reports that a socket is writable again only in a poll phase. A main-loop block inside an I/O callback can + -- make a ping tick run before that poll, so the tick finds no write progress although Claude read during the block. + -- A tick that would drop a client for a stall therefore arms a check handle, which libuv runs right after its next + -- poll, and only the judgement there drops the client. The stall count (`write_stalled`) is kept by the ticks: the + -- second tick in a row without progress asks for the confirmation. A wake tick, a gap above 1.5 intervals, starts a + -- fresh stall window, and the two ticks after it always count. + describe("when a stall is confirmed after the next poll", function() + local pong + + before_each(function() + pong = frame.create_frame(frame.OPCODE.PONG, "", true, true) + end) + + ---A client with a queued reply that Claude does not read, and one tick that saw the 101 reply as progress. + local function stalled_client() + local tick, raw_tick = start_ping_timer() + local client, writes, handle = connect_and_upgrade() + tcp.send_to_client(server, client.id, string.rep("x", 262144)) + handle._read_cb(nil, pong) + clock = clock + interval + tick() + return tick, raw_tick, client, handle, writes + end + + ---Claude answers, `gap` intervals pass, and `run` runs the tick. + local function answer_and(run, gap) + return function(handle) + handle._read_cb(nil, pong) + clock = clock + gap * interval + run() + end + end + + it("keeps a reading client until libuv can report the bytes it read during the block", function() + local tick, raw_tick, client, handle = stalled_client() + answer_and(tick, 1)(handle) -- no write progress for one interval: the first tick that counts + assert.is_true(client.write_stalled) + + -- A libuv I/O callback blocks for 1.3 intervals. Claude reads during the block. The late tick runs before the next + -- poll can account for the newly writable socket, while the last pong is still within its window. + answer_and(raw_tick, 1.3)(handle) + assert.is_not_nil(server.clients[client.id], "the late tick dropped a reader before libuv reported its progress") + poll(function() + handle.drain(1024) -- the poll phase reports the bytes that Claude read + end) + + assert.is_not_nil(server.clients[client.id], "the judgement after the poll dropped a reader that made progress") + assert.is_nil(client.write_stalled) + assert.spy(callbacks.on_disconnect).was_not_called() + end) + + it("keeps a reader whose bytes libuv reports only at the poll after a tick that came on time", function() + local tick, raw_tick, client, handle = stalled_client() + answer_and(tick, 1)(handle) + assert.is_true(client.write_stalled) + + answer_and(raw_tick, 1)(handle) -- a block ended just after this tick was due: on time, and before the poll + assert.is_not_nil(server.clients[client.id], "the tick dropped the client before the next poll") + poll(function() + handle.drain(1) + end) + + assert.is_not_nil(server.clients[client.id]) + assert.is_nil(client.write_stalled) + end) + + it("drops a client at the poll after its second tick in a row without progress, not at the tick", function() + local tick, raw_tick, client, handle, writes = stalled_client() + answer_and(tick, 1)(handle) + assert.is_true(client.write_stalled) + + answer_and(raw_tick, 1)(handle) -- the second tick without progress, and no poll phase yet + assert.is_not_nil(server.clients[client.id], "the tick dropped the client before a fresh poll") + assert.spy(callbacks.on_disconnect).was_not_called() + + poll() -- the poll reports nothing: the stall is confirmed + assert.is_nil(server.clients[client.id]) + assert.spy(callbacks.on_disconnect).was_called_with(client, 1006, "Write queue stalled") + assert.equals("Not reading", parse(writes[#writes]).payload:sub(3)) + end) + + it("does not ask for a confirmation at the first tick without progress", function() + local tick, _, client, handle = stalled_client() + + answer_and(tick, 1)(handle) + + assert.is_true(client.write_stalled) + assert.is_not_nil(server.clients[client.id]) + assert.equals(0, #checks, "the first tick armed a check handle") + end) + + it("clears the count when a tick finds progress", function() + local tick, _, client, handle = stalled_client() + answer_and(tick, 1)(handle) + assert.is_true(client.write_stalled) + + handle.drain(1024) + answer_and(tick, 1)(handle) + assert.is_nil(client.write_stalled) + + answer_and(tick, 1)(handle) -- the first tick without progress again + assert.is_not_nil(server.clients[client.id]) + assert.is_true(client.write_stalled) + end) + + it("still drops a stalled client at its second tick when every tick comes late", function() + local tick, _, client, handle = stalled_client() + + answer_and(tick, 1.06)(handle) -- 6% late, as under a loop that is always a little busy + assert.is_true(client.write_stalled) + answer_and(tick, 1.4)(handle) -- below the wake threshold of 1.5 intervals + + assert.is_nil(server.clients[client.id], "late ticks kept a stalled client") + assert.spy(callbacks.on_disconnect).was_called_with(client, 1006, "Write queue stalled") + end) + + it("judges only the clients that answered at the tick", function() + local tick, _, client, handle = stalled_client() + answer_and(tick, 1)(handle) + assert.is_true(client.write_stalled) + + -- No pong for two intervals: the pong rule gives its second chance, and the stall rule leaves the client alone. + client.last_pong = clock - 2 * interval + clock = clock + interval + tick() + + assert.is_not_nil(server.clients[client.id]) + assert.is_true(client.pong_overdue) + assert.equals(0, #checks, "a client that did not answer was confirmed for a stall") + end) + + it("makes no stall judgement at the tick that finds the machine woken from sleep, and clears the count", function() + local tick, _, client, handle = stalled_client() + answer_and(tick, 1)(handle) + assert.is_true(client.write_stalled) + + clock = clock + 600000 -- the machine slept for ten minutes + tick() + + assert.is_nil(client.write_stalled, "the wake kept the count") + assert.is_not_nil(server.clients[client.id], "the sleep counted against the client") + assert.equals(0, #checks) + end) + + it("counts the two ticks after a wake tick, so that a run of wake ticks keeps no stalled client", function() + local tick, _, client = stalled_client() + + clock = clock + 600000 -- a wake tick: it starts a fresh stall window + tick() + assert.is_nil(client.write_stalled, "the wake tick kept the count") + assert.is_not_nil(server.clients[client.id]) + + clock = clock + 600000 -- the loop was blocked again, for more than 1.5 intervals: no sleep this time + tick() + assert.is_true(client.write_stalled, "the second wake tick in a row did not count") + assert.is_not_nil(server.clients[client.id]) + + clock = clock + 600000 + tick() -- the second tick in a row without progress: the poll confirms it + assert.is_nil(server.clients[client.id], "a run of wake ticks kept a stalled client") + assert.spy(callbacks.on_disconnect).was_called_with(client, 1006, "Write queue stalled") + end) + + it("counts a wake tick that follows a tick in time, so that alternating ticks keep no stalled client", function() + local tick, _, client, handle = stalled_client() + + clock = clock + 600000 + tick() -- wake tick: a fresh stall window + answer_and(tick, 1)(handle) -- a tick in time: the first without progress + assert.is_true(client.write_stalled) + + clock = clock + 600000 + tick() -- a wake tick again, two ticks after the fresh window started: it counts + + assert.is_nil(server.clients[client.id], "alternating wake ticks kept a stalled client") + assert.spy(callbacks.on_disconnect).was_called_with(client, 1006, "Write queue stalled") + end) + + it("starts a fresh stall window at a wake tick three ticks after the last one", function() + local tick, _, client, handle = stalled_client() + + clock = clock + 600000 + tick() -- wake tick: a fresh stall window + handle.drain(1) + answer_and(tick, 1)(handle) + handle.drain(1) + answer_and(tick, 1)(handle) + + clock = clock + 600000 + tick() -- a wake tick three ticks later: the machine slept again + + assert.is_nil(client.write_stalled) + assert.equals(clock, client.last_drain, "the second sleep did not start a fresh window") + end) + + it( + "reports the time since the last progress, not since before the sleep, when it drops the client after a wake", + function() + local logger = require("claudecode.logger") + local lines = {} + local saved_one_line = logger.one_line + logger.one_line = function(_, _, text) + table.insert(lines, text) + end + local tick, _, client, handle = stalled_client() + + clock = clock + 600000 -- the machine slept for ten minutes + tick() + answer_and(tick, 1)(handle) -- the first tick without progress after the wake + answer_and(tick, 1)(handle) -- the second, and the poll confirms it + logger.one_line = saved_one_line + + assert.is_nil(server.clients[client.id]) + assert.same({ "Closed " .. client.id .. ": not reading for 60s, 256 KB queued" }, lines) + end + ) + + describe("with the check and idle handles of a confirmation", function() + ---Arm a confirmation: a client whose second tick in a row found no progress. + local function pending_confirmation() + local tick, raw_tick, client, handle, writes = stalled_client() + answer_and(tick, 1)(handle) + answer_and(raw_tick, 1)(handle) + assert.equals(1, #checks) + assert.equals(1, #idles) + return client, handle, writes + end + + ---Run `poll` with the info lines of the logger collected. + local function logged_info_during(run) + local logger = require("claudecode.logger") + local lines, saved_one_line = {}, logger.one_line + logger.one_line = function(_, _, text) + table.insert(lines, text) + end + local ok, err = pcall(run) + logger.one_line = saved_one_line + assert(ok, err) + return lines + end + + it("arms an idle handle with the check handle, and closes both once the check ran", function() + pending_confirmation() + assert.is_true(checks[1].active and idles[1].active, "the idle handle does not keep the poll from waiting") + + poll(function() end) + + assert.is_true(checks[1].closed and idles[1].closed, "a handle stays open after its check ran") + assert.is_false(checks[1].active or idles[1].active) + end) + + it("leaves no handle open after the client was dropped", function() + local client = pending_confirmation() + + poll() + + assert.is_nil(server.clients[client.id]) + assert.is_true(checks[1].closed and idles[1].closed) + end) + + it("leaves no handle open after the client made progress", function() + local _, handle = pending_confirmation() + + poll(function() + handle.drain(1) + end) + + assert.is_true(checks[1].closed and idles[1].closed) + end) + + it("removes the handles of a confirmation from the pending ones once it ran", function() + pending_confirmation() + assert.is_not_nil(next(server.pending_polls), "an armed confirmation is not listed") + + poll() + + assert.is_nil(next(server.pending_polls), "a finished confirmation stays listed until stop_server") + end) + + it("skips a client that broke the protocol in the poll of its confirmation, and logs no stall", function() + local client, handle, writes = pending_confirmation() + + local lines = logged_info_during(function() + poll(function() + handle._read_cb(nil, frame.create_frame(frame.OPCODE.BINARY, "x", true, true)) + end) + end) + + assert.same({}, lines, "a client that already left was logged as not reading") + assert.spy(callbacks.on_disconnect).was_called(1) + assert.equals("Client error: Binary messages not supported", callbacks.on_disconnect.calls[1].vals[3]) + assert.equals(1003, close_status(writes[#writes]), "the Close frame of the protocol error was replaced") + assert.is_false(handle:is_closing(), "the confirmation closed the socket before its Close frame was written") + assert.is_nil(server.clients[client.id]) + end) + + it("skips a listed peer-closing client until its scheduled close callback runs", function() + local client, handle = pending_confirmation() + scheduled = {} + local lines = logged_info_during(function() + poll(function() + handle._read_cb(nil, frame.create_frame(frame.OPCODE.CLOSE, string.char(3, 232) .. "Done", true, true)) + assert.equals("closing", client.state) + assert.equals(client, server.clients[client.id], "the peer Close must wait for vim.schedule") + end) + end) + assert.same({}, lines, "a peer-closing client was logged as stalled") + assert.equals(client, server.clients[client.id]) + assert.spy(callbacks.on_disconnect).was_not_called() + assert.is_true(checks[1].closed and idles[1].closed) + run_scheduled() + assert.spy(callbacks.on_disconnect).was_called(1) + assert.spy(callbacks.on_disconnect).was_called_with(client, 1000, "Done") + assert.is_nil(server.clients[client.id]) + end) + + it("skips a client that hung up in the poll of its confirmation, and logs no stall", function() + local client, handle = pending_confirmation() + + local lines = logged_info_during(function() + poll(function() + handle._read_cb(nil, nil) -- Claude hangs up: the poll reads the EOF + end) + end) + + assert.same({}, lines, "a client that already left was logged as not reading") + assert.spy(callbacks.on_disconnect).was_called(1) + assert.spy(callbacks.on_disconnect).was_called_with(client, 1006, "EOF") + end) + + it("closes the handles of a confirmation in stop_server, and the confirmation never runs", function() + local client = pending_confirmation() + + tcp.stop_server(server) + + assert.is_true(checks[1].closed and idles[1].closed, "stop_server left a handle open") + assert.is_nil(server.pending_polls) + poll() -- nothing runs, and nothing fails + assert.spy(callbacks.on_disconnect).was_not_called() + assert.is_not_nil(client) + tcp.stop_server(server) -- a second stop leaves nothing to close + end) + + it( + "closes both handles and leaves the client for the next tick when libuv cannot create the idle handle", + function() + vim.loop.new_idle = function() + return nil + end + local tick, raw_tick, client, handle = stalled_client() + answer_and(tick, 1)(handle) + + answer_and(raw_tick, 1)(handle) + + assert.equals(1, #checks) + assert.is_true(checks[1].closed, "the check handle was left open") + assert.is_not_nil(server.clients[client.id], "a failed handle dropped the client") + end + ) + end) + + describe("is_client_reading", function() + local client_manager + + before_each(function() + client_manager = require("claudecode.server.client") + end) + + local function silent_client() + local client, _, handle = connect_and_upgrade() + tcp.send_to_client(server, client.id, string.rep("x", 262144)) + client_manager.is_client_reading(client) -- sees the 101 reply + return client, handle + end + + it("says false only at the second check in a row that found no progress", function() + local client = silent_client() + assert.is_true(client_manager.is_client_reading(client)) + assert.is_false(client_manager.is_client_reading(client)) + end) + + it("clears the count when a check finds progress", function() + local client, handle = silent_client() + assert.is_true(client_manager.is_client_reading(client)) + assert.is_true(client.write_stalled) + handle.drain(1) + assert.is_true(client_manager.is_client_reading(client)) + assert.is_nil(client.write_stalled) + end) + + it("counts an empty queue as progress", function() + local client, handle = silent_client() + assert.is_true(client_manager.is_client_reading(client)) + handle.drain(math.huge) + assert.is_true(client_manager.is_client_reading(client)) + assert.is_nil(client.write_stalled) + end) + end) + end) + + -- `queued_bytes` must count every frame that the server writes, or a reader that makes progress looks stalled. A + -- reader that drains less than the frames queued behind its reads loses its progress in the count. + describe("when the server counts the frames it writes", function() + local reply = string.rep("x", 262144) + + it("counts the pings it sends, while a reader drains one byte per tick", function() + local tick = start_ping_timer() + local client, _, handle = connect_and_upgrade() + tcp.send_to_client(server, client.id, reply) + for _ = 1, 8 do + handle.drain(1) + handle._read_cb(nil, frame.create_frame(frame.OPCODE.PONG, "", true, true)) + clock = clock + interval + tick() + assert.is_not_nil(server.clients[client.id], "a reader was dropped despite one byte of progress") + end + end) + + it("counts the pongs it sends, while a reader drains less than the payload of a pong", function() + local tick = start_ping_timer() + local client, _, handle = connect_and_upgrade() + tcp.send_to_client(server, client.id, reply) + for _ = 1, 8 do + handle._read_cb(nil, frame.create_frame(frame.OPCODE.PING, string.rep("p", 125), true, true)) + handle.drain(64) + handle._read_cb(nil, frame.create_frame(frame.OPCODE.PONG, "", true, true)) + clock = clock + interval + tick() + assert.is_not_nil(server.clients[client.id], "a reader was dropped despite 64 bytes of progress") + end + end) + + it("reports the write error once, to the callback, when the last piece of a frame fails", function() + local client, writes = connect_and_upgrade() + local callback = spy.new(function() end) + tcp.send_to_client(server, client.id, string.rep("x", 2 * 1024 * 1024), callback) + for index = 2, #writes - 1 do + writes[index].done() + end + assert.spy(callback).was_not_called() + + writes[#writes].done("EPIPE") + + assert.spy(callback).was_called(1) + assert.spy(callback).was_called_with("EPIPE") + end) + + it("counts a broadcast to a client as failed when the last piece of its frame fails", function() + local _, writes = connect_and_upgrade() + local on_written = spy.new(function() end) + tcp.broadcast(server, string.rep("x", 2 * 1024 * 1024), on_written) + for index = 2, #writes - 1 do + writes[index].done() + end + assert.spy(on_written).was_not_called() + + writes[#writes].done("EPIPE") + + assert.spy(on_written).was_called(1) + assert.spy(on_written).was_called_with(0) + end) + end) + + -- libuv lowers the queue size of a socket by the bytes it wrote on Linux and macOS, but only by a whole write request + -- when that request is done on Windows. A large frame is queued in pieces, so progress shows after each piece. + describe("when a message is larger than one piece of a write", function() + local piece = 1024 * 1024 + + it("queues it in pieces that make up the frame, and calls back once, after the last piece", function() + local client, writes = connect_and_upgrade() + local message = string.rep("abcdefgh", 3 * piece / 8 + 5) -- three pieces and a short rest + local expected = frame.create_text_frame(message) + local callback = spy.new(function() end) + + tcp.send_to_client(server, client.id, message, callback) + + local pieces = {} + for index = 2, #writes do -- writes[1] is the 101 reply + table.insert(pieces, writes[index].data) + assert.is_true(#writes[index].data <= piece) + end + assert.equals(4, #pieces) + assert.is_true(expected == table.concat(pieces), "the pieces are not the frame") + for index = 2, #writes - 1 do + writes[index].done() + end + assert.spy(callback).was_not_called() + writes[#writes].done() + assert.spy(callback).was_called(1) + end) + + -- A server frame with a payload of 64 KiB or more has a 10-byte header, so a payload of `size - 10` bytes makes a + -- frame of exactly `size` bytes. A frame of a whole number of pieces plus one byte ends in a piece of one byte, and + -- only that piece may call back. + for _, case in ipairs({ + { size = piece, pieces = 1 }, + { size = piece + 1, pieces = 2 }, + { size = 2 * piece + 1, pieces = 3 }, + }) do + it("calls back once for a frame of " .. case.size .. " bytes", function() + local client, writes = connect_and_upgrade() + local callback = spy.new(function() end) + + tcp.send_to_client(server, client.id, string.rep("x", case.size - 10), callback) + + assert.equals(1 + case.pieces, #writes, "the 101 reply and the pieces") + local total = 0 + for index = 2, #writes do + total = total + #writes[index].data + end + assert.equals(case.size, total) + for index = 2, #writes - 1 do + writes[index].done() + end + assert.spy(callback).was_not_called() + writes[#writes].done() + assert.spy(callback).was_called(1) + end) + end + + it("queues a message of one piece or less as one write", function() + local client, writes = connect_and_upgrade() + + tcp.send_to_client(server, client.id, string.rep("x", piece - 16)) + + assert.equals(2, #writes) + end) + + it("keeps a client whose reader takes several intervals for a message, as libuv counts it on Windows", function() + local tick = start_ping_timer() + local client, writes, handle = connect_and_upgrade() + tcp.send_to_client(server, client.id, string.rep("x", 8 * piece)) + + -- Claude reads 2 MB per interval. A write counts as done once Claude read all of it. + local credit = 0 + for _ = 1, 6 do + credit = credit + 2 * piece + for _, write in ipairs(writes) do + if write.left > 0 then + if #write.data > credit then + break + end + credit = credit - #write.data + write.done() + end + end + handle._read_cb(nil, frame.create_frame(frame.OPCODE.PONG, "", true, true)) + clock = clock + interval + tick() + end + + assert.is_not_nil(server.clients[client.id], "a Claude that reads the message was dropped") + assert.spy(callbacks.on_disconnect).was_not_called() + end) + end) + + -- The ping timer sees only the clients in `server.clients`. A client that left it with its Close frame still queued + -- is covered by the deadline of that frame: the same window, and the same rule, no write progress for two intervals. + describe("when a Close frame waits in the queue of a Claude that does not read", function() + local reply = string.rep("x", 262144) + local function binary() + return frame.create_frame(frame.OPCODE.BINARY, "x", true, true) + end + + it( + "closes the socket of a client that broke the protocol once the queue made no progress for the window", + function() + start_ping_timer() + local client, _, handle = connect_and_upgrade() + tcp.send_to_client(server, client.id, reply) + + handle._read_cb(nil, binary()) + assert.is_nil(server.clients[client.id], "the ping timer would still see the client") + assert.is_false(handle:is_closing(), "the socket was closed before the Close frame was written") + local deadline = close_deadline() + + run_after_window(deadline) + + assert.is_true(handle:is_closing(), "the socket still holds the queue") + assert.equals("closed", client.state) + assert.is_true(deadline.stopped and deadline.closed, "the deadline keeps running") + assert.spy(callbacks.on_disconnect).was_called(1) + end + ) + + it( + "closes the socket of a client that the server stopped, once the queue made no progress for the window", + function() + start_ping_timer() + local client, _, handle = connect_and_upgrade() + tcp.send_to_client(server, client.id, reply) + + tcp.stop_server(server) + assert.is_false(handle:is_closing(), "the socket was closed before the Close frame was written") + local deadline = close_deadline() + + run_after_window(deadline) + + assert.is_true(handle:is_closing(), "the socket still holds the queue") + assert.equals("closed", client.state) + assert.is_true(deadline.stopped and deadline.closed, "the deadline keeps running") + end + ) + + it("waits while Claude reads, and the write callback then closes the socket and ends the deadline", function() + start_ping_timer() + local client, writes, handle = connect_and_upgrade() + tcp.send_to_client(server, client.id, reply) + handle._read_cb(nil, binary()) + local deadline = close_deadline() + + handle.drain(1024) -- Claude reads a little: that is progress + run_after_window(deadline) + assert.is_false(handle:is_closing(), "a Claude that reads was cut off") + assert.is_false(deadline.stopped) + + handle.drain(math.huge) -- Claude reads the rest, so the Close frame is written + assert.is_true(handle:is_closing()) + assert.is_true(deadline.stopped and deadline.closed, "the deadline outlives the write") + assert.equals(1003, close_status(writes[#writes])) + + -- A deadline that runs after the socket closed does nothing: closing a handle twice is an error in libuv. + deadline.callback() + assert.equals("closed", client.state) + end) + + it("closes the socket once when libuv cancels the Close frame's write after the deadline closed it", function() + start_ping_timer() + local client, writes, handle = connect_and_upgrade() + tcp.send_to_client(server, client.id, reply) + handle._read_cb(nil, binary()) + local deadline = close_deadline() + + run_after_window(deadline) + assert.is_true(handle:is_closing()) + + -- libuv runs the callback of each write still queued, with ECANCELED, once the handle closes. + writes[#writes].done("ECANCELED") + assert.equals("closed", client.state) + end) + + it("does not wait for a client that the ping timer dropped, whose socket closes at once", function() + local tick = start_ping_timer() + local client, _, handle = connect_and_upgrade() + tcp.send_to_client(server, client.id, reply) + + for _ = 1, 3 do + handle._read_cb(nil, frame.create_frame(frame.OPCODE.PONG, "", true, true)) + clock = clock + interval + tick() + end + + assert.is_nil(server.clients[client.id]) + assert.is_true(handle:is_closing()) + local deadline = close_deadline() + assert.is_true(deadline.stopped and deadline.closed, "the deadline outlives the drop") + deadline.callback() + end) + + it("arms nothing when the Close frame was written at once", function() + local handle = { closed = false } + function handle:write(_, callback) + if callback then + callback() + end + end + function handle:get_write_queue_size() + return 0 + end + function handle:close() + assert(not self.closed, "handle is already closing") + self.closed = true + end + function handle:is_closing() + return self.closed + end + local client = require("claudecode.server.client").create_client(handle, 2 * interval) + client.state, client.handshake_complete = "connected", true + + require("claudecode.server.client").close_client(client, 1000, "") + + assert.equals(0, #timers) + assert.is_true(handle.closed) + end) + + it("arms nothing for a client that has no keepalive window", function() + local client, _, handle = connect_and_upgrade() -- no ping timer was started, so the server has no window + tcp.send_to_client(server, client.id, reply) + + tcp.stop_server(server) + + assert.equals(0, #timers) + assert.is_false(handle:is_closing()) + end) + end) end) diff --git a/tests/unit/server/frame_spec.lua b/tests/unit/server/frame_spec.lua index 317eeac44..a514c0c14 100644 --- a/tests/unit/server/frame_spec.lua +++ b/tests/unit/server/frame_spec.lua @@ -189,6 +189,7 @@ describe("WebSocket client malformed frame handling", function() handshake_complete = true, last_ping = 0, last_pong = 0, + queued_bytes = 0, } return c, writes, handle diff --git a/tests/unit/server/message_spec.lua b/tests/unit/server/message_spec.lua index 42bddd363..6598ed422 100644 --- a/tests/unit/server/message_spec.lua +++ b/tests/unit/server/message_spec.lua @@ -26,7 +26,14 @@ describe("WebSocket client messages", function() function handle:is_closing() return self.closed end - c = { id = "test_client", tcp_handle = handle, state = "connected", buffer = "", handshake_complete = true } + c = { + id = "test_client", + tcp_handle = handle, + state = "connected", + buffer = "", + handshake_complete = true, + queued_bytes = 0, + } end) -- The client masks every frame it sends. diff --git a/tests/unit/server/send_utf8_spec.lua b/tests/unit/server/send_utf8_spec.lua index 5d19f2122..d7f106882 100644 --- a/tests/unit/server/send_utf8_spec.lua +++ b/tests/unit/server/send_utf8_spec.lua @@ -12,6 +12,7 @@ describe("WebSocket client send_message", function() local written = {} local client = { state = "connected", + queued_bytes = 0, tcp_handle = { write = function(_, data) table.insert(written, data) diff --git a/tests/unit/server/tcp_spec.lua b/tests/unit/server/tcp_spec.lua index 8141106a2..5ae0409b6 100644 --- a/tests/unit/server/tcp_spec.lua +++ b/tests/unit/server/tcp_spec.lua @@ -337,6 +337,9 @@ describe("TCP server disconnect handling", function() function handle:is_closing() return self.closed end + function handle:get_write_queue_size() + return 0 + end local client = client_manager.create_client(handle) client.state, client.handshake_complete, client.last_pong = "connected", true, vim.loop.now() - 61000 local server = { clients = { [client.id] = client }, on_disconnect = spy.new(function() end) } diff --git a/tests/unit/server_spec.lua b/tests/unit/server_spec.lua index aecd3acbf..e61307490 100644 --- a/tests/unit/server_spec.lua +++ b/tests/unit/server_spec.lua @@ -186,6 +186,7 @@ describe("WebSocket Server", function() server.state.server.clients[id] = { id = id, state = state, + queued_bytes = 0, tcp_handle = { write = function(_, data) table.insert(written, id) @@ -220,6 +221,7 @@ describe("WebSocket Server", function() server.state.server.clients[id] = { id = id, state = "connected", + queued_bytes = 0, tcp_handle = { write = function(_, _, callback) if sync_error ~= nil then