Skip to content

Commit 2ba7496

Browse files
Patel230Patel230
authored
refactor(tool): migrate BatchExec/Diagnostics/CodeSearch/CodeMatch/Outline to typed schemas (#324)
Single-source Schema()+Parameters() pattern, Execute parses via DecodeInput[T], wire-shape parity tests including enums, array items, and numeric bounds.
1 parent 2eb812e commit 2ba7496

6 files changed

Lines changed: 227 additions & 140 deletions

File tree

‎internal/tool/batch_exec.go‎

Lines changed: 28 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -26,49 +26,36 @@ func (BatchExecTool) Description() string {
2626
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."
2727
}
2828

29-
func (BatchExecTool) Parameters() map[string]interface{} {
30-
return map[string]interface{}{
31-
"type": "object",
32-
"properties": map[string]interface{}{
33-
"action": map[string]interface{}{
34-
"type": "string",
35-
"enum": []string{"submit", "poll", "wait"},
36-
"description": "submit: send prompts; poll: single status check; wait: poll until the batch reaches a terminal state (with backoff + Retry-After honoring).",
37-
},
38-
"prompts": map[string]interface{}{
39-
"type": "array",
40-
"items": map[string]interface{}{"type": "string"},
41-
"description": "Prompts to submit (action=submit).",
42-
},
43-
"model": map[string]interface{}{
44-
"type": "string",
45-
"description": "Model ID (default claude-sonnet-4-20250514).",
46-
},
47-
"batch_id": map[string]interface{}{
48-
"type": "string",
49-
"description": "Batch ID to poll/wait on.",
50-
},
51-
"max_tokens": map[string]interface{}{
52-
"type": "integer",
53-
"description": "Max output tokens per request (default 4096).",
54-
},
55-
"timeout_seconds": map[string]interface{}{
56-
"type": "integer",
57-
"description": "Max seconds to wait (default 600).",
58-
},
59-
"poll_interval_seconds": map[string]interface{}{
60-
"type": "integer",
61-
"description": "Initial poll interval in seconds (default 2).",
62-
},
29+
// Schema returns the typed input schema. Parameters() delegates to it so the
30+
// two cannot diverge.
31+
func (BatchExecTool) Schema() ToolSchema {
32+
return ToolSchema{
33+
Type: "object",
34+
Properties: map[string]SchemaProperty{
35+
"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"}},
36+
"prompts": {Type: "array", Description: "Prompts to submit (action=submit).", Items: &SchemaProperty{Type: "string"}},
37+
"model": {Type: "string", Description: "Model ID (default claude-sonnet-4-20250514)."},
38+
"batch_id": {Type: "string", Description: "Batch ID to poll/wait on."},
39+
"max_tokens": {Type: "integer", Description: "Max output tokens per request (default 4096)."},
40+
"timeout_seconds": {Type: "integer", Description: "Max seconds to wait (default 600)."},
41+
"poll_interval_seconds": {Type: "integer", Description: "Initial poll interval in seconds (default 2)."},
6342
},
64-
"required": []string{"action"},
43+
Required: []string{"action"},
6544
}
6645
}
6746

47+
func (BatchExecTool) Parameters() map[string]interface{} {
48+
return batchExecSchema.ToJSONSchema()
49+
}
50+
51+
// batchExecSchema is the single source of truth for BatchExec's input schema.
52+
var batchExecSchema = BatchExecTool{}.Schema()
53+
6854
var batchHTTP = &http.Client{Timeout: 5 * time.Minute}
6955

70-
// batchExecParams are the parsed parameters for all BatchExec actions.
71-
type batchExecParams struct {
56+
// BatchExecInput are the parsed parameters for all BatchExec actions.
57+
// BatchExecInput is the typed input for BatchExecTool.
58+
type BatchExecInput struct {
7259
Action string `json:"action"`
7360
Prompts []string `json:"prompts"`
7461
Model string `json:"model"`
@@ -102,9 +89,9 @@ type batchPollResult struct {
10289
}
10390

10491
func (BatchExecTool) Execute(ctx context.Context, input json.RawMessage) (string, error) {
105-
var p batchExecParams
106-
if err := json.Unmarshal(input, &p); err != nil {
107-
return "", fmt.Errorf("invalid input: %w", err)
92+
p, err := DecodeInput[BatchExecInput]("BatchExec", input)
93+
if err != nil {
94+
return "", err
10895
}
10996

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

133-
func batchSubmit(ctx context.Context, apiKey string, p batchExecParams) (string, error) {
120+
func batchSubmit(ctx context.Context, apiKey string, p BatchExecInput) (string, error) {
134121
if len(p.Prompts) == 0 {
135122
return "", fmt.Errorf("at least one prompt is required")
136123
}

‎internal/tool/code_match.go‎

Lines changed: 29 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,14 @@ import (
2121
// and strings cannot produce false positives.
2222
type CodeMatchTool struct{}
2323

24+
// CodeMatchInput is the typed input for CodeMatchTool.
25+
type CodeMatchInput struct {
26+
Pattern string `json:"pattern"`
27+
Path string `json:"path"`
28+
Language string `json:"language"`
29+
Limit int `json:"limit"`
30+
}
31+
2432
func (CodeMatchTool) Name() string { return "CodeMatch" }
2533
func (CodeMatchTool) RiskLevel() string { return "low" }
2634
func (CodeMatchTool) Aliases() []string { return []string{"code_match", "match_code"} }
@@ -29,34 +37,28 @@ func (CodeMatchTool) Description() string {
2937
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.`
3038
}
3139

32-
func (CodeMatchTool) Parameters() map[string]interface{} {
33-
return map[string]interface{}{
34-
"type": "object",
35-
"properties": map[string]interface{}{
36-
"pattern": map[string]interface{}{
37-
"type": "string",
38-
"description": "Tree-sitter query pattern (S-expression). Captures (@name) are returned per match.",
39-
},
40-
"path": map[string]interface{}{
41-
"type": "string",
42-
"description": "Project directory (default: session working directory).",
43-
},
44-
"language": map[string]interface{}{
45-
"type": "string",
46-
"enum": []string{"go", "python", "typescript", "tsx"},
47-
"description": "Restrict to one language (default: all supported).",
48-
},
49-
"limit": map[string]interface{}{
50-
"type": "integer",
51-
"minimum": 1,
52-
"maximum": 200,
53-
"description": "Maximum matches total (default 30).",
54-
},
40+
// Schema returns the typed input schema. Parameters() delegates to it so the
41+
// two cannot diverge.
42+
func (CodeMatchTool) Schema() ToolSchema {
43+
return ToolSchema{
44+
Type: "object",
45+
Properties: map[string]SchemaProperty{
46+
"pattern": {Type: "string", Description: "Tree-sitter query pattern (S-expression). Captures (@name) are returned per match."},
47+
"path": {Type: "string", Description: "Project directory (default: session working directory)."},
48+
"language": {Type: "string", Enum: []interface{}{"go", "python", "typescript", "tsx"}, Description: "Restrict to one language (default: all supported)."},
49+
"limit": {Type: "integer", Minimum: 1, Maximum: 200, Description: "Maximum matches total (default 30)."},
5550
},
56-
"required": []string{"pattern"},
51+
Required: []string{"pattern"},
5752
}
5853
}
5954

55+
func (CodeMatchTool) Parameters() map[string]interface{} {
56+
return codeMatchSchema.ToJSONSchema()
57+
}
58+
59+
// codeMatchSchema is the single source of truth for CodeMatch's input schema.
60+
var codeMatchSchema = CodeMatchTool{}.Schema()
61+
6062
// codeMatchHit is one structural match.
6163
type codeMatchHit struct {
6264
File string `json:"file"`
@@ -73,14 +75,9 @@ var codeMatchExtLang = map[string]string{
7375
}
7476

7577
func (CodeMatchTool) Execute(ctx context.Context, input json.RawMessage) (string, error) {
76-
var params struct {
77-
Pattern string `json:"pattern"`
78-
Path string `json:"path"`
79-
Language string `json:"language"`
80-
Limit int `json:"limit"`
81-
}
82-
if err := json.Unmarshal(input, &params); err != nil {
83-
return "", fmt.Errorf("invalid input: %w", err)
78+
params, err := DecodeInput[CodeMatchInput]("CodeMatch", input)
79+
if err != nil {
80+
return "", err
8481
}
8582
params.Pattern = strings.TrimSpace(params.Pattern)
8683
if params.Pattern == "" {

‎internal/tool/code_search.go‎

Lines changed: 29 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,14 @@ import (
1010
// CodeSearchTool searches the codebase semantically.
1111
type CodeSearchTool struct{}
1212

13+
// CodeSearchInput is the typed input for CodeSearchTool.
14+
type CodeSearchInput struct {
15+
Query string `json:"query"`
16+
Limit int `json:"limit"`
17+
Language string `json:"language"`
18+
Refresh bool `json:"refresh"`
19+
}
20+
1321
func (CodeSearchTool) Name() string { return "CodeSearch" }
1422
func (CodeSearchTool) RiskLevel() string { return "low" }
1523

@@ -19,40 +27,32 @@ func (CodeSearchTool) Description() string {
1927
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.`
2028
}
2129

22-
func (CodeSearchTool) Parameters() map[string]interface{} {
23-
return map[string]interface{}{
24-
"type": "object",
25-
"properties": map[string]interface{}{
26-
"query": map[string]interface{}{
27-
"type": "string",
28-
"description": "The semantic search query describing what you are looking for.",
29-
},
30-
"limit": map[string]interface{}{
31-
"type": "integer",
32-
"description": "Maximum number of results to return (default 5).",
33-
},
34-
"language": map[string]interface{}{
35-
"type": "string",
36-
"description": "Optional language filter (e.g. go, python, typescript).",
37-
},
38-
"refresh": map[string]interface{}{
39-
"type": "boolean",
40-
"description": "If true, refresh the code index before searching.",
41-
},
30+
// Schema returns the typed input schema. Parameters() delegates to it so the
31+
// two cannot diverge.
32+
func (CodeSearchTool) Schema() ToolSchema {
33+
return ToolSchema{
34+
Type: "object",
35+
Properties: map[string]SchemaProperty{
36+
"query": {Type: "string", Description: "The semantic search query describing what you are looking for."},
37+
"limit": {Type: "integer", Description: "Maximum number of results to return (default 5)."},
38+
"language": {Type: "string", Description: "Optional language filter (e.g. go, python, typescript)."},
39+
"refresh": {Type: "boolean", Description: "If true, refresh the code index before searching."},
4240
},
43-
"required": []string{"query"},
41+
Required: []string{"query"},
4442
}
4543
}
4644

45+
func (CodeSearchTool) Parameters() map[string]interface{} {
46+
return codeSearchSchema.ToJSONSchema()
47+
}
48+
49+
// codeSearchSchema is the single source of truth for CodeSearch's input schema.
50+
var codeSearchSchema = CodeSearchTool{}.Schema()
51+
4752
func (CodeSearchTool) Execute(ctx context.Context, input json.RawMessage) (string, error) {
48-
var params struct {
49-
Query string `json:"query"`
50-
Limit int `json:"limit"`
51-
Language string `json:"language"`
52-
Refresh bool `json:"refresh"`
53-
}
54-
if err := json.Unmarshal(input, &params); err != nil {
55-
return "", fmt.Errorf("invalid input: %w", err)
53+
params, err := DecodeInput[CodeSearchInput]("CodeSearch", input)
54+
if err != nil {
55+
return "", err
5656
}
5757

5858
if params.Query == "" {

‎internal/tool/diagnostics.go‎

Lines changed: 25 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,12 @@ import (
1212
// DiagnosticsTool runs lint/type/compile diagnostics for a file or project.
1313
type DiagnosticsTool struct{}
1414

15+
// DiagnosticsInput is the typed input for DiagnosticsTool.
16+
type DiagnosticsInput struct {
17+
Path string `json:"path"`
18+
Scope string `json:"scope"`
19+
}
20+
1521
func (DiagnosticsTool) Name() string { return "Diagnostics" }
1622
func (DiagnosticsTool) Aliases() []string { return []string{"diagnostics", "lint"} }
1723
func (DiagnosticsTool) Description() string {
@@ -23,31 +29,30 @@ func (DiagnosticsTool) Description() string {
2329
Run this after making code changes to catch errors before committing.`
2430
}
2531

26-
func (DiagnosticsTool) Parameters() map[string]interface{} {
27-
return map[string]interface{}{
28-
"type": "object",
29-
"properties": map[string]interface{}{
30-
"path": map[string]interface{}{
31-
"type": "string",
32-
"description": "File or directory path to diagnose",
33-
},
34-
"scope": map[string]interface{}{
35-
"type": "string",
36-
"description": "Scope of diagnostics: 'file' or 'project' (default: file)",
37-
"enum": []string{"file", "project"},
38-
},
32+
// Schema returns the typed input schema. Parameters() delegates to it so the
33+
// two cannot diverge.
34+
func (DiagnosticsTool) Schema() ToolSchema {
35+
return ToolSchema{
36+
Type: "object",
37+
Properties: map[string]SchemaProperty{
38+
"path": {Type: "string", Description: "File or directory path to diagnose"},
39+
"scope": {Type: "string", Enum: []interface{}{"file", "project"}, Description: "Scope of diagnostics: 'file' or 'project' (default: file)"},
3940
},
40-
"required": []string{"path"},
41+
Required: []string{"path"},
4142
}
4243
}
4344

45+
func (DiagnosticsTool) Parameters() map[string]interface{} {
46+
return diagnosticsSchema.ToJSONSchema()
47+
}
48+
49+
// diagnosticsSchema is the single source of truth for Diagnostics' input schema.
50+
var diagnosticsSchema = DiagnosticsTool{}.Schema()
51+
4452
func (DiagnosticsTool) Execute(ctx context.Context, input json.RawMessage) (string, error) {
45-
var args struct {
46-
Path string `json:"path"`
47-
Scope string `json:"scope"`
48-
}
49-
if err := json.Unmarshal(input, &args); err != nil {
50-
return "", fmt.Errorf("invalid input: %w", err)
53+
args, err := DecodeInput[DiagnosticsInput]("Diagnostics", input)
54+
if err != nil {
55+
return "", err
5156
}
5257
if args.Path == "" {
5358
return "", fmt.Errorf("path is required")

‎internal/tool/outline.go‎

Lines changed: 23 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -15,38 +15,43 @@ import (
1515
// reading the full content. Returns a compact outline with line numbers.
1616
type OutlineTool struct{}
1717

18+
// OutlineInput is the typed input for OutlineTool.
19+
type OutlineInput struct {
20+
FilePath string `json:"file_path"`
21+
FilePaths []string `json:"file_paths,omitempty"`
22+
}
23+
1824
func (OutlineTool) Name() string { return "Outline" }
1925
func (OutlineTool) RiskLevel() string { return "low" }
2026
func (OutlineTool) Aliases() []string { return []string{"outline"} }
2127
func (OutlineTool) Description() string {
2228
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."
2329
}
2430

25-
func (OutlineTool) Parameters() map[string]interface{} {
26-
return map[string]interface{}{
27-
"type": "object",
28-
"properties": map[string]interface{}{
29-
"file_path": map[string]interface{}{
30-
"type": "string",
31-
"description": "Path to a single file, relative to the project root (e.g. 'src/main.go')",
32-
},
33-
"file_paths": map[string]interface{}{
34-
"type": "array",
35-
"items": map[string]interface{}{"type": "string"},
36-
"description": "Paths to multiple files. Use instead of file_path to outline several files in one call (max 20).",
37-
},
31+
// Schema returns the typed input schema. Parameters() delegates to it so the
32+
// two cannot diverge.
33+
func (OutlineTool) Schema() ToolSchema {
34+
return ToolSchema{
35+
Type: "object",
36+
Properties: map[string]SchemaProperty{
37+
"file_path": {Type: "string", Description: "Path to a single file, relative to the project root (e.g. 'src/main.go')"},
38+
"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)."},
3839
},
3940
}
4041
}
4142

43+
func (OutlineTool) Parameters() map[string]interface{} {
44+
return outlineSchema.ToJSONSchema()
45+
}
46+
47+
// outlineSchema is the single source of truth for Outline's input schema.
48+
var outlineSchema = OutlineTool{}.Schema()
49+
4250
const outlineMaxFiles = 20
4351

4452
func (OutlineTool) Execute(ctx context.Context, input json.RawMessage) (string, error) {
45-
var in struct {
46-
FilePath string `json:"file_path"`
47-
FilePaths []string `json:"file_paths,omitempty"`
48-
}
49-
if err := json.Unmarshal(input, &in); err != nil {
53+
in, err := DecodeInput[OutlineInput]("Outline", input)
54+
if err != nil {
5055
return "", err
5156
}
5257

0 commit comments

Comments
 (0)