shishantbiswas/coder
1
0
Fork 0
mirror of https://github.com/coder/coder.git synced 2026-06-06 06:28:20 +00:00
Code Issues Packages Projects Releases Wiki Activity
Files
596e55b13630b43decdfd40df9dbcfc6507dabcb
coder/coderd/x/chatd/chattool/createworkspace_test.go
T
Kyle Carberry d11849d94a fix: re-fetch context files and skills from workspace on each turn (#24360) 
Context files (AGENTS.md) and skills were only fetched from the
workspace on the first turn or when the agent changed. On subsequent
turns, stale content from persisted messages was used. This meant that
if AGENTS.md or skills were modified on the workspace between turns, the
agent wouldn't see the changes until the user created a new chat.

## Changes

- Extract `fetchWorkspaceContext` from `persistInstructionFiles` to
allow fetching workspace context without persisting
- On subsequent turns, re-fetch fresh context from the workspace instead
of reading stale persisted content; falls back to persisted messages if
the workspace dial fails
- Update `ReloadMessages` callback to re-derive instruction and skills
from reloaded database messages after compaction, instead of using
captured closure variables
- Add `formatSystemInstructionsFromParts` helper to build system
instructions directly from agent parts without requiring separate
OS/directory params
- Add tests for the new helper

<details><summary>Implementation Notes</summary>

### Root cause

In `runChat`, the `else if hasContextFiles` branch (subsequent turns)
called `instructionFromContextFiles(messages)` which read stale content
from persisted DB messages. The `ReloadMessages` callback
(post-compaction) also used captured `instruction`/`skills` closure
variables from the start of the turn, never re-deriving them.

### Approach

1. **Extract `fetchWorkspaceContext`** — Pure refactor of the fetch-only
part of `persistInstructionFiles` (agent connection, context config
retrieval, content sanitization, metadata stamping). Returns parts +
skills without persisting.

2. **Subsequent turns**: Instead of reading from persisted messages,
launch a `g2` goroutine that calls `fetchWorkspaceContext` to get fresh
context from the workspace. Falls back gracefully to persisted messages
if the workspace is unreachable.

3. **ReloadMessages**: Re-derive `instruction` from
`instructionFromContextFiles(reloadedMsgs)` and `skills` from
`skillsFromParts(reloadedMsgs)` using the freshly loaded messages, with
fallback to captured values if the reloaded messages don't contain
context (e.g. compacted away).

</details>

> 🤖 Generated by Coder Agents
2026-04-15 16:41:15 -04:00

1265 lines
35 KiB
Go
Raw Blame History

package chattool //nolint:testpackage // Uses internal symbols.
import (
"context"
"database/sql"
"encoding/json"
"fmt"
"sync"
"testing"
"time"
"charm.land/fantasy"
"github.com/google/uuid"
"github.com/stretchr/testify/require"
"go.uber.org/mock/gomock"
"golang.org/x/xerrors"
"cdr.dev/slog/v3/sloggers/slogtest"
"github.com/coder/coder/v2/coderd/database"
"github.com/coder/coder/v2/coderd/database/dbmock"
"github.com/coder/coder/v2/coderd/util/ptr"
"github.com/coder/coder/v2/codersdk"
"github.com/coder/coder/v2/codersdk/workspacesdk"
)
func TestWaitForAgentReady(t *testing.T) {
t.Parallel()
t.Run("AgentConnectsAndLifecycleReady", func(t *testing.T) {
t.Parallel()
ctrl := gomock.NewController(t)
db := dbmock.NewMockStore(ctrl)
agentID := uuid.New()
// Mock returns Ready lifecycle state.
db.EXPECT().
GetWorkspaceAgentLifecycleStateByID(gomock.Any(), agentID).
Return(database.GetWorkspaceAgentLifecycleStateByIDRow{
LifecycleState: database.WorkspaceAgentLifecycleStateReady,
}, nil)
// AgentConnFn succeeds immediately.
connFn := func(ctx context.Context, id uuid.UUID) (workspacesdk.AgentConn, func(), error) {
return nil, func() {}, nil
}
result := waitForAgentReady(context.Background(), db, agentID, connFn)
require.Empty(t, result)
})
t.Run("AgentConnectTimeout", func(t *testing.T) {
t.Parallel()
ctrl := gomock.NewController(t)
db := dbmock.NewMockStore(ctrl)
agentID := uuid.New()
// AgentConnFn always fails - context will timeout.
connFn := func(ctx context.Context, id uuid.UUID) (workspacesdk.AgentConn, func(), error) {
return nil, nil, context.DeadlineExceeded
}
// Use a context that's already canceled to avoid waiting.
ctx, cancel := context.WithCancel(context.Background())
cancel()
result := waitForAgentReady(ctx, db, agentID, connFn)
require.Equal(t, "not_ready", result["agent_status"])
require.NotEmpty(t, result["agent_error"])
})
t.Run("AgentConnectsButStartupFails", func(t *testing.T) {
t.Parallel()
ctrl := gomock.NewController(t)
db := dbmock.NewMockStore(ctrl)
agentID := uuid.New()
// Mock returns StartError lifecycle state.
db.EXPECT().
GetWorkspaceAgentLifecycleStateByID(gomock.Any(), agentID).
Return(database.GetWorkspaceAgentLifecycleStateByIDRow{
LifecycleState: database.WorkspaceAgentLifecycleStateStartError,
}, nil)
connFn := func(ctx context.Context, id uuid.UUID) (workspacesdk.AgentConn, func(), error) {
return nil, func() {}, nil
}
result := waitForAgentReady(context.Background(), db, agentID, connFn)
require.Equal(t, "startup_scripts_failed", result["startup_scripts"])
require.Equal(t, "start_error", result["lifecycle_state"])
})
t.Run("NilAgentConnFn", func(t *testing.T) {
t.Parallel()
ctrl := gomock.NewController(t)
db := dbmock.NewMockStore(ctrl)
agentID := uuid.New()
// Mock returns Ready lifecycle state.
db.EXPECT().
GetWorkspaceAgentLifecycleStateByID(gomock.Any(), agentID).
Return(database.GetWorkspaceAgentLifecycleStateByIDRow{
LifecycleState: database.WorkspaceAgentLifecycleStateReady,
}, nil)
result := waitForAgentReady(context.Background(), db, agentID, nil)
require.Empty(t, result)
})
t.Run("NilDB", func(t *testing.T) {
t.Parallel()
connFn := func(ctx context.Context, id uuid.UUID) (workspacesdk.AgentConn, func(), error) {
return nil, func() {}, nil
}
result := waitForAgentReady(context.Background(), nil, uuid.New(), connFn)
require.Empty(t, result)
})
}
func TestCreateWorkspace_PrefersChatSuffixAgent(t *testing.T) {
t.Parallel()
ctrl := gomock.NewController(t)
db := dbmock.NewMockStore(ctrl)
ownerID := uuid.New()
orgID := uuid.New()
templateID := uuid.New()
workspaceID := uuid.New()
jobID := uuid.New()
buildID := uuid.New()
fallbackAgentID := uuid.New()
chatAgentID := uuid.New()
db.EXPECT().
GetAuthorizationUserRoles(gomock.Any(), ownerID).
Return(database.GetAuthorizationUserRolesRow{
ID: ownerID,
Roles: []string{},
Groups: []string{},
Status: database.UserStatusActive,
}, nil)
db.EXPECT().
GetTemplateByID(gomock.Any(), templateID).
Return(database.Template{
ID: templateID,
OrganizationID: orgID,
}, nil)
db.EXPECT().
GetChatWorkspaceTTL(gomock.Any()).
Return("0s", nil)
db.EXPECT().
GetWorkspaceBuildByID(gomock.Any(), buildID).
Return(database.WorkspaceBuild{
ID: buildID,
WorkspaceID: workspaceID,
JobID: jobID,
}, nil)
db.EXPECT().
GetProvisionerJobByID(gomock.Any(), jobID).
Return(database.ProvisionerJob{
ID: jobID,
JobStatus: database.ProvisionerJobStatusSucceeded,
}, nil)
db.EXPECT().
GetWorkspaceAgentsInLatestBuildByWorkspaceID(gomock.Any(), workspaceID).
Return([]database.WorkspaceAgent{
{ID: fallbackAgentID, Name: "dev", DisplayOrder: 0},
{ID: chatAgentID, Name: "dev-coderd-chat", DisplayOrder: 1},
}, nil)
db.EXPECT().
GetWorkspaceAgentLifecycleStateByID(gomock.Any(), chatAgentID).
Return(database.GetWorkspaceAgentLifecycleStateByIDRow{
LifecycleState: database.WorkspaceAgentLifecycleStateReady,
}, nil)
var connectedAgentID uuid.UUID
createFn := func(_ context.Context, _ uuid.UUID, req codersdk.CreateWorkspaceRequest) (codersdk.Workspace, error) {
return codersdk.Workspace{
ID: workspaceID,
Name: req.Name,
OwnerName: "testuser",
LatestBuild: codersdk.WorkspaceBuild{
ID: buildID,
},
}, nil
}
agentConnFn := func(_ context.Context, agentID uuid.UUID) (workspacesdk.AgentConn, func(), error) {
connectedAgentID = agentID
return nil, func() {}, nil
}
tool := CreateWorkspace(orgID, db, CreateWorkspaceOptions{
OwnerID: ownerID,
CreateFn: createFn,
AgentConnFn: agentConnFn,
WorkspaceMu: &sync.Mutex{},
Logger: slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}),
})
input := fmt.Sprintf(`{"template_id":%q,"name":"test-chat-agent"}`, templateID.String())
resp, err := tool.Run(context.Background(), fantasy.ToolCall{
ID: "call-1",
Name: "create_workspace",
Input: input,
})
require.NoError(t, err)
require.NotEmpty(t, resp.Content)
require.Equal(t, chatAgentID, connectedAgentID)
var result map[string]any
require.NoError(t, json.Unmarshal([]byte(resp.Content), &result))
require.Equal(t, buildID.String(), result["build_id"])
}
func TestCreateWorkspace_ReturnsSelectionErrorImmediately(t *testing.T) {
t.Parallel()
ctrl := gomock.NewController(t)
db := dbmock.NewMockStore(ctrl)
ownerID := uuid.New()
orgID := uuid.New()
chatID := uuid.New()
templateID := uuid.New()
workspaceID := uuid.New()
jobID := uuid.New()
buildID := uuid.New()
db.EXPECT().
GetChatByID(gomock.Any(), chatID).
Return(database.Chat{ID: chatID}, nil)
db.EXPECT().
GetAuthorizationUserRoles(gomock.Any(), ownerID).
Return(database.GetAuthorizationUserRolesRow{
ID: ownerID,
Roles: []string{},
Groups: []string{},
Status: database.UserStatusActive,
}, nil)
db.EXPECT().
GetTemplateByID(gomock.Any(), templateID).
Return(database.Template{
ID: templateID,
OrganizationID: orgID,
}, nil)
db.EXPECT().
GetChatWorkspaceTTL(gomock.Any()).
Return("0s", nil)
db.EXPECT().
GetWorkspaceBuildByID(gomock.Any(), buildID).
Return(database.WorkspaceBuild{
ID: buildID,
WorkspaceID: workspaceID,
JobID: jobID,
}, nil)
db.EXPECT().
GetProvisionerJobByID(gomock.Any(), jobID).
Return(database.ProvisionerJob{
ID: jobID,
JobStatus: database.ProvisionerJobStatusSucceeded,
}, nil)
db.EXPECT().
UpdateChatWorkspaceBinding(gomock.Any(), database.UpdateChatWorkspaceBindingParams{
ID: chatID,
WorkspaceID: uuid.NullUUID{UUID: workspaceID, Valid: true},
BuildID: uuid.NullUUID{UUID: buildID, Valid: true},
AgentID: uuid.NullUUID{},
}).
Return(database.Chat{
ID: chatID,
WorkspaceID: uuid.NullUUID{UUID: workspaceID, Valid: true},
}, nil)
db.EXPECT().
GetWorkspaceAgentsInLatestBuildByWorkspaceID(gomock.Any(), workspaceID).
Return([]database.WorkspaceAgent{
{ID: uuid.New(), Name: "alpha-coderd-chat", DisplayOrder: 0},
{ID: uuid.New(), Name: "beta-coderd-chat", DisplayOrder: 1},
}, nil)
tool := CreateWorkspace(orgID, db, CreateWorkspaceOptions{
OwnerID: ownerID,
ChatID: chatID,
CreateFn: func(_ context.Context, _ uuid.UUID, req codersdk.CreateWorkspaceRequest) (codersdk.Workspace, error) {
return codersdk.Workspace{
ID: workspaceID,
Name: req.Name,
OwnerName: "testuser",
LatestBuild: codersdk.WorkspaceBuild{
ID: buildID,
},
}, nil
},
AgentConnFn: func(context.Context, uuid.UUID) (workspacesdk.AgentConn, func(), error) {
t.Fatal("AgentConnFn should not be called when agent selection fails")
return nil, nil, xerrors.New("unexpected agent dial")
},
WorkspaceMu: &sync.Mutex{},
Logger: slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}),
})
input := fmt.Sprintf(`{"template_id":%q,"name":"test-selection-error"}`, templateID.String())
resp, err := tool.Run(context.Background(), fantasy.ToolCall{
ID: "call-1",
Name: "create_workspace",
Input: input,
})
require.NoError(t, err)
var result map[string]any
require.NoError(t, json.Unmarshal([]byte(resp.Content), &result))
require.Equal(t, true, result["created"])
require.Equal(t, "testuser/test-selection-error", result["workspace_name"])
require.Equal(t, "selection_error", result["agent_status"])
require.Contains(t, result["agent_error"], "multiple agents match the chat suffix")
require.Equal(t, buildID.String(), result["build_id"])
}
func TestCreateWorkspace_PostCreationBuildFailure(t *testing.T) {
t.Parallel()
ctrl := gomock.NewController(t)
db := dbmock.NewMockStore(ctrl)
ownerID := uuid.New()
orgID := uuid.New()
templateID := uuid.New()
workspaceID := uuid.New()
jobID := uuid.New()
buildID := uuid.New()
db.EXPECT().
GetAuthorizationUserRoles(gomock.Any(), ownerID).
Return(database.GetAuthorizationUserRolesRow{
ID: ownerID,
Roles: []string{},
Groups: []string{},
Status: database.UserStatusActive,
}, nil)
db.EXPECT().
GetTemplateByID(gomock.Any(), templateID).
Return(database.Template{
ID: templateID,
OrganizationID: orgID,
}, nil)
db.EXPECT().
GetChatWorkspaceTTL(gomock.Any()).
Return("0s", nil)
// waitForBuild fetches the build by ID.
db.EXPECT().
GetWorkspaceBuildByID(gomock.Any(), buildID).
Return(database.WorkspaceBuild{
ID: buildID,
WorkspaceID: workspaceID,
JobID: jobID,
}, nil)
// waitForBuild polls the provisioner job. Return Failed.
db.EXPECT().
GetProvisionerJobByID(gomock.Any(), jobID).
Return(database.ProvisionerJob{
ID: jobID,
JobStatus: database.ProvisionerJobStatusFailed,
Error: sql.NullString{String: "terraform apply failed", Valid: true},
}, nil)
createFn := func(_ context.Context, _ uuid.UUID, req codersdk.CreateWorkspaceRequest) (codersdk.Workspace, error) {
return codersdk.Workspace{
ID: workspaceID,
Name: req.Name,
OwnerName: "testuser",
LatestBuild: codersdk.WorkspaceBuild{
ID: buildID,
},
}, nil
}
tool := CreateWorkspace(orgID, db, CreateWorkspaceOptions{
OwnerID: ownerID,
ChatID: uuid.Nil,
CreateFn: createFn,
WorkspaceMu: &sync.Mutex{},
Logger: slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}),
})
input := fmt.Sprintf(`{"template_id":%q,"name":"test-build-fail"}`, templateID.String())
resp, err := tool.Run(context.Background(), fantasy.ToolCall{
ID: "call-1",
Name: "create_workspace",
Input: input,
})
require.NoError(t, err)
var result map[string]any
require.NoError(t, json.Unmarshal([]byte(resp.Content), &result))
require.Contains(t, result["error"], "workspace build failed")
require.Equal(t, buildID.String(), result["build_id"])
require.False(t, resp.IsError,
"buildToolResponse must not set IsError; chatprompt strips structured fields from error responses")
}
func TestCreateWorkspace_GlobalTTL(t *testing.T) {
t.Parallel()
tests := []struct {
name string
ttlReturn string
ttlErr error
wantTTLMs *int64
}{
{
name: "PositiveTTL",
ttlReturn: "2h",
wantTTLMs: ptr.Ref(int64(2 * time.Hour / time.Millisecond)),
},
{
name: "ZeroTTLUsesTemplateDefault",
ttlReturn: "0s",
wantTTLMs: nil,
},
{
name: "DBError_FallsBackToNil",
ttlReturn: "",
ttlErr: xerrors.New("db error"),
wantTTLMs: nil,
},
{
name: "InvalidStoredValue_FallsBackToNil",
ttlReturn: "not-a-duration",
wantTTLMs: nil,
},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
t.Parallel()
ctrl := gomock.NewController(t)
db := dbmock.NewMockStore(ctrl)
ownerID := uuid.New()
orgID := uuid.New()
templateID := uuid.New()
workspaceID := uuid.New()
jobID := uuid.New()
buildID := uuid.New()
db.EXPECT().
GetAuthorizationUserRoles(gomock.Any(), ownerID).
Return(database.GetAuthorizationUserRolesRow{
ID: ownerID,
Roles: []string{},
Groups: []string{},
Status: database.UserStatusActive,
}, nil)
db.EXPECT().
GetTemplateByID(gomock.Any(), templateID).
Return(database.Template{
ID: templateID,
OrganizationID: orgID,
}, nil)
db.EXPECT().
GetChatWorkspaceTTL(gomock.Any()).
Return(tc.ttlReturn, tc.ttlErr)
db.EXPECT().
GetWorkspaceBuildByID(gomock.Any(), buildID).
Return(database.WorkspaceBuild{
ID: buildID,
WorkspaceID: workspaceID,
JobID: jobID,
}, nil)
db.EXPECT().
GetProvisionerJobByID(gomock.Any(), jobID).
Return(database.ProvisionerJob{
ID: jobID,
JobStatus: database.ProvisionerJobStatusSucceeded,
}, nil)
db.EXPECT().
GetWorkspaceAgentsInLatestBuildByWorkspaceID(gomock.Any(), workspaceID).
Return([]database.WorkspaceAgent{}, nil)
var capturedReq codersdk.CreateWorkspaceRequest
createFn := func(_ context.Context, _ uuid.UUID, req codersdk.CreateWorkspaceRequest) (codersdk.Workspace, error) {
capturedReq = req
return codersdk.Workspace{
ID: workspaceID,
Name: req.Name,
OwnerName: "testuser",
LatestBuild: codersdk.WorkspaceBuild{
ID: buildID,
},
}, nil
}
tool := CreateWorkspace(orgID, db, CreateWorkspaceOptions{
OwnerID: ownerID,
ChatID: uuid.Nil,
CreateFn: createFn,
WorkspaceMu: &sync.Mutex{},
Logger: slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}),
})
input := fmt.Sprintf(`{"template_id":%q,"name":"test-ws-%s"}`, templateID.String(), tc.name)
resp, err := tool.Run(context.Background(), fantasy.ToolCall{
ID: "call-1",
Name: "create_workspace",
Input: input,
})
require.NoError(t, err)
require.NotEmpty(t, resp.Content)
var result map[string]any
require.NoError(t, json.Unmarshal([]byte(resp.Content), &result))
require.Equal(t, buildID.String(), result["build_id"])
if tc.wantTTLMs != nil {
require.NotNil(t, capturedReq.TTLMillis)
require.Equal(t, *tc.wantTTLMs, *capturedReq.TTLMillis)
} else {
require.Nil(t, capturedReq.TTLMillis)
}
})
}
}
func TestCreateWorkspace_RejectsCrossOrgTemplate(t *testing.T) {
t.Parallel()
ctrl := gomock.NewController(t)
db := dbmock.NewMockStore(ctrl)
ownerID := uuid.New()
chatOrgID := uuid.New()
templateOrgID := uuid.New() // Different org.
templateID := uuid.New()
chatID := uuid.New()
// Chat exists but has no workspace binding.
db.EXPECT().
GetChatByID(gomock.Any(), chatID).
Return(database.Chat{
ID: chatID,
WorkspaceID: uuid.NullUUID{},
}, nil)
db.EXPECT().
GetAuthorizationUserRoles(gomock.Any(), ownerID).
Return(database.GetAuthorizationUserRolesRow{
ID: ownerID,
Roles: []string{},
Groups: []string{},
Status: database.UserStatusActive,
}, nil)
db.EXPECT().
GetTemplateByID(gomock.Any(), templateID).
Return(database.Template{
ID: templateID,
OrganizationID: templateOrgID,
Name: "wrong-org-template",
}, nil)
createCalled := false
tool := CreateWorkspace(chatOrgID, db, CreateWorkspaceOptions{
OwnerID: ownerID,
ChatID: chatID,
CreateFn: func(context.Context, uuid.UUID, codersdk.CreateWorkspaceRequest) (codersdk.Workspace, error) {
createCalled = true
return codersdk.Workspace{}, nil
},
WorkspaceMu: &sync.Mutex{},
Logger: slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}),
})
input := fmt.Sprintf(`{"template_id":%q}`, templateID.String())
resp, err := tool.Run(context.Background(), fantasy.ToolCall{
ID: "call-1",
Name: "create_workspace",
Input: input,
})
require.NoError(t, err)
require.False(t, createCalled, "CreateFn must not be called for cross-org template")
require.Contains(t, resp.Content, "organization")
}
func TestCheckExistingWorkspace_ConnectedAgent(t *testing.T) {
t.Parallel()
ctrl := gomock.NewController(t)
db := dbmock.NewMockStore(ctrl)
chatID := uuid.New()
workspaceID := uuid.New()
jobID := uuid.New()
agentID := uuid.New()
now := time.Now().UTC()
expectExistingWorkspaceLookup(
db,
chatID,
workspaceID,
jobID,
"existing-workspace",
database.ProvisionerJobStatusSucceeded,
database.WorkspaceTransitionStart,
)
db.EXPECT().
GetWorkspaceAgentsInLatestBuildByWorkspaceID(gomock.Any(), workspaceID).
Return([]database.WorkspaceAgent{{
ID: agentID,
Name: "dev",
CreatedAt: now.Add(-time.Minute),
FirstConnectedAt: validNullTime(now.Add(-45 * time.Second)),
LastConnectedAt: validNullTime(now.Add(-5 * time.Second)),
}}, nil)
db.EXPECT().
GetWorkspaceAgentLifecycleStateByID(gomock.Any(), agentID).
Return(database.GetWorkspaceAgentLifecycleStateByIDRow{
LifecycleState: database.WorkspaceAgentLifecycleStateReady,
}, nil)
connFn := func(context.Context, uuid.UUID) (workspacesdk.AgentConn, func(), error) {
t.Fatalf("unexpected agent dial for connected workspace")
return nil, nil, xerrors.New("unexpected agent dial")
}
options := testCheckExistingWorkspaceOptions(chatID, connFn)
check := options.checkExistingWorkspace(context.Background(), db)
require.NoError(t, check.Err)
require.True(t, check.Done)
require.Equal(t, "already_exists", check.Result["status"])
require.Equal(t, "existing-workspace", check.Result["workspace_name"])
require.Equal(t, "workspace is already running and recently connected", check.Result["message"])
}
func TestCheckExistingWorkspace_InProgressBuildReturnsBuildID(t *testing.T) {
t.Parallel()
ctrl := gomock.NewController(t)
db := dbmock.NewMockStore(ctrl)
chatID := uuid.New()
workspaceID := uuid.New()
jobID := uuid.New()
buildID := uuid.New()
// GetChatByID returns a chat linked to a workspace.
db.EXPECT().
GetChatByID(gomock.Any(), chatID).
Return(database.Chat{
ID: chatID,
WorkspaceID: uuid.NullUUID{UUID: workspaceID, Valid: true},
}, nil)
// GetWorkspaceByID returns a non-deleted workspace.
db.EXPECT().
GetWorkspaceByID(gomock.Any(), workspaceID).
Return(database.Workspace{
ID: workspaceID,
Name: "building-workspace",
}, nil)
// GetLatestWorkspaceBuildByWorkspaceID is called once in
// checkExistingWorkspace. waitForBuild now uses
// GetWorkspaceBuildByID to track the specific build.
db.EXPECT().
GetLatestWorkspaceBuildByWorkspaceID(gomock.Any(), workspaceID).
Return(database.WorkspaceBuild{
ID: buildID,
WorkspaceID: workspaceID,
JobID: jobID,
Transition: database.WorkspaceTransitionStart,
}, nil)
db.EXPECT().
GetWorkspaceBuildByID(gomock.Any(), buildID).
Return(database.WorkspaceBuild{
ID: buildID,
WorkspaceID: workspaceID,
JobID: jobID,
Transition: database.WorkspaceTransitionStart,
}, nil)
// First GetProvisionerJobByID (in checkExistingWorkspace) returns
// Running, triggering waitForBuild. The second call (waitForBuild's
// first poll) returns Succeeded so the loop exits immediately.
firstJob := db.EXPECT().
GetProvisionerJobByID(gomock.Any(), jobID).
Return(database.ProvisionerJob{
ID: jobID,
JobStatus: database.ProvisionerJobStatusRunning,
}, nil)
db.EXPECT().
GetProvisionerJobByID(gomock.Any(), jobID).
Return(database.ProvisionerJob{
ID: jobID,
JobStatus: database.ProvisionerJobStatusSucceeded,
}, nil).
After(firstJob)
// The in-progress path now publishes the build ID before
// waitForBuild.
db.EXPECT().
UpdateChatWorkspaceBinding(gomock.Any(), database.UpdateChatWorkspaceBindingParams{
ID: chatID,
WorkspaceID: uuid.NullUUID{UUID: workspaceID, Valid: true},
BuildID: uuid.NullUUID{UUID: buildID, Valid: true},
AgentID: uuid.NullUUID{},
}).
Return(database.Chat{
ID: chatID,
WorkspaceID: uuid.NullUUID{UUID: workspaceID, Valid: true},
}, nil)
// After waitForBuild completes, checkExistingWorkspace fetches
// agents. Return empty to keep the test focused on build_id.
db.EXPECT().
GetWorkspaceAgentsInLatestBuildByWorkspaceID(gomock.Any(), workspaceID).
Return([]database.WorkspaceAgent{}, nil)
options := testCheckExistingWorkspaceOptions(chatID, nil)
check := options.checkExistingWorkspace(context.Background(), db)
require.NoError(t, check.Err)
require.True(t, check.Done)
require.Equal(t, false, check.Result["created"])
require.Equal(t, "already_exists", check.Result["status"])
require.Equal(t, buildID.String(), check.Result["build_id"])
require.Equal(t, "building-workspace", check.Result["workspace_name"])
require.Equal(t, "workspace build completed", check.Result["message"])
}
func TestCheckExistingWorkspace_InProgressBuildFailureReturnsBuildID(t *testing.T) {
t.Parallel()
ctrl := gomock.NewController(t)
db := dbmock.NewMockStore(ctrl)
chatID := uuid.New()
workspaceID := uuid.New()
jobID := uuid.New()
buildID := uuid.New()
db.EXPECT().
GetChatByID(gomock.Any(), chatID).
Return(database.Chat{
ID: chatID,
WorkspaceID: uuid.NullUUID{UUID: workspaceID, Valid: true},
}, nil)
db.EXPECT().
GetWorkspaceByID(gomock.Any(), workspaceID).
Return(database.Workspace{
ID: workspaceID,
Name: "failing-workspace",
}, nil)
db.EXPECT().
GetLatestWorkspaceBuildByWorkspaceID(gomock.Any(), workspaceID).
Return(database.WorkspaceBuild{
ID: buildID,
WorkspaceID: workspaceID,
JobID: jobID,
Transition: database.WorkspaceTransitionStart,
}, nil)
db.EXPECT().
GetWorkspaceBuildByID(gomock.Any(), buildID).
Return(database.WorkspaceBuild{
ID: buildID,
WorkspaceID: workspaceID,
JobID: jobID,
Transition: database.WorkspaceTransitionStart,
}, nil)
// First call returns Running (triggers waitForBuild), second
// returns Failed so waitForBuild returns an error.
firstJob := db.EXPECT().
GetProvisionerJobByID(gomock.Any(), jobID).
Return(database.ProvisionerJob{
ID: jobID,
JobStatus: database.ProvisionerJobStatusRunning,
}, nil)
db.EXPECT().
GetProvisionerJobByID(gomock.Any(), jobID).
Return(database.ProvisionerJob{
ID: jobID,
JobStatus: database.ProvisionerJobStatusFailed,
}, nil).
After(firstJob)
// The in-progress path publishes the build ID before
// waitForBuild.
db.EXPECT().
UpdateChatWorkspaceBinding(gomock.Any(), database.UpdateChatWorkspaceBindingParams{
ID: chatID,
WorkspaceID: uuid.NullUUID{UUID: workspaceID, Valid: true},
BuildID: uuid.NullUUID{UUID: buildID, Valid: true},
AgentID: uuid.NullUUID{},
}).
Return(database.Chat{
ID: chatID,
WorkspaceID: uuid.NullUUID{UUID: workspaceID, Valid: true},
}, nil)
options := testCheckExistingWorkspaceOptions(chatID, nil)
check := options.checkExistingWorkspace(context.Background(), db)
require.Error(t, check.Err)
require.Contains(t, check.Err.Error(), "existing workspace build failed")
require.Equal(t, buildID, check.FailedBuildID)
}
func TestCheckExistingWorkspace_ConnectingAgentWaits(t *testing.T) {
t.Parallel()
ctrl := gomock.NewController(t)
db := dbmock.NewMockStore(ctrl)
chatID := uuid.New()
workspaceID := uuid.New()
jobID := uuid.New()
agentID := uuid.New()
now := time.Now().UTC()
connectCalls := 0
expectExistingWorkspaceLookup(
db,
chatID,
workspaceID,
jobID,
"existing-workspace",
database.ProvisionerJobStatusSucceeded,
database.WorkspaceTransitionStart,
)
db.EXPECT().
GetWorkspaceAgentsInLatestBuildByWorkspaceID(gomock.Any(), workspaceID).
Return([]database.WorkspaceAgent{{
ID: agentID,
Name: "dev",
CreatedAt: now,
ConnectionTimeoutSeconds: 60,
}}, nil)
db.EXPECT().
GetWorkspaceAgentLifecycleStateByID(gomock.Any(), agentID).
Return(database.GetWorkspaceAgentLifecycleStateByIDRow{
LifecycleState: database.WorkspaceAgentLifecycleStateReady,
}, nil)
connFn := func(context.Context, uuid.UUID) (workspacesdk.AgentConn, func(), error) {
connectCalls++
return nil, func() {}, nil
}
options := testCheckExistingWorkspaceOptions(chatID, connFn)
check := options.checkExistingWorkspace(context.Background(), db)
require.NoError(t, check.Err)
require.True(t, check.Done)
require.Equal(t, 1, connectCalls)
require.Equal(t, "already_exists", check.Result["status"])
require.Equal(t, "workspace exists and the agent is still connecting", check.Result["message"])
}
func TestCheckExistingWorkspace_DeadAgentAllowsCreation(t *testing.T) {
t.Parallel()
tests := []struct {
name string
agent database.WorkspaceAgent
}{
{
name: "Disconnected",
agent: database.WorkspaceAgent{
ID: uuid.New(),
Name: "disconnected",
CreatedAt: time.Now().UTC().Add(-2 * time.Minute),
FirstConnectedAt: validNullTime(time.Now().UTC().Add(-2 * time.Minute)),
LastConnectedAt: validNullTime(time.Now().UTC().Add(-time.Minute)),
},
},
{
name: "TimedOut",
agent: database.WorkspaceAgent{
ID: uuid.New(),
Name: "timed-out",
CreatedAt: time.Now().UTC().Add(-2 * time.Second),
ConnectionTimeoutSeconds: 1,
},
},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
t.Parallel()
ctrl := gomock.NewController(t)
db := dbmock.NewMockStore(ctrl)
chatID := uuid.New()
workspaceID := uuid.New()
jobID := uuid.New()
expectExistingWorkspaceLookup(
db,
chatID,
workspaceID,
jobID,
"existing-workspace",
database.ProvisionerJobStatusSucceeded,
database.WorkspaceTransitionStart,
)
db.EXPECT().
GetWorkspaceAgentsInLatestBuildByWorkspaceID(gomock.Any(), workspaceID).
Return([]database.WorkspaceAgent{tc.agent}, nil)
options := testCheckExistingWorkspaceOptions(chatID, nil)
check := options.checkExistingWorkspace(context.Background(), db)
require.NoError(t, check.Err)
require.False(t, check.Done)
require.Nil(t, check.Result)
})
}
}
func TestWaitForBuild_CanceledJob(t *testing.T) {
t.Parallel()
ctrl := gomock.NewController(t)
db := dbmock.NewMockStore(ctrl)
ownerID := uuid.New()
orgID := uuid.New()
templateID := uuid.New()
workspaceID := uuid.New()
jobID := uuid.New()
buildID := uuid.New()
db.EXPECT().
GetAuthorizationUserRoles(gomock.Any(), ownerID).
Return(database.GetAuthorizationUserRolesRow{
ID: ownerID,
Roles: []string{},
Groups: []string{},
Status: database.UserStatusActive,
}, nil)
db.EXPECT().
GetTemplateByID(gomock.Any(), templateID).
Return(database.Template{
ID: templateID,
OrganizationID: orgID,
}, nil)
db.EXPECT().
GetChatWorkspaceTTL(gomock.Any()).
Return("0s", nil)
// waitForBuild fetches the build by ID.
db.EXPECT().
GetWorkspaceBuildByID(gomock.Any(), buildID).
Return(database.WorkspaceBuild{
ID: buildID,
WorkspaceID: workspaceID,
JobID: jobID,
}, nil)
// waitForBuild polls the provisioner job. Return Canceled.
db.EXPECT().
GetProvisionerJobByID(gomock.Any(), jobID).
Return(database.ProvisionerJob{
ID: jobID,
JobStatus: database.ProvisionerJobStatusCanceled,
}, nil)
createFn := func(_ context.Context, _ uuid.UUID, req codersdk.CreateWorkspaceRequest) (codersdk.Workspace, error) {
return codersdk.Workspace{
ID: workspaceID,
Name: req.Name,
OwnerName: "testuser",
LatestBuild: codersdk.WorkspaceBuild{
ID: buildID,
},
}, nil
}
tool := CreateWorkspace(orgID, db, CreateWorkspaceOptions{
OwnerID: ownerID,
ChatID: uuid.Nil,
CreateFn: createFn,
WorkspaceMu: &sync.Mutex{},
Logger: slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}),
})
input := fmt.Sprintf(`{"template_id":%q,"name":"test-build-cancel"}`, templateID.String())
resp, err := tool.Run(context.Background(), fantasy.ToolCall{
ID: "call-1",
Name: "create_workspace",
Input: input,
})
require.NoError(t, err)
var result map[string]any
require.NoError(t, json.Unmarshal([]byte(resp.Content), &result))
require.Contains(t, result["error"], "build was canceled")
require.Equal(t, buildID.String(), result["build_id"])
require.False(t, resp.IsError,
"buildToolResponse must not set IsError; chatprompt strips structured fields from error responses")
}
func TestCheckExistingWorkspace_StoppedWorkspace(t *testing.T) {
t.Parallel()
ctrl := gomock.NewController(t)
db := dbmock.NewMockStore(ctrl)
chatID := uuid.New()
workspaceID := uuid.New()
jobID := uuid.New()
expectExistingWorkspaceLookup(
db,
chatID,
workspaceID,
jobID,
"stopped-workspace",
database.ProvisionerJobStatusSucceeded,
database.WorkspaceTransitionStop,
)
options := testCheckExistingWorkspaceOptions(chatID, nil)
check := options.checkExistingWorkspace(context.Background(), db)
require.True(t, check.Done)
require.NoError(t, check.Err)
require.Equal(t, "stopped", check.Result["status"])
require.Contains(t, check.Result["message"], "start_workspace")
}
func TestCheckExistingWorkspace_DeletedWorkspace(t *testing.T) {
t.Parallel()
ctrl := gomock.NewController(t)
db := dbmock.NewMockStore(ctrl)
chatID := uuid.New()
workspaceID := uuid.New()
// Mock GetChatByID returns a chat linked to a workspace.
db.EXPECT().
GetChatByID(gomock.Any(), chatID).
Return(database.Chat{
ID: chatID,
WorkspaceID: uuid.NullUUID{UUID: workspaceID, Valid: true},
}, nil)
// Mock GetWorkspaceByID returns a soft-deleted workspace.
db.EXPECT().
GetWorkspaceByID(gomock.Any(), workspaceID).
Return(database.Workspace{
ID: workspaceID,
Deleted: true,
}, nil)
options := testCheckExistingWorkspaceOptions(chatID, nil)
check := options.checkExistingWorkspace(context.Background(), db)
require.NoError(t, check.Err)
require.False(t, check.Done, "should allow creation for deleted workspace")
require.Nil(t, check.Result)
}
func testCheckExistingWorkspaceOptions(
chatID uuid.UUID,
agentConnFn AgentConnFunc,
) CreateWorkspaceOptions {
return CreateWorkspaceOptions{
ChatID: chatID,
AgentConnFn: agentConnFn,
AgentInactiveDisconnectTimeout: 30 * time.Second,
}
}
func expectExistingWorkspaceLookup(
db *dbmock.MockStore,
chatID uuid.UUID,
workspaceID uuid.UUID,
jobID uuid.UUID,
workspaceName string,
jobStatus database.ProvisionerJobStatus,
transition database.WorkspaceTransition,
) {
db.EXPECT().
GetChatByID(gomock.Any(), chatID).
Return(database.Chat{
ID: chatID,
WorkspaceID: uuid.NullUUID{UUID: workspaceID, Valid: true},
}, nil)
db.EXPECT().
GetWorkspaceByID(gomock.Any(), workspaceID).
Return(database.Workspace{
ID: workspaceID,
Name: workspaceName,
}, nil)
db.EXPECT().
GetLatestWorkspaceBuildByWorkspaceID(gomock.Any(), workspaceID).
Return(database.WorkspaceBuild{
WorkspaceID: workspaceID,
JobID: jobID,
Transition: transition,
}, nil)
db.EXPECT().
GetProvisionerJobByID(gomock.Any(), jobID).
Return(database.ProvisionerJob{
ID: jobID,
JobStatus: jobStatus,
}, nil)
}
func TestCreateWorkspace_OnChatUpdatedFiresAfterBuild(t *testing.T) {
t.Parallel()
ctrl := gomock.NewController(t)
db := dbmock.NewMockStore(ctrl)
ownerID := uuid.New()
templateID := uuid.New()
workspaceID := uuid.New()
chatID := uuid.New()
jobID := uuid.New()
buildID := uuid.New()
// checkExistingWorkspace calls GetChatByID first. Return a chat
// with no workspace so the tool proceeds to creation.
db.EXPECT().
GetChatByID(gomock.Any(), chatID).
Return(database.Chat{
ID: chatID,
}, nil)
db.EXPECT().
GetAuthorizationUserRoles(gomock.Any(), ownerID).
Return(database.GetAuthorizationUserRolesRow{
ID: ownerID,
Roles: []string{},
Groups: []string{},
Status: database.UserStatusActive,
}, nil)
// Org check: GetTemplateByID returns a template in the
// same org (uuid.Nil matches our organizationID param).
db.EXPECT().
GetTemplateByID(gomock.Any(), templateID).
Return(database.Template{
ID: templateID,
OrganizationID: uuid.Nil,
}, nil)
db.EXPECT().
GetChatWorkspaceTTL(gomock.Any()).
Return("0s", nil)
// UpdateChatWorkspaceBinding — triggers first OnChatUpdated.
db.EXPECT().
UpdateChatWorkspaceBinding(gomock.Any(), gomock.Any()).
Return(database.Chat{
ID: chatID,
WorkspaceID: uuid.NullUUID{UUID: workspaceID, Valid: true},
}, nil)
// waitForBuild: fetch build, then poll job as completed.
db.EXPECT().
GetWorkspaceBuildByID(gomock.Any(), buildID).
Return(database.WorkspaceBuild{
ID: buildID,
WorkspaceID: workspaceID,
JobID: jobID,
}, nil)
db.EXPECT().
GetProvisionerJobByID(gomock.Any(), jobID).
Return(database.ProvisionerJob{
ID: jobID,
JobStatus: database.ProvisionerJobStatusSucceeded,
CompletedAt: validNullTime(time.Now()),
}, nil)
// GetChatByID — called after waitForBuild for second OnChatUpdated.
// GetChatByID — called after waitForBuild for second OnChatUpdated.
db.EXPECT().
GetChatByID(gomock.Any(), chatID).
Return(database.Chat{
ID: chatID,
WorkspaceID: uuid.NullUUID{UUID: workspaceID, Valid: true},
}, nil)
// Agent lookup after build completes — return empty so we skip
// agent selection and waitForAgentReady.
db.EXPECT().
GetWorkspaceAgentsInLatestBuildByWorkspaceID(gomock.Any(), workspaceID).
Return([]database.WorkspaceAgent{}, nil)
var mu sync.Mutex
var callbackChats []database.Chat
createFn := func(_ context.Context, _ uuid.UUID, req codersdk.CreateWorkspaceRequest) (codersdk.Workspace, error) {
return codersdk.Workspace{
ID: workspaceID,
Name: req.Name,
OwnerName: "testuser",
LatestBuild: codersdk.WorkspaceBuild{
ID: buildID,
},
}, nil
}
tool := CreateWorkspace(uuid.Nil, db, CreateWorkspaceOptions{
OwnerID: ownerID,
ChatID: chatID,
CreateFn: createFn,
WorkspaceMu: &sync.Mutex{},
Logger: slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}),
OnChatUpdated: func(chat database.Chat) {
mu.Lock()
callbackChats = append(callbackChats, chat)
mu.Unlock()
},
})
input := fmt.Sprintf(`{"template_id":%q,"name":"test-callback"}`, templateID.String())
resp, err := tool.Run(context.Background(), fantasy.ToolCall{
ID: "call-1",
Name: "create_workspace",
Input: input,
})
require.NoError(t, err)
require.False(t, resp.IsError)
mu.Lock()
defer mu.Unlock()
require.Len(t, callbackChats, 2,
"OnChatUpdated should fire twice: once on binding, once after build completes")
// Both callbacks should carry the workspace ID.
for i, chat := range callbackChats {
require.True(t, chat.WorkspaceID.Valid, "callback %d should have workspace ID", i)
require.Equal(t, workspaceID, chat.WorkspaceID.UUID)
}
}
func validNullTime(t time.Time) sql.NullTime {
return sql.NullTime{Time: t, Valid: true}
}
Reference in New Issue View Git Blame Copy Permalink