Skip to content

Commit e5cfeb3

Browse files
Patel230Patel230
committed
refactor(tool): migrate Agent, CodeGraph, PRGenerator, SmartCreate to typed schemas
1 parent e286b44 commit e5cfeb3

5 files changed

Lines changed: 200 additions & 165 deletions

File tree

‎internal/tool/agent.go‎

Lines changed: 44 additions & 72 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,22 @@ const (
2626

2727
type AgentTool struct{}
2828

29+
// AgentInput is the typed input for AgentTool.
30+
type AgentInput struct {
31+
Prompt string `json:"prompt"`
32+
Description string `json:"description"`
33+
SubagentType string `json:"subagent_type"`
34+
CapabilityMode string `json:"capability_mode"`
35+
Isolation string `json:"isolation"`
36+
Thoroughness string `json:"thoroughness"`
37+
CWD string `json:"cwd"`
38+
Model string `json:"model"`
39+
RunInBackground bool `json:"run_in_background"`
40+
AgentID string `json:"agent_id"`
41+
ResumeFrom string `json:"resume_from"`
42+
RetryOf string `json:"retry_of"`
43+
}
44+
2945
func (AgentTool) Name() string { return "Agent" }
3046
func (AgentTool) RiskLevel() string { return "medium" }
3147
func (AgentTool) Aliases() []string { return []string{"agent", "Task"} }
@@ -36,83 +52,39 @@ func (AgentTool) Description() string {
3652
"cwd, model, description, and run_in_background control spawn behavior."
3753
}
3854

39-
func (AgentTool) Parameters() map[string]interface{} {
40-
return map[string]interface{}{
41-
"type": "object",
42-
"properties": map[string]interface{}{
43-
"prompt": map[string]interface{}{
44-
"type": "string",
45-
"description": "Task description for the sub-agent",
46-
},
47-
"description": map[string]interface{}{
48-
"type": "string",
49-
"description": "Short human-readable label for the spawn (3–5 words).",
50-
},
51-
"subagent_type": map[string]interface{}{
52-
"type": "string",
53-
"description": "explore | plan | general-purpose (alias: general). Default: explore.",
54-
"enum": []string{"explore", "plan", "general-purpose", "general"},
55-
},
56-
"capability_mode": map[string]interface{}{
57-
"type": "string",
58-
"description": "read-only | read-write | execute | all. Defaults from subagent_type when omitted.",
59-
"enum": []string{"read-only", "read-write", "execute", "all"},
60-
},
61-
"isolation": map[string]interface{}{
62-
"type": "string",
63-
"description": "none | worktree. Mutually exclusive with cwd when worktree.",
64-
"enum": []string{"none", "worktree"},
65-
},
66-
"thoroughness": map[string]interface{}{
67-
"type": "string",
68-
"description": "Explore only: quick | medium | very-thorough.",
69-
"enum": []string{"quick", "medium", "very-thorough"},
70-
},
71-
"cwd": map[string]interface{}{
72-
"type": "string",
73-
"description": "Working directory for the sub-agent. Mutually exclusive with isolation=worktree.",
74-
},
75-
"model": map[string]interface{}{
76-
"type": "string",
77-
"description": "Optional model override for the sub-agent.",
78-
},
79-
"run_in_background": map[string]interface{}{
80-
"type": "boolean",
81-
"description": "If true, spawn asynchronously — results are collected when the main turn ends.",
82-
},
83-
"agent_id": map[string]interface{}{
84-
"type": "string",
85-
"description": "ID of a previous sub-agent to query status/result (legacy resume lookup).",
86-
},
87-
"resume_from": map[string]interface{}{
88-
"type": "string",
89-
"description": "Subagent ID to resume with full transcript (typed spawn).",
90-
},
91-
"retry_of": map[string]interface{}{
92-
"type": "string",
93-
"description": "ID of a failed sub-agent to retry. Spawns a new agent with the same request.",
94-
},
55+
// Schema returns the typed input schema. Parameters() delegates to it so the
56+
// two cannot diverge.
57+
func (AgentTool) Schema() ToolSchema {
58+
return ToolSchema{
59+
Type: "object",
60+
Properties: map[string]SchemaProperty{
61+
"prompt": {Type: "string", Description: "Task description for the sub-agent"},
62+
"description": {Type: "string", Description: "Short human-readable label for the spawn (3–5 words)."},
63+
"subagent_type": {Type: "string", Enum: []interface{}{"explore", "plan", "general-purpose", "general"}, Description: "explore | plan | general-purpose (alias: general). Default: explore."},
64+
"capability_mode": {Type: "string", Enum: []interface{}{"read-only", "read-write", "execute", "all"}, Description: "read-only | read-write | execute | all. Defaults from subagent_type when omitted."},
65+
"isolation": {Type: "string", Enum: []interface{}{"none", "worktree"}, Description: "none | worktree. Mutually exclusive with cwd when worktree."},
66+
"thoroughness": {Type: "string", Enum: []interface{}{"quick", "medium", "very-thorough"}, Description: "Explore only: quick | medium | very-thorough."},
67+
"cwd": {Type: "string", Description: "Working directory for the sub-agent. Mutually exclusive with isolation=worktree."},
68+
"model": {Type: "string", Description: "Optional model override for the sub-agent."},
69+
"run_in_background": {Type: "boolean", Description: "If true, spawn asynchronously — results are collected when the main turn ends."},
70+
"agent_id": {Type: "string", Description: "ID of a previous sub-agent to query status/result (legacy resume lookup)."},
71+
"resume_from": {Type: "string", Description: "Subagent ID to resume with full transcript (typed spawn)."},
72+
"retry_of": {Type: "string", Description: "ID of a failed sub-agent to retry. Spawns a new agent with the same request."},
9573
},
96-
"required": []string{"prompt"},
74+
Required: []string{"prompt"},
9775
}
9876
}
9977

78+
func (AgentTool) Parameters() map[string]interface{} {
79+
return agentSchema.ToJSONSchema()
80+
}
81+
82+
// agentSchema is the single source of truth for Agent's input schema.
83+
var agentSchema = AgentTool{}.Schema()
84+
10085
func (AgentTool) Execute(ctx context.Context, input json.RawMessage) (string, error) {
101-
var p struct {
102-
Prompt string `json:"prompt"`
103-
Description string `json:"description"`
104-
SubagentType string `json:"subagent_type"`
105-
CapabilityMode string `json:"capability_mode"`
106-
Isolation string `json:"isolation"`
107-
Thoroughness string `json:"thoroughness"`
108-
CWD string `json:"cwd"`
109-
Model string `json:"model"`
110-
RunInBackground bool `json:"run_in_background"`
111-
AgentID string `json:"agent_id"`
112-
ResumeFrom string `json:"resume_from"`
113-
RetryOf string `json:"retry_of"`
114-
}
115-
if err := json.Unmarshal(input, &p); err != nil {
86+
p, err := DecodeInput[AgentInput]("Agent", input)
87+
if err != nil {
11688
return "", err
11789
}
11890
if len(p.Prompt) > maxAgentPromptBytes {

‎internal/tool/codegraph.go‎

Lines changed: 40 additions & 59 deletions
Original file line numberDiff line numberDiff line change
@@ -15,77 +15,58 @@ import (
1515
// CodeGraphTool provides tree-sitter based code intelligence.
1616
type CodeGraphTool struct{}
1717

18+
// CodeGraphInput is the typed input for CodeGraphTool.
19+
type CodeGraphInput struct {
20+
Action string `json:"action"`
21+
Query string `json:"query"`
22+
NodeID string `json:"node_id"`
23+
MaxDepth int `json:"max_depth"`
24+
MaxNodes int `json:"max_nodes"`
25+
Root string `json:"root"`
26+
From string `json:"from"`
27+
To string `json:"to"`
28+
MaxFiles int `json:"max_files"`
29+
Dir string `json:"dir"`
30+
}
31+
1832
func (CodeGraphTool) Name() string { return "CodeGraph" }
1933
func (CodeGraphTool) RiskLevel() string { return "low" }
2034
func (CodeGraphTool) Aliases() []string { return []string{"cg", "graph"} }
2135
func (CodeGraphTool) Description() string {
2236
return "Query the code knowledge graph: search symbols, trace callers/callees, compute impact radius, build context."
2337
}
2438

25-
func (CodeGraphTool) Parameters() map[string]interface{} {
26-
return map[string]interface{}{
27-
"type": "object",
28-
"properties": map[string]interface{}{
29-
"action": map[string]interface{}{
30-
"type": "string",
31-
"enum": []string{"search", "callers", "callees", "impact", "context", "index", "sync", "trace", "explore", "files", "status", "stats", "pagerank", "centrality", "communities", "components", "deadcode", "coupling", "cross_repo", "semantic_search", "hybrid_search"},
32-
"description": "Action: search/find symbols, callers/who calls, callees/what it calls, impact/breakage radius, context/build task context, index/full re-index, sync/incremental update, trace/call path A→B, explore/multi-symbol source, files/list indexed, status/health check, stats/counts, pagerank/file importance, centrality/bridge files, communities/module clusters, components/isolated subsystems, deadcode/unused code, coupling/tightly coupled files, cross_repo/cross-repo dependencies, semantic_search/embedding-based search, hybrid_search/combined FTS5+semantic",
33-
},
34-
"query": map[string]interface{}{
35-
"type": "string",
36-
"description": "Search query, symbol name, or task description (for context/trace use 'from -> to')",
37-
},
38-
"node_id": map[string]interface{}{
39-
"type": "string",
40-
"description": "Node ID for callers/callees/impact",
41-
},
42-
"max_depth": map[string]interface{}{
43-
"type": "integer",
44-
"description": "Max traversal depth (default: 3)",
45-
},
46-
"max_nodes": map[string]interface{}{
47-
"type": "integer",
48-
"description": "Max nodes to return (default: 30)",
49-
},
50-
"root": map[string]interface{}{
51-
"type": "string",
52-
"description": "Project root directory (default: current dir)",
53-
},
54-
"from": map[string]interface{}{
55-
"type": "string",
56-
"description": "Source symbol for trace action",
57-
},
58-
"to": map[string]interface{}{
59-
"type": "string",
60-
"description": "Target symbol for trace action",
61-
},
62-
"max_files": map[string]interface{}{
63-
"type": "integer",
64-
"description": "Max files for explore action (default: 10)",
65-
},
66-
"dir": map[string]interface{}{
67-
"type": "string",
68-
"description": "Directory filter for files action",
69-
},
39+
// Schema returns the typed input schema. Parameters() delegates to it so the
40+
// two cannot diverge.
41+
func (CodeGraphTool) Schema() ToolSchema {
42+
return ToolSchema{
43+
Type: "object",
44+
Properties: map[string]SchemaProperty{
45+
"action": {Type: "string", Enum: []interface{}{"search", "callers", "callees", "impact", "context", "index", "sync", "trace", "explore", "files", "status", "stats", "pagerank", "centrality", "communities", "components", "deadcode", "coupling", "cross_repo", "semantic_search", "hybrid_search"}, Description: "Action: search/find symbols, callers/who calls, callees/what it calls, impact/breakage radius, context/build task context, index/full re-index, sync/incremental update, trace/call path A→B, explore/multi-symbol source, files/list indexed, status/health check, stats/counts, pagerank/file importance, centrality/bridge files, communities/module clusters, components/isolated subsystems, deadcode/unused code, coupling/tightly coupled files, cross_repo/cross-repo dependencies, semantic_search/embedding-based search, hybrid_search/combined FTS5+semantic"},
46+
"query": {Type: "string", Description: "Search query, symbol name, or task description (for context/trace use 'from -> to')"},
47+
"node_id": {Type: "string", Description: "Node ID for callers/callees/impact"},
48+
"max_depth": {Type: "integer", Description: "Max traversal depth (default: 3)"},
49+
"max_nodes": {Type: "integer", Description: "Max nodes to return (default: 30)"},
50+
"root": {Type: "string", Description: "Project root directory (default: current dir)"},
51+
"from": {Type: "string", Description: "Source symbol for trace action"},
52+
"to": {Type: "string", Description: "Target symbol for trace action"},
53+
"max_files": {Type: "integer", Description: "Max files for explore action (default: 10)"},
54+
"dir": {Type: "string", Description: "Directory filter for files action"},
7055
},
71-
"required": []string{"action"},
56+
Required: []string{"action"},
7257
}
7358
}
7459

60+
func (CodeGraphTool) Parameters() map[string]interface{} {
61+
return codeGraphSchema.ToJSONSchema()
62+
}
63+
64+
// codeGraphSchema is the single source of truth for CodeGraph's input schema.
65+
var codeGraphSchema = CodeGraphTool{}.Schema()
66+
7567
func (CodeGraphTool) Execute(ctx context.Context, input json.RawMessage) (string, error) {
76-
var p struct {
77-
Action string `json:"action"`
78-
Query string `json:"query"`
79-
NodeID string `json:"node_id"`
80-
MaxDepth int `json:"max_depth"`
81-
MaxNodes int `json:"max_nodes"`
82-
Root string `json:"root"`
83-
From string `json:"from"`
84-
To string `json:"to"`
85-
MaxFiles int `json:"max_files"`
86-
Dir string `json:"dir"`
87-
}
88-
if err := json.Unmarshal(input, &p); err != nil {
68+
p, err := DecodeInput[CodeGraphInput]("CodeGraph", input)
69+
if err != nil {
8970
return "", err
9071
}
9172

‎internal/tool/pr_generator.go‎

Lines changed: 25 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -504,33 +504,37 @@ func (t *PRGeneratorTool) Description() string {
504504
}
505505

506506
// Parameters returns the JSON schema for the tool's input.
507-
func (t *PRGeneratorTool) Parameters() map[string]interface{} {
508-
return map[string]interface{}{
509-
"type": "object",
510-
"properties": map[string]interface{}{
511-
"base_branch": map[string]interface{}{
512-
"type": "string",
513-
"description": "The base branch to compare against (e.g., 'main', 'develop')",
514-
"default": "main",
515-
},
516-
"project_dir": map[string]interface{}{
517-
"type": "string",
518-
"description": "The project directory (defaults to current directory)",
519-
},
507+
// Schema returns the typed input schema. Parameters() delegates to it so the
508+
// two cannot diverge.
509+
func (t *PRGeneratorTool) Schema() ToolSchema {
510+
return ToolSchema{
511+
Type: "object",
512+
Properties: map[string]SchemaProperty{
513+
"base_branch": {Type: "string", Default: "main", Description: "The base branch to compare against (e.g., 'main', 'develop')"},
514+
"project_dir": {Type: "string", Description: "The project directory (defaults to current directory)"},
520515
},
521-
"required": []string{},
516+
Required: []string{},
522517
}
523518
}
524519

520+
func (t *PRGeneratorTool) Parameters() map[string]interface{} {
521+
return prGeneratorSchema.ToJSONSchema()
522+
}
523+
524+
// prGeneratorSchema is the single source of truth for PRGenerator's input schema.
525+
var prGeneratorSchema = (&PRGeneratorTool{}).Schema()
526+
527+
// PRGeneratorInput is the typed input for PRGeneratorTool.
528+
type PRGeneratorInput struct {
529+
BaseBranch string `json:"base_branch"`
530+
ProjectDir string `json:"project_dir"`
531+
}
532+
525533
// Execute runs the PR generator tool.
526534
func (t *PRGeneratorTool) Execute(ctx context.Context, input json.RawMessage) (string, error) {
527-
var params struct {
528-
BaseBranch string `json:"base_branch"`
529-
ProjectDir string `json:"project_dir"`
530-
}
531-
532-
if err := json.Unmarshal(input, &params); err != nil {
533-
return "", fmt.Errorf("invalid input: %w", err)
535+
params, err := DecodeInput[PRGeneratorInput]("PRGenerator", input)
536+
if err != nil {
537+
return "", err
534538
}
535539

536540
if params.BaseBranch == "" {

‎internal/tool/schema_batch_test.go‎

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1232,3 +1232,74 @@ func TestTaskStopSchemaProvider(t *testing.T) {
12321232
t.Fatalf("required = %v, want [task_id]", TaskStopTool{}.Parameters()["required"])
12331233
}
12341234
}
1235+
1236+
func TestAgentSchemaProvider(t *testing.T) {
1237+
var _ SchemaProvider = AgentTool{}
1238+
props := schemaProps(t, AgentTool{}.Parameters())
1239+
if props["prompt"].(map[string]interface{})["type"] != "string" {
1240+
t.Fatal("prompt type wrong")
1241+
}
1242+
subEnum, ok := props["subagent_type"].(map[string]interface{})["enum"].([]interface{})
1243+
if !ok || len(subEnum) != 4 || subEnum[0] != "explore" || subEnum[3] != "general" {
1244+
t.Fatalf("subagent_type enum = %v, want 4 options", props["subagent_type"])
1245+
}
1246+
capEnum, ok := props["capability_mode"].(map[string]interface{})["enum"].([]interface{})
1247+
if !ok || len(capEnum) != 4 || capEnum[0] != "read-only" || capEnum[3] != "all" {
1248+
t.Fatalf("capability_mode enum = %v, want 4 options", props["capability_mode"])
1249+
}
1250+
if props["run_in_background"].(map[string]interface{})["type"] != "boolean" {
1251+
t.Fatal("run_in_background type wrong")
1252+
}
1253+
req, _ := AgentTool{}.Parameters()["required"].([]string)
1254+
if len(req) != 1 || req[0] != "prompt" {
1255+
t.Fatalf("required = %v, want [prompt]", AgentTool{}.Parameters()["required"])
1256+
}
1257+
}
1258+
1259+
func TestCodeGraphSchemaProvider(t *testing.T) {
1260+
var _ SchemaProvider = CodeGraphTool{}
1261+
props := schemaProps(t, CodeGraphTool{}.Parameters())
1262+
enum, ok := props["action"].(map[string]interface{})["enum"].([]interface{})
1263+
if !ok || len(enum) != 21 || enum[0] != "search" || enum[20] != "hybrid_search" {
1264+
t.Fatalf("action enum has %d options, want 21", len(enum))
1265+
}
1266+
if props["max_depth"].(map[string]interface{})["type"] != "integer" {
1267+
t.Fatal("max_depth type wrong")
1268+
}
1269+
if props["query"].(map[string]interface{})["type"] != "string" {
1270+
t.Fatal("query type wrong")
1271+
}
1272+
req, _ := CodeGraphTool{}.Parameters()["required"].([]string)
1273+
if len(req) != 1 || req[0] != "action" {
1274+
t.Fatalf("required = %v, want [action]", CodeGraphTool{}.Parameters()["required"])
1275+
}
1276+
}
1277+
1278+
func TestPRGeneratorSchemaProvider(t *testing.T) {
1279+
var inter SchemaProvider = (&PRGeneratorTool{})
1280+
_ = inter
1281+
props := schemaProps(t, (&PRGeneratorTool{}).Parameters())
1282+
base := props["base_branch"].(map[string]interface{})
1283+
if base["type"] != "string" {
1284+
t.Fatal("base_branch type wrong")
1285+
}
1286+
if base["default"] != "main" {
1287+
t.Fatalf("base_branch default = %v, want main", base["default"])
1288+
}
1289+
if props["project_dir"].(map[string]interface{})["type"] != "string" {
1290+
t.Fatal("project_dir type wrong")
1291+
}
1292+
}
1293+
1294+
func TestSmartCreateSchemaProvider(t *testing.T) {
1295+
var inter SchemaProvider = (&SmartCreateTool{})
1296+
_ = inter
1297+
props := schemaProps(t, (&SmartCreateTool{}).Parameters())
1298+
if props["path"].(map[string]interface{})["type"] != "string" {
1299+
t.Fatal("path type wrong")
1300+
}
1301+
req, _ := (&SmartCreateTool{}).Parameters()["required"].([]string)
1302+
if len(req) != 1 || req[0] != "path" {
1303+
t.Fatalf("required = %v, want [path]", (&SmartCreateTool{}).Parameters()["required"])
1304+
}
1305+
}

0 commit comments

Comments
 (0)