mirror of
https://github.com/coder/coder.git
synced 2026-06-02 20:48:20 +00:00
ef0151601e
## Summary When a workspace build fails because the user is over their group quota, the chat tools currently surface the failure as a bare `"workspace build failed: insufficient quota"` string with no machine-readable error code and no visibility into the user's current usage. Agents and the UI cannot distinguish quota failures from any other Terraform error, so users see an opaque message and have no clear path to recovery. This PR tags quota failures with a typed error code at the source and propagates it through the chat tool layer so callers can react to it explicitly. Relates to CODAGT-20 ## Changes **Provisioner runner** - Add `InsufficientQuotaErrorCode = "INSUFFICIENT_QUOTA"` and set it explicitly at the `commitQuota` failure site via a new `failedWorkspaceBuildfCode` helper, so `provisioner_jobs.error_code` is populated only on the genuine quota path. The substring matcher used for externally produced sentinels (e.g. `"missing parameter"`, `"required template variables"`) is intentionally not extended; provider errors that happen to mention "insufficient quota" stay classified as generic build failures. **SDK and API contract** - Add `JobErrorCodeInsufficientQuota` and a `JobIsInsufficientQuotaErrorCode` helper to `codersdk`. - Extend the swagger `enums` tag on `ProvisionerJob.ErrorCode` to include `INSUFFICIENT_QUOTA`. - Regenerate `coderd/apidoc`, `docs/reference/api/*`, and `site/src/api/typesGenerated.ts`. **chattool create_workspace / start_workspace** - `waitForBuild` now returns a typed `*workspaceBuildError` carrying both the message and the `JobErrorCode`, instead of a bare error string. - New `quotaerror.go` introduces a structured `quotaErrorResult` (with `error_code`, `title`, `message`, `build_id`, and optional `quota`) and a best-effort `workspaceQuotaDetails` lookup that wraps owner authorization internally and fetches `credits_consumed` and `budget` from the database. Quota lookup failures (including authorization failures) never block the failure payload. - On quota-coded build failures, both `create_workspace` and `start_workspace` now return the structured response (with the recovery guidance inlined into `message`) instead of the bare `"insufficient quota"` string. This applies to all three failure paths: post-creation, an in-progress existing build, and a freshly triggered start build. Non-quota build failures continue to use the existing `buildToolResponse` / `newBuildError` path. - Owner authorization is wrapped only on the call sites that need it (the `CreateFn` and `StartFn` invocations and the quota-detail lookup), so idempotent fast paths (already running, already in progress, existing-workspace early returns) do not pay for an extra RBAC round-trip or fail when role lookup is transient. ## Out of scope - No changes to quota math, allowances, or bypass behavior. - No automatic retries. - No new quota-inspection tools and no changes to MCP `coder_create_workspace` (which returns immediately and never observed the build outcome here). - No frontend UI changes; those will land in a follow-up PR that consumes the new `INSUFFICIENT_QUOTA` code.
1863 lines
53 KiB
Go
1863 lines
53 KiB
Go
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/dbauthz"
|
|
"github.com/coder/coder/v2/coderd/database/dbmock"
|
|
"github.com/coder/coder/v2/coderd/httpapi/httperror"
|
|
"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)
|
|
})
|
|
}
|
|
|
|
func TestCreateWorkspace_PrefersChatSuffixAgent(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()
|
|
fallbackAgentID := uuid.New()
|
|
chatAgentID := uuid.New()
|
|
|
|
db.EXPECT().
|
|
GetChatByID(gomock.Any(), chatID).
|
|
Return(database.Chat{ID: chatID}, nil)
|
|
|
|
db.EXPECT().
|
|
UpdateChatWorkspaceBinding(gomock.Any(), gomock.Any()).
|
|
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().
|
|
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(db, orgID, chatID, 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(db, orgID, chatID, CreateWorkspaceOptions{
|
|
OwnerID: ownerID,
|
|
|
|
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()
|
|
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().
|
|
UpdateChatWorkspaceBinding(gomock.Any(), gomock.Any()).
|
|
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)
|
|
|
|
// 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(db, orgID, chatID, CreateWorkspaceOptions{
|
|
OwnerID: ownerID,
|
|
|
|
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.NotContains(t, result, "error_code",
|
|
"generic build failures must not surface a quota error_code")
|
|
require.NotContains(t, result, "quota",
|
|
"generic build failures must not surface quota details")
|
|
require.False(t, resp.IsError,
|
|
"buildToolResponse must not set IsError; chatprompt strips structured fields from error responses")
|
|
}
|
|
|
|
func TestCreateWorkspace_PostCreationQuotaFailure(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().
|
|
UpdateChatWorkspaceBinding(gomock.Any(), gomock.Any()).
|
|
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.ProvisionerJobStatusFailed,
|
|
Error: sql.NullString{String: "insufficient quota", Valid: true},
|
|
ErrorCode: sql.NullString{
|
|
String: string(codersdk.InsufficientQuota),
|
|
Valid: true,
|
|
},
|
|
}, nil)
|
|
|
|
db.EXPECT().
|
|
GetQuotaConsumedForUser(gomock.Any(), database.GetQuotaConsumedForUserParams{
|
|
OwnerID: ownerID,
|
|
OrganizationID: orgID,
|
|
}).
|
|
Return(int64(40), nil)
|
|
db.EXPECT().
|
|
GetQuotaAllowanceForUser(gomock.Any(), database.GetQuotaAllowanceForUserParams{
|
|
UserID: ownerID,
|
|
OrganizationID: orgID,
|
|
}).
|
|
Return(int64(40), 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(db, orgID, chatID, CreateWorkspaceOptions{
|
|
OwnerID: ownerID,
|
|
CreateFn: createFn,
|
|
WorkspaceMu: &sync.Mutex{},
|
|
Logger: slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}),
|
|
})
|
|
|
|
input := fmt.Sprintf(`{"template_id":%q,"name":"test-quota-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.Equal(t, string(codersdk.InsufficientQuota), result["error_code"])
|
|
require.Equal(t, "Workspace quota reached", result["title"])
|
|
require.Contains(t, result["error"], "workspace build failed")
|
|
require.Contains(t, result["message"], "workspace quota is full")
|
|
require.Contains(t, result["message"], "Delete a workspace")
|
|
require.Contains(t, result["message"], "raise your group quota allowance")
|
|
require.NotContains(t, result, "next_steps")
|
|
require.Equal(t, buildID.String(), result["build_id"])
|
|
quota, ok := result["quota"].(map[string]any)
|
|
require.True(t, ok)
|
|
require.Equal(t, float64(40), quota["credits_consumed"])
|
|
require.Equal(t, float64(40), quota["budget"])
|
|
require.False(t, resp.IsError,
|
|
"quota responses must not set IsError; chatprompt strips structured fields from error responses")
|
|
}
|
|
|
|
func TestCreateWorkspace_ExistingBuildQuotaFailure(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().
|
|
GetAuthorizationUserRoles(gomock.Any(), ownerID).
|
|
Return(database.GetAuthorizationUserRolesRow{
|
|
ID: ownerID,
|
|
Roles: []string{},
|
|
Groups: []string{},
|
|
Status: database.UserStatusActive,
|
|
}, nil)
|
|
|
|
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: "existing-quota-workspace",
|
|
OrganizationID: orgID,
|
|
}, nil)
|
|
db.EXPECT().
|
|
GetLatestWorkspaceBuildByWorkspaceID(gomock.Any(), workspaceID).
|
|
Return(database.WorkspaceBuild{
|
|
ID: buildID,
|
|
WorkspaceID: workspaceID,
|
|
JobID: jobID,
|
|
Transition: database.WorkspaceTransitionStart,
|
|
}, nil)
|
|
firstJob := db.EXPECT().
|
|
GetProvisionerJobByID(gomock.Any(), jobID).
|
|
Return(database.ProvisionerJob{
|
|
ID: jobID,
|
|
JobStatus: database.ProvisionerJobStatusRunning,
|
|
}, 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().
|
|
GetWorkspaceBuildByID(gomock.Any(), buildID).
|
|
Return(database.WorkspaceBuild{
|
|
ID: buildID,
|
|
WorkspaceID: workspaceID,
|
|
JobID: jobID,
|
|
Transition: database.WorkspaceTransitionStart,
|
|
}, nil)
|
|
db.EXPECT().
|
|
GetProvisionerJobByID(gomock.Any(), jobID).
|
|
Return(database.ProvisionerJob{
|
|
ID: jobID,
|
|
JobStatus: database.ProvisionerJobStatusFailed,
|
|
Error: sql.NullString{String: "insufficient quota", Valid: true},
|
|
ErrorCode: sql.NullString{
|
|
String: string(codersdk.InsufficientQuota),
|
|
Valid: true,
|
|
},
|
|
}, nil).
|
|
After(firstJob)
|
|
ownerCtx := ownerContextMatcher{ownerID: ownerID}
|
|
db.EXPECT().
|
|
GetQuotaConsumedForUser(ownerCtx, database.GetQuotaConsumedForUserParams{
|
|
OwnerID: ownerID,
|
|
OrganizationID: orgID,
|
|
}).
|
|
Return(int64(40), nil)
|
|
db.EXPECT().
|
|
GetQuotaAllowanceForUser(ownerCtx, database.GetQuotaAllowanceForUserParams{
|
|
UserID: ownerID,
|
|
OrganizationID: orgID,
|
|
}).
|
|
Return(int64(40), nil)
|
|
|
|
tool := CreateWorkspace(db, orgID, chatID, CreateWorkspaceOptions{
|
|
OwnerID: ownerID,
|
|
CreateFn: func(context.Context, uuid.UUID, codersdk.CreateWorkspaceRequest) (codersdk.Workspace, error) {
|
|
t.Fatal("CreateFn should not be called when an existing build is in progress")
|
|
return codersdk.Workspace{}, nil
|
|
},
|
|
WorkspaceMu: &sync.Mutex{},
|
|
Logger: slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}),
|
|
})
|
|
|
|
input := fmt.Sprintf(`{"template_id":%q,"name":"test-existing-quota-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.Equal(t, string(codersdk.InsufficientQuota), result["error_code"])
|
|
require.Equal(t, "Workspace quota reached", result["title"])
|
|
require.Contains(t, result["error"], "existing workspace build failed")
|
|
require.Contains(t, result["message"], "could not start this workspace")
|
|
require.Contains(t, result["message"], "workspace quota is full")
|
|
require.Equal(t, buildID.String(), result["build_id"])
|
|
quota, ok := result["quota"].(map[string]any)
|
|
require.True(t, ok)
|
|
require.Equal(t, float64(40), quota["credits_consumed"])
|
|
require.Equal(t, float64(40), quota["budget"])
|
|
require.False(t, resp.IsError)
|
|
}
|
|
|
|
func TestCreateWorkspace_ResponderErrorPreservesStructuredFields(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()
|
|
|
|
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)
|
|
|
|
tool := CreateWorkspace(db, orgID, chatID, CreateWorkspaceOptions{
|
|
OwnerID: ownerID,
|
|
CreateFn: func(context.Context, uuid.UUID, codersdk.CreateWorkspaceRequest) (codersdk.Workspace, error) {
|
|
return codersdk.Workspace{}, httperror.NewResponseError(400, codersdk.Response{
|
|
Message: "missing required parameter",
|
|
Detail: "region must be set before the workspace can start",
|
|
Validations: []codersdk.ValidationError{{
|
|
Field: "region",
|
|
Detail: "region must be set before the workspace can start",
|
|
}},
|
|
})
|
|
},
|
|
WorkspaceMu: &sync.Mutex{},
|
|
Logger: slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}),
|
|
})
|
|
|
|
input := fmt.Sprintf(`{"template_id":%q,"name":"test-structured-error"}`, 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)
|
|
|
|
var result struct {
|
|
Error string `json:"error"`
|
|
Detail string `json:"detail"`
|
|
Validations []codersdk.ValidationError `json:"validations"`
|
|
}
|
|
require.NoError(t, json.Unmarshal([]byte(resp.Content), &result))
|
|
require.Equal(t, "missing required parameter", result.Error)
|
|
require.Equal(t, "region must be set before the workspace can start", result.Detail)
|
|
require.Equal(t, []codersdk.ValidationError{{
|
|
Field: "region",
|
|
Detail: "region must be set before the workspace can start",
|
|
}}, result.Validations)
|
|
}
|
|
|
|
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()
|
|
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().
|
|
UpdateChatWorkspaceBinding(gomock.Any(), gomock.Any()).
|
|
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(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(db, orgID, chatID, CreateWorkspaceOptions{
|
|
OwnerID: ownerID,
|
|
|
|
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(db, chatOrgID, chatID, CreateWorkspaceOptions{
|
|
OwnerID: ownerID,
|
|
|
|
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(connFn)
|
|
check := options.checkExistingWorkspace(context.Background(), db, chatID)
|
|
|
|
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(nil)
|
|
check := options.checkExistingWorkspace(context.Background(), db, chatID)
|
|
|
|
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(nil)
|
|
check := options.checkExistingWorkspace(context.Background(), db, chatID)
|
|
|
|
require.Error(t, check.BuildErr)
|
|
require.Contains(t, check.BuildErr.Error(), "existing workspace build failed")
|
|
require.Equal(t, buildID, check.BuildID)
|
|
require.Equal(t, buildFailureActionStart, check.BuildAction)
|
|
require.NoError(t, check.Err)
|
|
}
|
|
|
|
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(connFn)
|
|
check := options.checkExistingWorkspace(context.Background(), db, chatID)
|
|
|
|
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(nil)
|
|
check := options.checkExistingWorkspace(context.Background(), db, chatID)
|
|
|
|
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()
|
|
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().
|
|
UpdateChatWorkspaceBinding(gomock.Any(), gomock.Any()).
|
|
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)
|
|
|
|
// 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(db, orgID, chatID, CreateWorkspaceOptions{
|
|
OwnerID: ownerID,
|
|
|
|
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(nil)
|
|
check := options.checkExistingWorkspace(context.Background(), db, chatID)
|
|
|
|
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(nil)
|
|
check := options.checkExistingWorkspace(context.Background(), db, chatID)
|
|
|
|
require.NoError(t, check.Err)
|
|
require.False(t, check.Done, "should allow creation for deleted workspace")
|
|
require.Nil(t, check.Result)
|
|
}
|
|
|
|
func testCheckExistingWorkspaceOptions(
|
|
agentConnFn AgentConnFunc,
|
|
) CreateWorkspaceOptions {
|
|
return CreateWorkspaceOptions{
|
|
AgentConnFn: agentConnFn,
|
|
AgentInactiveDisconnectTimeout: 30 * time.Second,
|
|
}
|
|
}
|
|
|
|
type ownerContextMatcher struct {
|
|
ownerID uuid.UUID
|
|
}
|
|
|
|
func (m ownerContextMatcher) Matches(v any) bool {
|
|
ctx, ok := v.(context.Context)
|
|
if !ok {
|
|
return false
|
|
}
|
|
actor, ok := dbauthz.ActorFromContext(ctx)
|
|
return ok && actor.ID == m.ownerID.String()
|
|
}
|
|
|
|
func (ownerContextMatcher) String() string {
|
|
return "context with owner actor"
|
|
}
|
|
|
|
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.
|
|
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(db, uuid.Nil, chatID, CreateWorkspaceOptions{
|
|
OwnerID: ownerID,
|
|
|
|
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}
|
|
}
|
|
|
|
// createWorkspacePresetTestSetup holds common test dependencies
|
|
// for create_workspace preset tests.
|
|
type createWorkspacePresetTestSetup struct {
|
|
DB *dbmock.MockStore
|
|
OwnerID uuid.UUID
|
|
OrgID uuid.UUID
|
|
TemplateID uuid.UUID
|
|
ChatID uuid.UUID
|
|
WorkspaceID uuid.UUID
|
|
BuildID uuid.UUID
|
|
AgentID uuid.UUID
|
|
}
|
|
|
|
// setupCreateWorkspacePresetTest creates common mock expectations
|
|
// for preset-related create_workspace tests. It sets up RBAC,
|
|
// template lookup, TTL, and chat lookup.
|
|
func setupCreateWorkspacePresetTest(t *testing.T) createWorkspacePresetTestSetup {
|
|
t.Helper()
|
|
|
|
ctrl := gomock.NewController(t)
|
|
db := dbmock.NewMockStore(ctrl)
|
|
|
|
s := createWorkspacePresetTestSetup{
|
|
DB: db,
|
|
OwnerID: uuid.New(),
|
|
OrgID: uuid.New(),
|
|
TemplateID: uuid.New(),
|
|
ChatID: uuid.New(),
|
|
WorkspaceID: uuid.New(),
|
|
BuildID: uuid.New(),
|
|
AgentID: uuid.New(),
|
|
}
|
|
|
|
// RBAC.
|
|
db.EXPECT().
|
|
GetAuthorizationUserRoles(gomock.Any(), s.OwnerID).
|
|
Return(database.GetAuthorizationUserRolesRow{
|
|
ID: s.OwnerID,
|
|
Username: "testuser",
|
|
Status: "active",
|
|
}, nil)
|
|
|
|
// Template lookup.
|
|
db.EXPECT().
|
|
GetTemplateByID(gomock.Any(), s.TemplateID).
|
|
Return(database.Template{
|
|
ID: s.TemplateID,
|
|
OrganizationID: s.OrgID,
|
|
Name: "test-template",
|
|
ActiveVersionID: uuid.New(),
|
|
}, nil)
|
|
|
|
// Chat workspace TTL.
|
|
db.EXPECT().
|
|
GetChatWorkspaceTTL(gomock.Any()).
|
|
Return("", sql.ErrNoRows)
|
|
|
|
// Check for existing workspace (no existing).
|
|
db.EXPECT().
|
|
GetChatByID(gomock.Any(), s.ChatID).
|
|
Return(database.Chat{ID: s.ChatID}, nil)
|
|
|
|
return s
|
|
}
|
|
|
|
// expectSuccessfulBuild adds mock expectations for a successful
|
|
// build, agent lookup, and agent lifecycle check.
|
|
func (s createWorkspacePresetTestSetup) expectSuccessfulBuild() {
|
|
s.DB.EXPECT().
|
|
UpdateChatWorkspaceBinding(gomock.Any(), gomock.Any()).
|
|
Return(database.Chat{ID: s.ChatID}, nil)
|
|
|
|
s.DB.EXPECT().
|
|
GetWorkspaceBuildByID(gomock.Any(), s.BuildID).
|
|
Return(database.WorkspaceBuild{
|
|
ID: s.BuildID,
|
|
JobID: uuid.New(),
|
|
}, nil)
|
|
s.DB.EXPECT().
|
|
GetProvisionerJobByID(gomock.Any(), gomock.Any()).
|
|
Return(database.ProvisionerJob{
|
|
JobStatus: database.ProvisionerJobStatusSucceeded,
|
|
}, nil)
|
|
|
|
s.DB.EXPECT().
|
|
GetWorkspaceAgentsInLatestBuildByWorkspaceID(gomock.Any(), s.WorkspaceID).
|
|
Return([]database.WorkspaceAgent{{
|
|
ID: s.AgentID,
|
|
Name: "main",
|
|
}}, nil)
|
|
|
|
s.DB.EXPECT().
|
|
GetWorkspaceAgentLifecycleStateByID(gomock.Any(), s.AgentID).
|
|
Return(database.GetWorkspaceAgentLifecycleStateByIDRow{
|
|
LifecycleState: database.WorkspaceAgentLifecycleStateReady,
|
|
}, nil)
|
|
}
|
|
|
|
func TestCreateWorkspace_WithPresetID(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
s := setupCreateWorkspacePresetTest(t)
|
|
s.expectSuccessfulBuild()
|
|
|
|
presetID := uuid.New()
|
|
|
|
var capturedReq codersdk.CreateWorkspaceRequest
|
|
createFn := func(_ context.Context, _ uuid.UUID, req codersdk.CreateWorkspaceRequest) (codersdk.Workspace, error) {
|
|
capturedReq = req
|
|
return codersdk.Workspace{
|
|
ID: s.WorkspaceID,
|
|
Name: req.Name,
|
|
LatestBuild: codersdk.WorkspaceBuild{
|
|
ID: s.BuildID,
|
|
},
|
|
}, nil
|
|
}
|
|
|
|
agentConnFn := func(_ context.Context, _ uuid.UUID) (workspacesdk.AgentConn, func(), error) {
|
|
return nil, func() {}, nil
|
|
}
|
|
|
|
tool := CreateWorkspace(s.DB, s.OrgID, s.ChatID, CreateWorkspaceOptions{
|
|
OwnerID: s.OwnerID,
|
|
CreateFn: createFn,
|
|
AgentConnFn: agentConnFn,
|
|
WorkspaceMu: &sync.Mutex{},
|
|
Logger: slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}),
|
|
})
|
|
|
|
input := fmt.Sprintf(
|
|
`{"template_id":%q,"preset_id":%q,"name":"test-ws"}`,
|
|
s.TemplateID.String(), presetID.String(),
|
|
)
|
|
|
|
ctx := context.Background()
|
|
resp, err := tool.Run(ctx, fantasy.ToolCall{
|
|
ID: "call-preset",
|
|
Name: "create_workspace",
|
|
Input: input,
|
|
})
|
|
require.NoError(t, err)
|
|
require.False(t, resp.IsError, "unexpected error: %s", resp.Content)
|
|
|
|
require.Equal(t, presetID, capturedReq.TemplateVersionPresetID,
|
|
"expected preset ID to be set on CreateWorkspaceRequest")
|
|
}
|
|
|
|
func TestCreateWorkspace_InvalidPresetID(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
s := setupCreateWorkspacePresetTest(t)
|
|
|
|
tool := CreateWorkspace(s.DB, s.OrgID, s.ChatID, CreateWorkspaceOptions{
|
|
OwnerID: s.OwnerID,
|
|
CreateFn: func(_ context.Context, _ uuid.UUID, _ codersdk.CreateWorkspaceRequest) (codersdk.Workspace, error) {
|
|
t.Fatal("CreateFn should not be called with invalid preset_id")
|
|
return codersdk.Workspace{}, nil
|
|
},
|
|
WorkspaceMu: &sync.Mutex{},
|
|
Logger: slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}),
|
|
})
|
|
|
|
input := fmt.Sprintf(
|
|
`{"template_id":%q,"preset_id":"not-a-uuid","name":"test-ws"}`,
|
|
s.TemplateID.String(),
|
|
)
|
|
|
|
ctx := context.Background()
|
|
resp, err := tool.Run(ctx, fantasy.ToolCall{
|
|
ID: "call-bad-preset",
|
|
Name: "create_workspace",
|
|
Input: input,
|
|
})
|
|
require.NoError(t, err)
|
|
require.True(t, resp.IsError)
|
|
require.Contains(t, resp.Content, "invalid preset_id")
|
|
}
|
|
|
|
func TestCreateWorkspace_WithPresetAndParams(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
s := setupCreateWorkspacePresetTest(t)
|
|
s.expectSuccessfulBuild()
|
|
|
|
presetID := uuid.New()
|
|
|
|
var capturedReq codersdk.CreateWorkspaceRequest
|
|
createFn := func(_ context.Context, _ uuid.UUID, req codersdk.CreateWorkspaceRequest) (codersdk.Workspace, error) {
|
|
capturedReq = req
|
|
return codersdk.Workspace{
|
|
ID: s.WorkspaceID,
|
|
Name: req.Name,
|
|
LatestBuild: codersdk.WorkspaceBuild{
|
|
ID: s.BuildID,
|
|
},
|
|
}, nil
|
|
}
|
|
|
|
agentConnFn := func(_ context.Context, _ uuid.UUID) (workspacesdk.AgentConn, func(), error) {
|
|
return nil, func() {}, nil
|
|
}
|
|
|
|
tool := CreateWorkspace(s.DB, s.OrgID, s.ChatID, CreateWorkspaceOptions{
|
|
OwnerID: s.OwnerID,
|
|
CreateFn: createFn,
|
|
AgentConnFn: agentConnFn,
|
|
WorkspaceMu: &sync.Mutex{},
|
|
Logger: slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}),
|
|
})
|
|
|
|
input := fmt.Sprintf(
|
|
`{"template_id":%q,"preset_id":%q,"name":"test-ws","parameters":{"region":"us-east"}}`,
|
|
s.TemplateID.String(), presetID.String(),
|
|
)
|
|
|
|
ctx := context.Background()
|
|
resp, err := tool.Run(ctx, fantasy.ToolCall{
|
|
ID: "call-preset-params",
|
|
Name: "create_workspace",
|
|
Input: input,
|
|
})
|
|
require.NoError(t, err)
|
|
require.False(t, resp.IsError, "unexpected error: %s", resp.Content)
|
|
|
|
// Verify preset ID is set.
|
|
require.Equal(t, presetID, capturedReq.TemplateVersionPresetID,
|
|
"expected preset ID to be set")
|
|
|
|
// Verify parameters are also populated.
|
|
require.Len(t, capturedReq.RichParameterValues, 1,
|
|
"expected rich parameter values to be set")
|
|
require.Equal(t, "region", capturedReq.RichParameterValues[0].Name)
|
|
require.Equal(t, "us-east", capturedReq.RichParameterValues[0].Value)
|
|
}
|