From c0a323a751b82d470f75a894ca272a70e6dde3ac Mon Sep 17 00:00:00 2001 From: Ethan <39577870+ethanndickson@users.noreply.github.com> Date: Thu, 26 Mar 2026 00:12:05 +1100 Subject: [PATCH] fix(coderd): use DB liveness for chat workspace reuse (#23551) create_workspace could create a replacement workspace after a single 5s agent dial failed, even when the existing workspace agent had recently checked in. That made temporary reachability blips look like dead workspaces and let chatd replace a running workspace too aggressively. Use the workspace agent's DB-backed status with the deployment's AgentInactiveDisconnectTimeout before allowing replacement. Recently connected and still-connecting agents now reuse the existing workspace, while disconnected or timed-out agents still allow a new workspace. This also threads the inactivity timeout through chatd and adds focused coverage for the reuse and replacement branches. --- coderd/coderd.go | 25 +- coderd/x/chatd/chatd.go | 100 ++++---- coderd/x/chatd/chattool/createworkspace.go | 103 +++++---- .../x/chatd/chattool/createworkspace_test.go | 213 +++++++++++++++++- 4 files changed, 329 insertions(+), 112 deletions(-) diff --git a/coderd/coderd.go b/coderd/coderd.go index 7a8865bfe3..5f21cb40d2 100644 --- a/coderd/coderd.go +++ b/coderd/coderd.go @@ -777,18 +777,19 @@ func New(options *Options) *API { } api.chatDaemon = chatd.New(chatd.Config{ - Logger: options.Logger.Named("chatd"), - Database: options.Database, - ReplicaID: api.ID, - SubscribeFn: options.ChatSubscribeFn, - MaxChatsPerAcquire: int32(maxChatsPerAcquire), //nolint:gosec // maxChatsPerAcquire is clamped to int32 range above. - ProviderAPIKeys: chatProviderAPIKeysFromDeploymentValues(options.DeploymentValues), - AgentConn: api.agentProvider.AgentConn, - CreateWorkspace: api.chatCreateWorkspace, - StartWorkspace: api.chatStartWorkspace, - Pubsub: options.Pubsub, - WebpushDispatcher: options.WebPushDispatcher, - UsageTracker: options.WorkspaceUsageTracker, + Logger: options.Logger.Named("chatd"), + Database: options.Database, + ReplicaID: api.ID, + SubscribeFn: options.ChatSubscribeFn, + MaxChatsPerAcquire: int32(maxChatsPerAcquire), //nolint:gosec // maxChatsPerAcquire is clamped to int32 range above. + ProviderAPIKeys: chatProviderAPIKeysFromDeploymentValues(options.DeploymentValues), + AgentConn: api.agentProvider.AgentConn, + AgentInactiveDisconnectTimeout: api.AgentInactiveDisconnectTimeout, + CreateWorkspace: api.chatCreateWorkspace, + StartWorkspace: api.chatStartWorkspace, + Pubsub: options.Pubsub, + WebpushDispatcher: options.WebPushDispatcher, + UsageTracker: options.WorkspaceUsageTracker, }) gitSyncLogger := options.Logger.Named("gitsync") refresher := gitsync.NewRefresher( diff --git a/coderd/x/chatd/chatd.go b/coderd/x/chatd/chatd.go index 22a3aab186..e0018b5ed8 100644 --- a/coderd/x/chatd/chatd.go +++ b/coderd/x/chatd/chatd.go @@ -98,12 +98,13 @@ type Server struct { subscribeFn SubscribeFn - agentConnFn AgentConnFunc - createWorkspaceFn chattool.CreateWorkspaceFn - startWorkspaceFn chattool.StartWorkspaceFn - pubsub pubsub.Pubsub - webpushDispatcher webpush.Dispatcher - providerAPIKeys chatprovider.ProviderAPIKeys + agentConnFn AgentConnFunc + agentInactiveDisconnectTimeout time.Duration + createWorkspaceFn chattool.CreateWorkspaceFn + startWorkspaceFn chattool.StartWorkspaceFn + pubsub pubsub.Pubsub + webpushDispatcher webpush.Dispatcher + providerAPIKeys chatprovider.ProviderAPIKeys // chatStreams stores per-chat stream state. Using sync.Map // gives each chat independent locking — concurrent chats @@ -1433,22 +1434,23 @@ func shouldQueueUserMessage(status database.ChatStatus) bool { // Config configures a chat processor. type Config struct { - Logger slog.Logger - Database database.Store - ReplicaID uuid.UUID - SubscribeFn SubscribeFn - PendingChatAcquireInterval time.Duration - MaxChatsPerAcquire int32 - InFlightChatStaleAfter time.Duration - ChatHeartbeatInterval time.Duration - AgentConn AgentConnFunc - CreateWorkspace chattool.CreateWorkspaceFn - StartWorkspace chattool.StartWorkspaceFn - Pubsub pubsub.Pubsub - ProviderAPIKeys chatprovider.ProviderAPIKeys - WebpushDispatcher webpush.Dispatcher - UsageTracker *workspacestats.UsageTracker - Clock quartz.Clock + Logger slog.Logger + Database database.Store + ReplicaID uuid.UUID + SubscribeFn SubscribeFn + PendingChatAcquireInterval time.Duration + MaxChatsPerAcquire int32 + InFlightChatStaleAfter time.Duration + ChatHeartbeatInterval time.Duration + AgentConn AgentConnFunc + AgentInactiveDisconnectTimeout time.Duration + CreateWorkspace chattool.CreateWorkspaceFn + StartWorkspace chattool.StartWorkspaceFn + Pubsub pubsub.Pubsub + ProviderAPIKeys chatprovider.ProviderAPIKeys + WebpushDispatcher webpush.Dispatcher + UsageTracker *workspacestats.UsageTracker + Clock quartz.Clock } // New creates a new chat processor. The processor polls for pending @@ -1488,25 +1490,26 @@ func New(cfg Config) *Server { } p := &Server{ - cancel: cancel, - closed: make(chan struct{}), - db: cfg.Database, - workerID: workerID, - logger: cfg.Logger.Named("processor"), - subscribeFn: cfg.SubscribeFn, - agentConnFn: cfg.AgentConn, - createWorkspaceFn: cfg.CreateWorkspace, - startWorkspaceFn: cfg.StartWorkspace, - pubsub: cfg.Pubsub, - webpushDispatcher: cfg.WebpushDispatcher, - providerAPIKeys: cfg.ProviderAPIKeys, - instructionCache: make(map[uuid.UUID]cachedInstruction), - pendingChatAcquireInterval: pendingChatAcquireInterval, - maxChatsPerAcquire: maxChatsPerAcquire, - inFlightChatStaleAfter: inFlightChatStaleAfter, - chatHeartbeatInterval: chatHeartbeatInterval, - usageTracker: cfg.UsageTracker, - clock: clk, + cancel: cancel, + closed: make(chan struct{}), + db: cfg.Database, + workerID: workerID, + logger: cfg.Logger.Named("processor"), + subscribeFn: cfg.SubscribeFn, + agentConnFn: cfg.AgentConn, + agentInactiveDisconnectTimeout: cfg.AgentInactiveDisconnectTimeout, + createWorkspaceFn: cfg.CreateWorkspace, + startWorkspaceFn: cfg.StartWorkspace, + pubsub: cfg.Pubsub, + webpushDispatcher: cfg.WebpushDispatcher, + providerAPIKeys: cfg.ProviderAPIKeys, + instructionCache: make(map[uuid.UUID]cachedInstruction), + pendingChatAcquireInterval: pendingChatAcquireInterval, + maxChatsPerAcquire: maxChatsPerAcquire, + inFlightChatStaleAfter: inFlightChatStaleAfter, + chatHeartbeatInterval: chatHeartbeatInterval, + usageTracker: cfg.UsageTracker, + clock: clk, } //nolint:gocritic // The chat processor uses a scoped chatd context. @@ -3383,13 +3386,14 @@ func (p *Server) runChat( OwnerID: chat.OwnerID, }), chattool.CreateWorkspace(chattool.CreateWorkspaceOptions{ - DB: p.db, - OwnerID: chat.OwnerID, - ChatID: chat.ID, - CreateFn: p.createWorkspaceFn, - AgentConnFn: chattool.AgentConnFunc(p.agentConnFn), - WorkspaceMu: &workspaceMu, - Logger: p.logger, + DB: p.db, + OwnerID: chat.OwnerID, + ChatID: chat.ID, + CreateFn: p.createWorkspaceFn, + AgentConnFn: chattool.AgentConnFunc(p.agentConnFn), + AgentInactiveDisconnectTimeout: p.agentInactiveDisconnectTimeout, + WorkspaceMu: &workspaceMu, + Logger: p.logger, }), chattool.StartWorkspace(chattool.StartWorkspaceOptions{ DB: p.db, diff --git a/coderd/x/chatd/chattool/createworkspace.go b/coderd/x/chatd/chattool/createworkspace.go index f92f9ed4f6..18fdc07e52 100644 --- a/coderd/x/chatd/chattool/createworkspace.go +++ b/coderd/x/chatd/chattool/createworkspace.go @@ -35,9 +35,6 @@ const ( // agentAttemptTimeout is the timeout for a single connection // attempt to the workspace agent during the retry loop. agentAttemptTimeout = 5 * time.Second - // agentPingTimeout is the timeout for a single agent ping - // when checking whether an existing workspace is alive. - agentPingTimeout = 5 * time.Second // startupScriptTimeout is the maximum time to wait for the // workspace agent's startup scripts to finish after the agent // is reachable. @@ -62,13 +59,14 @@ type AgentConnFunc func( // CreateWorkspaceOptions configures the create_workspace tool. type CreateWorkspaceOptions struct { - DB database.Store - OwnerID uuid.UUID - ChatID uuid.UUID - CreateFn CreateWorkspaceFn - AgentConnFn AgentConnFunc - WorkspaceMu *sync.Mutex - Logger slog.Logger + DB database.Store + OwnerID uuid.UUID + ChatID uuid.UUID + CreateFn CreateWorkspaceFn + AgentConnFn AgentConnFunc + AgentInactiveDisconnectTimeout time.Duration + WorkspaceMu *sync.Mutex + Logger slog.Logger } type createWorkspaceArgs struct { @@ -116,17 +114,12 @@ func CreateWorkspace(options CreateWorkspaceOptions) fantasy.AgentTool { } // Check for an existing workspace on the chat. - if options.DB != nil && options.ChatID != uuid.Nil { - existing, done, existErr := checkExistingWorkspace( - ctx, options.DB, options.ChatID, - options.AgentConnFn, - ) - if existErr != nil { - return fantasy.NewTextErrorResponse(existErr.Error()), nil - } - if done { - return toolResponse(existing), nil - } + existing, done, existErr := options.checkExistingWorkspace(ctx) + if existErr != nil { + return fantasy.NewTextErrorResponse(existErr.Error()), nil + } + if done { + return toolResponse(existing), nil } ownerID := options.OwnerID @@ -251,17 +244,23 @@ func CreateWorkspace(options CreateWorkspaceOptions) fantasy.AgentTool { }) } -// checkExistingWorkspace checks whether the chat already has a usable -// workspace. Returns the result map and true if the caller should -// return early (workspace exists and is alive or building). Returns -// false if the caller should proceed with creation (workspace is dead -// or missing). -func checkExistingWorkspace( +// checkExistingWorkspace checks whether the configured chat already has +// a usable workspace. Returns the result map and true if the caller +// should return early (workspace exists and is alive or building). +// Returns false if the caller should proceed with creation (workspace +// is dead or missing). +func (o CreateWorkspaceOptions) checkExistingWorkspace( ctx context.Context, - db database.Store, - chatID uuid.UUID, - agentConnFn AgentConnFunc, ) (map[string]any, bool, error) { + if o.DB == nil || o.ChatID == uuid.Nil { + return nil, false, nil + } + + db := o.DB + chatID := o.ChatID + agentConnFn := o.AgentConnFn + agentInactiveDisconnectTimeout := o.AgentInactiveDisconnectTimeout + chat, err := db.GetChatByID(ctx, chatID) if err != nil { return nil, false, xerrors.Errorf("load chat: %w", err) @@ -327,32 +326,38 @@ func checkExistingWorkspace( }, true, nil } - // Build succeeded — check if agent is reachable. + // Build succeeded — use the agent's recent DB-backed + // connection status to decide whether the workspace is + // still usable. agents, agentsErr := db.GetWorkspaceAgentsInLatestBuildByWorkspaceID(ctx, ws.ID) - if agentsErr == nil && len(agents) > 0 && agentConnFn != nil { - pingCtx, cancel := context.WithTimeout(ctx, agentPingTimeout) - conn, release, connErr := agentConnFn(pingCtx, agents[0].ID) - cancel() - if connErr == nil { - release() - _ = conn - // Agent is reachable; wait for startup scripts. - result := map[string]any{ - "created": false, - "workspace_name": ws.Name, - "status": "already_exists", - "message": "workspace is already running and reachable", - } - // Pass nil for agentConnFn since we already confirmed connectivity. + if agentsErr == nil && len(agents) > 0 { + status := agents[0].Status(agentInactiveDisconnectTimeout) + result := map[string]any{ + "created": false, + "workspace_name": ws.Name, + "status": "already_exists", + } + + switch status.Status { + case database.WorkspaceAgentStatusConnected: + result["message"] = "workspace is already running and recently connected" for k, v := range waitForAgentReady(ctx, db, agents[0].ID, nil) { result[k] = v } return result, true, nil + case database.WorkspaceAgentStatusConnecting: + result["message"] = "workspace exists and the agent is still connecting" + for k, v := range waitForAgentReady(ctx, db, agents[0].ID, agentConnFn) { + result[k] = v + } + return result, true, nil + case database.WorkspaceAgentStatusDisconnected, + database.WorkspaceAgentStatusTimeout: + // Agent is offline or never became ready — allow + // creation. } - // Agent unreachable — workspace is dead, allow - // creation. } - // No agent ID or no conn func — allow creation. + // No agent ID or no agent status — allow creation. return nil, false, nil default: diff --git a/coderd/x/chatd/chattool/createworkspace_test.go b/coderd/x/chatd/chattool/createworkspace_test.go index f1174f1b87..6257393c67 100644 --- a/coderd/x/chatd/chattool/createworkspace_test.go +++ b/coderd/x/chatd/chattool/createworkspace_test.go @@ -2,6 +2,7 @@ package chattool //nolint:testpackage // Uses internal symbols. import ( "context" + "database/sql" "fmt" "sync" "testing" @@ -228,6 +229,160 @@ func TestCreateWorkspace_GlobalTTL(t *testing.T) { } } +func TestCheckExistingWorkspace_ConnectedAgent(t *testing.T) { + t.Parallel() + ctrl := gomock.NewController(t) + db := dbmock.NewMockStore(ctrl) + + chatID := uuid.New() + workspaceID := uuid.New() + jobID := uuid.New() + agentID := uuid.New() + now := time.Now().UTC() + + expectExistingWorkspaceLookup( + db, + chatID, + workspaceID, + jobID, + "existing-workspace", + database.ProvisionerJobStatusSucceeded, + database.WorkspaceTransitionStart, + ) + db.EXPECT(). + GetWorkspaceAgentsInLatestBuildByWorkspaceID(gomock.Any(), workspaceID). + Return([]database.WorkspaceAgent{{ + ID: agentID, + CreatedAt: now.Add(-time.Minute), + FirstConnectedAt: validNullTime(now.Add(-45 * time.Second)), + LastConnectedAt: validNullTime(now.Add(-5 * time.Second)), + }}, nil) + db.EXPECT(). + GetWorkspaceAgentLifecycleStateByID(gomock.Any(), agentID). + Return(database.GetWorkspaceAgentLifecycleStateByIDRow{ + LifecycleState: database.WorkspaceAgentLifecycleStateReady, + }, nil) + + connFn := func(context.Context, uuid.UUID) (workspacesdk.AgentConn, func(), error) { + t.Fatalf("unexpected agent dial for connected workspace") + return nil, nil, xerrors.New("unexpected agent dial") + } + + options := testCheckExistingWorkspaceOptions(db, chatID, connFn) + result, done, err := options.checkExistingWorkspace(context.Background()) + require.NoError(t, err) + require.True(t, done) + require.Equal(t, "already_exists", result["status"]) + require.Equal(t, "existing-workspace", result["workspace_name"]) + require.Equal(t, "workspace is already running and recently connected", result["message"]) +} + +func TestCheckExistingWorkspace_ConnectingAgentWaits(t *testing.T) { + t.Parallel() + ctrl := gomock.NewController(t) + db := dbmock.NewMockStore(ctrl) + + chatID := uuid.New() + workspaceID := uuid.New() + jobID := uuid.New() + agentID := uuid.New() + now := time.Now().UTC() + connectCalls := 0 + + expectExistingWorkspaceLookup( + db, + chatID, + workspaceID, + jobID, + "existing-workspace", + database.ProvisionerJobStatusSucceeded, + database.WorkspaceTransitionStart, + ) + db.EXPECT(). + GetWorkspaceAgentsInLatestBuildByWorkspaceID(gomock.Any(), workspaceID). + Return([]database.WorkspaceAgent{{ + ID: agentID, + CreatedAt: now, + ConnectionTimeoutSeconds: 60, + }}, nil) + db.EXPECT(). + GetWorkspaceAgentLifecycleStateByID(gomock.Any(), agentID). + Return(database.GetWorkspaceAgentLifecycleStateByIDRow{ + LifecycleState: database.WorkspaceAgentLifecycleStateReady, + }, nil) + + connFn := func(context.Context, uuid.UUID) (workspacesdk.AgentConn, func(), error) { + connectCalls++ + return nil, func() {}, nil + } + + options := testCheckExistingWorkspaceOptions(db, chatID, connFn) + result, done, err := options.checkExistingWorkspace(context.Background()) + require.NoError(t, err) + require.True(t, done) + require.Equal(t, 1, connectCalls) + require.Equal(t, "already_exists", result["status"]) + require.Equal(t, "workspace exists and the agent is still connecting", result["message"]) +} + +func TestCheckExistingWorkspace_DeadAgentAllowsCreation(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + agent database.WorkspaceAgent + }{ + { + name: "Disconnected", + agent: database.WorkspaceAgent{ + ID: uuid.New(), + CreatedAt: time.Now().UTC().Add(-2 * time.Minute), + FirstConnectedAt: validNullTime(time.Now().UTC().Add(-2 * time.Minute)), + LastConnectedAt: validNullTime(time.Now().UTC().Add(-time.Minute)), + }, + }, + { + name: "TimedOut", + agent: database.WorkspaceAgent{ + ID: uuid.New(), + CreatedAt: time.Now().UTC().Add(-2 * time.Second), + ConnectionTimeoutSeconds: 1, + }, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + ctrl := gomock.NewController(t) + db := dbmock.NewMockStore(ctrl) + + chatID := uuid.New() + workspaceID := uuid.New() + jobID := uuid.New() + + expectExistingWorkspaceLookup( + db, + chatID, + workspaceID, + jobID, + "existing-workspace", + database.ProvisionerJobStatusSucceeded, + database.WorkspaceTransitionStart, + ) + db.EXPECT(). + GetWorkspaceAgentsInLatestBuildByWorkspaceID(gomock.Any(), workspaceID). + Return([]database.WorkspaceAgent{tc.agent}, nil) + + options := testCheckExistingWorkspaceOptions(db, chatID, nil) + result, done, err := options.checkExistingWorkspace(context.Background()) + require.NoError(t, err) + require.False(t, done) + require.Nil(t, result) + }) + } +} + func TestCheckExistingWorkspace_DeletedWorkspace(t *testing.T) { t.Parallel() ctrl := gomock.NewController(t) @@ -252,10 +407,62 @@ func TestCheckExistingWorkspace_DeletedWorkspace(t *testing.T) { Deleted: true, }, nil) - result, done, err := checkExistingWorkspace( - context.Background(), db, chatID, nil, - ) + options := testCheckExistingWorkspaceOptions(db, chatID, nil) + result, done, err := options.checkExistingWorkspace(context.Background()) require.NoError(t, err) require.False(t, done, "should allow creation for deleted workspace") require.Nil(t, result) } + +func testCheckExistingWorkspaceOptions( + db *dbmock.MockStore, + chatID uuid.UUID, + agentConnFn AgentConnFunc, +) CreateWorkspaceOptions { + return CreateWorkspaceOptions{ + DB: db, + ChatID: chatID, + AgentConnFn: agentConnFn, + AgentInactiveDisconnectTimeout: 30 * time.Second, + } +} + +func expectExistingWorkspaceLookup( + db *dbmock.MockStore, + chatID uuid.UUID, + workspaceID uuid.UUID, + jobID uuid.UUID, + workspaceName string, + jobStatus database.ProvisionerJobStatus, + transition database.WorkspaceTransition, +) { + db.EXPECT(). + GetChatByID(gomock.Any(), chatID). + Return(database.Chat{ + ID: chatID, + WorkspaceID: uuid.NullUUID{UUID: workspaceID, Valid: true}, + }, nil) + db.EXPECT(). + GetWorkspaceByID(gomock.Any(), workspaceID). + Return(database.Workspace{ + ID: workspaceID, + Name: workspaceName, + }, nil) + db.EXPECT(). + GetLatestWorkspaceBuildByWorkspaceID(gomock.Any(), workspaceID). + Return(database.WorkspaceBuild{ + WorkspaceID: workspaceID, + JobID: jobID, + Transition: transition, + }, nil) + db.EXPECT(). + GetProvisionerJobByID(gomock.Any(), jobID). + Return(database.ProvisionerJob{ + ID: jobID, + JobStatus: jobStatus, + }, nil) +} + +func validNullTime(t time.Time) sql.NullTime { + return sql.NullTime{Time: t, Valid: true} +}