opensource/coder
3
0
Fork 0
mirror of https://github.com/coder/coder.git synced 2026-09-24 15:04:27 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix(coderd): use DB liveness for chat workspace reuse (#23551)

Browse Source 
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.
This commit is contained in:
Ethan
2026-03-26 00:12:05 +11:00
committed by GitHub
parent 4ba9986301
commit c0a323a751
4 changed files with 329 additions and 112 deletions
Download Patch File Download Diff File Expand all files Collapse all files
coderd/coderd.go
+13 -12
View File
@@ -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(
coderd/x/chatd/chatd.go
+52 -48
View File
@@ -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,
coderd/x/chatd/chattool/createworkspace.go
+54 -49
View File
@@ -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:
coderd/x/chatd/chattool/createworkspace_test.go
+210 -3
View File
@@ -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}
}