diff --git a/docs/llm_features.md b/docs/llm_features.md index 8bd3684..5908ad2 100644 --- a/docs/llm_features.md +++ b/docs/llm_features.md @@ -52,18 +52,19 @@ LLMs sometimes need the structural match but also a few surrounding lines of con `fast-mcp` already exists and exposes structural search and rewrite tools over stdio. - **Requirement**: Improve host integration guidance and expand tool coverage where needed. - **Tools to expose**: - - `search_ruby_ast(pattern, dir)`: Search for a RuboCop AST pattern natively, returning JSON results. - - `get_method(method_name, dir)`: Shortcut tool that extracts a specific method by name. - - `get_class(class_name, dir)`: Shortcut tool that extracts the body of a specific class. + - `search_ruby_ast(pattern, dir, offset, limit)`: Search for a RuboCop AST pattern natively, returning paginated JSON results (default limit: 20). + - `get_method(method_name, dir, class_name, offset, limit)`: Shortcut tool that extracts a specific method by name, with pagination. + - `get_class(class_name, dir, offset, limit)`: Shortcut tool that extracts the body of a specific class, with pagination. - `rewrite_ruby_file(file, pattern, replacement)`: Apply a Fast replacement to a Ruby file in-place. - `run_fast_experiment(name, lookup, search, edit, policy)`: Use Fast experiments to apply iterative code refactorings safely, validated automatically via test policies. - **Next gaps**: + - Pagination and metadata: Added `offset`, `limit`, `total`, and `has_more` to search tools. - Add examples for Codex, Claude Desktop, and other MCP-capable hosts. - Consider exposing resources/templates only if they add value beyond tools. - Improve scoping for `ruby_method_source` with `class_name` so it is lexical, not file-level. ## 5. Token Limit Awareness (`--max-tokens` / `--truncate`) -If a query matches an entire 3000-line class, it might blow out an LLM's context window. +The MCP server now supports `limit` and `offset` for search results to avoid hitting token limits in a single response. - **Requirement**: A flag `fast "(class ...)" --max-tokens=1000` that will purposefully truncate the body of large AST nodes (perhaps retaining signatures but omitting block bodies) to fit within token boundaries. ## 6. Auto-disable ANSI colors for non-TTY diff --git a/docs/mcp_tutorial.md b/docs/mcp_tutorial.md index 4f36a97..52aa7e8 100644 --- a/docs/mcp_tutorial.md +++ b/docs/mcp_tutorial.md @@ -65,17 +65,22 @@ printf '%s\n' \ | ruby -Ilib bin/fast-mcp ``` -This returns a JSON payload whose `result.content[0].text` is itself a JSON array of matches. Each match includes `file`, `line_start`, `line_end`, and the trimmed `code` snippet. +This returns a JSON payload whose `result.content[0].text` is itself a JSON object containing: +- `matches`: A JSON array of matches. Each match includes `file`, `line_start`, `line_end`, and the trimmed `code` snippet. +- `total`: Total number of matches found. +- `offset`: The current offset used. +- `limit`: The maximum number of matches returned. +- `has_more`: Boolean indicating if there are more results available. -### 3. Extract a method from a known class +### 3. Extract a method from a known class with pagination ```bash printf '%s\n' \ - '{"jsonrpc":"2.0","id":3,"method":"tools/call","params":{"name":"ruby_method_source","arguments":{"method_name":"run","class_name":"McpServer","paths":["lib"]}}}' \ + '{"jsonrpc":"2.0","id":3,"method":"tools/call","params":{"name":"ruby_method_source","arguments":{"method_name":"run","class_name":"McpServer","paths":["lib"],"limit":5,"offset":0}}}' \ | ruby -Ilib bin/fast-mcp ``` -This is often more token-efficient than returning the full class body. +This is often more token-efficient than returning the full class body. The `limit` defaults to 20. ### 4. Search with a raw AST pattern diff --git a/lib/fast/mcp_server.rb b/lib/fast/mcp_server.rb index dc5adab..e61dec9 100644 --- a/lib/fast/mcp_server.rb +++ b/lib/fast/mcp_server.rb @@ -29,7 +29,9 @@ class McpServer properties: { pattern: { type: 'string', description: 'Fast AST pattern, e.g. "(def match?)" or "(send nil :raise ...)".' }, paths: { type: 'array', items: { type: 'string' }, description: 'Files or directories to search.' }, - show_ast: { type: 'boolean', description: 'Include s-expression AST in results (default: false).' } + show_ast: { type: 'boolean', description: 'Include s-expression AST in results (default: false).' }, + offset: { type: 'integer', description: 'Offset for pagination (default: 0).' }, + limit: { type: 'integer', description: 'Maximum number of results to return (default: 20).' } }, required: ['pattern', 'paths'] } @@ -43,7 +45,9 @@ class McpServer method_name: { type: 'string', description: 'Method name, e.g. "initialize".' }, paths: { type: 'array', items: { type: 'string' }, description: 'Files or directories to search.' }, class_name: { type: 'string', description: 'Optional class name to restrict results, e.g. "Matcher".' }, - show_ast: { type: 'boolean', description: 'Include s-expression AST in results (default: false).' } + show_ast: { type: 'boolean', description: 'Include s-expression AST in results (default: false).' }, + offset: { type: 'integer', description: 'Offset for pagination (default: 0).' }, + limit: { type: 'integer', description: 'Maximum number of results to return (default: 20).' } }, required: ['method_name', 'paths'] } @@ -56,7 +60,9 @@ class McpServer properties: { class_name: { type: 'string', description: 'Class name to extract, e.g. "Rewriter".' }, paths: { type: 'array', items: { type: 'string' }, description: 'Files or directories to search.' }, - show_ast: { type: 'boolean', description: 'Include s-expression AST in results (default: false).' } + show_ast: { type: 'boolean', description: 'Include s-expression AST in results (default: false).' }, + offset: { type: 'integer', description: 'Offset for pagination (default: 0).' }, + limit: { type: 'integer', description: 'Maximum number of results to return (default: 20).' } }, required: ['class_name', 'paths'] } @@ -157,6 +163,8 @@ def handle_tool_call(id, params) tool_name = params['name'] args = params['arguments'] || {} show_ast = args['show_ast'] || false + offset = args['offset'] || 0 + limit = args['limit'] || 20 @gains = Gains.new("mcp:#{tool_name}") if args['pattern'] && !args['pattern'].start_with?('(', '{', '[') && !args['pattern'].match?(/^[a-z_]+$/) @@ -168,12 +176,12 @@ def handle_tool_call(id, params) when 'validate_fast_pattern' execute_validate_pattern(args['pattern']) when 'search_ruby_ast' - execute_search(args['pattern'], args['paths'], show_ast: show_ast) + execute_search(args['pattern'], args['paths'], show_ast: show_ast, offset: offset, limit: limit) when 'ruby_method_source' execute_method_search(args['method_name'], args['paths'], - class_name: args['class_name'], show_ast: show_ast) + class_name: args['class_name'], show_ast: show_ast, offset: offset, limit: limit) when 'ruby_class_source' - execute_class_search(args['class_name'], args['paths'], show_ast: show_ast) + execute_class_search(args['class_name'], args['paths'], show_ast: show_ast, offset: offset, limit: limit) when 'rewrite_ruby' execute_rewrite(args['source'], args['pattern'], args['replacement']) when 'rewrite_ruby_file' @@ -198,7 +206,7 @@ def execute_validate_pattern(pattern) { valid: false, error: e.message } end - def execute_search(pattern, paths, show_ast: false) + def execute_search(pattern, paths, show_ast: false, offset: nil, limit: nil) results = [] on_result = ->(file, matches) do @gains&.record_match(file) if matches.any? @@ -218,21 +226,61 @@ def execute_search(pattern, paths, show_ast: false) on_search = ->(file) { @gains&.record_search(file) } Fast.search_all(pattern, paths, parallel: false, on_result: on_result, on_search: on_search) - results + + matches = if offset || limit + results[offset || 0, limit || results.size] || [] + else + results + end + + { + matches: matches, + total: results.size, + offset: offset, + limit: limit, + has_more: (offset || 0) + (limit || results.size) < results.size + } end - def execute_method_search(method_name, paths, class_name: nil, show_ast: false) + def execute_method_search(method_name, paths, class_name: nil, show_ast: false, offset: nil, limit: nil) pattern = "(def #{method_name})" - results = execute_search(pattern, paths, show_ast: show_ast) - return results unless class_name + results = [] + on_result = ->(file, matches) do + @gains&.record_match(file) if matches.any? + matches.compact.each do |node| + next unless (exp = node_expression(node)) + next if class_name && !class_defined_in_file?(class_name, file) - # Filter: keep only methods whose file contains the class - results.select do |r| - class_defined_in_file?(class_name, r[:file]) + entry = { + file: file, + line_start: exp.line, + line_end: exp.last_line, + code: Fast.highlight(node, colorize: false) + } + entry[:ast] = Fast.highlight(node, show_sexp: true, colorize: false) if show_ast + results << entry + end end + on_search = ->(file) { @gains&.record_search(file) } + + Fast.search_all(pattern, paths, parallel: false, on_result: on_result, on_search: on_search) + + matches = if offset || limit + results[offset || 0, limit || results.size] || [] + else + results + end + + { + matches: matches, + total: results.size, + offset: offset, + limit: limit, + has_more: (offset || 0) + (limit || results.size) < results.size + } end - def execute_class_search(class_name, paths, show_ast: false) + def execute_class_search(class_name, paths, show_ast: false, offset: nil, limit: nil) # Use simple (class ...) pattern then filter by name — avoids nil/superclass edge cases results = [] on_result = ->(file, matches) do @@ -254,7 +302,20 @@ def execute_class_search(class_name, paths, show_ast: false) end on_search = ->(file) { @gains&.record_search(file) } Fast.search_all('(class ...)', paths, parallel: false, on_result: on_result, on_search: on_search) - results.select { |r| r[:file] } # already filtered above + + matches = if offset || limit + results[offset || 0, limit || results.size] || [] + else + results + end + + { + matches: matches, + total: results.size, + offset: offset, + limit: limit, + has_more: (offset || 0) + (limit || results.size) < results.size + } end def execute_rewrite(source, pattern, replacement) diff --git a/spec/fast/experiment_spec.rb b/spec/fast/experiment_spec.rb index f7e2191..c9ae8f9 100644 --- a/spec/fast/experiment_spec.rb +++ b/spec/fast/experiment_spec.rb @@ -16,18 +16,19 @@ describe Fast::ExperimentFile do let(:experiment_file) { Fast::ExperimentFile.new(spec, experiment) } - let(:spec) do - tempfile = Tempfile.new - tempfile.write <<~RUBY + let(:tempfile) do + tf = Tempfile.new + tf.write <<~RUBY let(:user) { create(:user) } let(:address) { create(:address) } let(:phone_number) { create(:phone_number) } let(:country) { create(:country) } let(:language) { create(:language) } RUBY - tempfile.close - tempfile.path + tf.close + tf end + let(:spec) { tempfile.path } describe '#filename' do it { expect(experiment_file.experimental_filename(1)).to include('experiment_1') } diff --git a/spec/fast/mcp_pagination_spec.rb b/spec/fast/mcp_pagination_spec.rb new file mode 100644 index 0000000..6ed0e8d --- /dev/null +++ b/spec/fast/mcp_pagination_spec.rb @@ -0,0 +1,121 @@ +# frozen_string_literal: true + +require 'spec_helper' +require 'fast/mcp_server' +require 'fileutils' +require 'json' + +RSpec.describe Fast::McpServer do + let(:server) { Fast::McpServer.new } + let(:temp_dir) { File.expand_path('tmp_mcp_pagination', Dir.pwd) } + let(:test_file) { File.join(temp_dir, 'sample.rb') } + + before do + FileUtils.mkdir_p(temp_dir) + File.write(test_file, <<~RUBY) + def first; end + def second; end + def third; end + RUBY + allow(server).to receive(:write_response) + end + + after do + FileUtils.rm_rf(temp_dir) + end + + describe 'pagination' do + it 'returns paginated results for search_ruby_ast' do + params = { + 'name' => 'search_ruby_ast', + 'arguments' => { + 'pattern' => '(def ...)', + 'paths' => [test_file], + 'limit' => 2 + } + } + + expect(server).to receive(:write_response) do |_id, response| + content = JSON.parse(response[:content].first[:text]) + expect(content['matches'].size).to eq(2) + expect(content['total']).to eq(3) + expect(content['offset']).to eq(0) + expect(content['limit']).to eq(2) + expect(content['has_more']).to be true + expect(content['matches'][0]['code']).to include('def first') + expect(content['matches'][1]['code']).to include('def second') + end + + server.send(:handle_tool_call, '1', params) + end + + it 'handles offset for search_ruby_ast' do + params = { + 'name' => 'search_ruby_ast', + 'arguments' => { + 'pattern' => '(def ...)', + 'paths' => [test_file], + 'offset' => 2, + 'limit' => 2 + } + } + + expect(server).to receive(:write_response) do |_id, response| + content = JSON.parse(response[:content].first[:text]) + expect(content['matches'].size).to eq(1) + expect(content['total']).to eq(3) + expect(content['offset']).to eq(2) + expect(content['limit']).to eq(2) + expect(content['has_more']).to be false + expect(content['matches'][0]['code']).to include('def third') + end + + server.send(:handle_tool_call, '1', params) + end + + it 'returns paginated results for ruby_method_source' do + params = { + 'name' => 'ruby_method_source', + 'arguments' => { + 'method_name' => '...', + 'paths' => [test_file], + 'limit' => 1 + } + } + + expect(server).to receive(:write_response) do |_id, response| + content = JSON.parse(response[:content].first[:text]) + expect(content['matches'].size).to eq(1) + expect(content['total']).to eq(3) + expect(content['has_more']).to be true + end + + server.send(:handle_tool_call, '1', params) + end + + it 'returns paginated results for ruby_class_source' do + class_file = File.join(temp_dir, 'classes.rb') + File.write(class_file, <<~RUBY) + class A; end + class A; end + RUBY + params = { + 'name' => 'ruby_class_source', + 'arguments' => { + 'class_name' => 'A', + 'paths' => [class_file], + 'limit' => 1 + } + } + + expect(server).to receive(:write_response) do |_id, response| + content = JSON.parse(response[:content].first[:text]) + expect(content['matches'].size).to eq(1) + expect(content['total']).to eq(2) + expect(content['has_more']).to be true + end + + server.send(:handle_tool_call, '1', params) + end + end +end