fix(handlers): repair instructions test compile

workspace-server/internal/handlers/a2a_proxy.go

@@ -97,28 +97,28 @@ const maxProxyResponseBody = 10 << 20
 //
 // Timeout model — three independent budgets, none of which gets in each other's way:
 //
-//   1. Client.Timeout — DELIBERATELY UNSET. Client.Timeout is a hard wall on
-//      the entire request including streamed body reads, and would pre-empt
-//      legitimate slow cold-start flows (Claude Code first-token over OAuth
-//      can take 30-60s on boot; long-running agent synthesis can stream
-//      tokens for minutes). Total-request budget is enforced per-request
-//      via context deadline (canvas = idle-only, agent-to-agent = 30 min ceiling).
+//  1. Client.Timeout — DELIBERATELY UNSET. Client.Timeout is a hard wall on
+//     the entire request including streamed body reads, and would pre-empt
+//     legitimate slow cold-start flows (Claude Code first-token over OAuth
+//     can take 30-60s on boot; long-running agent synthesis can stream
+//     tokens for minutes). Total-request budget is enforced per-request
+//     via context deadline (canvas = idle-only, agent-to-agent = 30 min ceiling).
 //
-//   2. Transport.DialContext — 10s connect timeout. When a workspace's EC2
-//      black-holes TCP connects (instance terminated mid-flight, security group
-//      flipped, NACL bug), the OS default is 75s on Linux / 21s on macOS — long
-//      enough that Cloudflare's ~100s edge timeout can fire first and surface
-//      a generic 502 page to canvas. 10s is well above realistic intra-region
-//      latencies and well below CF's edge timeout.
+//  2. Transport.DialContext — 10s connect timeout. When a workspace's EC2
+//     black-holes TCP connects (instance terminated mid-flight, security group
+//     flipped, NACL bug), the OS default is 75s on Linux / 21s on macOS — long
+//     enough that Cloudflare's ~100s edge timeout can fire first and surface
+//     a generic 502 page to canvas. 10s is well above realistic intra-region
+//     latencies and well below CF's edge timeout.
 //
-//   3. Transport.ResponseHeaderTimeout — 180s default. From request-body-end
-//      to response-headers-start. Configurable via
-//      A2A_PROXY_RESPONSE_HEADER_TIMEOUT (envx.Duration). Covers cold-start
-//      first-byte (30-60s OAuth flow above) with enough room for Opus agent
-//      turns (big context + internal delegate_task round-trips routinely exceed
-//      the old 60s ceiling). Body streaming after headers is governed by the
-//      per-request context deadline, NOT this timeout — so multi-minute agent
-//      responses still work fine.
+//  3. Transport.ResponseHeaderTimeout — 180s default. From request-body-end
+//     to response-headers-start. Configurable via
+//     A2A_PROXY_RESPONSE_HEADER_TIMEOUT (envx.Duration). Covers cold-start
+//     first-byte (30-60s OAuth flow above) with enough room for Opus agent
+//     turns (big context + internal delegate_task round-trips routinely exceed
+//     the old 60s ceiling). Body streaming after headers is governed by the
+//     per-request context deadline, NOT this timeout — so multi-minute agent
+//     responses still work fine.
 //
 // The point of (2) and (3) is to surface a *structured* 503 from
 // handleA2ADispatchError when the workspace agent is unreachable, so canvas
@@ -645,7 +645,7 @@ func (h *WorkspaceHandler) resolveAgentURL(ctx context.Context, workspaceID stri
 			// the caller can retry once the workspace is back online (~10s).
 			if status == "hibernated" {
 				log.Printf("ProxyA2A: waking hibernated workspace %s", workspaceID)
-				go h.RestartByID(workspaceID)
+				h.goAsync(func() { h.RestartByID(workspaceID) })
 				return "", &proxyA2AError{
 					Status:  http.StatusServiceUnavailable,
 					Headers: map[string]string{"Retry-After": "15"},

workspace-server/internal/handlers/a2a_proxy_helpers.go

@@ -194,7 +194,7 @@ func (h *WorkspaceHandler) maybeMarkContainerDead(ctx context.Context, workspace
 	}
 	db.ClearWorkspaceKeys(ctx, workspaceID)
 	h.broadcaster.RecordAndBroadcast(ctx, string(events.EventWorkspaceOffline), workspaceID, map[string]interface{}{})
-	go h.RestartByID(workspaceID)
+	h.goAsync(func() { h.RestartByID(workspaceID) })
 	return true
 }
 
@@ -241,7 +241,7 @@ func (h *WorkspaceHandler) preflightContainerHealth(ctx context.Context, workspa
 	}
 	db.ClearWorkspaceKeys(ctx, workspaceID)
 	h.broadcaster.RecordAndBroadcast(ctx, string(events.EventWorkspaceOffline), workspaceID, map[string]interface{}{})
-	go h.RestartByID(workspaceID)
+	h.goAsync(func() { h.RestartByID(workspaceID) })
 	return &proxyA2AError{
 		Status: http.StatusServiceUnavailable,
 		Response: gin.H{
@@ -262,8 +262,8 @@ func (h *WorkspaceHandler) logA2AFailure(ctx context.Context, workspaceID, calle
 		errWsName = workspaceID
 	}
 	summary := "A2A request to " + errWsName + " failed: " + errMsg
-	go func(parent context.Context) {
-		logCtx, cancel := context.WithTimeout(context.WithoutCancel(parent), 30*time.Second)
+	h.goAsync(func() {
+		logCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 30*time.Second)
 		defer cancel()
 		LogActivity(logCtx, h.broadcaster, ActivityParams{
 			WorkspaceID:  workspaceID,
@@ -277,7 +277,7 @@ func (h *WorkspaceHandler) logA2AFailure(ctx context.Context, workspaceID, calle
 			Status:       "error",
 			ErrorDetail:  &errMsg,
 		})
-	}(ctx)
+	})
 }
 
 // logA2ASuccess records a successful A2A round-trip and (for canvas-initiated
@@ -298,19 +298,19 @@ func (h *WorkspaceHandler) logA2ASuccess(ctx context.Context, workspaceID, calle
 	// silent workspaces. Only update when callerID is a real workspace (not
 	// canvas, not a system caller) and the target returned 2xx/3xx.
 	if callerID != "" && !isSystemCaller(callerID) && statusCode < 400 {
-		go func() {
+		h.goAsync(func() {
 			bgCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
 			defer cancel()
 			if _, err := db.DB.ExecContext(bgCtx,
 				`UPDATE workspaces SET last_outbound_at = NOW() WHERE id = $1`, callerID); err != nil {
 				log.Printf("last_outbound_at update failed for %s: %v", callerID, err)
 			}
-		}()
+		})
 	}
 	summary := a2aMethod + " → " + wsNameForLog
 	toolTrace := extractToolTrace(respBody)
-	go func(parent context.Context) {
-		logCtx, cancel := context.WithTimeout(context.WithoutCancel(parent), 30*time.Second)
+	h.goAsync(func() {
+		logCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 30*time.Second)
 		defer cancel()
 		LogActivity(logCtx, h.broadcaster, ActivityParams{
 			WorkspaceID:  workspaceID,
@@ -325,7 +325,7 @@ func (h *WorkspaceHandler) logA2ASuccess(ctx context.Context, workspaceID, calle
 			DurationMs:   &durationMs,
 			Status:       logStatus,
 		})
-	}(ctx)
+	})
 
 	if callerID == "" && statusCode < 400 {
 		h.broadcaster.BroadcastOnly(workspaceID, string(events.EventA2AResponse), map[string]interface{}{
@@ -510,8 +510,8 @@ func (h *WorkspaceHandler) logA2AReceiveQueued(ctx context.Context, workspaceID,
 		wsName = workspaceID
 	}
 	summary := a2aMethod + " → " + wsName + " (queued for poll)"
-	go func(parent context.Context) {
-		logCtx, cancel := context.WithTimeout(context.WithoutCancel(parent), 30*time.Second)
+	h.goAsync(func() {
+		logCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 30*time.Second)
 		defer cancel()
 		LogActivity(logCtx, h.broadcaster, ActivityParams{
 			WorkspaceID:  workspaceID,
@@ -523,7 +523,7 @@ func (h *WorkspaceHandler) logA2AReceiveQueued(ctx context.Context, workspaceID,
 			RequestBody:  json.RawMessage(body),
 			Status:       "ok",
 		})
-	}(ctx)
+	})
 }
 
 // readUsageMap extracts input_tokens / output_tokens from the "usage" key of m.

workspace-server/internal/handlers/a2a_proxy_preflight_test.go

@@ -54,6 +54,7 @@ func TestPreflight_ContainerRunning_ReturnsNil(t *testing.T) {
 	_ = setupTestDB(t)
 	stub := &preflightLocalProv{running: true, err: nil}
 	h := NewWorkspaceHandler(newTestBroadcaster(), nil, "http://localhost:8080", t.TempDir())
+	waitForHandlerAsyncBeforeDBCleanup(t, h)
 	h.provisioner = stub
 
 	if err := h.preflightContainerHealth(context.Background(), "ws-running-123"); err != nil {
@@ -186,8 +187,8 @@ func TestProxyA2A_Preflight_RoutesThroughProvisionerSSOT(t *testing.T) {
 	}
 
 	var (
-		callsIsRunning             bool
-		callsContainerInspectRaw   bool
+		callsIsRunning                  bool
+		callsContainerInspectRaw        bool
 		callsRunningContainerNameDirect bool
 	)
 	ast.Inspect(fn.Body, func(n ast.Node) bool {

workspace-server/internal/handlers/a2a_proxy_test.go

@@ -262,6 +262,7 @@ func TestProxyA2A_Upstream502_TriggersContainerDeadCheck(t *testing.T) {
 	allowLoopbackForTest(t)
 	broadcaster := newTestBroadcaster()
 	handler := NewWorkspaceHandler(broadcaster, nil, "http://localhost:8080", t.TempDir())
+	waitForHandlerAsyncBeforeDBCleanup(t, handler)
 	cp := &fakeCPProv{running: false}
 	handler.SetCPProvisioner(cp)
 
@@ -324,6 +325,7 @@ func TestProxyA2A_Upstream502_AliveAgent_PropagatesAsIs(t *testing.T) {
 	allowLoopbackForTest(t)
 	broadcaster := newTestBroadcaster()
 	handler := NewWorkspaceHandler(broadcaster, nil, "http://localhost:8080", t.TempDir())
+	waitForHandlerAsyncBeforeDBCleanup(t, handler)
 	cp := &fakeCPProv{running: true}
 	handler.SetCPProvisioner(cp)
 
@@ -513,6 +515,7 @@ func TestProxyA2A_AllowedSelf_SkipsAccessCheck(t *testing.T) {
 	allowLoopbackForTest(t)
 	broadcaster := newTestBroadcaster()
 	handler := NewWorkspaceHandler(broadcaster, nil, "http://localhost:8080", t.TempDir())
+	waitForHandlerAsyncBeforeDBCleanup(t, handler)
 
 	agentServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		w.Header().Set("Content-Type", "application/json")
@@ -661,18 +664,18 @@ func TestProxyA2A_CallerIDDerivedFromBearer(t *testing.T) {
 	//    (column order: workspace_id, activity_type, source_id, target_id, ...)
 	mock.ExpectExec("INSERT INTO activity_logs").
 		WithArgs(
-			"ws-target",                       // $1 workspace_id
-			"a2a_receive",                     // $2 activity_type
-			sqlmock.AnyArg(),                  // $3 source_id — *string("ws-caller"), checked below
-			sqlmock.AnyArg(),                  // $4 target_id
-			sqlmock.AnyArg(),                  // $5 method
-			sqlmock.AnyArg(),                  // $6 summary
-			sqlmock.AnyArg(),                  // $7 request_body
-			sqlmock.AnyArg(),                  // $8 response_body
-			sqlmock.AnyArg(),                  // $9 tool_trace
-			sqlmock.AnyArg(),                  // $10 duration_ms
-			sqlmock.AnyArg(),                  // $11 status
-			sqlmock.AnyArg(),                  // $12 error_detail
+			"ws-target",      // $1 workspace_id
+			"a2a_receive",    // $2 activity_type
+			sqlmock.AnyArg(), // $3 source_id — *string("ws-caller"), checked below
+			sqlmock.AnyArg(), // $4 target_id
+			sqlmock.AnyArg(), // $5 method
+			sqlmock.AnyArg(), // $6 summary
+			sqlmock.AnyArg(), // $7 request_body
+			sqlmock.AnyArg(), // $8 response_body
+			sqlmock.AnyArg(), // $9 tool_trace
+			sqlmock.AnyArg(), // $10 duration_ms
+			sqlmock.AnyArg(), // $11 status
+			sqlmock.AnyArg(), // $12 error_detail
 		).
 		WillReturnResult(sqlmock.NewResult(0, 1))
 
@@ -1716,7 +1719,6 @@ func TestDispatchA2A_RejectsUnsafeURL(t *testing.T) {
 	}
 }
 
-
 // --- handleA2ADispatchError ---
 
 func TestHandleA2ADispatchError_ContextDeadline(t *testing.T) {
@@ -1803,6 +1805,7 @@ func TestMaybeMarkContainerDead_CPOnly_NotRunning(t *testing.T) {
 	mock := setupTestDB(t)
 	setupTestRedis(t)
 	handler := NewWorkspaceHandler(newTestBroadcaster(), nil, "http://localhost:8080", t.TempDir())
+	waitForHandlerAsyncBeforeDBCleanup(t, handler)
 	cp := &fakeCPProv{running: false}
 	handler.SetCPProvisioner(cp)
 
@@ -1955,6 +1958,7 @@ func TestLogA2AFailure_Smoke(t *testing.T) {
 	mock := setupTestDB(t)
 	setupTestRedis(t)
 	handler := NewWorkspaceHandler(newTestBroadcaster(), nil, "http://localhost:8080", t.TempDir())
+	waitForHandlerAsyncBeforeDBCleanup(t, handler)
 
 	// Sync workspace-name lookup (called in the caller goroutine).
 	mock.ExpectQuery(`SELECT name FROM workspaces WHERE id =`).
@@ -1973,6 +1977,7 @@ func TestLogA2AFailure_EmptyNameFallback(t *testing.T) {
 	mock := setupTestDB(t)
 	setupTestRedis(t)
 	handler := NewWorkspaceHandler(newTestBroadcaster(), nil, "http://localhost:8080", t.TempDir())
+	waitForHandlerAsyncBeforeDBCleanup(t, handler)
 
 	// Empty name from DB → summary uses the workspaceID as the name.
 	mock.ExpectQuery(`SELECT name FROM workspaces WHERE id =`).
@@ -1989,6 +1994,7 @@ func TestLogA2ASuccess_Smoke(t *testing.T) {
 	mock := setupTestDB(t)
 	setupTestRedis(t)
 	handler := NewWorkspaceHandler(newTestBroadcaster(), nil, "http://localhost:8080", t.TempDir())
+	waitForHandlerAsyncBeforeDBCleanup(t, handler)
 
 	mock.ExpectQuery(`SELECT name FROM workspaces WHERE id =`).
 		WithArgs("ws-ok").
@@ -2005,6 +2011,7 @@ func TestLogA2ASuccess_ErrorStatus(t *testing.T) {
 	mock := setupTestDB(t)
 	setupTestRedis(t)
 	handler := NewWorkspaceHandler(newTestBroadcaster(), nil, "http://localhost:8080", t.TempDir())
+	waitForHandlerAsyncBeforeDBCleanup(t, handler)
 
 	mock.ExpectQuery(`SELECT name FROM workspaces WHERE id =`).
 		WithArgs("ws-err").

workspace-server/internal/handlers/handlers_test.go

@@ -62,6 +62,11 @@ func setupTestDB(t *testing.T) sqlmock.Sqlmock {
 	return mock
 }
 
+func waitForHandlerAsyncBeforeDBCleanup(t *testing.T, h *WorkspaceHandler) {
+	t.Helper()
+	t.Cleanup(h.waitAsyncForTest)
+}
+
 // setupTestRedis creates a miniredis instance and assigns it to the global db.RDB.
 func setupTestRedis(t *testing.T) *miniredis.Miniredis {
 	t.Helper()
@@ -361,6 +366,11 @@ func TestWorkspaceCreate(t *testing.T) {
 }
 
 func TestBuildProvisionerConfig_IncludesAwarenessSettings(t *testing.T) {
+	mock := setupTestDB(t)
+	mock.ExpectQuery(`SELECT digest FROM runtime_image_pins`).
+		WithArgs("claude-code").
+		WillReturnError(sql.ErrNoRows)
+
 	broadcaster := newTestBroadcaster()
 	handler := NewWorkspaceHandler(broadcaster, nil, "http://localhost:8080", "/tmp/configs")
 

workspace-server/internal/handlers/instructions_test.go

@@ -11,7 +11,6 @@ import (
 	"time"
 
 	"github.com/DATA-DOG/go-sqlmock"
-	"github.com/Molecule-AI/molecule-monorepo/platform/internal/db"
 	"github.com/gin-gonic/gin"
 )
 
@@ -193,7 +192,7 @@ func TestInstructionsHandler_Create_InvalidScope(t *testing.T) {
 
 	handler.Create(c)
 
-	if w.Code != http.BadRequest {
+	if w.Code != http.StatusBadRequest {
 		t.Fatalf("expected 400, got %d: %s", w.Code, w.Body.String())
 	}
 }
@@ -277,7 +276,7 @@ func TestInstructionsHandler_Create_WorkspaceScopeWithScopeTarget(t *testing.T)
 		"scope_target": wsID,
 		"title":        "WS rule",
 		"content":      "Use HTTPS",
-		"priority":      10,
+		"priority":     10,
 	})
 	w := httptest.NewRecorder()
 	c, _ := gin.CreateTestContext(w)
@@ -299,10 +298,9 @@ func TestInstructionsHandler_Create_WorkspaceScopeWithScopeTarget(t *testing.T)
 func TestInstructionsHandler_Update_Success(t *testing.T) {
 	mock := setupTestDB(t)
 	handler := NewInstructionsHandler()
-	title := "Updated title"
 
 	mock.ExpectExec(regexp.QuoteMeta("UPDATE platform_instructions SET\n\t\t\t\ttitle = COALESCE($2, title),\n\t\t\t\tcontent = COALESCE($3, content),\n\t\t\t\tpriority = COALESCE($4, priority),\n\t\t\t\tenabled = COALESCE($5, enabled),\n\t\t\t\tupdated_at = NOW()\n\t\t\t\tWHERE id = $1")).
-		WithArgs(&title, "inst-1").
+		WithArgs("inst-1", sqlmock.AnyArg(), nil, nil, nil).
 		WillReturnResult(sqlmock.NewResult(0, 1))
 
 	body, _ := json.Marshal(map[string]interface{}{"title": "Updated title"})
@@ -325,10 +323,9 @@ func TestInstructionsHandler_Update_Success(t *testing.T) {
 func TestInstructionsHandler_Update_NotFound(t *testing.T) {
 	mock := setupTestDB(t)
 	handler := NewInstructionsHandler()
-	title := "Updated title"
 
 	mock.ExpectExec(regexp.QuoteMeta("UPDATE platform_instructions SET\n\t\t\t\ttitle = COALESCE($2, title),\n\t\t\t\tcontent = COALESCE($3, content),\n\t\t\t\tpriority = COALESCE($4, priority),\n\t\t\t\tenabled = COALESCE($5, enabled),\n\t\t\t\tupdated_at = NOW()\n\t\t\t\tWHERE id = $1")).
-		WithArgs(&title, "nonexistent").
+		WithArgs("nonexistent", sqlmock.AnyArg(), nil, nil, nil).
 		WillReturnResult(sqlmock.NewResult(0, 0))
 
 	body, _ := json.Marshal(map[string]interface{}{"title": "Updated title"})

workspace-server/internal/handlers/org_helpers_security_test.go

@@ -93,7 +93,7 @@ func TestResolveInsideRoot_DotPathComponent(t *testing.T) {
 	if err != nil {
 		t.Fatalf("dot path component: unexpected error: %v", err)
 	}
-	if got[len(got)-14:] != "/subdir/file.txt" {
+	if !strings.HasSuffix(got, "/subdir/file.txt") {
 		t.Errorf("dot path component: got %q, want suffix /subdir/file.txt", got)
 	}
 }

workspace-server/internal/handlers/plugins_install_eic_test.go

@@ -342,6 +342,11 @@ func TestPluginInstall_InstanceLookupError_Returns503(t *testing.T) {
 // ---------- dispatch: uninstall ----------
 
 func TestPluginUninstall_SaaS_DispatchesToEIC(t *testing.T) {
+	mock := setupTestDB(t)
+	mock.ExpectExec("DELETE FROM workspace_plugins WHERE workspace_id").
+		WithArgs("ws-1", "browser-automation").
+		WillReturnResult(sqlmock.NewResult(0, 1))
+
 	stubReadPluginManifestViaEIC(t, func(ctx context.Context, instanceID, runtime, pluginName string) ([]byte, error) {
 		return []byte("name: browser-automation\nskills:\n  - browse\n"), nil
 	})

workspace-server/internal/handlers/plugins_test.go

@@ -629,6 +629,9 @@ func TestPluginInstall_RejectsUnknownScheme(t *testing.T) {
 }
 
 func TestPluginInstall_LocalSourceReachesContainerLookup(t *testing.T) {
+	mock := setupTestDB(t)
+	expectAllowlistAllowAll(mock)
+
 	base := t.TempDir()
 	pluginDir := filepath.Join(base, "demo")
 	_ = os.MkdirAll(pluginDir, 0o755)
@@ -955,14 +958,14 @@ func TestLogInstallLimitsOnce(t *testing.T) {
 
 func TestRegexpEscapeForAwk(t *testing.T) {
 	cases := map[string]string{
-		"my-plugin":                 `my-plugin`,
-		"# Plugin: foo /":           `# Plugin: foo \/`,
-		"# Plugin: a.b /":           `# Plugin: a\.b \/`,
-		"foo[bar]":                  `foo\[bar\]`,
-		"a*b+c?":                    `a\*b\+c\?`,
-		"path|with|pipes":           `path\|with\|pipes`,
-		`back\slash`:                `back\\slash`,
-		"":                          ``,
+		"my-plugin":       `my-plugin`,
+		"# Plugin: foo /": `# Plugin: foo \/`,
+		"# Plugin: a.b /": `# Plugin: a\.b \/`,
+		"foo[bar]":        `foo\[bar\]`,
+		"a*b+c?":          `a\*b\+c\?`,
+		"path|with|pipes": `path\|with\|pipes`,
+		`back\slash`:      `back\\slash`,
+		"":                ``,
 	}
 	for in, want := range cases {
 		got := regexpEscapeForAwk(in)
@@ -1247,7 +1250,7 @@ func TestPluginDownload_GithubSchemeStreamsTarball(t *testing.T) {
 		scheme: "github",
 		fetchFn: func(_ context.Context, _ string, dst string) (string, error) {
 			files := map[string]string{
-				"plugin.yaml":            "name: remote-plugin\nversion: 1.0.0\n",
+				"plugin.yaml":             "name: remote-plugin\nversion: 1.0.0\n",
 				"skills/x/SKILL.md":       "---\nname: x\n---\n",
 				"adapters/claude_code.py": "from plugins_registry.builtins import AgentskillsAdaptor as Adaptor\n",
 			}

workspace-server/internal/handlers/restart_signals.go

@@ -58,7 +58,7 @@ func (h *WorkspaceHandler) gracefulPreRestart(ctx context.Context, workspaceID s
 	// Non-blocking send — don't stall the restart cycle.
 	// Run in a detached goroutine so the caller (runRestartCycle) can
 	// proceed to stopForRestart without waiting.
-	go func() {
+	h.goAsync(func() {
 		signalCtx, cancel := context.WithTimeout(context.Background(), restartSignalTimeout)
 		defer cancel()
 
@@ -109,7 +109,7 @@ func (h *WorkspaceHandler) gracefulPreRestart(ctx context.Context, workspaceID s
 		} else {
 			log.Printf("A2AGracefulRestart: %s returned status %d — proceeding with stop", workspaceID, resp.StatusCode)
 		}
-	}()
+	})
 }
 
 // resolveAgentURLForRestartSignal returns the routable URL for the workspace

workspace-server/internal/handlers/restart_signals_test.go

@@ -271,6 +271,7 @@ func TestGracefulPreRestart_URLResolutionError(t *testing.T) {
 		WorkspaceHandler: newHandlerWithTestDeps(t),
 		errToReturn:      context.DeadlineExceeded,
 	}
+	waitForHandlerAsyncBeforeDBCleanup(t, hWrapper.WorkspaceHandler)
 
 	hWrapper.gracefulPreRestart(context.Background(), "ws-url-err-111")
 	time.Sleep(200 * time.Millisecond)

workspace-server/internal/handlers/terminal_test.go

@@ -340,6 +340,11 @@ func TestSSHCommandCmd_BuildsArgv(t *testing.T) {
 // a workspace must still be able to access its own terminal. The CanCommunicate
 // fast-path returns true when callerID == targetID.
 func TestTerminalConnect_KI005_AllowsOwnTerminal(t *testing.T) {
+	mock := setupTestDB(t)
+	mock.ExpectQuery("SELECT COALESCE").
+		WithArgs("ws-alice").
+		WillReturnRows(sqlmock.NewRows([]string{"instance_id"}).AddRow(""))
+
 	// CanCommunicate fast-path: callerID == targetID → returns true without DB.
 	prev := canCommunicateCheck
 	canCommunicateCheck = func(callerID, targetID string) bool { return callerID == targetID }
@@ -367,6 +372,11 @@ func TestTerminalConnect_KI005_AllowsOwnTerminal(t *testing.T) {
 // skip the CanCommunicate check entirely and fall through to the Docker auth path.
 // We assert they get the nil-docker 503 instead of 403.
 func TestTerminalConnect_KI005_SkipsCheckWithoutHeader(t *testing.T) {
+	mock := setupTestDB(t)
+	mock.ExpectQuery("SELECT COALESCE").
+		WithArgs("ws-any").
+		WillReturnRows(sqlmock.NewRows([]string{"instance_id"}).AddRow(""))
+
 	h := NewTerminalHandler(nil) // nil docker → 503 if reached
 	w := httptest.NewRecorder()
 	c, _ := gin.CreateTestContext(w)
@@ -439,6 +449,9 @@ func TestTerminalConnect_KI005_AllowsSiblingWorkspace(t *testing.T) {
 	mock.ExpectExec(`UPDATE workspace_auth_tokens SET last_used_at`).
 		WithArgs(sqlmock.AnyArg()).
 		WillReturnResult(sqlmock.NewResult(0, 1))
+	mock.ExpectQuery("SELECT COALESCE").
+		WithArgs("ws-dev").
+		WillReturnRows(sqlmock.NewRows([]string{"instance_id"}).AddRow(""))
 
 	h := NewTerminalHandler(nil)
 	w := httptest.NewRecorder()
@@ -463,7 +476,10 @@ func TestTerminalConnect_KI005_AllowsSiblingWorkspace(t *testing.T) {
 // introduced in GH#1885: internal routing uses org tokens which are not in
 // workspace_auth_tokens, so ValidateToken would always fail for them.
 func TestKI005_OrgToken_SkipsValidateToken(t *testing.T) {
-	setupTestDB(t) // no ValidateToken ExpectQuery — none should fire
+	mock := setupTestDB(t) // no ValidateToken ExpectQuery — none should fire
+	mock.ExpectQuery("SELECT COALESCE").
+		WithArgs("ws-target").
+		WillReturnRows(sqlmock.NewRows([]string{"instance_id"}).AddRow(""))
 	prev := canCommunicateCheck
 	canCommunicateCheck = func(callerID, targetID string) bool {
 		// Simulate platform agent → target workspace (same org).
@@ -544,4 +560,3 @@ func TestSSHCommandCmd_ConnectTimeoutPresent(t *testing.T) {
 			args)
 	}
 }
-

workspace-server/internal/handlers/workspace.go

@@ -15,6 +15,7 @@ import (
 	"os"
 	"path/filepath"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/Molecule-AI/molecule-monorepo/platform/internal/crypto"
@@ -73,6 +74,19 @@ type WorkspaceHandler struct {
 	// memory plugin). main.go sets this to plugin.DeleteNamespace
 	// when MEMORY_PLUGIN_URL is configured.
 	namespaceCleanupFn func(ctx context.Context, workspaceID string)
+	asyncWG            sync.WaitGroup
+}
+
+func (h *WorkspaceHandler) goAsync(fn func()) {
+	h.asyncWG.Add(1)
+	go func() {
+		defer h.asyncWG.Done()
+		fn()
+	}()
+}
+
+func (h *WorkspaceHandler) waitAsyncForTest() {
+	h.asyncWG.Wait()
 }
 
 func NewWorkspaceHandler(b events.EventEmitter, p *provisioner.Provisioner, platformURL, configsDir string) *WorkspaceHandler {

workspace-server/internal/handlers/workspace_dispatchers.go

@@ -111,11 +111,11 @@ func (h *WorkspaceHandler) provisionWorkspaceAuto(workspaceID, templatePath stri
 		"sync":         false,
 	})
 	if h.cpProv != nil {
-		go h.provisionWorkspaceCP(workspaceID, templatePath, configFiles, payload)
+		h.goAsync(func() { h.provisionWorkspaceCP(workspaceID, templatePath, configFiles, payload) })
 		return true
 	}
 	if h.provisioner != nil {
-		go h.provisionWorkspace(workspaceID, templatePath, configFiles, payload)
+		h.goAsync(func() { h.provisionWorkspace(workspaceID, templatePath, configFiles, payload) })
 		return true
 	}
 	// No backend wired — mark failed so the workspace doesn't linger in
@@ -275,13 +275,13 @@ func (h *WorkspaceHandler) RestartWorkspaceAutoOpts(ctx context.Context, workspa
 	if h.cpProv != nil {
 		h.cpStopWithRetry(ctx, workspaceID, "RestartWorkspaceAuto")
 		// resetClaudeSession is Docker-only — CP has no session state to clear.
-		go h.provisionWorkspaceCP(workspaceID, templatePath, configFiles, payload)
+		h.goAsync(func() { h.provisionWorkspaceCP(workspaceID, templatePath, configFiles, payload) })
 		return true
 	}
 	if h.provisioner != nil {
 		// Docker.Stop has no retry — see docstring rationale.
 		h.provisioner.Stop(ctx, workspaceID)
-		go h.provisionWorkspaceOpts(workspaceID, templatePath, configFiles, payload, resetClaudeSession)
+		h.goAsync(func() { h.provisionWorkspaceOpts(workspaceID, templatePath, configFiles, payload, resetClaudeSession) })
 		return true
 	}
 	// No backend wired — same shape as provisionWorkspaceAuto's no-backend

workspace-server/internal/handlers/workspace_provision_auto_test.go

@@ -144,6 +144,7 @@ func TestProvisionWorkspaceAuto_RoutesToCPWhenSet(t *testing.T) {
 	rec := &trackingCPProv{startErr: errors.New("simulated CP rejection")}
 	bcast := &concurrentSafeBroadcaster{}
 	h := NewWorkspaceHandler(bcast, nil, "http://localhost:8080", t.TempDir())
+	waitForHandlerAsyncBeforeDBCleanup(t, h)
 	h.SetCPProvisioner(rec)
 
 	wsID := "ws-routes-to-cp-0123456789abcdef"
@@ -595,6 +596,7 @@ func TestRestartWorkspaceAuto_RoutesToCPWhenSet(t *testing.T) {
 
 	// Mock DB so cpStopWithRetry can run without a real Postgres.
 	mock := setupTestDB(t)
+	waitForHandlerAsyncBeforeDBCleanup(t, h)
 	mock.MatchExpectationsInOrder(false)
 	// provisionWorkspaceCP runs in the goroutine and will hit secrets
 	// SELECTs + UPDATE workspace as failed (we make CP Start return
@@ -670,6 +672,7 @@ func TestRestartWorkspaceAuto_RoutesToDockerWhenOnlyDocker(t *testing.T) {
 
 	bcast := &concurrentSafeBroadcaster{}
 	h := NewWorkspaceHandler(bcast, nil, "http://localhost:8080", t.TempDir())
+	waitForHandlerAsyncBeforeDBCleanup(t, h)
 	stub := &stoppingLocalProv{}
 	h.provisioner = stub
 

workspace-server/internal/handlers/workspace_provision_test.go

@@ -2,6 +2,7 @@ package handlers
 
 import (
 	"context"
+	"database/sql"
 	"fmt"
 	"net/http"
 	"os"
@@ -634,6 +635,11 @@ func TestSeedInitialMemories_EmptyMemoriesNil(t *testing.T) {
 // ==================== buildProvisionerConfig ====================
 
 func TestBuildProvisionerConfig_BasicFields(t *testing.T) {
+	mock := setupTestDB(t)
+	mock.ExpectQuery(`SELECT COALESCE\(workspace_dir`).
+		WithArgs("ws-basic").
+		WillReturnRows(sqlmock.NewRows([]string{"workspace_dir", "workspace_access"}).AddRow("", "none"))
+
 	broadcaster := newTestBroadcaster()
 	tmpDir := t.TempDir()
 	handler := NewWorkspaceHandler(broadcaster, nil, "http://localhost:8080", tmpDir)
@@ -678,6 +684,14 @@ func TestBuildProvisionerConfig_BasicFields(t *testing.T) {
 }
 
 func TestBuildProvisionerConfig_WorkspacePathFromEnv(t *testing.T) {
+	mock := setupTestDB(t)
+	mock.ExpectQuery(`SELECT COALESCE\(workspace_dir`).
+		WithArgs("ws-env").
+		WillReturnError(sql.ErrNoRows)
+	mock.ExpectQuery(`SELECT digest FROM runtime_image_pins`).
+		WithArgs("claude-code").
+		WillReturnError(sql.ErrNoRows)
+
 	broadcaster := newTestBroadcaster()
 	handler := NewWorkspaceHandler(broadcaster, nil, "http://localhost:8080", t.TempDir())