diff --git a/coderd/autobuild/lifecycle_executor_test.go b/coderd/autobuild/lifecycle_executor_test.go index be6ccda582..630bbe14d8 100644 --- a/coderd/autobuild/lifecycle_executor_test.go +++ b/coderd/autobuild/lifecycle_executor_test.go @@ -18,6 +18,7 @@ import ( "github.com/coder/coder/v2/coderd/coderdtest" "github.com/coder/coder/v2/coderd/database" "github.com/coder/coder/v2/coderd/database/dbauthz" + "github.com/coder/coder/v2/coderd/database/dbfake" "github.com/coder/coder/v2/coderd/database/dbgen" "github.com/coder/coder/v2/coderd/database/dbtestutil" "github.com/coder/coder/v2/coderd/database/pubsub" @@ -29,6 +30,7 @@ import ( "github.com/coder/coder/v2/coderd/schedule/cron" "github.com/coder/coder/v2/coderd/util/ptr" "github.com/coder/coder/v2/codersdk" + "github.com/coder/coder/v2/codersdk/agentsdk" "github.com/coder/coder/v2/provisioner/echo" "github.com/coder/coder/v2/provisionersdk/proto" "github.com/coder/coder/v2/testutil" @@ -522,6 +524,96 @@ func TestExecutorAutostopExtend(t *testing.T) { assert.Equal(t, database.WorkspaceTransitionStop, stats.Transitions[workspace.ID]) } +func TestExecutorAutostopAIAgentActivity(t *testing.T) { + t.Parallel() + + var ( + ctx = context.Background() + tickCh = make(chan time.Time) + statsCh = make(chan autobuild.Stats) + client, db = coderdtest.NewWithDatabase(t, &coderdtest.Options{ + AutobuildTicker: tickCh, + IncludeProvisionerDaemon: true, + AutobuildStats: statsCh, + }) + ) + + // Given: we have a user with a task workspace. + user := coderdtest.CreateFirstUser(t, client) + r := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{ + OrganizationID: user.OrganizationID, + OwnerID: user.UserID, + }).WithTask(database.TaskTable{ + Name: "test-task", + Prompt: "AI agent activity test task", + }, &proto.App{Slug: "test-app"}).Do() + + // Given: template has activity bump enabled. + _, err := client.UpdateTemplateMeta(ctx, r.Template.ID, codersdk.UpdateTemplateMeta{ + DefaultTTLMillis: (2 * time.Hour).Milliseconds(), + ActivityBumpMillis: time.Hour.Milliseconds(), + }) + require.NoError(t, err) + + // Set deadline to past to meet 5% threshold for activity bump. + now := time.Now() + pastDeadline := now.Add(-30 * time.Minute) + err = db.UpdateWorkspaceBuildDeadlineByID(dbauthz.AsSystemRestricted(ctx), database.UpdateWorkspaceBuildDeadlineByIDParams{ + ID: r.Build.ID, + UpdatedAt: now, + Deadline: pastDeadline, + MaxDeadline: time.Time{}, + }) + require.NoError(t, err) + + // Given: agent reports "working" status. + agentClient := agentsdk.New(client.URL, agentsdk.WithFixedToken(r.AgentToken)) + err = agentClient.PatchAppStatus(ctx, agentsdk.PatchAppStatus{ + AppSlug: "test-app", + State: codersdk.WorkspaceAppStatusStateWorking, + Message: "AI agent is working", + }) + require.NoError(t, err) + + p, err := coderdtest.GetProvisionerForTags(db, time.Now(), r.Workspace.OrganizationID, nil) + require.NoError(t, err) + + // When: the autobuild executor ticks after the past deadline. + go func() { + tickTime := now.Add(30 * time.Minute) + coderdtest.UpdateProvisionerLastSeenAt(t, db, p.ID, tickTime) + tickCh <- tickTime + }() + + // Then: nothing should happen and the workspace should stay running. + stats := <-statsCh + require.Len(t, stats.Errors, 0) + require.Len(t, stats.Transitions, 0) + + // Given: agent reports "complete" status. + err = agentClient.PatchAppStatus(ctx, agentsdk.PatchAppStatus{ + AppSlug: "test-app", + State: codersdk.WorkspaceAppStatusStateComplete, + Message: "AI agent completed", + }) + require.NoError(t, err) + + // When: the autobuild executor ticks after the bumped deadline. + go func() { + tickTime := now.Add(time.Hour).Add(time.Minute) + coderdtest.UpdateProvisionerLastSeenAt(t, db, p.ID, tickTime) + tickCh <- tickTime + close(tickCh) + }() + + // Then: the workspace should be stopped. + stats = <-statsCh + require.Len(t, stats.Errors, 0) + require.Len(t, stats.Transitions, 1) + require.Contains(t, stats.Transitions, r.Workspace.ID) + require.Equal(t, database.WorkspaceTransitionStop, stats.Transitions[r.Workspace.ID]) +} + func TestExecutorAutostopAlreadyStopped(t *testing.T) { t.Parallel() diff --git a/coderd/workspaceagents.go b/coderd/workspaceagents.go index f97cfde85f..17a831188e 100644 --- a/coderd/workspaceagents.go +++ b/coderd/workspaceagents.go @@ -41,6 +41,7 @@ import ( "github.com/coder/coder/v2/coderd/telemetry" maputil "github.com/coder/coder/v2/coderd/util/maps" strutil "github.com/coder/coder/v2/coderd/util/strings" + "github.com/coder/coder/v2/coderd/workspacestats" "github.com/coder/coder/v2/coderd/wspubsub" "github.com/coder/coder/v2/codersdk" "github.com/coder/coder/v2/codersdk/agentsdk" @@ -396,6 +397,31 @@ func (api *API) patchWorkspaceAgentAppStatus(rw http.ResponseWriter, r *http.Req // Notify on state change to Working/Idle for AI tasks api.enqueueAITaskStateNotification(ctx, app.ID, latestAppStatus, req.State, workspace, workspaceAgent) + // Bump deadline when agent reports working or transitions away from working. + // This prevents auto-pause during active work and gives users time to interact + // after work completes. + shouldBump := false + newState := database.WorkspaceAppStatusState(req.State) + + // Bump if reporting working state. + if newState == database.WorkspaceAppStatusStateWorking { + shouldBump = true + } + + // Bump if transitioning away from working state. + if latestAppStatus.ID != uuid.Nil { + prevState := latestAppStatus.State + if prevState == database.WorkspaceAppStatusStateWorking { + shouldBump = true + } + } + if shouldBump { + // We pass time.Time{} for nextAutostart since we don't have access to + // TemplateScheduleStore here. The activity bump logic handles this by + // defaulting to the template's activity_bump duration (typically 1 hour). + workspacestats.ActivityBumpWorkspace(ctx, api.Logger, api.Database, workspace.ID, time.Time{}) + } + httpapi.Write(ctx, rw, http.StatusOK, nil) } diff --git a/coderd/workspaceagents_test.go b/coderd/workspaceagents_test.go index 101b2040e8..f055ff7c75 100644 --- a/coderd/workspaceagents_test.go +++ b/coderd/workspaceagents_test.go @@ -425,6 +425,174 @@ func TestWorkspaceAgentAppStatus(t *testing.T) { }) } +func TestWorkspaceAgentAppStatus_ActivityBump(t *testing.T) { + t.Parallel() + + client, db := coderdtest.NewWithDatabase(t, nil) + user := coderdtest.CreateFirstUser(t, client) + + tests := []struct { + name string + prevState *codersdk.WorkspaceAppStatusState // nil means no previous state + newState codersdk.WorkspaceAppStatusState + shouldBump bool + }{ + { + name: "FirstStatusBumps", + prevState: nil, + newState: codersdk.WorkspaceAppStatusStateWorking, + shouldBump: true, + }, + { + name: "WorkingToIdleBumps", + prevState: ptr.Ref(codersdk.WorkspaceAppStatusStateWorking), + newState: codersdk.WorkspaceAppStatusStateIdle, + shouldBump: true, + }, + { + name: "WorkingToCompleteBumps", + prevState: ptr.Ref(codersdk.WorkspaceAppStatusStateWorking), + newState: codersdk.WorkspaceAppStatusStateComplete, + shouldBump: true, + }, + { + name: "CompleteToIdleNoBump", + prevState: ptr.Ref(codersdk.WorkspaceAppStatusStateComplete), + newState: codersdk.WorkspaceAppStatusStateIdle, + shouldBump: false, + }, + { + name: "CompleteToCompleteNoBump", + prevState: ptr.Ref(codersdk.WorkspaceAppStatusStateComplete), + newState: codersdk.WorkspaceAppStatusStateComplete, + shouldBump: false, + }, + { + name: "FailureToIdleNoBump", + prevState: ptr.Ref(codersdk.WorkspaceAppStatusStateFailure), + newState: codersdk.WorkspaceAppStatusStateIdle, + shouldBump: false, + }, + { + name: "FailureToFailureNoBump", + prevState: ptr.Ref(codersdk.WorkspaceAppStatusStateFailure), + newState: codersdk.WorkspaceAppStatusStateFailure, + shouldBump: false, + }, + { + name: "CompleteToWorkingBumps", + prevState: ptr.Ref(codersdk.WorkspaceAppStatusStateComplete), + newState: codersdk.WorkspaceAppStatusStateWorking, + shouldBump: true, + }, + { + name: "FailureToCompleteNoBump", + prevState: ptr.Ref(codersdk.WorkspaceAppStatusStateFailure), + newState: codersdk.WorkspaceAppStatusStateComplete, + shouldBump: false, + }, + { + name: "WorkingToFailureBumps", + prevState: ptr.Ref(codersdk.WorkspaceAppStatusStateWorking), + newState: codersdk.WorkspaceAppStatusStateFailure, + shouldBump: true, + }, + { + name: "IdleToIdleNoBump", + prevState: ptr.Ref(codersdk.WorkspaceAppStatusStateIdle), + newState: codersdk.WorkspaceAppStatusStateIdle, + shouldBump: false, + }, + { + name: "IdleToWorkingBumps", + prevState: ptr.Ref(codersdk.WorkspaceAppStatusStateIdle), + newState: codersdk.WorkspaceAppStatusStateWorking, + shouldBump: true, + }, + } + + for _, tt := range tests { + tt := tt + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + // Create workspace with agent and app. + r := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{ + OrganizationID: user.OrganizationID, + OwnerID: user.UserID, + }).WithAgent(func(a []*proto.Agent) []*proto.Agent { + a[0].Apps = []*proto.App{{Slug: "test-app"}} + return a + }).Do() + + ctx := testutil.Context(t, testutil.WaitMedium) + + // Configure template with activity_bump to enable deadline bumping. + _, err := client.UpdateTemplateMeta(ctx, r.Template.ID, codersdk.UpdateTemplateMeta{ + ActivityBumpMillis: time.Hour.Milliseconds(), + }) + require.NoError(t, err) + + // Set the workspace build deadline to the past to ensure the 5% + // threshold is met for activity bumping. + pastDeadline := dbtime.Now().Add(-30 * time.Minute) + err = db.UpdateWorkspaceBuildDeadlineByID(dbauthz.AsSystemRestricted(ctx), database.UpdateWorkspaceBuildDeadlineByIDParams{ + ID: r.Build.ID, + UpdatedAt: dbtime.Now(), + Deadline: pastDeadline, + MaxDeadline: time.Time{}, + }) + require.NoError(t, err) + + agentClient := agentsdk.New(client.URL, agentsdk.WithFixedToken(r.AgentToken)) + + // If there's a previous state, report it first. + if tt.prevState != nil { + err := agentClient.PatchAppStatus(ctx, agentsdk.PatchAppStatus{ + AppSlug: "test-app", + State: *tt.prevState, + Message: "previous state", + }) + require.NoError(t, err) + + // Reset deadline to past again to meet 5% threshold for next bump. + err = db.UpdateWorkspaceBuildDeadlineByID(dbauthz.AsSystemRestricted(ctx), database.UpdateWorkspaceBuildDeadlineByIDParams{ + ID: r.Build.ID, + UpdatedAt: dbtime.Now(), + Deadline: pastDeadline, + MaxDeadline: time.Time{}, + }) + require.NoError(t, err) + } + + // Get the deadline before the new status report. + beforeBuild, err := db.GetWorkspaceBuildByID(dbauthz.AsSystemRestricted(ctx), r.Build.ID) + require.NoError(t, err) + beforeDeadline := beforeBuild.Deadline + + // Report the new state. + err = agentClient.PatchAppStatus(ctx, agentsdk.PatchAppStatus{ + AppSlug: "test-app", + State: tt.newState, + Message: "new state", + }) + require.NoError(t, err) + + // Check if deadline changed. + afterBuild, err := db.GetWorkspaceBuildByID(dbauthz.AsSystemRestricted(ctx), r.Build.ID) + require.NoError(t, err) + afterDeadline := afterBuild.Deadline + + didBump := afterDeadline.After(beforeDeadline) + if tt.shouldBump { + require.True(t, didBump, "wanted deadline to bump but it didn't") + } else { + require.False(t, didBump, "wanted deadline not to bump but it did") + } + }) + } +} + func TestWorkspaceAgentConnectRPC(t *testing.T) { t.Parallel()