fix

Author: slyt3

pkg/api/v1/workload_service.go

@@ -12,6 +12,7 @@ import (
 	"time"
 
 	nameref "github.com/google/go-containerregistry/pkg/name"
+
 	regtypes "github.com/stacklok/toolhive-core/registry/types"
 	groupval "github.com/stacklok/toolhive-core/validation/group"
 	httpval "github.com/stacklok/toolhive-core/validation/http"
@@ -435,7 +436,10 @@ func validateRuntimeConfig(runtimeConfig *templates.RuntimeConfig) error {
 	return nil
 }
 
-func runtimeConfigForImageBuild(req *createRequest, runtimeConfigOverride *templates.RuntimeConfig) (*templates.RuntimeConfig, error) {
+func runtimeConfigForImageBuild(
+	req *createRequest,
+	runtimeConfigOverride *templates.RuntimeConfig,
+) (*templates.RuntimeConfig, error) {
 	if runtimeConfigOverride == nil || req == nil {
 		return nil, nil
 	}

pkg/api/v1/workloads_types_test.go

@@ -10,10 +10,10 @@ import (
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
+	"github.com/stacklok/toolhive-core/permissions"
 	"github.com/stacklok/toolhive/pkg/auth"
 	"github.com/stacklok/toolhive/pkg/auth/remote"
 	"github.com/stacklok/toolhive/pkg/container/templates"
-	"github.com/stacklok/toolhive-core/permissions"
 	"github.com/stacklok/toolhive/pkg/runner"
 	"github.com/stacklok/toolhive/pkg/secrets"
 	"github.com/stacklok/toolhive/pkg/transport/types"

pkg/vmcp/composer/state_store_test.go

@@ -143,15 +143,15 @@ func TestInMemoryStateStore_ListActiveWorkflows(t *testing.T) {
 // TestInMemoryStateStore_Cleanup tests automatic cleanup of stale workflows.
 func TestInMemoryStateStore_Cleanup(t *testing.T) {
 	t.Parallel()
-	// Use very short intervals for testing but with sufficient margin
+	// Use short intervals with enough margin to reduce CI timing flakiness.
 	cleanupInterval := 50 * time.Millisecond
-	maxAge := 50 * time.Millisecond
+	maxAge := 100 * time.Millisecond
 
 	store := NewInMemoryStateStore(cleanupInterval, maxAge).(*inMemoryStateStore)
 	defer store.Stop()
 
 	// Create workflows directly in the store with specific timestamps
-	veryOldTime := time.Now().Add(-1 * time.Second) // Way older than maxAge
+	veryOldTime := time.Now().Add(-500 * time.Millisecond) // Well older than maxAge
 
 	store.mu.Lock()
 	// Old completed workflow - should be cleaned up
@@ -169,20 +169,15 @@ func TestInMemoryStateStore_Cleanup(t *testing.T) {
 	}
 	store.mu.Unlock()
 
-	// Wait for at least 2 cleanup cycles
-	time.Sleep(150 * time.Millisecond)
+	require.Eventually(t, func() bool {
+		store.mu.RLock()
+		oldExists := store.states["old-workflow"]
+		runningExists := store.states["running-workflow"]
+		store.mu.RUnlock()
 
-	// Verify cleanup results
-	store.mu.RLock()
-	oldExists := store.states["old-workflow"]
-	runningExists := store.states["running-workflow"]
-	store.mu.RUnlock()
-
-	// Old completed workflow should be cleaned up
-	assert.Nil(t, oldExists, "old completed workflow should be cleaned up")
-
-	// Running workflow should still exist (not a terminal state)
-	assert.NotNil(t, runningExists, "running workflow should not be cleaned up")
+		// Old completed workflow should be cleaned up, running workflow must remain.
+		return oldExists == nil && runningExists != nil
+	}, 1*time.Second, 25*time.Millisecond, "expected stale terminal workflow to be cleaned up")
 }
 
 // TestInMemoryStateStore_GetStats tests statistics retrieval.