Harden display control tab state handling

canvas/src/components/tabs/DisplayTab.tsx

@@ -1,6 +1,6 @@
 "use client";
 
-import { useEffect, useState } from "react";
+import { useEffect, useRef, useState } from "react";
 import { api } from "@/lib/api";
 
 interface DisplayStatus {
@@ -13,31 +13,86 @@ interface DisplayStatus {
   height?: number;
 }
 
+interface DisplayControlStatus {
+  controller: "none" | "user" | "agent";
+  controlled_by?: string;
+  expires_at?: string;
+}
+
 interface Props {
   workspaceId: string;
 }
 
 export function DisplayTab({ workspaceId }: Props) {
   const [status, setStatus] = useState<DisplayStatus | null>(null);
+  const [control, setControl] = useState<DisplayControlStatus | null>(null);
   const [error, setError] = useState<string | null>(null);
+  const [controlError, setControlError] = useState<string | null>(null);
+  const [controlBusy, setControlBusy] = useState(false);
+  const requestGeneration = useRef(0);
 
   useEffect(() => {
+    const generation = requestGeneration.current + 1;
+    requestGeneration.current = generation;
     let cancelled = false;
     setStatus(null);
+    setControl(null);
     setError(null);
-    api
-      .get<DisplayStatus>(`/workspaces/${workspaceId}/display`)
-      .then((data) => {
-        if (!cancelled) setStatus(data);
-      })
-      .catch((err) => {
-        if (!cancelled) setError(err instanceof Error ? err.message : "Display status unavailable");
-      });
+    setControlError(null);
+    setControlBusy(false);
+    async function load() {
+      try {
+        const displayStatus = await api.get<DisplayStatus>(`/workspaces/${workspaceId}/display`);
+        if (cancelled || requestGeneration.current !== generation) return;
+        setStatus(displayStatus);
+        if (displayStatus.reason === "display_not_enabled") return;
+        try {
+          const displayControl = await api.get<DisplayControlStatus>(`/workspaces/${workspaceId}/display/control`);
+          if (!cancelled && requestGeneration.current === generation) setControl(displayControl);
+        } catch (err) {
+          if (!cancelled && requestGeneration.current === generation) {
+            setControl(null);
+            setControlError("Display control unavailable");
+          }
+        }
+      } catch (err) {
+        if (!cancelled && requestGeneration.current === generation) setError("The display status could not be loaded.");
+      }
+    }
+    load();
     return () => {
       cancelled = true;
     };
   }, [workspaceId]);
 
+  const acquireControl = async () => {
+    const generation = requestGeneration.current;
+    const controlPath = `/workspaces/${workspaceId}/display/control`;
+    setControlBusy(true);
+    setControlError(null);
+    try {
+      const next = await api.post<DisplayControlStatus>(`${controlPath}/acquire`, {
+        controller: "user",
+        ttl_seconds: 300,
+      });
+      if (requestGeneration.current !== generation) return;
+      setControl(next);
+    } catch (err) {
+      if (requestGeneration.current !== generation) return;
+      setControlError("Failed to take control");
+      try {
+        const latest = await api.get<DisplayControlStatus>(controlPath);
+        if (requestGeneration.current !== generation) return;
+        setControl(latest);
+      } catch {
+        if (requestGeneration.current !== generation) return;
+        setControl(null);
+      }
+    } finally {
+      if (requestGeneration.current === generation) setControlBusy(false);
+    }
+  };
+
   if (error) {
     return (
       <div className="p-5">
@@ -81,12 +136,50 @@ export function DisplayTab({ workspaceId }: Props) {
             : "This workspace has display configuration, but the desktop session infrastructure is not configured yet."}
         </p>
         {!isNotEnabled && (
-          <dl className="mt-5 grid grid-cols-2 gap-x-4 gap-y-2 text-left text-[11px]">
-            <dt className="text-ink-mid">Mode</dt>
-            <dd className="font-mono text-ink">{status.mode || "unknown"}</dd>
-            <dt className="text-ink-mid">Status</dt>
-            <dd className="font-mono text-ink">{status.status || "unknown"}</dd>
-          </dl>
+          <>
+            <dl className="mt-5 grid grid-cols-2 gap-x-4 gap-y-2 text-left text-[11px]">
+              <dt className="text-ink-mid">Mode</dt>
+              <dd className="font-mono text-ink">{status.mode || "unknown"}</dd>
+              <dt className="text-ink-mid">Status</dt>
+              <dd className="font-mono text-ink">{status.status || "unknown"}</dd>
+            </dl>
+            <div className="mt-5 w-full max-w-xs border-t border-line/50 pt-4">
+              {control ? (
+                <div className="flex items-center justify-between gap-3 text-left">
+                  <div className="min-w-0">
+                    <p className="text-[11px] font-medium text-ink">
+                      {control.controller === "none"
+                        ? "No active controller"
+                        : `Controlled by ${displayControlActorLabel(control)}`}
+                    </p>
+                    {control.expires_at && (
+                      <p className="mt-1 truncate font-mono text-[10px] text-ink-mid">
+                        Until {new Date(control.expires_at).toLocaleTimeString()}
+                      </p>
+                    )}
+                    {controlError && <p className="mt-1 text-[10px] leading-snug text-red-200">{controlError}</p>}
+                  </div>
+                  {control.controller === "none" && (
+                    <button
+                      type="button"
+                      onClick={acquireControl}
+                      disabled={controlBusy}
+                      className="h-8 shrink-0 rounded border border-line bg-surface px-3 text-[11px] font-medium text-ink hover:bg-surface-elevated disabled:cursor-not-allowed disabled:opacity-60"
+                    >
+                      Take control
+                    </button>
+                  )}
+                </div>
+              ) : (
+                <div className="text-left">
+                  {!controlError && (
+                    <div className="h-8 rounded border border-line/40 bg-surface-sunken/30 motion-safe:animate-pulse" />
+                  )}
+                  {controlError && <p className="mt-2 text-[10px] leading-snug text-red-200">{controlError}</p>}
+                </div>
+              )}
+            </div>
+          </>
         )}
       </div>
     );
@@ -94,3 +187,10 @@ export function DisplayTab({ workspaceId }: Props) {
 
   return null;
 }
+
+function displayControlActorLabel(control: DisplayControlStatus): string {
+  if (control.controller === "agent") return "Agent";
+  if (control.controlled_by === "admin-token") return "Admin";
+  if (control.controlled_by?.startsWith("org-token:")) return "Automation";
+  return "User";
+}

canvas/src/components/tabs/__tests__/DisplayTab.test.tsx

@@ -1,12 +1,13 @@
 // @vitest-environment jsdom
 import { describe, it, expect, vi, beforeEach } from "vitest";
-import { render, screen, waitFor } from "@testing-library/react";
+import { cleanup, fireEvent, render, screen, waitFor } from "@testing-library/react";
 
-const { mockGet } = vi.hoisted(() => ({ mockGet: vi.fn() }));
+const { mockGet, mockPost } = vi.hoisted(() => ({ mockGet: vi.fn(), mockPost: vi.fn() }));
 
 vi.mock("@/lib/api", () => ({
   api: {
     get: mockGet,
+    post: mockPost,
   },
 }));
 
@@ -14,7 +15,9 @@ import { DisplayTab } from "../DisplayTab";
 
 describe("DisplayTab", () => {
   beforeEach(() => {
+    cleanup();
     mockGet.mockReset();
+    mockPost.mockReset();
   });
 
   it("renders unavailable state for non-display workspaces", async () => {
@@ -29,5 +32,219 @@ describe("DisplayTab", () => {
       expect(screen.getByText("Display is not enabled for this workspace.")).toBeTruthy();
     });
     expect(mockGet).toHaveBeenCalledWith("/workspaces/ws-no-display/display");
+    expect(mockGet).not.toHaveBeenCalledWith("/workspaces/ws-no-display/display/control");
+  });
+
+  it("renders control acquisition for display-configured workspaces", async () => {
+    mockGet
+      .mockResolvedValueOnce({
+        available: false,
+        reason: "display_session_unavailable",
+        mode: "desktop-control",
+        status: "not_configured",
+      })
+      .mockResolvedValueOnce({
+        controller: "none",
+      });
+    mockPost.mockResolvedValueOnce({
+      controller: "user",
+      controlled_by: "admin-token",
+      expires_at: "2026-05-23T08:48:27Z",
+    });
+
+    render(<DisplayTab workspaceId="ws-display" />);
+
+    await waitFor(() => {
+      expect(screen.getByRole("button", { name: "Take control" })).toBeTruthy();
+    });
+    expect(mockGet).toHaveBeenCalledWith("/workspaces/ws-display/display");
+    expect(mockGet).toHaveBeenCalledWith("/workspaces/ws-display/display/control");
+
+    fireEvent.click(screen.getByRole("button", { name: "Take control" }));
+
+    await waitFor(() => {
+      expect(screen.getByText("Controlled by Admin")).toBeTruthy();
+    });
+    expect(mockPost).toHaveBeenCalledWith("/workspaces/ws-display/display/control/acquire", {
+      controller: "user",
+      ttl_seconds: 300,
+    });
+  });
+
+  it("renders active display control locks as observe-only", async () => {
+    mockGet
+      .mockResolvedValueOnce({
+        available: false,
+        reason: "display_session_unavailable",
+        mode: "desktop-control",
+        status: "not_configured",
+      })
+      .mockResolvedValueOnce({
+        controller: "agent",
+        controlled_by: "sidecar",
+        expires_at: "2026-05-23T08:48:27Z",
+      });
+
+    render(<DisplayTab workspaceId="ws-display" />);
+
+    await waitFor(() => {
+      expect(screen.getByText("Controlled by Agent")).toBeTruthy();
+    });
+    expect(screen.queryByRole("button", { name: "Release" })).toBeNull();
+    expect(screen.queryByRole("button", { name: "Take control" })).toBeNull();
+    expect(mockPost).not.toHaveBeenCalled();
+  });
+
+  it("labels org-token display control locks as automation", async () => {
+    mockGet
+      .mockResolvedValueOnce({
+        available: false,
+        reason: "display_session_unavailable",
+        mode: "desktop-control",
+        status: "not_configured",
+      })
+      .mockResolvedValueOnce({
+        controller: "user",
+        controlled_by: "org-token:abc123",
+        expires_at: "2026-05-23T08:48:27Z",
+      });
+
+    render(<DisplayTab workspaceId="ws-display" />);
+
+    await waitFor(() => {
+      expect(screen.getByText("Controlled by Automation")).toBeTruthy();
+    });
+    expect(screen.queryByText("org-token:abc123")).toBeNull();
+    expect(screen.queryByRole("button", { name: "Take control" })).toBeNull();
+  });
+
+  it("refreshes display control state after failed acquisition", async () => {
+    mockGet
+      .mockResolvedValueOnce({
+        available: false,
+        reason: "display_session_unavailable",
+        mode: "desktop-control",
+        status: "not_configured",
+      })
+      .mockResolvedValueOnce({
+        controller: "none",
+      })
+      .mockResolvedValueOnce({
+        controller: "agent",
+        controlled_by: "sidecar",
+        expires_at: "2026-05-23T08:48:27Z",
+      });
+    mockPost.mockRejectedValueOnce(new Error("API POST /workspaces/ws-display/display/control/acquire: 409 conflict"));
+
+    render(<DisplayTab workspaceId="ws-display" />);
+
+    await waitFor(() => {
+      expect(screen.getByRole("button", { name: "Take control" })).toBeTruthy();
+    });
+
+    fireEvent.click(screen.getByRole("button", { name: "Take control" }));
+
+    await waitFor(() => {
+      expect(screen.getByText("Controlled by Agent")).toBeTruthy();
+    });
+    expect(screen.getByText("Failed to take control")).toBeTruthy();
+    expect(mockGet).toHaveBeenCalledWith("/workspaces/ws-display/display/control");
+    expect(mockGet).toHaveBeenCalledTimes(3);
+    expect(mockPost).toHaveBeenCalledWith("/workspaces/ws-display/display/control/acquire", {
+      controller: "user",
+      ttl_seconds: 300,
+    });
+  });
+
+  it("keeps display status visible without takeover actions when control status fails", async () => {
+    mockGet
+      .mockResolvedValueOnce({
+        available: false,
+        reason: "display_session_unavailable",
+        mode: "desktop-control",
+        status: "not_configured",
+      })
+      .mockRejectedValueOnce(new Error("API GET /workspaces/ws-display/display/control: 401 unauthorized"));
+
+    render(<DisplayTab workspaceId="ws-display" />);
+
+    await waitFor(() => {
+      expect(screen.getByText("Display session is not ready.")).toBeTruthy();
+    });
+    expect(screen.queryByRole("button", { name: "Take control" })).toBeNull();
+    expect(screen.getByText("Display control unavailable")).toBeTruthy();
+  });
+
+  it("does not render raw display status errors", async () => {
+    mockGet.mockRejectedValueOnce(new Error("API GET /workspaces/ws-display/display: 500 secret backend details"));
+
+    render(<DisplayTab workspaceId="ws-display" />);
+
+    await waitFor(() => {
+      expect(screen.getByText("Display status unavailable")).toBeTruthy();
+    });
+    expect(screen.queryByText(/secret backend details/)).toBeNull();
+  });
+
+  it("ignores stale acquire responses after workspace changes", async () => {
+    const acquire = deferred<{ controller: "user"; controlled_by: string; expires_at: string }>();
+    mockGet
+      .mockResolvedValueOnce({
+        available: false,
+        reason: "display_session_unavailable",
+        mode: "desktop-control",
+        status: "not_configured",
+      })
+      .mockResolvedValueOnce({
+        controller: "none",
+      })
+      .mockResolvedValueOnce({
+        available: false,
+        reason: "display_session_unavailable",
+        mode: "desktop-control",
+        status: "not_configured",
+      })
+      .mockResolvedValueOnce({
+        controller: "none",
+      });
+    mockPost.mockReturnValueOnce(acquire.promise);
+
+    const { rerender } = render(<DisplayTab workspaceId="ws-a" />);
+
+    await waitFor(() => {
+      expect(screen.getByRole("button", { name: "Take control" })).toBeTruthy();
+    });
+    fireEvent.click(screen.getByRole("button", { name: "Take control" }));
+
+    rerender(<DisplayTab workspaceId="ws-b" />);
+
+    await waitFor(() => {
+      expect(mockGet).toHaveBeenCalledWith("/workspaces/ws-b/display/control");
+    });
+    await waitFor(() => {
+      expect(screen.getByRole("button", { name: "Take control" })).toBeTruthy();
+    });
+
+    acquire.resolve({
+      controller: "user",
+      controlled_by: "admin-token",
+      expires_at: "2026-05-23T08:48:27Z",
+    });
+    await acquire.promise;
+
+    await waitFor(() => {
+      expect(screen.queryByText("Controlled by Admin")).toBeNull();
+    });
+    expect(screen.getByRole("button", { name: "Take control" })).toBeTruthy();
   });
 });
+
+function deferred<T>() {
+  let resolve!: (value: T) => void;
+  let reject!: (reason?: unknown) => void;
+  const promise = new Promise<T>((res, rej) => {
+    resolve = res;
+    reject = rej;
+  });
+  return { promise, resolve, reject };
+}

workspace-server/internal/memory/pgplugin/handlers.go

@@ -3,7 +3,6 @@ package pgplugin
 import (
 	"encoding/json"
 	"errors"
-	"log"
 	"net/http"
 	"strings"
 
@@ -247,9 +246,7 @@ func (h *Handler) forget(w http.ResponseWriter, r *http.Request, id string) {
 func writeJSON(w http.ResponseWriter, status int, body interface{}) {
 	w.Header().Set("Content-Type", "application/json")
 	w.WriteHeader(status)
-	if err := json.NewEncoder(w).Encode(body); err != nil {
-		log.Printf("pgplugin: JSON encode error: %v", err)
-	}
+	_ = json.NewEncoder(w).Encode(body)
 }
 
 func writeError(w http.ResponseWriter, status int, code contract.ErrorCode, message string, details map[string]interface{}) {