diff --git a/coderd/x/chatd/chatd.go b/coderd/x/chatd/chatd.go index 927c6f28a4..6677aa52c7 100644 --- a/coderd/x/chatd/chatd.go +++ b/coderd/x/chatd/chatd.go @@ -64,6 +64,9 @@ const ( instructionCacheTTL = 5 * time.Minute workspaceDialValidationDelay = 5 * time.Second workspaceMCPDiscoveryTimeout = 5 * time.Second + // defaultDialTimeout matches the timeout used by ~8 other + // server-side AgentConn callers. + defaultDialTimeout = 30 * time.Second // DefaultChatHeartbeatInterval is the default time between chat // heartbeat updates while a chat is being processed. DefaultChatHeartbeatInterval = 30 * time.Second @@ -116,7 +119,17 @@ const ( defaultSubagentInstruction = "You are running as a delegated sub-agent chat. Complete the delegated task and provide clear, concise assistant responses for the parent agent." ) -var errChatHasNoWorkspaceAgent = xerrors.New("workspace has no running agent: the workspace is likely stopped. Use the start_workspace tool to start it") +var ( + errChatHasNoWorkspaceAgent = xerrors.New("workspace has no running agent: the workspace is likely stopped. Use the start_workspace tool to start it") + errChatAgentDisconnected = xerrors.New( + "workspace agent is disconnected and cannot execute tools. " + + "The workspace may need to be restarted from the Coder dashboard", + ) + errChatDialTimeout = xerrors.New( + "connection to the workspace agent timed out. " + + "The workspace may need to be restarted from the Coder dashboard", + ) +) // Server handles background processing of pending chats. type Server struct { @@ -133,6 +146,7 @@ type Server struct { agentConnFn AgentConnFunc agentInactiveDisconnectTimeout time.Duration + dialTimeout time.Duration instructionLookupTimeout time.Duration createWorkspaceFn chattool.CreateWorkspaceFn startWorkspaceFn chattool.StartWorkspaceFn @@ -547,6 +561,16 @@ func (c *turnWorkspaceContext) getWorkspaceConnLocked() (workspacesdk.AgentConn, return nil, agentRelease } +// isAgentUnreachable reports whether the given agent row's +// status is disconnected or timed out. It uses timestamp +// arithmetic on the row. The "connecting" state is allowed +// through because it is normal after a fresh workspace build. +func isAgentUnreachable(now time.Time, agent database.WorkspaceAgent, inactiveTimeout time.Duration) bool { + status := agent.Status(now, inactiveTimeout) + return status.Status == database.WorkspaceAgentStatusDisconnected || + status.Status == database.WorkspaceAgentStatusTimeout +} + func (c *turnWorkspaceContext) getWorkspaceConn(ctx context.Context) (workspacesdk.AgentConn, error) { if c.server.agentConnFn == nil { return nil, xerrors.New("workspace agent connector is not configured") @@ -555,8 +579,30 @@ func (c *turnWorkspaceContext) getWorkspaceConn(ctx context.Context) (workspaces for attempt := 0; attempt < 2; attempt++ { c.mu.Lock() currentConn, staleRelease := c.getWorkspaceConnLocked() + // Capture agentID in the same lock section as + // currentConn to prevent a TOCTOU race with + // concurrent clearCachedWorkspaceState calls. + agentID := c.agent.ID c.mu.Unlock() + + // Status check on cache hit: re-fetch the agent + // row so we see the latest heartbeat rather than + // a potentially stale cached copy. if currentConn != nil { + if agentID != uuid.Nil { + freshAgent, err := c.server.db.GetWorkspaceAgentByID(ctx, agentID) + if err != nil { + c.server.logger.Warn(ctx, "failed to re-fetch agent for status check", + slog.F("agent_id", agentID), + slog.Error(err), + ) + // On DB error the check re-runs on the + // next tool call. + } else if isAgentUnreachable(c.server.clock.Now(), freshAgent, c.server.agentInactiveDisconnectTimeout) { + c.clearCachedWorkspaceState() + return nil, errChatAgentDisconnected + } + } return currentConn, nil } if staleRelease != nil { @@ -568,8 +614,20 @@ func (c *turnWorkspaceContext) getWorkspaceConn(ctx context.Context) (workspaces return nil, err } + // Status check on cache miss: the freshly fetched + // agent row may already show disconnected. + if isAgentUnreachable(c.server.clock.Now(), agent, c.server.agentInactiveDisconnectTimeout) { + c.clearCachedWorkspaceState() + return nil, errChatAgentDisconnected + } + + // Wrap the dial in a timeout to bound the time spent + // waiting for an unreachable agent. The timeout scopes + // only dialWithLazyValidation, not ensureWorkspaceAgent + // or the post-dial binding steps. + dialCtx, dialCancel := context.WithTimeoutCause(ctx, c.server.dialTimeout, errChatDialTimeout) dialResult, err := dialWithLazyValidation( - ctx, + dialCtx, agent.ID, chatSnapshot.WorkspaceID.UUID, DialFunc(c.server.agentConnFn), @@ -578,13 +636,21 @@ func (c *turnWorkspaceContext) getWorkspaceConn(ctx context.Context) (workspaces }, workspaceDialValidationDelay, ) + dialCancel() if err != nil { if xerrors.Is(err, errChatHasNoWorkspaceAgent) { c.clearCachedWorkspaceState() + return nil, err + } + // Surface the dial timeout sentinel only when the + // parent context is still alive. If the parent was + // canceled (e.g. ErrInterrupted), its error must + // propagate unchanged so the chatloop can detect it. + if ctx.Err() == nil && errors.Is(context.Cause(dialCtx), errChatDialTimeout) { + return nil, errChatDialTimeout } return nil, err } - agentConn := dialResult.Conn agentRelease := dialResult.Release if dialResult.WasSwitched { @@ -3353,6 +3419,7 @@ func New(cfg Config) *Server { subscribeFn: cfg.SubscribeFn, agentConnFn: cfg.AgentConn, agentInactiveDisconnectTimeout: cfg.AgentInactiveDisconnectTimeout, + dialTimeout: defaultDialTimeout, instructionLookupTimeout: instructionLookupTimeout, createWorkspaceFn: cfg.CreateWorkspace, startWorkspaceFn: cfg.StartWorkspace, diff --git a/coderd/x/chatd/chatd_internal_test.go b/coderd/x/chatd/chatd_internal_test.go index 974d5fb707..f00ebf70a8 100644 --- a/coderd/x/chatd/chatd_internal_test.go +++ b/coderd/x/chatd/chatd_internal_test.go @@ -1111,9 +1111,12 @@ func TestPersistInstructionFilesIncludesAgentMetadata(t *testing.T) { }, nil).AnyTimes() logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) server := &Server{ - db: db, - logger: logger, - instructionLookupTimeout: 5 * time.Second, + db: db, + logger: logger, + clock: quartz.NewReal(), + instructionLookupTimeout: 5 * time.Second, + agentInactiveDisconnectTimeout: 30 * time.Second, + dialTimeout: 30 * time.Second, agentConnFn: func(context.Context, uuid.UUID) (workspacesdk.AgentConn, func(), error) { return conn, func() {}, nil }, @@ -1276,9 +1279,12 @@ func TestPersistInstructionFilesSentinelWithSkills(t *testing.T) { }, nil).AnyTimes() logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) server := &Server{ - db: db, - logger: logger, - instructionLookupTimeout: 5 * time.Second, + db: db, + logger: logger, + clock: quartz.NewReal(), + instructionLookupTimeout: 5 * time.Second, + agentInactiveDisconnectTimeout: 30 * time.Second, + dialTimeout: 30 * time.Second, agentConnFn: func(context.Context, uuid.UUID) (workspacesdk.AgentConn, func(), error) { return conn, func() {}, nil }, @@ -1361,9 +1367,12 @@ func TestPersistInstructionFilesSentinelNoSkillsClearsColumn(t *testing.T) { }, nil).AnyTimes() logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) server := &Server{ - db: db, - logger: logger, - instructionLookupTimeout: 5 * time.Second, + db: db, + logger: logger, + clock: quartz.NewReal(), + instructionLookupTimeout: 5 * time.Second, + agentInactiveDisconnectTimeout: 30 * time.Second, + dialTimeout: 30 * time.Second, agentConnFn: func(context.Context, uuid.UUID) (workspacesdk.AgentConn, func(), error) { return conn, func() {}, nil }, @@ -1589,7 +1598,12 @@ func TestTurnWorkspaceContextGetWorkspaceConnLazyValidationSwitchesWorkspaceAgen conn.EXPECT().SetExtraHeaders(gomock.Any()).Times(1) var dialed []uuid.UUID - server := &Server{db: db} + server := &Server{ + db: db, + clock: quartz.NewReal(), + agentInactiveDisconnectTimeout: 30 * time.Second, + dialTimeout: 30 * time.Second, + } server.agentConnFn = func(_ context.Context, agentID uuid.UUID) (workspacesdk.AgentConn, func(), error) { dialed = append(dialed, agentID) if agentID == staleAgentID { @@ -1649,7 +1663,12 @@ func TestTurnWorkspaceContextGetWorkspaceConnFastFailsWithoutCurrentAgent(t *tes Return([]database.WorkspaceAgent{}, nil). Times(1) - server := &Server{db: db} + server := &Server{ + db: db, + clock: quartz.NewReal(), + agentInactiveDisconnectTimeout: 30 * time.Second, + dialTimeout: 30 * time.Second, + } server.agentConnFn = func(context.Context, uuid.UUID) (workspacesdk.AgentConn, func(), error) { return nil, nil, xerrors.New("dial failed") } @@ -3689,3 +3708,455 @@ func TestSafeSweepIdleStreams_RecoversFromPanic(t *testing.T) { server.safeSweepIdleStreams(context.Background()) }, "safeSweepIdleStreams must recover panics so the janitor loop keeps running") } + +func TestGetWorkspaceConn_StatusCheck(t *testing.T) { + t.Parallel() + + type testCase struct { + name string + agent database.WorkspaceAgent + cacheHit bool + dbError bool + wantErr error + wantDialCalled bool + wantReleaseCalled bool + } + + tests := []testCase{ + { + name: "DisconnectedAgentCacheMiss", + agent: database.WorkspaceAgent{ + FirstConnectedAt: sql.NullTime{ + Time: time.Now().Add(-10 * time.Minute), + Valid: true, + }, + LastConnectedAt: sql.NullTime{ + Time: time.Now().Add(-10 * time.Minute), + Valid: true, + }, + }, + wantErr: errChatAgentDisconnected, + }, + { + name: "DisconnectedAgentCacheHit", + agent: database.WorkspaceAgent{ + FirstConnectedAt: sql.NullTime{ + Time: time.Now().Add(-10 * time.Minute), + Valid: true, + }, + LastConnectedAt: sql.NullTime{ + Time: time.Now().Add(-10 * time.Minute), + Valid: true, + }, + }, + cacheHit: true, + wantErr: errChatAgentDisconnected, + wantReleaseCalled: true, + }, + { + name: "TimedOutAgentCacheMiss", + agent: database.WorkspaceAgent{ + CreatedAt: time.Now().Add(-10 * time.Minute), + ConnectionTimeoutSeconds: 60, + }, + wantErr: errChatAgentDisconnected, + }, + { + // A "connecting" agent (never connected, normal after + // fresh build) must NOT be blocked by the status check. + name: "ConnectingAgentProceeds", + agent: database.WorkspaceAgent{}, + wantDialCalled: true, + }, + { + name: "CacheHitHealthyAgent", + agent: database.WorkspaceAgent{ + FirstConnectedAt: sql.NullTime{ + Time: time.Now().Add(-5 * time.Minute), + Valid: true, + }, + LastConnectedAt: sql.NullTime{ + Time: time.Now(), + Valid: true, + }, + }, + cacheHit: true, + }, + { + // When GetWorkspaceAgentByID returns an error on + // cache hit, the cached connection should be returned. + name: "CacheHitDBError", + agent: database.WorkspaceAgent{ + FirstConnectedAt: sql.NullTime{ + Time: time.Now().Add(-5 * time.Minute), + Valid: true, + }, + LastConnectedAt: sql.NullTime{ + Time: time.Now(), + Valid: true, + }, + }, + cacheHit: true, + dbError: true, + }, + } + + for _, tc := range tests { + tc := tc + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + ctrl := gomock.NewController(t) + db := dbmock.NewMockStore(ctrl) + + workspaceID := uuid.New() + agentID := uuid.New() + chat := database.Chat{ + ID: uuid.New(), + WorkspaceID: uuid.NullUUID{ + UUID: workspaceID, + Valid: true, + }, + AgentID: uuid.NullUUID{ + UUID: agentID, + Valid: true, + }, + } + + // Stamp the agent with the generated ID. + agent := tc.agent + agent.ID = agentID + + // Set up the DB mock for GetWorkspaceAgentByID. + if tc.dbError { + db.EXPECT().GetWorkspaceAgentByID(gomock.Any(), agentID). + Return(database.WorkspaceAgent{}, xerrors.New("connection reset")). + Times(1) + } else { + db.EXPECT().GetWorkspaceAgentByID(gomock.Any(), agentID). + Return(agent, nil). + Times(1) + } + + var dialCalled bool + var releaseCalled bool + + // For ConnectingAgentProceeds the dial returns a real + // mock conn; for all others it should not be reached. + var dialConn *agentconnmock.MockAgentConn + if tc.wantDialCalled { + dialConn = agentconnmock.NewMockAgentConn(ctrl) + dialConn.EXPECT().SetExtraHeaders(gomock.Any()).Times(1) + } + + server := &Server{ + db: db, + logger: slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}), + clock: quartz.NewReal(), + agentInactiveDisconnectTimeout: 30 * time.Second, + dialTimeout: defaultDialTimeout, + } + server.agentConnFn = func(_ context.Context, id uuid.UUID) (workspacesdk.AgentConn, func(), error) { + dialCalled = true + if dialConn != nil { + return dialConn, func() {}, nil + } + return nil, nil, xerrors.New("should not be called") + } + + chatStateMu := &sync.Mutex{} + currentChat := chat + workspaceCtx := turnWorkspaceContext{ + server: server, + chatStateMu: chatStateMu, + currentChat: ¤tChat, + loadChatSnapshot: func(context.Context, uuid.UUID) (database.Chat, error) { + return database.Chat{}, nil + }, + } + defer workspaceCtx.close() + + // For cache-hit tests, pre-populate the cached + // connection state. + var cachedConn *agentconnmock.MockAgentConn + if tc.cacheHit { + cachedConn = agentconnmock.NewMockAgentConn(ctrl) + workspaceCtx.agent = agent + workspaceCtx.agentLoaded = true + workspaceCtx.conn = cachedConn + workspaceCtx.releaseConn = func() { releaseCalled = true } + workspaceCtx.cachedWorkspaceID = chat.WorkspaceID + } + + ctx := testutil.Context(t, testutil.WaitShort) + gotConn, err := workspaceCtx.getWorkspaceConn(ctx) + + switch { + case tc.wantErr != nil: + require.Nil(t, gotConn) + require.ErrorIs(t, err, tc.wantErr) + case tc.cacheHit: + require.NoError(t, err) + require.Same(t, cachedConn, gotConn) + default: + // Cache-miss success (ConnectingAgentProceeds). + require.NoError(t, err) + require.Same(t, dialConn, gotConn) + } + + require.Equal(t, tc.wantDialCalled, dialCalled, "dial called") + require.Equal(t, tc.wantReleaseCalled, releaseCalled, "release called") + + // For error cases on cache-miss, the cache should be + // cleared. + if tc.wantErr != nil && !tc.cacheHit { + workspaceCtx.mu.Lock() + defer workspaceCtx.mu.Unlock() + require.False(t, workspaceCtx.agentLoaded) + require.Nil(t, workspaceCtx.conn) + } + // For cache-hit disconnect, the cache should also be + // cleared. + if tc.wantErr != nil && tc.cacheHit { + workspaceCtx.mu.Lock() + defer workspaceCtx.mu.Unlock() + require.False(t, workspaceCtx.agentLoaded) + require.Nil(t, workspaceCtx.conn) + } + }) + } +} + +func TestGetWorkspaceConn_DialTimeout(t *testing.T) { + // When dialWithLazyValidation blocks beyond the dial + // timeout, getWorkspaceConn should return + // errChatDialTimeout instead of hanging indefinitely. + t.Parallel() + + ctrl := gomock.NewController(t) + db := dbmock.NewMockStore(ctrl) + + workspaceID := uuid.New() + agentID := uuid.New() + chat := database.Chat{ + ID: uuid.New(), + WorkspaceID: uuid.NullUUID{ + UUID: workspaceID, + Valid: true, + }, + AgentID: uuid.NullUUID{ + UUID: agentID, + Valid: true, + }, + } + + // Agent appears connected so the status check passes. + connectedAgent := database.WorkspaceAgent{ + ID: agentID, + FirstConnectedAt: sql.NullTime{ + Time: time.Now().Add(-1 * time.Minute), + Valid: true, + }, + LastConnectedAt: sql.NullTime{ + Time: time.Now(), + Valid: true, + }, + } + + db.EXPECT().GetWorkspaceAgentByID(gomock.Any(), agentID). + Return(connectedAgent, nil). + Times(1) + + server := &Server{ + db: db, + clock: quartz.NewReal(), + agentInactiveDisconnectTimeout: 30 * time.Second, + dialTimeout: 10 * time.Millisecond, + } + // Dial blocks forever (simulates unreachable agent). + server.agentConnFn = func(ctx context.Context, _ uuid.UUID) (workspacesdk.AgentConn, func(), error) { + <-ctx.Done() + return nil, nil, ctx.Err() + } + + chatStateMu := &sync.Mutex{} + currentChat := chat + workspaceCtx := turnWorkspaceContext{ + server: server, + chatStateMu: chatStateMu, + currentChat: ¤tChat, + loadChatSnapshot: func(context.Context, uuid.UUID) (database.Chat, error) { return database.Chat{}, nil }, + } + defer workspaceCtx.close() + + ctx := testutil.Context(t, testutil.WaitShort) + gotConn, err := workspaceCtx.getWorkspaceConn(ctx) + require.Nil(t, gotConn) + require.ErrorIs(t, err, errChatDialTimeout) +} + +func TestGetWorkspaceConn_DialTimeoutParentCanceled(t *testing.T) { + // When the parent context is canceled, the parent's error + // must propagate unchanged (not wrapped as a dial timeout). + // This is critical because the chatloop checks + // context.Cause(ctx) for ErrInterrupted. + t.Parallel() + + ctrl := gomock.NewController(t) + db := dbmock.NewMockStore(ctrl) + + workspaceID := uuid.New() + agentID := uuid.New() + chat := database.Chat{ + ID: uuid.New(), + WorkspaceID: uuid.NullUUID{ + UUID: workspaceID, + Valid: true, + }, + AgentID: uuid.NullUUID{ + UUID: agentID, + Valid: true, + }, + } + + connectedAgent := database.WorkspaceAgent{ + ID: agentID, + FirstConnectedAt: sql.NullTime{ + Time: time.Now().Add(-1 * time.Minute), + Valid: true, + }, + LastConnectedAt: sql.NullTime{ + Time: time.Now(), + Valid: true, + }, + } + + db.EXPECT().GetWorkspaceAgentByID(gomock.Any(), agentID). + Return(connectedAgent, nil). + Times(1) + + parentErr := xerrors.New("parent canceled") + ctx, cancel := context.WithCancelCause(testutil.Context(t, testutil.WaitShort)) + + server := &Server{ + db: db, + clock: quartz.NewReal(), + agentInactiveDisconnectTimeout: 30 * time.Second, + // Use a very long dial timeout so the parent cancel fires + // first. + dialTimeout: 10 * time.Minute, + } + // Signal when the dial goroutine has started so we can + // cancel the parent at the right time without time.Sleep. + dialStarted := make(chan struct{}) + server.agentConnFn = func(ctx context.Context, _ uuid.UUID) (workspacesdk.AgentConn, func(), error) { + close(dialStarted) + <-ctx.Done() + return nil, nil, ctx.Err() + } + + chatStateMu := &sync.Mutex{} + currentChat := chat + workspaceCtx := turnWorkspaceContext{ + server: server, + chatStateMu: chatStateMu, + currentChat: ¤tChat, + loadChatSnapshot: func(context.Context, uuid.UUID) (database.Chat, error) { return database.Chat{}, nil }, + } + defer workspaceCtx.close() + + // Cancel the parent after the dial starts. + go func() { + <-dialStarted + cancel(parentErr) + }() + + gotConn, err := workspaceCtx.getWorkspaceConn(ctx) + require.Nil(t, gotConn) + // The error must NOT be errChatDialTimeout. + require.NotErrorIs(t, err, errChatDialTimeout) + // The parent context's error should propagate. + require.Error(t, err) + require.ErrorIs(t, err, context.Canceled) +} + +func TestGetWorkspaceConn_DialErrorNotMisclassifiedAsTimeout(t *testing.T) { + // Regression test: a non-timeout dial error (e.g. auth + // failure) with the parent context still alive must NOT be + // converted to errChatDialTimeout. Before the fix, + // dialCancel() poisoned dialCtx.Err(), causing all errors + // to be misclassified. + t.Parallel() + + ctrl := gomock.NewController(t) + db := dbmock.NewMockStore(ctrl) + + workspaceID := uuid.New() + agentID := uuid.New() + chat := database.Chat{ + ID: uuid.New(), + WorkspaceID: uuid.NullUUID{ + UUID: workspaceID, + Valid: true, + }, + AgentID: uuid.NullUUID{ + UUID: agentID, + Valid: true, + }, + } + + connectedAgent := database.WorkspaceAgent{ + ID: agentID, + FirstConnectedAt: sql.NullTime{ + Time: time.Now().Add(-1 * time.Minute), + Valid: true, + }, + LastConnectedAt: sql.NullTime{ + Time: time.Now(), + Valid: true, + }, + } + + db.EXPECT().GetWorkspaceAgentByID(gomock.Any(), agentID). + Return(connectedAgent, nil). + Times(1) + // When the initial dial fails immediately, dialWithLazyValidation + // calls resolveFastFailure which validates the binding. Mock the + // validation to return the same agent, triggering a synchronous + // redial that also returns the error. + db.EXPECT().GetWorkspaceAgentsInLatestBuildByWorkspaceID(gomock.Any(), workspaceID). + Return([]database.WorkspaceAgent{connectedAgent}, nil). + AnyTimes() + + dialErr := xerrors.New("authentication failed") + server := &Server{ + db: db, + clock: quartz.NewReal(), + agentInactiveDisconnectTimeout: 30 * time.Second, + // Generous timeout so the dial error fires well before + // the timeout. + dialTimeout: defaultDialTimeout, + } + server.agentConnFn = func(context.Context, uuid.UUID) (workspacesdk.AgentConn, func(), error) { + // Return an error immediately (not a timeout). + return nil, nil, dialErr + } + + chatStateMu := &sync.Mutex{} + currentChat := chat + workspaceCtx := turnWorkspaceContext{ + server: server, + chatStateMu: chatStateMu, + currentChat: ¤tChat, + loadChatSnapshot: func(context.Context, uuid.UUID) (database.Chat, error) { return database.Chat{}, nil }, + } + defer workspaceCtx.close() + + ctx := testutil.Context(t, testutil.WaitShort) + gotConn, err := workspaceCtx.getWorkspaceConn(ctx) + require.Nil(t, gotConn) + // Must NOT be misclassified as a dial timeout. + require.NotErrorIs(t, err, errChatDialTimeout) + // The original dial error should propagate. + require.ErrorContains(t, err, "authentication failed") +}