Skip to content

refactor(tool): migrate BatchExec/Diagnostics/CodeSearch/CodeMatch/Outline to typed schemas #324

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

Merged
Patel230 merged 1 commit into from
Sep 16, 2026
Merged

refactor(tool): migrate BatchExec/Diagnostics/CodeSearch/CodeMatch/Outline to typed schemas #324

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
69 changes: 28 additions & 41 deletions internal/tool/batch_exec.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,49 +26,36 @@ func (BatchExecTool) Description() string {
return "Submit one or more prompts to the Anthropic Message Batches API for 50%-cost async execution. Returns a batch ID; poll with action=poll to check status."
}

func (BatchExecTool) Parameters() map[string]interface{} {
return map[string]interface{}{
"type": "object",
"properties": map[string]interface{}{
"action": map[string]interface{}{
"type": "string",
"enum": []string{"submit", "poll", "wait"},
"description": "submit: send prompts; poll: single status check; wait: poll until the batch reaches a terminal state (with backoff + Retry-After honoring).",
},
"prompts": map[string]interface{}{
"type": "array",
"items": map[string]interface{}{"type": "string"},
"description": "Prompts to submit (action=submit).",
},
"model": map[string]interface{}{
"type": "string",
"description": "Model ID (default claude-sonnet-4-20250514).",
},
"batch_id": map[string]interface{}{
"type": "string",
"description": "Batch ID to poll/wait on.",
},
"max_tokens": map[string]interface{}{
"type": "integer",
"description": "Max output tokens per request (default 4096).",
},
"timeout_seconds": map[string]interface{}{
"type": "integer",
"description": "Max seconds to wait (default 600).",
},
"poll_interval_seconds": map[string]interface{}{
"type": "integer",
"description": "Initial poll interval in seconds (default 2).",
},
// Schema returns the typed input schema. Parameters() delegates to it so the
// two cannot diverge.
func (BatchExecTool) Schema() ToolSchema {
return ToolSchema{
Type: "object",
Properties: map[string]SchemaProperty{
"action": {Type: "string", Description: "submit: send prompts; poll: single status check; wait: poll until the batch reaches a terminal state (with backoff + Retry-After honoring).", Enum: []interface{}{"submit", "poll", "wait"}},
"prompts": {Type: "array", Description: "Prompts to submit (action=submit).", Items: &SchemaProperty{Type: "string"}},
"model": {Type: "string", Description: "Model ID (default claude-sonnet-4-20250514)."},
"batch_id": {Type: "string", Description: "Batch ID to poll/wait on."},
"max_tokens": {Type: "integer", Description: "Max output tokens per request (default 4096)."},
"timeout_seconds": {Type: "integer", Description: "Max seconds to wait (default 600)."},
"poll_interval_seconds": {Type: "integer", Description: "Initial poll interval in seconds (default 2)."},
},
"required": []string{"action"},
Required: []string{"action"},
}
}

func (BatchExecTool) Parameters() map[string]interface{} {
return batchExecSchema.ToJSONSchema()
}

// batchExecSchema is the single source of truth for BatchExec's input schema.
var batchExecSchema = BatchExecTool{}.Schema()

var batchHTTP = &http.Client{Timeout: 5 * time.Minute}

// batchExecParams are the parsed parameters for all BatchExec actions.
type batchExecParams struct {
// BatchExecInput are the parsed parameters for all BatchExec actions.
// BatchExecInput is the typed input for BatchExecTool.
type BatchExecInput struct {
Action string `json:"action"`
Prompts []string `json:"prompts"`
Model string `json:"model"`
Expand Down Expand Up @@ -102,9 +89,9 @@ type batchPollResult struct {
}

func (BatchExecTool) Execute(ctx context.Context, input json.RawMessage) (string, error) {
var p batchExecParams
if err := json.Unmarshal(input, &p); err != nil {
return "", fmt.Errorf("invalid input: %w", err)
p, err := DecodeInput[BatchExecInput]("BatchExec", input)
if err != nil {
return "", err
}

apiKey := batchAPIKey()
Expand All @@ -130,7 +117,7 @@ func (BatchExecTool) Execute(ctx context.Context, input json.RawMessage) (string
}
}

func batchSubmit(ctx context.Context, apiKey string, p batchExecParams) (string, error) {
func batchSubmit(ctx context.Context, apiKey string, p BatchExecInput) (string, error) {
if len(p.Prompts) == 0 {
return "", fmt.Errorf("at least one prompt is required")
}
Expand Down
61 changes: 29 additions & 32 deletions internal/tool/code_match.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,14 @@ import (
// and strings cannot produce false positives.
type CodeMatchTool struct{}

// CodeMatchInput is the typed input for CodeMatchTool.
type CodeMatchInput struct {
Pattern string `json:"pattern"`
Path string `json:"path"`
Language string `json:"language"`
Limit int `json:"limit"`
}

func (CodeMatchTool) Name() string { return "CodeMatch" }
func (CodeMatchTool) RiskLevel() string { return "low" }
func (CodeMatchTool) Aliases() []string { return []string{"code_match", "match_code"} }
Expand All @@ -29,34 +37,28 @@ func (CodeMatchTool) Description() string {
return `Structural code search with tree-sitter query patterns. Matches the AST, not text: comments/strings cannot false-positive. Pattern is a tree-sitter query S-expression; use @captures to extract parts. Examples - Go functions: "(function_declaration name: (identifier) @name) @fn" | Go calls of one function: "(call_expression function: (identifier) @callee) @call" | Python defs: "(function_definition name: (identifier) @name) @fn". Language auto-detects per file; restrict with language=go|python|typescript|tsx.`
}

func (CodeMatchTool) Parameters() map[string]interface{} {
return map[string]interface{}{
"type": "object",
"properties": map[string]interface{}{
"pattern": map[string]interface{}{
"type": "string",
"description": "Tree-sitter query pattern (S-expression). Captures (@name) are returned per match.",
},
"path": map[string]interface{}{
"type": "string",
"description": "Project directory (default: session working directory).",
},
"language": map[string]interface{}{
"type": "string",
"enum": []string{"go", "python", "typescript", "tsx"},
"description": "Restrict to one language (default: all supported).",
},
"limit": map[string]interface{}{
"type": "integer",
"minimum": 1,
"maximum": 200,
"description": "Maximum matches total (default 30).",
},
// Schema returns the typed input schema. Parameters() delegates to it so the
// two cannot diverge.
func (CodeMatchTool) Schema() ToolSchema {
return ToolSchema{
Type: "object",
Properties: map[string]SchemaProperty{
"pattern": {Type: "string", Description: "Tree-sitter query pattern (S-expression). Captures (@name) are returned per match."},
"path": {Type: "string", Description: "Project directory (default: session working directory)."},
"language": {Type: "string", Enum: []interface{}{"go", "python", "typescript", "tsx"}, Description: "Restrict to one language (default: all supported)."},
"limit": {Type: "integer", Minimum: 1, Maximum: 200, Description: "Maximum matches total (default 30)."},
},
"required": []string{"pattern"},
Required: []string{"pattern"},
}
}

func (CodeMatchTool) Parameters() map[string]interface{} {
return codeMatchSchema.ToJSONSchema()
}

// codeMatchSchema is the single source of truth for CodeMatch's input schema.
var codeMatchSchema = CodeMatchTool{}.Schema()

// codeMatchHit is one structural match.
type codeMatchHit struct {
File string `json:"file"`
Expand All @@ -73,14 +75,9 @@ var codeMatchExtLang = map[string]string{
}

func (CodeMatchTool) Execute(ctx context.Context, input json.RawMessage) (string, error) {
var params struct {
Pattern string `json:"pattern"`
Path string `json:"path"`
Language string `json:"language"`
Limit int `json:"limit"`
}
if err := json.Unmarshal(input, &params); err != nil {
return "", fmt.Errorf("invalid input: %w", err)
params, err := DecodeInput[CodeMatchInput]("CodeMatch", input)
if err != nil {
return "", err
}
params.Pattern = strings.TrimSpace(params.Pattern)
if params.Pattern == "" {
Expand Down
58 changes: 29 additions & 29 deletions internal/tool/code_search.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,14 @@ import (
// CodeSearchTool searches the codebase semantically.
type CodeSearchTool struct{}

// CodeSearchInput is the typed input for CodeSearchTool.
type CodeSearchInput struct {
Query string `json:"query"`
Limit int `json:"limit"`
Language string `json:"language"`
Refresh bool `json:"refresh"`
}

func (CodeSearchTool) Name() string { return "CodeSearch" }
func (CodeSearchTool) RiskLevel() string { return "low" }

Expand All @@ -19,40 +27,32 @@ func (CodeSearchTool) Description() string {
return `Semantic code search over the local code index. Use this instead of Grep when you need to find implementations by meaning, not exact text. Start with limit=5; if results look relevant, use offset to paginate. Set refresh=true to update the index first.`
}

func (CodeSearchTool) Parameters() map[string]interface{} {
return map[string]interface{}{
"type": "object",
"properties": map[string]interface{}{
"query": map[string]interface{}{
"type": "string",
"description": "The semantic search query describing what you are looking for.",
},
"limit": map[string]interface{}{
"type": "integer",
"description": "Maximum number of results to return (default 5).",
},
"language": map[string]interface{}{
"type": "string",
"description": "Optional language filter (e.g. go, python, typescript).",
},
"refresh": map[string]interface{}{
"type": "boolean",
"description": "If true, refresh the code index before searching.",
},
// Schema returns the typed input schema. Parameters() delegates to it so the
// two cannot diverge.
func (CodeSearchTool) Schema() ToolSchema {
return ToolSchema{
Type: "object",
Properties: map[string]SchemaProperty{
"query": {Type: "string", Description: "The semantic search query describing what you are looking for."},
"limit": {Type: "integer", Description: "Maximum number of results to return (default 5)."},
"language": {Type: "string", Description: "Optional language filter (e.g. go, python, typescript)."},
"refresh": {Type: "boolean", Description: "If true, refresh the code index before searching."},
},
"required": []string{"query"},
Required: []string{"query"},
}
}

func (CodeSearchTool) Parameters() map[string]interface{} {
return codeSearchSchema.ToJSONSchema()
}

// codeSearchSchema is the single source of truth for CodeSearch's input schema.
var codeSearchSchema = CodeSearchTool{}.Schema()

func (CodeSearchTool) Execute(ctx context.Context, input json.RawMessage) (string, error) {
var params struct {
Query string `json:"query"`
Limit int `json:"limit"`
Language string `json:"language"`
Refresh bool `json:"refresh"`
}
if err := json.Unmarshal(input, &params); err != nil {
return "", fmt.Errorf("invalid input: %w", err)
params, err := DecodeInput[CodeSearchInput]("CodeSearch", input)
if err != nil {
return "", err
}

if params.Query == "" {
Expand Down
45 changes: 25 additions & 20 deletions internal/tool/diagnostics.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,12 @@ import (
// DiagnosticsTool runs lint/type/compile diagnostics for a file or project.
type DiagnosticsTool struct{}

// DiagnosticsInput is the typed input for DiagnosticsTool.
type DiagnosticsInput struct {
Path string `json:"path"`
Scope string `json:"scope"`
}

func (DiagnosticsTool) Name() string { return "Diagnostics" }
func (DiagnosticsTool) Aliases() []string { return []string{"diagnostics", "lint"} }
func (DiagnosticsTool) Description() string {
Expand All @@ -23,31 +29,30 @@ func (DiagnosticsTool) Description() string {
Run this after making code changes to catch errors before committing.`
}

func (DiagnosticsTool) Parameters() map[string]interface{} {
return map[string]interface{}{
"type": "object",
"properties": map[string]interface{}{
"path": map[string]interface{}{
"type": "string",
"description": "File or directory path to diagnose",
},
"scope": map[string]interface{}{
"type": "string",
"description": "Scope of diagnostics: 'file' or 'project' (default: file)",
"enum": []string{"file", "project"},
},
// Schema returns the typed input schema. Parameters() delegates to it so the
// two cannot diverge.
func (DiagnosticsTool) Schema() ToolSchema {
return ToolSchema{
Type: "object",
Properties: map[string]SchemaProperty{
"path": {Type: "string", Description: "File or directory path to diagnose"},
"scope": {Type: "string", Enum: []interface{}{"file", "project"}, Description: "Scope of diagnostics: 'file' or 'project' (default: file)"},
},
"required": []string{"path"},
Required: []string{"path"},
}
}

func (DiagnosticsTool) Parameters() map[string]interface{} {
return diagnosticsSchema.ToJSONSchema()
}

// diagnosticsSchema is the single source of truth for Diagnostics' input schema.
var diagnosticsSchema = DiagnosticsTool{}.Schema()

func (DiagnosticsTool) Execute(ctx context.Context, input json.RawMessage) (string, error) {
var args struct {
Path string `json:"path"`
Scope string `json:"scope"`
}
if err := json.Unmarshal(input, &args); err != nil {
return "", fmt.Errorf("invalid input: %w", err)
args, err := DecodeInput[DiagnosticsInput]("Diagnostics", input)
if err != nil {
return "", err
}
if args.Path == "" {
return "", fmt.Errorf("path is required")
Expand Down
41 changes: 23 additions & 18 deletions internal/tool/outline.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,38 +15,43 @@ import (
// reading the full content. Returns a compact outline with line numbers.
type OutlineTool struct{}

// OutlineInput is the typed input for OutlineTool.
type OutlineInput struct {
FilePath string `json:"file_path"`
FilePaths []string `json:"file_paths,omitempty"`
}

func (OutlineTool) Name() string { return "Outline" }
func (OutlineTool) RiskLevel() string { return "low" }
func (OutlineTool) Aliases() []string { return []string{"outline"} }
func (OutlineTool) Description() string {
return "Extract function/type/class signatures from one or more files. Returns a compact outline with line numbers — much cheaper than reading the full file."
}

func (OutlineTool) Parameters() map[string]interface{} {
return map[string]interface{}{
"type": "object",
"properties": map[string]interface{}{
"file_path": map[string]interface{}{
"type": "string",
"description": "Path to a single file, relative to the project root (e.g. 'src/main.go')",
},
"file_paths": map[string]interface{}{
"type": "array",
"items": map[string]interface{}{"type": "string"},
"description": "Paths to multiple files. Use instead of file_path to outline several files in one call (max 20).",
},
// Schema returns the typed input schema. Parameters() delegates to it so the
// two cannot diverge.
func (OutlineTool) Schema() ToolSchema {
return ToolSchema{
Type: "object",
Properties: map[string]SchemaProperty{
"file_path": {Type: "string", Description: "Path to a single file, relative to the project root (e.g. 'src/main.go')"},
"file_paths": {Type: "array", Items: &SchemaProperty{Type: "string"}, Description: "Paths to multiple files. Use instead of file_path to outline several files in one call (max 20)."},
},
}
}

func (OutlineTool) Parameters() map[string]interface{} {
return outlineSchema.ToJSONSchema()
}

// outlineSchema is the single source of truth for Outline's input schema.
var outlineSchema = OutlineTool{}.Schema()

const outlineMaxFiles = 20

func (OutlineTool) Execute(ctx context.Context, input json.RawMessage) (string, error) {
var in struct {
FilePath string `json:"file_path"`
FilePaths []string `json:"file_paths,omitempty"`
}
if err := json.Unmarshal(input, &in); err != nil {
in, err := DecodeInput[OutlineInput]("Outline", input)
if err != nil {
return "", err
}

Expand Down
Loading
Loading