mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
fix(coderd/chatd): exclude |& from background detection, add tests (#23313)
The ampersand detection treated bash's pipe-stderr operator (|&) as a trailing & for backgrounding, stripping it and producing a broken pipe command. Also adds tests for execute.go and chatloop context limit helpers, covering previously untested edge cases.
This commit is contained in:
Mathias Fredriksson
@@ -0,0 +1,399 @@
|
||||
package chatloop
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"testing"
|
||||
|
||||
"charm.land/fantasy"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// testProviderData implements fantasy.ProviderOptionsData so we can
|
||||
// construct arbitrary ProviderMetadata for extractContextLimit tests.
|
||||
type testProviderData struct {
|
||||
data map[string]any
|
||||
}
|
||||
|
||||
func (*testProviderData) Options() {}
|
||||
|
||||
func (d *testProviderData) MarshalJSON() ([]byte, error) {
|
||||
return json.Marshal(d.data)
|
||||
}
|
||||
|
||||
// Required by the ProviderOptionsData interface; unused in tests.
|
||||
func (d *testProviderData) UnmarshalJSON(b []byte) error {
|
||||
return json.Unmarshal(b, &d.data)
|
||||
}
|
||||
|
||||
func TestNormalizeMetadataKey(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
key string
|
||||
want string
|
||||
}{
|
||||
{name: "lowercase", key: "camelCase", want: "camelcase"},
|
||||
{name: "hyphens stripped", key: "kebab-case", want: "kebabcase"},
|
||||
{name: "underscores stripped", key: "snake_case", want: "snakecase"},
|
||||
{name: "uppercase", key: "UPPER", want: "upper"},
|
||||
{name: "spaces stripped", key: "with spaces", want: "withspaces"},
|
||||
{name: "empty", key: "", want: ""},
|
||||
{name: "digits preserved", key: "123", want: "123"},
|
||||
{name: "mixed separators", key: "Max_Context-Tokens", want: "maxcontexttokens"},
|
||||
{name: "dots stripped", key: "context.limit", want: "contextlimit"},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
got := normalizeMetadataKey(tt.key)
|
||||
require.Equal(t, tt.want, got)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestIsContextLimitKey(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
key string
|
||||
want bool
|
||||
skip bool
|
||||
}{ // Exact matches after normalization.
|
||||
{name: "context_limit", key: "context_limit", want: true},
|
||||
{name: "context_window", key: "context_window", want: true},
|
||||
{name: "context_length", key: "context_length", want: true},
|
||||
{name: "max_context", key: "max_context", want: true},
|
||||
{name: "max_context_tokens", key: "max_context_tokens", want: true},
|
||||
{name: "max_input_tokens", key: "max_input_tokens", want: true},
|
||||
{name: "max_input_token", key: "max_input_token", want: true},
|
||||
{name: "input_token_limit", key: "input_token_limit", want: true},
|
||||
|
||||
// Case and separator variations.
|
||||
{name: "Context-Window mixed case", key: "Context-Window", want: true},
|
||||
{name: "MAX_CONTEXT_TOKENS screaming", key: "MAX_CONTEXT_TOKENS", want: true},
|
||||
{name: "contextLimit camelCase", key: "contextLimit", want: true},
|
||||
|
||||
// Fallback heuristic: contains "context" + limit/window/length.
|
||||
{name: "model_context_limit", key: "model_context_limit", want: true},
|
||||
{name: "context_window_size", key: "context_window_size", want: true},
|
||||
{name: "context_length_max", key: "context_length_max", want: true},
|
||||
|
||||
// Fallback heuristic: starts with "max" + contains "context".
|
||||
// BUG(isContextLimitKey): "max_context_version" matches
|
||||
// because it contains "context" and starts with "max",
|
||||
// but a version field is not a context limit.
|
||||
// TODO: Fix the heuristic and remove this skip.
|
||||
{name: "max_context_version false positive", key: "max_context_version", want: false, skip: true}, // Non-matching keys.
|
||||
{name: "context_id no limit keyword", key: "context_id", want: false},
|
||||
{name: "empty string", key: "", want: false},
|
||||
{name: "unrelated key", key: "model_name", want: false},
|
||||
{name: "limit without context", key: "rate_limit", want: false},
|
||||
{name: "max without context", key: "max_tokens", want: false},
|
||||
{name: "context alone", key: "context", want: false},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
if tt.skip {
|
||||
t.Skip("known bug: isContextLimitKey false positive")
|
||||
}
|
||||
got := isContextLimitKey(tt.key)
|
||||
require.Equal(t, tt.want, got)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestNumericContextLimitValue(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
value any
|
||||
want int64
|
||||
wantOK bool
|
||||
}{
|
||||
// float64: the default numeric type from json.Unmarshal.
|
||||
{name: "float64 integer", value: float64(128000), want: 128000, wantOK: true},
|
||||
{name: "float64 fractional rejected", value: float64(128000.5), want: 0, wantOK: false},
|
||||
{name: "float64 zero rejected", value: float64(0), want: 0, wantOK: false},
|
||||
{name: "float64 negative rejected", value: float64(-1), want: 0, wantOK: false},
|
||||
|
||||
// int64
|
||||
{name: "int64 positive", value: int64(200000), want: 200000, wantOK: true},
|
||||
{name: "int64 zero rejected", value: int64(0), want: 0, wantOK: false},
|
||||
{name: "int64 negative rejected", value: int64(-1), want: 0, wantOK: false},
|
||||
|
||||
// int32
|
||||
{name: "int32 positive", value: int32(50000), want: 50000, wantOK: true},
|
||||
{name: "int32 zero rejected", value: int32(0), want: 0, wantOK: false},
|
||||
|
||||
// int
|
||||
{name: "int positive", value: int(50000), want: 50000, wantOK: true},
|
||||
{name: "int zero rejected", value: int(0), want: 0, wantOK: false},
|
||||
|
||||
// string
|
||||
{name: "string numeric", value: "128000", want: 128000, wantOK: true},
|
||||
{name: "string trimmed", value: " 128000 ", want: 128000, wantOK: true},
|
||||
{name: "string non-numeric rejected", value: "not a number", want: 0, wantOK: false},
|
||||
{name: "string empty rejected", value: "", want: 0, wantOK: false},
|
||||
{name: "string zero rejected", value: "0", want: 0, wantOK: false},
|
||||
{name: "string negative rejected", value: "-1", want: 0, wantOK: false},
|
||||
|
||||
// json.Number
|
||||
{name: "json.Number valid", value: json.Number("200000"), want: 200000, wantOK: true},
|
||||
{name: "json.Number invalid rejected", value: json.Number("invalid"), want: 0, wantOK: false},
|
||||
{name: "json.Number zero rejected", value: json.Number("0"), want: 0, wantOK: false},
|
||||
|
||||
// Unhandled types.
|
||||
{name: "bool rejected", value: true, want: 0, wantOK: false},
|
||||
{name: "nil rejected", value: nil, want: 0, wantOK: false},
|
||||
{name: "slice rejected", value: []int{1}, want: 0, wantOK: false},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
got, ok := numericContextLimitValue(tt.value)
|
||||
require.Equal(t, tt.wantOK, ok)
|
||||
require.Equal(t, tt.want, got)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestPositiveInt64(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
got, ok := positiveInt64(42)
|
||||
require.True(t, ok)
|
||||
require.Equal(t, int64(42), got)
|
||||
|
||||
got, ok = positiveInt64(0)
|
||||
require.False(t, ok)
|
||||
require.Equal(t, int64(0), got)
|
||||
|
||||
got, ok = positiveInt64(-1)
|
||||
require.False(t, ok)
|
||||
require.Equal(t, int64(0), got)
|
||||
}
|
||||
|
||||
func TestCollectContextLimitValues(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
t.Run("FlatMap", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
input := map[string]any{
|
||||
"context_limit": float64(200000),
|
||||
"other_key": float64(999),
|
||||
}
|
||||
var collected []int64
|
||||
collectContextLimitValues(input, func(v int64) {
|
||||
collected = append(collected, v)
|
||||
})
|
||||
require.Equal(t, []int64{200000}, collected)
|
||||
})
|
||||
|
||||
t.Run("NestedMaps", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
input := map[string]any{
|
||||
"provider": map[string]any{
|
||||
"info": map[string]any{
|
||||
"context_window": float64(100000),
|
||||
},
|
||||
},
|
||||
}
|
||||
var collected []int64
|
||||
collectContextLimitValues(input, func(v int64) {
|
||||
collected = append(collected, v)
|
||||
})
|
||||
require.Equal(t, []int64{100000}, collected)
|
||||
})
|
||||
|
||||
t.Run("ArrayTraversal", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
input := []any{
|
||||
map[string]any{"context_limit": float64(50000)},
|
||||
map[string]any{"context_limit": float64(80000)},
|
||||
}
|
||||
var collected []int64
|
||||
collectContextLimitValues(input, func(v int64) {
|
||||
collected = append(collected, v)
|
||||
})
|
||||
require.Len(t, collected, 2)
|
||||
require.Contains(t, collected, int64(50000))
|
||||
require.Contains(t, collected, int64(80000))
|
||||
})
|
||||
|
||||
t.Run("MixedNesting", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
input := map[string]any{
|
||||
"models": []any{
|
||||
map[string]any{
|
||||
"context_limit": float64(128000),
|
||||
},
|
||||
},
|
||||
}
|
||||
var collected []int64
|
||||
collectContextLimitValues(input, func(v int64) {
|
||||
collected = append(collected, v)
|
||||
})
|
||||
require.Equal(t, []int64{128000}, collected)
|
||||
})
|
||||
|
||||
t.Run("NonMatchingKey", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
input := map[string]any{
|
||||
"model_name": "gpt-4",
|
||||
"tokens": float64(1000),
|
||||
}
|
||||
var collected []int64
|
||||
collectContextLimitValues(input, func(v int64) {
|
||||
collected = append(collected, v)
|
||||
})
|
||||
require.Empty(t, collected)
|
||||
})
|
||||
|
||||
t.Run("ScalarIgnored", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
var collected []int64
|
||||
collectContextLimitValues("just a string", func(v int64) {
|
||||
collected = append(collected, v)
|
||||
})
|
||||
require.Empty(t, collected)
|
||||
})
|
||||
}
|
||||
|
||||
func TestFindContextLimitValue(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
t.Run("SingleCandidate", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
input := map[string]any{
|
||||
"context_limit": float64(200000),
|
||||
}
|
||||
limit, ok := findContextLimitValue(input)
|
||||
require.True(t, ok)
|
||||
require.Equal(t, int64(200000), limit)
|
||||
})
|
||||
|
||||
t.Run("MultipleCandidatesTakesMax", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
input := map[string]any{
|
||||
"a": map[string]any{"context_limit": float64(50000)},
|
||||
"b": map[string]any{"context_limit": float64(200000)},
|
||||
}
|
||||
limit, ok := findContextLimitValue(input)
|
||||
require.True(t, ok)
|
||||
require.Equal(t, int64(200000), limit)
|
||||
})
|
||||
|
||||
t.Run("NoCandidates", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
input := map[string]any{
|
||||
"model": "gpt-4",
|
||||
}
|
||||
_, ok := findContextLimitValue(input)
|
||||
require.False(t, ok)
|
||||
})
|
||||
|
||||
t.Run("NilInput", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
_, ok := findContextLimitValue(nil)
|
||||
require.False(t, ok)
|
||||
})
|
||||
}
|
||||
|
||||
func TestExtractContextLimit(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
t.Run("AnthropicStyle", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
metadata := fantasy.ProviderMetadata{
|
||||
"anthropic": &testProviderData{
|
||||
data: map[string]any{
|
||||
"cache_read_input_tokens": float64(100),
|
||||
"context_limit": float64(200000),
|
||||
},
|
||||
},
|
||||
}
|
||||
result := extractContextLimit(metadata)
|
||||
require.True(t, result.Valid)
|
||||
require.Equal(t, int64(200000), result.Int64)
|
||||
})
|
||||
|
||||
t.Run("OpenAIStyle", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
metadata := fantasy.ProviderMetadata{
|
||||
"openai": &testProviderData{
|
||||
data: map[string]any{
|
||||
"max_context_tokens": float64(128000),
|
||||
},
|
||||
},
|
||||
}
|
||||
result := extractContextLimit(metadata)
|
||||
require.True(t, result.Valid)
|
||||
require.Equal(t, int64(128000), result.Int64)
|
||||
})
|
||||
|
||||
t.Run("NestedDeeply", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
metadata := fantasy.ProviderMetadata{
|
||||
"provider": &testProviderData{
|
||||
data: map[string]any{
|
||||
"info": map[string]any{
|
||||
"context_window": float64(100000),
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
result := extractContextLimit(metadata)
|
||||
require.True(t, result.Valid)
|
||||
require.Equal(t, int64(100000), result.Int64)
|
||||
})
|
||||
|
||||
t.Run("MultipleCandidatesTakesMax", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
metadata := fantasy.ProviderMetadata{
|
||||
"a": &testProviderData{
|
||||
data: map[string]any{
|
||||
"context_limit": float64(50000),
|
||||
},
|
||||
},
|
||||
"b": &testProviderData{
|
||||
data: map[string]any{
|
||||
"context_limit": float64(200000),
|
||||
},
|
||||
},
|
||||
}
|
||||
result := extractContextLimit(metadata)
|
||||
require.True(t, result.Valid)
|
||||
require.Equal(t, int64(200000), result.Int64)
|
||||
})
|
||||
|
||||
t.Run("NoMatchingKeys", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
metadata := fantasy.ProviderMetadata{
|
||||
"openai": &testProviderData{
|
||||
data: map[string]any{
|
||||
"model": "gpt-4",
|
||||
"tokens": float64(1000),
|
||||
},
|
||||
},
|
||||
}
|
||||
result := extractContextLimit(metadata)
|
||||
assert.False(t, result.Valid)
|
||||
})
|
||||
|
||||
t.Run("NilMetadata", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
result := extractContextLimit(nil)
|
||||
assert.False(t, result.Valid)
|
||||
})
|
||||
|
||||
t.Run("EmptyMetadata", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
result := extractContextLimit(fantasy.ProviderMetadata{})
|
||||
assert.False(t, result.Valid)
|
||||
})
|
||||
}
|
||||
@@ -127,7 +127,7 @@ func executeTool(
|
||||
// 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, "&&") {
|
||||
if !background && strings.HasSuffix(trimmed, "&") && !strings.HasSuffix(trimmed, "&&") && !strings.HasSuffix(trimmed, "|&") {
|
||||
background = true
|
||||
args.Command = strings.TrimSpace(strings.TrimSuffix(trimmed, "&"))
|
||||
}
|
||||
|
||||
@@ -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").
|
||||
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,489 @@
|
||||
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/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").
|
||||
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").
|
||||
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").
|
||||
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)
|
||||
|
||||
// ProcessOutput always returns running. The poll loop
|
||||
// and the timeout-branch recovery call both hit this.
|
||||
mockConn.EXPECT().
|
||||
ProcessOutput(gomock.Any(), "proc-1").
|
||||
Return(workspacesdk.ProcessOutputResponse{
|
||||
Running: true,
|
||||
Output: "partial output",
|
||||
}, nil).
|
||||
AnyTimes()
|
||||
|
||||
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 before the 200ms poll interval,
|
||||
// so the context-done branch fires first.
|
||||
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)
|
||||
mockConn.EXPECT().
|
||||
ProcessOutput(gomock.Any(), "proc-1").
|
||||
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")
|
||||
})
|
||||
|
||||
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").
|
||||
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
|
||||
}
|
||||
Reference in New Issue
Block a user