diff --git a/coderd/chatd/chatd.go b/coderd/chatd/chatd.go index aee39187f9..f08b22a8de 100644 --- a/coderd/chatd/chatd.go +++ b/coderd/chatd/chatd.go @@ -2351,11 +2351,12 @@ func (p *Server) tryAutoPromoteQueuedMessage( } // trackWorkspaceUsage bumps the workspace's last_used_at via the -// usage tracker. If wsID is not yet valid, it re-reads the chat -// from the DB to pick up late associations (e.g. create_workspace -// linking a workspace mid-conversation). The caller should store -// the returned value so that subsequent calls skip the DB lookup -// once a workspace has been found. +// usage tracker and extends the workspace's autostop deadline. If +// wsID is not yet valid, it re-reads the chat from the DB to pick +// up late associations (e.g. create_workspace linking a workspace +// mid-conversation). The caller should store the returned value so +// that subsequent calls skip the DB lookup once a workspace has +// been found. func (p *Server) trackWorkspaceUsage( ctx context.Context, chatID uuid.UUID, @@ -2375,6 +2376,22 @@ func (p *Server) trackWorkspaceUsage( } if wsID.Valid { p.usageTracker.Add(wsID.UUID) + // Bump the workspace autostop deadline. We pass time.Time{} + // for nextAutostart since we don't have access to + // TemplateScheduleStore here. The activity bump logic + // defaults to the template's activity_bump duration + // (typically 1 hour). Chat workspaces are never prebuilds, + // so no prebuild guard is needed (unlike reporter.go). + // + // This fires every heartbeat (~30s) but the SQL only + // writes when 5% of the deadline has elapsed — most calls + // perform a read-only CTE lookup with no UPDATE. + // + // Scaling note: for 10,000 active chats, this could lead to + // approx. 333 CTE queries/second. A cheap fix for this could + // be to heartbeat every Nth query. Leaving as potential future + // low-hanging fruit if needed. + workspacestats.ActivityBumpWorkspace(ctx, logger.Named("activity_bump"), p.db, wsID.UUID, time.Time{}) } return wsID } diff --git a/coderd/chatd/chatd_test.go b/coderd/chatd/chatd_test.go index eb58889283..02f6b4b408 100644 --- a/coderd/chatd/chatd_test.go +++ b/coderd/chatd/chatd_test.go @@ -17,6 +17,7 @@ import ( "time" "github.com/google/uuid" + "github.com/prometheus/client_golang/prometheus" "github.com/stretchr/testify/require" "go.uber.org/mock/gomock" "golang.org/x/xerrors" @@ -30,9 +31,12 @@ import ( "github.com/coder/coder/v2/coderd/coderdtest" "github.com/coder/coder/v2/coderd/database" "github.com/coder/coder/v2/coderd/database/db2sdk" + "github.com/coder/coder/v2/coderd/database/dbauthz" "github.com/coder/coder/v2/coderd/database/dbgen" "github.com/coder/coder/v2/coderd/database/dbtestutil" + "github.com/coder/coder/v2/coderd/database/dbtime" dbpubsub "github.com/coder/coder/v2/coderd/database/pubsub" + "github.com/coder/coder/v2/coderd/rbac" "github.com/coder/coder/v2/coderd/util/slice" "github.com/coder/coder/v2/coderd/workspacestats" "github.com/coder/coder/v2/codersdk" @@ -2187,19 +2191,51 @@ func TestHeartbeatBumpsWorkspaceUsage(t *testing.T) { return chattest.OpenAIResponse{StreamingChunks: chunks} })) - // Create a workspace that will be linked to the chat later, - // simulating the normal flow where a chat is created first - // and then creates a workspace via create_workspace. + // Create a workspace with a full build chain so we can verify + // both last_used_at (dormancy) and deadline (autostop) bumps. org := dbgen.Organization(t, db, database.Organization{}) - tmpl := dbgen.Template(t, db, database.Template{ + _ = dbgen.OrganizationMember(t, db, database.OrganizationMember{ + UserID: user.ID, + OrganizationID: org.ID, + }) + tv := dbgen.TemplateVersion(t, db, database.TemplateVersion{ OrganizationID: org.ID, CreatedBy: user.ID, }) + tmpl := dbgen.Template(t, db, database.Template{ + OrganizationID: org.ID, + ActiveVersionID: tv.ID, + CreatedBy: user.ID, + }) + require.NoError(t, db.UpdateTemplateScheduleByID(ctx, database.UpdateTemplateScheduleByIDParams{ + ID: tmpl.ID, + UpdatedAt: dbtime.Now(), + AllowUserAutostop: true, + ActivityBump: int64(time.Hour), + })) ws := dbgen.Workspace(t, db, database.WorkspaceTable{ OwnerID: user.ID, OrganizationID: org.ID, TemplateID: tmpl.ID, + Ttl: sql.NullInt64{Valid: true, Int64: int64(8 * time.Hour)}, }) + pj := dbgen.ProvisionerJob(t, db, nil, database.ProvisionerJob{ + OrganizationID: org.ID, + CompletedAt: sql.NullTime{ + Valid: true, + Time: dbtime.Now().Add(-30 * time.Minute), + }, + }) + // Build deadline is 30 minutes in the past — close enough to + // be bumped by the default 1-hour activity bump. + build := dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{ + WorkspaceID: ws.ID, + TemplateVersionID: tv.ID, + JobID: pj.ID, + Transition: database.WorkspaceTransitionStart, + Deadline: dbtime.Now().Add(-30 * time.Minute), + }) + originalDeadline := build.Deadline // Set up a short heartbeat interval and a UsageTracker that // flushes frequently so last_used_at gets updated in the DB. @@ -2212,9 +2248,13 @@ func TestHeartbeatBumpsWorkspaceUsage(t *testing.T) { t.Cleanup(func() { tracker.Close() }) logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) + // Wrap the database with dbauthz so the chatd server's + // AsChatd context is enforced on every query, matching + // production behavior. + authzDB := dbauthz.New(db, rbac.NewStrictCachingAuthorizer(prometheus.NewRegistry()), slogtest.Make(t, nil), coderdtest.AccessControlStorePointer()) server := chatd.New(chatd.Config{ Logger: logger, - Database: db, + Database: authzDB, ReplicaID: uuid.New(), Pubsub: ps, PendingChatAcquireInterval: 10 * time.Millisecond, @@ -2227,7 +2267,11 @@ func TestHeartbeatBumpsWorkspaceUsage(t *testing.T) { }) // Create a chat WITHOUT a workspace, the normal starting state. - chat, err := server.CreateChat(ctx, chatd.CreateOptions{ + // In production, CreateChat is called from the HTTP handler with + // the authenticated user's context. Here we use AsChatd since + // the chatd server processes everything under that role. + chatCtx := dbauthz.AsChatd(ctx) + chat, err := server.CreateChat(chatCtx, chatd.CreateOptions{ OwnerID: user.ID, Title: "usage-tracking-test", ModelConfigID: model.ID, @@ -2285,6 +2329,22 @@ func TestHeartbeatBumpsWorkspaceUsage(t *testing.T) { require.NoError(t, err) require.True(t, updatedWs.LastUsedAt.After(ws.LastUsedAt), "workspace last_used_at should have been bumped") + + // Verify the workspace build deadline was also extended. + // The SQL only writes when 5% of the deadline has elapsed — + // most calls perform a read-only CTE lookup. Wider ±2 + // minute tolerance than activitybump_test.go because the bump + // happens asynchronously via the heartbeat goroutine. + testutil.Eventually(ctx, t, func(ctx context.Context) bool { + updatedBuild, buildErr := db.GetLatestWorkspaceBuildByWorkspaceID(ctx, ws.ID) + if buildErr != nil || !updatedBuild.Deadline.After(originalDeadline) { + return false + } + now := dbtime.Now() + return updatedBuild.Deadline.After(now.Add(time.Hour-2*time.Minute)) && + updatedBuild.Deadline.Before(now.Add(time.Hour+2*time.Minute)) + }, testutil.IntervalFast, + "workspace build deadline should have been bumped to ~now+1h") } func TestHeartbeatNoWorkspaceNoBump(t *testing.T) { diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go index 5a2a5981cb..6f9b2d4bf4 100644 --- a/coderd/database/dbauthz/dbauthz.go +++ b/coderd/database/dbauthz/dbauthz.go @@ -705,7 +705,7 @@ var ( DisplayName: "Chat Daemon", Site: rbac.Permissions(map[string][]policy.Action{ rbac.ResourceChat.Type: {policy.ActionCreate, policy.ActionRead, policy.ActionUpdate, policy.ActionDelete}, - rbac.ResourceWorkspace.Type: {policy.ActionRead}, + rbac.ResourceWorkspace.Type: {policy.ActionRead, policy.ActionUpdate}, rbac.ResourceDeploymentConfig.Type: {policy.ActionRead}, rbac.ResourceUser.Type: {policy.ActionReadPersonal}, }), diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go index cc76fe7c8d..6d7ea9932b 100644 --- a/coderd/database/dbauthz/dbauthz_test.go +++ b/coderd/database/dbauthz/dbauthz_test.go @@ -5697,12 +5697,16 @@ func TestAsChatd(t *testing.T) { require.NoError(t, err, "chat %s should be allowed", action) } - // Workspace read. - err := auth.Authorize(ctx, actor, policy.ActionRead, rbac.ResourceWorkspace) - require.NoError(t, err, "workspace read should be allowed") + // Workspace read + update (update needed for ActivityBumpWorkspace). + for _, action := range []policy.Action{ + policy.ActionRead, policy.ActionUpdate, + } { + err := auth.Authorize(ctx, actor, action, rbac.ResourceWorkspace) + require.NoError(t, err, "workspace %s should be allowed", action) + } // DeploymentConfig read. - err = auth.Authorize(ctx, actor, policy.ActionRead, rbac.ResourceDeploymentConfig) + err := auth.Authorize(ctx, actor, policy.ActionRead, rbac.ResourceDeploymentConfig) require.NoError(t, err, "deployment config read should be allowed") // User read_personal (needed for GetUserChatCustomPrompt). @@ -5713,16 +5717,12 @@ func TestAsChatd(t *testing.T) { t.Run("DeniedActions", func(t *testing.T) { t.Parallel() - // Cannot write workspaces. - for _, action := range []policy.Action{ - policy.ActionUpdate, policy.ActionDelete, - } { - err := auth.Authorize(ctx, actor, action, rbac.ResourceWorkspace) - require.Error(t, err, "workspace %s should be denied", action) - } + // Cannot delete workspaces. + err := auth.Authorize(ctx, actor, policy.ActionDelete, rbac.ResourceWorkspace) + require.Error(t, err, "workspace delete should be denied") // Cannot access users. - err := auth.Authorize(ctx, actor, policy.ActionRead, rbac.ResourceUser) + err = auth.Authorize(ctx, actor, policy.ActionRead, rbac.ResourceUser) require.Error(t, err, "user read should be denied") // Cannot access API keys.