Skip to content

Commit fea2d05

Browse files
mafredrimafredri
committed
feat(coderd): bump workspace deadline on AI agent activity
AI agents report status via patchWorkspaceAgentAppStatus, but this wasn't extending workspace deadlines. This prevented proper task auto-pause behavior, causing tasks to pause mid-execution when there were no human connections. Now we call ActivityBumpWorkspace when agents report status, using the same logic as SSH/IDE connections. We skip bumping only when transitioning from terminal states (complete/failure) to terminal or idle states, since the work is done. Closes coder/internal#1251
1 parent e79f1d0 commit fea2d05

3 files changed

Lines changed: 270 additions & 0 deletions

File tree

‎coderd/autobuild/lifecycle_executor_test.go‎

Lines changed: 92 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ import (
1818
"github.com/coder/coder/v2/coderd/coderdtest"
1919
"github.com/coder/coder/v2/coderd/database"
2020
"github.com/coder/coder/v2/coderd/database/dbauthz"
21+
"github.com/coder/coder/v2/coderd/database/dbfake"
2122
"github.com/coder/coder/v2/coderd/database/dbgen"
2223
"github.com/coder/coder/v2/coderd/database/dbtestutil"
2324
"github.com/coder/coder/v2/coderd/database/pubsub"
@@ -29,6 +30,7 @@ import (
2930
"github.com/coder/coder/v2/coderd/schedule/cron"
3031
"github.com/coder/coder/v2/coderd/util/ptr"
3132
"github.com/coder/coder/v2/codersdk"
33+
"github.com/coder/coder/v2/codersdk/agentsdk"
3234
"github.com/coder/coder/v2/provisioner/echo"
3335
"github.com/coder/coder/v2/provisionersdk/proto"
3436
"github.com/coder/coder/v2/testutil"
@@ -522,6 +524,96 @@ func TestExecutorAutostopExtend(t *testing.T) {
522524
assert.Equal(t, database.WorkspaceTransitionStop, stats.Transitions[workspace.ID])
523525
}
524526

527+
func TestExecutorAutostopAIAgentActivity(t *testing.T) {
528+
t.Parallel()
529+
530+
var (
531+
ctx = context.Background()
532+
tickCh = make(chan time.Time)
533+
statsCh = make(chan autobuild.Stats)
534+
client, db = coderdtest.NewWithDatabase(t, &coderdtest.Options{
535+
AutobuildTicker: tickCh,
536+
IncludeProvisionerDaemon: true,
537+
AutobuildStats: statsCh,
538+
})
539+
)
540+
541+
// Given: we have a user with a task workspace.
542+
user := coderdtest.CreateFirstUser(t, client)
543+
r := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{
544+
OrganizationID: user.OrganizationID,
545+
OwnerID: user.UserID,
546+
}).WithTask(database.TaskTable{
547+
Name: "test-task",
548+
Prompt: "AI agent activity test task",
549+
}, &proto.App{Slug: "test-app"}).Do()
550+
551+
// Given: template has activity bump enabled.
552+
_, err := client.UpdateTemplateMeta(ctx, r.Template.ID, codersdk.UpdateTemplateMeta{
553+
DefaultTTLMillis: (2 * time.Hour).Milliseconds(),
554+
ActivityBumpMillis: time.Hour.Milliseconds(),
555+
})
556+
require.NoError(t, err)
557+
558+
// Set deadline to past to meet 5% threshold for activity bump.
559+
now := time.Now()
560+
pastDeadline := now.Add(-30 * time.Minute)
561+
err = db.UpdateWorkspaceBuildDeadlineByID(dbauthz.AsSystemRestricted(ctx), database.UpdateWorkspaceBuildDeadlineByIDParams{
562+
ID: r.Build.ID,
563+
UpdatedAt: now,
564+
Deadline: pastDeadline,
565+
MaxDeadline: time.Time{},
566+
})
567+
require.NoError(t, err)
568+
569+
// Given: agent reports "working" status.
570+
agentClient := agentsdk.New(client.URL, agentsdk.WithFixedToken(r.AgentToken))
571+
err = agentClient.PatchAppStatus(ctx, agentsdk.PatchAppStatus{
572+
AppSlug: "test-app",
573+
State: codersdk.WorkspaceAppStatusStateWorking,
574+
Message: "AI agent is working",
575+
})
576+
require.NoError(t, err)
577+
578+
p, err := coderdtest.GetProvisionerForTags(db, time.Now(), r.Workspace.OrganizationID, nil)
579+
require.NoError(t, err)
580+
581+
// When: the autobuild executor ticks after the past deadline.
582+
go func() {
583+
tickTime := now.Add(30 * time.Minute)
584+
coderdtest.UpdateProvisionerLastSeenAt(t, db, p.ID, tickTime)
585+
tickCh <- tickTime
586+
}()
587+
588+
// Then: nothing should happen and the workspace should stay running.
589+
stats := <-statsCh
590+
require.Len(t, stats.Errors, 0)
591+
require.Len(t, stats.Transitions, 0)
592+
593+
// Given: agent reports "complete" status.
594+
err = agentClient.PatchAppStatus(ctx, agentsdk.PatchAppStatus{
595+
AppSlug: "test-app",
596+
State: codersdk.WorkspaceAppStatusStateComplete,
597+
Message: "AI agent completed",
598+
})
599+
require.NoError(t, err)
600+
601+
// When: the autobuild executor ticks after the bumped deadline.
602+
go func() {
603+
tickTime := now.Add(time.Hour).Add(time.Minute)
604+
coderdtest.UpdateProvisionerLastSeenAt(t, db, p.ID, tickTime)
605+
tickCh <- tickTime
606+
close(tickCh)
607+
}()
608+
609+
// Then: the workspace should be stopped.
610+
stats = <-statsCh
611+
require.Len(t, stats.Errors, 0)
612+
require.Len(t, stats.Transitions, 1)
613+
require.Contains(t, stats.Transitions, r.Workspace.ID)
614+
require.Equal(t, database.WorkspaceTransitionStop, stats.Transitions[r.Workspace.ID])
615+
}
616+
525617
func TestExecutorAutostopAlreadyStopped(t *testing.T) {
526618
t.Parallel()
527619

‎coderd/workspaceagents.go‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,7 +40,9 @@ import (
4040
"github.com/coder/coder/v2/coderd/rbac/policy"
4141
"github.com/coder/coder/v2/coderd/telemetry"
4242
maputil "github.com/coder/coder/v2/coderd/util/maps"
43+
"github.com/coder/coder/v2/coderd/util/slice"
4344
strutil "github.com/coder/coder/v2/coderd/util/strings"
45+
"github.com/coder/coder/v2/coderd/workspacestats"
4446
"github.com/coder/coder/v2/coderd/wspubsub"
4547
"github.com/coder/coder/v2/codersdk"
4648
"github.com/coder/coder/v2/codersdk/agentsdk"
@@ -396,6 +398,32 @@ func (api *API) patchWorkspaceAgentAppStatus(rw http.ResponseWriter, r *http.Req
396398
// Notify on state change to Working/Idle for AI tasks
397399
api.enqueueAITaskStateNotification(ctx, app.ID, latestAppStatus, req.State, workspace, workspaceAgent)
398400

401+
// Determine if we should bump workspace activity based on state transitions.
402+
// We bump the deadline to prevent auto-stop (task pause) during AI activity,
403+
// except when transitioning from terminal states (complete/failure) to
404+
// terminal or idle states, which indicates the work is done.
405+
shouldBump := true
406+
terminalStates := []database.WorkspaceAppStatusState{
407+
database.WorkspaceAppStatusStateComplete,
408+
database.WorkspaceAppStatusStateFailure,
409+
}
410+
if latestAppStatus.ID != uuid.Nil {
411+
prevState := latestAppStatus.State
412+
newState := database.WorkspaceAppStatusState(req.State)
413+
414+
// Skip bump for terminal to terminal or terminal to idle.
415+
if slice.Contains(terminalStates, prevState) &&
416+
(slice.Contains(terminalStates, newState) || newState == database.WorkspaceAppStatusStateIdle) {
417+
shouldBump = false
418+
}
419+
}
420+
if shouldBump {
421+
// We pass time.Time{} for nextAutostart since we don't have access to
422+
// TemplateScheduleStore here. The activity bump logic handles this by
423+
// defaulting to the template's activity_bump duration (typically 1 hour).
424+
workspacestats.ActivityBumpWorkspace(ctx, api.Logger, api.Database, workspace.ID, time.Time{})
425+
}
426+
399427
httpapi.Write(ctx, rw, http.StatusOK, nil)
400428
}
401429

‎coderd/workspaceagents_test.go‎

Lines changed: 150 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -425,6 +425,156 @@ func TestWorkspaceAgentAppStatus(t *testing.T) {
425425
})
426426
}
427427

428+
func TestWorkspaceAgentAppStatus_ActivityBump(t *testing.T) {
429+
t.Parallel()
430+
431+
client, db := coderdtest.NewWithDatabase(t, nil)
432+
user := coderdtest.CreateFirstUser(t, client)
433+
434+
tests := []struct {
435+
name string
436+
prevState *codersdk.WorkspaceAppStatusState // nil means no previous state
437+
newState codersdk.WorkspaceAppStatusState
438+
shouldBump bool
439+
}{
440+
{
441+
name: "FirstStatusBumps",
442+
prevState: nil,
443+
newState: codersdk.WorkspaceAppStatusStateWorking,
444+
shouldBump: true,
445+
},
446+
{
447+
name: "WorkingToIdleBumps",
448+
prevState: ptr.Ref(codersdk.WorkspaceAppStatusStateWorking),
449+
newState: codersdk.WorkspaceAppStatusStateIdle,
450+
shouldBump: true,
451+
},
452+
{
453+
name: "WorkingToCompleteBumps",
454+
prevState: ptr.Ref(codersdk.WorkspaceAppStatusStateWorking),
455+
newState: codersdk.WorkspaceAppStatusStateComplete,
456+
shouldBump: true,
457+
},
458+
{
459+
name: "CompleteToIdleNoBump",
460+
prevState: ptr.Ref(codersdk.WorkspaceAppStatusStateComplete),
461+
newState: codersdk.WorkspaceAppStatusStateIdle,
462+
shouldBump: false,
463+
},
464+
{
465+
name: "CompleteToCompleteNoBump",
466+
prevState: ptr.Ref(codersdk.WorkspaceAppStatusStateComplete),
467+
newState: codersdk.WorkspaceAppStatusStateComplete,
468+
shouldBump: false,
469+
},
470+
{
471+
name: "FailureToIdleNoBump",
472+
prevState: ptr.Ref(codersdk.WorkspaceAppStatusStateFailure),
473+
newState: codersdk.WorkspaceAppStatusStateIdle,
474+
shouldBump: false,
475+
},
476+
{
477+
name: "FailureToFailureNoBump",
478+
prevState: ptr.Ref(codersdk.WorkspaceAppStatusStateFailure),
479+
newState: codersdk.WorkspaceAppStatusStateFailure,
480+
shouldBump: false,
481+
},
482+
{
483+
name: "CompleteToWorkingBumps",
484+
prevState: ptr.Ref(codersdk.WorkspaceAppStatusStateComplete),
485+
newState: codersdk.WorkspaceAppStatusStateWorking,
486+
shouldBump: true,
487+
},
488+
{
489+
name: "FailureToCompleteNoBump",
490+
prevState: ptr.Ref(codersdk.WorkspaceAppStatusStateFailure),
491+
newState: codersdk.WorkspaceAppStatusStateComplete,
492+
shouldBump: false,
493+
},
494+
}
495+
496+
for _, tt := range tests {
497+
tt := tt
498+
t.Run(tt.name, func(t *testing.T) {
499+
t.Parallel()
500+
501+
// Create workspace with agent and app.
502+
r := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{
503+
OrganizationID: user.OrganizationID,
504+
OwnerID: user.UserID,
505+
}).WithAgent(func(a []*proto.Agent) []*proto.Agent {
506+
a[0].Apps = []*proto.App{{Slug: "test-app"}}
507+
return a
508+
}).Do()
509+
510+
ctx := testutil.Context(t, testutil.WaitMedium)
511+
512+
// Configure template with activity_bump to enable deadline bumping.
513+
_, err := client.UpdateTemplateMeta(ctx, r.Template.ID, codersdk.UpdateTemplateMeta{
514+
ActivityBumpMillis: time.Hour.Milliseconds(),
515+
})
516+
require.NoError(t, err)
517+
518+
// Set the workspace build deadline to the past to ensure the 5%
519+
// threshold is met for activity bumping.
520+
pastDeadline := dbtime.Now().Add(-30 * time.Minute)
521+
err = db.UpdateWorkspaceBuildDeadlineByID(dbauthz.AsSystemRestricted(ctx), database.UpdateWorkspaceBuildDeadlineByIDParams{
522+
ID: r.Build.ID,
523+
UpdatedAt: dbtime.Now(),
524+
Deadline: pastDeadline,
525+
MaxDeadline: time.Time{},
526+
})
527+
require.NoError(t, err)
528+
529+
agentClient := agentsdk.New(client.URL, agentsdk.WithFixedToken(r.AgentToken))
530+
531+
// If there's a previous state, report it first.
532+
if tt.prevState != nil {
533+
err := agentClient.PatchAppStatus(ctx, agentsdk.PatchAppStatus{
534+
AppSlug: "test-app",
535+
State: *tt.prevState,
536+
Message: "previous state",
537+
})
538+
require.NoError(t, err)
539+
540+
// Reset deadline to past again to meet 5% threshold for next bump.
541+
err = db.UpdateWorkspaceBuildDeadlineByID(dbauthz.AsSystemRestricted(ctx), database.UpdateWorkspaceBuildDeadlineByIDParams{
542+
ID: r.Build.ID,
543+
UpdatedAt: dbtime.Now(),
544+
Deadline: pastDeadline,
545+
MaxDeadline: time.Time{},
546+
})
547+
require.NoError(t, err)
548+
}
549+
550+
// Get the deadline before the new status report.
551+
beforeBuild, err := db.GetWorkspaceBuildByID(dbauthz.AsSystemRestricted(ctx), r.Build.ID)
552+
require.NoError(t, err)
553+
beforeDeadline := beforeBuild.Deadline
554+
555+
// Report the new state.
556+
err = agentClient.PatchAppStatus(ctx, agentsdk.PatchAppStatus{
557+
AppSlug: "test-app",
558+
State: tt.newState,
559+
Message: "new state",
560+
})
561+
require.NoError(t, err)
562+
563+
// Check if deadline changed.
564+
afterBuild, err := db.GetWorkspaceBuildByID(dbauthz.AsSystemRestricted(ctx), r.Build.ID)
565+
require.NoError(t, err)
566+
afterDeadline := afterBuild.Deadline
567+
568+
didBump := afterDeadline.After(beforeDeadline)
569+
if tt.shouldBump {
570+
require.True(t, didBump, "wanted deadline to bump but it didn't")
571+
} else {
572+
require.False(t, didBump, "wanted deadline not to bump but it did")
573+
}
574+
})
575+
}
576+
}
577+
428578
func TestWorkspaceAgentConnectRPC(t *testing.T) {
429579
t.Parallel()
430580

0 commit comments

Comments
 (0)