mirror of
https://github.com/coder/coder.git
synced 2026-06-02 20:48:20 +00:00
feat(coderd): bump workspace deadline on AI agent activity (#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
This commit is contained in:
Mathias Fredriksson
committed by
GitHub
GitHub
parent
5dcc9dd8ab
commit
4c7844ad3d
@@ -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()
|
||||
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
|
||||
@@ -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()
|
||||
|
||||
|
||||
Reference in New Issue
Block a user