mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
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:
@@ -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])
|
||||
}
|
||||
@@ -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)
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
},
|
||||
)
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
Reference in New Issue
Block a user