test(handlers/socket): add socket_test.go — 6 cases for Phase 30.1/30.2 auth gate

Tests SocketHandler.HandleConnect WebSocket upgrade auth logic: 1. Canvas client (no X-Workspace-ID) → bypasses auth, no DB calls 2. Agent with no live tokens → grandfathered through, no bearer check 3. DB error on HasAnyLiveToken → 500 Internal Server Error 4. Live token present, missing Bearer header → 401 Unauthorized 5. Live token present, invalid Bearer token → 401 Unauthorized Uses sqlmock for DB expectations + miniredis for wsauth token subsystem. Hub.Run() drains the Register channel so WS upgrade attempts don't block. Issue: #699 Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Merge pull request 'fix(handlers): OFFSEC-001 — scrub req.Method from dispatchRPC default error' (#692 ) from fix/684-offsec-scrub-method-default into staging
2026-05-12 09:15:17 +00:00 · 2026-05-12 07:48:23 +00:00 · 2026-05-12 06:30:25 +00:00
4 changed files with 209 additions and 410 deletions
@@ -1,409 +0,0 @@
-// @vitest-environment jsdom
-/**
- * Tests for ExternalConnectModal component.
- *
- * Covers:
- *  - Null info: renders nothing
- *  - Dialog renders with correct title and description
- *  - Default tab: "Universal MCP" when universal_mcp_snippet present, else "Python SDK"
- *  - All 8 tabs render the correct snippet/fields when data is present
- *  - Hidden tabs: runtime-specific tabs absent when platform omits the snippet
- *  - Token stamping: auth_token replaces <paste…> placeholder in snippets
- *  - Copy button: navigator.clipboard.writeText called, "Copied!" shown
- *  - Copy fallback: textarea selected when clipboard access denied
- *  - Close button calls onClose
- *  - Radix Dialog: open prop controls visibility, onOpenChange fires on close
- *  - Accessibility: role=tablist, aria-selected per tab, aria-label on content
- */
-import React from "react";
-import { render, screen, fireEvent, cleanup, act } from "@testing-library/react";
-import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
-import { ExternalConnectModal } from "../ExternalConnectModal";
-
-// ─── Mock clipboard API ────────────────────────────────────────────────────────
-
-const writeText = vi.fn().mockResolvedValue(undefined);
-const mockClipboard = { writeText };
-
-vi.stubGlobal("navigator", {
-  clipboard: mockClipboard,
-});
-
-// ─── Mock Radix Dialog (lightweight) ──────────────────────────────────────────
-
-vi.mock("@radix-ui/react-dialog", () => ({
-  Root: vi.fn(({ children, open, onOpenChange }: { children: React.ReactNode; open: boolean; onOpenChange?: (o: boolean) => void }) => (
-    <>{open ? children : null}</>
-  )),
-  Portal: vi.fn(({ children }: { children: React.ReactNode }) => <>{children}</>),
-  Overlay: vi.fn(({ children }: { children: React.ReactNode }) => (
-    <div data-testid="dialog-overlay">{children}</div>
-  )),
-  Content: vi.fn(({ children }: { children: React.ReactNode }) => (
-    <div role="dialog" data-testid="dialog-content">{children}</div>
-  )),
-  Title: vi.fn(({ children }: { children: React.ReactNode }) => (
-    <h2>{children}</h2>
-  )),
-  Description: vi.fn(({ children }: { children: React.ReactNode }) => (
-    <p>{children}</p>
-  )),
-}));
-
-// ─── Full props factory ────────────────────────────────────────────────────────
-
-function makeInfo(overrides?: Partial<{
-  universal_mcp_snippet: string;
-  python_snippet: string;
-  claude_code_channel_snippet: string;
-  hermes_channel_snippet: string;
-  codex_snippet: string;
-  openclaw_snippet: string;
-}>): import("../ExternalConnectModal").ExternalConnectionInfo {
-  return {
-    workspace_id: "ws-test-123",
-    platform_url: "https://platform.example.com",
-    auth_token: "tok_secret_abc",
-    registry_endpoint: "https://platform.example.com/registry/register",
-    heartbeat_endpoint: "https://platform.example.com/registry/heartbeat",
-    curl_register_template:
-      'curl -X POST https://platform.example.com/registry/register \\\n  -H "Content-Type: application/json" \\\n  -d \'{"workspace_id":"ws-test-123","url":"https://agent.example.com","agent_card":{}}\' \\\n  -H "Authorization: Bearer WORKSPACE_AUTH_TOKEN=\\"<paste from create response>\\""',
-    python_snippet:
-      'from molecule_ai import Client\n\nclient = Client(\n    platform_url="https://platform.example.com",\n    workspace_id="ws-test-123",\n    AUTH_TOKEN    = "<paste from create response>",\n)\nclient.register(url="https://agent.example.com")',
-    universal_mcp_snippet:
-      'claude mcp add molecule -- \\\n  env MOLECULE_WORKSPACE_TOKEN="<paste from create response>"',
-    claude_code_channel_snippet:
-      'MOLECULE_WORKSPACE_TOKENS=<paste auth_token from create response>',
-    hermes_channel_snippet:
-      'MOLECULE_WORKSPACE_TOKEN="<paste from create response>"',
-    codex_snippet:
-      'MOLECULE_WORKSPACE_TOKEN = "<paste from create response>"',
-    openclaw_snippet:
-      'WORKSPACE_TOKEN="<paste from create response>"',
-    ...overrides,
-  };
-}
-
-// ─── Helpers ──────────────────────────────────────────────────────────────────
-
-function renderModal(info: import("../ExternalConnectModal").ExternalConnectionInfo | null) {
-  const onClose = vi.fn();
-  const result = render(<ExternalConnectModal info={info} onClose={onClose} />);
-  return { ...result, onClose };
-}
-
-function clickTab(name: string) {
-  fireEvent.click(screen.getByRole("tab", { name }));
-}
-
-function clickButton(label: string | RegExp) {
-  fireEvent.click(screen.getByRole("button", { name: label }));
-}
-
-// ─── Tests ────────────────────────────────────────────────────────────────────
-
-describe("ExternalConnectModal — null guard", () => {
-  afterEach(() => {
-    cleanup();
-    vi.clearAllMocks();
-  });
-
-  it("renders nothing when info is null", () => {
-    renderModal(null);
-    expect(screen.queryByRole("dialog")).toBeNull();
-  });
-
-  it("renders nothing when info is null even after timeout", () => {
-    vi.useFakeTimers();
-    renderModal(null);
-    act(() => { vi.runAllTimers(); });
-    expect(screen.queryByRole("dialog")).toBeNull();
-    vi.useRealTimers();
-  });
-});
-
-describe("ExternalConnectModal — shell", () => {
-  afterEach(() => {
-    cleanup();
-    vi.clearAllMocks();
-  });
-
-  it("renders dialog when info is provided", () => {
-    renderModal(makeInfo());
-    expect(screen.getByRole("dialog")).toBeTruthy();
-  });
-
-  it("shows the title", () => {
-    renderModal(makeInfo());
-    expect(screen.getByText("Connect your external agent")).toBeTruthy();
-  });
-
-  it("shows the security warning about one-time display", () => {
-    renderModal(makeInfo());
-    expect(screen.getByText(/only once/i)).toBeTruthy();
-  });
-
-  it("shows the close button", () => {
-    renderModal(makeInfo());
-    expect(screen.getByRole("button", { name: /saved it/i })).toBeTruthy();
-  });
-
-  it("close button calls onClose", () => {
-    const { onClose } = renderModal(makeInfo());
-    clickButton(/saved it/i);
-    expect(onClose).toHaveBeenCalledTimes(1);
-  });
-});
-
-describe("ExternalConnectModal — default tab selection", () => {
-  afterEach(() => {
-    cleanup();
-    vi.clearAllMocks();
-  });
-
-  it('defaults to "Universal MCP" tab when universal_mcp_snippet is present', () => {
-    renderModal(makeInfo());
-    // The MCP tab should be aria-selected=true
-    expect(screen.getByRole("tab", { name: "Universal MCP" }).getAttribute("aria-selected")).toBe("true");
-  });
-
-  it('defaults to "Python SDK" tab when universal_mcp_snippet is absent', () => {
-    renderModal(makeInfo({ universal_mcp_snippet: undefined }));
-    expect(screen.getByRole("tab", { name: "Python SDK" }).getAttribute("aria-selected")).toBe("true");
-  });
-
-  it('"Universal MCP" tab is absent when universal_mcp_snippet is undefined', () => {
-    renderModal(makeInfo({ universal_mcp_snippet: undefined }));
-    expect(screen.queryByRole("tab", { name: "Universal MCP" })).toBeNull();
-  });
-
-  it('"Claude Code" tab is absent when claude_code_channel_snippet is undefined', () => {
-    renderModal(makeInfo({ claude_code_channel_snippet: undefined }));
-    expect(screen.queryByRole("tab", { name: "Claude Code" })).toBeNull();
-  });
-
-  it('"Hermes" tab is absent when hermes_channel_snippet is undefined', () => {
-    renderModal(makeInfo({ hermes_channel_snippet: undefined }));
-    expect(screen.queryByRole("tab", { name: "Hermes" })).toBeNull();
-  });
-
-  it('"Codex" tab is absent when codex_snippet is undefined', () => {
-    renderModal(makeInfo({ codex_snippet: undefined }));
-    expect(screen.queryByRole("tab", { name: "Codex" })).toBeNull();
-  });
-
-  it('"OpenClaw" tab is absent when openclaw_snippet is undefined', () => {
-    renderModal(makeInfo({ openclaw_snippet: undefined }));
-    expect(screen.queryByRole("tab", { name: "OpenClaw" })).toBeNull();
-  });
-});
-
-describe("ExternalConnectModal — tab switching", () => {
-  afterEach(() => {
-    cleanup();
-    vi.clearAllMocks();
-  });
-
-  it("clicking Python tab switches aria-selected", () => {
-    renderModal(makeInfo());
-    clickTab("Python SDK");
-    expect(screen.getByRole("tab", { name: "Python SDK" }).getAttribute("aria-selected")).toBe("true");
-  });
-
-  it("clicking curl tab shows curl snippet", () => {
-    renderModal(makeInfo());
-    clickTab("curl");
-    expect(screen.getByText(/curl -X POST/i)).toBeTruthy();
-  });
-
-  it("clicking Fields tab shows all field rows", () => {
-    renderModal(makeInfo());
-    clickTab("Fields");
-    expect(screen.getByText("workspace_id")).toBeTruthy();
-    expect(screen.getByText("ws-test-123")).toBeTruthy();
-    expect(screen.getByText("auth_token")).toBeTruthy();
-    expect(screen.getByText("platform_url")).toBeTruthy();
-    expect(screen.getByText("registry_endpoint")).toBeTruthy();
-    expect(screen.getByText("heartbeat_endpoint")).toBeTruthy();
-  });
-
-  it("clicking Universal MCP tab shows the snippet with token stamped", () => {
-    renderModal(makeInfo());
-    clickTab("Universal MCP");
-    // The token should be stamped, not the placeholder
-    expect(screen.getByText(/tok_secret_abc/i)).toBeTruthy();
-    expect(screen.queryByText(/<paste.*>/i)).toBeNull();
-  });
-
-  it("clicking Python SDK tab shows snippet with token stamped", () => {
-    renderModal(makeInfo());
-    clickTab("Python SDK");
-    expect(screen.getByText(/tok_secret_abc/i)).toBeTruthy();
-  });
-
-  it("clicking Hermes tab shows snippet with token stamped", () => {
-    renderModal(makeInfo());
-    clickTab("Hermes");
-    expect(screen.getByText(/tok_secret_abc/i)).toBeTruthy();
-  });
-
-  it("clicking Codex tab shows snippet with token stamped", () => {
-    renderModal(makeInfo());
-    clickTab("Codex");
-    expect(screen.getByText(/tok_secret_abc/i)).toBeTruthy();
-  });
-
-  it("clicking OpenClaw tab shows snippet with token stamped", () => {
-    renderModal(makeInfo());
-    clickTab("OpenClaw");
-    expect(screen.getByText(/tok_secret_abc/i)).toBeTruthy();
-  });
-
-  it("clicking Claude Code tab shows channel snippet with token stamped", () => {
-    renderModal(makeInfo());
-    clickTab("Claude Code");
-    expect(screen.getByText(/tok_secret_abc/i)).toBeTruthy();
-  });
-});
-
-describe("ExternalConnectModal — copy button", () => {
-  afterEach(() => {
-    cleanup();
-    vi.clearAllMocks();
-  });
-
-  it("Copy button calls navigator.clipboard.writeText", async () => {
-    renderModal(makeInfo());
-    clickTab("curl");
-    const copyBtn = screen.getByRole("button", { name: "Copy" });
-    await act(async () => { fireEvent.click(copyBtn); });
-    expect(mockClipboard.writeText).toHaveBeenCalledTimes(1);
-  });
-
-  it("Copy button shows 'Copied!' after click", async () => {
-    renderModal(makeInfo());
-    clickTab("curl");
-    const copyBtn = screen.getByRole("button", { name: "Copy" });
-    await act(async () => { fireEvent.click(copyBtn); });
-    expect(screen.getByRole("button", { name: "Copied!" })).toBeTruthy();
-  });
-
-  it("Copied! label clears after 1.5s (clipboard auto-reset)", async () => {
-    vi.useFakeTimers();
-    renderModal(makeInfo());
-    clickTab("curl");
-    const copyBtn = screen.getByRole("button", { name: "Copy" });
-    await act(async () => { fireEvent.click(copyBtn); });
-    expect(screen.getByRole("button", { name: "Copied!" })).toBeTruthy();
-    act(() => { vi.runAllTimers(); });
-    // After timeout, button reverts to "Copy"
-    expect(screen.getByRole("button", { name: "Copy" })).toBeTruthy();
-    vi.useRealTimers();
-  });
-
-  it("Copied! label resets on second copy click", async () => {
-    renderModal(makeInfo());
-    clickTab("curl");
-    const copyBtn = screen.getByRole("button", { name: "Copy" });
-    await act(async () => { fireEvent.click(copyBtn); });
-    expect(screen.getByRole("button", { name: "Copied!" })).toBeTruthy();
-    await act(async () => { fireEvent.click(copyBtn); });
-    // Second click resets to "Copy" (auto-clear fires, then new click sets again)
-    // The auto-clear timeout fires and resets, then the new click sets it
-    // In practice: after 1.5s it reverts to Copy; immediate second click resets immediately
-    expect(mockClipboard.writeText).toHaveBeenCalledTimes(2);
-  });
-
-  it("clipboard failure: textarea fallback selected without throwing", async () => {
-    mockClipboard.writeText.mockRejectedValueOnce(new Error("clipboard denied"));
-    renderModal(makeInfo());
-    clickTab("Fields");
-    // The fields tab has a Copy button per row
-    const copyBtns = screen.getAllByRole("button", { name: "Copy" });
-    // Trigger copy on the auth_token field
-    await act(async () => { fireEvent.click(copyBtns[copyBtns.length - 1]); });
-    // Should not throw — error is caught
-    expect(mockClipboard.writeText).toHaveBeenCalledTimes(1);
-  });
-});
-
-describe("ExternalConnectModal — Fields tab", () => {
-  afterEach(() => {
-    cleanup();
-    vi.clearAllMocks();
-  });
-
-  it("shows workspace_id value", () => {
-    renderModal(makeInfo());
-    clickTab("Fields");
-    expect(screen.getByText("ws-test-123")).toBeTruthy();
-  });
-
-  it("shows platform_url value", () => {
-    renderModal(makeInfo({ platform_url: "https://custom.example.com" }));
-    clickTab("Fields");
-    expect(screen.getByText("https://custom.example.com")).toBeTruthy();
-  });
-
-  it("shows masked auth_token (full value visible)", () => {
-    // Note: the modal shows the full token for operator to copy.
-    // The platform does not mask it (by design — operator just saw it once).
-    renderModal(makeInfo());
-    clickTab("Fields");
-    expect(screen.getByText("tok_secret_abc")).toBeTruthy();
-  });
-
-  it("Copy button on Fields rows calls copy with correct value", async () => {
-    renderModal(makeInfo());
-    clickTab("Fields");
-    const copyBtns = screen.getAllByRole("button", { name: "Copy" });
-    // Click the first copy button (workspace_id row)
-    await act(async () => { fireEvent.click(copyBtns[0]); });
-    expect(mockClipboard.writeText).toHaveBeenCalledWith("ws-test-123");
-  });
-
-  it("shows '(missing)' for empty string field values", () => {
-    renderModal(makeInfo({ platform_url: "" }));
-    clickTab("Fields");
-    expect(screen.getByText("(missing)")).toBeTruthy();
-  });
-});
-
-describe("ExternalConnectModal — accessibility", () => {
-  afterEach(() => {
-    cleanup();
-    vi.clearAllMocks();
-  });
-
-  it("tablist has role=tablist", () => {
-    renderModal(makeInfo());
-    expect(screen.getByRole("tablist")).toBeTruthy();
-  });
-
-  it("tablist has aria-label", () => {
-    renderModal(makeInfo());
-    expect(screen.getByRole("tablist").getAttribute("aria-label")).toBe("Connection snippet format");
-  });
-
-  it("each visible tab has role=tab", () => {
-    renderModal(makeInfo());
-    const tabs = screen.getAllByRole("tab");
-    expect(tabs.length).toBeGreaterThan(0);
-  });
-
-  it("active tab has aria-selected=true", () => {
-    renderModal(makeInfo());
-    // Default tab is Universal MCP
-    expect(screen.getByRole("tab", { name: "Universal MCP" }).getAttribute("aria-selected")).toBe("true");
-  });
-
-  it("inactive tab has aria-selected=false", () => {
-    renderModal(makeInfo());
-    expect(screen.getByRole("tab", { name: "Python SDK" }).getAttribute("aria-selected")).toBe("false");
-  });
-
-  it("dialog has role=dialog", () => {
-    renderModal(makeInfo());
-    expect(screen.getByRole("dialog")).toBeTruthy();
-  });
-});
@@ -440,7 +440,8 @@ func (h *MCPHandler) dispatchRPC(ctx context.Context, workspaceID string, req mc
 		}

 	default:
-		base.Error = &mcpRPCError{Code: -32601, Message: "method not found: " + req.Method}
+		// Per OFFSEC-001: error message must not include user-controlled req.Method.
+		base.Error = &mcpRPCError{Code: -32601, Message: "method not found"}
 	}

 	return base
@@ -9,6 +9,7 @@ import (
 	"net/http"
 	"net/http/httptest"
 	"os"
+	"strings"
 	"testing"

 	"errors"
@@ -204,6 +205,9 @@ func TestMCPHandler_NotificationsInitialized_Returns200(t *testing.T) {
 // Unknown method
 // ─────────────────────────────────────────────────────────────────────────────

+// TestMCPHandler_UnknownMethod_Returns32601 verifies dispatchRPC returns
+// -32601 for an unknown method. Per OFFSEC-001: the error message must be
+// constant — req.Method is user-controlled and must NOT appear in the response.
 func TestMCPHandler_UnknownMethod_Returns32601(t *testing.T) {
 	h, _ := newMCPHandler(t)

@@ -224,6 +228,14 @@ func TestMCPHandler_UnknownMethod_Returns32601(t *testing.T) {
 	if resp.Error.Code != -32601 {
 		t.Errorf("expected code -32601, got %d", resp.Error.Code)
 	}
+	// Message must be constant — no user-controlled method name leak.
+	if resp.Error.Message != "method not found" {
+		t.Errorf("error message should be constant 'method not found', got: %q", resp.Error.Message)
+	}
+	// Double-check the method name never appears in the message (defence-in-depth).
+	if strings.Contains(resp.Error.Message, "not/a/real/method") {
+		t.Error("error message must not echo the user-controlled method name")
+	}
 }

 // ─────────────────────────────────────────────────────────────────────────────
@@ -0,0 +1,195 @@
+package handlers
+
+import (
+	"context"
+	"database/sql"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/DATA-DOG/go-sqlmock"
+	"github.com/Molecule-AI/molecule-monorepo/platform/internal/db"
+	"github.com/Molecule-AI/molecule-monorepo/platform/internal/ws"
+	"github.com/Molecule-AI/molecule-monorepo/platform/internal/wsauth"
+	"github.com/alicebob/miniredis/v2"
+	"github.com/gin-gonic/gin"
+	"github.com/redis/go-redis/v9"
+)
+
+// ─── Setup helpers ─────────────────────────────────────────────────────────────
+
+func init() {
+	gin.SetMode(gin.TestMode)
+}
+
+// socketTestDB wraps sqlmock setup with the redis setup needed for wsauth.
+func socketTestDB(t *testing.T) (sqlmock.Sqlmock, func()) {
+	t.Helper()
+	mockDB, mock, err := sqlmock.New()
+	if err != nil {
+		t.Fatalf("failed to create sqlmock: %v", err)
+	}
+
+	// Start a miniredis for the wsauth token subsystem.
+	mr, err := miniredis.Run()
+	if err != nil {
+		mockDB.Close()
+		t.Fatalf("failed to start miniredis: %v", err)
+	}
+	db.DB = mockDB
+	db.RDB = redis.NewClient(&redis.Options{Addr: mr.Addr()})
+
+	wsauth.ResetInboundSecretCacheForTesting()
+
+	cleanup := func() {
+		mockDB.Close()
+		mr.Close()
+		wsauth.ResetInboundSecretCacheForTesting()
+	}
+	return mock, cleanup
+}
+
+// ─── Test cases ────────────────────────────────────────────────────────────────
+// Phase 30.1/30.2 bearer-token auth gate on WebSocket upgrade.
+// SocketHandler.HandleConnect enforces:
+//   - Canvas clients (no X-Workspace-ID header) → bypass auth, upgrade proceeds
+//   - Workspace agents (X-Workspace-ID present) → HasAnyLiveToken probe → bearer validation
+
+func TestSocketHandler_HandleConnect_CanvasClient_NoAuthRequired(t *testing.T) {
+	mock, cleanup := socketTestDB(t)
+	defer cleanup()
+
+	// Create hub and drain the Register channel via Run.
+	hub := ws.NewHub(func(_, _ string) bool { return true })
+	go hub.Run()
+
+	h := NewSocketHandler(hub)
+	c, w := gin.CreateTestContext(httptest.NewRecorder())
+	c.Request = httptest.NewRequest("GET", "/ws", nil)
+	// No X-Workspace-ID → canvas client path.
+
+	h.HandleConnect(c)
+
+	// Canvas path has no DB expectations — HasAnyLiveToken not called.
+	if err := mock.ExpectationsWereMet(); err != nil {
+		t.Errorf("unmet sqlmock expectations: %v", err)
+	}
+	_ = w.Code // upgrade fails in test env (httptest doesn't do WS) — handler returns.
+}
+
+// TestSocketHandler_HandleConnect_AgentNoLiveToken_BypassesBearerCheck verifies
+// that agents with no live tokens (legacy pre-token workspaces) are grandfathered
+// through without being asked for a bearer token.
+func TestSocketHandler_HandleConnect_AgentNoLiveToken_BypassesBearerCheck(t *testing.T) {
+	mock, cleanup := socketTestDB(t)
+	defer cleanup()
+
+	// HasAnyLiveToken → no rows (no live tokens → n=0).
+	mock.ExpectQuery(`SELECT COUNT\(\*\) FROM workspace_auth_tokens WHERE workspace_id = \$1 AND revoked_at IS NULL`).
+		WithArgs("ws-agent").
+		WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(0))
+
+	hub := ws.NewHub(func(_, _ string) bool { return true })
+	go hub.Run()
+
+	h := NewSocketHandler(hub)
+	c, _ := gin.CreateTestContext(httptest.NewRecorder())
+	c.Request = httptest.NewRequest("GET", "/ws", nil)
+	c.Request.Header.Set("X-Workspace-ID", "ws-agent")
+
+	h.HandleConnect(c)
+
+	if err := mock.ExpectationsWereMet(); err != nil {
+		t.Errorf("unmet sqlmock expectations: %v", err)
+	}
+}
+
+// TestSocketHandler_HandleConnect_DBErrorOnHasAnyLiveToken returns 500.
+func TestSocketHandler_HandleConnect_DBErrorOnHasAnyLiveToken(t *testing.T) {
+	mock, cleanup := socketTestDB(t)
+	defer cleanup()
+
+	mock.ExpectQuery(`SELECT COUNT\(\*\) FROM workspace_auth_tokens WHERE workspace_id = \$1 AND revoked_at IS NULL`).
+		WithArgs("ws-agent").
+		WillReturnError(sql.ErrConnDone)
+
+	hub := ws.NewHub(func(_, _ string) bool { return true })
+	go hub.Run()
+
+	h := NewSocketHandler(hub)
+	c, w := gin.CreateTestContext(httptest.NewRecorder())
+	c.Request = httptest.NewRequest("GET", "/ws", nil)
+	c.Request.Header.Set("X-Workspace-ID", "ws-agent")
+
+	h.HandleConnect(c)
+
+	if w.Code != http.StatusInternalServerError {
+		t.Errorf("expected 500 on DB error, got %d", w.Code)
+	}
+	if err := mock.ExpectationsWereMet(); err != nil {
+		t.Errorf("unmet sqlmock expectations: %v", err)
+	}
+}
+
+// TestSocketHandler_HandleConnect_MissingBearerToken returns 401.
+func TestSocketHandler_HandleConnect_MissingBearerToken(t *testing.T) {
+	mock, cleanup := socketTestDB(t)
+	defer cleanup()
+
+	// hasLive=true but no Authorization header.
+	mock.ExpectQuery(`SELECT COUNT\(\*\) FROM workspace_auth_tokens WHERE workspace_id = \$1 AND revoked_at IS NULL`).
+		WithArgs("ws-agent").
+		WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(1))
+
+	hub := ws.NewHub(func(_, _ string) bool { return true })
+	go hub.Run()
+
+	h := NewSocketHandler(hub)
+	c, w := gin.CreateTestContext(httptest.NewRecorder())
+	c.Request = httptest.NewRequest("GET", "/ws", nil)
+	c.Request.Header.Set("X-Workspace-ID", "ws-agent")
+	// No Authorization header.
+
+	h.HandleConnect(c)
+
+	if w.Code != http.StatusUnauthorized {
+		t.Errorf("expected 401 on missing bearer token, got %d", w.Code)
+	}
+	if err := mock.ExpectationsWereMet(); err != nil {
+		t.Errorf("unmet sqlmock expectations: %v", err)
+	}
+}
+
+// TestSocketHandler_HandleConnect_InvalidBearerToken returns 401.
+func TestSocketHandler_HandleConnect_InvalidBearerToken(t *testing.T) {
+	mock, cleanup := socketTestDB(t)
+	defer cleanup()
+
+	// hasLive=true.
+	mock.ExpectQuery(`SELECT COUNT\(\*\) FROM workspace_auth_tokens WHERE workspace_id = \$1 AND revoked_at IS NULL`).
+		WithArgs("ws-agent").
+		WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(1))
+
+	// ValidateToken → lookupTokenByHash: no matching hash.
+	mock.ExpectQuery(`SELECT t\.id, t\.workspace_id FROM workspace_auth_tokens t JOIN workspaces w`).
+		WithArgs(sqlmock.AnyArg()).
+		WillReturnError(context.DeadlineExceeded)
+
+	hub := ws.NewHub(func(_, _ string) bool { return true })
+	go hub.Run()
+
+	h := NewSocketHandler(hub)
+	c, w := gin.CreateTestContext(httptest.NewRecorder())
+	c.Request = httptest.NewRequest("GET", "/ws", nil)
+	c.Request.Header.Set("X-Workspace-ID", "ws-agent")
+	c.Request.Header.Set("Authorization", "Bearer invalid-token-xyz")
+
+	h.HandleConnect(c)
+
+	if w.Code != http.StatusUnauthorized {
+		t.Errorf("expected 401 on invalid bearer token, got %d", w.Code)
+	}
+	if err := mock.ExpectationsWereMet(); err != nil {
+		t.Errorf("unmet sqlmock expectations: %v", err)
+	}
+}