feat(coderd): add stop_workspace chatd tool and recovery classification (#24997)

## Summary

Adds a `stop_workspace` tool to chatd so the model can recover from the
"workspace running but agent dead" failure mode (e.g. an OOM that leaves
the workspace running but the agent unreachable) by stopping and then
starting the workspace.

<img width="924" height="742" alt="image"
src="https://github.com/user-attachments/assets/279dedb6-6e29-4fe1-8754-3a1f01e538bf"
/>



## What changed

**New `stop_workspace` chatd tool**
(`coderd/x/chatd/chattool/stopworkspace.go`). Mirrors `start_workspace`:
shares `WorkspaceMu` to serialize with create/start, waits for any
in-progress build before issuing a stop, and is idempotent only after a
successful Stop transition. Failed stop builds re-attempt rather than
reporting success.

**New `chatStopWorkspace` coderd hook** (`coderd/exp_chats.go`). Mirrors
`chatStartWorkspace` minus the `RequireActiveVersion` gate. Stop should
not be blocked by template version policy.

**Differentiated recovery sentinels** (`coderd/x/chatd/chatd.go`).
`errChatAgentDisconnected` instructs the model to call `stop_workspace`
then `start_workspace`. `errChatDialTimeout` instructs a single retry,
then user escalation if it repeats. The previous single message
conflated transient and persistent failures.

**Two-signal recovery gate.** Recovery is only surfaced when a tool call
times out *and* a fresh DB read of the latest workspace agent says
`Disconnected`. The previous draft escalated on the DB read alone, which
would fire on a 30-second heartbeat blip (e.g. agent respawn) and prompt
a destructive stop/start unnecessarily.

**Cache-hit disconnected handling** now clears the cache and retries a
fresh dial before escalating, rather than returning the recovery
sentinel immediately. Latest-agent classification uses
`GetWorkspaceAgentsInLatestBuildByWorkspaceID` instead of the chat's
bound `AgentID`, so stale bindings after a rebuild don't misclassify.

**Shared chattool helpers** in `coderd/x/chatd/chattool/chattool.go`:
`latestWorkspaceBuildAndJob`, `publishBuildBinding`,
`provisionerJobTerminal`. Applied to both `start_workspace` and
`stop_workspace`.

## Notes

- Reverts an earlier draft that widened `ask_user_question` to root
standard turns. Plan-mode-only behavior is restored.
- The `stop_workspace` tool currently renders via the generic chat
tool-call UI. A follow-up frontend PR will prettify the `stop_workspace`
tool and style it like the `start_workspace` tool.
- Never-connected (`Timeout` status) agents are intentionally excluded
from recovery. They indicate template or startup failure, not the
running-but-dead case this PR targets.

Closes CODAGT-315
This commit is contained in:
Ethan
2026-05-11 16:23:07 +10:00
committed by GitHub
parent cee504e8a0
commit bd6cc1aaf2
11 changed files with 1250 additions and 114 deletions
+63
View File
@@ -1,12 +1,16 @@
package chattool
import (
"context"
"encoding/json"
"unicode/utf8"
"charm.land/fantasy"
"github.com/google/uuid"
"golang.org/x/xerrors"
"cdr.dev/slog/v3"
"github.com/coder/coder/v2/coderd/database"
"github.com/coder/coder/v2/codersdk"
)
@@ -54,6 +58,65 @@ func responseErrorResult(resp codersdk.Response) map[string]any {
return result
}
func latestWorkspaceBuildAndJob(
ctx context.Context,
db database.Store,
workspaceID uuid.UUID,
) (database.WorkspaceBuild, database.ProvisionerJob, error) {
build, err := db.GetLatestWorkspaceBuildByWorkspaceID(ctx, workspaceID)
if err != nil {
return database.WorkspaceBuild{}, database.ProvisionerJob{}, xerrors.Errorf("get latest build: %w", err)
}
job, err := db.GetProvisionerJobByID(ctx, build.JobID)
if err != nil {
return database.WorkspaceBuild{}, database.ProvisionerJob{}, xerrors.Errorf("get provisioner job: %w", err)
}
return build, job, nil
}
func publishBuildBinding(
ctx context.Context,
db database.Store,
logger slog.Logger,
chatID uuid.UUID,
workspaceID uuid.UUID,
buildID uuid.UUID,
onChatUpdated func(database.Chat),
) {
updatedChat, bindErr := db.UpdateChatWorkspaceBinding(ctx, database.UpdateChatWorkspaceBindingParams{
ID: chatID,
WorkspaceID: uuid.NullUUID{UUID: workspaceID, Valid: true},
BuildID: uuid.NullUUID{
UUID: buildID,
Valid: buildID != uuid.Nil,
},
AgentID: uuid.NullUUID{},
})
if bindErr != nil {
logger.Error(ctx, "failed to persist build ID on chat binding",
slog.F("chat_id", chatID),
slog.F("build_id", buildID),
slog.Error(bindErr),
)
return
}
if onChatUpdated != nil {
onChatUpdated(updatedChat)
}
}
func provisionerJobTerminal(status database.ProvisionerJobStatus) bool {
switch status {
case database.ProvisionerJobStatusSucceeded,
database.ProvisionerJobStatusFailed,
database.ProvisionerJobStatusCanceled:
return true
default:
return false
}
}
func truncateRunes(value string, maxLen int) string {
if maxLen <= 0 || value == "" {
return ""
+5 -48
View File
@@ -56,7 +56,7 @@ func StartWorkspace(db database.Store, chatID uuid.UUID, options StartWorkspaceO
return fantasy.NewTextErrorResponse("workspace starter is not configured"), nil
}
// Serialize with create_workspace to prevent races.
// Serialize with create_workspace and stop_workspace to prevent races.
if options.WorkspaceMu != nil {
options.WorkspaceMu.Lock()
defer options.WorkspaceMu.Unlock()
@@ -86,18 +86,9 @@ func StartWorkspace(db database.Store, chatID uuid.UUID, options StartWorkspaceO
), nil
}
build, err := db.GetLatestWorkspaceBuildByWorkspaceID(ctx, ws.ID)
build, job, err := latestWorkspaceBuildAndJob(ctx, db, ws.ID)
if err != nil {
return fantasy.NewTextErrorResponse(
xerrors.Errorf("get latest build: %w", err).Error(),
), nil
}
job, err := db.GetProvisionerJobByID(ctx, build.JobID)
if err != nil {
return fantasy.NewTextErrorResponse(
xerrors.Errorf("get provisioner job: %w", err).Error(),
), nil
return fantasy.NewTextErrorResponse(err.Error()), nil
}
// If a build is already in progress, wait for it.
@@ -106,24 +97,7 @@ func StartWorkspace(db database.Store, chatID uuid.UUID, options StartWorkspaceO
database.ProvisionerJobStatusRunning:
// Publish the build ID to the frontend so it
// can start streaming logs immediately.
updatedChat, bindErr := db.UpdateChatWorkspaceBinding(ctx, database.UpdateChatWorkspaceBindingParams{
ID: chatID,
WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true},
BuildID: uuid.NullUUID{
UUID: build.ID,
Valid: build.ID != uuid.Nil,
},
AgentID: uuid.NullUUID{},
})
if bindErr != nil {
options.Logger.Error(ctx, "failed to persist build ID on chat binding",
slog.F("chat_id", chatID),
slog.F("build_id", build.ID),
slog.Error(bindErr),
)
} else if options.OnChatUpdated != nil {
options.OnChatUpdated(updatedChat)
}
publishBuildBinding(ctx, db, options.Logger, chatID, ws.ID, build.ID, options.OnChatUpdated)
if err := waitForBuild(ctx, db, build.ID); err != nil {
// newBuildError returns via toolResponse (IsError: false)
// rather than NewTextErrorResponse (IsError: true) so the
@@ -199,24 +173,7 @@ func StartWorkspace(db database.Store, chatID uuid.UUID, options StartWorkspaceO
// Persist the build ID on the chat binding so the
// frontend can stream logs without polling.
updatedChat, bindErr := db.UpdateChatWorkspaceBinding(ctx, database.UpdateChatWorkspaceBindingParams{
ID: chatID,
WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true},
BuildID: uuid.NullUUID{
UUID: startBuild.ID,
Valid: startBuild.ID != uuid.Nil,
},
AgentID: uuid.NullUUID{},
})
if bindErr != nil {
options.Logger.Error(ctx, "failed to persist build ID on chat binding",
slog.F("chat_id", chatID),
slog.F("build_id", startBuild.ID),
slog.Error(bindErr),
)
} else if options.OnChatUpdated != nil {
options.OnChatUpdated(updatedChat)
}
publishBuildBinding(ctx, db, options.Logger, chatID, ws.ID, startBuild.ID, options.OnChatUpdated)
if err := waitForBuild(ctx, db, startBuild.ID); err != nil {
return buildFailureToolResponse(
ctx,
+181
View File
@@ -0,0 +1,181 @@
package chattool
import (
"context"
"sync"
"charm.land/fantasy"
"github.com/google/uuid"
"golang.org/x/xerrors"
"cdr.dev/slog/v3"
"github.com/coder/coder/v2/coderd/database"
"github.com/coder/coder/v2/coderd/httpapi/httperror"
"github.com/coder/coder/v2/codersdk"
)
// StopWorkspaceFn stops a workspace by creating a new build with
// the "stop" transition.
type StopWorkspaceFn func(
ctx context.Context,
ownerID uuid.UUID,
workspaceID uuid.UUID,
req codersdk.CreateWorkspaceBuildRequest,
) (codersdk.WorkspaceBuild, error)
// StopWorkspaceOptions configures the stop_workspace tool.
type StopWorkspaceOptions struct {
OwnerID uuid.UUID
StopFn StopWorkspaceFn
WorkspaceMu *sync.Mutex
OnChatUpdated func(database.Chat)
Logger slog.Logger
}
type stopWorkspaceArgs struct{}
// StopWorkspace returns a tool that stops the workspace associated
// with the current chat. The tool is idempotent when the workspace is
// already stopped. db must not be nil and chatID must not be uuid.Nil.
func StopWorkspace(db database.Store, chatID uuid.UUID, options StopWorkspaceOptions) fantasy.AgentTool {
return fantasy.NewAgentTool(
"stop_workspace",
"Stop the chat's workspace and wait for the stop build to complete. "+
"If another workspace build is already in progress, this waits "+
"for that build first, then stops the workspace if needed. "+
"After waiting, this tool is idempotent if the workspace is "+
"already stopped or the in-progress build stopped it. Use "+
"this when the "+
"user explicitly asks to stop the workspace, or when a "+
"workspace-agent error tells you to stop and then start the "+
"workspace. Stopping a workspace terminates running processes "+
"and may discard unsaved in-memory state. This tool does not "+
"delete the workspace.",
func(ctx context.Context, _ stopWorkspaceArgs, _ fantasy.ToolCall) (fantasy.ToolResponse, error) {
if options.StopFn == nil {
return fantasy.NewTextErrorResponse("workspace stopper is not configured"), nil
}
// Serialize with create_workspace and start_workspace to
// prevent lifecycle races.
if options.WorkspaceMu != nil {
options.WorkspaceMu.Lock()
defer options.WorkspaceMu.Unlock()
}
chat, err := db.GetChatByID(ctx, chatID)
if err != nil {
return fantasy.NewTextErrorResponse(
xerrors.Errorf("load chat: %w", err).Error(),
), nil
}
if !chat.WorkspaceID.Valid {
return fantasy.NewTextErrorResponse(
"chat has no workspace; use create_workspace first",
), nil
}
ws, err := db.GetWorkspaceByID(ctx, chat.WorkspaceID.UUID)
if err != nil {
return fantasy.NewTextErrorResponse(
xerrors.Errorf("load workspace: %w", err).Error(),
), nil
}
if ws.Deleted {
return fantasy.NewTextErrorResponse(
"workspace was deleted; use create_workspace to make a new one",
), nil
}
build, job, err := latestWorkspaceBuildAndJob(ctx, db, ws.ID)
if err != nil {
return fantasy.NewTextErrorResponse(err.Error()), nil
}
// If a build is already in progress, wait for it before
// deciding whether a stop build is still needed.
switch job.JobStatus {
case database.ProvisionerJobStatusPending,
database.ProvisionerJobStatusRunning,
database.ProvisionerJobStatusCanceling:
publishBuildBinding(ctx, db, options.Logger, chatID, ws.ID, build.ID, options.OnChatUpdated)
waitErr := waitForBuild(ctx, db, build.ID)
// Re-read after waiting because another transition may
// have completed while this tool was blocked.
ws, err = db.GetWorkspaceByID(ctx, ws.ID)
if err != nil {
return fantasy.NewTextErrorResponse(
xerrors.Errorf("load workspace: %w", err).Error(),
), nil
}
if ws.Deleted {
return fantasy.NewTextErrorResponse(
"workspace was deleted; use create_workspace to make a new one",
), nil
}
build, job, err = latestWorkspaceBuildAndJob(ctx, db, ws.ID)
if err != nil {
return fantasy.NewTextErrorResponse(err.Error()), nil
}
// The fresh job row is authoritative. A wait error can
// be stale if the build reached a terminal state while the
// wait context was ending.
if waitErr != nil && !provisionerJobTerminal(job.JobStatus) {
return buildToolResponse(newBuildError(
xerrors.Errorf("waiting for in-progress build: %w", waitErr).Error(),
build.ID,
)), nil
}
}
if job.JobStatus == database.ProvisionerJobStatusSucceeded &&
build.Transition == database.WorkspaceTransitionStop {
result := map[string]any{
"stopped": true,
"workspace_name": ws.Name,
}
setNoBuild(result, uuid.Nil)
return toolResponse(result), nil
}
ownerCtx, ownerErr := asOwner(ctx, db, options.OwnerID)
if ownerErr != nil {
return fantasy.NewTextErrorResponse(ownerErr.Error()), nil
}
stopBuild, err := options.StopFn(ownerCtx, options.OwnerID, ws.ID, codersdk.CreateWorkspaceBuildRequest{
Transition: codersdk.WorkspaceTransitionStop,
})
if err != nil {
if responseErr, ok := httperror.IsResponder(err); ok {
_, resp := responseErr.Response()
return toolResponse(responseErrorResult(resp)), nil
}
return fantasy.NewTextErrorResponse(
xerrors.Errorf("stop workspace: %w", err).Error(),
), nil
}
publishBuildBinding(ctx, db, options.Logger, chatID, ws.ID, stopBuild.ID, options.OnChatUpdated)
if err := waitForBuild(ctx, db, stopBuild.ID); err != nil {
return buildToolResponse(newBuildError(
xerrors.Errorf("workspace stop build failed: %w", err).Error(),
stopBuild.ID,
)), nil
}
if options.OnChatUpdated != nil {
if latest, err := db.GetChatByID(ctx, chatID); err == nil {
options.OnChatUpdated(latest)
}
}
result := map[string]any{
"stopped": true,
"workspace_name": ws.Name,
}
setBuildID(result, stopBuild.ID)
return toolResponse(result), nil
})
}
@@ -0,0 +1,449 @@
package chattool_test
import (
"context"
"database/sql"
"encoding/json"
"sync"
"sync/atomic"
"testing"
"time"
"charm.land/fantasy"
"github.com/google/uuid"
"github.com/stretchr/testify/require"
"cdr.dev/slog/v3/sloggers/slogtest"
"github.com/coder/coder/v2/coderd/database"
"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/x/chatd/chattool"
"github.com/coder/coder/v2/codersdk"
"github.com/coder/coder/v2/testutil"
)
func TestStopWorkspace(t *testing.T) {
t.Parallel()
t.Run("NoWorkspace", func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitLong)
db, _ := dbtestutil.NewDB(t)
user := dbgen.User(t, db, database.User{})
modelCfg := seedModelConfig(t, db)
org := dbgen.Organization(t, db, database.Organization{})
_ = dbgen.OrganizationMember(t, db, database.OrganizationMember{
UserID: user.ID,
OrganizationID: org.ID,
})
chat := dbgen.Chat(t, db, database.Chat{
OrganizationID: org.ID,
OwnerID: user.ID,
LastModelConfigID: modelCfg.ID,
Title: "test-stop-no-workspace",
})
tool := chattool.StopWorkspace(db, chat.ID, chattool.StopWorkspaceOptions{
StopFn: func(_ context.Context, _ uuid.UUID, _ uuid.UUID, _ codersdk.CreateWorkspaceBuildRequest) (codersdk.WorkspaceBuild, error) {
t.Fatal("StopFn should not be called")
return codersdk.WorkspaceBuild{}, nil
},
WorkspaceMu: &sync.Mutex{},
})
resp, err := tool.Run(ctx, fantasy.ToolCall{ID: "call-1", Name: "stop_workspace", Input: "{}"})
require.NoError(t, err)
require.Contains(t, resp.Content, "use create_workspace first")
})
t.Run("DeletedWorkspace", func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitLong)
db, _ := dbtestutil.NewDB(t)
user := dbgen.User(t, db, database.User{})
modelCfg := seedModelConfig(t, db)
org := dbgen.Organization(t, db, database.Organization{})
_ = dbgen.OrganizationMember(t, db, database.OrganizationMember{
UserID: user.ID,
OrganizationID: org.ID,
})
wsResp := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{
OwnerID: user.ID,
OrganizationID: org.ID,
Deleted: true,
}).Seed(database.WorkspaceBuild{
Transition: database.WorkspaceTransitionDelete,
}).Do()
ws := wsResp.Workspace
chat := dbgen.Chat(t, db, database.Chat{
OrganizationID: org.ID,
OwnerID: user.ID,
WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true},
LastModelConfigID: modelCfg.ID,
Title: "test-stop-deleted-workspace",
})
tool := chattool.StopWorkspace(db, chat.ID, chattool.StopWorkspaceOptions{
StopFn: func(_ context.Context, _ uuid.UUID, _ uuid.UUID, _ codersdk.CreateWorkspaceBuildRequest) (codersdk.WorkspaceBuild, error) {
t.Fatal("StopFn should not be called for deleted workspace")
return codersdk.WorkspaceBuild{}, nil
},
WorkspaceMu: &sync.Mutex{},
})
resp, err := tool.Run(ctx, fantasy.ToolCall{ID: "call-1", Name: "stop_workspace", Input: "{}"})
require.NoError(t, err)
require.Contains(t, resp.Content, "workspace was deleted")
require.Contains(t, resp.Content, "create_workspace")
})
t.Run("AlreadyStopped", func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitLong)
db, _ := dbtestutil.NewDB(t)
user := dbgen.User(t, db, database.User{})
modelCfg := seedModelConfig(t, db)
org := dbgen.Organization(t, db, database.Organization{})
_ = dbgen.OrganizationMember(t, db, database.OrganizationMember{
UserID: user.ID,
OrganizationID: org.ID,
})
wsResp := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{
OwnerID: user.ID,
OrganizationID: org.ID,
}).Seed(database.WorkspaceBuild{
Transition: database.WorkspaceTransitionStop,
}).Do()
ws := wsResp.Workspace
chat := dbgen.Chat(t, db, database.Chat{
OrganizationID: org.ID,
OwnerID: user.ID,
WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true},
LastModelConfigID: modelCfg.ID,
Title: "test-stop-already-stopped",
})
tool := chattool.StopWorkspace(db, chat.ID, chattool.StopWorkspaceOptions{
OwnerID: user.ID,
StopFn: func(_ context.Context, _ uuid.UUID, _ uuid.UUID, _ codersdk.CreateWorkspaceBuildRequest) (codersdk.WorkspaceBuild, error) {
t.Fatal("StopFn should not be called for already-stopped workspace")
return codersdk.WorkspaceBuild{}, nil
},
WorkspaceMu: &sync.Mutex{},
})
resp, err := tool.Run(ctx, fantasy.ToolCall{ID: "call-1", Name: "stop_workspace", Input: "{}"})
require.NoError(t, err)
var result map[string]any
require.NoError(t, json.Unmarshal([]byte(resp.Content), &result))
require.Equal(t, true, result["stopped"])
require.Equal(t, ws.Name, result["workspace_name"])
require.Equal(t, true, result["no_build"])
require.Nil(t, result["build_id"])
})
t.Run("RunningWorkspaceStops", func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitLong)
db, _ := dbtestutil.NewDB(t)
user := dbgen.User(t, db, database.User{})
modelCfg := seedModelConfig(t, db)
org := dbgen.Organization(t, db, database.Organization{})
_ = dbgen.OrganizationMember(t, db, database.OrganizationMember{
UserID: user.ID,
OrganizationID: org.ID,
})
wsResp := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{
OwnerID: user.ID,
OrganizationID: org.ID,
}).Seed(database.WorkspaceBuild{
Transition: database.WorkspaceTransitionStart,
}).Do()
ws := wsResp.Workspace
chat := dbgen.Chat(t, db, database.Chat{
OrganizationID: org.ID,
OwnerID: user.ID,
WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true},
LastModelConfigID: modelCfg.ID,
Title: "test-stop-running-workspace",
})
var stopCalled atomic.Bool
var stopBuildID uuid.UUID
var seenBuildID uuid.UUID
var onChatUpdatedCalls atomic.Int32
tool := chattool.StopWorkspace(db, chat.ID, chattool.StopWorkspaceOptions{
OwnerID: user.ID,
StopFn: func(_ context.Context, _ uuid.UUID, wsID uuid.UUID, req codersdk.CreateWorkspaceBuildRequest) (codersdk.WorkspaceBuild, error) {
stopCalled.Store(true)
require.Equal(t, ws.ID, wsID)
require.Equal(t, codersdk.WorkspaceTransitionStop, req.Transition)
buildResp := dbfake.WorkspaceBuild(t, db, ws).Seed(database.WorkspaceBuild{
Transition: database.WorkspaceTransitionStop,
BuildNumber: 2,
}).Do()
stopBuildID = buildResp.Build.ID
return codersdk.WorkspaceBuild{ID: buildResp.Build.ID}, nil
},
WorkspaceMu: &sync.Mutex{},
OnChatUpdated: func(chat database.Chat) {
onChatUpdatedCalls.Add(1)
if chat.BuildID.Valid {
seenBuildID = chat.BuildID.UUID
}
},
})
resp, err := tool.Run(ctx, fantasy.ToolCall{ID: "call-1", Name: "stop_workspace", Input: "{}"})
require.NoError(t, err)
require.True(t, stopCalled.Load())
var result map[string]any
require.NoError(t, json.Unmarshal([]byte(resp.Content), &result))
require.Equal(t, true, result["stopped"])
require.Equal(t, ws.Name, result["workspace_name"])
require.Equal(t, stopBuildID.String(), result["build_id"])
require.Nil(t, result["no_build"])
require.GreaterOrEqual(t, onChatUpdatedCalls.Load(), int32(1))
require.Equal(t, stopBuildID, seenBuildID)
updatedChat, err := db.GetChatByID(ctx, chat.ID)
require.NoError(t, err)
require.True(t, updatedChat.BuildID.Valid)
require.Equal(t, stopBuildID, updatedChat.BuildID.UUID)
})
t.Run("InProgressBuildWaitsThenStops", func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitLong)
db, _ := dbtestutil.NewDB(t)
user := dbgen.User(t, db, database.User{})
modelCfg := seedModelConfig(t, db)
org := dbgen.Organization(t, db, database.Organization{})
_ = dbgen.OrganizationMember(t, db, database.OrganizationMember{
UserID: user.ID,
OrganizationID: org.ID,
})
wsResp := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{
OwnerID: user.ID,
OrganizationID: org.ID,
}).Seed(database.WorkspaceBuild{
Transition: database.WorkspaceTransitionStart,
}).Starting().Do()
ws := wsResp.Workspace
chat := dbgen.Chat(t, db, database.Chat{
OrganizationID: org.ID,
OwnerID: user.ID,
WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true},
LastModelConfigID: modelCfg.ID,
Title: "test-stop-in-progress-build",
})
jobRead := make(chan struct{}, 1)
wrappedDB := &jobInterceptStore{Store: db, jobRead: jobRead}
var stopCalled atomic.Bool
var stopBuildID uuid.UUID
var onChatUpdatedCalled atomic.Bool
tool := chattool.StopWorkspace(wrappedDB, chat.ID, chattool.StopWorkspaceOptions{
OwnerID: user.ID,
StopFn: func(_ context.Context, _ uuid.UUID, wsID uuid.UUID, req codersdk.CreateWorkspaceBuildRequest) (codersdk.WorkspaceBuild, error) {
stopCalled.Store(true)
require.Equal(t, ws.ID, wsID)
require.Equal(t, codersdk.WorkspaceTransitionStop, req.Transition)
buildResp := dbfake.WorkspaceBuild(t, db, ws).Seed(database.WorkspaceBuild{
Transition: database.WorkspaceTransitionStop,
BuildNumber: 2,
}).Do()
stopBuildID = buildResp.Build.ID
return codersdk.WorkspaceBuild{ID: buildResp.Build.ID}, nil
},
WorkspaceMu: &sync.Mutex{},
Logger: slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}),
OnChatUpdated: func(_ database.Chat) { onChatUpdatedCalled.Store(true) },
})
type toolResult struct {
resp fantasy.ToolResponse
err error
}
done := make(chan toolResult, 1)
go func() {
resp, err := tool.Run(ctx, fantasy.ToolCall{ID: "call-1", Name: "stop_workspace", Input: "{}"})
done <- toolResult{resp: resp, err: err}
}()
testutil.TryReceive(ctx, t, jobRead)
require.False(t, stopCalled.Load(), "StopFn must wait for the in-progress build")
now := time.Now().UTC()
require.NoError(t, db.UpdateProvisionerJobWithCompleteByID(ctx, database.UpdateProvisionerJobWithCompleteByIDParams{
ID: wsResp.Build.JobID,
UpdatedAt: now,
CompletedAt: sql.NullTime{Time: now, Valid: true},
}))
res := testutil.TryReceive(ctx, t, done)
require.NoError(t, res.err)
require.True(t, stopCalled.Load())
require.True(t, onChatUpdatedCalled.Load())
var result map[string]any
require.NoError(t, json.Unmarshal([]byte(res.resp.Content), &result))
require.Equal(t, true, result["stopped"])
require.Equal(t, stopBuildID.String(), result["build_id"])
})
t.Run("FailedLatestStopBuildStillStops", func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitLong)
db, _ := dbtestutil.NewDB(t)
user := dbgen.User(t, db, database.User{})
modelCfg := seedModelConfig(t, db)
org := dbgen.Organization(t, db, database.Organization{})
_ = dbgen.OrganizationMember(t, db, database.OrganizationMember{
UserID: user.ID,
OrganizationID: org.ID,
})
wsResp := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{
OwnerID: user.ID,
OrganizationID: org.ID,
}).Seed(database.WorkspaceBuild{
Transition: database.WorkspaceTransitionStop,
}).Do()
ws := wsResp.Workspace
now := time.Now().UTC()
require.NoError(t, db.UpdateProvisionerJobWithCompleteByID(ctx, database.UpdateProvisionerJobWithCompleteByIDParams{
ID: wsResp.Build.JobID,
UpdatedAt: now,
CompletedAt: sql.NullTime{Time: now, Valid: true},
Error: sql.NullString{String: "latest build failed", Valid: true},
}))
chat := dbgen.Chat(t, db, database.Chat{
OrganizationID: org.ID,
OwnerID: user.ID,
WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true},
LastModelConfigID: modelCfg.ID,
Title: "test-stop-failed-latest-build",
})
var stopCalled atomic.Bool
tool := chattool.StopWorkspace(db, chat.ID, chattool.StopWorkspaceOptions{
OwnerID: user.ID,
StopFn: func(_ context.Context, _ uuid.UUID, _ uuid.UUID, req codersdk.CreateWorkspaceBuildRequest) (codersdk.WorkspaceBuild, error) {
stopCalled.Store(true)
require.Equal(t, codersdk.WorkspaceTransitionStop, req.Transition)
buildResp := dbfake.WorkspaceBuild(t, db, ws).Seed(database.WorkspaceBuild{
Transition: database.WorkspaceTransitionStop,
BuildNumber: 2,
}).Do()
return codersdk.WorkspaceBuild{ID: buildResp.Build.ID}, nil
},
WorkspaceMu: &sync.Mutex{},
})
resp, err := tool.Run(ctx, fantasy.ToolCall{ID: "call-1", Name: "stop_workspace", Input: "{}"})
require.NoError(t, err)
require.True(t, stopCalled.Load())
var result map[string]any
require.NoError(t, json.Unmarshal([]byte(resp.Content), &result))
require.Equal(t, true, result["stopped"])
})
t.Run("StopTriggeredBuildFailure", func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitLong)
db, _ := dbtestutil.NewDB(t)
user := dbgen.User(t, db, database.User{})
modelCfg := seedModelConfig(t, db)
org := dbgen.Organization(t, db, database.Organization{})
_ = dbgen.OrganizationMember(t, db, database.OrganizationMember{
UserID: user.ID,
OrganizationID: org.ID,
})
wsResp := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{
OwnerID: user.ID,
OrganizationID: org.ID,
}).Seed(database.WorkspaceBuild{
Transition: database.WorkspaceTransitionStart,
}).Do()
ws := wsResp.Workspace
chat := dbgen.Chat(t, db, database.Chat{
OrganizationID: org.ID,
OwnerID: user.ID,
WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true},
LastModelConfigID: modelCfg.ID,
Title: "test-stop-triggered-build-failure",
})
var stopBuildJobID uuid.UUID
var stopBuildID uuid.UUID
stopFn := func(_ context.Context, _ uuid.UUID, wsID uuid.UUID, req codersdk.CreateWorkspaceBuildRequest) (codersdk.WorkspaceBuild, error) {
require.Equal(t, ws.ID, wsID)
require.Equal(t, codersdk.WorkspaceTransitionStop, req.Transition)
buildResp := dbfake.WorkspaceBuild(t, db, ws).Seed(database.WorkspaceBuild{
Transition: database.WorkspaceTransitionStop,
BuildNumber: 2,
}).Starting().Do()
stopBuildJobID = buildResp.Build.JobID
stopBuildID = buildResp.Build.ID
return codersdk.WorkspaceBuild{ID: buildResp.Build.ID}, nil
}
jobRead := make(chan struct{}, 2)
wrappedDB := &jobInterceptStore{Store: db, jobRead: jobRead}
tool := chattool.StopWorkspace(wrappedDB, chat.ID, chattool.StopWorkspaceOptions{
OwnerID: user.ID,
StopFn: stopFn,
WorkspaceMu: &sync.Mutex{},
Logger: slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}),
})
type toolResult struct {
resp fantasy.ToolResponse
err error
}
done := make(chan toolResult, 1)
go func() {
resp, err := tool.Run(ctx, fantasy.ToolCall{ID: "call-1", Name: "stop_workspace", Input: "{}"})
done <- toolResult{resp: resp, err: err}
}()
testutil.TryReceive(ctx, t, jobRead)
testutil.TryReceive(ctx, t, jobRead)
now := time.Now().UTC()
require.NoError(t, db.UpdateProvisionerJobWithCompleteByID(ctx, database.UpdateProvisionerJobWithCompleteByIDParams{
ID: stopBuildJobID,
UpdatedAt: now,
CompletedAt: sql.NullTime{Time: now, Valid: true},
Error: sql.NullString{String: "terraform destroy failed", Valid: true},
}))
res := testutil.TryReceive(ctx, t, done)
require.NoError(t, res.err)
var result map[string]any
require.NoError(t, json.Unmarshal([]byte(res.resp.Content), &result))
require.Contains(t, result["error"], "workspace stop build failed")
require.Equal(t, stopBuildID.String(), result["build_id"])
require.False(t, res.resp.IsError,
"buildToolResponse must not set IsError; chatprompt strips structured fields from error responses")
})
}