Skip to content

Commit 4c7844a

Browse files
mafredrimafredri
authored
feat(coderd): bump workspace deadline on AI agent activity (coder#21584)
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 bump when transitioning to or from the working state. Closes coder/internal#1251
1 parent 5dcc9dd commit 4c7844a

3 files changed

Lines changed: 286 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: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,7 @@ import (
4141
"github.com/coder/coder/v2/coderd/telemetry"
4242
maputil "github.com/coder/coder/v2/coderd/util/maps"
4343
strutil "github.com/coder/coder/v2/coderd/util/strings"
44+
"github.com/coder/coder/v2/coderd/workspacestats"
4445
"github.com/coder/coder/v2/coderd/wspubsub"
4546
"github.com/coder/coder/v2/codersdk"
4647
"github.com/coder/coder/v2/codersdk/agentsdk"
@@ -396,6 +397,31 @@ func (api *API) patchWorkspaceAgentAppStatus(rw http.ResponseWriter, r *http.Req
396397
// Notify on state change to Working/Idle for AI tasks
397398
api.enqueueAITaskStateNotification(ctx, app.ID, latestAppStatus, req.State, workspace, workspaceAgent)
398399

400+
// Bump deadline when agent reports working or transitions away from working.
401+
// This prevents auto-pause during active work and gives users time to interact
402+
// after work completes.
403+
shouldBump := false
404+
newState := database.WorkspaceAppStatusState(req.State)
405+
406+
// Bump if reporting working state.
407+
if newState == database.WorkspaceAppStatusStateWorking {
408+
shouldBump = true
409+
}
410+
411+
// Bump if transitioning away from working state.
412+
if latestAppStatus.ID != uuid.Nil {
413+
prevState := latestAppStatus.State
414+
if prevState == database.WorkspaceAppStatusStateWorking {
415+
shouldBump = true
416+
}
417+
}
418+
if shouldBump {
419+
// We pass time.Time{} for nextAutostart since we don't have access to
420+
// TemplateScheduleStore here. The activity bump logic handles this by
421+
// defaulting to the template's activity_bump duration (typically 1 hour).
422+
workspacestats.ActivityBumpWorkspace(ctx, api.Logger, api.Database, workspace.ID, time.Time{})
423+
}
424+
399425
httpapi.Write(ctx, rw, http.StatusOK, nil)
400426
}
401427

‎coderd/workspaceagents_test.go‎

Lines changed: 168 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -425,6 +425,174 @@ 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+
name: "WorkingToFailureBumps",
496+
prevState: ptr.Ref(codersdk.WorkspaceAppStatusStateWorking),
497+
newState: codersdk.WorkspaceAppStatusStateFailure,
498+
shouldBump: true,
499+
},
500+
{
501+
name: "IdleToIdleNoBump",
502+
prevState: ptr.Ref(codersdk.WorkspaceAppStatusStateIdle),
503+
newState: codersdk.WorkspaceAppStatusStateIdle,
504+
shouldBump: false,
505+
},
506+
{
507+
name: "IdleToWorkingBumps",
508+
prevState: ptr.Ref(codersdk.WorkspaceAppStatusStateIdle),
509+
newState: codersdk.WorkspaceAppStatusStateWorking,
510+
shouldBump: true,
511+
},
512+
}
513+
514+
for _, tt := range tests {
515+
tt := tt
516+
t.Run(tt.name, func(t *testing.T) {
517+
t.Parallel()
518+
519+
// Create workspace with agent and app.
520+
r := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{
521+
OrganizationID: user.OrganizationID,
522+
OwnerID: user.UserID,
523+
}).WithAgent(func(a []*proto.Agent) []*proto.Agent {
524+
a[0].Apps = []*proto.App{{Slug: "test-app"}}
525+
return a
526+
}).Do()
527+
528+
ctx := testutil.Context(t, testutil.WaitMedium)
529+
530+
// Configure template with activity_bump to enable deadline bumping.
531+
_, err := client.UpdateTemplateMeta(ctx, r.Template.ID, codersdk.UpdateTemplateMeta{
532+
ActivityBumpMillis: time.Hour.Milliseconds(),
533+
})
534+
require.NoError(t, err)
535+
536+
// Set the workspace build deadline to the past to ensure the 5%
537+
// threshold is met for activity bumping.
538+
pastDeadline := dbtime.Now().Add(-30 * time.Minute)
539+
err = db.UpdateWorkspaceBuildDeadlineByID(dbauthz.AsSystemRestricted(ctx), database.UpdateWorkspaceBuildDeadlineByIDParams{
540+
ID: r.Build.ID,
541+
UpdatedAt: dbtime.Now(),
542+
Deadline: pastDeadline,
543+
MaxDeadline: time.Time{},
544+
})
545+
require.NoError(t, err)
546+
547+
agentClient := agentsdk.New(client.URL, agentsdk.WithFixedToken(r.AgentToken))
548+
549+
// If there's a previous state, report it first.
550+
if tt.prevState != nil {
551+
err := agentClient.PatchAppStatus(ctx, agentsdk.PatchAppStatus{
552+
AppSlug: "test-app",
553+
State: *tt.prevState,
554+
Message: "previous state",
555+
})
556+
require.NoError(t, err)
557+
558+
// Reset deadline to past again to meet 5% threshold for next bump.
559+
err = db.UpdateWorkspaceBuildDeadlineByID(dbauthz.AsSystemRestricted(ctx), database.UpdateWorkspaceBuildDeadlineByIDParams{
560+
ID: r.Build.ID,
561+
UpdatedAt: dbtime.Now(),
562+
Deadline: pastDeadline,
563+
MaxDeadline: time.Time{},
564+
})
565+
require.NoError(t, err)
566+
}
567+
568+
// Get the deadline before the new status report.
569+
beforeBuild, err := db.GetWorkspaceBuildByID(dbauthz.AsSystemRestricted(ctx), r.Build.ID)
570+
require.NoError(t, err)
571+
beforeDeadline := beforeBuild.Deadline
572+
573+
// Report the new state.
574+
err = agentClient.PatchAppStatus(ctx, agentsdk.PatchAppStatus{
575+
AppSlug: "test-app",
576+
State: tt.newState,
577+
Message: "new state",
578+
})
579+
require.NoError(t, err)
580+
581+
// Check if deadline changed.
582+
afterBuild, err := db.GetWorkspaceBuildByID(dbauthz.AsSystemRestricted(ctx), r.Build.ID)
583+
require.NoError(t, err)
584+
afterDeadline := afterBuild.Deadline
585+
586+
didBump := afterDeadline.After(beforeDeadline)
587+
if tt.shouldBump {
588+
require.True(t, didBump, "wanted deadline to bump but it didn't")
589+
} else {
590+
require.False(t, didBump, "wanted deadline not to bump but it did")
591+
}
592+
})
593+
}
594+
}
595+
428596
func TestWorkspaceAgentConnectRPC(t *testing.T) {
429597
t.Parallel()
430598

0 commit comments

Comments
 (0)