From e5cfeb3d2d0ca0a5e9d5ad53c54ca6547a7cd167 Mon Sep 17 00:00:00 2001 From: Lakshman Patel Date: Wed, 16 Sep 2026 16:37:40 +0530 Subject: [PATCH] refactor(tool): migrate Agent, CodeGraph, PRGenerator, SmartCreate to typed schemas --- internal/tool/agent.go | 116 +++++++++++------------------ internal/tool/codegraph.go | 99 ++++++++++-------------- internal/tool/pr_generator.go | 46 ++++++------ internal/tool/schema_batch_test.go | 71 ++++++++++++++++++ internal/tool/smart_create.go | 33 ++++---- 5 files changed, 200 insertions(+), 165 deletions(-) diff --git a/internal/tool/agent.go b/internal/tool/agent.go index b641f28f7..45809d888 100644 --- a/internal/tool/agent.go +++ b/internal/tool/agent.go @@ -26,6 +26,22 @@ const ( type AgentTool struct{} +// AgentInput is the typed input for AgentTool. +type AgentInput struct { + Prompt string `json:"prompt"` + Description string `json:"description"` + SubagentType string `json:"subagent_type"` + CapabilityMode string `json:"capability_mode"` + Isolation string `json:"isolation"` + Thoroughness string `json:"thoroughness"` + CWD string `json:"cwd"` + Model string `json:"model"` + RunInBackground bool `json:"run_in_background"` + AgentID string `json:"agent_id"` + ResumeFrom string `json:"resume_from"` + RetryOf string `json:"retry_of"` +} + func (AgentTool) Name() string { return "Agent" } func (AgentTool) RiskLevel() string { return "medium" } func (AgentTool) Aliases() []string { return []string{"agent", "Task"} } @@ -36,83 +52,39 @@ func (AgentTool) Description() string { "cwd, model, description, and run_in_background control spawn behavior." } -func (AgentTool) Parameters() map[string]interface{} { - return map[string]interface{}{ - "type": "object", - "properties": map[string]interface{}{ - "prompt": map[string]interface{}{ - "type": "string", - "description": "Task description for the sub-agent", - }, - "description": map[string]interface{}{ - "type": "string", - "description": "Short human-readable label for the spawn (3–5 words).", - }, - "subagent_type": map[string]interface{}{ - "type": "string", - "description": "explore | plan | general-purpose (alias: general). Default: explore.", - "enum": []string{"explore", "plan", "general-purpose", "general"}, - }, - "capability_mode": map[string]interface{}{ - "type": "string", - "description": "read-only | read-write | execute | all. Defaults from subagent_type when omitted.", - "enum": []string{"read-only", "read-write", "execute", "all"}, - }, - "isolation": map[string]interface{}{ - "type": "string", - "description": "none | worktree. Mutually exclusive with cwd when worktree.", - "enum": []string{"none", "worktree"}, - }, - "thoroughness": map[string]interface{}{ - "type": "string", - "description": "Explore only: quick | medium | very-thorough.", - "enum": []string{"quick", "medium", "very-thorough"}, - }, - "cwd": map[string]interface{}{ - "type": "string", - "description": "Working directory for the sub-agent. Mutually exclusive with isolation=worktree.", - }, - "model": map[string]interface{}{ - "type": "string", - "description": "Optional model override for the sub-agent.", - }, - "run_in_background": map[string]interface{}{ - "type": "boolean", - "description": "If true, spawn asynchronously — results are collected when the main turn ends.", - }, - "agent_id": map[string]interface{}{ - "type": "string", - "description": "ID of a previous sub-agent to query status/result (legacy resume lookup).", - }, - "resume_from": map[string]interface{}{ - "type": "string", - "description": "Subagent ID to resume with full transcript (typed spawn).", - }, - "retry_of": map[string]interface{}{ - "type": "string", - "description": "ID of a failed sub-agent to retry. Spawns a new agent with the same request.", - }, +// Schema returns the typed input schema. Parameters() delegates to it so the +// two cannot diverge. +func (AgentTool) Schema() ToolSchema { + return ToolSchema{ + Type: "object", + Properties: map[string]SchemaProperty{ + "prompt": {Type: "string", Description: "Task description for the sub-agent"}, + "description": {Type: "string", Description: "Short human-readable label for the spawn (3–5 words)."}, + "subagent_type": {Type: "string", Enum: []interface{}{"explore", "plan", "general-purpose", "general"}, Description: "explore | plan | general-purpose (alias: general). Default: explore."}, + "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."}, + "isolation": {Type: "string", Enum: []interface{}{"none", "worktree"}, Description: "none | worktree. Mutually exclusive with cwd when worktree."}, + "thoroughness": {Type: "string", Enum: []interface{}{"quick", "medium", "very-thorough"}, Description: "Explore only: quick | medium | very-thorough."}, + "cwd": {Type: "string", Description: "Working directory for the sub-agent. Mutually exclusive with isolation=worktree."}, + "model": {Type: "string", Description: "Optional model override for the sub-agent."}, + "run_in_background": {Type: "boolean", Description: "If true, spawn asynchronously — results are collected when the main turn ends."}, + "agent_id": {Type: "string", Description: "ID of a previous sub-agent to query status/result (legacy resume lookup)."}, + "resume_from": {Type: "string", Description: "Subagent ID to resume with full transcript (typed spawn)."}, + "retry_of": {Type: "string", Description: "ID of a failed sub-agent to retry. Spawns a new agent with the same request."}, }, - "required": []string{"prompt"}, + Required: []string{"prompt"}, } } +func (AgentTool) Parameters() map[string]interface{} { + return agentSchema.ToJSONSchema() +} + +// agentSchema is the single source of truth for Agent's input schema. +var agentSchema = AgentTool{}.Schema() + func (AgentTool) Execute(ctx context.Context, input json.RawMessage) (string, error) { - var p struct { - Prompt string `json:"prompt"` - Description string `json:"description"` - SubagentType string `json:"subagent_type"` - CapabilityMode string `json:"capability_mode"` - Isolation string `json:"isolation"` - Thoroughness string `json:"thoroughness"` - CWD string `json:"cwd"` - Model string `json:"model"` - RunInBackground bool `json:"run_in_background"` - AgentID string `json:"agent_id"` - ResumeFrom string `json:"resume_from"` - RetryOf string `json:"retry_of"` - } - if err := json.Unmarshal(input, &p); err != nil { + p, err := DecodeInput[AgentInput]("Agent", input) + if err != nil { return "", err } if len(p.Prompt) > maxAgentPromptBytes { diff --git a/internal/tool/codegraph.go b/internal/tool/codegraph.go index 858935e36..a7651c420 100644 --- a/internal/tool/codegraph.go +++ b/internal/tool/codegraph.go @@ -15,6 +15,20 @@ import ( // CodeGraphTool provides tree-sitter based code intelligence. type CodeGraphTool struct{} +// CodeGraphInput is the typed input for CodeGraphTool. +type CodeGraphInput struct { + Action string `json:"action"` + Query string `json:"query"` + NodeID string `json:"node_id"` + MaxDepth int `json:"max_depth"` + MaxNodes int `json:"max_nodes"` + Root string `json:"root"` + From string `json:"from"` + To string `json:"to"` + MaxFiles int `json:"max_files"` + Dir string `json:"dir"` +} + func (CodeGraphTool) Name() string { return "CodeGraph" } func (CodeGraphTool) RiskLevel() string { return "low" } func (CodeGraphTool) Aliases() []string { return []string{"cg", "graph"} } @@ -22,70 +36,37 @@ func (CodeGraphTool) Description() string { return "Query the code knowledge graph: search symbols, trace callers/callees, compute impact radius, build context." } -func (CodeGraphTool) Parameters() map[string]interface{} { - return map[string]interface{}{ - "type": "object", - "properties": map[string]interface{}{ - "action": map[string]interface{}{ - "type": "string", - "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"}, - "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", - }, - "query": map[string]interface{}{ - "type": "string", - "description": "Search query, symbol name, or task description (for context/trace use 'from -> to')", - }, - "node_id": map[string]interface{}{ - "type": "string", - "description": "Node ID for callers/callees/impact", - }, - "max_depth": map[string]interface{}{ - "type": "integer", - "description": "Max traversal depth (default: 3)", - }, - "max_nodes": map[string]interface{}{ - "type": "integer", - "description": "Max nodes to return (default: 30)", - }, - "root": map[string]interface{}{ - "type": "string", - "description": "Project root directory (default: current dir)", - }, - "from": map[string]interface{}{ - "type": "string", - "description": "Source symbol for trace action", - }, - "to": map[string]interface{}{ - "type": "string", - "description": "Target symbol for trace action", - }, - "max_files": map[string]interface{}{ - "type": "integer", - "description": "Max files for explore action (default: 10)", - }, - "dir": map[string]interface{}{ - "type": "string", - "description": "Directory filter for files action", - }, +// Schema returns the typed input schema. Parameters() delegates to it so the +// two cannot diverge. +func (CodeGraphTool) Schema() ToolSchema { + return ToolSchema{ + Type: "object", + Properties: map[string]SchemaProperty{ + "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"}, + "query": {Type: "string", Description: "Search query, symbol name, or task description (for context/trace use 'from -> to')"}, + "node_id": {Type: "string", Description: "Node ID for callers/callees/impact"}, + "max_depth": {Type: "integer", Description: "Max traversal depth (default: 3)"}, + "max_nodes": {Type: "integer", Description: "Max nodes to return (default: 30)"}, + "root": {Type: "string", Description: "Project root directory (default: current dir)"}, + "from": {Type: "string", Description: "Source symbol for trace action"}, + "to": {Type: "string", Description: "Target symbol for trace action"}, + "max_files": {Type: "integer", Description: "Max files for explore action (default: 10)"}, + "dir": {Type: "string", Description: "Directory filter for files action"}, }, - "required": []string{"action"}, + Required: []string{"action"}, } } +func (CodeGraphTool) Parameters() map[string]interface{} { + return codeGraphSchema.ToJSONSchema() +} + +// codeGraphSchema is the single source of truth for CodeGraph's input schema. +var codeGraphSchema = CodeGraphTool{}.Schema() + func (CodeGraphTool) Execute(ctx context.Context, input json.RawMessage) (string, error) { - var p struct { - Action string `json:"action"` - Query string `json:"query"` - NodeID string `json:"node_id"` - MaxDepth int `json:"max_depth"` - MaxNodes int `json:"max_nodes"` - Root string `json:"root"` - From string `json:"from"` - To string `json:"to"` - MaxFiles int `json:"max_files"` - Dir string `json:"dir"` - } - if err := json.Unmarshal(input, &p); err != nil { + p, err := DecodeInput[CodeGraphInput]("CodeGraph", input) + if err != nil { return "", err } diff --git a/internal/tool/pr_generator.go b/internal/tool/pr_generator.go index 934755ea7..9d2e2aee4 100644 --- a/internal/tool/pr_generator.go +++ b/internal/tool/pr_generator.go @@ -504,33 +504,37 @@ func (t *PRGeneratorTool) Description() string { } // Parameters returns the JSON schema for the tool's input. -func (t *PRGeneratorTool) Parameters() map[string]interface{} { - return map[string]interface{}{ - "type": "object", - "properties": map[string]interface{}{ - "base_branch": map[string]interface{}{ - "type": "string", - "description": "The base branch to compare against (e.g., 'main', 'develop')", - "default": "main", - }, - "project_dir": map[string]interface{}{ - "type": "string", - "description": "The project directory (defaults to current directory)", - }, +// Schema returns the typed input schema. Parameters() delegates to it so the +// two cannot diverge. +func (t *PRGeneratorTool) Schema() ToolSchema { + return ToolSchema{ + Type: "object", + Properties: map[string]SchemaProperty{ + "base_branch": {Type: "string", Default: "main", Description: "The base branch to compare against (e.g., 'main', 'develop')"}, + "project_dir": {Type: "string", Description: "The project directory (defaults to current directory)"}, }, - "required": []string{}, + Required: []string{}, } } +func (t *PRGeneratorTool) Parameters() map[string]interface{} { + return prGeneratorSchema.ToJSONSchema() +} + +// prGeneratorSchema is the single source of truth for PRGenerator's input schema. +var prGeneratorSchema = (&PRGeneratorTool{}).Schema() + +// PRGeneratorInput is the typed input for PRGeneratorTool. +type PRGeneratorInput struct { + BaseBranch string `json:"base_branch"` + ProjectDir string `json:"project_dir"` +} + // Execute runs the PR generator tool. func (t *PRGeneratorTool) Execute(ctx context.Context, input json.RawMessage) (string, error) { - var params struct { - BaseBranch string `json:"base_branch"` - ProjectDir string `json:"project_dir"` - } - - if err := json.Unmarshal(input, ¶ms); err != nil { - return "", fmt.Errorf("invalid input: %w", err) + params, err := DecodeInput[PRGeneratorInput]("PRGenerator", input) + if err != nil { + return "", err } if params.BaseBranch == "" { diff --git a/internal/tool/schema_batch_test.go b/internal/tool/schema_batch_test.go index 515584a37..e8ef83c58 100644 --- a/internal/tool/schema_batch_test.go +++ b/internal/tool/schema_batch_test.go @@ -1232,3 +1232,74 @@ func TestTaskStopSchemaProvider(t *testing.T) { t.Fatalf("required = %v, want [task_id]", TaskStopTool{}.Parameters()["required"]) } } + +func TestAgentSchemaProvider(t *testing.T) { + var _ SchemaProvider = AgentTool{} + props := schemaProps(t, AgentTool{}.Parameters()) + if props["prompt"].(map[string]interface{})["type"] != "string" { + t.Fatal("prompt type wrong") + } + subEnum, ok := props["subagent_type"].(map[string]interface{})["enum"].([]interface{}) + if !ok || len(subEnum) != 4 || subEnum[0] != "explore" || subEnum[3] != "general" { + t.Fatalf("subagent_type enum = %v, want 4 options", props["subagent_type"]) + } + capEnum, ok := props["capability_mode"].(map[string]interface{})["enum"].([]interface{}) + if !ok || len(capEnum) != 4 || capEnum[0] != "read-only" || capEnum[3] != "all" { + t.Fatalf("capability_mode enum = %v, want 4 options", props["capability_mode"]) + } + if props["run_in_background"].(map[string]interface{})["type"] != "boolean" { + t.Fatal("run_in_background type wrong") + } + req, _ := AgentTool{}.Parameters()["required"].([]string) + if len(req) != 1 || req[0] != "prompt" { + t.Fatalf("required = %v, want [prompt]", AgentTool{}.Parameters()["required"]) + } +} + +func TestCodeGraphSchemaProvider(t *testing.T) { + var _ SchemaProvider = CodeGraphTool{} + props := schemaProps(t, CodeGraphTool{}.Parameters()) + enum, ok := props["action"].(map[string]interface{})["enum"].([]interface{}) + if !ok || len(enum) != 21 || enum[0] != "search" || enum[20] != "hybrid_search" { + t.Fatalf("action enum has %d options, want 21", len(enum)) + } + if props["max_depth"].(map[string]interface{})["type"] != "integer" { + t.Fatal("max_depth type wrong") + } + if props["query"].(map[string]interface{})["type"] != "string" { + t.Fatal("query type wrong") + } + req, _ := CodeGraphTool{}.Parameters()["required"].([]string) + if len(req) != 1 || req[0] != "action" { + t.Fatalf("required = %v, want [action]", CodeGraphTool{}.Parameters()["required"]) + } +} + +func TestPRGeneratorSchemaProvider(t *testing.T) { + var inter SchemaProvider = (&PRGeneratorTool{}) + _ = inter + props := schemaProps(t, (&PRGeneratorTool{}).Parameters()) + base := props["base_branch"].(map[string]interface{}) + if base["type"] != "string" { + t.Fatal("base_branch type wrong") + } + if base["default"] != "main" { + t.Fatalf("base_branch default = %v, want main", base["default"]) + } + if props["project_dir"].(map[string]interface{})["type"] != "string" { + t.Fatal("project_dir type wrong") + } +} + +func TestSmartCreateSchemaProvider(t *testing.T) { + var inter SchemaProvider = (&SmartCreateTool{}) + _ = inter + props := schemaProps(t, (&SmartCreateTool{}).Parameters()) + if props["path"].(map[string]interface{})["type"] != "string" { + t.Fatal("path type wrong") + } + req, _ := (&SmartCreateTool{}).Parameters()["required"].([]string) + if len(req) != 1 || req[0] != "path" { + t.Fatalf("required = %v, want [path]", (&SmartCreateTool{}).Parameters()["required"]) + } +} diff --git a/internal/tool/smart_create.go b/internal/tool/smart_create.go index 01d23c350..e7dc65d3b 100644 --- a/internal/tool/smart_create.go +++ b/internal/tool/smart_create.go @@ -646,7 +646,8 @@ type SmartCreateTool struct { } // smartCreateInput is the JSON input for the SmartCreate tool. -type smartCreateInput struct { +// SmartCreateInput is the typed input for SmartCreateTool. +type SmartCreateInput struct { Path string `json:"path"` } @@ -658,23 +659,29 @@ func (t *SmartCreateTool) Description() string { return "Creates a new file with appropriate boilerplate based on project conventions and file type." } -func (t *SmartCreateTool) Parameters() map[string]interface{} { - return map[string]interface{}{ - "type": "object", - "properties": map[string]interface{}{ - "path": map[string]interface{}{ - "type": "string", - "description": "Path of the file to create", - }, +// Schema returns the typed input schema. Parameters() delegates to it so the +// two cannot diverge. +func (t *SmartCreateTool) Schema() ToolSchema { + return ToolSchema{ + Type: "object", + Properties: map[string]SchemaProperty{ + "path": {Type: "string", Description: "Path of the file to create"}, }, - "required": []string{"path"}, + Required: []string{"path"}, } } +func (t *SmartCreateTool) Parameters() map[string]interface{} { + return smartCreateSchema.ToJSONSchema() +} + +// smartCreateSchema is the single source of truth for SmartCreate's input schema. +var smartCreateSchema = (&SmartCreateTool{}).Schema() + func (t *SmartCreateTool) Execute(ctx context.Context, input json.RawMessage) (string, error) { - var params smartCreateInput - if err := json.Unmarshal(input, ¶ms); err != nil { - return "", fmt.Errorf("invalid input: %w", err) + params, err := DecodeInput[SmartCreateInput]("SmartCreate", input) + if err != nil { + return "", err } if params.Path == "" {