From 9eefa5c47445a18cd2b336693eed93cbda15f81e Mon Sep 17 00:00:00 2001
From: hongming <hongmingwang@moleculesai.app>
Date: Sun, 24 May 2026 20:31:29 -0700
Subject: [PATCH] fix(display): allow browser sessions to take control

---
 canvas/src/components/tabs/DisplayTab.tsx     | 10 ++--
 .../handlers/workspace_display_control.go     |  7 ++-
 .../workspace_display_control_test.go         | 42 ++++++++++++++++
 .../internal/middleware/wsauth_middleware.go  |  8 ++++
 .../middleware/wsauth_middleware_test.go      | 48 +++++++++++++++++++
 5 files changed, 110 insertions(+), 5 deletions(-)

diff --git a/canvas/src/components/tabs/DisplayTab.tsx b/canvas/src/components/tabs/DisplayTab.tsx
index 5944db28a..540f2e358 100644
--- a/canvas/src/components/tabs/DisplayTab.tsx
+++ b/canvas/src/components/tabs/DisplayTab.tsx
@@ -265,6 +265,11 @@ function DisplayControlBar({
   onAcquire: () => void;
   onRelease: () => void;
 }) {
+  const userControl = control?.controller === "user";
+  const adminControl = userControl && control?.controlled_by === "admin-token";
+  const canAcquireUserControl = control?.controller === "none" || (userControl && !hasSession);
+  const canReleaseUserControl = adminControl || (userControl && hasSession);
+
   return (
     <div className="flex min-w-0 items-center gap-3">
       {control && (
@@ -282,8 +287,7 @@ function DisplayControlBar({
           {controlError && <p className="mt-0.5 text-[10px] text-red-200">{controlError}</p>}
         </div>
       )}
-      {(control?.controller === "none" ||
-        (control?.controller === "user" && control.controlled_by === "admin-token" && !hasSession)) && (
+      {canAcquireUserControl && (
         <button
           type="button"
           onClick={onAcquire}
@@ -293,7 +297,7 @@ function DisplayControlBar({
           Take control
         </button>
       )}
-      {control?.controller === "user" && control.controlled_by === "admin-token" && (
+      {canReleaseUserControl && (
         <button
           type="button"
           onClick={onRelease}
diff --git a/workspace-server/internal/handlers/workspace_display_control.go b/workspace-server/internal/handlers/workspace_display_control.go
index 6e23db348..8c159a2e4 100644
--- a/workspace-server/internal/handlers/workspace_display_control.go
+++ b/workspace-server/internal/handlers/workspace_display_control.go
@@ -337,11 +337,14 @@ func displayControlActor(c *gin.Context) (string, bool) {
 			return actorOrgTokenPrefix + s, true
 		}
 	}
+	if v, ok := c.Get("cp_session_actor"); ok {
+		if s, ok := v.(string); ok && s != "" {
+			return s, true
+		}
+	}
 	if displayControlIsAdminToken(c) {
 		return actorAdminToken, true
 	}
-	// Browser session auth is intentionally observe-only until AdminAuth
-	// exposes a stable per-user or per-session identity in gin.Context.
 	return "", false
 }
 
diff --git a/workspace-server/internal/handlers/workspace_display_control_test.go b/workspace-server/internal/handlers/workspace_display_control_test.go
index e137cec3e..0f15146d0 100644
--- a/workspace-server/internal/handlers/workspace_display_control_test.go
+++ b/workspace-server/internal/handlers/workspace_display_control_test.go
@@ -258,6 +258,48 @@ func TestWorkspaceDisplayControlAcquire_RejectsCoarseSessionActor(t *testing.T)
 	}
 }
 
+func TestWorkspaceDisplayControlAcquire_AcceptsVerifiedBrowserSessionActor(t *testing.T) {
+	mock := setupTestDB(t)
+	t.Setenv("DISPLAY_SESSION_SIGNING_SECRET", "display-session-test-secret")
+	handler := NewWorkspaceHandler(newTestBroadcaster(), nil, "http://localhost:8080", t.TempDir())
+	expiresAt := time.Date(2026, 5, 23, 18, 30, 0, 0, time.UTC)
+
+	mock.ExpectQuery(`SELECT COALESCE\(compute, '\{\}'::jsonb\) FROM workspaces WHERE id = \$1`).
+		WithArgs("ws-display").
+		WillReturnRows(sqlmock.NewRows([]string{"compute"}).AddRow(`{"display":{"mode":"desktop-control","protocol":"dcv","width":1920,"height":1080}}`))
+	mock.ExpectQuery(`INSERT INTO workspace_display_control_locks`).
+		WithArgs("ws-display", "user", "session:abc123", 300).
+		WillReturnRows(sqlmock.NewRows([]string{"controller", "controlled_by", "expires_at"}).
+			AddRow("user", "session:abc123", expiresAt))
+
+	w := httptest.NewRecorder()
+	c, _ := gin.CreateTestContext(w)
+	c.Params = gin.Params{{Key: "id", Value: "ws-display"}}
+	c.Request = httptest.NewRequest("POST", "/workspaces/ws-display/display/control/acquire", bytes.NewBufferString(`{"controller":"user","ttl_seconds":300}`))
+	c.Request.Header.Set("Content-Type", "application/json")
+	c.Set("cp_session_actor", "session:abc123")
+
+	handler.AcquireDisplayControl(c)
+
+	if w.Code != http.StatusOK {
+		t.Fatalf("expected status 200, got %d: %s", w.Code, w.Body.String())
+	}
+	var resp map[string]interface{}
+	if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
+		t.Fatalf("failed to parse response: %v", err)
+	}
+	if resp["controller"] != "user" || resp["controlled_by"] != "session:abc123" {
+		t.Fatalf("lock response = %#v, want user/session actor", resp)
+	}
+	sessionURL, ok := resp["session_url"].(string)
+	if !ok || !strings.HasPrefix(sessionURL, "/workspaces/ws-display/display/session/websockify#token=") {
+		t.Fatalf("session_url = %#v, want signed websockify URL fragment", resp["session_url"])
+	}
+	if err := mock.ExpectationsWereMet(); err != nil {
+		t.Errorf("unmet sqlmock expectations: %v", err)
+	}
+}
+
 func TestWorkspaceDisplayControlRelease_RemovesCallerLock(t *testing.T) {
 	mock := setupTestDB(t)
 	handler := NewWorkspaceHandler(newTestBroadcaster(), nil, "http://localhost:8080", t.TempDir())
diff --git a/workspace-server/internal/middleware/wsauth_middleware.go b/workspace-server/internal/middleware/wsauth_middleware.go
index a15302cbe..821da64f0 100644
--- a/workspace-server/internal/middleware/wsauth_middleware.go
+++ b/workspace-server/internal/middleware/wsauth_middleware.go
@@ -1,8 +1,10 @@
 package middleware
 
 import (
+	"crypto/sha256"
 	"crypto/subtle"
 	"database/sql"
+	"encoding/hex"
 	"errors"
 	"log"
 	"net/http"
@@ -196,6 +198,7 @@ func AdminAuth(database *sql.DB) gin.HandlerFunc {
 		// bearer-only path unchanged.
 		if cookieHeader := c.GetHeader("Cookie"); cookieHeader != "" {
 			if ok, _ := VerifiedCPSession(cookieHeader); ok {
+				c.Set("cp_session_actor", cpSessionActor(cookieHeader))
 				c.Next()
 				return
 			}
@@ -260,6 +263,11 @@ func AdminAuth(database *sql.DB) gin.HandlerFunc {
 	}
 }
 
+func cpSessionActor(cookieHeader string) string {
+	sum := sha256.Sum256([]byte(tenantSlug() + "\x00" + cookieHeader))
+	return "session:" + hex.EncodeToString(sum[:])[:16]
+}
+
 // CanvasOrBearer is a softer admin-auth variant used ONLY for cosmetic
 // canvas routes where forging the request has zero security impact (PUT
 // /canvas/viewport: worst case an attacker resets the shared viewport
diff --git a/workspace-server/internal/middleware/wsauth_middleware_test.go b/workspace-server/internal/middleware/wsauth_middleware_test.go
index 8371802c2..0b748a9b3 100644
--- a/workspace-server/internal/middleware/wsauth_middleware_test.go
+++ b/workspace-server/internal/middleware/wsauth_middleware_test.go
@@ -284,6 +284,54 @@ func TestAdminAuth_FailOpen_NoTokensGlobally(t *testing.T) {
 	}
 }
 
+func TestAdminAuth_VerifiedCPSession_SetsSessionActor(t *testing.T) {
+	resetSessionCache()
+	srv, hits := mockCPServer(t, http.StatusOK, `{"member":true,"user_id":"u_1","role":"owner","org_id":"org_1"}`)
+	t.Setenv("CP_UPSTREAM_URL", srv.URL)
+	t.Setenv("MOLECULE_ORG_SLUG", "acme")
+	t.Setenv("ADMIN_TOKEN", "admin-secret")
+
+	mockDB, mock, err := sqlmock.New()
+	if err != nil {
+		t.Fatalf("sqlmock.New: %v", err)
+	}
+	defer mockDB.Close()
+
+	mock.ExpectQuery(hasAnyLiveTokenGlobalQuery).
+		WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(1))
+
+	const cookie = "session=valid-browser-session"
+	r := gin.New()
+	r.GET("/workspaces", AdminAuth(mockDB), func(c *gin.Context) {
+		actor, ok := c.Get("cp_session_actor")
+		if !ok {
+			t.Fatalf("expected cp_session_actor in context")
+		}
+		if actor != cpSessionActor(cookie) {
+			t.Fatalf("cp_session_actor = %q, want stable hashed actor", actor)
+		}
+		if actor == cookie {
+			t.Fatalf("cp_session_actor must not expose the raw cookie")
+		}
+		c.JSON(http.StatusOK, gin.H{"ok": true})
+	})
+
+	w := httptest.NewRecorder()
+	req, _ := http.NewRequest(http.MethodGet, "/workspaces", nil)
+	req.Header.Set("Cookie", cookie)
+	r.ServeHTTP(w, req)
+
+	if w.Code != http.StatusOK {
+		t.Fatalf("expected 200 for verified CP session, got %d: %s", w.Code, w.Body.String())
+	}
+	if hits.Load() != 1 {
+		t.Fatalf("expected one CP verification request, got %d", hits.Load())
+	}
+	if err := mock.ExpectationsWereMet(); err != nil {
+		t.Errorf("unmet sqlmock expectations: %v", err)
+	}
+}
+
 // TestAdminAuth_C10_NoBearer_Returns401 — C10 critical path: when at least
 // one workspace has tokens, GET /admin/secrets without a bearer → 401.
 func TestAdminAuth_C10_NoBearer_Returns401(t *testing.T) {
-- 
2.52.0