Skip to content

Commit 60d3750

Browse files
Patel230Patel230
authored
refactor(tool): migrate SQL, StructuredEdit, VerifyPlan, Workflow, CodeGen, TicketCompliance to typed schemas (#331)
1 parent 19e4a41 commit 60d3750

7 files changed

Lines changed: 277 additions & 183 deletions

File tree

‎internal/tool/codegen.go‎

Lines changed: 27 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -229,50 +229,46 @@ func (t *CodeGenTool) Description() string {
229229
return "Generate common code patterns from templates. Supports Go, Python, and TypeScript templates for handlers, tests, middleware, CRUD operations, and more."
230230
}
231231

232-
func (t *CodeGenTool) Parameters() map[string]interface{} {
233-
return map[string]interface{}{
234-
"type": "object",
235-
"properties": map[string]interface{}{
236-
"action": map[string]interface{}{
237-
"type": "string",
238-
"enum": []string{"generate", "list", "preview", "suggest"},
239-
"description": "Action to perform",
240-
},
241-
"template": map[string]interface{}{
242-
"type": "string",
243-
"description": "Template name (e.g., go-handler, py-fastapi-endpoint)",
244-
},
245-
"variables": map[string]interface{}{
246-
"type": "object",
247-
"description": "Template variables as key-value pairs",
248-
},
249-
"language": map[string]interface{}{
250-
"type": "string",
251-
"description": "Filter templates by language (go, python, typescript)",
252-
},
253-
"description": map[string]interface{}{
254-
"type": "string",
255-
"description": "Natural language description for template suggestion",
256-
},
232+
// Schema returns the typed input schema. Parameters() delegates to it so the
233+
// two cannot diverge.
234+
func (t *CodeGenTool) Schema() ToolSchema {
235+
return ToolSchema{
236+
Type: "object",
237+
Properties: map[string]SchemaProperty{
238+
"action": {Type: "string", Enum: []interface{}{"generate", "list", "preview", "suggest"}, Description: "Action to perform"},
239+
"template": {Type: "string", Description: "Template name (e.g., go-handler, py-fastapi-endpoint)"},
240+
"variables": {Type: "object", Description: "Template variables as key-value pairs"},
241+
"language": {Type: "string", Description: "Filter templates by language (go, python, typescript)"},
242+
"description": {Type: "string", Description: "Natural language description for template suggestion"},
257243
},
258-
"required": []string{"action"},
244+
Required: []string{"action"},
259245
}
260246
}
261247

262-
type codeGenInput struct {
248+
func (t *CodeGenTool) Parameters() map[string]interface{} {
249+
return codeGenSchema.ToJSONSchema()
250+
}
251+
252+
// codeGenSchema is the single source of truth for CodeGen's input schema.
253+
var codeGenSchema = (&CodeGenTool{}).Schema()
254+
255+
// CodeGenInput is the typed input for CodeGenTool.
256+
type CodeGenInput struct {
263257
Action string `json:"action"`
264258
Template string `json:"template"`
265259
Variables map[string]string `json:"variables"`
266260
Language string `json:"language"`
267261
Description string `json:"description"`
268262
}
269263

264+
// codeGenInput is an unexported alias for backward compatibility with tests.
265+
type codeGenInput = CodeGenInput
266+
270267
func (t *CodeGenTool) Execute(ctx context.Context, input json.RawMessage) (string, error) {
271-
var in codeGenInput
272-
if err := json.Unmarshal(input, &in); err != nil {
273-
return "", fmt.Errorf("invalid input: %w", err)
268+
in, err := DecodeInput[CodeGenInput]("CodeGen", input)
269+
if err != nil {
270+
return "", err
274271
}
275-
276272
switch in.Action {
277273
case "generate":
278274
if in.Template == "" {

‎internal/tool/schema_batch_test.go‎

Lines changed: 97 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -994,3 +994,100 @@ func TestRefactorSchemaProvider(t *testing.T) {
994994
t.Fatalf("required = %v, want [action file]", RefactorTool{}.Parameters()["required"])
995995
}
996996
}
997+
998+
func TestSQLSchemaProvider(t *testing.T) {
999+
var _ SchemaProvider = SQLTool{}
1000+
props := schemaProps(t, SQLTool{}.Parameters())
1001+
enum, ok := props["driver"].(map[string]interface{})["enum"].([]interface{})
1002+
if !ok || len(enum) != 3 || enum[0] != "sqlite" || enum[2] != "mysql" {
1003+
t.Fatalf("driver enum = %v, want 3 options", props["driver"])
1004+
}
1005+
req, _ := SQLTool{}.Parameters()["required"].([]string)
1006+
if len(req) != 2 || req[0] != "dsn" || req[1] != "query" {
1007+
t.Fatalf("required = %v, want [dsn query]", SQLTool{}.Parameters()["required"])
1008+
}
1009+
}
1010+
1011+
func TestStructuredEditSchemaProvider(t *testing.T) {
1012+
var _ SchemaProvider = StructuredEditTool{}
1013+
props := schemaProps(t, StructuredEditTool{}.Parameters())
1014+
if props["path"].(map[string]interface{})["type"] != "string" {
1015+
t.Fatal("path type wrong")
1016+
}
1017+
blocks := props["blocks"].(map[string]interface{})
1018+
if blocks["type"] != "array" {
1019+
t.Fatalf("blocks type = %v, want array", blocks["type"])
1020+
}
1021+
items := blocks["items"].(map[string]interface{})
1022+
itemReq, ok := items["required"].([]string)
1023+
if !ok || len(itemReq) != 2 || itemReq[0] != "search" || itemReq[1] != "replace" {
1024+
t.Fatalf("items required = %v, want [search replace]", items["required"])
1025+
}
1026+
req, _ := StructuredEditTool{}.Parameters()["required"].([]string)
1027+
if len(req) != 2 || req[0] != "path" || req[1] != "blocks" {
1028+
t.Fatalf("required = %v, want [path blocks]", StructuredEditTool{}.Parameters()["required"])
1029+
}
1030+
}
1031+
1032+
func TestVerifyPlanExecutionSchemaProvider(t *testing.T) {
1033+
var _ SchemaProvider = VerifyPlanExecutionTool{}
1034+
props := schemaProps(t, VerifyPlanExecutionTool{}.Parameters())
1035+
steps := props["plan_steps"].(map[string]interface{})
1036+
if steps["type"] != "array" {
1037+
t.Fatalf("plan_steps type = %v, want array", steps["type"])
1038+
}
1039+
itemProps := steps["items"].(map[string]interface{})["properties"].(map[string]interface{})
1040+
if itemProps["description"].(map[string]interface{})["type"] != "string" {
1041+
t.Fatal("description type wrong")
1042+
}
1043+
if itemProps["expected"].(map[string]interface{})["type"] != "string" {
1044+
t.Fatal("expected type wrong")
1045+
}
1046+
req, _ := VerifyPlanExecutionTool{}.Parameters()["required"].([]string)
1047+
if len(req) != 1 || req[0] != "plan_steps" {
1048+
t.Fatalf("required = %v, want [plan_steps]", VerifyPlanExecutionTool{}.Parameters()["required"])
1049+
}
1050+
}
1051+
1052+
func TestWorkflowSchemaProvider(t *testing.T) {
1053+
var _ SchemaProvider = WorkflowTool{}
1054+
props := schemaProps(t, WorkflowTool{}.Parameters())
1055+
if props["workflow"].(map[string]interface{})["type"] != "string" {
1056+
t.Fatal("workflow type wrong")
1057+
}
1058+
if props["args"].(map[string]interface{})["type"] != "object" {
1059+
t.Fatal("args type wrong")
1060+
}
1061+
req, _ := WorkflowTool{}.Parameters()["required"].([]string)
1062+
if len(req) != 1 || req[0] != "workflow" {
1063+
t.Fatalf("required = %v, want [workflow]", WorkflowTool{}.Parameters()["required"])
1064+
}
1065+
}
1066+
1067+
func TestCodeGenSchemaProvider(t *testing.T) {
1068+
var inter SchemaProvider = (&CodeGenTool{})
1069+
_ = inter
1070+
props := schemaProps(t, (&CodeGenTool{}).Parameters())
1071+
enum, ok := props["action"].(map[string]interface{})["enum"].([]interface{})
1072+
if !ok || len(enum) != 4 || enum[0] != "generate" || enum[3] != "suggest" {
1073+
t.Fatalf("action enum = %v, want 4 options", props["action"])
1074+
}
1075+
req, _ := (&CodeGenTool{}).Parameters()["required"].([]string)
1076+
if len(req) != 1 || req[0] != "action" {
1077+
t.Fatalf("required = %v, want [action]", (&CodeGenTool{}).Parameters()["required"])
1078+
}
1079+
}
1080+
1081+
func TestTicketComplianceSchemaProvider(t *testing.T) {
1082+
var inter SchemaProvider = (&TicketComplianceTool{})
1083+
_ = inter
1084+
props := schemaProps(t, (&TicketComplianceTool{}).Parameters())
1085+
enum, ok := props["ticket_source"].(map[string]interface{})["enum"].([]interface{})
1086+
if !ok || len(enum) != 3 || enum[0] != "github" || enum[2] != "linear" {
1087+
t.Fatalf("ticket_source enum = %v, want 3 options", props["ticket_source"])
1088+
}
1089+
req, _ := (&TicketComplianceTool{}).Parameters()["required"].([]string)
1090+
if len(req) != 2 || req[0] != "ticket_content" || req[1] != "diff" {
1091+
t.Fatalf("required = %v, want [ticket_content diff]", (&TicketComplianceTool{}).Parameters()["required"])
1092+
}
1093+
}

‎internal/tool/sql.go‎

Lines changed: 30 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,15 @@ import (
2424
// dialect light up without further changes here.
2525
type SQLTool struct{}
2626

27+
// SQLInput is the typed input for SQLTool.
28+
type SQLInput struct {
29+
Driver string `json:"driver"`
30+
DSN string `json:"dsn"`
31+
Query string `json:"query"`
32+
AllowWrite bool `json:"allow_write"`
33+
MaxRows int `json:"max_rows"`
34+
}
35+
2736
func (SQLTool) Name() string { return "SQL" }
2837
func (SQLTool) Aliases() []string { return []string{"sql", "sql_query"} }
2938

@@ -40,38 +49,29 @@ func (SQLTool) Description() string {
4049
// can prompt when appropriate.
4150
func (SQLTool) RiskLevel() string { return "medium" }
4251

43-
func (SQLTool) Parameters() map[string]interface{} {
44-
return map[string]interface{}{
45-
"type": "object",
46-
"properties": map[string]interface{}{
47-
"driver": map[string]interface{}{
48-
"type": "string",
49-
"description": "Database dialect: one of sqlite, postgres, mysql. Defaults to sqlite.",
50-
"enum": []string{"sqlite", "postgres", "mysql"},
51-
},
52-
"dsn": map[string]interface{}{
53-
"type": "string",
54-
"description": "Data source name / connection string. For sqlite " +
55-
"this is the file path (or \":memory:\").",
56-
},
57-
"query": map[string]interface{}{
58-
"type": "string",
59-
"description": "The SQL statement to execute.",
60-
},
61-
"allow_write": map[string]interface{}{
62-
"type": "boolean",
63-
"description": "Allow destructive / mutating statements. When false " +
64-
"(the default) only read-only queries are permitted.",
65-
},
66-
"max_rows": map[string]interface{}{
67-
"type": "integer",
68-
"description": "Maximum number of rows to return (default 100).",
69-
},
52+
// Schema returns the typed input schema. Parameters() delegates to it so the
53+
// two cannot diverge.
54+
func (SQLTool) Schema() ToolSchema {
55+
return ToolSchema{
56+
Type: "object",
57+
Properties: map[string]SchemaProperty{
58+
"driver": {Type: "string", Enum: []interface{}{"sqlite", "postgres", "mysql"}, Description: "Database dialect: one of sqlite, postgres, mysql. Defaults to sqlite."},
59+
"dsn": {Type: "string", Description: `Data source name / connection string. For sqlite this is the file path (or ":memory:").`},
60+
"query": {Type: "string", Description: "The SQL statement to execute."},
61+
"allow_write": {Type: "boolean", Description: "Allow destructive / mutating statements. When false (the default) only read-only queries are permitted."},
62+
"max_rows": {Type: "integer", Description: "Maximum number of rows to return (default 100)."},
7063
},
71-
"required": []string{"dsn", "query"},
64+
Required: []string{"dsn", "query"},
7265
}
7366
}
7467

68+
func (SQLTool) Parameters() map[string]interface{} {
69+
return sqlSchema.ToJSONSchema()
70+
}
71+
72+
// sqlSchema is the single source of truth for SQL's input schema.
73+
var sqlSchema = SQLTool{}.Schema()
74+
7575
// driverName maps a user-facing dialect to its database/sql driver name. The
7676
// bool reports whether rho bundles a driver for that dialect.
7777
func sqlDriverName(dialect string) (driver string, bundled bool, err error) {
@@ -250,14 +250,8 @@ func firstSQLKeyword(query string) string {
250250
}
251251

252252
func (t SQLTool) Execute(ctx context.Context, input json.RawMessage) (string, error) {
253-
var p struct {
254-
Driver string `json:"driver"`
255-
DSN string `json:"dsn"`
256-
Query string `json:"query"`
257-
AllowWrite bool `json:"allow_write"`
258-
MaxRows int `json:"max_rows"`
259-
}
260-
if err := json.Unmarshal(input, &p); err != nil {
253+
p, err := DecodeInput[SQLInput]("SQL", input)
254+
if err != nil {
261255
return "", err
262256
}
263257
if strings.TrimSpace(p.DSN) == "" {

‎internal/tool/structured_edit.go‎

Lines changed: 32 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -21,41 +21,42 @@ func (StructuredEditTool) Description() string {
2121
return "Apply search-and-replace edits to a file. Provide one or more SEARCH/REPLACE blocks. Each block finds exact text and replaces it. The SEARCH text must match the file contents exactly, including whitespace."
2222
}
2323

24-
func (StructuredEditTool) Parameters() map[string]interface{} {
25-
return map[string]interface{}{
26-
"type": "object",
27-
"properties": map[string]interface{}{
28-
"path": map[string]interface{}{
29-
"type": "string",
30-
"description": "Path to the file to edit",
31-
},
32-
"blocks": map[string]interface{}{
33-
"type": "array",
34-
"description": "List of SEARCH/REPLACE blocks",
35-
"items": map[string]interface{}{
36-
"type": "object",
37-
"properties": map[string]interface{}{
38-
"search": map[string]interface{}{
39-
"type": "string",
40-
"description": "Exact text to find. Must match the file exactly, including whitespace.",
41-
},
42-
"replace": map[string]interface{}{
43-
"type": "string",
44-
"description": "Text to replace the search text with.",
45-
},
24+
// Schema returns the typed input schema. Parameters() delegates to it so the
25+
// two cannot diverge.
26+
func (StructuredEditTool) Schema() ToolSchema {
27+
return ToolSchema{
28+
Type: "object",
29+
Properties: map[string]SchemaProperty{
30+
"path": {Type: "string", Description: "Path to the file to edit"},
31+
"blocks": {
32+
Type: "array",
33+
Description: "List of SEARCH/REPLACE blocks",
34+
Items: &SchemaProperty{
35+
Type: "object",
36+
Properties: map[string]SchemaProperty{
37+
"search": {Type: "string", Description: "Exact text to find. Must match the file exactly, including whitespace."},
38+
"replace": {Type: "string", Description: "Text to replace the search text with."},
4639
},
47-
"required": []string{"search", "replace"},
40+
Required: []string{"search", "replace"},
4841
},
4942
},
50-
"auto_format": map[string]interface{}{
51-
"type": "boolean",
52-
"description": "If true, ignore insignificant whitespace differences when matching (default: false)",
53-
},
43+
"auto_format": {Type: "boolean", Description: "If true, ignore insignificant whitespace differences when matching (default: false)"},
5444
},
55-
"required": []string{"path", "blocks"},
45+
Required: []string{"path", "blocks"},
5646
}
5747
}
5848

49+
func (StructuredEditTool) Parameters() map[string]interface{} {
50+
return structuredEditSchema.ToJSONSchema()
51+
}
52+
53+
// structuredEditSchema is the single source of truth for StructuredEdit's input schema.
54+
var structuredEditSchema = StructuredEditTool{}.Schema()
55+
56+
// StructuredEditInput is the typed input for StructuredEditTool.
57+
// Aliased from the legacy unexported name.
58+
type StructuredEditInput = structuredEditInput
59+
5960
type structuredEditInput struct {
6061
Path string `json:"path"`
6162
Blocks []searchReplaceBlock `json:"blocks"`
@@ -68,11 +69,10 @@ type searchReplaceBlock struct {
6869
}
6970

7071
func (s StructuredEditTool) Execute(ctx context.Context, input json.RawMessage) (string, error) {
71-
var p structuredEditInput
72-
if err := json.Unmarshal(input, &p); err != nil {
73-
return "", fmt.Errorf("invalid input: %w", err)
72+
p, err := DecodeInput[structuredEditInput]("StructuredEdit", input)
73+
if err != nil {
74+
return "", err
7475
}
75-
7676
if p.Path == "" {
7777
return "", fmt.Errorf("path is required")
7878
}

0 commit comments

Comments
 (0)