From 14ed3e3644a67ca57dfa264ae83a8b436f68a24b Mon Sep 17 00:00:00 2001 From: Cian Johnston Date: Wed, 18 Mar 2026 19:07:21 +0000 Subject: [PATCH] feat: bump workspace last_used_at on chat heartbeat (#23205) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - coderd: Wires `options.WorkspaceUsageTracker` into the chatd config. - chatd: Adds `UsageTracker` and calls `UsageTracker.Add(workspaceID)` on each heartbeat tick - chatd: adds tests to verify `last_used_at` bump behaviour > 🤖 This PR was created with the help of Coder Agents, and will be reviewed by my human. 🧑‍💻 --- coderd/chatd/chatd.go | 57 ++++++++++- coderd/chatd/chatd_test.go | 194 +++++++++++++++++++++++++++++++++++++ coderd/coderd.go | 1 + 3 files changed, 250 insertions(+), 2 deletions(-) diff --git a/coderd/chatd/chatd.go b/coderd/chatd/chatd.go index a8c4b04c64..38e6f47417 100644 --- a/coderd/chatd/chatd.go +++ b/coderd/chatd/chatd.go @@ -31,6 +31,7 @@ import ( "github.com/coder/coder/v2/coderd/database/pubsub" coderdpubsub "github.com/coder/coder/v2/coderd/pubsub" "github.com/coder/coder/v2/coderd/webpush" + "github.com/coder/coder/v2/coderd/workspacestats" "github.com/coder/coder/v2/codersdk" "github.com/coder/coder/v2/codersdk/workspacesdk" "github.com/coder/quartz" @@ -46,7 +47,9 @@ const ( homeInstructionLookupTimeout = 5 * time.Second instructionCacheTTL = 5 * time.Minute - chatHeartbeatInterval = 60 * time.Second + // DefaultChatHeartbeatInterval is the default time between chat + // heartbeat updates while a chat is being processed. + DefaultChatHeartbeatInterval = 30 * time.Second maxChatSteps = 1200 // maxStreamBufferSize caps the number of events buffered // per chat during a single LLM step. When exceeded the @@ -96,10 +99,14 @@ type Server struct { instructionCacheMu sync.RWMutex instructionCache map[uuid.UUID]cachedInstruction + usageTracker *workspacestats.UsageTracker + clock quartz.Clock + // Configuration pendingChatAcquireInterval time.Duration maxChatsPerAcquire int32 inFlightChatStaleAfter time.Duration + chatHeartbeatInterval time.Duration } type cachedInstruction struct { @@ -1285,12 +1292,15 @@ type Config struct { 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 } // New creates a new chat processor. The processor polls for pending @@ -1314,6 +1324,16 @@ func New(cfg Config) *Server { maxChatsPerAcquire = DefaultMaxChatsPerAcquire } + chatHeartbeatInterval := cfg.ChatHeartbeatInterval + if chatHeartbeatInterval == 0 { + chatHeartbeatInterval = DefaultChatHeartbeatInterval + } + + clk := cfg.Clock + if clk == nil { + clk = quartz.NewReal() + } + workerID := cfg.ReplicaID if workerID == uuid.Nil { workerID = uuid.New() @@ -1336,6 +1356,9 @@ func New(cfg Config) *Server { pendingChatAcquireInterval: pendingChatAcquireInterval, maxChatsPerAcquire: maxChatsPerAcquire, inFlightChatStaleAfter: inFlightChatStaleAfter, + chatHeartbeatInterval: chatHeartbeatInterval, + usageTracker: cfg.UsageTracker, + clock: clk, } //nolint:gocritic // The chat processor uses a scoped chatd context. @@ -2230,6 +2253,35 @@ func (p *Server) tryAutoPromoteQueuedMessage( return &msg, remainingQueuedMessages, true, nil } +// 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. +func (p *Server) trackWorkspaceUsage( + ctx context.Context, + chatID uuid.UUID, + wsID uuid.NullUUID, + logger slog.Logger, +) uuid.NullUUID { + if p.usageTracker == nil { + return wsID + } + if !wsID.Valid { + latest, err := p.db.GetChatByID(ctx, chatID) + if err != nil { + logger.Warn(ctx, "failed to re-read chat for workspace association", slog.Error(err)) + return wsID + } + wsID = latest.WorkspaceID + } + if wsID.Valid { + p.usageTracker.Add(wsID.UUID) + } + return wsID +} + func (p *Server) processChat(ctx context.Context, chat database.Chat) { logger := p.logger.With(slog.F("chat_id", chat.ID)) logger.Info(ctx, "processing chat request") @@ -2248,7 +2300,7 @@ func (p *Server) processChat(ctx context.Context, chat database.Chat) { // worker is still alive. The goroutine stops when chatCtx is // canceled (either by completion or interruption). go func() { - ticker := time.NewTicker(chatHeartbeatInterval) + ticker := p.clock.NewTicker(p.chatHeartbeatInterval, "chatd", "heartbeat") defer ticker.Stop() for { select { @@ -2267,6 +2319,7 @@ func (p *Server) processChat(ctx context.Context, chat database.Chat) { cancel(chatloop.ErrInterrupted) return } + chat.WorkspaceID = p.trackWorkspaceUsage(chatCtx, chat.ID, chat.WorkspaceID, logger) } } }() diff --git a/coderd/chatd/chatd_test.go b/coderd/chatd/chatd_test.go index 61565a523e..3d732726a0 100644 --- a/coderd/chatd/chatd_test.go +++ b/coderd/chatd/chatd_test.go @@ -33,6 +33,7 @@ import ( "github.com/coder/coder/v2/coderd/database/dbtestutil" dbpubsub "github.com/coder/coder/v2/coderd/database/pubsub" "github.com/coder/coder/v2/coderd/util/slice" + "github.com/coder/coder/v2/coderd/workspacestats" "github.com/coder/coder/v2/codersdk" "github.com/coder/coder/v2/codersdk/workspacesdk" "github.com/coder/coder/v2/codersdk/workspacesdk/agentconnmock" @@ -1956,6 +1957,199 @@ func TestStartWorkspaceTool_EndToEnd(t *testing.T) { require.True(t, foundToolResultInSecondCall, "expected second streamed model call to include start_workspace tool output") } +func TestHeartbeatBumpsWorkspaceUsage(t *testing.T) { + t.Parallel() + + db, ps := dbtestutil.NewDB(t) + + ctx := testutil.Context(t, testutil.WaitLong) + user, model := seedChatDependencies(ctx, t, db) + setOpenAIProviderBaseURL(ctx, t, db, chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { + if !req.Stream { + return chattest.OpenAINonStreamingResponse("ok") + } + // Block until the request context is canceled so the chat + // stays in a processing state long enough for heartbeats + // to fire. + chunks := make(chan chattest.OpenAIChunk) + go func() { + defer close(chunks) + <-req.Context().Done() + }() + 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. + org := dbgen.Organization(t, db, database.Organization{}) + tmpl := dbgen.Template(t, db, database.Template{ + OrganizationID: org.ID, + CreatedBy: user.ID, + }) + ws := dbgen.Workspace(t, db, database.WorkspaceTable{ + OwnerID: user.ID, + OrganizationID: org.ID, + TemplateID: tmpl.ID, + }) + + // Set up a short heartbeat interval and a UsageTracker that + // flushes frequently so last_used_at gets updated in the DB. + flushTick := make(chan time.Time) + flushDone := make(chan int, 1) + tracker := workspacestats.NewTracker(db, + workspacestats.TrackerWithTickFlush(flushTick, flushDone), + workspacestats.TrackerWithLogger(slogtest.Make(t, nil)), + ) + t.Cleanup(func() { tracker.Close() }) + + logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) + server := chatd.New(chatd.Config{ + Logger: logger, + Database: db, + ReplicaID: uuid.New(), + Pubsub: ps, + PendingChatAcquireInterval: 10 * time.Millisecond, + InFlightChatStaleAfter: testutil.WaitLong, + ChatHeartbeatInterval: 100 * time.Millisecond, + UsageTracker: tracker, + }) + t.Cleanup(func() { + require.NoError(t, server.Close()) + }) + + // Create a chat WITHOUT a workspace, the normal starting state. + chat, err := server.CreateChat(ctx, chatd.CreateOptions{ + OwnerID: user.ID, + Title: "usage-tracking-test", + ModelConfigID: model.ID, + InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("hello")}, + }) + require.NoError(t, err) + + // Wait for the chat to start processing and at least one + // heartbeat to fire. + testutil.Eventually(ctx, t, func(ctx context.Context) bool { + fromDB, listErr := db.GetChatByID(ctx, chat.ID) + if listErr != nil { + return false + } + return fromDB.Status == database.ChatStatusRunning && + fromDB.HeartbeatAt.Valid && + fromDB.HeartbeatAt.Time.After(fromDB.CreatedAt) + }, testutil.IntervalFast, + "chat should be running with at least one heartbeat") + + // Flush the tracker and verify nothing was tracked yet + // (no workspace linked). + testutil.RequireSend(ctx, t, flushTick, time.Now()) + count := testutil.RequireReceive(ctx, t, flushDone) + require.Equal(t, 0, count, + "expected no workspaces to be flushed before association") + + // Link the workspace to the chat in the DB, simulating what + // the create_workspace tool does mid-conversation. + _, err = db.UpdateChatWorkspace(ctx, database.UpdateChatWorkspaceParams{ + WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, + ID: chat.ID, + }) + require.NoError(t, err) + + // The heartbeat re-reads the workspace association from the DB + // on each tick. Wait for the tracker to pick it up. + testutil.Eventually(ctx, t, func(ctx context.Context) bool { + select { + case flushTick <- time.Now(): + case <-ctx.Done(): + return false + } + select { + case c := <-flushDone: + return c > 0 + case <-ctx.Done(): + return false + } + }, testutil.IntervalMedium, + "expected usage tracker to flush the late-associated workspace") + + // Verify the workspace's last_used_at was actually updated. + updatedWs, err := db.GetWorkspaceByID(ctx, ws.ID) + require.NoError(t, err) + require.True(t, updatedWs.LastUsedAt.After(ws.LastUsedAt), + "workspace last_used_at should have been bumped") +} + +func TestHeartbeatNoWorkspaceNoBump(t *testing.T) { + t.Parallel() + + db, ps := dbtestutil.NewDB(t) + + ctx := testutil.Context(t, testutil.WaitLong) + user, model := seedChatDependencies(ctx, t, db) + setOpenAIProviderBaseURL(ctx, t, db, chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { + if !req.Stream { + return chattest.OpenAINonStreamingResponse("ok") + } + chunks := make(chan chattest.OpenAIChunk) + go func() { + defer close(chunks) + <-req.Context().Done() + }() + return chattest.OpenAIResponse{StreamingChunks: chunks} + })) + + // Set up UsageTracker with manual tick/flush. + usageTickCh := make(chan time.Time) + flushCh := make(chan int, 1) + tracker := workspacestats.NewTracker(db, + workspacestats.TrackerWithTickFlush(usageTickCh, flushCh), + workspacestats.TrackerWithLogger(slogtest.Make(t, nil)), + ) + t.Cleanup(func() { tracker.Close() }) + + logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) + server := chatd.New(chatd.Config{ + Logger: logger, + Database: db, + ReplicaID: uuid.New(), + Pubsub: ps, + PendingChatAcquireInterval: 10 * time.Millisecond, + InFlightChatStaleAfter: testutil.WaitLong, + ChatHeartbeatInterval: 100 * time.Millisecond, + }) + t.Cleanup(func() { + require.NoError(t, server.Close()) + }) + + // Create a chat WITHOUT linking a workspace. + chat, err := server.CreateChat(ctx, chatd.CreateOptions{ + OwnerID: user.ID, + Title: "no-workspace-test", + ModelConfigID: model.ID, + InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("hello")}, + }) + require.NoError(t, err) + + // Wait for the chat to be acquired and at least one heartbeat + // to fire. + testutil.Eventually(ctx, t, func(ctx context.Context) bool { + fromDB, listErr := db.GetChatByID(ctx, chat.ID) + if listErr != nil { + return false + } + return fromDB.Status == database.ChatStatusRunning && + fromDB.HeartbeatAt.Valid && + fromDB.HeartbeatAt.Time.After(fromDB.CreatedAt) + }, testutil.IntervalFast, + "chat should be running with at least one heartbeat") + + // Flush the tracker. Since no workspace was linked, count + // should be 0. + testutil.RequireSend(ctx, t, usageTickCh, time.Now()) + count := testutil.RequireReceive(ctx, t, flushCh) + require.Equal(t, 0, count, "expected no workspaces to be flushed when chat has no workspace") +} + func newTestServer( t *testing.T, db database.Store, diff --git a/coderd/coderd.go b/coderd/coderd.go index 15984d7170..99a2c6cbbd 100644 --- a/coderd/coderd.go +++ b/coderd/coderd.go @@ -787,6 +787,7 @@ func New(options *Options) *API { StartWorkspace: api.chatStartWorkspace, Pubsub: options.Pubsub, WebpushDispatcher: options.WebPushDispatcher, + UsageTracker: options.WorkspaceUsageTracker, }) gitSyncLogger := options.Logger.Named("gitsync") refresher := gitsync.NewRefresher(