@@ -4,6 +4,7 @@ import type { ResolvedGatewayAuth } from "../auth.js";
44import { MAX_BUFFERED_BYTES } from "../server-constants.js" ;
55import {
66 attachGatewayWsForTest ,
7+ createGatewayWsTestLogger ,
78 createGatewayWsTestRequestContext ,
89 createGatewayWsTestSocket ,
910 createResolvedGatewayTokenAuth ,
@@ -47,18 +48,20 @@ async function connectTestWs(
4748 options ?: Partial < Parameters < typeof attachGatewayWsConnectionHandler > [ 0 ] > ;
4849 } = { } ,
4950) {
51+ const logWsControl = createGatewayWsTestLogger ( ) ;
5052 const connected = attachGatewayWsForTest ( {
5153 attach : attachGatewayWsConnectionHandler ,
5254 clients : params . clients ,
5355 headers : params . headers ,
5456 host : params . host ,
55- options : params . options ,
57+ options : { ... params . options , logWsControl : logWsControl as never } ,
5658 socket : params . socket ,
5759 } ) ;
5860 await waitForLazyMessageHandler ( ) ;
5961
6062 return {
6163 clients : connected . clients ,
64+ logWsControl,
6265 socket : connected . socket ,
6366 passed : firstAttachedHandlerParams ( ) ,
6467 } ;
@@ -197,6 +200,76 @@ describe("attachGatewayWsConnectionHandler", () => {
197200 expect ( socket . close ) . toHaveBeenCalledWith ( 1008 , "slow consumer" ) ;
198201 } ) ;
199202
203+ it ( "keeps handshake phase advancement monotonic" , async ( ) => {
204+ const { socket, logWsControl, passed } = await connectTestWs ( ) ;
205+ const handlerParams = passed as {
206+ advanceHandshakePhase : ( phase : string ) => void ;
207+ } ;
208+
209+ handlerParams . advanceHandshakePhase ( "auth_credentials_received" ) ;
210+ handlerParams . advanceHandshakePhase ( "auth_validated" ) ;
211+ handlerParams . advanceHandshakePhase ( "auth_credentials_received" ) ;
212+ socket . emit ( "close" , 1006 , Buffer . from ( "client disappeared" ) ) ;
213+
214+ const [ message , context ] = logWsControl . warn . mock . calls [ 0 ] as [ string , { phase ?: string } ] ;
215+ expect ( message ) . toContain ( "phase=auth_validated" ) ;
216+ expect ( context ) . toMatchObject ( { phase : "auth_validated" } ) ;
217+ } ) ;
218+
219+ it ( "includes the last completed handshake phase in pre-connect close logs" , async ( ) => {
220+ const { socket, logWsControl } = await connectTestWs ( ) ;
221+
222+ socket . emit ( "close" , 1006 , Buffer . from ( "client disappeared" ) ) ;
223+
224+ expect ( logWsControl . warn ) . toHaveBeenCalled ( ) ;
225+ const [ message , context ] = logWsControl . warn . mock . calls [ 0 ] as [ string , { phase ?: string } ] ;
226+ expect ( message ) . toContain ( "closed before connect" ) ;
227+ expect ( message ) . toContain ( "phase=ws_upgrade_started" ) ;
228+ expect ( context ) . toMatchObject ( { phase : "ws_upgrade_started" } ) ;
229+ } ) ;
230+
231+ it ( "includes the last completed handshake phase on preauth timeout logs" , async ( ) => {
232+ vi . useFakeTimers ( ) ;
233+ const { logWsControl } = await connectTestWs ( {
234+ options : { preauthHandshakeTimeoutMs : 100 } ,
235+ } ) ;
236+
237+ vi . advanceTimersByTime ( 150 ) ;
238+
239+ expect ( logWsControl . warn ) . toHaveBeenCalledWith ( expect . stringContaining ( "handshake timeout" ) ) ;
240+ expect ( logWsControl . warn ) . toHaveBeenCalledWith (
241+ expect . stringContaining ( "phase=ws_upgrade_started" ) ,
242+ ) ;
243+ } ) ;
244+
245+ it ( "omits handshake phase metadata after the connection is ready" , async ( ) => {
246+ const { socket, logWsControl, passed } = await connectTestWs ( ) ;
247+ const handlerParams = passed as {
248+ advanceHandshakePhase : ( phase : string ) => void ;
249+ setClient : ( client : never ) => boolean ;
250+ setHandshakeState : ( state : "pending" | "connected" | "failed" ) => void ;
251+ } ;
252+
253+ handlerParams . advanceHandshakePhase ( "auth_credentials_received" ) ;
254+ handlerParams . advanceHandshakePhase ( "auth_validated" ) ;
255+ expect (
256+ handlerParams . setClient ( {
257+ socket,
258+ connect : { client : { id : "openclaw-control-ui" , mode : "webchat" } } ,
259+ connId : "ready-client" ,
260+ usesSharedGatewayAuth : false ,
261+ } as never ) ,
262+ ) . toBe ( true ) ;
263+ handlerParams . setHandshakeState ( "connected" ) ;
264+ handlerParams . advanceHandshakePhase ( "session_attached" ) ;
265+ handlerParams . advanceHandshakePhase ( "hello_payload_prepared" ) ;
266+ handlerParams . advanceHandshakePhase ( "ready" ) ;
267+
268+ socket . emit ( "close" , 1000 , Buffer . from ( "done" ) ) ;
269+
270+ expect ( logWsControl . warn ) . not . toHaveBeenCalled ( ) ;
271+ } ) ;
272+
200273 it ( "skips node presence disconnects for stale reconnected sockets" , async ( ) => {
201274 const unregister = vi . fn ( ( ) => null ) ;
202275 const { socket } = attachGatewayWsForTest ( {
0 commit comments