From 6afc1bac0bc1a04a7b61623d40e1c4c854f0c106 Mon Sep 17 00:00:00 2001 From: Mathias Fredriksson Date: Fri, 20 Mar 2026 00:18:23 +0200 Subject: [PATCH] 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. --- .../chatloop/contextlimit_internal_test.go | 399 ++++++++++++++ coderd/chatd/chattool/execute.go | 2 +- .../chatd/chattool/execute_internal_test.go | 100 ++++ coderd/chatd/chattool/execute_test.go | 489 ++++++++++++++++++ 4 files changed, 989 insertions(+), 1 deletion(-) create mode 100644 coderd/chatd/chatloop/contextlimit_internal_test.go create mode 100644 coderd/chatd/chattool/execute_internal_test.go create mode 100644 coderd/chatd/chattool/execute_test.go diff --git a/coderd/chatd/chatloop/contextlimit_internal_test.go b/coderd/chatd/chatloop/contextlimit_internal_test.go new file mode 100644 index 0000000000..85e490f425 --- /dev/null +++ b/coderd/chatd/chatloop/contextlimit_internal_test.go @@ -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) + }) +} diff --git a/coderd/chatd/chattool/execute.go b/coderd/chatd/chattool/execute.go index d22e65bea0..d7b2203ff4 100644 --- a/coderd/chatd/chattool/execute.go +++ b/coderd/chatd/chattool/execute.go @@ -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, "&")) } diff --git a/coderd/chatd/chattool/execute_internal_test.go b/coderd/chatd/chattool/execute_internal_test.go new file mode 100644 index 0000000000..e9100f3051 --- /dev/null +++ b/coderd/chatd/chattool/execute_internal_test.go @@ -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 +} diff --git a/coderd/chatd/chattool/execute_test.go b/coderd/chatd/chattool/execute_test.go new file mode 100644 index 0000000000..c89086c219 --- /dev/null +++ b/coderd/chatd/chattool/execute_test.go @@ -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 +}