mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
feat(chatd): add start_workspace tool to agent flow (#22646)
## Summary When a chat's workspace is stopped, the LLM previously had no way to start it — `create_workspace` would either create a duplicate workspace or fail. This adds a dedicated `start_workspace` tool to the agent flow. ## Changes ### New: `start_workspace` tool (`coderd/chatd/chattool/startworkspace.go`) - Detects if the chat's workspace is stopped and starts it via a new build with `transition=start` - Reuses the existing `waitForBuild` and `waitForAgent` helpers (shared logic) - Shares the workspace mutex with `create_workspace` to prevent races - Idempotent: returns immediately if the workspace is already running or building - Returns a `no_agent` / `not_ready` status if the agent isn't available yet (non-fatal) ### Updated: `create_workspace` stopped-workspace hint - `checkExistingWorkspace` now returns a `stopped` status with message `"use start_workspace to start it"` when it detects the chat's workspace is stopped, instead of falling through to create a new workspace ### Wiring - `chatd.Config` / `chatd.Server`: new `StartWorkspace` / `startWorkspaceFn` field - `coderd/chats.go`: new `chatStartWorkspace` method that calls `postWorkspaceBuildsInternal` with proper RBAC context - `coderd/coderd.go`: passes `chatStartWorkspace` into chatd config - Tool registered alongside `create_workspace` for root chats only (not subagents) ### Tests (`startworkspace_test.go`) - `NoWorkspace`: error when chat has no workspace - `AlreadyRunning`: idempotent return for workspace with successful start build - `StoppedWorkspace`: verifies StartFn is called, build is waited on, and success response returned
This commit is contained in:
@@ -66,6 +66,7 @@ type Server struct {
|
||||
|
||||
agentConnFn AgentConnFunc
|
||||
createWorkspaceFn chattool.CreateWorkspaceFn
|
||||
startWorkspaceFn chattool.StartWorkspaceFn
|
||||
pubsub pubsub.Pubsub
|
||||
webpushDispatcher webpush.Dispatcher
|
||||
providerAPIKeys chatprovider.ProviderAPIKeys
|
||||
@@ -852,6 +853,7 @@ type Config struct {
|
||||
InFlightChatStaleAfter time.Duration
|
||||
AgentConn AgentConnFunc
|
||||
CreateWorkspace chattool.CreateWorkspaceFn
|
||||
StartWorkspace chattool.StartWorkspaceFn
|
||||
Pubsub pubsub.Pubsub
|
||||
ProviderAPIKeys chatprovider.ProviderAPIKeys
|
||||
WebpushDispatcher webpush.Dispatcher
|
||||
@@ -887,6 +889,7 @@ func New(cfg Config) *Server {
|
||||
subscribeFn: cfg.SubscribeFn,
|
||||
agentConnFn: cfg.AgentConn,
|
||||
createWorkspaceFn: cfg.CreateWorkspace,
|
||||
startWorkspaceFn: cfg.StartWorkspace,
|
||||
pubsub: cfg.Pubsub,
|
||||
webpushDispatcher: cfg.WebpushDispatcher,
|
||||
providerAPIKeys: cfg.ProviderAPIKeys,
|
||||
@@ -2210,6 +2213,14 @@ func (p *Server) runChat(
|
||||
AgentConnFn: chattool.AgentConnFunc(p.agentConnFn),
|
||||
WorkspaceMu: &workspaceMu,
|
||||
}),
|
||||
chattool.StartWorkspace(chattool.StartWorkspaceOptions{
|
||||
DB: p.db,
|
||||
OwnerID: chat.OwnerID,
|
||||
ChatID: chat.ID,
|
||||
StartFn: p.startWorkspaceFn,
|
||||
AgentConnFn: chattool.AgentConnFunc(p.agentConnFn),
|
||||
WorkspaceMu: &workspaceMu,
|
||||
}),
|
||||
)
|
||||
tools = append(tools, p.subagentTools(func() database.Chat {
|
||||
return chat
|
||||
|
||||
@@ -1133,6 +1133,162 @@ func TestCreateWorkspaceTool_EndToEnd(t *testing.T) {
|
||||
require.True(t, foundToolResultInSecondCall, "expected second streamed model call to include create_workspace tool output")
|
||||
}
|
||||
|
||||
func TestStartWorkspaceTool_EndToEnd(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
ctx := testutil.Context(t, testutil.WaitSuperLong)
|
||||
deploymentValues := coderdtest.DeploymentValues(t)
|
||||
deploymentValues.Experiments = []string{string(codersdk.ExperimentAgents)}
|
||||
client := coderdtest.New(t, &coderdtest.Options{
|
||||
DeploymentValues: deploymentValues,
|
||||
IncludeProvisionerDaemon: true,
|
||||
})
|
||||
user := coderdtest.CreateFirstUser(t, client)
|
||||
|
||||
version := coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, &echo.Responses{
|
||||
Parse: echo.ParseComplete,
|
||||
ProvisionPlan: echo.PlanComplete,
|
||||
ProvisionApply: echo.ApplyComplete,
|
||||
})
|
||||
coderdtest.AwaitTemplateVersionJobCompleted(t, client, version.ID)
|
||||
template := coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID)
|
||||
|
||||
// Create a workspace, then stop it so start_workspace has
|
||||
// something to start. We intentionally skip starting a test
|
||||
// agent — the echo provisioner creates new agent rows for each
|
||||
// build, so an agent started for build 1 cannot serve build 3.
|
||||
// The tool handles the no-agent case gracefully.
|
||||
workspace := coderdtest.CreateWorkspace(t, client, template.ID)
|
||||
coderdtest.AwaitWorkspaceBuildJobCompleted(t, client, workspace.LatestBuild.ID)
|
||||
workspace = coderdtest.MustTransitionWorkspace(
|
||||
t, client, workspace.ID,
|
||||
codersdk.WorkspaceTransitionStart, codersdk.WorkspaceTransitionStop,
|
||||
)
|
||||
|
||||
var streamedCallCount atomic.Int32
|
||||
var streamedCallsMu sync.Mutex
|
||||
streamedCalls := make([][]chattest.OpenAIMessage, 0, 2)
|
||||
|
||||
openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse {
|
||||
if !req.Stream {
|
||||
return chattest.OpenAINonStreamingResponse("Start workspace test")
|
||||
}
|
||||
|
||||
streamedCallsMu.Lock()
|
||||
streamedCalls = append(streamedCalls, append([]chattest.OpenAIMessage(nil), req.Messages...))
|
||||
streamedCallsMu.Unlock()
|
||||
|
||||
if streamedCallCount.Add(1) == 1 {
|
||||
return chattest.OpenAIStreamingResponse(
|
||||
chattest.OpenAIToolCallChunk("start_workspace", "{}"),
|
||||
)
|
||||
}
|
||||
return chattest.OpenAIStreamingResponse(
|
||||
chattest.OpenAITextChunks("Workspace started and ready.")...,
|
||||
)
|
||||
})
|
||||
|
||||
_, err := client.CreateChatProvider(ctx, codersdk.CreateChatProviderConfigRequest{
|
||||
Provider: "openai-compat",
|
||||
APIKey: "test-api-key",
|
||||
BaseURL: openAIURL,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
contextLimit := int64(4096)
|
||||
isDefault := true
|
||||
_, err = client.CreateChatModelConfig(ctx, codersdk.CreateChatModelConfigRequest{
|
||||
Provider: "openai-compat",
|
||||
Model: "gpt-4o-mini",
|
||||
ContextLimit: &contextLimit,
|
||||
IsDefault: &isDefault,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
// Create a chat with the stopped workspace pre-associated.
|
||||
chat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{
|
||||
Content: []codersdk.ChatInputPart{
|
||||
{
|
||||
Type: codersdk.ChatInputPartTypeText,
|
||||
Text: "Start the workspace.",
|
||||
},
|
||||
},
|
||||
WorkspaceID: &workspace.ID,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
var chatWithMessages codersdk.ChatWithMessages
|
||||
require.Eventually(t, func() bool {
|
||||
got, getErr := client.GetChat(ctx, chat.ID)
|
||||
if getErr != nil {
|
||||
return false
|
||||
}
|
||||
chatWithMessages = got
|
||||
return got.Chat.Status == codersdk.ChatStatusWaiting || got.Chat.Status == codersdk.ChatStatusError
|
||||
}, testutil.WaitSuperLong, testutil.IntervalFast)
|
||||
|
||||
if chatWithMessages.Chat.Status == codersdk.ChatStatusError {
|
||||
lastError := ""
|
||||
if chatWithMessages.Chat.LastError != nil {
|
||||
lastError = *chatWithMessages.Chat.LastError
|
||||
}
|
||||
require.FailNowf(t, "chat run failed", "last_error=%q", lastError)
|
||||
}
|
||||
|
||||
// Verify the workspace was started.
|
||||
require.NotNil(t, chatWithMessages.Chat.WorkspaceID)
|
||||
updatedWorkspace, err := client.Workspace(ctx, workspace.ID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, codersdk.WorkspaceTransitionStart, updatedWorkspace.LatestBuild.Transition)
|
||||
|
||||
// Verify start_workspace tool result exists in the chat messages.
|
||||
var foundStartWorkspaceResult bool
|
||||
for _, message := range chatWithMessages.Messages {
|
||||
if message.Role != "tool" {
|
||||
continue
|
||||
}
|
||||
for _, part := range message.Content {
|
||||
if part.Type != codersdk.ChatMessagePartTypeToolResult || part.ToolName != "start_workspace" {
|
||||
continue
|
||||
}
|
||||
var result map[string]any
|
||||
require.NoError(t, json.Unmarshal(part.Result, &result))
|
||||
started, ok := result["started"].(bool)
|
||||
require.True(t, ok)
|
||||
require.True(t, started)
|
||||
foundStartWorkspaceResult = true
|
||||
}
|
||||
}
|
||||
require.True(t, foundStartWorkspaceResult, "expected start_workspace tool result message")
|
||||
|
||||
// Verify the LLM received the tool result in its second call.
|
||||
require.GreaterOrEqual(t, streamedCallCount.Load(), int32(2))
|
||||
streamedCallsMu.Lock()
|
||||
recordedStreamCalls := append([][]chattest.OpenAIMessage(nil), streamedCalls...)
|
||||
streamedCallsMu.Unlock()
|
||||
require.GreaterOrEqual(t, len(recordedStreamCalls), 2)
|
||||
|
||||
var foundToolResultInSecondCall bool
|
||||
for _, message := range recordedStreamCalls[1] {
|
||||
if message.Role != "tool" {
|
||||
continue
|
||||
}
|
||||
if !json.Valid([]byte(message.Content)) {
|
||||
continue
|
||||
}
|
||||
var result map[string]any
|
||||
if err := json.Unmarshal([]byte(message.Content), &result); err != nil {
|
||||
continue
|
||||
}
|
||||
started, ok := result["started"].(bool)
|
||||
if ok && started {
|
||||
foundToolResultInSecondCall = true
|
||||
break
|
||||
}
|
||||
}
|
||||
require.True(t, foundToolResultInSecondCall, "expected second streamed model call to include start_workspace tool output")
|
||||
}
|
||||
|
||||
func newTestServer(
|
||||
t *testing.T,
|
||||
db database.Store,
|
||||
|
||||
@@ -288,6 +288,17 @@ func checkExistingWorkspace(
|
||||
return result, true, nil
|
||||
|
||||
case database.ProvisionerJobStatusSucceeded:
|
||||
// If the workspace was stopped, tell the model to use
|
||||
// start_workspace instead of creating a new one.
|
||||
if build.Transition == database.WorkspaceTransitionStop {
|
||||
return map[string]any{
|
||||
"created": false,
|
||||
"workspace_name": ws.Name,
|
||||
"status": "stopped",
|
||||
"message": "workspace is stopped; use start_workspace to start it",
|
||||
}, true, nil
|
||||
}
|
||||
|
||||
// Build succeeded — check if agent is reachable.
|
||||
agents, agentsErr := db.GetWorkspaceAgentsInLatestBuildByWorkspaceID(ctx, ws.ID)
|
||||
if agentsErr == nil && len(agents) > 0 && agentConnFn != nil {
|
||||
|
||||
@@ -0,0 +1,176 @@
|
||||
package chattool
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"sync"
|
||||
|
||||
"charm.land/fantasy"
|
||||
"github.com/google/uuid"
|
||||
"golang.org/x/xerrors"
|
||||
|
||||
"github.com/coder/coder/v2/coderd/database"
|
||||
"github.com/coder/coder/v2/codersdk"
|
||||
)
|
||||
|
||||
// StartWorkspaceFn starts a workspace by creating a new build with
|
||||
// the "start" transition.
|
||||
type StartWorkspaceFn func(
|
||||
ctx context.Context,
|
||||
ownerID uuid.UUID,
|
||||
workspaceID uuid.UUID,
|
||||
req codersdk.CreateWorkspaceBuildRequest,
|
||||
) (codersdk.WorkspaceBuild, error)
|
||||
|
||||
// StartWorkspaceOptions configures the start_workspace tool.
|
||||
type StartWorkspaceOptions struct {
|
||||
DB database.Store
|
||||
OwnerID uuid.UUID
|
||||
ChatID uuid.UUID
|
||||
StartFn StartWorkspaceFn
|
||||
AgentConnFn AgentConnFunc
|
||||
WorkspaceMu *sync.Mutex
|
||||
}
|
||||
|
||||
// StartWorkspace returns a tool that starts a stopped workspace
|
||||
// associated with the current chat. The tool is idempotent: if the
|
||||
// workspace is already running or building, it returns immediately.
|
||||
func StartWorkspace(options StartWorkspaceOptions) fantasy.AgentTool {
|
||||
return fantasy.NewAgentTool(
|
||||
"start_workspace",
|
||||
"Start the chat's workspace if it is currently stopped. "+
|
||||
"This tool is idempotent — if the workspace is already "+
|
||||
"running, it returns immediately. Use create_workspace "+
|
||||
"first if no workspace exists yet.",
|
||||
func(ctx context.Context, _ struct{}, _ fantasy.ToolCall) (fantasy.ToolResponse, error) {
|
||||
if options.StartFn == nil {
|
||||
return fantasy.NewTextErrorResponse("workspace starter is not configured"), nil
|
||||
}
|
||||
|
||||
// Serialize with create_workspace to prevent races.
|
||||
if options.WorkspaceMu != nil {
|
||||
options.WorkspaceMu.Lock()
|
||||
defer options.WorkspaceMu.Unlock()
|
||||
}
|
||||
|
||||
if options.DB == nil || options.ChatID == uuid.Nil {
|
||||
return fantasy.NewTextErrorResponse("start_workspace is not properly configured"), nil
|
||||
}
|
||||
|
||||
chat, err := options.DB.GetChatByID(ctx, options.ChatID)
|
||||
if err != nil {
|
||||
return fantasy.NewTextErrorResponse(
|
||||
xerrors.Errorf("load chat: %w", err).Error(),
|
||||
), nil
|
||||
}
|
||||
if !chat.WorkspaceID.Valid {
|
||||
return fantasy.NewTextErrorResponse(
|
||||
"chat has no workspace; use create_workspace first",
|
||||
), nil
|
||||
}
|
||||
|
||||
ws, err := options.DB.GetWorkspaceByID(ctx, chat.WorkspaceID.UUID)
|
||||
if err != nil {
|
||||
if xerrors.Is(err, sql.ErrNoRows) {
|
||||
return fantasy.NewTextErrorResponse(
|
||||
"workspace was deleted; use create_workspace to make a new one",
|
||||
), nil
|
||||
}
|
||||
return fantasy.NewTextErrorResponse(
|
||||
xerrors.Errorf("load workspace: %w", err).Error(),
|
||||
), nil
|
||||
}
|
||||
|
||||
build, err := options.DB.GetLatestWorkspaceBuildByWorkspaceID(ctx, ws.ID)
|
||||
if err != nil {
|
||||
return fantasy.NewTextErrorResponse(
|
||||
xerrors.Errorf("get latest build: %w", err).Error(),
|
||||
), nil
|
||||
}
|
||||
|
||||
job, err := options.DB.GetProvisionerJobByID(ctx, build.JobID)
|
||||
if err != nil {
|
||||
return fantasy.NewTextErrorResponse(
|
||||
xerrors.Errorf("get provisioner job: %w", err).Error(),
|
||||
), nil
|
||||
}
|
||||
|
||||
// If a build is already in progress, wait for it.
|
||||
switch job.JobStatus {
|
||||
case database.ProvisionerJobStatusPending,
|
||||
database.ProvisionerJobStatusRunning:
|
||||
if err := waitForBuild(ctx, options.DB, ws.ID); err != nil {
|
||||
return fantasy.NewTextErrorResponse(
|
||||
xerrors.Errorf("waiting for in-progress build: %w", err).Error(),
|
||||
), nil
|
||||
}
|
||||
return waitForAgentAndRespond(ctx, options.DB, options.AgentConnFn, ws)
|
||||
|
||||
case database.ProvisionerJobStatusSucceeded:
|
||||
// If the latest successful build is a start
|
||||
// transition, the workspace should be running.
|
||||
if build.Transition == database.WorkspaceTransitionStart {
|
||||
return waitForAgentAndRespond(ctx, options.DB, options.AgentConnFn, ws)
|
||||
}
|
||||
// Otherwise it is stopped (or deleted) — proceed
|
||||
// to start it below.
|
||||
|
||||
default:
|
||||
// Failed, canceled, etc — try starting anyway.
|
||||
}
|
||||
|
||||
// Set up dbauthz context for the start call.
|
||||
ownerCtx, ownerErr := asOwner(ctx, options.DB, options.OwnerID)
|
||||
if ownerErr != nil {
|
||||
return fantasy.NewTextErrorResponse(ownerErr.Error()), nil
|
||||
}
|
||||
|
||||
_, err = options.StartFn(ownerCtx, options.OwnerID, ws.ID, codersdk.CreateWorkspaceBuildRequest{
|
||||
Transition: codersdk.WorkspaceTransitionStart,
|
||||
})
|
||||
if err != nil {
|
||||
return fantasy.NewTextErrorResponse(
|
||||
xerrors.Errorf("start workspace: %w", err).Error(),
|
||||
), nil
|
||||
}
|
||||
|
||||
if err := waitForBuild(ctx, options.DB, ws.ID); err != nil {
|
||||
return fantasy.NewTextErrorResponse(
|
||||
xerrors.Errorf("workspace start build failed: %w", err).Error(),
|
||||
), nil
|
||||
}
|
||||
|
||||
return waitForAgentAndRespond(ctx, options.DB, options.AgentConnFn, ws)
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
// waitForAgentAndRespond looks up the first agent in the workspace's
|
||||
// latest build, waits for it to become reachable, and returns a
|
||||
// success response.
|
||||
func waitForAgentAndRespond(
|
||||
ctx context.Context,
|
||||
db database.Store,
|
||||
agentConnFn AgentConnFunc,
|
||||
ws database.Workspace,
|
||||
) (fantasy.ToolResponse, error) {
|
||||
agents, err := db.GetWorkspaceAgentsInLatestBuildByWorkspaceID(ctx, ws.ID)
|
||||
if err != nil || len(agents) == 0 {
|
||||
// Workspace started but no agent found — still report
|
||||
// success so the model knows the workspace is up.
|
||||
return toolResponse(map[string]any{
|
||||
"started": true,
|
||||
"workspace_name": ws.Name,
|
||||
"agent_status": "no_agent",
|
||||
}), nil
|
||||
}
|
||||
|
||||
result := map[string]any{
|
||||
"started": true,
|
||||
"workspace_name": ws.Name,
|
||||
}
|
||||
for k, v := range waitForAgentReady(ctx, db, agents[0].ID, agentConnFn) {
|
||||
result[k] = v
|
||||
}
|
||||
return toolResponse(result), nil
|
||||
}
|
||||
@@ -0,0 +1,213 @@
|
||||
package chattool_test
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
"charm.land/fantasy"
|
||||
"github.com/google/uuid"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/coder/coder/v2/coderd/chatd/chattool"
|
||||
"github.com/coder/coder/v2/coderd/database"
|
||||
"github.com/coder/coder/v2/coderd/database/dbfake"
|
||||
"github.com/coder/coder/v2/coderd/database/dbgen"
|
||||
"github.com/coder/coder/v2/coderd/database/dbtestutil"
|
||||
"github.com/coder/coder/v2/codersdk"
|
||||
"github.com/coder/coder/v2/codersdk/workspacesdk"
|
||||
"github.com/coder/coder/v2/testutil"
|
||||
)
|
||||
|
||||
func TestStartWorkspace(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
t.Run("NoWorkspace", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
db, _ := dbtestutil.NewDB(t)
|
||||
|
||||
user := dbgen.User(t, db, database.User{})
|
||||
modelCfg := seedModelConfig(ctx, t, db, user.ID)
|
||||
|
||||
chat, err := db.InsertChat(ctx, database.InsertChatParams{
|
||||
OwnerID: user.ID,
|
||||
LastModelConfigID: modelCfg.ID,
|
||||
Title: "test-no-workspace",
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
tool := chattool.StartWorkspace(chattool.StartWorkspaceOptions{
|
||||
DB: db,
|
||||
ChatID: chat.ID,
|
||||
StartFn: func(_ context.Context, _ uuid.UUID, _ uuid.UUID, _ codersdk.CreateWorkspaceBuildRequest) (codersdk.WorkspaceBuild, error) {
|
||||
t.Fatal("StartFn should not be called")
|
||||
return codersdk.WorkspaceBuild{}, nil
|
||||
},
|
||||
WorkspaceMu: &sync.Mutex{},
|
||||
})
|
||||
|
||||
resp, err := tool.Run(ctx, fantasy.ToolCall{ID: "call-1", Name: "start_workspace", Input: "{}"})
|
||||
require.NoError(t, err)
|
||||
require.Contains(t, resp.Content, "no workspace")
|
||||
})
|
||||
|
||||
t.Run("AlreadyRunning", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
db, _ := dbtestutil.NewDB(t)
|
||||
|
||||
user := dbgen.User(t, db, database.User{})
|
||||
modelCfg := seedModelConfig(ctx, t, db, user.ID)
|
||||
org := dbgen.Organization(t, db, database.Organization{})
|
||||
_ = dbgen.OrganizationMember(t, db, database.OrganizationMember{
|
||||
UserID: user.ID,
|
||||
OrganizationID: org.ID,
|
||||
})
|
||||
wsResp := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{
|
||||
OwnerID: user.ID,
|
||||
OrganizationID: org.ID,
|
||||
}).Seed(database.WorkspaceBuild{
|
||||
Transition: database.WorkspaceTransitionStart,
|
||||
}).Do()
|
||||
ws := wsResp.Workspace
|
||||
|
||||
chat, err := db.InsertChat(ctx, database.InsertChatParams{
|
||||
OwnerID: user.ID,
|
||||
WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true},
|
||||
LastModelConfigID: modelCfg.ID,
|
||||
Title: "test-already-running",
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
agentConnFn := func(_ context.Context, _ uuid.UUID) (workspacesdk.AgentConn, func(), error) {
|
||||
return nil, func() {}, nil
|
||||
}
|
||||
|
||||
tool := chattool.StartWorkspace(chattool.StartWorkspaceOptions{
|
||||
DB: db,
|
||||
OwnerID: user.ID,
|
||||
ChatID: chat.ID,
|
||||
AgentConnFn: agentConnFn,
|
||||
StartFn: func(_ context.Context, _ uuid.UUID, _ uuid.UUID, _ codersdk.CreateWorkspaceBuildRequest) (codersdk.WorkspaceBuild, error) {
|
||||
t.Fatal("StartFn should not be called for already-running workspace")
|
||||
return codersdk.WorkspaceBuild{}, nil
|
||||
},
|
||||
WorkspaceMu: &sync.Mutex{},
|
||||
})
|
||||
|
||||
resp, err := tool.Run(ctx, fantasy.ToolCall{ID: "call-1", Name: "start_workspace", Input: "{}"})
|
||||
require.NoError(t, err)
|
||||
|
||||
var result map[string]any
|
||||
require.NoError(t, json.Unmarshal([]byte(resp.Content), &result))
|
||||
started, ok := result["started"].(bool)
|
||||
require.True(t, ok)
|
||||
require.True(t, started)
|
||||
})
|
||||
|
||||
t.Run("StoppedWorkspace", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
db, _ := dbtestutil.NewDB(t)
|
||||
|
||||
user := dbgen.User(t, db, database.User{})
|
||||
modelCfg := seedModelConfig(ctx, t, db, user.ID)
|
||||
org := dbgen.Organization(t, db, database.Organization{})
|
||||
_ = dbgen.OrganizationMember(t, db, database.OrganizationMember{
|
||||
UserID: user.ID,
|
||||
OrganizationID: org.ID,
|
||||
})
|
||||
// Create a completed "stop" build so the workspace is stopped.
|
||||
wsResp := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{
|
||||
OwnerID: user.ID,
|
||||
OrganizationID: org.ID,
|
||||
}).Seed(database.WorkspaceBuild{
|
||||
Transition: database.WorkspaceTransitionStop,
|
||||
}).Do()
|
||||
ws := wsResp.Workspace
|
||||
|
||||
chat, err := db.InsertChat(ctx, database.InsertChatParams{
|
||||
OwnerID: user.ID,
|
||||
WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true},
|
||||
LastModelConfigID: modelCfg.ID,
|
||||
Title: "test-stopped-workspace",
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
var startCalled bool
|
||||
startFn := func(_ context.Context, _ uuid.UUID, wsID uuid.UUID, req codersdk.CreateWorkspaceBuildRequest) (codersdk.WorkspaceBuild, error) {
|
||||
startCalled = true
|
||||
require.Equal(t, codersdk.WorkspaceTransitionStart, req.Transition)
|
||||
require.Equal(t, ws.ID, wsID)
|
||||
|
||||
// Simulate start by inserting a new completed "start" build.
|
||||
dbfake.WorkspaceBuild(t, db, ws).Seed(database.WorkspaceBuild{
|
||||
Transition: database.WorkspaceTransitionStart,
|
||||
BuildNumber: 2,
|
||||
}).Do()
|
||||
return codersdk.WorkspaceBuild{}, nil
|
||||
}
|
||||
|
||||
agentConnFn := func(_ context.Context, _ uuid.UUID) (workspacesdk.AgentConn, func(), error) {
|
||||
return nil, func() {}, nil
|
||||
}
|
||||
|
||||
tool := chattool.StartWorkspace(chattool.StartWorkspaceOptions{
|
||||
DB: db,
|
||||
OwnerID: user.ID,
|
||||
ChatID: chat.ID,
|
||||
StartFn: startFn,
|
||||
AgentConnFn: agentConnFn,
|
||||
WorkspaceMu: &sync.Mutex{},
|
||||
})
|
||||
|
||||
resp, err := tool.Run(ctx, fantasy.ToolCall{ID: "call-1", Name: "start_workspace", Input: "{}"})
|
||||
require.NoError(t, err)
|
||||
require.True(t, startCalled, "expected StartFn to be called")
|
||||
|
||||
var result map[string]any
|
||||
require.NoError(t, json.Unmarshal([]byte(resp.Content), &result))
|
||||
started, ok := result["started"].(bool)
|
||||
require.True(t, ok)
|
||||
require.True(t, started)
|
||||
})
|
||||
}
|
||||
|
||||
// seedModelConfig inserts a provider and model config for testing.
|
||||
func seedModelConfig(
|
||||
ctx context.Context,
|
||||
t *testing.T,
|
||||
db database.Store,
|
||||
userID uuid.UUID,
|
||||
) database.ChatModelConfig {
|
||||
t.Helper()
|
||||
|
||||
_, err := db.InsertChatProvider(ctx, database.InsertChatProviderParams{
|
||||
Provider: "openai",
|
||||
DisplayName: "OpenAI",
|
||||
APIKey: "test-key",
|
||||
BaseUrl: "",
|
||||
ApiKeyKeyID: sql.NullString{},
|
||||
CreatedBy: uuid.NullUUID{UUID: userID, Valid: true},
|
||||
Enabled: true,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
model, err := db.InsertChatModelConfig(ctx, database.InsertChatModelConfigParams{
|
||||
Provider: "openai",
|
||||
Model: "gpt-4o-mini",
|
||||
DisplayName: "Test Model",
|
||||
CreatedBy: uuid.NullUUID{UUID: userID, Valid: true},
|
||||
UpdatedBy: uuid.NullUUID{UUID: userID, Valid: true},
|
||||
Enabled: true,
|
||||
IsDefault: true,
|
||||
ContextLimit: 128000,
|
||||
CompressionThreshold: 70,
|
||||
Options: json.RawMessage(`{}`),
|
||||
})
|
||||
require.NoError(t, err)
|
||||
return model
|
||||
}
|
||||
@@ -932,6 +932,51 @@ func (api *API) chatCreateWorkspace(
|
||||
return workspace, nil
|
||||
}
|
||||
|
||||
// chatStartWorkspace starts a stopped workspace by creating a new
|
||||
// build with the "start" transition. It mirrors chatCreateWorkspace
|
||||
// but for the start path.
|
||||
func (api *API) chatStartWorkspace(
|
||||
ctx context.Context,
|
||||
ownerID uuid.UUID,
|
||||
workspaceID uuid.UUID,
|
||||
req codersdk.CreateWorkspaceBuildRequest,
|
||||
) (codersdk.WorkspaceBuild, error) {
|
||||
actor, _, err := httpmw.UserRBACSubject(ctx, api.Database, ownerID, rbac.ScopeAll)
|
||||
if err != nil {
|
||||
return codersdk.WorkspaceBuild{}, xerrors.Errorf("load user authorization: %w", err)
|
||||
}
|
||||
ctx = dbauthz.As(ctx, actor)
|
||||
|
||||
workspace, err := api.Database.GetWorkspaceByID(ctx, workspaceID)
|
||||
if err != nil {
|
||||
return codersdk.WorkspaceBuild{}, xerrors.Errorf("get workspace: %w", err)
|
||||
}
|
||||
|
||||
// Build a synthetic API key so postWorkspaceBuildsInternal can
|
||||
// record the correct initiator.
|
||||
syntheticKey := database.APIKey{
|
||||
UserID: ownerID,
|
||||
}
|
||||
|
||||
apiBuild, err := api.postWorkspaceBuildsInternal(
|
||||
ctx,
|
||||
syntheticKey,
|
||||
workspace,
|
||||
req,
|
||||
func(action policy.Action, object rbac.Objecter) bool {
|
||||
// Authorization is handled by dbauthz on the context.
|
||||
authErr := api.HTTPAuth.Authorizer.Authorize(ctx, actor, action, object.RBACObject())
|
||||
return authErr == nil
|
||||
},
|
||||
audit.WorkspaceBuildBaggage{},
|
||||
)
|
||||
if err != nil {
|
||||
return codersdk.WorkspaceBuild{}, xerrors.Errorf("create workspace build: %w", err)
|
||||
}
|
||||
|
||||
return apiBuild, nil
|
||||
}
|
||||
|
||||
func chatWorkspaceAuditStatus(err error) int {
|
||||
if responder, ok := httperror.IsResponder(err); ok {
|
||||
status, _ := responder.Response()
|
||||
|
||||
@@ -767,6 +767,7 @@ func New(options *Options) *API {
|
||||
ProviderAPIKeys: chatProviderAPIKeysFromDeploymentValues(options.DeploymentValues),
|
||||
AgentConn: api.agentProvider.AgentConn,
|
||||
CreateWorkspace: api.chatCreateWorkspace,
|
||||
StartWorkspace: api.chatStartWorkspace,
|
||||
Pubsub: options.Pubsub,
|
||||
WebpushDispatcher: options.WebPushDispatcher,
|
||||
})
|
||||
|
||||
Reference in New Issue
Block a user