chore: move chatd and related packages to /x/ subpackage (#23445)

- Moves `coderd/chatd/`, `coderd/gitsync/`, `enterprise/coderd/chatd/`
under `x/` parent directories to signal instability
- Adds `Experimental:` glue code comments in `coderd/coderd.go`

> 🤖 This PR was created with the help of Coder Agents, and was
reviewed by my human. 🧑‍💻
This commit is contained in:
Cian Johnston
2026-03-23 17:34:43 +00:00
committed by GitHub
parent 86d8b6daee
commit 80a172f932
64 changed files with 92 additions and 90 deletions
+33
View File
@@ -0,0 +1,33 @@
package chattool
import (
"encoding/json"
"unicode/utf8"
"charm.land/fantasy"
)
// toolResponse builds a fantasy.ToolResponse from a JSON-serializable
// result payload.
func toolResponse(result map[string]any) fantasy.ToolResponse {
data, err := json.Marshal(result)
if err != nil {
return fantasy.NewTextResponse("{}")
}
return fantasy.NewTextResponse(string(data))
}
func truncateRunes(value string, maxLen int) string {
if maxLen <= 0 || value == "" {
return ""
}
if utf8.RuneCountInString(value) <= maxLen {
return value
}
runes := []rune(value)
if maxLen > len(runes) {
maxLen = len(runes)
}
return string(runes[:maxLen])
}
+211
View File
@@ -0,0 +1,211 @@
package chattool
import (
"context"
"fmt"
"time"
"charm.land/fantasy"
fantasyanthropic "charm.land/fantasy/providers/anthropic"
"github.com/coder/coder/v2/codersdk/workspacesdk"
"github.com/coder/quartz"
)
const (
// ComputerUseModelProvider is the provider for the computer
// use model.
ComputerUseModelProvider = "anthropic"
// ComputerUseModelName is the model used for computer use
// subagents.
ComputerUseModelName = "claude-opus-4-6"
)
// computerUseTool implements fantasy.AgentTool and
// chatloop.ToolDefiner for Anthropic computer use.
type computerUseTool struct {
declaredWidth int
declaredHeight int
getWorkspaceConn func(ctx context.Context) (workspacesdk.AgentConn, error)
providerOptions fantasy.ProviderOptions
clock quartz.Clock
}
// NewComputerUseTool creates a computer use AgentTool that delegates to the
// agent's desktop endpoints. declaredWidth and declaredHeight are the
// model-facing desktop dimensions advertised to Anthropic and requested for
// screenshots.
func NewComputerUseTool(
declaredWidth, declaredHeight int,
getWorkspaceConn func(ctx context.Context) (workspacesdk.AgentConn, error),
clock quartz.Clock,
) fantasy.AgentTool {
return &computerUseTool{
declaredWidth: declaredWidth,
declaredHeight: declaredHeight,
getWorkspaceConn: getWorkspaceConn,
clock: clock,
}
}
func (*computerUseTool) Info() fantasy.ToolInfo {
return fantasy.ToolInfo{
Name: "computer",
Description: "Control the desktop: take screenshots, move the mouse, click, type, and scroll.",
Parameters: map[string]any{},
Required: []string{},
}
}
// ComputerUseProviderTool creates the provider-defined Anthropic computer-use
// tool definition using the declared model-facing desktop geometry.
func ComputerUseProviderTool(declaredWidth, declaredHeight int) fantasy.Tool {
return fantasyanthropic.NewComputerUseTool(
fantasyanthropic.ComputerUseToolOptions{
DisplayWidthPx: int64(declaredWidth),
DisplayHeightPx: int64(declaredHeight),
ToolVersion: fantasyanthropic.ComputerUse20251124,
},
)
}
func (t *computerUseTool) ProviderOptions() fantasy.ProviderOptions {
return t.providerOptions
}
func (t *computerUseTool) SetProviderOptions(opts fantasy.ProviderOptions) {
t.providerOptions = opts
}
func (t *computerUseTool) Run(ctx context.Context, call fantasy.ToolCall) (fantasy.ToolResponse, error) {
input, err := fantasyanthropic.ParseComputerUseInput(call.Input)
if err != nil {
return fantasy.NewTextErrorResponse(
fmt.Sprintf("invalid computer use input: %v", err),
), nil
}
conn, err := t.getWorkspaceConn(ctx)
if err != nil {
return fantasy.NewTextErrorResponse(
fmt.Sprintf("failed to connect to workspace: %v", err),
), nil
}
declaredWidth, declaredHeight := t.declaredActionDimensions()
// For wait actions, sleep then return a screenshot.
if input.Action == fantasyanthropic.ActionWait {
d := input.Duration
if d <= 0 {
d = 1000
}
timer := t.clock.NewTimer(time.Duration(d)*time.Millisecond, "computeruse", "wait")
defer timer.Stop()
select {
case <-ctx.Done():
case <-timer.C:
}
screenshotAction := workspacesdk.DesktopAction{
Action: "screenshot",
ScaledWidth: &declaredWidth,
ScaledHeight: &declaredHeight,
}
screenResp, sErr := conn.ExecuteDesktopAction(ctx, screenshotAction)
if sErr != nil {
return fantasy.NewTextErrorResponse(
fmt.Sprintf("screenshot failed: %v", sErr),
), nil
}
return fantasy.NewImageResponse(
[]byte(screenResp.ScreenshotData), "image/png",
), nil
}
// For screenshot action, use ExecuteDesktopAction.
if input.Action == fantasyanthropic.ActionScreenshot {
screenshotAction := workspacesdk.DesktopAction{
Action: "screenshot",
ScaledWidth: &declaredWidth,
ScaledHeight: &declaredHeight,
}
screenResp, sErr := conn.ExecuteDesktopAction(ctx, screenshotAction)
if sErr != nil {
return fantasy.NewTextErrorResponse(
fmt.Sprintf("screenshot failed: %v", sErr),
), nil
}
return fantasy.NewImageResponse(
[]byte(screenResp.ScreenshotData), "image/png",
), nil
}
// Build the action request.
action := workspacesdk.DesktopAction{
Action: string(input.Action),
ScaledWidth: &declaredWidth,
ScaledHeight: &declaredHeight,
}
if input.Coordinate != ([2]int64{}) {
coord := [2]int{int(input.Coordinate[0]), int(input.Coordinate[1])}
action.Coordinate = &coord
}
if input.StartCoordinate != ([2]int64{}) {
coord := [2]int{int(input.StartCoordinate[0]), int(input.StartCoordinate[1])}
action.StartCoordinate = &coord
}
if input.Text != "" {
action.Text = &input.Text
}
if input.Duration > 0 {
d := int(input.Duration)
action.Duration = &d
}
if input.ScrollAmount > 0 {
s := int(input.ScrollAmount)
action.ScrollAmount = &s
}
if input.ScrollDirection != "" {
action.ScrollDirection = &input.ScrollDirection
}
// Execute the action.
_, err = conn.ExecuteDesktopAction(ctx, action)
if err != nil {
return fantasy.NewTextErrorResponse(
fmt.Sprintf("action %q failed: %v", input.Action, err),
), nil
}
// Take a screenshot after every action (Anthropic pattern).
screenshotAction := workspacesdk.DesktopAction{
Action: "screenshot",
ScaledWidth: &declaredWidth,
ScaledHeight: &declaredHeight,
}
screenResp, sErr := conn.ExecuteDesktopAction(ctx, screenshotAction)
if sErr != nil {
return fantasy.NewTextErrorResponse(
fmt.Sprintf("screenshot failed: %v", sErr),
), nil
}
return fantasy.NewImageResponse(
[]byte(screenResp.ScreenshotData), "image/png",
), nil
}
func (t *computerUseTool) declaredActionDimensions() (declaredWidth, declaredHeight int) {
if t.declaredWidth <= 0 || t.declaredHeight <= 0 {
geometry := workspacesdk.DefaultDesktopGeometry()
return geometry.DeclaredWidth, geometry.DeclaredHeight
}
return t.declaredWidth, t.declaredHeight
}
// computeScaledScreenshotSize preserves the historical helper name while using
// the shared declared-geometry selection logic.
func computeScaledScreenshotSize(width, height int) (scaledWidth int, scaledHeight int) {
geometry := workspacesdk.NewDesktopGeometry(width, height)
return geometry.DeclaredWidth, geometry.DeclaredHeight
}
@@ -0,0 +1,85 @@
package chattool
import (
"testing"
"github.com/stretchr/testify/assert"
)
func TestComputeScaledScreenshotSize(t *testing.T) {
t.Parallel()
tests := []struct {
name string
width, height int
wantW, wantH int
}{
{
name: "1920x1080_prefers_standard_1280x720",
width: 1920,
height: 1080,
wantW: 1280,
wantH: 720,
},
{
name: "1280x800_no_scaling",
width: 1280,
height: 800,
wantW: 1280,
wantH: 800,
},
{
name: "3840x2160_prefers_standard_1280x720",
width: 3840,
height: 2160,
wantW: 1280,
wantH: 720,
},
{
name: "1568x1000_prefers_standard_1280x816",
width: 1568,
height: 1000,
wantW: 1280,
wantH: 816,
},
{
name: "100x100_small_display",
width: 100,
height: 100,
wantW: 100,
wantH: 100,
},
{
name: "4000x3000_prefers_standard_1024x768",
width: 4000,
height: 3000,
wantW: 1024,
wantH: 768,
},
{
name: "1920x1200_prefers_standard_1280x800",
width: 1920,
height: 1200,
wantW: 1280,
wantH: 800,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
t.Parallel()
gotW, gotH := computeScaledScreenshotSize(tt.width, tt.height)
assert.Equal(t, tt.wantW, gotW)
assert.Equal(t, tt.wantH, gotH)
// Invariant: results must respect Anthropic constraints.
const maxLongEdge = 1568
const maxTotalPixels = 1_150_000
longEdge := max(gotW, gotH)
assert.LessOrEqual(t, longEdge, maxLongEdge,
"long edge %d exceeds max %d", longEdge, maxLongEdge)
assert.LessOrEqual(t, gotW*gotH, maxTotalPixels,
"total pixels %d exceeds max %d", gotW*gotH, maxTotalPixels)
})
}
}
+226
View File
@@ -0,0 +1,226 @@
package chattool_test
import (
"context"
"testing"
"charm.land/fantasy"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"go.uber.org/mock/gomock"
"golang.org/x/xerrors"
"github.com/coder/coder/v2/coderd/x/chatd/chattool"
"github.com/coder/coder/v2/codersdk/workspacesdk"
"github.com/coder/coder/v2/codersdk/workspacesdk/agentconnmock"
"github.com/coder/quartz"
)
func TestComputerUseTool_Info(t *testing.T) {
t.Parallel()
geometry := workspacesdk.DefaultDesktopGeometry()
tool := chattool.NewComputerUseTool(geometry.DeclaredWidth, geometry.DeclaredHeight, nil, quartz.NewReal())
info := tool.Info()
assert.Equal(t, "computer", info.Name)
assert.NotEmpty(t, info.Description)
}
func TestComputerUseProviderTool(t *testing.T) {
t.Parallel()
geometry := workspacesdk.DefaultDesktopGeometry()
def := chattool.ComputerUseProviderTool(geometry.DeclaredWidth, geometry.DeclaredHeight)
pdt, ok := def.(fantasy.ProviderDefinedTool)
require.True(t, ok, "ComputerUseProviderTool should return a ProviderDefinedTool")
assert.Contains(t, pdt.ID, "computer")
assert.Equal(t, "computer", pdt.Name)
assert.Equal(t, int64(geometry.DeclaredWidth), pdt.Args["display_width_px"])
assert.Equal(t, int64(geometry.DeclaredHeight), pdt.Args["display_height_px"])
}
func TestComputerUseProviderTool_PrefersDeclaredGeometry(t *testing.T) {
t.Parallel()
geometry := workspacesdk.NewDesktopGeometry(1920, 1080)
def := chattool.ComputerUseProviderTool(geometry.DeclaredWidth, geometry.DeclaredHeight)
pdt, ok := def.(fantasy.ProviderDefinedTool)
require.True(t, ok, "ComputerUseProviderTool should return a ProviderDefinedTool")
assert.Equal(t, int64(1280), pdt.Args["display_width_px"])
assert.Equal(t, int64(720), pdt.Args["display_height_px"])
}
func TestComputerUseTool_Run_Screenshot(t *testing.T) {
t.Parallel()
ctrl := gomock.NewController(t)
mockConn := agentconnmock.NewMockAgentConn(ctrl)
geometry := workspacesdk.DefaultDesktopGeometry()
mockConn.EXPECT().ExecuteDesktopAction(
gomock.Any(),
gomock.AssignableToTypeOf(workspacesdk.DesktopAction{}),
).DoAndReturn(func(_ context.Context, action workspacesdk.DesktopAction) (workspacesdk.DesktopActionResponse, error) {
require.NotNil(t, action.ScaledWidth)
require.NotNil(t, action.ScaledHeight)
assert.Equal(t, geometry.DeclaredWidth, *action.ScaledWidth)
assert.Equal(t, geometry.DeclaredHeight, *action.ScaledHeight)
return workspacesdk.DesktopActionResponse{
Output: "screenshot",
ScreenshotData: "base64png",
ScreenshotWidth: geometry.DeclaredWidth,
ScreenshotHeight: geometry.DeclaredHeight,
}, nil
})
tool := chattool.NewComputerUseTool(geometry.DeclaredWidth, geometry.DeclaredHeight, func(_ context.Context) (workspacesdk.AgentConn, error) {
return mockConn, nil
}, quartz.NewReal())
call := fantasy.ToolCall{
ID: "test-1",
Name: "computer",
Input: `{"action":"screenshot"}`,
}
resp, err := tool.Run(context.Background(), call)
require.NoError(t, err)
assert.Equal(t, "image", resp.Type)
assert.Equal(t, "image/png", resp.MediaType)
assert.Equal(t, []byte("base64png"), resp.Data)
assert.False(t, resp.IsError)
}
func TestComputerUseTool_Run_LeftClick(t *testing.T) {
t.Parallel()
ctrl := gomock.NewController(t)
mockConn := agentconnmock.NewMockAgentConn(ctrl)
geometry := workspacesdk.DefaultDesktopGeometry()
mockConn.EXPECT().ExecuteDesktopAction(
gomock.Any(),
gomock.AssignableToTypeOf(workspacesdk.DesktopAction{}),
).DoAndReturn(func(_ context.Context, action workspacesdk.DesktopAction) (workspacesdk.DesktopActionResponse, error) {
require.NotNil(t, action.Coordinate)
assert.Equal(t, [2]int{100, 200}, *action.Coordinate)
require.NotNil(t, action.ScaledWidth)
require.NotNil(t, action.ScaledHeight)
assert.Equal(t, geometry.DeclaredWidth, *action.ScaledWidth)
assert.Equal(t, geometry.DeclaredHeight, *action.ScaledHeight)
return workspacesdk.DesktopActionResponse{Output: "left_click performed"}, nil
})
mockConn.EXPECT().ExecuteDesktopAction(
gomock.Any(),
gomock.AssignableToTypeOf(workspacesdk.DesktopAction{}),
).DoAndReturn(func(_ context.Context, action workspacesdk.DesktopAction) (workspacesdk.DesktopActionResponse, error) {
assert.Equal(t, "screenshot", action.Action)
require.NotNil(t, action.ScaledWidth)
require.NotNil(t, action.ScaledHeight)
assert.Equal(t, geometry.DeclaredWidth, *action.ScaledWidth)
assert.Equal(t, geometry.DeclaredHeight, *action.ScaledHeight)
return workspacesdk.DesktopActionResponse{
Output: "screenshot",
ScreenshotData: "after-click",
ScreenshotWidth: geometry.DeclaredWidth,
ScreenshotHeight: geometry.DeclaredHeight,
}, nil
})
tool := chattool.NewComputerUseTool(geometry.DeclaredWidth, geometry.DeclaredHeight, func(_ context.Context) (workspacesdk.AgentConn, error) {
return mockConn, nil
}, quartz.NewReal())
call := fantasy.ToolCall{
ID: "test-2",
Name: "computer",
Input: `{"action":"left_click","coordinate":[100,200]}`,
}
resp, err := tool.Run(context.Background(), call)
require.NoError(t, err)
assert.Equal(t, "image", resp.Type)
assert.Equal(t, []byte("after-click"), resp.Data)
}
func TestComputerUseTool_Run_Wait(t *testing.T) {
t.Parallel()
ctrl := gomock.NewController(t)
mockConn := agentconnmock.NewMockAgentConn(ctrl)
geometry := workspacesdk.DefaultDesktopGeometry()
mockConn.EXPECT().ExecuteDesktopAction(
gomock.Any(),
gomock.AssignableToTypeOf(workspacesdk.DesktopAction{}),
).DoAndReturn(func(_ context.Context, action workspacesdk.DesktopAction) (workspacesdk.DesktopActionResponse, error) {
require.NotNil(t, action.ScaledWidth)
require.NotNil(t, action.ScaledHeight)
assert.Equal(t, geometry.DeclaredWidth, *action.ScaledWidth)
assert.Equal(t, geometry.DeclaredHeight, *action.ScaledHeight)
return workspacesdk.DesktopActionResponse{
Output: "screenshot",
ScreenshotData: "after-wait",
ScreenshotWidth: geometry.DeclaredWidth,
ScreenshotHeight: geometry.DeclaredHeight,
}, nil
})
tool := chattool.NewComputerUseTool(geometry.DeclaredWidth, geometry.DeclaredHeight, func(_ context.Context) (workspacesdk.AgentConn, error) {
return mockConn, nil
}, quartz.NewReal())
call := fantasy.ToolCall{
ID: "test-3",
Name: "computer",
Input: `{"action":"wait","duration":10}`,
}
resp, err := tool.Run(context.Background(), call)
require.NoError(t, err)
assert.Equal(t, "image", resp.Type)
assert.Equal(t, "image/png", resp.MediaType)
assert.Equal(t, []byte("after-wait"), resp.Data)
assert.False(t, resp.IsError)
}
func TestComputerUseTool_Run_ConnError(t *testing.T) {
t.Parallel()
geometry := workspacesdk.DefaultDesktopGeometry()
tool := chattool.NewComputerUseTool(geometry.DeclaredWidth, geometry.DeclaredHeight, func(_ context.Context) (workspacesdk.AgentConn, error) {
return nil, xerrors.New("workspace not available")
}, quartz.NewReal())
call := fantasy.ToolCall{
ID: "test-4",
Name: "computer",
Input: `{"action":"screenshot"}`,
}
resp, err := tool.Run(context.Background(), call)
require.NoError(t, err)
assert.True(t, resp.IsError)
assert.Contains(t, resp.Content, "workspace not available")
}
func TestComputerUseTool_Run_InvalidInput(t *testing.T) {
t.Parallel()
geometry := workspacesdk.DefaultDesktopGeometry()
tool := chattool.NewComputerUseTool(geometry.DeclaredWidth, geometry.DeclaredHeight, func(_ context.Context) (workspacesdk.AgentConn, error) {
return nil, xerrors.New("should not be called")
}, quartz.NewReal())
call := fantasy.ToolCall{
ID: "test-5",
Name: "computer",
Input: `{invalid json`,
}
resp, err := tool.Run(context.Background(), call)
require.NoError(t, err)
assert.True(t, resp.IsError)
assert.Contains(t, resp.Content, "invalid computer use input")
}
+523
View File
@@ -0,0 +1,523 @@
package chattool
import (
"context"
"errors"
"fmt"
"strings"
"sync"
"time"
"charm.land/fantasy"
"github.com/google/uuid"
"golang.org/x/xerrors"
"cdr.dev/slog/v3"
"github.com/coder/coder/v2/coderd/database"
"github.com/coder/coder/v2/coderd/util/namesgenerator"
"github.com/coder/coder/v2/codersdk"
"github.com/coder/coder/v2/codersdk/workspacesdk"
)
const (
// buildPollInterval is how often we check if the workspace
// build has completed.
buildPollInterval = 2 * time.Second
// buildTimeout is the maximum time to wait for a workspace
// build to complete before giving up.
buildTimeout = 10 * time.Minute
// agentConnectTimeout is the maximum time to wait for the
// workspace agent to become reachable after a successful build.
agentConnectTimeout = 2 * time.Minute
// agentRetryInterval is how often we retry connecting to the
// workspace agent.
agentRetryInterval = 2 * time.Second
// agentAttemptTimeout is the timeout for a single connection
// attempt to the workspace agent during the retry loop.
agentAttemptTimeout = 5 * time.Second
// agentPingTimeout is the timeout for a single agent ping
// when checking whether an existing workspace is alive.
agentPingTimeout = 5 * time.Second
// startupScriptTimeout is the maximum time to wait for the
// workspace agent's startup scripts to finish after the agent
// is reachable.
startupScriptTimeout = 10 * time.Minute
// startupScriptPollInterval is how often we check the agent's
// lifecycle state while waiting for startup scripts.
startupScriptPollInterval = 2 * time.Second
)
// CreateWorkspaceFn creates a workspace for the given owner.
type CreateWorkspaceFn func(
ctx context.Context,
ownerID uuid.UUID,
req codersdk.CreateWorkspaceRequest,
) (codersdk.Workspace, error)
// AgentConnFunc provides access to workspace agent connections.
type AgentConnFunc func(
ctx context.Context,
agentID uuid.UUID,
) (workspacesdk.AgentConn, func(), error)
// CreateWorkspaceOptions configures the create_workspace tool.
type CreateWorkspaceOptions struct {
DB database.Store
OwnerID uuid.UUID
ChatID uuid.UUID
CreateFn CreateWorkspaceFn
AgentConnFn AgentConnFunc
WorkspaceMu *sync.Mutex
Logger slog.Logger
}
type createWorkspaceArgs struct {
TemplateID string `json:"template_id"`
Name string `json:"name,omitempty"`
Parameters map[string]string `json:"parameters,omitempty"`
}
// CreateWorkspace returns a tool that creates a new workspace from a
// template. The tool is idempotent: if the chat already has a
// workspace that is building or running, it returns the existing
// workspace instead of creating a new one. A mutex prevents parallel
// calls from creating duplicate workspaces.
func CreateWorkspace(options CreateWorkspaceOptions) fantasy.AgentTool {
return fantasy.NewAgentTool(
"create_workspace",
"Create a new workspace from a template. Requires a "+
"template_id (from list_templates). Optionally provide "+
"a name and parameter values (from read_template). "+
"If no name is given, one will be generated. "+
"This tool is idempotent — if the chat already has a "+
"workspace that is building or running, the existing "+
"workspace is returned.",
func(ctx context.Context, args createWorkspaceArgs, _ fantasy.ToolCall) (fantasy.ToolResponse, error) {
if options.CreateFn == nil {
return fantasy.NewTextErrorResponse("workspace creator is not configured"), nil
}
templateIDStr := strings.TrimSpace(args.TemplateID)
if templateIDStr == "" {
return fantasy.NewTextErrorResponse("template_id is required; use list_templates to find one"), nil
}
templateID, err := uuid.Parse(templateIDStr)
if err != nil {
return fantasy.NewTextErrorResponse(
xerrors.Errorf("invalid template_id: %w", err).Error(),
), nil
}
// Serialize workspace creation to prevent parallel
// tool calls from creating duplicate workspaces.
if options.WorkspaceMu != nil {
options.WorkspaceMu.Lock()
defer options.WorkspaceMu.Unlock()
}
// Check for an existing workspace on the chat.
if options.DB != nil && options.ChatID != uuid.Nil {
existing, done, existErr := checkExistingWorkspace(
ctx, options.DB, options.ChatID,
options.AgentConnFn,
)
if existErr != nil {
return fantasy.NewTextErrorResponse(existErr.Error()), nil
}
if done {
return toolResponse(existing), nil
}
}
ownerID := options.OwnerID
// Set up dbauthz context for DB lookups.
if options.DB != nil {
ownerCtx, ownerErr := asOwner(ctx, options.DB, ownerID)
if ownerErr != nil {
return fantasy.NewTextErrorResponse(ownerErr.Error()), nil
}
ctx = ownerCtx
}
var ttlMs *int64
if options.DB != nil {
raw, err := options.DB.GetChatWorkspaceTTL(ctx)
if err != nil {
options.Logger.Error(ctx, "failed to read chat workspace TTL setting, using template default",
slog.Error(err),
)
} else {
d, parseErr := codersdk.ParseChatWorkspaceTTL(raw)
if parseErr != nil {
options.Logger.Warn(ctx, "invalid chat workspace TTL setting, using template default",
slog.F("raw", raw),
slog.Error(parseErr),
)
} else if d > 0 {
ms := d.Milliseconds()
ttlMs = &ms
}
}
}
createReq := codersdk.CreateWorkspaceRequest{
TemplateID: templateID,
TTLMillis: ttlMs,
}
// Resolve workspace name.
name := strings.TrimSpace(args.Name)
if name == "" {
seed := "workspace"
if options.DB != nil {
if t, lookupErr := options.DB.GetTemplateByID(ctx, templateID); lookupErr == nil {
seed = t.Name
}
}
name = generatedWorkspaceName(seed)
} else if err := codersdk.NameValid(name); err != nil {
name = generatedWorkspaceName(name)
}
createReq.Name = name
// Map parameters.
for k, v := range args.Parameters {
createReq.RichParameterValues = append(
createReq.RichParameterValues,
codersdk.WorkspaceBuildParameter{Name: k, Value: v},
)
}
workspace, err := options.CreateFn(ctx, ownerID, createReq)
if err != nil {
return fantasy.NewTextErrorResponse(err.Error()), nil
}
// Wait for the build to complete and the agent to
// come online so subsequent tools can use the
// workspace immediately.
if options.DB != nil {
if err := waitForBuild(ctx, options.DB, workspace.ID); err != nil {
return fantasy.NewTextErrorResponse(
xerrors.Errorf("workspace build failed: %w", err).Error(),
), nil
}
}
// Look up the first agent so we can link it to the chat.
workspaceAgentID := uuid.Nil
if options.DB != nil {
agents, agentErr := options.DB.GetWorkspaceAgentsInLatestBuildByWorkspaceID(ctx, workspace.ID)
if agentErr == nil && len(agents) > 0 {
workspaceAgentID = agents[0].ID
}
}
// Persist workspace + agent association on the chat.
if options.DB != nil && options.ChatID != uuid.Nil {
if _, err := options.DB.UpdateChatWorkspace(ctx, database.UpdateChatWorkspaceParams{
ID: options.ChatID,
WorkspaceID: uuid.NullUUID{
UUID: workspace.ID,
Valid: true,
},
}); err != nil {
options.Logger.Error(ctx, "failed to persist chat workspace association",
slog.F("chat_id", options.ChatID),
slog.F("workspace_id", workspace.ID),
slog.Error(err),
)
}
}
// Wait for the agent to come online and startup scripts to finish.
if workspaceAgentID != uuid.Nil {
agentStatus := waitForAgentReady(ctx, options.DB, workspaceAgentID, options.AgentConnFn)
result := map[string]any{
"created": true,
"workspace_name": workspace.FullName(),
}
for k, v := range agentStatus {
result[k] = v
}
return toolResponse(result), nil
}
return toolResponse(map[string]any{
"created": true,
"workspace_name": workspace.FullName(),
}), nil
})
}
// checkExistingWorkspace checks whether the chat already has a usable
// workspace. Returns the result map and true if the caller should
// return early (workspace exists and is alive or building). Returns
// false if the caller should proceed with creation (workspace is dead
// or missing).
func checkExistingWorkspace(
ctx context.Context,
db database.Store,
chatID uuid.UUID,
agentConnFn AgentConnFunc,
) (map[string]any, bool, error) {
chat, err := db.GetChatByID(ctx, chatID)
if err != nil {
return nil, false, xerrors.Errorf("load chat: %w", err)
}
if !chat.WorkspaceID.Valid {
return nil, false, nil
}
ws, err := db.GetWorkspaceByID(ctx, chat.WorkspaceID.UUID)
if err != nil {
return nil, false, xerrors.Errorf("load workspace: %w", err)
}
// Workspace was soft-deleted — allow creation.
if ws.Deleted {
return nil, false, nil
}
// Check the latest build status.
build, err := db.GetLatestWorkspaceBuildByWorkspaceID(ctx, ws.ID)
if err != nil {
// Can't determine status — allow creation.
return nil, false, nil
}
job, err := db.GetProvisionerJobByID(ctx, build.JobID)
if err != nil {
return nil, false, nil
}
switch job.JobStatus {
case database.ProvisionerJobStatusPending,
database.ProvisionerJobStatusRunning:
// Build is in progress — wait for it instead of
// creating a new workspace.
if err := waitForBuild(ctx, db, ws.ID); err != nil {
return nil, false, xerrors.Errorf(
"existing workspace build failed: %w", err,
)
}
result := map[string]any{
"created": false,
"workspace_name": ws.Name,
"status": "already_exists",
"message": "workspace build completed",
}
agents, agentsErr := db.GetWorkspaceAgentsInLatestBuildByWorkspaceID(ctx, ws.ID)
if agentsErr == nil && len(agents) > 0 {
for k, v := range waitForAgentReady(ctx, db, agents[0].ID, agentConnFn) {
result[k] = v
}
}
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 {
pingCtx, cancel := context.WithTimeout(ctx, agentPingTimeout)
conn, release, connErr := agentConnFn(pingCtx, agents[0].ID)
cancel()
if connErr == nil {
release()
_ = conn
// Agent is reachable; wait for startup scripts.
result := map[string]any{
"created": false,
"workspace_name": ws.Name,
"status": "already_exists",
"message": "workspace is already running and reachable",
}
// Pass nil for agentConnFn since we already confirmed connectivity.
for k, v := range waitForAgentReady(ctx, db, agents[0].ID, nil) {
result[k] = v
}
return result, true, nil
}
// Agent unreachable — workspace is dead, allow
// creation.
}
// No agent ID or no conn func — allow creation.
return nil, false, nil
default:
// Failed, canceled, etc — allow creation.
return nil, false, nil
}
}
// waitForBuild polls the workspace's latest build until it
// completes or the context expires.
func waitForBuild(
ctx context.Context,
db database.Store,
workspaceID uuid.UUID,
) error {
buildCtx, cancel := context.WithTimeout(ctx, buildTimeout)
defer cancel()
ticker := time.NewTicker(buildPollInterval)
defer ticker.Stop()
for {
build, err := db.GetLatestWorkspaceBuildByWorkspaceID(
buildCtx, workspaceID,
)
if err != nil {
return xerrors.Errorf("get latest build: %w", err)
}
job, err := db.GetProvisionerJobByID(buildCtx, build.JobID)
if err != nil {
return xerrors.Errorf("get provisioner job: %w", err)
}
switch job.JobStatus {
case database.ProvisionerJobStatusSucceeded:
return nil
case database.ProvisionerJobStatusFailed:
errMsg := "build failed"
if job.Error.Valid {
errMsg = job.Error.String
}
return xerrors.New(errMsg)
case database.ProvisionerJobStatusCanceled:
return xerrors.New("build was canceled")
case database.ProvisionerJobStatusPending,
database.ProvisionerJobStatusRunning,
database.ProvisionerJobStatusCanceling:
// Still in progress — keep waiting.
default:
return xerrors.Errorf("unexpected job status: %s", job.JobStatus)
}
select {
case <-buildCtx.Done():
return xerrors.Errorf(
"timed out waiting for workspace build: %w",
buildCtx.Err(),
)
case <-ticker.C:
}
}
}
// waitForAgentReady waits for the workspace agent to become
// reachable and for its startup scripts to finish. It returns
// status fields suitable for merging into a tool response.
func waitForAgentReady(
ctx context.Context,
db database.Store,
agentID uuid.UUID,
agentConnFn AgentConnFunc,
) map[string]any {
result := map[string]any{}
// Phase 1: retry connecting to the agent.
if agentConnFn != nil {
agentCtx, agentCancel := context.WithTimeout(ctx, agentConnectTimeout)
defer agentCancel()
ticker := time.NewTicker(agentRetryInterval)
defer ticker.Stop()
var lastErr error
for {
attemptCtx, attemptCancel := context.WithTimeout(agentCtx, agentAttemptTimeout)
conn, release, err := agentConnFn(attemptCtx, agentID)
attemptCancel()
if err == nil {
release()
_ = conn
break
}
lastErr = err
select {
case <-agentCtx.Done():
result["agent_status"] = "not_ready"
result["agent_error"] = lastErr.Error()
return result
case <-ticker.C:
}
}
}
// Phase 2: poll lifecycle until startup scripts finish.
if db != nil {
scriptCtx, scriptCancel := context.WithTimeout(ctx, startupScriptTimeout)
defer scriptCancel()
ticker := time.NewTicker(startupScriptPollInterval)
defer ticker.Stop()
var lastState database.WorkspaceAgentLifecycleState
for {
row, err := db.GetWorkspaceAgentLifecycleStateByID(scriptCtx, agentID)
if err == nil {
lastState = row.LifecycleState
switch lastState {
case database.WorkspaceAgentLifecycleStateCreated,
database.WorkspaceAgentLifecycleStateStarting:
// Still in progress, keep polling.
case database.WorkspaceAgentLifecycleStateReady:
return result
default:
// Terminal non-ready state.
result["startup_scripts"] = "startup_scripts_failed"
result["lifecycle_state"] = string(lastState)
return result
}
}
select {
case <-scriptCtx.Done():
if errors.Is(scriptCtx.Err(), context.DeadlineExceeded) {
result["startup_scripts"] = "startup_scripts_timeout"
} else {
result["startup_scripts"] = "startup_scripts_unknown"
}
return result
case <-ticker.C:
}
}
}
return result
}
func generatedWorkspaceName(seed string) string {
base := codersdk.UsernameFrom(strings.TrimSpace(strings.ToLower(seed)))
if strings.TrimSpace(base) == "" {
base = "workspace"
}
suffix := strings.ReplaceAll(uuid.NewString(), "-", "")[:4]
if len(base) > 27 {
base = strings.Trim(base[:27], "-")
}
if base == "" {
base = "workspace"
}
name := fmt.Sprintf("%s-%s", base, suffix)
if err := codersdk.NameValid(name); err == nil {
return name
}
return namesgenerator.NameDigitWith("-")
}
@@ -0,0 +1,261 @@
package chattool //nolint:testpackage // Uses internal symbols.
import (
"context"
"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_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()
templateID := uuid.New()
workspaceID := uuid.New()
jobID := uuid.New()
db.EXPECT().
GetAuthorizationUserRoles(gomock.Any(), ownerID).
Return(database.GetAuthorizationUserRolesRow{
ID: ownerID,
Roles: []string{},
Groups: []string{},
Status: database.UserStatusActive,
}, nil)
db.EXPECT().
GetChatWorkspaceTTL(gomock.Any()).
Return(tc.ttlReturn, tc.ttlErr)
db.EXPECT().
GetLatestWorkspaceBuildByWorkspaceID(gomock.Any(), workspaceID).
Return(database.WorkspaceBuild{
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",
}, nil
}
tool := CreateWorkspace(CreateWorkspaceOptions{
DB: db,
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)
if tc.wantTTLMs != nil {
require.NotNil(t, capturedReq.TTLMillis)
require.Equal(t, *tc.wantTTLMs, *capturedReq.TTLMillis)
} else {
require.Nil(t, capturedReq.TTLMillis)
}
})
}
}
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)
result, done, err := checkExistingWorkspace(
context.Background(), db, chatID, nil,
)
require.NoError(t, err)
require.False(t, done, "should allow creation for deleted workspace")
require.Nil(t, result)
}
+50
View File
@@ -0,0 +1,50 @@
package chattool
import (
"context"
"charm.land/fantasy"
"github.com/coder/coder/v2/codersdk/workspacesdk"
)
type EditFilesOptions struct {
GetWorkspaceConn func(context.Context) (workspacesdk.AgentConn, error)
}
type EditFilesArgs struct {
Files []workspacesdk.FileEdits `json:"files"`
}
func EditFiles(options EditFilesOptions) fantasy.AgentTool {
return fantasy.NewAgentTool(
"edit_files",
"Perform search-and-replace edits on one or more files in the workspace."+
" Each file can have multiple edits applied atomically.",
func(ctx context.Context, args EditFilesArgs, _ fantasy.ToolCall) (fantasy.ToolResponse, error) {
if options.GetWorkspaceConn == nil {
return fantasy.NewTextErrorResponse("workspace connection resolver is not configured"), nil
}
conn, err := options.GetWorkspaceConn(ctx)
if err != nil {
return fantasy.NewTextErrorResponse(err.Error()), nil
}
return executeEditFilesTool(ctx, conn, args)
},
)
}
func executeEditFilesTool(
ctx context.Context,
conn workspacesdk.AgentConn,
args EditFilesArgs,
) (fantasy.ToolResponse, error) {
if len(args.Files) == 0 {
return fantasy.NewTextErrorResponse("files is required"), nil
}
if err := conn.EditFiles(ctx, workspacesdk.FileEditRequest{Files: args.Files}); err != nil {
return fantasy.NewTextErrorResponse(err.Error()), nil
}
return toolResponse(map[string]any{"ok": true}), nil
}
+561
View File
@@ -0,0 +1,561 @@
package chattool
import (
"context"
"encoding/json"
"fmt"
"regexp"
"strings"
"time"
"charm.land/fantasy"
"github.com/coder/coder/v2/codersdk/workspacesdk"
)
const (
// defaultTimeout is the default timeout for command
// execution.
defaultTimeout = 10 * time.Second
// maxOutputToModel is the maximum output sent to the LLM.
maxOutputToModel = 32 << 10 // 32KB
// snapshotTimeout is how long a non-blocking fallback
// request is allowed to take when retrieving a process
// output snapshot after a blocking wait times out.
snapshotTimeout = 30 * time.Second
)
// nonInteractiveEnvVars are set on every process to prevent
// interactive prompts that would hang a headless execution.
var nonInteractiveEnvVars = map[string]string{
"GIT_EDITOR": "true",
"GIT_SEQUENCE_EDITOR": "true",
"EDITOR": "true",
"VISUAL": "true",
"GIT_TERMINAL_PROMPT": "0",
"NO_COLOR": "1",
"TERM": "dumb",
"PAGER": "cat",
"GIT_PAGER": "cat",
}
// fileDumpPatterns detects commands that dump entire files.
// When matched, a note is added suggesting read_file instead.
var fileDumpPatterns = []*regexp.Regexp{
regexp.MustCompile(`^cat\s+`),
regexp.MustCompile(`^(rg|grep)\s+.*--include-all`),
regexp.MustCompile(`^(rg|grep)\s+-l\s+`),
}
// ExecuteResult is the structured response from the execute
// tool.
type ExecuteResult struct {
Success bool `json:"success"`
Output string `json:"output,omitempty"`
ExitCode int `json:"exit_code"`
WallDurationMs int64 `json:"wall_duration_ms"`
Error string `json:"error,omitempty"`
Truncated *workspacesdk.ProcessTruncation `json:"truncated,omitempty"`
Note string `json:"note,omitempty"`
BackgroundProcessID string `json:"background_process_id,omitempty"`
}
// ExecuteOptions configures the execute tool.
type ExecuteOptions struct {
GetWorkspaceConn func(context.Context) (workspacesdk.AgentConn, error)
DefaultTimeout time.Duration
}
// ProcessToolOptions configures a process management tool
// (process_output, process_list, or process_signal). Each of
// these tools only needs a workspace connection resolver.
type ProcessToolOptions struct {
GetWorkspaceConn func(context.Context) (workspacesdk.AgentConn, error)
}
// ExecuteArgs are the parameters accepted by the execute tool.
type ExecuteArgs struct {
Command string `json:"command" description:"The shell command to execute."`
Timeout *string `json:"timeout,omitempty" description:"Timeout duration (e.g. '30s', '5m'). Default is 10s. Only applies to foreground commands."`
WorkDir *string `json:"workdir,omitempty" description:"Working directory for the command."`
RunInBackground *bool `json:"run_in_background,omitempty" description:"Run this command in the background without blocking. Use for long-running processes like dev servers, file watchers, or builds that run longer than 5 seconds. Do NOT use shell & to background processes — it will not work correctly. Always use this parameter instead."`
}
// Execute returns an AgentTool that runs a shell command in the
// workspace via the agent HTTP API.
func Execute(options ExecuteOptions) fantasy.AgentTool {
return fantasy.NewAgentTool(
"execute",
"Execute a shell command in the workspace. Use run_in_background=true for long-running processes (dev servers, file watchers, builds). Never use shell '&' for backgrounding. If the command fails or times out, the response may include a background_process_id; use process_output with that ID to retrieve the result.",
func(ctx context.Context, args ExecuteArgs, _ fantasy.ToolCall) (fantasy.ToolResponse, error) {
if options.GetWorkspaceConn == nil {
return fantasy.NewTextErrorResponse("workspace connection resolver is not configured"), nil
}
conn, err := options.GetWorkspaceConn(ctx)
if err != nil {
return fantasy.NewTextErrorResponse(err.Error()), nil
}
return executeTool(ctx, conn, args, options.DefaultTimeout), nil
},
)
}
func executeTool(
ctx context.Context,
conn workspacesdk.AgentConn,
args ExecuteArgs,
optTimeout time.Duration,
) fantasy.ToolResponse {
if args.Command == "" {
return fantasy.NewTextErrorResponse("command is required")
}
// Build the environment map for the process request.
env := make(map[string]string, len(nonInteractiveEnvVars)+1)
env["CODER_CHAT_AGENT"] = "true"
for k, v := range nonInteractiveEnvVars {
env[k] = v
}
background := args.RunInBackground != nil && *args.RunInBackground
// Detect shell-style backgrounding (trailing &) and promote to
// background mode. Models sometimes use "cmd &" instead of the
// run_in_background parameter, which causes the shell to fork
// and exit immediately, leaving an untracked orphan process.
trimmed := strings.TrimSpace(args.Command)
if !background && strings.HasSuffix(trimmed, "&") && !strings.HasSuffix(trimmed, "&&") && !strings.HasSuffix(trimmed, "|&") {
background = true
args.Command = strings.TrimSpace(strings.TrimSuffix(trimmed, "&"))
}
var workDir string
if args.WorkDir != nil {
workDir = *args.WorkDir
}
if background {
return executeBackground(ctx, conn, args.Command, workDir, env)
}
return executeForeground(ctx, conn, args, optTimeout, workDir, env)
}
// executeBackground starts a process in the background and
// returns immediately with the process ID.
func executeBackground(
ctx context.Context,
conn workspacesdk.AgentConn,
command string,
workDir string,
env map[string]string,
) fantasy.ToolResponse {
resp, err := conn.StartProcess(ctx, workspacesdk.StartProcessRequest{
Command: command,
WorkDir: workDir,
Env: env,
Background: true,
})
if err != nil {
return errorResult(fmt.Sprintf("start background process: %v", err))
}
result := ExecuteResult{
Success: true,
BackgroundProcessID: resp.ID,
}
data, err := json.Marshal(result)
if err != nil {
return fantasy.NewTextErrorResponse(err.Error())
}
return fantasy.NewTextResponse(string(data))
}
// executeForeground starts a process and waits for its
// completion, enforcing the configured timeout.
func executeForeground(
ctx context.Context,
conn workspacesdk.AgentConn,
args ExecuteArgs,
optTimeout time.Duration,
workDir string,
env map[string]string,
) fantasy.ToolResponse {
timeout := optTimeout
if timeout <= 0 {
timeout = defaultTimeout
}
if args.Timeout != nil {
parsed, err := time.ParseDuration(*args.Timeout)
if err != nil {
return fantasy.NewTextErrorResponse(
fmt.Sprintf("invalid timeout %q: %v", *args.Timeout, err),
)
}
timeout = parsed
}
cmdCtx, cancel := context.WithTimeout(ctx, timeout)
defer cancel()
start := time.Now()
resp, err := conn.StartProcess(cmdCtx, workspacesdk.StartProcessRequest{
Command: args.Command,
WorkDir: workDir,
Env: env,
Background: false,
})
if err != nil {
return errorResult(fmt.Sprintf("start process: %v", err))
}
result := waitForProcess(cmdCtx, ctx, conn, resp.ID, timeout)
result.WallDurationMs = time.Since(start).Milliseconds()
// Add an advisory note for file-dump commands.
if note := detectFileDump(args.Command); note != "" {
result.Note = note
}
data, err := json.Marshal(result)
if err != nil {
return fantasy.NewTextErrorResponse(err.Error())
}
return fantasy.NewTextResponse(string(data))
}
// truncateOutput safely truncates output to maxOutputToModel,
// ensuring the result is valid UTF-8 even if the cut falls in
// the middle of a multi-byte character.
func truncateOutput(output string) string {
if len(output) > maxOutputToModel {
output = strings.ToValidUTF8(output[:maxOutputToModel], "")
}
return output
}
// waitForProcess waits for process completion using the
// blocking process output API instead of polling.
// waitForProcess blocks until the process exits or the context
// expires. On any error (timeout or transport), it tries a
// non-blocking snapshot to recover. Total wall time may exceed
// timeout by up to snapshotTimeout if recovery is needed.
func waitForProcess(
ctx context.Context,
parentCtx context.Context,
conn workspacesdk.AgentConn,
processID string,
timeout time.Duration,
) ExecuteResult {
// Block until the process exits or the context is
// canceled.
resp, err := conn.ProcessOutput(ctx, processID, &workspacesdk.ProcessOutputOptions{
Wait: true,
})
if err != nil {
origErr := err
timedOut := ctx.Err() != nil
// Fetch a snapshot with a fresh context. The blocking
// request may have failed due to a context timeout or
// a transport error (e.g. the server's WriteTimeout
// killed the connection). Either way, the process may
// still have output available.
bgCtx, bgCancel := context.WithTimeout(
parentCtx,
snapshotTimeout,
)
defer bgCancel()
resp, err = conn.ProcessOutput(bgCtx, processID, nil)
if err != nil {
errMsg := fmt.Sprintf("get process output: %v; use process_output with ID %s to retry", origErr, processID)
if timedOut {
errMsg = fmt.Sprintf("command timed out after %s; failed to get output: %v", timeout, err)
}
return ExecuteResult{
Success: false,
ExitCode: -1,
Error: errMsg,
BackgroundProcessID: processID,
}
}
// Snapshot succeeded. If the process finished, return
// its real result (transparent recovery).
if !resp.Running {
exitCode := 0
if resp.ExitCode != nil {
exitCode = *resp.ExitCode
}
output := truncateOutput(resp.Output)
return ExecuteResult{
Success: exitCode == 0,
Output: output,
ExitCode: exitCode,
Truncated: resp.Truncated,
}
}
// Process still running, return partial output.
output := truncateOutput(resp.Output)
errMsg := fmt.Sprintf("command timed out after %s", timeout)
if !timedOut {
errMsg = fmt.Sprintf("get process output: %v (process still running, use process_output to check later)", origErr)
}
return ExecuteResult{
Success: false,
Output: output,
ExitCode: -1,
Error: errMsg,
Truncated: resp.Truncated,
BackgroundProcessID: processID,
}
}
// The server-side wait may return before the
// process exits if maxWaitDuration is shorter than
// the client's timeout. Retry if our context still
// has time left.
if resp.Running {
if ctx.Err() == nil {
// Still within the caller's timeout, retry.
return waitForProcess(ctx, parentCtx, conn, processID, timeout)
}
output := truncateOutput(resp.Output)
return ExecuteResult{
Success: false,
Output: output,
ExitCode: -1,
Error: fmt.Sprintf("command timed out after %s", timeout),
Truncated: resp.Truncated,
BackgroundProcessID: processID,
}
}
exitCode := 0
if resp.ExitCode != nil {
exitCode = *resp.ExitCode
}
output := truncateOutput(resp.Output)
return ExecuteResult{
Success: exitCode == 0,
Output: output,
ExitCode: exitCode,
Truncated: resp.Truncated,
}
}
// errorResult builds a ToolResponse from an ExecuteResult with
// an error message.
func errorResult(msg string) fantasy.ToolResponse {
data, err := json.Marshal(ExecuteResult{
Success: false,
Error: msg,
})
if err != nil {
return fantasy.NewTextErrorResponse(msg)
}
return fantasy.NewTextResponse(string(data))
}
// detectFileDump checks whether the command matches a file-dump
// pattern and returns an advisory note, or empty string if no
// match.
func detectFileDump(command string) string {
for _, pat := range fileDumpPatterns {
if pat.MatchString(command) {
return "Consider using read_file instead of " +
"dumping file contents with shell commands."
}
}
return ""
}
const (
// defaultProcessOutputTimeout is the default time the
// process_output tool blocks waiting for new output or
// process exit before returning. This avoids polling
// loops that waste tokens and HTTP round-trips.
defaultProcessOutputTimeout = 10 * time.Second
)
// ProcessOutputArgs are the parameters accepted by the
// process_output tool.
type ProcessOutputArgs struct {
ProcessID string `json:"process_id"`
WaitTimeout *string `json:"wait_timeout,omitempty" description:"Override the default 10s block duration. The call blocks until the process exits or this timeout is reached. Set to '0s' for an immediate snapshot without waiting."`
}
// ProcessOutput returns an AgentTool that retrieves the output
// of a background process by its ID.
func ProcessOutput(options ProcessToolOptions) fantasy.AgentTool {
return fantasy.NewAgentTool(
"process_output",
"Retrieve output from a background process. "+
"Use the process_id returned by execute with "+
"run_in_background=true or from a timed-out "+
"execute's background_process_id. Blocks up to "+
"10s for the process to exit, then returns the "+
"output and exit_code. If still running after "+
"the timeout, returns the output so far. Use "+
"wait_timeout to override the default 10s wait "+
"(e.g. '30s', or '0s' for an immediate snapshot).",
func(ctx context.Context, args ProcessOutputArgs, _ fantasy.ToolCall) (fantasy.ToolResponse, error) {
if options.GetWorkspaceConn == nil {
return fantasy.NewTextErrorResponse("workspace connection resolver is not configured"), nil
}
if args.ProcessID == "" {
return fantasy.NewTextErrorResponse("process_id is required"), nil
}
conn, err := options.GetWorkspaceConn(ctx)
if err != nil {
return fantasy.NewTextErrorResponse(err.Error()), nil
}
timeout := defaultProcessOutputTimeout
if args.WaitTimeout != nil {
parsed, err := time.ParseDuration(*args.WaitTimeout)
if err != nil {
return fantasy.NewTextErrorResponse(
fmt.Sprintf("invalid wait_timeout %q: %v", *args.WaitTimeout, err),
), nil
}
timeout = parsed
}
var opts *workspacesdk.ProcessOutputOptions
// Save parent context before applying timeout.
parentCtx := ctx
if timeout > 0 {
opts = &workspacesdk.ProcessOutputOptions{
Wait: true,
}
var cancel context.CancelFunc
ctx, cancel = context.WithTimeout(ctx, timeout)
defer cancel()
}
resp, err := conn.ProcessOutput(ctx, args.ProcessID, opts)
if err != nil {
// The blocking request may have failed due to a
// context timeout or a transport error (e.g.
// server WriteTimeout). Try a non-blocking
// snapshot if the parent context is still alive.
if parentCtx.Err() != nil {
return errorResult(fmt.Sprintf("get process output: %v", err)), nil
}
bgCtx, bgCancel := context.WithTimeout(parentCtx, snapshotTimeout)
defer bgCancel()
resp, err = conn.ProcessOutput(bgCtx, args.ProcessID, nil)
if err != nil {
return errorResult(fmt.Sprintf("get process output: %v", err)), nil
}
// Fall through to normal response handling below.
}
output := truncateOutput(resp.Output)
exitCode := 0
if resp.ExitCode != nil {
exitCode = *resp.ExitCode
}
result := ExecuteResult{
Success: !resp.Running && exitCode == 0,
Output: output,
ExitCode: exitCode,
Truncated: resp.Truncated,
}
if resp.Running {
// Process is still running, success is not
// yet determined.
result.Success = true
result.Note = "process is still running"
}
data, err := json.Marshal(result)
if err != nil {
return fantasy.NewTextErrorResponse(err.Error()), nil
}
return fantasy.NewTextResponse(string(data)), nil
},
)
}
// ProcessList returns an AgentTool that lists all tracked
// processes on the workspace agent.
func ProcessList(options ProcessToolOptions) fantasy.AgentTool {
return fantasy.NewAgentTool(
"process_list",
"List all tracked processes in the workspace. "+
"Returns process IDs, commands, status (running or "+
"exited), and exit codes. Use this to discover "+
"background processes or check which processes are "+
"still running.",
func(ctx context.Context, _ struct{}, _ fantasy.ToolCall) (fantasy.ToolResponse, error) {
if options.GetWorkspaceConn == nil {
return fantasy.NewTextErrorResponse("workspace connection resolver is not configured"), nil
}
conn, err := options.GetWorkspaceConn(ctx)
if err != nil {
return fantasy.NewTextErrorResponse(err.Error()), nil
}
resp, err := conn.ListProcesses(ctx)
if err != nil {
return errorResult(fmt.Sprintf("list processes: %v", err)), nil
}
data, err := json.Marshal(resp)
if err != nil {
return fantasy.NewTextErrorResponse(err.Error()), nil
}
return fantasy.NewTextResponse(string(data)), nil
},
)
}
// ProcessSignalArgs are the parameters accepted by the
// process_signal tool.
type ProcessSignalArgs struct {
ProcessID string `json:"process_id"`
Signal string `json:"signal"`
}
// ProcessSignal returns an AgentTool that sends a signal to a
// tracked process on the workspace agent.
func ProcessSignal(options ProcessToolOptions) fantasy.AgentTool {
return fantasy.NewAgentTool(
"process_signal",
"Send a signal to a background process. "+
"Use \"terminate\" (SIGTERM) for graceful shutdown "+
"or \"kill\" (SIGKILL) to force stop. Use the "+
"process_id returned by execute with "+
"run_in_background=true or from process_list.",
func(ctx context.Context, args ProcessSignalArgs, _ fantasy.ToolCall) (fantasy.ToolResponse, error) {
if options.GetWorkspaceConn == nil {
return fantasy.NewTextErrorResponse("workspace connection resolver is not configured"), nil
}
if args.ProcessID == "" {
return fantasy.NewTextErrorResponse("process_id is required"), nil
}
if args.Signal != "terminate" && args.Signal != "kill" {
return fantasy.NewTextErrorResponse(
"signal must be \"terminate\" or \"kill\"",
), nil
}
conn, err := options.GetWorkspaceConn(ctx)
if err != nil {
return fantasy.NewTextErrorResponse(err.Error()), nil
}
if err := conn.SignalProcess(ctx, args.ProcessID, args.Signal); err != nil {
return errorResult(fmt.Sprintf("signal process: %v", err)), nil
}
data, err := json.Marshal(map[string]any{
"success": true,
"message": fmt.Sprintf(
"signal %q sent to process %s",
args.Signal, args.ProcessID,
),
})
if err != nil {
return fantasy.NewTextErrorResponse(err.Error()), nil
}
return fantasy.NewTextResponse(string(data)), nil
},
)
}
@@ -0,0 +1,100 @@
package chattool
import (
"context"
"encoding/json"
"strings"
"testing"
"unicode/utf8"
"charm.land/fantasy"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"go.uber.org/mock/gomock"
"github.com/coder/coder/v2/codersdk/workspacesdk"
"github.com/coder/coder/v2/codersdk/workspacesdk/agentconnmock"
"github.com/coder/coder/v2/testutil"
)
func TestTruncateOutput(t *testing.T) {
t.Parallel()
t.Run("EmptyOutput", func(t *testing.T) {
t.Parallel()
result := runForegroundWithOutput(t, "")
assert.Empty(t, result.Output)
})
t.Run("ShortOutput", func(t *testing.T) {
t.Parallel()
result := runForegroundWithOutput(t, "short")
assert.Equal(t, "short", result.Output)
})
t.Run("ExactlyAtLimit", func(t *testing.T) {
t.Parallel()
output := strings.Repeat("a", maxOutputToModel)
result := runForegroundWithOutput(t, output)
assert.Equal(t, maxOutputToModel, len(result.Output))
assert.Equal(t, output, result.Output)
})
t.Run("OverLimit", func(t *testing.T) {
t.Parallel()
output := strings.Repeat("b", maxOutputToModel+1024)
result := runForegroundWithOutput(t, output)
assert.Equal(t, maxOutputToModel, len(result.Output))
})
t.Run("MultiByteCutMidCharacter", func(t *testing.T) {
t.Parallel()
// Build output that places a 3-byte UTF-8 character
// (U+2603, snowman ☃) right at the truncation boundary
// so the cut falls mid-character.
padding := strings.Repeat("x", maxOutputToModel-1)
output := padding + "☃" // ☃ is 3 bytes, only 1 byte fits
result := runForegroundWithOutput(t, output)
assert.LessOrEqual(t, len(result.Output), maxOutputToModel)
assert.True(t, utf8.ValidString(result.Output),
"truncated output must be valid UTF-8")
})
}
// runForegroundWithOutput runs a foreground command through the
// Execute tool with a mock that returns the given output, and
// returns the parsed result.
func runForegroundWithOutput(t *testing.T, output string) ExecuteResult {
t.Helper()
ctrl := gomock.NewController(t)
mockConn := agentconnmock.NewMockAgentConn(ctrl)
mockConn.EXPECT().
StartProcess(gomock.Any(), gomock.Any()).
Return(workspacesdk.StartProcessResponse{ID: "proc-1"}, nil)
exitCode := 0
mockConn.EXPECT().
ProcessOutput(gomock.Any(), "proc-1", gomock.Any()).
Return(workspacesdk.ProcessOutputResponse{
Running: false,
ExitCode: &exitCode,
Output: output,
}, nil)
tool := Execute(ExecuteOptions{
GetWorkspaceConn: func(_ context.Context) (workspacesdk.AgentConn, error) {
return mockConn, nil
},
})
ctx := testutil.Context(t, testutil.WaitMedium)
resp, err := tool.Run(ctx, fantasy.ToolCall{
ID: "call-1",
Name: "execute",
Input: `{"command":"echo test"}`,
})
require.NoError(t, err)
var result ExecuteResult
require.NoError(t, json.Unmarshal([]byte(resp.Content), &result))
return result
}
+581
View File
@@ -0,0 +1,581 @@
package chattool_test
import (
"context"
"encoding/json"
"testing"
"charm.land/fantasy"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"go.uber.org/mock/gomock"
"golang.org/x/xerrors"
"github.com/coder/coder/v2/coderd/x/chatd/chattool"
"github.com/coder/coder/v2/codersdk/workspacesdk"
"github.com/coder/coder/v2/codersdk/workspacesdk/agentconnmock"
"github.com/coder/coder/v2/testutil"
)
func TestExecuteTool(t *testing.T) {
t.Parallel()
t.Run("EmptyCommand", func(t *testing.T) {
t.Parallel()
ctrl := gomock.NewController(t)
mockConn := agentconnmock.NewMockAgentConn(ctrl)
tool := newExecuteTool(t, mockConn)
resp, err := tool.Run(context.Background(), fantasy.ToolCall{
ID: "call-1",
Name: "execute",
Input: `{"command":""}`,
})
require.NoError(t, err)
assert.True(t, resp.IsError)
assert.Contains(t, resp.Content, "command is required")
})
t.Run("AmpersandDetection", func(t *testing.T) {
t.Parallel()
tests := []struct {
name string
command string
runInBackground *bool
wantCommand string
wantBackground bool
wantBackgroundResp bool // true if the response should contain a background_process_id
comment string
}{
{
name: "SimpleBackground",
command: "cmd &",
wantCommand: "cmd",
wantBackground: true,
wantBackgroundResp: true,
comment: "Trailing & is correctly detected and stripped.",
},
{
name: "TrailingDoubleAmpersand",
command: "cmd &&",
wantCommand: "cmd &&",
wantBackground: false,
wantBackgroundResp: false,
comment: "Ends with &&, excluded by the && suffix check.",
},
{
name: "NoAmpersand",
command: "cmd",
wantCommand: "cmd",
wantBackground: false,
wantBackgroundResp: false,
},
{
name: "ChainThenBackground",
command: "cmd1 && cmd2 &",
wantCommand: "cmd1 && cmd2",
wantBackground: true,
wantBackgroundResp: true,
comment: "Ends with & but not &&, so it gets promoted " +
"to background and the trailing & is stripped. " +
"The remaining command runs in background mode.",
},
{
// "|&" is bash's pipe-stderr operator, not
// backgrounding. It must not be detected as a
// trailing "&".
name: "BashPipeStderr",
command: "cmd |&",
wantCommand: "cmd |&",
wantBackground: false,
wantBackgroundResp: false,
},
{
name: "AlreadyBackgroundWithTrailingAmpersand",
command: "cmd &",
runInBackground: ptr(true),
wantCommand: "cmd &",
wantBackground: true,
wantBackgroundResp: true,
comment: "When run_in_background is already true, " +
"the stripping logic is skipped, preserving " +
"the original command.",
},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
t.Parallel()
ctrl := gomock.NewController(t)
mockConn := agentconnmock.NewMockAgentConn(ctrl)
var capturedReq workspacesdk.StartProcessRequest
mockConn.EXPECT().
StartProcess(gomock.Any(), gomock.Any()).
DoAndReturn(func(_ context.Context, req workspacesdk.StartProcessRequest) (workspacesdk.StartProcessResponse, error) {
capturedReq = req
return workspacesdk.StartProcessResponse{ID: "proc-1"}, nil
})
// For foreground cases, ProcessOutput is polled.
exitCode := 0
mockConn.EXPECT().
ProcessOutput(gomock.Any(), "proc-1", gomock.Any()).
Return(workspacesdk.ProcessOutputResponse{
Running: false,
ExitCode: &exitCode,
}, nil).
AnyTimes()
tool := newExecuteTool(t, mockConn)
input := map[string]any{"command": tc.command}
if tc.runInBackground != nil {
input["run_in_background"] = *tc.runInBackground
}
inputJSON, err := json.Marshal(input)
require.NoError(t, err)
ctx := testutil.Context(t, testutil.WaitMedium)
resp, err := tool.Run(ctx, fantasy.ToolCall{
ID: "call-1",
Name: "execute",
Input: string(inputJSON),
})
require.NoError(t, err)
assert.False(t, resp.IsError, "response should not be an error")
assert.Equal(t, tc.wantCommand, capturedReq.Command,
"command passed to StartProcess")
assert.Equal(t, tc.wantBackground, capturedReq.Background,
"background flag passed to StartProcess")
var result chattool.ExecuteResult
require.NoError(t, json.Unmarshal([]byte(resp.Content), &result))
if tc.wantBackgroundResp {
assert.NotEmpty(t, result.BackgroundProcessID,
"expected background_process_id in response")
} else {
assert.Empty(t, result.BackgroundProcessID,
"expected no background_process_id")
}
})
}
})
t.Run("ForegroundSuccess", func(t *testing.T) {
t.Parallel()
ctrl := gomock.NewController(t)
mockConn := agentconnmock.NewMockAgentConn(ctrl)
var capturedReq workspacesdk.StartProcessRequest
mockConn.EXPECT().
StartProcess(gomock.Any(), gomock.Any()).
DoAndReturn(func(_ context.Context, req workspacesdk.StartProcessRequest) (workspacesdk.StartProcessResponse, error) {
capturedReq = req
return workspacesdk.StartProcessResponse{ID: "proc-1"}, nil
})
exitCode := 0
mockConn.EXPECT().
ProcessOutput(gomock.Any(), "proc-1", gomock.Any()).
Return(workspacesdk.ProcessOutputResponse{
Running: false,
ExitCode: &exitCode,
Output: "hello world",
}, nil)
tool := newExecuteTool(t, mockConn)
ctx := testutil.Context(t, testutil.WaitMedium)
resp, err := tool.Run(ctx, fantasy.ToolCall{
ID: "call-1",
Name: "execute",
Input: `{"command":"echo hello"}`,
})
require.NoError(t, err)
assert.False(t, resp.IsError)
var result chattool.ExecuteResult
require.NoError(t, json.Unmarshal([]byte(resp.Content), &result))
assert.True(t, result.Success)
assert.Equal(t, 0, result.ExitCode)
assert.Equal(t, "hello world", result.Output)
assert.Empty(t, result.BackgroundProcessID)
assert.Equal(t, "true", capturedReq.Env["CODER_CHAT_AGENT"])
})
t.Run("ForegroundNonZeroExit", func(t *testing.T) {
t.Parallel()
ctrl := gomock.NewController(t)
mockConn := agentconnmock.NewMockAgentConn(ctrl)
mockConn.EXPECT().
StartProcess(gomock.Any(), gomock.Any()).
Return(workspacesdk.StartProcessResponse{ID: "proc-1"}, nil)
exitCode := 42
mockConn.EXPECT().
ProcessOutput(gomock.Any(), "proc-1", gomock.Any()).
Return(workspacesdk.ProcessOutputResponse{
Running: false,
ExitCode: &exitCode,
Output: "something failed",
}, nil)
tool := newExecuteTool(t, mockConn)
ctx := testutil.Context(t, testutil.WaitMedium)
resp, err := tool.Run(ctx, fantasy.ToolCall{
ID: "call-1",
Name: "execute",
Input: `{"command":"exit 42"}`,
})
require.NoError(t, err)
assert.False(t, resp.IsError)
var result chattool.ExecuteResult
require.NoError(t, json.Unmarshal([]byte(resp.Content), &result))
assert.False(t, result.Success)
assert.Equal(t, 42, result.ExitCode)
assert.Equal(t, "something failed", result.Output)
})
t.Run("BackgroundExecution", func(t *testing.T) {
t.Parallel()
ctrl := gomock.NewController(t)
mockConn := agentconnmock.NewMockAgentConn(ctrl)
mockConn.EXPECT().
StartProcess(gomock.Any(), gomock.Any()).
DoAndReturn(func(_ context.Context, req workspacesdk.StartProcessRequest) (workspacesdk.StartProcessResponse, error) {
assert.True(t, req.Background)
return workspacesdk.StartProcessResponse{ID: "bg-42"}, nil
})
tool := newExecuteTool(t, mockConn)
ctx := testutil.Context(t, testutil.WaitMedium)
resp, err := tool.Run(ctx, fantasy.ToolCall{
ID: "call-1",
Name: "execute",
Input: `{"command":"sleep 999","run_in_background":true}`,
})
require.NoError(t, err)
assert.False(t, resp.IsError)
var result chattool.ExecuteResult
require.NoError(t, json.Unmarshal([]byte(resp.Content), &result))
assert.True(t, result.Success)
assert.Equal(t, "bg-42", result.BackgroundProcessID)
})
t.Run("Timeout", func(t *testing.T) {
t.Parallel()
ctrl := gomock.NewController(t)
mockConn := agentconnmock.NewMockAgentConn(ctrl)
mockConn.EXPECT().
StartProcess(gomock.Any(), gomock.Any()).
Return(workspacesdk.StartProcessResponse{ID: "proc-1"}, nil)
// First call (blocking wait) returns context error
// because the 50ms timeout expires.
mockConn.EXPECT().
ProcessOutput(gomock.Any(), "proc-1", gomock.Any()).
DoAndReturn(func(ctx context.Context, _ string, _ *workspacesdk.ProcessOutputOptions) (workspacesdk.ProcessOutputResponse, error) {
<-ctx.Done()
return workspacesdk.ProcessOutputResponse{}, ctx.Err()
})
// Second call (snapshot fallback) returns partial output.
mockConn.EXPECT().
ProcessOutput(gomock.Any(), "proc-1", gomock.Any()).
Return(workspacesdk.ProcessOutputResponse{
Running: true,
Output: "partial output",
}, nil)
tool := newExecuteTool(t, mockConn)
ctx := testutil.Context(t, testutil.WaitMedium)
resp, err := tool.Run(ctx, fantasy.ToolCall{
ID: "call-1",
Name: "execute",
// 50ms timeout expires during the blocking wait.
Input: `{"command":"sleep 999","timeout":"50ms"}`,
})
require.NoError(t, err)
assert.False(t, resp.IsError)
var result chattool.ExecuteResult
require.NoError(t, json.Unmarshal([]byte(resp.Content), &result))
assert.False(t, result.Success)
assert.Equal(t, -1, result.ExitCode)
assert.Contains(t, result.Error, "timed out")
assert.Equal(t, "partial output", result.Output)
})
t.Run("StartProcessError", func(t *testing.T) {
t.Parallel()
ctrl := gomock.NewController(t)
mockConn := agentconnmock.NewMockAgentConn(ctrl)
mockConn.EXPECT().
StartProcess(gomock.Any(), gomock.Any()).
Return(workspacesdk.StartProcessResponse{}, xerrors.New("connection lost"))
tool := newExecuteTool(t, mockConn)
ctx := testutil.Context(t, testutil.WaitMedium)
resp, err := tool.Run(ctx, fantasy.ToolCall{
ID: "call-1",
Name: "execute",
Input: `{"command":"echo hi"}`,
})
require.NoError(t, err)
// Errors from StartProcess are returned as a JSON body
// with success=false, not as a ToolResponse error.
assert.False(t, resp.IsError)
var result chattool.ExecuteResult
require.NoError(t, json.Unmarshal([]byte(resp.Content), &result))
assert.False(t, result.Success)
assert.Contains(t, result.Error, "connection lost")
})
t.Run("ProcessOutputError", func(t *testing.T) {
t.Parallel()
ctrl := gomock.NewController(t)
mockConn := agentconnmock.NewMockAgentConn(ctrl)
mockConn.EXPECT().
StartProcess(gomock.Any(), gomock.Any()).
Return(workspacesdk.StartProcessResponse{ID: "proc-1"}, nil)
// First call: blocking wait fails.
mockConn.EXPECT().
ProcessOutput(gomock.Any(), "proc-1", gomock.Any()).
Return(workspacesdk.ProcessOutputResponse{}, xerrors.New("agent disconnected"))
// Second call: snapshot fallback also fails.
mockConn.EXPECT().
ProcessOutput(gomock.Any(), "proc-1", gomock.Any()).
Return(workspacesdk.ProcessOutputResponse{}, xerrors.New("agent disconnected"))
tool := newExecuteTool(t, mockConn)
ctx := testutil.Context(t, testutil.WaitMedium)
resp, err := tool.Run(ctx, fantasy.ToolCall{
ID: "call-1",
Name: "execute",
Input: `{"command":"echo hi"}`,
})
require.NoError(t, err)
assert.False(t, resp.IsError)
var result chattool.ExecuteResult
require.NoError(t, json.Unmarshal([]byte(resp.Content), &result))
assert.False(t, result.Success)
assert.Contains(t, result.Error, "agent disconnected")
// Snapshot fallback should provide the process ID
// so the agent can retry manually.
assert.Equal(t, "proc-1", result.BackgroundProcessID)
})
t.Run("TransportErrorRecoveryProcessDone", func(t *testing.T) {
t.Parallel()
ctrl := gomock.NewController(t)
mockConn := agentconnmock.NewMockAgentConn(ctrl)
exitCode := 0
mockConn.EXPECT().
StartProcess(gomock.Any(), gomock.Any()).
Return(workspacesdk.StartProcessResponse{ID: "proc-1"}, nil)
// Blocking wait fails with transport error.
mockConn.EXPECT().
ProcessOutput(gomock.Any(), "proc-1", gomock.Any()).
Return(workspacesdk.ProcessOutputResponse{}, xerrors.New("EOF"))
// Snapshot fallback finds the process completed.
mockConn.EXPECT().
ProcessOutput(gomock.Any(), "proc-1", gomock.Any()).
Return(workspacesdk.ProcessOutputResponse{
Output: "hello\n",
Running: false,
ExitCode: &exitCode,
}, nil)
tool := newExecuteTool(t, mockConn)
ctx := testutil.Context(t, testutil.WaitMedium)
resp, err := tool.Run(ctx, fantasy.ToolCall{
ID: "call-1",
Name: "execute",
Input: `{"command":"echo hello"}`,
})
require.NoError(t, err)
assert.False(t, resp.IsError)
var result chattool.ExecuteResult
require.NoError(t, json.Unmarshal([]byte(resp.Content), &result))
// Transparent recovery: success with real output.
assert.True(t, result.Success)
assert.Equal(t, 0, result.ExitCode)
assert.Equal(t, "hello\n", result.Output)
assert.Empty(t, result.BackgroundProcessID)
})
t.Run("TransportErrorProcessStillRunning", func(t *testing.T) {
t.Parallel()
ctrl := gomock.NewController(t)
mockConn := agentconnmock.NewMockAgentConn(ctrl)
mockConn.EXPECT().
StartProcess(gomock.Any(), gomock.Any()).
Return(workspacesdk.StartProcessResponse{ID: "proc-1"}, nil)
// Blocking wait fails with transport error.
mockConn.EXPECT().
ProcessOutput(gomock.Any(), "proc-1", gomock.Any()).
Return(workspacesdk.ProcessOutputResponse{}, xerrors.New("EOF"))
// Snapshot fallback: process still running.
mockConn.EXPECT().
ProcessOutput(gomock.Any(), "proc-1", gomock.Any()).
Return(workspacesdk.ProcessOutputResponse{
Output: "partial output",
Running: true,
}, nil)
tool := newExecuteTool(t, mockConn)
ctx := testutil.Context(t, testutil.WaitMedium)
resp, err := tool.Run(ctx, fantasy.ToolCall{
ID: "call-1",
Name: "execute",
Input: `{"command":"sleep 60"}`,
})
require.NoError(t, err)
assert.False(t, resp.IsError)
var result chattool.ExecuteResult
require.NoError(t, json.Unmarshal([]byte(resp.Content), &result))
assert.False(t, result.Success)
assert.Contains(t, result.Error, "process still running")
assert.Contains(t, result.Error, "process_output")
assert.Equal(t, "partial output", result.Output)
assert.Equal(t, "proc-1", result.BackgroundProcessID)
})
t.Run("GetWorkspaceConnNil", func(t *testing.T) {
t.Parallel()
tool := chattool.Execute(chattool.ExecuteOptions{
GetWorkspaceConn: nil,
})
ctx := testutil.Context(t, testutil.WaitMedium)
resp, err := tool.Run(ctx, fantasy.ToolCall{
ID: "call-1",
Name: "execute",
Input: `{"command":"echo hi"}`,
})
require.NoError(t, err)
assert.True(t, resp.IsError)
assert.Contains(t, resp.Content, "not configured")
})
t.Run("GetWorkspaceConnError", func(t *testing.T) {
t.Parallel()
tool := chattool.Execute(chattool.ExecuteOptions{
GetWorkspaceConn: func(_ context.Context) (workspacesdk.AgentConn, error) {
return nil, xerrors.New("workspace offline")
},
})
ctx := testutil.Context(t, testutil.WaitMedium)
resp, err := tool.Run(ctx, fantasy.ToolCall{
ID: "call-1",
Name: "execute",
Input: `{"command":"echo hi"}`,
})
require.NoError(t, err)
assert.True(t, resp.IsError)
assert.Contains(t, resp.Content, "workspace offline")
})
}
func TestDetectFileDump(t *testing.T) {
t.Parallel()
tests := []struct {
name string
command string
wantHit bool
}{
{
name: "CatFile",
command: "cat foo.txt",
wantHit: true,
},
{
name: "NotCatPrefix",
command: "concatenate foo",
wantHit: false,
},
{
name: "GrepIncludeAll",
command: "grep --include-all pattern",
wantHit: true,
},
{
name: "RgListFiles",
command: "rg -l pattern",
wantHit: true,
},
{
name: "GrepRecursive",
command: "grep -r pattern",
wantHit: false,
},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
t.Parallel()
ctrl := gomock.NewController(t)
mockConn := agentconnmock.NewMockAgentConn(ctrl)
mockConn.EXPECT().
StartProcess(gomock.Any(), gomock.Any()).
Return(workspacesdk.StartProcessResponse{ID: "proc-1"}, nil)
exitCode := 0
mockConn.EXPECT().
ProcessOutput(gomock.Any(), "proc-1", gomock.Any()).
Return(workspacesdk.ProcessOutputResponse{
Running: false,
ExitCode: &exitCode,
Output: "output",
}, nil)
tool := newExecuteTool(t, mockConn)
ctx := testutil.Context(t, testutil.WaitMedium)
input, err := json.Marshal(map[string]any{
"command": tc.command,
})
require.NoError(t, err)
resp, err := tool.Run(ctx, fantasy.ToolCall{
ID: "call-1",
Name: "execute",
Input: string(input),
})
require.NoError(t, err)
var result chattool.ExecuteResult
require.NoError(t, json.Unmarshal([]byte(resp.Content), &result))
if tc.wantHit {
assert.Contains(t, result.Note, "read_file",
"expected advisory note for %q", tc.command)
} else {
assert.Empty(t, result.Note,
"expected no note for %q", tc.command)
}
})
}
}
// newExecuteTool creates an Execute tool wired to the given mock.
func newExecuteTool(t *testing.T, mockConn *agentconnmock.MockAgentConn) fantasy.AgentTool {
t.Helper()
return chattool.Execute(chattool.ExecuteOptions{
GetWorkspaceConn: func(_ context.Context) (workspacesdk.AgentConn, error) {
return mockConn, nil
},
})
}
func ptr[T any](v T) *T {
return &v
}
+148
View File
@@ -0,0 +1,148 @@
package chattool
import (
"context"
"database/sql"
"sort"
"strings"
"charm.land/fantasy"
"github.com/google/uuid"
"golang.org/x/xerrors"
"github.com/coder/coder/v2/coderd/database"
"github.com/coder/coder/v2/coderd/database/dbauthz"
"github.com/coder/coder/v2/coderd/httpmw"
"github.com/coder/coder/v2/coderd/rbac"
)
const listTemplatesPageSize = 10
// ListTemplatesOptions configures the list_templates tool.
type ListTemplatesOptions struct {
DB database.Store
OwnerID uuid.UUID
}
type listTemplatesArgs struct {
Query string `json:"query,omitempty"`
Page int `json:"page,omitempty"`
}
// ListTemplates returns a tool that lists available workspace templates.
// The agent uses this to discover templates before creating a workspace.
// Results are ordered by number of active developers (most popular first)
// and paginated at 10 per page.
func ListTemplates(options ListTemplatesOptions) fantasy.AgentTool {
return fantasy.NewAgentTool(
"list_templates",
"List available workspace templates. Optionally filter by a "+
"search query matching template name or description. "+
"Use this to find a template before creating a workspace. "+
"Results are ordered by number of active developers (most popular first). "+
"Returns 10 per page. Use the page parameter to paginate through results.",
func(ctx context.Context, args listTemplatesArgs, _ fantasy.ToolCall) (fantasy.ToolResponse, error) {
if options.DB == nil {
return fantasy.NewTextErrorResponse("database is not configured"), nil
}
ctx, err := asOwner(ctx, options.DB, options.OwnerID)
if err != nil {
return fantasy.NewTextErrorResponse(err.Error()), nil
}
filterParams := database.GetTemplatesWithFilterParams{
Deleted: false,
Deprecated: sql.NullBool{
Bool: false,
Valid: true,
},
}
query := strings.TrimSpace(args.Query)
if query != "" {
filterParams.FuzzyName = query
}
templates, err := options.DB.GetTemplatesWithFilter(ctx, filterParams)
if err != nil {
return fantasy.NewTextErrorResponse(err.Error()), nil
}
// Look up active developer counts so we can sort by popularity.
templateIDs := make([]uuid.UUID, len(templates))
for i, t := range templates {
templateIDs[i] = t.ID
}
ownerCounts := make(map[uuid.UUID]int64)
if len(templateIDs) > 0 {
rows, countErr := options.DB.GetWorkspaceUniqueOwnerCountByTemplateIDs(ctx, templateIDs)
if countErr == nil {
for _, row := range rows {
ownerCounts[row.TemplateID] = row.UniqueOwnersSum
}
}
}
// Sort by active developer count descending.
sort.SliceStable(templates, func(i, j int) bool {
return ownerCounts[templates[i].ID] > ownerCounts[templates[j].ID]
})
// Paginate.
page := args.Page
if page < 1 {
page = 1
}
totalCount := len(templates)
totalPages := (totalCount + listTemplatesPageSize - 1) / listTemplatesPageSize
if totalPages == 0 {
totalPages = 1
}
start := (page - 1) * listTemplatesPageSize
end := start + listTemplatesPageSize
if start > totalCount {
start = totalCount
}
if end > totalCount {
end = totalCount
}
pageTemplates := templates[start:end]
items := make([]map[string]any, 0, len(pageTemplates))
for _, t := range pageTemplates {
item := map[string]any{
"id": t.ID.String(),
"name": t.Name,
}
if display := strings.TrimSpace(t.DisplayName); display != "" {
item["display_name"] = display
}
if desc := strings.TrimSpace(t.Description); desc != "" {
item["description"] = truncateRunes(desc, 200)
}
if count, ok := ownerCounts[t.ID]; ok && count > 0 {
item["active_developers"] = count
}
items = append(items, item)
}
return toolResponse(map[string]any{
"templates": items,
"count": len(items),
"page": page,
"total_pages": totalPages,
"total_count": totalCount,
}), nil
},
)
}
// asOwner sets up a dbauthz context for the given owner so that
// subsequent database calls are scoped to what that user can access.
func asOwner(ctx context.Context, db database.Store, ownerID uuid.UUID) (context.Context, error) {
actor, _, err := httpmw.UserRBACSubject(ctx, db, ownerID, rbac.ScopeAll)
if err != nil {
return ctx, xerrors.Errorf("load user authorization: %w", err)
}
return dbauthz.As(ctx, actor), nil
}
+74
View File
@@ -0,0 +1,74 @@
package chattool
import (
"context"
"charm.land/fantasy"
"github.com/coder/coder/v2/codersdk/workspacesdk"
)
type ReadFileOptions struct {
GetWorkspaceConn func(context.Context) (workspacesdk.AgentConn, error)
}
type ReadFileArgs struct {
Path string `json:"path"`
Offset *int64 `json:"offset,omitempty"`
Limit *int64 `json:"limit,omitempty"`
}
func ReadFile(options ReadFileOptions) fantasy.AgentTool {
return fantasy.NewAgentTool(
"read_file",
"Read a file from the workspace. Returns line-numbered content. "+
"The offset parameter is a 1-based line number (default: 1). "+
"The limit parameter is the number of lines to return (default: 2000). "+
"For large files, use offset and limit to paginate.",
func(ctx context.Context, args ReadFileArgs, _ fantasy.ToolCall) (fantasy.ToolResponse, error) {
if options.GetWorkspaceConn == nil {
return fantasy.NewTextErrorResponse("workspace connection resolver is not configured"), nil
}
conn, err := options.GetWorkspaceConn(ctx)
if err != nil {
return fantasy.NewTextErrorResponse(err.Error()), nil
}
return executeReadFileTool(ctx, conn, args)
},
)
}
func executeReadFileTool(
ctx context.Context,
conn workspacesdk.AgentConn,
args ReadFileArgs,
) (fantasy.ToolResponse, error) {
if args.Path == "" {
return fantasy.NewTextErrorResponse("path is required"), nil
}
offset := int64(1) // 1-based line number default
limit := int64(0) // 0 means use server default (2000)
if args.Offset != nil {
offset = *args.Offset
}
if args.Limit != nil {
limit = *args.Limit
}
resp, err := conn.ReadFileLines(ctx, args.Path, offset, limit, workspacesdk.DefaultReadFileLinesLimits())
if err != nil {
return fantasy.NewTextErrorResponse(err.Error()), nil
}
if !resp.Success {
return fantasy.NewTextErrorResponse(resp.Error), nil
}
return toolResponse(map[string]any{
"content": resp.Content,
"file_size": resp.FileSize,
"total_lines": resp.TotalLines,
"lines_read": resp.LinesRead,
}), nil
}
+130
View File
@@ -0,0 +1,130 @@
package chattool
import (
"context"
"encoding/json"
"strings"
"charm.land/fantasy"
"github.com/google/uuid"
"golang.org/x/xerrors"
"github.com/coder/coder/v2/coderd/database"
)
// ReadTemplateOptions configures the read_template tool.
type ReadTemplateOptions struct {
DB database.Store
OwnerID uuid.UUID
}
type readTemplateArgs struct {
TemplateID string `json:"template_id"`
}
// ReadTemplate returns a tool that retrieves details about a specific
// template, including its configurable rich parameters. The agent
// uses this after list_templates and before create_workspace.
func ReadTemplate(options ReadTemplateOptions) fantasy.AgentTool {
return fantasy.NewAgentTool(
"read_template",
"Get details about a workspace template, including its "+
"configurable parameters. Use this after finding a "+
"template with list_templates and before creating a "+
"workspace with create_workspace.",
func(ctx context.Context, args readTemplateArgs, _ fantasy.ToolCall) (fantasy.ToolResponse, error) {
if options.DB == nil {
return fantasy.NewTextErrorResponse("database is not configured"), nil
}
templateIDStr := strings.TrimSpace(args.TemplateID)
if templateIDStr == "" {
return fantasy.NewTextErrorResponse("template_id is required"), nil
}
templateID, err := uuid.Parse(templateIDStr)
if err != nil {
return fantasy.NewTextErrorResponse(
xerrors.Errorf("invalid template_id: %w", err).Error(),
), nil
}
ctx, err = asOwner(ctx, options.DB, options.OwnerID)
if err != nil {
return fantasy.NewTextErrorResponse(err.Error()), nil
}
template, err := options.DB.GetTemplateByID(ctx, templateID)
if err != nil {
return fantasy.NewTextErrorResponse("template not found"), nil
}
params, err := options.DB.GetTemplateVersionParameters(ctx, template.ActiveVersionID)
if err != nil {
return fantasy.NewTextErrorResponse(
xerrors.Errorf("failed to get template parameters: %w", err).Error(),
), nil
}
templateInfo := map[string]any{
"id": template.ID.String(),
"name": template.Name,
"active_version_id": template.ActiveVersionID.String(),
}
if display := strings.TrimSpace(template.DisplayName); display != "" {
templateInfo["display_name"] = display
}
if desc := strings.TrimSpace(template.Description); desc != "" {
templateInfo["description"] = desc
}
paramList := make([]map[string]any, 0, len(params))
for _, p := range params {
param := map[string]any{
"name": p.Name,
"type": p.Type,
"required": p.Required,
}
if display := strings.TrimSpace(p.DisplayName); display != "" {
param["display_name"] = display
}
if desc := strings.TrimSpace(p.Description); desc != "" {
param["description"] = truncateRunes(desc, 300)
}
if p.DefaultValue != "" {
param["default"] = p.DefaultValue
}
if p.Mutable {
param["mutable"] = true
}
if p.Ephemeral {
param["ephemeral"] = true
}
if p.FormType != "" {
param["form_type"] = string(p.FormType)
}
if len(p.Options) > 0 && string(p.Options) != "null" && string(p.Options) != "[]" {
var opts []map[string]any
if err := json.Unmarshal(p.Options, &opts); err == nil && len(opts) > 0 {
param["options"] = opts
}
}
if p.ValidationRegex != "" {
param["validation_regex"] = p.ValidationRegex
}
if p.ValidationMin.Valid {
param["validation_min"] = p.ValidationMin.Int32
}
if p.ValidationMax.Valid {
param["validation_max"] = p.ValidationMax.Int32
}
paramList = append(paramList, param)
}
return toolResponse(map[string]any{
"template": templateInfo,
"parameters": paramList,
}), nil
},
)
}
+175
View File
@@ -0,0 +1,175 @@
package chattool
import (
"context"
"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 {
return fantasy.NewTextErrorResponse(
xerrors.Errorf("load workspace: %w", err).Error(),
), nil
}
if ws.Deleted {
return fantasy.NewTextErrorResponse(
"workspace was deleted; use create_workspace to make a new one",
), nil
}
build, 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,258 @@
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/database"
"github.com/coder/coder/v2/coderd/database/dbfake"
"github.com/coder/coder/v2/coderd/database/dbgen"
"github.com/coder/coder/v2/coderd/database/dbtestutil"
"github.com/coder/coder/v2/coderd/x/chatd/chattool"
"github.com/coder/coder/v2/codersdk"
"github.com/coder/coder/v2/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)
})
t.Run("DeletedWorkspace", func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitLong)
db, _ := dbtestutil.NewDB(t)
user := dbgen.User(t, db, database.User{})
modelCfg := seedModelConfig(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 workspace that has been soft-deleted.
wsResp := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{
OwnerID: user.ID,
OrganizationID: org.ID,
Deleted: true,
}).Seed(database.WorkspaceBuild{
Transition: database.WorkspaceTransitionDelete,
}).Do()
ws := wsResp.Workspace
chat, err := db.InsertChat(ctx, database.InsertChatParams{
OwnerID: user.ID,
WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true},
LastModelConfigID: modelCfg.ID,
Title: "test-deleted-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 for deleted 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)
require.Contains(t, resp.Content, "workspace was deleted")
})
}
// 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
}
+51
View File
@@ -0,0 +1,51 @@
package chattool
import (
"context"
"strings"
"charm.land/fantasy"
"github.com/coder/coder/v2/codersdk/workspacesdk"
)
type WriteFileOptions struct {
GetWorkspaceConn func(context.Context) (workspacesdk.AgentConn, error)
}
type WriteFileArgs struct {
Path string `json:"path"`
Content string `json:"content"`
}
func WriteFile(options WriteFileOptions) fantasy.AgentTool {
return fantasy.NewAgentTool(
"write_file",
"Write a file to the workspace.",
func(ctx context.Context, args WriteFileArgs, _ fantasy.ToolCall) (fantasy.ToolResponse, error) {
if options.GetWorkspaceConn == nil {
return fantasy.NewTextErrorResponse("workspace connection resolver is not configured"), nil
}
conn, err := options.GetWorkspaceConn(ctx)
if err != nil {
return fantasy.NewTextErrorResponse(err.Error()), nil
}
return executeWriteFileTool(ctx, conn, args)
},
)
}
func executeWriteFileTool(
ctx context.Context,
conn workspacesdk.AgentConn,
args WriteFileArgs,
) (fantasy.ToolResponse, error) {
if args.Path == "" {
return fantasy.NewTextErrorResponse("path is required"), nil
}
if err := conn.WriteFile(ctx, args.Path, strings.NewReader(args.Content)); err != nil {
return fantasy.NewTextErrorResponse(err.Error()), nil
}
return toolResponse(map[string]any{"ok": true}), nil
}