Skip to content

fix(mcp): bound and cancel Windows code search #1608

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ertankucukoglu wants to merge 1 commit into
base: main
Choose a base branch
from
+240 −4
Open

fix(mcp): bound and cancel Windows code search #1608

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
109 changes: 105 additions & 4 deletions src/mcp/mcp.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -123,6 +123,7 @@ enum {
* growing getline buffers without bound through ignored extension headers. */
#define MCP_MAX_MESSAGE_SIZE ((size_t)10U * 1024U * 1024U)
#define MCP_MAX_HEADER_SIZE ((size_t)8U * 1024U)
#define MCP_SEARCH_OUTPUT_MAX ((size_t)64U * 1024U * 1024U)

/* ── Helpers ────────────────────────────────────────────────────── */

Expand Down Expand Up @@ -1567,6 +1568,7 @@ struct cbm_mcp_server {
void *quarantine_test_context;
cbm_mcp_command_test_hook_fn command_test_hook;
void *command_test_context;
size_t search_output_limit_override;
cbm_thread_t autoindex_tid;
bool autoindex_active; /* true if auto-index thread was started */

Expand Down Expand Up @@ -1865,6 +1867,12 @@ void cbm_mcp_server_set_command_test_hook(cbm_mcp_server_t *srv, cbm_mcp_command
srv->command_test_context = context;
}

void cbm_mcp_server_set_search_output_limit_for_test(cbm_mcp_server_t *srv, size_t limit) {
if (srv) {
srv->search_output_limit_override = limit;
}
}

/* ── Cache dir + project DB path helpers ───────────────────────── */

/* Returns the cache directory. Writes to buf, returns buf for convenience. */
Expand Down Expand Up @@ -9634,6 +9642,32 @@ static bool compile_path_filter(const char *filter, cbm_regex_t *re) {
return cbm_regcomp(re, filter, CBM_REG_EXTENDED | CBM_REG_NOSUB) == CBM_REG_OK;
}

static int mcp_run_shell_command_cancellable_bounded(cbm_mcp_server_t *srv, const char *command,
char output_path[CBM_SZ_2K],
size_t output_limit,
bool *output_limit_exceeded,
cbm_proc_result_t *result_out);

#ifdef _WIN32
static char *search_code_scan_error(search_scratch_t *scratch, const char *output_path,
bool has_path_filter, cbm_regex_t *path_regex, char *root_path,
char *pattern, char *project, char *file_pattern,
const char *message) {
if (output_path && output_path[0]) {
(void)cbm_unlink(output_path);
}
search_scratch_close(scratch);
if (has_path_filter) {
cbm_regfree(path_regex);
}
free(root_path);
free(pattern);
free(project);
free(file_pattern);
return cbm_mcp_text_result(message, true);
}
#endif

static char *handle_search_code(cbm_mcp_server_t *srv, const char *args) {
char *pattern = cbm_mcp_get_string_arg(args, "pattern");
char *project = get_project_arg(args);
Expand Down Expand Up @@ -9810,6 +9844,44 @@ static char *handle_search_code(cbm_mcp_server_t *srv, const char *args) {
build_grep_cmd(cmd, sizeof(cmd), use_regex, scoped, file_pattern, tmpfile, filelist,
root_path);

#ifdef _WIN32
char output_path[CBM_SZ_2K] = {0};
cbm_proc_result_t scan_result = {0};
bool scan_output_exceeded = false;
size_t scan_output_limit = srv->search_output_limit_override
? srv->search_output_limit_override
: MCP_SEARCH_OUTPUT_MAX;
int scan_run = mcp_run_shell_command_cancellable_bounded(
srv, cmd, output_path, scan_output_limit, &scan_output_exceeded, &scan_result);
if (scan_output_exceeded) {
char message[CBM_SZ_128];
snprintf(message, sizeof(message),
"search failed: output exceeded the %zu-byte safety limit", scan_output_limit);
return search_code_scan_error(&scratch, output_path, has_path_filter, &path_regex,
root_path, pattern, project, file_pattern, message);
}
bool scan_cancelled = scan_result.cancellation_requested || mcp_request_cancelled(srv);
if (scan_cancelled) {
return search_code_scan_error(&scratch, output_path, has_path_filter, &path_regex,
root_path, pattern, project, file_pattern,
"search_code cancelled for this request");
}
if (scan_run != 0) {
return search_code_scan_error(
&scratch, output_path, has_path_filter, &path_regex, root_path, pattern, project,
file_pattern, "search failed: the contained command could not complete");
}
FILE *fp = cbm_fopen(output_path, "rb");
if (!fp) {
return search_code_scan_error(&scratch, output_path, has_path_filter, &path_regex,
root_path, pattern, project, file_pattern,
"search failed: contained output could not be read");
}
gm = collect_grep_matches(fp, root_path, strlen(root_path), has_path_filter, &path_regex,
grep_limit, &gm_count);
(void)fclose(fp);
(void)cbm_unlink(output_path);
#else
FILE *fp = cbm_popen(cmd, "r");
if (!fp) {
search_scratch_close(&scratch);
Expand All @@ -9823,6 +9895,7 @@ static char *handle_search_code(cbm_mcp_server_t *srv, const char *args) {
gm = collect_grep_matches(fp, root_path, strlen(root_path), has_path_filter, &path_regex,
grep_limit, &gm_count);
cbm_pclose(fp);
#endif
/* Both scratch files and the private directory go here — unlike the old
* code, the file list is removed even when the scan was not scoped. */
search_scratch_close(&scratch);
Expand Down Expand Up @@ -9976,10 +10049,16 @@ static bool mcp_resolve_windows_cmd(char out[CBM_SZ_4K]) {
}
#endif

static int mcp_run_shell_command_cancellable(cbm_mcp_server_t *srv, const char *command,
char output_path[CBM_SZ_2K],
cbm_proc_result_t *result_out) {
if (!srv || !command || !output_path || !result_out || !mcp_command_output_path(output_path)) {
static int mcp_run_shell_command_cancellable_bounded(cbm_mcp_server_t *srv, const char *command,
char output_path[CBM_SZ_2K],
size_t output_limit,
bool *output_limit_exceeded,
cbm_proc_result_t *result_out) {
if (output_limit_exceeded) {
*output_limit_exceeded = false;
}
if (!srv || !command || !output_path || !result_out ||
(output_limit > 0 && !output_limit_exceeded) || !mcp_command_output_path(output_path)) {
return -1;
}
/* Internal test seam: rejecting after output allocation exercises the same
Expand Down Expand Up @@ -10018,10 +10097,18 @@ static int mcp_run_shell_command_cancellable(cbm_mcp_server_t *srv, const char *
}

cbm_proc_poll_t state;
bool limit_exceeded = false;
for (;;) {
if (mcp_request_cancelled(srv)) {
(void)cbm_subprocess_request_cancel(process);
}
if (!limit_exceeded && output_limit > 0) {
int64_t output_size = cbm_file_size(output_path);
if (output_size > 0 && (uint64_t)output_size > output_limit) {
limit_exceeded = true;
(void)cbm_subprocess_request_cancel(process);
}
}
state = cbm_subprocess_poll(process, result_out);
if (state != CBM_PROC_POLL_RUNNING) {
break;
Expand All @@ -10031,9 +10118,23 @@ static int mcp_run_shell_command_cancellable(cbm_mcp_server_t *srv, const char *
bool contained = state == CBM_PROC_POLL_TERMINAL && result_out->tree_quiesced &&
!result_out->supervision_failed;
cbm_subprocess_destroy(process);
if (!limit_exceeded && output_limit > 0) {
int64_t final_size = cbm_file_size(output_path);
limit_exceeded = final_size > 0 && (uint64_t)final_size > output_limit;
}
if (output_limit_exceeded) {
*output_limit_exceeded = limit_exceeded;
}
return contained ? 0 : -1;
}

static int mcp_run_shell_command_cancellable(cbm_mcp_server_t *srv, const char *command,
char output_path[CBM_SZ_2K],
cbm_proc_result_t *result_out) {
return mcp_run_shell_command_cancellable_bounded(srv, command, output_path, 0, NULL,
result_out);
}

/* Does `node`'s line range overlap any recorded hunk for `file`? Used to scope
* seed detection to the actually-changed lines rather than the whole file.
* Non-static (declared in mcp_internal.h) so tests can exercise the overlap
Expand Down
1 change: 1 addition & 0 deletions src/mcp/mcp_internal.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@ void cbm_mcp_server_set_quarantine_test_hook(cbm_mcp_server_t *srv,
cbm_mcp_quarantine_test_hook_fn hook, void *context);
void cbm_mcp_server_set_command_test_hook(cbm_mcp_server_t *srv, cbm_mcp_command_test_hook_fn hook,
void *context);
void cbm_mcp_server_set_search_output_limit_for_test(cbm_mcp_server_t *srv, size_t limit);

/* Release only the constructor-created pristine in-memory store. Public
* cbm_mcp_server_new(NULL) semantics remain unchanged; daemon sessions use
Expand Down
134 changes: 134 additions & 0 deletions tests/test_mcp.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -162,6 +162,21 @@ typedef struct {
int merge_base_calls;
} mcp_command_hook_probe_t;

#ifdef _WIN32
typedef struct {
cbm_mcp_server_t *server;
bool cancel_on_call;
bool cancel_accepted;
int calls;
char command[CBM_SZ_4K];
} mcp_search_command_probe_t;

typedef struct {
char path[512];
char *saved_cache;
} mcp_search_cache_t;
#endif

static bool mcp_quarantine_hook_probe(void *context, const char *step) {
mcp_quarantine_hook_probe_t *probe = context;
if (!probe || !step) {
Expand All @@ -187,6 +202,45 @@ static bool mcp_command_hook_probe(void *context, const char *command) {
return true;
}

#ifdef _WIN32
static bool mcp_search_command_hook_probe(void *context, const char *command) {
mcp_search_command_probe_t *probe = context;
if (!probe || !command) {
return false;
}
probe->calls++;
snprintf(probe->command, sizeof(probe->command), "%s", command);
if (probe->cancel_on_call && probe->server) {
probe->cancel_accepted = cbm_mcp_server_cancel_active(probe->server);
}
return true;
}

static bool mcp_search_cache_open(mcp_search_cache_t *cache, const char *prefix) {
memset(cache, 0, sizeof(*cache));
snprintf(cache->path, sizeof(cache->path), "%s/%s-XXXXXX", cbm_tmpdir(), prefix);
if (!cbm_mkdtemp(cache->path)) {
return false;
}
const char *saved_cache = getenv("CBM_CACHE_DIR");
cache->saved_cache = saved_cache ? strdup(saved_cache) : NULL;
if ((saved_cache && !cache->saved_cache) || cbm_setenv("CBM_CACHE_DIR", cache->path, 1) != 0) {
free(cache->saved_cache);
cache->saved_cache = NULL;
(void)th_rmtree(cache->path);
return false;
}
return true;
}

static bool mcp_search_cache_close(mcp_search_cache_t *cache) {
restore_cache_dir(cache->saved_cache);
free(cache->saved_cache);
cache->saved_cache = NULL;
return th_rmtree(cache->path) == 0;
}
#endif

typedef struct {
const char *name;
char *value;
Expand Down Expand Up @@ -4501,6 +4555,84 @@ TEST(search_code_path_filter_matches_nothing) {
PASS();
}

TEST(search_code_windows_cancel_cleans_supervised_scan) {
#ifdef _WIN32
mcp_search_cache_t cache;
ASSERT_TRUE(mcp_search_cache_open(&cache, "cbm-search-cancel"));

char tmp[512], src_path[768], vendor_path[768];
cbm_mcp_server_t *srv = setup_prefilter_server(tmp, sizeof(tmp), src_path, sizeof(src_path),
vendor_path, sizeof(vendor_path));
ASSERT_NOT_NULL(srv);
mcp_search_command_probe_t probe = {
.server = srv,
.cancel_on_call = true,
};
cbm_mcp_server_set_command_test_hook(srv, mcp_search_command_hook_probe, &probe);

char *response =
cbm_mcp_handle_tool(srv, "search_code",
"{\"pattern\":\"HandleRequest\",\"project\":\"prefilter-search\","
"\"file_pattern\":\"*.go\"}");
ASSERT_NOT_NULL(response);
ASSERT_TRUE(probe.cancel_accepted);
ASSERT_NOT_NULL(strstr(response, "cancelled"));
ASSERT_NOT_NULL(strstr(response, "\"isError\":true"));

char logs[640];
snprintf(logs, sizeof(logs), "%s/logs", cache.path);
ASSERT_EQ(mcp_count_directory_entries_with_prefix(logs, ".mcp-command-"), 0);

free(response);
cbm_mcp_server_free(srv);
cleanup_prefilter_dir(tmp, src_path, vendor_path);
ASSERT_TRUE(mcp_search_cache_close(&cache));
PASS();
#else
SKIP_PLATFORM("supervised Select-String cancellation runs on Windows");
#endif
}

TEST(search_code_windows_output_limit_fails_closed_and_cleans_scan) {
#ifdef _WIN32
mcp_search_cache_t cache;
ASSERT_TRUE(mcp_search_cache_open(&cache, "cbm-search-limit"));

char tmp[512], src_path[768], vendor_path[768];
cbm_mcp_server_t *srv = setup_prefilter_server(tmp, sizeof(tmp), src_path, sizeof(src_path),
vendor_path, sizeof(vendor_path));
ASSERT_NOT_NULL(srv);
cbm_mcp_server_set_search_output_limit_for_test(srv, 512);

FILE *source = cbm_fopen(src_path, "ab");
ASSERT_NOT_NULL(source);
for (int i = 0; i < 256; i++) {
ASSERT_GT(fprintf(source, "func HandleRequest%d() error { return nil }\n", i), 0);
}
ASSERT_EQ(fclose(source), 0);

char *response =
cbm_mcp_handle_tool(srv, "search_code",
"{\"pattern\":\"HandleRequest\",\"project\":\"prefilter-search\","
"\"file_pattern\":\"*.go\"}");
ASSERT_NOT_NULL(response);
ASSERT_NOT_NULL(strstr(response, "output exceeded"));
ASSERT_NOT_NULL(strstr(response, "\"isError\":true"));

char logs[640];
snprintf(logs, sizeof(logs), "%s/logs", cache.path);
ASSERT_EQ(mcp_count_directory_entries_with_prefix(logs, ".mcp-command-"), 0);

free(response);
cbm_mcp_server_free(srv);
cleanup_prefilter_dir(tmp, src_path, vendor_path);
ASSERT_TRUE(mcp_search_cache_close(&cache));
PASS();
#else
SKIP_PLATFORM("supervised Select-String output limit runs on Windows");
#endif
}

/* issue #283: search_code with regex=true and a syntactically invalid pattern
* must return an explicit error, not an empty result indistinguishable from a
* legitimate no-match. */
Expand Down Expand Up @@ -10722,6 +10854,8 @@ SUITE(mcp) {
#endif
RUN_TEST(search_code_path_filter_prefilter_keeps_matches);
RUN_TEST(search_code_path_filter_matches_nothing);
RUN_TEST(search_code_windows_cancel_cleans_supervised_scan);
RUN_TEST(search_code_windows_output_limit_fails_closed_and_cleans_scan);
RUN_TEST(search_code_invalid_regex_errors_issue283);
RUN_TEST(search_code_literal_pipe_warns_issue282);
RUN_TEST(search_code_ampersand_accepted_issue272);
Expand Down
Loading