feat(websocketspec): fire BeforeDisconnect/AfterDisconnect hooks on close

Hooks fire from Connection.Close() once per connection. ConnectionManager
Shutdown now closes connections outside its lock so hooks can call back
into the manager. Update the single-transaction audit plan status.

pkg/common/tx_guard_test.go

@@ -102,10 +102,7 @@ func TestSpecHandlersDoNotQueryThePoolDirectly(t *testing.T) {
 // Anything else defined in a spec's hooks.go must have an Execute call site: an
 // unwired hook silently disables whatever is registered on it (resolvespec's
 // AfterRead skipped column-level security masking until it was wired).
-var unwiredHooks = map[string]string{
-	"websocketspec/BeforeDisconnect": "connection close is not hooked yet",
-	"websocketspec/AfterDisconnect":  "connection close is not hooked yet",
-}
+var unwiredHooks = map[string]string{}
 
 var hookConstRE = regexp.MustCompile(`(?m)^\s*([A-Z][A-Za-z0-9]*)\s+HookType\s*=`)
 

pkg/websocketspec/connection.go

@@ -193,12 +193,17 @@ func (cm *ConnectionManager) GetConnection(id string) (*Connection, bool) {
 func (cm *ConnectionManager) Shutdown() {
 	cm.cancel()
 
-	// Close all connections
-	cm.mu.Lock()
+	// Close all connections outside the lock: disconnect hooks may call back
+	// into the manager (Count, GetConnection).
+	cm.mu.RLock()
+	conns := make([]*Connection, 0, len(cm.connections))
 	for _, conn := range cm.connections {
+		conns = append(conns, conn)
+	}
+	cm.mu.RUnlock()
+	for _, conn := range conns {
 		conn.Close()
 	}
-	cm.mu.Unlock()
 }
 
 // ReadPump reads messages from the WebSocket connection
@@ -300,6 +305,14 @@ func (c *Connection) SendJSON(v interface{}) error {
 // Close closes the connection
 func (c *Connection) Close() {
 	c.closedOnce.Do(func() {
+		// Disconnect hooks fire once per connection, however it closes. They
+		// are best-effort: a failure is logged and never blocks the close.
+		hookCtx := c.disconnectHookContext()
+		c.logDisconnectHook(BeforeDisconnect, c.executeHook(func(r *HookRegistry) error { return r.Execute(BeforeDisconnect, hookCtx) }))
+		defer func() {
+			c.logDisconnectHook(AfterDisconnect, c.executeHook(func(r *HookRegistry) error { return r.Execute(AfterDisconnect, hookCtx) }))
+		}()
+
 		if c.cancel != nil {
 			c.cancel()
 		}
@@ -321,6 +334,34 @@ func (c *Connection) Close() {
 	})
 }
 
+// disconnectHookContext builds the hook context for the disconnect hooks. The
+// context is detached from the connection's cancellation so hooks that touch
+// the database still run after the connection context is cancelled.
+func (c *Connection) disconnectHookContext() *HookContext {
+	parent := c.ctx
+	if parent == nil {
+		parent = context.Background()
+	}
+	return &HookContext{
+		Context:    context.WithoutCancel(parent),
+		Handler:    c.handler,
+		Connection: c,
+	}
+}
+
+func (c *Connection) executeHook(run func(*HookRegistry) error) error {
+	if c.handler == nil || c.handler.hooks == nil {
+		return nil
+	}
+	return run(c.handler.hooks)
+}
+
+func (c *Connection) logDisconnectHook(hook HookType, err error) {
+	if err != nil {
+		logger.Warn("[WebSocketSpec] %s hook failed for connection %s: %v", hook, c.ID, err)
+	}
+}
+
 // AddSubscription adds a subscription to this connection
 func (c *Connection) AddSubscription(sub *Subscription) {
 	c.mu.Lock()

pkg/websocketspec/connection_test.go

@@ -594,3 +594,57 @@ func TestConnectionManager_CompleteLifecycle(t *testing.T) {
 	_, exists = cm.GetConnection("conn-1")
 	assert.False(t, exists)
 }
+
+func TestConnection_Close_FiresDisconnectHooksOnce(t *testing.T) {
+	h := &Handler{hooks: NewHookRegistry()}
+	var order []string
+	var ctxErrAfter error
+	h.hooks.Register(BeforeDisconnect, func(hc *HookContext) error {
+		order = append(order, "before")
+		assert.NoError(t, hc.Context.Err())
+		return nil
+	})
+	h.hooks.Register(AfterDisconnect, func(hc *HookContext) error {
+		order = append(order, "after")
+		ctxErrAfter = hc.Context.Err()
+		return nil
+	})
+
+	conn := createTestConnection("c1")
+	conn.handler = h
+
+	conn.Close()
+	conn.Close() // second close must not re-fire
+
+	assert.Equal(t, []string{"before", "after"}, order)
+	assert.NoError(t, ctxErrAfter, "hook context must survive the connection cancel")
+}
+
+func TestConnection_Close_HookErrorDoesNotBlockClose(t *testing.T) {
+	h := &Handler{hooks: NewHookRegistry()}
+	h.hooks.Register(BeforeDisconnect, func(*HookContext) error { return assert.AnError })
+
+	conn := createTestConnection("c2")
+	conn.handler = h
+	conn.Close()
+
+	assert.Error(t, conn.ctx.Err(), "connection must still be cancelled")
+}
+
+func TestConnectionManager_Shutdown_HookMayCallManager(t *testing.T) {
+	cm := NewConnectionManager(context.Background())
+	h := &Handler{hooks: NewHookRegistry(), connManager: cm}
+	h.hooks.Register(BeforeDisconnect, func(*HookContext) error { cm.Count(); return nil })
+
+	conn := createTestConnection("c3")
+	conn.handler = h
+	cm.connections[conn.ID] = conn
+
+	done := make(chan struct{})
+	go func() { cm.Shutdown(); close(done) }()
+	select {
+	case <-done:
+	case <-time.After(2 * time.Second):
+		t.Fatal("Shutdown deadlocked on a hook calling the manager")
+	}
+}